Recovering Latent Structure from Sequences with
Auto-regressive Models
Sarah Dean, Cornell University
IFDS Research Expo, August 2026
"Next token prediction"
Auto-regressive architectures trained with next element prediction loss have had great success in many domains
$$ (o_0, o_1, o_2, o_3, o_4, o_5 ..., o_t) \rightarrow o_{t+1} $$
Eyjolfsdottir & Branson, Agent-Centric Animal Pose Forecasting
Haffner et al., Mastering Diverse Domains through World Models
Radford et al., Improving Language Understanding by Generative Pre-Training
natural language
robotics
biology
Many researchers are interested in intermediate network activations, whether for use in downstream tasks or for scientific understanding
Inside the black box...
Levy et al., Simulation Distillation: Pretraining World Models
in Simulation for Rapid Real-World Adaptation
Eyjolfsdottir & Branson, Agent-Centric Animal Pose Forecasting
robotics
biology
1. In Theory
Outline: recovering latent structure
2. In the Wild
Partially observed stochastic dynamical system
$$ x_{t+1} = f(x_t,u_t,w_t),\quad y_t = g(x_t, v_t) $$
- Inputs \(u_t\) and outputs \(y_t\) are observed
- States \(x_t\), process noise \(w_t\), measurement
noise \(v_t\) are unobserved
Setting: stochastic dynamics
ex: UAV
\(x\)
\(f\)
\(g\)
\(w\)
\(v\)
Irreducible uncertainty in state is described by the posterior distribution
$$P(x_t | y_0, u_0, y_1, u_1,...,y_t) $$
\(u\)
\(y\)
$$x_0\sim \mathcal D_0, w_t\sim \mathcal D_w, v_t\sim \mathcal D_v$$
Computing this posterior is called filtering $$P(x_t|y_0,...,y_t) \propto P(x_{t-1}|y_0,...,y_{t-1}) P(x_t|x_{t-1} , u_t, y_{t})$$ but for general distributions it can be hard
Setting with no models, only observations
$$ (y_0, u_0, y_1, u_1, \dots, y_T, u_T) $$
- Idea: train a deep auto-regressive model on this sequence, predicting \(y_t\) from \(y_{0:t-1},u_{0:t-1}\)
- Question: will model activations recover something about the latent state?
Unknown dynamics/noise models
ex: UAV
\(x\)
\(?\)
\(?\)
\(w\)
\(v\)
For any similarity transform, \(\tilde x = S(x)\) is an equivalent state space representation with identical input/output behavior
$$ \tilde f = S \circ f \circ S^{-1}, \qquad \tilde g = g \circ S^{-1} $$
\(u\)
\(y\)
future
\(h_t \approx x_t??\)
past
\(\exists S ~\text{s.t.}~ S(h_t)\approx \arg\max_{x_t} P(x_t|\text{past})??\)
posterior?
\(h_t \approx \arg\max_{x_t} P(x_t|\text{past})??\)
posterior, up to similarity
transform?
Linear-Gaussian Systems
ex: UAV
Kinematic equations:
\(f_t = m a_t\)
\(v_{t+1} = v_t + a_t\)
\(p_{t+1} = p_t + v_t\)
\( \begin{bmatrix}p_{t+1}\\ v_{t+1} \end{bmatrix} = \begin{bmatrix} 1 & 1\\ 0 & 1\end{bmatrix} \begin{bmatrix}p_t\\ v_t\end{bmatrix} + \begin{bmatrix}0\\ 1/m\end{bmatrix}f_t\)
\((u_t+w_t)\)
\(y_t = \begin{bmatrix} 1 & 0 \end{bmatrix} x_t + v_t\)
\(=: x_t\)
Consider linear dynamics and measurement functions: $$ x_{t+1} = A x_t + B u_t + w_t, \qquad y_t = C x_t + v_t $$
Guassian initial state and noise: $$x_0\sim \mathcal N(0,\Sigma_0), \quad w_t\sim\mathcal N(0,\Sigma_w),\quad v_t\sim \mathcal N(0,\Sigma_v)$$
\(x\)
\(A, B\)
\(C\)
\(w\)
\(v\)
\(u\)
\(y\)
Kalman Filter
For linear-Gaussian systems, efficient posterior computation via Kalman filter:
- Predict: \(\hat x_{t|t-1} = A \hat x_{t-1} + B u_{t-1}\) and \(\Sigma_{t|t-1} = A \Sigma_{t-1} A^\top + \Sigma_w\)
- Update: for \( L_{t} = \Sigma_{t|t-1} C^\top (C \Sigma_{t|t-1} C^\top + R)^{-1} \) $$\hat x_{t} = \hat x_{t|t-1} + L_{t}\big(y_{t} - C\hat x_{t|t-1}\big),\qquad \Sigma_{t} = (I - L_{t} C)\, \Sigma_{t|t-1}$$
- We have that \(P(x_t \mid y_{0:t}, u_{0:t})\) = \(\mathcal N(\hat x_{t}, \Sigma_{t})\)
The steady-state Kalman filter (SS-KF) has \(\Sigma_t=\Sigma\) and \(L_t=L\) for $$A\Sigma A^\top - \Sigma + \Sigma_w - A\Sigma C^\top (C\Sigma C^\top + \Sigma_v)^{-1} C\Sigma A^\top = 0.$$
SS-KF estimate: \(\hat x_{t+1|t} = (A-ALC) \hat x_{t|t-1} + B u_{t} + AL y_{t} \)
\(\Sigma_t\) does not depend on inputs or outputs
\(\hat x_t\) is the mean/MLE
Auto-regressive model
Idea: train an auto-regressive model on \((y_0, u_0, y_1, u_1, \dots, y_T, u_T)\)
\(L\) past inputs/outputs
\(H\) future outputs
two layer linear network
Question: will model activations recover the Kalman filter state estimate, up to similarity transform?
squared loss \( \displaystyle \min \frac{1}{2T} \sum_{t=1}^T \Vert{}y_{t:t+H-1} - G_2 G_1 {z}_t\Vert{}_{2}^2\)
\(z_t:=\)
Theorem (Informal). Let \(d_h=d_x\) and \((\hat{G}_1, \hat{G}_2)\) be the global minimizer, and $$H \geq d_x,\qquad L \gtrsim \log(T) / (1 - \rho),\qquad T \gtrsim L \left(d_{{z}} + \log(T/\delta)\right)\:.$$ Then for any \(\{u_k, y_k\}_{k=0}^{t-1}\), with probability at least \(1 - \delta\),
$$\min_{S\in\text{GL}_{d_x}(\mathbb R)}\Vert{}Sh_t-\hat{x}_{t|t-1}\Vert{}_{2}^2 \lesssim \frac{H}{T} \left(d_x (H d_{{y}} + d_{{z}}) + \log \left(\frac{T}{\delta}\right)\right) \Vert{}{z}_t\Vert{}_{2}^2.$$
Main result
\(\begin{bmatrix} y_{t-1} \\ \vdots \\ y_{t-L} \\ u_{t-1} \\ \vdots \\ u_{t-L} \end{bmatrix}\)
SS-KF on \(\{u_k, y_k\}_{k=0}^{t-1}\)
Assumptions: Linear-Gaussian system with
- Non-explosive, i.e. \(\rho(A) = \max |\lambda_i(A)| \leq 1\)
- Observable, i.e. \(\begin{bmatrix}C^\top & (CA)^\top & \dots &(CA^{d_x-1})^\top\end{bmatrix}^\top\) is full rank
- Persistent excitation, \(\Sigma_w, \Sigma_v\succ 0\) and training \(\{u_t\}_{t=0}^{T-1} \sim \mathcal N(0,I)\)
\(\hat{G}_1 z_t\)
model dimension
Proposition (Informal). The optimization landscape is benign.
Yahya Sattar
Numerical Experiments
Proof outline
\(\|S h_t - \hat x_t \|^2 = \|S \hat G_1 z_t - G^\star_1 z_t + G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
- \( \leq 2\|S \hat G_1 - G_1^\star \|^2\|z_t\|^2 +2\|G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
- \(\leq \frac{10\|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\ + 2\|G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
Lemma: Procrustes distance bound holds if \(\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\) small
- Recall that for the steady state Kalman filter, $$\hat x_{t+1|t} = \underbrace{(A-ALC)}_{\bar A} \hat x_{t|t-1} + B u_{t} + \underbrace{AL}_{F} y_{t} ,\qquad y_t = C \hat x_{t|t-1} + C(x_t-\hat x_{t|t-1}) + v_t$$
- As a consequence, $$\hat{x}_{t|t-1} = \underbrace{\begin{bmatrix}F & \dots & \bar{A}^{L-1}F & B & \dots & \bar{A}^{L-1}B \end{bmatrix}}_{G_1^\star} {z}_t + \bar{A}^L \hat{x}_{t-L},\qquad y_{t:t+H-1} =\underbrace{ \begin{bmatrix} C \\ CA \\ \vdots \\ CA^{H-1} \end{bmatrix}}_{G_2^\star }\hat{x}_{t|t-1} + \xi_t $$
Proof outline
$$\hat{x}_{t|t-1} = {G_1^\star} {z}_t + \bar{A}^L \hat{x}_{t-L},\qquad y_{t:t+H-1} ={G_2^\star }\hat{x}_{t|t-1} + \xi_t $$
\(\|S h_t - \hat x_t \|^2 = \|S \hat G_1 z_t - G^\star_1 z_t + G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
- \( \leq 2\|S \hat G_1 - G_1^\star \|^2\|z_t\|^2 +2\| \bar{A}^L \hat{x}_{t-L} \|^2\)
- \(\leq \frac{10\|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\ + 2\| \bar{A}^L \hat{x}_{t-L}\|^2\)
Lemma: Procrustes distance bound holds if \(\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\) small - \(\leq\frac{ 10 \|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\frac{1}{\lambda_{\min}(\sum_{t=1}^T z_tz_t^\top)} \sum_{t=1}^T \| (\hat G_2 \hat G_1 - G_2^\star G_1^\star)z_t\|^2 + 2\| \bar{A}^L \hat{x}_{t-L} \|^2\)
properties of optimal KF
truncation error, decays geometrically in \(L\)
Theorem: in-sample prediction error $$\lesssim H\left(d_x (H d_{{y}} + d_{{z}}) + \log \left(\frac{T}{\delta}\right)\right)$$
Theorem: persistent excitation $$\sum_{t=1}^T z_t z_t^\top \succsim T I$$
Proof outline
$$\hat{x}_{t|t-1} = {G_1^\star} {z}_t + \bar{A}^L \hat{x}_{t-L},\qquad y_{t:t+H-1} ={G_2^\star }\hat{x}_{t|t-1} + \xi_t $$
\(\|S h_t - \hat x_t \|^2 = \|S \hat G_1 z_t - G^\star_1 z_t + G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
- \( \leq 2\|S \hat G_1 - G_1^\star \|^2\|z_t\|^2 +2\| \bar{A}^L \hat{x}_{t-L} \|^2\)
- \(\leq \frac{10\|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\ + 2\| \bar{A}^L \hat{x}_{t-L}\|^2\)
- \(\leq\frac{ 10 \|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\frac{1}{\lambda_{\min}(\sum_{t=1}^T z_tz_t^\top)} \sum_{t=1}^T \| (\hat G_2 \hat G_1 - G_2^\star G_1^\star)z_t\|^2 + 2\| \bar{A}^L \hat{x}_{t-L} \|^2\)
- \(\lesssim \frac{H}{T} \left(d_x (H d_{{y}} + d_{{z}}) + \log \left(\frac{T}{\delta}\right)\right) \Vert{}{z}_t\Vert{}_{2}^2 \) as long as $$L \gtrsim \log(T) / (1 - \rho),\qquad T \gtrsim L \left(d_{{z}} + \log(T/\delta)\right)\:.$$
- Opportunity for new provable learning and control algorithms based on (deep) end-to-end training
- Hidden state (posterior) recovery is limited by data
- cannot recover latent quantities absent from data, whether due to lack of observability or excitation
- Open question: expand theory beyond linear-Gaussian?
- Key challenge: tractability of optimal filter (nonlinear)
- Idea: study best filter within function class, e.g. class of linear filters
- Open question: expand theory to multi-task setting (simultaneous identification and estimation)
Implications and open questions
1. In Theory
Outline: recovering latent structure
2. In the Wild
- Setting: generate sequence with arbitrary HMM \((A,B)\). Provide LLM with sequence only $$(o_1,o_2,...,o_t)$$
- Observation: as LLMs observe longer sequences, their predictions converge towards that of the true posterior
Synthetic data: HMMs
Yijia Dai
Dai et al., Pre-trained Large Language Models Learn Hidden Markov Models In-context
Tokenization matters
Entropy matters, sometimes
Similar structure to stochastic dynamics:
- Latent: \(A,B\) (fixed) and \(h_t\) (time-varying, according to \(A\))
- Observed: \(o_t\) (time-varying, according to \(h_t\), \(B\))
Hypothesis
Hypothesis: the LLM forward pass approximates a learning (filtering/identification) algorithm
- What algorithm? Can we find evidence of latent quantities in the activations?
Related work: in-context learning (ICL) of Markov chains by induction heads
Edelman, et al., The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains
$$ (o_0, o_1, o_2, o_3, o_4, o_5 ..., o_t) \rightarrow o_{t+1} $$
- Baum-Welch
- Alternates between fitting \(\hat A,\hat B\)
and estimate most likely \(\hat h_{0:t}\) - Predict via posterior computation
- Alternates between fitting \(\hat A,\hat B\)
- Soft n-gram
- \(\hat W = \arg\min \sum_{k=1}^t \ell(W\phi(o_{k-1:k-n}), o_k)\)
- Predict with \( \hat W\phi(o_{t:t-n+1})\)
- Spectral learning
- Posterior in terms of observable operators $$1^\top AB[o_t]AB[o_{t-1}]\dots A B[o_0] \mu $$
- Estimate operators in low-dim subspace with method of moments
Possible ICL Algorithms
Results
1. Prediction performance: Baum Welch behaves qualitatively different, Soft n-gram and Spectral are plausible
2. Toy setting: 2-layer transformer
trained on single HMM setting appears to encode \(\phi(o_{t-1:t-n})\) in activation and \(W_\star\) in weights
3. Mechanistic Interpretability:
probing and patching models for algorithm "belief states" supports soft n-gram algorithm
Best guess: understand LLM forward pass as performing gradient descent on \(\sum_{k=1}^t \ell_{CE}(W\phi(o_{k-1:k-n}), o_k)\)
Discussion
Domain-specific GPTs for inference and discovery
from sequential data?
Pretrained LLMs as black-box sequence predictors?
- Two-Layer Linear Auto-Regressive Models Estimate Latent States at ICML26 (arxiv:2606.12691) with Yahya Sattar, Sunmook Choi, Leo Maynard-Zhang, Yassir Jedra, Maryam Fazel
- Pre-trained Large Language Models Learn Hidden Markov Models In-context at NeuRIPS25 (arxiv:2506.07298) with Yijia Dai, Zhaolin Gao, Yahya Sattar, Jennifer J. Sun
- Extracting Algorithms in Pre-trained LLMs: A Case on Hidden Markov Models (arxiv:2607.22646) with Yijia Dai, Zhaolin Gao, Yahya Sattar, Jennifer J. Sun
Thanks! Questions?
Sunmook Choi
Yahya Sattar
Yassir Jedra
Maryam Fazel
Leo Maynard-Zhang
Yijia Dai
Zhaolin Gao
Jennifer Sun