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: This paper studies improving the modeling of physical dynamics with equivariant neural networks. In particular, this paper focuses on a new type of data governed by physical models. Several special symmetry groups are considered to better characterize the system, including uniform motion equivariance, resolution-independent scaling, and resolution-dependent scaling, etc. Simulation results show that the proposed equivariant model yields better accuracy and physical consistency than the non-equivariant models even with data augmentation, given the type of distributional shift is known. Results on the real-world data show some of the equivariant models can generalize better than the non-equivariant models.