Vectorized Conditional Neural Fields: A Framework for Solving Time-dependent Parametric Partial Differential Equations

The simulation of physical systems such as weather forecasting and fluid dynamics relies on solving Partial Differential Equations (PDEs). Nowadays, Machine Learning (ML) is increasingly used for solving PDEs due to several advantages compared to classical numerical solvers. For instance, ML surrogate models enable faster simulation times than classical numerical PDE solvers, are differentiable, and they are able to be used even when the underlying PDEs are not known exactly. However, if knowledge about the PDEs is available, it can be added to the model as in Physics-Informed Neural Networks (PINNs). The following figure shows a neural surrogate for solving PDEs that takes the initial value (solution for time \(t = 0\)) as input and predicts the the solutions for future times. Intuitively, the problem of solving PDEs can be interpreted as an image-to-video prediction task.

Despite recent advances in neural architectures for PDE solving, current methods lack several of the characteristics of an ideal PDE solver: (i) generalization to different Initial Conditions (ICs), (ii) PDE parameters, (iii) support for 1D, 2D, and 3D PDEs, (iv) stability over long rollouts, (v) temporal extrapolation, (vi) spatial and temporal super- resolution capabilities, all with affordable cost, high speed, and accuracy. Towards developing a model that encompasses these ideal characteristics, we propose Vectorized Conditional Neural Field (VCNeF), a linear transformer-based conditional neural field that solves PDEs continuously in time, endowing the model with temporal as well as spatial Zero-Shot Super-Resolution (ZSSR) capabilities. The model introduces a new mechanism to condition a neural field on Initial Conditions (ICs) and PDE parameters to achieve generalization to both ICs and PDE parameter values not seen during training. While modeling the solution using neural fields such as PINNs naturally provide temporal and spatial ZSSR, these methods are inefficient since we need to query them separately for every temporal and spatial location in the domain. We achieve faster training and inference by vectorizing these computations on GPUs. Moreover, the proposed method explicitly models dependencies between multiple simultaneous spatio-temporal queries to the model.

Vectorized Conditional Neural Fields. Typically, a (conditional) neural field generates the output quantities for all input spatial and temporal coordinates in multiple and independent forward passes. The training and inference times can be improved by processing multiple inputs in parallel on the GPU, which is possible since all forward passes are independent. However, there are spatial dependencies between different input spatial coordinates, particularly for solving PDEs, that will not be exploited with CNeFs or by processing multiple inputs of CNeFs in parallel. Consequently, we propose extending CNeFs to

Hence, we name our proposed model Vectorized Conditional Neural Field since it implicitly generates a vectorization of the input spatial coordinates for a given time \( t \). The VCNeF model represents a function \[ f_\theta: (\mathbb{R}_{+} \times \mathbb{R}^{s \times d}) \rightarrow \mathbb{R}^{s \times c} \\ \text{ with } (t, \boldsymbol{X}) \mapsto u(t, \boldsymbol{X}) = \begin{pmatrix} u(t, \boldsymbol{x_1}) \\ \vdots \\ u(t, \boldsymbol{x_s}) \end{pmatrix} \] where \( u(t, \boldsymbol{x_i}) \) denotes the PDE solution for the spatial coordinates \( \boldsymbol{x_i} \). Note that we do not impose a structure on the spatial coordinates \( \boldsymbol{x_i} \) and that the number of spatial points (i.e., \( s \) ) can be arbitrary. The model can process multiple timesteps \( t \) in parallel on the GPU to further improve the training and inference time since VCNeF does not exploit dependencies between the temporal coordinates.

VCNeF Properties. The design of the proposed VCNeF model allows for the following properties. (i) Generalization to PDE parameters not seen during training, (ii) spatial- and temporal zero-shot super-resolution (i.e., increasing the spatial and temporal resolution after training), (iii) accelerated training and inference due to the vectorization, and (iv) allows including a physics-aware loss function.

Neural Architecture. VCNeF consists of a Linear Transformer to produce an attention-refined representation of the initial condition and PDE parameters. In the modulation blocks, the previously computed attention-refined latent-representations are used to modulate a vectorized conditional neural field. The vectorized conditional neural field can be queried with temporal and spatial coordinates and uses a self-attention mechanism that allows leveraging spatial dependencies between the generated solution points. To reduce the computational costs, the spatial domains are divided into non-overlapping patches as in Vision Transformers (ViTs) for 2D and 3D PDEs. However, unlike a traditional ViT, our patch generation has two branches: patches of a smaller size (\(p_S = 4 \) or \(4 \times 4 \)) and of a larger size (\(p_S = 16 \) or \(16 \times 16 \)) since we aim to capture the dynamics accurately at multiple scales.

Here, we compare visualizations of the predictions vs ground truth for 1D Burgers, Advection, and 2D Compressible Navier-Stokes PDEs. The 2D CNS dataset has four channels, namely density velocity-x and velocity-y, and pressure, and we visualize the predictions of our VCNeF model with the ground truth data.