Tensor-Var: Efficient Four-Dimensional Variational Data Assimilation
Authors: Yiming Yang
Xiaoyuan Cheng
Daniel Giles
Sibo Cheng
Yi He
Xiao Xue
Boli Chen
Yukun Hu
Venue: ICML 2025 poster (2025)
Keywords: Data assimilation
Dynamical system
Primary area: applications->chemistry_physics_and_earth_sciences

Abstract: Variational data assimilation estimates the dynamical system states by minimizing a cost function that fits the numerical models with the observational data. Although four-dimensional variational assimilation (4D-Var) is widely used, it faces high computational costs in complex nonlinear systems and depends on imperfect state-observation mappings. Deep learning (DL) offers more expressive approximators, while integrating DL models into 4D-Var is challenging due to their nonlinearities and lack of theoretical guarantees in assimilation results. In this paper, we propose \textit{Tensor-Var}, a novel framework that integrates kernel conditional mean embedding (CME) with 4D-Var to linearize nonlinear dynamics, achieving convex optimization in a learned feature space. Moreover, our method provides a new perspective for solving 4D-Var in a linear way, offering theoretical guarantees of consistent assimilation results between the original and feature spaces. To handle large-scale problems, we propose a method to learn deep features (DFs) using neural networks within the Tensor-Var framework. Experiments on chaotic systems and global weather prediction with real-time observations show that Tensor-Var outperforms conventional and DL hybrid 4D-Var baselines in accuracy while achieving a 10- to 20-fold speed improvement.

Tensor-Vars
Efficient Four-Dimensional Variational Data Assimilation
Yiming Yang Xiaoyuan Cheng" Daniel Giles'
Sibo Cheng²
Yi He
Xiao Xue
Boli Chen 1
Yukun Hu
Abstract
Tensor-Var
Variational data assimilation estimates the dynam- ical system states by minimizing a cost function that fits the numerical models with the obser- vational data. Although four-dimensional vari- ational assimilation (4D-Var) is widely used, faces high computational costs in complex non- linear systems and depends imperfect state- observation mappings. Deep learning (DL) offers more expressive approximators, while integrating DL models into 4D-Var is challenging due to their nonlinearities and lack of theoretical guarantees in assimilation results. In this paper, we propose Tensor-Var, novel framework that integrates ker- nel conditional mean embedding (CME) with 4D- Var to linearize nonlinear dynamics, achieving convex optimization in learned feature space. Moreover, our method provides a new perspec- tive for solving 4D- Var a linear way. offering theoretical guarantees of consistent assimilation results between the original and feature spaces. To handle large-scale problems, method to learn deep features (DFs) using neural networks within the Tensor- Var framework. Ex- periments on chaotic systems and global weather prediction with real-time observations show that Tensor-V Varr outperforms conventional and DL hy- brid 4D-Var baselines in accuracy while achieving 10- 20-fold speed improvement.
Feature Space
(5xr1) - Cru((s))
Hi
4D VD-Vart Cont
Capo
(Linear)
(Linear)
Ha
(Convex)
Original Space
S11 F(5))
40
(Non-linear)
(Non-linear)
(Non-convex)
Figure I. Demonstration Tensor- Var: system with non- linear dynamical model, observation model, and non-convex cost function (bottom) can be represented linearly in feature space, re- sulting convex cost function (top).
(Coveney, 2024). In this regard, data assimilation (DA) (Law al., 2015; Asch al., 2016) uses observation data to continuously calibrate models, improving forecast accu- racy. Various DA methods have been proposed to deal with different types of observation data and system dynamics. Among these methods, 4D variational (4D-Vat Var) data assimi- lation has been considered cutting-edge and used effectively in real-world applications such as numerical weather pre- diction (NWP) (Browne al., 2019; Milan al., 2020). 4D-Van minimizes quadratic cost function that finds the op- timal match between system states and observations (Asch et al., 2016). While effective in NWP, there are two criti- cal limitations for their applications: (1) Numerical models for complex, nonlinear systems are often inefficient for real- time assimilation and forecasting. (2) Observations are often noisy, incomplete representations of the states, out known state-observation mapping, posing challenges in utilizing the observations.
propose a
with-
1. Introduction
Forecasting of dynamical systems is an initial value problem of practical significance. Many real-world systems, such the ocean and atmosphere, are chaotic, which means that minor errors in current estimations in computational models can lead to rapid divergence and substantial forecast errors
Efforts have been made to integrate DL models to learn an observation (or inverse) model (Frerix et al., 2021; Wang et al., 2022; Liang et al., 2023), addressing the imperfect knowledge observation models 4D-Var. While these approaches improve observation utilization, they remain constrained by the complexities of numerical models and learned observation mappings, whereas our approach simpli- fies them by finding their linear representations. To improve computation efficiency, state-of-the--a DL models (Vaswani et al., 2017: Chen et al., 2018; Li et al., 2020; Kovachki et al., 2023; Cheng et al., 2025; He et al., 2025) are capable of con-
'University College London. London, United Kingdom *CEREA, ENPC and EDF R&D, Institut Polytechnique de Paris, Paris, France. Correspondence to: Yiming Yang <ze- ahyyl @ucl.ac.uk>. Yukun Hu
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Tensor-Vart Efficient Four-Dimensional Variational Data Assimilation
structing highly nonlinear mappings to surrogate dynamical systems and achieve notable successes in NWP (Bi al., 2022; Lam et al., 2022; Kurth et al., 2023; Chen et al., 2023; Conti, 2024; Vaughan et al., 2024). However, integrating such models into optimization-aast remains challenging due to their inherent non-linearities (Janner et al., 2021; Bocquet, 2023; Bocquet et al., 2024). Using dittoriifferentiationn (AD) of DL models in DD-Var can reduce computational costs and has shown success in simple examples (Geer, 2021; Dong et al., 2022; Cheng et al., 2024). However, concerns remain regarding the ac- curacy of AD-derived derivatives, and their computational complexity increases with the model scale (Baydin et al., 2018). Recently, (Xiao et al., 2024) applied AD of a pre- trained weather forecasting model in 4D-Var, forming self- contained DA framework for Global NWP. This approach demonstrates promising performance in large-scale settings, but relies on a well-designed pre-trained model and may not generalize easily to other domains. Latent DA (Peyron et al., 2021; Fablet et al., 2021; Meline & Zaplotnik, 2023; Cheng et al., 2023a; Fablet et al., 2023) addresses these challenges by performing DA in a learned low-dimensional latent space. Although efficient, these approaches lack theo- retical guarantees for the consistency of 4D-Var solutions between the latent and original spaces.
and o taking values in the spaces S R and O R where ny and ?% denote the dimensions of these spaces. The distribution of S is denoted by Ps. the joint distribution over (S, O) by Pso, and the conditional distribution of S given O by Ps A sequence of states over time steps from to t is denoted by $1:: = (s1)
tasks, such as 41)- Var,
2.1. 4D variational data assimilation
Consider a dynamical system in discrete-time comprising a dynamical model and observation model:
St = Fs + e and or = G(s,) + + (1)
in which F is the dynamical model that advances the state St to st+1. and G is the observation model that maps the 8t the observation oz. The noise components ef, ef such that ei ~ (N, Q) and ei (N R). The objective of 4D-variational DA is to minimize cost function:
(So.T) = ||s0 - -bbll
11: G(s))l121
(2)
T
+ 11sc - F(s-))) 1
Σ
in which si is a prior guess for the initial state 80, with B as the background covariance matrix representing the un- certainty, i.c., 80 ~ (s B). The second and third terms account for errors the observation and dynamical models in (1). Solving the minimization (2) in a high-dimensional nonlinear dynamical system with incomplete observations can be challenging. In this paper, we aim to characterize the nonlinear dynamical and observation models in (1) using lin- ear models, allowing efficient optimization and forecasting. In this sense, we use the CME that transforms the nonlinear system (2) into a linear system in a feature space.
In this paper, we introduce Tensor-Var, a framework that lincarizes nonlinear dynamics via kernel embedding, en- abling convex optimization in feature space. Our approach addresses efficiency and non-convex challenges in existing variational DA methods. To our knowledge, Tensor-Var is the first attempt to identify the linear feature space for DA systems with a convex cost function, enabling an efficient solution to the 4D-Var problem. To handle incomplete ob- servations, we derive an inverse observation operator that in- corporates historical observations to infer the system states. Moreover, we provide a theoretical analysis that demon- strates the existence of a linear representation of the system under kernel features and the consistency of the 4D-Var solution across original and feature spaces. A key challenge in extending the kernel embedding to practical variational DA is scalability. To overcome this, our approach learns adaptive deep features (DFs) that map data into a fixed- dimensional feature space, reducing computational com- plexities. Our experiments on two chaotic systems and two global NWP applications show that Tensor- Var outperforms both conventional and MI.-hybrid variational DA baselines in accuracy and computational efficiency, demonstrating the advantages of linearizing DA systems.
2.2. Kernel conditional mean embedding
A Reproducing Kernel Hilbert Space (RKHS) His on S with kernel kş is a Hilbert space, satisfying the reproduc- ing property (Schölkopf & Smola, 2002). Let kg and ko be positive-definite (pd) kernels on S and O with RKHS Hs and lilo. We denote the kernel-induced features as (o))) = ko and (st))) = (s )) referring the Hs and Ho as feature spaces respectively.
Kernel mean embedding of distribution Pss is defined as the expectation of the feature (S()) : μ2s = [s] e Hs and always exists for bounded kernels (Fukumizu al., 2011). With the reproducing property of RKHS, ex- pectations of the function J can easily be computed as Es(((S) =
2. Preliminary
for any f E lHl.. For charac- teristic kernels, these embeddings are injective, uniquely determining the probability distribution (Berlinet & Thomas- Agnan, 201 1). In addition to mean embedding, we will need
Notation. Let S and O be random variables representing the state and observation, respectively, with their realizations
2

---

Tensor-Var Var: Efficient Four-Dimensional Variational Data Assimilation
the (uncentered) cross-covariance operator (Baker, 1973) Cso: Ho Hss, defined as Cso = Eso Esolds(S) (S) (O)), where @ denotes the tensor product. This operator always exists for bounded kernels and can also be viewed as the embedding of the joint distribution Pso. represented as an element in the tensor product space lis @ Hll.. Simi- larly, the covariance operator Css: Itls + Hy is defined as Css = EE (s) @ (s)]) These operators extend the con- cepts of covariance matrices from fnite-dimensional spaces to infinite kernel feature spaces.
models of the system in (I).
CME of dynamical model. Let S be the one-step for- ward of S. Given the dynamical model F' in (1) and the ker- nel feature Φs, the CME operator Cs+ can be recognized as the best linear approximation in the feature space Hig that minimizes the regression residual EE] - Given finite data {( { sliced from the system tra- jectory S1:N+1+ we can obtain the empirical estimate C5+is as (4) with theoretical support of convergence (Fukumizu et al., 2013; Klus et al., 2020). The CME operator Cs+1s is effectively characterizes the system dynamics as a linear model and simplifies the 4D-Var as a convex optimization in the feature space H.s-
To represent the dynami- cal and observation models that arise in (1), the conditional mean embedding (CME) plays an important role. For a conditional distribution Iss of S. The CME of Psjo is defined as
Conditional mean embedding.
CME of inverse observation model. Analogous to the dynamical model, the observation model G in (1) can be linearized by the CME operator Cols, which has been used as observation models for filtering algorithms (Song et al., 2009; Fukumizu et al., 2013; Kanagawa et al., 2016; Geb- hardt et al., 2019). Most approaches assume a complete observation setting, where the observations can fully deter- mine the state. In practice, observations are often incom- plete representations of states, with ng > nos leading to underdetermined systems (Liu et al., 2022). In such sys- tems, the lack of a bijective mapping between the state and observation spaces means that observations cannot uniquely determine the system's 's state. As a result, the optimization problem in 4D- 4DrVar may produce suboptimal solutions (Asch et al., 2016). Is theoretically challenging to distinguish any two mixtures of states based on single-step observation if n. no. By introducing the past consecutive obser- vations as history hN = 0t-m-1a-1 E mR the joint information from history and current observation is enough to estimate the system state. The choice of history length is critical: Too short lacks sufficient information, while too long is inefficient. Empirically, we performed an ablation study to assess the effects of history length, detailed in the Subsection 4.4.
Ma = Es ( = o] for any E O,
(3)
requiring an operator CsJo: Ho + Hss such that (1) MPas (o), and (2) (f.HEs:))
(s) = ol for any f Hy. Under standard as- sumptions the CME operator can be expressed as Cs10 = CsoCoo (Fukumizu et al., 2004; Song et al., 2009). Given ii.d samples {(s,, o,) } ~ Pso. an empirical estimate of the CME operator can be obtained as
Csjo = Cso + A1) & + A (4)
where A is the regularization parameter, Ko is the Gram ma- (s,))
trix [Ko] ol (o, and &g = (obs(1) and to = |do (01), .... (ox) are the feature matrices stacked by columns.
3. Method
In this section, we introduce our Tensor-Var Var approach, which embeds 4D-Var into the kernel feature space and provides a theoretical analysis demonstrating the existence of linear dynamies with consistent convergence between the original and feature space solutions. To effectively address incomplete observations, we propose inverse observation operator that leverages consecutive historical observations. Additionally, we propose a method to learn adaptive deep features (DFs) using neural networks within the Tensor- Var framework, improving real-worl, applicability. Finally, we analyze optimization performance of Tensor- Var compared to existing 4D-Va Var based methods, demonstrating its effi- ciency.
To effectively incorporate history, we introduce another kernel feature Φи with an induced RKHS H The space How = Ho @ HH is called a tensor product RKHS on Ox H with associated kernel feature @ou (o, h) = Φo(() @ ΦH (H) As shown in (Song et al., 2013; Muandet al., be extended to high-order
2017), the CME operator features, allowing us to embed the joint distributions over 8t, o,, and h, into the feature space (Song et al., 2009; 2013). The tensor product feature Δoн captures the high-order dependencies between observations and the history. The CME operator Csjoh is the linear inverse observation model that minimizes the state estimation error given observations and history. Given dataset ((si, o, h.) b the empirical counterpart CS oh = C + AT) follows
3.1. CME of 4D-Var in RKHS
Having introduced the necessary tools for manipulating kerne embeddings, we now focus on learning the linearized
((1) Css is injective, and (2) [S)] o] H. for any
f E g and o e O.
3

---

Tensor-Vane Var: Efficient Four-Dimensional Variational Data Assimilation
Table 1. Comparison of Tensor-Var with other algorithms. Here, e is the error threshold and 7 is the length of the assimilation window.
Space Transformation
Non-Space Transformation
Methods
Tensor Var
Latent 4D-
4DrVar (wl w/ Adjoint)
Frerix et al. (2021)
Convexity Convergence Rate
X
x
Linear
Sublinear O(1/e)
Sublinear
Sublinear O (1/2)
Iteration Time
O (log (1/c))
O (1/22)
Computational Complexity
w/o Adjoint: O ((Tn))?/e)) w/ Adjoint: O (Tn,/ee)
O (Td, log (1/c))
O ((7)))%))
O ((Tn))//E)
the same way as (4). The Eson is the empirical high order N tensor, e.g., CSOH = - (s;) @ (o())) @ 6н (h.).
feature Ø8g- R. RRd with Cs+ s by minimizing the loss L(0s) minge L(C, 0s) = [[lloo,(s']
( Predictions using Es+ are in the fea- ture space; however, for DA problems, reconstruction to the original state space is required, which is known as the preim- age problem (Honeine & Richard, 201 1). Here, we learn an inverse feature 0% to solve the preimage problem during training, avoiding repeated optimization whenever comput- ing preimages. The final training loss is the combination of the two terms as
Feature space Var. Using the kernel features, we lin- earize the original nonlinear dynamics and observations in the feature space Hs.. This transformation enables us to reformulate the Var optimization objective (2) into the feature space and optimize over a sequence of elements 20:T in feature space Hsg:
min 1 os( 20.7
+
1 - Cs+ Cs+1s2ell@=+
L(Og. 0'z) B[lloss (s+) (s (s)
1=
+ w B ||+ - $bz (dog (s)))
2 - (u))
where w E (0, 1] is a weighting coefficient, and Cs+1s is computed the CME over training batches. Note that using linear kernel in the learned feature space, where køg (si, 83) -og (s:) 6ag (s,).
(5)
where the B, R. and Q are the covariance operators for the background error, observation error, and model error in the feature space. In this work, we estimate the three operators as the empirical error covariance matrices from the training dataset (the explicit estimation can be found in Appendix D.I). We present the seeudo-algoritmm in Appendix D.2, as shown in Algorithms 2 and 3.
DFs corresponds
Learning the observation and history features. Sim- ilar to learning the state feature, CsJoh is the minimizer of the regression problem mapping the tensor product of observation and history features to the state feature. In this phase, we learn the DFs for observation de : R. Rd, history Φon R Rd,, and Csjon jointly with the loss function:
3.2. Learning the deep features within Tensor-Vat
Using the pre-determined kernel features has theoreti- cal guarantees. However, it maps data into an infinite- dimensional feature space, e.g., radial basis function, mak- ing estimation challenging due to polynomial scaling with sample size. On the other hand, these feature maps strug- gle with irregular or nigh-dimensional data, often resulting in poor performance. Learned deep features (DFs) have emerged as alternatives to generic kernel features (Xu al., 2022; Kostic et al., 2023; Shimizu et al., 2024), projecting the data into a dinesonieesioonl feature space. To improve scalability, we integrate DFs with the Tensor Var framework and validate their effectiveness through experiments.
I(Oo, 0n) = B[ll (s) ČsjoH (0) @ Φo,, [h)]]
where ČOJOH is computed over training batches in parallel, with the @ denoting the Kronecker product in practice.
Improved 4D-Var Optimization. Tensor- Var improves 4D-Var optimization by transforming the nonlinear prob- lem into a linear convex optimization, enabling global op- timality and faster convergence in feature space. Unlike conventional 4D-Var, which relies on costly Hessian compu- tations and suffers from sublinear convergence, Tensor-Vas achieves linear convergence through lincarized dynamics ž = Cs+1s4. Table 1 presents a theoretical comparison of optimization properties between Tensor- Var and exist- ing 4D-Var- Var-based methods, including MI-hhbrrd 4D-Var
Learning the state feature. Recall that the CME oper- ator Cs-1s is the best linear approximation of system dy- namics in feature space. We propose to jointly learn the

---

Tensor-Van Var: Efficient Four-Dimensional Variational Data Assimilation
(Frerix et al., 2021), latent 4D- (Cheng et al., 2024), and traditional 4Dr with or without an adjoint model. As shown in Table 1, due to its linear dynamics, its iteration time scales as (log(1 1 / with error threshold which is more efficient than ()) the complexity of most 4D-Vat algorithms, as supported by empirical results in Section 4.1. Notably, there is an additional error source from the feature mapping, which is not explicitly discussed here; more in-depth analysis can be found in Appendix B.
ing (1) Lorenz 96 system (Lorenz, 1996) with n,, 40 and
80. (2) Auramoto-Sivashinsky (KS) equation: fourth-order nonlinear PDE system (Papagcorgiou & Smyrlis, 1991) with n, = 128 and 256, representing different spatial resolution. For both systems, we use a nonlinear observation model o = G(s) = 5arctan(s7, (s7/10) + E, where is white noise with a standard deviation of 0.01 times the standard devia- tion of the state variable distribution, and only 20% states can be observed. To assess the practical applicability of Tensor- we evaluate its performance in global medium- range weather forecasting (i.e. 3-5 days) by using a subset of the ECMWF Reanalysis v5 (ERA5) dataset for training and testing, with further details in Subsection 4.2 and (Rasp et al., 2024). Moreover, we incorporate observation loca- tions extracted from the real-time weather satellite track into the NWP experiment with higher spatial-resolution in Subsection 4.3.
3.3. Theoretical analysis.
In Section 3.1, we discuss how the dynamical system F in (1) can be embedded as a linear system in the feature space. However, two important questions remain: 1) Does such a linear dynamical system exist? 2) Are the solutions of the original and feature space 4D-Va consistent? In this section, we provide affirmative answers to both questions using the theory of Kazantzis-Ke Sravaris/Luenberger (KKL) observers (Andrieu & Praly, 2006). We give a road map of the theoretical analysis with the main results and refer the reader Appendix C for details.
Baselines. We compare our method against several base- line approaches: (1) 3D-Var with a known observation model, (2) a model-based 4D- Var algorithm that assumes known dynamical and observation models, (3) a learned in- verse observation model with the known dynamical model, as proposed by (Frerix et al., 2021). In the two NWP prob- lems, we include another three baselines: Latent 3D-Var and Latent 4D-Var (Cheng et al., 2024), which perform variational data assimilation in a latent space learned via an autoencoder, with Latent 4D- Var also modeling latent- space dynamics; and Fengwu 4D-Var (Xiao et al., 2024), an Al-embedded 4D-Var framework that integrates an ad- vanced neural weather forecasting model into the DA pro- cess. These competitive baselines cover both operational VA and ML-hybrid Var-DA methods.
Under the mild assumptions that (1) the dynamical model F'r is first-order differentiable and (2) the kernel features are all first-order differentiable, the kernel feature 6s satisfies the necessary conditions as the state transformation in the KKL observer framework. This transformation enables us to represent the nonlinear dynamical system as a linear system in a digherrdimensional feature space, as established by KKL observer theory (Tran & Bernard, 2023). This result confirms the existence of such a linear system, thus answering question 1). The KKL observer theory provides a theoretical foundation for our approach, bridging nonlinear dynamies and linear 4D-Var methods (Andrieu & Praly, 2006). A detailed derivation proving that ds satisfies the conditions as the state transformation of the KKL observer can be found in Appendix C.
4.1. Evaluation and results.
In each experiment, we measure the quality of assimilation using the Normalized Root Mean Square Error (NRMSH) -.= 1| sell2/(smmx Smin) over the assimilation window, where Smax, Smin are the maximum and minimum state values in the training dataset. These results and the average evaluation time for each algorithm are reported in Table 2. All metrics are evaluated 20 times with different initial conditions, reporting the mean and standard devia- tion. For Tensor-Var, the history length m is selected using a cross-validation approach as an ablation study in Subsec- tion 4.4, and the objective function in the feature space is minimized using quadratic programming, implemented via CVXPY (Diamond & Boyd, 2016). The baselines use the L-BFGS method (Nocedal & Wright, 2006) with 10 his- tory vectors for the Hessian approximation. For 4D-Var baselines, we consider two cases (with and without adjoint models). The background state st is set to the mean state of the training set.
We consider the nonlinear system in (1) within a compact state space and assume that the cost function in 2 has a unique solution. Given that Φs is a state transformation in the KKL observer, the system in the original state space can be represented linearly in the feature space. The solution in the feature space has consistent convergence to the unique solution of the original 4D- 4D-Var problem, minimizing with respect to the cost function 5, answering question 2). A formal theorem with detailed proofs can be found in Theorem C.9 in Appendix C.2.
4. Experiment
To evaluate our proposed method, the comparison is con- ducted on series of benchmark domains, representing the optimization problem (2) of increasing complexity, includ-
5

---

Tensor-Vart Efficient Four-Dimensional Variational Data Assimilation
Table Comparison of DA performances, All baseline methods use the strong-constraint 4D-Van objective, while our approach uses the weak-constraiig Var objective. Evaluation times and iterations reported as each assimilation window's s mean and standard deviation. Our method consistently outperforms the baselines across all benchmark domains.
Domain
Algorithm
NRMSE (%)
Evaluation time (10 *s)
Iteration
3D- Var
14.17 0.93
12.59
13 =
Lorenz 96
4D-Var
12.27 + 1.41
210.52 1 3.87
17 1 6
40
4D-Var by Adjoint
12.18 1 1.46
33.78 2.11
16 1
The 8
Frerix al. (2021)
9.89 1.63
167.43 + 1.33
11
Ours
0.87
12.51 1.97
8
3D-Var
15.19
19.38 + 0.37
14
Lorenz 96
4D-Var
12.38 1.18
322.21 5.73
22 + 8
n. 80
4D-Var Adjoint
12.44
18.07 2.59
22 + 8
no 16
Frerix (2021)
286.11 1 2.43
14
Ours
9.04 1.32
21.09 = 0.79
9 1
3D-Van Var
1..64
16.48 1.17
13 1340 0
Kuramoto Sivasiii--
4D-Var
15.46 1.07
94.83 - 3.89
77 4
Ti, 128
4D-Van by Adjoint
15.43 0.87
22.14 - 4.01
72
No 32
Frerix al. (2021)
10.25 + 1.34
63.28 = 1.91
28 1 1
Ours
9.69 + 1.56
19.58 + 1.23
18 + 2
3D-Var
16.66 1 0.69
18.81 1 0.92
15 1
Kuramoto- Kuramoto-Sivashinsky
4D-Var
0.62
99.68 + 2.35
Ttg = 256
4D- by Adjoint
10.23 0.21
25.77 + 2.13
81 1
no 64
Frerix et al. (2021)
8.87 + 0.55
71.39 = 1.23
31 + 2
Ours
4.31 + 0.19
21.37 = 1.36
23 =
Latent 3D-Var
Latent 4D-Yar
FengWu 4D-Var
Tensor-Var
z500
t850
q700
u850
v850
Evaluation time
Densiby
Time (s)
0
5
2.5 5.0 7.5 NRMSE (%)
10 20 NRMSE (%)
10
2.5 5.0 7.5 NRMSE (%)
Latent so.var 40 TangWuu 40 va-
Tomson Van
NRMSE (%)
NRMSE (%)
Figure 2. Comparison of distribution of NRMSE (%) across different atmospheric variables (z500, 1850, q700, u850, v850) for Latent
Var. Latent HD-Var, Fengwu 4D-Var, and Tensor- Var. The rightmost bar plot shows evaluation times and error bars indicate the standard deviation for evaluation time.
As shown in Table 2, our method consistently outperforms the other baseline methods all metrics. In all four tasks, Tensor- Var achieves the lowest mean and standard deviation of NRMSE for assimilation accuracy. demonstrating strong generalization from the Lorenz-96 systems the more com- plex KS systems. The better performances are attributed to the linearization of the dynamics and the history-augmented inverse observation operator, which makes the JD-Var op- timization convex with more reliable convergence to the global optimum the feature space. The larger NRMSE errors observed in the model-based 3D-Var and 4D-Var methods are due the fact that incomplete observations lead to uncontrolled errors in the unobserved state dimen- sions. Although the ML-hybrid 4D-Var method (Frerix et al., 2021) employs learned inverse observation model, it strug- gles with generalization due to the incomplete observations, leading to poor performance on test data. By incorporating historical information, our approach effectively controls es- timation errors in the unobserved state dimensions. This
improvement demonstrated in Figures 7 and 8 in Ap- pendix 1.1.2, which qualitatively compares the assimilation performance in the Lorenz-96 and KS systems.
In terms of computational efficiency, the linear structure and convexity of the Tensor enable linear optimization con- vergence rate and much lower computational complexity, as shown in Table 1. Our results confirm these improvements by showing lower iteration times and overall evaluation times compared to other baselines. Moreover, the computa- tional cost of feature mapping negligible relative to the 4D-Var optimization, and the lowest NRMSE validates con- sistency of the optimization across both the original and feature spaces.
4.2. Global NWP
Next, we consider global NWP problem. The European Centre for Medium-Range Weather Forecasts (ECMWF) Atmospheric Reanalysis (ERA5) dataset (Hersbach et al.,
6

---

Tensor-Var Efficient Four-Dimensional Variational Data Assimilation
Lalent 3D-Var
Latent 4D-Var
FengWu 4Dr
Tensor-Var
z500
t850
q700
u850
v850
NRMSE (%)
5
8
15.0
3
12.5
3
10.0
10
20
0
10
20
0
10
20
0
10
20
0
10
20
Time
Time
Time
Time
Time
Figure 3. Comparison of assimilation NRMSE (%) over 7-day horizon for five atmospheric variables from Latent 3D-Var, Latent 4D-Var, Fengwu 4D-Var, and Tensor-Var. Each time-step represents a 6-hour interval.
2020) provides the best estimate of the dynamics of the atmosphere covering the period from 1940 to present. The 500hPa geopotential, 850 hPa temperature, 700 hPa humid- ity, and 850 hPa wind speed (meridional and zonal direc- tions) at 64 32 resolution are considered here. The data is sourced from the WeatherBench2 repository (Rasp et al., 2024). Observations are sampled randomly from the grid with a 15% spatial coverage and with additive noise (0.01 times the standard deviation the state variable) (see Figure 9). Latent dimensions of the two baseline methods are set to match the feature dimension de in our approach. The evalu- ation metric is the area- wei-hhigg RMSE over grid points (see more details in Appendix E.3). We trained all models in ERA5 data from 1979-01-01 to 2016-01-01 and tested on data p0st-2018, 8, with a qualitative evaluation shown for 2018-01-01 00:00 in Figure 9.
4.3. Assimilation from Satellite Observations
In the final experiment, we consider observations from satel- lite track locations instead of random observation points in Section 4.2, addressing more realistic DA problem in global NWP. Weather satellites offer crucial observations, but their dynamic positions and spatial-tempores sparsity make data assimilation more challenging. The spatial reso- lution of the grid increases to 240 x 121 in this case.
Weather Satellites One Hour Tracks
Figure (left 1-5) shows the distributions of NRMSE (%) across different atmospheric variables (z500, t850, 4700, u850, v850) for Latent 3D-, 4D- 4ar,ar, and Tensor- Var, eval- uated on a test data set consisting of two years of data from 2018-01-01 00:00 to 2020-01-01 00:00. Our approach consistently achieves the lowest mean and standard devia- tion of NRMSE for all variables, demonstrating improved performance in both assimilation accuracy and robustness. The latent spaces of 3D- and 4D-Var are learned by the autoencoder purely based on the reconstruction loss with- out considering the system dynamics. This weakens the forecasting abilities of the dynamical systems, introducing extra errors in the 4D-Var optimization compared to the 3D)- Var optimization. In contrast, our approach jointly learns feature space representation and linear system dynamics, resulting in more accurate forecasting and improved 4D- performances. The rightmost barplot in Figure 2 shows the evaluation times Var achieves faster computation than Latent 3D-Var and Latent 4D-Var because its deep features are used only to map data into feature space rather than being directly in- volved in gradient-based optimization via AD. In addition to the assimilation results, we evaluate the forecast quality of Tensor- Var based the assimilated state; quantitative results can be found in Appendix E.3 Figure 10.
1273
127
Satellte--ee NOAA 15
Setei NETEOSATO IMSG-2)
Satelltt DMSSP 5D-3 F15 (USA 172:
Satelits EWS G1 (GDES 131
Sanel lite NOAA
Observation locations
Figure 4. Selected satellite tracks over one-hour horizon. with observations (black triangles) sampled at half-hour intervals.
An example of satellite tracks and observation distribution is shown in Figure 4. We extract satellite track data (lat- itude and longitude coordinates) from CelesTrak2 for the same periods Section 4.2, matching it with ERA5 data generate practical observations. These observations include satellite locations within two hours before the assimilation time, sampled at half-hour intervals, with an average cover- age of approximately 6%, see Figure 5. Other experimental settings, such as data volume, training/testing periods, vari- ables, and the 4D-Var window, align with Section 4.2. In this experiment, we evaluate our proposed method in an online operational scenario, where the DA is continuously applied. To assess the uncertainties, all methods are evalu- ated 10 times by randomly sampling sequences from the test dataset. We present the mean and standard deviation of the NRMSEs in Figure 3 over a 7-day time horizon. Tensor- Var
Nvidia RTX-4090 GPU. Tensor-
'CelesTrak provides public orbital data for a wide range of satellites, including those with meteorological sensors at www. celestrak. com. The data include positional details and tem- poral information, allowing accurate real-time satellite tracking.
7

---

Tensor-Vart Efficient Four-Dimensional Variational Data Assimilation
2018-01-01
2018-01-02 00:00
2018-01-03 00:00
2018-01-04 00:00
2018-01-05
00:00
00:00
True State
1300
MCC3
40200
3000
2000
MUCE
4033
56000
4XXC
4003
53300
MACO
40.30
1246
209
Tensor-Va Var Result
4400
4000
25000
30000
Absolute Error
150
10
26
63
X
100
Observations
Figure Visualization continuous assimilation results, absolute errors, and observation locations z500 (geopotential), starting from 2018-01-01 00:00. The observation coverage, defined as the ratio of the number of observations to the number of grid points, is 6.37%.
consistently achieves the lowest mean and standard devia- tion, demonstrating its robustness in a large-scale system with practical observations. Figure 5 presents the qualitative assimilation results for the variable z500 with observations; results for other variables are provided in Appendix E.4. Dy- namic satellite observations impact assimilation accuracy, with clustered observations near the equator from geosyn- chronous satellites substantially reducing errors the corre- sponding regions. The results demonstrate that Tensor- Var excels in accuracy and robustness when handling large-scale systems with practical observations, showing its potential for applications in operational DA and forecasting.
to be the same as the experiments Subsection 4.1. Table (right) shows that incorporating history (C 0) signifi- cantly improves the accuracy of the state estimation, with the NRMSE decreasing C' increases. However, improve- ments become marginal when increasing C beyond certain point (around C' = 4). This indicates a tradc-off, where in- creasing the history length beyond a certain threshold yields little additional benefit in state estimation.
Effect of DFs. This experiment compares DFs with vary- ing dimensions d, = 20, 40, 60, 80 to Gaussian kernel feature. To scale up the Gaussian kernel, we apply Nyström approximation and kernel PCA (See Appendix E.5 for de- tails). All three features are Gaussian kernels. The preimage of the system state is learned using kernel ridge regression on the low-dimensional representations. As shown Table (left), Gaussian kernel performance is close to the best DFs in n. 40 but degrades in n, 80. In lower-dimensional problem, the Gaussian kernel features are more robust than DFs, but their performance can degrade with increasing system dimensionality and complexity. DFs with d, = 60 and 120 perform consistently well, whereas de = 20 and 40 underperform, indicating trade-offs between data size, model complexity, and system dimension.
4.4. Ablation study
To support our empirical results, we perform ablation stud- ies on the 40- and 80-dimensional Lorenz-96 system to investigate (1) the effect of history length in learning the operator ESJOH. and explore (2) the effect of DFs Φs with different feature dimensions.
Effect of history length. We explore the effect of history length mn on learning the inverse observation operator and its impact on state estimation accuracy. According to the theory in (Liu et al., 2022), the history length can be cho- sen as m 1og(n.)². The ablation study conducted by scaling m proportionally m C' Clogn, n, where the con- stant C' was adjusted. The dimensions day do, dy are fixed
5. Conclusion
In this paper, we propose Tensor- Var, a framework that uses kernel conditional mean embedding to linearize non- linear models in 4D-Van Var, making optimization tractable and
Please note that the result omits the class of systems with exponential dependency on the history length.
8

---

Tensor-Vart Efficient Four-Dimensional Variational Data Assimilation
Asch, M., Bocquet, M., and Nodet, M. Data assimilation: methods, algorithms, and applications. SIAM, 2016.
Table Comparison of different features and history lengths with NRMSE (%) as the metric for state estimation accuracy. GK
denotes the Gaussian kernel.
Baker, C. R. Joint measures and cross-covariance operators. Transactions of the American Mathematical Society, 186: 273 289, 1973.
Feature Dimension (NRMSE %)
History Length (NRMSE %)
d,
40
80
C
40
80
16.7
17.3
11.7
3.4
10.8
40
14.1
+ 1.3
16.8 : 2.2
1
9.3
2.8
1.7
0.9
0.9
Baydin, A. G., Pearlmutter, B. A., Radul, A. A., and Siskind, J. M. Automatic differentiation in machine learning: survey. Journal of machine learning research, 18(153):
120
9.7 0.7
10.9 + 0.9
4
7.7 0.9
7.9
GK
8.4.1.
0.5
11.7 1
1.4
8
7.510.
7.5
0.5
7.0 7.0 0.8
1-43, 2018
efficient. By learning adaptive deep features, Tensor- Var addresses the scalability typically associated with traditional kernel methods. Our inverse observation operator, which incorporates historical observations, improves accuracy and robustness with incomplete observations. Experiments on two chaotic systems and global weather forecasting show that Tensor Var outperforms state-of-the-art hybrid ML-DA models in both accuracy and efficiency.
Berg, C., Christensen, J. P. R., and Ressel, P. Harmonic analysis semigroups, volume 100. Maringer-Verlaag New York, 1984.
Berlinet, A. and Thomas-Agnan, C. Reproducing kernel Hilbert spaces in probability and statistics. Springer Science & Business Media, 2011.
Bi, K. Xie, L., Zhang, H., Chen, X., Gu, X., and Tian, Q. Pangu-weather: A 3d high-resolution model for fast and accurate global weather forecast. arXiv preprint arXiv:2211.02556, 2022.
Limitations and future work. Tensor- Var requires access to system states to learn dynamics and observation models, which may not be feasible in practice. Future work will focus on learning these models directly from observations and calibrating dynamics in feature space. Additionally, the simplified error covariance used 4D- Var may underesti- mate system correlations; refining these within Tensor- Var could improve assimilation and performance.
Bobrowski, A. Convergence of one-parameter operator semigroups, volume 30. Cambridge University Press, 2016.
Bocquet, M. Surrogate modeling for the climate sciences dy- namics with machine learning and data assimilation. Fron- tiers in Applied Mathematics and Statistics, 9:1133226. 2023.
Acknowledgments
The authors thank the reviewers and the ICML 2025 pro- gram committee for their insightful feedback and construc- tive suggestions that have helped improve this work. We also acknowledge the support from the University College London (Dean's Prize, Chadwick Scholarship), the Engi- neering and Physical Sciences Research Council projects (EP/W007762/1), the United Kingdom Atomic Energy Authority (NEPTUNE 2057701-TN-03), and the Royal Academy of Engineering (1F-2425-19-A1165),
Bocquet, M., Farchi, A., Finn, S., Durand, C., Cheng, S., Chen, Y., Pasmans, I., and Carrassi, A. Deep learning- based sequential data assimilation for chaotic dynamics identifies local instabilities from single state forecasts. arXiv preprint arXiv:2408.047333 2024.
Boyd, S. P. and Vandenberghe, L. Convex optimization. Cambridge university press, 2004.
Impact Statement
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary,
C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne S., and Zhang, Q. JAX: Composable transformations of Python+NumPy programs. http:// github. com/ coon/google/// jax, 2018. Version 0.3.13.
This paper aims advance machine learning applications in studying the long-term properties of large-scale chaotic systems. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Browne, P. A., De Rosnay, P., Zuo, H., Bennett, A., and Dawson, A. Weakly coupled ocean -ecan-atmosphere data assimilation in the ecmwf nwp system. Remote Sensing, 11(3):234, 2019.
References
Andrieu, V. and Praly, L. On the existence of a kazantzis- kravaris/luenberger observer. SIAM Journal on Control and Optimization, 45(2):332-454 2006.
Butcher, J. C. A history of runge-kutta methods. Applied numerical mathematics, 20(3):247-660, 1996.

---

Tensor-Van Efficient Four-Dimensional Variational Data Assimilation
Chen, K., Han, T., Gong, J., Bai, L., Ling, F., Luo, J.-J, Chen, X., Ma, I.., Zhang, T., Su, R., et al. Fengwu: Push- ing the skillful global medium-range weather forecast beyond 10 days lead. arXiv preprint arXiv:2304.02948, 2023.
Fablet, R., Chapron, B., Drumetz, L.,, Mémin, E., Pannek- oucke, O., and Rousseau, F. Learning variational data assimilation models and solvers. Journal of Advances in Modeling Earth Systems, 13(10):2022MMS00552. 2021.
Fablet, R., Febvre, Q. and Chapron, B. Multimodal 4dvar- nets for the reconstruction of sea surface dynamies from sst-ssh synergies. IEEE Transactions on Geoscience and Remote Sensing, 61: 61:1-14, 2023.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Frerix, T., Kochkov, D., Smith, J., Cremers, D., Brenner, M., and Hoyer, S. Variational data assimilation with a learned inverse observation operator. International Conference Machine Learning, pp. 3449 -458.. PMLR, 2021.
Cheng, S., Quilodrán-Casas, C., Ouala, S., Farchi, A., Liu, C., Tandeo, P., Fablet, R., Lucor, D., Iooss, B., Brajard, J., et al. Machine learning with data assimilation and uncertainty quantification for dynamical systems: re- view. IEEE/CAA Journal of Automatica Sinica, 10(6): 2023a.
Fukumizu, K., Bach, F. R., and Jordan, M. I. Dimensionality reduction for supervised learning with reproducing kernel hilbert spaces. Journal of Machine Learning Research, (Jan):73-99, 2004.
Cheng, S., Min, J., Liu, C.,, and Arcucci, R. Torchda: A python package for performing data assimilation with deep learning forward and transformation functions. Com- puter Physics Commmicatiinns, pp. 109359, 2024.
Fukumizu, K., Song, L.,, and Gretton, A. Kernel bayes' rule. Advances in neural information processing systems, 24, 2011.
Cheng, X., Yang. Y., Jiang, W., and Hu, Y. Keec: Embed to control on an equivariant geometry. arXiv preprint arXiv:2312.01544, 2023b.
Fukumizu, K., Song. L., and Gretton, A. Kernel bayes' rule: Bayesian inference with positive definite kemels. The Journal of Machine Learning Research, 14(1):3753- 3783, 2013.
Cheng. X., He, Y.. Yang, Y., Xue, X., Cheng, S., Giles, D., Tang, X., and Hu, Y. Learning chaos in a linear way. In The Thirteenth International Conference on Learning Representations, 2025.
Gajic, Z. and Qureshi, M. T. Lyapunov matrix equation in system stability and control. Courier Corporation, 2008.
Conti, S. Artificial intelligence for weather forecasting. Nature Reviews Electrical Engineering, 1(1):888, 2024.
Gebhardt, G. H., Kupcsik, A., and Neumann, G. The kernel kalman rule: Efficient nonparametric inference by recur- sive least-squares and subspace projections. Machine Learning, (08(12):21332115) 2019.
Coveney, P. V. Sharkovskii's theorem and the limits of digital computers for the simulation of chaotic dynamical systems. Journal of Computational Science, 83:102449, 2024.
Geer, A. J. Learning earth system models from observations: machine learning or data assimilation? Philosophical Transactions of the Royal Society A, 4,39922199):22200888 2021.
Cox, M. and Matthews, C. Exponential time differenc- ing for stiff systems. Journal of Computational Physics, 176(2):430-455 455, 2002.
Gejadze, I. Y., Dimet, F.-X. L., and Shutyaev, V. On analysis eттor covariances in variational data assimilation. SIAM Journal on Scientific Computing, 10(4):1847-187)) 2008.
Cvitanović, P., Davidchack, R. I.., and Siminos, E. On the state space geometry of the kuramoto-sivashinge flow in periodic domain. SIAM Journal on Applied Dynamicul Systems, 9(1):1-33, 2010.
Givoli, D. A tutorial on the adjoint method for inverse problems. Computer Methods in Applied Mechanics and Engineering, 380: 111111101 10, 2021.
Diamond, S. and Boyd, S. Cvxpy: A python-embedder modeling language for convex optimization. Journal of Machine Learning Research, 17(83):1-5, 2016.
He, Y., Yang, Y., Cheng, X., Wang, H., Xue, X., Chen, B., and Hu, Y. Chaos meets attention: Transformers for large- scale dynamical prediction. In Forty-second International Conference Machine Learning, 2025. URL https: / // //operrevi ew, net forum? wit.piiddoxxxxeeeemm.m
Dong, R., Leng, H., Zhao, J., Song, J., and Liang, S. A framework for four-din dinensinnnional variational data assim- ilation based on machine learning. Entropy. 24(2):264, 2022.
Hersbach, H., Bell, B., Berrisford, P., Hirahara, S., Horányi, A., Sabaeer, J., Nicolas, J., Peubey, C., Radu, R.,
10

---

Tensor-Van Efficient Four-Dimensional Variational Data Assimilation
Schepers, D., Simmons, A., Soci, C., Abdalla, S., Abel- lan, X., Balsamo, G., Bechtold, P., Biavati, G., Bidlot, J. Bonavita, M., De Chiara, G., Dahlgren, P., Dee, D., Diamantakis, M., Dragani. R., Flemming. J., Forbes, R., Fuentes, M., Geer, A., Haimberger, I.,, Healy, S., Hogan, R. J. Hólm, E., Janisková, M., Keeley, S., Laloyaux, P., Lopez, P., Lupu, C., Radnoti, G., de Rosnay, P., Rozum, I, Vamborg. F., Villaume, S., and Thépaut, J.-N. The ERA5 global reanalysis. Quarterly Journal of the Royal Meteorological Society, 146(730):1999-2049 2020. doi: 10. 10.1002/gj.3803.
medium-range global weather forecasting. arXiv preprint arXiv:22/2./2794, 2022.
Law, K., Stuart, A., and Zygalakis, K. Data assimilation.
Cham, Switzerland: Springer, 214:52, 2015.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. arXiv preprint arXiv:2010.08893 2020.
Liang. J., Terasaki, K., and Miyoshi, T. A machine learning approach to the observation operator for satellite radiance data assimilation. Journal of the Meteorological Society of Japan. Ser. II, 101(1):79-95, 2023.
Honeine, P. and Richard, C. Preimage problem in kernel- based machine learning. IEEE Signal Processing Maga- zine, 28(2):77-55) (8, 2011.
Liu, Q. Chung, A., Szepesvári, C., and Jin, C. When is partially observable reinforcement learning not scary? In Conference on Learning Theory, pp. 5175-5220, PMLR, 2022.
Janner, M., Li, Q.. and Levine, S. Offline reinforcement learning as one big sequence modeling problem. Ad- vances in nearal information processing systems, 34: 1273-1286, 2021.
Lorenz, E. N. Predictability: A problem partly solved. In Proc. Seminar on predictability, volume 1. Reading, 1996.
Kanagawa, M., Nishiyama, Y., Gretton, A., and Fukumizu, K. Filtering with state-observation examples via kernel monte carlo filter. Neural computation, 28(2):382-444 2016.
Manchester, I. R. and Slotine, J.-J. E. Control contraction metrics: Convex and intrinsic criteria for nonlinear feed- back design. IEEE Transactions on Automatic Control, 6266))3046 2017.
Kingma, D. P. Adam: A method for stochastic optimization. arXiv preprint arXiv: pet::69128668 2014.
Klus, S., Schuster, I., and Muandet, K. Eigendecomposi- tions of transfer operators in reproducing kernel hilbert spaces. Journal of Nonlinear Science, 30:283-315, 2020.
McShane, E. J. Extension range of functions. 1934.
Meline, B. and Zaplotnik, Ž. Neural-network data assim- ilation using variational autoencoder. arXiv preprint arXiv:2308.16073 2023.
Kostic, R., Novelli, P., Grazzi, R., Lounici, K., and Pontil, M. Deep projection networks for learning time-homogeneous dynamical systems. arXiv preprini arXiv:2307.09991, 2023.
Milan, M., Macpherson, B., Tubbs, R., Dow, G., Inver- arity, G., Mittermaier, M., Halloran, G., Kelly, G., Li, D., Maycock, A., et al. Hourly 4d-var in the met office ukv operational forecast model. Quarterly Journal of the Royal Meteorological Society, 146(728):1281-133) 2020.
Kovachki, N., Li, Z., Liu, B., Azizzadenesheli, K., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Learning maps between function spaces with applications to pdes. Journal of Machine Learning Re- search, 24(89):1-97, 2023.
Muandet, K., Fukumizu, K., Sriperumbudur, B., Schölkopf, B., et al. Kernel mean embedding of distributions: review and beyond. Foundations and Trends@ in Machine Learning, 10(1-2):1-141, 2017.
Kurth, T., Subramanian, S., Harrington, P., Pathak, J., Mardani, M., Hall, D., Miele, A., Kashinath, K., and Anandkumar, A. Fourcastnet: Accelerating global high- resolution weather forecasting using adaptive fourier neu- ral operators. In Proceedings of the platform for advanced scientific computing conference, pp. 1-11, 2023.
Nguyen, D. and Fablet, R. A transformer network with sparse augmented data representation and cross entropy loss for ais-based vessel trajectory prediction. IEEE Ac- cess, 2024.
Kwok, J.-Y. and Tsang, I.-H. The pre-image problem in kernel methods. IEEE transactions on neural networks, 15(6)1511-11555, 2004.
Nocedal, J. and Wright, S. J. Quadratic programming. Nu- merical optimization, pp. 448 492, 2006.
Lam, R., Sanchez-Gonzale A., Willson, M., Wirnsberger, P., Fortunato, M., Alet, F., Ravuri, S., Ewalds, T., Eaton- Rosen, Z., Hu, W. et al. Graphcast: Learning skillful
Papageorgiou, D. T. and Smyrlis, Y. S. The route to chaos for the kuramoto kuramoto-sivashinsky equation. Theoretical and Computational Fluid Dynamics, 3(1):15-42 1991.

---

Tensor-Van Var: Efficient Four-Dimensional Variational Data Assimilation
Parks, P. C. A new proof of the Routh-Hurwitz stability criterion using the second method of Liapunov. In Math- ematical Proceedings of the Cambridge Philosophical Society, volume 58, pp. 694-702. Cambridge University Press, 1962.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, t.,, and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Vaughan, A., Markou, S., Tebbutt, W., Requeima, J., Bru- insma, W. P., Andersson, T. R., Herzog, M., Lane, N. D., Hosking, J. S., and Turner, R. E. Aardvark weather: end- to-end data-driven weather forecasting, arXiv preprint T1f1::2004.0444 2024.
Paulin, D., Jasra, A., Beskos, A., and Crisan, D. A 4d-var method with flow-dependent background covariances for the shallow-water equations. Statisties and Computing, 32(4):65, 2022.
Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., et al. Scikit-leam: Machine learning in python. the Journal of machine Learning research, 12:2825-2833 2830, 2011.
Wang, Y., Shi, X., Lei. L.,, and Fung. J. C.-H. Deep learning augmented data assimilation: Reconstructing missing information with convolutional autoencoders. Monthly Weather Review, 15008):1777-199, 2022.
Wright, S. Numerical optimization, 2006.
Peyron, M., Fillion, A., Gürol, S., Marchais, V., Gratton, S., Boudier, P., and Goret, G. Latent space data assimilation by using deep learning. Quarterly Journal of the Royal Meteorological Society, 147(740):7599-3777 2021.
Xiao, Y., Bai, L.,, Xue, W., Chen, H., Chen, K., Chen, K., Han, T., and Ouyang, W. Towards a self-contained data- driven global weather forecasting framework. In Pro- ceedings of the 41st International Conference on Ma- chine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 54255-54275. PMLR, 21-27 Jul 2024. htthttps:// WtLLhttps:/ppoceeeeeeee httpp//roceengnnee/ee mlr press/ v235 2/35/xia24a..... html.
Rasp, S., Hoyer, S., Merose, A., Langmore, I., Battaglia, P., Russell, T., Sanchez-Gonzale,, A., Yang, V., Carver, R., Agrawal, S., et al. Weatherbench 2: A benchmark for the next generation of data-driven global weather models. Journal of Advances in Modeling Earth Systems, 16(6):
e2023MS004019, 2024.
Xu, L.,, Chen, Y., Doucet, A., and Gretton, A. Importance weighted kernel bayes' rule. In International Conference on Machine Learning, pp. 24524-24538, PMLR, 2022.
Romanoff, N. On one-parameter groups of linear trans- formations. Annals of Mathematics, 48(2)3216-233, 1947.
Zamir, S. W., Arora, Khan, S., Hayat, M.. Khan, F. S.. and Yang, M.-H. Restormer: Efficient transformer for high-resolution image restoration. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 5728-5739, 2022.
Sastry, S. Nonlinear systems: analysis, stability, and control, volume 10. Springer Science & Business Media, 2013.
Schölkopf. and Smola, J. Learning with kernels: support vector machines, regularization, optimization, and beyond. MIT press, 2002.
Shimizu, E., Fukumizu, K., and Sejdinovic, D. Neural- kernel conditional mean embeddings. arXiv preprint arXiv:2403. 10859, 2024.
Song, L., Huang, J., Smola, A., and Fukumizu, K. Hilbert space embeddings of conditional distributions with appli- cations dynamical systems. In Proceedings of the 26th Annual International Conference on Machine Learning, pp. 961-968, 2009.
Song, L.., Fukumizu, K., and Gretton, A. Kernel embeddings of conditional distributions: unified kernel framework for nonparametric inference in graphical models. IEEE Signal Processing Magazine, 30(4):98-11) 2013.
Tran, G. Q. B. and Bernard, P. Arbitrarily fast robust kkl ob- server for nonlinear varying discrete systems. IEEE Transactions Automatic Control, 2023.
12

---

Tensor-Vart Efficient Four-Dimensional Variational Data Assimilation
A. Table of Notations
Notations
Meaning
S,O S,O,H H
random variables of state, observation, and history. samples of state, observation, and history.
8, 0, h S
domain of state space.
O
domain observation space.
n P
dimension of original spaces with specified subscripts.
probability distribution with specified subscripts.
Φs(s) (s(s,) ))
kernel feature map for state.
(o ko (0,
kernel feature map for observation.
ØF (h) = KH (h,
kernel feature map for history.
H.s, Ho, H
induced kernel feature space for S, O and H.
=
indicate the elements in feature space.
C
Conditional mean embedding (CME) operators with specified subscripts. error matrices for background, observation, and dynamical model in original space. error matrices for background, observation, and dynamical model in feature space. deep feature for state with neural network parameters Os.
Q,
B, B,R, Q
60g
inverse feature for state with neural network parameters 0s. deep feature for state with neural network parameters 8o. deep feature for state with neural network parameters 0μ- dimension of feature spaces with specified subscripts.
680
den d
13

---

Tensor- Var: Efficient Four-Dimensional Variational Data Assimilation
B. Variational Data Assimilation Optimization
A sequence { generated by an optimization algorithm converges linearly to x² there exists C E (0, 1) such that || - 2*I| X C/Iæa - I*I implying || z01 = O(C*). Convergence sublinear if || - ==' - 0 as h + co but slower than linear, often expressed as || - I*1| = O(1/k²) for > O. Superlinear convergence occurs when limg->o | x² [|| z=1l = O, with quadratic convergence special case where || x= || - C 2 x= x=11?.
B.1. Analysis of Optimization
Tensor-Var. Our approach based on CME theoretically linear in the learned feature space. Under this condition, it formulates a convex quadratic programming (QP) problem with linear dynamic constraints, illustrated in Equation (5). From the convergence rates of the Newton optimization method (Wright, 2006), the algorithm exhibits a linear convergence rate in worst-case and superlinear convergence rate in best-case. Consequently, the number of iterations required to achieve an c-error ranges between (1og(1 /c)) and (log log( (e))). ) In terms of computational complexity, the QP problem features an invariant Hessian matrix. This invariance allows us to compute the Hessian matrix only once. Incorporating the iteration complexity, the total computational complexity of Tensor- Var is efficiently bounded between O (Td, log(1/e)) /e)) and O (Td log log (1g))).). This demonstrates that Tensor- Var achieves efficient computation by minimizing redundant Hessian calculations and leveraging the convergence properties of the Newton method. It should be noted that the optimization error in Tensor-Va is in the feature space, not the state space. The true error is quantified as (e + eir) it decoding back, where eir represents the reconstruction error.
Latent 4D-Var. Unlike our linearized CME model, the latent 4D-Var approach (Cheng et al., 2024) directly constructs a neural network-based surrogate model using an autoencoder, thereby no longer depending on computationally intensive numerical models. In this framework, the inherent nonlinearity and non-convexity of the optimization problem remain unchanged. According to optimization theory, the L-BFGS algorithm achieves a worst-case sublinear convergence rate and a best-case superlinear convergence rate (Boyd& & Vandenberghe, 2004; Wright, 2006). The best computational complexity can only be achieved only if the problem exhibits strong local convexity around the initial point. Consequently, the number of iterations required to achieve an c-error ranges between (1/22) and (log log(1 /c)). Regarding computational complexity, latent 4D-Var does not decompose the quadratic programming problem into T separate subproblems. Instead, it concatenates the T features into high-dimensional vector of size Td.. At each time step. the Hessian matrix is updated using the BFGS algorithm. Considering the complexity of the iteration, the total computational complexity ranges from ((Td.))/ee to (((Td)) log log( (1()) which makes it computationally expensive. Since optimization is performed in the latent feature space, the true error is quantified as (e + eîr) upon decoding, where eir represents the reconstruction error.
4D- The original 4D-Var, when solved using the Newton method, exhibits a similar convergence rate, ranging from sublinear to superlinear. However, the number of iterations scales with the state dimension n. rather than the latent space dimension de, leading to a total complexity of ((TTn)22e) or ((Tn,,)² 2 log log()/))) to achieve an e- error. When the adjoint model is available, the total computational complexities can be further reduced to OTTn//e²) or OTn, log lo]( through forward and adjoint computations, each requiring (n,,) (Givoli, 2021; Wright, 2006). Since the 4D- 4D-Varin directly solved on the original space, there is no reconstruction error that needs to be counted.
Frerix et al., 2021. Frerix et al. (2021) propose learning an approximate inverse observation operator using deep learning and transform this problem into the physical state space. The 4D-Var problem is then solved using a standard Newton solver, resulting in convergence properties and computational complexities similar to those of the standard 4D-Van -D- without an adjoint model such as ((TTn))² /e²) or O ((Tn,)² log log(L/))). In practice, (Frerix al., 2021) argues that solving 4DrVar in physical state space with inverse observation operator leads to better initialization, inverse mapping provides an improved starting point for optimization. This improved initialization results in better-conditioned optimization landscape, i.e., strongly local convex and achieves superlinear convergence rate. Our numerical experiments further confirm that their method achieves faster convergence with fewer iterations compared to the standard 4D-Var model (with or without an adjoint model), as shown in Table 2.
comprehensive summary of these analyses and findings is provided in Table 4.
14

---

Tensor-Vart Efficient Four-Dimensional Variational Data Assimilation
Table Full comparison of Tensor-Va Var with other algorithms. Here, is the error threshold and is the length of the assimilation window. The best computational complexity can only be achieved only if the problem exhibits strong local convexity around the initial point.
Space Transformation
Non-Space Transformation
Methods
Tensor-Ve Var
Latent 4D-
4D-Var (w/o or w/ Adjoint)
Frerix et al. (2021)
Model-based
Convexity
X
X
X sublinear to linear O (1/e²)
Convergence Rate
linear to superlinear O (log (1/e))
sublinear to linear O(1/e))
sublinear to linear
Iteration Time
O (1/e²)
Worst Computational Complexity
w/o Adjoint: O ((Tn.)) (2)
O (T log (1/c))
((Td)))/e)
O ((Tn.))/E)
w/ Adjoint: O (Tn,/e")
w/o Adjoint:
Best Computational Complexity
(I n.) log log (1/c)
O (Td, log log (1/c)) O ((I',)² log log(o(()/))
O ((In.)² log (1/e))
w/ Adjoint:
O (Tn, log log (1/c))
C. Theoretical Analysis
In this section, we provide a theoretical convergence analysis of Tensor-Var, drawing on concepts from control theory and contraction analysis. We begin by introducing comparison functions, including class K. and K.C functions, as well as Lyapunov functions. Using these tools, we examine the convergence of Tensor- Var through differential equation, demonstrating monotonic contraction based the Lyapunov direct method. Furthermore, by using comparison functions, we show that contraction in the feature space implies contraction in the original space.
To perform a formal convergence analysis of Tensor- Var, we make mild assumption about the first-order differentiabilityy of the dynamical system F in (1) with the time derivative = f(s), a standard assumption in convergence studies (Sastry, 2013).
Assumption 1.
Assumption 2. We require that the kernel possess well-defined first-order derivative, as the convergence analysis is performed in the feature function space. This assumption is common in kernel methods and satisfied by many widely-used kernels, such the Gaussian, Fourier, Matérn, and Laplace kernels (Berg et al., 1984; Schölkopf & Smola, 2002).
C.1. Notations and Technical Lemmas
Definition C.1 (Class IC function (Gajic Qureshi, 2008).). A continuous function a a0 + [0, oo) said belong to class K. if it is strictly increasing and (0)) 0. It is of class Ko if a(co) oo and air)
co.
Definition C.2 (Class KCC function (Gajic & Qureshi, 2008).). A continuous function B: [0, a] x [0, oo) + I0, oo) is said belong to class K.C if for each fixed t, the mapping 3(r, t) belonging to class K. with respect r and, for each fixed a the mapping (aa 1) is decreasing with respect f and B(aa t) - as /. co.
The and KLL are two classes of comparison functions, and we can use the comparison function to analyze the monotone contraction in both spatial and temporal horizons.
Definition C.3 (Lyapunov Stability (Gajic Qureshi, 2008)). If the Lyapunov function V is globally positive definite, radially unbounded, the equilibrium is isolated and the time derivative of the Lyapunov function is globally negative definite:
dV
(x) <
O,
VR" \ {0}.
(6)
dt
then the equilibrium is proven to be globally asymptotically stable. The Lyapunov function is a class K. function, which satisfies the condition as follows
e1([==]) < V(x) S a2(([=]]), Vrr E [0, 00).
15

---

Tensor-Var Efficient Four-Dimensional Variational Data Assimilation
Lemma C.4 (Hurwitz stability criterion (Parks, 1962)). A square matrix A E IRux is said to be Hurwitz stable if all the eigenvalues of A have strictly negative real parts, L.e, for every eigenvalue X of A,
Re(X) < O.
(8)
In other words, the real part of each eigenvalue of the matrix must lie in the left half of the complex plane.
Lemma C.5 (Hurwitz stability criterion via Lyapunov function (Sastry, 2013)). Give a candidate Lyapunov function for linear dynamics as
de = Ar.
V(x)= = x xx,
99)
dt where symmetric, positive definite matrix; A governs the evolution of dynamics. For the system to stable, the time derivative dV must be negative definite, i.e., aV. 0 for all O. This means that:
dt
dV d (5f Px) (A" <0
dt (1)
(10)
dt
with
P
(11)
where (AT + P.A) is negative definite.
Based on the convergence analysis and Lyapunov theory, more generalized concept contraction metric is needed to support our paper.
Definition C.6 (Contraction (Manchester & Slotine, 2017)). Given the system bounded metric Mxx, t) (positive definite) such that
da = [ar, (), there exists a uniformly
dt
d.M as (x, M 05 (x,t) -CM, 0,
(12)
di dr
then we call the system contracting, and Mx, t) is contraction metric.
C.2. Proof of Key Theorems
To adopt the convergence analysis of the DA problem, we give mild assumption on the smoothness of dynamics, the first-order derivative exists in the system
Theorem C.7 (Embedding and Consistent convergence). Here, consider system 8t4 Far(s:) = = 5 f(t f(x
(1) first order differentiable with derivative des [(e) for s E Rhd with unique equilibrium point as 8.. If there
di
ddz exists a embedding as (s(s)) = [s) ..... 0g (st))T with that d, E co satisfying the following properties:
a. (embedding) For a finite de. the 2x (s.) is fuulllcoluma rank; when dg infinite, is assumed to be rank-d, countably 0o)
infinite, iee. (Vos(s;)} is full-column rank with Vás(s;) = 1 (st).
003
84 (st)....
b. (convergence) There exists Hurwitz matrix A verifying
døs (s:) = Abs(st).
(13)
dt
Then, the equilibrium 8. and (s(8)) are global asymptotic convergence.
Proof.
(Embedding.) The embedding property follows from RKHS theory. Since we restrict our analysis to separable Hilbert space, it has countable basis either finite or infinite. Thus, the dimension of the RKHS can be infinite, but the rank the embedding determined by countable set of basis functions (Schölkopf Smola, 2002).
16

---

Tensor-Vac Var: Efficient Four-Dimensional Variational Data Assimilation
(convergence.) According to the differential equation in Hilbert space, have
døg (st)
dt
dos dst
(8t) dl. (
(14)
0st
-sbs(st).
The second line of (14) follows from the chain rule, the final line represents the time derivative in Hilbert space, where A is a linear operator that governs the dynamics. Following the work from (Romanoff, 1947; Bobrowski, 2016; Cheng et al., 2023b), we define A as:
Cs - Id
lim +
(15)
where Cs+is the conditional covariance operator between future and current states in the RKHS. Given the smoothness of the kernel function and the differentiability of the system dynamics, the linear operator A exists and is well-defined in this context.
Since 8. is the equilibrium point, we can derive a natural result that f(x.) = 0. Invoking (14) (third line) we have des dt (s+) = 0. Thus (s.) is also a local equilibrium point in RKHS. From the embedding property of RKHS, we have a local injective map Φ IR. + Rd,, which ensures that the convergence properties of the system in the original space are preserved in the feature space. For the neighbourhood around the equilibrium point s., there exist class I. functions or and 022 as
an (1Is; - s.ll) 2 [lØ(s)) 6(s-)ll2 S az(1ls; -s.l), Ys, E ((s,,e).
(16)
B(s,,e) <) is denoted as b-ball centred at 8.. The smoothness of the kemne function and regularity of the dynamics ensure that the system remains well-behaved in feature space, and the convergence properties of the original system carry over to the feature space.
Thus, when the system is locally stable in the original space, the corresponding system in Hilbert space is also locally stable. According to the Hurwitz stability criterion in C.4, the linear operator A has only negative real part in its eigenvalues, guaranteeing exponential convergence. If s, is global equilibrium, then (s(s.) is also global equilibrium in feature space. Remark C.8. Theorem C.7 (a) implies the existence of a global coordinate. In many situations, when the embedding space chosen properly, we can have a stronger result that the existence of left inverse such that of (Ø(s) = s. This result can be naturally connected the Kazantzis Kravaris/Luenberger (KKL) observers (Tran & Bernard, 2023). The embedding corresponds to the injectivity of the state. Theorem C.7 (b) corresponds to the convergence in the KKL observer. When the embedding space is uniformly injective, the dynamics in feature space become rectifiable dynamics, yielding stable trajectory if the original system is stable.
Before entering the last main theorem, we need to introduce the KKL observer. Based on the theory of the KKL observer, will link the convergence problem in feature space with the KKL observer. Please note that assume that the state space S is a bounded set and system (1) is forward complete within this bounded set S.
Consider the nonlinear dynamical system in (1) with time derivative as
ds:
f(s);;
or G(st).
(17)
The design of the KKL observer is as follows:
Find an embedding map T S dd that transforms (17) to new coordinates T(s) as
0T (s)f(s) = AT(s) - BG(s),
(18)
Os
where A E [Rd.xd xd, Hurwitz matrix and B E Rd, such that the system (A, B) is controllable..
R = A2 BB, A (6-1 has full row rank (i.e. rank R) = d.)
17

---

Tensor- Var: Efficient Four-Dimensional Variational Data Assimilation
Since T is injective, the left inverse T exists, i.e., (T) The KKL observer is then given by
(19)
There are certain conditions (Andrieu & Praly, 2006) that (17) needs to satisfy in order to ensure the existence KKL
observer (18) in the sense that a |6; injective if there exists a class IC function a in C.1 such that, for every $1, $2 E S, satisfying || - s2ll s ((T(s)) (s2)) Our embedding properties and convergence conditions (as shown in C.7) are satisfied by the conditions, thus assert the existence of KKL observer.
- O. On the other hand, the map T Rd. need be uniformly
In this paper, we state the existence of a global linear dynamical system in feature space. We provide theoretical guarantee that the embedding property of dis can derive the equivalent convergence in the feature space. However, there are two parts that have not been proven: 1). Does the global linear dynamical system exist? 2) the embedding space in C.7 properly defined? We consider the nonlinear system in (1) within compact space and the cost function in has a unique solution. According to the condition of KKL observer, we guarantee (1) the existence of such linear system and (2) the solutions in the original space and feature space has consistent convergence properties, with respect the cost functions (2) and 5, and convergent exponentially to the unique solution.
Theorem C.9. Let S be bounded set in R If there exists C¹ function T S - Rd which satisfies the following two
conditions:
T solution of the partial differential equation
OT
AT(s)) BG(s), E
(20)
Ös
where A is Hurwitz matrix, and (A, B) is controllable;
There exists Lipschitz constant such that for all (81, 82) in S S, the following inequality holds:
[8] 82 L|T(s)) T(8)));
(21)
then there exists a continuous fontinuousfunctionn T mR + S such for all (s, T()) ESX
(T(+)) - sell S eL T((o) -T(s)) exp max (A)t), VE E (0,
(22)
and where (st, T(s)) is the solution of system in (1) and (17) at time t; mazz(A) is the largest eigenvalue of matrix A.
Proof.
possible to defined the left inverse function T RRd, S and this one satisfies,
|| s2ll S (s)) - ((22)l,, V(s1, 82) E SSSX
(23)
yields that the function T-1 Rdd, + S is global Lipschitz. Hence, the function T RRd. our problem is Lipschitz extension on the set S of this function. For more convenience, we denoted (s)). Following the approach
the Mc-Shane formula in (McShane, 1934), select the Tt as the function defined by
(1w) inf ((T-1(()))))) + T.I|2 w]]]-
(24)
The function is such that for all E S.
(T((s)) =s
(25)
and for all (w, 8) E dRd x S,
[T7(10) >|1 s Vd,L.l T())
(26)
This implies that along the trajectory (st, 24) of the system satisfying the following result
[T()2)) sell S Vd, - T(:))ll Vt E [0, oc).
(27)
18

---

Tensor-Vart Efficient Four-Dimensional Variational Data Assimilation
On the another hand, the function T is solution of the partial differential equation in (18), consequently, this implies that along the trajectory of system (2t, 82) Red S, have
2 (st) exp T((s))
Vt E [0, oo).
(28)
Note, since A is Hurwitz matrix, with 0 max (A) O, it can derive that equation (22) holds and concludes the proof of the theorem.
Remark C.10. The KKL observer asserts that the linear representation of the nonlinear system. After establishing the KKL observer, theorem C.9 asserts the convergence of the estimated trajectory to the equilibrium trajectory. Since the embedding property, we can derive that the existence of left inverse based on the Mc-Shane formula in (24). Meanwhile, the convergence holds when pulling back the state to the original space S. Thus can assert the feature Φs our framework aligns the KKL observer, and the coordinate transformation T is just our feature Øs.
D. Estimation of error matrices and Pseudo Algorithm
D.1. Estimation of error matrices.
To estimate these error covariance operators in the feature space, we empirically estimate these error matrices in the feature spaces from the training dataset (Gejadze et al., 2008). For R and Q, we estimated the covariances as R. ,/ Σi= ry 6 ry and Q nv N 9: 9s. where т; (s(s,) Cs (o;, h,) and q, os(s,) Cs+ s (s;) are the regression residuals, quantifying the errors of the empirical operators Čşjoh and Cs+ Similarly, we compute the background covariance as the empirical variance over an average of { This error should decay monotonically over time and stabilize after a sufficiently long time horizon. This is strongly related to the covariance estimation in Kalman filtering (see Chapter 6.7 (Asch et al., 2016)). We leave the investigation of such design for future efforts.
D.2. Pseudo algorithm.
In this section, we provide the pseudo-algorithm for training Tensor- Var with traditional kernel features in algorithm 2 and training with deep features in algorithm 1. The kernel feature map ø + Hys such that k(s,, sy) (δ(s;), d(s)))ss may not necessarily have an explicit form (e.g., RBF and Matérn kernels), long as the (:, ")n is an valid inner product. For clarity, we use the polynomial kernel with degree two and constant e as an example:
Explicit form, d(s) = k(s8)) )) = -07, 22. /21182,... 288, 18n, C)
Inner product, k(s;, 8,) (o(s:), ((;))) (sT's) cc²
Algorithm outlines the procedure of performing data assimilation with trained models.
'For eyclic application of Tensor-Var better design for B should be time-dependent, reflecting the hetween the estimated system state and the true state, e.g. (Paulin al., 2022).
19

---

Tensor-Vart Efficient Four-Dimensional Variational Data Assimilation
Algorithm 1 Tensor-Vas Var training with deep feature
Require: Data D (8t; oz h4 of K==i Initialized deep features 08g 90a Dag the inverse feature of
training epoch K,
learning rate a, batch size NB
for l,... K do
Random sample batch data Dunch C.
Cs+is. Csjoh Algorithm by using batch data Dontch and deep features
Compute loss (0s) |[Cs+ (s) - Φgg (st) (s+)
Compute loss (@o. (H) ||Csjon (deo (o) Φen (h)] 60 (s)l1
Compute loss [(8s, (/g) = 11ob/ (0gs (s)) s||?
Update the deep features.
8g As + a V I(0s);
0o. Fr = 9o + a V no 1000 On), Hr + VV oo 100, 100, 0n): 0g, 0''g = 0s + aV 02)), 0's + a aV 1(0s, 0'y)
end for
Compute Cs+ Esion = Algorithm 2 by using the whole dataset D and trained deep features 0gg. 200: 08 Compute the error covariance matrices B, R. Q from Subsection D.1 Return 0gg 09. 0a Cs+1s. Csjon. B, and Q
Algorithm 2 Tensor-Var training with kernel feature
Require: Dataset D fsa O₁, h4, 82 }. kernel features (s(s) ks(s, (0))) ko -1. ku (h, Compute the Gram matrix Ks where [Ks]ij = ks(s1,8j) Compute the Gram matrix Koн where [Kon]ı kon (o: @ h; 0j & h2) = ko oz) ) (h: h1) If N is too large, say N 2 10000, using the Nystrom approximation to select subset D> = {si, Oi, his 8 Hi=1 Compute the feature matrix 2 =
Compute the feature matrix 5+ = (ss( (5.t))
Compute the feature matrix ºoz = (0. 11), 20 (On kn)
CME for the system dynamics. C5+15 = Ds (Ks XI) - 4g
CME for the inverse observation model. Eyjon 4 (Kon XI) OH
Compute the eπror covariance matrices 13 R. Q from Subsection D.I.
Fit the projection matrix for pre-image. Epej = S(Ks + AT) Φ5 where S = (81,...,8n)
Return Es+1s. Č3 Eyooj, B, R, and Q
Algorithm 3 Vensor-Var assimilationn
Require: assimilation window for, by TT background state ss; leading time kernel features ás do, Φи (or trained deep features 00g, 60,, Φo,, ): CME operators Cs+1s and ČςjON; Error covariance matrices B, R. Q.
Perform Quadratic Programming with objective
min ||:0 65 (x.)
+
2 - fo))
Σ
||441
Project back to original space with & = Eproy (or using learned inverse feature š = 1b/ (zo))
for i 1,.., T do Predict 2++ Cs+|s2T+L1
Project back to original space with 8, = Gprojze (or inverse feature êT = obgg (2T+1))
end for
20

---

Tensor-Van Efficient Four-Dimensional Variational Data Assimilation
E. Experiment Settings
li=1 and Dohs ([(s) of o The DFs are trained in two steps using these datasets. First, the state DFs Φ05z Oog are trained on Daya and we store the estimated operator Cs+1s- Next, with the state features fixed, the observation DF Φbo and history DF don are trained on Dats according (3.2), storing the estimated operator ĆSJOH. The baseline method (Frerix et al., 2021) is trained on Dats. excluding history. All models are trained with the Adam optimizer (Kingma, 2014) for 200 epochs, using batch sizes from 256 to 1024 for stable operator estimation. Additional details on the DFs, baselines, and training procedures can be found in Appendix E.
Training details. Given the generated data, we construct two datasets: Däyn = {( s 1441) 1
Implementation details. For all baseline methods, we employ the L-BFGS algorithm for Variational Data Assimilation (Var-DA) optimization, implemented in JAX (Bradbury et al., 2018). The 4D-Var baselines use numerical dynamical models based on the 8th-order Runge-Kure Kutta method and the 4th-order Exponential Time Differencing Runge- Kutta (ETDRK) method (Cox& & Matthews, 2002) for the Lorenz-96 and KS systems. For Tensor- Var, we apply interior-poing quadratic programming solve the linearized 4D- 4ar optimization, using CVXPY (Diamond & Boyd, 2016). All training is conducted on workstation with a 48-core AMD 7980X CPU and an Nvidia GeForce 4090 GPU.
E.1. Lorenz 96
First, consider the single-level Lorenz-96 system, which is introduced in (Lorenz, 1996) low-order model of atmospheric circulation along latitude circle. The system state is [Si,... Sk] representing atmospheric velocity at K evenly spaced locations and is evolved according to the governing equation:
dSk = SS--1 - Sk+1) Sx F
dt
with periodic boundary conditions T = Tk. The first term models advection, and the second term represents a linear damping with magnitude F. In general, the dynamics become more turbulent/chaotic as F increases. We choose the number of variables K 40, 80 and the external forcing F 10, where the system is chaotic with Lyapunov time of approximately 0.6 time units. As observation model for the following experiments, we randomly observe 25% states (e.g. 10 in K 40). Our models are trained on dataset D of N = 100 trajectories, each trajectory consisting of = 5000 time steps, generated by integrating the system from randomly sampled initial conditions.
E.I.I. DATA GENERATION.
To generate the dataset, we use the 8th-order Runge-Kute (Butcher, 1996) method to numerically integrate the Lorenz-96 systems with sample step 0.1 and the integration step At size is set 0.01. The system is integrated from randomly sampled initial conditions, and data is collected once the system reaches stationary distribution. For an observation operator, we use subsampling which every 5th and 10th variable for 40 and 80-dimensional system observed via the nonlinear mapping 5 arctan( / 10) + with noise ~ (O, 0.1) (see Figure 6 in Appendix E. for an example). The arctan: : R + [- 2 squeezes the state variable S& into [- 2 I1, which is difficult for inverse estimation. We integrate the Lorenz96 system with observation interval At. = 0.1. The history length is set as 10 such that h4 =
1.2. ADDITIONAL EXPERIMENT RESULTS.
We provide qualitative results in Figure for the Lorenz 96 system at two different dimensions: 40 (left) and 80 (right). Each subplot illustrates the normalized absolute error for various methods, including 3D-VAR, 4D-VAR, (Frerix al., 2021), and Tensor- Var, compared to the ground truth. The assimilation window length is set to 5 (indicated by the red dashed line), with forecasts extended for an additional 100 steps based on the assimilated results. Tensor- Var generally outperforms the other methods in both dimensions and maintains stable long-term forecasts, comparable to other model-based approaches. For 3D- and AD- with partially observed models, the observed states show minimized errors (indicated by the dark lines), while the errors of unobserved states remain uncontrollable, as clearly shown in Figure 7.
21

---

Tensor-Var Efficient Four-Dimensional Variational Data Assimilation
10
State
Observation
4-
0-
-2
-4-
0
10
15
20
25
30
35
40
State Coordinate
Figure 6. Nonlinear observation model in the Lorenz 96 system: the state values are represented by the solid blue curve, with the observed grid points indicated red dots.
Ground Truth
Ground Truth
-AR Error
3D VAR Error
0.75
0.50
0.50
4D- VAR Error
4D- 4D-VAR Error
0.50
0.50
0.25
0.25
Frerix et al. (2021) Error
Frerix et al. (2021) Error
20
0.50
0.25
Tensor Var Error
Tensor Error
0.75
0.50
0.25
60
100
(a) 40 dimensional Lorenz 96 system
(b) 80 dimensional Lorenz 96 system
Figure 7. Qualitative enror comparison for the Lorenz 96 system at (a) 40 dimensions and (b) 80 dimensions. The plots show the normalized absolute errors various methods, including VAR, VAR, Frerix (2021), and Tensor- Var, compared to ground truth. The assimilation window length is set to 5 (indicated by the red dashed line), with forecasts extended for an additional 100 steps based on assimilated results.
22

---

Tensor-Van Var: Efficient Four-Dimensional Variational Data Assimilation
E.2. Kuramoto-Sivashinsky
Next, we consider the Kuramoto-Sivashinke (KS) equation, a nonlinear PDE system known for its chaotic behaviour and widely used to study instability in fluid dynamies and plasma physies (Papageorgiou & Smyrlis, 1991). The dynamics in spatial domain u(x, is given by.
du
Du
8u 2x
dau + dx1
+
= 0,
Ot
Oz
where 2 E [0, L] with periodic boundary conditions. We set the domain length L = 32x, large enough to induce complex patterns and temporal chaos due to high-order term interactions (Cvitanović et al., 2010). The system state u(x, t) was discretized into n, 128 and n, 256. The observation model is the same as Lorenz-96, where 25% states can be observed. In this case, our models are trained on dataset D consisting of N 100 trajectories, each with L 5000 time steps and discretization of At 0.01, sampled from the stationary distribution with different initial conditions.
E. DATA GENERATION.
To generate the dataset, we use the exponential time differencing Runge-Kutta method (ETDRK), which has proven effective computing nonlinear partial differential equation (Cox & Matthhes, 2002) with an integration step At = 0.001 and sample step 0.01. The system is integrated from randomly sampled initial conditions, and data is collected once reaches a stationary distribution. 256-dimensional systems, we use subsampling which every 8th for both 128 and 256-dimensional system are observed with noise (0, 1), and 5 arctan( /0) nonlinear mapping (see Figure in Appendix E. for an example). The history length is set to 10 as well.
observations, we use subsampling, observing every 8th state both 128- and
E.2.2. ADDITIONAL EXPERIMENT RESULTS.
Ground Truth
Ground Truth
320
100
200
3D-VAR Error
3D-VAR Error
0.75
0.75
50-
318
0.50
300
VAR Error
4D-VAR Error
0.75
0.75
0.50
0.50
3.25
0.25
Frerix et al. (2021) Error
Frerix et al. (2021) Error
3.50
320
0.25
0.55
228
Tensor-Var Error
Tensor-Va Var Error
3.75
0.75
50
322
100
0.25
328
(a) 128 dimensional KS system
(b) 256 dimensional KS system
Figure 8. Qualitative error comparison for the KS system at (a) 128 dimensions and (b) 256 dimensions. The plots show the normalized absolute errors various methods, including 3D- VAR. AD-.AR. Frerix et (2021), Tensor- Var, compared the ground truth. The assimilation window length is set to (indicated by the red dashed line), with forecasts extended for an additional 100 steps based on the assimilated results.
23

---

Tensor-Van Var: Efficient Four-Dimensional Variational Data Assimilation
We provide qualitative results in Figure 8 for the KS systems at different dimensions: 128 (left) and 256 (right). Each subplot illustrates the normalized absolute error for various methods, including 3D-VAR, 4D-VAR, Frerix et al. (2021), and Tensor-Vat compared to the ground truth. The assimilation window length is set to 5 (indicated by the red dashed line), with forecasts extended for an additional 100 steps based on the assimilated results.
Compared to the Lorenz-96 system, the KS system is more complex, being governed by partial differential equations (PDEs) that account for spatial evolution. In both dimensions, Tensor-Var Var consistently outperforms other methods, particularly in capturing chaotic dynamies during the initial forecast phase. also maintains long-tern stability in more complex PDE system, comparable other model-based approaches. In contrast, 3D-VAR struggles with assimilation, especially in the 356-dimensional case, due its inability to capture temporal evolution, leading to rapid error divergence. This underscores the critical importance of temporal modeling in chaotic systems. A similar pattern of error between observed and unobserved states is evident in Figure 8.
E.3. Global NWP
We consider global numerical weather prediction (NWP) problem using the European Centre for Medium- Weather Forecasts (ECMWF) Atmospheric Reanalysis (ERAS) dataset (Hersbach et al., 2020). This dataset provides high-resolution atmospheric reanalysis data from 1940 to the present, offering the most comprehensive estimate of atmospheric dynamics. For our proof of concept, we focus on five upper level physical variables: 500 hPa geopotential height, 850 hPa temperature, 700 hPa humidity, and 850 hPa wind speed (meridional and zonal components) at a spatial resolution of 64x32.
The data is sourced from the WeatherBench2 repository (Rasp et al., 2024). From this dataset, we randomly sample grid points with 15% spatial coverage. The sampled observations include additive noise equivalent to 0.01 times the standard deviation of the state variable, ensuring robustness against observational uncertainty (see Figure 9). For model training, we use ERA5 data from 1979-01-01 to 2016-01-01, separating data from post-2018 for testing. There are 51, consecutive system states with generated observations for training and 2,920 data for testing.
Background State
Observation
True State
Latent-3DVar error
Latent-4DVa error
Tensor- error
Z500
110
1201
1850
8700
1850
v850
31
Figure 9. Visualization of assimilation results for five variables from FRA5 data time 2018-01-01 00:00. Each column (from left right) displays the background state, observations, true state, and errors for Latent-3DVa: Latent-4DVar, and Tensor- Var. The reported error was weighted absolute error each pixel; see Appendix E.3 for more details.
In addition to the results presented in the main experiments, we evaluate the forecasting quality of Tensor- Var based on the assimilated state. Figure 10 shows the mean RMSE weighted by latitude (Rasp et al., 2024) for five variables predicted by Tensor- Var at various lead times 7, where T = 0 represents the assimilation error in the final state of the assimilation window.
24

---

Tensor-Var Efficient Four-Dimensional Variational Data Assimilation
Geopotential
Temperature
humidity
wind speed - longituda direction
wind speed latitude direction
100
1.0
RMSEE
RMSE
PMAS
RMEE
0.8
100
60
022
1.0
0.4
10
30
20
10
20
10
23
30
10
30
Leading (per hours)
Leading (per hours)
Leading time 6 hours)
Leading hours)
Leading time hours)
Figure 10. The (non-cyclic) forecasting quality of Tensor-Var in NWP experiments with leading time zero as the final state in the assimilation window, evaluated across different experiments. The five sub-figures display the NWP forecast for 5 variables (15-day in total).
Area-weighting Root mean squared error (RMSE). The error is defined for each variable and level
(w())rris S1,
which is area-weighting over grid points. This is because on an equiangular latitude latiuudelongitude grid, grid cells at the poles have a much smaller area compared to grid cells at the equator. Weighing all cells equally would result in an inordinate bias towards the polar regions. The latitude weights wii) are computed as:
sin (A sin (8)
w ii = E (sin (0"*) sin (0)))
0% and of indicate upper and lower latitude bounds, respectively, for the grid cell with latitude index i.
25

---

Tensor-Van Var: Efficient Four-Dimensional Variational Data Assimilation
E.4. ASSIMILATION FROM SATELLITE OBSERVATIONS
E. DATA GENERATION.
We collect the weather satellite track data from https: / celestzaa. org/NORAD/ wwtaa//eeeeeeetee for the period 1979- 01-01 00:00:00 to 2020-01-01 00:00:00. Observations are matched to the high-resolution ERA5 dataset (240 x 121 grid) by identifying the nearest neighborhood grid points along the satellite track to generate observations. Furthermore, we use an observation frequency of up to every half hour within the 2 hours before each assimilation time and add white noise with a standard deviation of 1% of the standard deviation of the corresponding state variables.
E.4.2. ANOMALY CORRELATION COEFFICIENT EVALUATION
In addition to the NRMSE results reported Subsection 4.3, we also evaluate the assimilated fields using the Anomaly Correlation Coefficient (ACC), which measures the pattern similarity between anomalies in the analysis (or forecast) and reference field relative climatological mean. This metric is widely used in numerical weather prediction assess the accuracy of large-scale spatial patterns. We compare the ACC performance of our method with the SOTA data-driven baseline, FengWu 4D- -ar (Xiao et al., 2024). ACC is formally defined as:
(, (5: 8) VE, -3 V2 2, 8)2"
ACC
where & the result of assimilation, 8 is the corresponding reference value (e.g. ERA5), and is the climatological mean. Figure 11 reports the mean and standard deviation of ACC for five atmospheric variables over a 7-day forecast horizon, with assimilation window of length 5. The results indicate that Tensor- Var consistently achieves higher ACC values than FengWu Wu 4D- Var over long horizon.
1.00 ACC 2500
ACC- 1850
ACC- Q700
M
0.504
8550
0.95
8550
Cos15
50c
8.50
Fengivu 40-Var
8.00
Feng/Wu 4D-Vir
FengWiu 4D-ar lenson
Verson-er
20
Time
ACC u850
ACC v850
0.30
0.00
3825
0.32
FengWru Var
FengWuu 4D
Var
Tansor-ar
25
10
25
Timo
Time
Figure 11. Comparison of ACC over 7-day horizon for five atmospheric variables from Fengwu 4D- and Tensor Var. Each time-step represents 6-hour interval.
E.4.3. LONG ROLL-OUT STABILITY TEST
To assess the long-term stability of our method, we conduct a one-year roll-out experiment under the following settings: spatial resolution of 240 121; meteorological variables (z500, t850, 9700, u850, v850]; a 6-hour time step; assimilation window length of 5 steps; and forecast lead time of days. The experiment covers the full year from January I to December 31, 2018, with training and evaluation settings consistent with Section 4.3.
Figure 12 illustrates the long-term behavior of Tensor- Var in comparison to FengWu Var. The results show that Tensor- Var effectively controls the estimation emror over extended periods and achieves lower error magnitudes than FengWu 4D-Van throughout most of the year. However, we observe that Tensor- Var exhibits instability after approximately 800 assimilation steps, whereas Feng Wu maintains stable performance.
26

---

Tensor-Var Efficient Four-Dimensional Variational Data Assimilation
Z500
t850
30
25
Forecast Fenglity 60 var
Forecast: Fenglity 40
Forecest Fengiuu
Forecast Fenglely
25
Foreced Temmrrder
20
I'reoasi-OA. Beraceller
NRMSE (90
20
NRMSE (%)
15
1400
800
1000
1200
Time
Time
q700
u850
Forerast- Tengttu 4D Var
Forecast Tengli: 4D
Torecast DA Fengiku VD-Ver
Terecas-OA- Fengliti 4D-War
Forecast Throo-Warr
25
Forecast Temag-Var
Forecant-Ok DA lenazriar
Forecasz DA. Tarserlbrr
NRASE (96)
NRMSE (%)
15
Parhh
800 Time
v850
Tersorkar Forecas:
NRMSF (96)
1.5
200
400
600
800
1000
1200
1400
Time
Figure 12. Long-ter roll-out stability test for the five metcorological variables one one-year horizon. The plot compares the NRMSE of forecasts (dashed) and forecast-DA (solid) outputs from Tensor- Var and FengWu 4D-Var.
The observed instability can be attributed to the linear dynamical structure of Tensor- Var, which is able to handle short- term forecasting but can be insufficient for capturing nonlinear long-term dynamics. In contrast, the transformer- based architecture of Feng Wu 4D-Var appears to more robust for extended temporal integration. These findings highlight a trade-off between assimilation efficiency and long-horizon stability. Future work will focus on improving Tensor- Var's dynamical expressiveness while preserving its computational advantages.
27

---

Tensor-Vart Efficient Four-Dimensional Variational Data Assimilation
E.4.4. ADDITIONAL QUALITATIVE RESULTS
2018-01-01
2018-01-02
2018-01-03 00:00
2018-01-04 00:00
2018-01-05
00:00
00:00
00:00
True State
Tensor-Var
Result
Absolute Error
SOOOO
Figure 1.3. Visualization of continuous assimilation results, absolute errors, and observation locations for 1850 (temperature)) starting from
2018-01-01 00:00.
2018-01-01
2018-01-02 00:00
2018-01-03
2018-01-04
2018-01-05
00:00
00:00
00:00
00:00
True State
2306
0250
2315
0.03
6.05
6.30
0.03
C1.0
0005
621
2315
0065
LEG
3015
0022
6400
4315
0003
Tensor-Vay Vaa Result
C003 2410 2013
0020
0.03 2310
0:03
0400
0003
520
2313
0.02
Absolute Error
0.00 5.00 1.32
535
Son
6230 CASL
C431
ES1
E522
ESS EIEL
DES
DA0E 0000 6.02
5009 530 San DAGE D0C 6002 5301 5.00 150
Observations
Figure 14. Visualization of continuous assimilation results, absolute errors, and observation locations for q700 (humidity), starting from 2018-01-00 00:00.
28

---

Tensor-Var Efficient Four-Dimensional Variational Data Assimilation
2018-01-01
2018-01-02 00:00
2018-01-03 00:00
2018-01-04 00:00
2018-01-05
00:00
00:00
True State
Tensor-Var
Result
BOCOO 
Absolute Error
Observations
Figure 15. Visualization of continuous assimilation results, absolute errors, and observation locations for u850 (meridional wind speed), starting from 1018-01-01 00:00.
2018-01-01
2018-01-02 00:00
2018-01-03 00:00
2018-01-04 00:00
2018-01-05
00:00
00:00
EebiBer TSNIN Pebiaee TRRIDCE
Tensor-Var Result
Absolute Error
COOOO
Figure 16. Visualization of continuous assimilation results, absolute error,, and observation locations for v850 (zonal wind speed), starting from 2018-01-01 00:00.
29

---

Tensor-Van Var: Efficient Four-Dimensional Variational Data Assimilation
E.5. Ablation study
In this section, we provide the details of the ablation studies of (1) the history length m, (2) the dimensions of the feature dimension and comparison with standard kernel functions, and (3) the effects of the estimated error matrices. All ablation studies are conducted on Lorenz-96 systems with n. 40 and 80. We fix the remaining hyperparameters consistent with the main experiments and vary only the parameters under investigation.
(1) Effect of the history length m. We examine the impact of the history length m on learning the inverse observation operator and its effect on state estimation accuracy. The feature dimensions d,, do. and d₁, remain constant, while the history length is varied by adjusting the size of the final linear layer (see Table 5). The state feature dimensions dg are set to and 120 for the two system dimensions.
(2) Effect DFs. We implement Tensor- Var with Gaussian kernel using kernel PCA projected the first 60 and 120 cigen-coordinates in scikit-learn (Pedregosa al., 2011) by aligning with the dimension of the used DFs. For the Gaussian kernel, approach the pre-image problem by fitting a projection operator (Kwok & Tsang, 2004). The space S C R together with the linear kernel k(s;, 83) of 8, also forms RKHS. Therefore, we can define the projection operator a CME operator that maps from the feature space Hs the original space S as Eproj S (Ks + - where = [1 sn]. By applying Gmeoj now to the state, we can obtain the mean estimation the kernel mean embedding in such that Croot/p Es (Cpoojtss(S) (8)] Es[S]. We normalize the dataset to standard Gaussian distribution and select the length scale = 1.0 by performing cross-validation on (0.5, 0.75, 1.0, 1.25, 1)
E.6. Model architecture
Table Deep feature architecture for chaotic systems with dimensions na: No. 7h.
m X n. and feature dimension da, do, da
Components
Layer
Weight size
Bias size
Activation
Fully Connected
ns 4n,
4n
Tanh
605
Fully Connected
4n, x 2n.
2n,
Tanh
Fully Connected
2n, x d,
d,
Fully Connected
d, x 2n
2n
Tanh
0%/
Fully Connected
2n, x 4n,
4n,
Tanh
Fully Connected
4n, x nss
nss
Fully Connected
no x 4n
4no
Tanh
doo
Fully Connected
4n x 2no
2no
Tanh
Fully Connected
2n, do
d,
Convolution 1D Max Pooling (size=2)
m 2m x 5
2m
Tanh
Convolution ID
2m 4m X
4m
Tanh
Max Pooling (size=2)
Flatten
Fully Connected
mn, x dn
dp
30

---

Tensor-Var Efficient Four- Your-Dimensional Variational Data Assimilation
Table Model architecture for Global NWP with input dimension (H, W, C) with physical variables and spatial resolution H x W. The implementation of vision Transformer (ViT) block follows (Zamir al., 2022) with applications in DA followed by (Nguyen Fablet, 2024).
Components
Layer
Layer number
C. W)
Activation
Convolution2d
4C,C (H, W)
Transformer Block
2
4C 4C, ( W 2
ReLU ReLU ReLU
60.
Transformer Block
3
4C - 8C, H W 1 H W 8C - SC, 8
Transformer Block
Flatten
(8C, 8 CHW
Fully Connected
1
CHW + do
Fully Connected
d, + CHW 8
Transpose
CHW 8 (8C, H1 )
00's
Transformer Block
3
8C + ( (8) 8
H W
ReLU
Transformer Block
3
8C + 4C, W
1. 4C C, (H,
ReLU ReLU
Transformer Block
2
40 + 4C4 H W
Convolution2d
1
Convolution2d
-C, (H
Transformer Block
2
2C - 2C, (H W
ReLU ReLU
(112) 20 40, (4, W
69o
Transformer Block
3
Flatten
(4C, H W
CHW
81 8
16
Fully Connected
CHW 16
+ d
Convolution2d
mC - 2C, (H,W)
Transformer Block
2
20 + -0, 21
H W
ReLU ReLU
690
Transformer Block
2C + 4C, ( W
Flatten
(4C, 11 M.
CHW
Fully Connected
CHW + d/ 16
31


Reviews:
[{"review_id": "bF48ziopcj", "reviewer": "Reviewer_Qu2S", "summary": "This paper introduces 4d Tensor-Var, a framework for performing data assimilation with learnable features. Building upon CME to parameterize the state space model, the authors propose two versions of constructing the features space.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank reviewer Qu2S for your thoughtful review and constructive feedback on our manuscript. Your comments have been invaluable in improving the clarity and presentation of our work. Below, we address each of your comments.\n\n- **Reorganize the method section**\n  \n    Thanks reviewer’s suggestion, we have improved clarity in the methods section as follows:\n\n    - Added a dedicated Related Work section to clearly distinguish the literature review from our methodological contributions.\n    - Included a diagram illustrating the training procedure for CME opeartor $\\mathcal{C}_{S^+|S}$ (can be accessed in the [anonymous Google Slides](https://docs.google.com/presentation/d/1SWXKhJL_4mAKVQlY4HAxQTl_79C94m-2TwgtkYIvuNw/edit?usp=sharing))\n    - Moved the algorithm (previously Algorithm 1) to the main text (end of Section 3.1), revising it to highlight key steps.\n  \n- **Clarification of Notations**\n  \n   Thank you for carefully reviewing our manuscript. We would like to clarify the following:\n   \n   - We do not state the CME operator as a mapping from the observation space $\\mathcal{O}$ to latent space $\\mathcal{S}$. Instead, we defined it as $\\mathcal{C}_{S|O}\\colon\\mathbb{H}_O\\rightarrow\\mathbb{H}_S$, where $\\mathbb{H}$ with subscript denotes the corresponding feature (latent) spaces. \n \n      The forward CME operator is consistently expressed as $\\mathbb{H}_{S^+}\\rightarrow\\mathbb{H}_{S}$ later in the manuscript.\n\n   - Algorithm 2 (Appendix C) focuses on Tensor-Var with predetermined kernel features, without learnable parameters $\\theta$ as in deep features.\n\n   - We realize the notations in our manuscript may be unclear. To improve clarity, we have added a dedicated table in the Appendix of the updated manuscript. You can also find it in the shared Google Slides above.\n\n- **Predetermined kernel features compare to deep features**\n\n    We would like to clarify that we conducted an ablation study to compare different feature dimensions and the Gaussian kernel (Section 4.4, Table 3). Details are in Appendix D.5. Additionally, we tested other kernel features, including Matern-1.5 and Matern-2.5.\n\n  **Table: Comparison of features using NRMSE (%) for state estimation accuracy.**\n  | Features/Experiments  | Lorenz-96 $n_s=40$ | Lorenz-96 $n_s=80$|\n  |-----------------------|----------------------|----------------------|\n  | $d_s=20$            | 16.7±2.1              | 17.3±2.6              |\n  | $d_s=40$            | 14.1±1.3              | 16.8±2.2              |\n  | $d_s=60$            | 8.3±0.9               | 11.4±0.9              |\n  | $d_s=120$          | 9.7±0.7               | **9.0**±0.9              |\n  | Gaussian Kernel       | 8.4±0.5               | 14.3±1.4              |\n  | Matern-1.5            | 9.6±0.8               | 13.2±1.1              |\n  | Matern-2.5            | **7.2±0.2**           | 9.7±**0.7**           |\n\n  In $n_s=40$, Matern-2.5 shows the lowest mean and standard deviation (std). In $n_s=80$, the deep feature with $d_s=120$ outperforms other methods, despite a slightly higher std than Matern-2.5. These results indicate fixed kernels are more reliable at low dimensions, while deep features excel as dimensionality increases. For high-dimensional weather forecasting, fixed kernel features often scale poorly and struggle with structured data. Thus, we propose deep features to extend Tensor-Var for broader applications.\n\n- **Comparing Training Time Across Methods**\n\n  We thank reviewer for the interest in training time. To address this, we measured the training time (per epoch) under identical configurations (NVIDIA RTX 4090) for the experiments in Section 4.2. These results are summarised in the following table and are included in the revised manuscript (Appendix D).\n\n  |Methods|Latent 3D-Var|Latent 4D-Var|FengWu 4D-Var|Tensor-Var|\n  |-|-|-|-|-|\n  |Training time (min)|0.75±0.06|0.88±0.41|11.05±1.65|5.32±1.07 (forward), 4.17±0.31 (inverse)|\n\n    Latent 3D and 4D-Var are computationally efficient as they only consist of convolutional and fully connected layers, requiring only one forward pass in training. FengWu 4D-Var, with its transformer-based encoder-decoder, is more time-consuming due to self-attention operations. Tensor-Var’s forward and inverse models are slower than latent DA due to solving least square problems in feature spaces, while the inverse model is faster as no decoding process.\n\n- **Suggested References**\n  \n  We appreciate suggested references. After careful review, we agree that they are relevant to our study. We have properly cited them in the updated manuscript.\n\nWe thank reviewer Qu2S again for invested efforts and constructive feedback, which helped improve the manuscript. We have carefully addressed all concerns and made substantial revisions to improve clarity. We look forward to any further feedback.", "claims_and_evidence": "I would recommend reorganizing the presentation: currently, the method section is a mix of literature review and modification based on past works. This format creates a barrier to accessing the method. A diagram or an algorithm block (like the ones in Appendix C with more clarification) would be ideal to present the method.\n\nNotations are also confusing: For example, the authors defined the CME (conditional mean embedding) as a mapping from the observation space ($\\mathcal{O}$) to latent space ($\\mathcal{S}$) on page 3, but later on, when they introduced state feature, the notation they adopted seem to indicate that the CME embedding is from one-step forward feature space $\\mathcal{S}^{+}$ to $\\mathcal{S}$. This inconsistency without explanation greatly impacts the presentation quality.\n\nAnother example we could find is regarding the definition of feature space $z_t$, and the feature mapping function $\\phi_{\\theta_{\\mathcal{S}}}$ which maps from $\\mathbb{R}^{n_s}$ to $\\mathbb{R}^{d_s}$. The authors did not discuss at all what $n_s$ and $d_s$ mean, and the feature $z_t$ is not introduced and suddenly applied to the loss function in Eqn. 5.", "methods_and_evaluation_criteria": "The dataset looks good to me: The authors conducted experiments on chaotic PDE data (KS, Lorenz systems) as well as real weather data (ERA5).\n\nMoreover, I would be curious to see how the training time compares between different methods.", "theoretical_claims": "Yes, I've checked their theoretical claims. However, because of the ambiguity in the method, I am not sure if I am capable of fully accessing their theoretical claims.", "experimental_designs_or_analyses": "I do like the empirical studies they perform, as well as the intuitive diagram they provided in Figure 1.\n\nBesides the current results, I am curious about the comparison between Algorithms 1 and 2: how do fixed kernel features compare to deep features?", "relation_to_broader_scientific_literature": "I think the field of research is important.", "essential_references_not_discussed": "I will recommend adding a discussion with the following literature:\n\nKalmanNet: Neural Network Aided Kalman Filtering for Partially Known Dynamics (https://ieeexplore.ieee.org/document/9733186)\n\nData Assimilation with Machine Learning Surrogate Models: A Case Study with FourCastNet (https://arxiv.org/abs/2405.13180)", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Please check the \"claims\" part. I do think improving the clarity of this draft is important.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "xydlSTliuU", "reviewer": "Reviewer_dYtv", "summary": "This paper proposes Tensor-Var, a framework combining kernel conditional mean embedding (CME) with 4D-Var to linearize nonlinear dynamics and enable convex optimization in a learned feature space. It addresses the high computational costs and theoretical limitations of traditional 4D-Var and deep learning hybrids. Tensor-Var incorporates deep features (DFs) for scalability and provides theoretical guarantees for consistent assimilation. Experiments on chaotic systems and global weather prediction show Tensor-Var achieves higher accuracy and a 10- to 20-fold speedup over conventional and DL hybrid 4D-Var methods, offering an efficient and theoretically sound approach for data assimilation.\n\n## Update after rebuttal\nI appreciate the authors' thorough efforts in addressing my feedback. The explanations and revisions have resolved my concerns, and I maintain my current evaluation of weak accept.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate valuable and constructive suggestions from Reviewer dYtv. Your efforts have significantly contributed to improving the quality of our paper. Below, we provide detailed responses to each of your comments.\n\n- **Higher-resolution experiment**\n  \n  We would like to clarify that the experiments in Section 4.3 are conducted at a resolution of **240×121**, which is roughly 14 times higher than the 64×32 resolution used in Section 4.2.\n\n- **Compare with sota ML-based DA methods & Baseline selections**\n  \n   We understand the reviewer’s concern about comparing with state-of-the-art ML-based DA methods, particularly FengWu-4DVar. To address this, we conducted additional evaluations of FengWu-4DVar in two weather data assimilation experiments. We used the official implementation from the [GitHub repository](https://github.com/OpenEarthLab/FengWu-4DVar). Training and evaluation settings follow our implementation as described in Sections 4.2 and 4.3.\n  - We have summarized the updated results for Figure 2 (page 7) for Section 4.2 into the following table for clarity. \n  - For the experiment with higher resolution, we have also added the additional results of FengWu-4DVar into Figure 3 (page 7). Both updated figures (2 and 3) are available in the [anonymous Google Slides](https://docs.google.com/presentation/d/1SWXKhJL_4mAKVQlY4HAxQTl_79C94m-2TwgtkYIvuNw/edit?usp=sharing), providing a clearer demonstration of the results.\n\n      **Table: Assimilation NRMSE (%) across variables and evaluation time (s).**\n      | Method|z500| t850| q700| u850| v850| Evaluation Time |\n      |-|-|-|-|-|-|-|\n      |Latent 3D-Var|3.10±0.28|3.83±0.28|11.4±0.65|6.40±0.38|6.36±0.36|1.84±0.06|\n      |Latent 4D-Var|3.02±1.12|3.89±1.29|13.79±3.17|6.57±1.07|6.31±0.83|4.90±0.26|\n      |FengWu 4D-Var|2.80±0.23|2.88±0.16|5.00±0.43|5.51±0.19|5.40±0.24|1.81±0.11|\n      |Tensor-Var|**1.79±0.11**|**1.94±0.09**|**4.77±0.19**|**3.19±0.14**|**3.47±0.15**|**0.83±0.02**|\n\n      The results demonstrate that **Tensor-Var outperforms FengWu-4DVar in terms of both accuracy and computational efficiency**. Similar results are observed in higher resolution experiments. We provide result figures along with detailed discussions in the shared Google Slides for the reviewer's potential interest.\n\n- **Prediction Model Details**\n  \n  In Tensor-Var, predictions are made using the fitted linear operator $\\mathcal{C}_{S^+|S}$ without the need for additional learnable models. To improve clarity, we have included a demonstrative figure of the fitting procedure in the revised manuscript's Appendix C and made it available in the Google Slides.\n\n- **Stability Test**\n  \n  We thank reviewer for highlighting the importance of long-term stability. In response, we added a long-term rollout stability test with the following key settings:\n\n  - Spatial Resolution: 240×121.\n  - Meteorological Variables: z500, t850, q700, u850, v850.\n  - Time Step: 6 hours per step.\n  - Assimilation Window Length: 5 steps.\n  - Test Horizon: One year (January 1, 2018 - December 31, 2018).\n  - Leading Time: 7 days.\n  - others: Training and evaluation settings follow Section 4.3.\n  \n  The result figures can be found in the Google Slides. These results demonstrate that Tensor-Var consistently controls the estimation error over one-year horizon. To benchmark performance, we also include the results of FengWu 4D-Var. Our experiments show that **Tensor-Var achieves better error control with DA compared to FengWu 4D-Var**. However, Tensor-Var rollouts unstablely after 800 steps, where FengWu demonstrates more stable long-term rollouts.\n  \n  Tensor-Var’s long-term instability likely stems from linear dynamics, which only capture short-term evolution accurately. In contrast, transformer model in FengWu better handles long-term evolution. We value the reviewer’s insight and will work on improvng Tensor-Var’s stability while maintaining DA efficiency and accuracy.\n\n- **Different metrics** \n \n   We thank reviewer for suggesting using different metrics. To address this, we evaluated FengWu 4D-Var (considered cutting-edge) and Tensor-Var using the anomaly correlation coefficient (ACC) at a resolution of $240*121$. Other settings follow Section 4.3 Figure 3. The results demonstrate that **Tensor-Var outperforms FengWu 4D-Var in ACC** as well. The corresponding figures can be found in the shared Google Slides.\n\n\n- **Clarification of Figure 10**\n  \n  We thank reviewer for carefully reviewing our manuscript. The reported metric is *normalized RMSE (NRMSE)*, and we have clarified this in the revised manuscript in Appendix D.3.\n\n- **Suggested References** \n  \n  We thank for the suggested references. After careful review, we have properly cited them in the revised manuscript and discussed their relevance.\n\nIn the end, we would like to thank reviewer dYtv again for the valuable feedback and invested efforts.  We believe that the revisions address your concerns and improve the overall confidence to our work.", "claims_and_evidence": "Claim 1: Efficiency and Accuracy Improvements:\nThe paper asserts that Tensor-Var outperforms conventional and ML-hybrid variational DA baselines in both accuracy and computational efficiency. This claim is supported by experimental results on chaotic systems and global NWP applications, showing significant improvements in accuracy and a notable reduction in computational cost. However, the comparison with state-of-the-art ML-based DA methods (e.g., FengWu-4DVar) is not provided, which could strengthen the claim of superiority.\n\nClaim 2: Handling Incomplete Observations:\nThe authors claim that Tensor-Var improves robustness and accuracy by deriving an inverse observation operator that incorporates historical data to infer system states. While this is a promising approach, the evidence for its effectiveness is limited to the presented experiments. Broader validation on diverse datasets with varying levels of observation incompleteness would strengthen this claim.\n\nClaim 3: Scalability via Deep Features (DFs):\nThe paper addresses scalability by learning adaptive DFs to map data into a fixed-dimensional feature space. While this is demonstrated on the 64×32 ERA5 dataset in the experiments, the scalability of the approach for even larger-scale problems (e.g., higher-resolution data) is not thoroughly explored, leaving room for further validation.\n\n[1] FengWu-4DVar: Coupling the Data-driven Weather Forecasting Model with 4D Variational Assimilation", "methods_and_evaluation_criteria": "Yes, the methods and evaluation criteria make sense. Tensor-Var uses kernel embedding and deep features to linearize nonlinear dynamics, addressing computational and non-convex challenges in 4D-Var. Experiments on chaotic systems and global NWP applications demonstrate improved accuracy and efficiency compared to conventional and ML-hybrid baselines", "theoretical_claims": "Yes, the theoretical seems correct.", "experimental_designs_or_analyses": "Evaluation Setup:\nAll metrics are evaluated 20 times with different initial conditions. However, 20 initial fields are insufficient for data assimilation tasks, as it is common to evaluate over an entire year to ensure robustness and reliability.\n\nBaseline Selection:\nThe choice of baselines could be more diverse. For example, DABench [1] provides a comprehensive set of advanced AI-based and numerical assimilation ensemble methods that could serve as stronger baselines for comparison.\n\nEvaluation Metrics:\nThe evaluation relies solely on NRMSE. Including additional metrics, such as latitude-weighted RMSE and ACC, would provide a more comprehensive comparison with state-of-the-art DA methods and offer readers better insights into the performance.\n\nPrediction Model Details:\nThe details of the prediction model used in the assimilation process are not provided. A clear description of the model's architecture, training, and implementation is necessary for reproducibility and thorough evaluation.\n\n[1] DABench: A Benchmark Dataset for Data-Driven Weather Data Assimilation", "relation_to_broader_scientific_literature": "The Tensor-Var framework advances the broader scientific literature by integrating kernel embeddings and deep features to linearize nonlinear dynamics in variational data assimilation (DA), addressing challenges like non-convex optimization and scalability. It builds on prior work in ML-enhanced DA but introduces theoretical guarantees and scalable deep features, achieving significant improvements in efficiency (10- to 20-fold speedup) and accuracy. This bridges gaps in traditional and ML-hybrid DA methods, offering a robust solution for real-time applications like numerical weather prediction.", "essential_references_not_discussed": "[1] Verma, Y., Heinonen, M., & Garg, V. (2024). ClimODE: Climate and weather forecasting with physics-informed neural ODEs. arXiv preprint arXiv:2404.10024.\n\n[2] Vaughan, A., Markou, S., Tebbutt, W., Requeima, J., Bruinsma, W. P., Andersson, T. R., ... & Turner, R. E. (2024). Aardvark weather: end-to-end data-driven weather forecasting. arXiv preprint arXiv:2404.00411.\n\n[3] Han, Tao, et al. \"Fengwu-ghr: Learning the kilometer-scale medium-range global weather forecasting.\" arXiv preprint arXiv:2402.00059 (2024).", "other_strengths_and_weaknesses": "Lack of Stability Testing: The authors do not conduct stability tests for the assimilation and forecasting cycle over extended periods. For example, it is unclear whether the proposed DA method can achieve stable predictions in a year-long rolling forecast scenario. This is a critical gap, as long-term stability is essential for operational DA systems.\nLowe \nResolution Limitations: All experiments are conducted at a 32x64 resolution, and it is unclear whether the conclusions scale to higher resolutions. Adding experiments at higher resolutions (e.g., 0.25°) would make the study more solid and practically applicable, as higher resolutions are critical for capturing fine-scale weather phenomena and operational forecasting.", "other_comments_or_suggestions": "See each part above.", "questions_for_authors": "Higher-Resolution Experiments:\nCan Tensor-Var's performance be validated at higher resolutions (e.g., 0.25°)?\nSuccessful high-resolution results would strengthen practical relevance and increase my score.\n\nStability Testing for Rolling Forecasts:\nCan the authors demonstrate Tensor-Var's stability in year-long rolling forecasts?\nEvidence of long-term stability would address a critical limitation and improve my evaluation.\n\nComparison with State-of-the-Art Methods:\nCan Tensor-Var be compared to state-of-the-art methods like FengWu-4DVar?\nSuperior or comparable performance would validate the method's significance and increase my score.\n\nAddressing these points would significantly enhance the paper's contributions, and I would like to rise my score.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XtCd0CkL4H", "reviewer": "Reviewer_wd4a", "summary": "This paper discusses Four-Dimensional Variational Data Assimilation (4D-Var), which is widely used in weather forecasting and dynamic system state estimation. Traditional methods struggle to properly model the nonlinear relationships between observational data and numerical models. The authors propose a novel approach that leverages Kernel Conditional Mean Embedding (CME) to linearize 4D-Var, transforming the nonlinear optimization problem into a convex optimization problem, thereby reducing the computational cost of optimization. Additionally, by incorporating past observational data, the proposed method enables robust estimation even in missing observations. Furthermore, the use of Deep Features (DFs) enhances scalability. Through experiments on two benchmark systems, as well as data assimilation experiments using global Numerical Weather Prediction (NWP) data and satellite observational data, the authors demonstrate that the proposed method achieves lower computational cost and reduced error compared to conventional 4D-Var methods.\n\n## update after rebuttal\n\nThank you for the detailed response. I find the integration with AI-based forecasting models particularly interesting. I also appreciate the addition of the comparison method. Overall, I believe the contribution is significant, so I will keep my evaluation as it is. I mistakenly posted my comment in the Official Comment section. My apologies.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate reviewer's constructive feedback and recognition of the strengths of our work. Below, we address each of the questions in detail.\n\n- **Compare with sota ML-based DA methods.** \n  \n  We thank reviewer for raising the questions regarding the comparison with recent ML-based methods.\n\n  - Most existing ML-based methods, such as GraphCast [Lam et al. (2023)], primarily focus on weather forecasting rather than data assimilation. Therefore, direct comparisons between our DA framework and these forecasting methods are not straightforward. \n  \n  - Nevertheless, our proposed DA framework can be seamlessly integrated with advanced AI-based weather forecasting models. By incorporating real-time observations, our approach consistently corrects the AI-based weather forecasting models, providing a more reliable starting point for subsequent forecasts. Furthermore, GNN-based approaches can be directly integrated into our proposed DA framework as kernel features, enabling the effective assimilation of spatially irregular data. This integration would further improve the model's ability to handle complex spatial information within the DA context.\n  \n  - We recognise the importance of comparing our method with state-of-the-art ML-based DA approaches, as also highlighted by Reviewer dYtv. To address this concern, our revised manuscript has added the new results from a state-of-the-art machine learning-based data assimilation method, FengWu 4D-Var [Xiao et al. (2023)]. The results and comparison details are provided in our response to Reviewer dYtv under the question *Compare with sota ML-based DA methods & Baseline selections* and are also presented in the accessible [anonymous Google Slides](https://docs.google.com/presentation/d/1SWXKhJL_4mAKVQlY4HAxQTl_79C94m-2TwgtkYIvuNw/edit?usp=sharing). We believe this addition improves the credibility of our study and provides a more comprehensive evaluation.\n\nWe hope this response adequately addresses the reviewer’s concerns. We sincerely appreciate the reviewer’s invested time and effort in carefully reviewing our work and providing valuable, constructive feedback.\n\n**References**\n  - Lam, Remi, et al. \"Learning skillful medium-range global weather forecasting.\" Science 382.6677 (2023): 1416-1421.\n  - Xiao, Yi, et al. \"Fengwu-4dvar: Coupling the data-driven weather forecasting model with 4d variational assimilation.\" arXiv preprint arXiv:2312.12455 (2023).", "claims_and_evidence": "The authors propose, for the first time, a method that utilizes Kernel Conditional Mean Embedding (CME) to address the high computational cost and inability to properly handle nonlinear relationships in 4D-Var. Furthermore, by leveraging DFs, they achieve a practically feasible computational cost. This idea is highly logical and well-founded. Additionally, the effectiveness of the proposed method is properly demonstrated through extensive experimental data.", "methods_and_evaluation_criteria": "The proposed method is compared against conventional 3D-Var, 4D-Var, and ML-based 4D-Var (Frerix et al., 2021). These are appropriate baseline methods for demonstrating the advancements of the proposed approach. Additionally, the datasets used for evaluation include chaotic nonlinear systems, making them suitable for assessing the advantages of the proposed method.", "theoretical_claims": "The specific theoretical framework, including the estimation of error covariance matrices using kernel covariance matrices, is detailed in the Supplementary Material. A broad review of the content confirms that the key points are well covered.", "experimental_designs_or_analyses": "As mentioned in Methods and Evaluation Criteria, appropriate comparison methods are used, and the datasets chosen for evaluation are also suitable.", "relation_to_broader_scientific_literature": "The proposed method has broad applicability in modeling large-scale spatiotemporal data. In addition to the weather and satellite observation data evaluated in this paper, it can also be applied to various domains, such as spatiotemporal estimation of air pollution levels, probabilistic forecasting of earthquakes, time series prediction of stock prices and exchange rates, and estimation of biosignals (e.g., EEG and ECG). This approach has the potential to drive significant breakthroughs across a wide range of scientific fields.", "essential_references_not_discussed": "To the best of my knowledge, prior studies related to 4D-Var, including the well-known GraphCast, have been appropriately cited.", "other_strengths_and_weaknesses": "As mentioned above, this paper has the potential to be widely applied to large-scale spatiotemporal data modeling and could lead to significant breakthroughs in this field.\n\nRegarding weaknesses, while recent machine learning-based methods, such as GraphCast and GNN-based approaches, are cited, there is little direct comparison with machine learning-based methods. This may be due to the high computational cost of directly learning nonlinear relationships, making comparisons difficult. However, providing at least some level of comparison or an explanation of the challenges in conducting such comparisons could enhance the credibility of the study.", "other_comments_or_suggestions": "I have no additional comments.", "questions_for_authors": "Could you explain the lack of comparison with recent machine learning-based methods, such as GraphCast and GNN-based approaches?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nB507puonc", "reviewer": "Reviewer_wFvT", "summary": "The paper introduces Tensor-Var, a novel framework for variational data assimilation, which integrates kernel conditional mean embedding (CME) with four-dimensional variational assimilation (4D-Var). Traditional 4D-Var methods are computationally expensive and struggle with nonlinear dynamics and imperfect state-observation mappings. Tensor-Var addresses these challenges by linearizing nonlinear dynamics in a learned feature space, making the optimization convex while ensuring theoretical guarantees for consistent assimilation results.\n\nTo enhance scalability, the framework incorporates deep features (DFs) learned via neural networks, allowing it to handle large-scale problems efficiently. Experiments on chaotic systems and global weather prediction demonstrate that Tensor-Var significantly outperforms conventional and deep learning (DL)-based 4D-Var methods in accuracy while achieving a 10- to 20-fold speedup.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank reviewer wFvT for these insightful questions. Below, we address each point in detail.\n\n- **Choice of kernel functions**\n  \n   Thank you for raising the question about choices of kernel functions in Tensor-Var.  We would like to clarify that we have included an ablation study in Section 4.4 Table 3  on the Lorenz96 systems with $n_s=40$ and $n_s=80$. To further address the reviewer's comments, we also tested other kernel features, including Matern-1.5 and Matern-2.5. The results are provided in the following Table: \n\n  **Table: Comparison of different features with NRMSE (%) as the metric for state estimation accuracy.**\n  | Features/Experiments  | Lorenz-96 $n_s=40$ | Lorenz-96 $n_s=80$|\n  |-----------------------|----------------------|----------------------|\n  | $d_s=20$            | 16.7±2.1              | 17.3±2.6              |\n  | $d_s=40$            | 14.1±1.3              | 16.8±2.2              |\n  | $d_s=60$            | 8.3±0.9               | 11.4±0.9              |\n  | $d_s=120$          | 9.7±0.7               | **9.0**±0.9              |\n  | Gaussian Kernel       | 8.4±0.5               | 14.3±1.4              |\n  | Matern-1.5            | 9.6±0.8               | 13.2±1.1              |\n  | Matern-2.5            | **7.2±0.2**           | 9.7±**0.7**           |\n\nIn a low dimensional system ($n_s=40$), Matern-2.5 achieves the lowest mean and standard deviation. In the higher-dimensional setting ($n_s=80$), the deep feature with $d_s=120$ outperforms the other approaches, albeit with a slightly higher standard deviation compared to Matern-2.5. In general, the results suggest that fixed kernel features may be more reliable at lower dimensions, but deep features become advantageous as dimensionality increases. \n\nFor high-dimensional weather forecasting problems, fixed kernel features tend to be poorly scaled and are inadequate for processing structured data. Thus, we propose deep features to address scalability challenges and extend Tensor-Var for broad applications. \n\n\n- **Computational Complexity and Scalability**\n  \n     We thank the reviewer for their insightful question regarding the scalability of Tensor-Var. The computational complexity of Tensor-Var primarily depends on the feature dimension rather than the system dimension, as DA optimisation is performed within the feature space. As shown in Table 1 (Section 4.1), the evaluation time exhibits only a slight increase with the increasing system dimension.\n\n     Moreover, the experiments presented in Section 4.2 and Section 4.3 differ in spatial resolution, with Section 4.3 using a resolution of $240*121$, approximately 14 times higher than that of Section 4.2. Nonetheless, the computation time does not increase substantially.\n\n     One potential bottleneck arises from solving the least square problems to compute the two CME operators in high-dimensional feature spaces during training. Incorporating more efficient solvers, such as the Krylov subspace method [Liesen et al. (2013)], could address this challenge in future work.\n\nWe hope these responses adequately address the reviewer’s concerns, and we look forward to any further comments you may have.\n\n\n**References**\n\n - Liesen, Jörg, and Zdenek Strakos. Krylov subspace methods: principles and analysis. Numerical Mathematics and Scie, 2013.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. How sensitive is the performance of Tensor-Var to the choice of the kernel function in the conditional mean embedding (CME)? Have you tested different kernels, and how do they impact generalization across different dynamical systems?\n2. While Tensor-Var achieves a 10- to 20-fold speedup, how does its computational complexity scale with increasing data dimensions and observation frequencies? Are there potential bottlenecks in large-scale real-world applications?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "NA"}}]

Author rebuttal:

Meta-review: This paper introduces Tensor-Var, a framework that uses kernel CME to linearize the dynamics in 4D-Var (a data assimilation approach common in meteorology) to yield a more tractable and efficient optimization problem. Reviewers were generally impressed with the methodological contributions, highlighting the potential for impact. While reviewers raised several concerns about the manuscript, almost all of these were addressed during the rebuttal period. These include updated experiments comparing with SOTA ML-based methods, stability tests, comparisons of different kernel functions, and comparisons of training times. I recommend acceptance of this work assuming all of these additions are included in the final revision, and _strongly recommend_ that the comparisons to other SOTA ML methods are included in the main body of the paper.
Decision: Accept (poster)