CALM-PDE: Continuous and Adaptive Convolutions for Latent Space Modeling of Time-dependent PDEs

TL;DR

CALM-PDE introduces a discretization-agnostic framework for time-dependent PDEs that encodes solutions into a fixed latent space using continuous, adaptive convolutions with learnable query points and a locality bias. Dynamics are evolved entirely in the latent space via a Transformer-based processor (latent time-stepping) and then decoded back to physical space at arbitrary query locations. The approach achieves competitive accuracy across regular and irregular meshes, while offering memory and speed advantages over Transformer-based baselines, demonstrating the practicality of continuous convolutions for PDE surrogates. The work also suggests future directions such as dynamic query points and broader applicability beyond PDEs to domains like chemistry.

Abstract

Solving time-dependent Partial Differential Equations (PDEs) using a densely discretized spatial domain is a fundamental problem in various scientific and engineering disciplines, including modeling climate phenomena and fluid dynamics. However, performing these computations directly in the physical space often incurs significant computational costs. To address this issue, several neural surrogate models have been developed that operate in a compressed latent space to solve the PDE. While these approaches reduce computational complexity, they often use Transformer-based attention mechanisms to handle irregularly sampled domains, resulting in increased memory consumption. In contrast, convolutional neural networks allow memory-efficient encoding and decoding but are limited to regular discretizations. Motivated by these considerations, we propose CALM-PDE, a model class that efficiently solves arbitrarily discretized PDEs in a compressed latent space. We introduce a novel continuous convolution-based encoder-decoder architecture that uses an epsilon-neighborhood-constrained kernel and learns to apply the convolution operator to adaptive and optimized query points. We demonstrate the effectiveness of CALM-PDE on a diverse set of PDEs with both regularly and irregularly sampled spatial domains. CALM-PDE is competitive with or outperforms existing baseline methods while offering significant improvements in memory and inference time efficiency compared to Transformer-based methods.

Figures (25)

• Figure 1: CALM-PDE encodes the arbitrarily discretized PDE solution into a fixed latent space $\mathbb{R}^{l \times d}$, computes the dynamics in the latent space, and decodes the solution for the given query points.
• Figure 2: Encode-process-decode architecture of CALM-PDE. The encoder reduces the spatial dimension and increases the channel dimension. It is based on multiple CALM layers, which perform continuous convolution on learnable query points constrained to an epsilon neighborhood.
• Figure 3: Learned query positions of CALM-PDE. The model samples more query points in the region where the cylinders are located (red rectangle) by moving query points to this region.
• Figure 4: Comparison of inference times on 2D Navier-Stokes with 200 trajectories. Prediction steps are increased to evaluate the scaling. FNO, PIT, and CALM-PDE achieve similar times.
• Figure 5: The continuous solution function $u(\cdot, x, y)$ has to be discretized along the spatial dimension with a suitable mesh or grid into a discrete representation $\{u(\cdot, (x, y)_n)\}_{n=1}^N$ of the function.