Structure-Preserving Learning Improves Geometry Generalization in Neural PDEs

TL;DR

The paper tackles the challenge of geometry generalization for neural PDE surrogates by enforcing physical structure and discretization-consistent conservation through Finite Element Exterior Calculus. It introduces General-Geometry Neural Whitney Forms (Geo-NeW), which learns a geometry-conditioned reduced FE space and a compatible nonlinear operator, solved implicitly via a PDE-constrained objective. A Lipschitz-bounded flux model guarantees well-posedness and differentiability, enabling stable training and exact enforcement of boundary conditions. Empirically, Geo-NeW achieves state-of-the-art performance on standard steady-state PDE benchmarks and shows strong out-of-distribution generalization to unseen geometries, outperforming transformer- and point-cloud-based baselines. This work advances geometry-aware, physics-grounded foundation models that generalize across complex domains and hold promise for real-time engineering design on realistic meshes.

Abstract

We aim to develop physics foundation models for science and engineering that provide real-time solutions to Partial Differential Equations (PDEs) which preserve structure and accuracy under adaptation to unseen geometries. To this end, we introduce General-Geometry Neural Whitney Forms (Geo-NeW): a data-driven finite element method. We jointly learn a differential operator and compatible reduced finite element spaces defined on the underlying geometry. The resulting model is solved to generate predictions, while exactly preserving physical conservation laws through Finite Element Exterior Calculus. Geometry enters the model as a discretized mesh both through a transformer-based encoding and as the basis for the learned finite element spaces. This explicitly connects the underlying geometry and imposed boundary conditions to the solution, providing a powerful inductive bias for learning neural PDEs, which we demonstrate improves generalization to unseen domains. We provide a novel parameterization of the constitutive model ensuring the existence and uniqueness of the solution. Our approach demonstrates state-of-the-art performance on several steady-state PDE benchmarks, and provides a significant improvement over conventional baselines on out-of-distribution geometries.

Figures (13)

• Figure 1: Geo-NeW provides improved generalization capability for strongly out-of-distribution geometries. Here, models trained on square domains with randomly circular obstacles, but evaluated on a domain including a variable angle step $\theta$; increasing $\theta$ provides a continuous measure of departure from the training distribution. We demonstrate reduced error compared to Transolver, which fails to produce meaningful predictions beyond $\theta=20^\circ$.
• Figure 2: Geo-NeW pipeline for geometry-generalizable PDE surrogate modeling. A geometry encoder maps mesh-derived features to a latent context $z$ that conditions both reduced finite element spaces and a nonlinear flux model. Geometry enters both through the learned encoding and explicitly as a basis for the finite element model. The resulting learned finite element system explicitly encodes mesh topology and metric structure, enforces conservation and boundary conditions, and is solved implicitly to produce the PDE solution. Teal indicates a learnable component.
• Figure 3: We demonstrate the adaptability of the learned finite element basis functions to variable geometries for airfoils with shocks. The shape functions track the location of the geometry-dependent shock, allowing accurate representations even at very low reduced dimension.
• Figure 4: Geo-NeW's physical and geometric biases improve prediction of steady fluid flow. For models trained on a flow past obstacles in a square domain, we gently perturb the domain to demonstrate out-of-distribution geometry's effect on pointwise error for velocity $(u,v)$ and pressure $p$. Row 1. Target solutions. Row 2. Prediction by Transolver gives large concentrations of pointwise error. Row 3. Scaling the coordinate of positional embedding for Transolver to match unit cube of training data yields lower error but with hallucinated obstacles. Row 4. Geo-NeW predicts under extrapolation despite only having trained on a unit square.
• Figure 5: We consider six example problems consisting of canonical steady state PDEs across heat transfer(a.), solid mechanics (b.), and fluid mechanics (c.-f.). In four of these problems, Geo-NeW provides the best performance out of our baselines and those available from previous work.