On learning linear dynamical systems in context with attention layers
Authors: Maria-Luiza Vladarean
Xuhui Zhang
Suvrit Sra
Venue: ICLR 2026 Poster (2026)
Keywords: in-context learning; linear attention; linear dynamical systems; kalman filter; time series
Primary area: foundation or frontier models, including LLMs
TL;DR: This paper studies how linear attention layers in-context learn linear dynamical systems and shows the optimal weight construction implements one step of Gradient Descent relative to an autoregression objective of window size one.

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.

ON LEARNING LINEAR DYNAMICAL SYSTEMS IN CONTEXT WITH ATTENTION LAYERS
Maria-Luiza Vladarean" maria... veria.vladare Aaaaareeantta tum.de
Xuhui Zhang xuhui zhang@tum.de
Suvrit Sra 8 w.s.ratuuum.. de
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 se- quences of mexacc observations produced by noise-corrupted LDSs, with all per- turbations being Gaussian. Importantly, this non-i.dd. data setting is a significant step towards modeling real-world scenarios. We provide the optimal weight con- struction for a single linear-attention layer and show its equivalence to one step autoregression objective of window size Guided by experiments, we uncover a connection to a generalization of the Pre- conditioned 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.
of Gradient Descent relative
INTRODUCTION
This paper contributes towards understanding transformers' expressive power when learning from non-i.i.d data. In particular, consider sequences produced by a time-invariant linear dynamical system (LDS) doubly-corruteed by Gaussian noise
Ø1+1 Ax. + wts
vt c xe + Ut,
where we Me NO and 44 NO with mutually independent wi and Us Our starting point is the well-known ability of transformer-based Large Language Models to perform in-context learning (ICL.) (Brown et al., 2020).
ICL boils down to accurately answering query based on a set of examples given as a textual prefix ("in context") (Brown et al., 2020). This behaviour is desirable, as it dampens the requirement for expensive data collection and fine-tuning stages (Liu et al., 2023). Present research on ICL spans the spectrum of directions between enhancing it through specialized training and prompt engineering, and building a mechanistic understanding of it - see the comprehensive review of Dong et al. (2022). Our work is aligned with the latter direction.
Broadly, there are two theoretical perspectives on ICL mechanics: a Bayesian view in which trans- formers recover latent concepts from prompts, thereby performing implicit Bayesian inference (Wang et al., 2023; Jiang, 2023; Wies et al., 2023; Xie et al., 2021, and subsequent works in this direction), and an algorithmic view in which they implement iterative optimization methods in context (Von Os- wald et al., 2023a; Giannou et al., 2023; Akyürek et al., 2022: Garg et al., 2022; Ahn et al., 2023; Mahankali et al.. 2023; Sander Peyré, 2024: Von Oswald et al., 2023b; Sander et al., 2024). While the latter view does not account for the "emergent" aspect of "in-the-win -hild" ICL (Shen et al., 2023). it provides conerete expressions for transformers' modelling power and identifies the minimal func- tional unit wielding it - a single, causally-masked linear attention layer, without positional encoding. Works adopting the algorithmic view investigate transformers' ability to perform linear regression
"qqual contribution.
School of Computation, Information and Technology, Technical University of Munich

---

in the forward pass, confirming this through both empirical evidence and formal analysis. The i.i.d. setting is well understood, with results showing that for specific tokenizations, data distributions, and architectural choices, the optimal transformer weights implement gradient- based optimization relative to a context-dependent least-squares loss (Von Oswald et al., 2023a; Mahankali et al., 2023; Ahn et al., 2023; Von Oswald al., 2023b; Sander et al., 2024). In contrast the i.i.d. case, our theoretical grasp of the more realistic non-i.i.d setting is weak. The main hurdle in analyzing this scenario is handling tokens' statistical dependence on their preceding context. Our work takes the first steps towards unraveling this difficulty.
Specifically, we study the ability of a single linear attention layer to predict observation yr based on context of past observations f generated by LDS (1). This setting is well-motivated: firstly, temporal scaffold similar to that of iaguage-induueed tokens, which stands in stark contrast to the i.i.d. setup predominantly addressed by prior works (with few exceptions discussed shortly). Notably, dynamical systems have already been proposed conveniently flexible models for grammatical sentence formation (Elman, 1995; Tabor et al., 1996; Beim Graben et al., 2004; Belanger & Kakade, 2015). thus making setting (1) particularly relevant. Secondly, considering LDS-produced data places the algorithmic and Bayesian views on closer footing, since LDSs are subclass of the Hidden Markov Models (Minka, 1999) used in the Bayesian formulation of ICL. (Xie et al., 2021). Finally, empirical observations highlight the strong performance of transformers relative to the Kalman Filter (KF) (Kalman, 1960) for predicting yr. even regimes where KF is provably optimal (Du et al., 2023). To our knowledge, the underlying mechanism is yet to be understood.
sequence {y=) is built
Our goal is to characterize the optimal single linear self-attention layer trained to predict yr based on the context {y: ) We begin by defining context-dependent loss for time-series data via the improper learning approach to system identification, whereby processes of type (1) are well approximated by autoregressive models. We then identify and interpret the structure of optimally trained linear attention layers algorithmic steps on the corresponding autoregressive loss. Our contributions are the following.
C1. In Theorem 4.1, we prove that for order-one autoregressive approximation of (1), the optimal linear attention layer implements a step of Gradient Descent on the associated least-squares loss. To our knowledge, this is the first optimality result for LDS data.
C2. In Lemma 4. 1, we identify a salient banded pattern of the matrices involved in the stationarity condition for generic order-s approx of (1). We further define class of parameters that satisfy this structural constraint and empirically observe that minimizers obey it, thereby narrowing the search for the provably optimal lincar attention layer when > 2.
C3. In Section 5, we provide numerical experiments verifying our theory for order-one au- toregressive approximations. Furthermore, we connect the tiling pattern of empirically determined minimizers for order-s approximations, 8 2, with a generalization of the Preconditioned Conjugate Gradient method, thus further highlighting the view of ICL as on-the-fly optimization. To our knowledge, this is the first interpretation of the in-context algorithm for general order-s autoregression.
C4. Conceptually, we provide theoretical grounding for interpreting ICL as implicit optimization in the LDS setting, bridging system identification theory and empirical observations of transformer-Ko parity.
RELATED LITERATURE
We review the studies treating ICL as in-context optimization, together with works on filtering and system identification. Further comparisons discussed in Section 4.1.
ICL for linear regression with i.i.d data. This line of work studies whether transformers trained on few-shot learning objective can perform linear regression in-context, and how. Garg et al. (2022); Akyürek et al. (2022); Von Oswald et al. (2023a) provide empirical results in the affirmative, along with possible architecture constructions implementing Gradient Descent (GD) steps relative to a context- context-induced least squares loss. Through this lens, ICL reduces to on-the-fly ffy optimization executed in the transformer's forward pass. Mahankali al. (2023); Zhang et al. (2024); Ahn et al.
2

---

(2023) complement these findings by proving that the one-layer linear self- implementing (preconditioned) GD step is global minimizer of the pretraining loss when covariates are Gaussian and i..dd Finally, Zhang et al. (2024) complete the picture by proving that Gradient Flow converges to these global minimizers. Our results extend this line of work to the non-i.i.d. setting.
ICL and system identification. Different from the prior section, the following papers use the standard causal pretraining objective (minimizing prediction error over all sequence positions) and, unless stated otherwise, the results concern a single layer of linear self-attention. Von Oswald et al. (2023b) give a construction implementing a GD step on L(W) = E:11 Wy yi++ |I in parallel for all positions f, for a specific token augmentation. Sander et al. (2024), further characterize the global minimizers of this causal pretraining loss relative to the noiseless data y = Ay:. with A sampled from the set of commuting orthogonal matrices. Notably, under the same token augmentation, this characterization coincides with the construction of Von Oswald et al. (2023b). Sander et al. (2024) also describe minimizers for architectures using positional encoding, albeit under the restriction of diagonal weight matrices. Zheng et al. (2024) complement these results by showing that, for a diagonal weight initialization and a controlled distribution of yoo Gradient Flow (GF) recovers the aforementioned GD- GD-implementing optimum. Finally, Sander & Peyré (2024) extend these results to arbitrary orthogonal As via an infinite-depth attention-only transformer that correctly predicts yr in the limit T - co. This result also extends to the standard softmax activation.
Moving away from the noiseless settings above, Cole et al. (2025) establish approximation theoretic results for deep attention-only transformers predicting the sequence yt++1 = Ay: + w. with we ~ (00 "w= I) and A E sd 4 They establish the existence of log( (T')-dephh transformer that achieves uniform-over- A error of order log(T") for predicting E/xr+1 æt. A], and derive lower bound on the prediction accuracy attainable by a single linear attention layer. Regarding transformers' capacity, Ziemann et al. (2024) establish that achieving a uniform-in-time error bound for next-step prediction requires a number of parameters at least quadratic in the algebraic multiplicities of A A's unstable eigenvalues, as well as a context length least logarithmic in T.
In summary, these works either study transformers' ICL performance relative to simplified LDSs or do not address the question of weight optimality. In contrast, we study fully-fledged systems (1) with the aim of characterizing the pretraining loss minimizers in the few-shot training setting.
Transformers and linear filtering. The classical model-based prediction tool for systems of type (1) is the Kalman Filter (KF) (Kalman, 1960). Using knowledge of system parameters, the KF gives the minimum expected squared error estimates à;; of the hidden states æ;; as linear combinations of the past observations yi- Transformers as potential implementers of KF were studied by Goel & Bartlett (2024), who prove that a softmax causal attention layer is an arbitrarily good approximator for a given type (1) system. Akram & Vikalo (2024) further construct a transformer emulating the KF. Finally, Du et al. (2023) provide empirical evidence that GPT-2 architecture (Radford et al., 2019) competes in accuracy with the KF for predicting the next observation in contetex,, though the mechanism remains unstudied. We partially fill this gap with our present work.
3 PROBLEM FORMULATION & ASSUMPTIONS
Notation. Vectors and matrices are denoted by bold, lowercase and uppercase letters, respectively, with lowercase letters reserved for scalars. Symbols 1& and 0₂2 are the all-ones and all-zeros vectors of dimension d, and 1dxm and 0axmm are the analogous matrices. Subscripts of the form Aij. identify the submatrix formed by rows from i to j and columns from p to q. Unless stated otherwise, ||- denotes the Euclidean norm of vectors and the spectral norm of matrices. We denote by Tr (-) the trace of a matrix, { >> the inner product, by || || its Frobenius norm, and by p(-) its spectral radius. We use e; for the th canonical basis vector and I for the identity matrix. The notation Sd +++) semidffiit (-definite) matrices in Rdxd respectively. We use gd-1 1 to denote the unit sphere in Rdd. We use @ to denote the Hadamard product. Finally, we use [n] when referencing the set of integers 21,2,... .... We write w.p. to mean "with probability".
defines the cone of symmetric
3

---

The big picture:
filtering, system identification, and linear regression. computes the optimal estimates σ; of x; through the system of recursions
The KF (Kalman, 1960)
Predict:
B(+1|t = Ad, P(IIIc AP,AT + En
Gain:
kt41 Po+1jec (c PetijiC+ o.)
Update: ω(+1 E/+1/t + k4t1 (16+1 c cT
P+1 = (Id k4410' Ptt2le
where do and the eπror covariance estimate Po are inputs. Under the Gaussian errors assumption, the state prediction satisfies 2; E6 [x, yts y/ and, consequently, the forward observation prediction follows #++1 = c Ad, = E 3t++1 Ut. yn]- The fast, constant-time KF predictions, however, require knowing the LDSS parameters
condition generally not satisfied in practice.
Consequently, "proper learning" approaches seek to reconstruct the underlying model by first estimat- ing A, C, Ew,, o, through costly parameter identification techniques, and then producing forward observation predictions using the KF (Hamilton, 1995). In contrast, "improper learning" methods eschew structural constraints and solely seek to achieve low error relative the underlying data distribution and the learning objective (Kozdoba et al., 2019, and references therein). For LDSs, this boils down expressing the next observation as a linear function of the recent past. Not only does the latter approach avoid parameter estimation, but it also benefits from convex formulations, thus being amenable to classical optimization techniques. Most importantly, for certain LDS classes, improper learning methods can closely track [E[y++1 ...... y1 yt)] as we describe next.
The data-generating process (1) can be rewritten via the KF quantities in expression (2). In particular, Tsiamis & Pappas (2019) express future observations as a function of their past s predecessors,
|o+11....
.......---1 1-
c [[A CT((ATTAA
kc ) -
k.... k] li..... .U....9T---. 9T-8-11
+ c
TT(A (A kc ) ST + [S++1,...T-i).
where D. = [yt. 3t.1..... 11+s-1] l is the steady-state gain, and e, E are i.idd, zero-mean Gaussian errors. Under KF convergence conditions, quantity p(A - kc ) makes the second term vanish exponentially in s and thus renders negligible. We are now in the familiar setting of noisy linear regression, albeit with non-i.i.d. data. The resulting order-s autoregressive process (AR(s)) is associated with the optimization objective
T-
1
min (w) = 2(T - 1) WER-
Σ t=1
(Jets w T )
(4)
This simplification is the crux of improper learning approaches to system identification (Kozdoba et al., 2019) and becomes relevant in conjunction with the view of transformers optimizers of context-induced least-squares objective. Should this view withstand scrutiny in the non-i.idd setting, it would suggest that transformers can learn LDS-based time series in context to arbitrary accuracy as a function of s. This our incentive for characterizing the few-shot pretraining loss minimizers.
To ensure the validity of the autoregressive process approximation, make the following assumption.
Assumption 3.1 (System assumptions). LDS (1) has strictly positive definite noise covariances En and a, > O. The transition matrix A E Rdxdd is marginally stable, with A) S 1, and the pair (A, c) observable, meaning that
0 [c, [c, (A4-1)T
(5)
has column rank of d.
Assumption 3.1 is standard in the literature, and ensures KF convergence (Harrison, 1997) along with the exponential vanishing of the bias term in (3). Furthermore, it ensures the closeness of forward observation predictions given by the KF with those produced by a linear autoregressive predictor determined by expression (4) (Kozdoba al., 2019).
4

---

Transformer architecture. Transformers (Vaswani et al., 2017) are neural architectures performing sequence- se-uuueece mapping. For input tokens SN = SN. E RN'P,, the transformer produces a corresponding SN = _ NN] E R x2 xP by dynamically mixing tokens via its attention mechanism. An L-layer transformer To Nx + RNNx parametrized by [0:] is composition of blocks TL = To, To. Each To, is a sequeence-t-sequeeee function given by
(S) (MLP 0UILP Agpe)
where MLP oA a multilayer perceptron and Agues is the attention mapping. This paper studies the simplified block Ta( := As(s), with L = 1 and MLP opuP ld.
The causal h--eaded attention block with residual connections is given by
H
Ao(S) S
Σ a
M
1 sw (Wk) TT ) Sw
h=1
where the parameters [W W W Wol are the query, key, value, and projection matrices, respectively; is scaling constant; a is the softmax normalizing function applied row-wise; and M R NxN, with M.; if 2 j and -c otherwise, causal mask.
Similar to prior works (Von Oswald et al., 2023a; Ahn et al., 2023; Mahankali et al., 2023), we restrict our study the analytically tractable setting of single-headed 2020). Without loss of expressivity, we drop the projection matrix Wo and consider the W W as a single matrix WoK & Rexp, Since we're working in the few-shot scenario, we're concerned solely with predicting the final position
headed linear attention (Katharopoulos et al.,
N-1 Σ
N = TO(S)N = SN +
W
8,8; W WQ
N
where we set 7 = N - and omit the last sum element due to a token asymmetry discussed next.
Token construction. We use the token construction approach of Von Oswald et (2023a); Ahn
et al. (2023); Mahankali et al. (2023). The input matrix Yo built for AR(s) data (4) is
y1
U2
Fr-s-1 VT-1
ÿr-s 0
Yo
ys+1
Mo+2
where 2 is the window size of the AR process and Yo ((++1)x(T-s), The corresponding scaling constant r in (6) becomes T - --1 1. The last column represents the "test" token, whose last entry is filled in the transformer's forward pass by the estimate ỘT. This asymmetry motivates the last term's S removal in (7).
Lemma 3.1 below ensures the existence (by construction) of linear attention layer producing Yo from the raw sequence (ye ( Its proof is deferred to Appendix C. due to space constraints.
Lemma 3.1. For a given s >= there exists an + 1-headed linear attention layer with positional encoding which transforms input sequences [V1. V2..... into
U1
UT
LUs+1
0
0
The latter quantity is essentially equivalent to Y, as defined in equation (8).
Data distribution, loss function, and training paradigm. We consider trajectories {yi (y sampled from system (1), where each trajectory corresponds to different parameters A, c and mp. We assume 2D (O ), and impose the following condition on the distributions of A and c.
Assumption 3.2 (LDS family). The system matrix A E gaxd is sampled from centrally symmetric distribution supported on { I E Idd p(M) S 1}. for which holds that
P ( Bi E [d], s.t. Aa(A) {₂(A))) 0.
In other words, has a simple spectrum almost surely. The observation vector c E Rd is sampled independently, from distribution that is absolutely continuous w.r.l. the Lebesgue measure over Te.
5

---

Except for the central symmetry condition, Assumption 3.2 simply ensures that Assumption 3.1 holds w.p. for every sampled LDS (proven in Appendix C.2). The central symmetry of A's distribution, on the other hand, a technical requirement for proving our main result.
Data generation procceds two steps: we sample A, C, and Vo independently and observe the evolution of system (1) for T steps. Note the noises wy and i/ are independent of A, C, and Zo- We then construct Y. (8) for a fixed 8, and train our model to minimize the in-context loss
( E
yr yr)
(10)
where the expectation here and throughout is taken A, c, æu, {we), {ve)- The subscript (s 1, T s) isolates the rightmost element of the final output token.
4 OPTIMAL PARAMETER CONFIGURATIONS
This section presents our formal analysis and implications relative to prior literature. Our theoretical contribution is two-fold. First, in Lemma 4.1 reveal salient structure within the first-order optimality condition, which plays an important role in finding optimum configurations relative to the in-context loss (10) for any 1. Second, in Theorem 4.1 we prove that the weight configuration implementing one step of GD a global minimizer of loss (10) for AR(I)-type tokens.
Unlike the i..dd case, each token generated by the LDS depends the entire history. This causes high-order data moments to appear in the in-context loss, which can only be handled by unrolling the process to its initial state æ4- general approach dealing with this statistical dependence given in Appendix D.3. We now describe the structure emerging within the first-order optimality condition.
Following Ahn et al. (2023), we use basic algebra (Appendix D.3) rewrite loss (10) as
[(b, .,) = E
Y bu /UT +-1h
(1I)
T
R=1
where W.T [0 (++1)xxs b) W QK [2
........ 2s, 0s+1] and Y = Yo, a.r--1 Y
The
pero-padding of both matrices comes from predicting solely the last position of the final token. Consequently, parameters satisfying for all j E [s] that
E (Y, (Y,
bo yT 14 1 I Y
(T -1)E UTVT- DTTTT---144 8]
(12)
critical points of the loss. Notably, the RHS of (12) has banded structure induced by the distributional symmetry of A, whose pattern alternates depending on the parity of 8 +. In particular, E 22, and RHS (12) (s if + j 2% + 1, with
RHS (12) B) if
0 0
* 0
0
0
Bo(s, =
and (s, =
(13)
0
0
0
*
*
where is placeholder for structurally unconstrained elements which can evaluate to arbitrary reals,
not necessarily the same different positions (see Appendix D.3). Note that Bn(s, j) is constrained constrained to
to zero at positions (r, () [s x [s such that r + 2N, and B (s, j) whenever r + 2N + 1. Lemma 4.1 below formalizes a class of parameters ensuring that LHS (12) respects the structural sparsity pattern of RHS (12).
Lemma 4.1. For any 2 J. E [s], the parameters WOk and Wy having the structure
(1
WeK
W
[((+1)xx T24e mod 2:2(
(14)
0 1+

---

* 0 with R = ( X
0 and r = r @ [0]
ensure that LHS (12),,, = whenever
r rl 2N and s + E 22, or E 2N and
22 + 1.
Note that these parameters recover the zero-structure of Bo(s. and B1 (s.j). Lemma 4.1 can be understood as structure-basse -base narrowing of the parameter class likely to hold minimizers of loss (10). The experiments in Section 5 give empirical support for this claim.
Our second step to use structure (14) to identify a global minimizer of loss (10) for A-(p)-type tokens. The result is stated in Theorem 4.1, whose proof we defer to Appendix D.4.
Theorem 4.1. Let Yo encode the input tokens for s = 1. Then, the optimal parameters 0* (W Qr W of single linear self-altention layer with respect to loss L(0) are
(T (T-2)E ur -MT Eisr Mc30+1
0 0 A
W
(15)
0
up to rescaling with nonzero constant.
Broadly, the proof of Theorem 4.1 encounters difficulties compared the i.i.d. case: the number of terms that need to be matched satisfying the first-order optimality condition, and the full- history dependence of the data. We address the first obstacle by using the structural simplification of Lemma 4. 1. and the second by invoking Isserlis' theorem (Isserlis, 1918), which provides a tractable decomposition of the higher-order moments of the data. Details are given in Appendix D.3.
Notably, a forward pass using the optimal parameters (15) amounts to the prediction given after one GD step on C CARRI)(U (w) starting from wg We thus recover the oottiiiiitii view upheld by works in the i.i.d. setting (Ahn et al., 2023; Mahankali et al., 2023) but for LDS -Do-prrddeed data.
4.1 DISCUSSION
To our knowledge, the only other architecture proposed for handling noisy observations yr of type (I) is given by Cole al. (2025). Theirs is part of a proof of existence by construction and, as such, is not accompanied by confirming experimental evidence. Different from us, they propose an attention-only transformer that unrolls modified Richardson iteration meant estimate
(+EL, (+ ET wxwe] -1
for simpler LDS with direct state access. Their construction extends to the setting of objective (4) via the work of Tsiamis & Pappas (2019), who give a high under our assumptions. However, their transformer has a minimum of two layers, of which the first is fixed, therefore providing no guarantee that training will recover it. Our results take first step towards filling this gap.
probability result for the existence of (ET++ प.प."
Tangentially, Akram & Vikalo (2024) construct a transformer emulating the KF, contingent on knowledge of the system parameters and an elaborate token augmentation scheme. While this architecture is capable of computing the forward KF observation Dr. it relies on ideal knowledge of LDS (I) which rarely encountered in practice.
Theorem 4.1 sets forth a plausible hypothesis for prior experiments (Du et al., 2023, Fig. 2) using a GPT-2 architecture trained autoregressively with data (1) for stable A e 5d 54 Their results highlight the transformer's 's competitive performance relative to the KF for predicting the next observation of a previously unseen sequence, in-context. These experiments suggest an implicit form of system identification might be executed in the forward pass, though the mechanism remains unstudied. Through the -o-aa-oopiimizaiion lens, we can interpret the high accuracy of GPT-2's in-context predictions possible consequence of Theorem 2 of Kozdoba et al. (2019). Importantly, the latter result implies that for an arbitrary. finite family S of LDSs (1) and an E > O, there exists a window-lengre s(E) such that the optimal AR(s(E)) predictor ineurs average error that is least as good, up to e, as that of the forward observation prediction D of the best KF in S. Our results take the first step in formally exploring this hypothesis.
7

---

-6.00
1.36
-0
6.15
0.4
04
0.2
02
-6.00
0.00
0.00
0.0
0.00
-1.34
6625
6000
(a) Training loss for AR(1) tokens.
(b) Wax for AR(1).
(c) Optimal Wy for AR(1).
6.3
1:91
-0.01
-2.35
1
0.00
0.00
1:31
-6.4
1.0
0.8
6.5
0.6
0.6
-6.6
-6
0.00
0.00
0.00
3-
0.00
0.00
-1.94
2000
1000 kiltaralions
6000
8000
(d) Training loss for AR(2) tokens.
(c) WeR for AR(2).
(f) W. for AR(2).
-72
0.45
0.00
0.21
0.00
0.00
0.00
0.00
0.00
7.3
-0.01
1.50
-0.00
-1.99
2
0.00
0.00
0.00
1.46
7.4
-7.5
3-
0.62
0.00
-1.16
-0.00
3+
0.00
0.00
0.00
0.01
0.0
0.0
-7.6
4-
0.00
0.00
0.00
0.00
4:
0.00
0.00
0.00
2.28
2000
4000
6000
8000
[(leeraioon)
(g) Training loss for AR(3) tokens.
(h) Wax for AR(3).
(i) W for AR(3).
1
77.8
0.00
0.41
0.00
-0.70
-0.00
2
0.00
0.00
0.00
0.00
-0.00
0.6
88.0
2.58
0.00
0..4
0.00
0.00
2.44
0.4
-.1
1.55
4+
-8.2
5
0.00
0.00
0.00
0.00
0.00
5:
0.00
0.00
0.00
0.00
0.89
2000
8000
kilterations)
() Training loss for AR(4) tokens.
(k) W for AR(4).
(I) W for AR(4).
Figure Experimental results AR(1-1) tokens showing optimally-trained attention parameters.
EXPERIMENTS
We now present numerical evidence supporting our theory. All experiments were implemented in Python 3.12 and run ThinkPad T14p with 32 GB RAM and 22-core Intel TM Ultra 185H processor. The code is available at https: /ttthbbttttttt//g//g/ com/ htmm//HZhangC1.. -ddo-aads
h
We train architecture (7) on sequences { T = 30, each sampled from a different LDS of type (1) with hidden state dimension d = 5. The number of training iterations is 8000 for all cases,
8

---

with increase in the batch size for every increase in window size, starting from 3000 for AR(I). A fresh batch of LDSs is sampled at every iteration. The experiments cover the following four settings.
(a) For each sequence, we sample A's diagonal entries uniformly in the interval [-1, 1] and set c = 1d- The noise covariances are set to Ene = 1e-21 and of = 1e-2. The results are depicted in Figure for AR(1-4) tokens.
(b) For each sequence, we sample a vector independently, and set A = Q diag(v)Q. We further sample c ~ Unif([-2, 2, 2| indepen- dently. The noise covariances are set to Σn 1e-2/ and o 1e-2. Experiments for AR(1-))type tokens are provided in Figure 3 in Appendix B.3.
Unif([-1, 114) and a matrix Q ~~ Haar(O
(c) For each sequence, we sample a vector ~ Unif([-1, 114) and a matrix Q ~ Haar(O independently, and set A Q di) We further sample
Unif([-0.5, 0.5)
independently. We fix the covariance En Qw T e for all systems, for arbitrary Qw ~ Haar (dd)) We fix og = 1e-2. Experiments for AR(1-3)-tye tokens are provided in Figure 4 in Appendix B.4. 1ni and matrix = pp.ll with Pi. i.d. 1 ([-1, IJ), and set A d-idiag(v)... We further sample C Uni The noise covariances are set to En, 1e-21 and of = 1e-2. Experiments for AR(1-2)- tokens are provided in Figure in Appendix B.5.
0.95, 1.0, 1.05, 1.1]) Qw
(d) For each sequence, we independently sample a vector
All the settings above have 35 ~ (O I) oj 1e-2. 2. Note that we could have used any other centrally symmetric distribution with marginals supported on [-1,1] for the sampling of the diagonal v. e.g., Unif(S*-1) - uniform on the unit sphere: Unif({ (x R ||>||2 S 1}} uniform inside the unit ball, etc. These sampling schemes obey Assumption 3.2, as proven in Appendix E. All results are averaged over 3 random seeds. The weights are learned using Adam (Loshchilov & Hutter, 2017) with gradient clipping and a learning rate schedule consisting of linear warm-up phase followed by cosine annealing (Loshchilov Hutter, 2016). A full list of hyperparameters provided in Appendix B.
Subplots (b,c) in Figures 1, 3, 4 and show optima conforming to Theorem 4. for AR(1)-type tokens. Moreover, subplots (e,f.h,i,k.I) in Figures 1 and 3. subplots (e.f.h.i) in Figure 4 and subplots (c.f) in Figure confirm the pattern dictated by Lemma 4.1 for general s 2 2. Finally, we empirically show that in setting (a), the weights converge the sparsity pattern predicted by Lemma 4.1 in terms of the Jaccard distance of their supports. Setup details and results are given in Appendix B.2.
Interpreting the sparsity pattern for AR(s) 2. A quick calculation of the forward pass reveals that weights trained to optimality with AR(s) tokens for 2 do not implement standard GD in the forward pass, but an iteration reminiscent of the Preconditioned Conjugate Gradient method (PCG) (Hestenes et al., 1952; Shewchuk et al., 1994), or more generally, to Krylov subspace methods (Saad, 2003, chapters and 7).
We start by observing that the forward pass factor T- Y (with Y defined in (11)) has meaningful block structure involving the gradient at zero and the Hessian of (AR(s)w),
V'Lan()) (O)
T Y
(16)
where 7 T--1 Σx V2 E R and V2LAR(s) E R denotes the constant Hessian.
Together with the parameter structure of Lemma 4.1 and the experiments, we use expression (16) to rewrite the transformer-induee after two PCG steps on LAR(s) (w) starting from wo = 0. We describe the case for even s, with the
predictor in manner that highlights its resemblance to that obtained
odd case following similarly. Let 2k, k E N and N = (++1)*+1
Consider weights
2
C1
Ck+1
CN-
Ck+2
0
0 Osxs
CN -++
Wox
E R 1xs+1
Wv
E <R+tixs+1
CN
0
CN

---

where e; RR, Vi E [V]. Renaming the top left x block of Wok P, the top right 1 block as P. and the top right s x 1 block of Wv as q, the consformer-indueee linear predictor T-s-T Wok Y WW w is
PV + ()PVCAR))(O) + E2p,
(17)
where E = cN E R and 6 = q VL)R()) (0) + CNT E R.
Comparatively, the PCG (Shewchuk et al., 1994, p. 51) predictor obtained after two steps on loss LAR(s) with preconditioner p-1 starting from Wo = is (see Appendix E.2)
PV PVL + T2 PVL))))))
(18)
where TI. F2 E R are serration-deennden constants, Note that we have suspended the requirement for P's symmetry here - we will address this drawback shortly.
We make few observations on the similarities and differences between the two predictors. The second terms of (17) and (18) coincide up to scaling. The first terms of these predictors represent directions mapped by PV AARe) which is typical of Krylov subspace methods. Moreover, we initialize the conjugate direction of PCG to the vector q. the first terms also coincide up to scaling. This prompts the natural question of whether directions
PV (0) and LARs) are aligned, or at least significantly so. As a preliminary test, we compute this alignment in terms of cosine similarity when replacing (AR and L AAR(s) (0) with empirical estimates of their expectations. Using the weights reported for the AR(4) case in settings a) and b) (Figures and 3), batches of 16000 sequences for estimating the expectations and averaging the result over five random seeds, we obtain cosine similarity of 0.88 + 0.05 for setting a) and 0.93 = 0.01 for setting b). The values suggest a strong alignment, which gives weight to this line of interpretation.
In terms of predictor dissimilarities, the set of directions whose linear combination yields (17) contains the additional vector p. This structure is typical of so-called augmented Krylov methods, whereby "correction directions" are added to the standard Krylov search space to compensate for ill-behaved modes of PV AAR(s) (Carpenter et al., 2010, and references therein). Broadly, they can be seen as PCG generalizations, and have the added benefit of accommodating asymmetric Ps, thus resolving the earlier drawback. We test whether p acts correction by evaluating the residual obtained prior to and after incorporating p. We use the same setup as above, and compute the cosine of the angle between the preconditioned residual evaluated q (which coincides with w in PCG with non- stansaandard initial conjugate direction), and the direction PVLLAR(s)) suggest that p moves the predictor in residual reduction direction. Indeed, we observe cosine of -0.99 + 75-5 for the AR(4) case in setting a), and -0.99 + 35 for the AR(4) case in setting b). Taken together, these preliminary results warrant further exploration of this interpretation.
Anti-alignment would
Finally, we remark that our interpretation of the AR(s) case does not contradict the plain GD step observed for AR(I), since PCG variants collapse GD for one-dimensional covariates.
LIMITATIONS & FUTURE DIRECTIONS
We have sketched a path to understanding how attention layers learn LDSs in context by leveraging results from the improper learning literature on system identification. Our study fully characterizes the learning of order-one autoregressive LDS approximations (Theorem 4.1) and narrows the class of weights that hold minimizers for higher-order approximations (Lemma 4.1), with all findings experimentally confirmed. These contributions provide the building blocks for closing further gaps, the most pressing of which is describing the optima of loss (10) for s 2 2. A meaningful intermediate result here is quantifying the approximation power of optima obtained in the stationary LDS regime, where the problem structure further simplifies due to Y being Toeplitz. Such assumption is reasonable, since systems with p(A) < 1 reach stationarity exponentially fast. Further valuable directions concern bridging the theory-practice gap, such as extending the analysis to typical causal pretraining objectives, and empirically probing the parallel between Krylov subspace methods and the predictor modeled by multi-layer transformer.
10

---

ACKNOWLEDGMENTS
The authors are grateful Xiang Cheng and Anastasios Tsiamis for helpful discussions throughout the development of this work. The authors also sincerely thank the anonymous reviewers for their constructive feedback which helped improve and clarify this manuscript. All the authors of this work were generously supported by the Alexander von Humboldt foundation.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 013456144456555 2023. 1, 2, 6, 7, 2.
Usman Akram and Haris Vikalo. Can transformers in-context learn behavior of a linear dynamical system?, 2024. URL https: / arxiv. org/ abs / 22200.10002 16446. 3,7
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm in-context learning? investigations with linear models. arXiv preprint arXiv:2211. 2022. 2
Peter Beim Graben, Bryan Jurish, Douglas Saddy, and Stefan Frisch. Language processing by dynamical systems. International Journal of Bifurcation and Chaos, 14(02))599-621, 2004. 2
David Belanger and Sham Kakade. A linear dynamical system model for text. In International Conference on Machine Learning, pp. 833-842. PMLR, 2015.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020. I
Mark H Carpenter, C Vuik, Peter Lucas, Martin vanGijzen, and Hester Bijl. A general algorithm for reusing krylov subspace information. i. unsteady navier-stokes. Technical report, 2010. 10
Frank Cole, Yulong Lu, Tianhao Zhang, and Yuxuan Zhao. In-context learning of linear dynamical systems with transformers: Error bounds and depth-separation. arXiv preprint arXiv:2502.08136, 2025. 3,7
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Tianyu Liu, et al. A survey on in-context learning. arXiv preprint xxXiv:2301.000344 2022.1 1
Zhe Du, Haldun Balim, Samet Oymak, and Necmiye Ozay. Can transformers learn optimal filtering for unknown systems? IEEE Control Systems Letters, 7:3525-3530, 2023. 2, 3,
Jeffrey L Elman. Language dynamical system. Mind motion: Explorations in the dynamics of cognition, pp. 195-223, 1995.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022. 1.2
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. International Conference on Machine Learning, pp. 13988-11442. PMLR, 2023.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings. 2010. 15, 16, 17
Gautam Goel and Peter Bartlett. Can transformer represent kalman filter? In 6th Annual Learning for Dynamics & Control Conference, pp. 1502-1512. PMLR, 2024. 3
James D Hamilton. Time series analysis, 1995. 4

---

P Jeff Harrison. Convergence and the constant dynamic linear model. Journal of Forecasting, 16(5): 287-292, 1997.
Magnus R Hestenes, Eduard Stiefel, et al. Methods of conjugate gradients for solving linear systems. Journal of research of the National Bureau of Standards, 49(6):409 436, 1952.
Leon Isserlis. On formula for the product-moment coefficient of any order of a normal frequency distribution in any number of variables. Biometrika, 12(1/2):134-139 1918. 7, 23
Paul Jaccard. Étude comparative de distribution florale dans une portion des alpes et des jura. Bull Soc Vaudoise Sci Nat, 37:547-579, 1901. 17
Hui Jiang. A latent space theory for emergent abilities in large language models. arXiv preprint NSKV:23304099660 2023.
Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. Journal of basic Engineering. 82(1):35-45, 1960. 2, 3,4
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rons: Fast autoregressive transformers with linear attention. In International conference on machine learning. pp. 5156-5165. PMLR, 2020. 5
Mark Kozdoba, Jakub Marecek, Tigran Tchrakian, and Shie Mannor. On-line learning of linear dy- namical systems: Exponential forgetting in kalman filters. Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 4098-4105 2019. 4,7
Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM computing surveys, 55(9):1-35, 2023.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint 1.XX::1608.03983,, 2016. 9, 15, 16, 17
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. 9, 15, 16, 17
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03377 2023. 1, 5,
Thomas P Minka. From hidden markov models linear dynamical systems. Technical report, Citeseer, 1999. 2
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019. 3
H. L. Royden and P. M. Fitzpatrick. Real Analysis. Prentice Hall, 4th edition, 2010. ISBN 978- 0131437470. 22
Yousef Saad. Iterative methods for sparse linear systems. SIAM, 2003.
Michael E Sander and Gabriel Peyré. Towards understanding the universality of transformers for next-token prediction. arXiv preprint arXiv:2410.03011, 2024. 1,
Michael E Sander, Raja Giryes, Taiji Suzuki, Mathieu Blondel, and Gabriel Peyré. How do trans- formers perform in-context autoregressive learning? arXiv preprint arX1v:2402.05787, 2024. 2,
Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Do pretrained transformers learn in-context by gradient descent? arXiv preprint arXiv:2310.08544) 2023. 1
Jonathan Richard Shewchuk et al. An introduction to the conjugate gradient method without the agonizing pain. Technical report, Camegie-MMlloo University. Department Computer Science Pittsburgh, 1994. 9, 10, 33
12

---

Whitney Tabor, Christopher Juliano, and Michael Tenenhaus. A dynamical system for language processing. In Proceedings of the Anmual Meeting of the Cognitive Science Sociery, volume 18, 1996. URL httpsss//sccooosssssstttieeeeeee
it ht 78r6h0cg.. 2
Anastasios Tsiamis and George Pappas. Finite sample analysis of stochastic system identification. In 2019 IEEE 58th Conference on Decision and Control (CDC), pp. 3648-3654, IEEE, 2019. 4, 7
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017. 5
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pp- 35151-35174. PMLR, 2023a. 2, 21
Johannes Von Oswald, Maximilian Schlegel, Alexander Meulemans. Seijin Kobayashi, Eyvind Niklasson, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, et al. Uncovering mesa-optiiiiztiin algorithms in transformers. arXiv preprint arXiv:2309.05858. 2023b. 1, 3,
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are latent variable models: Explaining and finding good demonstrations for in-context learning. Advances in Neural Information Processing Systems, 11111114-56636 2023. 1
Noam Wies, Yoav Levine, and Amnon Shashua. The learnaillity of in-context learning. Advances in
Neural Information Processing Systems, 36:36637-36651, 2023. 1
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in- in-context
learning as implicit bayesian inference. arXiv preprint arXiv:2111.02085 2021. 1, 2
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024. 2, 3
Chenyu Zheng, Wei Huang, Rongzhen Wang, Guoqiang Wu, Jun Zhu, and Chongxuan Li. On mesa-optimization in autoregressively trained transformers: Emergence and capability, Advances in Neural Information Processing Systems, 17:49081-49129, 2024. 3
Ingvar Ziemann, Nikolai Matni, and George Pappas. State space models, emergence, and ergodicity: How many parameters are needed for stable predictions? arXiv preprint arXiv:2409. 13421, 2024.
13

---

A LLM USAGE DISCLOSURE
LLMs were used in elaborating this paper as follows:
Finding related work.
Computing the result of polynomial multiplications.
Generating LaTeX tables and tikz figures.
Transferring proofs from pen-and-paper papeeppperrrr format into LaTeX automatically using the online tool Manus https: / /manus....... im/.
14

---

B EXPERIMENTS FURTHER DETAILS
B.1 HYPERPARAMETERS
Table Training hyperparameters of setting (a) in Section
Hyperparameter
Value
Xavier normal distribution (Glorot & Bengio, 2010) with
Weight initialization
gain 1e-6
AdamW (Loshchilov & Hutter, 2017) with B1 0.9 and 82 0.98 for all AR(1-4), Ie-9
Optimizer
Weight decay
5e-3 for AR(2 4) and 1e-2 for AR(1 & 3)
2e-2 for AR(I), 3e-2 for AR(2), 3e-3 for AR(3),
Learning rate (i.e., max. val.)
4e-3 for AR(4)
Min. learning rate
le 4 for AR(1 2) and 1e-5 for AR(3 & 4)
Linear warmup
800 iter.
Decay schedule
Cosine annealing (Loshchilov & Hutter, 2016)
Max. decay steps
7200 iter.
Max. grad norm (clipping)
300
Random seeds
(666013, 1, 1,0}
3000 for AR(1), 4000 for AR(2), 10000 for AR(3),
Batch size / iter.
18000 for AR(4) (4)
Total iter.
8001
15

---

Table Training hyperparameteers setting (b) in Section
Hyperparameter
Value
Xavier normal distribution (Glorot & Bengio, 2010) with
Weight initialization
gain 1e-6
AdamW (Loshchilov & Hutter, 2017) with 31 = 0.98 for
AR(1), 0.96 for AR(2), 0.92 for AR(3), 0.88 for AR(4), 82 = 0.99 for AR(I), 0.98 for AR(2), 0.96 for AR(3),
Optimizer
0.92 for AR(4), 1e-9
Weight decay
1e-2 for AR(I), 1e-3 for AR(2), 1e-4 for AR(3 & 4)
1e-3 for AR(1), 4e-3 for AR(2), 6e-3 for AR(3),
Learning rate (i.e., max. val.)
8e-3 for AR(4)
Min. learning
1e-5
Linear warmup
800 iter.
Decay schedule
Cosine annealing (Loshchilov & Hutter, 2016)
Max. decay steps
7200 iter.
Max. grad norm (clipping)
300
Random seeds
(666013, 1, 1,0}
4000 for AR(1), 8000 for AR(2), 14000 for AR(3),
Batch size iter.
18000 for AR(4)
Total iter.
8001
Table 3: Training hyperparameters of setting (c) in Section 5
Hyperparameter
Value
Xavier normal distribution (Glorot Bengio, 2010) with
Weight initialization
gain 1e-6
A (Loshchilov & Hutter, 2017) with 31 = 0.98 for AR(1), 0.96 for AR(2), 0.92 for AR(3), 83 = 0.99 for
Optimizer
AR(I), 0.98 for AR(2), 0.96 for AR(3), 1e-9
Weight decay
1e-2 for AR(1), 5e-3 for AR(2), 1e-3 for AR(3)
Learning rate (i.e., max. val.)
2e-3 for AR(1), 3e-3 for AR(2), 4e-3 for AR(3)
Min. learning rate
le-5
Linear warmup
800 iter.
Decay schedule
Cosine annealing (Loshchilov & Hutter, 2016)
Max. decay steps
7200 iter.
Max. grad norm (clipping)
300
Random seeds
(666013, 1, 1,0}
Batch size / iter.
4000 for AR(1), 8000 for AR(2). 16000 for AR(3)
Total iter.
8001
16

---

Table Training hyperparameteess setting (d) in Section
Hyperparameter
Value
Xavier normal distribution (Glorot & Bengio, 2010) with
Weight initialization
gain 1e-6
AdamW W (Loshchilov & Hutter, 2017) with $1 = 0.92 and ;2 0.96 for all AR(1 4), 1e-9
Optimizer
Weight decay
5e-3 for AR(I), 1e-3 for AR(2)
Learning rate (i.e., max. val.)
3e-4 for AR(1), 5e-4 for AR(2)
Min. learning rate
1e-5
Linear warmup
800 iter.
Decay schedule
Cosine annealing (Loshchilov & Hutter, 2016)
Max. decay steps
7200 iter.
Max. grad norm (clipping)
300
Random seeds
(666013, 1, 1,0}
Batch size iter.
4000 for AR(1), 10000 for AR(2)
Total iter.
8001
B.2 EXPERIMENTS SHOWING CONVERGENCE TO CHECKERBOARD PATTERN DURING TRAINING
This set of experiments serves illustrate that parameters WΩK and Wv converge to the checker- board pattern across iterations. Since the non-zero values of these parameters are of different magnitudes and we not have their theoretical expressions for window-sizes greater than we shall only consider their non-zero support, as follows.
Definition B.1. For matrix M Rdxm its support is defined as the collection of positions corresponding to non-zero values
SMpp(M) ((i,)) [d] x [m] ai.j + 0}.
(19)
Additionally, the support-inducee mask is binary matrix with unit entries on the support
iidjimm
mask(M) = [(AJ))ssup((M] (M)
(20)
ki=i
where 1 1 if condition is true and O otherwise, is the indicator function centered
We rely on the Jaccard distance (Jaccard, 1901) adapted to binary matrices A, B
1 Ens Di.bii. Ei, (eaxx(arjj. bij)
J (A, B)
(21)
to track whether the support-induced masks of our parameters during training converge to the predicted
(for AR(1)) hypothesized (for AR(s) 2 2) sparsity patterns of Lemma 4.1. Our experiments employ a tolerance level of 1e-1 when computing the masks of Wv and WoK. meaning that any entry below this value is considered zero. The results are depicted in Figure and its subplots for varying window sizes, where Ming and Mise represent the masks posited in Lemma 4.1 for a null tolerance level. The illustrations empirically confirm that our parameters' supports converge to the identified in Lemma 4.1.
17

---

1.0-
1.0
0.8
0.8
d.c. (..amkWWal (1.
bxc/mmsk(VVv).... Mi
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0-
0
2000
4000
6000
8000
2000
4000
6000
8000
(iterations)
(iterations)
(a) The Jaccard distance of WeR for AR(1).
(b) The Jaccard distance of W for AR(1).
1.0
1.0
0.8-
dwcc/mark(Mo).. Mpor)
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0-
0.0
0
2000
4000
6000
8000
2000
4000
6000
8000
(Iterations)
(Iteraaions)
(c) The Jaccard distance of Wox for AR(2).
(d) The Jaccard distance of Wy for AR(2).
1.0
1.0
0.8
0.8
(mask(()W)))) Mige))
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
2000
4000
6000
8000
2000
4000
6000
8000
k (iterations)
(Iterations)
(c) The Jaccard distance of Wox for AR(3).
(f) The Jaccard distance of Wy for AR(3).
1.0
1.0
0.8
(mask( M
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0-
2000
4000
6000
8000
2000
4000
6000
8000
(Iterations)
(iterations)
(g) The Jaccard distance of Wox for AR(4).
(h) The Jaccard distance of Wy for AR(4).
Figure The experiment results Jaceard distance between the M and W Wo and Jaceard distance between the Mi and Wy for AR(1 4). Both converge 0 the end of the training.
18

---

3.50
5.00
1021-200
08
04
-6.2
0.00
9.00
0.00
6.00
00
6.25
2000
4000
6000
8000
(a) Training loss for AR(1) tokens.
(b) Wax for AR(I).
(c) Optimal wP for AR(1).
6.7
1:38
0.01
-1.06
6.8
08
1400300
M/M.M
2-
0.4
-7.0
2000
4000
5000
8000
(arelaan)
(d) Training loss for AR(2) tokens.
WaR for AR(2).
(f) WV for AR(2).
1 9.36
0.30
0.14
0..00
0.00
0.03
0.09
-8.00
9.81
08
0.0
2:30
0.00
02
0.02
0.02
4-
3.00
3.20
0.00
0.00
0.00
0.03
0.03
3.00
7.7
4000
(g) Training loss for AR(3) tokens.
(h) WÓx for AR()).
(1) W for AR().
7.5
2: 0.09
0.92
0.00
0.35
0.00
0.00
0.00
0.00
0.00
0.50
0.
-766
0.00
0.25
0.00
0.31
06
0.00
0.00
0.02
2.24
0.6
7.7
-7.8
400
0.0
5- 5.00
0.00
0.00
0.00
0.00
54 0.00
0.00
5.00
0.03
0.10
2000
1000
5000
8000
(j) Training loss for AR(1) tokens.
(k) W for AR(1).
(1) W. for AR(1).
Figure 3: Experimental results for AR(1-4 with non-diagonal, symmetric A. random and isotropic D...
setting (b) in Section which align with the Lemma 4.1.
B.3 EXPERIMENTS WITH NON-DIAGONAL, SYMMETRIC A, RANDOM & AND ISOTROPIC Σ,,
The LDS which generates the training data is follows. For each sequence, sample Unif([-1, 1]d), sample Q ~ (d)r(O(d)) and set A = Q diag(v) Sample c Unif([-2, 2()) The noise covariances set to Σn, 1e-2I and of 1e-2.
19

---

B.4 EXPERIMENTS WITH NON-DIAGONAL, SYMMETRIC A, RANDOM AND NON-DIAGONAL, ANISOTROPIC Σng
The LDS which generates the training data as follows. For each sequence, sample v
Uni[[[-1, 1J)); sample Q Haar(O(d)) and set A = Q diag(v)Q; sam- ple e Unif([-0.3) ). Set the process noise covariance Ew = Qw diag(1e-2
(0.9. 0.95, 1.0, 1.05, 1. where Qw Haar(O(d)). Set oe 1e-2.
4,25
1
1-
0.8
-4.35
0.66
0.6
-4.40
-4.45
-4.0
2000
4000
6000
8000
ema)
(a) Training loss for AR(1) tokens.
(b) Wax for AR(1).
(c) Optimal W for AR(1).
0.00
0.53
5.0
1.0
1.0
m/
-.11
2
0.00
-0.23
-0.00
2+
0.00
0.00
.0.00
-5.2-
0.2
02
-.00
0.0
-5.3-
2000
4000
6000
8000
Iitiarabe kitaaraiion)
1
(d) Training loss for AR(2) tokens.
(c) W for AR(2).
(f) Optimal W for AR(2).
-5.6
-0.16
-5.7-
0.54
5.88
0.4
04
-5.9
0000
0.00
-0.80
-0.00
0.2
0.00
0.00
0.00
0.00
02
0.0
6.0-
0.00
0.00
0.00
0.00
4.
0.00
0.00
0.00
.0.86
4000
6000
(g) Training loss for AR(3) tokens.
(h) Wox for AR(3).
(i) Optimal Wi for AR(3).
Figure Experimental results for AR(I) with non-diagonal, symmetric random and non-diagonal, anisptropic Ew, setting (c) in Section 5, which align with the Lemma 4.1.
B.5 EXPERIMENTS WITH NON-DIAGONAL, NON-SYMMETRIC A, RANDOM e AND ISOTROPIC Σ,
The LDS which generates the training data is as follows.
For each sequence, sample d Unif([-1,1]"), sample P P= [p..], by sampling Pt, i.i.d. from ([-1, 1]), and set = P- P diddia())) Sample e ~ Unif([-1, 1)d). The noise covariances are set to E,, 1e-21 2I and o 1e-2.
In practice, we need to guarantee is well conditioned. After sampling Pı,j i.i.d. from / ([-1, 1)). decompose P Q Rt where Q is orthogonal matrix and R is upper-triangle matrix from the QR decomposition of P. We modify the diagonals of R manually to make sure max, R., 2 and right multiply Q with the modified R to have the well conditioned P.
min; R.;
20

---

-6.05
-3.00
5.00
6.10
08
04
-6.15
0.00
9.00
0.00
6.00
0.17
0.0
-6.20
2000
4000
6000
8000
(a) Training loss for AR(1) tokens.
(b) Wax for AR(1).
(c) Optimal wP for AR(1).
-710
-7115
4.13
0.
4.18
2.00
0.03
-.15
-7.20
10
1.0
7.25
4.02
26
7.00
0.03
46
-730
42
-7.00
602
0.00
3.00
0
2000
4000
60000
8000
(d) Training loss for AR(2) tokens.
(c) WŌw for AR(2).
(1) WV for AR(2).
Figure 5: Experimental results for AR(1-) with non-diagonal, non-symmetric A, random e and isotropic Eag.
setting (d) in Section 5, which align with the Lemma 4.1.
C SECTION 3 PROOFS
C.I PROOF OF TOKEN CONSTRUCTION LEMMA
1, there exists + 1-headed linear attention layer with positional encoding which transforms input sequences [y) ,,.....,37 into
Lemma 3.1. For given >=
91
ÚT
((-1)x(1---1)
Lys+1
0
0
The latter quantity is essentially equivalent to Yo as defined in equation (8).
Proof. We first define a matrix right-shift operator, which shifts each row one position to the right,
padding the first column with zeros. Let >; mmxx + R be > (M) M R, where
0
0
(22)
On-1
We follow Von Oswald et al. (2023a) in using the one-hot positional encodings, concatenated to the input sequence to obtain tokens {[y,, e.] We define 1 attention heads given by
Define Wq
T WK
ERT T and Wy E RT+Ixs as follows:
of Ir
/ =
Vh E |s + 1]
Oo:
(.. >> (Tr) ...)
(W
times
Wh
h
s+1
Vh E
[s
(23)
00+1
Or
Each head then computes the following
21

---

y1
1
0
07
3/1
Y2
ys
yr
y1
1
0
0
2/2
1
0
0
y/2
0
1
W
(W T
Wv
Lyr
I--
04-1xh-1
-
T
0
1/2
0
Uh
U4++1
yT
Opxa+1
Summing over the outputs of all heads, get equivalent representation to (8).
C.2 PROOF THE ALMOST SURE OBSERVABILITY OF THE LDS
We seek to show that Assumption 3.2 ensures LDS (1) observability w.p. 1. Note that the central symmetry of the distribution is irrelevant for this statement, and only relevant for the proofs in Section 4. We repeat Assumption 3.2 below for convenience.
Assumption 3.2 (LDS family). The system matrix A E gdxd is sampled from centrally symmetric distribution supported { M E Rdddd p(M) S 1}. for which holds that
P (((A Rij [d], s.l. ((AA = {y(A)))
In other words, A has a simple spectrum almost surely. The observation vector e E Rd is sampled independently, from distribution that is absolutely continuous w.r.l. the Lebesgue meusure over R.d.
Lemma C.1. Assumption 3.2 ensures the pair (A, c) observable w.p.
Proof. Since A has distinct eigenvalues w.p. 1 (the simple spectrum condition), it is (block) diagonalizable almost surely, and its eigenvectors (vi.... va) are linearly independent. Therefore, observability is ensured if v, + 0 almost surely for all i [d].
Since is sampled from distribution that absolutely continuous w.r.t. the Lebesgue measure in Rdd we want prove that the set
= UcEE Rd c 0}
has zero Lebesgue measure in the ambient Rdd Each collection {ce Rd c v; 0} forms proper subspace of Edd with dimension at most d - (it can be less, for complex V.). Therefore, its Lebesgue measure is null (see, e.g., (Royden & Fitzpatrick, 2010, pg. 435)).
Since 14 is finite union of measure zero sets, it is itself measure zero. Hence, observability holds w.p. 1.
22

---

D SECTION 4 PROOFS
D.1 PRELIMINARIES
Since we're dealing with data generated from stochastic processes, our proofs will heavily rely on taking expectations conditioned randomness up certain point in the process. In what follows, we formalize the natural filtrations with respect to process (1).
We denote the natural filtration associated with (1). (Fe where
Fr (A,, C, æ0, Wo..... Wr-1, 00, ......1 Uz-1),
t20 0.
(24)
By convention, when 0 the of noise variables are empty, and we define
Fo (A, c, zo),
(25)
illustrate that A and are sampled once at time and then remain fixed.
It follows that
(a) F W20
(b) D; is J.-measurable for all 2 0.
(c) Yr Fe+1 meesuurbee (since yr depends on v,)
(d) The noise at time independent on the respective filtration: w, 11 Fr, Us 11. Fr, for all O.
D.2 AUXILIARY RESULTS AND TECHNICAL LEMMATA
Theorem D.1 (Isserlis (1918)). Let y [yn, fy,.yy...-5 yu] (O, E) be an n-dimensional, zero multivariate normal vector. Then, for any even integer n,
D0 y1 ... yn]
Σ
II Elgerr.
PEPP(n) (2,r)Ep
where PP(n) denotes the set of all painwise partitions of [n] into disjoint pairs. If n. is odd, then Elyn12 yo] 0.
Lemma D.1. Given random vectors 2, w. q E RRd and assuming that w is independent of 2, q and the relevant integrability conditions hold, then
ww q
E/ww [q]
(26)
Proof. We use the towering property of expectations,
E 2 ww z]
=
TE
unw
2 q] q q]
E
Tww "]
where the last line follows from the quantities' independence.
Lemma D.2. Let the sequence {ys be generated by LDS (1) sampled according Assump- tion 3.2. For time indices < i j. it holds that
E [u. = E A'Em, (A yc] +
(A + 10m (27)
k=D
where 1 {{=5} takes the value ifi = and otherwise.
23

---

Proof. For process (1) it holds that
3-1
33 = A A'''r;
an +
A wu
k=
and therefore
j-1
yj = c A'-'a; +
Σ
c
A + vj.
A
k=
The product of scalars yiU, therefore takes the form yıU, = 3
T
= (c oi + u;)
c AT''''r; a
+
C. A A0 +
+Dj Djj
k=
j-1
= e' a,a0y (A
1-+++++ c +
c
u
(A c + c. Deyj
k=
+ u
T (A T y
M T (A U. c +
U.Uy.
(28)
Now, observing that E[y;U5] = E [ WiUs [F]] and remembering that F;, A, e are F-measurable, and that for all i and p 2 i, up 11 Fi and Up 1I Fs, and wp IL Bqg vp,q q 2 O, we have
E [c
x (A c|
= =c (ATTTAA
j-1
E
xz w (A (A k Fi
Σ
I w
(AT)-I-* c= 0,
k=
E [[TTx,0y] Fi]
= c cJ,,BBuy] = 0,
E
[ (A T g Fil
= Bjuy] x.. (A - = 0.
E w (A T p Fi
Ele; E w (A 0.
E [oeu,
Fi]
= Elu,yy] = 1 16=z)10.
Therefore,
Elyxvs] = EE [[[xv [Fi]] = E [TT,
a (A (AT) cc + 1
(29)
Noting that x; in (29) and get
A A'xo- + Σi= 11-1-k wwk we further unroll the first term inside the expectation
i-1
c. T (AT)i-cc
c. A'xo +
Σ
c
A
1-
We
T (ATye+ + Σ w (A (AT
0=0
R=
c A' T
(AT
c A'a T (A (A
k=0
1-1
1-1 (A Σ - w (A (AT ccc
T A--i1 1 2
k,t=
(30)
24

---

Using up 11 Fu C Fs Vp 2 0 and wy 1T wy: Vp + q in conjunction with (30) and Lemma D.1
we get
E x (A y c = c' A o (A yct
A A 2
Σ
(31)
k=0
Furthermore, noting that oAA, c) Fo. we have that
E [e , (A y-ic A, = c A'Eno (A ye+ Σ A
(A
(32)
k=
Taking full expectation in (32), and plugging everything back into (29), we get the stated result
[ = E Ic A'E (A yc + Σ A
y + 1i=j!
Lemma D.3. {y: 320 sequence of observations generated by an LDS (1) sampled according
to Assumption 3.2. Then,
(a) ifi j 2p + for some pe N4, [Ba 0:
(b) ifi + j k + 2p + for some p E N4. I [yCyuue
(c) ifi + + + + 2p + 1 for some p E N+, E 183 =
Note that there is no condition on the indices being pairwise distinct.
Proof. To prove point (a), we start from the expression derived in Lemma D.2.
[y, E [c A (ATyc]
k=0
Clearly, since + is odd, it holds that * j and hence the third term zero. Furthermore, since A has centrally symmetric distribution, we have that
E [ A'E (A y [c (-A) (- (-A yc (--1)++ [c' A'E,, (A (c].
(33)
implying that E [c A'Y (AT)c) 0. We apply similar reasoning for the other term and obtain that
E [y
thus proving the first point.
For both points (b) and (c), we will rely on Isserlis's theorem, which we replicate in The- orem D.1 for convenience. Note that conditioned, on A and c, the vectors [Н4Н, A, cj and (h A, c] are jointly Gaussian since they are linear transformations of the
jointly Gaussian vectors ri
a WT
k.. Uo.... (max(i.j.k.) and r2
2J wT .... w. Tax (..,..,,,,. No. ....... (maxx
respectively. We can therefore apply the
towering property along with Isserlis's 's result to get
E [yansuam]
E E (::b,и: A, c]]
E
E j::
A, E WaE
A, c] + E [veu A, cl E [y,3 A, c]
A, cl ].
+ E [ave
A, C E Uy
(34)
25

---

since PP( = (((i,)) (k, ((i, (())))))))))) (i))) Since + + k + is odd, the two pairs inside any E p. k, must have different parities (i.e., one even, one odd). W.l.o.g, we analyze the first term in (34), assuming 0
siisi S k From (32), we know that
E [y:0 AA cj E 1A AA
c A' (A yct
Σ
c
1 (A - c+ (i=1))
t=
k-1
C A (A
Σ
c
A
1
(A (A (k=l) 00
(35)
Assume w.l.o.g that + j is even, and k + 4 is odd. This implies that [(k= = 0. Taking full expectation on both sides and developing the product, we get
E E
A,
c) E
c,c
E
A'E
(Ayce A A0 (c]
k
E T A A'Eno (A y'ce A 1k-1-15pp (5-1 (ATT-
+
E
c T A
1-1-1 m
(AT) cc A (AT)'c]
t=
E [c]
A
11-1-s (A (A c A k-1- En (A (A
1=
+ 1
|1=1
E
c
A A (AT)'))
4-1
+ +
E
|
c
A
(A (A
(36)
t=0
Using the index parity assumptions and the reasoning based on the central symmetry of A's a distribu- tion from (33), get that all the terms on the RHS of (36) are zero. We treat the remaining terms in (34) similarly to get the final result in (b).
Finally, point (c) follows a similar path. We have
P, L, m)) = (((i.)), (ks, (m, (4)) (l,m)), {(i,)), (J, (4,, (l. m),)),
{(i,4), (j, k)), (,))))))))((())))))))))))))(.(.) (j,n), (1,)), ((i,m). (j. m), (k.
(k, m)),
((i,n), (j, k), (1, m)).
For the parity hypothesis to be satisfied, not that inside set q E PP k. m, n)), at least one pair must have an odd parity, while the other two must be of the same parity (either even or odd). W.o.l.g let 0 < < j < k < l < m I n. pick the first set in PP( j, k. I, m. n}) above (the rest follow the same logic) and assume that m + n is odd. By the same logic as before, we have that
1 (m=n) = 0 and
E E [У:
A, c) E (aa
A, c] E (ym A, c]1
E
T A' (ATye+
c. A -1-1 (A
t=
k-1
A k Emo
(A (A
1-
En (AT)--1-1 e + 1
c
c
0
k-1
C
TAA "mED (A ()
Σ
c
A
1
E
(A Tm--te
t=0
26

---

Without computing, whose sum is odd. Therefore, using the centrally symmetric property of A's distribution, all the terms evaluate to zero, and point (c) proven.
that every term in the expanded product will have powers of A
D.3 PROOF OF LEMMA 4.1
Lemma 4.1. For any j. E (s), the parameters WOK and Wy having the structure
(
WoK+
Wv (++1))x
12+s mod 2:2[ +
(14)
0 _
with R 1
+
0
*
* and 1 r+11 ensure that LIS (12),,, whenever
r rl 2N and s + E 2Z, or r + lE 2N 1 and 2Z + 1.
Proof. Recall the in-context loss in (10) with general coosStuuudt input token matrix
31
---
Dr-s is defined as VT-1
Yo
Ms+1
Us+2
) -
((0) = E (To (Yo) (Yo)
(37)
For equations (38) to (42) below, we use the same reformulations as Ahn et al. (2023). The last column of the transformer's output above can be written
yT 1
yr
T W.
Σ
y:y: y
Dili+ts Vi
W QK
(38)
Vi++
where the summation comes from the causal mask. Therefore, the transformer's prediction of yr. To (Y0)++17-- can be written as
T-8-1
y:uk y
Dillet vita
[2122 2z.] 1..62...2.. Br]yr-as
(39)
T 8
Vi+s
VERG+4Xx(+4))
where TT E RRx((++1 is the last row of WT and zj E [R(++)) is the jth column of W So the in-context loss L(Wv, Wok) can be rewritten function of TT and Zz [zy]]=1
[T-
2 )
)²}
C(b, E
b YZZZZr---a VT
(40)
T
Plugging in the expression of ÿT- the in-context loss is
2
yT
0r-++1
((b = E
T y
21 21 ...
yr
T
8T-1
27

---

E
Y 2T:/T 148 YTT
k=1
2
E
Tr(Y Tr(Yzub) bb WTT --7---148 yT
E
Y bz /T- - yr
(41)
k=1
We reparametrize the in-context loss using X == bzk
((X =E T
X
- yr
(42)
Note that the gradient of the in-context loss with respect to each X, is
Vx, ((X = 2E
Y. X X2)UT UT -- 14h yr) UT---1458
(43)
k=1
The gradient Vx, xe((ss) is a sum of two terms, Vx, (( replacing Y we have
TX, Tx where,
Tx,
2
E
Y., -1
y:y.
Diliits Vi+
(44)
T
y Vi++
1=1
T-s-1
T
2)
-2E
VT
yi Vit+
Deli+ts
yē+
(45)
Each matrix element of !2) X,
has the form
T-s-1
2E BTVT-
(46)
with E [1, m E s0 and E s).
The sum of y's indices in (46) for each term in the above sum 27 2i + (m - j).
The parity is determined by that of m + s - + j and is independent of the sum counter (i.c.. the same for all terms). According to Lemma D.3, (46) if (m + s-1+ odd, and of
arbitrary value if it is even. So general matrix element of (2)) X, is 0 if (m + n - -- 1 + j) is odd and of arbitrary value if (m
1 j) is even.
28

---

For given AR(s)-type token (s is fixed) and a specific j, whether matrix element of T 0 only depends on m (its position in the matrix). So,
0 * 0
* 0
*
if j 8 is even;
*
0
T
*
0
0
0
0
if [j 1| is odd.
now turn TQ)) X; with the end goal of finding parameter configuration that matches the sparsity
pattern of (2) For this section, assume is odd (the other case follows similarly). First, let
[k)] and unfold the expression of the matrix inner product
X = L 1f=1
Xx)
(k) Uptryptls
(47)
P=0 t=0
p=0
where r, I are the indices traversing Y.
Furthermore, each matrix element of T X, inside the expectation has the form
2 T
Σ
Σ (Y, Xx) yT---1-8 -- Witt Di+ms
(48)
T-
where m E 10, (.. In what follows, we'll use the notion of matrix's 's support, as follows.
8) are the indices traversing Y.
Definition D.1 (Matrix support). Consider matrix B E Rmxn, The support of B is the set of index puirs corresponding its nonzero entries,
sBpp(B) ((,,) [m] [n] B O0.
(49)
Assume that j is fixed and odd (we discuss the even case afterwards). Note that the sparsity of each +)) where, when s, jodd, the respective
(2) position in x, dictated by the parity of (m element is zero whenever + is even. Notice that except for the contribution of the matrix inner product, the sum of indices for the y-factors in (48) is 2(T We distinguish two cases:
1 + i) + k + j + n + m so the
parity is determined by that of l: +
(a) when & is even, + j odd, and we wish that the term zeroes out for even m + n. This means that X must select in (47) only pairs Bptrlpo4 for which r + e is even and zero-out the others. In other words, the support of X should satisfy
xupp(X.) ((r,)) + [s+1] r r+le22
(50)
29

---

Such an X₁ may look like
(R)
11
2 13
1,8
0
(k)
(k)
22
2,++1
(k)
(R)
(k)
431
3 33
X
(51)
(k)
T
#,1
0 s,3
(R)
(R) a+ ++1,++1
1+1,2
with arbitrary (possibly also zero) values for constants 2i. Note that, consequence, the entries corresponding odd m are not forced to zero by the distributional symmetry of A and can generally take arbitrary values.
(b) when k is odd, & + j is even, and we wish that the term zeroes out for even m + n. This means that X₂ must select in (47) only pairs Up+rup++ for which r + is odd and zero-out the others. In other words, the support of Xx should satisfy
supp Xx) S ((r,)) E [s + 11 1 rte 22 + +
(52)
Such an X may look like
112
0
x 1,s+1
(R)
(k)
(k)
T22
2,,
(k)
(b)
32
0
31
X
(53)
D: 8.2
0
D s
(R)
(k)
(k)
I 1s+1,1
+,133
sts
with arbitrary (possibly also zero) values for constants x Note that, as consequence, the not forced to zero by the distributional symmetry of
entries corresponding odd A and can generally take arbitrary values.
These patterns need to be coherent with the case of j-even. Note that in T X,, when s-odd, j-even, the respective element is zero whenever m + n is odd. We again distinguish two cases:
when & is even, + j even, and we wish that the term zeroes for odd m + n. This means that X3 must select in (47) only pairs Up+rop++ for which + I is even and zero-out the others. Notice that the pattern of X₁ in (51) for even satisfies this requirement and we have coherence.
(b) when & is odd, means that X₁ must select in (47) only pairs Up+rup++ for which r + is odd and zero-out the others. Notice that the pattern of X₁ in (53) for odd is satisfies this requirement and have coherence.
is odd, and wish that the term zeroes out for odd m + This
The same approach goes through for even window size To sum up, weights satisfying equations (50) and (52) for the appropriate parity of triplet (s,j, k) ensure the symmetry-indee -nduuced sparsity pattern complies with that of (2)
X,
Finally, recall that X := bz. It can easily shown that vectors b and 2 yielding an X₃ whose support has maximum cardinality are those whose support consists of indices of opposing parities (e.g., even indices for and odd indices for 2t. or vice versa) in the case of (52) or the same parities in the case of (50). For case of odd window sizes 8, the sparsity pattern of b and 2, yielding
30

---

complying X is
[0,2) ,...,,, 0,,,1 if k is even
(k)
(k)
b
1/
(54)
2,,0,00 ,, .0,..... 20 .0] if is odd 21
(M)
bst1
For even window size the patterns are
0
(k) .0] if k is even 0, ..... 0,
(K)
62
b
2/
(55)
...................... 20) 20........2 0, 2,, 20)
if odd
bs+1
Arranging these vectors inside WoK and Wv gives the stated result.
D.4 PROOF OF THEOREM 4.1
Theorem 4.1. Let Yo encode the input tokens for 1. Then, the optimal parameters 0* (Wōk: Wt) of single linear self-attention layer with respect to loss [(0) are
(T-2)E
1MT ET-2 M
0 0 W
W
(15)
QK
0 1
up to rescaling with nonzero constant.
Proof. For the transformer parameters in (15), the corresponding b = and the corresponding
(f-2)
pr- (Vr- ΣT wmr+1]
F F= (c 0j. where
Ep Mf (E 0:844:)
So X = bff bFT
The gradient the in-context loss Vx Vx((X) is
2
T T
ED [(Y, X X)v?
T-2
2
ED
T-2 v²
yr+1Vr
yr
Vr+1
0
C
0
0
//T
2
T
2
1i+13/
1=1
2
ED
vi
yey/+11
T --2
1/+1//
Vi+1
2
ED
T T-2
11--11111111
yeli+1
(56)
T 2
0
According to Lemma D.3, the two diagonal elements in (56) Ep
IT=
T-2
and E, E v-
T-2 vi
are O, since the sums of y indices are both odd.
T-2
T x,
2
VT
Y
T-
Vr -2
Σ
9il6+11
1i++19
31

---

T-
= D
UT UT
0
Mili+1
(57)
Uit1ys
According to Lemma D.3, the two diagonal elements in (57) ED yryr-1 Ei= v/
and
FD E 7- v are O, since the sums of y indices are both odd.
Plugging in the expression of c, it can be casily found that
Vx = Tx, + +T', =0
(58)
Since the in-context loss is convex in X and the X resulting from the A and Wa above makes V ((X) = 0, the W and Wyx above is a global minimizer for the in-context loss.
32

---

E PROOFS FOR SECTION 5
E. PROOF THAT OUR EXPERIMENTS' SAMPLING SCHEMES OBEY ASSUMPTION 3.2
All our experiments use sampling schemes whose generalization the following:
(a) constructed by sampling v ~ P, where P is centrally symmetric and absolutely con- tinuous w.r.t the Lebesgue measure R with marginals supported on [-1, 1), and independently sampling P, whose every entry is drawn i.i.d. from any absolutely continuous distribution w.r.t. Lebesgue measure in R. Matrix A is then formed as Pdiag(u)P
e sampled from an absolutely continuous distribution w.r.t. Lebesgue measure in R or otherwise fixed with e * Og-
We need show that
(a) A's distribution is centrally symmetrie, i.e., that A _ -A;
(b) A's 's spectrum is simple w.p. 1;
observability still holds when is fixed according to the above condition.
The first point achieved since, by the central symmetry of v's distribution,
-A = -i-gddda dia 4 - diag(v) A.
(59)
The second point is ensured by v's distribution being absolutely continuous w.r.t. the Lebesgue measure Rd, and hence the probability v belonging to (d 1)-dimensional subspaces (and lower) such as { e Rd j E [d] s.t. x; = zy} } is null. In conjunction with the above, when sample e from a continuous distribution in Rdd Assumption (3.2) is satisfied.
However, our proofs and experiments go through even if c is fixed, as follows. First, the theoretical results rest on A's distributional symmetry and are invariant to the linear transformation induced by Second, observability is ensured since det (O) in expression (5) is not w.p. as follows.
We use (OP) * 0 w.p.
de * w.p. 1, since det P) * w.p.
dct(OP)
det([z; diag( v) z; diag( d--12)) 21)
(60)
det (zibagg())dett
U2
0/
(61)
Ud
Since P's entries are drawn i.i.d. from an absolutely continuous distribution w.r.t. Lebesgue measure in R, it holds that * 0 w.p. 1. The remaining matrix is Vandermonde with e; * Uj, Vi,j [d] w.p. Hence, the determinant nonzero w.p. and observability holds almost surely.
E.2 RELATION OF TRANSFORMER FORWARD PASS WITH PCG
For convenience, reproduce below the PCG iteration of Shewchuk al. (1994) for minimizing objective
1 w T Aw + b w + fww We compute the first two steps of the algorithm with respect to the loss (4), which can be rewritten as
2
T
(w) = 2(T -
(ye4s w T )2
(62)
33

---

Algorithm I Preconditioned Conjugate Gradient
Input: preconditioner H, wo, ro Auro do H -2- tro knew ro do
2: for i = U......d do
3: 2; Ad,
4:
6ym d. z:
5:
wi+1 w) + o,
6:
Ti-1 ri 0;22
7:
u+11 H r
8:
8old Önew, dnew I 1 1441
9:
B++1 8ad
Suc.
10: d;+1 Vi+1 + Bs+1d;
11: end for
12: return fT
T-8-1
1
un J. w 20e+s T D.. + v2..
(63)
2(T 8 1-1)
1)
t=1
2uh
[22RAR(e) w w (LR(()) (0) + 12+t
(64)
Two iterations of Algorithm 1 starting from wa = 0 and using H P-1 yield the following predictor.
wt andg
Hn 'ro
U22
wi + aidd
agdo 01 [Pri [,
on (0) P(ro- PV
.n PVCAR.)))) + ((VCAR)s) (0) ao V L A + B, PV (AR(s
00 PVL (0) + [PVVLAR() an V CAR()) PVLAR() + 31 PVL PVLAR())
(00 + 0n + 31 PV - an PV (A
34


Reviews:
[{"review_id": "o7731fKa1J", "reviewer": "Reviewer_JMPY", "summary": "This paper investigates in-context learning (ICL) for non iid sequential data, addressing a critical gap in the theoretical understanding of transformers. This paper takes the first step toward sequential data by studying Linear Dynamical Systems (LDS). The authors prove that for order-one autoregressive (AR(1)) approximations, the optimal linear attention layer implements exactly one step of Gradient Descent on the context-dependent least-squares loss. The key insight is that AR(s) models can closely approximate Kalman Filter predictions when the window size s is sufficiently large. The authors identify structural patterns for higher-order cases and provide experimental validation. The connection to language modeling is still distant, but this result represents the first optimality result for ICL with non-i.i.d. data extending the optimization framework from independence to temporal dependence.", "questions": "You state  that AR(s≥2) optimal weights do not \nimplement gradient descent. However, you observe an intriguing \nrecursive pattern  where AR(s+1) structure embeds \nforward-shifted AR(s) structure. Can you clarify your current undestanding ?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The paper is well written, well organized.I am far from being a theoretical specialist in the field, but the problem is very clearly understandable, as are the SOTA and the paper's positioning in the literature.\n* Extends ICL theory to noise-corrupted dynamical systems with non iid data, more realistic scenario.\n*  Provides Mechanistic Explanation for Transformer-Kalman Filter connection\n* Solid theoretical contribution for AR(1) with Noise (theorem 4.1)\n* structural insight for higher orders (banded structures etc).\n* limitations are well discussed.", "weaknesses": "* in my opinion, some assumptions are highly restrictive (diagonal A for LDS, independance of the evolution of each dimension), with no clear insights how to generalize beyond the assumption.\n* The  proof is only for AR(1), the results for s>=1 are only empirical.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "dcygjrIkS2", "reviewer": "Reviewer_2rfD", "summary": "This paper investigates the theoretical behavior of linear attention layers when provided sequences generated by linear dynamical systems (LDS), aiming to understand how transformers can perform in-context learning under temporally correlated data. It was shown that for first-order autoregressive processes, the optimal weights of a linear attention layer perform a single gradient descent step on the least-squares estimation objective, essentially mirroring Kalman filtering updates. It is further shown which structured pattern the learned weights exhibit for higher-order AR models. Experiments on synthetic LDS data support these claims, suggesting that transformers can emulate filtering-like inference dynamics without explicit supervision.", "questions": "It is stated in the paper that Assumption 3.2 can be \"relaxed to symmetry at the expense of added complexity in the later data sampling step.\" This should be elaborated upon: what would be the path forward for more general A and c?\n\nThe core analytical result is for AR(1) / first-order LDS. Can you clarify precisely which parts of the argument fail for AR(p) with p > 1? Is it a technical difficulty (algebra gets intractable) or a fundamental obstacle (the correspondence to a single GD step is no longer true)? This matters for how general the claimed mechanism actually is.\n\nIs the linear attention update / Kalman-style filtering analogy exact (i.e., the attention layer implements the same update under the assumptions), or qualitative (as in \"both reduce prediction error using recent context\")? It would help if this is stated explicitly, i.e., if  exact equivalence under stated assumptions is precisely established... otherwise, one might get an impression that attention simply learns to mimic the effect of an adaptive linear predictor under LDS-generated data. Either way, precision would be beneficial.\n\nThe Discussion section makes several statements about prior work that appear overstated. For instance, lines 358–359 claim this is the only work addressing noisy LDSs apart from Cole et al. '25, yet Goel and Bartlett '24 (already cited) analyze a similar and more general LDS model without assuming Assumption 3.2. Likewise, the statement in lines 368–371 that prior KF-emulating transformers require complete knowledge of system parameters seems inconsistent with published claims of robustness to partial parameter availability in that approach.\n\nWhile the focus of the paper is on LDS + Gaussian inputs, an insight into what happens if the noise is not Gaussian but the system is still linear would be appreciated. Likewise, even a brief discussion on potential directions and/or expectation in non-linear + non-Gaussian case would be helpful.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a clear theoretical analysis linking linear attention updates to gradient descent steps in linear dynamical systems. The extension of the understanding of in-context learning ICL beyond i.i.d. data to temporally correlated, non-i.i.d. sequences is conceptually \nsignificant.", "weaknesses": "The scope is somewhat narrow: the main result holds only for AR(1) processes, with higher-order cases supported empirically rather than analytically, and the experimental validation remains limited to small synthetic settings.\n\nAssumption 3.2 (lines 207-208) appears rather limiting -- diagonal A, combined with the assumptions on the noise covariance matrix stated as part of Assumption 3.1, implies that the dynamics of the components of the state vector in (1) are completely separable. The assumption c = 1_d (part of Assumption 3.2) is also limiting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "miyrqFao7X", "reviewer": "Reviewer_jZgQ", "summary": "The paper investigates the expressive power of a single linear attention layer for in-context learning of Linear Dynamical Systems (LDS) using non-i.i.d. data. The main result demonstrates that, for a first-order autoregressive approximation, the optimal linear attention layer effectively performs one step of gradient descent (GD) on the corresponding least-squares loss.", "questions": "- The optimality result in Theorem 4.1 is restricted to the order-one autoregressive approximation of the Linear Dynamical System. Do the authors expect their findings to extend to higher-order cases, and if so, under what assumptions?\n\n- How sensitive are your results to relaxing Assumption 3.2, for example in systems with coupled (non-diagonal) dynamics or correlated, non-Gaussian noise?\n\n- Could the authors elaborate on how the key hyperparameters were selected and discuss how sensitive the empirical results are to changes in these parameters?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The main strength of the paper lies in its consideration of the challenging non-i.i.d. setting and its rigorous theoretical contribution. The authors formally prove that, for an order-one autoregressive approximation of a Linear Dynamical System, the optimal single linear attention layer corresponds to one step of gradient descent on the associated least-squares loss. Overall, the theoretical analysis is well-motivated and provides an interesting and valuable contribution to the understanding of the expressive power of linear attention mechanisms.", "weaknesses": "One of the main weaknesses of the paper lies in its presentation, which is at times convoluted and difficult to follow. Furthermore, since the paper’s primary contribution is theoretical, placing the key proofs in the appendix somewhat weakens its overall clarity and impact. Regarding the theoretical results, an additional limitation that reduces the paper’s significance is that the optimality result (Theorem 4.1), which demonstrates that the transformer performs one step of Gradient Descent on the least-squares loss, is restricted to the order-one autoregressive approximation of the Linear Dynamical System. Another substantial limitation lies in Assumption 3.2, which imposes a highly simplified structure on the linear dynamical system by restricting the matrix A to be diagonal with distinct, nonzero entries. While this choice facilitates analytical tractability, it eliminates interactions between state variables and therefore limits the applicability of the results to realistic systems exhibiting coupled dynamics. Moreover, the assumption that both the process and observation noise are Gaussian and isotropic further simplifies the analysis but fails to capture the correlated or non-Gaussian noise commonly encountered in real-world scenarios. The authors should also consider citing and discussing related recent work, such as Zhang, Yedi, et al. “Training dynamics of in-context learning in linear attention.” arXiv preprint arXiv:2501.16265 (2025), which similarly analyzes the expressive power and mechanistic basis of ICL within minimal transformer architectures.\nRegarding the empirical study, the rationale behind the proposed experimental setup is not entirely convincing. It remains unclear how several hyperparameters (such as T, d, and the number of iterations) were selected and how their choice may have influenced the results. For instance, the noise margin, which appears to play a crucial role, is fixed, yet no explanation is provided for how its value was determined. The authors should clarify the sensitivity of the analysis and the observed recursive patterns to these parameter choices.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "MMczhzSK4V", "reviewer": "Reviewer_hQkJ", "summary": "The paper investigates how a single linear self-attention layer can perform in-context learning (ICL) on data generated by linear dynamical systems (LDS)—a non-i.i.d. setting relevant to time-series and control problems. The authors derive theoretical results showing that the optimal linear-attention weights correspond to a gradient descent (GD) step on an autoregressive loss for first-order systems, extending prior work on ICL beyond the i.i.d. regime. They also identify structural patterns in optimal parameters for higher-order autoregressive models and support their findings with numerical experiments.", "questions": "N/A", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- This paper povides one of the first formal analyses of transformers’ in-context learning behavior on non-i.i.d. time-dependent data (LDS), bridging ICL theory with system identification literature.\n\n- Experimental results confirm the theoretical findings and reveal recurring structural patterns in optimal weights for higher-order autoregressive tasks.", "weaknesses": "- The assumptions in this paper suggest that the non-i.i.d setting seems to be overly simplified, each sample seems to be just a sum from previous ones after some scalings. This setting leads to a significant mismatch with real-life data. Addition to this, a single linear attention layer is studied. This model is also entirely far from a transformer often used in practice. Those facts lead to concern about the significance of their result to understand transformer's expressive power.\n- The optimality proof is limited to first-order systems; higher-order cases are only supported empirically without a full theoretical guarantee.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "0FYO0EWoTt", "reviewer": "Reviewer_yeiD", "summary": "The paper addresses the task of learning dynamical models in-context with a linear transformer. The authors show how the transformer can be optimised to fit a linear, size-one, autoregressive model in-context. Under certain assumptions, they compute the attention parameters that guarantee one in-context step of gradient descent of a Least Squares objective.", "questions": "- Why is the Gaussianity assumption necessary? Wouldn't requiring the noise to have a zero mean be enough?\n- Once data matrices are formed, the problem reduces to standard linear regression. Is the extra challenge the sum over $k$ in Equation 9? Aren't $Y_t|Y_{t-1}$, $t=1, \\dots, T$ i.i.d. when the model is stationary?\n- What do you mean by *\"temporal scaffold closer in nature to that of language-induced tokens, in stark contrast to the i.i.d.\"*?\n- What happens if the state-transition matrix is not diagonal? Would $W_{QK}$ and $W_V$ still have the structure shown in Lemma 4.1? \n- Does the idea generalise to i) multi-step GD and ii) larger sliding windows?\n- In the discussion, you say that the model does not implement GD for $s\\geq 2$. Does this falsify the theoretical equivalence claimed in the abstract? Is the *\"recursive pattern of the optimal weights that builds on top of the GD inducing parameters recovered for AR(1)*\" predicted by the theory?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Learning a dynamical model in context is challenging and new. \n- The state-of-the-art is well-presented, in the introduction and throughout the presentation.", "weaknesses": "- The performance of linear self-attention models does not match that of more expensive LLMs. \n- After the reduction to improper learning, the problem resembles linear regression. The authors should comment on the specific challenges coming from the non-i.i.d. setup.\n- The assumptions on the diagonality of the state transition matrix seem strong. The authors comment on the effects of relaxing it for the data-generating process. Is the assumption also crucial for Lemma and Theorem 4.1? \n- The authors write *\"However, a quick calculation of the forward pass reveals that weights trained to optimality with AR(s) tokens (7) for $s \\geq 2$ do not implement GD in the forward pass\"*. Does it mean the suggested procedure cannot be verified empirically? \n- The experimental results focus on the attention parameters. Except for the *quick calculation of the forward pas* mentioned above, the authors do not provide any evidence that the learned model performs gradient descent in context, which, to my understanding, is the paper's main claim.\n- The length of the sequences used in the experiments is 30. Is the same length used for in-context testing?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": {}}]

Author rebuttal:

Decision: Accept (Poster)