On Learning Linear Dynamical Systems in Context With Attention Layers

MCML Authors

Maria-Luiza Vladarean

Dr.

→ Group Suvrit Sra
 Resource Aware Machine Learning

Xuhui Zhang

→ Group Suvrit Sra
 Resource Aware Machine Learning

Suvrit Sra

Prof. Dr.

Core PI

Resource Aware Machine Learning

Abstract

This paper studies the expressive power of linear attention layers for in-context learning (ICL) of linear dynamical systems (LDS). We consider training on sequences of inexact observations produced by noise-corrupted LDSs, with all perturbations being Gaussian. Importantly, this non-i.i.d. data setting is a significant step towards modeling real-world scenarios. We provide the optimal weight construction for a single linear-attention layer and show its equivalence to one step of Gradient Descent relative to an autoregression objective of window size one. Guided by experiments, we uncover a connection to a generalization of the Preconditioned Conjugate Gradient method for larger window sizes. We back our findings with numerical evidence. These results add to the existing understanding of transformers' expressivity as in-context learners and offer plausible hypotheses for recent observations that place their performance on par with that of the Kalman Filter --- the optimal model-dependent learner for this setting.

inproceedings VZS+26

---

ICLR 2026

14th International Conference on Learning Representations. Rio de Janeiro, Brazil, Apr 23-27, 2026. To be published. Preprint available.

---

Authors

M.-L. Vladarean • X. Zhang • S. Sra

---

Links

URL

---

Research Area

 A2 | Mathematical Foundations