Equivariant Neural Fields For Symmetry Preserving Continous PDE Forecasting

Authors: david m knigge
david r wessels
riccardo  valperga
samuele  papa
efstratios  gavves
erik j bekkers
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-09 19:03:27

Abstract: Recently, Neural Fields (NeFs) have emerged as a powerful modelling paradigm to represent discretely-sampled continuous signals. As such, novel work has explored the use of Conditional NeFs to model PDEs,  by learning continuous flows in the latent space of the Conditional NeF. Although this approach benefits from favourable properties of neural fields such as grid-agnosticity and space-time-continuous dynamics modelling, it does not make use of important geometric information about the domain of the PDE being modelled -- such as information on symmetries of the PDE -- in favour of modelling flexibility.
Instead, we propose a NeF parameterization that preserves geometric information in the latent space of the Conditional NeF: \textit{Equivariant Neural Fields}. Using this representation, we construct a framework for space-time continuous PDE modelling that preserves known symmetries of the PDE. We experimentally validate our model and show it readily generalizes to arbitrary locations, as well as geometric transformations of the initial conditions - where other NeF-based PDE forecasting methods fail.

Main paper: 
 introduction: 
 Partial Differential Equations (PDEs) are a foundational tool in modelling and understanding spatio-temporal dynamics across diverse scientific domains. Classically, PDEs are solved using numerical methods like finite elements, finite volumes, or spectral methods. In recent years, Deep Learning (DL) methods have emerged as promising alternatives due to abundance of observed and simulated data as well as the accessibility to computational resources, with applications ranging from fluid simulations and weather modelling (Yin et al., 2022; Brandstetter et al., 2022a) to biology Moser et al. (2023) . The systems modelled by PDEs often have underlying symmetries. For example, heat diffusion or fluid dynamics can be modeled with differential operators which are rotation equivariant, i.e. given a solution to the system of PDEs, its rotation is also a valid solution.In such scenarios it is sensible, and even desirable, to design neural networks that incorporate and preserve such symmetries by design to improve generalization and data-efficiency (Cohen & Welling, 2016; Weiler & Cesa, 2019; Bekkers, 2019) .Crucially, DL-based approaches rely on data sampled on a grid, without the inherent ability to generalize outside of it, which is restrictive in many scenarios (Prasthofer et al., 2022) . To this end, Yin et al. (2022) propose to use Neural Fields (NeFs) for modelling and forecasting PDE dynamics. However, they fail to leverage the well-known equivariances present in the data.We introduce a new class of SE(n)-equivariant conditional neural fields. We condition the NeF through a point cloud of latent feature vectors defined over the symmetry group of interest, disentangling pose and appearance in the latent space of functions. We subsequently propose a framework for solving PDEs by learning traversals of the latent space of this equivariant NeF with an equivariant Neural ODE, see Fig. 1 for an illustration.  
background and related work: 
 We describe Deep Learning based approaches to data-driven dynamics modelling, and specifically highlight the recent use of neural fields (NeFs) in this setting. We highlight the concept of symmetry conservation in general DL model design and argue for the inclusion of such inductive biases in DL-based PDE surrogates, which we achieve by introducing equivariant conditional NeFs.DL approaches to dynamics modelling Most DL methods for solving PDEs attempt to directly replace solvers with mappings between finite-dimensional Euclidean spaces, i.e. through the use of CNNs (Guo et al., 2016; Ayed et al., 2020) or GNNs (Pfaff et al., 2020; Brandstetter et al., 2022b) often applied autoregressively to an observed (discretized) PDE state. Instead, the Neural Operator (NO) (Kovachki et al., 2021) paradigm attempts to learn infinite-dimensional operators, i.e. mappings between function spaces, with limited success. Fourier Neural Operator (FNO) (Li et al., 2020) extends this method by performing convolutions in the spectral domain. FNO obtains much improved performance, but due to its reliance on FFT is limited to data on regular grids. 2016), generalizations of CNNs that respect symmetries of the data -such as dilations and continuous rotations (Weiler & Cesa, 2019; Finzi et al., 2020; Knigge et al., 2022) . In the context of dynamics modelling, equivariant architectures have been employed to incorporate various properties of physical systems in the modelling process, examples of such properties are the symplectic structure (Jin et al., 2020) , discrete symmetries such reversing symmetries (Valperga et al., 2022) and energy conservation (Greydanus et al., 2019; Hernández et al., 2021) .Neural Fields in Dynamics Modelling Neural fields (NeFs) are a class of coordinate-based neural networks, often trained to reconstruct discretely-sampled input signals in a continuous way. More specifically, a neural field f θ : R n → R d is a field -parameterized by a neural network with parameters θthat maps n-dimensional input coordinates x ∈ R n in the data domain, to Dupont et al. (2022) showed the viability of using the latents z f as representations for downstream tasks (e.g. classification, generation) proposing a framework for learning on neural fields. This framework inherits desirable properties of neural fields, such as inherent support for sparsely and/or irregularly sampled data, and independence to signal resolution. Yin et al. (2022) propose to use conditional NeFs for PDE modelling by learning a continuous flow in the latent space of a conditional neural field. In particular, a set of latents {z ν i } T i=1 are obtained by fitting a conditional neural field to a given set of observations {ν i } T i=1 at timesteps 1, ..., T ; simultaneously, a neural ODE (Chen et al., 2018) is trained with those latents such that solutions correspond to the trajectories traced by the learned latents.d-dimensional signal values f (x) ∈ R d . By associating a conditioning latent z f ∈ R c to each signal f , a single conditional NeF f θ : R n × R c → R d can learn to represent families D of continuous signals such that ∀ f ∈ D : f (x) ≈ f θ (x; z f ).Figure 2 :The proposed Equivariant Neural Field (top) under a rototranslation g (bottom)., i.e., ∀ g∈SE(n) : L g ν t ⇔ l g Z ν t . 3. The neural ODE is equivariant, i.e., ∀ g∈SE(n) : F ψ (l g z ν t ) = l g F ψ (z ν t ).The next section describes how we solve the necessary conditions (1.) and (2.) by defining the latent space to a space of attributed point clouds in SE(n) and introducing a new class of SE(n)equivariant conditional neural fields. Condition (3.) will be satisfied by parametrizing the neural ODE F ψ with an equivariant graph neural network (Bekkers et al., 2023) .We define NeFs that are conditioned on attributed point clouds over SE(n). That is, as a conditioning variable we consider a set Z := {(p i , a i )} n i=1 of n tuples that consist of a pose p i ∈ SE(n) and appearance a i ∈ R c . In clearer terms, each field f will be associated with a latent set Z through the NeF such that f (x) ≈ f θ (x; Z f ). The latent space has a well defined action l g given by l g Z = {(g • p i , a i )} n i , with g • p i being simply the group operation of SE(n).We propose a new class of equivariant neural fields (Fig. 2 ) that posses the following steerability property ∀g ∈ SE(n) :EQUATIONThrough this property, a roto-translation of the field can be obtained by a roto-translation of the latent point cloud. We define the equivariant neural field through cross-attention between relative coordinates p -1 i x = R T i (x -x i ) and the appearance vectors a i viaEQUATIONin which q is a learnable relative coordinate embedding function (we use SIREN from Sitzmann et al. ( 2020)), and k, v are learnable linear transformations. The softmax is over the latent set, and d k represents the dimensionality of the embeddings. Note that q can be any type of function that is not necessarily equivariant since the quantity p -1 i x is bi-invariant under the group action via∀g ∈ SE(n) : (p i , x) → (g p i , g x) ⇔ p -1 i x → (g p i ) -1 g x = p -1 i g -1 g x = p -1 i x .(4) This also proves the steerability property (2) since by ( 4) the tuple g -1 x, p i is equivalent to x, g p i . 
method: 
 Mathematical preliminaries Given a group G with identity element e ∈ G, and a set X, the group action is a map L g : G × X → X such that L e (x) = x and L gh (x) = L g (L h (x)). In particular, we are interested in the Special Euclidean group SE(n) = T n ⋊ SO(n). SE(n) is the roto-translation group consisting of elements g = (t, R) with group operation g g ′ = (t, R) (t ′ , R ′ ) = (Rt ′ + t, RR ′ ); their action on function spaces is defined by L g f (x) = f (g -1 x) = f (R -1 (x -t)).EQUATIONwhere F ψ is a neural network that parametrizes the neural ODE, that is, it maps latent z ν τ ∈ R c to its temporal derivative:dz ν τ dτ =F ψ (z ν τ ).Here z ν t denotes the latent associated with ν t and is inferred during training. Note that compared to Yin et al. (2022) where two independent objectives are minimised simultaneously (one for z ν t and θ and one for ψ), we train our model with a single objective. Intuitively, this is to force the latents to not only reconstruct observations, but also to be arranged in such a way that they can be fitted by a continuous flow.Our objective is to leverage the inductive bias of equivariance to preserve SE(n) symmetries that a PDE may possess in our neural surrogate. This requires the following necessary conditions:1. The latent space Z is equiped with a well-defined group action l g . 
experiments: 
 We validate our model on two different PDEs. First, to experimentally validate the equivariance properties of our framework, we train our models on 2D diffusion dc dt = D∇ 2 c, where c is a scalar field, and D is the diffusivity. Specifically, we create a benchmark where for initial conditions we insert a pulse at a random location in the grid x = (x 1 , x 2 ) ∈ R 2 s.t. -1<x 1 <1, 0<x 2 <1 for the training data and -1<x 1 <1, -1<x 2 <0 for the test data (intuitively, train and test sets contain pulses under different group actions respectively). We also evaluate on 2D Navier Stokes (Stokes et al., 1851) corresponding to an incompressible fluid with dynamics dv dt = -u∇v + v∆µ + f, v = ∇ × u, ∇u = 0, where u is the velocity field, v the vorticity, and µ the viscosity. For f we choose a symmetric forcing term. For dataset details see Appx. A.1, for experimental details see Appx. A.2.The models are trained and evaluated with the MSE on intervals of T =20 steps split into t IN ∈ {1, ..., 10} and t OUT ∈ {10, ..., 20}. Only t IN is used for supervision with the loss described in Eq. 1. Following Yin et al. ( 2022), we additionally provide results for randomly subsampled initial states ν 1 , where we show very little degradation even when 95% of values are removed from the observed initial state of a trajectory. A visualization of a predicted train and test trajectory for both datasets may be found in Appx. A.3.On the diffusion dataset (Tab. 1), our model obtains good performance on both train and test sets. In this experiment, DINo fails to generalize to the test set due to its lack of symmetry-preservation. On Navier-Stokes with symmetric forcing (Tab. 2), we similarly obtain good results on both train and test sets, although performance degrades for timesteps t out . We conclude that symmetry-preservation improves space-time continuous PDE solving. 
conclusion: 
 In this paper we proposed a model which utilises pointclouds as conditioning terms for conditional neural fields. Using an equivariant neural ode parameterised with a SE(n)-equivariant GNN we showed that utilising these conditional local latent point clouds within PDE-modeling outperforms the baseline DINo in two different experiments.  
 
 Appendix:  


Decision: Dear Authors,

Thank you for submitting the draft.

Both reviewers agree that the presented work presents some interesting strengths. However, both reviewers do also raise some major points of concern, especially regarding the clarity of the presentation. It is expected that authors will be addressing comments by the reviewers in the final draft. 

regards

AC

Reviews: review: title: title: Equivariant latent fields
main_review: main_review: The authors propose using an equivariant point cloud representation for latent variables in a neural PDE formulation. The idea is tested on 2D diffusion and 2D vorticity equations with rotationally symmetric forcing. 

The numerical results show improvement based on the baseline, but the baseline DINo is never cited or defined in text. The authors do not provide details on the ground truth solver used or how the initial conditions for vorticity equation was generated. They do not provide value for diffusivity D and do not state what the time steps used in evaluation are, making it impossible to judge the results. Finally, there are 2 typos in the statement of 2d NS equation: it is \mu \Delta v and not \nabla u = 0 but \nabla \cdot u = 0. 

In addition to lack of detail, the experiments are very contrived: what happens if forcing f is not rotationally symmetric, as it is most practical applications?
title: title: Introduces SE(n) equivariant model to augment CNeFs with geometric properties such as roto-translation to achieve better generalization.
main_review: main_review: Existing neural fields architecture for PDE solving lack components that preserve symmetries in data. And this work introduces an SE(n) equivariant approach to augment CNeFs with geometric properties such as roto-translation.

**Strengths**:
- generalization to geometric transformations of ICs
- learning equivariant latent evolution
- Roto-Translation in the latent space

**Weaknesses**:
- Limited set of experiments.
- No separate validation set. Hence the model might have been tuned on the test set.

**Questions**:
- How was the 50% and 5% observed points selected?
- How to decide which PDEs possess SE(n) symmetries?
- How to choose the optimal number of latents?

**Suggestions for improvement**:
- The importance of steerability on high dimensional PDEs can be emphasized.