Incorporating Symmetry into Deep Dynamics Models for Improved Generalization

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 $\circ\colon G\times G\to G$. The composition map is required to be associative and have an identity $1\in 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 $\cdot\colon G\times S\to 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.

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 $\bm{w}(\bm{x},t)$ be the velocity vector field of a flow. The field $\bm{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,

where $H(\bm{x},t)$ is temperature, $p$ is pressure, $\kappa$ is the heat conductivity, $\rho_{0}$ is initial density, $\alpha$ is the coefficient of thermal expansion, $\nu$ is the kinematic viscosity, and $f$ is the buoyant force.

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=\mathbb{R}^{2}$ and 3-dimensional spacetime $\hat{X}=\mathbb{R}^{3}$. Let $V=\mathbb{R}^{d}$ be a $G$-representation. Denote the set of all $V$-fields on $\hat{X}$ as $\hat{\mathcal{F}}_{V}=\{\bm{w}\colon\hat{X}\to V:\bm{w}\text{ smooth}\}.$ Define $\mathcal{F}_{V}$ similarly to be $V$-fields on $X$. Then $G$ has an induced action on $\hat{\mathcal{F}}_{V}$ by $(g\bm{w})(x,t)=g(\bm{w}(g^{-1}x,g^{-1}t))$ and on $\mathcal{F}_{V}$ analogously.

Consider a system of differential operators $\mathcal{D}$ acting on $\hat{\mathcal{F}}_{V}$. Denote the set of solutions $\mathrm{Sol}(\mathcal{D})\subseteq\hat{\mathcal{F}}_{V}.$ We say $G$ is a symmetry group of $\mathcal{D}$ if $G$ preserves $\mathrm{Sol}(\mathcal{D})$. That is, if $\varphi$ is a solution of $\mathcal{D}$, then for all $g\in G$, $g(\varphi)$ is also. In order to forecast the evolution of a system $\mathcal{D}$, we model the forward prediction function $f$. Let $\bm{w}\in\mathrm{Sol}(\mathcal{D})$. The input to $f$ is a collection of $k$ snapshots at times $t-k,\ldots,t-1$ denoted $\bm{w}_{t-i}\in\mathcal{F}_{d}$. The prediction function $f\colon\mathcal{F}_{d}^{k}\to\mathcal{F}_{d}$ is defined $f(\bm{w}_{t-k},\ldots,\bm{w}_{t-1})=\bm{w}_{t}$. It predicts the solution at a time $t$ based on the solution in the past. Let $G$ be a symmetry group of $\mathcal{D}$. Then for $g\in G$, $g(\bm{w})$ is also a solution of $\mathcal{D}$. Thus $f(g\bm{w}_{t-k},\ldots,g\bm{w}_{t-1})=g\bm{w}_{t}$. Consequently, $f$ is $G$-equivariant.

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_{\bm{c}}^{\mathrm{sp}}\bm{w}(\bm{x},t)=\bm{w}(\bm{x-c},t)$,   $\bm{c}\in\mathbb{R}^{2}$,

• •

  Time translation:  $T_{\tau}^{\mathrm{time}}\bm{w}(\bm{x},t)=\bm{w}(\bm{x},t-\tau)$,   $\tau\in\mathbb{R}$,

• •

  Uniform motion:  $T_{\bm{c}}^{\mathrm{um}}\bm{w}(\bm{x},t)=\bm{w}(\bm{x},t)+\bm{c}$,   $\bm{c}\in\mathbb{R}^{2}$,

• •

  Rotation/Reflection:   $T_{R}^{\mathrm{rot}}\bm{w}(\bm{x},t)=R\bm{w}(R^{-1}\bm{x},t),\;R\in O(2)$,

• •

  Scaling:   $T_{\lambda}^{sc}\bm{w}(\bm{x},t)=\lambda\bm{w}(\lambda\bm{x},\lambda^{2}t)$,   $\lambda\in\mathbb{R}_{>0}$.

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 $\mathcal{F}_{\mathbb{R}^{d_{\mathrm{in}}}}\to\mathcal{F}_{\mathbb{R}^{d_{\mathrm{out}}}}$ with kernel $K$ from a $\mathbb{R}^{d_{\mathrm{in}}}$-field to a $\mathbb{R}^{d_{\mathrm{out}}}$-field. Let $\mathbb{R}^{d_{\mathrm{in}}}$ and $\mathbb{R}^{d_{\mathrm{out}}}$ be $G$-representations with action maps $\rho_{\mathrm{in}}$ and $\rho_{\mathrm{out}}$ respectively. Cohen et al. [11, Theorem 3.3] prove the network is $G$-equivariant if and only if

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.

CNNs are time translation-equivariant as long as we predict in an autoregressive manner. Convolutional layers are also naturally space translation-equivariant (if cropping is ignored). Any activation function which acts identically pixel-by-pixel is equivariant.

To incorporate rotational symmetry, we model $f$ using $\mathrm{SO}(2)$-equivariant convolutions and activations within the E(2)-CNN framework of Weiler and Cesa [49]. In practice, we use the cyclic group $G=C_{n}$ instead of $G=\mathrm{SO}(2)$ as for large enough $n$ the difference is practically indistinguishable due to space discretization. We use powers of the regular representation $\rho=\mathbb{R}[C_{n}]^{m}$ for hidden layers. The representation $\mathbb{R}[C_{n}]$ has basis given by elements of $C_{n}$ and $C_{n}$-action by permutation matrices. It has good descriptivity since it contains all irreducible representations of $C_{n}$, and it is compatible with any activation function applied channel-wise.

By the proposition 3 above that is proved in Appendix B.3, within ResNet, residual mappings should be invariant, not equivariant, to uniform motion. That is, the skip connection $f^{(i,i+2)}=I$ is equivariant and the residual function $f^{(i,i+1)}$ should be invariant. Hence, for the first layer in each residual block, we omit adding the mean back to the output $\bm{\mathcal{Q}_{ij}}$. In the case of Unet, when upscaling, we pad with the mean to preserve the overall mean.

Scale equivariance in dynamics is unique as the physical law dictates the scaling of magnitude, space and time simultaneously. This is very different from scaling in images regarding resolutions [51]. For example, the Navier-Stokes equations are preserved under a specific scaling ratio of time, space, and velocity given by the transformation

where $\lambda\in\mathbb{R}_{>0}$. We implement two different approaches for scale equivariance, depending on whether we tie the physical scale with the resolution of the data.

Resolution Independent Scaling. We fix the resolution and scale the magnitude of the input by varying the discretization step size. An input $\bm{w}\in\mathcal{F}_{\mathbb{R}^{2}}^{k}$ with step size $\Delta_{x}(\bm{w})$ and $\Delta_{t}(\bm{w})$ can be scaled $\bm{w}^{\prime}=T_{\lambda}^{sc}(\bm{w})=\lambda\bm{w}$ by scaling the magnitude of vector alone, provided the discretization constants are now assumed to be $\Delta_{x}(\bm{w}^{\prime})=1/\lambda\Delta_{x}(\bm{w})$ and $\Delta_{t}(\bm{w}^{\prime})=1/\lambda^{2}\Delta_{t}(\bm{w})$. We refer to this as magnitude equvariance hereafter.

To obtain magnitude equivariance, we divide the input tensor by the MinMax scaler (the maximum of the tensor minus the minimum) and scale the output back after convolution and activation per sliding block. We found that the standard deviation and mean L2 norm may work as well but are not as stable as the MinMax scaler. Specifically, using the same notation as in Section 3.4,

Resolution Dependent Scaling. If the physical scale of the data is fixed, then scaling corresponds to a change in resolution and time step size. To achieve this, we replace the convolution layers with group correlation layers over the group $G=(\mathbb{R}_{>0},\cdot)\ltimes(\mathbb{R}^{2},+)$ of scaling and translations. In convolution, we translate a kernel $K$ across an input $\bm{w}$ as such $\bm{v}(\bm{p})=\sum_{\bm{q}\in\mathbb{Z}^{2}}\bm{w}(\bm{p}+\bm{q})K(\bm{q}).$ The $G$-correlation upgrades this operation by both translating and scaling the kernel relative to the input,

where $s$ and $t$ denote the indices of output and input channels respectively. We add an axis to the tensors corresponding the scale factor $\mu$. Note that we treat the channel as a time dimension both with respective to our input and scaling action. As a consequence, as the number of channels increases in the lower layers of Unet and ResNet, the temporal resolution increases, which is analogous to temporal refinement in numerical methods [24, 31]. For the input $\tilde{\bm{w}}$ of first layer where $\tilde{\bm{w}}$ has no levels originally, $\bm{w}(p,s,\lambda)=\lambda\tilde{\bm{w}}(\lambda p,\lambda^{2}s)$.

Our model builds on the methods of Worrall and Welling [51], but with important adaptations for the physical domain. Our implementation of group correlation equation 5 directly incorporates the physical scaling law equation 3 of the system equation $\mathcal{D}_{\mathrm{NS}}$. This affects time, space, and magnitude. (For heat, we drop the magnitude scaling.) The physical scaling law dictates our model should be equivariant to both up and down scaling and by any $\lambda\in\mathbb{R}_{>0}$. Practically, the sum is truncated to 7 different $1/3\leq\lambda\leq 3$ and discrete data is continuously indexed using interpolation. Note equation 3 demands we scale anisotropically, i.e. differently across time and space.

The equivariance errors can be defined as $\mathrm{EE}_{T}(x)=|T(f(x))-f(T(x))|$, where $x$ is an input, $f$ is a neural net, $T$ is a transformation from a symmetry group. We empirically measure the equivariance errors of all equivariant models we have designed. Table 1 shows the equivariance errors of ResNet and Equ-ResNet. The transformation $T$ is sampled in the same way as we generated the transformed Rayleigh-Bénard Convection test sets. See more details in Appendix B.5.

Data Description. Rayleigh-Bénard Convection occurs in a horizontal layer of fluid heated from below and is a major feature of the El Niño dynamics. The dataset comes from two-dimensional turbulent flow simulated using the Lattice Boltzmann Method [8] with Rayleigh number $2.5\times 10^{8}$. We divide each 1792 $\times$ 256 image into 7 square subregions of size 256 $\times$ 256, then downsample to 64 $\times$ 64 pixels. To test the models’ generalization ability, we generate additional four test sets : 1) UM: added random vectors drawn from $U(-1,1)$; 2) Mag: multiplied by random values sampled from $U(0,2)$; 3) Rot: randomly rotated by the multiples of $\pi/2$; 4) Scale: scaled by $\lambda$ sampled from $U(1/5,2)$. Due to lack of a fixed reference frame, real-world data would be transformed relative to training data. We use transformed data to mimic this scenario.

Prediction Performance. Table 2 shows the prediction RMSE and ESE on the original and four transformed test sets by the non-equivariant ResNet(Unet) and four Equ-ResNets(Unets). Augm is ResNet(Unet) trained on the augmented training set with additional samples with random transformations applied from the relevant symmetry group. The augmented training set contains additional transformed samples and is three times the size of the original training set. Each column contains the prediction errors by the non-equivariant and equivariant models on each test set. On the original test set, all models have similar RMSE, yet the equivariant models have lower ESE. This demonstrates that incorporating symmetries preserves the representation powers of CNNs and even improves models’ physical consistency.

On the transformed test sets, we can see that ResNet(Unet) fails, while Equ-ResNets(Unets) performs even much better than Augm-ResNets(Unets). This demonstrates the value of equivariant models over data augmentation for improving generalization. Figure 2 shows the ground truth and the predicted velocity fields at time step $1$, $5$ and $10$ by the ResNet and four Equ-ResNets on the four transformed test samples.