🤖 machine learning

GAIA: Geometry-Adaptive Operator Learning for Forward and Inverse Problems

The paper introduces GAIA, a geometry-adaptive operator learning model that unifies forward and inverse partial differential equation solving on arbitrary domains by encoding boundary and interior features into tokens to condition integral transforms, achieving state-of-the-art accuracy without retraining or iterative optimization.

Original authors: Meenakshi Krishnan, Pranav Pulijala, Ke Chen, Haizhao Yang, Ramani Duraiswami

Published 2026-07-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Meenakshi Krishnan, Pranav Pulijala, Ke Chen, Haizhao Yang, Ramani Duraiswami

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a complex puzzle, like figuring out the weather inside a storm cloud just by looking at the rain on the ground, or determining the shape of a hidden object by how electricity flows around it. In the world of science and engineering, these are called Partial Differential Equations (PDEs).

For decades, computers have been great at solving these puzzles, but they are slow. If you want to solve the same puzzle 1,000 times with slightly different shapes or conditions (like designing a new airplane wing or a medical scanner), the computer has to start from scratch every time. It's like baking a cake from scratch for every single guest at a party.

Recently, scientists tried using AI to act as a "shortcut" or a "surrogate" that learns the rules of the puzzle so it can solve it instantly. However, most of these AI shortcuts have a major flaw: they are rigid. They only work if the puzzle pieces (the geometry) are in a perfect grid, like a chessboard. If the puzzle is an irregular shape—like a twisted piece of metal or a human organ—the AI gets confused and fails.

Enter GAIA: The Shape-Shifting AI

The paper introduces GAIA (Geometry-Adaptive Operator Learning). Think of GAIA not as a rigid robot, but as a master chef who can cook the same delicious dish regardless of the shape of the pan.

Here is how GAIA works, broken down into simple concepts:

1. The "Token" Translator

Most AI models struggle when the input (what you measure) and the output (what you want to predict) are in different places. For example, in a medical scan, you might measure electricity on the skin (boundary) but want to know what's happening inside the body.

The Analogy: Imagine you are trying to describe a house to a friend. You could just list the bricks (the raw data), or you could give them a "summary card" that says "It's a two-story Victorian with a red door."
GAIA's Trick: It creates two types of these "summary cards" (called tokens):
- Boundary Tokens: These summarize the shape of the container (the skin, the edge of the wing).
- Slice Tokens: These summarize the physics happening inside (the temperature, the pressure).
- GAIA reads these cards to understand the "vibe" of the geometry before it even starts solving the math.

2. The "Smart Lens" (Cross-Attention)

Once GAIA has its summary cards, it needs to solve the math. Traditional AI uses a "one-size-fits-all" lens. GAIA uses a smart, adjustable lens.

The Analogy: Imagine looking at a landscape through a camera lens. A normal lens blurs the background if you focus on the foreground. GAIA's lens is special: as it looks at a specific point in the solution, it asks its "summary cards," "Hey, what does the shape look like right here?"
This allows the AI to adapt its math locally. If the shape is curved here and flat there, the AI changes its calculation method instantly to match that specific spot. This is called cross-attention.

3. The "One-Pass" Magic

Old methods for solving these "inverse problems" (figuring out the cause from the effect) were like playing a game of "Guess Who?" where you have to ask, "Is it this shape? No? Okay, try that shape," over and over again. This is slow and expensive.

GAIA's Superpower: It solves the puzzle in one single pass. It looks at the boundary measurements and immediately predicts the interior solution without needing to guess, check, and re-calculate. It's like looking at a shadow and instantly knowing the exact 3D object casting it, without rotating the object to check.

What Did They Test?

The researchers didn't just talk about theory; they tested GAIA on seven different "puzzles" (benchmarks), including some brand new ones they created:

Electrical Impedance Tomography (EIT): Figuring out what's inside a body by measuring electricity on the skin (like a medical scanner).
Optical Tomography: Seeing through tissue using light.
Airfoil Reconstruction: Reconstructing the airflow around a wing from just a few noisy sensor points.
Mechanical Components: Solving stress problems on weirdly shaped gears, nuts, and bolts.

The Results

GAIA didn't just participate; it dominated.

Accuracy: On the medical and wing-reconstruction tasks, GAIA was 27% to 64% more accurate than the next best AI method. It could see fine details (like a shockwave on a wing) that other models blurred out or missed entirely.
Flexibility: Unlike other models that break when you change the resolution (the number of points you look at), GAIA stays accurate whether you look at the puzzle with a magnifying glass or a wide-angle lens.
Speed: It solved these complex problems in a fraction of a second, making it fast enough for real-time use.

The Bottom Line

GAIA is a new type of AI that finally learned how to handle irregular shapes and different input/output locations without needing to be retrained for every new shape. It acts like a universal translator that can read the "geometry" of a problem and instantly solve it, whether you are designing a new airplane part or trying to see inside a human body.

Limitations mentioned in the paper:

It currently only works for steady-state problems (things that aren't changing over time, like a static pressure map). It doesn't yet handle movies or time-evolving storms.
For very simple, grid-based problems, it trades a tiny bit of peak accuracy for the huge benefit of being able to handle any shape.
It requires a lot of memory if the 3D shapes get extremely huge (millions of points).

Technical Summary: GAIA - Geometry-Adaptive Operator Learning for Forward and Inverse Problems

1. Problem Statement

Partial differential equations (PDEs) are fundamental to modeling scientific and engineering phenomena, yet traditional solvers are computationally expensive, particularly in multi-query settings such as inverse problems, optimal design, and uncertainty quantification. While neural operator learning methods (e.g., FNO, DeepONet) have emerged to learn mappings between function spaces directly, they face significant limitations when applied to arbitrary geometries.

Existing geometry-adaptive operators (e.g., GINO, GAOT, Transolver) have made progress but remain primarily designed for forward problems where inputs and outputs share the same spatial domain. They struggle with Boundary Value Problems (BVPs) and inverse problems, where inputs (e.g., boundary measurements) and outputs (e.g., interior fields) reside on different point sets or domains. Furthermore, many existing methods rely on iterative optimization loops for inversion, require expensive graph construction per sample, or fail to generalize across varying geometries without retraining. There is a need for a unified architecture that can handle arbitrary, varying geometries for both forward and inverse tasks in a single pass without iterative optimization.

2. Methodology: The GAIA Architecture

The authors propose Geometry-Adaptive Integral Autoencoder (GAIA), a unified model that conditions learned integral kernels on domain geometry via cross-attention. GAIA builds upon the Integral Autoencoder (IAE-Net) framework but introduces explicit geometric context to overcome the geometry-agnostic limitations of previous integral operators.

Core Components

Dual-Pathway Geometry Tokenization:
GAIA encodes the domain geometry into "tokens" using two complementary pathways:
- Boundary Tokenizer: Uses a PointNet-style encoder with shared MLPs and global max-pooling to process boundary coordinates ( $X_{bnd}$ ). This produces tokens that summarize the global domain shape in a permutation-invariant manner.
- State Slice Tokenizer: Adapts soft-clustering aggregation (inspired by Transolver) to jointly embed spatial coordinates and field values ( $X, U$ ). This produces tokens that capture the spatially varying physics of the interior field distribution.
  These tokens are refined through transformer encoder layers to update the geometric context at each block.
Geometry-Conditioned Integral Transforms:
Unlike standard IAE-Net where kernels depend only on local coordinates, GAIA modifies the encoding and decoding integral transforms. The learnable kernel $K(x, z)$ is conditioned on the geometry tokens $T^{(l)}$ via cross-attention.
- Spatial coordinates are lifted using Fourier feature encodings.
- A context vector $C(x)$ is computed by attending to the geometry tokens.
- The kernel is parameterized by an MLP taking concatenated spatial features and the geometric context: $K(x, z) = \text{MLP}([\gamma(x), \gamma(z), U(x), C(x)])$ .
  This allows the kernel to adapt locally to geometric features at each integration point.
Unified Forward and Inverse Processing:
- Forward Problems: The network applies stacks of Adaptive IAE blocks with dense skip connections, performing volume integrals over the domain mesh.
- Inverse Problems & BVPs: GAIA employs a specialized Observation-to-Domain Decoder. This module bridges the gap between sparse boundary measurements and the dense interior mesh. It first lifts boundary data to a fixed set of latent coordinates via a geometry-conditioned integral, processes them, and then projects them back onto the interior mesh using a second conditioned integral. This decoupling of query points (output) from integration points (input) enables single-pass solutions for problems where input and output domains differ.

3. Key Contributions

Unified Architecture: GAIA is the first model to provide a single architecture capable of solving both forward and inverse problems on arbitrary, varying geometries in a single pass, without retraining, iterative optimization, or per-sample graph construction.
Explicit Geometric Conditioning: By using cross-attention to condition integral kernels on explicit geometry tokens (boundary and state slices), GAIA captures the dependence of solution operators (like Green's functions) on domain shape, a capability missing in prior geometry-agnostic integral operators.
New Benchmarks: The authors introduce and release new benchmarks for varying-geometry inverse and BVP tasks, including:
- Electrical Impedance Tomography (EIT): Conductivity reconstruction on varying star-shaped domains.
- Optical Tomography (OT): Scattering coefficient recovery on random convex pentagons.
- Airfoil Reconstruction: Sparse-to-full Mach field recovery on deformed NACA profiles.
- 3D Poisson BVP: Solving on mechanical components (gears, nuts, screws) from the MCB dataset.

4. Experimental Results

GAIA was evaluated on seven 2D and 3D benchmarks against five state-of-the-art geometry-adaptive baselines (GINO, CORAL, Transolver, GAOT, LNO) and specialized solvers (NGF).

Inverse Problems & BVPs: GAIA sets new state-of-the-art results on all inverse and BVP tasks.
- EIT: Reduced median relative $L_2$ error by 27% compared to the next best amortized method (LNO: 0.97% vs. GAIA: 0.71%).
- Airfoil: Reduced error by 64% compared to the next best (CORAL: 1.60% vs. GAIA: 0.58%).
- MCB (3D Poisson): Outperformed all baselines across all shape categories (gears, nuts, fittings, screws), significantly beating specialized solvers like NGF on general-purpose tasks.
- GAIA successfully recovered fine-scale features (e.g., transonic shocks, oscillatory conductivity patches) that baselines either smeared or missed.
Forward Problems: GAIA remains competitive with specialized solvers.
- Achieved the lowest error on Poisson-Gauss (0.71%).
- On 3D Darcy flow, while Transolver showed lower error at the exact training resolution, GAIA demonstrated superior discretization invariance. When evaluated at coarser resolutions, Transolver and GAOT degraded significantly, whereas GAIA maintained stable accuracy.
Efficiency & Robustness:
- GAIA achieved the highest throughput (365 samples/s) among geometry-adaptive methods on EIT, with competitive latency (~16ms) and modest memory usage.
- The model exhibited graceful degradation under increasing noise levels (up to 10%), consistent with integral-transform architectures that average noise rather than amplifying it.

5. Significance and Claims

The paper claims that GAIA addresses a critical gap in scientific machine learning: the inability of current operator learning methods to efficiently and accurately handle inverse problems and BVPs on arbitrary, varying geometries.

Single-Pass Efficiency: By eliminating the need for iterative optimization loops (common in PINNs or differentiable surrogates) and per-sample graph construction (common in GINO), GAIA offers a computationally efficient solution for multi-query settings.
Generalization: The architecture demonstrates robust generalization across different resolutions and noise levels, making it suitable for real-world applications where mesh discretizations may vary.
Architectural Trade-off: The authors acknowledge that for forward problems with fixed, shared discretizations, GAIA may trade a small amount of in-distribution peak accuracy for substantially improved discretization invariance. However, for applications involving varying resolutions or geometries, this robustness is presented as a preferable trade-off.

The work establishes that conditioning integral kernels on explicit, tokenized geometric representations via cross-attention is a viable and effective strategy for unifying the solution of forward and inverse PDE problems on complex, varying domains.