PR-Smoother: Simulator-Preserving Non-Gaussian Smoothing for Data Assimilation¶
Conference: NeurIPS2026
arXiv: 2609.26890
Code: https://github.com/Yuta-Tarumi/PRSmoother_Neurips2026
Area: Physics / Scientific Computing
Keywords: data assimilation, non-Gaussian smoothing, normalizing flow, physical simulator, joint parameter estimation
TL;DR¶
PR-Smoother learns a non-Gaussian initial-state posterior and future-conditioned stepwise corrections in physical space while retaining the prescribed simulator, using an observation-only variational objective to jointly estimate states, physical parameters, and sensor biases, with validation on nonlinear Lorenz–96 and a 16,384-dimensional fluid system.
Background & Motivation¶
Data assimilation is not ordinary time-series forecasting: physical equations and observation operators are already available, and the task is to infer hidden physical states from partial, noisy measurements, potentially while calibrating equation parameters. This paper studies windowed smoothing, which uses all observations in a window to estimate a trajectory posterior, rather than filtering with observations available only up to the current time. 4D-Var retains the equations but primarily produces a windowed maximum a posteriori estimate; ensemble Kalman methods also retain the physical model, but their updates rely on local Gaussian approximations. When measurements map positive and negative states to the same value or saturate at large magnitudes, these representations can lose genuine multimodality.
Flexible generative models offer another route, but learning replacement dynamics, latent transitions, or a trajectory prior changes the scientific problem: the calibrated quantities need not remain those of the original physical equations. When training data contain only observation windows, without true state labels, generating supervised trajectories using some choice of unknown parameters introduces further difficulties. The author therefore requires non-Gaussian physical-space posteriors, high-dimensional scalability, observation-only training, and a prescribed simulator that can be calibrated directly. This is not a claim that variational state-space inference did not previously exist, but a specific construction satisfying these combined requirements.
The approach separates the equations that generate the world from the posterior that explains observations: the physical simulator remains the generative transition, while neural networks learn only observation-conditioned posterior degrees of freedom. A deterministic trajectory is fully determined by its initial state, so a large generative network need not reconstruct the entire trajectory; with process noise, future-observation-driven corrections are added around each physical rollout step. Core Idea: use a normalizing flow for a non-Gaussian initial state, the prescribed simulator for temporal dependence, and future-conditioned corrections only in the posterior, jointly trained and calibrated through an ELBO retaining physical transitions and observation likelihoods.
Method¶
Overall Architecture¶
Inputs comprise an observation window, prescribed differentiable dynamics and observation operators, and an initial-state prior. The output is not a single trajectory, but an approximate physical trajectory posterior that supports sampling and evaluation of the relevant density terms; unknown physical parameters and persistent sensor biases are learned as point estimates shared across windows. Training and deployment target repeated windows from the same system, not arbitrary physical parameters or observation layouts without adaptation.
“Non-Gaussian Initial State” encodes the entire observation window and draws an initial state through a conditional normalizing flow; “Simulator Rollout” advances samples through the original physical equations. The deterministic variant already yields trajectory samples at this point, whereas the noisy variant adds “Future-Conditioned Corrections” at each step. During training, “Observation-Only Joint Learning” evaluates these trajectories through the original observation operator and constructs an objective comparing the physical model with the variational posterior; inference on a new window does not reoptimize the network for that window.
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
Y["Observation window"] --> A["Non-Gaussian Initial State"]
A --> B["Simulator Rollout"]
B -->|Noisy variant| C["Future-Conditioned Corrections"]
Y -->|Current and future observations| C
C -->|Continue at next time| B
B -->|Deterministic variant| O["Trajectory posterior samples"]
C --> O
O -.->|Training only: original observation likelihood| D["Observation-Only Joint Learning"]
Y -.->|Training only: observation supervision| D
D -.->|Update posterior and shared parameters| AThe correction loop represents sequential sampling within a window, not retraining at deployment. Training also updates physical parameters, biases, and applicable noise scales; the original dynamics and observation operator remain in the objective rather than being replaced by the posterior networks in the diagram.
Key Designs¶
1. Non-Gaussian Initial State: place complex uncertainty at the entrance to a physical trajectory
A diagonal-Gaussian initial posterior can compress a sign-ambiguous state into one peak. PR-Smoother first encodes the entire observation window into context, then maps a standard Gaussian sample to the initial state through conditional RealNVP. A nonlinear invertible transformation can represent multiple density peaks while providing the exact initial log density through the change-of-variables formula, as needed for the entropy term and initial-state KL in the ELBO. Normalizing flows are selected not only for sampling quality, but also to avoid solving a separate density-tracking ODE for every sample.
The random input is not an additional learned latent dynamical state: the mapping outputs the physical initial state used by the original simulator. Low-dimensional Lorenz–96 uses a periodic convolutional encoder, the 40-dimensional variant adds attention along the spatial axis, and the Kolmogorov variant uses a two-dimensional conditional coupling flow. Encoders can vary with grid structure while retaining the same interface: observation-conditioned initial states, explicit densities, and physical-space outputs.
2. Simulator Rollout: obtain temporal dependence from the original physical equations
The deterministic variant samples only the initial state and calls the prescribed simulator at every subsequent step. One initial state therefore determines the full trajectory, preserving temporal dependence without generating each time independently. Nonlinear rollout can also make later state marginals non-Gaussian; Gaussian components in some conditional terms do not imply a Gaussian trajectory posterior.
Strictly, the deterministic trajectory distribution is a pushforward measure of the initial posterior and need not have an ordinary full-dimensional density in trajectory space. The paper first constructs the generative model and variational family with identical positive-noise transition kernels, allowing their log-transition terms to cancel, then takes the vanishing-noise limit. This avoids treating Dirac transitions as ordinary Gaussian densities, leaving an objective that constrains observation fit against the initial-state prior.
3. Future-Conditioned Corrections: modify posterior sampling rather than learn replacement dynamics
With process noise, one initial sample followed by deterministic rollout cannot explain random deviations at every step. The author exploits the Markov property of the state-space model: conditioned on the previous state, the current smoothing conditional requires only current and future observations. The forward variational factorization follows this structure, with stepwise conditionals:
The correction mean and covariance in the general construction can depend on the previous state, future observations, and time encoding. An anti-causal encoder reverses the observation sequence, processes it with a causally masked Transformer, and reverses the outputs back, so each context contains only current and future measurements; a correction head uses this context to predict deviations around the physical rollout. Although each conditional is Gaussian, the non-Gaussian initial state, nonlinear simulator, and state-dependent corrections together need not yield a Gaussian joint posterior.
The essential boundary is that these corrections belong only to the variational posterior; the generative transition remains centered on the original simulator. Stepwise KL terms penalize deviations from the prescribed process-noise model, preventing the network from explaining observations through arbitrary cost-free jumps. Setting correction means to zero and matching variational and process covariances recovers the nested construction without additional future-conditioned corrections; the further vanishing-process-noise limit yields the deterministic variant.
The main text and the unified settings in Appendix I specify isotropic scales shared across time for the actual noisy experiments, rather than arbitrary full covariances. The proposition that the family contains the exact smoother in the linear-Gaussian limit therefore relies on idealized assumptions about flow expressivity and the correction head's ability to represent the required means and covariances; it is not a guarantee that the concrete isotropic implementation is exact for every linear-Gaussian system.
4. Observation-Only Joint Learning: place state inference and scientific calibration in one objective
The model draws trajectory samples, produces predicted measurements through the original observation operator, and compares them with actual measurements through the likelihood. Unknown dynamical parameters enter the original simulator, and sensor biases enter the original observation model; both are shared across training windows, whereas window-specific states are inferred through the amortized posterior. These are not unconstrained repairs for a single trajectory: shared parameters must explain many different windows to be distinguishable from window-specific state freedom.
The evidence lower bound (ELBO) lower-bounds the log evidence of the observations. Its general form can be written below; maximizing it improves observation fit while constraining the posterior against the initial prior and physical process model:
This is the conditional-KL equivalent of main-text Eq. (14); after corresponding transition terms cancel in the deterministic construction, Eq. (9) provides its limiting objective. Gaussian conditional-transition KLs are available in closed form, and the initial flow density is directly evaluable, so training needs no true trajectory labels. Test ELBOs are summed over full windows and should be compared only within the same window length and observation pattern.
A Worked Example¶
Consider a 50-step window of 40-dimensional Lorenz–96 with componentwise measurements \(h(x)=\min(x^2,5)\) and observation-noise standard deviation 1. A sufficiently large positive state and its negative counterpart produce the same saturated value, so one measurement cannot determine the sign.
The full-window encoder first provides context to the initial flow, allowing different samples to retain multiple plausible explanations. Each initial-state sample is advanced through the original Lorenz–96 equations rather than choosing signs independently at each time; later measurements constrain whether the entire dynamical history is consistent. With process noise, anti-causal context additionally helps select stepwise corrections consistent with future measurements, subject to their KL cost.
Multiple observations thus do more than smooth an estimate: they rule out trajectories unable to produce subsequent measurements. The model returns a collection of trajectory samples; applications requiring a point estimate can summarize the posterior or use its output to initialize per-window 4D-Var refinement. The latter is a proposed application direction, not a reported combined experiment.
Loss & Training¶
Training maximizes minibatch observation-window ELBOs, with Adam optimizing the networks and applicable shared generative parameters. Learning-rate warmup covers the first 1% of Lorenz–96 updates and 5% of Kolmogorov updates, followed by exponential decay to 0.1 times the maximum rate. The 40-dimensional initial-state Flow uses a maximum rate of \(2\times10^{-4}\), Noisyflow uses \(4\times10^{-4}\), and all Kolmogorov variational variants use \(3\times10^{-4}\).
The main 4-dimensional and 40-dimensional training sets each contain approximately \(1.0\times10^7\) windows; the Kolmogorov summary table rounds its size to \(2.9\times10^6\), whereas the data-scaling experiment specifies the full size as \(2.88\times10^6\). Kolmogorov also uses a spectral Gaussian initial prior in the Fourier basis, with prior hyperparameters optimized jointly with other generative parameters by default.
The source has two inconsistent implementation descriptions: parts of Appendix I.2 describe diagonal, state-dependent variances in correction heads, whereas Section 5.2, Appendix F, and the unified settings in Appendix I specify shared isotropic scales; I.2 also ends by stating there are no trainable model parameters, whereas Appendix F explicitly learns process-noise scales jointly. This note interprets experiments according to the unified settings while retaining these discrepancies, rather than presenting them as verified code behavior.
Key Experimental Results¶
Main Results¶
Unless otherwise stated, values are means and standard deviations over five independent random seeds. The 4-dimensional experiment uses squared observations with noise standard deviation 3 and a bootstrap particle filter with \(5\times10^6\) particles as an evaluation reference, not as trajectory supervision for training.
Energy distance (ED) compares average cross-distribution sample distances while subtracting within-distribution distance terms; the 1-Wasserstein distance \(W_1\) is the minimum expected transport distance needed to move one distribution into the other. Smaller values indicate closer agreement with the particle reference. The cached text does not specify whether ED takes a square root, distance normalization, or numerical solution details for \(W_1\), so no exact implementation formula is invented and these values are not treated as errors against an analytical posterior.
The following entries from Figure 2 most directly distinguish initial-posterior expressivity from trajectory coupling:
| Method | 1-step ED | 1-step \(W_1\) | 10-step ED | 10-step \(W_1\) |
|---|---|---|---|---|
| Flow | 0.20 ± 0.04 | 0.90 ± 0.09 | 0.13 ± 0.04 | 0.22 ± 0.02 |
| Gauss | 1.38 ± 0.24 | 3.63 ± 0.67 | 0.19 ± 0.02 | 0.29 ± 0.01 |
| MF | 1.37 ± 0.24 | 3.62 ± 0.67 | 0.57 ± 0.05 | 0.68 ± 0.07 |
| EnKF | 0.70 ± 0.09 | 2.63 ± 0.31 | 1.42 ± 0.26 | 2.57 ± 0.58 |
The high-dimensional Kolmogorov system uses a \(128\times128\) grid, 10-step windows, and observation-noise standard deviation 3, jointly estimating global forcing, the Reynolds number, and sensor biases persistent across windows. The next table selects state RMSEs from main-text Table 1; oracle baselines receive true physical parameters and biases in advance and therefore do not have the same information as joint-estimation methods.
| Method | Full | Half | Sparse | Realistic |
|---|---|---|---|---|
| Flow | 0.28 ± 0.00 | 2.31 ± 0.09 | 1.92 ± 0.12 | 1.60 ± 0.02 |
| Noisyflow | 0.65 ± 0.01 | 2.16 ± 0.10 | 2.18 ± 0.11 | 1.84 ± 0.07 |
| Gauss | 0.32 ± 0.00 | 2.16 ± 0.10 | 4.95 ± 0.10 | 2.99 ± 0.13 |
| IEnKS | 4.04 ± 0.12 | 4.13 ± 0.18 | 4.10 ± 0.18 | 4.02 ± 0.19 |
| 4D-Var (oracle) | 0.31 ± 0.00 | 2.42 ± 0.10 | 3.81 ± 0.14 | 3.06 ± 0.16 |
In the Realistic setting, Table 2 reports Flow parameter errors of \(|\Delta F|=(2.3\pm1.6)\times10^{-3}\) and \(|\Delta\log Re|=0.31\pm0.01\), with bias error \(0.57\pm0.02\). The true values are 1.0 for \(F\) and 1000 for \(Re\), initialized at 0.1 and 100 respectively; this is not merely state reconstruction with known correct parameters.
Ablation Study¶
The 40-dimensional experiment compares combinations of Flow/Gauss initial posteriors and enabled/disabled future-conditioned corrections, alongside classical baselines; lower state RMSE is better. Columns respectively cover linear observations, nonlinear observations with deterministic dynamics, and nonlinear observations with process-noise standard deviation 0.1 per step, from Figure 3.
| Config | Linear | Nonlinear | Nonlinear + process noise |
|---|---|---|---|
| Flow | 0.516 ± 0.006 | 0.738 ± 0.011 | 0.909 ± 0.015 |
| Noisyflow | 0.287 ± 0.001 | 0.549 ± 0.084 | 0.498 ± 0.003 |
| Gauss | 0.597 ± 0.027 | 0.827 ± 0.009 | 1.020 ± 0.014 |
| NoisyGauss | 0.343 ± 0.002 | 0.797 ± 0.019 | 0.958 ± 0.006 |
| MF | 0.413 ± 0.001 | 3.170 ± 0.004 | 3.019 ± 0.030 |
| 4D-Var | 0.142 ± 0.004 | 6.488 ± 0.202 | 3.516 ± 0.024 |
| IEnKS | 0.176 ± 0.004 | 5.675 ± 0.175 | 5.603 ± 0.177 |
Variants differ in learning rates, batch sizes, and epoch counts, so this is not a strictly single-module-removal ablation with identical compute. The results support the importance of variational expressivity, but not attribution of every difference to an individual module's independent causal contribution.
Key Findings¶
- Single-step squared measurements in four dimensions allow at most \(2^4\) sign combinations in principle; the illustrated marginal is bimodal, which should not be confused with a claim that the full four-dimensional joint posterior has only two modes. With ten steps, dynamics help resolve sign ambiguity and Flow ED decreases from 0.20 to 0.13.
- 4D-Var and IEnKS lead in the linear 40-dimensional setting, so PR-Smoother is not best in every regime. Under nonlinear observations and process noise, Noisyflow achieves RMSE 0.498, below Flow's 0.909 and NoisyGauss's 0.958.
- Appendix F, Table 5 reports Noisyflow learning \(\ln\sigma_Q=-2.29\pm0.04\), close to the true value -2.30, compared with \(-4.61\pm0.16\) for NoisyGauss. Nonzero learned scales in columns without true process noise are effective model-error or regularization scales, not discoveries of physical noise.
- Mixed Kolmogorov coverage provides parameter identification through dense regions and global coverage through sparse regions; the entirely unobserved half cannot be reliably recovered. Under Sparse coverage, Flow obtains state RMSE 1.92 but \(|\Delta\log Re|=2.47\pm0.02\), showing that reconstructable states do not imply identifiable parameters.
- Under approximately matched training wall-clock budgets in the high-dimensional deterministic experiment, Flow outperforms Noisyflow. The author interprets this as slower convergence in a larger noisy search space, not universal ineffectiveness of corrections.
Highlights & Insights¶
- Placing physical equations in both the generative objective and posterior sampling structure is more direct than adding a physical penalty only to a loss. The model need not relearn temporal relationships and can concentrate its learned capacity on observation-induced uncertainty.
- Non-Gaussian expressivity and temporal coupling serve different purposes. Flow/Gauss comparisons primarily test initial-density expressivity, whereas MF comparisons expose the consequences of removing trajectory coupling; calling all of them “Gaussian” obscures this distinction.
- Amortization is valuable for repeated windows rather than eliminating compute costs. Appendix Table 12 reports trained Flow inference times of 0.37 seconds for linear 40-dimensional observations, 0.38 seconds for nonlinear observations, and 1.89 seconds for mixed-coverage Kolmogorov flow, but requires approximately 11–18 and 19–20 GPU-hours of offline training respectively.
Limitations & Future Work¶
- Differentiable simulators and observation operators are required; existing adjoint methods may help connect legacy solvers, but efficient training has not been validated for all such simulators. Experiments remain synthetic physical benchmarks rather than operational weather systems.
- Shared physical parameters and sensor biases are point estimates, with trajectory posteriors conditional on these estimates rather than including full parameter uncertainty. A global parameter variational factor and prior are natural extensions.
- Noisy stepwise conditionals remain restricted to a Gaussian family, with simplified covariances in the actual implementation. Expressivity in two tractable limits proves neither exactness for arbitrary nonlinear noisy systems nor successful optimization.
- Repeated windows from one system suit amortization, whereas changes in parameter distributions or observation layouts may require adaptation. Out-of-distribution predictive checks, posterior coverage, and calibration should supplement RMSE evaluation.
- Main Kolmogorov IEnKS results use 16 ensemble members, with appendix checks at 16, 32, and 48 and calibrated localization radii; this has a compute-budget rationale but does not establish superiority over every larger-ensemble configuration. Runtime comparisons exclude offline training and use oracle high-dimensional baselines.
Related Work & Insights¶
- vs 4D-Var / IEnKS: Both retain the physical model, but the former mainly optimizes per-window point estimates and the latter relies on ensemble Gaussian structure; PR-Smoother learns a reusable non-Gaussian trajectory posterior. Classical methods retain accuracy advantages with informative, near-linear observations.
- vs latent state-space models / LEVDA: These approaches ease inference through learned representations or dynamics surrogates; this paper operates directly in physical space and keeps original equation parameters in the calibration objective. Its advantage depends on the original simulator supporting gradient computation.
- vs score-based DA / FlowDAS: These approaches learn trajectory priors or transitions before conditioning on observations; PR-Smoother does not use supervised state trajectories to learn a replacement prior, retaining the prescribed simulator as the generative transition.
- vs AFSF: The source identifies this as contemporaneous conditional-flow filtering/smoothing work trained on simulated state–observation trajectories; PR-Smoother emphasizes an observation-only ELBO and a prescribed simulator. More flexible stepwise posteriors retaining explicit densities and physical-parameter identifiability are worth studying, but this extension is not evaluated here.
Rating¶
- Novelty: 4/5 — The contribution is structured non-Gaussian smoothing under prescribed-simulator constraints, not the invention of variational state-space inference.
- Experimental Thoroughness: 4/5 — Multimodality, process noise, and high-dimensional calibration are supported, but operational systems and full uncertainty calibration remain missing.
- Writing Quality: 4/5 — The problem boundaries and two tractable limits are clear, with conflicting appendix descriptions of noise parameters and covariance structure.
- Value: 4/5 — Relevant to scientific computing with differentiable physical models, repeated assimilation windows, and nonlinear observations.