Improving Ensemble Filters with Flow Matching

1 Introduction

• For stochastic learned updates, the source can encode selected information about current forecast uncertainty before the learned transformation, rather than using the forecast only as conditioning information.This source choice anchors the analysis update to the forecast ensemble available at each assimilation time.Source: §1 · p. 2
• Across sparsely observed Lorenz–96, Kuramoto–Sivashinsky, and Kolmogorov flow, FlowEF yields lower RMSE and marginal CRPS than tuned classical filters and, where directly compared, learned DA baselines.The reported gains are smaller for already strong baseline analyses and larger in several challenging finite-ensemble settings.Source: §1 · p. 2
• FlowEF formulates ensemble analysis as conditional transport from a source constructed from the current forecast ensemble, with the classical-filter analysis as additional conditioning information.Its localized-covariance Gaussian source and evolving observation residual are also part of the contribution.Source: §1 · p. 2
• The paper establishes a Wasserstein error bound that separates principal approximation errors in the proposed analysis update.The bound distinguishes losses from static conditioning and finite-ensemble estimation, velocity-field approximation and integration, source mismatch, and calibration inflation.Source: §1 · p. 2

2.1 State-space model

• The state-space model uses latent state x_t and observation y_t, with dynamics f and observation operator h; unless stated otherwise, the noise terms are additive Gaussian.The observation likelihood is written p_t(y_t | x).Source: §2.1 · p. 3

2.2 Bayesian forecast and analysis laws

• The forecast and analysis laws are defined at each time, and outside the linear-Gaussian setting their mean and covariance generally do not determine the predictive distribution.The text introduces Bayes’ rule for the forecast and analysis laws.Source: §2.2 · p. 3

2.3 Ensembles and analysis rules

• Ensemble filters represent forecast and analysis laws with finite particle collections, propagate members through the dynamics, and apply an analysis rule that may be stochastic or deterministic.Ensemble Kalman methods base updates on empirical moments, while the full ensemble retains information beyond its first two moments.Source: §2.3 · pp. 3–4

Classical ensemble and variational methods

• Classical ensemble filters form analysis updates from forecast moments and anomalies, while localization and inflation help stabilize finite-ensemble covariance estimates in spatially extended systems.Variational methods instead optimize prior-plus-observation mismatch, with ensemble and hybrid variants using ensemble subspaces or covariances.Source: §3 · pp. 4–5

Learned analysis rules

• Deterministic learned methods modify parts of the conventional forecast–analysis update or learn the analysis map, with uncertainty represented through the ensemble or resulting distributional quantities.Stochastic learned rules instead sample analysis particles from auxiliary random variables.Source: §3 · p. 5

Stochastic and transport methods

• Score and diffusion methods use conditional scores or guided sampling, while normalizing-flow filters learn invertible conditional transformations.Other transport approaches learn continuous-time dynamics or construct maps from couplings or optimality criteria.Source: §3 · pp. 5–6
• FlowEF constructs its source from the current forecast ensemble and transports particles through a learned flow conditioned on the forecast, corresponding baseline analysis, and observation.The source choice determines where forecast information enters the stochastic analysis rule: before transport or through the learned transformation.Source: §3 · p. 6

4 FlowEF

• FlowEF augments a classical filter by transporting forecast-informed starting particles to an analysis ensemble with a learned velocity field conditioned on baseline features and a dynamic observation field.The transport endpoints form a raw analysis ensemble that is subsequently calibrated by inflation.Source: §4 · p. 6

4.1 Conditional flow matching

• Conditional flow matching trains the velocity field on straight-line interpolants from source particles toward analysis targets, and inference integrates the learned ODE with K forward Euler steps.The interpolation target is a straight-line velocity.Source: §4.1 · pp. 7–8

4.2 Forecast ensemble starting laws

• Training draws source particles from a localized Gaussian based on the baseline forecast covariance, while inference transports each corresponding forecast member exactly once without resampling.The Gaussian uses the baseline filter's tuned Gaspari–Cohn localization radius.Source: §4.2 · p. 8

4.3 Conditioning variables

• Static conditioning describes forecast–observation discrepancy and baseline-filter response, while dynamic conditioning supplies the current observation residual and transport stage.The observation CNN maps residuals, observation weights, and a pseudo-time embedding into a dynamic field; conditioning on the baseline response lets the model learn relative to an existing update.Source: §4.3 · pp. 8–9

4.4 Training

• Training data are generated from a classical baseline before FlowEF fitting, so a time-t training tuple does not require a FlowEF analysis at time t −1.Training samples use a source particle and a uniformly sampled interpolation time.Source: §4.4 · p. 10

4.5 Inference and calibration

• At inference, test-time transports are computed independently, and multiplicative spread inflation is tuned separately on validation data by minimizing marginal ensemble CRPS.The score averages scalar CRPS over state coordinates rather than using a multivariate energy score.Source: §4.5 · pp. 10–11

4.6 Convergence analysis

• Under the stated assumptions, a Wasserstein error bound separates static-conditioning information loss and estimation, velocity approximation, numerical integration, source mismatch, and calibration-inflation error.The source term reflects the difference between forecast-member initialization at inference and the localized Gaussian training source.Source: §4.6 · p. 11

5.1 Experimental setup

• Experiments use EnKF, ETKF, DEnKF, and LETKF across Lorenz–96, Kuramoto–Sivashinsky, and Kolmogorov flow, with results averaged over 10 seeds.For learned-method comparisons, the validation-selected baselines are DEnKF for the first two systems and LETKF for Kolmogorov flow.Source: §5 · pp. 11–12, §5.1 · p. 12

Metrics

• RMSE measures ensemble-mean point accuracy, while CRPS evaluates predictive accuracy and concentration; SSR and Cov95 assess dispersion and marginal interval coverage.SSR is referenced to 1, and Cov95 to 0.95; values below or above those references indicate under- or over-dispersion or coverage.Source: Metrics · pp. 12–13

5.2 Lorenz–96 model

• FlowEF leads learned-method benchmarks on all four metrics for low-dimensional Lorenz–96 and Kolmogorov flow; on Kuramoto–Sivashinsky it leads RMSE, CRPS, and SSR, while ADEnKF leads Cov95.For high-dimensional Lorenz–96, FlowEF has the lowest RMSE and CRPS and the closest SSR and Cov95 to their reference values among reported methods.Source: §5.2 · pp. 13–14, §5.3 · p. 16, §5.4 · p. 18
• With dx=100 and N=30 in Lorenz–96, FlowEF improves all four reported metrics over each baseline, with larger RMSE and CRPS reductions than at dx=40.The sample covariance rank is at most N−1, making covariance estimation more challenging as state dimension increases.Source: §5.2 · p. 13, Table 3 · p. 14

5.3 Kuramoto–Sivashinsky equation

• In Kuramoto–Sivashinsky, FlowEF improves all four metrics over each baseline and reduces run-to-run variability; gains are narrowest over DEnKF.On a representative seed, its RMSE and CRPS remain below DEnKF across the trajectory and its rank histogram is flatter.Source: §5.3 · p. 16, Table 4 · p. 16

5 Experiments

• For FlowEF atop each of four baselines, RMSE and CRPS improve across 10-seed averages, with the largest gains for ETKF and the smallest for EnKF.In low-dimensional Lorenz–96, FlowEF brings SSR and Cov95 closer to their reference values for all four baselines.Source: §5.2 · p. 13, §5.4 · p. 18

5.4 Kolmogorov flow

• In Kolmogorov flow, FlowEF improves RMSE and CRPS for all baselines, with the largest gains over ETKF and smallest over EnKF.SSR improves for ETKF and LETKF, while Cov95 improves for DEnKF and LETKF; observations cover only 1.6% of the state.Source: §5.4 · p. 18

6 Conclusion

• FlowEF transports forecast-derived particles through a learned conditional flow to form an analysis ensemble, allowing nonlinear updates beyond classical filters’ affine structure.Its velocity field is conditioned on baseline forecast and analysis ensembles and the current observation.Source: §6 · pp. 18–19
• Across the reported systems, FlowEF reduces RMSE and marginal CRPS relative to tuned classical baselines and outperforms evaluated learned DA alternatives where direct comparisons are available.Improvements are smaller for already strong baseline analyses and larger in several challenging finite-ensemble settings.Source: §6 · p. 19
• FlowEF generally improves spread–skill ratio and marginal coverage, although empirical 95% coverage remains below nominal in several settings.The results indicate improved point accuracy and probabilistic performance while retaining information from a classical ensemble filter.Source: §6 · p. 20
• FlowEF is trained separately for each baseline filter and dynamical system, limiting transfer without retraining; training also depends on simulated trajectories representing inference regimes.Numerical integration adds inference cost beyond the classical analysis, and the conditioning construction is designed primarily for the observation models considered.Source: §6 · p. 20
• Future work could pursue transfer across filters, regimes, and observation configurations, extend conditioning to nonlinear or misspecified observation models, and reduce inference cost.Other proposed directions include jointly learning localization and calibration with the transport.Source: §6 · p. 20

Appendix A. Implementation details

• Implementation details cover data generation and splits, localized Gaussian sampling, network architecture, training, inference and calibration, and evaluation metrics.The corresponding material is organized across Appendices A.1–A.6.Source: Appendix A · p. 25
• Algorithm 1 summarizes the complete FlowEF training and analysis-update procedure.Source: Appendix A · p. 25
• Table 6 collects the implementation and hyperparameter settings used in the experiments.Source: Appendix A · p. 25

A.1 Data generation and splits

• Each system uses a 25,000-step simulated trajectory, with disjoint windows of 10,000 training, 2,000 validation, and 2,000 test assimilation times.The pipeline is repeated over 10 data seeds, regenerating trajectories, splits, and other listed randomness.Source: §A.1 · p. 25
• The validation set tunes baseline inflation and localization before FlowEF training, then tunes FlowEF inflation after training.An initial spin-up lets the trajectory and baseline filter settle from arbitrary initialization.Source: §A.1 · p. 25

A.2 Localized Gaussian sampling

• Localization uses ring distance for one-dimensional systems and torus distance for Kolmogorov flow, with a Gaspari–Cohn taper of compact support ρ≤2rloc.Source: §A.2 · pp. 25–26
• The taper covariance is sampled through an eigendecomposition and truncated low-rank factor; setting rloc=0 gives a diagonal Gaussian source that preserves marginal spread but discards cross-coordinate correlations.Dense or randomized low-rank eigensolvers are used depending on state dimension and estimated rank.Source: §A.2 · p. 26

A.3 Architecture

• Two periodic CNNs extract features from observation residuals and pseudo-time, then compute velocity from the transported particle and static and dynamic conditioning channels.FiLM re-injects pseudo-time into each residual block so the velocity can vary smoothly with pseudo-time at every depth.Source: §A.3 · pp. 28–29

A.4 Training

• Training samples 32 assimilation times per minibatch and updates velocity-network parameters with Adam for 100 epochs, cosine learning-rate decay, and gradient clipping at norm 5.Source: §A.4 · pp. 30–31

A.5 Inference and calibration

• Inference integrates the flow with 10 Euler steps, recomputing dynamic conditioning at each step; a validation-set search selects inflation to minimize CRPS.The dynamic conditioning tracks the transported particle’s distance from the observation during intermediate steps.Source: §A.5 · p. 31

• RMSE measures point accuracy of the analysis ensemble mean against the ground truth.The metrics are computed on the test set and aggregated over 10 seeds.Source: §A.6 · pp. 33–34
• CRPS jointly measures ensemble accuracy and concentration around the truth, and is minimized when the predictive distribution matches the data-generating distribution.Source: §A.6 · p. 34
• SSR compares ensemble spread with actual error: 1 is the reference, below 1 indicates underdispersion, and above 1 indicates overdispersion.Source: §A.6 · p. 34
• Cov95 is coordinatewise marginal coverage of the truth by the ensemble’s 95% interval; 0.95 is the reference, with lower values indicating undercoverage and higher values overcoverage.Source: §A.6 · pp. 34–35

Appendix B. Benchmark comparison

• Many adapted benchmark methods perform worse than the tuned classical baseline in RMSE and CRPS; these are common-framework comparisons, not exact reproductions or evaluations in native settings.The experiments use tuned localized and inflated classical filters and a linear, sparse observation model; the comparisons concern analysis updates from fixed baseline forecasts, not recursive filtering performance.Source: Appendix B · pp. 35, 38

Appendix C. Convergence theory of FlowEF

• The error bound decomposes FlowEF’s error into six contributions, and exact recovery requires matching training and inference sources, an exact velocity, and an exact ODE solve.Source: Appendix C · p. 38

C.1 Notations

• The notation section uses σ-algebras to represent information carried by random variables; one generated by Z is contained in F when each variable in Z is measurable.Source: §C.1 · p. 38

C.1.1 Filtering laws notations

• The filtering-law notation distinguishes exact forecast and analysis laws, baseline ensemble empirical laws, and FlowEF output laws and ensembles.Source: §C.1.1 · pp. 38–39

C.1.2 Conditional variables notations

• Conditioning notation separates static baseline information from dynamic features recomputed from the transported state and observation residual.The static channel uses coordinatewise first and second moments of the forecast and analysis ensembles; the dynamic feature is a τ-dependent function of the transported state.Source: §C.1.2 · pp. 39–40

C.1.3 Training and inference notations

• Training and inference notation distinguishes the conditional training source from the inference source, which is exactly the baseline empirical forecast measure.Source: §C.1.3 · pp. 40–41

C.2.1 Observation and baseline assumptions

• The setup assumes known linear Gaussian observations, time-invariant observation design with fixed noise covariance, a statistically steady time-uniformly bounded baseline, common marginal law across train, validation, and test cycles, and shared trained parameters.Under suitable regularity conditions, finite-ensemble error relative to infinite-ensemble counterparts can satisfy O(N^-1/2); the bound's constant depends on assimilation time.Source: §C.2.1 · p. 42

C.2.2 Regularity assumptions

• The regularity assumptions require Wasserstein-stable conditioning, concentration of static features, and Lipschitz stability of the ideal transport; baseline bias can leave a nonvanishing conditioning offset.The concentration condition distinguishes sampling fluctuation from systematic offset.Source: §C.2.2 · p. 43

C.3.1 Induced velocity field

• The induced velocity formulation gives the exact CFM loss and Euler inference update, and its excess loss equals integrated mean-squared velocity error.The result is stated for the induced velocity field of the algorithm.Source: §C.3.1 · p. 44

C.3.2 Static and dynamic conditioning

• Under Assumption 3, the induced parameterization represents the ideal velocity, leaving flow-matching error attributable to estimation and optimization rather than approximation.The dynamic feature passes the observation record to the velocity network.Source: §C.3.2 · p. 47

C.3.3 Transport endpoint

• When induced velocity error is zero and training and inference share a source distribution, the learned flow has the exact conditional endpoint law.The result is stated under the endpoint lemma’s conditions.Source: §C.3.3 · pp. 48–49

C.3.4 Conditioning bias (εbias)

• The conditioning-bias analysis assumes the conditioned law satisfies a uniform Talagrand T2(κ_t) transport–entropy inequality.The supplied passage states this condition for the law given the conditioning variables.Source: §C.3.4 · p. 49

C.3.5 Finite-ensemble features (εfeat)

• Under Assumption 8, the endpoint conditional laws satisfy the stated conditional stability result.The supplied passage identifies the result as Lemma 15 but does not include its bound.Source: §C.3.5 · p. 49
• For finite-ensemble features, the endpoint laws are K_t(c, o) and K_t(c′, o), yielding the stated W2 bound.The accompanying assumption sets β_t = 1 and defines the continuous limit as 1 when L_u,z,t = 0.Source: §C.3.5 · pp. 49–50

C.3.6 Flow matching and Euler integration (εfm, εode)

• Under Assumptions 6 and 7 with β_t = 1, the results bound population conditional error, exact-flow CFM error, and K-step Euler error.The flow-matching and Euler results apply to the learned velocity field; the Euler result is stated for source laws supported in N_t.Source: §C.3.6 · pp. 50–51

C.3.7 Source mismatch (εsrc)

• Under Assumptions 5 and 9, the source-mismatch result includes an O(r_f,N) baseline-sampling term.The lemma labels this contribution as baseline sampling.Source: §C.3.7 · p. 52

C.3.8 Output calibration (εcal,N)

• Calibration error ε_cal,N is zero if and only if α_flow = 1; over the Table 6 search grid [0.80, 2.00], |α_flow − 1| ≤ 1.The result concerns inflation of the ensemble anomalies about an unchanged mean.Source: §C.3.8 · pp. 53–54

C.4 Main results

• Table 12 organizes each error term with its corresponding Algorithm 1 step and the condition under which it vanishes.The theorem’s bound is stated under Assumptions 3–7 and 9, with β_t = 1.Source: §C.4 · p. 54, Table 12 · p. 55
• Theorem 20 decomposes FlowEF’s error into six terms: conditioning bias, feature estimation, flow matching, Euler discretization, source mismatch, and inflation discrepancy.The source mismatch is between the forecast ensemble and the localized Gaussian training source.Source: §C.4 · pp. 54–55
• Under assumptions that remove the first four error terms, only the fifth term remains nonzero in the resulting bound.The proof identifies the endpoint law and applies the bound for that remaining term.Source: §C.4 · p. 57