Technical Summary: GRIFDIR

Problem Statement

Score-based diffusion models in infinite-dimensional function spaces offer a theoretically principled framework for modeling function-valued data, such as solutions to partial differential equations (PDEs) or spatial fields. These models promise resolution invariance (generalizing across discretizations) and the ability to handle irregular domain geometries and unstructured meshes. However, practical implementations have struggled to realize these benefits. Existing backbones face a dichotomy:

• Convolutional U-Nets: Perform well on fixed, regular grids but lack resolution invariance and fail on irregular geometries.
• Fourier Neural Operators (FNOs): Provide resolution invariance but are inherently biased toward regular grids and struggle with complex, non-regular domain topologies.

There remains a gap between the theoretical guarantees of infinite-dimensional diffusion models and practical architectures capable of operating on arbitrary, unstructured meshes without sacrificing resolution invariance.

Methodology: GRIFDIR

The authors propose GRIFDIR (Graph Resolution-Invariant FEM Diffusion Models), a novel architecture designed to serve as a backbone for score-based diffusion models in infinite-dimensional function spaces. The architecture is built upon four key desiderata: multiscale processing, resolution/discretization invariance, domain flexibility, and the ability to capture global structure.

Core Components

1. Resolution-Invariant FEM Convolution:
   Instead of learning kernels in the coordinate frame of the domain (which varies with mesh resolution), GRIFDIR generalizes CNN convolutions to unstructured meshes by parameterizing the convolutional filter as a Finite Element Method (FEM) function on a fixed reference patch in physical space.

   • For a node $x_i$, neighboring nodes are projected into a fixed reference patch $\hat{\Omega}$ via a rescaling map $\Phi_{x_i}$.
   • The filter is defined as a linear combination of $Q1$ bilinear basis functions on a $P \times P$ grid within this reference patch.
   • The convolution integral is approximated by a weighted sum over neighbors, where weights are determined by the basis function values at the projected neighbor locations.
   • Advantage: This ensures the physical receptive field remains constant regardless of mesh refinement, achieving strict resolution invariance. It also inherits the inductive biases of standard CNNs while extending them to arbitrary geometries.

2. Multiscale Mesh Hierarchy:
   To capture both local fine-scale features and global structures, the FEM convolution is embedded within a hierarchical mesh architecture (a Graph U-Net).

   • The model operates on a hierarchy of meshes $\{T^{(\ell)}\}$, from fine to coarse.
   • The physical radius of the FEM filter scales with the mesh coarseness ($r^{(\ell)} = \mu \bar{d}^{(\ell)}$), mimicking the receptive field expansion of standard U-Nets.
   • Transitions between levels use learned restriction and prolongation maps.

3. Global Latent Representation:
   To address the limitation of standard message-passing GNNs in capturing long-range dependencies, the coarsest level of the hierarchy is connected to a Diffusion Transformer (DiT).

   • A cross-attention mechanism maps the variable-size coarse-mesh tokens to a fixed-size latent space.
   • This latent space is processed by Transformer blocks with AdaLN-Zero time conditioning, serving as a global memory for the model.
   • The decoder mirrors the encoder, mapping the latent representation back to the finest mesh.

4. Conditional Sampling Framework:
   For inverse problems, the model utilizes a diffusion prior and guidance paradigm. The conditional score is approximated using either:

   • Fun-DPS: A heuristic guidance method using Tweedie estimates to approximate the guidance term.
   • Fun-DAPS: A decoupled annealing approach that performs Langevin updates in the denoised space, avoiding backpropagation through the score model.

Key Contributions

• Novel Architecture: Introduction of GRIFDIR, the first architecture to combine FEM-based convolutions with graph neural operators and transformers for infinite-dimensional diffusion models.
• Resolution Invariance: The architecture is mathematically designed to be discretization-invariant, allowing models trained on coarse meshes to generalize to finer resolutions without retraining.
• Geometric Flexibility: The method natively supports unstructured meshes and complex domain topologies (e.g., non-convex, multiply-connected domains with holes) without requiring domain masking or deformation maps.
• Conditional Sampling: Demonstration of the framework's ability to solve Bayesian inverse problems (sparse sensor reconstruction and PDE-constrained inverse problems) on irregular domains.

Experimental Results

The authors evaluated GRIFDIR on two datasets: Gaussian Blobs (conductivity fields) and the Pinball dataset (fluid flow around rotating obstacles).

Resolution Invariance

• Gaussian Blobs: Models were trained on a coarse mesh (2,048 elements) and tested on a fine mesh (8,192 elements).
  • Baselines (CNN, MP-GNN, GAOT) failed to generalize, producing artifacts or incorrect feature counts (e.g., generating four times as many blobs when resolution quadrupled).
  • GRIFDIR and FNO maintained high fidelity at the finer resolution. However, FNO is restricted to regular grids, whereas GRIFDIR succeeded on irregular domains.
  • MMD Scores: GRIFDIR achieved the lowest Maximum Mean Discrepancy (0.201) at the inference resolution, outperforming all baselines.

Irregular Domains

• GRIFDIR was trained on eight different domain geometries (including L-shapes, E-shapes, and domains with holes).
• The model successfully generated plausible unconditional samples for all geometries, demonstrating its ability to handle discontinuities and complex boundaries.

Conditional Sampling (Inverse Problems)

• Sparse Sensors: The model reconstructed fields from sparse point measurements (10–100 sensors). Performance (RMSE and Energy Score) improved as sensor density increased, with Fun-DPS generally outperforming Fun-DAPS in accuracy, though Fun-DAPS offered lower computational cost.
• Poisson Inverse Problem: The model solved the inverse problem of inferring a source term from a Poisson equation solution. It achieved an RMSE of $0.0223 \pm 0.0137$.
• Pinball Dataset: The model successfully reconstructed fluid flow fields from sparse sensors and performed conditional generation based on velocity parameters ($\mu$) and time ($\tau$), capturing the parametric dependence of the advection-diffusion dynamics.

Significance and Claims

The paper claims that GRIFDIR bridges the gap between the theoretical promises of infinite-dimensional diffusion models and practical implementation. By utilizing FEM-based convolutions, the architecture achieves resolution invariance and domain flexibility simultaneously, properties that existing methods (U-Nets, FNOs, standard GNNs) fail to provide in unison.

The authors emphasize that their approach:

1. Enables practical function-space diffusion modeling across varying domains and discretizations.
2. Maintains high fidelity in capturing functional distributions on non-trivial geometries.
3. Provides a scalable backbone for solving both forward generative tasks and Bayesian inverse problems on irregular meshes.

The work positions itself as a step toward fully realizing the benefits of infinite-dimensional modeling, where convergence guarantees hold independently of discretization resolution, and where the model can naturally incorporate the geometric structure of the data.