SBNO : Schrödinger Bridge Neural Operator for the Forward and Inverse Problems under Missing Information

AI summary

Overview

Research area: Scientific Machine Learning (SciML) for partial differential equations — specifically diffusion/Schrödinger-bridge-based generative models for operator learning under partial observations.

Technical level: Advanced. The paper assumes familiarity with stochastic differential equations, path measures, Girsanov's theorem, optimal transport, and diffusion bridge matching, and its main text works directly with drift regressions and probability-flow ODEs.

Scope: The paper introduces SBNO, an equation-free, paired-data diffusion bridge matching framework that learns both forward and inverse PDE operators from sparse observations using a unified mask-aware endpoint coupling.

What This Paper Is About

Learning-based PDE surrogate models usually assume dense observations or access to the governing equations, but real sensor deployments give sparse, partial measurements and the physics may be unknown or only partly specified. This paper proposes a single framework that reconstructs either the full solution field from a partial input (forward) or the full input field from a partial solution (inverse) when only 5% of the field is observed. The goal is accurate reconstruction that is also fast and tuning-free at inference time, rather than requiring guidance-weight search as guided-diffusion PDE solvers do.

Key Contributions

1. A unified fixed-coupling bridge framework for both PDE directions. SBNO keeps the empirical endpoint coupling induced by paired PDE data fixed (rather than optimizing over couplings as in the classical Schrödinger bridge problem) and mixes the forward and inverse task couplings with equal weight (1/2 each) so that one trained model supports both directions.

2. Mask-aware lifted state representation. Each field is lifted by concatenating the masked values with the binary observation mask, giving a state in R^(2D). This explicitly distinguishes observed values from zero-filled missing entries.

3. Tuning-free, equation-free deterministic inference. Reconstruction is produced by integrating the probability-flow ODE built from the learned forward/backward drift pair, requiring no inference-time observation or PDE guidance-weight tuning and no access to the governing PDE during training or inference.

4. Efficiency and dispersion claims. Against diffusion-based PDE solvers, SBNO is reported to be at least 9.7 times faster and to use at least 4.3 times less GPU memory, while also showing lower test-instance error dispersion.

Main Findings

• Accuracy across ten settings: On five PDE benchmarks under 5% observations (819 points on a 128 × 128 grid), SBNO achieves the best or second-best mean relative L2 error in all ten forward and inverse benchmark settings.

• Selected numbers (mean ± std of relative L2 error): Darcy flow forward: SBNO 0.0359 ± 0.0133 versus FunDPS 0.0461 ± 0.0207, DiffusionPDE 0.0552 ± 0.0849, FNOPE 0.2219 ± 0.0893, GNOT 0.0510 ± 0.0148, FNO 0.0582 ± 0.0135. Darcy flow inverse: SBNO 0.0269 ± 0.0125 versus FunDPS 0.0835 ± 0.0380, DiffusionPDE 0.0499 ± 0.0338, FNOPE 0.2241 ± 0.1110, GNOT 0.0779 ± 0.0367, FNO 0.0946 ± 0.0409.

• Poisson and Helmholtz cases: Poisson forward: SBNO 0.0424 ± 0.0219 versus FunDPS 0.0775 ± 0.0473, DiffusionPDE 0.0986 ± 0.1696, GNOT 0.0929 ± 0.0428, FNO 0.1458 ± 0.0518. Helmholtz forward: SBNO 0.0519 ± 0.0315 versus FunDPS 0.0563 ± 0.0310, DiffusionPDE 0.0902 ± 0.1412, GNOT 0.0926 ± 0.0393, FNO 0.2191 ± 0.0854. In the two inverse rows visible in the text, DiffusionPDE reports the lowest error (Poisson inverse 0.2015 ± 0.0792 versus SBNO 0.2167 ± 0.0345; Helmholtz inverse 0.1893 versus SBNO 0.2310 ± 0.0501), with the Helmholtz inverse row truncated in the provided content.

• Lower error dispersion: SBNO reduces the average error dispersion by 44.91%, from 0.0416 for FunDPS to 0.0229, indicating more consistent per-test-instance accuracy.

• Efficiency: SBNO is reported as at least 9.7 times faster and using at least 4.3 times less GPU memory than guided-diffusion baselines at inference.

• Qualitative capability differences (Table 1): SBNO is mask-aware, supports unified PDE-task training, requires no inference-time hyperparameters, is equation-free, and can use both stochastic and deterministic inference. FunDPS and DiffusionPDE use only observation conditioning, support unified training, require inference-time hyperparameters, use PDE guidance, and are stochastic. FNOPE is stochastic and equation-free but not unified or mask-aware. GNOT and FNO are deterministic and equation-free, but not unified; FNO is also mask-aware.

• Reproduction caveat: The FunDPS numbers are the authors' reproduced configurations; they state they did not recover FunDPS's originally reported accuracy and document an unresolved reproduction discrepancy in Appendix E.

• An acknowledged limitation: The probability-flow ODE matches only the one-time marginals of the Markovian projected fixed-coupling bridge. It does not reproduce the full stochastic path law of the bridge mixture, does not preserve the empirical endpoint coupling, and its endpoint is a deterministic point reconstruction that is not necessarily the conditional mean and cannot represent multimodal uncertainty.

Methodology in Plain English

SBNO treats each PDE sample as a pair of an input field and its solution field. From a pair drawn from the dataset, the method builds two possible "trips": a forward trip that starts from a partially observed input and ends at the fully observed solution, and an inverse trip that starts from a fully observed input and ends at a partially observed solution. Each field is first lifted by stacking its masked values with the mask itself, which tells the network exactly which entries are real observations and which are simply zero-filled holes.

Training data pairs are kept exactly as given — the method does not search for an optimal matching between the two endpoint distributions. Instead, forward and inverse trips are mixed with probability one half each, forming a fixed coupling. Intermediate states between the two endpoints are sampled from a reference diffusion bridge, and two separate neural networks (a time-conditioned U-Net with independently optimized parameters for each direction) are trained by regressing onto the bridge's conditional score targets. Because both directions live in the same lifted state space, one pair of networks serves both tasks.

At test time, the method does not run noisy stochastic sampling. Instead it integrates the probability-flow ODE associated with the learned drifts using a fourth-order Runge-Kutta solver with 32 integration steps, producing a single deterministic reconstruction. During training, the matching objective is computed only at the discrete time points that coincide with this 32-step RK4 trajectory, and sensor locations are resampled every iteration while the observation ratio stays fixed at 5%.

Why This Matters

Impact on research: The paper reframes PDE surrogate learning under missing data as learning bridge dynamics under a fixed empirical coupling rather than solving for an optimal coupling, and it shows that a single model can cover forward and inverse problems without physics constraints, guidance terms, or inference-time hyperparameter search. It also highlights error dispersion across test instances as a metric alongside mean accuracy, which is unusual in this literature.

Real-world applications:

• Subsurface and porous media flow: inferring permeability or coefficient fields from sparse pressure or saturation measurements.
• Wave propagation and Helmholtz-type problems: reconstructing fields from limited sensor arrays.
• Fluid dynamics: completing flow fields from partial instrumentation.
• Design optimization, control, and inverse modeling workflows where PDEs must be solved repeatedly and classical numerical solvers become expensive.

Industry relevance: The combination of deterministic, tuning-free inference with reported speedups (at least 9.7 times) and memory reductions (at least 4.3 times) makes deployment more practical than guided-diffusion solvers that require per-instance guidance-weight tuning. Equation-free operation matters where the governing physics is only partially known or where calibrated simulators are unavailable.

Future Directions

• Reconciling deterministic inference with uncertainty. Because the probability-flow ODE cannot represent multimodal inverse solutions, the paper's own analysis points to stochastic sampling or predictive-uncertainty diagnostics as a complementary mode; the authors report stochastic inference experiments in Appendix I.

• Robustness to noisy observations. The main experiments use noiseless observations, with noisy partial-observation experiments deferred to Appendix J; extending and quantifying performance under sensor noise remains open.

• Closing the FunDPS reproduction gap. The authors report they did not recover FunDPS's originally reported accuracy, so the baseline comparison protocol warrants independent verification.

• Broadening the coupling and PDE coverage. The empirical coupling is fixed by design, which leaves open whether optimizing the coupling, handling irregular or adaptive sensor patterns, and scaling beyond the evaluated benchmarks (including the three smaller-scale datasets in Appendix H and the time-dependent setting covered by the notation) would extend the framework's reach.

Target Audience

Researchers and practitioners in Scientific Machine Learning, generative modeling, and computational physics who work on neural operators, diffusion models for PDEs, or Schrödinger bridge and diffusion bridge matching methods. It is most useful to readers already comfortable with SDEs, path measures, and operator learning — beginners will find the preliminary section and appendices necessary before the method becomes accessible. Engineers deploying PDE surrogates under sparse sensing constraints, and reviewers evaluating efficiency and dispersion claims for diffusion-based solvers, are also direct beneficiaries.

Authors’ abstract

Partial differential equation (PDE) surrogate modeling often relies on either dense observations or access to governing equations/physics constraints. In practice, observations can be sparse and partial, and the underlying physics may be partially specified or unavailable. We introduce SBNO, an equation-free paired-data diffusion bridge matching framework that learns conditional operators for both forward and inverse PDE problems under partial observations. SBNO uses diffusion Schrödinger bridge matching to learn a pair of time-dependent drifts from a mask-aware endpoint coupling that preserves empirical input-solution pairing during training and supports both PDE-task directions within a single framework. At inference time, SBNO produces a deterministic point reconstruction by integrating the probability-flow ODE constructed from the learned drift pair, without inference-time observation or PDE guidance-weight tuning. On five PDE benchmarks with 5% observations, SBNO achieves superior or comparable accuracy to baselines while demonstrating lower test-instance error dispersion. Furthermore, SBNO is at least $9.7 \times$ faster and reduces GPU memory by at least $4.3 \times$ compared to guided-diffusion baselines.