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

TL;DR

This work tackles the challenge of efficiently solving time-dependent parametric PDEs with strong generalization and super-resolution capabilities. It introduces Vectorized Conditional Neural Fields (VCNeF), a time-continuous transformer-based neural field that conditions on initial conditions and PDE parameters and evaluates multiple spatio-temporal query points in parallel. By vectorizing spatial queries and employing a FiLM-like conditioning mechanism within a linear-transformer–based architecture, VCNeF achieves robust generalization to unseen parameters and enables spatial/temporal zero-shot super-resolution, often outperforming state-of-the-art baselines on Burgers, Advection, and compressible Navier–Stokes PDEs. The method offers practical impact for accelerated simulations and has flexible handling of time discretization and grid resolution, with potential extensions to physics-informed losses and adaptive time-stepping.

Abstract

Transformer models are increasingly used for solving Partial Differential Equations (PDEs). Several adaptations have been proposed, all of which suffer from the typical problems of Transformers, such as quadratic memory and time complexity. Furthermore, all prevalent architectures for PDE solving lack at least one of several desirable properties of an ideal surrogate model, such as (i) generalization to PDE parameters not seen during training, (ii) spatial and temporal zero-shot super-resolution, (iii) continuous temporal extrapolation, (iv) support for 1D, 2D, and 3D PDEs, and (v) efficient inference for longer temporal rollouts. To address these limitations, we propose Vectorized Conditional Neural Fields (VCNeFs), which represent the solution of time-dependent PDEs as neural fields. Contrary to prior methods, however, VCNeFs compute, for a set of multiple spatio-temporal query points, their solutions in parallel and model their dependencies through attention mechanisms. Moreover, VCNeF can condition the neural field on both the initial conditions and the parameters of the PDEs. An extensive set of experiments demonstrates that VCNeFs are competitive with and often outperform existing ML-based surrogate models.

Figures (28)

• Figure 1: Conditional Neural Field (CNeF) vs proposed Vectorized CNeF (VCNeF) for solving parameterized PDEs.
• Figure 2: An illustration of the VCNeF architecture for solving parametric time-dependent 2D PDEs. Latent representations of ICs are generated with a multi-scale patching mechanism multi-scale-vit-chen. A modulation block consists of self-attention, activation function $\sigma$, and a modulated neural field that uses the scaling of FiLM film-conditioning-perez:2018 to condition the spatio-temporal coordinates on ICs.
• Figure 3: Example predictions for density of 2D CNS.
• Figure 4: Error distribution of samples in the test set of 1D CNS. Boldfaced are the unseen PDE parameter values.
• Figure 5: Inference times of models trained on Burgers with $s = 256$, predicting different numbers of timesteps in the future.