summaries: The authors propose an algorithm that takes several neural networks and outputs an embedding for such networks and a meta-model. This meta-model can take any embedding of a network, and emulate its hidden states and outputs. The authors suggest that these embeddings can measure models similarity, and show how to interpolate and extrapolate between and from the models in the training set.

summaries: This paper proposes an approach to combine together multiple (pretrained) neural networks into a single model that can simulate each of the considered neural networks, both in terms of output and hidden states. The setup here relies on all the considered networks to operate on the same task, and in particular the same input data type (but perhaps with different sampled training sets). The combined meta-model then considers as inputs not only the input data, but also an encoding (which is learned during its training) of each model. This allows model-dependent latent features to be computed by the meta-model, which can then be decoded to hidden states of each network via specialized decoders that can be thought of as auxiliary computational pathways that operate in parallel to the one producing the main output of the meta-model. The model encoding are then used to visualize and explore the relations between different networks (or instantiations), and furthermore, to further tune the combined model towards an interpolation or extrapolation that goes beyond (and seems to outperform) each of the individual networks used to train the meta-model.

summaries: This paper introduces a new approach to studying the population of trained neural networks. Specifically, the author's design and train a meta-model which when given a vector, the meta-model emulates the model associated with the given vector. The authors present a few applications of their framework, including clustering and semi-supervised learning.

summaries: The manuscript proposes to analyze the similarity of different neural networks trained to solve the same task by developing a meta-model capable of reproducing/emulating both the output and the hidden representation of each individual network.  This meta-model is parametrized by a set of parameters which is common to all the networks, and by other parameters which are network-specific (theta). In this way, each network will correspond to a different theta, and one can analyze the global features of the set of networks by unsupervised manifold learning in the  theta-space.

abstractText: A core challenge in the interpretation of deep neural networks is identifying commonalities between the underlying algorithms implemented by distinct networks trained for the same task. Motivated by this problem, we introduce DYNAMO, an algorithm that constructs low-dimensional manifolds where each point corresponds to a neural network model, and two points are nearby if the corresponding neural networks enact similar high-level computational processes. DYNAMO takes as input a collection of pre-trained neural networks and outputs a meta-model that emulates the dynamics of the hidden states as well as the outputs of any model in the collection. The specific model to be emulated is determined by a model embedding vector that the meta-model takes as input; these model embedding vectors constitute a manifold corresponding to the given population of models. We apply DYNAMO to both RNNs and CNNs, and find that the resulting model embedding spaces enable novel applications: clustering of neural networks on the basis of their high-level computational processes in a manner that is less sensitive to reparameterization; model averaging of several neural networks trained on the same task to arrive at a new, operable neural network with similar task performance; and semi-supervised learning via optimization on the model embedding space. Using a fixed-point analysis of meta-models trained on populations of RNNs, we gain new insights into how similarities of the topology of RNN dynamics correspond to similarities of their high-level computational processes.

authors:
Kyle Aitken
Vinay V Ramasesh
Ankush Garg
Yuan Cao
David Sussillo
Niru Maheswaranathan

title: The geometry of integration in text classification rnns

venue: arXiv preprint arXiv:2010.15114,

authors:
Piotr Bojanowski
Armand Joulin
David Lopez-Pas
Arthur Szlam

title: Optimizing the latent space of generative networks

venue: In International Conference on Machine Learning,

authors:
Ekin D Cubuk
Barret Zoph
Jonathon Shlens
Quoc V Le

title: Randaugment: Practical automated data augmentation with a reduced search space

venue: In IEEE Conference on Computer Vision and Pattern Recognition Workshops,

authors:
Chelsea Finn
Pieter Abbeel
Sergey Levine

title: Model-agnostic meta-learning for fast adaptation of deep networks

authors:
Takashi Fukuda
Masayuki Suzuki
Gakuto Kurata
Samuel Thomas
Jia Cui
Bhuvana Ramabhadran

title: Efficient knowledge distillation from an ensemble of teachers

authors:
Matthew D Golub
David Sussillo

title: Fixedpointfinder: A tensorflow toolbox for identifying and characterizing fixed points in recurrent neural networks

venue: Journal of Open Source Software,

authors:
Geoffrey Hinton
Oriol Vinyals
Jeffrey Dean

title: Distilling the knowledge in a neural network

venue: In NIPS Deep Learning and Representation Learning Workshop,

authors:
P Izmailov
AG Wilson
D Podoprikhin
D Vetrov
T Garipov

title: Averaging weights leads to wider optima and better generalization

venue: In Conference on Uncertainty in Artificial Intelligence,

authors:
Anatole Katok
Boris Hasselblatt

title: Introduction to the modern theory of dynamical systems. Number 54

authors:
Simon Kornblith
Mohammad Norouzi
Honglak Lee
Geoffrey Hinton

title: Similarity of neural network representations revisited

authors:
Karel Lenc
Andrea Vedaldi

title: Understanding image representations by measuring their equivariance and equivalence

venue: In IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Yixuan Li
Jason Yosinski
Jeff Clune
Hod Lipson
John E Hopcroft

title: Convergent learning: Do different neural networks learn the same representations

venue: In International Conference on Learning Representations,

authors:
Ilya Loshchilov
Frank Hutter

title: Decoupled weight decay regularization

authors:
Andrew Maas
Raymond E Daly
Peter T Pham
Dan Huang
Andrew Y Ng
Christopher Potts

title: Learning word vectors for sentiment analysis

venue: In Annual Meeting of the Association for Computational Linguistics,

authors:
Niru Maheswaranathan
Alex H Williams
Matthew D Golub
Surya Ganguli
David Sussillo

title: Reverse engineering recurrent networks for sentiment classification reveals line attractor dynamics

venue: Advances in Neural Information Processing Systems,

title: Universality and individuality in neural dynamics across large populations of recurrent networks

authors:
Ari Morcos
Maithra Raghu
Samy Bengio

title: Insights on representational similarity in neural networks with canonical correlation

authors:
Maithra Raghu
Justin Gilmer
Jason Yosinski
Jascha Sohl-Dickstein

title: SVCCA: Singular Vector Canonical Correlation Analysis for Deep Learning Dynamics and Interpretability

venue: In International Conference on Neural Information Processing Systems,

authors:
Sachin Ravi
Hugo Larochelle

title: Optimization as a model for few-shot learning

title: bengio, université de montréal, adriana romero, nicolas ballas, samira ebrahimi kahou, antoine chassang, carlo gatta, and yoshua bengio. fitnets: Hints for thin deep nets

authors:
Andrei A Rusu
Dushyant Rao
Jakub Sygnowski
Oriol Vinyals
Razvan Pascanu
Simon Osindero
Raia Hadsell

title: Meta-Learning with Latent Embedding Optimization

authors:
Adam Santoro
Sergey Bartunov
Matthew Botvinick
Daan Wierstra
Timothy Lillicrap

title: Metalearning with memory-augmented neural networks

authors:
Karen Simonyan
Andrea Vedaldi
Andrew Zisserman

title: Deep inside convolutional networks: Visualising image classification models and saliency maps

venue: arXiv preprint arXiv:1312.6034,

authors:
David Sussillo
Omri Barak

title: Opening the black box: low-dimensional dynamics in highdimensional recurrent neural networks

venue: Neural Computation,

authors:
David Sussillo
Mark M Churchland
Matthew T Kaufman
Krishna V Shenoy

title: A neural network that finds a naturalistic solution for the production of muscle activity

venue: Nature Neuroscience,

authors:
Daniel LK Yamins
James J DiCarlo

title: Using goal-driven deep learning models to understand sensory cortex

authors:
Matthew D Zeiler
Rob Fergus

title: Visualizing and understanding convolutional networks

venue: In European Conference on Computer Vision,

authors:
Xiang Zhang
Junbo Zhao
Yann LeCun

title: Character-level convolutional networks for text classification

heading: 1 INTRODUCTION

text:
A crucial feature of neural networks with fixed network architecture is that they form a manifold by virtue of their continuously tunable weights, which underlies their capability to be trained by gradient descent. However, this conception of the space of neural networks is inadequate for understanding the computational processes they perform. For example, two neural networks trained to perform the same task may have vastly different weights, and yet implement the same high-level algorithms and computational processes (Maheswaranathan et al., 2019b).
In this paper, we construct an algorithm which provides alternative parametrizations of the space of RNNs and CNNs with the goal of endowing a geometric structure that is more compatible with the high-level computational processes performed by neural networks. In particular, given a set of neural networks with possibly different architectures (and possibly trained on different tasks), we find a parametrization of a low-dimensional submanifold of neural networks which approximately interpolates between these chosen “base models”, as well as extrapolates beyond them. We can use such model embedding spaces to cluster neural networks and even compute model averages of neural networks. A key feature is that two points in model embedding space are nearby if they correspond to neural networks which implement similar high-level computational processes, in a manner to be described later. In this way, two neural networks may correspond to nearby points in model embedding space even if those neural networks have distinct weights or even architectures.
The model embedding space is parametrized by a low-dimensional parameter θ ∈ Rd, and each base model is assigned a value of θ in the space. This allows us to apply traditional ideas from clustering and interpolation to the space of neural networks. Moreover, each model embedding space has an associated meta-model which, upon being given a θ, is rendered into an operable neural network. If a base model is mapped to some θ, then the meta-model, upon being given that θ, will emulate the
corresponding base model. See Figure 1 for a diagrammatic depiction. An interesting application is that given two models assigned to parameters θ1 and θ2, we can consider the averaged model corresponding to the value (θ1 + θ2)/2. We find that this averaged model performs similarly on the task for which the base models in the average were trained. We also use the model embedding space to extrapolate outside the space of base models, and find cases in which the model embedding manifold specifies models that perform better on a task than any trained base model. Later on we will explain how this can be regarded as a form of semi-supervised learning.
The rest of the paper is organized as follows. We first provide the mathematical setup for our algorithmic construction of model embedding spaces. After reviewing related work, we then present results of numerical experiments which implement the algorithm and explore clustering, model averaging, and semi-supervised learning on the model embedding space. We further examine how topological features of the dynamics of RNNs in model embedding spaces are reflective of classes of high-level computation processes. Finally, we conclude with a discussion.

heading: 2 DYNAMICAL MODEL EMBEDDING

heading: 2.1 MATHEMATICAL SETUP

text:
In this section, we provide the mathematical setup for construction of model embedding spaces. We treat this in the RNN setting, and relegate the CNN setting to Section A of the Appendix.
Algorithm 1 DYNAMO algorithm Input: Base models {(Fn, Gn)}Nn=1, loss functions {(Ln,hidden,Ln,output)}Nn=1,
datasets {Dn}Nn=1, output loss weight λ ≥ 0 Output: Meta-model (F̃ , G̃), state maps {Vn}Nn=1, model embeddings {θn}Nn=1 Initialize networks F̃ , G̃, {Vn}Nn=1 Initialize model embeddings θn = 0 for τ = 1, . . . , τmax do
Sample base model (Fi, Gi) uniformly at random Sample input {xt} from dataset Di Compute base model hidden states {ht} and outputs {Gi(ht)} Compute meta-model hidden states {hθi,t} and outputs {G̃(hθi,t)} Compute mapped meta-model states {Vi(hθi,t)} Compute loss L̂ = Li,hidden + λLi,output Update F̃ , G̃, Vi, θi using∇L̂
end for return (F̃ , G̃), {Vn}, {θn}
Notation. Let an RNN be denoted by F (x, h) where x is the input and h is the hidden state. We further denote the hidden-to-output map byG(h). We consider a collection ofN RNNs {(Fn, Gn)}Nn=1 we call the base models which may each have distinct dimensions for their hidden states, but all have the same dimension for their inputs as well as the same dimension for their outputs. A sequence of inputs is notated as {xt}Tt=0 which induces a sequence of hidden states by ht+1 = F (xt, ht) where the initial hidden state h0 is given. A collection of sequences of inputs, possibly each with different maximum lengths T , is denoted by D which we call an input data set. We suppose that each base model RNN has an associated input data set.

heading: 2.2 META-MODELS FOR RNNS

text:
Given a collection base model RNNs, we would like to construct a meta-model which emulates the behavior of each of the base models. In this case, the meta-model is itself an RNN with one additional input θ ∈ Rd and a corresponding map F̃ (θ, x, h) whose output is the next hidden state. Given a sequence of input states {xt}Tt=0, we have a corresponding sequence of output states hθ,t+1 = F̃ (θ, xt, hθ,t) starting from an initial hidden state h0 (which we suppose does not depend on θ). The meta-model also includes a hidden-to-output map G̃(h) that is independent of θ.
For the meta-model (F̃ , G̃) to emulate a particular base model (Fn, Gn) with respect to its corresponding data set Dn, we consider the following criteria: there is some θn for which
1. G̃(hθn,t) ≈ Gn(ht) for all t > 0 and all input sequences in the data set
2. Vn(hθ,t) ≈ ht for all t > 0 and all input sequences in the data set,
where Vn is a transformation of a meta-model’s hidden activity. We emphasize that θn and Vn depend on the particular base model under consideration. The first criterion means that at some particular θn, the outputs of the meta-model RNN dynamics are close to the outputs of the base model RNN dynamics. The second criterion means that at the same θn, there is a time-independent transformation Vn (i.e., Vn does not depend on t) such the transformed hidden state dynamics of the meta-model are close to the hidden state dynamics of the base model. See Figure 2 for a visualization. As depicted in the figure, it is convenient to regard the meta-model RNN as having input states (θn, x). As such, a sequence of inputs {xt}Tt=1 is appended by θn to become {(θn, xt)}Tt=1.
The desired properties of the meta-model are enforced in the loss function. Defining the functions
Loutput[F̃ , G̃, θn] := 1
T T∑ t=1 d(G̃(hθn,t), Gn(ht)) (1)
Lhidden[F̃ , θn, Vn] := 1
T T∑ t=1 ‖Vn(hθn,t)− ht‖ 2 2 (2)
where d is some suitable distance or divergence, we can construct the loss function E{xt}∼Dn [ Lhidden[F̃ , θn, Vn] + λLoutput[F̃ , G̃, θn] ] (3)
where we average over the choice of sequence {xt} coming from the input data set Dn. Above, λ is a hyperparameter. Our aim is to minimize the above over a suitable class of F̃ , G̃, Vn, as well as θn; this can be implemented computationally via the DYNAMO algorithm (see Algorithm 1). As a side remark, it appears that a suitable alternative choice to equation 2 would be 1 T ∑T t=1 ‖hθn,t −Wn(ht)‖ 2 2 where here Wn is a map from the hidden states of the base model to the hidden states of the meta-model. However, this would be problematic since minimization may pressure Wn to be the the zero map (or otherwise have outputs which are small in norm) and accordingly pressure the dynamics of the meta-model to be the trivial (or have small norm). As such, we opt to formulate Lhidden as it is written in equation 2. Suppose we want the meta-model to be able to emulate an entire collection of base models {(Fn, Gn)}Nn=1. In particular, the meta-model will attempt to assign to the nth base model a θn and Vn so that the two criteria listed above are satisfies for that base model. These desiderata can be implemented by minimizing the loss function
L[F̃ , G̃, {θn}Nn=1, {Vn}Nn=1] := 1
N N∑ n=1 E{xt}∼Dn [ Ln,hidden[F̃ , θn, Vn] + λLn,output[F̃ , G̃, θn] ] .
(4)
In some circumstances, we may want to consider base models with distinct dimensions for their output states. For instance, suppose half of the base models perform a task with outputs in Rd1 and the other half perform a task without outputs in Rd2 . To accommodate for this, we can have two hidden-to-output maps G̃1, G̃2 for the meta-model, where the maps have outputs in Rd1 and Rd2 respectively. The loss function is slightly modified so that we use G̃1 where compare the metamodel to the first kind of base model, and G̃2 when we compare the meta-model to the second kind of base model. This construction generalizes to the setting where the base models can be divided up into k groups with distinct output dimensions; this would necessitate k hidden-to-output functions G̃1, ..., G̃k for the meta-model.

heading: 3 RELATED WORK

text:
There is a substantial body of work on interpreting the computational processes implemented neural networks by studying their intermediate representations. Such analyses can be performed either on individual models (Simonyan et al., 2013; Zeiler & Fergus, 2014; Lenc & Vedaldi, 2015) or on collections of models (Li et al., 2016; Raghu et al., 2017; Morcos et al., 2018; Kornblith et al., 2019). Prior work in this latter category has focused on pairwise comparisons between models. For example, SVCCA (Raghu et al., 2017) uses canonical correlation analysis to measure the representational similarity between pairs of models. While these methods can also be used to derive model representations by embedding the pairwise distance matrix, our approach does not require the Θ(N2) computational cost of comparing all pairs of base models. Moreover, DYNAMO yields an executable meta-model that can be called with model embedding vectors other than those corresponding to the base models.
There is a related body of work in the field of computational neuroscience. CCA based techniques and representational geometry are standard for comparing neural networks to the neural activations of animals performing vision tasks (Yamins & DiCarlo, 2016) as well as motor tasks. In an example
of the latter, the authors of (Sussillo et al., 2015) used CCA techniques to compare brain recordings to those of neural networks trained to reproduce the reaching behaviors of animals, while the authors of (Maheswaranathan et al., 2019b) used fixed point analyses of RNNs to consider network similarity from a topological point of view.
DYNAMO can be viewed as a form of knowledge distillation (Hinton et al., 2015), since the outputs of the base models serve as targets in the optimization of the meta-model. However, unlike typical instances of knowledge distillation involving an ensemble of teacher networks (Hinton et al., 2015; Fukuda et al., 2017) where individual model predictions are averaged to provide more accurate target labels for the student, our approach instead aims to preserve the dynamics and outputs of each individual base model. FitNets (Romero, 2015) employ a form a knowledge distillation using maps between hidden representations to guide learning; this is similar to our use of hidden state maps in training the meta-model.
Our treatment of the model embedding vectors {θn} as learnable parameters is similar to the approach used in Generative Latent Optimization (Bojanowski et al., 2018), which jointly optimizes the image generator network and the latent vectors corresponding to each image in the training set. Bojanowski et al. (2018) find that the principal components of the image representation space found by GLO are semantically meaningful. We likewise find that the principal components of the model embeddings found by DYNAMO are discriminative between subsets of models.
Unlike methods such as hypernetworks (Ha et al., 2017) and LEO (Rusu et al., 2019) that use a model to generate parameters for a separate network, our approach does not attempt to reproduce the parameters of the base models in the collection. Instead, the meta-model aims to reproduce only the hidden states and outputs of a base model when conditioned on the corresponding embedding vector.
The core focus of this work also differs from that of the meta-learning literature (Santoro et al., 2016; Ravi & Larochelle, 2017; Finn et al., 2017; Munkhdalai & Yu, 2017), which primarily concerns the problem of few-shot adaptation when presented with data from a new task. Our empirical study centers on the post-hoc analysis of a given collection of models, which may or may not have been trained on different tasks. However, we remark that our exploration of optimization in lowdimensional model embedding space is related to LEO (Rusu et al., 2019), where a compressed model representation is leveraged for efficient meta-learning.

heading: 4 EMPIRICAL RESULTS

text: In this section, we describe the results of our empirical study of meta-models trained using DYNAMO on collections of RNNs trained on NLP tasks, and collections of CNNs trained for image classification. For RNN base models, we parameterize the meta-model as a GRU where the model embedding vector θ is presented as an additional input in each time step. For CNN base models, we use a ResNet meta-model where θ is an additional input for each ResNet block. In all our experiments, we hold out half the available training data for use as unlabeled data for training the meta-model; the base models were trained on the remaining training data (or a fraction thereof).1 By default, we set the output loss hyperparameter λ to 1. We defer further details on model architectures and training to the Appendix.

heading: 4.1 VISUALIZING MODEL SIMILARITY IN EMBEDDING SPACE

text:
The base model embeddings {θn}Nn=1 can be used for cluster analysis to evaluate the similarity structure of a collection of models. We illustrate this use case of DYNAMO via a series of example applications on NLP and vision tasks. Figure 3 shows model embeddings learned by DYNAMO on collections of RNNs. In these plots, we observe a clear separation of these networks according to the size of the available training data, the RNN model architecture used, and the specific NLP task used for training. By computing the eigenvalues of the covariance matrix corresponding to the N model embeddings, we obtain a measure of the intrinsic dimensionality of the corresponding collections of
1For example, our IMDB sentiment base models trained on 100% of the available training data were trained on 12,500 examples, with the remaining 12,500 examples used as unlabeled data for training the meta-model.
models. In Figure 3, we find that 2 to 6 components are sufficient to explain 95% of the variance in the model embeddings.
By tuning the output loss hyperparameter λ in equation 4, we can adjust the degree of emphasis placed on reproducing the hidden dynamics of the base networks versus their outputs. We demonstrate this effect in Figure 4: with λ = 1, we observe clear separation of ResNet-34 models trained on CIFAR-100 with different data augmentation policies (“weak”, with shifts and horizontal flips vs. “strong”, with RandAugment (Cubuk et al., 2020)), and with different training set sizes. In contrast, with λ = 0, we observe a weak clustering effect with differing data augmentation, and no detectable separation with different training set sizes. We infer that the change in data augmentation policy re-
sults in a larger difference in the learned feature representations than the change in training set size. In Appendix B, we illustrate the effect of setting λ = 0 for RNN models.
We additionally compare the embeddings obtained using DYNAMO to those derived from SVCCA (Raghu et al., 2017), a pairwise comparison technique that aligns the representations produced by a pair of networks using canonical correlation analysis (CCA). In Figure 5, we plot the 2D embeddings obtained using multidimensional scaling (MDS) on the pairwise distance matrix computed using SVCCA (modified to output L2 distances instead of correlations). Unlike the principal components of the DYNAMO model embeddings plotted in Figure 3, the MDS coordinates are not semantically interpretable. Additionally, the cluster structure of the collection of GRUs trained with varying training set sizes is less apparent in this representation.
Lastly, we note that DYNAMO allows for flexibility in defining the metric used to compare the hidden states and outputs of the base models with those of the meta-model. In Appendix B.3, we demonstrate the benefit of using the L1 distance for clustering CNN representations.

heading: 4.2 EXTRAPOLATION BEYOND BASE MODEL EMBEDDINGS

text:
We study the model embedding space corresponding to a trained meta-model by conditioning it on model embedding vectors θ other than those assigned to the set of base models. Figure 6 visualizes the landscape of test accuracies for two meta-models: (i) a meta-model for 10 GRUs trained with
50% of the IMDB training data, and 10 GRUs trained with 25% of the data, and (ii) a meta-model for 10 IMDB sentiment GRUs and 10 AG News classification GRUs.
We note two particularly salient properties of these plots. First, the test accuracy varies smoothly when interpolating θ between pairs of base model embeddings—we would in general not observe this property when interpolating the parameters of the base GRUs directly, since they were trained with different random initializations and orderings of the training examples. Second, we observe that the embedding vector that realizes the highest test accuracy lies outside the convex hull of the base model embeddings. This is perhaps surprising since typical training and inference protocols involve the use of convex combinations of various objects: for instance, averaging of predictions in model ensembles, and averaging of model parameters during training (e.g., using exponential moving averages, or with Stochastic Weight Averaging (Izmailov et al., 2018)). This extrapolatory phenomenon suggests that DYNAMO is able to derive a low-dimensional manifold of models that generalizes beyond the behavior of the base models used for training the meta-model.

heading: 4.3 SEMI-SUPERVISED LEARNING IN MODEL EMBEDDING SPACE

text:
The existence of model embeddings that improve on the accuracy of the base models suggests a natural semi-supervised learning (SSL) procedure involving a trained meta-model. In particular, we minimize the loss incurred by the meta-model on a small set of additional labeled examples by optimizing the value of θ. This is done by backpropagating gradients through the meta-model, with the meta-model parameters held fixed. Figure 7 shows the result of this procedure on an IMDB sentiment meta-model (previously depicted in Figure 6) with a set of additional labeled examples (disjoint from the test set) of size equal to 1% of the full training set. This procedure successfully finds a θ that improves on the test accuracy of the best base model by 6% (86.4% vs. 80.3%).
We observe that this SSL procedure achieves lower accuracy when we train the meta-model using fewer base models. In particular, a metamodel only on the 10 GRUs trained with 50% of the training data yields a test accuracy of
85.4%, and a meta-model on only a single GRU out of these 10 yields 81.6%. This result suggests that a diversity of base models helps improve the accuracy achievable by the meta-model.

heading: 5 DYNAMICS OF META-MODELS FOR RNNS

text:
In this section we perform an analysis of the dynamical features of the dynamics generated by the meta-models trained on base models that perform the sentiment classification task. The sentiment classification task has a well-understood dynamical structure (Sussillo & Barak, 2013; Maheswaranathan et al., 2019a;b; Aitken et al., 2020) that we can use as a basis for understanding the behavior of a corresponding meta-model. To a first approximation, the sentiment analysis task can be solved by a simple integrator that accumulates sentiment corresponding to each word (positive words such as ‘good’ or ‘fantastic’ adding positive sentiment, and negative words such as ‘bad’ or ‘terrible’ adding negative sentiment). It has been shown that simple sentiment analysis models approximate this integrator by constructing a line attractor in the space of hidden states. For instance, for the zero input x∗ = ~0, it has been observed that the dynamical system generated by the map Fx∗(h) = F (x
∗, h) has a tubular region with very little movement, in the sense that ‖Fx∗(h)− h‖2 is very small for hidden states h in this region.
To investigate the dynamical behavior of the meta-model space, we trained a meta-model on a set of 20 base models which were themselves trained on the IMDB sentiment analysis task. Of these 20 base models, 10 were trained with 50% of the available training data and the remaining 10 were trained on 100% of the training data. The θ points corresponding to these base models clustered in the model embedding space according to the amount of training data. In Figure 8 we perform a fixed-point analysis of several models corresponding to points of interest in the model embedding space.
The fixed-point analysis was run according to the procedure described in (Golub & Sussillo, 2018). First we selected a set of candidate hidden states hj by running the model on a typical batch of inputs. For each hidden state hj obtained in this way, we used gradient descent on the loss ‖F (x∗, h)−h‖22 to find the nearest approximate fixed point.
An interesting finding is that the meta-model found line attractor structures that were very geometrically similar for models within a cluster. An interpretation of this result pertaining to topological conjugacy in dynamical systems theory is discussed in Section C of the Supplementary Materials. Moreover, we find that the meta-model finds a continuous interpolation between line attractors that are relatively fat and thin (corresponding to models trained on 50% of the data), and models that are elongated (corresponding to models trained on 100% of the data).

heading: 6 DISCUSSION

text:
We have introduced the algorithm DYNAMO, which maps a set of neural network base models to a low dimensional feature space. Our results show that the model embeddings provided by DYNAMO capture relevant computational features of the base models. Moreover, the model embedding spaces produced by DYNAMO are a sufficiently smooth that model averaging can be performed, and even model extrapolation to reach new models with better performance than any base model. Indeed, in our experiments where the base models were trained on the sentiment analysis task, the model embedding space describes a space of line attractors which vary smoothly in the parameter θ.
We have demonstrated that DYNAMO can be applied broadly to neural networks that have a dynamical structure; for example, we used the demarcation of layers of convolutional neural networks as a proxy for a dynamical time variable. This also suggests possible scientific applications of DYNAMO to dynamical systems arising in nature. A present limitation of DYNAMO is the need for all base models to have the same input structure. For example, one cannot presently utilize DYNAMO
to compare language models trained with different encodings (character-based vs. word-based, for example).

heading: A META-MODELS FOR CNNS

text:
Our setup in Section 2.2 above can be readily adapted to CNNs, or feedforward neural networks more broadly. For instance, in the case of ResNets, there is a single input x0 followed by a sequence of blocks which operate on different numbers of channels. We let the meta-model likewise be a ResNet with the same block structure. Moreover, we consider the output bt of the tth block in place of theGn(ht)’s in equation 1, and take the hidden states between the t and t+1 blocks to be the ht’s in equation 2. Since each block of a ResNet is distinct, we consider a family of maps Vn,t which depend on the base model n and the block layer t. Suppose B is the total number of blocks. Then we can more explicitly write
LCNNoutput[F̃ , θ] := 1
B B∑ t=1 d(bθn,t, bt) (5)
LCNNhidden[F̃ , θn, {Vn,t}Bt=1] := 1
B B∑ t=1 ‖Vn,t(hθn,t)− ht‖ 2 2 . (6)
Then the total loss function accounting for all N of the ResNet base models is
LCNN[F̃ , {θn}Nn=1, {{Vn,t}Bt=1}Nn=1]
:= 1
N N∑ n=1 Ex0∼Dn [ LCNNn,hidden[F̃ , θn, {Vn,t}Bt=1] + λLCNNn,output[F̃ , θn] ] (7)
where we note that in the expectation value we are only sampling over x0’s in Dn since x0’s are the only form of input data.

heading: B ADDITIONAL EXPERIMENTAL DETAILS

text:
In this section, we provide further details on our experiments as well as additional empirical results.
B.1 META-MODEL ARCHITECTURES
RNN architecture. We parameterize the meta-model for RNNs as a GRU that takes the model embedding vector as an additional input at each time step. Specifically, the meta-model GRU takes as input the vector [xt ; θ] at each time step, where xt is an input token embedding and [ · ; · ] denotes concatenation. The model embedding vector θ therefore serves as a time-independent bias for the meta-model.
CNN architecture. We parameterize the meta-model for CNNs with a modified ResNet architecture. In each residual block, we use a linear transformation W to map the model embedding vector θ to the corresponding channel dimension of convolutional layer. We then add the vector Wθ to the channels at each spatial location of the feature map. This design emulates the approach used in our parametrization of the RNN meta-model, with the model embedding θ serving as a bias term in each residual block. We reuse the weight matrix W for all residual blocks with the same channel dimension.
Algorithm 2 Meta-Model Residual Block Input: features x, model embedding vector θ z = BatchNorm(Conv(x)) z = ReLU(z+Wθ) z = BatchNorm(Conv(x)) z = ReLU(x+ z) return z
The standard ResNet architecture consists of a sequence of four layers, with each layer consisting of a sequence of residual blocks. To compute the hidden state loss Lhidden in DYNAMO, we compute distances between the output representations of each of these four layers, averaging over the number of channels and spatial locations in each set of features.
B.2 TRAINING DETAILS
Table 1 lists the hyperparameters used for training our RNN base models and meta-models, and Table 2 lists the hyperparameters used for our CNN base models and meta-models. By default, we use a model embedding dimension of 16. In the case of visualizing the training trajectories of IMDB sentiment GRUs (Figure 3, second column), we use a model embedding dimension of 32 due to the relatively larger number of base models.
B.3 SUPPLEMENTARY EMPIRICAL RESULTS
Effect of output loss weight. Figure 9 shows the effect of setting the output loss weight λ = 0 for meta-models on RNNs. These plots illustrate that model clustering can be performed on the basis of comparing hidden state dynamics alone. We note that the change in the λ hyperparameter results in qualitative changes in the resulting clustering. In particular, the model embeddings for GRUs trained on AG News classification (rightmost column of Figure 9) are much more tightly coupled relative to the GRUs trained on the IMDB dataset when λ = 0. This indicates that the dynamics implemented by the AG News GRUs are much more similar than those implemented by the IMDB sentiment GRUs.
Clustering with other loss functions. As noted in Sec. 4.1, the loss functions used to compare hidden states and outputs can be easily changed to better match the characteristics of the base models under consideration. We demonstrate the benefit of this additional flexibility by replacing the L2 distance in Lhidden (equation 6) with the L1 distance for comparing the intermediate representations of ResNets. This choice is motivated by the observation that the use of the ReLU nonlinearity results in sparse representations, which suggests the use of the L1 metric. Figure 10 shows a comparison between these two distance functions in the case of ResNet-34 models trained on CIFAR-100 with “weak” data augmentation (random shifts and horizontal flips) and with “strong” data augmentation (RandAugment). Here, we parameterized the meta-model with a ResNet-50 architecture. The use of the L1 distance results in a clearer separation between the two sets of models. This is reflected qualitatively in the distribution of the base model embeddings in model embedding space, and quantitatively in the relative scale of the variance captured by the first principal component.
B.4 FURTHER INVESTIGATION OF THE DYNAMICS OF SENTIMENT ANALYSIS
Recall that our trained RNN’s implemented sentiment analysis via line attractor dynamics, in which inputted words kicked the hidden state in the ‘positive’ or ‘negative’ direction along the line attractor according to the valence of the word (i.e. how positive or negative the word is). Figure 11 investigates how valences assigned to words change as we scan across model embedding space. Given an RNN corresponding to a value θ, we first find a fixed point h∗ with neutral readout G(h∗) ≈ 0. For this fixed point, we compute G(F (x, h∗)) for a variety of word inputs x. To produce a “score” for the model, we compute
Score(h∗) = ∑
x∈Wpositive
G(F (x, h∗))− ∑
x∈Wnegative
G(F (x, h∗))− ∑
x∈Wneutral
|G(F (x, h∗))|, (8)
where the set of positive words Wpositive, negative words Wnegative, and neutral words Wneutral are listed in Table 3. Note that here we are not analyzing context effects, for instance how the string ‘not terrible’ would be rendered into a net-positive valence.

heading: C META-MODELS AND TOPOLOGICAL CONJUGACY

text:
In this section we describe the notion of topological conjugacy, its relationship to the loss function Lhidden , and provide speculation as to the interpretation of the results from Section 5. A topological conjugacy (Katok & Hasselblatt, 1997) between two dynamical systems defined by maps F : X → X and E : Y → Y is a homeomorphism Φ : X → Y satisfying
F (x) = (Φ−1 ◦ E ◦ Φ)(x). (9) Note that the dynamical systems F and E can have distinct domains. The significance of the relationship in equation 9 is that dynamics obtained from iterated applications of the map F and E are
related to each other by the formula
Fn(x) = (Φ−1 ◦ En ◦ Φ)(x).
Thus, if F and E are topologically conjugate, then their iterates are also topologically conjugate and this means that the dynamics are related by a change of variables. Notice that topological conjugacy is an equivalence relation; as such the transitive property tells us that if F ∼ E and E ∼ D then F ∼ D. The notion of topological conjugacy is an important motivation for defining the loss functionLhidden , which we recall is given by
Lhidden[F , θ, V ] := 1
T T∑ t=1 ‖V (hθ,t)− ht‖22,
where the hidden states ht are dynamics obtained from a base model F , the hidden states hθt are dynamics obtained from the meta-model F̃ with parameter θ, and V is the map from the meta-model hidden states to the base model hidden states. One way of obtaining zero loss is to find a topological conjugacy with map V between the meta-model F̃ (at fixed θ) and the base model F , meaning a relationship of the form
F (x, h) = (V ◦ F̃ )(θ, x, V −1h).
It is convenient to define the notation Fx(h) := F (x, h) and F̃θ,x(h) := F̃ (θ, x, h). Then we have
ht = (Fxt ◦ Fxt−1 ◦ · · · ◦ Fx1)(h0) = (V ◦ F̃θ,xt ◦ F̃θ,xt−1 ◦ · · · ◦ F̃θ,x1)(V −1h0) = V hθ,t
with hθ,0 = V −1h0 , so that Lhidden[F , θ, V ] = 0 would hold. As depicted in Figure 2, to obtain zero loss it would actually suffice to find a weaker relationship of the form
F (x, V h) = (V ◦ F̃ )(θ, x, h).
The difference here is that V need not be invertible.
Using the language of topological conjugacy, we can describe a speculative but plausible interpretation of the results of Section 5. In that section we observed that models from the same cluster had very similar dynamical features and performed similarly to the model average of the cluster. This suggests that for each model Fn in the same cluster, we have
Fn(x, Vnh) ≈ (Vn ◦ F )(θ, x, h)
where θ is the centroid of the cluster to which Fn belongs. Note that here we replaced θn with θ, thus assuming both that Lhidden is small and that θn is sufficiently close to θ. Second, making the hypothesis that there exists an inverse V −1n to the map Vn, the map Vn may provide a topological conjugacy between the base model Fn and the meta-model F̃θ evaluated at θ. Assuming further that our assumptions hold for all models in the cluster, using the transitivity of topological conjugacy we would conclude that base models belonging to the same cluster are topologically conjugate to one another. This would justify the intuition suggested by Figure 8 that DYNAMO clusters models according to commonalities of topological structures of dynamics.


summaries: This paper aims to add explicit/human-defined constraints to learning-based simulation frameworks, where a learned constraint function implicitly regularizes the dynamics, and future predictions are generated via a constraint solver. The authors built the framework on top of graph neural networks (GNNs) to capture the compositionality of the underlying system and enforce the constraint using an implicit constraint function optimized via gradient descent.  The authors tested the proposed method in four physical simulation environments, including rope, bouncing balls, bouncing rigids, and BoxBath. Experiment results show that the proposed C-GNS has a competitive or better performance compared to prior learned simulators. In addition, they have also demonstrated that C-GNS can generalize to unseen, hand-designed constraints by applying more solver iterations than experienced during training to improve the accuracy on larger systems.

summaries: This paper presents a neural-network simulator that learns to solve constraints inspired by classic physics-based simulation. The proposed simulator uses a graph neural network to encode a constraint solver and shows results in a number of simulation environments.  The main contribution of the paper is its idea of using graph networks to encode constraint-based simulation.

summaries: This paper presents to simulate physics via a constraint-based approach instead of direct prediction. In particular, the authors first employ GNN taking as input the history positions and dynamics, to capture the interaction between different particles within the system, whose output is considered as the constraint satisfaction scalar. Then, the gradient with respect to the constraint function is applied as the update of the dynamics over a certain number of solver iterations. The experiments are conducted on a variety of challenging physical domains, including simulated ropes, bouncing balls, colliding irregular shapes and splashing fluids.

summaries: This manuscript proposes to learn the numerical solutions to Lagrangian physical simulation with a constraint-based inference method on graph neural networks. This method involves an iterative update during inference and thus enables test-time dynamical correction. The contributions are: (1) this manuscript builds a scalar predictor to indicate how well the constraint is agreed. (2) this manuscript proposes using the graph neural networks as the backbone to deal with a variable length of the physical domain. They also examine the effectiveness with a bunch of experiments including the state prediction experiment on four different environments and multiple ablation studies towards different hyper-parameters.

abstractText: In the rapidly advancing area of learned physical simulators, nearly all methods train a forward model that directly predicts future states from input states. However, many traditional simulation engines use a constraint-based approach instead of direct prediction. Here we present a framework for constraint-based learned simulation, where a scalar constraint function is implemented as a trainable function approximator, and future predictions are computed as the solutions to a constraint satisfaction problem. We implement our method using a graph neural network as the constraint function and gradient descent as the constraint solver. The architecture can be trained by standard backpropagation. We test the model on a variety of challenging physical domains, including simulated ropes, bouncing balls, colliding irregular shapes and splashing fluids. Our model achieves better or comparable performance to top learned simulators. A key advantage of our model is the ability to generalize to more solver iterations at test time to improve the simulation accuracy. We also show how hand-designed constraints can be added at test time to satisfy objectives which were not present in the training data, which is not possible with forward approaches. Our constraint-based framework is applicable to any setting in which forward learned simulators are used, and more generally demonstrates key ways that learned models can leverage popular methods in numerical methods.

authors:
Brandon Amos
J Zico Kolter

title: Optnet: Differentiable optimization as a layer in neural networks

authors:
Shaojie Bai
J Zico Kolter
Vladlen Koltun

title: Deep equilibrium models

venue: arXiv preprint arXiv:1909.01377,

authors:
Shaojie Bai
Vladlen Koltun
J Zico Kolter

title: Multiscale deep equilibrium models

venue: arXiv preprint arXiv:2006.08656,

title: Fast contact force computation for nonpenetrating rigid bodies

venue: In Proceedings of the 21st annual conference on Computer graphics and interactive techniques,

authors:
Sergey Bartunov
Vinod Nair
Peter Battaglia
Tim Lillicrap

title: Continuous latent search for combinatorial optimization

venue: In Learning Meets Combinatorial Algorithms at NeurIPS2020,

authors:
P. Battaglia
Razvan Pascanu
Matthew Lai
Danilo Jimenez Rezende
K. Kavukcuoglu

title: Interaction networks for learning about objects, relations and physics

authors:
Peter W Battaglia
Jessica B Hamrick
Victor Bapst
Alvaro Sanchez-Gonzalez
Vinicius Zambaldi
Mateusz Malinowski
Andrea Tacchetti
David Raposo
Adam Santoro
Ryan Faulkner

title: Relational inductive biases, deep learning, and graph networks

venue: arXiv preprint arXiv:1806.01261,

authors:
Sofien Bouaziz
Sebastian Martin
Tiantian Liu
Ladislav Kavan
Mark Pauly

title: Projective dynamics: Fusing constraint projections for fast simulation

venue: ACM transactions on graphics (TOG),

authors:
Michael M Bronstein
Joan Bruna
Yann LeCun
Arthur Szlam
Pierre Vandergheynst

title: Geometric deep learning: going beyond euclidean data

venue: IEEE Signal Processing Magazine,

authors:
Ricky TQ Chen
Yulia Rubanova
Jesse Bettencourt
David Duvenaud

title: Neural ordinary differential equations

venue: arXiv preprint arXiv:1806.07366,

authors:
Zhengdao Chen
Jianyu Zhang
Martin Arjovsky
Léon Bottou

title: Symplectic recurrent neural networks

venue: arXiv preprint arXiv:1909.13334,

title: The numerical solution of the navier-stokes equations for an incompressible fluid

venue: Bulletin of the American Mathematical Society,

authors:
Filipe De Avila Belbute-Peres
Thomas Economon
Zico Kolter

title: Combining differentiable PDE solvers and graph neural networks for fluid flow prediction

venue: Proceedings of the 37th International Conference on Machine Learning,

authors:
Priya L Donti
David Rolnick
J Zico Kolter

title: Dc3: A learning method for optimization with hard constraints

venue: arXiv preprint arXiv:2104.12225,

authors:
David Duvenaud
Zico Kolter
Matt Johnson

title: Deep implicit layers - neural odes, deep equilibrium models, and beyond, 2020

authors:
Marc Finzi
Ke Alexander Wang
Andrew Gordon Wilson

title: Simplifying hamiltonian and lagrangian neural networks via explicit constraints

venue: arXiv preprint arXiv:2010.13581,

authors:
Rony Goldenthal
David Harmon
Raanan Fattal
Michel Bercovier
Eitan Grinspun

title: Efficient simulation of inextensible cloth

venue: In ACM SIGGRAPH 2007 papers, pp. 49–es

authors:
Samuel Greydanus
Misko Dzamba
Jason Yosinski

title: Hamiltonian neural networks

authors:
George Em Karniadakis
Ioannis G Kevrekidis
Lu Lu
Paris Perdikaris
Sifan Wang
Liu Yang

title: Physics-informed machine learning

venue: Nature Reviews Physics,

authors:
Dmitrii Kochkov
Jamie A Smith
Ayya Alieva
Qing Wang
Michael P Brenner
Stephan Hoyer

title: Machine learning–accelerated computational fluid dynamics

venue: Proceedings of the National Academy of Sciences,

authors:
Yunzhu Li
Jiajun Wu
Russ Tedrake
Joshua B Tenenbaum
Antonio Torralba

title: Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids

venue: In ICLR,

authors:
João Loula
Kelsey Allen
Tom Silver
Josh Tenenbaum

title: Learning constraint-based planning models from demonstrations

venue: In 2020 IEEE/RSJ International Conference on Intellitgent Robots and Systems (IROS),

authors:
Michael Lutter
Christian Ritter
Jan Peters

title: Deep lagrangian networks: Using physics as model prior for deep learning

authors:
Miles Macklin
Matthias Müller
Nuttapong Chentanez
Tae-Yong Kim

title: Unified particle physics for real-time applications

venue: ACM Transactions on Graphics (TOG),

authors:
Damian Mrowca
Chengxu Zhuang
Elias Wang
Nick Haber
Li Fei-Fei
Joshua B Tenenbaum
Daniel LK Yamins

title: Flexible neural representation for physics prediction

venue: arXiv preprint arXiv:1806.08047,

authors:
Matthias Müller
Bruno Heidelberger
Marcus Hennix
John Ratcliff

title: Position based dynamics

venue: Journal of Visual Communication and Image Representation,

authors:
Tobias Pfaff
Meire Fortunato
Alvaro Sanchez-Gonzalez
Peter Battaglia

title: Learning mesh-based simulation with graph networks

authors:
Yulia Rubanova
Ricky TQ Chen
David Duvenaud

title: Latent odes for irregularly-sampled time series

venue: In Proceedings of the 33rd International Conference on Neural Information Processing Systems,

authors:
Alvaro Sanchez-Gonzalez
Nicolas Heess
Jost Tobias Springenberg
Josh Merel
Martin Riedmiller
Raia Hadsell
Peter Battaglia

title: Graph networks as learnable physics engines for inference and control

authors:
Alvaro Sanchez-Gonzalez
Victor Bapst
Kyle Cranmer
Peter Battaglia

title: Hamiltonian graph networks with ode integrators

venue: arXiv preprint arXiv:1909.12790,

authors:
Alvaro Sanchez-Gonzalez
Jonathan Godwin
Tobias Pfaff
Rex Ying
Jure Leskovec
Peter Battaglia

title: Learning to simulate complex physics with graph networks

authors:
Bernhard Thomaszewski
Simon Pabst
Wolfgang Strasser

title: Continuum-based strain limiting

venue: In Computer Graphics Forum,

authors:
Nils Thuerey
Konstantin Weißenow
Lukas Prantl
Xiangyu Hu

title: Deep learning methods for reynolds-averaged navier–stokes simulations of airfoil flows

venue: AIAA Journal,

authors:
Emanuel Todorov
Tom Erez
Yuval Tassa

title: Mujoco: A physics engine for model-based control

venue: In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems,

authors:
Po-Wei Wang
Priya Donti
Bryan Wilder
Zico Kolter

title: Satnet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver

authors:
Rui Wang
Karthik Kashinath
Mustafa Mustafa
Adrian Albert
Rose Yu

title: Towards physicsinformed deep learning for turbulent flow prediction

venue: In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining,

authors:
Jin-Long Wu
Heng Xiao
Eric Paterson

title: Physics-informed machine learning approach for augmenting turbulence models: A comprehensive framework

venue: Physical Review Fluids,

authors:
Shuqi Yang
Xingzhe He
Bo Zhu

title: Learning physical constraints with neural projections

authors:
Yao Zhang
Woong Je Sung
Dimitri N Mavris

title: Application of convolutional neural network to predict airfoil lift coefficient

venue: AIAA/ASCE/AHS/ASC Structures, Structural Dynamics, and Materials Conference,

title: rather than raw velocity to better match the approach in Sanchez-Gonzalez et al

title: 6We use five for BOXBATH to match the model in Sanchez-Gonzalez et al. (2020) 7For BOXBATH, we also provide the norm of the relative distances as an additional vector to match features

authors: ConstraintMLP For the ConstraintMLP we use a similar setup to as in the original paper by Yang

title: We concatenate the features for each node into a single vector and run an MLP to produce a scalar constraint output

venue: The node features include the absolute positions, velocities and distances to the walls for each node for the optimized state and for the past three time points. Note that this model does not have access to the position differences between the nodes (this information

title: The MLPs for node and edge processing

heading: 1 INTRODUCTION

text:
Consider a bowling ball colliding with a bowling pin. You might explain this event as involving a pair of forces being generated, one which causes the pin to move, and the other which causes the ball to careen away with a different direction and speed. This kind of intuitive cause-and-effect approach is analogous to physical simulators that apply an explicit forward model to calculate a future state directly from the current one, such as when numerically integrating discretized equations of motion.
An alternative, but equally valid, way to explain the collision is in terms of constraint satisfaction: the ball and pin cannot occupy the same location at the same time, and their combined energies and momenta must be conserved, so the post-collision trajectories are the only way the future can unfold without violating these constraints. This constraint-based approach is analogous to physical simulators that use an implicit function to model a system of constraints over the current and future states, and which generate a prediction by searching for a future state that respects all constraints.
Both families of simulators—those based on explicit, forward functions versus those which define the dynamics implicitly, via constraints—are widely used in physics, engineering, and graphics. In principle they can model the same types of dynamics, however they differ in how their respective predictions are computed and in practice strike different trade-offs that determine why one or the other is preferred in different domains. For example, explicit methods are popular for large systems with (mostly) independent local effects whose space and time derivatives are relatively smooth, and their accuracy can often be increased by discretizing space and time more finely. Implicit approaches are often preferred for systems with strong interactions, such rigid and stiff dynamics, and more accurate solutions can often be found by using more sophisticated constraint solvers or by increasing the computational budget (e.g., solver iterations) allocated to searching for solutions. In machine learning (ML), there have been rapid advances recently in methods for learning to simulate complex dynamic processes, however almost all (e.g., Sanchez-Gonzalez et al. (2020); Pfaff et al. (2021)) have focused on explicit forward model approaches, with few exceptions (Yang et al., 2020).
Here we present a framework for learning to simulate complex dynamics via constraint satisfaction. Our “Constraint-based Graph Network Simulator” (C-GNS) defines a single scalar-valued constraint function that represents whether a future state satisfies the physical constraints, conditioned on the current and previous states. The constraint function is implemented as a Graph Neural Network (GNN) (Bronstein et al., 2017; Battaglia et al., 2018), which can model systems with rich compositional structure—multiple bodies, complex meshes, etc. To predict the next state via the constraint function’s implicit representation of the dynamics, a gradient-based solver finds a proposed state which satisfies the constraints. We train it through the solver by backpropagation. We also introduce a hybrid approach that proposes and refines the future state using an explicit iterative predictor, rather than solving for learned constraints.
We tested the C-GNS on a variety of challenging physical simulation domains generated by several different simulation engines: simulated rope, bouncing balls, and bouncing irregular rigid shapes (MuJoCo, Todorov et al. (2012)) and splashing fluids (Flex, Macklin et al. (2014)). We found that the C-GNS’s simulated rollouts were more accurate than a state-of-the-art Graph Net Simulator (GNS) (Sanchez-Gonzalez et al., 2020) with comparable number of parameters. At test time, the C-GNS could use additional solver iterations to improve its predictive accuracy, striking desired speed-accuracy trade-offs. It could also satisfy new, hand-designed constraints jointly alongside its learned constraints. Neither of these capabilities are possible in explicit forward-style approaches.

heading: 2 BACKGROUND AND RELATED WORK

text:
Constraint solvers are central to many physics simulators. Most rigid-body and game engines use constraints to model joints, collision and contact (Baraff, 1994). They are used for limiting strain in realistic cloth simulation (Thomaszewski et al., 2009), and are a core component in Eulerian incompressible fluid solvers to solve for pressure (Chorin, 1967). Recently, position-based (Müller et al., 2007) and projective dynamics methods (Bouaziz et al., 2014) have become very popular for interactive simulation. These methods express dynamics purely as constraints, and can simulate a wide range of physical systems from rigids over soft-bodies to fluids (Macklin et al., 2014).
Machine learning methods for accelerating scientific simulation of complex systems, such as turbulence (Kochkov et al., 2021; Wang et al., 2020) and aerodynamics (Thuerey et al., 2020; Zhang et al., 2018), have grown rapidly in recent years. GNN-based learned simulators, in particular, are a very flexible approach which can model a wide range of systems, from articulated dynamics (Sanchez-Gonzalez et al., 2018) to particle-based physics (Mrowca et al., 2018; Li et al., 2019; Sanchez-Gonzalez et al., 2020) and mesh-based continuum systems (Pfaff et al., 2021; De Avila Belbute-Peres et al., 2020), and generalize well to unseen scenarios. Combining learning algorithms with principles from physics and numerical methods, such as auxiliary loss terms and rich inductive biases, can improve sample complexity, computational efficiency, and generalization (Wu et al., 2018; Karniadakis et al., 2021; Chen et al., 2018; Rubanova et al., 2019). Imposing Hamiltonian (Greydanus et al., 2019; Sanchez-Gonzalez et al., 2019; Chen et al., 2019) and Lagrangian (Lutter et al., 2019; Cranmer et al., 2020; Finzi et al., 2020) mechanics in learned simulators offers unique speed/accuracy tradeoffs and can preserve symmetries more effectively.
Recent methods have been proposed for learning constraint functions and solving them in a model’s forward pass (Duvenaud et al. (2020)’s “Deep Implicit Layers” tutorial is an excellent hands-on survey). Such models can play games (Amos & Kolter, 2017; Wang et al., 2019), optimize power flow (Donti et al., 2021), support robotic planning (Loula et al., 2020), and perform combinatorial optimization (Bartunov et al., 2020). Solvers such as gradient descent and Newton’s method are differentiable, and support training by backpropagation, but this can be computationally expensive, so approaches such as Deep Equilibrium Models (DEM) (Bai et al., 2019; 2020) use implicit differentiation to compute gradients only at the solution point.
Despite the popularity of constraint-based traditional simulators, only a single simulator which uses learned constraints has been reported (Yang et al., 2020). Their “Neural Projections” method, based on Goldenthal et al. (2007), iteratively proposes a future state with an Euler step, then projects the proposal onto a learned constraint manifold, implemented as a multilayer perceptron (MLP). Crucially, their constraint function only measures how much an individual state violates the learned constraints, and thus is not an implicit representation of the dynamics. It is suitable for quasi-static regimes, but not scenarios such as the elastic collisions in the bowling ball example described above.
Under review as a conference paper at ICLR 2021
<latexit sha1_base64="pfoxs8rsC/moP79HBUVBdU/MhtI=">AAAB6nicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkUI9FLx4r2g9oQ9lsJ+3SzSbsboQS+hO8eFDEq7/Im//GbZuDtj4YeLw3w8y8IBFcG9f9dgobm1vbO8Xd0t7+weFR+fikreNUMWyxWMSqG1CNgktsGW4EdhOFNAoEdoLJ7dzvPKHSPJaPZpqgH9GR5CFn1FjpoTvwBuWKW3UXIOvEy0kFcjQH5a/+MGZphNIwQbXueW5i/Iwqw5nAWamfakwom9AR9iyVNELtZ4tTZ+TCKkMSxsqWNGSh/p7IaKT1NApsZ0TNWK96c/E/r5ea8NrPuExSg5ItF4WpICYm87/JkCtkRkwtoUxxeythY6ooMzadkg3BW315nbSvql6tWruvVRo3eRxFOINzuAQP6tCAO2hCCxiM4Ble4c0Rzovz7nwsWwtOPnMKf+B8/gDdJ42H</latexit>
<latexit sha1_base64="t5Z8j6uw1dDrul9BiPoeucvIxm0=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2A9oQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgobm1vbO8Xd0t7+weFR+fikreNUMWyxWMSqG1CNgktsGW4EdhOFNAoEdoLJ3dzvPKHSPJYPZpqgH9GR5CFn1FipqQflilt1FyDrxMtJBXI0BuWv/jBmaYTSMEG17nluYvyMKsOZwFmpn2pMKJvQEfYslTRC7WeLQ2fkwipDEsbKljRkof6eyGik9TQKbGdEzVivenPxP6+XmvDGz7hMUoOSLReFqSAmJvOvyZArZEZMLaFMcXsrYWOqKDM2m5INwVt9eZ20r6perVpr1ir12zyOIpzBOVyCB9dQh3toQAsYIDzDK7w5j86L8+58LFsLTj5zCn/gfP4A4OuM/g==</latexit>
<latexit sha1_base64="HJ/awZ5XFr1cDfruVAeJQcOIixk=">AAAB8HicbVDLSgNBEOz1GeMr6tHLYBA8hV0J6DHoxWME85BkCbOTSTJkZnaZ6RXCkq/w4kERr36ON//GSbIHTSxoKKq66e6KEiks+v63t7a+sbm1Xdgp7u7tHxyWjo6bNk4N4w0Wy9i0I2q5FJo3UKDk7cRwqiLJW9H4dua3nrixItYPOEl4qOhQi4FgFJ302B1RzNrTHvZKZb/iz0FWSZCTMuSo90pf3X7MUsU1Mkmt7QR+gmFGDQom+bTYTS1PKBvTIe84qqniNszmB0/JuVP6ZBAbVxrJXP09kVFl7URFrlNRHNllbyb+53VSHFyHmdBJilyzxaJBKgnGZPY96QvDGcqJI5QZ4W4lbEQNZegyKroQguWXV0nzshJUK9X7arl2k8dRgFM4gwsI4ApqcAd1aAADBc/wCm+e8V68d+9j0brm5TMn8Afe5w8R35CX</latexit>
<latexit sha1_base64="ajLqfCZov3o6FTnfb54h4xEyGI0=">AAAB9HicbVBNS8NAEJ3Ur1q/qh69LBZBEEoiBT0WvXisYGuhDWWz3bZLN5u4OymUkN/hxYMiXv0x3vw3btsctPXBwOO9GWbmBbEUBl332ymsrW9sbhW3Szu7e/sH5cOjlokSzXiTRTLS7YAaLoXiTRQoeTvWnIaB5I/B+HbmP064NiJSDziNuR/SoRIDwShaye+OKKbtrJfihZf1yhW36s5BVomXkwrkaPTKX91+xJKQK2SSGtPx3Bj9lGoUTPKs1E0Mjykb0yHvWKpoyI2fzo/OyJlV+mQQaVsKyVz9PZHS0JhpGNjOkOLILHsz8T+vk+Dg2k+FihPkii0WDRJJMCKzBEhfaM5QTi2hTAt7K2EjqilDm1PJhuAtv7xKWpdVr1at3dcq9Zs8jiKcwCmcgwdXUIc7aEATGDzBM7zCmzNxXpx352PRWnDymWP4A+fzB7oYkhM=</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="pmdXaQtx/RkbjEEj0JjG94undSA=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2FZpQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgpr6xubW8Xt0s7u3v5B+fCoreNUMWyxWMTqPqAaBZfYMtwIvE8U0igQ2AnGNzO/84RK81jemUmCfkSHkoecUWOl5kO/XHGr7hxklXg5qUCORr/81RvELI1QGiao1l3PTYyfUWU4Ezgt9VKNCWVjOsSupZJGqP1sfuiUnFllQMJY2ZKGzNXfExmNtJ5Ege2MqBnpZW8m/ud1UxNe+RmXSWpQssWiMBXExGT2NRlwhcyIiSWUKW5vJWxEFWXGZlOyIXjLL6+S9kXVq1VrzVqlfp3HUYQTOIVz8OAS6nALDWgBA4RneIU359F5cd6dj0VrwclnjuEPnM8fuYOM5A==</latexit>
<latexit sha1_base64="Jxyp7Y07VUyiyNEiRyqcT+XYQDk=">AAAB9HicbVDLSgNBEJz1GeMr6tHLYBA8hV0J6DGoB71FMA9IljA76U2GzD6c6Q2GZb/DiwdFvPox3vwbJ8keNLGgoajqprvLi6XQaNvf1srq2vrGZmGruL2zu7dfOjhs6ihRHBo8kpFqe0yDFCE0UKCEdqyABZ6Elje6nvqtMSgtovABJzG4ARuEwhecoZFcv5d2EZ4wvbnLsl6pbFfsGegycXJSJjnqvdJXtx/xJIAQuWRadxw7RjdlCgWXkBW7iYaY8REbQMfQkAWg3XR2dEZPjdKnfqRMhUhn6u+JlAVaTwLPdAYMh3rRm4r/eZ0E/Us3FWGcIIR8vshPJMWIThOgfaGAo5wYwrgS5lbKh0wxjianognBWXx5mTTPK061Ur2vlmtXeRwFckxOyBlxyAWpkVtSJw3CySN5Jq/kzRpbL9a79TFvXbHymSPyB9bnDziakmY=</latexit>
<latexit sha1_base64="rQRE9g+W9jgLwJeS3fY3jHXSBJo=">AAAB83icbVDLSgNBEJyNrxhfUY9eBoPgKexKQI9BPXiMYB6QXcLspDcZMvtgplcMy/6GFw+KePVnvPk3TpI9aGJBQ1HVTXeXn0ih0ba/rdLa+sbmVnm7srO7t39QPTzq6DhVHNo8lrHq+UyDFBG0UaCEXqKAhb6Erj+5mfndR1BaxNEDThPwQjaKRCA4QyO5wSBzEZ4wu83zQbVm1+056CpxClIjBVqD6pc7jHkaQoRcMq37jp2glzGFgkvIK26qIWF8wkbQNzRiIWgvm9+c0zOjDGkQK1MR0rn6eyJjodbT0DedIcOxXvZm4n9eP8XgystElKQIEV8sClJJMaazAOhQKOAop4YwroS5lfIxU4yjialiQnCWX14lnYu606g37hu15nURR5mckFNyThxySZrkjrRIm3CSkGfySt6s1Hqx3q2PRWvJKmaOyR9Ynz+gO5IT</latexit>
<latexit sha1_base64="kZkYF8749t+4eNJpcL8A/9X3kSM=">AAAB83icbVBNS8NAEJ3Ur1q/qh69BIvgqSRS0GOxF48V7Ac0oWy2m3bpZhN2J2IJ+RtePCji1T/jzX/jts1BWx8MPN6bYWZekAiu0XG+rdLG5tb2Tnm3srd/cHhUPT7p6jhVlHVoLGLVD4hmgkvWQY6C9RPFSBQI1gumrbnfe2RK81g+4CxhfkTGkoecEjSSFw4zD9kTZq08H1ZrTt1ZwF4nbkFqUKA9rH55o5imEZNIBdF64DoJ+hlRyKlgecVLNUsInZIxGxgqScS0ny1uzu0Lo4zsMFamJNoL9fdERiKtZ1FgOiOCE73qzcX/vEGK4Y2fcZmkyCRdLgpTYWNszwOwR1wximJmCKGKm1ttOiGKUDQxVUwI7urL66R7VXcb9cZ9o9a8LeIowxmcwyW4cA1NuIM2dIBCAs/wCm9War1Y79bHsrVkFTOn8AfW5w+etZIS</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="FgDtDkQujOEF3ZM1C6wBi95/Q5Y=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXAnoMevEYwTwkWcPspJMMmZ1dZmaFsOQjvHhQxKvf482/cZLsQRMLGoqqbrq7glhwbVz328mtrW9sbuW3Czu7e/sHxcOjpo4SxbDBIhGpdkA1Ci6xYbgR2I4V0jAQ2ArGNzO/9YRK80jem0mMfkiHkg84o8ZKrYfHtMzPp71iya24c5BV4mWkBBnqveJXtx+xJERpmKBadzw3Nn5KleFM4LTQTTTGlI3pEDuWShqi9tP5uVNyZpU+GUTKljRkrv6eSGmo9SQMbGdIzUgvezPxP6+TmMGVn3IZJwYlWywaJIKYiMx+J32ukBkxsYQyxe2thI2ooszYhAo2BG/55VXSvKh41Ur1rlqqXWdx5OEETqEMHlxCDW6hDg1gMIZneIU3J3ZenHfnY9Gac7KZY/gD5/MHvsCPMA==</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="FgDtDkQujOEF3ZM1C6wBi95/Q5Y=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXAnoMevEYwTwkWcPspJMMmZ1dZmaFsOQjvHhQxKvf482/cZLsQRMLGoqqbrq7glhwbVz328mtrW9sbuW3Czu7e/sHxcOjpo4SxbDBIhGpdkA1Ci6xYbgR2I4V0jAQ2ArGNzO/9YRK80jem0mMfkiHkg84o8ZKrYfHtMzPp71iya24c5BV4mWkBBnqveJXtx+xJERpmKBadzw3Nn5KleFM4LTQTTTGlI3pEDuWShqi9tP5uVNyZpU+GUTKljRkrv6eSGmo9SQMbGdIzUgvezPxP6+TmMGVn3IZJwYlWywaJIKYiMx+J32ukBkxsYQyxe2thI2ooszYhAo2BG/55VXSvKh41Ur1rlqqXWdx5OEETqEMHlxCDW6hDg1gMIZneIU3J3ZenHfnY9Gac7KZY/gD5/MHvsCPMA==</latexit>
<latexit sha1_base64="nk/O4KvoHG7nWgCKsKo+m2fcKfM=">AAAB/nicbVBNS8NAEN3Ur1q/ouLJy2IRPJVECnos9uKxgm2VJoTNdtMu3WzC7kQsoeBf8eJBEa/+Dm/+G7dtDtr6YODx3gwz88JUcA2O822VVlbX1jfKm5Wt7Z3dPXv/oKOTTFHWpolI1F1INBNcsjZwEOwuVYzEoWDdcNSc+t0HpjRP5C2MU+bHZCB5xCkBIwX2kSdJKEhwj6Mg94A9Qt6cTAK76tScGfAycQtSRQVagf3l9ROaxUwCFUTrnuuk4OdEAaeCTSpepllK6IgMWM9QSWKm/Xx2/gSfGqWPo0SZkoBn6u+JnMRaj+PQdMYEhnrRm4r/eb0Moks/5zLNgEk6XxRlAkOCp1ngPleMghgbQqji5lZMh0QRCiaxignBXXx5mXTOa269Vr+pVxtXRRxldIxO0Bly0QVqoGvUQm1EUY6e0St6s56sF+vd+pi3lqxi5hD9gfX5A3C7lc8=</latexit>
<latexit sha1_base64="vQ3/3W0LApHcMPwu9rwFnhCUcHI=">AAAB73icbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ie0oWw2k3bpZhN3N0Ip/RNePCji1b/jzX/jts1BWx8MPN6bYWZekAqujet+O4W19Y3NreJ2aWd3b/+gfHjU0kmmGDZZIhLVCahGwSU2DTcCO6lCGgcC28HoZua3n1Bpnsh7M07Rj+lA8ogzaqzU6YUoDCUP/XLFrbpzkFXi5aQCORr98lcvTFgWozRMUK27npsaf0KV4UzgtNTLNKaUjegAu5ZKGqP2J/N7p+TMKiGJEmVLGjJXf09MaKz1OA5sZ0zNUC97M/E/r5uZ6MqfcJlmBiVbLIoyQUxCZs+TkCtkRowtoUxxeythQ6ooMzaikg3BW355lbQuql6tWrurVerXeRxFOIFTOAcPLqEOt9CAJjAQ8Ayv8OY8Oi/Ou/OxaC04+cwx/IHz+QOaD4+w</latexit>
<latexit sha1_base64="fAxixLzqPP0ibRTKMP9FMpa2phA=">AAACB3icbZDLSsNAFIZPvNZ6i7oUZLAIlUJJpKAboejGZQV7o41lMpm0QycXZiZCCd258VXcuFDEra/gzrdxmnahrT8MfPznHM6c3405k8qyvo2l5ZXVtfXcRn5za3tn19zbb8goEYTWScQj0XKxpJyFtK6Y4rQVC4oDl9OmO7ye1JsPVEgWhXdqFFMnwP2Q+Yxgpa2eedS+T4usZJ+O0SXKWFMJdT3KFUbtnlmwylYmtAj2DAowU61nfnW9iCQBDRXhWMqObcXKSbFQjHA6zncTSWNMhrhPOxpDHFDppNkdY3SiHQ/5kdAvVChzf0+kOJByFLi6M8BqIOdrE/O/WidR/oWTsjBOFA3JdJGfcKQiNAkFeUxQovhIAyaC6b8iMsACE6Wjy+sQ7PmTF6FxVrYr5cptpVC9msWRg0M4hiLYcA5VuIEa1IHAIzzDK7wZT8aL8W58TFuXjNnMAfyR8fkD8P+W0w==</latexit>
<latexit sha1_base64="fAxixLzqPP0ibRTKMP9FMpa2phA=">AAACB3icbZDLSsNAFIZPvNZ6i7oUZLAIlUJJpKAboejGZQV7o41lMpm0QycXZiZCCd258VXcuFDEra/gzrdxmnahrT8MfPznHM6c3405k8qyvo2l5ZXVtfXcRn5za3tn19zbb8goEYTWScQj0XKxpJyFtK6Y4rQVC4oDl9OmO7ye1JsPVEgWhXdqFFMnwP2Q+Yxgpa2eedS+T4usZJ+O0SXKWFMJdT3KFUbtnlmwylYmtAj2DAowU61nfnW9iCQBDRXhWMqObcXKSbFQjHA6zncTSWNMhrhPOxpDHFDppNkdY3SiHQ/5kdAvVChzf0+kOJByFLi6M8BqIOdrE/O/WidR/oWTsjBOFA3JdJGfcKQiNAkFeUxQovhIAyaC6b8iMsACE6Wjy+sQ7PmTF6FxVrYr5cptpVC9msWRg0M4hiLYcA5VuIEa1IHAIzzDK7wZT8aL8W58TFuXjNnMAfyR8fkD8P+W0w==</latexit>
<latexit sha1_base64="9UY9Tzo0xSguDEfPuYbWHGsRnh0=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXAnoMevEYwTwkWcPsZJIMmZ1dZnqFsOQjvHhQxKvf482/cZLsQRMLGoqqbrq7glgKg6777eTW1jc2t/LbhZ3dvf2D4uFR00SJZrzBIhnpdkANl0LxBgqUvB1rTsNA8lYwvpn5rSeujYjUPU5i7od0qMRAMIpWaj08pmX3fNorltyKOwdZJV5GSpCh3it+dfsRS0KukElqTMdzY/RTqlEwyaeFbmJ4TNmYDnnHUkVDbvx0fu6UnFmlTwaRtqWQzNXfEykNjZmEge0MKY7MsjcT//M6CQ6u/FSoOEGu2GLRIJEEIzL7nfSF5gzlxBLKtLC3EjaimjK0CRVsCN7yy6ukeVHxqpXqXbVUu87iyMMJnEIZPLiEGtxCHRrAYAzP8ApvTuy8OO/Ox6I152Qzx/AHzucPZ+qO9w==</latexit>
<latexit sha1_base64="4l4tyu1wQNISUpBmeK1HI2A9PJA=">AAAB7nicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkUI9FLx4r2A9pQ9lsN+3SzSbsToQS+iO8eFDEq7/Hm//GbZuDtj4YeLw3w8y8IJHCoOt+O4WNza3tneJuaW//4PCofHzSNnGqGW+xWMa6G1DDpVC8hQIl7yaa0yiQvBNMbud+54lrI2L1gNOE+xEdKREKRtFKnf6YYvY4G5QrbtVdgKwTLycVyNEclL/6w5ilEVfIJDWm57kJ+hnVKJjks1I/NTyhbEJHvGepohE3frY4d0YurDIkYaxtKSQL9fdERiNjplFgOyOKY7PqzcX/vF6K4bWfCZWkyBVbLgpTSTAm89/JUGjOUE4toUwLeythY6opQ5tQyYbgrb68TtpXVa9Wrd3XKo2bPI4inME5XIIHdWjAHTShBQwm8Ayv8OYkzovz7nwsWwtOPnMKf+B8/gCDZ4+x</latexit>
<latexit sha1_base64="ErhthiE4wcyOR5LhKX4LnvmTkNM=">AAAB+HicbVBNS8NAEN34WetHox69LBahXkoiBb0IRS+epIL9oo1ls922SzebsDsRasgv8eJBEa/+FG/+G7dtDtr6YODx3gwz8/xIcA2O822trK6tb2zmtvLbO7t7BXv/oKHDWFFWp6EIVcsnmgkuWR04CNaKFCOBL1jTH19P/eYjU5qH8h4mEfMCMpR8wCkBI/XsQndEIGmnl+2HpHR7mvbsolN2ZsDLxM1IEWWo9eyvbj+kccAkUEG07rhOBF5CFHAqWJrvxppFhI7JkHUMlSRg2ktmh6f4xCh9PAiVKQl4pv6eSEig9STwTWdAYKQXvan4n9eJYXDhJVxGMTBJ54sGscAQ4mkKuM8VoyAmhhCquLkV0xFRhILJKm9CcBdfXiaNs7JbKVfuKsXqVRZHDh2hY1RCLjpHVXSDaqiOKIrRM3pFb9aT9WK9Wx/z1hUrmzlEf2B9/gAd6pK9</latexit>
<latexit sha1_base64="ErhthiE4wcyOR5LhKX4LnvmTkNM=">AAAB+HicbVBNS8NAEN34WetHox69LBahXkoiBb0IRS+epIL9oo1ls922SzebsDsRasgv8eJBEa/+FG/+G7dtDtr6YODx3gwz8/xIcA2O822trK6tb2zmtvLbO7t7BXv/oKHDWFFWp6EIVcsnmgkuWR04CNaKFCOBL1jTH19P/eYjU5qH8h4mEfMCMpR8wCkBI/XsQndEIGmnl+2HpHR7mvbsolN2ZsDLxM1IEWWo9eyvbj+kccAkUEG07rhOBF5CFHAqWJrvxppFhI7JkHUMlSRg2ktmh6f4xCh9PAiVKQl4pv6eSEig9STwTWdAYKQXvan4n9eJYXDhJVxGMTBJ54sGscAQ4mkKuM8VoyAmhhCquLkV0xFRhILJKm9CcBdfXiaNs7JbKVfuKsXqVRZHDh2hY1RCLjpHVXSDaqiOKIrRM3pFb9aT9WK9Wx/z1hUrmzlEf2B9/gAd6pK9</latexit>
<latexit sha1_base64="9UY9Tzo0xSguDEfPuYbWHGsRnh0=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXAnoMevEYwTwkWcPsZJIMmZ1dZnqFsOQjvHhQxKvf482/cZLsQRMLGoqqbrq7glgKg6777eTW1jc2t/LbhZ3dvf2D4uFR00SJZrzBIhnpdkANl0LxBgqUvB1rTsNA8lYwvpn5rSeujYjUPU5i7od0qMRAMIpWaj08pmX3fNorltyKOwdZJV5GSpCh3it+dfsRS0KukElqTMdzY/RTqlEwyaeFbmJ4TNmYDnnHUkVDbvx0fu6UnFmlTwaRtqWQzNXfEykNjZmEge0MKY7MsjcT//M6CQ6u/FSoOEGu2GLRIJEEIzL7nfSF5gzlxBLKtLC3EjaimjK0CRVsCN7yy6ukeVHxqpXqXbVUu87iyMMJnEIZPLiEGtxCHRrAYAzP8ApvTuy8OO/Ox6I152Qzx/AHzucPZ+qO9w==</latexit>
<latexit sha1_base64="pmdXaQtx/RkbjEEj0JjG94undSA=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2FZpQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgpr6xubW8Xt0s7u3v5B+fCoreNUMWyxWMTqPqAaBZfYMtwIvE8U0igQ2AnGNzO/84RK81jemUmCfkSHkoecUWOl5kO/XHGr7hxklXg5qUCORr/81RvELI1QGiao1l3PTYyfUWU4Ezgt9VKNCWVjOsSupZJGqP1sfuiUnFllQMJY2ZKGzNXfExmNtJ5Ege2MqBnpZW8m/ud1UxNe+RmXSWpQssWiMBXExGT2NRlwhcyIiSWUKW5vJWxEFWXGZlOyIXjLL6+S9kXVq1VrzVqlfp3HUYQTOIVz8OAS6nALDWgBA4RneIU359F5cd6dj0VrwclnjuEPnM8fuYOM5A==</latexit>
<latexit sha1_base64="YET3kkOS2zY8mMv2K4bjQDXERhw=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2A9oQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgobm1vbO8Xd0t7+weFR+fikreNUMWyxWMSqG1CNgktsGW4EdhOFNAoEdoLJ3dzvPKHSPJYPZpqgH9GR5CFn1FipyQblilt1FyDrxMtJBXI0BuWv/jBmaYTSMEG17nluYvyMKsOZwFmpn2pMKJvQEfYslTRC7WeLQ2fkwipDEsbKljRkof6eyGik9TQKbGdEzVivenPxP6+XmvDGz7hMUoOSLReFqSAmJvOvyZArZEZMLaFMcXsrYWOqKDM2m5INwVt9eZ20r6perVpr1ir12zyOIpzBOVyCB9dQh3toQAsYIDzDK7w5j86L8+58LFsLTj5zCn/gfP4AyKuM7g==</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="t5Z8j6uw1dDrul9BiPoeucvIxm0=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2A9oQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgobm1vbO8Xd0t7+weFR+fikreNUMWyxWMSqG1CNgktsGW4EdhOFNAoEdoLJ3dzvPKHSPJYPZpqgH9GR5CFn1FipqQflilt1FyDrxMtJBXI0BuWv/jBmaYTSMEG17nluYvyMKsOZwFmpn2pMKJvQEfYslTRC7WeLQ2fkwipDEsbKljRkof6eyGik9TQKbGdEzVivenPxP6+XmvDGz7hMUoOSLReFqSAmJvOvyZArZEZMLaFMcXsrYWOqKDM2m5INwVt9eZ20r6perVpr1ir12zyOIpzBOVyCB9dQh3toQAsYIDzDK7w5j86L8+58LFsLTj5zCn/gfP4A4OuM/g==</latexit>
<latexit sha1_base64="NkuDH4GMVAhCg8VuQG4AALMi2Sw=">AAACNHicbVBNSxxBEO3RxJhNNKsec2myBFaIy0xYMBdB9BLIxUBWd9jZDDU9NdrY0zN214jLMD/Kiz8kFwnkEJFc8xvSu+4hfhQ0vH6vXnXXS0olLfn+T29h8dnzpRfLL1uvXq+svmmvrR/aojICB6JQhRkmYFFJjQOSpHBYGoQ8UXiUnO5P9aNzNFYW+htNShzncKxlJgWQo+L2lyhFRcBDvsO3IuWMKfBIQ6IgDnkW1xHhBdX7TdMduovCM07NBx5u8si4uRTX4U74ve7KzaaJ2x2/58+KPwbBHHTYvA7i9o8oLUSVoyahwNpR4Jc0rsGQFAqbVlRZLEGcwjGOHNSQox3Xs6Ub/t4xKc8K444mPmP/d9SQWzvJE9eZA53Yh9qUfEobVZR9GtdSlxWhFncPZZXiVPBpgjyVBgWpiQMgjHR/5eIEDAhyObdcCMHDlR+Dw4+9oN/rf+13dvfmcSyzt+wd67KAbbNd9pkdsAET7JJds9/sxrvyfnm33p+71gVv7tlg98r7+w94eaot</latexit>

heading: 3 MODEL FRAMEWORK

text:
Simulation basics A physical trajectory, measured at discrete time intervals, is a sequence of states, (X1, . . . , XT ), where Xt represents properties such as the positions, velocities, masses, etc, of elements of the system. A physical simulator, s, is a function that maps current and/or previous state(s), which we term the context, X≤t, to a predicted future state, X̂t+1 = s(X≤t) (see Figure 1a)1. A simulated physical trajectory, termed a rollout, (Xt, X̂t+1, X̂t+2, . . . ), can be generated by repeatedly applying s to its own predicted state, X̂t+1 = s(X̂≤t).
Simulators are often comprised of a PREDICTOR mechanism which maps the context X≤t to an update value Ŷ , that represents information about the system’s temporal evolution at the current time. Then Ŷ is used by an UPDATER mechanism to update the current state to the next state: X̂t+1 = UPDATER(X≤t, Ŷ ), e.g., updating current positions and velocities represented by Xt with new velocities and accelerations represented by Ŷ , to predict the next state.
Explicit simulators Across science, engineering, and graphics, a popular class of simulators are defined explicitly: the state update Ŷ is predicted directly fromX≤t using an explicit forward function, Ŷ = fD(X≤t), as illustrated in Figure 1b. Among the rapidly growing family of learned simulators, the forward function fD is typically implemented using a neural network (Sanchez-Gonzalez et al., 2020; Pfaff et al., 2021).
Constraint-based implicit simulators Here we explore learned simulators based on implicit formulations of the dynamics. Rather than predicting the desired state directly, as in explicit formulations, our implicit simulator uses a differentiable constraint function, c = fC(X≤t, Ŷ ), where c is a scalar that quantifies how well a proposed state update Ŷ agrees withX≤t. A future prediction is generated by applying a solver, such as an optimization or zero-finding algorithm, to find a Ŷ that satisfies the constraint function, and applying the UPDATER to update Xt to X̂t+1. The fC can represent all the physical constraints in the system, including the time dynamics.
1Despite that physics is Markovian, we use X≤t as input because our framework can also apply to dynamic processes which are non-Markovian. Providing previous states can also often be helpful when there are hidden properties of the system which are only identifiable over a sequence of observed states, and when a state does not represent velocity or momentum information.
As illustrated in Figure 1d, we formulate our constraint-solving procedure via an iterative method that starts with an initial proposal, Y (0). On the i-th iteration, the solver uses the gradient of fC w.r.t. Y at the current proposal to compute a change to the proposal, δY = −λ ∇Y fC(X≤t, Y )|Y=Y (i) . This δY is then used to revise the proposal to, Y (i+1) = Y (i) + δY . This process repeats for N steps, and the final proposal value is treated as the PREDICTOR’s output, Ŷ = Y (N).
Our constraint-based model’s fC is defined as a trainable function approximator which is real-valued and lower bounded at zero, and uses gradient descent to find Ŷ that minimizes it, where λ is a fixed step size. This induces the semantics that the desired Ŷ = arg minY fC(X≤t, Y ).
We also explore a second constraint-solving procedure, inspired by Yang et al. (2020)’s Neural Projections’ use of “fast projection” Goldenthal et al. (2007). Specifically, λ = − fC(X≤t,Y (i))
‖∇Y fC(X≤t,Y )| Y =Y (i)
‖2 . Unlike gradient descent, fast projection is a zero-finding algorithm, so in
this case fC is not lower bounded. This induces the semantics that fC(X≤t, Ŷ ) = 0.
This general formulation of constraint-based learned simulation can be trained by backpropagating loss gradients through the solver loop2. The computational budget of the forward pass can be varied via the number of solver iterations N .
Explicit iterative simulators As a hybrid between forward and constraint-based simulators, we introduced a model which iteratively refines a proposed state update, like in the constraint-based approach described above, but using an explicit function to directly output a δY at each iteration, rather than solving a constraint function (see Figure 1c). See Section 4.3 for details.

heading: 4 EXPERIMENTS

heading: 4.1 EXPERIMENTAL TASK DOMAINS

text: We test our framework on a variety of physical environments, shown in Figure 2: ROPE, BOUNCING BALLS and BOUNCING RIGIDS, whose ground truth training and test data were generated by the MuJoCo physics simulator, as well as BOXBATH from Li et al. (2019). These environments demonstrate a diverse set of physical constraints: ‘hard’ constraints (preserving the shape of the rigid object and resolving collisions), and ‘soft’ constraints on fluid movement, handling gravity and preserving the momentum of the rope and bouncing balls. See the Supplementary Materials for details.

heading: 4.2 MODEL IMPLEMENTATIONS

text:
Representing the physical system Our experimental domains are physical systems comprised of sets of interacting point-like elements, e.g., objects, particles, mesh vertices, etc. We represent the state as Xt = (p j t )
j=1...|Xt|, where |Xt| is the number of elements, and pjt is the j-th element’s position at time t. There are also other static properties of the physical elements, e.g., masses, material types, etc., which we represent with Z to keep it distinct from the dynamic state information represented by Xt. The input context is X≤t = (Z,Xt−3, Xt−2, Xt−1, Xt).
2Implicit differentiation at the solution point should be applicable as well, and potentially offer computational benefits as mentioned in the Section 2, though we do not explore that here.
In our implementation, Ŷ , represents the predicted changes in position (i.e., the “average velocity” across the time step)3, ŷj = ∆p̂jt+1 = p̂ j t+1 − pjt . The UPDATER then computes X̂t+1 using p̂jt+1 = p j t + ∆p̂ j t+1, where p j t is provided in the input X≤t.
Constructing the input graph Our implementations of the fD, fDI, and fC use GNNs as the function approximators, so we need to pack the context, X≤t, and (for the fDI and fC) the proposed state update information, Y (i), into an input graph, Gt = (Vt, Et). The edges Et represent possible interactions among the elements, such as fully connected edges to represent collisions and rigid attachments in BOUNCING BALLS and BOUNCING RIGIDS, spring constraints in ROPE, and interactions among particles within a fixed connectivity radius in BOXBATH.
We enforced translation-invariance by construction, by never providing absolute positions as input to the models. Instead, the j-th input node’s features are the static properties, and a sequence of the three most recent position changes (i.e. average velocities), vjt = [z j ,∆pjt−2,∆p j t−1,∆p j t ], where, ∆pjt = p j t − pjt−1. For fDI and fC, which also take the solver’s current proposed Y (i), we also concatenate the proposed average velocity from the i-th solver iteration, yj,(i) − pjt , as input. For the input edge feature for an edge that connects from node j to k, we also provide the relative displacement vector between the nodes’ positions, ejkt = p k t − pjt .
GNN-based Encode-Process-Decode core We implemented fD, fDI, and fC using Graph Networks (GN) (Battaglia et al., 2018), arranged in the Encode-Process-Decode architecture, similar to previous work on GN-based learned simulators (Sanchez-Gonzalez et al., 2018; 2020; Pfaff et al., 2021). The Encoder uses two MLPs to encode node and edge features into high-dimensional latent vectors. The Processor applies multiple GNs, with unshared weights, in sequence, with node and edge residual connections at each step. We do not use global updates for the GNs. The Decoder uses an MLP to produce an output for each node.
The fD directly returns Ŷ . The fDI returns a change to the proposed update δY for the current iteration. The fC’s Decoder returns a scalar for each node to produce a constraint value per node {cj |j = 1 . . . |V |}. These node-wise constraint values are averaged to compute a single scalar c constraint for the entire system, c = fC(X≤t, Ŷ ) = 1|V | ∑|V | j=1 c j .
Solving the constraint For fDI and fC we initialize Y (0) = ∆pjt to the most recent average velocity4. We used auto-differentiation in JAX to compute the gradient function,∇Y fC, and the step size λwas specific to the model variant, as described below. During training we used N = 5 solver iterations.

heading: 4.3 MODEL VARIANTS

text:
The key questions in this work are whether constraint-based learned simulators can compete with explicit, forward learned simulators, whether implementing the constraint function with GNNs is more effective than with MLPs, and how minima-based constraint functions solved by gradient descent compare to constraints defined as the zeros of a function which are solved by fast projection (Goldenthal et al., 2007). The following model variants allow us to answer these questions.
Forward GNN This is an explicit, forward GNN-based learned simulator based on the GNS models from Sanchez-Gonzalez et al. (2020); Pfaff et al. (2021). It directly predicts the state update Ŷ from the past time points X≤t.
C-GNS Gradient Descent (C-GNS-GD) and C-GNS-Fast Projections (C-GNS-FP) These are our proposed constraint-based GNN models. For the C-GNS-GD, the scalar per-node output cj was squared, to force the overall fC to be non-negative, and a gradient descent solver with a fixed step size, λ = 0.001, was used to minimize it. For C-GNS-FP, the λ was based on “fast projection” (Goldenthal et al., 2007; Yang et al., 2020), as described in Section 3. Supplementary Figure B.5(c-d) shows ablations.
3For BOXBATH we vary a number of modelling choices to best match those in Sanchez-Gonzalez et al. (2020). The major difference is that we set Ŷ to be the average acceleration rather than average velocity. See Supplementary Materials for other differences.
4To ensure analogous information is provided downstream of fD, the update rule also includes the previous average velocity: p̂jt+1 = p j t + ∆p j t + ŷ j
Iterative GNN We implemented a hybrid between the Forward GNN and C-GNS, as shown in Figure 1c. It was identical to the C-GNS models, except its fDI directly predicted proposed state updates as in fD, rather than being computed via the gradients as was done with fC.
ConstraintMLP Gradient Descent (ConstraintMLP-GD) and ConstraintMLP-Fast Projections (ConstraintMLP-FP) These were MLP-based constraint models, which, rather than using GNNs to implement fC, instead concatenated the embeddings of all the input nodes into a single vector and passed them to an MLP implementation of fC. By default, these models cannot handle variable-length inputs, so we padded smaller states with zeros up to the maximum state size. The ConstraintMLP-FP was the MLP analog to our C-GNS-FP, and was similar to Neural Projections (Yang et al., 2020). The ConstraintMLP-GD used gradient descent, and was the MLP analog to our C-GNS-GD. We omit the results for the ConstraintMLP models on BOXBATH (1024 nodes), as MLPs do not generally work well on physical systems with more than a few particles (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2018).

heading: 4.4 TRAINING AND EVALUATION

text:
We trained the models to make next-step predictions, by computing the L2 loss between the predicted X̂t+1 and the corresponding ground truth Xt+1, averaged over nodes. All model weights and biases were trained using standard backpropagation with the Adam optimizer.
At test time, we compute 1-step metrics by evaluating the 1-step errors along each point of the ground truth trajectory. We also evaluate rollout errors by iteratively applying the learned model starting from an initial state, over 160 rollout steps, and computing the error between the predicted and ground truth trajectories.

heading: 5 RESULTS

text:
Predictive accuracy5 Our experimental results show that our C-GNS-GD’s performance was generally better than the other model variants. Figure 3 compares the different models on 1-step and rollout position MSE (see Supplementary Table B.1 for numerical results). For each dataset, we
5Videos of the model rollouts are available at sites.google.com/view/constraint-based-simulator
used the same number of message-passing steps (MP) for all GN-based models. We used 2 MPs for the ROPE dataset, and 1 MP for all other tasks.
The C-GNS-GD has lower 1-step MSE between the ground truth and predicted positions than other models across all datasets. Qualitatively, we observed that for Forward GNN with a single messagepassing step, the box in BOXBATH “melts” over time, as the forward model cannot preserve its rigid shape (see Videos). The comparable C-GNS-GD, by contrast, maintains the rigidity more effectively. These quantitative results suggest that constraint-based learned simulators are competitive alternative to explicit, forward learned simulators. We generally found that the Iterative GNN was fairly competitive with the C-GNS-GD in overall performance and better than the Forward GNN.
We also found that the C-GNS-FP was generally less stable across seeds, and not as accurate as the C-GNS-GD. The same conclusion holds for ConstraintMLP-FP versus ConstraintMLP-GD. We speculate that the fast projection algorithm may make training challenging because the step size λ is proportional to fC, which may cause poor zero-finding early in training when the fC is not yet informative. Additionally, we find that C-GNS-FP algorithm becomes unstable in the areas with shallow constraint gradients, perhaps because its λ depends on the inverse of the gradient’s norm.
We explored how varying the message-passing steps and solver iterations (N ) influenced the relative performance among the models in our ROPE dataset. Figure 4 shows that the C-GNS-GD generally required fewer parameters and message-passing steps to achieve comparable 1-step MSE to the other models. Supplementary Figure B.3 shows similar results for the rollout MSE. For most combinations of message-passing steps and number of solver iterations, C-GNS-GD (green) outperforms the Iterative GNN (yellow), C-GNS-FP (purple) as well as the Forward GNN (blue) with the same number of MPs (the Forward GNN is not iterative model, so we plot it as a single bar). We hypothesize that the solver iterations in the C-GNS and Iterative GNN may play a similar role to message passing with shared weights .
Interpreting the learned constraints To better understand the learned fC functions in the C-GNSGD, Figure 5 visualizes the node-wise constraint values as a function of Y (proposed average velocity) for different nodes in the ROPE dataset while holding the other nodes’ proposed update Y fixed.
We also overlay the sequence of five points that represent the proposed Y (i) steps from the solver where all nodes were jointly optimized. The figure shows the learned fC has a minimum near the ground truth Y, which the gradient descent steps are able to reach.
Incorporating novel constraints at test time We next explored a unique advantage of the constraint-based model: because the fC measures the degree the physical constraints are violated, we can incorporate additional, hand-designed constraints at test time, and use the model to potentially satisfy them. For the ROPE dataset, we designed three constraint functions that return positive values which increase quadratically as the rope enters different “forbidden” regions of the space: a vertical wall, a horizontal floor, and a disk-shaped region. We weighted these constraint terms by a coefficient hyperparameter and added each of the hand-designed constraints to the learned fC term of C-GNS-GD and ran the forward evaluation of the model.
As shown in Figure 6, the model was able to simulate the dynamics in a way that the corresponding forbidden region was avoided. In some cases, satisfying the joint constraint resulted in unintuitive behaviors, such as the rope links changing in length to adapt to the obstacle (Videos). However, this is to be expected, as the minimum of the joint constraint may not overlap with the minimum of the learned constraint, which is the one that would otherwise guarantee length preservation. For this example we added a further hand-designed constraint that incentivizes maintaining relative distances between nodes. In general this is a powerful example of how constraint-based models can generalize outside their training data, and solve both for the learned dynamics and arbitrary desired constraints.
Generalizing to larger systems via increased solver iterations In principle, iterative and constraint-based simulators should find more accurate solutions by increasing the number of solver iterations, N . We investigated whether the C-GNS-GD and Iterative GNN trained on ROPE could generalize from Ntrain = 5 on which they were trained, to Ntest ∈ [0, 15]. We also analyzed whether increased solver iterations could improve generalize performance from training on ropes with 5−10 nodes, to test ropes with 20 nodes.
Figure 7a (top row) shows that for test ropes that match the 5−10 nodes experienced during training, the Iterative GNN (light blue) overfits very heavily to Ntest = Ntrain = 5: error increases abruptly for N ≤ 4 and N ≥ 6. By contrast, the C-GNS-GD (light red) generalizes much better to different Ntest. Figure 7a (bottom row) shows that for test ropes with 20 nodes, the Iterative GNN again overfits, while the C-GNS-GD can generalize well to longer ropes if Ntest is increased.
We also trained the Iterative GNN and C-GNS-GD with additional loss terms that were applied to the Y (i) on each solver iteration, not only the final one, Ŷ = Y (N). We used an exponential decay factor, α = 0.25, which downweighted this additional loss term more heavily for earlier solver proposals. The dark blue and red curves in Figure 7a show how this additional loss further improves generalization to more solver iterations and larger systems as test time for the Iterative GNN, but especially the C-GNS-GD. Figure 7b visualizes how increasing the solver iterations systematically improves the quality of the long-term rollout accuracy in the ROPE dataset.
Together these results show the C-GNS-GD is effective in making use of additional resources at test time. This opens the exciting possibility of training on small, simple systems, and testing on large, complex systems. See Supplementary Figure B.2 for further details.

heading: 6 DISCUSSION

text:
We presented a general-purpose framework for constraint-based learned simulation, where a learned constraint function implicitly represents the dynamics, and future predictions are generated via a constraint solver. We implemented our framework using GNNs as the constraint function and gradient descent as the constraint solver, and tested it in a variety of challenging physical simulation problems. Our results showed that our C-GNS has competitive or better performance compared to previous learned simulators. We demonstrated unique abilities to generalize to novel, hand-designed constraints, and use more solver iterations than experienced during training to improve the accuracy on larger systems.
We can hypothesize about the relationship between explicit, forward learned simulators and implicit, constraint-based ones in terms of the sharing schemes of these architectures. The C-GNS has a stronger inductive bias than the Forward GNN. The transformation of fC in C-GNS effectively ties the parameters in the resulting ∇Y fC function, and the solver iterations are analogous to how a recurrent neural network’s parameters are shared over iterations. In contrast, the message-passing steps in the Forward GNN used in our work are unshared. In principle, the fD of the Forward GNN is more expressive because if given enough depth, after training it could learn to take parameter values that are equivalent to the shared parameters of C-GNS. Our results shown in Figure 4 supports this possibility: the Forward GNN with many more message-passing steps eventually approaches the C-GNS’s performance. Moreover, we speculate the C-GNS’s inductive biases contribute to its advantages in terms of incorporating novel hand-designed constraints and generalizing to more solver iterations and larger systems.
More broadly, the performance, generality and unique advantages of constraint-based learned simulation make it an important new direction in the advancement of machine learning methods for complex simulation problems in science and engineering.

heading: 7 REPRODUCIBILITY STATEMENT

text: We are committed to open-source the model code after the paper is accepted. Also, we are going to open-source the MuJoCo datasets that we generated for this paper. We provide more details on the model implementation as well as the hyperparameters used for each model in the Supplementary Material.