PaperlayerView exact revision on arXiv

Source-linked AI summary

A Unified Framework for Bayesian Data Assimilation with Generative Models and Observation Interpolants

Nikolaj T. Mücke, Benjamin Sanderse

arXiv:2610.03396v1cs.LGcs.CE

TL;DR

Bayesian data assimilation needs posterior samples that combine model forecasts with noisy observations, but high-dimensional, non-Gaussian posteriors are challenging. This paper conditions pretrained stochastic interpolant, flow-matching, and diffusion samplers using a shared observation-interpolant score correction and practical Gaussian approximations. In evaluations, the full inflated-covariance samplers recover the linear-Gaussian posterior closely, while field-case performance varies by variable and observation regime.

  • Problem

    Bayesian data assimilation combines model predictions with noisy observations, but existing ensemble Kalman and particle-filter approaches face Gaussian-assumption or sample-cost constraints, while deterministic surrogates require difficult-to-specify model-error distributions.

    Source: §1 · p. 1

  • Method

    The framework conditions pretrained interpolant-based models on observations through a shared likelihood-score correction, approximated in practice by a Gaussian surrogate derived from the observation interpolant.

    Source: §1 · p. 2, §4 · p. 4, §4.1 · p. 4

  • Results

    In the linear-Gaussian test, all three samplers with full inflated covariance converge to KL ≈10^-4, while field-case performance varies across samplers, variables, and observation densities.

    Source: §5.1 · p. 7, §5.2 · p. 7, §5.3 · p. 9

  • Takeaways & Limitations

    With an exact intermediate likelihood score, the conditioned stochastic and deterministic sampling dynamics target the same posterior distribution.

    Source: §1 · p. 2, §4 · p. 4

    The practical method approximates the likelihood score and assumes linear Gaussian observations; its autoregressive one-step formulation neither reweights past states nor uses future observations.

    Source: §6 · p. 10

Abstract

from arXiv · show

Bayesian data assimilation combines model forecasts with noisy observations, but sampling high-dimensional, non-Gaussian posteriors remains challenging. We introduce an observation-interpolant framework that turns pretrained stochastic interpolant, flow matching, and diffusion models into posterior samplers without retraining. Conditioning the interpolant path on observations yields a shared likelihood-score correction to the drift or velocity, unifying stochastic and deterministic posterior sampling. The resulting SDEs and ODEs sample the exact posterior when the intermediate likelihood score is known. For practical computation, we approximate this score using a closed-form Gaussian surrogate with a bias-corrected mean and covariance inflated by the model's source covariance. Jacobian-free and ensemble-shared approximations make the method tractable in high dimensions. We evaluate the framework on linear-Gaussian dynamics, stochastic two-dimensional Navier-Stokes, and urban airflow with up to $O(10^4)$ degrees of freedom.

1 INTRODUCTION

Bayesian data assimilation must represent uncertainty in high-dimensional, nonlinear settings, while existing approaches either impose Gaussian assumptions or require many samples. This paper unifies stochastic interpolants, flow matching, and diffusion models as posterior samplers, with a shared observation-based score correction and practical Gaussian and covariance approximations.

  • Ensemble Kalman filters rely on Gaussian assumptions, while particle filters typically require many more samples; deterministic neural surrogates also need difficult-to-specify model-error distributions.Source: §1 · p. 1
  • Existing posterior-sampling methods for stochastic interpolants, flow matching, and diffusion have been developed model by model, tied to stochastic or deterministic generators.Source: §1 · p. 1
  • Conditioning a pretrained interpolant on observations supplies a shared likelihood-score correction to its drift or velocity, targeting the posterior without retraining.Different diffusion choices yield stochastic or deterministic samplers while preserving the conditioned interpolant’s distribution at each pseudo-time.Source: §1 · pp. 1–2
  • A Gaussian likelihood-score surrogate uses the observation interpolant’s conditional mean and covariance, including source-noise bias correction and covariance inflation.Source: §1 · p. 2
  • Jacobian-free covariance and ensemble-shared Jacobian approximations make the method feasible in sparse-observation, field-scale settings.Source: §1 · p. 2

2 RELATED WORK

Generative posterior samplers condition pretrained flow-based models by modifying their sampling dynamics, but existing methods differ in their likelihood-score approximations. The paper extends observation interpolation into a shared framework for diffusion, flow matching, and stochastic interpolants.

3 PRELIMINARIES

The preliminaries formulate assimilation as sequential Bayesian updating under a Markov transition model and introduce interpolant-based generators for its transition prior. Flow matching, diffusion, and stochastic interpolants provide related sampling dynamics with different stochasticity choices.

  • At each Markov step, the state and observation are linked by an observation matrix, the transition prior and noise-induced likelihood define a one-step posterior, and posterior samples advance assimilation autoregressively.Source: §3 · p. 3
  • The general interpolant xτ = ατx0 + βτx1 + στ ε connects conditioning state x0 to target x1, with pseudo-time schedules and endpoint constraints specifying the path.Source: §3 · p. 3
  • Flow matching connects a Gaussian source to the transition prior and samples by integrating a learned marginal velocity; with exact velocity and well-posed dynamics, the terminal sample follows the transition prior.Source: §3 · p. 3
  • For differentiable schedules, flow matching and diffusion are equivalent representations, while the diffusion term controls sampling dynamics and setting it to zero recovers the flow-matching ODE.Source: §3 · p. 3
  • Stochastic interpolants can connect the conditioning state directly to the transition prior using Brownian source noise; their drift shares the velocity–score structure, with schedules determining the learning objective.Source: §3 · pp. 3–4

4 UNIFIED POSTERIOR SAMPLER

The unified sampler conditions a pretrained interpolant on observations to produce posterior-sampling ODEs or SDEs across stochastic interpolants, flow matching, and diffusion models. A practical Gaussian score approximation and computational shortcuts support implementation, including autoregressive assimilation.

4 UNIFIED POSTERIOR SAMPLER

  • Under the theorem’s stated regularity and well-posedness assumptions, the conditioned path yields posterior marginals and an exact terminal posterior sampler across stochastic interpolants, flow matching, and diffusion models.The framework modifies a pretrained model’s drift, velocity, or score without retraining.Source: §4 · p. 4

4.1 A CLOSED-FORM LIKELIHOOD SCORE FROM OBSERVATION INTERPOLATION

  • Interpolating observations alongside the state makes the otherwise intractable intermediate likelihood score usable by matching observations to each sampler state.The interpolated likelihood is defined as the conditional expectation of the terminal likelihood given the intermediate state and source state.Source: §4.1 · pp. 4–5, §4 · p. 4

4.2 PRACTICAL APPROXIMATIONS

  • A Gaussian surrogate uses conditionally derived moments; its bias-corrected mean and inflated covariance account for source-noise contributions and retain prior correlations.The closure is exact in the linear–Gaussian regime; retaining correlations allows observations to inform unobserved degrees of freedom, particularly under sparse sensing.Source: §4.2 · p. 5, §4.1 · p. 5
  • The ensemble-shared Jacobian reduces per-step Jacobian work to O(Ny) network passes and one factorisation, and refreshing it every k steps further reduces evaluation cost by a factor of k.Assimilation starts with the prior drift at the first step, then applies the correction; each posterior sample initializes the next forecast in the autoregressive rollout.Source: §4.2 · p. 5, §4.3 · p. 6, Figure 2 · p. 6

5 RESULTS

The experiments evaluate posterior samplers on a linear–Gaussian system, stochastic 2D Navier–Stokes, and urban airflow. Results show strong exact-posterior recovery in the analytical case, while covariance quality and observed-versus-unobserved variables distinguish performance in the field cases.

5 RESULTS

  • The evaluation uses synthetic data for three cases, with field experiments assimilated autoregressively using 64-member ensembles.For field cases, the reported metrics include RMSE, CRPS, and runtime per assimilation step; the linear–Gaussian case uses KL divergence to the exact posterior.Source: §5 · pp. 6–7

5.1 SIMPLE ANALYTICAL TEST CASE

  • KL ≈10^-4: all three samplers with full inflated covariance converge to the exact posterior, while the Jacobian-free variant yields KL ≈0.12–0.24.Only FlowDAS+SURGE among the baselines also recovers the exact posterior; its error diverges as SDE step count increases.Source: §5.1 · p. 7, Figure 3 · p. 7

5.2 STOCHASTIC NAVIER–STOKES

  • 2–4× lower SI-SDE RMSE with shared inflated covariance than with the Jacobian-free variant in sparse Navier–Stokes; shared SI-SDE also has best or tied-best CRPS in three of four scenarios.The Jacobian-free variant and Guided FM oversmooth toward the prior mean; D-Flow SGLD costs about 700 s per step, roughly an order of magnitude more than shared SI-SDE.Source: §5.2 · pp. 7–8, Table 1 · p. 8

5.3 URBAN AIRFLOW

  • Urban airflow results split by variable: DM-SDE has the lowest velocity RMSE and CRPS among the proposed samplers, while SI-SDE has the lowest temperature RMSE and CRPS at both sensor densities.With shared covariance, SI-SDE temperature RMSE improves from 0.38/0.43 to 0.36/0.42; temperature is unobserved and its errors accumulate over the long rollout.Source: §5.3 · p. 9, Figure 5 · p. 9

6 CONCLUSION

The appendix establishes how velocity and score representations yield marginal-preserving samplers, how trained networks supply the quantities they use, and why conditioning targets the posterior at the terminal time. The conclusion qualifies exactness with practical approximations and rollout limits.

6 CONCLUSION

  • Table 2 reports rollout-averaged velocity and temperature RMSE and CRPS over 4–5 held-out urban-airflow trajectories, using 64 ensemble members and 50 sampler steps.Bold identifies the best value in each column, including ties to two decimal places.Source: Table 2 · p. 9
  • The method’s exactness assumes an exact prior and likelihood score, while practical Gaussian, stopped-gradient, and Jacobian approximations can introduce bias; one-step assimilation may accumulate errors over long rollouts.The formulation also assumes linear Gaussian observations, and shared covariance grows costly with observation dimension.Source: §6 · p. 10

A VELOCITY–SCORE DUALITY AND THE MARGINAL-PRESERVING SDE

  • Under well-posedness and uniqueness assumptions, the SDE family preserves the interpolant marginals, while setting its diffusion to zero gives the probability-flow ODE.The assumptions allow the point-mass initial law used by stochastic interpolants.Source: A VELOCITY–SCORE DUALITY AND THE… · p. 14

A.1 READING THE VELOCITY AND SCORE FROM THE TRAINED NETWORK

  • Proposition A.2 converts trained model outputs into the velocity–score pair used by the samplers; flow matching directly trains velocity, while stochastic interpolants train an SDE drift.For rectified flow, the score relation degenerates as pseudo-time approaches 1, motivating an endpoint-vanishing diffusion for the DM-SDE lift.Source: §A.1 · pp. 14–15

B CONDITIONING THE INTERPOLANT PATH: PROOFS

  • Conditioning on observations changes the endpoint law to the posterior while retaining the interpolant representation and its source noise.The source noise remains independent under conditioning, so only the endpoint distribution changes.Source: §B.1 · p. 15

B.1 POSTERIOR PATH

  • The posterior path retains the interpolant form, with its endpoint drawn from the posterior distribution.Bayes’ rule and the velocity–score identity relate the conditioned path quantities to the prior path.Source: §B.1 · p. 15

B.2 PROOF OF THEOREM 4.1

  • The sampler starts from the unchanged source law and ends with the posterior law, because the conditioned interpolant equals the endpoint state at pseudo-time 1.FM/DM start from N(0, I), while SI starts from δx0.Source: §B.2 · p. 16

C PROOF OF LEMMA 4.2 AND LEMMA 4.3

The proof derives conditional moments for the interpolated observation by separating observation noise from source noise and applying interpolant Tweedie identities.

D IMPLEMENTATION DETAILS

Implementation uses fixed interpolant schedules, Euler-based pseudo-time integration, conditional U-Net priors, and shared pretrained checkpoints across assimilation samplers.

  • Stochastic-interpolant priors use quadratic schedules, while flow-matching priors use linear schedules with a noise source term.The SI schedules are ατ=1−τ, βτ=τ^2, γτ=1−τ; FM uses ατ=0, βτ=τ, στ=1−τ.Source: D IMPLEMENTATION DETAILS · p. 16
  • Sampling uses M uniform pseudo-time steps, with Euler–Maruyama for SDE samplers and forward Euler for the guided ODE.Source: D IMPLEMENTATION DETAILS · p. 17
  • Field-case priors are conditional U-Nets that receive pseudo-time and recent physical states, with architecture and channels tailored to Navier–Stokes or urban fields.Urban networks use four field channels and exclude solid cells with a building mask.Source: D IMPLEMENTATION DETAILS · p. 17
  • Priors are trained by simulation-free regression, and each case reuses the same trained FM checkpoint across FM/DM samplers and the same SI checkpoint across SI samplers without assimilation fine-tuning.Source: D IMPLEMENTATION DETAILS · p. 17

E EVALUATION METRICS

The evaluation defines ensemble forecast metrics and diagnostics, and notes that finite-ensemble estimation contributes sampling error to the reported analytical-case KL.

  • Ensemble-mean RMSE measures error of the ensemble mean rather than averaging member-level RMSE, which would additionally count ensemble spread.Source: E EVALUATION METRICS · p. 17
  • CRPS uses an unbiased ensemble estimator, with its spread-term normalization chosen to make the estimate of E|X−X′| unbiased.Source: E EVALUATION METRICS · p. 17
  • Spread–skill reports the absolute distance of the corrected spread-to-skill ratio from 1, so zero indicates perfect calibration.Source: E EVALUATION METRICS · p. 17
  • The reported Gaussian KL is generally nonzero even for exact-posterior samples because ensemble moments are estimated from a finite ensemble; more integration steps do not remove this sampling error.Source: E EVALUATION METRICS · p. 18
  • Sliced W1 averages Wasserstein distances over 128 random one-dimensional projections, while cost is wall-clock seconds per assimilation step at matched ensemble size and hardware.Source: E EVALUATION METRICS · p. 18

F METHOD DESCRIPTIONS AND COMPARISON TO ALTERNATIVES

The section compares alternative guided samplers and describes their costs, approximations, and tuning. It also gives a linear–Gaussian setup and identifies the diagnostics used to compare samples with the exact posterior.

F METHOD DESCRIPTIONS AND COMPARISON TO ALTERNATIVES

  • FlowDAS estimates likelihood scores by Monte Carlo through one-step predictions, whereas the described method uses a closed-form observation interpolant and Gaussian surrogate.FlowDAS uses a tuned scenario-dependent step size and fixed J = 25 samples; its guidance can become unstable with overly large step sizes.Source: F METHOD DESCRIPTIONS AND COMPARISON… · p. 18
  • FIG guides a deterministic flow-matching ODE with residual-based inner corrections, but approximates xτ ≈ x1 and uses no covariance.Its correction scale is tuned per scenario and sampler-step count; the number of inner steps depends on the observation operator.Source: F METHOD DESCRIPTIONS AND COMPARISON… · p. 18
  • SDA uses Tweedie-based Gaussian-likelihood guidance on a diffusion prior, while its single-window autoregressive form does not propagate cross-window correlations as a joint smoother would.Its isotropic covariance carries no model source information, and the method is tied to the native reverse diffusion.Source: F METHOD DESCRIPTIONS AND COMPARISON… · p. 18
  • D-Flow SGLD backpropagates through a full generative rollout at every chain step, costing 12 forwards and 12 backwards per step in the reported settings.Its tuned parameters are the SGLD step size and source penalty, with the step size dominant in the reported sweeps.Source: F METHOD DESCRIPTIONS AND COMPARISON… · p. 19
  • SURGE wraps FlowDAS or SDA guidance with importance weights targeting the exact terminal likelihood asymptotically as particle count grows, but particle degeneracy limits it in high dimensions.Particles resample below ESS < N/2; the implementation adds one denoiser forward per step and no backward pass.Source: F METHOD DESCRIPTIONS AND COMPARISON… · p. 19

G.1 SETUP

  • The analytical case uses a two-dimensional identity-operator linear–Gaussian system with unit Gaussian noise; closed-form SI drift and FM velocity drive the corresponding samplers.For the FM velocity, the described schedules are linear; the same velocity also drives the DM-SDE through the score identity.Source: §G.1 · pp. 19–20

G.2 ADDITIONAL RESULTS

  • The analytical comparison reports Gaussian-fit KL divergence, sliced-W1 distance to the exact posterior, and wall-clock cost, with Figure 6 visualizing the posterior and closed-form-drift samples.Table 4 reports metrics at M = 50 and M = 500; the supplied passage does not state the resulting values.Source: §G.2 · p. 20, Figure 6 · p. 20

H NAVIER–STOKES TEST CASE: SETUP AND ADDITIONAL RESULTS

The appendix details the stochastic Navier–Stokes data and assimilation setup, then presents covariance-approximation ablations and diagnostics of reconstructed fields. The diagnostics include pointwise errors and posterior-mean enstrophy spectra across observation scenarios.

H NAVIER–STOKES TEST CASE: SETUP AND ADDITIONAL RESULTS

  • The Navier–Stokes appendix reports five held-out trajectories, 64 ensemble members, 15 assimilation steps, and observation noise η ∼ N(0, 0.05^2 I).The simulation uses a 256^2 grid and is down-sampled to 128^2 for training and assimilation; retained snapshots follow a spin-up period.Source: H NAVIER–STOKES TEST CASE: SETUP AND… · p. 20, §H.1 · p. 21

H.1 FORCING AND DATA GENERATION

  • The stochastic forcing is temporally white, acts on four Fourier modes, and uses forcing amplitude ε = 1.The trajectories are generated with JAX-CFD using a pseudo-spectral discretization, viscosity ν = 10^-3, and linear drag α = 0.1.Source: §H.1 · p. 21

H.2 ABLATIONS

  • The covariance ablation compares ensemble-shared inflated covariance at Jacobian refresh cadences k ∈ {1, 5, 10} with the Jacobian-free approximation.Table 5 reports the same metrics and scenarios as Table 1, plus per-step cost, for E = 64 and M = 50.Source: §H.2 · p. 21, Table 5 · p. 30

H.3 ENSTROPHY SPECTRA

  • The qualitative figures show true vorticity and observations alongside posterior means, ensemble standard deviations, and absolute errors for super-resolution and sparse-sensor cases.The sparse-sensor examples include 5% and 1.5625% observation densities; the captions identify the final assimilation step.Source: Figure 7 · p. 22, Figure 8 · p. 23, Figure 9 · p. 24
  • Figure 14 compares posterior-mean radially averaged enstrophy spectra with the true spectrum across four observation scenarios, with zoomed panels at k ≥ 30.Curves below truth indicate over-smoothing, while curves above it indicate spurious small-scale enstrophy; spectra average over 15 assimilated steps.Source: §H.3 · p. 21, Figure 14 · p. 29

I URBAN AIRFLOW TEST CASE: SETUP AND ADDITIONAL RESULTS

The urban-airflow appendix describes a masked, sparse-sensor evaluation over 145 forecast steps and 49 posterior updates. It presents qualitative reconstructions and forecast-step trajectories for error, probabilistic score, and spread–skill deviation.

I URBAN AIRFLOW TEST CASE: SETUP AND ADDITIONAL RESULTS

  • The urban-airflow run uses one held-out trajectory, 64 ensemble members, 50 sampler steps, and sensor densities of 1.5625% and 0.78125% of fluid cells.Observations are assimilated every third forecast step for 49 updates; metrics are evaluated over fluid cells with buildings masked.Source: I URBAN AIRFLOW TEST CASE: SETUP AND… · p. 30
  • The data come from uDALES simulations of coupled momentum, temperature, and building-surface energy in a periodic 128 × 128 × 128 m^3 domain at 1 m resolution.The solver uses a staggered Cartesian grid, the Vreman subgrid-stress model, and an immersed-boundary representation of buildings.Source: I URBAN AIRFLOW TEST CASE: SETUP AND… · p. 30
  • Figures 15–18 show true fields, observations, reconstructions, uncertainty, and error for velocity magnitude and potential temperature at both sensor densities.The appendix retains SURGE versions of FlowDAS and SDA in these figures and omits DM-SDE variants.Source: I URBAN AIRFLOW TEST CASE: SETUP AND… · p. 30, Figure 15 · p. 31, Figure 16 · p. 32
  • Figures 19–21 plot fluid-cell RMSE, CRPS, and spread–skill deviation by forecast step for velocity and temperature at both sensor densities.The captions specify that observations are assimilated every third step; they do not state comparative outcomes.Source: Figure 19 · p. 35, Figure 20 · p. 35, Figure 21 · p. 36

Related papers

Loading 2610.03396v1…