Incorporating Symmetry into Deep Dynamics Models for Improved Generalization

1 INTRODUCTION . Modeling dynamical systems in order to forecast the future is of critical importance in a wide range of fields including , e.g. , fluid dynamics , epidemiology , economics , and neuroscience [ 2 ; 21 ; 45 ; 22 ; 14 ] . Many dynamical systems are described by systems of non-linear differential equations that are difficult to simulate numerically . Accurate numerical computation thus requires long run times and manual engineering in each application . Recently , there has been much work applying deep learning to accelerate solving differential equations [ 46 ; 6 ] . However , current approaches struggle with generalization . The underlying problem is that physical data has no canonical frame of reference to use for data normalization . For example , it is not clear how to rotate samples of fluid flow such that they share a common orientation . Thus real-world out-of-distribution test data is difficult to align with training data . Another limitation of current approaches is low physical accuracy . Even when mean error is low , errors are often spatially correlated , producing a different energy distribution from the ground truth . We propose to improve the generalization and physical accuracy of deep learning models for physical dynamics by incorporating symmetries into the forecasting model . In physics , Noether ’ s Law gives a correspondence between conserved quantities and groups of symmetries . By building a neural network which inherently respects a given symmetry , we thus make conservation of the associated quantity more likely and consequently the model ’ s prediction more physically accurate . ∗Equal contribution A function f is equivariant if when its input x is transformed by a symmetry group g , the output is transformed by the same symmetry , f ( g · x ) = g · f ( x ) . See Figure 1 for an illustration . In the setting of forecasting , f approximates the underlying dynamical system . The set of valid transformations g is called the symmetry group of the system . Figure 1 : Illustration of equivariance of e.g . f ( x ) = 2x with respect to T = rot ( π/4 ) . By designing a model that is inherently equivariant to transformations of its input , we can guarantee that our model generalizes automatically across these transformations , making it robust to distributional shift . The symmetries we consider , translation , rotation , uniform motion , and scale , have different properties , and thus we tailor our methods for incorporating each symmetry . Specifically , for scale equivariance , we replace the convolution operation with group correlation over the group G generated by translations and rescalings . Our method builds on that of Worrall and Welling [ 51 ] , with significant novel adaptations to the physics domain : scaling affecting time , space , and magnitude ; both up and down scaling ; and scaling by any real number . For rotational symmetries , we leverage the key insight of Cohen and Welling [ 9 ] that the input , output , and hidden layers of the network are all acted upon by the symmetry group and thus should be treated as representations of the symmetry group . Our rotation-equivariant model is built using the flexible E ( 2 ) -CNN framework developed by Weiler and Cesa [ 49 ] . In the case of a uniform motion , or Galilean transformation , we show the above methods are too constrained . We use the simple but effective technique of convolutions conjugated by averaging operations . Research into equivariant neural networks has mostly been applied to tasks such as image classification and segmentation [ 27 ; 50 ; 49 ] . In contrast , we design equivariant networks in a completely different context , that of a time series representing a physical process . Forecasting high-dimensional turbulence is a significant step for equivariant neural networks compared to the low-dimensional physics examples and computer vision problems treated in other works . We test on a simulated turbulent convection dataset and on real-world ocean current and temperature data . Ocean currents are difficult to predict using numerical methods due to unknown external forces and complex dynamics not fully captured by simplified mathematical models . These domains are chosen as examples , but since the symmetries we focus on are pervasive in almost all physics problems , we expect our techniques will be widely applicable . Our contributions include : • We study the problem of improving the generalization capability and physical accuracy of deep learning models for learning complex physical dynamics such as turbulence and ocean currents . • We design tailored methods with theoretical guarantees to incorporate various symmetries , including uniform motion , rotation , and scaling , into convolutional neural networks . • When evaluated on turbulent convection and ocean current prediction , our models achieve significant improvement on generalization of both predictions and physical consistency . • For different symmetries , our methods have an average 31 % and maximum 78 % reduction in energy error when evaluated on turbulent convection with no distributional shift . 2 MATHEMATICAL PRELIMINARIES . 2.1 SYMMETRY GROUPS AND EQUIVARIANT FUNCTIONS . Formal discussion of symmetry relies on the concept of an abstract symmetry group . We give a brief overview , for a more formal treatment see Appendix A , or Lang [ 28 ] . A group of symmetries or simply group consists of a set G together with a composition map ◦ : G×G→ G. The composition map is required to be associative and have an identity 1 ∈ G. Most importantly , composition with any element of G is required to be invertible . Groups are abstract objects , but they become concrete when we let them act . A group G has an action on a set S if there is an action map · : G× S → S which is compatible with the composition law . We say further that S is a G-representation if the set S is a vector space and the group acts on S by linear transformations . Definition 1 ( invariant , equivariant ) . Let f : X → Y be a function and G be a group . Assume G acts on X and Y . The function f is G-equivariant if f ( gx ) = gf ( x ) for all x ∈ X and g ∈ G. The function f is G-invariant if f ( gx ) = f ( x ) for all x ∈ X and g ∈ G . 2.2 PHYSICAL DYNAMICAL SYSTEMS . We investigate two dynamical systems : Rayleigh–Bénard convection and real-world ocean current and temperature . These systems are governed by Navier-Stokes equations . 2D Navier-Stokes ( NS ) Equations . Let w ( x , t ) be the velocity vector field of a flow . The field w has two components ( u , v ) , velocities along the x and y directions . The governing equations for this physical system are the momentum equation , continuity equation , and temperature equation , ∂w ∂t = − ( w · ∇ ) w − 1 ρ0 ∇p+ ν∇2w + f ; ∇ ·w = 0 ; ∂H ∂t = κ∆H − ( w · ∇ ) H , ( DNS ) where H ( x , t ) is temperature , p is pressure , κ is the heat conductivity , ρ0 is initial density , α is the coefficient of thermal expansion , ν is the kinematic viscosity , and f is the buoyant force . 2.3 SYMMETRIES OF DIFFERENTIAL EQUATIONS . By classifying the symmetries of a system of differential equations , the task of finding solutions is made far simpler , since the space of solutions will exhibit those same symmetries . Let G be a group equipped with an action on 2-dimensional space X = R2 and 3-dimensional spacetime X̂ = R3 . Let V = Rd be a G-representation . Denote the set of all V -fields on X̂ as F̂V = { w : X̂ → V : w smooth } . Define FV similarly to be V -fields on X . Then G has an induced action on F̂V by ( gw ) ( x , t ) = g ( w ( g−1x , g−1t ) ) and on FV analogously . Consider a system of differential operatorsD acting on F̂V . Denote the set of solutions Sol ( D ) ⊆ F̂V . We sayG is a symmetry group ofD ifG preserves Sol ( D ) . That is , ifϕ is a solution ofD , then for all g ∈ G , g ( ϕ ) is also . In order to forecast the evolution of a systemD , we model the forward prediction function f . Let w ∈ Sol ( D ) . The input to f is a collection of k snapshots at times t− k , . . . , t− 1 denoted wt−i ∈ Fd . The prediction function f : Fkd → Fd is defined f ( wt−k , . . . , wt−1 ) = wt . It predicts the solution at a time t based on the solution in the past . Let G be a symmetry group of D. Then for g ∈ G , g ( w ) is also a solution of D. Thus f ( gwt−k , . . . , gwt−1 ) = gwt . Consequently , f is G-equivariant . 2.4 SYMMETRIES OF NAVIER-STOKES EQUATIONS . The Navier-Stokes equations are invariant under the following five different transformations . Individually each of these types of transformations generates a group of symmetries of the system . The full list of symmetry groups of NS equations and Heat equations are shown in Appendix B.6 . • Space translation : T spc w ( x , t ) = w ( x− c , t ) , c ∈ R2 , • Time translation : T timeτ w ( x , t ) = w ( x , t− τ ) , τ ∈ R , • Uniform motion : T umc w ( x , t ) = w ( x , t ) + c , c ∈ R2 , • Rotation/Reflection : T rotR w ( x , t ) = Rw ( R−1x , t ) , R ∈ O ( 2 ) , • Scaling : T scλ w ( x , t ) = λw ( λx , λ2t ) , λ ∈ R > 0 . 3 METHODOLOGY . We prescribe equivariance by training within function classes containing only equivariant functions . Our models can thus be theoretically guaranteed to be equivariant up to discretization error . We incorporate equivariance into two state-of-the-art architectures for dynamics prediction , ResNet and U-net [ 48 ] . Below , we describe how we modify the convolution operation in these models for different symmetries G to form four EquG-ResNet and four EquG-Unet models . 3.1 EQUIVARIANT NETWORKS . The key to building equivariant networks is that the composition of equivariant functions is equivariant . Hence , if the maps between layers of a neural network are equivariant , then the whole network will be equivariant . Note that both the linear maps and activation functions must be equivariant . An important consequence of this principle is that the hidden layers must also carry a G-action . Thus , the hidden layers are not collections of scalar channels , but vector-valued G-representations . Equivariant Convolutions . Consider a convolutional layer FRdin → FRdout with kernel K from a Rdin-field to a Rdout-field . Let Rdin and Rdout be G-representations with action maps ρin and ρout respectively . Cohen et al . [ 11 , Theorem 3.3 ] prove the network is G-equivariant if and only if K ( gv ) = ρ−1out ( g ) K ( v ) ρin ( g ) for all g ∈ G. ( 1 ) A network composed of such equivariant convolutions is called a steerable CNN . Equivariant ResNet and U-net . Equivariant ResNet architectures appear in [ 9 ; 10 ] , and equivariant transposed convolution , a feature of U-net , is implemented in [ 49 ] . We prove in general that adding skip connections to a network does not affect its equivariance with respect to linear actions and also give a condition for ResNet or Unet to be equivariant in Appendix B.2 . Relation to Data Augmentation . To improve generalization , equivariant networks offer a better performing alternative to the popular technique of data augmentation [ 13 ] . Large symmetry groups normally require augmentation with many transformed examples . In contrast , for equivariant models , we have following proposition . ( See Appendix B.1 for proof . ) Proposition 1 . G-equivariant models with equivariant loss learn equally ( up to sample weight ) from any transformation g ( s ) of a sample s. Thus data augmentation does not help during training .

Summary: The paper demonstrates that incorporating equivariance (i.e. symmetries) into model for predicting fluid dynamics improves its performance, especially when the test distribution is transformed by those symmetry groups. Leveraging the recent literature on equivariant CNNs, the paper proposes a CNN model that is equivariant with respect to known symmetries of the Navier-Stokes equations (time/space translation, rotation, uniform motion, and scaling). This approached is validated on 2 datasets on fluid dynamics: a synthetic dataset on Rayleigh-Benard convection and a real-world ocean dynamics dataset. On the synthetic dataset, the proposed models demonstrate better performance under distribution shift. On the real-world dataset, the models yields predictions that are more accurate and physically consistent.