Skip to content

Sparsely-Supervised Data Assimilation via Physics-Informed Schrödinger Bridge

Conference: NeurIPS 2026
arXiv: 2603.22319
OpenReview: https://openreview.net/forum?id=6WETopIg2t
Area: Physics
Keywords: Data Assimilation, Schrödinger Bridge, Physics-Informed Neural Networks, Multi-Fidelity Modeling, Partial Differential Equations

TL;DR

Addressing the bottleneck of unavailable full-field high-fidelity training supervision in partial differential equation systems, Physics-Informed Conditional Schrödinger Bridge (PICSB) introduces periodic surrogate-endpoint refresh and hard observation projection throughout sampling, training the transport velocity purely via physical PDE residuals to achieve sub-second zero-shot full-field assimilation.

Background & Motivation

In operational numerical weather prediction, ocean dynamics modeling, and real-time hydrological forecasting, reconstructing continuous high-resolution spatiotemporal fields from sparse sensor measurements—a process formalised as data assimilation (DA)—is a fundamental computational requirement. Direct high-fidelity (HF) numerical simulation of governing partial differential equations (PDEs) incurs prohibitive computational latency, rendering it unsuitable for time-critical, rapid-update operational workflows. In response, practitioners and researchers widely rely on low-fidelity (LF) physical models, coarse spatial discretizations, or sparse observation interpolations as provisional priors. However, such LF representations suffer from systematic physical biases, posing a fundamental dilemma: how to effectively eliminate these biases and reconstruct physically faithful full-field HF states when high-accuracy observational data is restricted to sparse spatial locations.

Existing algorithmic paradigms face acute trade-offs between computational latency and supervision requirements. On one hand, per-instance test-time optimization methods, exemplified by Physics-Informed Neural Networks (PINNs) and classical variational adjoint formulations, must repeatedly optimize network parameters from scratch for every newly encountered observation instance by trading off observation fitting and PDE residual penalties. This optimization loop typically consumes hundreds to thousands of seconds (or even hours) per sample, entirely preventing deployment in real-time operational assimilation cycles. On the other hand, recent amortized generative approaches based on conditional diffusion models and flow matching achieve faster zero-shot inference, but they overwhelmingly assume access to fully observed, dense HF fields during offline training—a luxury rarely available in real-world physical and experimental systems. Furthermore, standard conditional diffusion baselines require hundreds to thousands of iterative denoising steps along with gradient-based guidance at inference time to enforce physical constraints, leading to minute-level latency and frequent gradient noise amplification.

Schrödinger Bridge (SB) formulations establish an optimal stochastic transport between arbitrary endpoint distributions, offering an intrinsically favorable geometry to warp an informative LF prior directly into the desired HF posterior with far fewer sampling steps than noise-initialized diffusion models. Nevertheless, conventional SB matching strictly mandates paired full-field endpoints to specify reference Brownian bridges, directly clashing with the sparse-supervision reality of PDE inverse problems. The core idea of this paper is to formulate sparsely supervised data assimilation as a conditional Schrödinger bridge that transports a low-fidelity prior toward an observation-conditioned high-fidelity posterior, utilizing periodic refresh of a static surrogate model to establish reference bridge states, enforcing sparse observations via hard projection across all sampling steps, and updating the trainable velocity network exclusively through PDE residual loss to achieve guidance-free, amortized sub-second field reconstruction.

Method

Overall Architecture

PICSB formulates multi-fidelity field reconstruction as an amortized conditional transport map: given a full-grid low-fidelity field \(X_0 \in \mathbb{R}^n\) and sparse high-fidelity observations \(Y = \mathcal{H}(X_1) \in \mathbb{R}^m\) (\(m \ll n\)), the model integrates a controlled diffusion process over bridge time \(\tau \in [0, 1]\) to generate an observation-consistent and physically plausible HF reconstruction \(\hat{x}_1\). Because ground-truth full-field endpoints \(X_1\) are absent during training, PICSB deploys a dual-network configuration: a slowly updated static velocity network \(v_{\tilde{\theta}}\) acts as a reference generator to sample surrogate terminal endpoints and instantiate intermediate Brownian bridge states; concurrently, a trainable velocity network \(v_\theta\) drives trajectories from these intermediate states to terminal predictions, optimized directly by back-propagating discrete PDE residual energy.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input Data<br/>Full-grid LF field X0 + Sparse observations Y with mask M"] --> B["Surrogate Endpoint Periodic Refresh<br/>Static network generates surrogate endpoints and samples Brownian bridge"]
    B --> C["Hard Observation Projection Throughout Sampling<br/>Coordinate overwriting enforces constraints and eliminates guidance"]
    C --> D["Guidance-Free Training via Physics Residuals<br/>Discrete PDE energy loss back-propagates to update velocity network"]
    D --> E["Output Physically Consistent HF Full Field x1"]

Key Designs

1. Surrogate Endpoint Periodic Refresh: Self-consistent bridge construction without dense supervision

Standard Schrödinger bridge matching frameworks (such as Diffusion Schrödinger Bridge Matching) rely essentially on pre-paired endpoint distributions \((x_0, x_1)\) to construct closed-form Brownian bridge conditional drifts. Under sparse HF supervision, the complete target field \(x_1\) is unavailable, breaking standard bridge-matching objectives. To overcome this limitation, PICSB draws inspiration from the target-network stabilization paradigm in deep reinforcement learning and introduces a periodic refresh mechanism with a slowly updated static network \(v_{\tilde{\theta}}\). Within each refresh round \(k\), the static model parameters are updated with the current trainable weights (\(\tilde{\theta} \leftarrow \theta\)). For each training instance \((x_0, y)\), the static sampler generates a surrogate HF candidate \(\tilde{x}_1 \sim \tilde{\pi}_1^{(k)}(\cdot|y)\). An intermediate bridge state is then analytically drawn via the Brownian bridge closed form: $\(X_\tau \mid (x_0, \tilde{x}_1) \sim \mathcal{N}\big((1-\tau)x_0 + \tau \tilde{x}_1, \; \varepsilon \tau (1-\tau) I\big)\)$ at a uniform random time \(\tau \sim \text{Unif}(0, 1)\), serving as the starting state for the trainable sampler \(v_\theta\). Theoretical analysis (Proposition 4.2) guarantees that at a fixed point of this alternating process, the surrogate distribution is self-consistent and minimizes the expected PDE residual under the fixed reference measure.

2. Hard Observation Projection Throughout Sampling: Eliminating inference guidance and condition drift

In conventional conditional diffusion models, steering an unconditioned or biased prior to respect sparse measurements requires computing test-time gradients of measurement fidelity (\(\nabla_{x_t} \|y - \mathcal{H}(x_t)\|^2\)) at every diffusion step, substantially inflating computational latency and inducing potential manifold collapse under severe noise. PICSB adopts an exact projection geometry by defining the feasible observation subspace \(\mathcal{C}_y = \{x \in \mathbb{R}^n : \mathcal{H}(x) = y\}\). For coordinate subsampling operators, it applies an exact coordinate overwriting projection: $\(\Pi_{\mathcal{C}_y}(x) := \mathbf{M} \odot y + (1 - \mathbf{M}) \odot x\)$ where \(\mathbf{M} \in \{0, 1\}^n\) is a binary sensor location mask. Hard projection is strictly applied after every drift update during surrogate endpoint generation, intermediate Brownian bridge sampling, and both training and test-time Euler integration steps. By forcing intermediate states to lie exactly within \(\mathcal{C}_y\) prior to subsequent velocity evaluations, PICSB completely removes the need for inference-time observation guidance gradients, enabling fully autonomous, fast feedforward sampling.

3. Guidance-Free Training via Physics Residuals: Direct energy minimization over generative transport

Because full-field ground-truth fields are absent, conventional mean-squared error supervision cannot be applied. PICSB instead leverages the governing physical laws as an implicit energy landscape, employing a discretized PDE residual operator \(\mathcal{R}_h: \mathbb{R}^n \to \mathbb{R}^r\) to supervise velocity learning. Starting from an intermediate projected bridge state \(\tilde{x}_\tau\), the trainable velocity network integrates the remaining bridge trajectory to produce a candidate full field \(\hat{x}_1 = \text{Sample}_\theta(\tilde{x}_\tau \mid y, \tau)\). The training loss is defined strictly as the expected norm of the PDE residual: $\(\mathcal{L}(\theta) := \mathbb{E}_{(x_0, y) \sim \mathcal{D}_{\text{train}}} \mathbb{E}_{\tau \sim \text{Unif}(0, 1)} \big[ \|\mathcal{R}_h(\hat{x}_1)\| \big]\)$ Back-propagating this loss end-to-end through the multi-step sampling trajectory penalizes non-physical artifacts (e.g., blocky interpolation steps or numerical vortex dissipation). Under localized PDE identifiability (Proposition 4.1), minimizing this discrete residual under hard observation constraints guarantees an explicit error bound against the true physical field: \(\|\hat{x}_1 - x_1^\star\| \le c(\delta - \delta^\star(y; h)) + \eta(h)\), theoretically validating the consistency of purely residual-driven bridge training.

Loss & Training

The velocity network backbone is an AdaIN-conditioned U-Net adapted from unpaired Schrödinger bridge architectures. Diffusion time \(\tau_t\) is encoded via a 2-layer MLP (hidden dimension 16) into time embeddings that modulate U-Net feature maps through channel-wise Adaptive Instance Normalization, maintaining full spatial convolutionality. - Discretization & Solver: Sampling uses a fixed uniform discretization of \(T=10\) bridge-time steps (step size \(s=0.1\)) with constant noise scale \(\varepsilon = 10^{-2}\) across both training and inference. - Refresh Schedule: The static network is refreshed every \(C=100\) iterations (\(\tilde{\theta} \leftarrow \theta\)) over a total of 100k training iterations. - Optimization & Gradient Clipping: Optimization uses Adam. For 1D Burgers and 2D Kolmogorov benchmarks, learning rates are set to \(10^{-3}\) with gradient norm clipping at 1.0. For the 2D Darcy benchmark, governed by an elliptic operator and spatially heterogeneous permeability, the learning rate is scaled down to \(10^{-6}\) with strict gradient clipping at \(10^{-5}\) to ensure numerical stability.

Key Experimental Results

Main Results

Evaluation was conducted on three canonical fluid mechanics benchmarks: 1D viscous Burgers' equation, 2D steady-state Darcy flow, and 2D turbulent Kolmogorov flow (incompressible Navier-Stokes in vorticity form). All methods were evaluated under 10% sparse observations across three sensor regimes: resampled random sensors (R1), trajectory-fixed random sensors (R2), and a globally fixed sensor network (R3). Baselines comprise test-time re-optimization via PINNs and inference-time guidance over an LF-pretrained diffusion prior (Guidance). Accuracy is quantified by relative \(L^2\) reconstruction error (RelError %).

Dataset Regime PINNs (Optimization) Guidance (Diffusion Guided) PICSB (Ours) Relative Advantage
1D Burgers R1 (Resampled) 8.29% 0.99% 2.59% Outperforms PINNs by 68.8%
1D Burgers R2 (Sample-fixed) 13.48% 19.04% 14.05% Outperforms Guidance by 26.2%
1D Burgers R3 (Globally fixed) 11.79% 10.66% 14.25% Competitive with Guidance
2D Darcy R1 (Random) 5.69% 12.02% 0.29% Error reduced by 97.6%
2D Darcy R3 (Globally fixed) 5.64% 27.39% 0.28% Error reduced by 95.0%
2D Kolmogorov R1 (Resampled) 64.47% 8.99% 7.32% Clear best performance
2D Kolmogorov R2 (Sample-fixed) 65.20% 10.62% 8.66% Clear best performance
2D Kolmogorov R3 (Globally fixed) 65.06% 10.61% 8.37% Clear best performance

Inference Efficiency & Ablation/Robustness Study

Computational latency represents the primary practical hurdle in operational data assimilation. The table below reports the average inference wall-clock time per sample under regime R3, alongside key ablation and robustness findings.

Configuration / Method Burgers Wall-time (s) Darcy Wall-time (s) Kolmogorov Wall-time (s) Structural Characteristic & Note
PINNs (Test-time Optimization) 1733 s (~29 min) 2682 s (~45 min) 10680 s (~3.0 h) 100k optimization steps; converges to over-smoothed states
Guidance (Diffusion Guidance) 174 s (~3 min) 173 s (~3 min) 801 s (~13 min) 2000 diffusion steps with online gradient back-propagation
PICSB (Full Model, T=10) 0.023 s 0.023 s 0.134 s Amortized zero-shot inference, 6000x - 80000x speedup
PICSB (w/o Surrogate Refresh) Diverged / Failed Diverged / Failed Diverged / Failed Lacking moving reference causes bridge collapse
PICSB (w/o Hard Projection, soft loss) 4.82% (R1) 1.15% (R1) 12.45% (R1) Inexact observation fit leads to substantial error spikes
PICSB (Compressed to T=5 steps) 2.84% (R1) 0.35% (R1) 8.10% (R1) 2x faster (~0.01s), retaining competitive accuracy

Key Findings

  • Decisive gains on complex flows and spatial artifact removal: In 2D Darcy flow, the LF prior is constructed from nearest-neighbor interpolation, introducing severe piecewise-constant blocky artifacts. Guidance diffusion remains trapped within this coarse manifold (error reaching 27.39%), whereas PICSB completely erases interpolation artifacts via residual training, achieving a near-zero error of 0.28% (a 20x improvement). In 2D turbulent Kolmogorov flow, PINNs collapse to overly smooth solutions missing critical vortices (error >64%), while PICSB faithfully recovers fine-scale turbulent eddies.
  • Sub-second assimilation resolves real-time latency bottlenecks: For high-dimensional Kolmogorov dynamics across 40 time frames, PINNs require 3 hours and Guidance requires 13 minutes, while PICSB finishes full-field spatiotemporal assimilation in 0.134 seconds, unlocking feasibility for operational forecast integration.
  • Robustness against observation noise: When sparse observations are corrupted with up to 15% Gaussian noise, Guidance suffers from severe error amplification caused by back-propagating noisy gradients, whereas PICSB demonstrates graceful, gradual error degradation due to the stabilizing effect of hard projection within the transport geometry.

Highlights & Insights

  • Closing the loop on unsupervised Schrödinger bridges: While Schrödinger bridges traditionally require paired full endpoints and PINNs suffer from sluggish optimization, PICSB uses delayed target-network refreshes to define self-consistent surrogate endpoints, steering the transport manifold using physics residual energy while retaining fast 10-step bridge convergence.
  • Hard projection over gradient guidance: Replacing complex test-time gradient guidance with exact coordinate overwriting (\(\Pi_{\mathcal{C}_y}\)) eliminates the heavy computational footprint of reverse-mode automatic differentiation during sampling, while guaranteeing exact matching at sensor locations.
  • Transferable multi-sparsity adaptation: Additional experiments show that under extreme sensor sparsity (e.g., dropping from 10% to 1%), an integrated finetuning strategy across nested sparsity subsets enables simultaneous adaptation across heterogeneous observational densities.

Limitations & Future Work

  • Reliance on known analytical PDE operators: PICSB requires explicit mathematical forms for the discrete residual operator \(\mathcal{R}_h\); extending it to complex unparameterized or partially known systems (e.g., subgrid turbulence parameterizations) remains an open challenge.
  • Hard projection vulnerability under heavy sensor noise: Enforcing exact coordinate overwriting assumes reliable sensor readings; under extreme measurement noise, hard projection may inject unphysical local shocks, pointing to the need for noise-aware soft projection or confidence gating.
  • Dependency on informative LF priors: Under extreme observation sparsity where LF proxies provide degraded structural hints, reconstruction accuracy degrades; incorporating self-supervised neural operator backbones could further strengthen representation capacity.
  • vs PINNs (Raissi et al., 2019): PINNs optimize an implicit coordinate MLP per instance over thousands of iterations, making them computationally impractical for real-time applications and prone to over-smoothing; PICSB amortizes computation into offline training, performing zero-shot sub-second inference with superior accuracy on turbulent Kolmogorov flows (8.37% vs 65.06%).
  • vs DiffusionPDE / Guidance-based Diffusion (Huang et al., 2024): DiffusionPDE relies on thousands of diffusion steps and test-time gradient guidance; PICSB replaces noise-based diffusion with LF-to-HF Schrödinger transport requiring only 10 steps and no inference guidance, accelerating sampling by 3-4 orders of magnitude.
  • vs I2SB / Conditional Schrödinger Bridges (Liu et al., 2023; Li et al., 2025a): Existing SB formulations in vision and scientific modeling require paired full-field targets; PICSB introduces the periodic surrogate-endpoint refresh scheme, pioneering Schrödinger bridge training under sparse localized supervision.

Rating

  • Novelty: ⭐⭐⭐⭐⭐ Seamlessly marries physics-informed residuals with endpoint-free Schrödinger bridges via target-network-style surrogate refreshes.
  • Experimental Thoroughness: ⭐⭐⭐⭐⭐ Comprehensive evaluations across 1D Burgers, 2D Darcy, and 2D Navier-Stokes flow with rigorous sensor configurations, noise testing, and sparsity finetuning.
  • Writing Quality: ⭐⭐⭐⭐⭐ Strong theoretical grounding with mathematical proofs for identifiability error bounds and self-consistent fixed points.
  • Value: ⭐⭐⭐⭐⭐ Demonstrates sub-second zero-shot generative assimilation, offering immense practical value for real-time numerical weather prediction and physical fluid systems.