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: **Summary**:  This paper is concerned with the problem of automatic symmetry discovery, by meta-learning conserved quantities that are useful for predictions. They propose Noether Networks, where the learned symmetries take the form of a meta-learned conservative loss optimised inside the prediction function during test time. They demonstrate experimentally that the proposed approach can improve prediction quality and discover useful inductive biases in sequential problems.  **Contributions**: - This paper proposes a new algorithm that automatically learns its inductive biases in the form of a meta-learned conservative loss optimised inside the prediction function. - (In a simplified setting) They theoretically show when and why enforcing conservative laws brings benefits. - They demonstrate through experiments that the proposed algorithm can recover known physical laws, and improve prediction performance for video data.