PDEInvBench: A Comprehensive Dataset and Design Space Exploration of Neural Networks for PDE Inverse Problems

This paper introduces PDEInvBench, a comprehensive benchmark dataset for PDE inverse problems, and uses it to explore neural network design spaces, revealing that a two-stage training procedure combining parameter supervision with test-time residual fine-tuning, along with PDE derivative inputs and diverse initial conditions, significantly improves parameter estimation performance.

The Gist

Imagine you are a detective trying to figure out the rules of a game just by watching the players play.

In the world of physics, these "rules" are called Partial Differential Equations (PDEs). They describe how things like heat, water, or light move and change. Usually, scientists know the rules (the parameters, like how thick the water is) and use computers to predict what will happen (the solution). This is the "forward problem."

But what if you only have a video of the water moving, and you need to figure out how thick the water is? That is the Inverse Problem. It's like looking at a finished cake and trying to guess the exact recipe, or watching a car crash and trying to deduce the speed of the car before the impact.

This paper, PDEInvBench, is a massive new toolkit designed to help Artificial Intelligence (AI) get better at solving these "reverse engineering" puzzles. Here is a breakdown of what they did and what they found, using simple analogies.

1. The Problem: No Map for the Reverse Journey

Until now, researchers had plenty of maps for the "forward" journey (predicting the future from known rules), but very few for the "inverse" journey (figuring out the rules from the future). Existing AI benchmarks were like driving a car with a map that only showed how to get to the destination, but gave no clues on how to figure out where you started based on where you ended up.

The authors created PDEInvBench, a comprehensive "training gym" for AI. It contains simulations of five different physical systems (like fluid flow, chemical reactions, and wave motion) with thousands of different scenarios. It's a massive library of "videos" (solution fields) paired with the "secret recipes" (physical parameters) that created them.

2. The Experiment: Testing Three Key Ingredients

The researchers didn't just build the dataset; they used it to test three main ways to train AI, asking: What makes the best detective?

A. The Training Method (Optimization)

• The Old Way: Just show the AI the video and the answer, and say, "Memorize this." (Supervised Learning).
• The Physics Way: Don't give the answer. Instead, tell the AI, "Guess the rules, then check if your guess makes sense according to the laws of physics." (Self-Supervised).
• The Hybrid Way (The Winner): First, teach the AI the answers. Then, right before the final test, let the AI "think" for a moment using the laws of physics to refine its guess.
• The Finding: The best strategy is a two-step process. First, learn from the data (memorize the patterns). Second, right before you need to solve a new problem, do a quick "check-up" using the physics equations to fine-tune your answer. It's like studying your flashcards, then doing a quick mental rehearsal of the rules right before the exam.

B. The Tools (Problem Representation)

• The Question: Should the AI be given just the video, or should we also hand it a "cheat sheet" showing how fast things are changing (derivatives)?
• The Finding: Giving the AI the derivatives (the rate of change) as extra input features is like giving a detective a magnifying glass. It consistently helps the AI solve the puzzle faster and more accurately, even if the AI is smart enough to theoretically figure it out on its own.
• The Architecture: For moving systems (like flowing water), a specific type of AI called FNO (Fourier Neural Operator) worked best. It's like a specialized lens that is great at seeing waves and smooth patterns. However, for static systems (like water sitting still in a sponge), a standard image-recognition style AI (ResNet) actually worked better.

C. The Data Diet (Scaling)

• The Question: If you have a limited amount of computer power, should you generate data with more different recipes (more parameters) or more different starting points (more initial conditions) for the same recipe?
• The Finding: It is better to show the AI many different starting points for the same recipe.
• The Analogy: Imagine you are trying to learn how a specific engine works. You will learn more by watching that same engine run on a flat road, a steep hill, and a bumpy track (different starting conditions) than by watching five different engines run on a flat road. Seeing how the system reacts to different inputs teaches the AI the underlying rules better than just seeing more variations of the rules.

3. The Big Takeaways

The paper distills their findings into four practical rules for anyone building AI to solve physics puzzles:

1. Train in two stages: Learn from data first, then use physics laws to polish the answer right before you make a prediction.
2. Hand over the derivatives: Don't make the AI guess how fast things are changing; give it that information explicitly.
3. Pick the right tool: Use "wave-specialist" AI (FNO) for moving fluids, but "image-specialist" AI (ResNet) for static problems.
4. Diversity over quantity: When generating training data, it's better to have the same physical rules play out in many different scenarios than to have many different rules play out in the same scenario.

Summary

PDEInvBench is the first major step in standardizing how we teach AI to reverse-engineer the laws of physics. It shows that by combining data learning with physics checks, and by feeding the AI the right kind of diverse data, we can build much smarter systems for understanding the physical world. The authors have made their dataset and code public so other scientists can use this "gym" to train their own AI detectives.

Technical Summary

Technical Summary: PDEInvBench

Problem Definition

Partial Differential Equation (PDE) inverse problems involve inferring unknown physical parameters (e.g., viscosity, diffusivity, Reynolds number) of a system from observed spatiotemporal solution fields. While machine learning, particularly Neural Operators (NOs), has achieved significant success in the forward problem (mapping parameters to solutions), the inverse problem remains under-explored. The inverse setting presents distinct challenges, including ill-posedness (non-existence, non-uniqueness, and instability with respect to noise) and the need to map complex function spaces back to scalar or spatially varying parameters. Currently, there is a lack of comprehensive benchmarks and datasets dedicated to systematically evaluating neural network approaches for PDE inverse problems.

Methodology

The authors introduce PDEInvBench, a comprehensive benchmark dataset and evaluation framework designed to fill this gap. The methodology involves three primary components:

1. Dataset Construction

PDEInvBench comprises numerical simulations for five distinct PDE systems, covering elliptic, parabolic, and hyperbolic types across various physical behaviors (turbulence, steady-states, laminar flow, Turing patterns, and diffusion). The datasets include:

• 2D Reaction-Diffusion (RD): Fitzhugh-Nagumo equations with varying activator/inhibitor diffusion coefficients and reaction thresholds.
• Unforced 2D Navier-Stokes (NS): Laminar flow regimes governed by viscosity.
• Forced 2D Navier-Stokes (TF): Turbulent flow regimes with Kolmogorov forcing.
• Korteweg-De Vries (KdV): 1D shallow water waves with dispersive strength parameters.
• Darcy Flow (DF): Time-independent flow through porous media with spatially varying diffusion coefficients.

The dataset includes diverse spatial resolutions (e.g., 64x64 to 2048x2048) and evaluation splits designed to test In-Distribution (ID) generalization, Out-of-Distribution (OOD) Non-Extreme (middle parameter ranges omitted from training), and OOD Extreme (tail parameter values).

2. Design Space Exploration

The authors systematically investigate the design space of neural networks for inverse problems along three axes:

• Optimization Procedures: Comparing purely data-driven supervised learning (minimizing parameter prediction error) against self-supervised learning (minimizing PDE residuals). They also evaluate Test-Time Training (TTT), where models are fine-tuned at inference using the PDE residual loss to adapt to specific test samples.
• Problem Representation & Inductive Bias: Benchmarking four architectures with different inductive biases:
  • Fourier Neural Operator (FNO): Spectral representation.
  • ResNet: Local convolutional patterns.
  • Scalable Operator Transformer (scOT): Global attention mechanisms.
  • DeepONet: Branch-trunk decomposition.
    The study also examines the impact of explicitly conditioning inputs on PDE partial derivatives and the number of temporal frames provided.
• Scaling Properties: Analyzing the effects of scaling model size (parameter count) and data quantity. Data scaling experiments vary the number of initial conditions per parameter, the number of unique PDE parameters, and the temporal horizon of training trajectories.

3. Evaluation Metrics

Performance is measured using relative error for parameter prediction. Additionally, the authors employ the Negative Slope of the Line of Best Fit (NLS) to quantify scaling laws (how error decreases as data/model size increases). A physical consistency check is also performed by evolving simulations with predicted parameters and comparing energy spectra against reference solutions.

Key Results

The experiments yield several practical insights regarding the design of neural networks for PDE inverse problems:

1. Optimization Strategy: Purely data-driven supervised training consistently outperforms purely self-supervised (residual-based) training. However, a two-stage training approach is optimal: initial supervised training followed by Test-Time Training (TTT) using the PDE residual. TTT provides significant performance gains, particularly for OOD initial conditions and extreme parameter regimes, without degrading performance on in-distribution data.
2. Input Conditioning: Explicitly providing PDE partial derivatives as input features consistently improves accuracy across all architectures and datasets, suggesting that while networks can learn derivatives implicitly, explicit conditioning aids learning efficiency and generalization.
3. Architectural Performance:
   • For time-dependent PDEs, the FNO architecture generally outperforms ResNet, scOT, and DeepONet, likely due to its spectral inductive bias being well-suited for continuous solution fields.
   • For time-independent PDEs (e.g., Darcy Flow), ResNet and scOT significantly outperform FNO, as the problem reduces to spatial modeling where convolutional or attention-based local/global patterns are more effective.
4. Scaling Insights:
   • Initial Conditions vs. Parameters: Increasing the diversity of initial conditions for a fixed set of PDE parameters yields greater performance gains than increasing the number of unique PDE parameters. This suggests that observing the evolution of diverse initial states provides more information about the underlying parameter-to-solution mapping than simply sampling more parameter values.
   • Model Scaling: Increasing model size generally improves in-distribution and OOD non-extreme performance, though extrapolation to extreme OOD regimes does not always scale linearly with model size.

Significance and Claims

The paper claims to provide the first comprehensive benchmark dedicated to PDE inverse problems, addressing a critical gap where existing benchmarks (e.g., PDEBench, PDEArena) focus primarily on forward problems. By releasing the dataset and codebase, the authors aim to:

• Establish a standardized environment for evaluating and comparing ML methods for inverse problems.
• Derive actionable guidelines for practitioners, such as the preference for two-stage training and the value of derivative conditioning.
• Facilitate future research into the design space of neural operators, specifically regarding how inductive biases and scaling laws affect inverse problem solving.

The authors maintain a modest stance, noting that their findings are specific to the chosen datasets and architectures, and that further investigation is needed to understand the full scope of architectural factors and their impact on different physical regimes. They explicitly state that while TTT is recommended, it should be used with consideration for computational constraints.