Latent Diffusion Transformer with Local Neural Field as PDE Surrogate Model

Authors: louis  serrano
jean-noël  vittaut
patrick  gallinari
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-09 16:18:40

Abstract: We introduce a diffusion transformer architecture, AROMA (Attentive Reduced Order Model with Attention), for dynamics modeling of complex systems. By employing a discretization-free encoder and a local neural field decoder, we construct a latent space that accurately captures spatiality without requiring traditional space discretization. The diffusion transformer models the dynamics in this latent space conditioned on the previous state. It refines the predictions providing enhanced stability compared to traditional transformers and then enabling longer rollouts. AROMA demonstrates superior performance over existing neural field methods in simulating 1D  and 2D equations, highlighting the effectiveness of our approach in capturing complex dynamical behaviors.

Main paper: 
 introduction: 
 In recent years, many deep learning (DL) surrogate models have been introduced to approximate solutions to partial differential equations (PDEs) (Lu et al., 2021; Li et al., 2021; Brandstetter et al., 2022; Stachenfeld et al., 2022) . Amongst them, the family of neural operators and operator networks has been extensively adopted and tested in various scientific domains, showing the potential of datacentric DL models in science (Pathak et al., 2022; Vinuesa & Brunton, 2022) .In essence, neural operators adopt the functional formulation of a PDE problem, yet they are implemented with architectures that exhibit constraints with respect to the discretization of the data and the geometry of the domain. To alleviate these issues, neural-field based methods have been proposed to model the dynamics of a system that follows PDEs (Yin et al., 2022; Chen et al., 2022) or more generally to learn operators between functions (Serrano et al., 2023) . These models are discretization free, they can be applied on regular and irregular grids, with different samplings. However, they still face some challenges in applicability: (i) they work with auto-decoding which can be unstable during training and computationally intensive at inference, (ii) they unroll the dynamics in a latent-space over a trajectory of global codes that ignore the spatial information of the data.In this paper, we propose a discretization-free model that offers a solution to both of these issues by taking a different approach. Our hypothesis is that considering spatiality is essential to modeling spatiotemporal phenomena, and therefore the latent space should conserve a spatial structure. For instance, state-of-the-art latent U-Net (Rombach et al., 2022) and transformer (Peebles & Xie, 2023) diffusion models operate respectively on latent images or sequences of smaller dimensions that keep a spatial property. However, in the literature these latent representations are mostly obtained with traditional convolutional neural networks (CNNs) (He et al., 2016) which are not compatible with our discretization-free constraint. For the same reason, patching observations Lee et al. (2023) is not aligned with our objective. Therefore, there is no existing architecture capable of encoding without discretization constraints function values into a latent representation that conserves a spatial dimension.• A novel architecture for latent-space modeling of dynamical systems, combining local neural fields with a diffusion transformer for enhanced stability.• A discretization-free encoder that captures spatially-aware tokens from data, bypassing traditional patching, and a decoder that uses cross-attention with local neural fields for accurate, continuous data reconstruction.• An innovative method for modeling dynamics in latent space, utilizing a diffusion transformer for iterative refinement.• Detailed validation of our model's superior performance over existing frameworks, evidenced by results on two datasets and supported by a perturbation analysis showcasing the spatial significance of our latent tokens. 
problem setting: 
 In this paper, we focus on time-dependent PDEs defined over a spatial domain Ω (with boundary ∂Ω) and temporal domain [0, T ]. In the general form, their solutions u(x, t) satisfy the following constraints :EQUATIONEQUATIONEQUATIONWhere ν represents a set of PDE coefficients, Equations ( 2) and (3) represent the constraints with respect to the boundary and initial conditions. We aim to learn from solutions data obtained with classical solvers the evolution operator G that predicts the state of the system at the next time step: u t+∆t = G(u t ). We have access to training trajectories obtained with different initial conditions, and we want to generate accurate trajectory rollouts for new initial conditions at test time. A rollout is obtained by the iterative application of the evolution operator u m∆t = G m (u 0 ). 
model description: 
 Our model comprises key components to enhance the latent-space modeling of dynamical systems.• Encoder : u t X → Z t . The encoder takes input values over a grid X and reduces the dimension to a set of M multivariate normal statistics E w (u t X ) = (µ t j , σ t j ) M j=1 from which latent tokens can be sampled z t j ∼ N (µ t j , (σ t j ) 2 ). We note the mean, standard deviation and latent tokensµ t = (µ t 1 , • • • , µ t M ), σ t = (σ t 1 , • • • , σ t M ), Z t = (z t 1 , • • • , z t M ). Each token is d-dimensional. • Latent refiner : Z t → Ẑt+∆t .We model the dynamics in the latent space through a latent conditional diffusion model r θ . The diffusion model r θ is used to predict the next set of latent tokens by denoising them in K steps. Starting from pure gaussian noise Zt+∆t K ∼ N (0, I), we reverse the diffusion process:Zt+∆t k-1 = DENOISE( Zt+∆t k , r θ (Z t , Zt+∆t k , k))for 1 ≤ k ≤ K and set the prediction as Ẑt+∆t = Zt+∆t 0 . We note R θ this denoising operation. The diffusion model enables a "refinement process" throughout the denoising, improving the prediction at each denoising step. This is inspired by Lippe et al. (2023) , as we experimentally observed that a probabilistic model was more robust than a deterministic one for temporal extrapolation.• Decoder : Ẑt+∆t → ût+∆t . The decoder uses the latent tokens Ẑt+∆t to approximate the function value ût+∆t (x) = D ψ (x, Ẑt+∆t ) for each query coordinate x ∈ Ω. We therefore note ût+∆t = D ψ (Z t+∆t ) the predicted function, and ût+∆t X = D ψ (Z t+∆t , X ) the predicted values over the grid X . Inference We encode the initial condition and unroll the dynamics in the latent space by successive denoisings:ûm∆t = D ψ • R m θ • E w (u 0 ).We then decode along the trajectory to get the reconstructions. We outline the full inference pipeline in Figure 1 Training We do a two-stage training: first train the encoder and decoder, second train the refiner.Encoder The encoder architecture efficiently processes input functions irrespective of their discretization (Figure 5 ). It shares some similarities with the Perceiver architecture (Jaegle et al., 2021) but is uniquely designed for our setting. (i) context embedding creation: Given an input sequence of coordinate-value pairs (x 1 , u(x 1 )), . . . , (x N , u(x N )), we construct high-dimensional context embeddings. Each embedding c l is obtained as c l = f u (u(x l )) + f x (x l ) , where f u and f x are feedforward neural networks. (ii) information aggregation: The sequence of context embeddings c = (c 1 , . . . , c N ) is aggregated onto a smaller set of learnable query tokens T = (T 1 , . . . , T M ). This is achieved using a multihead cross-attention layer and a feedforward (FF) network, expressed asT ′ = FF(CrossAttention(Q = W Q T, K = W K c, V = W V c)). (iii) channel dimension com- pression:Finally, the information in the channel dimension of T ′ is compressed using a bottleneck linear layer, yielding the mean and standard deviation tokens µ = W µ T ′ and σ = W σ T ′ .Decoder The decoder architecture proceeds as follows (Figure 5 ). (i) latent tokens processing:The decoder first lifts the latent tokens Z -sampled from a multivariate gaussian with parameters µ and σ -to a higher channel dimension. (ii) self-attention layers: Several layers of self-attention are then applied to get processed tokens Z ′ . (iii) cross-attention: The decoder applies cross-attention to obtain a modulation vector that depends on the query coordinate ϕu (x) = CrossAttention(Q = W Q γ q (x), K = W K Z ′ , V = W V Z ′ ), where γ q is a feed forward neural network with fourier features of bandwidth ω q . (iv) local implicit neural representation (INR) decoding: Finally, an INR is used to decode the local modulation vector to obtain the output of the neural field. As in Lee et al. (2023) this INR employs a two-level decoding with two frequency bandwidths ω 1 , ω 2 such that ω q ≤ ω 2 ≤ ω 1 . We refer the reader to additional details regarding its implementation in Appendix B. In contrast with existing neural field methods for dynamics modeling, the modulation vector here is local, i.e. it does not modify the weights of the INR for the whole coordinate space.Training The encoder and decoder are jointly optimized as a variational autoencoder (VAE) (Kingma & Welling, 2013) to minimize the following objective :L = L recon + β • L KL ; where L recon = MSE(u t X , ût X )is the reconstruction loss between the input and D ψ (Z t , X ), with Z t ∼ N (µ t , (σ t ) 2 ) and µ t , σ t = E w (u t X ). The KL loss L KL = D KL (N (µ t , (σ t ) 2 ) || N (0, I)) helps regularizing the network and prevents overfitting.As the latent tokens Z t have a sequential structure, we use a conditional diffusion transformer architecture close to Peebles & Xie (2023) for r θ . At diffusion step k, the input to the network is a sequence stacking the tokens and the current noisy target estimate (Z t , Zt+∆t k ). See Appendix B, Figure 3 and Figure 2 for more details. To train the refiner r θ , we freeze the encoder and decoder, and use the encoder to sample pairs of successive latent tokens (Z t , Z t+∆t ). We employ the "v-predict" formulation of DDPM (Salimans & Ho ( 2022 Baselines We compare our model to FNO (Li et al., 2021) , and choose DINo (Yin et al., 2022) and CORAL (Serrano et al., 2023) as the neural field baselines. We also include a deterministic AROMA without diffusion as an ablaton study.During training, we use only the data fromn the training horizon (In-t). At test time, we evaluate the models to unroll the dynamics for new initial conditions in the training horizon (In-t) and for temporal extrapolation (Out-t).Results Table 1 showcases the performance of our model against various models on both datasets. The results are evaluated using the mean squared error (MSE) metric. Notably, AROMA demonstrates superior performance for Navier-Stokes dataset, achieving the lowest MSE scores in both In-t and Out-t scenarios. It also shows competitive results on Burgers, outperforming other INR-based models and closely following FNO. AROMA without diffusion obtains good results in the training horizon but is particularly less robust on longer rollouts (Figure 7 , Figure 8 , Figure 9 ). .68e-4 8.95e-3 3.76e-4 1.08e-3 DINo 1.27e-3 1.11e-2 5.82e-2 6.50e-2 CORAL 1.86e-4 1.02e-3 8.24e-4 2.56e-3 AROMA 6.95e-5 1.02e-3 3.96e-4 1.29e-3 AROMA no refiner 8.21e-5 1.70e-3 5.79e-4 2.04e-3 
conclusion: 
 AROMA offers a novel approach to neural field modeling, incorporating attention mechanisms and a latent diffusion transformer for spatio-temporal dynamics. It has shown improved performance on small-size datasets and holds potential for effective scaling to larger datasets, thanks to its use of transformer blocks. A RELATED WORKOur model differs from existing models in the field of operator learning and more broadly from existing neural field architectures. The works most related to ours are the following.Transformer for PDE Hao et al. (2023) introduced a versatile transformer architecture specifically tailored for operator learning. Their design incorporates an attention mechanism and employs a mixture of experts strategy, effective in addressing multi-scale challenges. However, their architecture resorts to linear attention as they do not reduce spatial dimensions, and each attention layer has a linear complexity in the sequence size, which can be prohibitive for very deep networks. 
 
 Appendix: We are not the first work that proposes to leverage locality to improve the design of neural fields. In a different approach, Bauer et al. (2023) proposed a grid-based latent space where the modulation function ϕ is dependent on the query coordinate x. This concept enables the application of architectures with spatial inductive biases for generation on the latent representations, such as a U-Net Denoiser for diffusion processes. Similarly, Lee et al. (2023) 2021). Their approach also leverages cross-attention instead of patching but their architecture differs from ours as it is designed for processing point-cloud shapes while our blocks can be used with arbitrary functions.We illustrate how our diffusion transformer is trained and used at inference in Figure 2 and Figure 3 . We provide the diffusion step k which acts as a conditioning input for the diffusion model. We use an exponential decrease for the noise level as in Lippe et al. (2023) i.e.α k = 1 -σ k/K min .Encoder-Decoder We provide a more detailed description of the encoder-decoder pipeline in 1D Burgers' Equation (Burgers): This equation models shock waves without diffusion, expressed as ∂u ∂t + 0.5 ∂u 2 ∂x = f . We take the same dataset as in Brandstetter et al. (2022) to guarantee periodicity of the initial conditions and forcing term f . Our training and test sets consist of 2048 and 128 trajectories, respectively, reduced to temporal and spatial resolutions of (250, 100). For training, we create 10 sub-trajectories of size 25, from which we keep only the first 20 timestamps. For evaluation we use the 25 timestamps. training and evaluate over 40 timestamps. The original data is at a 256 × 256 resolution, which we sub-sample to achieve a frame size of 64 × 64.On Burgers, we test the capacity of the models to unroll the dynamics from the 50th timestamp for 200 steps up to timestamp 250. Note that none of the models have been trained for this task, and that the models have not seen 20% of the data within this horizon. As showcased in Table 1 the performance on this dataset quickly deteriorates in extrapolation as the intensity and location of the shocks are hard to predict, and therefore the performance of all models is relatively poor. Still, AROMA performs the best with a test MSE over the whole trajectory of 0.119 compared to 0.142 for FNO. Most notably this analysis helps in motivating the modeling choice of our diffusion transformer. While the predictions with FNO and AROMA are relatively close to each other in Figure 6 and Figure 7 , AROMA without diffusion quickly diverges from the ground truth (Figure 8 ) highlighting the increased stability with the refinement process. We show in Figure 9 that AROMA's and FNO's predictions are still highly correlated at the end of the trajectory while AROMA without diffusion cannot cope with long rollouts.For Navier-Stokes, we show an example of test trajectory in the training horizon (Figure 10 ) and in extrapolation (Figure 11 ). To validate the spatial interpretation of our latent tokens, we establish a baseline code Z 0 , and introduce perturbations by sequentially replacing the j-th token, z 0 j , with subsequent tokens along the trajectory, denoted as z 1 j , z 2 j , . . . , z t j . Thus, the perturbed tokens mirrors Z 0 in all aspects except for the j-th token, which evolves according to the true token dynamics. We show reconstruction visualizations of the perturbed tokens in Figure 12 , Figure 13 , Figure 14 , Figure 15 , Figure 16 , Figure 17 , Figure 18 , Figure 19 . As evidenced in these figures, the perturbation of a token changes only locally the reconstructed field, which validates the spatial structure of our tokens. Notice in Figure 14 and in Figure 19 how these token affect the reconstruction in regions close to the boundary conditions. Our encoder-decoder has discovered from data and without explicit supervision that the solutions had periodic boundary conditions.For Navier-Stokes, we show how the mean (Figure 20 ) and standard deviation tokens (Figure 21 ) evolve over time for a given test trajectory. We show the predicted trajectory in the latent space in Figure 22 .  


Decision: This paper proposes to combine diffusion model and a cross-attention layer for dynamics modeling of complex systems. The concerns raised by the reviewer 7rSp are valid. I strongly command the author to include the training time of their model vs the baseline to present a more comprehensive picture.

Reviews: review: title: title: Advanced architecture for neural PDEs
main_review: main_review: The authors suggest using a combination of diffusion model and a cross-attention layer to learn latent variables in a neural PDE formulation. Despite a more complicated architecture that must have resulted in more GPU-hours than the baselines, the numerical results were not noticeably different. The authors should provide a comparison on training time of their model vs the baseline and further numerical justification for their idea.
title: title: official review
main_review: main_review: The premise of the paper is very interesting-- leveraging discretization-free INR representation without losing the benefits of a (somewhat) local representation. 
Method-wise there's a lot going on in this paper-- a latent-space transformer, a diffusion model for next-step prediction and an INR decoder. This is good for a workshop, there's a lot of food for thought and things to discuss. 

There isn't much validation of the papers' ideas so far, experiments are basically one comparison study. So currently we don't learn too much from the method choices-- I would definitely encourage the authors to keep exploring this space, drilling down on the individual components, and running ablations. E.g. what does setting this up as a diffusion model buy? Or maybe you get really good results for a diffusion transformer with a plain grid-based VAE decoder.
Focussing on one key element that works well, and really studying it could make an excellent full conference paper.