Noether Networks: meta-learning useful conserved quantities

1 Introduction . The clever use of inductive biases to exploit symmetries has been at the heart of many landmark achievements in machine learning , such as translation invariance in CNN image classification [ 35 ] , permutation invariance in Graph Neural Networks [ 52 ] for drug design [ 57 ] , or roto-translational equivariance in SE3-transformers [ 25 ] for protein structure prediction [ 33 ] . However , for data distributions of interest , there may be exploitable symmetries that are either unknown or difficult to describe with code . Progress has been made in automatically discovering symmetries for finite groups [ 74 ] , but meta-learning and exploiting general continuous symmetries has presented a major challenge . In part , this is because symmetries describe the effect of counterfactuals about perturbations to the data , which are not directly observable . In this work , we propose to exploit symmetries in sequential prediction problems indirectly . We take inspiration from Noether ’ s theorem [ 44 ] , which loosely states the following : For every continuous symmetry property of a dynamical system , there is a corresponding quantity whose value is conserved in time . For example , consider a system of planets interacting via gravity : this system is translation invariant in all three cardinal directions ( i.e . translating the entire system in the x , y , or z axis conserves the laws of motion ) . Noether ’ s theorem asserts there must be a conserved quantity for each of these symmetries ; in this case , linear momentum . Similarly , the system has a time-invariance ( i.e . the laws of motion are the same today as they will be tomorrow ) . In this case , the corresponding conserved quantity is the total energy of the system . ∗Equal contribution . Our code is publicly available at https : //lis.csail.mit.edu/noether . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Inspired by this equivalence , we propose that approximate conservation laws are a powerful paradigm for meta-learning useful inductive biases in sequential prediction problems . Whereas symmetries are difficult to discover because they are global properties linked to counterfactuals about unobserved perturbations of the data , conserved quantities can be directly observed in the true data . This provides an immediate signal for machine learning algorithms to exploit . Our approach involves meta-learning a parametric conservation loss function which is useful to the prediction task . We leverage the tailoring framework [ 4 ] , which proposes to encode inductive biases by fine-tuning neural networks with hand-designed unsupervised losses inside the prediction function . Whereas traditional auxiliary losses are added to the task loss during training , tailoring losses are optimized inside the prediction function both during training and testing , customizing the model to each individual query . In doing so , we ensure there is no generalization gap for the conservation loss . We propose to meta-learn the self-supervised loss function and parameterize it in the form of a conservation loss ; i.e . L ( x0 , x̃1 : T ; gφ ) = �T t=1 |gφ ( x0 ) − gφ ( x̃t ) |2 . This conservation form encodes a meta-inductive bias ( inductive bias over inductive biases ) which narrows the search space exponentially ( in T ) and simplifies the parameterization . Figure 1 demonstrates the Noether Network pipeline for a video prediction task . The main contribution of this paper is the Noether Network , an architecture class and algorithm that can automatically learn its own inductive biases in the form of meta-learned conservation loss functions . We theoretically characterize the advantages of such conservation laws as effective regularizers that constrain the learning to lower-dimensional output manifolds , lowering the generalization gap . Empirically , we first find that , when the meta-learned conservation loss takes the form of a synthesized program , Noether Networks recover known conservation laws from raw physical data . Second , we find that Noether Networks can learn conservation losses on raw videos , modestly improving the generalization of a video prediction model , especially on longer-term predictions . 2 Theoretical Advantages of Enforcing Conservation Laws . In this section , we demonstrate principled advantages of enforcing conservation laws of the form gφ ( fθ ( x ) ) = gφ ( x ) by considering a special case where preimages under gφ form affine subspaces . Let input x and target y be x , y ∈ Rd , and let the Noether embedding be gφ : Rd → P where P = { gφ ( x ) : x ∈ Rd } . We consider a restricted class of models , parameterized by θ ∈ Θ , of the form fθ ( x ) = x + vθ for vθ ∈ Rd such that for all x , the preimage of gφ is g−1φ [ { gφ ( x ) } ] = { x + Az : z ∈ Rm } , A ∈ Rd×m . Here , m ≤ d is the dimensionality of the preimages of gφ . We denote by C the smallest upper bound on the loss value as L ( fθ ( x ) , y ) ≤ C ( for all x , y and θ ) . Define ψ ( v ) = Ex , y [ L ( f ( x , v ) , y ) ] − 1n �n i=1 L ( f ( xi , v ) , yi ) , with Lipschitz constant ζ . Therefore , ζ ≥ 0 is the smallest real number such that , for all v and v� in V , |ψ ( v ) −ψ ( v� ) | ≤ ζ�v−v��2 , where V = { vθ ∈ Rd : θ ∈ Θ } . Finally , we define R = supθ∈Θ �vθ�2 , and the generalization gap by G ( θ ) = Ex , y [ L ( fθ ( x ) , y ) ] − 1 n n� i=1 L ( fθ ( xi ) , yi ) . Theorem 1 shows that enforcing conservation laws of gφ ( fθ ( x ) ) = gφ ( x ) is advantageous when the dimension of the preimages under gφ is less than the dimension of x ; that is , when m < d. Theorem 1 . Let ρ ∈ N+ . Then , for any δ > 0 , with probability at least 1− δ over an iid draw of n examples ( ( xi , yi ) ) ni=1 , the following holds for all θ ∈ Θ : G ( θ ) ≤ C � ξ ln ( max ( √ ξ , 1 ) ) + ξ ln ( 2R ( ζ1−1/ρ ) √ n ) + ln ( 1/δ ) 2n + � { ξ ≥ 1 } � ζ2/ρ n . ( 1 ) where ξ = m if gφ ( fθ ( x ) ) = gφ ( x ) for any x ∈ Rd and θ ∈ Θ , and ξ = d otherwise . The proof is presented in Appendix A . In Theorem 1 , when we enforce the conservation laws , d is replaced by m , the dimension of the preimage . We now discuss various cases for the values of m : • ( Case of m = d ) Let us consider an extreme scenario where the function gφ maps all x ∈ Rd to one single point . In this scenario , the dimensionality of preimages under gφ is maximized as m = d. Accordingly , the bounds with and without enforcing the conservation laws become the same . Indeed , in this scenario , the conservation laws of gφ ( fθ ( x ) ) = gφ ( x ) give us no information , because they always hold for all x ∈ Rd , even without imposing them . • ( Case of m = 0 ) Let us consider another extreme scenario where the function gφ is invertible . In this scenario , the dimensionality of preimages under gφ is zero as m = 0 . Thus , imposing the condition of gφ ( fθ ( x ) ) = gφ ( x ) makes the bound in Theorem 1 to be very small . Indeed , in this scenario , the condition of gφ ( fθ ( x ) ) = gφ ( x ) implies that fθ ( x ) = x : i.e. , x is not moving , and thus it is easy to generalize . • ( Case of 0 < m < d ) From the above two cases , we can see that the benefit of enforcing conservation laws comes from more practical cases in-between these , with 0 < m < d. In Theorem 1 , the function gφ can differ from the true function g∗φ underlying the system . This is because we analyze a standard generalization gap : i.e. , the difference between the expected loss and the training loss . The cost of not using the true g∗φ is captured in the training loss ; i.e. , the training loss can be large with the function gφ that differs significantly from the true g∗φ . Even in this case , the generalization gap can be small . For example , in the case of m = 0 , the generalization bound is small , whereas the training loss will be large unless xt+1 = xt . Therefore , our analysis gives us the insight on the trade-off between the training loss and the dimensionality of preimages under gφ . 3 Noether Networks . Leveraging tailoring to encode inductive biases . We perform a prediction-time optimization to encourage outputs to follow conservation laws using the tailoring framework [ 4 ] . Tailoring encodes inductive biases in the form of unsupervised losses optimized inside the prediction function . In doing so , tailoring fine-tunes the model to each query to ensure that it minimizes the unsupervised loss for that query . For example , we may optimize for energy conservation in a physics prediction problem . In meta-tailoring , we train the model to do well on the task loss after the tailoring step has fine-tuned its parameters . In contrast to auxiliary losses , which would optimise the conservation loss only for the training points , tailoring allows us to ensure conservation at test time . Since we aim to build in the conservation in the architecture , we want to ensure it is also satisfied for unseen test samples . Another advantage of tailoring losses is that they are easier to meta-learn . Auxiliary losses are pooled over all examples and training epochs and their effect is only known at validation/test time . We would need to use implicit gradients [ 39 , 47 ] to know their eventual effect on the final weights at the end of training . With tailoring , we can directly measure the effect of the meta-learned update on the same sample . A limitation of tailoring framework is that the tailoring loss must be user-specified . This is acceptable in domains where the desired inductive bias is both known and easily encoded , but problematic in general — we address this issue with Noether Networks . Our approach can be seen as a generalization of tailoring where the unsupervised loss is meta-learned and takes the form of a conservation loss . Noether Networks : discovering useful conservation laws . We build on the paradigm of conservation laws ; i.e . quantities that are conserved over time . This has the following two challenges : 1 . Real data often has noise , breaking any exact conservation law . Moreover , many conservation laws are only approximately satisfied given only partial information of the system . For example , conservation of energy is not fully satisfied in real dissipative systems and estimating conservation of mass from pixels is bound to be inexact in the case of occlusions . 2 . Optimizing for conservation alone can lead to trivial quantities , such as predicting a constant value gφ ( x ) = C independent of x . We search for useful conservation laws , whose ( approximate ) enforcing brings us closer to the true data manifold for the current sequence . Note that a useful conserved quantity doesn ’ t need to be exactly conserved in the training data to improve prediction performance . We only need its conservation to be informative of the data manifold . This allows us to discover conservation laws even in imperfect , noisy environments . Conversely , tailoring with a trivial conserved quantity can not improve the final predictions ( see App . B for a detailed explanation ) . Finally , viewing conserved quantities as useful inductive biases aligns well with their use , since we often care about biases only insofar as they help improve task performance . The search for such useful conserved quantities can take many forms . In this work , we focus on two : first , in the beginning of section 4.1 , we use a combination of program synthesis and gradient descent to generate a large , but finite , number of candidate parametric formulas for physical system prediction . We then try meta-tailoring with each formula on the training set and pick the conservation formula with the best loss . Formulas can be a useful , interpretable description in scientific domains given the appropriate state-space . However , for general descriptions from raw data , we would like to describe the loss with a neural network , as we describe in the next section . Meta-learning a neural loss function . We propose Noether Networks , an architecture class for sequential prediction that consists of a base predictor fθ : x̃t �→ x̃t+1 and a meta-learned tailoring loss , parameterized as the conservation of a neural embedding gφ : x̃t �→ Rk . This embedding takes raw predictions as input ( such as images in the case of video prediction ) . The conservation form induces a meta-inductive bias over potential learned tailoring losses . The Noether loss is formulated as LNoether ( x0 , x̃1 : T ; gφ ) = T� t=1 |gφ ( x0 ) − gφ ( x̃t ) |2 � �� � ( a ) ≈ T� t=1 |gφ ( x̃t−1 ) − gφ ( x̃t ) |2 � �� � ( b ) ( 2 ) where x0 is the ground-truth input , the x̃t = fθ ( x̃t−1 ) are the model ’ s predictions , and x̃0 � x0 . Expressions 2 ( a ) and 2 ( b ) are equivalent if we fully enforce the conservation law , but they differ if conservation is not fully enforced . When not fully conserving , 2 ( a ) propagates information from ground truth more directly to the prediction , but 2 ( b ) may be a useful approximation which better handles imperfectly conserved quantities , where the quantity should be Lipschitz but not exactly constant . In both cases , if we tailor θ with a single gradient step ; the gradient update takes the form θ ( x0 ; φ ) = θ − λin∇θLNoether ( x0 , x̃1 : T ( θ ) ; gφ ) . ( 3 ) We compute final predictions as x̂t = fθ ( x0 ; φ ) ( x̂t−1 ) . We can now backpropagate from Ltask ( x1 : T , x̂1 : T ) back to φ , which will be optimized so that the unsupervised adaptation θ �→ θ ( x0 ; φ ) helps lower Ltask.The optimization requires second-order gradients to express how φ affects Ltask through θ ( x0 ; φ ) . This is similar to MAML [ 23 ] , as well as works on meta-learning loss functions for few-shot learning [ 5 ] and group distribution shift [ 71 ] . Algorithm 1 provides pseudo-code . For deep learning frameworks that allow per-example weights , such as JAX [ 9 ] , the loop over sequences in Alg . 1 can be efficiently parallelized . To parallelize it for other frameworks we use the CNGRAD algorithm [ 4 ] , which adapts only the Conditional Normalization ( CN ) layers in the inner loop . Similar to BN layers , CN only performs element-wise affine transformations : y = x� γ + β , which can be efficiently parallelized in most deep learning frameworks even with per-example γ , β . Finally , even though we use the two-layer optimization typical of meta-learning , we are still in the classical single-task single-distribution setting . Noether Networks learn to impose their own inductive biases via a learned loss and to leverage it via an adaptation of its parameters θ . This is useful as we often do not have access to meta-partitions that distinguish data between distributions or tasks . Algorithm 1 Prediction and training procedures for Noether Networks with neural conservation loss Given : predictive model class f ; embedding model class g ; prediction horizon T ; training dist . Dtrain ; batch size N ; learning rates λin , λout , λemb ; task loss Ltask ; Noether loss LNoether 1 : procedure PREDICTSEQUENCE ( x0 ; θ , φ ) 2 : x̃0 , x̂0 ← x0 , x0 3 : x̃t ← fθ ( x̃t−1 ) ∀t ∈ { 1 , . . . , T } � Initial predictions 4 : θ ( x0 ; φ ) ← θ − λin∇θLNoether ( x0 , x̃1 : T ; gφ ) � Inner step with Noether loss 5 : x̂t ← fθ ( x0 ; φ ) ( x̂t−1 ) ∀t ∈ { 1 , . . . , T } � Final prediction with tailored weights 6 : return x̂1 : T 7 : procedure TRAIN 8 : φ ← randomly initialized weights � Initialize weights for Noether embedding g 9 : θ ← randomly initialized weights � Initialize weights for predictive model f 10 : while not done do 11 : Sample batch x ( 0 ) 0 : T , . . . , x ( N ) 0 : T ∼ Dtrain 12 : for 0 ≤ n ≤ N do 13 : x̂ ( n ) 1 : T ← PREDICTSEQUENCE ( x ( n ) 0 ; θ , φ ) 14 : φ ← φ− λemb∇φ �N n=0 Ltask ( x̂ ( n ) 1 : T , x ( n ) 1 : T ) � Outer step with task loss for embedding 15 : θ ← θ − λout∇θ �N n=0 Ltask ( x̂ ( n ) 1 : T , x ( n ) 1 : T ) � Outer step for predictive model 16 : return φ , θ

Summary: The authors tackle the challenge of automatically discovering symmetries in dynamical systems. As continuous symmetries are equivalently described by conserved quantities--as stated by Noether's theorem--they introduce a method for learning (approximately) conserved quantities. They do so by building on the `meta-tailoring' framework which allows to fine-tune the model according to the learned conservation laws at test time. For problems with few and interpretable coordinates (e.g. position, momentum, etc) they propose a symbolic approach that generate many valid (i.e. with compatible physical units) expressions that are parametrized with learnable weights. For raw-pixel data, they parametrize and learn the function to be conserved with a neural network. In both cases, the learned conservation function is used to guide the model towards physically 'realistic' solutions. They empirically assess their method on pendulum and spring systems, but also on real data from the Physics 101 dataset.