Two-Layer Linear Auto-Regressive Models Estimate Latent States

AI summary

Authors: Yahya Sattar, Sunmook Choi, Leo Maynard-Zhang, Yassir Jedra, Maryam Fazel, Sarah Dean (Cornell; U Washington & Amazon; Imperial College London) arXiv: 2606.12691v2 [cs.LG], revised 08 Aug 2026; License CC BY 4.0

Overview

Research area: Machine learning theory at the intersection of deep learning, representation learning, and control/dynamical systems — specifically, statistical learning theory for partially observed linear dynamical systems and non-convex optimization for matrix factorization.

Technical level: Advanced. The paper uses state-space control theory (Kalman filtering, observability/controllability, Riccati equations), matrix factorization theory, and finite-sample statistical learning bounds.

One-sentence scope: The paper proves that a two-layer linear auto-regressive model trained end-to-end by empirical risk minimization on input-output data from a partially observed linear dynamical system learns hidden activations that coincide, up to a similarity transformation, with steady-state Kalman filter state estimates.

What This Paper Is About

In a partially observed linear dynamical system, a learner sees only control inputs and measured outputs, while the latent state and the true system matrices stay hidden. Classical system identification handles this by first fitting a linear model and then decomposing it (for example, via Ho-Kalman factorization) to extract the state — a two-step pipeline that differs from how deep models are trained in practice.

This paper asks whether the end-to-end paradigm can do the job directly: if you train a two-layer linear auto-regressive network with a squared next-output prediction loss, does the network's hidden layer implicitly represent the best-possible state estimate (the Kalman filter) even though the model has no explicit knowledge of the dynamics or state?

Key Contributions

1. A two-layer auto-regressive model for latent state estimation. The authors formulate a new architecture that maps a window of L past outputs and inputs to a horizon of H future outputs through a linear hidden layer of dimension h, and train it with squared loss via empirical risk minimization (Section 3).
2. A benign optimization landscape. Despite the loss being non-convex, the authors show all local minima are global minima and every saddle point is a strict saddle (Section 4.2, Proposition 2).
3. Finite-sample statistical guarantees. They bound in-sample prediction error (Theorem 2) and parameter estimation error (Theorem 3), with sample complexity and explicit error rates in terms of trajectory length T, dimensions, and failure probability δ.
4. Latent state recovery. They show the parameter error implies that the learned hidden representation matches the Kalman filter state estimate up to a similarity transform, with an error decaying at rate Õ(1/√T) in the amount of training data (Theorem 4, Section 4.4).

Main Findings

• The Kalman filter is well approximated by an auto-regressive model with bounded truncation error. Expanding the steady-state Kalman predictor gives x̂_t = 𝒞 z̄_t + Ā^L x̂_{t-L}, where 𝒞 is an extended controllability matrix. Proposition 1 bounds the truncation error as C_ρ² ρ^{2L} ||Σ[x̂_t]|| (n + log(1/δ)) with probability at least 1 − δ, which can be made arbitrarily small by choosing the history length L ≳ β log(T) for some β > 0.

• The non-convex landscape is benign. Under Assumptions 1 and 2, and with L ≳ β log(C_ρ T ||C|| √(C_z̄ C_x̂))/(1 − ρ) and a sufficiently long trajectory, Proposition 2 states that (i) any local minimum is a global minimum and (ii) any saddle point is a strict saddle, with probability at least 1 − δ. The proof builds on Theorem 3 of prior work on matrix factorization plus persistence-of-excitation properties of the training data.

• In-sample prediction error is near-optimal in trajectory length and dimensionality. For the global minimizer (Ĝ_1, Ĝ_2), Theorem 2 gives (1/T) Σ_t ||(Ĝ_2 Ĝ_1 − 𝒪𝒞) z̄_t||² ≲ (||Σ[ξ_1]|| H / T)(r(d_ȳ + d_z̄) log(TΛ) + log(T/δ)), where 𝒪 is the extended observability matrix, 𝒞 the extended controllability matrix, and ξ_t collects future-input and future-innovation terms. This implies a prediction error of order Õ(1/T), and bounds the number of samples needed: T ≳ L(d_z̄ + log(T/δ)).

• Parameter estimation error inherits the same rate, scaled by excitation. Theorem 3 bounds ||Ĝ_2 Ĝ_1 − 𝒪𝒞||_F² ≲ (||Σ[ξ_1]|| H / (λ_min(Σ[z̄_L]) T))(r(d_ȳ + d_z̄) log(TΛ) + log(T/δ)), requiring the covariates to persistently excite all modes of the Hankel matrix ℋ = 𝒪𝒞, i.e., λ_min(Σ_{t=1}^T z̄_t z̄_tᵀ) ≥ Õ(λ_min(Σ[z̄_L]) T).

• The hidden layer recovers the latent state. Theorem 4 states that if the extended observability matrix has full column rank, the extended controllability matrix has full row rank, n̂ = n, the factorization (𝒪, 𝒞) is balanced (𝒪ᵀ𝒪 = 𝒞𝒞ᵀ), and the robustness condition 2||Ĝ_2 Ĝ_1 − 𝒪𝒞|| ≤ σ_n(𝒪𝒞) =: σ_n holds, then there exists a similarity transform S such that ||x̂_t − S Ĝ_1 z̄_t||² ≲ (||Σ[ξ_1]|| H / (λ_min(Σ[z̄_L]) σ_n T)) ||z̄_t||² (r(d_ȳ + d_z̄) log(TΛ) + log(T/δ)) with probability at least 1 − δ. In short, the trained weights are close to (𝒞, 𝒪) up to a similarity transform, and the hidden activations are state estimates.

• A horizon trade-off. The rank conditions are implied by Assumption 2 whenever H ≥ n and L ≥ n. If H is chosen smaller than n, 𝒪 loses full column rank and latent state recovery may be impossible; if H grows beyond n, the bound loosens because ||Σ[ξ_1]|| grows polynomially in H. This is an explicit trade-off between prediction horizon length and latent recovery accuracy.

• Similarity transformations are unavoidable. Because any invertible S yields the same input-output statistics with state Sx, the guarantee is necessarily stated up to a similarity transform rather than in absolute coordinates.

• Numerical simulations support the theory. The paper reports simulations demonstrating that auto-regressive models automatically represent latent state estimates (Section 5). The truncated content provided does not include the specific simulation settings, dataset sizes, hyperparameters, or numerical error values.

Methodology in Plain English

The authors build a single trajectory of observed input-output pairs from an unknown linear dynamical system and slice it into training examples. Each example feeds the most recent L outputs and inputs (the covariate z̄_t, dimension d_z̄ = (m + p)L) into a two-layer linear network and predicts the next H outputs (dimension d_ȳ = mH). The hidden layer has dimension h, and the product of the two weight matrices G_2 G_1 acts as a linear map.

Training minimizes an ordinary squared error over the training examples. The network weights are norm-bounded, which the authors note is equivalent to a regularized objective with penalty corresponding to a constant c_0 — the same effect deep learning practitioners get from weight decay. The model is fit by searching over the hidden dimension h and running gradient descent on G_1 and G_2 for each fixed h.

To prove the results, the authors take three steps. First, they algebraically expand the Kalman filter's predictor form to show the true state estimate is essentially a linear function of the recent input-output window, with a ρ^{2L} truncation error where ρ bounds the spectral radius of the filter's closed loop. Second, they translate the learning problem into the language of matrix factorization with a Hankel matrix ℋ = 𝒪𝒞 and import existing landscape results for that problem, combined with excitation properties of the data. Third, they bound the generalization error of the empirical minimizer using a self-normalized Gaussian complexity argument tailored to correlated time-series data, then convert the parameter error into a latent state error using singular value conditions and a balanced factorization.

Why This Matters

Impact on research. The paper unites two research traditions that usually operate apart: classical system identification, which estimates parameters and then factorizes, and end-to-end deep learning, which trains a network and hopes latent structure emerges. It gives a theoretical reason to expect that hidden activations in trained auto-regressive models carry interpretable state information, and it does so with a completely non-convex learning procedure — rather than special-purpose regularization such as nuclear norm penalties or Ho-Kalman factorization. It also contributes a rare finite-sample latent-recovery guarantee for an explicitly non-convex objective.

Real-world applications (as suggested by the paper's framing):

• Large language models and sequence models: understanding whether models that predict the next token internally track the underlying state of a process, such as a chessboard configuration given a sequence of moves.
• World models for robotics: models trained on robot video and action streams, where knowing whether the model tracks true physical state matters for downstream planning.
• Sequential decision-making and control: the optimal policy for quadratic-cost linear control decomposes into a Kalman filter composed with linear state feedback, so a filter learned directly from data could feed control design.
• Time-series forecasting and recommendation-style factorization tasks: settings where the long-standing goal is to extract compact latent representations from observed sequences, in the spirit of word2vec for language or matrix factorization for movie recommendation.

Industry relevance. The result speaks to a common practical question: when a

Authors’ abstract

Auto-regressive models have emerged as powerful tools for sequential data, from language to video. Understanding how and why these models learn latent representations remains an open theoretical question. In this work, we demonstrate that when trained by empirical risk minimization on data from partially observed linear dynamical systems, two-layer linear auto-regressive models naturally learn to approximate Kalman filtering. In particular, we show that the learned hidden representation coincides, up to a similarity transformation, with the state estimates produced by the optimal (Kalman) filter, even though the model has no explicit knowledge of the underlying dynamics or state. The result follows from three main insights. First, we establish that the Kalman filter is well approximated by an auto-regressive model with bounded truncation error. Second, we show that despite non-convexity, the two-layer optimization landscape is benign, i.e., all stationary points are either strict saddles or global minima. Finally, as our main contributions, we provide finite-sample guarantees on prediction error, parameter estimation error, and latent state recovery. Numerical simulations support the theoretical results and demonstrate that the latent representations of auto-regressive models recover state estimates.