Structure-Preserving Learning Improves Geometry Generalization in Neural PDEs

This paper introduces General-Geometry Neural Whitney Forms (Geo-NeW), a data-driven finite element method that jointly learns differential operators and compatible reduced spaces to preserve physical conservation laws and achieve superior generalization to unseen geometries in solving Partial Differential Equations.

The Gist

Imagine you are trying to teach a computer to predict how water flows around rocks, how heat spreads through a metal plate, or how a bridge bends under weight. These are problems governed by Partial Differential Equations (PDEs). Traditionally, solving these requires massive, slow simulations that act like a digital wind tunnel or a virtual stress test.

Recently, scientists have tried to train "AI models" to act as shortcuts, predicting the answers instantly. However, most of these AI shortcuts have a major flaw: they are like students who memorized the answers for a specific set of test questions but fail completely when the test paper changes shape. If you train an AI on a square room, it often gets confused when asked to solve the problem for a room with a weirdly shaped corner or a circular obstacle.

This paper introduces a new method called Geo-NeW (General-Geometry Neural Whitney Forms). Think of it as teaching the AI not just the answers, but the rules of the game and how to adapt those rules to any shape.

Here is how it works, using simple analogies:

1. The Problem: The "Rigid Mold" vs. The "Clay"

Most current AI models for physics are like rigid plastic molds. They are trained on a specific shape (like a square). If you try to pour the physics into a different shape (like a circle), the mold doesn't fit, and the result is a mess. They try to guess the answer based on patterns they've seen before, but they don't truly understand the geometry.

2. The Solution: The "Smart, Shape-Shifting Net"

Geo-NeW is different. Instead of a rigid mold, it builds a smart, shape-shifting net (a mathematical mesh) that fits perfectly around whatever shape you give it, whether it's a square, a circle, or a complex airfoil.

• The Mesh as a Skeleton: Imagine the shape of your object is a skeleton. Geo-NeW builds a flexible net over this skeleton. This net isn't just a grid; it's a "Whitney Form." In plain English, this is a special type of mathematical net designed to respect the laws of physics (like conservation of mass or energy) no matter how you stretch or twist the net.
• The "Teacher" (The Transformer): The AI uses a "teacher" (a Transformer network) to look at the shape of the skeleton. It asks: "What does this shape look like? Where are the walls? Where are the holes?"
• The "Student" (The Solver): Based on the teacher's description, the AI instantly reshapes its net and recalculates the physics rules for that specific shape. It doesn't just guess the answer; it sets up a mini-math problem that is guaranteed to have a correct, stable solution.

3. The "Inductive Bias": Teaching the Rules, Not Just the Answers

The paper claims that by forcing the AI to use this special net structure, it gains a powerful "inductive bias."

• Analogy: Imagine teaching a child to bake a cake.
  • Old AI: You show them a photo of a chocolate cake. They memorize the photo. If you ask for a strawberry cake, they are lost.
  • Geo-NeW: You teach them the recipe (the conservation laws) and how to adjust the ingredients based on the size of the pan (the geometry). Even if you give them a pan shaped like a star, they know exactly how to bake the cake because they understand the rules, not just the picture.

4. Why It's Better at "Unseen" Shapes

The paper tested this on shapes the AI had never seen before (Out-of-Distribution).

• The Test: They trained the AI on a square room with round obstacles. Then, they tested it on a room with a sharp, angled step (a shape it had never seen).
• The Result: Other AI models (like Transolver) failed completely, producing nonsense or "hallucinations" (imaginary obstacles). Geo-NeW, however, successfully predicted the flow of air or water around the new shape.
• Why? Because the math behind Geo-NeW is built on "Finite Element Exterior Calculus." This is a fancy way of saying the math is structurally sound. It guarantees that if you put a wall here, the flow stops there. It preserves the "physics" even when the "geometry" changes.

5. The "Black Box" vs. The "Transparent Box"

Many AI models are "black boxes"—you put data in, and an answer comes out, but you don't know if the answer makes sense.
Geo-NeW is more like a transparent box. Because it solves a simplified version of the actual physics equations, we can mathematically prove that a solution exists and that it is unique. It's not just guessing; it's solving a well-posed puzzle every time.

Summary of Claims

• What it does: It creates a physics solver that works on any 2D shape (geometry) without needing to be retrained for every new shape.
• How it does it: It combines a deep learning "encoder" (to understand the shape) with a specialized "solver" (to calculate the physics) that respects conservation laws.
• The Result: It is significantly more accurate than other AI models when asked to solve problems on shapes it has never seen before.
• The Trade-off: It is slightly slower than the fastest "guessing" AI models because it actually solves a math problem, but it is still much faster than traditional physics simulations and far more reliable.

In short, Geo-NeW teaches the AI to understand the shape of the world and the rules of physics simultaneously, allowing it to solve problems on any terrain, not just the ones it memorized.

Technical Summary

Technical Summary: General-Geometry Neural Whitney Forms (Geo-NeW)

Problem Statement

Scientific foundation models aim to provide real-time solutions to Partial Differential Equations (PDEs) across diverse physical systems and parameter regimes. However, a central barrier to deploying these models in realistic engineering settings is geometry generalization. Existing methods often restrict attention to rectilinear domains or domains diffeomorphic to them, utilizing Cartesian representations. While recent operator learning approaches (e.g., Fourier Neural Operators, Transformers) have advanced, they typically treat geometry as an input modality (via coordinate warping, graph connectivity, or point-cloud tokens) rather than embedding it into the fundamental structure of the solver. Consequently, these models struggle to generalize to unseen, complex geometries without retraining, failing to preserve the topological and conservation structures inherent to physical laws.

The authors posit that for scientific foundation models to be viable for engineering design, they must possess the capability to generalize across arbitrary geometries while strictly preserving physical structure (conservation laws, boundary conditions, and stability).

Methodology: Geo-NeW

The paper introduces General-Geometry Neural Whitney Forms (Geo-NeW), a data-driven finite element method that jointly learns a differential operator and compatible reduced finite element spaces defined on the underlying geometry. Unlike standard operator learning which regresses directly from inputs to solutions, Geo-NeW formulates the problem as identifying a reduced-order finite element model that links coordinate-free representations of physics and geometry.

Core Components

1. Reduced Finite Element Spaces (Whitney Forms):
   The method constructs reduced finite element spaces $W^k_g$ as subspaces of low-order Whitney forms. These spaces preserve the discrete exterior calculus structure (exact sequence property), ensuring that discrete divergence, gradient, and curl operations satisfy conservation laws (e.g., generalized Stokes theorems) at the discrete level.

   • A conditional neural field $W(x, z)$ maps fine-scale basis functions to a reduced basis, conditioned on a latent geometry encoding $z$.
   • This ensures the partition-of-unity property and maintains the exact sequence property ($\nabla W^0_g \subseteq W^1_g$), which is critical for numerical stability and conservation.

2. Geometry Conditioning and Encoding:
   Geometry enters the model through a geometry encoder that maps mesh-derived features to a latent context $z$. The input features include:

   • Heat Kernel Signature (HKS): An intrinsic spectral descriptor invariant to rigid motion.
   • Harmonic Coordinates (HC): Intrinsic coordinates constructed from boundary partitions.
   • Signed Distance Function (SDF): Measures closeness to boundaries.
   • Patch Labels: One-hot encoding of boundary types (e.g., inlet, wall).
     The encoder utilizes an inducing-point anchor transformer, which samples a small set of anchor tokens to capture global geometric context with sub-quadratic scaling, producing per-node embeddings that condition the physics model.

3. Learned Physics and Conservation Laws:
   The PDE is expressed as a reduced finite element system. For a flux $J$ and nonlinearity $F_\theta$, the conservation law $\nabla \cdot J = f$ is discretized as:
   $$\varepsilon \delta_0^\top M_1 \delta_0 u + \delta_0^\top F_\theta(u, z) - M_0 f_\theta(z) = 0$$
   Here, $\varepsilon \delta_0^\top M_1 \delta_0$ acts as a reduced Laplacian (artificial viscosity) providing a stable linear backbone, while $F_\theta$ is a learned nonlinear flux conditioned on geometry.

   • Structure Preservation: The flux $F_\theta$ is parameterized as an antisymmetric edge function to ensure discrete conservation.
   • Stability Guarantees: The authors enforce a uniform Lipschitz bound on the flux term with respect to the state variables. This ensures the nonlinear system is coercive, guaranteeing the existence and uniqueness of the solution (well-posedness) and enabling stable implicit differentiation during training.

4. Implicit Operator Learning:
   Training is performed via PDE-constrained optimization. Instead of regressing the solution $u$, the model learns the operator $G_\theta(u, \alpha) = 0$. Prediction requires solving this nonlinear system implicitly. Gradients are computed via the adjoint method using implicit differentiation through the Newton solver, allowing the model to learn the governing operator rather than just the solution map.

Key Contributions

1. First Geometry-Generalizing Structure-Preserving Model: Geo-NeW is the first model to enforce boundary conditions and conservation laws exactly while generalizing across unseen geometries.
2. Extension of FEEC Framework: It extends Finite Element Exterior Calculus (FEEC) to enable learning physics on distinct geometries without retraining, utilizing a geometry-conditioned reduced basis.
3. Invariant Geometric Embeddings: The combination of HKS, HC, and SDF with a transformer architecture yields translation and rotation-invariant geometric embeddings and a discretization-invariant prescription of partitions.
4. Provable Solvability: The paper provides a novel parameterization of the constitutive model (flux) that is convex in solution variables but not geometric variables, allowing for a proof of existence and uniqueness of the learned model's solution.

Results

The authors evaluate Geo-NeW on several steady-state PDE benchmarks, including pipe flow, linear elasticity, and Navier-Stokes flow.

• Standard Benchmarks: On in-distribution tasks (NACA airfoils, elasticity, pipe flow), Geo-NeW achieves state-of-the-art performance, outperforming previous ML-based methods (e.g., Geo-FNO, GNOT, Transolver) with improvements of up to 71% in error reduction for pipe flow and 29% for elasticity.
• Out-of-Distribution (OOD) Generalization:
  • Poly-Poisson: Trained on circular domains with polygonal cutouts ($n < 5$), the model successfully generalizes to polygons with $n > 5$, significantly outperforming baselines.
  • Navier-Stokes (NS2d-c++): Trained on unit squares with circular obstacles ($n < 4$), the model is evaluated on geometries with variable obstacle counts, square obstacles, extended domain lengths, and angled steps. Geo-NeW demonstrates robust generalization where baselines (like Transolver) fail to produce meaningful predictions or hallucinate obstacles. For instance, on angled steps, Geo-NeW maintains low error while baselines degrade significantly.
• Boundary Condition Satisfaction: Unlike unconstrained baselines that often violate boundary conditions, Geo-NeW enforces Dirichlet boundary conditions exactly by construction, resulting in near-zero boundary error.

Significance and Claims

The paper claims that Geo-NeW represents a significant step toward physics and geometry-aware foundation models. By explicitly connecting the underlying geometry and boundary conditions to the solution through a structure-preserving finite element framework, the model provides a powerful inductive bias that improves generalization to unseen domains.

The authors emphasize that while the current work is restricted to 2D steady-state problems, the framework naturally extends to 3D. They position this approach as essential for engineering design applications, where the geometries of interest are precisely those not available during training. The method balances the computational cost of solving a reduced nonlinear system with the benefits of real-time inference (significantly faster than traditional FEM solvers like COMSOL) and superior generalization capabilities compared to unconstrained neural operators.

The work concludes that preserving physical structure (conservation, stability, boundary enforcement) is not merely a constraint but a mechanism that enables robust learning on complex, unseen geometries, addressing a critical gap in the deployment of scientific foundation models.