PaperlayerView exact revision on arXiv

Source-linked AI summary

Score-based Data Assimilation

François Rozet, Gilles Louppe

arXiv:2306.10574v2cs.LGstat.ML

TL;DR

Data assimilation seeks plausible trajectories from incomplete or noisy observations, while computational costs can constrain inference for complex systems. The paper introduces score-based data assimilation, using local trajectory scores and observation guidance at inference. It reports accurate approximations of whole trajectory posteriors, with limits on scaling and on jointly fitting dynamics parameters.

  • Problem

    Data assimilation must infer plausible trajectories from imperfect observations, while computationally costly inference through physical dynamics can restrict data use in complex systems.

    Source: §1 · pp. 1–2

  • Method

    Score-based data assimilation learns local trajectory scores and uses the observation model only at inference to guide simultaneous generation of whole trajectories.

    Source: §1 · p. 2, §3.1 · p. 4, §5 · p. 9

  • Results

    The authors report accurate approximations of the whole trajectory posterior, with accuracy improving as k and the number of LMC corrections C increase.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Local score decomposition supports generating and inferring arbitrarily long Markov chains from short training segments, with parallelizable inference.

    Source: §5 · p. 9

    SDA assumes shared dynamics across trajectories and therefore cannot also fit system parameters without further development.

    Source: §5 · p. 9

Abstract

from arXiv · show

Data assimilation, in its most comprehensive form, addresses the Bayesian inverse problem of identifying plausible state trajectories that explain noisy or incomplete observations of stochastic dynamical systems. Various approaches have been proposed to solve this problem, including particle-based and variational methods. However, most algorithms depend on the transition dynamics for inference, which becomes intractable for long time horizons or for high-dimensional systems with complex dynamics, such as oceans or atmospheres. In this work, we introduce score-based data assimilation for trajectory inference. We learn a score-based generative model of state trajectories based on the key insight that the score of an arbitrarily long trajectory can be decomposed into a series of scores over short segments. After training, inference is carried out using the score model, in a non-autoregressive manner by generating all states simultaneously. Quite distinctively, we decouple the observation model from the training procedure and use it only at inference to guide the generative process, which enables a wide range of zero-shot observation scenarios. We present theoretical and empirical evidence supporting the effectiveness of our method.

1 Introduction

Data assimilation infers plausible trajectories from imperfect observations, but computationally costly methods that differentiate through transition dynamics restrict the data they can use. The paper proposes score-based assimilation that approximates whole trajectory posteriors without simulating or differentiating through the physical model.

1 Introduction

  • Data assimilation infers plausible state trajectories from imperfect, sparse, intermittent, or noisy observations of a dynamical system.The goal is to estimate the trajectory posterior given observations.Source: §1 · p. 1
  • Variational methods often seek a maximum-a-posteriori trajectory by gradient ascent, but differentiating through physical dynamics makes inference costly and limits assimilated data.Operational forecasts use only a small volume of available satellite data despite data assimilation accounting for a significant share of weather-prediction computation.Source: §1 · pp. 1–2
  • Score-based data assimilation learns from short trajectory segments, uses observations only at inference, and approximates the whole posterior without simulating or differentiating through physical dynamics.The observation-model decoupling supports zero-shot observation scenarios.Source: §1 · p. 2

2 Background

Score-based generative modeling perturbs samples with a diffusion process and uses a learned score to reverse that process. For inverse problems, the posterior score combines a prior score with a likelihood score, allowing observations to guide generation without training a separate conditional network.

2 Background

  • A continuous-time linear stochastic differential equation progressively perturbs data, with drift, diffusion, and Brownian motion determining the process.The Gaussian perturbation is chosen so the final state is approximately noise and the initial state approximately data.Source: §2 · p. 2
  • Reverse-time diffusion generates samples from noise by removing it, provided the score of the perturbed data distribution is available.A neural score network is trained with denoising score matching to approximate that score.Source: §2 · pp. 2–3
  • For inverse problems, Bayes’ rule decomposes the posterior score into a prior score and a likelihood score, avoiding the need to train a new conditional score network for each observation process.The cited likelihood approximation evaluates a Gaussian likelihood at a denoised estimate, enabling zero-shot likelihood-score computation.Source: §2 · p. 3

3 Score-based data assimilation

Score-based data assimilation combines local approximations to the trajectory prior score, an observation likelihood score, and predictor-corrector sampling. This design addresses the cost of modeling long, high-dimensional trajectories while supporting observation-guided posterior inference.

3 Score-based data assimilation

  • SDA decomposes trajectory posterior scoring into prior and likelihood terms, combining local prior scores, a new likelihood approximation, and predictor-corrector sampling.The approach targets the impractical network and data requirements of scoring long, high-dimensional trajectories directly.Source: §3 · p. 3

3.1 How is your blanket?

  • For dynamical systems, a narrow pseudo-blanket window lets a local score network train on short trajectory segments rather than the full chain.The window size k is problem-dependent and argued to be much smaller than trajectory length L; applying the local network across segments combines outputs into an approximate full score.Source: §3.1 · p. 4

3.2 Stable likelihood score

  • A perturbed-likelihood approximation addresses instability in the earlier likelihood estimate and permits computing the likelihood score without rescaling it or ignoring observation covariance.The method uses a covariance approximation based on the measurement Jacobian; the matrix term can often be simplified to a constant diagonal matrix.Source: §3.2 · p. 5

3.3 Predictor-Corrector sampling

  • Interleaving reverse-SDE steps with Langevin Monte Carlo corrections limits error accumulation from score approximations and discretization.In the limit of infinitely many sufficiently small LMC steps, samples follow the distribution defined by the approximate posterior score at each time; the authors find few steps necessary in practice.Source: §3.3 · p. 5

4 Results

The experiments evaluate SDA on Lorenz 1963 and Kolmogorov flow, using posterior comparisons for the former and primarily qualitative assessment for the latter. They report accurate and multimodal posterior inference in Lorenz, and trajectory recovery and physical consistency in Kolmogorov flow.

4.2 Kolmogorov flow

  • The Kolmogorov-flow experiment uses 64 × 64 velocity snapshots from a 256 × 256 Navier–Stokes grid, and its results are assessed primarily qualitatively because reliable assessment tools are unavailable.The fluid is two-dimensional and subject to Kolmogorov forcing.Source: §4 · p. 6, §4.2 · p. 7

4.1 Lorenz 1963

  • Across 64 observations from low- and high-frequency processes, increasing k and correction count C brings approximate posterior statistics closer to ground truth with diminishing returns; results for k ≥3 and C ≥2 are almost equivalent.The comparison used bootstrap-particle-filter ground-truth posterior samples and expected log-prior, expected log-likelihood, and trajectory-space Wasserstein distance.Source: Figure 2 · p. 6, §4.1 · p. 7
  • SDA identifies several plausible Lorenz trajectories, including four posterior modes consistent with the observations, rather than returning only a point estimate.The multimodal example uses k = 4 and C = 2.Source: Figure 3 · p. 7, §4.1 · p. 7

4.2 Kolmogorov flow

  • Under coarse, intermittent, noisy observations, SDA closely recovers the original trajectory, while the DPS likelihood-score approximation yields trajectories inconsistent with the observations.The Kolmogorov-flow observations are taken every four steps, coarsened to 8 × 8 resolution, and perturbed with moderate Gaussian noise.Source: Figure 4 · p. 8, §4.2 · p. 8
  • For an unlikely final-state vorticity observation, SDA generates a consistent trajectory whose initial state reproduces an almost identical trajectory when run through the physical model.The observation specifies uniform positive vorticity in a circle-shaped sub-domain of the final state.Source: Figure 5 · p. 8, §4.2 · p. 9

5 Conclusion

SDA uses local score models to infer arbitrarily long trajectories, with a Gaussian-likelihood approximation supporting posterior inference. Its computational limits, approximation error, scaling, and shared-dynamics assumption qualify its scope.

5 Conclusion

  • SDA uses local scores to generate arbitrarily long trajectories, while a perturbed-likelihood approximation supports inference under linear or nonlinear Gaussian observation models.The conclusion also identifies technical contributions to score-based generative modeling.Source: §5 · p. 9
  • The perturbed-likelihood approximation yields accurate posterior inference without stability tricks, and the authors note adaptations to inpainting, deblurring, super-resolution, and scientific inverse problems.The stated assumption is a linear or nonlinear Gaussian likelihood.Source: §5 · p. 9
  • Inference remains limited by the speed of simulating the reverse SDE, despite avoiding simulation or differentiation through the physical model.The authors suggest exploring faster score-model sampling methods.Source: §5 · p. 9
  • The impact of approximation error on the posterior remains unquantified, and scaling from tens of thousands of dimensions to operational systems with millions remains an open question.The authors describe such scaling as a serious engineering challenge.Source: §5 · p. 9
  • SDA assumes all trajectories share the same dynamics, so fitting system parameters alongside assimilation requires further development.With synthetic training data, accuracy is limited by the generating model; robustness to misspecification or distribution shift remains future work.Source: §5 · p. 9
  • The discussion contrasts full-trajectory inference with simpler state estimation and relates SDA to learned aids for variational methods and non-autoregressive sequence generation.It also notes that adapting cascaded video upsampling to assimilation could be difficult with arbitrary observation processes.Source: §5 · pp. 9–10

A The pseudo-blanket approximation is unbiased

The pseudo-blanket score for a state and its blanket equals the expectation of the full-trajectory score over the remaining variables, making it unbiased in expectation.

A The pseudo-blanket approximation is unbiased

  • The local score for a state and its pseudo-blanket is the expected full-trajectory score over variables outside that blanket, so it is unbiased in expectation.This holds regardless of the elements or size of the pseudo-blanket.Source: A The pseudo-blanket approximation is… · p. 15

B On the covariance of p(x | x(t))

For a Gaussian prior, the paper examines the covariance of the conditional distribution given a perturbed state and proposes an approximation based on its eigendecomposition.

B On the covariance of p(x | x(t))

  • Under a Gaussian prior with covariance Σx, the conditional distribution p(x | x(t)) has covariance Σ̂.The approximation uses the eigendecomposition of Σx, whose central diagonal term is observed to be close to Λ(Λ + I)^-1 for most perturbation times.Source: B On the covariance of p(x | x(t)) · p. 15

C Algorithms

The algorithms estimate the posterior score by adding prior and likelihood scores, then use predictor-corrector sampling to generate conditional samples.

C Algorithms

  • The posterior-score algorithm returns the sum of the prior score and the likelihood score.The supplied algorithm excerpt identifies this return step.Source: C Algorithms · p. 16
  • The predictor-corrector sampler initializes from a Gaussian noise distribution and iterates backward through the specified time steps.At each step, it updates the state and performs C correction iterations using noise and the conditional score.Source: C Algorithms · p. 16

D Experiment details

The experiments use GPU resources and train score networks for Lorenz and Kolmogorov flow with distinct architectures.

D Experiment details

  • Score networks were trained and evaluated concurrently, each on one GPU with at least 11 GB of memory.Source: D Experiment details · p. 17
  • Lorenz uses fully connected local score networks for k ≤ 4 and fully convolutional networks for k > 4, trained for 1024 epochs with AdamW and a linearly decreasing learning rate.Both architectures use residual blocks, SiLU activations, and layer normalization.Source: D Experiment details · p. 17, Table 1 · p. 17
  • Kolmogorov flow uses a local U-Net motivated by Navier–Stokes locality, trained for 1024 epochs with AdamW and a linearly decreasing learning rate.The U-Net uses residual blocks, SiLU activations, and layer normalization.Source: D Experiment details · p. 17, Table 2 · p. 17

E Assimilation examples

The examples show SDA recovering observed states while generating plausible trajectories under insufficient, sparse, and nonlinear observations; DPS comparisons are less consistent with observations in the reported cases.

E Assimilation examples

  • With insufficient center-field observations every three steps at 8 × 8 resolution, SDA identifies observed states and generates plausible states elsewhere, while DPS is less observation-consistent.The observation has small Gaussian noise, Σy = 0.012I.Source: Figure 6 · p. 18
  • For a nonlinear saturating observation of vorticity, SDA identifies the original state where data is sufficient and generates plausible states elsewhere, while DPS yields an inconsistent trajectory.Source: Figure 7 · p. 18
  • Across spatial subsampling factors n, SDA closely recovers the original trajectory, while DPS trajectories become progressively less consistent with observations as n increases.Source: Figure 8 · p. 19, E Assimilation examples · p. 19
  • The zero-observation initial-to-final difference enforces a looping trajectory, but this unrealistic setup produces physically implausible trajectories.The example uses a long trajectory with L = 127.Source: Figure 9 · p. 20
  • With spatial subsampling by a factor of 16 and medium Gaussian noise, SDA generates trajectories similar to the original with physically plausible variations.Source: Figure 11 · p. 21
Loading 2306.10574v2…