Question: represent the article for finding a single-sentence tldr

Question: Many challenging prediction problems, from molecular optimization to program synthesis, involve creating complex structured objects as outputs. However, available training data may not be sufficient for a generative model to learn all possible complex transformations. By leveraging the idea that evaluation is easier than generation, we show how a simple, broadly applicable, iterative target augmentation scheme can be surprisingly effective in guiding the training and use of such models. Our scheme views the generative model as a prior distribution, and employs a separately trained filter as the likelihood. In each augmentation step, we filter the model's outputs to obtain additional prediction targets for the next training epoch. Our method is applicable in the supervised as well as semi-supervised settings. We demonstrate that our approach yields significant gains over strong baselines both in molecular optimization and program synthesis. In particular, our augmented model outperforms the previous state-of-the-art in molecular optimization by over 10% in absolute gain. Deep architectures are becoming increasingly adept at generating complex objects such as images, text, molecules, or programs. Many useful generation problems can be seen as translation tasks, where the goal is to take a source (precursor) object such as a molecule and turn it into a target satisfying given design characteristics. Indeed, molecular optimization of this kind is a key step in drug development, though the adoption of automated tools remains limited due to accuracy concerns. We propose here a simple, broadly applicable meta-algorithm to improve translation quality. Translation is a challenging task for many reasons. Objects are complex and the available training data pairs do not fully exemplify the intricate ways in which valid targets can be created from the precursors. Moreover, precursors provided at test time may differ substantially from those available during training -a scenario common in drug development. While data augmentation and semisupervised methods have been used to address some of these challenges, the focus has been on either simple prediction tasks (e.g., classification) or augmenting data primarily on the source side. We show, in contrast, that iteratively augmenting translation targets significantly improves performance on complex generation tasks in which each precursor corresponds to multiple possible outputs. Our iterative target augmentation approach builds on the idea that it is easier to evaluate candidate objects than to generate them. Thus a learned predictor of target object quality (a filter) can be used to effectively guide the generation process. To this end, we construct an external filter and apply it to the complex generative model's sampled translations of training set precursors. Candidate translations that pass the filter criteria become part of the training data for the next training epoch. The translation model is therefore iteratively guided to generate candidates that pass the filter. The generative model can be viewed as an adaptively tuned prior distribution over complex objects, with the filter as the likelihood. For this reason, it is helpful to apply the filter at test time as well, or to use the approach transductively 1 to adapt the generation process to novel test cases. The approach is reminiscent of self-training or reranking approaches employed with some success for parsing . However, in our case, it is the candidate generator that is complex while the filter is relatively simple and remains fixed during the iterative process. We demonstrate that our meta-algorithm is quite effective and consistent in its ability to improve translation quality in the supervised setting. On a program synthesis task , under the same neural architecture, our augmented model outperforms their MLE baseline by 8% and their RL model by 3% in top-1 generalization accuracy (in absolute measure). On molecular optimization (a), their sequence to sequence translation baseline, when combined with our target data augmentation, achieves a new state-of-the-art and outperforms their graph based approach by over 10% in success rate. Their graph based methods are also improved by iterative target augmentation with more than 10% absolute gain. The reflect the difficulty of generation in comparison to evaluation; indeed, the gains persist even if the filter quality is reduced somewhat. Source side augmentation with unlabeled precursors (the semi-supervised setting) can further improve , but only when combined with the filter in the target data augmentation framework. We provide ablation experiments to empirically highlight the effect of our method and also offer some theoretical insights for why it is effective. Molecular Optimization The goal of molecular optimization is to learn to modify compounds so as to improve their chemical properties.;; used reinforcement learning approaches, while Jin et al. (2019a; b) formulated this problem as graphto-graph translation and significantly outperformed previous methods. However, their performance remains imperfect due to the limited size of given training sets. Our work uses property prediction models to check whether generated molecules have desired chemical properties. Recent advances in graph convolutional networks have provided effective solutions to predict those properties in silico. In this work, we use an off-the-shelf property prediction model to filter proposed translation pairs during data augmentation. Program Synthesis Program synthesis is the task of generating a program (using domain-specific language) based on given input-output specifications (; ;). One can check a generated program's correctness by simply executing it on each input and verifying its output.; leverage this idea in their respective decoding procedures, while also using structural constraints on valid programs. Semi-supervised Learning Our method is related to various approaches in semi-supervised learning. In image and text classification, data augmentation and label guessing are commonly applied to obtain artificial labels for unlabeled data. In machine translation, sample new targets from a stationary distribution in order to match the model distribution to the exponentiated payoff distribution centered at a single target sentence. Back-translation creates extra translation pairs by using a backward translation system to translate unlabeled sentences from a target language into a source language. In contrast, our method works in the forward direction because many translation tasks are not symmetric. Moreover, our data augmentation is carried out over multiple iterations, in which we use the augmented model to generate new data for the next iteration. In syntactic parsing, our method is closely related to self-training . They generate new parse trees from unlabeled sentences by applying an existing parser followed by a reranker, and then treat the ing parse trees as new training targets. However, their method is not iterative, and their reranker is explicitly trained to operate over the top k outputs of the parser; in contrast, our filter is independent of the generative model. In addition we show that our approach, which can be viewed as iteratively combining reranking and self-training, is theoretically motivated and can improve the performance of highly complex neural models in multiple domains. Co-training and tri-training also augment a parsing dataset by adding targets on which multiple baseline models agree. Instead of using multiple learners, our method uses task-specific constraints to select correct outputs. Our iterative target augmentation framework can be applied to any conditional generation task with task-specific constraints. For example, molecular optimization (a; b) is the task of transforming a given molecule X into another compound Y with improved chemical properties, Figure 1: Illustration of our data generation process in the program synthesis setting. Given an input-output specification, we first use our generation model to generate candidate programs, and then select correct programs using our external filter. Images of input-output specification and the program A are from. for attempt in 1,..., C do Train model on D. 14: for epoch in 1,..., n 2 do Iterative target augmentation 15: while constraining Y to remain similar to X. Program synthesis is the task of generating a program Y satisfying input specification X; for example, X may be a set of input-output test cases which Y must pass. Without loss of generality, we formulate the generation task as a translation problem. For a given input X, the model learns to generate an output Y satisfying the constraint c. The proposed augmentation framework can be applied to any translation model M trained on an existing dataset As illustrated in Figure 1, our method is an iterative procedure in which each iteration consists of the following two steps: Step: Let D t be the training set at iteration t. To construct each next training set D t+1, we feed each input X i ∈ D (the original training set, not D t) into the translation model up to C times to sample C candidate translations We take the first K distinct translations for each X i satisfying the constraint c and add them to D t+1. When we do not find K distinct valid translations, we simply add the original translation Y i to D t+1. • Training Step: We continue to train the model M t over the new training set D t+1 for one epoch. The above training procedure is summarized in Algorithm 1. As the constraint c is known a priori, we can construct an external filter to remove generated outputs that violate c during the augmentation step. At test time, we also use this filter to screen predicted outputs. To propose the final translation of a given input X, we have the model generate up to L outputs until we find one satisfying the constraint c. If all L attempts fail for a particular input, we just output the first of the failed attempts. Finally, as an additional improvement, we observe that the augmentation step can be carried out for unlabeled inputs X that have no corresponding Y. Thus we can further augment our training dataset in the transductive setting by including test set inputs during the augmentation step, or in the semi-supervised setting by simply including unlabeled inputs. We provide here some theoretical motivation for our iterative target augmentation framework. For simplicity, we consider an external filter c X,Y that is a binary indicator function representing whether output Y satisfies the desired constraint in relation to input X. In other words, we would like to generate Y such that Y ∈ B(X) = {Y |c X,Y = 1}. If the initial translation model P (Y |X) serves as a reasonable prior distribution over outputs, we could simply "invert" the filter and use as the ideal translation model. While this posterior calculation is typically not feasible but could be approximated through samples, it relies heavily on the appropriateness of the prior (model prior to augmentation). Instead, we go a step further and iteratively optimize our parametrically defined prior translation model P θ (Y |X). Note that the ing prior can become much more concentrated around acceptable translations. We maximize the log-likelihood that candidate translations satisfy the constraints implicitly encoded in the filter In many cases there are multiple viable outputs for any given input X. The training data may provide only one (or none) of them. Therefore, we treat the output structure Y as a latent variable, and expand the inner term of Eq. as Since the above objective involves discrete latent variables Y, we propose to maximize Eq. using the standard EM algorithm , especially its incremental, approximate variant. The target augmentation step in our approach is a sampled version of the E-step where the posterior samples are drawn with rejection sampling guided by the filter. The number of samples K controls the quality of approximation to the posterior. 3 The additional training step based on the augmented targets corresponds to a generalized M-step. More precisely, let P (t) θ (Y |X) be the current translation model after t epochs of augmentation training. In epoch t + 1, the augmentation step first samples C different candidates for each input X using the old model P (t) parameterized by θ (t), and then removes those which violate the constraint c, interpretable as samples from the current posterior As a , the training step maximizes the EM auxiliary objective via stochastic gradient descent: We train the model with multiple iterations and show empirically that model performance indeed keeps improving as we add more iterations. The EM approach is likely to converge to a different and better-performing translation model than the initial posterior calculation discussed above. We demonstrate the broad applicability of iterative target augmentation by applying it to two tasks of different domains: molecular optimization and program synthesis. Figure 2: Illustration of molecular optimization. Molecules can be modeled as graphs, with atoms as nodes and bonds as edges. Here, the task is to train a translation model to modify a given input molecule into a target molecule with higher drug-likeness (QED) score. The constraint has two components: the output Y must be highly drug-like, and must be sufficiently similar to the input X. The goal of molecular optimization is to learn to modify molecules so as to improve their chemical properties. As illustrated in Figure 2, this task is formulated as a graph-to-graph translation problem. Similar to machine translation, the training set is a set of molecular pairs {(X, Y)}. X is the input molecule (precursor) and Y is a similar molecule with improved chemical properties. Each molecule in the training set D is further labeled with its property score. Our method is well-suited to this task because the target molecule is not unique: each precursor molecule can be modified in many different ways to optimize its properties. External Filter The constraint for this task contains two parts: 1) the chemical property of Y must exceed a certain threshold β, and 2) the molecular similarity between X and Y must exceed a certain threshold δ. The molecular similarity sim(X, Y) is defined as Tanimoto similarity on Morgan fingerprints , which measures structural overlap between two molecules. In real world settings, ground truth values of chemical properties are often evaluated through experimental assays, which are too expensive and time-consuming to run for iterative target augmentation. Therefore, we construct an in silico property predictor F 1 to approximate the true property evaluator F 0. To train this property prediction model, we use the molecules in the training set and their labeled property values. The predictor F 1 is parameterized as a graph convolutional network and trained using the Chemprop package . During data augmentation, we use F 1 to filter out molecules whose predicted property is under the threshold β. We follow the evaluation setup of Jin et al. (2019b) for two molecular optimization tasks: 1. QED Optimization: The task is to improve the drug-likeness (QED) of a given compound X. The task is to optimize biological activity against the dopamine type 2 receptor (DRD2). The similarity constraint is sim(X, Y) ≥ 0.4 and the property constraint is DRD2(Y) ≥ 0.5, where DRD2(Y) ∈ is the predicted probability of biological activity given by the model from. We treat the output of the in silico evaluators from and as ground truth, and we use them only during test-time evaluation. Evaluation Metrics. During evaluation, we are interested both in the probability that the model will find a successful modification for a given molecule, as well as the diversity of the successful modifications when there are multiple. Table 1: Performance of different models on QED and DRD2 optimization tasks. Italicized models with + are augmented with iterative target augmentation. We emphasize that iterative target augmentation remains critical to performance in the semi-supervised and transductive settings; data augmentation without an external filter instead decreases performance. 1. Success: The fraction of molecules X for which any of the outputs Y 1... Y Z meet the required similarity and property constraints (specified previously for each task). This is our main metric. 2. Diversity: For each molecule X, we measure the average Tanimoto distance (defined as 1 − sim(Y i, Y j)) between pairs within the set of successfully translated compounds among Y 1... Y Z. If there are one or fewer successful translations then the diversity is 0. We average this quantity across all test molecules. We consider the following two model architectures from Jin et al. (2019a) to show that our augmentation scheme is not tied to specific neural architectures. 1. VSeq2Seq, a sequence-to-sequence translation model generating molecules by their SMILES string . 2. HierGNN, a hierarchical graph-to-graph architecture that achieves state-of-the-art performance on the QED and DRD2 tasks, outperforming VSeq2Seq by a wide margin. We apply our iterative augmentation procedure to the above two models, generating up to K = 4 new targets per precursor during each epoch of iterative target augmentation. Additionally, we evaluate our augmentation of VSeq2Seq in a transductive setting, as well as in a semi-supervised setting where we provide 100K additional source-side precursors from the ZINC database . Full hyperparameters are in Appendix A. As shown in Table 1, our iterative augmentation paradigm significantly improves the performance of VSeq2Seq and HierGNN. On both datasets, the translation success rate increases by over 10% in absolute terms for both models. In fact, VSeq2Seq+, our augmentation of the simple VSeq2Seq model, outperforms the non-augmented version of HierGNN. This strongly confirms our hypothesis about the inherent challenge of learning translation models in data sparse scenarios. Moreover, we find that adding more precursors during data augmentation further improves the VSeq2Seq model. On the QED dataset, the translation success rate improves from 89.0% to 92.6% by just adding test set molecules as precursors (VSeq2Seq+, transductive). When instead adding 100K presursors from the external ZINC database, the performance further increases to 95.0% (VSeq2Seq+, semisupervised). We observe similar improvements for the DRD2 task as well. Beyond accuracy gain, our augmentation strategy also improves the diversity of generated molecules. For instance, on the DRD2 dataset, our approach yields 100% relative gain in terms of output diversity. Although the property predictor used in data augmentation is different from the ground truth property evaluator used at test time, the difference in evaluators does not derail the overall training process. Here we analyze the influence of the quality of the property predictor used in data augmentation. Specifically, we rerun our experiments using less accurate predictors in the property-predicting component of our external filter. We obtain these less accurate predictors by undertraining Chemprop and decreasing its hidden dimension. For comparison, we also report with the oracle property predictor which is the ground truth property evaluator. As shown in Figure 3, on the DRD2 dataset, we are able to maintain strong performance despite using predictors that deviate significantly from the ground truth. This implies that our framework can potentially be applied to other properties that are harder to predict. On the QED dataset, our Table 2: Ablation analysis of filtering at training and test time. "Train" indicates a model whose training process uses data augmentation according to our framework. "Test" indicates a model that uses the external filter at prediction time to discard candidate outputs which fail to pass the filter. The evaluation for VSeq2Seq(no-filter) is conducted after 10 augmentation epochs, as the best validation set performance only decreases over the course of training. method is less tolerant to inaccurate property prediction because the property constraint is much tighter -it requires the QED score of an output Y to be in the range [0.9, 1.0]. Importance of External Filtering Our full model (VSeq2Seq+) uses the external filter during both training and testing. We further experiment with Vseq2seq(test), a version of our model trained without data augmentation but which uses the external filter to remove invalid outputs at test time. As shown in Table 2, VSeq2Seq(test) performs significantly worse than our full model trained under data augmentation. Similarly, a model VSeq2Seq(train) trained with the data augmentation but without the prediction time filtering also performs much worse than the full model. In addition, we run an augmentation-only version of the model without an external filter. This model (referred to as VSeq2Seq(no-filter) in Table 2 ) augments the data in each epoch by simply using the first K distinct candidate translations for each precursor X in the training set, without using the external filter at all. In addition, we provide this model with the 100K unlabeled precursors from the semi-supervised setting. Nevertheless, we find that the performance of this model steadily declines from that of the bootstrapped starting point with each data augmentation epoch. Thus the external filter is necessary to prevent poor targets from leading the model training astray. In program synthesis, the source is a set of input-output specifications for the program, and the target is a program that passes all test cases. Our method is suitable for this task because the target program is not unique. Multiple programs may be consistent with the given input-output specifications. The external filter is straightforward for this task: we simply check whether the generated output passes all test cases. Note that at evaluation time, each instance contains extra held-out input-output test cases; the program must pass these in addition to the given test cases in order to be considered correct. When we perform prediction time filtering, we do not use held-out test cases in our filter. Top-1 Generalization MLE 71.91 MLE + RL + Beam Search 77.12 80.17 Table 3: Model performance on Karel program synthesis task. MLE+ is our augmented version of the MLE model . Our task is based on the educational Karel programming language used for evaluation in and. Commands in the Karel language guide a robot's actions in a 2D grid, and may include for loops, while loops, and conditionals. Figure 1 contains an example. We follow the experiment setup of. Evaluation Metrics. The evaluation metric is top-1 generalization. This metric measures how often the model can generate a program that passes the input-output test cases on the test set. At test time, we use our model to generate up to L candidate programs and select the first one to pass the input-output specifications (not including held-out test cases). Models and Baselines. Our main baseline is the MLE baseline from. This model consists of a CNN encoder for the input-output grids and a LSTM decoder along with a handcoded syntax checker. It is trained to maximize the likelihood of the provided target program. Our model is the augmentation of this MLE baseline by our iterative target augmentation framework. As with molecular optimization, we generate up to K = 4 new targets per precursor during each augmentation step. Additionally, we compare against the best model from , which finetunes the same MLE architecture using an RL method with beam search to estimate gradients. We use the same hyperparameters as the original MLE baseline; see Appendix A for details. Table 3 shows the performance of our model in comparison to previous work. Our model (MLE+) outperforms the base MLE model in model by a wide margin. Moreover, our model outperforms the best reinforcement learning model (RL + Beam Search) in , which was trained to directly maximize the generalization metric. This demonstrates the efficacy of our approach in the program synthesis domain. Since our augmentation framework is complementary to architectural improvements, we hypothesize that other techniques, such as execution based synthesis , can benefit from our approach as well. In this work, we have presented an iterative target augmentation framework for generation tasks with multiple possible outputs. Our approach is theoretically motivated, and we demonstrate strong empirical on both the molecular optimization and program synthesis tasks, significantly outperforming baseline models on each task. Moreover, we find that iterative target augmentation is complementary to architectural improvements, and that its effect can be quite robust to the quality of the external filter. Finally, in principle our approach is applicable to other domains as well. Our augmented models share the same hyperparameters as their baseline counterparts in all cases. For the VSeq2Seq model we use batch size 64, embedding and hidden dimension 300, VAE latent dimension 30, and an LSTM with depth 1 (bidirectional in the encoder, unidirectional in the decoder). For models using iterative target augmentation, n 1 is set to 5 and n 2 is set to 10, while for the baseline models we train for 20 epochs (corresponding to n 1 = 20, n 2 = 0). The HierGNN model shares the same hyperparameters as in Jin et al. (2019a). For the training time and prediction time filtering parameters, we set K = 4, C = 200, and L = 10 for both the QED and DRD2 tasks. For the Karel program synthesis task, we use the same hyperparameters as the MLE baseline model in. We use a beam size of 64 at test time, the same as the MLE baseline, but simply sample programs from the decoder distribution when running iterative target augmentation during training. The baseline model is trained for 100 epochs, while for the model employing iterative target augmentation we train as normal for n 1 = 15 epochs followed by n 2 = 50 epochs of iterative target augmentation. Due to the large size of the full training dataset, in each epoch of iterative augmentation we use 1 10 of the dataset, so in total we make 5 passes over the entire dataset. For the training time and prediction time filtering parameters, we set K = 4, C = 50, and L = 10. In Table 4 we provide the training, validation, and test set sizes for all of our tasks. For each task we use the same splits as our baselines. Training Set Validation Set Test Set QED 88306 360 800 DRD2 34404 500 1000 Karel 1116854 2500 2500 Table 4: Number of source-target pairs in training, validation, and test sets for each task. In Figure 5, we provide the validation set performance per iterative target augmentation epoch for our VSeq2Seq+ model on both the QED and DRD2 tasks. The corresponding figure for the MLE+ model on the Karel task is in the main text in Figure 4. In our molecular optimization tasks, we experiment with the effect of modifying K, the number of new targets added per precursor during each training epoch. In all other experiments we have used K = 4. Since taking K = 0 corresponds to the base non-augmented model, it is unsurprising that performance may suffer when K is too small. However, as shown in Table 5, at least in molecular optimization there is relatively little change in performance for K much larger than 4. We also experiment with a version of our method which continually grows the training dataset by keeping all augmented targets, instead of discarding new targets at the end of each epoch. We chose the latter version for our main experiments due to its closer alignment to our EM motivation. However, we demonstrate in Table 6 that performance gains from continually growing the dataset are small to insignificant in our molecular optimization tasks. Table 6: Performance of our proposed augmentation scheme, VSeq2Seq+, compared to an alternative version (VSeq2Seq+, keep-targets) which keeps all generated targets and continually grows the training dataset. In Table 7 we provide the same ablation analysis that we provided in the main text for molecular optimization, demonstrating that both training time iterative target augmentation as well as prediction time filtering are beneficial to model performance. However, we note that even MLE(train), our model without prediction time filtering, outperforms the best RL method from. Table 7: Ablation analysis of filtering at training and test time. "Train" indicates a model whose training process uses data augmentation according to our framework. "Test" indicates a model that uses the external filter at prediction time to discard candidate outputs which fail to pass the filter. Note that MLE and MLE(test) are based on an MLE checkpoint which underperforms the published from by 1 point, due to training for fewer epochs.

We improve generative models by proposing a meta-algorithm that filters new training data from the model's outputs.

Represent the TLDR for finding the corresponding scientific article

Pseudo-labeling has shown to be a weak alternative for semi-supervised learning. We, conversely, demonstrate that dealing with confirmation bias with several regularizations makes pseudo-labeling a suitable approach.


Question: represent the text to find a very short summary (1-2 sentences max)

Question: The Boltzmann distribution is a natural model for many systems, from brains to materials and biomolecules, but is often of limited utility for fitting data because Monte Carlo algorithms are unable to simulate it in available time. This gap between the expressive capabilities and sampling practicalities of energy-based models is exemplified by the protein folding problem, since energy landscapes underlie contemporary knowledge of protein biophysics but computer simulations are challenged to fold all but the smallest proteins from first principles. In this work we aim to bridge the gap between the expressive capacity of energy functions and the practical capabilities of their simulators by using an unrolled Monte Carlo simulation as a model for data. We compose a neural energy function with a novel and efficient simulator based on Langevin dynamics to build an end-to-end-differentiable model of atomic protein structure given amino acid sequence information. We introduce techniques for stabilizing backpropagation under long roll-outs and demonstrate the model's capacity to make multimodal predictions and to, in some cases, generalize to unobserved protein fold types when trained on a large corpus of protein structures. Many natural systems, such as cells in a tissue or atoms in a protein, organize into complex structures from simple underlying interactions. Explaining and predicting how macroscopic structures such as these arise from simple interactions is a major goal of science and, increasingly, machine learning. The Boltzmann distribution is a foundational model for relating local interactions to system behavior, but can be difficult to fit to data. Given an energy function U ✓ [x], the probability of a system configuration x scales exponentially with energy as DISPLAYFORM0 where the (typically intractable) constant Z normalizes the distribution. Importantly, simple energy functions U ✓ [x] consisting of weak, local interactions can collectively encode complex system behaviors, such as the structures of materials and molecules or, when endowed with latent variables, the statistics of images, sound, and text BID0 BID17. Unfortunately, learning model parameters✓ and generating samples x ⇠ p ✓ (x) of the Boltzmann distribution is difficult in practice, as these procedures depend on expensive Monte Carlo simulations that may struggle to mix effectively. These difficulties have driven a shift towards generative models that are easier to learn and sample from, such as directed latent variable models and autoregressive models .The protein folding problem provides a prime example of both the power of energy-based models at describing complex relationships in data as well as the challenge of generating samples from them. Decades of research in biochemistry and biophysics support an energy landscape theory of An unrolled simulator as a model for protein structure. NEMO combines a neural energy function for coarse protein structure, a stochastic simulator based on Langevin dynamics with learned (amortized) initialization, and an atomic imputation network to build atomic coordinate output from sequence information. It is trained end-to-end by backpropagating through the unrolled folding simulation.protein folding , in which the folds that natural protein sequences adopt are those that minimize free energy. Without the availability of external information such as coevolutionary information or homologous structures (Martí-) to constrain the energy function, however, contemporary simulations are challenged to generate globally favorable low-energy structures in available time. How can we get the representational benefits of energy-based models with the sampling efficiency of directed models? Here we explore a potential solution of directly training an unrolled simulator of an energy function as a model for data. By directly training the sampling process, we eschew the question'when has the simulator converged' and instead demand that it produce a useful answer in a fixed amount of time. Leveraging this idea, we construct an end-to-end differentiable model of protein structure that is trained by backpropagtion through folding (FIG0). NEMO (Neural energy modeling and optimization) can learn at scale to generate 3D protein structures consisting of hundreds of points directly from sequence information. Our main contributions are:• Neural energy simulator model for protein structure that composes a deep energy function, unrolled Langevin dynamics, and an atomic imputation network for an end-to-end differentiable model of protein structure given sequence information• Efficient sampling algorithm that is based on a transform integrator for efficient sampling in transformed coordinate systems• Stabilization techniques for long roll-outs of simulators that can exhibit chaotic dynamics and, in turn, exploding gradients during backpropagation• Systematic analysis of combinatorial generalization with a new dataset of protein sequence and structure Protein modeling Our model builds on a long history of coarse-grained modeling of protein structure (; Figure 2 : A neural energy function models coarse grained structure and is sampled by internal coordinate dynamics. (A) The energy function is formulated as a Markov Random Field with structure-based features and sequence-based weights computed by neural networks FIG2 ). (B) To rapidly sample low-energy configurations, the Langevin dynamics simulator leverages both (i) an internal coordinate parameterization, which is more effective for global rearrangements, and (ii) a Cartesian parameterization, which is more effective for localized structural refinement. (C) The base features of the structure network are rotationally and translationally invariant internal coordinates (not shown), pairwise distances, and pairwise orientations.2016; BID2 BID5. Structured Prediction Energy Networks (SPENs) with unrolled optimization BID6 are a highly similar approach to ours, differing in terms of the use of optimization rather than sampling. Additional methodologically related work includes approaches to learn energy functions and samplers simultaneously (; BID25 ; BID20 BID8, to learn efficient MCMC operators BID20), to build expressive approximating distributions with unrolled Monte Carlo simulations BID18 BID23, and to learn the parameters of simulators with implicitly defined likelihoods 1 BID10 BID24. Overview NEMO is an end-to-end differentiable model of protein structure X conditioned on sequence information s consisting of three components (FIG0): (i) a neural energy function U ✓ [x; s] for coarse grained structure x given sequence, (ii) an unrolled simulator that generates approximate samples from U via internal coordinate Langevin dynamics (§ 2.3), and (iii) an imputation network that generates an atomic model X from the final coarse-grained sample x (T) (§ 2.4). All components are trained simultaneously via backpropagation through the unrolled process. Proteins Proteins are linear polymers (sequences) of amino acids that fold into defined 3D structures. The 20 natural amino acids have a common monomer structure [-(N-H)-(C-R)-(C=O)-] with variable side-chain R groups that can differ in properties such as hydrophobicity, charge, and ability to form hydrogen bonds. When placed in solvent (such as water or a lipid membrane), interactions between the side-chains, backbone, and solvent drive proteins into particular 3D configurations ('folds'), which are the basis for understanding protein properties such as biochemical activity, ligand binding, and interactions with drugs. Coordinate representations We predict protein structure X in terms of 5 positions per amino acid: the four heavy atoms of the backbone (N, C ↵, and carbonyl C=O) and the center of mass of the side chain R group. While it is well-established that the locations of C ↵ carbons are sufficient to reconstruct a full atomic structure , we include these additional positions for evaluating backbone hydrogen bonding (secondary structure) and coarse side-chain placement. Internally, the differentiable simulator generates an initial coarse-grained structure (1-position-peramino-acid) with the loss function targeted to the midpoint of the C ↵ carbon and the side chain center of mass. Sequence conditioning We consider two modes for conditioning our model on sequence information: 1-seq, in which s is an L ⇥ 20 matrix containing a one-hot encoding of the amino acid sequence, and Profile, in which s is an L ⇥ 40 matrix encoding both the amino acid sequence and a profile of evolutionarily related sequences (§ B.7).Internal coordinates In contrast to Cartesian coordinates x, which parameterize structure in terms of absolute positions of points x i 2 R 3, internal coordinates z parameterize structure in terms of relative distances and angles between points. We adopt a standard convention for internal coordinates of chains BID13 where each point x i is placed in a spherical coordinate system defined by the three preceding points x i 1, x i 2, x i 3 in terms of a radius (bond length 2) b i 2, a polar angle (bond angle) a i 2 [0, ⇡), and an azimuthal angle (dihedral angle) d i 2 [0, 2⇡) (Figure 2B). We define z i = {b i,ã i, d i}, whereb i,ã i are unconstrained parameterizations of b i and a i (§ A.1). The transformation x = F(z) from internal coordinates to Cartesian is then defined by the recurrence DISPLAYFORM0 ||ûi 1 ⇥ûi|| is a unit vector normal to each bond plane. The inverse transformation z = F 1 (x) is simpler to compute, as it only involves local (and fully parallelizable) calculations of distances and angles (§ A.1). Deep Markov Random Field We model the distribution of a structure x conditioned on a sequence s with the Boltzmann distribution, p ✓ (x|s) = 1. Internal coordinates z All internal coordinates except 6 are invariant to rotation and translation 3 and we mask these in the energy loss.2. Distances D ij = kx i x j k between all pairs of points. We further process these by 4 radial basis functions with (learned) Gaussian kernels.3. Orientation vectorsv ij, which are unit vectors encoding the relative position of point x j in a local coordinate system of x i with base vectorsû DISPLAYFORM0 kûi û i+1 k,n i+1, and the cross product thereof. Langevin dynamics The Langevin dynamics is a stochastic differential equation that aymptotically samples from the Boltzmann distribution (Equation 1). It is typically simulated by a first-order discretization as DISPLAYFORM0 Internal coordinate dynamics The efficiency with which Langevin dynamics explores conformational space is highly dependent on the geometry (and thus parameterization) of the energy landscape U (x). While Cartesian dynamics are efficient at local structural rearrangement, internal coordinate dynamics much more efficiently sample global, coherent changes to the topology of the fold (Figure 2B). We interleave the Cartesian Langevin dynamics with preconditioned Internal Coordinate dynamics, DISPLAYFORM1 where C is a preconditioning matrix that sets the relative scaling of changes to each degree of freedom. For all simulations we unroll T = 250 time steps, each of which is comprised of one Cartesian step followed by one internal coordinate step (Equation 9, § A.3).Transform integrator Simulating internal coordinate dynamics is often computationally intensive as it requires rebuilding Cartesian geometry x from internal coordinates z with F(z) BID13 which is an intrinsically sequential process. Here we bypass the need for recomputing coordinate transformations at every step by instead computing on-the-fly transformation integration (Figure 3). The idea is to directly apply coordinate updates in one coordinate system to another by numerically integrating the Jacobian. This can be favorable when the Jacobian has a simple structure, such as in our case where it requires only distributed cross products. Local reference frame reconstruction The imputation network builds an atomic model X from the final coarse coordinates x (T). Each atomic coordinate X i,j of atom type j at position i is placed in a local reference frame as DISPLAYFORM0 where e i,j (z; ✓) and r i,j (z; ✓) are computed by a 1D convolutional neural network (FIG2). We train and evaluate the model on a set of ⇠67,000 protein structures (domains) that are hierarchically and temporally split. The model is trained by gradient descent using a composite loss that combines terms from likelihood-based and empirical-risk minimization-based training. Compute forces fz = @x @z DISPLAYFORM0 Compute forces fz = @x @z DISPLAYFORM1 Figure 3: A transform integrator simulates Langevin dynamics in a more favorable coordinate system (e.g. internal coordinates z) directly in terms of the untransformed state variables (e.g. Cartesian x). This exchanges the cost of an inner-loop transformation step (e.g. geometry construction F(z)) for an extra Jacobian evaluation, which is fully parallelizable on modern hardware (e.g. GPUs). Structural stratification There are several scales of generalization in protein structure prediction, which range from predicting the structure of a sequence that differs from the training set at a few positions to predicting a 3D fold topology that is absent from training set. To test these various levels of generalization systematically across many different protein families, we built a dataset on top of the CATH hierarchical classification of protein folds BID11. CATH hierarchically organizes proteins from the Protein Data Bank BID7 into domains (individual folds) that are classified at the levels of Class, Architecture, Topology, and Homologous superfamily (from general to specific). We collected protein domains from CATH releases 4.1 and 4.2 up to length 200 and hierarchically and temporally split this set (§ B.1) into training (⇠35k folds), validation (⇠21k folds), and test sets (⇠10k folds). The final test set is subdivided into four subsets: C, A, T, and H, based on the level of maximal similarity between a given test domain and domains in the training set. For example, domains in the C or A sets may share class and potentially architecture classifications with train but will not share topology (i.e. fold). Likelihood The gradient of the data-averaged log likelihood of the Boltzmann distribution is DISPLAYFORM0 which, when ascended, will minimize the average energy of samples from the data relative to samples from the model. In an automatic differentiation setting, we implement a Monte Carlo estimator for (the negative of) this gradient by adding the energy gap, DISPLAYFORM1 to the loss, where? is an identity operator that sets the gradient to zero 4.Empirical Risk In addition to the likelihood loss, which backpropagates through the energy function but not the whole simulation, we developed an empirical risk loss composing several measures of protein model quality. It takes the form schematized in FIG2. Our combined loss sums all of the terms L = L ER + L ML without weighting. DISPLAYFORM2 We found that the long roll-outs of our simulator were prone to chaotic dynamics and exploding gradients, as seen in other work (; BID12 . Unfortunately, when chaotic dynamics do occur, it is typical for all gradients to explode (across learning steps) and standard techniques such as gradient clipping BID14 are unable to rescue learning (§ B.5). To stabilize training, we developed two complimentary techniques that regularize against chaotic simulator dynamics while still facilitating learning when they arise. They are• Lyapunov regularization We regularize the simulator time-step function (rather than the energy function) to be approximately 1-Lipschitz. (If exactly satisfied, this eliminates the possibility of chaotic dynamics.)• Damped backpropagation through time We exponentially decay gradient accumulation on the backwards pass of automatic differentiation by multiplying each backwards iteration by a damping factor. We adaptively tune to cancel the scale of the exploding gradients. This can be thought of as a continuous relaxation of and a quantitatively tunable alternative to truncated backpropagation through time. Figure 5: Examples of fold generalization at topology and architecture level. These predicted structures show a range of prediction accuracy for structural generalization (C and A) tasks, with the TM-score comparing the top ranked 3D-Jury pick against the target. The largest clusters are the three most-populated clusters derived from 100 models per domain with a within-cluster cutoff of TM > 0.5. CATH IDs: 2oy8A03; 5c3uA02; 2y6xA00; 3cimB00; 4ykaC00; 2f09A00; 3i5qA02; 2ayxA01. For each of the 10,381 protein structures in our test set, we sampled 100 models from NEMO, clustered them by structural similarity, and selected a representative structure by a standard consensus algorithm . For evaluation of performance we focus on the TM-Score BID27, a measure of structural similarity between 0 and 1 for which TM > 0.5 is typically considered an approximate reconstruction of a fold. Calibrated uncertainty We find that, when the model is confident (i.e. the number of distinct structural clusters is low ⇠1-3), it is also accurate with some predictions having average TM > 0.5 FIG1. Unsurprisingly, the confidence of the model tends to go with the difficulty of generalization, with the most confident predictions from the H test set and the least confident from C. Structural generalization However, even when sequence identity is low and generalization difficulty is high FIG1, center), the model is still able to make some accurate predictions of 3D structures. Figure 5 illustrates some of these successful predictions at C and A levels, specifically 4ykaC00, 5c3uA02 and beta sheet formation in 2oy8A03. We observe that the predictive distribution is multimodal with non-trivial differences between the clusters representing alternate packing of the chain. In some of the models there is uneven distribition of uncertainty along the chain, which sometimes corresponded to loosely packed regions of the protein. Comparison to an end-to-end baseline We constructed a baseline model that is a non-iterative replica of NEMO which replaces the coarse-grained simulator module (and energy function) with a two-layer bidirectional LSTM that directly predicts coarse internal coordinates z (followed by transformation to Cartesian coordinates with F). We trained this baseline across a range of hyperparameter values and found that for difficult C, A, and T tasks, NEMO generalized more effectively than the RNNs TAB1. For the best performing 2x300 architecture, we trained two additional replicates and report the averaged perfomance in FIG1 (right).Additionally, we report the of a sequence-only NEMO model in TAB1. Paralleling secondary structure prediction BID16 ), we find that the availability of evolutionary information has significant impact on prediction quality. This work presents a novel approach for protein structure prediction that combines the inductive bias of simulators with the speed of directed models. A major advantage of the approach is that model sampling (inference) times can be considerably faster than conventional approaches to protein structure prediction TAB5 ). There are two major disadvantages. First, the computational cost of training and sampling is higher than that of angle-predicting RNNs FIG0 ) such as our baseline or. Consequently, those methods have been scaled to larger datasets than ours (in protein length and diversity) which are more relevant to protein structure prediction tasks. Second, the instability of backpropagating through long simulations is unavoidable and only partially remedied by our approaches of Lipschitz regularization and gradient damping. These approaches can also lead to slower learning and less expressive energy functions. Methods for efficient (i.e. subquadratic) N -body simulations and for more principled stabilization of deep networks may be relevant to addressing both of these challenges in the future. We described a model for protein structure given sequence information that combines a coarse-grained neural energy function and an unrolled simulation into an end-to-end differentiable model. To realize this idea at the scale of real proteins, we introduced an efficient simulator for Langevin dynamics in transformed coordinate systems and stabilization techniques for backpropagating through long simulator roll-outs. We find that that model is able to predict the structures of protein molecules with hundreds of atoms while capturing structural uncertainty, and that the model can structurally generalize to distant fold classifications more effectively than a strong baseline. (MPNN, bottom left), and outputs energy function weights l as well as simulator hyperparameters (top center). Second, the simulator iteratively modifies the structure via Langevin dynamics based on the gradient of the energy landscape (Forces, bottom center). Third, the imputation network constructs predicted atomic coordinates X from the final simulator time step x (T). During training, the true atomic coordinates X (Data), predicted atomic coordinates X, simulator trajectory x,..., x (T), and secondary structure predictions SS (Model) feed into a composite loss function (Loss, bottom right), which is then optimized via backpropagation. Inverse transformation The inverse transformation z = F 1 (x) involves fully local computations of bong lengths and angles. DISPLAYFORM0 Jacobian The Jacobian @x @z defines the infinitesimal response of the Cartesian coordinates x to perturbations of the internal coordinates z. It will be important for both converting Cartesian forces into angular torques and bond forces as well as the development of our transform integrator. It is defined element-wise as DISPLAYFORM1 The Jacobian has a simple form that can be understood by imagining the protein backbone as a robot arm that is planted at x 0 (Figure 2B). Increasing or decreasing the bond length b i extends or retracts all downstream coordinates along the bonds axis, moving a bond angle a i drives circular motion of all downstream coordinates around the bond normal vectorn i centered at x i 1, and moving a dihedral angle d i drives circular motion of downstream coordinate x j around bond vectorû i 1 centered at x i 1.Unconstrained representations Bond lengths and angles are subject to the constraints b i > 0 and 0 < a i < ⇡. We enforce these constraints by representing these degrees of freedom in terms of fully unconstrained variablesb i andã i via the transformations b i = log ⇣ 1 + eb i ⌘ and a i = ⇡ 1+e ã i. All references to the internal coordinates z and Jacobians @x @z will refer to the use of fully unconstrained representations TAB3. FIG2 provides an overall schematic of the model, including the components of the energy function. CNN primitives All convolutional neural network primitives in the model schematic (FIG2) follow a common structure consisting of stacks of residual blocks. Each residual block includes DISPLAYFORM0 consists of a layer of channel mixing (1x1 convolution), a variable-sized convolution layer, and a second layer of channel mixing. We use dropout with p = 0.9 and Batch Renormalization on all convolutional layers. Batch Renormalization rather than Normalization was necessary rather owing to the large variation in sizes of the structures of the proteins and ing large variation in mini-batch statistics. Why sampling vs. optimization Deterministic methods for optimizing the energy U (x; s) such as gradient descent or quasi-Newton methods can effectively seek local minima of the energy surface, but are challenged to optimize globally and completely ignore the contribution of the widths of energy minima (entropy) to their probability. We prefer sampling to optimization for three reasons: (i) noise in sampling algorithms can facilitate faster global conformational exploration by overcoming local minima and saddle points, (ii) sampling generates populations of states that respect the width (entropy) of wells in U and can be used for uncertainty quantification, and (iii) sampling allows training with an approximate Maximum Likelihood objective (Equation 5).Langevin Dynamics The Langevin dynamics are a stochastic dynamics that sample from the canonical ensemble. They are defined as a continuous-time stochastic differential equation, and are simulated in discrete time with the first order discretization DISPLAYFORM0 Each time step of ✏ involves a descent step down the energy gradient plus a perturbation of Gaussian noise. Importantly, as time tends toward to infinity, the time-distribution of the Langevin dynamics converges to the canonical ensemble. Our goal is to design a dynamics that converge to an approximate sample in a very short period of time. The comparison of this algorithm with naive integration is given in Figure 8. The corrector step is important for eliminating the large second-order errors that arise in curvilinear motions caused by angle changes (Figure 2B and Figure 8). In principle higher-order numerical integration methods or more time steps could increase accuracy at the cost of more evaluations of the Jacobian, but we found that second-order effects seemed to be the most relevant on our timescales. Mixed integrator Cartesian dynamics favor local structural rearrangements, such as the transitioning from a helical to an extended conformation, while internal coordinate dynamics favor global motions such as the change of the overall fold topology. Since both kinds of structural rearrangements are important to the folding process, we form a hybrid integrator (Algorithm 3) by taking one step with each integrator per force evaluation. Translational and rotational detrending Both Cartesian and Internal coordinates are overparameterized with 3L degrees of freedom, since only 3L 6 degrees of freedom are necessary to encode a centered and un-oriented structure 5. As a consequence, a significant fraction of the per time-step changes x can be explained by rigid translational and rotational motions of the entire structure. We isolate and remove these components of motion by treating the system {x 1, . . ., x L} as a set of particles with unit mass, and computing effective structural translational and rotational velocities by summing point-wise momenta. The translational component of motion is simply the average displacement across positions x Trans i = h x i i. For rotational motion around the center of mass, it is convenient to define the non-translational motion as x i = x i x Trans i and the centered Cartesian coordinates asx i = x i hx i i. The point-wise angular momentum is then l i =x i ⇥ x i and we define a total angular velocity of the structure! by summing these and dividing by the moment of inertia as! = (DISPLAYFORM1 Input :Initial state DISPLAYFORM0 Speed clipping We found it helpful to stabilize the model by enforcing a speed limit on overall structural motions for the internal coordinate steps. This prevents small changes to the energy function during learning from causing extreme dynamics that in turn produce a non-informative learning signal. To accomplish this, we translationally and rotationally detrend the update of the predictor step x and compute a hypothetical time step✏ z that would limit the fastest motion to 2 Angstroms per iteration. We then compute modified predictor and corrector steps subject to this new, potentially slower, time step. While this breaks the asymptotics of Langevin dynamics, (i) it is unlikely on our timescales that we achieve stationarity and (ii) it can be avoided by regularizing the dynamics away from situations where clipping is necessary. In the future, considering non-Gaussian perturbations with kinetic energies similar to Relativistic Monte Carlo might accomplish a similar goal in a more principled manner. The final integrator combining these ideas is presented in Figure 3. B APPENDIX B: TRAINING B.1 DATA For a training and validation set, we downloaded all protein domains of length L  200 from Classes ↵,, and ↵/ in CATH release 4.1, and then hierarchically purged a randomly selected set of A, T, and H categories. This created three validation sets of increasing levels of difficulty: H, which contains domains with superfamilies that are excluded from train (but fold topologies may be present), T, which contains fold topologies that were excluded from train (fold generalization), and A which contains secondary structure architectures that were excluded from train. For a test set, we downloaded all folds that were new to CATH release 4.2, which (due to a propensity of structural biology to make new structures of previously solved folds), provided 10,381 test domains. We further stratified this test set into C, A, T, and H categories based on their nearest CATH classification in the training set. We also analyzed test set stratifications based on nearest neighbors in both training and validation in FIG0. We note that the validation set was not explicitly used to tune hyperparameters due to the large cost of training (2 months on 2 M40 GPUs), but we did keep track of validation statistics during training. We optimized all models for 200,000 iterations with Adam . We optimize the model using a composite loss containing several terms, which are detailed as follows. Distance loss We score distances in the model with a contact-focused distance loss DISPLAYFORM0 where the contact-focusing weights are DISPLAYFORM1 is the sigmoid function. Angle loss We use the loss DISPLAYFORM2 where DISPLAYFORM3 T are unit length feature vectors that map the angles {a i, d i} to the unit sphere. Other angular losses, such as the negative log probability of a Von-Mises Fisher distribution, are based on the inner product of the feature vectors H(z a) · H(z b) rather than the Euclidean distance ||H(z a) H(z b)|| between them. It is worth noting that these two quantities are directly related DISPLAYFORM4 Taking z a as fixed and z b as the argument, the Euclidean loss has a cusp at z a whereas the Von-Mises Fisher loss is smooth around z a. This is analogous to the difference between L 1 and L 2 losses, where the cusped L 1 loss favors median behavior while the smooth L 2 loss favors average behavior. Trajectory loss In a further analogy to reinforcement learning, damped backpropation through time necessitates an intermediate loss function that can criticize transient states of the simulator. We compute this by featurizing the per time step coordinates as the product D ijvij (Figure 2C) and doing the same contact-weighted averaging as the distance loss. Template Modelling (TM) Score The TM-score BID27, DISPLAYFORM5 is a measure of superposition quality between two protein structures on that was presented as an approximately length-independent alternative to RMSD. The TM-score is the best attainable value of the preceding quantity for all possible superpositions of two structures, where DISPLAYFORM6 ||. This requires iterative optimization, which we implemented with a sign gradient descent with 100 iterations to optimally superimpose the model and target structure. We backpropagate through this unrolled optimization process as well as that of the simulator. Hydrogen bond loss We determine intra-backbone hydrogen bonds using the electrostatic model of DSSP . First, we place virtual hydrogens at 1 Angstroms along the negative angle bisector of the C i 1 N i C↵ i bond angle. Second, we compute a putative energy U h-bond ij (in kcal/mol) for each potential hydrogen bond from an amide donor at i to a carbonyl acceptor at j as DISPLAYFORM7 = 0.084 DISPLAYFORM8 where D ab = ||X i,a X j,b || is the Euclidean distance between atom a of residue i and atom b of residue j. We then make hard assignments of hydrogen bonds for the data with DISPLAYFORM9 Published as a conference paper at ICLR 2019We'predict' the probabilities of hydrogen bonds of the data given the model via logisitic regression of soft model assignments as DISPLAYFORM10 where a, b, c are learned parameters with the softplus parameterizations enforcing a, b > 0 and (u) = 1/(1 + exp( u) is the sigmoid function. The final hydrogen bond loss is the cross-entropy between these predictions and the data, DISPLAYFORM11 Secondary Structure Prediction We output standard 8-class predictions of secondary structure and score them with a cross-entropy loss. The combination of energy function, simulator, and refinement network can build an atomic level model of protein structure from sequence, and our goal is to optimize (meta-learn) this entire procedure by gradient descent. Before going into specifics of the loss function, however, we will discuss a challenges and solutions for computing gradients of unrolled simulations in the face of chaos. Gradient-based learning of iterative computational procedures such as Recurrent Neural Networks (RNNs) is well known to be subject to the problems of exploding and vanishing gradients BID14. Informally, these occur when the sensitivities of model outputs to inputs become either extremely large or extremely small and the gradient is no longer an informative signal for optimization. We find that backpropagation through unrolled simulations such as those presented is no exception to this rule. Often we observed that a model would productively learn for tens of thousands of iterations, only to suddenly and catastrophically exhibit diverging gradients from which the optimizer could not recover -even when the observed simulation dynamics exhibited no obvious qualitative changes to behavior and the standard solutions of gradient clipping BID14 were in effect. Similar phenomena have been observed previously in the context of meta-learning and are explored in detail in a concurrent work BID12.In FIG5, we furnish a minimal example that illustrates how chaos can lead to irrevocable loss of learning. We see that for even a simple particle-in-a-well, some choices of system parameters (such as too large a time step) can lead to chaotic dynamics which are synonymous with explosive gradients. This example is hardly contrived, and is in fact a simple model of the distance potentials between coordinates in our simulations. Moreover, it is important to note that chaos may not be easy to diagnose: for learning rates ↵ 2 [1.7, 1.8] the position of the particle x remains more or less confined in the well while the sensitivities diverge to 10 200. It seems unlikely that meta-learning would be able to recover after descending into chaos. The view per time step Exploding gradients and chaotic dynamics involve the same mechanism: a multiplicative accumulation of sensitivities. In dynamical systems this is frequently phrased as'exponentially diverging sensitivity to initial conditions'. Intuitively, this can be understood by examining how the Jacobian of an entire trajectory decomposes into a product of Jacobians as DISPLAYFORM0 When the norms of the per time-step Jacobians DISPLAYFORM1 @x (t 1) are typically larger than 1, the sensitivity || @x (T) @x || will grow exponentially with T. Ideally, we would keep these norms well-behaved which is the rationale recent work on stabilization of RNNs (; BID9 . Next we will offer a general-purpose regularizer to approximately enforce this goal for any differentiable computational iteration with continuous state. dx (bottom). When the step size ↵ is small, these dynamics converge to a periodic orbit over 2 k values where 0  k < 1. After some critical step size, the dynamics undergo a period-doubling bifurcation , become chaotic, and the gradients regularly diverge to huge numbers. Approximate Lipschitz conditions One condition that guarantees that a deterministic map F: R N! R N, x t = F (x t 1, ✓) cannot exhibit exponential sensitivity to initial conditions is the condition of being non-expansive (also known as 1-Lipschitz or Metric). That is, for any two input points x a, x b 2 R N, iterating the map cannot increase the distance between them as |F (x a, ✓) DISPLAYFORM2 Repplying the map to the bound immediately implies DISPLAYFORM3 for any number of iterations t. Thus, two initially close trajectories iterated through a non-expansive mapping must remain at least that close for arbitrary time. We approximately enforce non-expansivity by performing an online sensitivity analysis within simulations. At randomly selected time-steps, the current time step x (t) is rolled back to the preceding state and re-executed with small Gaussian perturbations to the state ⇠ N (0, 10 4 I) 6. We regularize the sensitivity by adding DISPLAYFORM4 to the loss. Interestingly, the stochastic nature of this approximate regularizer is likely a good thing -a truly non-expansive map is quite limited in what it can model. However, being'almost' non-expansive seems to be incredibly helpful for learning. Damped Backpropagation through Time The approximate Lipschitz conditions (or Lyapunov regularization) encourage but do not guarantee stable backpropagation. When chaotic phasetransitions or otherwise occur we need a fall-back plan to be able to continue learning. At the same time, we would like gradient descent to proceed in the usual manner when simulator dynamics To reduce our reliance on alignments and the generation of profiles for inference of new sequences while still leveraging evolutionary sequence data, we augmented our training set by dynamically spiking in diverse, related sequence into the model during training. Given a set of M sequences in the alignment we sample a sequence t based on its normalized Hamming distance d t with probability DISPLAYFORM5 where EDA is a scaling parameter that we set to 5. When the alternate sequence contains gaps, we construct a chimeric sequence that substitutes those sites with the query. This strategy increased the number of available sequence-structure pairs by several orders of magnitude, and we used it for both profile and 1-seq based training. C APPENDIX C: For each sequence from the CATH release 4.2 dataset, 100 structures were generated from both the profile and sequence-only models, while a single structure was generated from the RNN baseline models. The reported TM-scores were calculated using Maxcluster BID19. A single representative structure was chosen from the ensemble of 100 structures using 3D-Jury . A pairwise distance matrix of TM-scores was calculated for all of the 100 structures in the ensemble. Clusters were determined by agglomerative hierarchical clustering with complete linkage using a TM-score threshold of 0.5 to determine cluster membership. NEMO RNN, 300 hidden units FIG0: Sampling speed. Per-protein sampling times for various batch sizes across NEMO and one of the RNN baselines on a single Tesla M40 GPU with 12GB memory and 20 cores. For all in the main paper, 100 models were sampled per protein followed by consensus clustering with 3D-jury, adding an additional factor of 10 2 cost between NEMO and the RNN.

Represent the short summary for finding the corresponding scientific article (long)!

We use an unrolled simulator as an end-to-end differentiable model of protein structure and show it can (sometimes) hierarchically generalize to unseen fold topologies.

represent the short summary for finding the corresponding scientific article (long)

We use continuous time dynamics to define a generative model with exact likelihoods and efficient sampling that is parameterized by unrestricted neural networks.


Question: Represent the text to find a very short summary (1-2 sentences max)

Question: Progress in understanding how individual animals learn requires high-throughput standardized methods for behavioral training and ways of adapting training. During the course of training with hundreds or thousands of trials, an animal may change its underlying strategy abruptly, and capturing these changes requires real-time inference of the animal’s latent decision-making strategy. To address this challenge, we have developed an integrated platform for automated animal training, and an iterative decision-inference model that is able to infer the momentary decision-making policy, and predict the animal’s choice on each trial with an accuracy of ~80\%, even when the animal is performing poorly. We also combined decision predictions at single-trial resolution with automated pose estimation to assess movement trajectories. Analysis of these features revealed categories of movement trajectories that associate with decision confidence. We designed our automated training system as follows (

Represent this text to find a much much longer article about it with say hundreds-thousands of words

Automated mice training for neuroscience with online iterative latent strategy inference for behavior prediction

Represent the TLDR for finding the corresponding scientific article.

Inspired by trial-to-trial variability in the brain that can result from multiple noise sources, we introduce variability through noise in the knowledge distillation framework and studied their effect on generalization and robustness.


Question: Represent the text to find a very short summary (1-2 sentences max)

Question: Recurrent neural networks (RNNs) are a powerful tool for modeling sequential data. Despite their widespread usage, understanding how RNNs solve complex problems remains elusive. Here, we characterize how popular RNN architectures perform document-level sentiment classification. Despite their theoretical capacity to implement complex, high-dimensional computations, we find that trained networks converge to highly interpretable, low-dimensional representations. We identify a simple mechanism, integration along an approximate line attractor, and find this mechanism present across RNN architectures (including LSTMs, GRUs, and vanilla RNNs). Overall, these demonstrate that surprisingly universal and human interpretable computations can arise across a range of recurrent networks. Recurrent neural networks (RNNs) are a popular tool for sequence modelling tasks. These architectures are thought to learn complex relationships in input sequences, and exploit this structure in a nonlinear fashion. RNNs are typically viewed as black boxes, despite considerable interest in better understanding how they function. Here, we focus on studying how recurrent networks solve document-level sentiment analysis-a simple, but longstanding benchmark task for language modeling BID6 BID13. We demonstrate that popular RNN architectures, despite having the capacity to implement high-dimensional and nonlinear computations, in practice converge to low-dimensional representations when trained against this task. Moreover, using analysis techniques from dynamical systems theory, we show that locally linear approximations to the nonlinear RNN dynamics are highly interpretable. In particular, they all involve approximate low-dimensional line attractor dynamics-a useful dynamical feature that can be implemented by linear dynamics and used to store an analog value BID10. Furthermore, we show that this mechanism is surprisingly consistent across a range of RNN architectures. We trained four RNN architectures-LSTM BID3, GRU BID0, Update Gate RNN (UGRNN) BID1, and standard (vanilla) RNNs-on binary sentiment classifcation tasks. We trained each network type on each of three datasets: the IMDB movie review dataset, which contains 50,000 highly polarized reviews BID7; the Yelp review dataset, which contained 500,000 user reviews BID14; and the Stanford Sentiment Treebank, which contains 11,855 sentences taken from movie reviews BID11. For each task and architecture, we analyzed the best performing networks, selected using a validation set (see Appendix A for details). We analyzed trained networks by linearizing the dynamics around approximate fixed points. Approximate fixed points are state vectors {h 1 N h − F (h, 0) 2 2, and then minimizing q with respect to hidden states, h, using standard auto-differentiation methods BID2. We ran this optimization multiple times starting from different initial values of h. These initial conditions were sampled randomly from the state activation visited by the trained network, which was done to intentionally sample states related to the operation of the RNN. For brevity, we explain our approach using the working example of the LSTM trained on the Yelp dataset FIG0 ). We find similar for all architectures and datasets, these are shown in Appendix B. As an initial exploratory analysis step, we performed principal components analysis (PCA) on the RNN states concatenated across 1,000 test examples. The top 2-3 PCs explained ∼90% of the variance in hidden state activity FIG0, black line). The distribution of hidden states visited by untrained networks on the same test set was much higher dimensional FIG0, gray line), suggesting that training the networks stretched the geometry of their representations along a lowdimensional subspace. We then visualized the RNN dynamics in this lowdimensional space by forming a 2D histogram of the density of RNN states colored by the sentiment label FIG0. This visualization is reminiscent of integration dynamics along a line attractor-a well-studied mechanism for evidence accumulation in simple recurrent networks BID10 BID8 )-and we reasoned that similar dynamics may be used for sentiment classification. The hypothesis that RNNs approximate line attractor dynamics during sentiment classification makes four specific predictions, which we investigate and confirm in subsequent sections. First, the fixed points form an approximately 1D manifold that is aligned/correlated with the readout weights (Section 3.2). Second, all fixed points are attracting and marginally stable. That is, in the absence of input (or, perhaps, if a string of neutral/uninformative words are encountered) the RNN state should rapidly converge to the closest fixed point and then should not change appreciably (Section 3.4). Third, locally around each fixed point, inputs representing positive vs. negative evidence should produce linearly separable effects on the RNN state vector along some dimension (Section 3.5). Finally, these instantaneous effects should be integrated by the recurrent dynamics along the direction of the 1D fixed point manifold (Section 3.5). We numerically identified the location of ∼500 RNN fixed points using previously established methods BID12 BID2. We then projected these fixed points into the same low-dimensional space used in FIG0. Although the PCA projection was fit to the RNN hidden states, and not the fixed points, a very high percentage of variance in fixed points was captured by this projection FIG0, suggesting that that the RNN states remain close to the manifold of fixed points. We call the vector that describes the main axis of variation of the 1D manifold m. Consistent with the line attractor hypothesis, the fixed points appeared to be spread along a 1D curve when visualized in PC space, and furthermore the principal direction of this curve was aligned with the readout weights FIG0. We further verified that this low-dimensional approximation was accurate by using locally linear embedding (LLE; BID9 to parameterize a 1D manifold of fixed points in the raw, high-dimensional data. This provided a scalar coordinate, θ i ∈ [−1, 1], for each fixed point, which was well-matched to the position of the fixed point manifold in PC space (coloring of points in FIG0). We next aimed to demonstrate that the identified fixed points were marginally stable, and thus could be used to preserve accumulated information from the inputs. To do this, we used a standard linearization procedure BID4 ) to obtain an approximate, but highly interpretable, description of the RNN dynamics near the fixed point manifold. Given the last state h t−1 and the current input x t, the approach is to locally approximate the update rule with a first-order Taylor expansion: DISPLAYFORM0 where ∆h t−1 = h t−1 − h * and ∆x t = x t − x *, and {J rec, J inp} are Jacobian matrices of the system: DISPLAYFORM1 We choose h * to be a numerically identified fixed point and x * = 0, thus we have F (h *, x *) ≈ h * and ∆x t = x t. Under this choice, equation FORMULA0 reduces to a discrete-time linear dynamical system: DISPLAYFORM2 It is important to note that both Jacobians depend on which fixed point we choose to linearize around, and should thus be thought of as functions of h *; for notational simplicity we do not denote this dependence explicitly. By reducing the nonlinear RNN to a linear system, we can analytically estimate the network's response to a sequence of T inputs. In this approximation, the effect of each input x t is decoupled from all others; that is, the final state is given by the sum of all individual effects. 1 We can restrict our focus to the effect of a single input, x t (i.e. a single term in this sum). Let k = T − t be the number of time steps between x t and the end of the document. The total effect of x t on the final RNN state becomes: DISPLAYFORM3 1 We consider the case where the network has closely converged to a fixed point, so that h0 = h * and thus ∆h0 = 0.where L = R −1, the columns of R (denoted r a) contain the right eigenvectors of J rec, the rows of L (denoted a) contain the left eigenvectors of J rec, and Λ is a diagonal matrix containing complex-valued eigenvalues, λ 1 > λ 2 >... > λ N, which are sorted based on their magnitude. From equation 3 we see that x t affects the representation of the network through N terms (called the eigenmodes or modes of the system). The magnitude of each mode after k steps is given by the λ k a; thus, the size of each mode either reduces to zero or diverges exponentially fast, with a time constant given by: τ a = 1 log(|λa|). This time constant has units of tokens (or, roughly, words) and yields an interpretable number for the effective memory of the system. FIG1 plots the eigenvalues and associated time constants and shows the distribution of all eigenvalues at three representative fixed points along the fixed point manifold FIG1. In FIG1, we plot the decay time constant of the top three modes; the slowest decaying mode persists after ∼1000 time steps, while the next two modes persists after ∼100 time steps, with lower modes decaying even faster. Since the average review length for the Yelp dataset is ∼175 words, only a small number of modes could represent information from the beginning of the document. Overall, these eigenvalue spectra are consistent with our observation that RNN states only explore a low-dimensional subspace when performing sentiment classification. RNN activity along the majority of dimensions is associated with fast time constants and is therefore quickly forgotten. Restricting our focus to the top eigenmode for simplicity (there may be a few slow modes of integration), the effect of a single input, x t, on the network activity (equation 3) becomes: r 1 1 J inp x, where we have dropped the dependence on t since λ 1 ≈ 1, so the effect of x is largely insensitive to the exact time it was input to system. Using this expression, we separately analyzed the effects of specific words. We first examined the term J inp x for various choices of x (i.e. various word tokens). This quantity represents the instantaneous linear effect of x on the RNN state vector and is shared across all eigenmodes. We projected the ing vectors onto the same low-dimensional subspace shown in FIG0. We see that positive and negative valence words push the hidden state in opposite directions. Neutral words, in contrast, exert much smaller effects on the RNN state FIG2.While J inp x represents the instantaneous effect of a word, only the features of this input that overlap with the top few eigenmodes are reliably remembered by the network. The scalar quantity 1 J inp x, which we call the input projection, captures the magnitude of change induced by x along the eigenmode associated with the longest timescale. Again we observe that the valence of x strongly correlates with this quantity: neutral words have an input projection near zero while positive and negative words produced larger magnitude responses of opposite sign. Furthermore, this is reliably observed across all fixed points. FIG2 shows the average input projection for positive, negative, and neutral words; the histogram shows the distribution of these average effects across all fixed points along the line attractor. Finally, if the input projection onto the top eigenmode is non-negligible, then the right eigenvector r 1 (which is normalized to unit length) represents the direction along which x is integrated. If the RNN implements an approximate line attractor, then r 1 (and potentially other slow modes) should align with the principal direction of the manifold of fixed points, m. We indeed observe a high degree of overlap between r 1 and m both visually in PC space FIG2 and quantitatively across all fixed points FIG2. In this work we applied dynamical systems analysis to understand how RNNs solve sentiment analysis. We found a simple mechanismintegration along a line attractorpresent in multiple architectures trained to solve the task. Overall, this work provides preliminary, but optimistic, evidence that different, highly intricate network models can converge to similar solutions that may be reduced and understood by human practitioners. An RNN is defined by a nonlinear update rule h t = F (h t−1, x t), which is applied recursively from an initial state h 0 over a sequence of inputs x 1, x 2,..., x T. Let N and M denote the dimensionality of the hidden states and the input vectors, so that h t ∈ R N and x t ∈ R M. In sentiment classification, T represents the number of word tokens in a sequence, which can vary on a document-by-document basis. To process word sequences for a given dataset, we build a vocabulary and encode words as one-hot vectors. These are fed to a dense linear embedding layer with an embedding size of M = 128 (x t are the embeddings in what follows). The word embeddings were trained from scratch simultaneously with the RNN. We considered four RNN architectures-LSTM BID3, GRU BID0, UGRNN BID1, and the vanilla RNN (VRNN)-each corresponding to a separate nonlinear update rule, F (·, ·). For the LSTM architecture, h t consists of a concatenated hidden state vector and cell state vector so that N is twice the number of computational units; in all other architectures N is equal to the number of units. The RNN prediction is evaluated at the final time step T, and is given byŷ = w h T + b, where we call w ∈ R N the readout weights. In the LSTM architecture, the cell state vector is not read out, and thus half of the entries in w are enforced to be zero under this notation. We examined three benchmark datasets for sentiment classification: the IMDB movie review dataset, which contains 50,000 highly polarized reviews BID7; the Yelp review dataset, which contained 500,000 user reviews BID14; and the Stanford Sentiment Treebank, which contains 11,855 sentences taken from movie reviews BID11. The Stanford Sentiment Treebank also contains short phrases with labeled sentiments; these were not analyzed. For each task and architecture, we performed a randomized hyper-parameter search and selected the best networks based on a validation set. All models were trained using Adam BID5 ) with a batch size of 64. The hyper-parameter search was performed over the following ranges: the initial learning rate (10 −5 to 10 −1), learning rate decay factor (0 to 1), gradient norm clipping (10 −1 to 10), 2 regularization penalty (10 −3 to 10 −1), the β 1 (0.5 to 0.99), and β 2 (0.9999 to 0.99) parameters of the Adam optimization routine. We additionally trained a bag-of-words model (logistic regression trained with word counts), as a baseline comparison. The accuracies of our final models on the held-out test set are summarized in Table 1. We analyzed the best performing models for each combination of architecture type and dataset. For each model, we numerically identified a large set of fixed points {h * BID12 . Briefly, we accomplished this by first defining the loss function q = 1 N h−F (h, 0) 2 2, and then minimizing q with respect to hidden states, h, using standard auto-differentiation methods BID2. We ran this optimization multiple times starting from different initial values of h. These initial conditions were sampled randomly from the state activation during the operation of the trained network, which was done to intentionally sample states related to the operation of the RNN. We varied the stopping tolerance for q using 9 points logarithmically spaced between 10 −9 and 10 −5 running the optimization from 1000 different initial conditions for each tolerance. This allowed us to find approximate fixed points of varying speeds. Values of q at numerical zero are true fixed points, while small but non-zero values are called slow points. Slow points are often reasonable places to perform a linearization, assuming that √ q, which is akin to speed, is slow compared to the operation of the network. DISPLAYFORM0 Below we provide figures summarizing the linear integration mechanism for each combination of architecture (LSTMs, GRUs, Update Gate RNNs, Vanilla RNNs) and dataset (Yelp, IMDB, and Stanford Sentiment). Note that the first figure, LSTMs trained on Yelp, reproduces the figures in the main text-we include it here for completeness. The description of each panel is given in FIG0, note that these descriptions are the same across all figures. We find that these mechanisms are remarkably consistent across architectures and datasets. (lower left) Instantaneous effect of word inputs, J inp x, for positive (green), negative (red), and neutral (cyan) words. Blue arrows denote 1, the top left eigenvector. The PCA projection is the same as FIG1, but centered around each fixed point.(lower middle left) Average of 1 J inp x over 100 different words, shown for positive, negative, neutral words.(lower middle right) Same plot as in FIG1, with an example fixed point highlighted (approximate fixed points in grey). Blue arrows denote r1, the top right eigenvector.(lower right) Distribution of r 1 m (overlap of the top right eigenvector with the fixed point manifold) over all fixed points. Null distribution is randomly generated unit vectors of the size of the hidden state.

Represent the short summary for finding the corresponding scientific article (long)

We analyze recurrent networks trained on sentiment classification, and find that they all exhibit approximate line attractor dynamics when solving this task.

Represent this text to find a much much longer article about it with say hundreds-thousands of words

We show shortcut connections should be placed in patterns that minimize between-layer distances during backpropagation, and design networks that achieve log L distances using L log(L) connections.