PR-Smoother: Simulator-Preserving Non-Gaussian Smoothing for Data Assimilation

• The target regime combines non-Gaussian physical-space smoothing, high-dimensional scalability, observation-only training, and calibration of prescribed simulator parameters.Existing approaches often trade off simulator preservation against richer learned posteriors.Source: §1 · p. 2
• PR-Smoother learns future-conditioned corrections around the prescribed rollout, keeping the simulator explicit in both the ELBO and variational family.A conditional normalizing flow models the initial state, while transitions are Gaussian corrections centered on the simulator rollout.Source: §1 · p. 2
• The variational family contains the exact smoother under deterministic dynamics and in linear-Gaussian state-space models.The linear-Gaussian case assumes the flow can represent any smooth initial-state density and the Gaussian corrections can realize the required means and covariances.Source: §1 · p. 2
• Experiments cover multimodality in 4-dimensional Lorenz–96, ambiguous nonlinear observations and process noise in 40-dimensional Lorenz–96, and joint state–parameter–bias inference in 16,384-dimensional Kolmogorov flow.Source: §1 · p. 2

• The state-space model has prescribed, differentiable dynamics and observation operators, with unknown physical or bias parameters potentially shared across windows.Each window has its own latent trajectory, while unknown shared parameters are optimized jointly with the inference network.Source: §2 · pp. 2–3
• Data assimilation seeks filtering or smoothing posteriors over latent states from observations, with parameter estimates when applicable.Source: §2 · p. 2
• Nonlinear dynamics and observation operators can produce smoothing posteriors without closed-form solutions, including non-Gaussian or multimodal cases.Source: §2 · p. 3
• State trajectories are unavailable for training, so inference must learn from observation windows alone.Source: §2 · p. 3

• Classical 4D-Var and weak-constraint methods solve windowed MAP problems [Dimet and Talagrand, 1986, Lewis and Derber, 1985], while ensemble Kalman methods propagate Gaussian analyses [Evensen, 1994, Bocquet and Sakov, 2014].Source: §3 · p. 3
• Latent-space DA and latent SSM methods gain posterior flexibility through learned representations or auxiliary dynamics, emissions, and recognition models [Girin et al., 2021, Krishnan et al., 2015, 2017, Karl et al., 2017, Fraccaro et al., 2017, Doerr et al., 2018, Becker et al., 2019].The paper contrasts these with direct physical-space smoothing that retains prescribed operators.Source: §3 · p. 3
• Learned-prior methods train trajectory priors or transition models on simulated trajectories, whereas PR-Smoother keeps the simulator as the generative model and places flexibility in the posterior [Cranmer et al., 2020].Source: §3 · p. 3
• PR-Smoother uses a forward, future-conditioned factorization that pairs an expressive initial-state flow with transitions centered on the physical rollout, unlike reverse-time smoothing decompositions [Campbell et al., 2021, Chagneux et al., 2024].It trains from observation windows through a prescribed-simulator ELBO.Source: §3 · p. 4

4 Method

• PR-Smoother models the initial-state posterior with a conditional normalizing flow and each transition with future-conditioned Gaussian corrections centered on the prescribed rollout.Source: §4 · p. 4

4.1 Objective

• The observation-only ELBO jointly trains the amortized smoother and point estimates of unknown shared physical and sensor parameters.Source: §4.1 · p. 4

4.2 Simulator-Preserving Smoother Construction and ELBO

• The construction uses nested simulator-preserving families: deterministic rollout transitions or rollout-centered transitions with future-conditioned Gaussian corrections.Source: §4.2 · p. 4

4.2.1 Initial-Time Posterior qϕ(X1 |Y1:T )

• A conditional normalizing flow maps a standard-normal latent variable to the initial-state posterior, with observation context summarizing the full window.RealNVP provides direct access to log qϕ for the ELBO.Source: §4.2.1 · p. 5

4.2.2 Deterministic Dynamics Model

• In the deterministic limit, the trajectory posterior is the pushforward of the initial-state posterior through the prescribed dynamics, and transition terms cancel from the ELBO.Source: §4.2.2 · p. 5

4.2.3 Noisy Dynamics Model

• With process noise, future-conditioned Gaussian corrections augment each rollout-centered transition; setting their mean to zero and covariance to Q recovers the deterministic family.The generative transition remains unchanged, and the per-step Gaussian KL terms are closed-form.Source: §4.2.3 · pp. 5–6

4.3 Consistency Checks in Tractable Regimes

• The family contains the exact smoothing posterior for deterministic dynamics and for linear-Gaussian state-space models under the stated flow and correction-capacity assumptions.In the latter case, the Gaussian transitions can reproduce the Kalman-RTS smoother conditionals.Source: §4.3 · p. 6

5 Experiments

• Training uses observation windows alone, with offline costs ranging from 1 GPU-hour for 4D Lorenz-96 to 19–20 GPU-hours for Kolmogorov flow.The 40D Lorenz-96 experiment takes 11–18 GPU-hours offline.Source: §5 · p. 6

5.1 Multimodality Test: 4-Dimensional Lorenz-96

• The flow reproduces the PF’s one-step bimodality and ten-step concentration as accumulated observations and dynamics resolve the sign ambiguity.The one-step ambiguity arises because componentwise-square observations cannot distinguish x from −x.Source: Figure 2 · p. 7, §5.1 · p. 7

5.2 Nonlinear Observation Operator and Dynamics Noise Test: 40-Dimensional Lorenz-96

• Under saturating many-to-one observations, PR-Smoother performs well and Noisyflow has the lowest RMSE across settings; classical methods perform best with linear observations.The 40D experiment also finds that noisy-transition variants improve on noiseless counterparts even with deterministic data.Source: §5.2 · pp. 7–8, Figure 3 · p. 8

5.3 High State Dimension and Model Parameter Training Test: Kolmogorov Flow

• Under “full” observations, PR-Smoother’s state RMSE was well below observation noise and matched oracle 4D-Var (true parameters and bias field; 10 restarts); IEnKS (16 members) was less competitive.Source: §5.3 · p. 8
• With half coverage, neither family reliably recovers the unobserved half; sparse sensing leaves log Re biased, while mixed coverage supports state and parameter performance.Sparse point measurements do not resolve the small-scale structures controlled by Reynolds number.Source: §5.3 · pp. 8–9
• In realistic coverage, deterministic variants match or outperform noisy-transition variants at matched wall-clock budget, and dense regions identify F and log Re.The sparse network supplies global coverage, allowing calibrated dynamics to propagate information across the domain.Source: §5.3 · p. 9

6 Practical Scope and Limitations

• Undertraining or distribution shift can reduce amortized inference accuracy, although the method remains competitive with oracle 4D-Var under a fixed compute budget with N = 288 Kolmogorov-flow windows.Most gains occur by N ≃105–106.Source: §6 · p. 9
• Offline training costs are amortized across repeated windows from the same system, and adjoint routines can support gradient-based training for legacy solvers while retaining their forward models.The approach targets prescribed differentiable dynamics and observation models.Source: §6 · p. 9
• For applications needing a point estimate, PR-Smoother could initialize per-window refinement with 4D-Var.Source: §6 · p. 10
• Reported state posteriors condition on point estimates of shared physical and sensor parameters, leaving parameter uncertainty for a proposed global variational factor qψ(θ).Observation coverage matters for parameter identifiability; changed parameters or observation layouts may require adaptation, with held-out predictive checks supporting deployment assessment.Source: §6 · p. 10
• Evaluation on operational numerical weather prediction systems remains an important next step.Source: §6 · p. 10

7 Conclusion

• Across Lorenz–96 and 16,384-dimensional Kolmogorov flow, PR-Smoother captures multimodal posteriors and scales to high-dimensional nonlinear assimilation.Scientific deployment still requires domain-specific validation and calibration.Source: §7 · p. 10

A On the Variational Family Including the True Posterior in Two Limiting Cases

• Under the stated expressivity assumptions, the variational family contains the true smoother for deterministic dynamics and linear-Gaussian state-space models.The conditional flow must represent smooth initial-state posteriors, and the shock head must represent Gaussian means and covariances.Source: A On the Variational Family Including… · p. 15

A.1 Deterministic Dynamics

• In the deterministic limit, the initial-state posterior is the only nontrivial posterior degree of freedom; subsequent states follow the prescribed rollout.The conditional flow can represent that initial-state posterior exactly.Source: §A.1 · p. 16

A.2 Linear-Gaussian SSMs

• In linear-Gaussian SSMs, exact smoothing conditionals are Gaussian, matching the family’s Gaussian transition form and Gaussian initial-state posterior.The setting also admits Kalman filtering and Rauch–Tung–Striebel smoothing.Source: §A.2 · p. 16

A.2.1 Gaussian Backward Messages

• Gaussian backward messages remain Gaussian because each recursion multiplies Gaussian factors and integrates out the next state.The proof establishes this by induction from the constant terminal message.Source: §A.2.1 · pp. 16–17

A.2.2 The Conditional p(Xt | Xt−1, Yt:T ) is Gaussian

• The exact future-conditioned transition is Gaussian, and its mean depends affinely on the previous state, matching the variational transition family.The shock head can reproduce the exact conditional’s mean and covariance.Source: §A.2.2 · p. 17

A.2.3 Putting the Pieces Together

• Combining the Gaussian initial-state posterior with exact future-conditioned transitions makes the linear-Gaussian smoother posterior-inclusive; the shock head can reproduce RTS backward messages.The variational factorization uses the same future-conditioned structure as the exact smoother.Source: §A.2.3 · p. 18

B Algorithm

• The training algorithm samples observation-window minibatches, estimates the ELBO, and updates the model and smoother parameters by gradient ascent.It repeats the minibatch updates until convergence.Source: B Algorithm · p. 18

C Architecture for the Anti-Causal Encoder

• Reversing the observations before causal encoding yields per-step contexts that depend only on current and future observations and condition Gaussian corrections around the simulator rollout.The shock head uses the previous state, context, and normalized time; transitions and likelihoods still use the prescribed simulator and observation operator.Source: C Architecture for the Anti-Causal… · pp. 18–19

D Use of NF over DM or FM

• Diffusion and flow-matching alternatives require ODE-based density evaluation, with diffusion adding divergence integration and flow matching also requiring ODE sampling.The passage characterizes these computations as substantially heavier than a single normalizing-flow evaluation unless very small-NFE integrators are used.Source: D Use of NF over DM or FM · p. 19

E Full ELBO Tables

• Tables 3–4 report test-set ELBOs, with higher values better; estimates use one sample and five seeds, and magnitudes are comparable only at matched window lengths and observation patterns.Each estimate is for a full test sequence and is reported as mean ± standard deviation.Source: Table 3 · p. 19, Table 4 · p. 19, E Full ELBO Tables · p. 19

Model and training

• Both noisy-transition variants use isotropic generative process noise and learn σQ and the variational scale σS jointly by maximizing the ELBO.The variational transition uses a mean correction around f(xt) and an isotropic covariance shared across time.Source: F Trained σQ Values in Noisyflow and… · pp. 19–20

Interpretation

• In deterministic regimes, σQ is not strictly identifiable and primarily controls per-step KL regularization, so it is an effective model-error scale rather than a physical noise estimate.The deterministic data-generating setting has σTrueQ = 0.Source: F Trained σQ Values in Noisyflow and… · p. 20

Trained σQ Values

• ln σQ = −2.29 ± 0.04 for Noisyflow, versus −4.61 ± 0.16 for NoisyGauss, against the planted ln σQ = −2.30 in noisy dynamics.The passage says this diagnostic complements state-RMSE results and indicates more accurate process-noise calibration for Noisyflow.Source: F Trained σQ Values in Noisyflow and… · p. 20

Parameter estimation errors

• Table 6 reports |∆F|, |∆log Re|, and observation-bias errors for Full and Half observation patterns.Lower error is better.Source: Table 6 · p. 20

Reporting scope

• Results are marked N/A when none of five initial seeds succeeds in training or optimization; oracle models are excluded because their F, log Re, and bias are prescribed.The tables therefore report estimated-parameter errors for non-oracle models only.Source: Table 6 · p. 20

Compute-budgeted data scaling

• At N = 288 windows, Flow training remains stable, recovers coherent vortical structure, and keeps state RMSE below the oracle 4D-Var reference.The qualitative reconstructions are shown in Figure 5, and the state-RMSE comparison is in Figure 6(a).Source: Figure 5 · p. 21, Figure 6 · p. 21, H Compute-budgeted data scaling · p. 22
• With gradient updates held fixed across dataset sizes, performance is close to the main experiment at about 10^5 windows, with the most noticeable state-RMSE increase below about 10^4.Smaller datasets receive more epochs; metrics use a held-out test set.Source: H Compute-budgeted data scaling · pp. 20, 22, Figure 6 · p. 21

Parameter estimation

• Forcing amplitude F is robustly estimated across dataset sizes and bias-field RMSE varies modestly, while log Re shows weaker sensitivity and mild non-monotonicity.The passage identifies the corresponding comparisons as Figure 6(c), 6(e), and 6(d), respectively.Source: Figure 6 · p. 21, H Compute-budgeted data scaling · p. 22

I Experimental Settings

• Across experiments, training uses Adam and a warm-up-then-exponential-decay learning-rate schedule; noisyFlow and noisyGauss use learned isotropic noise scales and nested residual corrections.The learning rate warms from 10^-3 of its maximum for 1% of Lorenz-96 iterations or 5% of Kolmogorov iterations, then decays to 0.1 of maximum; setting the correction to zero and residual covariance equal to Q recovers the corresponding flow/Gauss model.Source: I Experimental Settings · p. 22

I.1 Lorenz-96, 4D

• The 4D Lorenz-96 multimodality test uses deterministic dynamics, squared noisy observations, and T = 1 or T = 10 windows, with a 5 × 10^6-particle filter as the evaluation reference.The initial state comes directly from a Gaussian prior with coordinate mean 0 and standard deviation 2; the test compares flow, diagonal-Gaussian, and mean-field variational families.Source: §I.1 · p. 23

I.2 Lorenz-96, 40D

• The 40D Lorenz-96 experiments use noisy dynamics and either linear or nonlinear observations, with seven smoother and baseline models evaluated on spin-up-discarded trajectories.Dynamics noise has variance 0.12 per observation interval, observation noise variance is 1.02, and the nonlinear operator is h(Xt) = min(Xt^2, 5).Source: §I.2 · p. 25

I.3 Kolmogorov Flow

• Kolmogorov flow uses 128 × 128 grids, 10-step windows, four observation-coverage patterns, unknown sensor biases, and joint learning of state, forcing, and Reynolds number.The dataset contains approximately 2.9 × 10^6 windows; true F = 1.0 and Re = 1000, while algorithms initialize at F = 0.1 and Re = 100.Source: §I.3 · pp. 27–28
• Kolmogorov flow uses a spectral Gaussian prior on the initial vorticity field, with precision diagonal in the 2D discrete Fourier basis.The experiment uses a periodic 128 × 128 grid.Source: §I.3 · pp. 27–28
• Kolmogorov variational baselines vary initial-state posteriors and stochastic corrections, spanning flow/noisyflow, Gauss/noisyGauss, and a factorized, dynamics-agnostic trajectory model.Flow uses a conditional normalizing flow and deterministic rollout; noisyflow adds Gaussian corrections. Gauss uses a diagonal-Gaussian initial posterior, noisyGauss combines that posterior with noisy transitions, and mean-field factorizes the trajectory posterior.Source: §I.3 · pp. 28–29
• The 4D-Var baseline jointly estimates the 16,384-dimensional flow state, sensor biases, and physical parameters; its main Kolmogorov results use strong constraints, with a weak-constraint option optimizing the trajectory.An oracle variant fixes the true parameters and biases and optimizes only the flow state.Source: §I.3 · p. 29
• Kolmogorov IEnKS uses 16 ensemble members and up to 100 iterations, with localization radius selected by state RMSE and an ensemble-size check spanning 16, 32, and 48 members.It augments the state with biases and physical parameters outside the oracle setting; runtime growth motivates using 16 members.Source: §I.3 · p. 30, Figure 7 · p. 30, Figure 8 · p. 30

I.4 Runtime and Amortization

• Training takes about 1 GPU-hour for 4D Lorenz-96, 11–18 GPU-hours for 40D Lorenz-96, and 19–20 GPU-hours for Kolmogorov flow.The costs are paid once per trained model and can be amortized across future assimilation windows.Source: §I.4 · p. 31
• After one-time offline training, PR-Smoother’s per-window cost stays fixed as windows accumulate, whereas 4D-Var and IEnKS require a new solve for each window.Inference uses a single forward pass of the encoder and variational networks plus a physical-model rollout.Source: §I.4 · p. 31

J Inference Examples in the Lorenz-96, 4D Experiment

• Histograms show Lorenz-96 4D inference results for the PF, flow, and gauss models.The passages identify the plotted models but do not report histogram outcomes.Source: J Inference Examples in the Lorenz-96,… · p. 31, Figure 9 · p. 32, Figure 10 · p. 33
• Figures 12, 13, and 14 provide histograms for the mean-field, EnKF, and ETKF models, respectively; the cited captions identify each model but do not report comparative results.Source: Figure 12 · p. 35, Figure 13 · p. 36, Figure 14 · p. 37

NeurIPS Paper Checklist

• The checklist identifies differentiability of (f, h), amortization gaps from undertraining or distribution shift, sparse-observation identifiability, and point-estimate treatment of θ as limitations.Source: §2 · p. 38
• Reproduction materials specify experimental setups, architectures, baselines, hyperparameters, hardware, runtimes, and supplementary code with scripts and commands.Source: §2 · p. 38
• The checklist identifies improved uncertainty-aware state estimation and parameter calibration as potential benefits, and over-trust without domain-specific validation and calibration as a potential risk.Source: §2 · p. 39