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

• Sparse observations can degrade conventional neural operators, while unknown or partially specified physics limits equation-based constraints and calibrated simulators.Source: §1 · p. 2
• SBNO preserves paired-data endpoint couplings, conditions on masked values and their binary mask, and learns bridge dynamics for both task directions.Inference uses the learned drifts in a probability-flow ODE, without PDE access or inference-time guidance-weight tuning.Source: §1 · p. 2
• SBNO achieves better or comparable reconstruction accuracy to prior baselines while exhibiting lower error dispersion across test samples.Source: §1 · p. 2
• Against diffusion-based PDE solvers, SBNO is at least 9.7× faster and uses at least 4.3× less GPU memory at inference.Source: §1 · p. 2

• PDE inverse problems infer unknown coefficients, sources, or boundary conditions, but sparse or irregular observations can leave the mapping ill-posed.Prior approaches include PDE-constrained optimization, Bayesian methods, PINNs, and operator learning.Source: §2 · pp. 2–3
• Function-space generative methods model distributions over fields, while recent diffusion PDE methods often use stochastic inference or hand-designed observation guidance.Such sampling or guidance can be costly or sensitive to conditioning design.Source: §2 · p. 3
• SBNO uses diffusion Schrödinger bridge matching machinery with a PDE-specific, mask-aware endpoint coupling for forward and inverse reconstruction.The related-work passage identifies this coupling and its unified use as SBNO’s contribution.Source: §2 · p. 3

3.1 Path Measures and a Reference Diffusion

• A path measure describes the law of a continuous-time trajectory, and the reference diffusion supplies the comparison dynamics for path-space KL.When measures share a diffusion coefficient, Girsanov’s theorem expresses the divergence through expected drift-difference kinetic energy, connecting it to stochastic optimal control.Source: §3.1 · pp. 3–4

3.2 Schrödinger Bridge and Entropy-regularized Optimal Transport

• A Schrödinger bridge minimizes path-space KL subject to fixed endpoint marginals, equivalently selecting an entropy-regularized endpoint coupling relative to the reference dynamics.The entropy-cost trade-off depends on the reference dynamics; under Brownian diffusion the induced cost is quadratic, and greater noise yields higher-entropy couplings.Source: §3.2 · p. 4
• SBNO instead preserves the empirical endpoint coupling from paired PDE data and learns bridge dynamics for forward and inverse reconstruction.Its overview describes mixing the paired task directions in one coupling and learning forward and backward drifts.Source: §3.2 · p. 4, Figure 2 · p. 5

4.1 Problem Definition

• The inverse reconstruction target may be non-unique under partial observations, and exact recovery of its endpoint conditional by learned dynamics is not guaranteed.The framework learns bridge dynamics and evaluates deterministic point reconstructions for this potentially non-unique target.Source: §4.1 · p. 6, §4.2 · p. 6

4.2 State Representation and Bridge Endpoints

• A masked observation and its binary mask are lifted into a shared R2D state space, allowing both partial observations and full fields to use the same drift representation.For a fully observed field, the mask is all ones and the lifted state is E1(x) = [x; 1].Source: §4.2 · p. 6

4.3 Schrödinger-Bridge-Inspired Fixed-Coupling Formulations

• SBNO keeps the empirical endpoint coupling from paired PDE samples fixed and learns Markovian drift fields for the resulting reciprocal bridge mixture.Unlike the classical Schrödinger bridge formulation, this construction does not optimize over endpoint couplings with prescribed marginals.Source: §4.3 · pp. 6–7

4.4 Unified Bridge Matching and Training Objective for Forward and Inverse Operators

• A fixed mixture of forward and inverse couplings preserves the empirical input-solution pairing while training time-dependent drifts for both reconstruction directions.The same drift pair supports both partial-to-full directions through the shared lifted state space.Source: §4.4 · p. 7
• Given a partial input or solution, the corresponding learned forward or backward drift process produces a full-field reconstruction.The forward process predicts the full solution from a partial input; the backward process predicts the full input from a partial solution.Source: §4.4 · p. 7

4.5 Deterministic Inference via Probability-flow ODE

• Deterministic inference integrates the learned probability-flow ODE, matching the projected bridge's one-time marginals but not its full path law or empirical endpoint coupling.The resulting endpoint is a deterministic point reconstruction, not necessarily a conditional mean, and cannot represent multimodal uncertainty.Source: §4.5 · p. 8

5 Experiments

• The main experiments evaluate forward and inverse tasks on five PDE benchmarks, using noiseless observations and deterministic SBNO inference throughout.The datasets are discretized on 128 × 128 grids, with 5% of observations available at inference.Source: §5 · p. 8, §5.1 · p. 8

5.1 Experimental Setup

• Baselines span guided diffusion, deterministic neural operators, and probabilistic neural operators; applicable methods without a unified objective use separate forward and inverse models.FunDPS results use reproduced configurations, and the authors report not recovering its originally reported accuracy.Source: §5.1 · pp. 8–9

5.2 Results

• SBNO achieves the best or second-best mean relative L2 error in all ten forward and inverse benchmark settings, while reducing average error dispersion by 44.91%, from 0.0416 for FunDPS to 0.0229.The reported results cover five benchmarks under 5% partial observations; SBNO also outperforms deterministic and probabilistic operator baselines in most settings.Source: §5.2 · p. 9, Table 2 · p. 9
• SBNO is at least 9.7× faster in wall-clock time per sample and consumes at least 4.3× less peak GPU memory than guided-diffusion baselines.The inference-efficiency comparison is restricted to FunDPS and DiffusionPDE; results are averaged over five runs.Source: §5.2 · p. 9, Table 3 · p. 9
• The reported accuracy-cost trade-off is relative to the reproduced guided-diffusion baselines in the evaluated setting.The experiments also examine reconstruction-performance versus inference-cost trade-offs by varying ODE solver steps.Source: §5.2 · pp. 9–10

6 Conclusion

• Across five PDE benchmarks with 5% point-wise observations, SBNO achieves competitive or superior accuracy, reduces test-instance error dispersion, and is at least 9.7× faster with at least 4.3× less peak GPU memory than guided diffusion solvers.It uses one trained drift pair for forward and inverse reconstruction and deterministic probability-flow ODE inference without observation or PDE guidance-weight tuning.Source: §6 · p. 10

A.1 Markovian and Reciprocal Projections

• The Markovian projection replaces non-Markovian history dependencies with a drift conditioned on the current state, and minimizes reverse KL divergence from the mixture to the Markov manifold.Source: §A.1 · p. 14
• The Schrödinger bridge is characterized as the intersection of Markovian and reciprocal path-measure classes, with the reciprocal projection preserving bridge dynamics and endpoint coupling.The Markovian projection enforces Markovianity but may distort the bridge; the reciprocal projection can preserve bridge dynamics while breaking Markovianity.Source: §A.1 · p. 14

A.2 Diffusion Schrödinger Bridge Matching

• Diffusion Schrödinger bridge matching makes high-dimensional projection scalable by replacing the intractable Markovian projection with regression-based bridge matching and updating forward and backward drifts.The reciprocal update reuses the current endpoint coupling.Source: §A.2 · p. 15

B Training Algorithm

• The unified training algorithm samples paired PDE data, chooses forward or inverse endpoint constructions, samples bridge states, and updates two time-conditioned drift networks.Source: B Training Algorithm · p. 15

B.1 Analytic Bridge Sampling and Regression Targets

• For a zero-drift Gaussian reference with deterministic scalar diffusion, bridge states and conditional-score regression targets are analytically computable, so no separately trained score networks are needed.A small time cutoff avoids evaluating singular targets at the endpoints.Source: §B.1 · pp. 15–16

C Evaluation

• Evaluation reports mean and standard deviation of test-sample relative L2 error, plus inference time per sample and peak GPU memory.Time is measured with each method’s largest feasible batch size and is not batch-one request latency; memory is measured at batch size one.Source: C Evaluation · p. 16

D Datasets

• The main experiments comprise five PDE datasets: four from DiffusionPDE [10] and one Gray-Scott reaction-diffusion dataset from The Well [27]; the four DiffusionPDE datasets include three static PDEs and one time-dependent Navier-Stokes dataset.The four DiffusionPDE datasets are defined on the unit square and discretized on 128 × 128 grids.Source: D Datasets · p. 16

D.2.2 Gray-Scott Equation

• Gray-Scott adds a two-channel reaction-diffusion forecasting dataset, constructed from stored trajectory frames into one-step input-output pairs.It contributes 10,000 training pairs and 500 test pairs from unseen initial conditions.Source: D Datasets · p. 16, §D.2.2 · p. 18

D.2.3–D.2.5 Additional PDE Benchmarks

• Three APEBench datasets—Korteweg-De Vries, Kuramoto-Sivashinsky, and Kolmogorov flow—provide complementary time-dependent PDE evaluations outside the main benchmark tables.Each uses forecasting pairs with the target ten time steps after the input; each dataset has 2,000 training samples and 500 test samples.Source: D Datasets · p. 17, §D.2.3 · p. 18, §D.2.4 · p. 19

D.1 Time-independent Static PDEs

• The three static benchmarks are Darcy flow, Poisson, and Helmholtz, constructed from Gaussian random fields and solved on 128 × 128 grids.Darcy permeability is formed by thresholding a Gaussian random field; Poisson and Helmholtz use that field as the coefficient.Source: §D.1 · p. 17, §D.1.1 · p. 17, §D.1.2 · p. 17

D.2.1 Incompressible Navier-Stokes Equation

• The time-dependent Navier-Stokes benchmark models 2D incompressible flow in vorticity form on a periodic unit torus and stores fields on a 128 × 128 grid.The viscosity is 10^-3, and the dynamics are simulated with a pseudo-spectral solver.Source: §D.2.1 · pp. 17–18

D.3 Data Quality Issue in the Accessed Bounded Navier-Stokes Release

• The accessed bounded Navier-Stokes dataset from FunDPS shows anomalies near a cylindrical obstacle, possibly from preprocessing or boundary handling.The authors exclude this dataset from their experiments because of the identified quality issue.Source: D Datasets · p. 17, §D.3 · p. 19

E Baseline Details

• DiffusionPDE and FunDPS baselines use guided diffusion sampling; the main FunDPS experiments use multi-resolution inference, with guidance weights selected on validation data.The authors did not reproduce FunDPS’s originally reported accuracy and state that their comparisons do not establish superiority over its original performance.Source: E Baseline Details · pp. 19–20
• FNO and GNOT are deterministic operator baselines, while FNOPE is probabilistic; unlike SBNO, these operator baselines are not trained with one unified forward/inverse objective.FNO receives masked-grid inputs, while GNOT can process partial observations directly; separate task models are trained where needed.Source: E Baseline Details · pp. 20–21

F Hyperparameter Sensitivity of FunDPS in Inference

• FunDPS is highly sensitive to guidance weight ζ across datasets and tasks; forward problems generally favor moderate guidance, while inverse problems often need stronger weights and show greater instability.Overly large weights can increase error or cause divergence; Figure 4 marks nonconvergent runs with ×.Source: F Hyperparameter Sensitivity of FunDPS… · pp. 21–22, Figure 4 · p. 21

G.1 Fixed-resolution FunDPS

• Fixed-resolution FunDPS generally performs better than multi-resolution FunDPS and sometimes outperforms SBNO, but its 4.61-second inference time is 1.7× slower than multi-resolution FunDPS and 16.5× slower than SBNO.Source: §G.1 · p. 22, Table 5 · p. 22

G.2 Sensitivity to Inference Seeds

• Across datasets and both task directions, SBNO’s seed-level variation in mean relative L2 error is small compared with the corresponding test-set error level.The authors report that the main results are not driven by a particular inference seed.Source: §G.2 · p. 22, Table 6 · p. 23

G.3 High-frequency Spectral Evaluation

• 23.8%, 8.7%, 54.3%, and 20.8% lower high-frequency relative spectral error than the next-best baseline on Darcy forward, Darcy inverse, Helmholtz forward, and Helmholtz inverse, respectively.SBNO also exhibits excess high-frequency energy in both forward tasks, so the results do not support uniform energy preservation.Source: §G.3 · pp. 23–24, Table 7 · p. 23

G.4 Mask-channel Ablation

• 61.3% and 10.4% higher mean error follow mask removal on Darcy forward and inverse reconstruction, while the mask-aware model has lower mean error in all ten settings.The ablation removes mask coordinates while retaining observed values and zero filling; without the mask, observed zeros and missing entries are not explicitly distinguished.Source: §G.4 · p. 24, Table 9 · p. 24

G.5 PDE-task Coupling Ablation

• 12.3% and 44.2% lower retained-task mean errors follow task-specific training, but excluded-task errors rise to 0.5063 for inverse and 0.5420 for forward reconstruction.In this Darcy experiment, unified training offers one checkpoint for both PDE tasks, whereas specialization improves the selected task at the expense of the other capability.Source: §G.5 · pp. 24–25, Table 10 · p. 25

G.6 Forward-drift-only SDE Evaluation

• Approximately 34% lower mean error follows specializing the forward drift to the forward-task coupling in the 60-path SDE sample-mean comparison.This evaluates a sample mean, not individual SDE samples against deterministic ODE predictions; the ODE requires both drift networks.Source: §G.6 · p. 25, Table 11 · p. 25

H Additional Experiments on Diverse Datasets

• SBNO achieves the best mean error on four of six tasks across three additional datasets, and second-best mean error on both Korteweg–de Vries tasks.DiffusionPDE has the best Korteweg–de Vries mean error; the datasets contain fewer trajectories and training pairs than the primary benchmarks.Source: H Additional Experiments on Diverse… · p. 26, Table 12 · p. 26

I Uncertainty Quantification with Stochastic Sampling

• Stochastic bridge sampling uses 60 samples per test case to report sample-mean reconstruction, uncertainty magnitude, calibration and sharpness, and error–uncertainty relations.These metrics complement deterministic probability-flow ODE inference, which is preferable when a fast, stable point reconstruction is desired.Source: I Uncertainty Quantification with… · pp. 26–27, Table 13 · p. 27
• Individual SDE predictions retain the overall pattern but appear more diffuse and less sharp than ODE reconstructions in the selected examples.These visualizations do not establish calibrated uncertainty or exact recovery of the target conditional distribution.Source: I Uncertainty Quantification with… · pp. 27–28

J Robustness to Noisy Partial Observations

• SBNO remains competitive across tested Gaussian-noise levels, with slight quality changes under low-to-moderate noise and greater degradation in the inverse problem as noise rises.It continues to compare favorably with the considered baselines throughout the evaluated range; other corruption types and non-Gaussian noise remain future-work directions.Source: J Robustness to Noisy Partial… · pp. 28–29, Table 14 · p. 28

K Sensitivity to the ODE Solver and Number of Inference Steps

• Modest gains from additional ODE steps quickly saturate relative to their inference cost, and adaptive DOPRI offers no meaningful advantage on the challenging inverse tasks.The remaining performance gap is attributed primarily to ambiguity in sparse inverse mapping rather than numerical integration accuracy.Source: K Sensitivity to the ODE Solver and… · pp. 29–30, Table 15 · p. 29

L Sensitivity to the Observation Ratio

• Performance is best at the trained 5% observation ratio and worsens as observations decrease, with sharper degradation below the training condition, especially for inverse reconstruction.The passages identify robustness across substantially different observation ratios without retraining or adaptation as a limitation.Source: L Sensitivity to the Observation Ratio · p. 30, Table 16 · p. 30

M Dynamic-ratio Training

• At 0.5%, dynamic training reduces forward/inverse errors from 0.2727/0.3362 to 0.1790/0.1710, while at 10% it reduces them from 0.1266/0.0810 to 0.0664/0.0442.At 5%, the fixed-density model remains more accurate, with errors of 0.0359/0.0269 versus 0.0836/0.0474; dynamic training was not evaluated at 15% or 20%.Source: M Dynamic-ratio Training · pp. 30–31, Table 17 · p. 30

N Closed-loop Forward/Inverse Evaluation

• In an evaluation-only closed-loop experiment, predictions from one inference direction were reused as partial observations for the opposite direction on held-out pairs with independently sampled masks, without closed-loop training.Source: N Closed-loop Forward/Inverse Evaluation · p. 31
• Darcy forward error changes only slightly under looping, while inverse error degrades more; Helmholtz shows stronger forward degradation and more moderate inverse change.Across harder settings, reused predictions can yield smoother reconstructions with less fine detail and increased error.Source: N Closed-loop Forward/Inverse Evaluation · p. 31

O.1 Reproducibility

• The authors report that the code and datasets will accompany the supplementary material and become publicly available upon publication.Source: §O.1 · p. 32

O.2 Environments

• Main experiments used four NVIDIA B200 GPUs for 200 training epochs and one B200 GPU for inference; Appendix H experiments used four A6000 GPUs for 100 epochs.Source: §O.2 · p. 32

O.3 Hyperparameters

• SBNO used 32 RK4 steps by default, while DiffusionPDE and FunDPS used original-paper settings with inference-time guidance-weight search.Hyperparameters were selected on validation data by lowest relative L2 error; Table 19 summarizes the main-experiment settings.Source: §O.3 · p. 32, Table 19 · p. 33

P Impact Statement

• The impact statement identifies potential benefits in climate and geophysical simulation, biomedical modeling, and engineering, while advising validation and governance for high-stakes use.Source: P Impact Statement · p. 33