Paper deep dive
Rethinking Forward Processes for Score-Based Data Assimilation in High Dimensions
Eunbi Yoon, Donghan Kim, Dae Wook Kim
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 98%
Last extracted: 4/10/2026, 2:02:21 AM
Summary
The paper introduces the Measurement-Aware Score-based Filter (MASF), a novel approach for high-dimensional data assimilation. By constructing a forward process directly from the measurement equation, MASF makes the likelihood score analytically tractable, avoiding the heuristic approximations common in existing score-based filters. This allows for more accurate and stable posterior estimation through reverse-time sampling, as demonstrated in experiments on Lorenz-63, Lorenz-96, and Kolmogorov flow datasets.
Entities (5)
Relation Signals (2)
MASF → improves → Data Assimilation
confidence 100% · Numerical experiments covering a range of settings, including high-dimensional datasets, demonstrate improved accuracy and stability over existing score-based filters.
MASF → utilizes → Score-based Generative Models
confidence 95% · We propose a measurement-aware score-based filter (MASF) that defines a measurement-aware forward process directly from the measurement equation.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Data assimilation is the process of estimating the time-evolving state of a dynamical system by integrating model predictions and noisy observations. It is commonly formulated as Bayesian filtering, but classical filters often struggle with accuracy or computational feasibility in high dimensions. Recently, score-based generative models have emerged as a scalable approach for high-dimensional data assimilation, enabling accurate modeling and sampling of complex distributions. However, existing score-based filters often specify the forward process independently of the data assimilation. As a result, the measurement-update step depends on heuristic approximations of the likelihood score, which can accumulate errors and degrade performance over time. Here, we propose a measurement-aware score-based filter (MASF) that defines a measurement-aware forward process directly from the measurement equation. This construction makes the likelihood score analytically tractable: for linear measurements, we derive the exact likelihood score and combine it with a learned prior score to obtain the posterior score. Numerical experiments covering a range of settings, including high-dimensional datasets, demonstrate improved accuracy and stability over existing score-based filters.
Tags
Links
- Source: https://arxiv.org/abs/2604.02889v1
- Canonical: https://arxiv.org/abs/2604.02889v1
Trouble viewing inline? Open PDF directly →
Full Text
75,974 characters extracted from source content.
Expand or collapse full text
Rethinking Forward Processes for Score-Based Data Assimilation in High Dimensions Eunbi Yoon1, Donghan Kim2, and Dae Wook Kim1 1 Department of Brain and Cognitive Sciences, KAIST, Daejeon, South Korea 2 Department of Mathematical Sciences, KAIST, Daejeon, South Korea Abstract Data assimilation is the process of estimating the time-evolving state of a dynamical system by integrating model predictions and noisy observations. It is commonly formulated as Bayesian filtering, but classical filters often struggle with accuracy or computational feasibility in high dimensions. Recently, score-based generative models have emerged as a scalable approach for high-dimensional data assimilation, enabling accurate modeling and sampling of complex distributions. However, existing score-based filters often specify the forward process independently of the data assimilation. As a result, the measurement-update step depends on heuristic approximations of the likelihood score, which can accumulate errors and degrade performance over time. Here, we propose a measurement-aware score-based filter (MASF) that defines a measurement-aware forward process directly from the measurement equation. This construction makes the likelihood score analytically tractable: for linear measurements, we derive the exact likelihood score and combine it with a learned prior score to obtain the posterior score. Numerical experiments covering a range of settings, including high-dimensional datasets, demonstrate improved accuracy and stability over existing score-based filters. 1 Introduction Figure 1: Schematic comparison of likelihood score. (a) Existing approaches specify the forward process independently of the measurement equation, which makes the likelihood intractable. (b) Our approach aligns the forward process with the measurement equation, so the likelihood score becomes tractable. Data assimilation estimates latent states from partial and noisy observations by combining dynamical model predictions with measurement information over time Evensen (2009a); Reich and Cotter (2015). It is often formulated as Bayesian filtering, which alternates between a time update and a measurement update Särkkä (2013). In the time update, the current state is propagated under the state equation to produce a prediction Asch et al. (2016). This is then corrected in the measurement update using observations via the measurement equation. Such filtering problems arise broadly in domains where time-evolving dynamics must be inferred from incomplete information, e.g., geophysical forecasting, and biological processes Chipilski et al. (2020); Aksoy et al. (2009); Cogan et al. (2021). Although the optimal Bayesian filters are well defined, exact computation is rarely feasible in high-dimensional nonlinear settings, since both updates involve integrals that typically admit no closed-form expressions Särkkä (2013); Doucet et al. (2001). In practice, two widely used families of Bayesian filters are Kalman filters and particle filters Kalman (1960); Evensen (2009b). Kalman filters, such as the ensemble Kalman filter (EnKF), approximate the posterior through Gaussian moments and update them recursively Kalman (1960); Whitaker and Hamill (2002). They are computationally efficient, but their accuracy can degrade when the posterior is non-Gaussian or when the state and measurement equations are highly nonlinear Asch et al. (2016); Houtekamer and Mitchell (1998). Particle filters, such as the auxiliary particle filter, represent the posterior with weighted samples and can capture non-Gaussian structure more faithfully Gordon et al. (1993); Andrieu et al. (2010); however, they often suffer from severe weight degeneracy in high dimensions Arulampalam et al. (2002); Snyder et al. (2008). Figure 2: Pipeline of the proposed method, MASF. The forward process is constructed by interpolating between the identity and the measurement operator, so that the state is progressively degraded toward the measurement. The reverse-time process samples state trajectories from the posterior. Recently, there has been growing interest in score-based generative models as a tool for representing complex high-dimensional distributions Song and Ermon (2019); Song et al. (2021); Dhariwal and Nichol (2021). These models learn the score function, the gradient of the log-density, and enable sampling by running a reverse-time process Hyvärinen (2005); Vincent (2011). Motivated by this, recent work in data assimilation trains score models to estimate the prior score and include measurements at sampling time through a likelihood score Rozet and Louppe (2023); Bao et al. (2024a, b); Ding et al. (2025). A representative approach is the score-based filter (SF), which applies a standard forward process to generate perturbed states and trains a score model to estimate the corresponding prior score Bao et al. (2024a). In what follows, we use score-based filtering to refer broadly to this family of methods; SF specifically denotes the algorithm of Bao et al. (2024a). At the measurement-update step, SF incorporates measurements by approximating the likelihood score with respect to the perturbed state, and guides reverse-time sampling using the sum of the learned prior score and the approximated likelihood score. However, this likelihood score approximation is not theoretically justified and can lead to errors that accumulate over sequential updates. Score-based Sequential Langevin Sampling (SSLS) instead adopts score matching with Annealing Langevin Monte Carlo Song and Ermon (2019), where sampling only requires the score of the target distribution Ding et al. (2025). In this framework, the exact likelihood score can be derived from the measurement equation, avoiding the need to approximate a perturbed likelihood score. Nevertheless, Langevin-based sampling typically relies on annealing over noise levels, which can substantially increase the number of sampling steps and make inference computationally expensive Song et al. (2021). A key challenge is to retain the advantages of score-based filters while removing the main bottleneck of existing approaches. This is a problem in data assimilation where the measurement updates are performed sequentially over time; as a result, even small approximation errors can accumulate and degrade performance over time, unlike in conditional generation. The main obstacle is that sampling evolves on perturbed states produced by a forward process, whereas conditioning requires evaluating the likelihood score at those same perturbed states. When the forward process is chosen independently of the measurement equation, the likelihood induced on perturbed states is generally intractable, forcing a likelihood score approximation at every update. To address this problem, we propose a Measurement-Aware Score-based Filter (MASF) that defines a measurement-aware forward process directly from the measurement equation. This method ensures that the likelihood score along the perturbed trajectory remains analytically feasible as in Fig. 1. In the linear measurement case, we construct the forward process by interpolating between the identity and the measurement operator, progressively mapping from state space to the measurement space over time; see Fig. 2 for an illustration. This yields a closed-form expression for the likelihood score along the perturbed trajectory, which we combine with a learned prior score to obtain the posterior score and derive reverse-time sampling without ad hoc approximations and annealing. Experiments on Lorenz–63, Lorenz–96, and Kolmogorov flow demonstrate that the proposed method consistently improves accuracy and stability over existing score-based filters over a broad range of settings, including high-dimensional datasets, supporting the approach both theoretically and experimentally. 2 Background 2.1 Bayesian Filtering Consider a continuous-time latent state process Xτ∈ℝdX_τ ^d governed by the stochastic differential equation (SDE) dXτ=f(Xτ,τ)dτ+g(Xτ,τ)dBτ, dX_τ=f(X_τ,τ)\,dτ+g(X_τ,τ)\,dB_τ, (1) where f:ℝd×ℝ→ℝdf:R^d×R ^d and g:ℝd×ℝ→ℝd×dg:R^d×R ^d× d denote the drift and diffusion terms. Bτ≥0\B_τ\_τ≥ 0 is a standard Brownian motion. Let τkk=1K\ _k\_k=1^K be the discrete measurement times. The corresponding linear measurement equation is Zk=AXk+σϵk,ϵk∼(0,I), Z_k=AX_k+σ\, ε_k, ε_k (0,I), (2) where Xk≐XτkX_k X_ _k, A∈ℝd×dA ^d× d is the measurement operator, and σ>0σ>0 is the noise scale. In the main text, we focus on linear measurements; nonlinear measurement equations are discussed in Section 6. The goal of Bayesian filtering is to estimate the posterior of the state given measurements up to time τk _k Särkkä (2013): p(k∣1:k):=p(Xk=τk|Z1=1,…,Zk=k), p(x_k _1:k):=p\! (X_k=x_ _k\, |\,Z_1=z_1,…,Z_k=z_k ), where 1:k=(1,…,k)z_1:k=(z_1,…,z_k). The posterior can be computed recursively by alternating a time-update step and a measurement-update step. Time-update step. Given the posterior at time τk−1 _k-1, the state SDE (1) induces the transition density p(k∣k−1):=p(Xτk=k|Xk−1=k−1). p(x_k _k-1):=p\! (X_ _k=x_k\, |\,X_k-1=x_k-1 ). (3) Then the prior in time τk _k is obtained by the Chapman–Kolmogorov equation Law et al. (2015): p(k∣1:k−1)=∫p(k∣k−1)p(k−1∣1:k−1)xk−1. p(x_k _1:k-1)= p(x_k _k-1)\,p(x_k-1 _1:k-1)\,dx_k-1. Measurement-update step. The posterior satisfies Bayes’ rule: p(k∣1:k)⏟Posterior∝p(k∣1:k−1)⏟Priorp(k∣k)⏟Likelihood. p(x_k _1:k)_ Posterior\; \; p(x_k _1:k-1)_ Prior\; p(z_k _k)_ Likelihood. (4) At time τk _k, the new measurement ZkZ_k is incorporated through the likelihood term p(k∣k)p(z_k _k), which is specified by the measurement equation. Taking logs and gradients with respect to kx_k then gives the additive decomposition of (4): ∇klogp(k∣1:k) _x_k p(x_k _1:k) (5) =∇klogp(k∣1:k−1)+∇klogp(k∣k). = _x_k p(x_k _1:k-1)+ _x_k p(z_k _k). We refer to the gradients of the posterior, prior, and likelihood log-densities as the posterior, prior, and likelihood scores, respectively. 2.2 Score-based generative models We consider the linear SDEs commonly used in score-based generative modeling Song et al. (2021): dXt=F(t)Xtdt+G(t)dBt,t∈[0,1], dX_t=F(t)X_t\,dt+G(t)\,dB_t, t∈[0,1], (6) where F(t)∈ℝd×dF(t) ^d× d and G(t)∈ℝd×dG(t) ^d× d are the time-dependent drift and diffusion terms, respectively. A widely used instance is the variance-preserving (VP) SDE, F(t)=−12β(t)I,G(t)=β(t)I, F(t)=- 12β(t)I, G(t)= β(t)\,I, (7) with an increasing function β(t)β(t) Nichol and Dhariwal (2021). The solution of the VP SDE has a closed-form: Xt=a(t)X0+γ(t)ϵ, X_t=a(t)X_0+γ(t)\, ε, (8) where ϵ∼(0,I) ε (0,I) and dtloga(t)=−12β(t),γ2(t)=1−a2(t) ddt a(t)=- 12β(t), γ^2(t)=1-a^2(t) (9) with a(0)=1a(0)=1 and a(1)=0a(1)=0 Song et al. (2021); Ho et al. (2020). This solution (8) implies that the conditional score ∇tlogp(t∣0) _x_t p(x_t _0) is linear in tx_t, enabling an efficient denoising-score-matching objective Hyvärinen (2005); Vincent (2011). 3 Methods 3.1 Forward process from State to Measurement space We introduce a time-dependent linear operator and an isotropic covariance A(t) A(t) =(1−a(t))A+a(t)I, =(1-a(t))A+a(t)I, (10) Σ(t) (t) =σ2γ2(t)I, =σ^2γ^2(t)\,I, (11) where a(t)a(t) and γ(t)γ(t) follow (9), A∈ℝd×dA ^d× d is fixed, and I is the identity matrix. We assume that A(t)A(t) is invertible for all t∈[0,1)t∈[0,1), while the endpoint operator A may be singular. We define the forward process by Xt=A(t)X0+Σ(t)12ϵ,ϵ∼(0,I). X_t=A(t)X_0+ (t) 12 ε, ε (0,I). (12) Equivalently, the conditional moments are (Xt∣X0=0) \! (X_t X_0=x_0 ) =A(t)0, =A(t)x_0, (13) Cov(Xt∣X0=0) \! (X_t X_0=x_0 ) =Σ(t). = (t). (14) Moment-matching SDE. We construct a linear SDE (6) whose solution matches Eq. (12). For (6), the conditional mean u(t):=(Xt∣X0=0)u(t):=E(X_t X_0=x_0) satisfies u˙(t)=F(t)u(t),u(0)=0. u(t)=F(t)u(t), u(0)=x_0. (15) Requiring u(t)=A(t)0u(t)=A(t)x_0 for all 0x_0 implies F(t)=A˙(t)A(t)−1, F(t)= A(t)\,A(t)^-1, (16) which is well-defined on t∈[0,1)t∈[0,1). Similarly, the conditional covariance v(t):=Cov(Xt∣X0=0)v(t):=Cov(X_t X_0=x_0) satisfies the Lyapunov equation Kloeden and Platen (1992) v˙(t)=F(t)v(t)+v(t)F(t)+G(t)G(t). v(t)=F(t)v(t)+v(t)F(t) T+G(t)G(t) T. (17) Imposing v(t)=Σ(t)v(t)= (t) then implies G(t)G(t)=Σ˙(t)−F(t)Σ(t)−Σ(t)F(t). G(t)G(t) T= (t)-F(t) (t)- (t)F(t) T. (18) Eq. (16) and (18) ensure that the resulting SDE matches the moments (13)–(14). See Appendix A for details. Transition law and likelihood score. The linear SDE with solution (12) induces Gaussian transition kernels: Xt∣Xs∼(Ms→tXs,Σs→t), X_t X_s \! (M_s→ tX_s, _s→ t ), (19) where Ms→t M_s→ t =A(t)A(s)−1, =A(t)A(s)^-1, (20) Σs→t _s→ t =Σ(t)−Ms→tΣ(s)Ms→t,0≤s<t≤1. = (t)-M_s→ t\, (s)\,M_s→ t T, 0≤ s<t≤ 1. Since Z=X1Z=X_1, the transition kernel from t to 11 leads to Z=Mt→1Xt+Σt→112ϵ,ϵ∼(0,I), Z=M_t→ 1X_t+ _t→ 1 12 ε, ε (0,I), (21) for t∈[0,1)t∈[0,1). Therefore, the likelihood score is ∇tlogp(∣t)=Mt→1Σt→1−1(−Mt→1t). _x_t p(z _t)=M_t→ 1 T\, _t→ 1^-1 (z-M_t→ 1x_t ). (22) This form clarifies how the measurement influence varies with t. As t→1t→ 1, the uncertainty Σt→1 _t→ 1 shrinks and the precision Σt→1−1 _t→ 1^-1 increases, amplifying the residual −Mt→1tz-M_t→ 1x_t. Consequently, in reverse-time sampling from near 11 to 0, the likelihood term is typically most influential near the beginning of the trajectory. See Appendix B for details. Reverse-time SDE. The reverse-time SDE associated with (6) is given by Anderson (1982): dXt dX_t =(F(t)Xt−G(t)G(t)∇logpt(Xt))dt = (F(t)X_t-G(t)G(t) T _x p_t(X_t) )dt (23) +G(t)dB¯t, +G(t)\,d B_t, where B¯t B_t is a Brownian motion in reverse time, and ptp_t denotes the marginal density of XtX_t. When conditioning on a measurement z, we replace the prior score by the posterior score via ∇tlogpt(t∣)=∇tlogpt(t)+∇tlogp(∣t). _x_t p_t(x_t )= _x_t p_t(x_t)+ _x_t p(z _t). Since ∇tlogpt(t) _x_t p_t(x_t) is generally intractable, we approximate it with a learned score model. For a perturbed state tx_t, the conditional score is ∇tlogp(t∣0) _x_t p(x_t _0) =−Σ(t)−1(t−A(t)0) =- (t)^-1 (x_t-A(t)x_0 ) (24) =−Σ(t)−12ϵ. =- (t)^- 12\, ε. The likelihood score ∇tlogp(∣t) _x_t p(z _t) is available in closed form, and we form the posterior score by adding it to the learned prior score. Reverse-time Sampling. A discretization of the reverse-time SDE (23) is t≈Ms→ts−Σs→t∇logps(s∣)+Σs→t12ϵ, _t≈ M_s→ tx_s- _s→ t\, _x p_s(x_s )+ _s→ t 12 ε, (25) for t<st<s. See Appendix B for details. In summary, we take the measurement equation to construct the forward process tailored for data assimilation. This allows us to obtain a moment-matching SDE, an exact likelihood score ∇tlogp(∣t) _x_t p(z _t), and a reverse-time SDE that enables sampling from the posterior. Applying to data assimilation. Measurement examples include additive noise with the identity operator and pixel-wise masking. In both cases, the interpolation (10) is invertible for all t∈[0,1)t∈[0,1) since the measurement operator A has a nonnegative spectrum; see Appendix C. Such measurements are standard in data assimilation Law et al. (2015); Asch et al. (2016); Carrassi et al. (2018). At each measurement time τk _k, we generate perturbed states by applying the measurement-aware forward process to a prior sample kx_k: k,t=A(t)k+Σ(t)12ϵ,ϵ∼(0,I). _k,t=A(t)x_k+ (t) 12 ε, ε (0,I). (26) Training on these perturbed states via denoising score matching provides a score model that approximates the prior score ∇k,tlogpk,t(k,t) _x_k,t p_k,t(x_k,t). Bayesian filtering is implemented by alternating (i) time update, which propagates the current posterior ensemble through the state dynamics to form a prior, and (i) measurement update, which learns the prior score from the propagated ensemble, combines it with the closed-form likelihood score to obtain a posterior score, and performs reverse-time sampling to produce the next posterior ensemble. 3.2 Training via Denoising Score Matching At each measurement time τk _k, we learn a score model θk(k,t,t)≈∇k,tlogpt(k,t)S_ _k(x_k,t,t)≈ _x_k,t p_t(x_k,t). For each k, we minimize the denoising score-matching objective ℒθk(t)=ϵ∼(0,I)[‖θk(k,t,t)+Σ(t)−12ϵ‖22]. _ _k(t)=E_ ε (0,I) [ \|S_ _k(x_k,t,t)+ (t)^- 12 ε \|_2^2 ]. (27) We update θk _k by minimizing t∼U(0,1)[ℒθk(t)]E_t (0,1)\! [L_ _k(t) ], where t is sampled from the uniform distribution U(0,1)U(0,1). 4 Experimental Setup Dataset construction. To generate ground-truth trajectories, we integrate the state equation (1) over [τ0,τR][ _0, _R] on a uniform grid τr=τ0+rΔτ _r= _0+r τ with Δτ=(τR−τ0)/R τ=( _R- _0)/R, for r=0,…,Rr=0,…,R, and set r:=τrx_r:=x_ _r. Let ⊂0,…,RK⊂\0,…,R\ denote the measurement steps, ||=K|K|=K; for each k∈k , we generate kz_k by applying the measurement equation (2) to kx_k. We initialize the prior by sampling N ensemble ^0(i)i=1N\ x^(i)_0\_i=1^N from a user-specified distribution. Unless otherwise stated, ^0(i)∼(0,I) x^(i)_0 (0,I) and we use N=100N=100. Training at a fixed measurement step. Given the prior ensemble ^k(i)i=1N\ x^(i)_k\_i=1^N, we generate perturbed states ^k,t(i) x^(i)_k,t by applying the forward process (12). We train a score model θk(⋅,t)S_ _k(·,t) with the loss ℒθk(t)L_ _k(t) in (27). In principle, a separate parameter set θk _k is required for each measurement step k. To reduce computational cost, we fully train the model at the first measurement step and, for subsequent steps, update it by fine-tuning only a subset of parameters. Algorithm 1 MASF algorithm 1: Input: measurement step set ⊂0,…,RK⊂\0,…,R\, measurements (k)k∈(z_k)_k , particles N, epochs E, nfe 2: Output: state estimates (r)r=0R(x_r)_r=0^R 3: (0) Initialization: sample (^0(i))i=1N∼p0(0)( x^(i)_0)_i=1^N p_0(x_0); set 0←1N∑i=1N^0(i)x_0← 1N _i=1^N x^(i)_0 4: for r=1r=1 to R do 5: (1) Time-update step: 6: ^r(i)←Transition(^r−1(i)),i=1,…,N x^(i)_r \! ( x^(i)_r-1 ), i=1,…,N 7: if r∈r then 8: (2) Train prior score at measurement step r: 9: for ℓ=1 =1 to E do 10: Sample t∼(0,1)t (0,1) and ϵ(i)∼(0,I) ε^(i) (0,I) 11: ^r,t(i)←(t)^r(i)+Σ12(t)ϵ(i) x^(i)_r,t (t) x^(i)_r+ 12(t) ε^(i) 12: L←1N∑i=1N‖r(^r,t(i),t)+Σ−12(t)ϵ(i)‖2L← 1N _i=1^N \|S_ θ_r( x^(i)_r,t,t)+ ^- 12(t) ε^(i) \|^2 13: Update r θ_r by minimizing L 14: end for 15: (3) Measurement-update step: 16: Initialize r,1−eps(i)x^(i)_r,1- eps by forward process on ^r(i) x^(i)_r 17: times←linspace(1−eps,0,nfe+1) times← linspace(1- eps,0, nfe+1) 18: for j=0j=0 to nfe−1 nfe-1 do 19: Sampleϵ∼(0,I)Sample\,\, ε (0,I) 20: s←times[j],t←times[j+1]s← times[j], t← times[j+1] 21: Guidance←Ms→1Σs→1−1(r−Ms→1r,s(i)) Guidance← M_s→ 1 T _s→ 1^-1 (z_r-M_s→ 1x^(i)_r,s ) 22: Score←θr(r,s(i),s)+Guidance Score _ _r(x^(i)_r,s,s)+ Guidance 23: r,t(i)←Ms→tr,s(i)−Σs→tScore+Σs→t12ϵx^(i)_r,t← M_s→ tx^(i)_r,s- _s→ t Score+ _s→ t 12 ε 24: end for 25: r←1N∑i=1Nr,0(i)x_r← 1N _i=1^Nx^(i)_r,0 26: Set posterior ensemble: ^r(i)←r,0(i) x^(i)_r ^(i)_r,0 27: end if 28: end for Measurement update via reverse-time sampling. After training θkS_ _k, we perform the measurement update by running reverse-time sampling initialized from the perturbed prior at t=1−epst=1- eps for a small eps>0 eps>0. During sampling, we combine the learned prior score with the likelihood score (22) induced by the measurement kz_k to obtain posterior samples. Unless stated otherwise, we use and set the Number of Function Evaluations (NFE), nfe=500 nfe=500. Time update between measurement steps. After the measurement update at step k, we propagate the posterior ensemble forward under Eq. (1) up to the next measurement step in K, giving the prior ensemble for the next update. We used the Euler–Maruyama method Kloeden and Platen (1992), as in prior work Bao et al. (2024a). Evaluation. We report the estimated trajectory as the ensemble mean at each time step and evaluate accuracy using the root mean squared error (RMSE) Willmott and Matsuura (2005) or structural similarity index measure (SSIM) between the estimated and ground-truth trajectories Wang et al. (2004). The overall procedure is summarized in Algorithm 1. Figure 3: State trajectories for the Lorenz–63 system with measurement gap 100. Each panel shows the reference trajectory and the assimilated trajectory produced by one of the considered methods: (a) EnKF, (b) SF, (c) SSLS, and (d) MASF. The title of each subplot reports the trajectory RMSE for a representative run (seed 1), followed by the mean ± standard deviation of RMSE computed over five random seeds. Overall, MASF achieves consistently lower RMSE compared to the baselines. 5 Experimental Results We evaluated MASF on three benchmark datasets: Lorenz–63, Lorenz–96, and Kolmogorov flow. For the ordinary differential equation (ODE) benchmarks (Lorenz–63 and Lorenz–96), we compared against EnKF, SF, and SSLS; for Kolmogorov flow, we compared against SF and SSLS. Lorenz-63. Lorenz–63 is a three-dimensional nonlinear ODE system, originally introduced as a simplified model of atmospheric convection Lorenz (1963): x˙ x =σ(y−x),y˙=x(ρ−z)−y,z˙=xy−βz =σ(y-x),\, y=x(ρ-z)-y,\, z=xy-β z (28) where σ=10σ=10, β=8/3β=8/3, and ρ=28ρ=28. We integrated (28) using the Euler-Maruyama method with step size dt=0.01dt=0.01. The measurement equation is given by k=Ak+σϵkϵk∼(0,I). _k=Ax_k+σ ε_k ε_k (0,I). (29) where A=IA=I, σ=1σ=1, with measurements taken every 100 steps. Under this configuration, we compared the performance of MASF with that of EnKF, SF, and SSLS over the time interval from step 2000 to 2500. All methods used the same MLP architecture Bishop (2006); Perez et al. (2018); SF and MASF additionally incorporate time embeddings. Detailed architectural and training configurations are provided in Appendix D.1. Fig. 3 illustrates representative trajectories and RMSE values averaged over five seeds. All filters showed reasonable performance. However, the two score-based filters, SF and SSLS, performed slightly worse than EnKF (Fig. 3a–c). In contrast, MASF tracks the trajectory more faithfully than EnKF, with smaller accumulated error (Fig. 3d). Overall, MASF achieves the lowest RMSE and exhibits reduced variance across random seeds, indicating more stable state estimation. Figures for additional seeds are shown in Fig. 7. Figure 4: Performance on the Lorenz–96 system across state dimension, chaoticity, and measurement sparsity. Panels (a)–(b) vary the state dimension, (c)–(d) vary the forcing parameter, and (e)–(f) vary the measurement gap, with the remaining parameters fixed as indicated in each panel title. Across all three sweeps, MASF achieves consistently lower RMSE and shows robust performance under variations in dimension, forcing, and measurement gap. The mean ± standard deviation of RMSE computed over five random seeds. Lorenz-96. Lorenz–96 is a d-dimensional nonlinear dynamical system on a one-dimensional periodic lattice Lorenz (1996). The data assimilation difficulty scales with the state dimension d, and the forcing parameter F controls the degree of instability. The state equation is x˙i x_i =(xi+1−xi−2)xi−1−xi+F, =(x_i+1-x_i-2)x_i-1-x_i+F, (30) for i=1,…,di=1,…,d with cyclic indexing xi+d=xix_i+d=x_i. The measurement equation follows (29). We evaluated performance over steps 2525 to 100100 with step size dt=0.01dt=0.01, comparing EnKF, SSLS, SF, and MASF under a common experimental configuration. All methods used the same 1D U-Net architecture Stoller et al. (2018); Perslev et al. (2019); SF and MASF additionally incorporate time embeddings. Detailed architectural and training configurations are provided in Appendix D.2. We swept three factors that determine data assimilation difficulty—state dimension, chaoticity, and sparsity—as shown in Fig. 4. Specifically, we varied the state dimension d∈(256,512,1024,2048)d∈(256,512,1024,2048), the forcing parameter F∈(8,12,16,20,24)F∈(8,12,16,20,24) (default: 88), and the measurement gap Gap∈(5,10,15,20,25)Gap∈(5,10,15,20,25). As the state dimension increased, the performance of the EnKF degraded more steeply than that of the other methods (Fig. 4a). This became more pronounced under sparser measurements (Fig. 4b). Importantly, MASF showed improved performance across a range of dimensions under both fine and sparse measurements. As the chaoticity increased with the forcing parameter, the accuracy of the EnKF decreased dramatically, whereas the score-based filters were affected more moderately (Fig. 4c). MASF exhibited the greatest robustness to increasing chaoticity. This outperformance of MASF became more pronounced under sparser measurements (Fig. 4d). Finally, as the measurement gap increased, MASF showed the lowest RMSE and the most stable behavior compared to the other methods (Fig. 4e–f). Taken together, these results indicate that a measurement-aware design of the forward process (12), which allows for exact likelihood computation (22), is crucial for accurate high-dimensional data assimilation in challenging settings with strong chaoticity and sparse measurements. Figure 5: Performance on the Kolmogorov flow. (a) RMSE as a function of the measurement gap. Points show the mean over 5 random seeds and error bars indicate ± standard deviation across seeds. (b,c) RMSE over time for representative runs at gap=5=5 (b) and gap=25=25 (c) with seed 0. Open circles denote measurement-update steps; numbers in parentheses report the time-averaged RMSE for each method on the shown trajectory. Across gaps, MASF achieves the lowest mean RMSE compared to the baselines. Kolmogorov flow. Kolmogorov flow is a two-dimensional incompressible fluid benchmark in which each state is a velocity field t=(t)∈ℝ2×H×Wx_t=u(t) ^2× H× W (two channels for (u,v)(u,v)) on a periodic grid Meshalkin and Sinai (1961); Chandler and Kerswell (2013); Kochkov et al. (2021a). The state follows the incompressible Navier–Stokes equations with external forcing: ∂t _tu =−(⋅∇)+1Re∇2−1ρ∇p+, =-(u·∇)u+ 1Re∇^2u- 1ρ∇ p+f, ∇⋅ ∇·u =0, =0, (31) where u is the velocity field, p is pressure, ρ is density, f is external forcing, and ReRe is the Reynolds number. We used a periodic domain [0,2π]2[0,2π]^2 with ρ=1ρ=1 and Re=2000Re=2000, and simulated trajectories on a 64×6464× 64 grid using the JAX-CFD solver Kochkov et al. (2021b, a). We set the step size to dt=0.2dt=0.2. The measurement equation is given by k=⊙k+σϵkϵk∼(0,I), _k=M _k+σ ε_k ε_k (0,I), (32) where ∈0,11×1×H×HM∈\0,1\^1× 1× H× H is a pixel-wise mask with H=64H=64 and ⊙ is element-wise multiplication. With stride s, we set M:,:,i,j=1M_:,:,i,j=1 if i≡0(mods)i≡ 0\ (mod\ s) and j≡0(mods)j≡ 0\ (mod\ s), and 0 otherwise. We fixed s=5s=5 and σ=0.1σ=0.1, and varied the measurement gap by setting gap∈(5,10,15,25) gap∈(5,10,15,25). All methods used the same 2D U-Net architecture Ronneberger et al. (2015); SF and MASF additionally incorporate time embeddings. Figure 6: Estimated system state on Kolmogorov flow (gap==15, seed 0). Vorticity fields are shown at three representative time indices (τ=15,30,45τ=15,30,45). Top to bottom: reference state, sparse measurement, and reconstructions by SF, SSLS, and MASF. Numbers in each reconstruction panel report the per-frame SSIM with respect to the reference at the same τ. Row labels (e.g., MASF(0.9765)) indicate the average SSIM over the three displayed time points. MASF yields the highest SSIM and most faithful spatial structures across the shown times. We focused on comparing the performance of the score-based filters, SF, SSLS, and MASF, since EnKF is not well suited for high-dimensional problems (Fig. 4a–b). Specifically, we compared the performance of the three score-based filters on the 64×6464× 64 Kolmogorov flow benchmark while increasing the measurement gap (5/10/15/25) (Fig. 5). SF performed poorly across all gaps, likely due to an inappropriate likelihood approximation. In contrast, SSLS and MASF achieved comparable performance for gap 5, but MASF showed improved performance as the gap increased. This indicates that MASF is more robust to long-range prediction and sparse temporal supervision, which is necessary for real-world data assimilation problems. Fig. 6 visualizes the estimated system state at measured time points. Consistent with the quantitative trends in Fig. 5, MASF produced cleaner and more structurally faithful flow fields, whereas SF exhibited noticeable artifacts and SSLS showed increasing blur or distortion as the prediction horizon grew. This was also captured in SSIM, where MASF maintained higher similarity to the reference and appeared visually cleaner at the measurement-conditioned frames. Figures for other gaps are shown in Fig. 8 and Fig. 9. 6 Limitations and Future Works MASF has several limitations that suggest directions for future work. First, the current formulation assumes the same dimensions for the state and measurement spaces, although dimensional mismatch often occurs in real-world settings. A natural extension to address this is to combine representation learning, enabling filtering in a shared latent space Amendola et al. (2020); Fan et al. (2025); Pasmans et al. (2025). Second, the well-posedness of our moment-matching SDE relies on the measurement operator having a nonnegative spectrum. When this assumption is violated, the interpolation used to define the forward process may be invalid. Thus, alternative interpolation schemes may be required, such as extensions to the complex domain Higham (1986); Gawlik and Leok (2018). Third, extending the idea for moment-matching SDE to nonlinear measurements is not straightforward Solin and Särkkä (2019). In the linear case, closed-form moments yield a mean ODE and a Lyapunov equation for the covariance, which together specify a consistent forward process. For nonlinear measurements, moment constraints generally do not uniquely determine a globally consistent drift, and a single SDE that matches prescribed moments over time may not exist Faedo et al. (2021); Varona et al. (2019). Fourth, current implementations of score-based filters can be computationally demanding because the prior score may need to be retrained at each measurement step as the state distribution evolves. Designing shared-parameter models is a next step Becker et al. (2019). Finally, MASF can be viewed as a new conditional generation framework that combines an analytic likelihood score with a learned prior score. This perspective suggests broader applicability to conditional inference problems. 7 Conclusion We proposed a new score-based filter that explicitly merges the measurement equation into the forward process. This construction yields an exact likelihood along the perturbed trajectory, enabling construction of the posterior score from a learned prior score and an analytically computed likelihood score. As a result, the proposed method performs sequential measurement updates without ad hoc likelihood-score approximations. We showed that this theoretically well-grounded approach outperformed baselines, across chaotic and high-dimensional benchmarks with sparse measurements. Our results demonstrate the capability for robust state tracking and reconstruction in challenging real-world systems. References A. Aksoy, D. Dowell, and C. Snyder (2009) A multicase comparative assessment of the ensemble kalman filter for assimilation of radar observations. part i: storm-scale analyses. Monthly Weather Review 137, p. 1805–1824. Cited by: §1. M. Amendola, R. Arcucci, L. Mottet, C. Q. Casas, S. Fan, C. Pain, P. Linden, and Y. Guo (2020) Data assimilation in the latent space of a neural network. arXiv preprint arXiv:2012.12056. External Links: Document Cited by: §6. B. D. O. Anderson (1982) Reverse-time diffusion equation models. Stochastic Processes and their Applications. Cited by: Appendix B, §3.1. C. Andrieu, A. Doucet, and R. Holenstein (2010) Particle markov chain monte carlo methods. Journal of the Royal Statistical Society: Series B 72, p. 269–342. Cited by: §1. M. S. Arulampalam, S. Maskell, N. Gordon, and T. Clapp (2002) A tutorial on particle filters for online nonlinear/non-Gaussian Bayesian tracking. IEEE Transactions on Signal Processing. Cited by: §1. M. Asch, M. Bocquet, and M. Nodet (2016) Data assimilation: methods, algorithms, and applications. SIAM. Cited by: §1, §1, §3.1. F. Bao, Z. Zhang, and G. Zhang (2024a) A score-based filter for nonlinear data assimilation. Journal of Computational Physics. Cited by: §1, §1, §4. F. Bao, Z. Zhang, and G. Zhang (2024b) An ensemble score filter for tracking high-dimensional nonlinear dynamical systems. Computer Methods in Applied Mechanics and Engineering 432, p. 117447. External Links: Document Cited by: §1. P. Becker, H. Pandya, G. Gebhardt, C. Zhao, J. Taylor, and G. Neumann (2019) Recurrent kalman networks: factorized inference in high-dimensional deep feature spaces. In International Conference on Machine Learning (ICML), Cited by: §6. C. Bishop (2006) Pattern recognition and machine learning. Springer. Cited by: §D.1, §D.3, §5. A. Carrassi, M. Bocquet, L. Bertino, and G. Evensen (2018) Data assimilation in the geosciences: an overview of methods, issues, and perspectives. Wiley Interdisciplinary Reviews: Climate Change. Cited by: §3.1. G. J. Chandler and R. R. Kerswell (2013) Invariant recurrent solutions embedded in a turbulent two-dimensional kolmogorov flow. Journal of Fluid Mechanics 722, p. 554–595. External Links: Document Cited by: §D.3, §5. H. G. Chipilski, X. Wang, and D. B. Parsons (2020) Impact of assimilating pecan profilers on the prediction of bore-driven nocturnal convection: a multiscale forecast evaluation for the 6 july 2015 case study. Monthly Weather Review 148, p. 1147–1175. Cited by: §1. N. Cogan, F. Bao, R. Paus, and A. Dobreva (2021) Data assimilation of synthetic data as a novel strategy for predicting disease progression in alopecia areata. Mathematical Medicine and Biology. Cited by: §1. P. Dhariwal and A. Nichol (2021) Diffusion models beat gans on image synthesis. External Links: 2105.05233, Link Cited by: §1. Z. Ding, C. Duan, Y. Jiao, J. Z. Yang, C. Yuan, and P. Zhang (2025) Nonlinear assimilation via score-based sequential Langevin sampling. External Links: 2411.13443, Link Cited by: §1, §1. A. Doucet, N. de Freitas, and N. Gordon (Eds.) (2001) Sequential monte carlo methods in practice. Springer. Cited by: §1. G. Evensen (2009a) Data assimilation: the ensemble kalman filter. Springer. Cited by: §1. G. Evensen (2009b) The ensemble kalman filter for combined state and parameter estimation: monte carlo techniques for data assimilation in large systems. IEEE Control Systems Magazine 29, p. 83–104. Cited by: §1. N. Faedo, G. Scarciotti, Member, A. Astolfi, and J. V. Ringwood (2021) On the approximation of moments for nonlinear systems. IEEE. Cited by: §6. H. Fan, Y. Liu, Z. Huo, Y. Liu, Y. Shi, and Y. Li (2025) A novel latent space data assimilation framework with autoencoder-observation to latent space. Monthly Weather Review. Cited by: §6. E. S. Gawlik and M. Leok (2018) Interpolation on symmetric spaces via the generalized polar decomposition. Foundations of Computational Mathematics. Cited by: §6. N. Gordon, D. Salmond, and A. Smith (1993) Novel approach to nonlinear/non-gaussian bayesian state estimation. IEE Proceedings F 140, p. 107–113. Cited by: §1. N. J. Higham (1986) Computing the polar decomposition—with applications. SIAM Journal on Scientific and Statistical Computing. Cited by: §6. J. Ho, A. Jain, and P. Abbeel (2020) Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, Vol. 33, p. 6840–6851. Cited by: §2.2. P. L. Houtekamer and H. L. Mitchell (1998) Data assimilation using an ensemble kalman filter technique. Monthly Weather Review. Cited by: §1. A. Hyvärinen (2005) Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research. Cited by: §1, §2.2. R. E. Kalman (1960) A new approach to linear filtering and prediction problems. Journal of Basic Engineering. Cited by: §1. P. E. Kloeden and E. Platen (1992) Numerical solution of stochastic differential equations. Springer. Cited by: §3.1, §4. D. Kochkov, J. A. Smith, A. Alieva, Q. Wang, M. P. Brenner, and S. Hoyer (2021a) Machine learning–accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences 118 (21), p. e2101784118. External Links: Document Cited by: §5, §5. D. Kochkov, J. A. Smith, P. Norgaard, G. Dresdner, A. Alieva, and S. Hoyer (2021b) JAX-cfd: computational fluid dynamics in jax. Note: https://github.com/google/jax-cfdAccessed 2026-01-28 Cited by: §5. K. J. H. Law, A. M. Stuart, and K. C. Zygalakis (2015) Data assimilation: a mathematical introduction. Springer. Cited by: §2.1, §3.1. E. N. Lorenz (1963) Deterministic nonperiodic flow. Journal of the Atmospheric Sciences 20 (2), p. 130–141. External Links: Document Cited by: §D.1, §5. E. N. Lorenz (1996) Predictability: a problem partly solved. In Proceedings of the Seminar on Predictability, Vol. I, Cited by: §D.2, §5. L. D. Meshalkin and I. G. Sinai (1961) Investigation of the stability of a stationary solution of a system of equations for the plane movement of an incompressible viscous liquid. Journal of Applied Mathematics and Mechanics 25 (6), p. 1700–1705. External Links: Document Cited by: §D.3, §5. A. Nichol and P. Dhariwal (2021) Improved denoising diffusion probabilistic models. arXiv preprint arXiv:2102.09672. Cited by: §2.2. I. Pasmans, Y. Chen, T. S. Finn, M. Bocquet, and A. Carrassi (2025) Ensemble kalman filter in latent space using a variational autoencoder pair. arXiv preprint arXiv:2502.12987. Cited by: §6. E. Perez, F. Strub, H. de Vries, V. Dumoulin, and A. Courville (2018) FiLM: visual reasoning with a general conditioning layer. In Proceedings of the AAAI Conference on Artificial Intelligence, Cited by: §D.1, §D.3, §5. M. Perslev, M. H. Jensen, S. Darkner, P. J. Jennum, and C. Igel (2019) U-time: a fully convolutional network for time series segmentation applied to sleep staging. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: §D.2, §5. S. Reich and C. Cotter (2015) Probabilistic forecasting and bayesian data assimilation. Cambridge University Press. Cited by: §1. O. Ronneberger, P. Fischer, and T. Brox (2015) U-Net: convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention (MICCAI), Cited by: §5. F. Rozet and G. Louppe (2023) Score-based data assimilation. In Thirty-seventh Conference on Neural Information Processing Systems, External Links: Link Cited by: §1. S. Särkkä (2013) Bayesian filtering and smoothing. Cambridge University Press. Cited by: §1, §2.1. C. Snyder, T. Bengtsson, P. Bickel, and J. Anderson (2008) Obstacles to high-dimensional particle filtering. Technical report Mathematical Advances in Data Assimilation. Cited by: §1. A. Solin and S. Särkkä (2019) Applied stochastic differential equations. Cambridge University Press. Cited by: §6. Y. Song and S. Ermon (2019) Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: §1, §1. Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole (2021) Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), Cited by: §1, §1, §2.2, §2.2. D. Stoller, S. Ewert, and S. Dixon (2018) Wave-u-net: a multi-scale neural network for end-to-end audio source separation. In Proceedings of the 19th International Society for Music Information Retrieval Conference (ISMIR), Cited by: §D.2, §5. M. C. Varona, R. Gebhart, J. Suk, and B. Lohmann (2019) Practicable simulation-free model order reduction by nonlinear moment matching. arXiv preprint arXiv:1901.10750. Cited by: §6. P. Vincent (2011) A connection between score matching and denoising autoencoders. Neural Computation. Cited by: §1, §2.2. Z. Wang, A. C. Bovik, H. R. Sheikh, and E. P. Simoncelli (2004) Image quality assessment: from error visibility to structural similarity. IEEE Transactions on Image Processing 13 (4), p. 600–612. External Links: Document Cited by: §4. J. S. Whitaker and T. M. Hamill (2002) Ensemble data assimilation without perturbed observations. Monthly Weather Review. Cited by: §1. C. J. Willmott and K. Matsuura (2005) Advantages of the mean absolute error (mae) over the root mean square error (rmse) in assessing average model performance. Climate Research. Cited by: §4. Appendix A Derivation of the moment-matching SDE Problem setup. Fix t∈[0,1]t∈[0,1]. Suppose that, for every ∈ℝdx ^d, the conditional law of XtX_t given X0=X_0=x is prescribed as Xt∣(X0=)∼(A(t),Σ(t)), X_t (X_0=x)\; \;N\! (A(t)x,\, (t) ), (33) where A:[0,1]→ℝd×dA:[0,1] ^d× d is differentiable with A(0)=IA(0)=I and A(1)=A(1)=A. We assume A(t)A(t) is invertible for all t∈[0,1)t∈[0,1) and the endpoint A(1)A(1) need not be invertible. Moreover, Σ:[0,1]→ℝd×d :[0,1] ^d× d is differentiable, symmetric, and positive semidefinite, with Σ(0)=0 (0)=0 and Σ(1)=σ2I (1)=σ^2I. We seek a linear SDE dXt=F(t)Xtdt+G(t)dBt, dX_t=F(t)X_t\,dt+G(t)\,dB_t, (34) whose solution matches (33). Theorem A.1 (Moment-matching SDE). Let A(⋅)A(·) and Σ(⋅) (·) be as in (33), with A(t)A(t) invertible for all t∈[0,1)t∈[0,1). Consider (34) with F:[0,1)→ℝd×dF:[0,1) ^d× d and G:[0,1)→ℝd×dG:[0,1) ^d× d satisfying, for all t∈[0,1)t∈[0,1), F(t) F(t) =A˙(t)A(t)−1, = A(t)\,A(t)^-1, (35) G(t)G(t) G(t)G(t) T =Σ˙(t)−F(t)Σ(t)−Σ(t)F(t),G(t)G(t)is symmetric and⪰0. = (t)-F(t) (t)- (t)F(t) T, G(t)G(t) T\ is symmetric and\ 0. (36) Assume additionally that F and G are locally bounded on [0,1)[0,1) (e.g., continuous on [0,T][0,T] for every T<1T<1). Then for every T<1T<1, the SDE (34) admits a unique strong solution on [0,T][0,T], and for every ∈ℝdx ^d and all t∈[0,T]t∈[0,T], Xt∣(X0=)∼(A(t),Σ(t)). X_t (X_0=x) \! (A(t)x,\, (t) ). (37) Moreover, if the limits A:=limt→1−A(t)A:= _t→ 1^-A(t) and Σ:=limt→1−Σ(t) := _t→ 1^- (t) exist, then Xt∣(X0=)⇒(A,Σ)as t→1−, X_t (X_0=x) \! (Ax, ) t→ 1^-, (38) where ⇒ denotes weak convergence. Equivalently, one may define X1:=AX0+Σ12ϵX_1:=AX_0+ 12 ε with ϵ∼(0,I) ε (0,I). Before proving Theorem A.1, we establish two lemmas. Lemma A.2 (Variation-of-constants formula). Let Φ:[0,1)→ℝd×d :[0,1) ^d× d be the fundamental matrix solving Φ˙(t)=F(t)Φ(t),Φ(0)=I, (t)=F(t) (t), (0)=I, (39) and assume Φ(t) (t) is invertible for all t∈[0,1)t∈[0,1). Then the unique strong solution of (34) satisfies, for all t∈[0,1)t∈[0,1), Xt=Φ(t)X0+∫0tΦ(t)Φ(s)−1G(s)Bs. X_t= (t)X_0+ _0^t (t) (s)^-1G(s)\,dB_s. (40) If Φ(1−):=limt→1−Φ(t) (1^-):= _t→ 1^- (t) exists and the Itô integral in (40) converges in L2L^2 as t→1−t→ 1^-, then the representation extends to t=1t=1: X1=Φ(1−)X0+∫01Φ(1−)Φ(s)−1G(s)Bs. X_1= (1^-)X_0+ _0^1 (1^-) (s)^-1G(s)\,dB_s. (41) Proof. Define Yt:=Φ(t)−1XtY_t:= (t)^-1X_t. Differentiating Φ(t)Φ(t)−1=I (t) (t)^-1=I gives dtΦ(t)−1=−Φ(t)−1F(t) ddt (t)^-1=- (t)^-1F(t). Applying Itô’s formula to Yt=Φ(t)−1XtY_t= (t)^-1X_t and using dXt=F(t)Xtdt+G(t)dBtdX_t=F(t)X_t\,dt+G(t)\,dB_t yields dYt dY_t =d(Φ(t)−1)Xt+Φ(t)−1dXt=−Φ(t)−1F(t)Xtdt+Φ(t)−1(F(t)Xtdt+G(t)dBt) =d( (t)^-1)\,X_t+ (t)^-1\,dX_t=- (t)^-1F(t)X_t\,dt+ (t)^-1 (F(t)X_t\,dt+G(t)\,dB_t ) =Φ(t)−1G(t)dBt. = (t)^-1G(t)\,dB_t. (42) Integrating from 0 to t gives Yt=X0+∫0tΦ(s)−1G(s)BsY_t=X_0+ _0^t (s)^-1G(s)\,dB_s. Multiplying by Φ(t) (t) yields (40). The extension to t=1t=1 follows by the stated limits. ∎ Lemma A.3 (Matching the conditional mean). Fix 0∈ℝdx_0 ^d and define m(t):=[Xt∣X0=0]m(t):=E[X_t X_0=x_0]. Assume A:[0,1)→ℝd×dA:[0,1) ^d× d is differentiable and invertible for all t∈[0,1)t∈[0,1) with A(0)=IA(0)=I. If F(t)=A˙(t)A(t)−1F(t)= A(t)A(t)^-1 on [0,1)[0,1), then m(t)=A(t)0m(t)=A(t)x_0 for all t∈[0,1)t∈[0,1). Moreover, the fundamental matrix Φ solving (39) satisfies Φ(t)=A(t) (t)=A(t) on [0,1)[0,1). Proof. Conditioning on X0=0X_0=x_0 and taking conditional expectation in (34) yields m˙(t)=F(t)m(t),m(0)=0. m(t)=F(t)m(t), m(0)=x_0. (43) Let Φ be the fundamental matrix from LemmaA.2. Then m(t)=Φ(t)0m(t)= (t)x_0. Define Ψ(t):=A(t)−1Φ(t) (t):=A(t)^-1 (t). Using dtA(t)−1=−A(t)−1A˙(t)A(t)−1 ddtA(t)^-1=-A(t)^-1 A(t)A(t)^-1 and Φ˙(t)=F(t)Φ(t) (t)=F(t) (t), Ψ˙(t) (t) =dt(A(t)−1Φ(t))=−A(t)−1A˙(t)A(t)−1Φ(t)+A(t)−1Φ˙(t) = ddt (A(t)^-1 (t) )=-A(t)^-1 A(t)A(t)^-1 (t)+A(t)^-1 (t) =−A(t)−1A˙(t)A(t)−1Φ(t)+A(t)−1F(t)Φ(t)=0. =-A(t)^-1 A(t)A(t)^-1 (t)+A(t)^-1F(t) (t)=0. (44) Thus Ψ(t)≡Ψ(0)=A(0)−1Φ(0)=I (t)≡ (0)=A(0)^-1 (0)=I, so Φ(t)=A(t) (t)=A(t) and m(t)=A(t)0m(t)=A(t)x_0. ∎ Proposition A.4 (Lyapunov equation for the conditional covariance). Fix ∈ℝdx ^d and define the centered process Yt:=Xt−m(t)Y_t:=X_t-m(t), where m(t)=[Xt∣X0=]m(t)=E[X_t X_0=x]. Then the conditional covariance ΣX(t):=Cov(Xt∣X0=)=(YtYt∣X0=) _X(t):=Cov(X_t X_0=x)=E\! (Y_tY_t T X_0=x ) (45) satisfies the matrix ODE Σ˙X(t)=F(t)ΣX(t)+ΣX(t)F(t)+G(t)G(t),ΣX(0)=0. _X(t)=F(t) _X(t)+ _X(t)F(t) T+G(t)G(t) T, _X(0)=0. (46) Proof. Since dm(t)=F(t)m(t)dtdm(t)=F(t)m(t)\,dt, subtracting dm(t)dm(t) from (34) yields dYt=F(t)Ytdt+G(t)dBt,Y0=0. dY_t=F(t)Y_t\,dt+G(t)\,dB_t, Y_0=0. (47) Apply Itô’s product rule to YtYtY_tY_t T: d(YtYt)=(dYt)Yt+Yt(dYt)+(dYt)(dYt). d(Y_tY_t T)=(dY_t)Y_t T+Y_t(dY_t) T+(dY_t)(dY_t) T. (48) Using dBtdBt=IdtdB_t\,dB_t T=I\,dt and taking conditional expectations given X0=X_0=x eliminates the local martingale terms, yielding (46). ∎ Lemma A.5 (Matching the covariance). Assume G is chosen so that (36) holds. Then Cov(Xt∣X0=)=Σ(t),∀t∈[0,1). (X_t X_0=x)= (t), ∀\,t∈[0,1). (49) Proof. By PropositionA.4, ΣX _X satisfies (46). If (36) holds, then Σ satisfies the same ODE with the same initial condition: Σ˙(t)=F(t)Σ(t)+Σ(t)F(t)+G(t)G(t),Σ(0)=0. (t)=F(t) (t)+ (t)F(t) T+G(t)G(t) T, (0)=0. (50) Uniqueness of solutions to this matrix ODE implies ΣX(t)=Σ(t) _X(t)= (t) for all t∈[0,1)t∈[0,1). ∎ Corollary A.6 (Moment-matching with linear interpolation). Fix a matrix A∈ℝd×dA ^d× d and define the interpolation A(t):=(1−a(t))A+a(t)I,t∈[0,1), A(t):=(1-a(t))A+a(t)I, t∈[0,1), (51) and assume that A(t)A(t) is invertible for all t∈[0,1)t∈[0,1). Define F(t):=A˙(t)A(t)−1,t∈[0,1), F(t):= A(t)\,A(t)^-1, t∈[0,1), (52) and let Φ(t) (t) be the fundamental matrix solving Φ˙(t)=F(t)Φ(t),Φ(0)=I. (t)=F(t) (t), (0)=I. (53) Then Φ(t)=A(t) (t)=A(t) for all t∈[0,1)t∈[0,1). In particular, if a(t)→0a(t)→ 0 as t→1t→ 1, then Φ(1−):=limt→1Φ(t)=A, (1-):= _t→ 1 (t)=A, (54) XtX_t admits a limit in distribution as t→1t→ 1, and the terminal random variable X1X_1 may be defined by Xt⇒X1X_t X_1 as t→1t→ 1, with representation X1=AX0+Σ12ϵ, X_1=AX_0+ 12 ε, (55) where ϵ∼(0,I) ε (0,I) and Σ12(Σ12)=Σ 12( 12) T= . Proof. Since A(t)A(t) is invertible on [0,1)[0,1), F(t)F(t) is well-defined. Moreover, F(t)A(t)=A˙(t)A(t)−1A(t)=A˙(t), F(t)A(t)= A(t)A(t)^-1A(t)= A(t), (56) so A(t)A(t) solves the linear matrix ODE A˙(t)=F(t)A(t),A(0)=I. A(t)=F(t)A(t), A(0)=I. (57) By definition, Φ(t) (t) solves Φ˙(t)=F(t)Φ(t),Φ(0)=I. (t)=F(t) (t), (0)=I. (58) Hence A(t)A(t) and Φ(t) (t) satisfy the same ODE with the same initial condition, and by uniqueness of solutions to M˙(t)=F(t)M(t) M(t)=F(t)M(t) with M(0)=IM(0)=I, we have Φ(t)=A(t) (t)=A(t) for all t∈[0,1)t∈[0,1). If a(t)→0a(t)→ 0 as t→1t→ 1, then limt→1A(t)=limt→1((1−a(t))A+a(t)I)=A, _t→ 1A(t)= _t→ 1 ((1-a(t))A+a(t)I )=A, (59) and therefore Φ(1−)=limt→1Φ(t)=A (1-)= _t→ 1 (t)=A. Finally, if Σ:=limt→1Σ(t) := _t→ 1 (t) exists, then the conditional laws Xt∣(X0=)∼(A(t),Σ(t))X_t (X_0=x) (A(t)x, (t)) converge as t→1t→ 1 to (A,Σ)N(Ax, ), and one may represent the limit by X1=AX0+Σ12ϵ, X_1=AX_0+ 12 ε, (60) with ϵ∼(0,I)ε (0,I) and Σ12(Σ12)=Σ 12( 12) T= . ∎ Corollary A.7 (Moment-matching linear with affine transformation). Let A:[0,1]→ℝd×dA:[0,1] ^d× d and Σ:[0,1]→ℝd×d :[0,1] ^d× d be as in (33), with A(0)=IA(0)=I, Σ(0)=0 (0)=0, and A(t)A(t) invertible for all t∈[0,1)t∈[0,1). Let b:[0,1]→ℝdb:[0,1] ^d be absolutely continuous with b(0)=0b(0)=0. Consider the linear SDE dXt=F(t)Xtdt+f(t)dt+G(t)dBt, dX_t=F(t)X_t\,dt+f(t)\,dt+G(t)\,dB_t, (61) where F:[0,1)→ℝd×dF:[0,1) ^d× d, f:[0,1)→ℝdf:[0,1) ^d, and G:[0,1)→ℝd×dG:[0,1) ^d× d are measurable. Assume that, for all t∈[0,1)t∈[0,1), F(t) F(t) =A˙(t)A(t)−1, = A(t)\,A(t)^-1, (62) f(t) f(t) =b˙(t)−F(t)b(t), = b(t)-F(t)b(t), (63) G(t)G(t) G(t)G(t) T =Σ˙(t)−F(t)Σ(t)−Σ(t)F(t). = (t)-F(t) (t)- (t)F(t) T. (64) Then the unique strong solution of (61) satisfies, for every ∈ℝdx ^d and all t∈[0,1)t∈[0,1), Xt∣(X0=)∼(A(t)+b(t),Σ(t)). X_t (X_0=x) \! (A(t)x+b(t),\, (t) ). (65) Moreover, if A:=limt→1A(t)A:= _t→ 1A(t), b:=limt→1b(t)b:= _t→ 1b(t), and Σ:=limt→1Σ(t) := _t→ 1 (t) exist, then Xt∣(X0=)⇒(A+b,Σ)as t→1, X_t (X_0=x) \! (Ax+b,\, ) t→ 1, (66) It can be defined as X1:=AX0+b+Σ12ϵ, X_1:=AX_0+b+ 12 ε, (67) where ϵ∼(0,I) ε (0,I) and Σ12(Σ12)=Σ 12( 12) T= . Appendix B Transition law and likelihood score Theorem B.1 (Gaussian transition of the moment-matching SDE). Consider the linear SDE dXt=F(t)Xtdt+G(t)dBt, dX_t=F(t)X_t\,dt+G(t)\,dB_t, (68) with fundamental matrix Φ solving Φ˙(t)=F(t)Φ(t) (t)=F(t) (t) and Φ(0)=I (0)=I. Assume that the conditional law given X0X_0 is Gaussian with Xt∣X0∼(A(t)X0,Σ(t)), X_t X_0 \! (A(t)X_0, (t) ), (69) and that A(t)A(t) is invertible for all t∈[0,1)t∈[0,1). Then for any 0<s<t≤10<s<t≤ 1, the transition is Gaussian: Xt∣Xs∼(Ms→tXs,Σs→t), X_t X_s \! (M_s→ tX_s, _s→ t ), (70) where Ms→t M_s→ t =A(t)A(s)−1, =A(t)A(s)^-1, (71) Σs→t _s→ t =Σ(t)−Ms→tΣ(s)Ms→t. = (t)-M_s→ t\, (s)\,M_s→ t T. (72) The proof proceeds by introducing supporting lemmas and then combining them to conclude (70)–(72). Lemma B.2 (Variation-of-constants formula). Let Φ be the fundamental matrix of (68). Then for any 0≤s<t<10≤ s<t<1, Xt=Φ(t)Φ(s)−1Xs+∫stΦ(t)Φ(u)−1G(u)Bu. X_t= (t) (s)^-1X_s+ _s^t (t) (u)^-1G(u)\,dB_u. (73) Proof. Apply LemmaA.2 on [0,t][0,t] and rewrite the resulting expression conditionally on time s; equivalently, apply the same argument to the shifted process on [s,t][s,t]. ∎ Lemma B.3 (Identification of the linear operator). Let 0≤s<t<10≤ s<t<1 and define Ms→t:=Φ(t)Φ(s)−1M_s→ t:= (t) (s)^-1. If F(t)=A˙(t)A(t)−1F(t)= A(t)A(t)^-1 and A(0)=IA(0)=I, then Φ(t)=A(t) (t)=A(t) for all t∈[0,1)t∈[0,1) and hence Ms→t=A(t)A(s)−1. M_s→ t=A(t)A(s)^-1. (74) Proof. Both Φ and A solve the matrix ODE Y˙(t)=F(t)Y(t) Y(t)=F(t)Y(t) with the same initial condition Y(0)=IY(0)=I. Uniqueness of solutions to linear ODEs implies Φ(t)=A(t) (t)=A(t) on [0,1)[0,1). ∎ Lemma B.4 (Gaussian increment and transition covariance). Fix 0≤s<t<10≤ s<t<1 and let Φ be the fundamental matrix of (68). Define Ms→t M_s→ t :=Φ(t)Φ(s)−1, := (t) (s)^-1, (75) ηs→t _s→ t :=∫stΦ(t)Φ(u)−1G(u)Bu. := _s^t (t) (u)^-1G(u)\,dB_u. (76) Then ηs→t _s→ t is Gaussian with [ηs→t]=0E[ _s→ t]=0 and is independent of σ(Xr:r≤s)σ(X_r:\,r≤ s). Moreover, under (69), Cov(ηs→t)=Σ(t)−Ms→tΣ(s)Ms→t. ( _s→ t)= (t)-M_s→ t (s)M_s→ t T. (77) Proof. Since the integrand in (76) is deterministic, ηs→t _s→ t is an Itô integral of a deterministic function against Brownian motion and is therefore Gaussian with mean zero. It depends only on the increment Bu−Bs:u∈[s,t]\B_u-B_s:\,u∈[s,t]\, hence it is independent of ℱs:=σ(Br:r≤s)F_s:=σ(B_r:\,r≤ s), and thus independent of σ(Xr:r≤s)⊆ℱsσ(X_r:\,r≤ s) _s. By LemmaB.2, Xt=Ms→tXs+ηs→tX_t=M_s→ tX_s+ _s→ t. Taking conditional covariance given X0X_0 and using the stated independence gives Cov(Xt∣X0)=Ms→tCov(Xs∣X0)Ms→t+Cov(ηs→t). (X_t X_0)=M_s→ t\,Cov(X_s X_0)\,M_s→ t T+Cov( _s→ t). (78) Substituting Cov(Xt∣X0)=Σ(t)Cov(X_t X_0)= (t) and Cov(Xs∣X0)=Σ(s)Cov(X_s X_0)= (s) from (69) yields (77). ∎ Proof. From Theorem B.1, the conditional density is p(∣t) p(z _t) =(2π)−d/2|Σt→1|−12exp(−12(−Mt→1t)Σt→1−1(−Mt→1t)). =(2π)^-d/2\,| _t→ 1|^- 12 \! (- 12(z-M_t→ 1x_t) T _t→ 1^-1(z-M_t→ 1x_t) ). (79) Hence logp(∣t) p(z _t) =−12(−Mt→1t)Σt→1−1(−Mt→1t)−12log|Σt→1|−d2log(2π). =- 12(z-M_t→ 1x_t) T _t→ 1^-1(z-M_t→ 1x_t)- 12 | _t→ 1|- d2 (2π). (80) The last two terms do not depend on tx_t. For the quadratic term, using ∇t(−Mt→1t)=−Mt→1 _x_t(z-M_t→ 1x_t)=-M_t→ 1 and the symmetry of Σt→1−1 _t→ 1^-1, we obtain ∇tlogp(∣t) _x_t p(z _t) =−12∇t[(−Mt→1t)Σt→1−1(−Mt→1t)] =- 12\, _x_t [(z-M_t→ 1x_t) T _t→ 1^-1(z-M_t→ 1x_t) ] (81) =−12[(−Mt→1)Σt→1−1(−Mt→1t)+(−Mt→1)Σt→1−1(−Mt→1t)] =- 12 [(-M_t→ 1) T _t→ 1^-1(z-M_t→ 1x_t)+(-M_t→ 1) T _t→ 1^-1(z-M_t→ 1x_t) ] (82) =Mt→1Σt→1−1(−Mt→1t). =M_t→ 1 T _t→ 1^-1 (z-M_t→ 1x_t ). (83) ∎ Theorem B.5 (Reverse-time Sampling). Let ptp_t be the marginal density of XtX_t from (6). The reverse-time SDE (run backward from s to t with t<st<s) is given by (23): dXt dX_t =(F(t)Xt−G(t)G(t)∇logpt(Xt))dt+G(t)dB¯t. = (F(t)X_t-G(t)G(t) T _x p_t(X_t) )\,dt+G(t)\,d B_t. (84) Conditioning on a measurement z replaces the score by the posterior score ∇logpt(∣) _x p_t(x ) =∇logpt()+∇logp(∣). = _x p_t(x)+ _x p(z ). (85) A one-step reverse-time sampling approximation over the step s→ts→ t is Xt X_t ≈Ms→tXs−Σs→t∇logps(Xs∣)+Σs→t12ϵ,ϵ∼(0,I), ≈ M_s→ tX_s- _s→ t\, _x p_s(X_s )+ _s→ t 12 ε, ε (0,I), (86) Proof. The reverse-time SDE (23) is the standard result of AndersonAnderson [1982]. We derive the time-reverse sampling (25) from the solution of (23). By variation of constants, for t<st<s the solution of (23) satisfies Xt X_t =Ms→tXs−∫stMu→tG(u)G(u)∇logpu(Xu∣)u+∫stMu→tG(u)B¯u, =M_s→ tX_s- _s^tM_u→ t\,G(u)G(u) T\, _x p_u(X_u )\,du+ _s^tM_u→ t\,G(u)\,d B_u, (87) Define P(t) P(t) :=A(t)−1Σ(t)A(t)−,equivalentlyΣ(t)=A(t)P(t)A(t). :=A(t)^-1 (t)A(t)^- T, (t)=A(t)P(t)A(t) T. (88) Differentiating Σ(t)=A(t)P(t)A(t) (t)=A(t)P(t)A(t) T and using A′(t)=F(t)A(t)A (t)=F(t)A(t) yields Σ˙(t) (t) =F(t)Σ(t)+Σ(t)F(t)+A(t)P˙(t)A(t). =F(t) (t)+ (t)F(t) T+A(t) P(t)A(t) T. (89) On the other hand, the covariance of the linear SDE (6) satisfies Σ˙(t) (t) =F(t)Σ(t)+Σ(t)F(t)+G(t)G(t). =F(t) (t)+ (t)F(t) T+G(t)G(t) T. (90) Comparing (89) and (90) gives A(t)P˙(t)A(t) A(t) P(t)A(t) T =G(t)G(t),henceP˙(t)=A(t)−1G(t)G(t)A(t)−. =G(t)G(t) T, P(t)=A(t)^-1G(t)G(t) TA(t)^- T. (91) We state the approximation that is central for the one-step sampler: over a single step s→ts→ t, we freeze the score term at time s and evaluate it at the current iterate XsX_s, ∇logpu(Xu∣)≈∇logps(Xs∣),u∈[t,s]. _x p_u(X_u )\;≈\; _x p_s(X_s ), u∈[t,s]. (92) Substituting (92) into the drift integral in (87) gives ∫stMu→tG(u)G(u)∇logpu(Xu∣)u≈∫stMu→tG(u)G(u)u∇logps(Xs∣) _s^tM_u→ t\,G(u)G(u) T\, _x p_u(X_u )\,du≈ _s^tM_u→ t\,G(u)G(u) T\,du\; _x p_s(X_s ) (93) =(∫stA(t)A(u)−1G(u)G(u)A(u)−A(t)u)∇logps(Xs∣) = ( _s^tA(t)A(u)^-1\,G(u)G(u) T\,A(u)^- TA(t) T\,du )\, _x p_s(X_s ) (94) =A(t)(∫stA(u)−1G(u)G(u)A(u)−u)A(t)∇logps(Xs∣) =A(t) ( _s^tA(u)^-1\,G(u)G(u) T\,A(u)^- T\,du )A(t) T\, _x p_s(X_s ) (95) =A(t)(P(t)−P(s))A(t)∇logps(Xs∣) =A(t) (P(t)-P(s) )A(t) T\, _x p_s(X_s ) (96) =Σs→t∇logps(Xs∣), = _s→ t\, _x p_s(X_s ), (97) Applying the substitutions (93) to (87) yields Xt X_t ≈Ms→tXs−Σs→t∇logps(Xs∣)+Σs→t12ϵ,ϵ∼(0,I), ≈ M_s→ tX_s- _s→ t\, _x p_s(X_s )+ _s→ t 12 ε, ε (0,I), (98) which is (25). ∎ Appendix C Invertibility of A(t)A(t) Sufficient spectral condition. Assume throughout that a(t)∈(0,1]a(t)∈(0,1] with a(0)=1a(0)=1 and a(1)=0a(1)=0, and define A(t):=(1−a(t))A+a(t)I. A(t)\;:= (1-a(t) )A+a(t)I. (99) A convenient sufficient condition for invertibility on t∈[0,1)t∈[0,1) is σ(A)⊂[0,∞), σ(A)⊂[0,∞), (100) i.e., all eigenvalues of A are real and nonnegative. Lemma C.1 (Invertibility of A(t)A(t) for t∈[0,1)t∈[0,1)). Assume (100) and a(t)∈(0,1]a(t)∈(0,1] for t∈[0,1)t∈[0,1). Then A(t)A(t) is invertible for every t∈[0,1)t∈[0,1). Moreover, its spectrum satisfies σ(A(t))⊂[a(t),∞). σ\! (A(t) )⊂[\,a(t),\,∞). (101) Proof. Let λ∈σ(A)λ∈σ(A) and let v≠0v≠ 0 satisfy Av=λvAv=λ v. Then by (99), A(t)v A(t)v =(((1−a(t))A+a(t)I)v = (( (1-a(t) )A+a(t)I )v (102) =(1−a(t))λv+a(t)v = (1-a(t) )λ v+a(t)v (103) =((1−a(t))λ+a(t))v. = ( (1-a(t) )λ+a(t) )v. (104) Hence μ(t):=(1−a(t))λ+a(t)μ(t):= (1-a(t) )λ+a(t) is an eigenvalue of A(t)A(t). Since λ≥0λ≥ 0 and a(t)∈(0,1]a(t)∈(0,1], μ(t)≥a(t)> 0,∀t∈[0,1). μ(t)\;≥\;a(t)\;>\;0, ∀\,t∈[0,1). (105) Therefore all eigenvalues of A(t)A(t) are strictly positive, and A(t)A(t) is invertible for every t∈[0,1)t∈[0,1). The spectral inclusion σ(A(t))⊂[a(t),∞)σ(A(t))⊂[\,a(t),\,∞) follows from the same bound. ∎ Example: grid mask. A grid mask keeps some coordinates and zeros out the rest. Thus it can be written as A=diag(m1,…,md),mi∈0,1. A=diag(m_1,…,m_d), m_i∈\0,1\. (106) Consequently, the eigenvalues of A are exactly its diagonal entries: σ(A)=mi:i=1,…,md⊂0,1⊂[0,∞). σ(A)=\m_i:\;i=1,…,m_d\⊂\0,1\⊂[0,∞). (107) Hence the sufficient condition (100) holds, and by Lemma C.1, A(t)A(t) is invertible for every t∈[0,1)t∈[0,1). Appendix D Experimental details D.1 Lorenz-63: Configuration Details Dynamics and data generation. We generate trajectories from the Lorenz–63 system (28) Lorenz [1963] with parameters σ=10σ=10, β=8/3β=8/3, and ρ=28ρ=28. The state dimension is d=3d=3 with the step size dt=0.01dt=0.01. Initial states are sampled from (0,I)N(0,I), and an additional Gaussian perturbation with standard deviation 0.10.1 is applied. For each trajectory, we use the time steps from step 20002000 to 25002500. Measurements are taken every gap=100 gap=100. Measurement equation. We use a linear observation model of the form τ=τ+σϵ,ϵ∼(0,I), _τ=x_τ+σ ε, ε (0,I), (108) where τ,τ∈ℝ3x_τ,z_τ ^3 and σ=1σ=1. We use tdefault=0.992t_default=0.992 for the terminal time. Model architecture. For Lorenz–63, we use a time-conditioned MLP Bishop [2006], Perez et al. [2018] that maps (,t)↦^(x,t) x with input, output dimensions 3→33→ 3, hidden width 6464, depth 33, and dropout 0.00.0. Training setup. We train for 500500 epochs with batch size 32 and learning rate 3×10−43× 10^-4, using a validation split of 0.20.2. D.2 Lorenz-96: Configuration Details Dynamics and data generation. We generate trajectories from the Lorenz–96 system (30) Lorenz [1996] with forcing F=8F=8 with the step size dt=0.01dt=0.01. Initial states are sampled from (0,I)N(0,I), and an additional Gaussian perturbation with standard deviation 1.01.0 is applied. For each trajectory, we use the time step from step 2525 to 100100. Measurement equation. We use a linear observation model τ=τ+σϵ,ϵ∼(0,I), _τ=x_τ+σ ε, ε (0,I), (109) where τ,τ∈ℝ64x_τ,z_τ ^64 and σ=1σ=1. We use tdefault=0.992t_default=0.992 for the terminal time. Model architecture. For Lorenz–96, we use a 1D U-Net Stoller et al. [2018], Perslev et al. [2019] operating on ∈ℝ1×dx ^1× d. The model uses base width 6464 with multiscale channels dim_mults=(1,2,4) dim\_mults=(1,2,4), time conditioning via sinusoidal positional embeddings, and attention with 44 heads of dimension 3232. We set dropout to 0.00.0 and do not use self-conditioning or learned variance. Training setup. We train for 500500 epochs and learning rate 3×10−43× 10^-4, using a validation split of 0.10.1. D.3 Kolmogorov Flow: Configuration Details Dynamics and data generation. We generate 2D trajectories from Kolmogorov flow Meshalkin and Sinai [1961], Chandler and Kerswell [2013] on the 64×6464× 64 grid. Each state is a velocity field t∈ℝ2×64×64x_t ^2× 64× 64. We simulate trajectories with Reynolds number Re=2000Re=2000 using the step size dt=0.2dt=0.2 from 5050 to 100100. Measurement equation. We use a grid-masked measurement equation with additive Gaussian noise: τ=⊙τ+σϵ,ϵ∼(0,I), _τ=M _τ+σ ε, ε (0,I), (110) where ∈0,11×1×64×64M∈\0,1\^1× 1× 64× 64 is a pixel-wise mask and ⊙ denotes element-wise multiplication. We use a regular mask with stride s=5s=5, i.e., M:,:,i,j=1M_:,:,i,j=1 if i≡0(mods)i≡ 0\ (mod\ s) and j≡0(mods)j≡ 0\ (mod\ s) and 0 otherwise. We set the observation noise to σ=0.1σ=0.1. Model architecture. For Kolmogorov flow, we use a 2D U-Net Bishop [2006], Perez et al. [2018] with time conditioning and attention. The model takes 22 input channels and outputs 22 channels, with base width model_channels=32 model\_channels=32 and channel_mult=(1,2,4) channel\_mult=(1,2,4). We use 22 residual blocks per resolution level, attention at resolution 44, dropout 0.00.0. Training setup. We train for 400400 epochs with batch size 1616 and learning rate 3×10−43× 10^-4, using a validation split of 0.10.1. Appendix E Additional Results This section summarizes additional experimental results that are not included in the main text. Figure 7: Additional qualitative results for different random seeds. Lorenz–63 state trajectories with measurement gap 100 for (a) EnKF, (b) SF, (c) SSLS, and (d) MASF. Each row corresponds to a different random seed (0, 2, 3, 4), showing the reference trajectory and the assimilated trajectory; subplot titles report the trajectory RMSE for each run. Figure 8: Estimated system state on Kolmogorov flow (seed 0) for two measurement gaps. Vorticity fields are shown at the indicated time indices τ. (a) gap=5=5 and (b) gap=10=10. Top to bottom: reference state, sparse measurement, and reconstructions by SF, SSLS, and MASF. Numbers in each reconstruction panel report the per-frame SSIM with respect to the reference at the same τ. Figure 9: Estimated system state on Kolmogorov flow (seed 0) for two measurement gaps. Vorticity fields are shown at the indicated time indices τ. (a) gap=15=15 and (b) gap=25=25. Top to bottom: reference state, sparse measurement, and reconstructions by SF, SSLS, and MASF. Numbers in each reconstruction panel report the per-frame SSIM with respect to the reference at the same τ.