Two-Layer Linear Auto-Regressive Models Estimate Latent States

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.

• The paper demonstrates that two-layer linear AR models can recover latent state estimates, approximating those produced by the Kalman filter using only observable data.
• It establishes that the loss landscape has a strict saddle property, enabling gradient-based methods to efficiently reach global optima with finite-sample guarantees.
• Empirical results confirm that, with sufficient trajectory data, the hidden layer activations align closely with the true latent states, ensuring reliable state recovery.

Two-Layer Linear Auto-Regressive Models Estimate Latent States: Technical Summary

Problem Setting and Theoretical Motivation

This paper addresses the statistical and optimization-theoretic underpinnings of two-layer linear auto-regressive (AR) models in the context of learning latent states from time series data generated by partially observed linear dynamical systems (LDS). Given only input-output sequences (no access to ground-truth state or system parameters), the goal is to recover—up to similarity—optimal state estimates as would be produced by the Kalman filter.

The prevalent approach in system identification has traditionally relied on learning a shallow AR model and applying a post hoc factorization (e.g., Ho–Kalman, nuclear norm regularization) to extract latent representations. In contrast, this work demonstrates that two-layer AR models, trained end-to-end via ERM on observable sequences, learn internal representations in the hidden layer that align (modulo similarity) with Kalman filter state estimates—without explicit knowledge of model structure or parameters.

Model Architecture and Main Theoretical Contributions

The modeling framework considers a standard LDS: $$x_{t+1} = A x_t + B u_t + w_t,\quad y_t = C x_t + v_t$$ where $x_t$ is the latent state, $u_t$ is the input, $y_t$ is the observed output, and $w_t, v_t$ are Gaussian noise sources. Crucially, the two-layer AR model maps concatenated histories of past $L$ outputs and inputs (the covariate $z_t$) through a linear hidden layer to a linearly decoded prediction of $H$ future outputs.

The function class is $f(z) = G_2 G_1 z$ with hidden dimension $h$, trained to minimize squared prediction error over an input-output trajectory. Both the history length $x_t$0 and model dimension $x_t$1 are treated as hyperparameters. The authors offer the following key theoretical advances:

• Kalman Filter Approximation: The Kalman filter state estimate at time $x_t$2 can be approximated as a linear function of a sufficiently long history of past outputs and inputs, with the error decaying as $x_t$3 where $x_t$4 relates to the spectral radius of the closed-loop estimator.
• Benign Non-Convexity: The empirical risk minimization problem over $x_t$5, despite its non-convexity (due to factorization), enjoys a benign landscape: all stationary points are either strict saddles or global minima. This ensures gradient-based methods can find global optimizers under reasonable data/initialization regimes.
• Finite-Sample Guarantees: The work provides non-asymptotic, finite-sample statistical bounds (in trajectory length $x_t$6, problem dimensions, and noise) for prediction error, parameter estimation, and, critically, for error in latent state recovery (i.e., the distance between the hidden layer representation and true state, up to similarity).

The essential technical finding is that, under classical control-theoretic assumptions (stabilizability, observability, non-explosiveness), and given sufficient trajectory length and history $x_t$7, the hidden activations of the trained model encode the Kalman state estimate, up to linear transformation. This effect is robust across diverse system classes.

Optimization Landscape and Statistical Learning Results

The authors rigorously show that:

• The two-layer factorized loss landscape satisfies the “strict saddle” property. All non-global stationary points correspond to strict saddles; thus, local optimization with random initialization and suitable perturbations will not get stuck in suboptimal points. This property is proven through an adaptation of existing matrix factorization non-convexity analyses and direct application to the AR loss.
• The sample complexity for predicting outputs and recovering latent states is explicit: as $x_t$8 increases, prediction and estimation errors scale as $x_t$9 (with problem-dependent constants and logarithmic dependencies).
• The alignment and sample complexity of latent recovery depends on full-rank conditions on extended observability and controllability matrices, as well as sufficient history length $u_t$0 and prediction horizon $u_t$1.

Figure 2: Sample complexity plot by varying the history length $u_t$2, with fixed future window length $u_t$3.

Empirical Results: Recovery of Latent States

Empirical studies support the theory using both synthetic and standard control benchmark (ControlGym) systems. The key findings are:

• When trained with $u_t$4 (dimension of the latent state), the hidden activations of the model exhibit extremely high $u_t$5 with the Kalman filter estimates after appropriate linear transformation. This holds consistently across random linear systems and practical control benchmarks, as validated by quantitative $u_t$6 metrics.
• The success of latent recovery does not depend on system particulars, confirming the generality of the phenomenon.
• The observed decline in latent recovery error, as trajectory length $u_t$7 increases, empirically matches the theoretical sample complexity rates.

Figure 3: $u_t$8.

These results are robust to various settings of history $u_t$9 and prediction window $y_t$0, with minor degradation for suboptimally chosen hyperparameters. Theoretical trade-offs are confirmed empirically: larger $y_t$1 improves approximation at the cost of more parameters, while the choice of $y_t$2 impacts the rank and identifiability of the underlying observability map.

Implications and Future Directions

These results have several theoretical and practical implications:

• Theoretical Implications: The findings unify perspectives from system identification and deep learning: end-to-end representation learning in AR models naturally (and provably) recovers the best linear sufficient statistic (Kalman state) without any explicit supervision or additional regularization, provided sufficient data and model size.
• Statistical Guarantees for Deep Architectures: This work provides a template for deeper or nonlinear architectures, suggesting that similar latent recovery phenomena may occur in higher-capacity or recurrent models—although further work is required to extend these results beyond the linear AR setting.
• Practical Relevance: In reinforcement learning and control, where interpretable latent representations are crucial, these guarantees provide a foundation for using simple AR models in pipeline designs, as well as benchmarking state recovery against classical methods in the presence of only observed input-output data.

Looking forward, future research is indicated in several directions, including:

• Extending statistical and landscape guarantees to non-Gaussian or adversarial noise settings, where Kalman filtering is not optimal and error correlations complicate the analysis.
• Investigating the emergence of state-encoding representations in deep or nonlinear AR and RNN architectures, and their alignment with classical observers or sufficient statistics.
• Understanding meta-generalization and “in-context learning” in high-capacity AR models trained on data from multiple systems, seeking conditions under which state recovery generalizes across system classes.

Conclusion

This paper provides a rigorous statistical and optimization-theoretic analysis showing that two-layer linear auto-regressive models, trained on observed input-output sequences from partially observed linear dynamical systems, automatically learn hidden representations that closely align with optimal state estimates produced by Kalman filtering, up to similarity. This effect is supported by finite-sample guarantees and robust empirical validation, bridging a key gap between modern representation learning methodology and classical control-theoretic state estimation. The analysis sets the stage for systematic investigation of representation and state recovery in more general architectures and data regimes.

Reference: "Two-Layer Linear Auto-Regressive Models Estimate Latent States" (2606.12691)