Source-linked AI summary
Incorporating Symmetry into Deep Dynamics Models for Improved Generalization
Rui Wang, Robin Walters, Rose Yu
TL;DR
Deep models for physical dynamics are limited by physical inaccuracies and weak generalization under distributional shift. The paper builds convolutional networks with symmetry-specific equivariance methods and evaluates them on convection and ocean dynamics. The resulting models improve prediction generalization and physical consistency, including 31% average and 78% maximum energy-error reductions on turbulent convection without distributional shift.
Problem
Deep-learning models for physical dynamics have limited physical accuracy and struggle to generalize when test data undergoes distributional shifts.
Method
The paper designs convolutional neural networks with tailored equivariance methods for translation, rotation, uniform motion, and scaling symmetries.
Results
The models significantly improve generalization and physical consistency on turbulent convection and ocean-current prediction, with 31% average and 78% maximum energy-error reductions on turbulent convection without distributional shift.
Takeaways & Limitations
Equivariant networks provide a theoretically and experimentally supported approach for forecasting complex physical dynamics under symmetry-group distributional shifts.
Takeaways & Limitations
Requiring every CNN layer to enforce uniform-motion equivariance severely limits the model, motivating a relaxed shifted-distribution construction.
Abstract
from arXiv · showhide
Recent work has shown deep learning can accelerate the prediction of physical dynamics relative to numerical solvers. However, limited physical accuracy and an inability to generalize under distributional shift limit its applicability to the real world. We propose to improve accuracy and generalization by incorporating symmetries into convolutional neural networks. Specifically, we employ a variety of methods each tailored to enforce a different symmetry. Our models are both theoretically and experimentally robust to distributional shift by symmetry group transformations and enjoy favorable sample complexity. We demonstrate the advantage of our approach on a variety of physical dynamics including Rayleigh Bénard convection and real-world ocean currents and temperatures. Compared with image or text applications, our work is a significant step towards applying equivariant neural networks to high-dimensional systems with complex dynamics. We open-source our simulation, data, and code at \url{https://github.com/Rose-STL-Lab/Equivariant-Net}.
1 INTRODUCTION
Deep models for physical dynamics face slow numerical baselines, limited physical accuracy, and poor generalization under distributional shift. This paper incorporates tailored symmetries into convolutional networks and evaluates them on turbulent convection and ocean data.
- Numerical simulation of nonlinear dynamical systems can require long runtimes and manual engineering for accurate computation.
- Existing deep-learning approaches struggle with distributional shift because physical data lacks a canonical frame for normalization and alignment.Low mean error can still produce spatially correlated errors and incorrect energy distributions.
- The paper incorporates translation, rotation, uniform-motion, and scale symmetries into convolutional neural networks using methods tailored to each symmetry.For scale, group correlation replaces convolution; rotational designs treat inputs, outputs, and hidden layers as group representations.
- The work extends equivariant neural networks from primarily image and segmentation tasks to high-dimensional physical time-series dynamics.
- The models improve generalization and physical consistency on turbulent convection and real-world ocean-current prediction.These experiments cover simulated turbulent convection and real-world ocean current and temperature data.
- 31% average and 78% maximum energy-error reductions are achieved on turbulent convection without distributional shift across different symmetries.
2 MATHEMATICAL PRELIMINARIES
The paper formulates physical forecasting through symmetry groups acting on fields and solutions of differential equations. Equivariance then follows for the forecasting map, while Navier–Stokes supplies concrete transformation groups and field equations.
- A symmetry group is a set with composition, an identity, and invertible composition elements.
- A group action maps group elements and set elements to the set compatibly with composition; a representation is a vector space with linear group actions.
- A function is equivariant when transforming its input transforms its output correspondingly, whereas an invariant function produces the same output after transformation.
- 2.2 PHYSICAL DYNAMICAL SYSTEMS: The study models Navier–Stokes dynamics, whose fields include velocity and temperature governed by momentum, continuity, and temperature equations.The formulation names pressure, heat conductivity, density, thermal expansion, viscosity, and buoyant force as physical quantities.
- 2.3 SYMMETRIES OF DIFFERENTIAL EQUATIONS: A differential-equation symmetry group preserves its solution set, so transformed solutions remain valid solutions of the system.
- 2.3 SYMMETRIES OF DIFFERENTIAL EQUATIONS: Forecasting from past snapshots is G-equivariant when transformed input histories produce the correspondingly transformed future state.
- 2.3 SYMMETRIES OF DIFFERENTIAL EQUATIONS: Navier–Stokes admits space and time translations, uniform motion, rotation/reflection, and scaling as transformation types.
3 METHODOLOGY
The methodology builds CNNs whose layers and hidden representations respect selected physical symmetries, providing equivariance up to discretization error. It tailors transformations for translation, rotation, uniform motion, and scale, including specialized relaxations for uniform motion and physically coupled scaling.
- General equivariant architectures: Equivariant ResNet and U-net models prescribe symmetry-respecting function classes, with theoretical guarantees up to discretization error.Equivariance is maintained through equivariant layer maps, activations, and hidden-layer group representations.
- General equivariant architectures: G-equivariant models with equivariant loss learn equally, up to sample weight, from transformed samples, so augmentation does not help during training.This is presented as a proposition contrasting equivariant models with conventional data augmentation.
- Rotation: Rotation is implemented with SO(2)-equivariant convolutions and activations, approximated in practice by the cyclic group Cn under spatial discretization.Hidden layers use powers of the regular representation of Cn.
- Uniform motion: Uniform-motion equivariance severely constrains alternating convolutional CNNs: if every combined layer is equivariant, the resulting network is affine.The method relaxes this requirement by centering each local input block and restoring its mean after convolution and activation.
- Scale: Scale methods either vary discretization through magnitude scaling or replace convolution with group correlation over translations and positive rescalings.The scale-aware representation adds a scale axis; channels are treated as a time dimension, increasing temporal resolution in deeper layers.
- Scale: Scale equivariance couples magnitude, space, and time according to the physical law, supporting both upscaling and downscaling for any λ ∈ R>0.The implementation truncates the scale sum to seven values between 1/3 and 3 and uses interpolation for discrete data.
4 RELATED WORK
Prior equivariant-neural-network research largely targeted image recognition, while this work applies equivariance to forecasting high-dimensional physical time series. Related physics and video-prediction studies address adjacent problems but differ in task or physical structure.
- Equivariant neural networks: Earlier equivariant neural networks were mainly developed for image classification and segmentation, not physical-dynamics forecasting.The paper positions high-dimensional turbulence forecasting as a substantially different application context.
- Equivariant neural networks: The paper differs from representation-learning approaches that seek invariance to specified factors, because it targets physical symmetries.The cited prior approaches are described as distinct from the symmetries modeled here.
- Physics-informed deep learning: Physics-informed deep-learning studies include CFD–U-net prediction, unsupervised turbulent-flow generation, and rotationally covariant molecular modeling.The paper distinguishes its forecasting setting from these neighboring applications.
- Video prediction: Video-prediction models forecast future frames from observations, whereas this work predicts physical processes without modeling object or camera motions.Natural-video models also face difficulty explicitly incorporating physical principles from unknown processes.
5 EXPERIMENTS
Experiments evaluate equivariant ResNet and U-net variants on simulated convection and real-world ocean dynamics, using prediction accuracy, physical consistency, equivariance, and distribution-shift tests. Across transformed and future-domain data, equivariant models generally outperform non-equivariant baselines and data-augmentation models.
- Experimental setup: Experiments use Rayleigh-Bénard convection and real-world ocean current data to evaluate equivariant ResNet and U-net models.Models predict raw velocity and temperature fields autoregressively, with errors reported over five random runs.
- Evaluation metrics: RMSE measures pixelwise prediction accuracy, while ESE measures the RMSE of log energy spectra as an indicator of physical consistency.ESE assesses whether predictions preserve fluid statistical distributions and obey energy conservation.
- Equivariance evaluation: Equivariance errors quantify the discrepancy between transforming a prediction and predicting from a transformed input.The study measures these errors for ResNet and Equ-ResNet models using transformations sampled as in the transformed convection test sets.
- Rayleigh-Bénard convection: On original convection data, models have similar RMSE, while equivariant models achieve lower ESE; on transformed sets, Equ-ResNets outperform ResNet and augmented ResNets.The comparison uses four transformed test sets and an augmented baseline trained with a dataset three times the original size.
- Distribution-shift generalization: EquScal-ResNet remains robust across scale factors from 1.5 to 1, whereas ResNet does not generalize across the corresponding distribution shifts.Figure 3 reports RMSE and ESE over five runs on upscaled convection test sets.
- Real-world ocean dynamics: Equivariant models outperform non-equivariant baselines on future and different-domain ocean data; EquScal-ResNet has the lowest RMSE, while EquRot models have significantly lower ESE than others.Equivariant models also outperform data-augmentation baselines, whose ESE is uniformly worse than that of the baselines.
6 CONCLUSION AND FUTURE WORK
The paper develops equivariant neural networks for physical dynamics and reports improved physical consistency and generalization under transformed test data. It also identifies combining full Navier–Stokes equivariances as an ongoing challenge.
- Conclusion: The authors develop equivariant neural networks that incorporate multiple symmetries for learning physical dynamics.The paper also states that these designs have theoretical and experimental support on 2D time-series prediction.
- Conclusion: The models demonstrate superior performance and improved physical consistency compared with non-equivariant counterparts.The reported comparison concerns both theoretical and experimental results on physical-dynamics forecasting.
- Conclusion: The proposed equivariant models generalize significantly better on transformed test data than non-equivariant models.
- Future work: All equivariant models can be combined and extended to three-dimensional cases.
- Future work: No single model currently provides equivariance to the full symmetry group of the Navier–Stokes equations.Combining different equivariances is described as possible but non-trivial and remains future work.
B.1 EQUIVARIANT NETWORKS AND DATA AUGMENTATION
For equivariant models, transformed training examples can be redundant because equivariance makes their learning signal equivalent up to loss weighting. Skip connections preserve equivariance when the layer mappings and activations satisfy the required representation structure.
- Data augmentation: Data augmentation offers no advantage for a perfectly equivariant parameterized function because transformed samples provide equivalent learning signals.For scale, the transformed sample can differ only by the symmetry-dependent sample weight.
- Data augmentation: For RMSE under rotation or uniform motion, original and transformed samples have equal gradients because the loss is invariant.
- Data augmentation: For scale transformations, the loss is equivariant and transformed samples are equivalent with weight χ(g).The paper specifies G = (R>0, ·) and χ(λ) = λ.
- Skip connections: ResNet and U-net remain G-equivariant when their layers are G-representations and convolutional mappings and activations are G-equivariant.This follows because skip connections act as identity mappings, which are equivariant for any G.
- Skip connections: This construction supports equivariant convolutional networks for rotational and scaling transformations, which act linearly.
B.3 RESULTS ON UNIFORM MOTION EQUIVARIANCE
Uniform-motion equivariance sharply constrains pointwise CNN architectures: individual convolution-activation layers must be affine, while residual blocks can retain equivariance through invariant residual mappings.
- Uniform motion: A convolutional layer equivariant to arbitrary uniform additions must satisfy a kernel constraint on constant tensors.The supplied proof passages show that the kernel response to a constant field is determined by the kernel sum.
- Uniform motion: Uniform-motion-equivariant pointwise activation functions must be translations.The paper derives σ(x) = a + x from the condition σ(x + c) = σ(x) + c.
- Uniform motion: A convolution followed by a piecewise differentiable activation is uniform-motion equivariant only when the composition is affine.Its form is ϕ(X) = K′ ∗X + b.
- Uniform motion: A CNN alternating convolutions and pointwise activations is affine if every combined layer is uniform-motion equivariant.This follows from composing affine functions.
- Limitation: More descriptive networks may be possible with activation functions spanning multiple channels, because the treatment considers only pointwise activations.
- Residual blocks: A residual block f(x) + x is uniform-motion equivariant when its residual connection f is uniform-motion invariant.This design avoids requiring the residual mapping itself to be equivariant.
B.4 RESULTS ON SCALE EQUIVARIANCE
Scale-equivariant CNNs are constrained by standard convolutional assumptions, motivating group-correlation designs and error analyses for transformed predictions. The paper also uses the turbulence energy spectrum to assess whether predictions preserve physical statistics.
- Scale equivariance: A scale-invariant CNN under the stated formulation is extremely limited.
- Scale equivariance: Under zero padding, a scale-equivariant convolutional kernel must be 1x1.
- Error analysis: The transformed test error is bounded using the untransformed test error and equivariance error.The paper defines equivariance error as EET(x) = |T(f(x)) − f(T(x))| and transformed test error as TTE = |T(ŷ) − f(T(x))|.
- Error analysis: For uniform motion, transformed test error satisfies TTE ≤ EE + TE.
- Physical symmetries: The Navier–Stokes equations have five listed symmetry transformations that collectively form a 7-dimensional symmetry group.The transformations include space and time translation, uniform motion, reflection/rotation, and scaling.
- Turbulence spectrum: Energy Spectrum Error indicates whether predictions preserve the correct statistical distribution and obey energy conservation.It can reveal smoothing of small or medium eddies even when accumulated RMSE is not high.
C HEAT DIFFUSION
The heat-diffusion experiments simulate temperature fields under varied initial conditions and diffusivities, then evaluate prediction accuracy and thermal-energy consistency under transformed tests.
- The 2D heat equation models temperature diffusion and is relevant to heat transfer, Brownian motion, and particle diffusion.
- The authors generated 6k scalar temperature fields using finite differences across varied initial conditions and thermal diffusivities.Figure 6 illustrates diffusion from a hotter interior region with thermal diffusivity 4.
- Thermal-energy consistency is evaluated with L1 loss because energy conservation requires each field’s temperature sum to remain consistent over time.
- Equ-CNNs consistently outperform CNNs in prediction RMSE and thermal-energy L1 loss across three transformed test sets.
D.1 DATASETS DESCRIPTION
The experiments cover simulated turbulent convection and real-world ocean currents, using transformed tests and matched convolutional architectures to study equivariant dynamics prediction.
- Rayleigh-Bénard convection: The Rayleigh-Bénard convection dataset contains 1792×256 velocity-field images from 2D turbulent flow simulated with the Lattice Boltzmann Method.Each image is divided into seven 256×256 regions and downsampled to 64×64 pixels; the Rayleigh number is 2.5 × 10^8.
- Transformed tests: Transformed test sets apply uniform-motion shifts, magnitude changes, rotations, and scale changes to evaluate symmetry-related distribution shifts.
- Ocean currents: Ocean-current experiments use 64×64 subregions from Atlantic, Indian, and North Pacific Ocean reanalysis data generated by NEMO from 01/01/2016 to 08/18/2017.
- Architectures: The study evaluates equivariant modifications of 18-layer ResNet and 13-layer U-net architectures for dynamics prediction.
- Prediction setup: Models use the past l velocity-field frames as input and predict one velocity field, with l tuned to approximately 25 and longer forecasts generated autoregressively.
- Training setup: Training uses MSE over the next k steps, a 60%-20%-20% temporal split, and approximately matched parameter counts across models.The best-performing k values were 3 or 4, and validation data was used for hyperparameter tuning.
E ADDITIONAL RESULTS
Additional results report temperature RMSEs and visualize predicted velocity fields for transformed samples and future-time ocean-current forecasts across regular and equivariant architectures.
- Temperature predictions: Table 8 reports test-data RMSEs for temperature predictions, with equivariant-model cells listing ResNet and U-net values separately.
- Transformed test samples: Figure 8 compares ground truth, Equ-Unet predictions, and U-net predictions at time steps 1, 5, and 10 under four transformed test conditions.The transformations are uniform-motion shift, magnitude scaling, 90-degree rotation, and 3× upscaling.
- Future-time ocean currents: Figure 9 compares ground-truth ocean-current velocity norms with predictions from regular ResNet and four Equ-ResNets at future time steps 5 and 10.