Optimality and NP-Hardness of Transformers in Learning  Markovian Dynamical Functions
Authors: Yanna Ding
Songtao Lu
Yingdong Lu
Tomasz J Nowicki
Jianxi Gao
Venue: NeurIPS 2025 poster ()
Keywords: In-Context Learning
Markovian Dynamics
Transformers
Optimization
Primary area: theory
TL;DR: We study the expression and limitations of transformers in performing in-context learning of functions governed by Markovian dynamics.

Abstract: Transformer architectures can solve unseen tasks based on input-output pairs in a given prompt due to in-context learning (ICL). Existing theoretical studies on ICL have mainly focused on linear regression tasks, often with i.i.d. inputs. To understand how transformers express in-context learning when modeling dynamics-driven functions, we investigate Markovian function learning through a structured ICL setup, where we characterize the loss landscape to reveal underlying optimization behaviors. Specifically, we (1)   provide the closed-form expression of the global minimizer (in an enlarged parameter space) for a single-layer linear self-attention (LSA) model; (2)  prove that recovering transformer parameters that realize the optimal solution is NP-hard in general, revealing a fundamental limitation of one-layer LSA in representing structured dynamical functions; and (3)  supply a novel interpretation of a multilayer LSA as performing preconditioned gradient descent to optimize multiple objectives beyond the square loss. These theoretical results are numerically validated using simplified transformers.

Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions
Yanna Ding' Songtao Lu² Yingdong Lu², Tomasz Nowicki', Jianxi Gao Department of Computer Science, Rensselaer Polytechnic Institute Department of Computer Science and Engineering, The Chinese University of Hong Kong IBM Research
{dingy6, gaarv.gaaajarrpppp edu, stlu@cse. cuhk. edu. hk, (yingdong, (oowicki)auu ibm. com
Abstract
Transformer architectures solve unseen tasks based on input-output pairs in a given prompt due to in-context learning (ICL). Existing theoretical studies ICL have mainly focused linear regression tasks, often with i.i.d. inputs. To understand how transformers express ICL when modeling dynamics-driven functions, we investigate Markovian function learning through a structured ICL setup, where we characterize the loss landscape to reveal underlying optimization behaviors. Specifically, we (1) provide the closed-form expression of the global minimizer (in an enlarged parameter space) for a single-layer linear self-attention (LSA) model: (2) prove that recovering transformer parameters that realize the optimal solution is NP-hard in general, revealing fundamental limitation of one- layer LSA representing structured dynamical functions; and (3) supply a novel interpretation of multilayer LSA as performing preconditioned gradient descent to optimize multiple objectives beyond the square loss. These theoretical results are numerically validated using simplified transformers.
Introduction
Transformer-based language models have demonstrated remarkable in-context learning (ICL) capa- bilities, predicting outputs for unseen inputs using only examples provided in the prompt, without parameter updates [ 2 5 7). This phenomenon has motivated growing body of theo- retical work aiming to understand the mechanisms underlying ICL. Much of this work focuses regression tasks with i.i.d. Gaussian inputs [8]
20 (1],, showing that transformers can emulate classical algorithms like gradient descent. Others have transformer expressivity. especially under structured in- puts [22] 23 24] [25] 25 26 27 28 29], revealing that even moderate-depth transformers struggle
[2 13
14
[5]
16 17
begun to explore the limitations
with certain algorithmic tasks.
A complementary learning scenario involves predicting dynamical functions from temporally struc- tured data, such as those governed by Markovian processes. Recent studies have made progress in understanding how transformers perform next-token prediction in structured englecsequence settings,
including Markov chains [24] 22 23), autoregressive models [30], and causal sequences [28] 29]. In
contrast, we focus on different task formulation: learning latent transition dynamics from multiple in-context sequences, where the model must integrate structural information across examples to generalize to new sequence.
Yanna Ding, Songtao Lu, Yingdong Lu, Tomasz Nowicki contributed equally to this work.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

In this work, we study ICL for Markovian dynamical functions, where each in-context example is trajectory sampled from a discrete-time Markov chain, and the task is predict the next token of a query sequence by leveraging shared transition structure across the examples. This setting introduces new challenges that go beyond classical ICI. tasks with i.i.d. inputs. In the standard regression case, each input-label pair corresponds to single token with simple statistical structure, typically zero-mean and independent. In contrast, each in-context example in our setting is a full sequence sampled from Markov chain. While a convex reparameterization enables us to characterize the global optimum in an enlarged parameter space, mapping it back to transformer parameters highly nontrivial due to dense parameter interactions induced by the underlying dynamics. These challenges call for a deeper understanding of how transformers generalize from dynamics-driit driven functions in context.
Our contributions. To this end, we study how transformers learn Markovian dynamical functions in-context through the lens of optimization. Given the challenges posed by non-convexity and stochasticity, focus binary Markov chains with first-order memory our first step. which are classical model for statistical language modeling [3 32 22). The major contributions of this work are highlighted as follows.
We establish an analytical framework for understanding ICL of Markovian dynamical functions, and characterize the global minimum of the loss landscape for 1-layer LSA under a tractable case of length-2 chains with both independent and correlated initial conditions. This result reveals how the optimal solution adapts to the Markovian dynamics, exhibiting denser structure compared to the i.i.d. linear regression case (Theorems 3.3
Going beyond the length-2 case, we identity fundamental gap between the global optimum of a structured prediction objective and what a 1-layer LSA can realize for arbitrary-lengtt Markov chains. Even when the optimal solution is well-defined and analytically characterized in an extended space, recovering the corresponding transformer parameters is NP-hard. This realizability gap reveals structural limitation of 1-layer LSA when learning Markovian dynamical functions in-context (Theorem 3.5).
Finally, we advance the understanding of multilayer transformer expressivity by exploring parameter subspace that mirrors the structure of the derived global minimum for Markovian dynamics. Our results show that the forward pass of the multilayer linear transformers is equivalent to solving a multi-objective optimization problem. This problem minimizes a squared loss while simultaneously maximizing multiple linear objectives (Theorem
Related work. The capability of transformers to perform ICL
33
4 3
has inspired extensive
work aiming to uncover the underlying mechanisms 34
10
35
36
37
38 39
21 30 42).
These studies can be broadly categorized into two groups based on the nature of the in-context task.
(1) Regression tasks. This line of work typically focuses on linear regression, where each in-context token represents pair consisting of an input i.i.d. sampled from Gaussian distribution and its corresponding output generated by a wround-truth function. From the expressiveness perspective, [,1 demonstrate that transformers trained such prompts can implicitly perform classical algorithms such as gradient descent, ridge regression, and algorithm selection. From the optimization theory angle, 20 [9] analyze the loss landscape of trained LSA. networks and show that they emulate preconditioned gradient descent.
43 extends this by proving that the global minimizer corresponds
to multi-step preconditioned gradient descent in looped transformer architectures. From the viewpoint of training dynamics. [15 16 prove convergence to the global optimum under mild distribution shift by leveraging the Polyak-L.ojasewicc (PL.) condition.
Departing from the standard i.i.d. setup, our work considers onput-output samples as realizations from Markov chain governed by shared kernel. We aim characterize how transformers behave when the input distribution exhibits structured temporal dependencics and the prediction task is governed by the same latent dynamics.
(2) Sequence generation. Another line of work investigates transformer behavior in sequence generation tasks, particularly under structured or Markovian data. [22] analyzes transformers trained on first-order Markov data and shows that shallow models often converge to oversimplified solutions that ignore sequential dependencies, while deeper architectures can capture bigram transitions. [ further demonstrates that the convergence hehavior in this setting is highly sensitive to initialization and its alignment with the data structure.
124 25 study transformers trained on higher-order

---

Related Work
This Work
31,1
11,2
11,,d
3/1
#1
22
Zd
Jd
Compared ICL for Structured Sequence
32211
21,2
...
Fi,d
y/
Attention Flow Unknown
Sa-1a Jnoia
Smold
Un+1
Compared to ICL for Linear Regression
a.d. N(O,E) 3;
1 ~ Markov (T,, Proo
(0,
Inilial Prob. Transl ion Prob.
y/
x;
Figure I: Comparison between sequence-level in-context Markovian data based attention struc- tures and the existing works. Top: Compared ICL for structured sequences, the key difference that the exiting studies of the attention mechanism [22] 30 24 28) adopted on a token-level, whereas our study studies sequence-level attention. Bottom: While prior work for linear regression samples in-context input and task vectors independently from some given Gaussian distribution
we consider input vectors generated through Markovian transition kernel with parameters Pol, P10 from given initial distributions.
19 15],
Markov sequences, with a focus on how model depth and tokenization affect the ability to model such processes, respectively. While setup also involves fixed-length substrings input units, similar to (25). we do not study the role of tokenization. Instead, we focus on how attention organizes information across correlated examples to enable prediction for unseen inputs. On the other hand, several works [26] 27] have shown that transformers can learn to implement the induction head mechanism [4 45 from an empirical perspective. Additionally, (28) proves that a two-
layer transformer trained with a sequential layer-wise scheme can recover the empirical transition probabilities of single-parent Markov chains. [29 shows that gradient descent n-gram data yields layered specialization: the first attention layer copies parent tokens, the feedforward network selects relevant features, and the final attention layer predicts the next token. Although our study also involves structured sequences, our formulation aligns more closely with regression-style ICL tasks: each token represents full training sample, and the model learns interpolate from these structured examples to predict a new query.
Preliminaries
2.1 In-Context Learning
ICL refers to the operation prompt consisting of input-output pairs and query input p ([(x;, Tx+1) where y; = h((a)), Vi [n + for some unknown function h H, and æs, y: belong to some input space A' and output space )/ respectively. ICL aims to form an output th+1 for the query input Fn+1 that approximates its true label in+1 R3 hi(2++1). The function remains the same within single prompt yet varies across prompts.
y
Prior works have focused on linear function space H: h(x) = y = w for some w X. Under such a construction, is deterministic once x is provided. Despite being commonly encountered in many
applications, the case where stochastic remains largely unexplored. For example, h can represent text generation mechanism that provides descriptions revolving given topic. Then the token generated in the next step is associated with probability based on previously generated words [46]. To approach the ICL for such scenarios, we consider simplified setting of next token prediction for Markov chains. The state space resembles vocabulary and the transition probability akin to the conditional probability of the next word given the previous text.

---

2.2 Markov Chains
The evolution of a Markov chain s of order l on state space S depends solely on the k most recent states. For time step т E Z21, we let s, denote the rth state in the sequence 8, the probability of observing state j E S at time step + is: P(5 = $1:T) = P( where STim,, denotes the subsequence from time step to 72. For first-order Markov chains, the dynamics are determined by the transition probabilities Pij = [8++1 = 8 = 1), which indicate the probability of transitioning from state E S to state E S. These probabilities constitute the Markov kernel P = (piz) [0. |] For a binary state space S = {0. 1), the transition matrix is represented as P(Por. P1o) = [1 P01. Poi P10. 1- Pro]- Let n, E [0. 1||s| denote the marginal probability at the rth time step. binary Markov chain 8 P(po1 P1o)) can be generated
sT-k+1:7) )
by starting with an initial distribution I and iteratively applying P(P01: P1o) to update the state
probabilities at each time step.
2.3 Data Formalism
We introduce the input embedding matrix formulation used for our theoretical results. For a Markov chain 8 with length d 1, we take its first d states to be the input x 81:d and the final state to be the output St. The input and output space are X' = St and y = S. We use subscripts to denote the indices of in-context samples, such that x, represents the first time steps the ith in-context Markov chain, while ys denotes its final state. To form an input embedding matrix Zo E g(d+1)x(n-))
stack (æi, 9:) E gd+1 as the first n columns and let the last column be 1I 1,,0) 0), inspired by [15
7
x1
32
In In+1 0
where 2; = [2,; ~ (T1, Pao)) for initial probability mass function Ti = - p. pl with
p E (0, 1) and transition probabilities Pot: P10 ~ U0, 1). We use double subscripts (zx,,) to indicate the jthh entry of the ith in-context input sequence. The Markov kernel varies for each prompt, while the initial probability p remains constant across all prompts. Let TF denote ransformer-based autoregressive model. The goal of ICL is to learn a model TF that can accurately predict the label of the query input: - TF(Zo) R Va+1-
2.4 Model and Training Objective
We consider transformers with LSA layers 1119. We recall a single-head self-attention layer 47 parameterized by key, queue, value weight matrices are defined as follows:
Attnw a. (Z) = W Z.ZM- softmax (Z WT Wgz)
(2)
M =
I Laxx
0
E TR (+-
(3)
where W., Wg. W, E ((d++))(d+11 are the (key, queue, value) weight matrices and Inx denotes the identity matrix. The attention scores are normalized by the softmax operator. The mask matrix M reflects the asymmetric prompt due to the absence of the label for (=++1), Motivated by [19 15), we simplify the architecture by (i) removing the softmax nonlinearity and (ii) reorganizing the weights as P == W. and Q == Wx W,, merging the query and key matrices into a single matrix:
Attn (1a) (Z) PZM(Z QZ).
Despite its simplicity, LSA demonstrates ICL capability for linear functions 15 and has been shown implement gradient descent [8] and preconditioned gradient descent 19 to solve linear regression in-context. We will prove in Sec. that certain parameter configuration implements preconditioned gradient descent for multi-objecttvv optimization problem that includes linear regression. Finally, we consider architecture consists of L-layer LSA modules. Let Zi denote the output of the ith layer attention, we have
1 (lin)
Z141 =2, + Attn P,Q, (Zi)
(5)
n1
1
= 2i + P.ZIM (Z, Q,Zi)
(LSA)
n

---

for 1 = .....,L 1. The normalizing factor n averages the attention weights gathered from the in-context examples. We consider the output of the transformer be the bottom-right entry of the Ith layer, i.e., TF Pt Q:)1=.......-11 optimize the following population loss in the limit of infinite number of training prompts such that each prompt corresponds to distinct Markov kernel (pt hijes:
[zr.] (d+1).xx+1)- To train the in-context learner, we
((P),Q0)) = (TF Pi Q1)) - Mn++)
(6)
where E fo....,L -- 1} and the expectation taken over Zu. (ptj hi,es. This objective function formulates the in-context task as last-token prediction for in-context sequences. Once trained. the model can autoregressively predict the next token of the query sequence. Unlike ICL for linear regression tasks [19 [5]. where the in-context input distribution assumes zero-mean Gaussian and the input-output relationship is linear, our setting involves a Markovian input and a stochastic input-output relationship. Furthermore, compared ICL for single sequence (28 24 attention is applied across sequences rather than being restricted to local tokens within a single sequence.
3 Global Minimizers and Expressive Limits of 1-Layer LSA
In this section, analyze the loss landscape of the in-context objective function f. To address its nonconvexity, we introduce a reparameterized objective in an expanded parameter space, resulting in strictly convex formulation with closed- global minimum. This analysis shows that achieving optimality in f requires more parameters than in linear tasks.
Parameter space. For single-layer LSA, only the last row of Po and the first d columns of Qo affect the output. Therefore, we consider optimization over the following subset Po and Qo:
Po = (dx((+1))
Qo = [A 0a+1]
where b E TR++1, A E R(d+1)xd Throughout this section, we assume that Po and Qo follow the above format and refer to them as P and Q for simplicity.
Reparameterized objective. We define a reparameteriization 6 which maps from LSA parameter space to 1dd where m (d+2)ld-1).
2
x, = = A), = [bAAk.j b. + A
i k; otherwise
(8)
where 15-11mm- -(d +(d+1)++ k 1). Here d(-), is the rth entry of the resulting vector in Rm and Ak.j denotes the (k:, j)-th entry of and b, denotes the ith element of b. For clarity, we use X to represent @(b, A). Let f Rddm - R denote the reparameterized objective s.t. f(o(b, A)) Q).
We collect the unique elements in the symmetric data matrix EL 2. into a vector g. Then can be expressed as square loss of a linear model parameterized by X:
j(x) = E ((#m+) @ Tg x Un+1
(9)
The equivalence of the objective function before and after reparameterization is verified in Ap- pendix C.1 C.3 The reparameterized objective, f(X), exhibits the following desired property:
Lemma 3.1 (Strict convexity). Suppose the initial probability of the Markov chains is ग = [1 [1-p.p)
with p (0, 1) and the transition probabilities are sampled from (0,1) Then f (Eg. 9 is strictly
convex w.r.t. X E gdm.
The proof, provided Appendix C.2 leverages the nonzero transition probabilities and the properties of Markov chains to establish the positive definiteness of the Hessian of f. Consequently, f is strictly convex, ensuring the existence and uniqueness of its global minimizer, denoted as X*. We derive its expression by solving for the zero of the gradient of ff below.
5

---

Lemma 3.2 (Global minimum for reparameterized objective). Consider the in-context learning of length-d +1 (d 1) Markov chains {(2,, uk) (E E {0.1}) with transition kemel
P
Poo
Poi
(0,1)². Suppose the initial states I2 are i..dd sampled from Bernoulli(p) for
P1o P11
some constani E (0.1). Consider indices i,j E [d] J W, E [d + 1] with j' ''','' S V'. We
denote t $ 12 S t3 $ t1 as the sorted version of J W, 1'). Define H E Rdmxdmm as
Hn.
E
(p(P")
111 p)((p))(1)))) )o1) " 12 (n (p)4-
R [D(P)' 111 + (1 - p) )01) (p'-")) (p(p) 11 (1 -p)(pt'-1)))) 1001) (P"-R")) )as
7
where r (i 1)m + j Er=z d+ + 1 - T, c= (j 1)m + 1' I + -. Define
b E dm as
b. =E [(M(P- "11 (1 (PP )oi) (pati-)))1
(p(P"-)) (1 pp ))))) (p' (p]]
forr (j- (j-1)m++ j - d+ The global minimum described in Eq. eequals X H-
X T of the objective functionn
The full derivation, given in Appendix
utilizes the first-order dependence in the in-context sequences evaluate the expected value of each token. Specifically, the expectation of each token is expressed as the probability of the first token multiplied by successive powers of the transition kernel.
Global minimum of the ICL objective. We now characterize the global minimum the original ICL loss. It suffices to find (b, A) such that @(b, A) = X However, since & maps into an expanded parameter space, an inverse mapping is not always guaranteed. We derive an analytic solution where possible and provide approximation for the more general case.
The following result presents the analytic solution in length-2 Markov chains.
Theorem 3.3 (Global minima for i.i.d. in-context initial states). Consider the in-context learning of length-2 Markov chains {(I, B= (I4, ys E {0, 1}) with transition probabilities P01. ~ (0, 1). Suppose the initial states 24 are i.i.d. sampled from Bernoullip)) for some p E (0,
Let X* == -H--H1 [D2/2 p²/3 /²/12 + p/4] T, where H is a symmetric matrix defined as follows (repeating entries in the lower half triangle are omitted)
2
(n
1)p
P.
(n- 1)p?
g
7
71
2n
2n
p
P.
+
(n (n-1)p²
P
+
(n
1)
P
p
2n
3n
2n
11
(n
1)
1
+
2n
Then the following choice of parameters
0
Xi X1
P=
XI 4x7 X
X
Q=
X
X 4x
(10)
1
Xi
2
is global minimizer of f(P Q), where X is the ith element of X
The proof is given in Appendix C.I We observe that the key LSA parameters are nontrivial in the optimizer, unlike in-context linear tasks with zero-mean Gaussian feature and task vectors, which
result in sparser structure where the first dl entries of b and the last row of A is zero [19 [6 [5].
The independence assumption on the initial states in Theorem 3.3
can be relaxed, and the global minima of [PP Q) still maintain the same structure as the independent case.
Theorem 3.4 (Global minima for generalized in-context initial states distribution). Consider the in-context learning of length-2 Markov chuins {(x,, 11)) Pi (X), yi E {))1)) with transition probabilities pou, Pii UO, 1). Suppose the initial states 2; are sampled from Bernoulli()) for some constant E (0, 1). Let C1 = - E[I Di= Lj= E[0

---

Define X as H-1 1c1/2n
c1/3n p/4 + 61/ c1/12n], and H is symmetric matrix defined as follows
+ +
2n
C1
C2
(n+1)er 11+1
2n
3n
An
12n
(2n++-10
6n
6n
692
(repeating entries in the lower half triangle are omitted)
Then, substitutiin X into Eq. 10 gives global minimizer of f(P.Q)-
The proof is deferred to Appendix C.2 of the query and in-context samples.
The global minimizer is determined by the joint expectation
While the global minimizer X of our reparameterized loss function f(X) be characterized analytically by convexity, it remains unclear whether such a solution can be realized by the actual transformer parameters (b, A). In particular, the correspondence between X and feasible transformer configurations is highly nontrivial as the input space exhibits structured dependencies. This motivates fundamental question: given optimal representation X in the extended parameter space, is computationally feasible to recover any compatible transformer parameters that achieve it? The following result addresses this question and reveals structural limitation of 1-layer LSA transformers.
Theorem 3.5 (NP-hardness of transformer parameter reconstruction from reparameteriattoon) Let d E Z21 denote the dimension of the in-context data (the length of the Markov chain minus one), and let (d+1)d+2) Given X Rddm representing the global minimizer of the reparameterized loss, solving for 1-layer LSA parameters (b, A) that satisfy the reparameterization equation Eq. is NP-hard with respect to d.
Proof. (sketch) We show that solving the reparameterization equation for transformer parameters (b, A) belongs to a general class of bilinear feasibility problems. Specifically, we recast the system as a collection of bilinear constraints b D(r) A.j = X, (r E [dm]), where each matrix D(r) encodes the structure of the corresponding term in σbb, A). This allows us further express the problem as bilinear program with an objective and constraints over (b, A).
To establish hardness, reduce from the bilinear separability problem, formulated in system (13) of Theorem 3.1 in [8), which asks whether two point sets in I2 can be strictly separated by pair of hyperplanes such that set occupies exactly three of the four induced regions. This problem is known to be NP-complete.
We construct variant of our bilinear program through two standard transformations: variable splitting (to impose nonnegativity) and variable fixing (to introduce linear structure). subset of transformer parameters is assigned to represent the decision variables in the separability problem. Other variables are fixed to constants to encode the geometric structure and auxiliary terms of the separability constraints. Through this construction, each constraint in the separability problem is converted to a constraint in our system.
Our bilinear program subsumes the bilinear separability formulation from 48 First, the objective functions are equivalent: we replicate their bilinear objective by selecting corresponding set of (b, A) variables and constructing a matrix (r) that enforces the same multiplicative interaction. Second, for every constraint in the separability problem, there exists corresponding constraint in our system constructed via combination of variable assignment and fixing. Finally, our program includes additional constraints beyond those in the separability setting, arising from the full reparameterization of the transformer system. This establishes that our recovery problem is at least as hard as bilinear separability, and therefore NP-hard.
The complete proof is provided in Appendix C.4 where show that the dimension d scales polynomially with the parameters of the separability program. This indicates that the problem remains NP-hard with respect to its dimension, as the reduction preserves the scaling of problem size.
In prior work [19, the reparameterized loss under an i.i.d. linear regression setting admits a structural simplification. Under Gaussian inputs, the optimal solution X* becomes sparse: each group of reparameterized variables associated with a column of contains only a single nonzero entry. This sparsity allows the full bilinear system to reduce d independent constraints, each involving single
7

---

bilinear term between b and A.j which can be analytically solved. In our setting with Markovian input sequences, the reparameterized solution X" does not exhibit such sparsity. The temporal structure of the data introduces statistical dependencies and nonzero mean patterns that persist across in-context samples, leading to dense coupling between variables in the bilinear system. This increased structural complexity motivates the hardness result above.
Therefore, even if the global optimum can be attained in an extended parameter space, it may not correspond to any realizable configuration in the original transformer parameter space. Moreover, since the reparameterization defines an onto mapping from the transformer's 's native parameter space Q) to the enlarged space X, the optimal value f == minx f(X) provides a lower bound for the best performance achievable by the transformer, i.e., f S min f(P, Q). This reveals an inherent architectural limitation: the I-layer transformer's s representational capacity may be insufficient to express dynamical functions induced by structured inputs, such as those governed by shared Markovian dynamics. Example [C.66 in Appendix [C.3] illustrates the absence of a corresponding LSA parameter for X in longer Markov chains. Nevertheless, the optimal configuration of (P, Q) maintains the same structure in the length-2 case.
4 Multilayer LSA Implements Preconditioned Multi-objective Optimization
We now demonstrate that the forward pass of an L-layer LSA can be interpreted preconditioned multi-objective optimization when trained within the parameter space with the following structure:
Oax(++1) bi
A
02+1
P
Q
Q=
(11)
04
0
We identify two groups of objective functions, R1, R2 Rd - Rd+1, for the linear model w Tax (w E IR²),, such that TF, performs gradient descent on these objectives, preconditioned by b,, At, an. Notably, this result does not rely on taking the expectation of the objective and holds for any given prompt instance.
Theorem 4.1 (Forward pass as minimizing multiple objectives). Consider the L-layer transformer
A
parameterzed by bt, A =
where by E pd+1, AA E Rdxd, aa E IRd &R I E [L]. Let Vn+1 be
(1)
0i
the bottom- right entry of the lth layer output. Then ( = (wi, In+1) where wi is optimizing two multi-objectiive problems R1, Rs Rdd + RRd+1, iteratively defined as follows: Wp and
wTI wi b/ VR (W At + (aa(w))
[0weppxd] where Ri (w)
mz 6o (w, (ww,,,) 3t5)
ai
(walvs (wid xz).d=1) 1)) w3 2,,)) 30, (Wا w T 2z}}
if #0
R2(w)) Σ
(walus 1 1) ag
if =
-(y; (wid- 3, Wu
The full derivation is provided Appendix D
In the above expression, W3d-1, denotes the first d entries of w and a., respec- tively. The first d terms in R₁ (Æyx=" I, for k E [d) capture the interaction between the k-th state of I and the linear combi- nation of states determined by w. These terms assign weight to each state's contribution to the overall objective, emphasizing their individual roles within the sequence. The final term,
Objo
Olj
0bj2
Obj
Objective (R1)
0.5
Objective (R2)
0.0
0.8
0.0
0.6
-0.2
0.4
-1.0
10
# Layer
Layer
Layer
Preconditioned multi-objective optimization perfor- mance of trained LSA across layers.
Figure 2:
(a) Values for the six
objectives in the first group (R1), (b) Values for the six objectives in the second group (RI), and (c) Generational Distance (GD) measuring deviation from the Pareto front. The model is a 10- layer LSA traiged on 100 in-context Markov chains of length 5.

---

(w (w"zz - - 9 ensures alignment between the linear model's prediction and the target state Yj-
R2 specifically emphasizes the role of wa, aligning with the structure of the last preconditioning matrix, which focuses updates on the last parameter. The first d terms in R2 scale the target state yy by
w²
T.j and Wa, with additional quadratic terms like alignment between y/ and partial prediction components of the imput sequence. The final objective is in cubic penalty form 3æ,,d (v/ w x2)²
2 Fjad- These terms capture the influence of wds the -1) and the value of Fj,d, emphasizing key
when J.dd * 0, emphasizing sequences with smaller Fj,d- When 3 = 0, the penalty changes to a quadratic term - (g) (w-d- (0:d=1, Fz::d=)) u focusing solely on aligning y/ with the partial prediction based on w. This adaptation ensures that the optimization prioritizes the appropriate components of the input sequence depending on the presence or absence of zj,d. Furthermore, when "idd = 0, the final objective becomes convex if wa and concave otherwise.
Empirical validation. We evaluate the role of transformer weights in preconditioned multi- objective optimization, as proved in Theorem 4.1 We train a 10-layer LSA model within the parameter space specified by Eq. oT 100 in-context Markov chains of length 5, sampled with initial probability of 0.5. As the forward pass progresses through deeper layers, we track the values of the multi-objectives and measure Generational Distance (GD) to quantify deviation from the Pareto front.
Fig. demonstrates that initially, LSA weights move closer to the Pareto front, indicating effective multi-objective optimization. However, beyond a certain depth, all objective values begin to increase simultaneously, suggesting that the optimization process deteriorates rather than balancing competing objectives. This behavior implies that while TF initially optimizes multiple objectives, deeper layers may prevent sustained improvement, potentially due to the nature of the restricted parameter space.
To investigate how model depth and data complexity affect performance, we conduct synthetic experiments varying the number of layers and the size of the state space. Prompts are constructed from Markov chains with [S] E [2, 3, 4), each consisting of 10 in-context sequences and one query of length 6. We train LSA models with
0.9
[S]
2
[S]
0.8
[S]
Accuracy
0.7
E {1, 2, 3} layers using gradient descent, and report mean accuracy over 5 seeds (see Fig. 3. The results show I-layer LSA performs poorly, close to random guessing, regardless of the state space size. As the number of layers increases, accuracy improves significantly. For instance, in the [S] 3 case, the mean accuracy increases from 0.55 (1-layer) to 0.98 (3-layer). This aligns with our theoreti- cal findings: while single-layer transformer struggles to capture structured dynamics, slightly deeper architectures allow the forward computation approximately reduce squared loss, as revealed by our multi-objective interpre- tation. Additional experiments using self-altention models and GPT-2 architectures under varied data and architectural configurations are provided in Ap-
0.6
0.5
0.4
1.5
3.0
Laaers
Accuracy of LSA models trained on Markovian prompts w.r.t number of layers. Each curve corre- sponds a different state space size.
Figure 3:
pendix A
5 Conclusion
In this work, we study ICL of Markovian dynamical functions using a LSA model. Focusing on one-layer transformers, we analyze the loss landscape induced by prompts constructed from first-order binary Markov chains. Our results show that the global minimum adapts to the underlying dynamics and deviates significantly from the sparsity structures typically observed i.i.d. regression tasks.
Despite the existence of an analytically characterized optimum in an extended parameter space, we prove that recovering corresponding transformer parameters is NP-hard. This establishes fundamental limitation: one-layer LSA transformers may be unable to realize optimal solutions,

---

even when those solutions are simple and fully specified. The result reveals representational and computational gap between what is learnable in principle and what the model can express.
Finally, we interpret the forward pass of multilayer LSA models as performing a form of multi- objective optimization, which includes squared loss one of the components. This formulation offers insight into why deeper architectures empirically outperform shallow ones when learning structured dynamical patterns.
Limitation. Our theoretical analysis focuses on l-layer LSA models and first-order Markov chains. These choices allow to isolate key structural effects and obtain analytically tractable results, such as characterizing the global minimum and proving NP-hardness of parameter recovery. However, this framework does not capture the full expressivity of modern transformers, which typically involve nonlinear attention mechanisms, multi-layer architectures, and learned positional encodings. Moreover, real-world dynamical processes often involve higher-order dependencies or more complex causal structures beyond simple Markovian assumptions. Extending our framework to handle richer classes of causal data, deeper networks, and nonlinear components such as softmax attention and MLP blocks remains important direction for future work.
6
Acknowledgements
This work was supported by the IBM- -eM-Rensselaer Future of Computing Research Collaboration and the U.S. National Science Foundation project 2047488.
References
[1] Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
[2] Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446. 2021.
[3] Shivam Garg, Dimitris Tsipras, Percy S 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.
[4] 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.
[5] Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean- Cean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.1188 2023.
[6] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
[7] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.. 13971, 2023.
[8] 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 on Machine Learning, pages 35151-35174,
PMLR, 2023.
[9] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023.
10

---

[10] Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning, pages 11398-11442. 11422, 2023.
[11] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In International Conference Machine Learning, pages 19565-19594, 2023.
[12] Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? language models secretly perform gradient descent as meta-optimizerss In Findings of the Association for Computational Linguistics, pages 4005-4015 2023.
[13] Haoyu Zhao, Abhishek Panigrahi, Rong Ge. and Sanjeev Arora. Do transformers parse while predicting the masked word? In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 16513 16542.. Association for Computational Linguistics, 2023.
[14] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisti- cians: Provable in-context learning with in-context algorithm selection. Advances in Neural Information Processing Systems, 2024.
[15] Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context. J. Mach. Learn. Res., 25:49:1 149:55,, 2024.
[16] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05229, 2023.
[17] Karthik Duraisamy. Finite sample analysis and bounds of generalization error of gradient descent in in-context linear regression. CoRR, hbs/2405.02466) 2024.
[18] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. How do nonlinear transformers learn and generalize in in-context learning? In Forty-first International Conference on Machine Learning (ICML), 2024.
[19] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Advances in Neural Information Processing Systems, 2023.
[20] Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations (ICLR), 2024.
[21] Yingcong Li, Ankit Singh Rawat, and Samet Oymak. Fine-grained analysis of in-context linear estimation: Data, architecture, and beyond. CoRR, abs/2407.100005 2024.
[22] Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A framework for principled analysis of transformers via markov chains. CoRR, 1bs/2402.04161, 2024.
[23] Ashok Vardhan Makkuva, Marco Bondaschi, Chanakya Ekbote, Adway Girish, Alliot Nagle, Hyeji Kim, and Michael Gastpar. Local to global: Learning dynamics and effect of initialization for transformers. arXiv preprint arXiv:2406.00 03072, 2024.
[24] Nived Rajaraman. Marco Bondaschi. Ashok Vardhan Makkuva, Kannan Ramchandran, and Michael Gastpar. Transformers on markov data: Constant depth suffices. Advances in Neural Information Processing Systems, 37:137521-1375566 2024.
[25] Nived Rajaraman, Jiantao Jiao, and Kannan Ramchandran. Toward theory of tokenization in lmss. arXiv preprint arXiv:24444.08335, 2024.
[26] Ezra Edelman, Nikolaos Tsilivis, Benjamin Edelman, Eran Malach, and Surbhi Goel. The evolution of statistical induction heads: In-context learning markov chains. Advances Neural Information Processing Systems, 37:64273-64311, 2024.

---

[27] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer. memory viewpoint. Advances in Neural Information Processing Systems, 36: 560-1588, 2023.
[28] Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In Forty-first International Conference on Machine Learning (ICML), 2024.
[29] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Prov- able training dynamics and feature learning in transformers. arXiv preprint arXiv:2409.1055% 2024.
[30] Michael Eli Sander, Raja Giryes, Taiji Suzuki, Mathieu Blondel, and Gabriel Peyré. How do transformers perform in-context autoregressive learning ? In International Conference on Machine Learning (ICML), 2024.
[31] Claude Elwood Shannon. A mathematical theory of communication. The Bell system technical journal, 2773):379-423, 1948.
[32] Claude Shannon. Prediction and entropy of printed english. Bell system technical journal, 30(1):50-64, 1951.
[33] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
[34] Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37836, 2022.
[35] Tianyu Guo, Wei Hu, Song Mei, Huan Wang. Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with representations. arXiv preprint urXiv:2310. 10616, 2023.
[36] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference on Machine Learning, pages 19689-19729, PMLR, 2023.
[37] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. Advances in Neural Information Processing Systems, 66.36677-33707,
[38] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarithmic depth. arXiv preprint N1t::2202.0926666 2024.
[39] Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Revisiting the equivalence of in-context learning and gradient descent: The impact of data distribution. In JEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 7410-7414. IEEE, 2024.
[40] Dake Bu, Wei Huang, Andi Han, Atsushi Nitanda, Taiji Suzuki, Qingfu Zhang, and Hau-San Wong. Provably transformers harness multi-concept word semantics for efficient in-context learning. Advances in Neural Information Processing Systems, 37:63342-63405, 2024.
[41] Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of- thought: Compositionality through in-context filtering and learning. Advances in Neural Information Processing Systems, 36, 2024.
[42] Jianhao Huang, Zixuan Wang, and Jason Lee. Transformers learn to implement multi-step gradient descent with chain of thought. arXiv preprint arXiv:2502.212111 2025.
[43] Khashayar Gatmiry, Nikunj Saunshi, Sashank J. Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learning? In International Conference on Machine Learning (ICML), 2024.
[44] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 121): 2021.
12

---

[45] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. /895, 2022.
[46] Jan Chorowski and Navdeep Jaitly. Towards better decoding and language model integration in sequence to sequence models. arXiv preprint 1612. 2016.
[47] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Informa- tion Processing Systems, pages 5998-6008, 2017.
[48] Kristin P Bennett and Olvi L Mangasarian. Bilinear separation of two sets in n-space. Compu- tational optimization and applications, 2(3):277-227, 1993.
13

---

Appendix
Comparative Analysis of Setups
14
Additional Experiments
14
B.1 Limitation I-Layer LSA
14
B.2 Capacity of Multi-Layer Transtormers
15
Proofs for Loss Landscape Characterization
18
Global Optimum for Length-2 Markov Chains (1.i.d. Initialization; Theorem 3.3
18
Global Optimum for Length-2 Markov Chains (Correlated Initialization; Theorem 3.3)
24
Reparameterized Optimum in the General Case (Lemma 3.2)
27
Optimization Limitation of Single-Layer LSA (Theorem) 3.5)
32
D Proof for Forward Pass as Multi-Objective Optinizzionn (Theorem) 4.1)
39
Additional Related Works
42
Comparative Analysis of Setups
We train two variants of single-layer self-attention models to in-context learn length-2 Markov chains using gradient descent over 10K random prompts.
1. Variant (LSA P.Q
(spanse),
): LSA (Eq.
parameterized by sparse P,Q (Eq.
Z1 = Zo +
PZM QZ)
0
0
a1
P.Q E {(
aa bj E R)
b1
b2
a2
0
2. Variant 2 (LSAp.Q): LSA (Bq. 4
parameterized by dense P,Q
Zi Zo + (PZM(ZZ
P.Q R2x2
To justify the choice the sparse parameter space, we plot the training loss curve of the above three variants in Fig. 4 The loss value is computed as the mean squared error for the query sequence averaged over random prompts. We set B 100 and use 30 in-context examples for each prompt. The in-context sequences are Markov chains with initial probability 0.3 and transition probabilities P01, P10 sampled from U (0, 1). The results demonstrate that the loss curves under variant and converge to nearly the same value, indicating that the sparse and dense parameter matrices perform cquivalently for LSA.
B Additional Experiments
In this section, first illustrate the limitation of 1-layer LSAA the problem dimension d increases, consistent with the NP-hardness analysis in Theorem 3.5 We then empirically explore the expressivity of deeper transformer architectures by evaluating GPT-2 models with varying numbers of layers. The results show that performance improves as the number of layers increases, suggesting that additional depth helps mitigate the limitations observed in shallow models.
B.1 Limitation of L-Layer LSA
Theoretically, we showed that while the reparameterized model admits an analytically computable global minimizer, recovering the corresponding LSA parameters is NP-hard due bilinear coupling.
14

---

0,354
LSA_ (P.Q) Sparse
LSA_<P.) (P.Q)
0,352
0.350
Loss
0.348
0.346
0.344
0.342
0
5000
10000
15000 20000 25000
Epoch
Figure Training loss w.r.t epochs for variants of the self-attention models, evaluated on 100 random prompts, each containing 30 in-context samples and a query sequence.
The hardness stems from bilinear feasibility program whose dimension scales polynomially with the dimension d of the input Markov chain. To empirically verify our theoretical findings on the limitations of 1 I-layer LSA, we measure how the performance gap between the reparameterized model and the LSA model evolves with increasing dimension d.
For data, we construct each prompt with n + = 101 sequences, where each sequence represents a binary Markov chain of length d+ 1. The first n sequences serve as in-context examples, and the final sequence In- is a query with its last token 1 masked. The prediction task is to infer Vn++1 based on the shared temporal dynamics among chains. The initial states I4,1 are independently sampled from Bernouli(O.3)
All transitions within a prompt are governed by shared transition kernel E R2x2 sampled from the same distribution the previous experiment. For each dimension d (1,2..... 10), we generate a batch of B 1000 prompts. The LSA model is trained using the Adam optimizer for 1000 iterations with a learning rate of 0.01, minimizing the mean squared error between predicted and true labels. As reference, we analytically compute the prediction from the reparameterized linear model using least squares regression. Accuracy computed by rounding each predicted value to the nearest integer (0 or 1) and comparing it to the binary label. The experiment is repeated across random seeds.
As shown in Fig. 5 the reparameterized model consistently outperforms the LSA model, with the accuracy gap grows as the dimension d increases. Although both models attempt fit the same underlying process, the hardness of parameter recovery suggests degraded performance at higher dimensions. In other words, the computational intractability has practical implications for expressivity and learning effectiveness in 1-layer LSA models.
B.2 Capacity of Multi-Layer Transformers
To further investigate ICL for Markovian dynamics learning for practical models, we conduct empirical studies using GPT-2-based models to answer the following research questions"
RQI. Does nonlinear attention transformers functionally approximate classical models?
RQ2. How do attention patterns in transformers evolve across layers?
We adopt architectures based GPT-2-blocks. We consider three configurations (embedding dimension, number of transformer blocks, number of heads), inspired by [41] (i) tiny: (64, 3, 2). (ii) small: (128, 4). (iii) standard: (256, 12, The models are optimized by Adam over 50K epochs with learning rate 0.0001. For each epoch, we randomly generate 64 data samples to train the model parameters. To ensure high prediction performance given any length-n' prompt (n' E [n]). we train on the average of the error over different prompt lengths from through n and update n from 26 to 101 during training. To generate data, we use the process described in Section B. We
Our code is available at suthttps:///ggg //gt:/u///ttttb/// wttptttuu..ommminnnntne com www.///ingananaallllloo.. Markov-M Dynooiic
15

---

0.95
Reparameterized
0.90
LSA
0.85
Accuracy
0.75
0.70
0.65
0.60
Figure 5: Accuracy versus dimension d for the reparameterized model (blue) and the 1-layer LSA model (orange). Results show mean and standard deviation computed over random seeds. The growing performance gap highlights the increasing difficulty for LSA to recover the optimal solution as the dimension increases.
(b)0
80
0.75
Test Accuracy
75
Test Accuracy
9.
0:70
cnns Gemalh oeandirt)) Lngiette
0.63
TF cony
TF Gamall
0.65
0.60
(atandard)
Logistic
Unner
0.55
Unear
0,50
20
60
60
1.00
20
40
60
100
in-context samples
in-context samples
Figure 6: Testing accuracy for three model configurations, compared baseline learning algorithms for (a) independent and (b) correlated initial conditions, respectively.
measure prediction accuracy by assigning the integer in {0, 13 closest to the transformer's output as the predicted state. The experiments conducted NVIDIA A40 GPU.
Deeper transformers outperform classical models in predicting the next token for query quence in-context (RQI). We investigate the performance of trained transformer compared to baseline learning algorithms, including logistic regression, linear regression, 3- 3-Nearest Neighbors (3-NN), and Support Vector Machine (SVM), when the number of in-context samples vary from to 100. 6 demonstratee the test accuracy for independent and correlated initial states. The accuracy averaged over 1280 prompts, where the shaded region denotes 90% confidence intervals computed using 1000 bootstraps. The result implies that the trained transformers with small standard size have comparable performance with SVM and logistic regression and better than the simple bascline 3-NN, while the test performance for tiny is slightly worse than its larger counterparts.
Transformers capture similarities between in-context sequences (RQ2). To investigate whether the attention mechanism captures structural similarities in in-context sequences with distinct transition kernels, we visualize the attention matrix and compare it to the sequence similarity matrix. We gener- ate prompts by partitioning the in-context sequences into & groups of equal size, each corresponding a distinct transition kernel. Within each group, sequences are sampled from the same kernel, ensuring shared transition dynamics. To quantify sequence similarity, we compute a similarity score Hamming distance
Fig. 7 presents the attention matrices across layers, illustrating progressive emergence of structure that aligns with the sequence similarity matrix. In particular, by the third layer, sequences governed by the same transition kernel exhibit significantly stronger mutual attention, indicating that the model increasingly attends to structurally similar sequences. This
as: Similarity = 1
Sequence length
16

---

(a)
Attention Map
1 Kernel
Kernels
Kernels
0.032270
0.032265
15
15
0.032260
20
20
20
25
25
25
0.032255
10 15
20
25
10 15 20 25
10
15 20
25
Key
Key
Key
(b)
0.032260
5
10
10
10
0.032255
15
15
15
0.032250
20
20
0.032245
0.032240
10 15 20
25
10 15 20
25
10 15 20 25
Key
Key
Key
(c)
5
5-
0.03228
10-
10
10
15
15
15
0.03227
0.03226
0
10 15 20 25
0
5
10
20 25 Key
15
0
5 10 15 Key
20 25
Key
(d)
Sequence Similarity
1 Kernel
2 Kernels
3 Kemels
ques
10-
15
15
15
20
20
20
25
25
25
15 Key
Key
Key
Figure 7: Attention patterns across 3 transformer layers and sequence similarity..))-cc) Attention maps from shallow to deeper layers. (d) Sequence similarity. In each subfigure, the three columns correspond datasets with increasing numbers of transition kernels in the input prompt. The emergence of block-wise patterns becomes more pronounced in deeper layers.
layer-wise sharpening of attention suggests that transformers can progressively reflect the in-context Markovian dependencies.
In the following sections, we analyze the optimization landscape of the 1-layer LSA model for ICL over Markovian dynamics. We begin by characterizing the global minima of LSA through reparameterized objective in low-dimensional settings (Sections C. C.2p. deriving exact solutions for both i.i.d. and correlated initial-state cases. In Section reparameterized objective to arbitrary-length sequences and general state spaces. Section C.4 then establishes fundamental limitation that recovering transformer parameters from the global minimizer is NP-hard. Finally, Section D offers an interpretation of the forward computation in multi-layer LSA models as multi-objective optimization process.
extend the analysis of the
The theoretical results and their relationships are organized as follows:
1. In Section C.I we assume that initial states are independently sampled.
(a) Lemma derives global minimizer
reparameterized convex objective under this
i...d. assumption.
(b) Theorem C.2 (Theorem 3.3 restated) builds on Lemma C.T characterize the global mini- mizer of the original 1-layer LSA objective.
17

---

2 Independent Initial States
Section C.1
Lemma C.1 = Theorem C.2 (Thearem 3.3 restated)
goba (i min.
grocall minimum. 1-lywe LBA
reparr reterrere model
Generaiion int.
Length-2 Correlated Initial States Gerwaine
Section C.2
Lemma C.3
Theorem C.4 (Theorem 3.4 restated)
gibbal
poobal minimum.
rpparameto
Arbitrary-length
Lemma C.6 (Lemma 3.1 restated)
atrir comnediyy repar amelerlad midel
Lemma C.7 (Lemma 3.2 restated)
Section C.3
giobal nam resersndeeead model
General State Space
Lemma C.8
model
Opnimizttio geo
Limitation
Section C.4
Lemma C.12 b bllenee formudation
Theorem C.14 (Theorem 3.5 restated)
layer
NP-rardreee
Seyond
Section D
Multi-layer LSA
Theorem D.1 (Theorem 4.1 restated)
Eowwerd muul--heeeceee optimiration
Figure 8: Structure of theoretical results this section. Each row pair corresponds to subsection, with the left showing the modeling case and the right listing the corresponding lemmas and theorems. Results build from short Markov chains with i.i.d. initial states to general, long sequences with general state space. This leads to a limitation result, showing that recovering transformer parameters from the global optimum is NP-hard in general. Finally, we provide separate interpretation of the multi-layer LSA forward pass as multi-objective optimization process.
In Section we extend the analysis the setting where initial states may be correlated. (a) Lemma generalizes the convex minimization result to accommodate arbitrary initial-state distributions.
(b) Theorem C.4
(Theorem 3.4 restated) then applies this generalization to characterize the global minimizer of the 1-layer LSA objective in the correlated setting.
3. In Section C.3
we consider Markov chains of arbitrary length and general (binary or nonbinary)
state spaces.
(a) Lemma C.6 (Lemma 3.1 restated) establishes that the reparameterized objective remains
strictly convex this general setting.
(b) Lemma C (Lemma 3.2 restated) and Lemma C.8 provide closed-form expressions for the
global minimizer of the reparameterized objective for binary and nonbinary state spaces, respectively.
We then investigate whether such optimal solutions can be realized by any feasible choice of 1-layer transformer parameters in Section C.4
(a) Lemma C.12 12 recasts this recovery question as a bilinear feasibility problem.
(b) Theorem C.14 (Theorem 3.5 restated) shows that this problem is NP-hard, revealing
fundamental limitation: 1-layer LSA may unable to represent optimal solutions even when
they exist in the reparameterized space.
5. Section recasss the forward pass of a multi-layer LSA as multi-objective optimization involving losses beyond squared error.
The logical flow the main theoretical results is summarized in Figure S
C Proofs for Loss Landscape Characterization
C.1 Global Optimum for Length-2 Markov Chains (i.idd. Initialization; Theorem 3.3
We begin by rewriting the loss by keeping parameters that affect the output prediction for the query
Int1.
18

---

The input prompt is formatted as a (d + 1) x (n 1) matrix:
Zo
21
Zn
2n+1 0
91
We assume 2; ~ Bernoulli(p) and let Pij denote the transition probability from state to (i,j X' (0, 13). We define the label y to be the next state. By definition of Markov chain, the expected value of Vi given is
4.1.d.
Ely Pol,
Pol Pii]
(1 x:)por T.P Poi + (P11 )01)
(12)
Rewriting the objective function. defined
The in-context objective function for the single layer case is
[((0 + Attmp,o(Zo)
27
f(P. Q) = E (n. jus Po1.P11
Mn++1
(13)
d41,++1
In 0
By definition of attention (Eq. 4) (here M
E [(a++)x(++D) is the mask matrix),
0 0
1
Zu AAttn::(Zoo) Zu
P (Z) QZo)
Zo
P(Z. ZJ JQZ
n
1/1
Zo + - P
(n)
Fn+11
In
0
:
QZo
y1 ...
yn
In
Yn
y/
Zu +
QZu
1/1
Yn
Yn
Ln++1
0
Zo + P
2
acys ys
QZo
Diys
2
1=1
The last column of above matrix can written as
Fn+1
1 PGQ
En+1 0
For the binary input case, d and P.Q E 2x2 Let b = b2) (b E IR2) be the last row of
and [a 02] 2 be the first column of Q. The bottom-right entry of Zo (ttt can be expressed bT Gar Since f(P. Q) only depends on parameters b, a, we rewrite the objective function as
S(P. Q) E (n P01P11
[b Ga - yn+1)²]
(14)
Reparameterization.
We further expand the term b Ga
b
b2]
x
rzli
a1
Filk
yE
as
bnb
1 Σ
x + (arby + azb1)
1 Σ
xays + azb2
1
Σ yE
km
im
19

---

Let Gga, Gzy: Gyy denote the top-left, top-right, and bottom-right entry, respectively. For any vector
X = 11: x2 X3] in IR's we consider the following loss function
f(x) E (ma) P11
[(XiGr + X2Gny + X3G In+1 Mn++1
(15)
We first derive the unique global minimum the reparameterized loss function (Eq. 15) and then find the set of global minima for the original loss function (Eq. [3) over the space of P.Q.
Lemma C.1. Consider the in-context learning of length-2 Markov chains {(1,, yo b))) (2 Us {0, 13) with transition probabilities POl. P11 ~ (0, 1). Suppose the initial states 24 are i.idd
sampled from Bernoulli(p) for some constant p E (0, 1). Consider the reparameterized objective
f(x) = 18 (1 (2..0.)
pnd PP1.PP Po1.P1
((Xx)ez X2Gry X3 Gyg) In-1 |n++)²]
(16)
where X = x2,, E RS and y/ = (1 Fn+1 Poo1)Poo + Int1P11 denotes the conditional probability observing the next state given the current state.
(1) The objective function f is strictly convex.
(2) The global minimum X is given as X = H-1 102/2 p²13 p²/12 p/4] symmetric matrix defined follows
where is
p
n-1 p² [n 1'p"
H
+ n-1
4 + 1
(17)
2n
n-1
P + x
2n
(omitting repeating entries in the lower half triangle).
Proof. We defer the proof of (7) to Lemma 3.1 Since f(X) is strictly convex, it has unique global minimum that sets the gradient (f(X) to zero. To show (2), we first set up the equation to evaluate the minimizer.
Setting up equations to solve for minimizer. The gradient of f w.r.t. X can be expressed as:
E
#++1
(G X1 + Gay Gaa Xo + Gyy G.. X3)
En+ 19++ G
vf(x= =
E
#n+1
Gzz Gzy X1 + G
By X2 + Guy Xz)
- 2n++n+1 18n+11 Gay.
(18)
E
(Gen Gyn X1 + Gzy Gyy + yy Xa) - Int 18 Gya
The global minimizer X is the solution the following system:
E
G2
E Gza
Gzy] all my
E 1ut1 Gev Gyv Gay Gyv)
Xi X;
Int + Gra
G
Gry.
E
E
E Fn+ 10 Gay
n G
Gyy
E
#1 Gay
Gyy]]
E
2
my
[In+ 18a+1 Gay
(19)
Next, we compute the expected values in the linear system 199. The reason for each key derivation step marked with labels such (i), (ii), etc., and further explained at the end of the derivation.
Computing RHS of Eq. 19
We evaluate the three elements in RHS separately below.
20

---

1. For the first element, we have
E (0r. 32 Por
Fn+ 18h++1 Gazz]
E (1 1a PP1.P1
- In+ 18n+1
E
(21:2 not pon Por.p1 pm
11+
71
independeent of
E 2u
1.844 1Pa P11 (2n+19n+14
E Ex Bx PP1.
E yn+1 (En-1ln+11))
24, 2n+1 P01. Pı1]
E Ex P.1.1.111
[x] In+1 In+1Bgnns [N++1 [1 En+1: p01 Pn]]
(8)
E Ex An Po1.P1a
[ In 1 Byngs [Mr++1 Px+1, Pil
R aut P11
[n] (po1 + (P11 Pa1)=
(11)
Rutt 1. 111. [x] P11Tn+1]
E pps [pn)]
Ex, [x]]] Exsts [xn+1]
(iv) 1
2dd
(20)
2. Similarly, for the second element, we have
E (z.,0i Hil P01.
[In+18h+1Ge Gzyl
E e, Me #n Po [ethhn+tVe++
Eyy [3%]] Ø4, 041 Pıa] Int Byn.s 12++1 Øn+1 P01 P11]]
aat Por Pm1
E anta Por Pm
[Po1 + (P11 po1)=?) (porent1 + (P11 Po1)
(4X)
E
[P113: P112n++1
(888)
Pp
[pi
Ezy [3re]
Be. in [mn+1]
(iv)
3 3
(21)
The third element can be expanded as follows.
E
[In+18++11 Gyyl
(2re 1i Pos P11
Br Bu..y (1.. Pot P1 (1n+10n+1164)
E
Byg [v? Ie Po.,
In+ By [aa41 Fnt1, Pra]]
S. PP1 .P11
E
Byy [vs [z,, Po. P11] (P11#n+1))
E aati
(Pp1 + (P11 Po1)zz) (p114441)
E m..
[Po1P119041 pi pir0:4111 Po1P12e7n1)]
(868)
E pox [Po1] Epn [P1]] Exaea [xn+1]
+ (Epm
[r]]]
Epos [Poi] Epi. [Pn1) Er,, [1:] Exner [fati]
(v)
1
4 p + P
(22)
Computing LHS of Eq. E[G]].
We evaluate the expectation of the covariance of in-context examples:
21

---

1.
E (1 P11 [
=E (arie) P01.P11
1 Σ
x
E
[ 20
+
σ 2zn +
In 21
In nn
n²
1
nBm [2;" d7 + ninn 1) [2
n
1
(np n(n 1)Ex, [x:3] Eee, (=]]])
n
(60)
(np + n(n 1)p²)
n
p
n
p
(23)
72
2.
(1y puit Porer GGr Gry)
1 Por PP1
x
Tilli
1=1
i=1
(8)
E ( P01-P11
1
E
Σ
P01, P11
i=1
1=1
1
p: + (p11 - pon)
[B - Po1-P11
iml
im
1
[B. Po1
Po1
i=1
i=1
1
(p1
(po))
po1)
Σ
7m1
im
(84)
1
E
P01
P11
i=1
i=1
(o
[2 g
2
p
11
+
p²
(24)
2n
2n
22

---

3.
E (or Min Po1pr
GG Gyy]
=B(o Par
B
Σ
Σ
si
(ct)
1=1
i=1
(10 Por Por-
E Bi
Σ
[x:)
PO1.P11
1=1
= ( PP1
24
Σ
(po1 + (p11 po1)
1=1
1=1
27
= -B 1i 1 PM
Po1
Iz
+ (P11 -oo)
It
==1
i=1
(681)
Epo [poi] + Epon 1p11 poi)l )
(se) p
(25)
2
4.
E
(2..n 1n4 Pm
[G
1
E (e Po1-P11
rtth
1=1
i=1
n (1..)) 12 Porspis
E
x
y/3
n2a
E
(1r No
por.pr
Σ
Σ
1=1
E
(i)
n
Big,
(x.
Him Por P11
Σ P1ı
n2
E
(z.
Pi P11.P11
Σ
J=1,f$1
p
(n -1)p²
+
(26)
2n
3n
5.
(e ao 1841 Por Po1
E
[G Gyyl
(o PoriPL
Filli
t=1
n²
1
E
n++1 Po1.P11
Σ
+
n
1
B
(m
12 P01
Σ
Σ
i==l
i=1 j=1,j1
(3)
n²
E
(z.
Po1.P11
1=1
+ n²
E
(a P011P111
Σ
Σ
ppo + (pn - po))az))
i=1 j=ljj
(su)
P
n
1
p p
(27)
2n
4
12
23

---

6.
E (21.0)) 1i 1o1...1 p01.p11
G
yyl
1:001.11111 [
yr
Σ
vt
n
1 n2
E
(a.:)))
mit PP1-PP1
Σ
M/i
+ n²a
(m P01--11
Σ
n
Σ Nēus j=1,fş
1
n
E
Σ
Pui + (P11 p0i)az
n
2
P01.P11
1
n
E
(P01 + (P11 - Po1) + +P11 po1)zz)
n²
P11
(50) 1
1
2n
+
6
(28)
Throughout the derivation, (1) uses the fact that {2, 83 and {ay, My G + j) are conditionally
independent given Por,Prr: : (Hi) holds since Tús Yt are binary random variables and x = wss yf = y for any integer k; (Mii follows from the fact that pot, Pm and xr Gj E [n+1]) are jointly independent; (iv) holds because the kth moments of uniform distribution U(0,1) 1) and Bernoulli distribution Bernoulli(p) are E+I and p, respectively.
Since In+1 and H; (i [n] are independent, we have E[z G.]
1 pE[G?]
Hence we have the expression for H.
Since f(X) strictly convex, Eq. 19 a unique solution X H-1 [p² p²/3 12/12 p/4]
Theorem C.2 (Theorem 3.3 restated). Consider the in-context learning of length-2 Markov chains
(Es, ye) Pi (It, ys E (0. 1 with transition probabilities P01- Po1-P11 ~ UO, 1). Suppose the initial
states 24 Lid. sampled from Bernoulli(p) (p) for some constant p (0,
Let X = H 122/2
p²/3
p²/12 + p/4] T, where H is a symmetric matrix defined as follows
2
(n-1)p²
(n-1)p
2n
2n
H =
2n
(a-
3m
1ip
2n
2-
g
12
(29)
R-1
2n
6
Then the following choice of parameters
0
Xi
P=
N X 4X
Q
XI v xx-
(30)
X
is a global minimizer of S(P.Q).
C.2 Global Optimum for Length-2 Markov Chains (Correlated Initialization; Theorem 3.3)
Lemma C.3. Consider the in-context learning of length-2 Markov chains ((1,, w:)) (2t, Wi
{0. 13) with transition probabilities P01,P11 ~ U(,, 1). Suppose the initial states In are sam-
pled from Bernoulli(p) for some constant p (0,1). Let ei En= E[p,r4+1, ET E B(I4z5In+1)
Consider the reparameterized objective
j(x) = E [ 1n+1 PP1.P1 [((X] G.a + X2Gmy + X X3G xn-1 1n+1)
(31)
24

---

where X = Xxx] E Rd and ys (1
In+1 Po1 + In+1P11 denotes the conditional
probability observing I the next state given the current state. Then a global minimum is given as
C1 2n
=
C1 /3
(32)
Lp/ + C1 12n
where
+
2n
2n
H
(n-+-))o
(33)
12n
(2n++11p
10
6m
672
6n
(omitting repeating entries in the lower half triangle).
Proof. Since the objective function remains the same, the derivation for the equations follows from the independent in-context example case (Eq, [9]. Similarly, we label the key steps with (i), (ii), etc., and defer the explanation at the end of the derivation.
Computing RHS of Eq. 19 w/o assuming independence of (Is (zj)t))1
1. For the first element, we have
E (msw) Po1
1 Ua++1 Grar]
72 E
12 Mn++1 In+ Po1 Pu]]
(4)
E Fn 1i1 (pLIsEnt1))
(Gi)
Epis [pn)] Eagonrr [1)0+4]]
2n C1
(34)
2. Similarly, for the second element, have
E
P01.P1L
E Baga Po1-P11
(in+1 Ey [yn+1 P01,1, Erg [yr Poi pr1]]
E a m1 Po1 ( ((n+)1))((spp)))
i=1
(184
1
E[0,0n+1|
C1-
(35)
3n
3n
The third element can be expanded follows.
E
[n++10++ G
Po1.P
Emg M. Bh PP411/++MM
(4))
E m. 1...1.PP00 Po1
En+ Eynra [Na+1 1.141,P01,Pi1 Eyg [yn
we Por, Pn]]
Ex, nn PM ((xm41p1)(00) + (p1 (oo1):)))
71
1=1
p
12n C1.
(36)
25

---

The derivation holds due the following facts: (i) the states are binary, (ti) My and y: are conditionally independent for j * (itt) independence between I; and P01.P11-
Computing LHS of Eq. 19 w/o assuming independence of ( Jeei++1l- We directly present the results for the other terms, their derivation is similar to that of the RHS in the independent case.
n+11
200
n
E B
+
n
B Blaayynnn
1
n²
+
n
22
(37)
E #n++1 Gaz Gay
El
+
B BirgzyTn41
2n
2n
1=1
1=
1
2n
+
2n 02
(38)
1
E [ Gea Guy]
E [n+n+11]
2n
1
C1
(39)
2n
1
1
E 11 Gry
B| E[243n+11
+
E B[zz0yan+11
2n
3n2
1
2n
C1
+
3n2
C2
(40)
1
n1
12
1
E
Gzy Gyu]
E[2,8n+11
+
2 +
E E0,
2n2
n
12
j
71
+
C1
+
C2
(41)
4n
12n
p
(n 1)p
n
n
E
+
2n
3n
n²
i 3m0
E
[0 1]
E
12
12
22n + 1)p
n
1
6n
6n
C1 +
6n 02.
(42)
L
Theorem C.4 (Theorem 3.4 restated). Consider the in-context learning of length-2 Markov
chains {1r,, 2 (Ii, Yi E {0,1}) with transition probabilities po1;P11 U 1). Sup-
pose the initial states H; are sampled from Bernoulli(p) for some constant E (0,1). Let
C1 -E Elxgen+1) = -Ei Ej=1ie B1r
We define X* as X H-1 [c1/2n c1/3n p/4+ / 11n)] where H a symmetric matrix defined as follows
2n
2nd
2m
C1
(z 1
(43)
2n
dn
(2n-
CO 6
6n
6n2
(repeating entries in the lower half triangle are omitted).
Then by substituting X" into Eq. 10
gives global minimizer of f(P.Q).
26

---

Example C.5. Suppose (ntt--Berooull(p) and 24 In+1 ~ Bernoulligg(xxtt) for some func-
tion g {0, 1} > [0, 1]. For example, when g(z) (x p)?, the expected values can be computed as follows.
For [n]] .J E|a = Bz,.1 [nat1 Be,, (2t [1n++]]]
= Eza-- [out] 2px + p pn1n11]
p 2p² +p²
Therefore c1 =npp-pp 22p p²).
C2
n 1)
Bz
In+ Eeg (xt Fn+1 Euy, [zi] [ [n+1]]
B
[n+1(In+1 - p)2()) (En+1 - p)²
Baat1
[n+1 (1 2p + p²² ) 2ppnn+1 +p²)]
ni) Eunt1
|n+1 ((1 - 2nt + +2²)
B
[[] - 4p + 4p"))= + +
2p)p²x==+ + +p xn+1]
4p² + 4px - -2p" 4p
4p +p²
pd 4p + 6y² 4p² +p.
The above derivation holds because (i) 04, Ty are conditionally ind. given #n+1, and (ii) the states of Markov chain are binary.
C.3 Reparameterized Optimum in the General Case (Lemmal 3.2
We recall (xs,, y:) form binary Markov chain of length d+ 1. Assuming the initial states sampled from Bernoullipp), the probability of x4,1 being is For 1 < S d, the probability of day being 1, given 2t,-1, is Pi1 (1 z11-1)po1- The probability y, being given Fi,d, is Piidid + +1 P01-
Reparameterization. For general d 2 the projection matrix P and attention weight matrix Q are size (d 1) x (d 1). We write
0ax(++1)
p
= [A 0a+1]:
(44)
where bT E [d+(+)1) denote the last row of P and A E R(d+1)xd (j E [d]) represent the first columns of Q. The objective function can be rewritten as:
GAAAA - yo+1 -Wat1 )
((, = E (x..0) ) PoL-PI PM
(45)
where In+1.5 (j (d) denotes the jth element of #n+1 and A. denote the jth column of A. The i-j entry of G (Gr,j has the following expression:
1/n k= Ik.
if i, [d] if [d.j = d 1 loi
G
1/n Tk.Mk
d d+ [d]
(46)
1/ Ek=1
if = d+1
Since G is symmetric, to obtain objective function with unique global minimum, collect model parameters that share the same coefficients Gú, = GG,.- We introduce a reparameterizzaion which maps from the model parameter space to Rddm where (d+2)(d+1),
6(P.Q), =X, = (Anpby + Ay.br Aryby
for [d [d+ 11,j' for i' E |d + 11,j =
(47)
27

---

Here of-), is the rth entry of the resulting vector. with r (j 1)m + ((d++ + + and Av, denotes the (i.j)-th entry of A and denotes the ith element of b.
To simplify notation, we collapse the unique elements in G into vector:
g g= [G,11 G1,2
G1,6+1
22.2
Ga
Ga,d+1 Go+1,4+1]
(48)
We concatenate the parameters X [d] into vector X [X()),... y.... :(()) E Rdm and consider the following reparameterized objective function
f(x) = (m.. P01.P11 ((En+1 @ g)
T yn+1)
(49)
Building on the formulation of the reparameterized objective for Sigitrary-leggth Markov chains, we next show that this objective is strictly convex.
Lemma C.6 Lemma 3.1 restated). Suppose the initial probability of the Markov chains is π = [1- p. with p E (0, 1) and the transition probabilities are sampled from (0, 1). The reparameterized objective function Eq. 49 strictly convex w.r.l. X B.dm.
Proof. We show the Hessian of f w.r.t X, B[ 1+1 gg J. is positive definite. Let + 0 be arbitrary nontrivial vector in Rdm, Let Un++1 g- Then for any Fn+11 {0. 1] and g [0, 1jmh w B[ @ gg Ju w [x
@ g) (Zn+1 @ g) Jw w 22
|=T2 212 2 O. Since # 0, at least of its entry is nonzero and this entry is multiplied by one of (En+1.5Gr.f" j [d..''3
[d + 1]} in the expression w 2. Take j a, 3. y Then suffices to find specific {ag, } and Øn+1 s.t. In+1[a]G > with positive probability, ie.. P( are nonzero, by definition of Markov chains, PI G3,y] is the product p (or - p) and PıjS and therefore is nonzero. Now because w" (22 )w 2 0 for all & in its support and there exists at least one Rdm s.t. (=z T Yw > and IP(=] > O, we have w T[(z (izzTJw > Hence the matrix 5[ n+1 X gg] is positive definite and it follows that is strictly convex.
O. Since the initial probability p (0,1) and the transition probabilities Pij
Leveraging the strict convexity of the objective, we proceed to derive its global minimizer in the case of binary Markov chains with arbitrary length.
Lemma C.7 (Lemma 3.2 restated). Consider the in-context learning of length-d + 1 (d 2 1) Markov
chains (:)) (ate, yi {0, 1}) with transition kernel P Poo
Poi
E (0, 1)2. Suppose the
Pio P11
initial states xi are i.i.d. sampled from Bernoulli(p) for some constant p E (0, 1). Consider indices
i E [d]. i'jj", ,, K,
E [d+ 1] with i' S j', ',,''= S l.. We denote 4 S t2 S ts S l4 as the sorted version of (i', I k, ("). Define H E Rddxxmm as
H
)o1 (pta-t1) ") (P's (1 (p In + (1 (p)((p)-)))))
n
n-1
(p(P"
n1 - (p (p)-)))
E
n
(p(pt' nn (1 (p )01) (p" (p")
(50)
where r (i 1)m j E d + T, c (j 1i-1)m- m+ l' Ek'-2 d
Define b dmm
b (1-1)mt d+1-; =E [p((pii nn + (1- (p))))p--))o)) (pd+1-j)
(p(P"" -J (1- )011 (p )m])
(51)
The global minimum X E dm of the objective function described in Eq. 49 equals x = H-1b.
Proof. Setting the gradient of Eq. 49 w.r.t. to zero, we have
[ & gg
X E [y++1 (#n+1 & g)]
(52)
where & denotes the Kronecker product.
28

---

Evaluating LHS of Eq. 52: E #1,1 Gr. with i.j E S. E [d + 1],
and i' Sk S. Let P
Poo
Poi
denote the transition probability matrix and π = [l-p.) P]
Pio
Pi1
the initial marginal probability. Further, (P*) (i,j E {0, ()) denotes the specific entry of P raised to the power of k. Then
E 1nn+1, Ge Gg =E 1a41,ea411, E [G Gar,
(53)
due to the fact that 2; (i E [n|) and #n+1 are independent and G contains in-context samples only. We then evaluate the two terms [In+1,54a11,] [Gy Ggr. separately.
For i S j. the probability of Tn+1.; = In+1, = equivalent to that of In+1, = conditioned 1,+11, multiplied by the probability of 2n+1, = 1 conditioned In+1,i 1. Therefore,
E 1a+1,iia+11
E
P 1n+ = En+1.j = =]
E (n(Pi-1) )1 (1 (Pip)) (o1) (PI-))
(54)
We temporarily let #k.d-1 = yk for E [n]]
For ,', K, I,'' ld + 1] and 1' s j', k S L', have
[G y GR
=E
Fk Tk
n2
m
E
Fk DK
Trckk Brrll
n?
n12
Te Tk.y Tk
n12
(55)
n
k= Kol
The summands in the first term, in the of jr S W, has the following form. The remaining orderings j R., 1 can be computed similar manner as follows.
E
T
E2
IP
Tk,
Ik.s'
Ik.K
=n
(p(pi'
J11 + (1 (p
D''''' (P sn]
(56)
Fach summand in the second term Eq. 55 contains a product coactivations from two distinct chains, 2か and Zne (k o x). Because the chains are independently drawn from the Markov kemel,
29

---

the terms 2k.j and Ex, are statistically independent.
E
k= k=1,
E
Plana 2k,"" (2x, xn, = 11
k=
n(n 1)B ( (p(p)' (1 p)(PP'
(1-D)(p''-)) )101) (po'
(p(pl'-)) In + +11 ()(('''1))) (p"'-N )11
(57)
Evaluating RHS of Eq.52
52 E El 1n ( with i' S
1
21 Grarl
E 2h
n1
k=1
1
E P [2n = Un+1 = 1| [I Kk.j = 1]]
E (pp())-1 11 + p)((P=-))))) (po+1->)
(p(pi' (1 ()p (p
(58)
We next extend the global minimization result to the case of a non-binary state space. The derivation follows similar strategy as the binary case, but now incorporates. Dirichlet prior over the Markov transition kernel P. where each row P..: (s S) is independently sampled from Dir(a 1s). The concentration parameter a > 0 controls the variability of the transitions: smaller values encourage sparse, deterministic dynamics, while larger values yield more uniform behavior.
Lemma C.8. Assumptions:
Each row of P is sampled from the Dirichlet distribution of order S with parameters o, with $ E S= {O....,lo |S))
The initial states are independently and uniformly sampled from S.
Define H Rdm
H
1
ΣΣ
Σ
88' ss'(pi-i) (pi-l))))
S
L x' &'ES S"
1
E
Σ
$152838455(P"'
Jore "20'-))so (p )'es"
Desos (1f -
S
n
Σ $18288(P''11) )e (P5' )eses
n/ S|S]
$1:82.83ES
10 (pl' Yoges ].
Σ
818282 (P''')
(59)
$1.82,83ES
where (i 1)m + + E d+
+1' I' 1m
-K'-2 d 1-7.
30

---

Define E Rm as
n (i--)m+t d+1-r = E
Σ
818283 $18283(PD11 (Pi (put1-s) ) )020
$1,82:33ES
(P3253(Pr )oro) (PD'
)220s
(60)
$1.42:33ES
The global minimum X E Rm of the objective function described Eq. 49 equals X H-1h.
Proof. Evaluating entries in H: E (n+11, Gy Ga v with i.j E [d, i S j, k. (d + 11, and < k' l'. Since #n+1 and xy (1 E (n)) are independent, as before, we evaluate E [f 1.j] and E [G Ger,r] separately.
The first term captures the expected coactivation between positions in the query chain.
E 11 11
E
ss' P In+l,i 8, Fnt1j = s')
LES #'ES
E
88' (((D)t)) Yas PuePPlnn11 1n+ =s]
LoES e'ES
ΣΣ
Σ 88 (5'(Pi-))ss (asx
(61)
S
LESS E'ES'ES**E EEES
We proceed to evaluate the second term, which denotes the aggregated coactivation over the in-context samples.
1
[G Gr
Fke x 2
n12
1
E
Σ
n2
(62)
The summands in the first term, in the case of j' W, has the following form. The remaining orderings of i', i', k can be computed in similar manner.
Tk T
Σ
$192535155
"eooo (PD (P
-0' De
(p'es )ssss (p''- )asso
(63)
S
In the summands of the second term in Eq. 62 contains the product of the coactivation from two distinct chains 2k, 2x (k # k). Since the two chains are independently sampled from the Markovian kernel, the two products Th and Fx are independent.
n
E
Σ
Σ
k=
k=1,
nnn 1)E
Σ
818288 (Pi-
(P )ozs
1223
$1.82133ES
(1525y(pk'-))
)
Jezes
(64)
31

---

Evaluating entries in h: 11 Gr with S We utilize the independence of the query and in-context chains to separate the expectation of 1.+1,8n+1 and Gz.z. Specifically, we have
El Ge
E (1t4
[In+1J3h+tl [ap
Σ
Σ
818283 (Pi2))(P111))) (pd+1-3) 82203
$1,82,838S
Σ
818283 (Pi' 00 (P5'
(65)
Although obtain closed-form expression for the global minimum = H- the entries of both H and h involve high-order interactions across the kernel. In particular, each entry entails expectations of products of entries from powers of random transition matrix P. where each row of P is independently sampled from a Dirichlet prior. Evaluating such expectations is analytically challenging due the nonlinear dependencies introduced by matrix multiplication, especially when powers of P couple multiple rows. However, for fixed values of the Dirichlet concentration parameter a, these expectations are tractable via Monte Carlo sampling, enabling empirical evaluation of H and h.
Moreover, the functional form of these entries suggests that H is dense in general, most entries involve sums over all state triplets and nontrivially depend on all entries of P. As a consequence, the inverse H-1 is also expected to be dense, indicating that the global optimum X" integrates information across all positions.
C.4 Optimization Limitation of Single-Layer LSA (Theorem 3.5
We have found the global minimizer of the reparameterized objective F. The goal is to find transformer parameter (P, Q) such that Φ(P. Q) X*. We show that in general this problem is NP-hard in Theorem 3.3 We formally state the original problem to solve the reparameterization
equation for transformer parameters (b, A) below.
Problem C.9 (Transformer parameter reconstruction.). Let d E Zz1 denote the dimension of the in-context data, i.e., the length of the in-context Markov chain minus one. Given X E RRd (m (d + 1) + 2)/2), solve the following system for E Rd+1, A E RR(d+1xxd
[b.Aks = X, ifi k b.Au + b = Xx i
(66)
where i, [1;j E [d]ir (j 1)m + i(d 1) + Er
We introduce a broader class of bilinear feasibility problem that incorporates the above problem as follows.
Problem C.10 (Bilinear feasibility.). Let d Zz1 denote the dimension of the in-context data, i.e., the length the in-context Markov chain minus one. Define (d+1)(d+2)/2 Let X E Rd and D(r) IR (dd+ ()x(1)) (r [dm]] be given. Solve the following system b E gd++1 A (dd+1)xd
b ((r() Anj X
(67)
where [d];- = (j 1)m i(d + + k- Ev 1) with k [d (
Next, we define bilinear program and establish its equivalence with the above bilinear feasibility problem. Each equality in Eq. 67 equivalent to two inequalities. We hold one inequality out and treat it as the objective. We state the bilinear program below.
32

---

Problem C.11 (Bilinear Program). Let X be some given vector in Rdm where d E B21 is the problem dimension and m (d + 1) + 2)/2. Let D E (R((+1)x((-)) with [dm] be some given malrices. Find the optimal solution b' E R A" E R ldd+ for the following problem:
min BER AEEROTTX
b D A. X1
S.1.
bt D A. 2 X1,
b ((D)) A. Xr:
b D(r)) A.ji Xr: +
(68)
where j E r (j 1)m
( '' with i E [d 1].
We now show the equivalence between solving the bilinear feasibility problem and finding a zero minimum in the bilinear program.
Lemma C.12 (Equivalence between bilinear feasibility and bilinear program). The bilinear feasibility problem (67) has solution if and only if the bilinear program (68) has a feasible solution whose objective function value is zero.
Proof. Before showing the equivalence, we note that the minimal objective function value achievable by the feasible variables for the program (68) is greater than equal zero. Because for solution to be feasible, it must satisfy 5t ((1) A X1 and hence b" ((1) A. X1 must be greater than
equal to zero.
(=) Suppose the bilinear feasibility problem has solution b, A. Then they must satisfy the inequality constraints the program (68). Moreover, the objective function value would also evaluate to zero
by definition of b, A. Therefore, A is an optimal solution of the bilinear program (68)
(<=) Suppose b" A is the optimal solution of the bilinear program (68) whose objective function
value zero. Then they must satisfy all inequality constraints, which are equivalent to dm 1 equations in the bilinear system (67). Additionally, the optimal solution guarantees the objective being zero and hence the remaining one equation also holds.
Bilinear separability program. By Lemma C. 12 it suffices to show finding a feasible solution achieving zero minimum for the bilinear program (68) is NP-hard. To do this, we reduce from bilinear program (13) in Theorem 3.1 in [8). This bilinear program has been shown to be NP-hard by reducing from the problem of bilinear separability: determining whether
disjoint sets of points can be strictly separated by two planes such that one set occupies exactly three of the four regions defined by the planes. Let n' denote dimension of the points and m', k the number of points in the two sets. Let A' E Th xn' and E R store the points the two disjoint sets, where each row corresponds to point. Let denote a vector of ones with arbitrary dimension. We follow their notation, and cases of conflict, we use a prime symbol indicate their variables. Their original system to determine the bilinear separability is:
212²
min #1 2,, aw 20 n2
A'w + eeete
50, 0,
A'w²2 + + He
< 0
s.t.
Bw e'ete $
Bw Pete
(69)
0
< 1z1
Here, 2² E RK decision variables, w w² E R² and 71 2 R are the remaining auxiliary
variables; is some given constant vector in R= The above system has zero minimum if and only if the two sets are separable in the above mentioned way.
Program 669 is equivalent to the following problem after splitting the decision variables into two nonnegative parts.
(2) =1) - 22)
min
2 ?
A'w' + + te
< 0,
A'w² + 2 te
< 0,
s.C.
Bwd - -
< 2
Bw 22
2
2
(70)
0
< 2 2
0
< 4
2
33

---

21
3/
7/2
3
A'
3
B
Figure 9: How each variable in the LSA parameter recovery problem is utilized to construct feasi- bility program that contains the objective and all constraints in the bilinear separability program (70) for the case m' k' 2, n' 2, and
The reduction idea is to map decision variables to the first few entries of (b, A., auxiliary variables entries in b, and many constants as possible entries in A. We first give reduction example (C.13) and then formalize the NP-hardness of solving the bilinear program (68) corresponding to recovering transformer parameters from reparameterization.
Example C.13. As an example, we consider the case for m' 2, the variable assignment below and visualize it in Fig. 9
and d 8. We detail
In this example, we explicitly map each constant, decision, and auxiliary variable from the bilinear separability program (70) to the corresponding parameter entries in the bilinear feasibility pro- gram (68). We assign variables to the following matrices and vectors in the separability program:
(i) constant: A' E R B E R2x2 c [1 _]_ (ii) decision variables: 24. 2_, 24. x2 E RR? (iii) auxiliary variables: w 2 E 7 IR?.
For indexing, we use by: to denote the subvector consisting of entries through of b, and A to denote the corresponding rows i through of the jth column of A. Additionally, Aas' ja denotes the submatrix of A spanning rows and columns j to j', inclusive. We use (=4); to denote the ith entry of 1
Decision variables and objective for m' -- We map the decision variables and objective from the separability program into the feasibility program framework. To do so, split bı.2, A1:2,1 into two nonnegative components and keep the rest of the variables variable:
Let d denote the resulting
b= b
bj
b - by
by
by
by
b6
by by bo]
(71)
Likewise, we denote partially split version of the first column of A as follows
A=
:- [ AT,
At A21
As.1
A41 As.1
A
A7.
As,1
(72)
Variable assignment: We map decision variables as follows:
2 bi
z1 + bT2.
2 + A
22
AT
(73)
D(')'s 's construction: We set D entry to at the following positions and 0 otherwise:
{ i) [2]] Then the objective in the feasibility program becomes
(74)
min (bit2 bT:)) (A AT.21) X,.
(75)
34

---

Since X1 is fixed, minimizing the above expression is equivalent to minimizing the objective in (70). Moreover, because we apply the same variable splitting to the corresponding entries of and A, the nonnegativity constraints on 24, =4,=_,2, x are preserved in the feasibility formulation.
Linear constraints for m' = W = n' =. We then encode the linear constraints into the feasibility program setup.
Variable assignment: 1. We map the given constant matrices A', B to sub-matrices of A. Since first column of A is reserved for decision variables, we begin assigning the remaining constants starting from the second column.
A' + - 1:2,2:3
B - A 3:1,233
(76)
2. We assign the entries in after the decision variables to the auxiliary variables:
w ba:4,
w² + bs:6; 7 + by, 72 + bs.
(77)
3. We further require few entries to be fixed at 0 or 1:
bg A5,23 := 1, A6,233 1= O.
(78)
Counting constraints both programs: the separability program, the constraints -A'w" + 7 e S and A''wa 2 + S each consist of 2 scalar inequalities. Similarly, the constraints Bw - - + S =1 - and Bw 72 + S 2 =² also contribute inequalities each. In total, there are 8 scalar constraints in the separability program, and we construct 8 corresponding (D)) matrices in the feasibility program to match these.
In the feasibility program, the right-hand side index ranges from I to dm = 8-445 = 360. Due to the indexing convention, incrementing the column index j of increases by 45. Specifically, for each j. the set E {(i-1)m 1, (j-1)m + 2,...... jm } = (45(j - 1) + 45 - 1) + 2, 45j} corresponds to 45 constraints involving the expression b D A..i.
p(r)'s 's construction: For each linear constraint in the separability program, we proceed as follows: (a) express the constraint in its full matrix form; (b) encode it using the variables in the feasibility program based on the earlier assignments (Eqs. 76 78); and (c) configure the corresponding D to implement the constraint within the fcasibility setup. We set each D(r) based the terms b,AAk.; that appear in the constraint written in step (b). For instance, a term b,Ak.j appears with coefficient set the corresponding entry (( 1.
-A'w' + + e (a) Full matrix:
A4 11.2
[
Ab
Ab.2
W22
+ +y
+
<
(79)
(b) Mapping to feasibility program's 's constraints:
A1,
A2,2
b3
A5.2
A A5,3
A6.2
+
by
+ by
< by
(80)
A1,3
A2.3
b4
A5.3
A6.3
(c) We the first index corresponding to the first corresponding to jo to encode the second row. This leads us to construct (6)) and (99)), Based on the terms in (80). we set the following entries of (445), D(91) as 1 and others
encode the first row of inequality (80) and
(4.))) ((9,))(()0))))))))))))))))1) ((9,6))
(81)
2. =A'w3 + +eee+es 0.
(a) Full matrix:
A4
A
0
12
1
+
+
<
(82)
35

---

(b) Mapping to feasibility program's 's constraints:
A1,2
22,2
bbs
As,2
As,2 A5,3
A6.2 A6.3
+ bg
+ bg
I b9
(83)
A1,3
22,3
bs
A5.3
(c) We use the second index r corresponding to j = to encode the first row of inequality
[3) and the second corresponding to j 3 to encode the second row, which leads us
to construct D(47) and ( Based on the terms in (83 we set the following entries of D(47) (002) as 1 and others 0:
((5, (6, 2)) U ((8,))) U ((. U (()
(84)
3. Bw - - # =L.
(a) Full matrix:
B1,1
B1.2
H
(2
11
B2,1
B2,2
i
+
(24)2 - (2_)a]
(85)
(b) Mapping to feasibility program's 's constraints:
A3,2
A4.2
ba
As.2 A5.3
As.2 As.3
[bt - A5.2 As.
A3.3
A4.3
ba
by
+ bg
- by A
(86)
(c) We use the third index corresponding to j = 2 to encode the first row of inequality (86). and the third r corresponding to j 3 to encode the second row. We set the following entries of D(18) (93) as
(33.), 4)) ((9,5)),
(87)
the following entries to -1
((7,5)) U {(j - 1,5)},
(88)
and the rest to O where j 2 for ( and j 3 for (93))
4. Bw? - - e S 2 x2.
(a) Full matrix:
B1.1
B1.2
(89)
B2,1
B2
2
[23]
(b) Mapping to feasibility program's constraints:
Aa,2
A4.2
bsl
A5,2 AA53
A5.2
(b
- AG, b Ass33
6
by
+ by
(90)
A3,3
A4.3
A5.3
(b)
(c) We use the forth index r corresponding to = 2 to encode the first row of inequality (0). and the forth r corresponding jo 3 to encode the second row. We set the following entries of (449) (94) as 1
(55,3), (6, 4)) U ((9,5)),
(91)
the following entries to 1
((8,5)) U (() - 1,5)),
(92)
and the rest to 0, where j = 2 for D(4 and j j= 3 for ((44)
In this example, we matched each constraint by configuring the corresponding D(r) matrices along with the variable assignments for and A. Additionally, we explicitly set X46, X91 X47, X92
X48, X93, X49, X94 to zero.
Building on the illustrative reduction above, we now formalize the general reduction. The variable assignment patterns and matrix constructions introduced earlier naturally extend to arbitrary problem sizes, as detailed below.
Theorem C.14 (Theorem 3.5 restated). Let E Z,1 denote the dimension of the in-context data
(the length of the Markov chain minus one), and let m= = (d+1)(d+2) Given X E &Rdm representing the global minimizer of the reparameterized loss, solving for I-layer LSA parameters (b, A) that satisfy the reparaneterization equation Eq. 8 is NP- hard with respect to d.
36

---

Proof. We map each constant, decision and auxiliary variable from the bilinear program for the separability problem (70) the variables that for the bilinear feasibility problem (68). By freezing
the variables that were assigned a constant, we obtain program for bilinear feasibility that subsumes the objective and constraints in the separability program.
Specifically, we make assignment for the following entities in the separability program 770 (i) con-
stant: A' E IRm xn' B E xn'xn", 1 whose shape is dependent on context; (ii) decision variables: 24, 2__ 27 22 E IR;; (iii) auxiliary variables: w w² E RR", 7 7227 E R"
Decision variables and objective. We split bl Andk, 1 into two nonnegative components and keep the rest of the variables as is. Let b denote the resulting variable:
b b
-
bx+1
ba+1]
(93)
Likewise, we denote partially split version of the first column of A as follows
t. :- [A - Al,
A A AV+LL
Ast1.]
(94)
Variable assignment: We map decision variables as follows:
2 + bige 21_ + bie
+ A
22
A
(95)
D 's construction: We set D entry to 1 at the following positions and 0 otherwise:
((i, : i [K]]]
(96)
Then the objective in the feasibility program becomes
min (6 ) (AT...
A ATv,1 X1.
(97)
Since X₁ is constant by construction, minimizing the bilinear term in the feasibility program directly corresponds to minimizing the objective in the separability program (70). The same variable splitting
applied to the relevant entries of and A ensures that the nonnegativity constraints on = =1 2,, =2 are preserved.
Linear constraints. To encode each linear constraint in program (70). we construct ((r)s s and make appropriate assignments for constants and auxiliary variables.
Variable assignment: 1, We embed the constant matrices A' and B into submatrices of A, starting from the second column.
A' + - 2:14m B AT 11:21
(98)
2. We assign the entries in à after the decision variables to the auxiliary variables:
7² + b w- + by +1'k' 2² + by +1 7 + b
(99)
We further require a few entries to be fixed at or 1:
bk 43 Azn' 1,2: mmx JR) == 1, Azn 2,22+ == 0.
(100)
Counting constraints in both programs: For constraints AA'w ty ete and - ete S
0, they each contain m' inequalities. Similarly, for constraints Bw - e'e + e S 2 21 and
Bu? - 22 + e S 2 -2,, they each contain ke' inequalities. In total, we set 2(m' + k') ( to encode the linear constraints.
In the feasibility program, the index for the RHS entries (r) ranges from to dmn By the relationship between and the index for b and A, if the column index of A is incremented by 1, would increase by m. particular, any r E {(j 1)m + 1, [j 1)m + 2 2..., fm] corresponds to constraint in the feasibility program with LHS of the form D(r) A.
Dir''s ' construction; For each D(), we set four parts of its entries to nonzero, in order to match with each term in the linear constraint.
37

---

Take the first group of constraints -A'w' + + e S 0 as an example, the first row is
-Ai. w + + s 0.
(101)
By the variable assignment, the desired inequality in the feasibility program would utilize partial entries in the first column of A:
+ Aim' 2 + 42442n++1 A2n/+12 + 4 S -by+2m'43 A22 (102)
The above inequality be achieved by setting n' + entries in D(im+1) accordingly, where blue and red indicate the row and column to be set to nonzero, respectively. The rest of the rows in the first group of constraints utilizes the next m' - columns of A. Since moving column index corresponds to shifting by m, the corresponding inequalities are of the form: b Dam+1 A.J+1 S Xjm+11 for j E [m'].
Using the above method, we map each linear constraint below.
1. - A'w' + - + & O. For j E [m'], we set the following entries D(m++) as and others 0:
{(K' + k. k) & E [n'I} U {(K' + 2n' 1,22 + 1)}
U [(K' + 2n' + 2n' + 1)} U {(K' + 2n' + 2n' + +2)-
(103)
2. -A'w² + +~ + e < 0. For j E [m'], we set the following entries of D(m++) as and others 0:
[(K + + k) & E [m']} U {(K' + 2m' + 2n' + 1)}
U [(K' + 2n' 3 2n' + 1)} U {(K' + 2n' 3,, 2n' + 2)}-
(104)
3. Bw - - < #4 - 2Σ. For j E [KT we set the following entries of D(im+3) as
[(K' k, n' + k) & [n'13 {(k' + 2m' + 2n' + 1)}.
(105)
the following entries to -1
{(K' + 2n' + 1,2n' +
- -,-,,n + 1)};
1,2n'
(106)
and the remaining entries to 0. 4. Bw - 7 + < 2 - z2. For j E [k', we set the following entries of Dlim+4) as
((K' + nl ki nl + k) E [n']} u {(V' + 2m' + +3, 3,,21 2n' + 1)},
(107)
the following entries to -1
{(K' + 2n' + 2n' 2n' 11) - 1,2n' + +1)}
(108)
and the remaining entries to 0.
We solely rely on the configuration of D and the b, A to match with the constraints and set X Xjm+1, Xim+2, Xymm43, Xym+ to zero for all j E [mn'] and j' [k].
The dimension d of the feasibility problem depends polynomially the problem dimension of the separability program parameters m', n' k'. Specifically, the reduction requires that
d 2 max (k' + 2n' + 2, 1 + max(m', k)))
This reduction preserves the scaling of the problem size, and thus the feasibility problem remains NP-hard with respect to its dimension d. In particular, the computational hardness increases with the problem dimension. This trend is also reflected empirically in Fig. 5 where the performance gap between LSA and the reparameterized model widens as d increases.
Since it is potentially impossible to find transformer parameter that achieves the global minimum of the reparameterized results, we observe the following performance boundaries for single-layer LSA.
38

---

Remark C.15. We define mapping y that projects X E Rm to the parameter space:
(X) = argmin -
(109)
Here, v finds a parameter set that maps to the closest point to X under Φ. 1/ (X) is the preimage of X under Φ, if such a preimage exists. Let f be the global minimum of f. Then f(x)) S I" S [()()*)) Let P*, Q" denote the global minimizer corresponding to f". Since f is strictly convex w.r. X dd it follows that f(X*) is the lower bound for any f(o(P, Q)), including
- Q*) f(p(P ))) Therefore f(X*) S I Similarly, since f is smaller than any I(,,)), we have l S ())()))
Example C.16. As an example, for d = 2, gg
becomes
G
G
G11 G13
G11 G22
G G23
G G33
G
G12G13
G12
G G23
G12
G
G13
G13 G23
G1a
(110)
G
G22 G23
G22 G33
G
G2n G33
G
(omitting the index-separating comma and the repeating entries in the lower half triangle).
After reparameterization, the objective function can be rewritten as
2
)
27
f(x) (2. Jitt Po1.P11
Σ
x0imnt Un++1
where X E 12 denotes the concatenation of the two vectors x(()xx2)) R The gradient of w.r.t. X is
vf(X) E
(1n+1,1) )
1n 1,288
X Elyn+1 (2n+1 @g)
[In+1,22n+1,288 1,
(2n+1.2)
We obtain the global minimizer X by solving V f((**) = 0:
x1)) 0-15 0.39
0.15
0.12
2.40
-0.09]
x [0.07
-.19
-0.07
0.06 -.22 0.04]
We project (1), (2) into the model parameter space.
Since the entires of (1) are nonzero, we have by f b2.
To verify the derivation, we plot the loss function w.r.t X,, indicating the global optimizer X; using red dashed line in Fig. 10 The theoretical global minimizer aligns with the lowest error.
expected global minimizer
1.5
1.25
1.00
1-
3
2.5 0.0
2.5
-2.5 0.0 2.5
2.5 2.5
-2.5 0.0 2.5
2.5 0.0 2.5
-2.5 0.0 2.5
X-1
X.2
X.3
X.4
X.5
X.6
Figure 10: Loss function w.r.t. the first six parameters after reparametcrization.
D Proof for Forward Pass as Multi-Objective Optimization (Theorem 4.1
To demonstrate the equivalence between the forward pass and preconditioned gradient descent, we aim to express the iterative definition of L.SA as an update of weight vectors, drawing inspiration from [19.. However, unlike their approach, our proof diverges because the update formula for L.SA cannot be simplified due to the presence of nonzero entries in bj.
39

---

Theorem D.1 (Theorem 4.1 restated). Consider the L-layer transformer parameterzed by bt, A, A
where by E Rd+1, A Iddd ot E IRd for E [L]. The forward pass of the transformer is equivalent to optimizing two groups of objectives R1, R2 IR + ga+1, Specifically, let Vu+1 be the bottom-right entry of the Ith layer output. Then (1) (wi, In+1) where wy is iteratively defined as follows: wh = and
Wi+1 =wi- -b bi
(u)) (R)(wi)
(111)
ai
where Ru(w) 1
sigaary 0.5(1s (w, 23)))
Σ
(112)
n j=1
(-walys (wid 1, Fj.dd 12) + 2 2 )33
ifaj #
R2(w)=
1
Σ
n
3øj,d (U) w xz))²
1
(113)
72 j=1
(-walys (w
2 2
- - (w q j.:d ) 3
In the above expression, W.d-1, W d-:11 denote the first d 1 entries of and z respectively.
(k) denote the (d++ 1)-i entry of the embdding Zu and (k) is the first d entries of the ith column in Zz. Since the first d rows of P is zero, the first rows of Zg is the same as Zo. Therefore
Proof. Let y;
(R()) = _ Fg, Vi E [n 1].
We define mapping to represent applying k transformer layers the bottom right entry of an embedding matrix Zo with Zooa-11,441 = yf g(x, y, k) R2 R x Z + R. When x = In+ = (0)
Va+1 = 9(x, y, k) = g( 0, R) In+t- We establish two claims for g(x, y, k)) when En+1-
(k:)
Claim 1: k)) (2, 0, y. y. The equation implies that applying the transfomer & times on Zo with 2o, = y is equivalent to applying the transformer times on Z/b with [2o]a41.441 = 0 and then add the resulting bottom-right entry with y.
By definition of LSA, the iterative definition of yi
(k)
(i In + 1J) is given by:
(k)
(k+1) Vi
(N)
ZJZ
yj a y
Vi
bk
(k)
(k)2
A
(114)
n
x
Cik)
(k') Since vAa) is independent of Un+11 for any k" and Vat1 depends on v additively, one can show inductively that g(Ø, k) and g(x, 0, k) always differ by i.c., g( y, k) = g(x, 0, k) +
(k)
(k)
Claim 2: g(x, 0, k) linear in x. We prove the claim inductively. When & 0, ((x,0,0)
Unt1 b% (A)) Agent1 is linear in In+1- For k 2 0. suppose g(v, 0, k) is linear in then (k+1)
g(x, 0, h: + 1) = Mn+1 Mn+1 b (()) A&I = g(2, 0, k) - b% G(k) Ag The first term
(k)
g(x, 0, k)) linear y. The term yن with j A n + does not depend on #n+1 according to Eq. 114
Hence 6%. (1+) AςT also linear in In+1.
Combining the two claims, we have
g(r,y,k) g(2, + y= (fk, x²) + y.
(115)
for some 0g &dd with 6 0. One can copy the values in the ith column to the + 1th column and adopt the previous arguments to show that g( yr, (0c, xz) + vr. By substituting
40

---

9: = (0x. xzi + into Eq.
114
we have
Tz T
(0j + (3)
yntl Mnt1
(l+11
(k)
bk
(y5 + O )an
(U3 0% xx)²
Ag
(116)
n
100+1,7a441 (Ox, xn+
Σ
bi
zzz
(uy + 0 2))
Autnttl (117)
(yj + 0 2z)
(2) + x,)²
Since the above equation holds for any In+1, we obtain
(yj 0 ) (y1 + 0g 2z)²
8k+1
T
=0p²
le
(My + 2z)
Ak
1 n1
1 n
[yg + 0 xr )ay
0k
bk
(リا + 0 43 22)
A
Σ
(3/ + 0 xr,) )2
ak
11
72 j
GIER(441xd
EQú+1
T5Z] T (35 + + 2z)a
Old lid- ag
=0p bp
A by
G2 cR(at1ied
(118)
where if 2j + 0, then
2z On atg
0k 0 jdd- jzj
(Mj + 0 2z 2z)
G2 (y/ + 1 x3) )
(119)
(0) + 0% 2z) #s.d
(Us + O x2)²
Izid
otherwise if Ij,d + 0, then
Jz. W1.10
2j.d-1kk
(yj + 2g
G2 2(y; + 0 n. (,,
2(y; + 0 x. (r,))),,a--- Ox,dd
(y) + 02 2z))
(120)
In the above expression, 0k,5, denote the ith element of 0k. By, respectively. We treat
be: Ak: [xxd as preconditioners. We construct two sets of muli-objectives Rí, Rs Rd
ax
Rd+1 as follows. We drop the index for layer on I and use 6g to denote the dth entry of 8. Let
Tj 1, O denote the first d 1 entries of the vectors Fj, 0 respectively.
15 15 xz
3 Izj
(0,z))z, 0.5((0.z)) + )
(1)
(121)
1 IJ
0.5(0 3j 03)
(Oa(y) + (B.a11 F.:dd 25.dd1)) 1)) + 0
2 22,a)
if #
3E (3) + xz)²
m
R2( - n Σ
(Dalys + P..d1, 1)) + 82
3 5.
(122)
(vj + (0.--1,
if af =0
The derivation of (0) for the case Fd utilizes the fact that the terms marked by (i) cquivalent: 20)) = (0.a- 25,dd when zd = O.
Then (#) G, and (R2 G2 yield
Oa- 0
041 = 0 - R ( - by (R2(0)
a
o - b V Ri (0) A + (i2
O(d-1)xd
(123)
By letting w -0 and (,(w) R-(-0)
0) (i
[2]), obtain the desired result.
41

---

The first d terms in R1 (1) (=Ij, 2, for lo [d] capture the interaction between the k-th state of zj and the linear combination states determined by w, emphasizing their individual roles within the sequence. The final term, (w Iy y3) ensures alignment between the linear model's prediction and the target state Vj-
(w(w) emphasizes the role of was since the first d - rows of the last preconditioning matrix zeros. The first d objectives in R2 (w) scale the target state yj by æ and wa, with additional quadratic wi zjdd- These terms capture the influence Wa, j. the alignment between Vj and partial terms like 2
prediction (w:d-1 1)d-11) The final objective is in cubic penalty form 3zy.a (y) - w x,)² when Fjdd * 0, magnifying sequences with smaller 2j.d- When žj 0, the penalty changes
quadratic
term- -w; (Wid 1; -.j.d-1)) focusing solely on aligning Dj with the partial prediction based on 1- Furthermore, when Fj,d = 0, the final objective becomes convex if wa < and concave otherwise.
E Additional Related Works
A growing line of research investigates the structural and computational limitations of shallow transformer models. Rajaraman et al. [24] examine ICL on Markovian data and show that while l-layer transformers may struggle to capture high-order dependencies, adding one two layers enables accurate modeling of k-th order Markov processes. Olsson et al. 45 demonstrate that induction heads do not arise in 1-layer transformers, and that their emergence and the corresponding gains in ICL performance are only observed in models with two or more layers. Separately, Sanford et al. 37 38 study the representational and computational limits of transformers from theoretical perspective. While not directly focused on ICI., they characterize tasks that require deeper models or higher attention capacity solve, and establish lower bounds by connecting transformer depth to parallel computation. Our work complements these studies from an optimization perspective, where we show that even when global minimizer exists in relaxed formulation, recovering feasible parameterization in the 1 -laayer transformer space is NP-hard.
42

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The three main claims are supported by Theorems 3.3 and 3.4 Theorem
Theorem
3.5 and
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the
results can be expected to generalize to other settings.
is fine to include aspirational goals motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The limitation is discussed in Section 5
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model well-specification,
asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't t acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer:| Yes]
Justification: The assumptions of each theoretical result is stated within the result. The results are followed by proofs detailed in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
43

---

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The access of code and data is detailed in Section B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
example, if the contribution
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e.g.. a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g.. with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome to describe the particular way they provide for reproducibility. In the case of closed-sous source models, it may be that access to the model is limited in some way (c.g.. to registered users). but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have provided the data and source code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (httts::/nips.cc/ cctpsbiic/
for more details.
While we encourage the release code and data, we understand that this might not be possible,
"No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed run reproduce the results. See the NeurlPS code and data submission guidelines (https: / /wpp. cttpsccccpbcccccccc SesiidesCodeSSur
onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
44

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why. submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits. hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experimental details are introduced in Section 4 and Appendix A Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is
necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The experiments are repeated over 5 random seeds and the mean and standard deviation of the metries are reported.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors). It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of eтrors is not verified.
For asymmetric distributions, the authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
Experiments compute resources
Question: For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time of execution) needed reproduce the experi- ments?
Answer: [Yes]
Justification: The device detail is included in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud
provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e-g., preliminary or failed experiments that didn't 't make into the paper).
Code of ethics
45

---

Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Question:
Answer: [Yes]
Justification: The research presented in this paper adheres to the NeurIPS Code of Ethics. It does not involve human subjects or sensitive personal data.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require a deviation
from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there a special consideration due
to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work focuses on theoretical and empirical analysis of in-context learning using synthetic data and simplified model architectures. does not involve real-world deployment, user data, or application-specific systems.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to par- ticular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, is legitimate to point out that improvement in the quality of generative models could be used generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: This paper does not release any models or datasets, and the research is purely theoretical in nature. It focuses on the analysis of loss landscapes and representational limits of simplified transformer models under synthetic settings. such, poses no foreseeable safety or misuse risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
46

---

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: The paper simulates synthetic data according to certain distributions and does not use existing datasets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version the asset used and, if possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information
available online, the authors are encouraged reach out the asset's
creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation
provided alongside the assets?
Answer: [NA]
Justification: The paper does not create new synthetic or real dataset.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their sub-
missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification; This paper does include any crowdsourcing experiments or research involving human subjects. All experiments are based on synthetic data.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should included in the main paper.
According to the NeurlPS Code Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This study does not include experiments involving human subjects, participants, or user studies that would necessitate IRB approval. The research is entirely computational, relying on synthetic data
Guidelines:
47

---

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equívalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble). such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non- standard component of the core methods in this research? Note that if the L.M is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required. Answer:[NA]
Justification: LLMs
involved in the methodology, experimental design, scientific contributions of this work. The use of LLMs was restricted to writing support and had no influence on the research's originality or scientific integrity. Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, or non-standard components. Please refer our L.L.M [ttpge///beerrroccccc////eeerrre should or should not be described.
for what
48


Reviews:
[{"review_id": "mgIMYlgYTX", "reviewer": "Reviewer_MTmD", "summary": "The paper investigates how transformers exhibit in-context learning (ICL) when trained on data generated from Markovian processes. It also identifies the global minima for such data sources in the case of a one-layer linear self-attention model. Additionally, the authors demonstrate that recovering the transformer's parameters could be NP-hard. Finally, they show that multiple layers can be interpreted as preconditioned gradient descent operating with multiple objectives.", "questions": "**Questions:**\n\n1. Could the authors please clarify what the solutions in Theorem 3.3 and subsequently in Theorem 3.4 represent? For example, assuming the data generation process described between lines 133–139, it seems that estimating the correct downstream output, even in a regression settings, should ideally involve predicting the mean (if I’m not wrong). Is that what the solution actually learns? If not, could the discrepancy stem from the formulation, the nature of the solution, or some other factor? A clearer interpretation would help in understanding the practical implications of these results.\n\n2. In the proof of Theorem 3.3 (and by extension Theorem 3.4), the authors simplify the LSA considerably and eventually reduce the problem to one of linear regression. While this helps in obtaining a tractable solution, the actual optimization landscape of $P$ and $Q$ appears quite different from the simplified setting (even if the global optima coincide). Understanding the behavior of the original optimization, particularly in terms of $P$ and $Q$, could provide additional insights. Would it be possible to characterize the loss landscape with respect to $P$ and $Q$ more directly? Moreover, how might the optimization dynamics change in the original setup compared to the simplified one?\n\n3. I found the discussion around the problem being NP-hard a bit difficult to follow. It would be helpful if the authors could reiterate the main arguments more clearly. Also, please correct me if I'm misunderstanding, but the proof seems to present a generalization of Theorem 3.3 with a larger number of parameters. If that's the case, is the idea that we can reduce the problem to solving a linear regression, much like in Theorem 3.3, but in order to recover $P$ and $Q$, we ultimately need to solve a system of polynomial equations (and solving such systems is known to be NP-hard, as discussed in [E])? If this interpretation is correct, wouldn’t that offer a more straightforward sketch of the hardness argument?\n\n4. As a follow-up to Problem 3, I remain a bit skeptical about the result. If my interpretation of why Problem 3 is NP-hard is accurate, then it seems that the hardness stems from the specific way the authors define optimality. I'm asking this because many problems in machine learning are technically NP-hard due to non-convexity [F], yet we still manage to find approximate solutions using gradient-based methods in practice. So even if the problem is NP-hard in the worst case, is it possible to converge to a reasonable solution using standard descent algorithms? Some clarification on this point would be helpful.\n\n5. Could the authors elaborate on what the multiple objectives in Theorem 4.1 are, and what they intuitively represent? Understanding the nature and interpretation of these objectives would help clarify the broader implications of the result. Additionally, it would be helpful if the authors could walk through the intuitions behind the explanations outlined in lines 328–332\n\n\n**Note:** I find the analysis in this work to be interesting. However, my main concern is regarding the applicability of the results. Much of the analysis appears to depend on very simplified settings, and the assumptions made are quite strong and seem to diverge significantly from practical scenarios. Based on these concerns, I am currently assigning a lower score. That said, I am open to revisiting my evaluation during the rebuttal phase, depending on the authors' clarifications and the views expressed by other reviewers.\n\n**Minor Typos:**\n\n1. Line before line 996 - the underbace should not be over the $1/n$ term.\n\n2. Line 1045 - missing bracket in $U(0, 1)$\n\n---\n\n**References:**\n\n[A] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a Transformer: A Memory Viewpoint. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.\n\n[B] Ezra Edelman, Nikolaos Tsilivis, Benjamin Edelman, Eran Malach, and Surbhi Goel. The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains. In Advances in Neural Information Processing Systems, 37:64273–64311, 2024.\n\n[C] Eshaan Nichani, Alex Damian, and Jason D. Lee. How Transformers Learn Causal Structure with Gradient Descent. In Forty-first International Conference on Machine Learning, 2024. URL\n\n[D] Nived Rajaraman, Marco Bondaschi, Ashok Vardhan Makkuva, Kannan Ramchandran, and Michael Gastpar. Transformers on Markov Data: Constant Depth Suffices. In Advances in Neural Information Processing Systems, 37:137521–137556, 2024.\n\n[E] https://bpb-us-e1.wpmucdn.com/sites.harvard.edu/dist/7/800/files/2024/10/spring2012-st12-scribe-lect14.pdf\n\n[F] Blum, Avrim, and Ronald Rivest. \"Training a 3-node neural network is NP-complete.\" Advances in neural information processing systems 1 (1988).", "limitations": "None.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n\n1. The paper tries to tackle a very interesting problem which is ICL on Markov chains.  \n\n2. The authors have tried to contextualize their work very well with regards to other work. \n\n3. The authors have also tried to extend the conventional regression-based theory for ICL tasks to a new domain which is Markov chains and this indeed is a challenging problem.\n\n**Weaknesses:**\n1. I feel somewhat uncertain about the applicability of this work. The problem setting is certainly interesting, especially in light of prior studies such as [A], [B], [C], and [D], which explore similar questions using the actual architectures employed by transformers in practice. In contrast, the current setup comes across as a bit contrived. The paper appears to build on the idea that some existing studies interpret in-context learning (ICL) through regression-based tasks and then explores how this perspective might extend to Markov chains. However, when focusing on Markov chains, particularly in the context of ICL, the primary objective typically involves next-token prediction. In that regard, a more natural setup already appears in works like [A], [B], [C], and [D], which the authors have appropriately cited.\n\n2. The authors have clearly invested substantial effort in characterizing optimal solutions. However, they introduce several simplifications that make it harder to assess how applicable the findings are in practice. Additionally, the paper does not clearly convey the intuitive meaning of the global solution. Please see the questions for a follow-up\n\n3. The paper seems to leave out the intuitions behind many of its theorems and propositions. While this is not necessarily a shortcoming, it gives the impression that the focus leans more toward mathematical formalism than toward building a qualitative understanding of the results. For instance, in Theorem 4, prior work such as [A], [B], [C], and [D] might lead one to expect that adding multiple transformer layers would yield a counting-based estimator (even though the settings are different). In this setup, however, it's not immediately clear why multiple layers of LSA would result in solving a multi-objective optimization problem. It would be helpful to develop more intuition around why this behavior arises and what gives rise to the presence of multiple objectives.", "was_revised": true, "final_justification": "I thank the authors for their comments and efforts throughout the rebuttal process. After carefully weighing the pros and cons, I have decided to increase my rating from **3** to **4**.\n\n**Pros:**  \n1. The paper is generally well-written and includes a substantial number of experimental results.  \n2. The authors did an excellent job during the rebuttal in clarifying my concerns, particularly regarding the proofs and the NP-hardness result.\n\n**Cons:**  \n1. The setting feels somewhat artificial and does not align well with common practices for training on sequential data.  \n2. To the best of my knowledge, the paper does not address learning dynamics, which is a concern given the reliance on the NP-hardness argument.  \n3. The multi-objective argument supporting Theorem 4 is relatively weak, and the preconditioner argument is also unconvincing.\n\n**Overall Assessment:**  \n\nHaving said this, making progress in theoretical machine learning is inherently challenging and is also important if we are to deepen our understanding of large language models in the future. Moreover, to reiterate, the authors did a phenomenal job during the rebuttal, addressing concerns with clarity and professionalism. Hence, I recommend acceptance.", "rebuttal": "We sincerely thank the reviewer for the thoughtful and detailed feedback. We aim to clarify the concerns including the motivation behind our design choices and the interpretation of our results below.\n\n## 1. Setup clarification \n> **W1**: Concerns about the applicability of the problem setup compared to prior works [A] [B] [C] [D] that study next-token prediction for Markov chains.\n- **Relationship with prior studies** \n    - **Extending regression-style tasks.** We extend the ICL framework applied in [1–7] to study a new class of functions that has not been previously considered. In these works, the underlying function class is often linear and the input data is assumed to be i.i.d. and zero-centered. In contrast, we focus on settings where the input-output relationship is governed by a dynamical function with non-iid input. \n    - **Differences within Markov chain studies.** While studies in [A-D] (as cited by the reviewer) focus on next token prediction for probabilistic sequences, they also extend to model dependence between many tokens and within sequences with different Markov chain mechanisms, such as higher order Markov chains and Markov chains with causality.   \n    - **Incomparable approaches.** We note that [1–7] and [A–D] approach the ICL problem from distinct perspectives, and their setups are not directly comparable. Specifically, [A–D] aim to predict the next token in a sequence whereas our goal is to recover a function using in-context examples. \n- **Reusable framework for probing expressivity.**  We offer a general approach to quantify and understand the limitation of 1-layer LSA. We reparameterize the LSA objective $f$ into a linear model $\\tilde{f}$ over an enlarged input space (Eq. 8) and characterize its global optimum X* (Lemma 3.2). We then show that bridging the gap between 1-layer LSA and the best achievable predictor under $\\tilde{f}$ is computationally hard (Theorem 3.5). This analysis is applicable to other ICL function classes that induce a dense structure on the optimum X*, where most entries in X* are nonzero, which might lead to NP-hardness in finding the corresponding optimal LSA parameters. This insight helps clarify what kinds of structured inference tasks that are challenging for shallow LSAs.\n \n## 2. Global minima and reparameterization \n\n> **W2**: Unclear interpretation of global minimum and its relation to mean prediction. \n> **Q1**: What the solutions in Theorem 3.3 and 3.4 represent? \n- **Role of global min derivation.**  Theorems 3.3 and 3.4 characterize the global minimum of the 1-layer LSA in the tractable case $d=1$. They show that mapping from a dense reparameterized optimum X* to transformer parameters (b*,A*) results in dense solutions, in contrast to the sparse structures typical in ICL for linear tasks, thereby motivating the construction of a broader parameter class (Eq.11) for which we derive the multi-layer LSA interpretation.\n- **Empirically,** we showed in Fig. 5 that for binary state space with 100 in-context examples and $d=1$, the accuracy of this global minimum is ~60%. This relatively low accuracy is due to the limited input dimension, and it improves as $d$ increases.\n\n> **W2**:  Introduced several simplifications (reparameterization) that make it harder to assess how applicable the findings are in practice. \n> **Q2**: Is it possible to characterize the loss landscape more directly? Moreover, how might the optimization dynamics change in the original setup compared to the simplified one?\n- **Usage of reparameterization.** The reparameterization simplifies the optimization while preserving expressivity, as the linear model fully captures the 1-layer LSA. Such technique has been applied in recent studies [5,8] to study the loss landscape for  ICL of linear tasks.  \n    - **New discovery under this reparameterization.** Under linear tasks, the global minimum of the reparameterization can be mapped back to the transformer parameter space. For more complex function class and non-iid input data, finding the corresponding transformer parameters becomes NP-hard, which sets our task apart from their setup.\n- **Impact on training dynamics.** The reparameterized objective $\\tilde{f}$ is strictly convex (Lemma 3.1), ensuring convergence to its unique global optimum via gradient descent. In contrast, the original LSA objective $f(b, A)$ is non-convex due to bilinear parameter coupling, which can lead to suboptimal convergence depending on initialization. Empirically, we observe that training over $\\tilde{f}$ reaches lower loss more quickly, likely because it avoids saddle points and plateaus present in $f$’s landscape.\n\n \n## 3. NP-Hardness interpretation  \n> **Q3**: NP-hard proof unclear; appears to generalize Theorem 3.3 with large $d$; why not directly reduce to known NP-hard polynomial system (e.g., [E])?\n\nWe refer the reviewer to our response to W1 for a full walkthrough of the reparameterization-to-hardness chain.\n- **Reduction justification.** The NP-hardness result in [E] (as cited by the reviewer) is established for general polynomial systems. Since the bilinear system in Theorem 3.5  is a subclass of the general system, this result cannot be directly inherited. Therefore, we construct a dedicated reduction specific to our setting.\n    \n> **Q4**: Is NP-hardness due to how optimality is defined? Many problems are NP-hard due to non-convexity and gradient-based methods help find approximate solutions in practice. Can standard training provide a reasonable solution even if the problem is NP-hard in the worst-case?\n- **Theoretical perspective:**\n    - The NP-hardness does not arise from the way optimality is defined, but from the non-convexity of the 1-layer LSA. The difficulty of solving the equation $\\phi(b, A) = X$ stems from $\\phi$ is bilinear, i.e., it couples the parameters b and A multiplicatively. This bilinear structure arises directly from the 1-layer LSA architecture, where the forward computation depends linearly on the product terms $b_i A_{j,k}$, rather than separately on $b_i$ and $A_{j,k}$.\n- **Empirical perspective:**\n    - **For 1-layer LSA:** As shown in Fig. 5 (Section B.1), as the sequence length d increases, the performance gap between the 1-layer LSA trained with gradient descent with 1000 epochs and the reparameterized optimum X* becomes larger. This aligns with our NP-hardness result, which suggests that recovering the global optimum becomes intractable as $d$ increases. \n    - **For deeper LSA:** Fig. 3 in the main text shows that adding layers improves accuracy to 80~90%, suggesting that increased expressivity helps overcome the limitations we identify for shallow architectures. \n    - **Modern LLMs.** Empirical works [9] have observed the emergence of ICL behavior in deeper transformers trained on language datasets, up to 40 layers and 13B parameters, while such capabilities are largely absent in shallow models with fewer than 8 layers. Our work provides a theoretical lens for this type of observation by framing expressivity as an optimization problem, explaining why shallow models struggle to represent certain function classes in context.\n\n## 4.  Multi-Layer behavior  \n> **W3**: Lack of intuitive explanation for the multi-layer forward pass and the emergence of multiple objectives in Theorem 4.\n> **Q5**: What do the  objectives in Theorem 4.1 represent intuitively, and how does the observed depth-dependent optimization behavior arise?\n- **Emergence of multiple objectives.** The parameter assumption in Eq. 11 requires certain entries in the LSA weights $P, Q$ to be nonzero, with their locations matching the global minimum structure in Theorems 3.3&3.4, though their values remain unrestricted. This construction allows us to interpret the forward computation as iterative steps of preconditioned gradient descent, where the emergence of multiple objectives arises from the structured sparsity in $P$ and $Q$. \n- **Interpretation.** At a high level, $R_1$ captures standard regression behavior, encouraging the linear model to align with the training data through squared error and weighted contributions from each input coordinate. $R_2$ introduces a cubic-order correction that places additional emphasis on the last input dimension.  \n- **Explanation of depth-dependent behavior (Fig.2).** To verify the multi-layer interpretation, we fixed the LSA parameter obtained by gradient descent training. The objectives are evaluated for this fixed LSA at each layer.  We observed that initially, the model makes progress on both $R_1$ and $R_2$, but beyond a certain depth, optimization deteriorates, all losses begin increasing. The non-monotonic decreasing is likely due to deeper LSA preconditioners  negatively affecting the performance. \n\n**Minor typos.**\nThank you for pointing out these typos. We will correct the misplaced underbrace and add the missing bracket in the final version.\n\n---\n**References**\n\n[1] Johannes von Oswald et al., Transformers Learn In-Context by Gradient Descent. ICML 2023 \n\n[2] Ruiqi Zhang, et al., Trained Transformers Learn Linear Models In-Context. J. Mach. Learn. Res. 25: 49:1-49:55 2024\n\n[3] Arvind V. Mahankali et al., One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention. ICLR 2024\n\n[4] Kwangjun Ahn et al., Linear attention is (maybe) all you need (to understand Transformer optimization). ICLR 2024 \n\n[5] Kwangjun Ahn et al., Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023\n\n[6] Yu Bai et al., Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection. NeurIPS 2023\n\n[7] Hongkang Li et al., How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? ICML 2024\n \n[8] Yedi Zhang, et al., Training Dynamics of In-Context Learning in Linear Attention, ICML 2025.\n\n[9] Catherine Olsson et al. In-context Learning and Induction Heads. 2022"}, {"review_id": "RZsAz4EI1z", "reviewer": "Reviewer_SbDi", "summary": "The paper studies how do transformers in-context learning markovian chains by proving that there exists (reparametrized) transformer parameters which are global optima of the objective loss function. Additionally, it shows that for an actual one-layer LSA to achieve the reparametrized parameter is NP-hard. Moreover, it demonstrates that an n L-layer LSA can be interpreted as preconditioned multi-objective optimization.", "questions": "Q1: I do not fully comprehend the motivation behind studying the ICL ability on Markov chains. If the authors want to study non IID tasks, dynamical systems would have been enough, so could the authors clarify on the necessity to study Markov chains?\n\nQ2: While Theorem 3.3 & 3.4 states the global minima under the reparametrized objective, Theorem 3.5 admits the NP-hardness to reconstruct transformer parameters from reparametrization. Does it mean that a one-layer transformer can or cannot learn Markov chain?\n\nQ3: For the multi-layer setting (Theorem 4.1), why does the $(x,y)$ pair suddenly change to linear regression, and what does the two objectives $R_1$ and $R_2$ mean from a high level?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\n1. The paper studies the ICL ability of transformers on Markov chains which distinguishes itself with previous works on IID tasks like linear regression.\n \n2. The theoretical aspect of the paper is solid, and some experiments are provided to support the conclusions despite being a theoretical paper.\n\n## Weaknesses\n\nPlease see Questions.", "was_revised": true, "final_justification": "During the rebuttal period my main questions are solved and it appears that I have misunderstood the data representation in Q3 and thank you for the clarification. Also I have surveyed works on the explanation of ICL Markovian process and admits its significance.  Overall the paper is a good and natural extension from previous works of in-context gradient descent on linear regression to Markov chains.", "rebuttal": "We are grateful to the reviewer for raising insightful questions. We address the key concerns about the role of Markov chains,  the interpretation of our NP-hardness result and multi-layer analysis in the following responses.\n\n## 1. Setup clarification \n> **Q1**: Why study Markov chains instead of general dynamical systems for non-i.i.d. ICL \n\nThank you for your suggestion on general dynamical system models. Indeed, it has been our general purpose to replace IID model with a more realistic and structured stochastic dynamical system model for analyzing transforms whose inputs and outputs can be characterized as probability measures.  \n- **Modeling power of Markov models.**  Markov processes, due to its strength in modelling capability/flexibility and extensive reservoirs of theory and technical methods developed, has been an important model for stochastic dynamical systems, see e.g. references [1] and [2] for their treatment in widely used textbooks in the area. Hence, it is natural to introduce a Markov chain, a Markov processes with discrete time and finite state space, in this study. This choice provides a controlled way to examine how Transformers approximate latent temporal structure, making it a natural step toward understanding ICL on more general function classes.\n \n## 2. Global minima to NP-hard results  \n> **Q2**: While Theorem 3.3 and 3.4 states the global minima under the reparametrized objective, Theorem 3.5 admits the NP-hardness to reconstruct transformer parameters from reparametrization. Does NP-hardness mean 1-layer transformers can’t learn Markov chains?\n\nThe NP-hardness result does not mean 1-layer transformer can’t learn Markov chains, but finding the best possible parameters is computationally hard in the worst-case.\n- **Tractable vs NP-hard case.** Theorems 3.3 and 3.4 provide a closed-form solution for recovering (b, A) from X* when d = 1, but this tractability breaks down for larger d, where the bilinear structure of $\\phi(b, A)$ induced by the 1-layer LSA becomes more complex.\n- **NP-Hardness' implication for performance.** As shown in Fig. 5 (Section B.1), when the sequence length $d$ increases, the performance gap between the 1-layer LSA trained with gradient descent (1K epochs) and the reparameterized optimum $X^*$ becomes larger. While the reparameterized model achieved 95% accuracy for $d=10$, the corresponding 1-layer LSA trained with 1K epochs only reached ~70%. This aligns with our NP-hardness result, which suggests that recovering the global optimum becomes intractable as $d$ increases. \n \n## 3. Multi-objective interpretation  \n> **Q3**: In Theorem 4.1, why does the forward pass reduce to a linear regression problem, and what do the two objectives R_1 and R_2 represent at a high level?\n- **Linear model origin.** As detailed in Appendix Section D, we show that for an $L$-layer LSA, the final prediction $\\hat{y}\\_{n+1}^{(L)}$ can be expressed as a linear function of $x\\_{n+1}$, i.e., $\\hat{y}\\_{n+1}^{(l)} = \\langle w\\_{l}, x\\_{n+1} \\rangle$. This holds for any prompt instance and follows directly from unrolling the forward pass. Given this linear form, we analyze how $w_l$ evolves across layers and show that it follows a preconditioned gradient descent path on two sets of objectives, $R_1$ and $R_2$.\n- **Objective interpretation.** At a high level, $R_1$ captures standard regression behavior, encouraging the linear model to align with the training data through squared error and weighted contributions from each input coordinate. $R_2$ introduces a cubic-order correction that places additional emphasis on the last input dimension. \n\n**References**\n\n[1] L. Arnold, Random dynamical systems, Springer, 1998.\n\n[2] R. Bhattacharya and M. Majumdar, Random dynamical systems: theory and applications, CUP, 2007."}, {"review_id": "yb1x8uluAA", "reviewer": "Reviewer_Fq3M", "summary": "This paper investigates how transformers perform in-context learning when the data arises from structured Markovian dynamics, rather than i.i.d. distributions. Focusing on linear self-attention (LSA), the authors derive a closed-form global minimizer for short Markov sequences and show that finding transformer parameters that realize this optimum is NP-hard in general—highlighting a fundamental limitation of one-layer LSA models. They further show that multilayer LSA models can overcome this expressivity barrier by implicitly performing a preconditioned multi-objective optimization. The theoretical results are supported by empirical validation, offering a novel perspective on how architectural depth enables transformers to generalize across structured, temporally dependent inputs.", "questions": "I would consider increasing my score if most of the following questions are clearly addressed. Please also refer to the weaknesses section.\n\nQuestions:\n\n- Your analysis assumes linear self-attention and binary first-order Markov chains. How do your results generalize to softmax attention, multi-head architectures, or higher-order/non-binary dynamics?\n- The NP-hardness result is worst-case. Do you expect this computational hardness to appear in typical learning scenarios? Are there cases where parameter recovery remains tractable?\n- The multi-layer analysis relies on a structured parameter space. Do real-world trained transformers actually follow this structure? Can you justify that the multi-objective interpretation holds more generally?\n- The experiments focus on simple synthetic tasks. Can you share more details or actually main-text results for the GPT-2 experiments, and can you evaluate on more realistic sequence data and modern models?", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The paper provides a rigorous analytical characterization of in-context learning for Markovian dynamical tasks.\n- The authors prove that recovering the transformer parameters that achieve the optimal in-context performance on general Markovian functions is an NP-hard problem (Theorem 3.5). This is a notable finding that reveals a structural limitation of single-layer transformers, as there exists a gap between the optimal structured predictor and what a one-layer LSA can represent or learn in practice. \n- The paper presents a novel interpretation of multi-layer transformers as performing a preconditioned multi-objective optimization in their forward pass (Theorem 4.1). In essence, a deeper LSA model is shown to implicitly minimize a main squared-loss objective while simultaneously maximizing several linear objectives.\n- Rigor and clarity of the analysis overall.\n\nWeaknesses:\n\n- The theoretical analysis is confined to a highly simplified setting – one-layer linear self-attention models and first-order binary Markov chains. While this abstraction is necessary for analytical tractability, it limits the direct applicability of the results.\n- The NP-hardness result, while theoretically important, is a worst-case complexity statement. It shows that exactly recovering the optimal parameter solution in the general Markovian setting is intractable, but this doesn’t necessarily mean that learning such tasks is impossible in practice, for example heuristic training or approximate solutions might perform well on typical instances. The paper could discuss more about the practical implications of this hardness.\n- The multi-layer LSA analysis (Theorem 4.1) is conducted under a special parameterization that “mirrors” the single-layer optimum’s structure. Essentially, the authors restrict the model to a certain subspace of parameters (Eq. 11 in the paper) to derive the multi-objective interpretation. This raises a question of how general the insight is as real transformers trained without such constraints might deviate from this idealized behavior.\n- While the paper does include experiments, they are relatively limited in scope. The main text only reports synthetic data results for small state spaces, and points to additional experiments with standard transformer architectures (GPT-2) in the appendix.", "was_revised": true, "final_justification": "After considering the rebuttal and discussion, I have decided to maintain my original evaluation. This paper offers a rigorous and well-presented theoretical analysis of in-context learning for Markovian dynamics, including an NP-hardness result for recovering the optimal one-layer LSA parameters and an interpretation of multi-layer models as performing preconditioned multi-objective optimization. The authors provided thoughtful clarifications and additional experiments. Nonetheless, several points remain only partially addressed: the generalization of the analysis beyond one-layer LSA and binary Markov chains is left for future work, the practical implications of the NP-hardness result are not fully explored, and the multi-layer interpretation depends on a structured parameterization that may not emerge in real-world settings. The contributions are technically solid and of theoretical interest, but the remaining scope limitations lead me to keep my borderline accept rating.", "rebuttal": "We thank the reviewer for the constructive feedback and careful reading of our theoretical results. Below, we discuss the  relevance of our analysis to more general settings and  clarify the intuition behind our NP-hardness result and multi-objective formulation.\n\n## 1. Restricted assumptions in problem setup  \n> **W1**: Analysis assumes 1-layer linear attention and binary first-order Markov chains with limiting direct applicability to richer architectures or dynamics.\n\n- **Theoretical tractability.** We would like to point out that one-layer linear self-attention models are widely adopted architectures in prior ICL studies [1,2,3], enabling clear theoretical insights.\n- **Generalizability.** Moreover, the limitation result of 1-layer LSA is not restricted to the current data distribution, i.e., binary state space or memory order of Markov chains. As shown in Lemma C.8, the reparameterized global optimum exists for Markov chains with arbitrary state spaces and trajectory lengths, yet shallow Transformers still fail to represent this optimum.\n \n> **Q1**: How do results generalize to softmax attention, multi-head architectures, or higher-order/non-binary dynamics?\n- **Scope of our study.** Our analysis focuses on LSA and first-order binary Markov dynamics, which offer an analytically tractable setting to study expressivity and optimization barriers in ICL. \n- **Novelty of our analysis.** To the best of our knowledge, this work takes an initial step toward understanding how ICL models handle structured function classes. Our work provides a theoretical lens by framing expressivity as an optimization problem, explaining why shallow models struggle to learn certain function classes in context. \n- **Extensibility.** This line of analysis offers a potential starting point for understanding expressivity limitations in broader function classes.  Extending our framework to cover these richer settings is an exciting direction we hope to pursue.\n \n## 2. NP-hard result interpretation\n> **W2 & Q2**: Heuristic training or approximate solutions may work well in typical cases. Will computational hardness show up in typical learning scenarios? Any tractable cases?\n- **For 1-layer LSA:** As shown in Fig. 5 (Section B.1), as the sequence length d increases, the performance gap between the 1-layer LSA trained with gradient descent with 1000 epochs and the reparameterized optimum X* becomes larger. This aligns with our NP-hardness result, which suggests that recovering the global optimum becomes intractable as $d$ increases. \n- **For deeper LSA:** Fig. 3 in the main text shows that adding layers improves accuracy, suggesting that increased expressivity helps overcome the limitations we identify for shallow architectures. For the same task and number of training epochs, standard training can succeed in practice when the model class is sufficiently expressive to begin with.\n- **Tractable cases:** Theorems 3.3 and 3.4 show that for d=1, the inverse mapping (b, A) that achieves a given X* can be computed in closed form. However, for larger d, $\\phi(b, A)$ becomes increasingly entangled due to bilinear interactions, and recovering parameters consistent with a given X* becomes intractable in general (Theorem 3.5).\n\n## 3. Generality of multi-Layer analysis  \n> **W3**: Limited generality of the multi-layer analysis due to structured parameterization. \n> **Q3**: Do real-world trained transformers actually follow this structure? Can you justify that the multi-objective interpretation holds more generally?\n- **Parameter assumption justification.** The parameter assumption in Eq. 11 requires certain entries in the LSA weights $P, Q$ to be nonzero, with their locations matching the global minimum structure in Theorems 3.3 and 3.4, though their values remain unrestricted. This construction allows us to interpret the forward computation as iterative steps of preconditioned gradient descent, where the emergence of multiple objectives arises from the structured sparsity in $P$ and $Q$.\n- **Generality.** Our analysis provides an analytically tractable case that highlights how deeper layers can improve performance by coordinating optimization across objectives using each layer's parameters. While we do not claim that real-world trained transformers follow this specific parameterization, we hope this controlled setting offers some intuition that may transfer to more realistic architectures.\n\n## 4. Experimental scope  \n> **W4**: Limited experimental scope and missing main-text results for larger models\n> **Q4**: Can you share more details or actually main-text results for the GPT-2 experiments, and can you evaluate on more realistic sequence data and modern models?\n\nDue to space constraints, we included the GPT-2 experiments in Appendix B.2. These experiments are designed to complement our theoretical analysis by evaluating how deeper, nonlinear transformer models perform on the same ICL task.\n- **GPT-2 Experiments.** We compare GPT-2 variants (with 3, 6, and 12 layers, respectively) on our synthetic benchmark. Fig. 6 shows deeper models consistently achieve higher accuracy in predicting the final token of a Markov chain, confirming the theoretical insight that multi-layer attention improves performance. Notably, performance with 6 and 12 layers becomes comparable to classical models like SVM and logistic regression and outperforms 3-NN baselines. \n    - **Attention maps.** To probe model behavior, we visualize attention maps across layers. Fig. 7 shows the emergence of block-wise attention, with deeper layers increasingly focusing on structurally similar in-context sequences governed by the same transition kernel.\n- **Additional experiments.** We extended our experiments to deeper models and observed that performance continues to improve up to 6 layers. We fix the total number of prompts as 1000, number of in-context examples as 100 and consider binary Markov chains with prefix length $d=5$. We perform gradient descent for 5000 epochs for all layers. The accuracy and MSE w.r.t. # layers is shown below. These results support the intuition that additional layers can mitigate the expressivity limitation. We hypothesize that the decrease at 5 layers may be due to optimization instability. Nonetheless, the overall trend supports the benefit of increased depth.\n\n| # Layers | Accuracy   | MSE  |\n| --- | --- | --- |\n| 1 | 0.681 | 0.225 |\n| 2 | 0.687 | 0.216 |\n| 3 | 0.713 | 0.200 |\n| 4 | 0.747 | 0.188 |\n| 5 | 0.681 | 0.218 |\n| 6 | 0.784  | 0.166 |\n\n- **Relationship to modern LLMs phenomenon.** Empirical works [4] have observed the emergence of ICL behavior in deeper transformers trained on language datasets, up to 40 layers and 13B parameters, while such capabilities are largely absent in shallow models with fewer than 8 layers. Our work provides a theoretical lens for this type of observation by framing expressivity as an optimization problem, explaining why shallow models struggle to represent certain function classes in context.\n  \n \n**References**\n\n[1] Kwangjun Ahn, et al., Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.\n\n[2] Ruiqi Zhang, et al., Trained Transformers Learn Linear Models In-Context. J. Mach. Learn. Res. 25: 49:1-49:55 2024.\n\n[3] Jianhao Huang, et al., Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought. ICLR 2025.\n\n[4] Catherine Olsson et al., In-context Learning and Induction Heads. 2022"}, {"review_id": "4xSaXxQshL", "reviewer": "Reviewer_88Bq", "summary": "This paper studies the in-context learning (in)ability of transformers for Markovian dynamical functions where the first $d$ $x_i$ variables are drawn from a Markov process and the $y_i$ is the $d+1$ element of the Markov chain. Each context is a distinct set of transition probabilities. The authors show that under a reparameterization of one layer linear self attention, they can identify a global minimizer for the ICL population risk. They then show that it is NP hard for one to find a one layer transformer configuration that achieves a given reparameterized configuration $X_\\star$. They do this by mapping the reparameterization to an NP hard bilinear optimization problem. They then show that deeper transformers with linear self attention can be interpreted as gradient descent on a mixture of two objective functions for any given context. They conclude with empirical simulations showing that performance improves as the depth of a model increases.", "questions": "1. Would the improvements to the performance on this task continue if you keep increasing the depth? Have the authors tried training models deeper than 3 layers? Is the global minimizing loss / Bayes error /  best possible accuracy computed? \n2. Could the authors clarify how what the NP hardness result indicates about learnability of this ICL problem in one layer LSA? The authors state \"the 1-layer transformer’s representational capacity may be insufficient to express dynamical functions induced by structured inputs, such as those governed by shared Markovian dynamics.\" I understand that finding a $P,Q$ for a given $X$ may be NP hard, but couldn't GD find a solution that may be very close to $X_\\star$ in performance in principle? I am not an expert in this style of result so any additional clarification would be useful to me.", "limitations": "The authors do a good job addressing potential limitations.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths\n\n1. This paper introduces a very novel and interesting setting for ICL that involves dense correlations across tokens. \n2. The authors provide several interesting results about LSA transformers on this problem. \n3. They also provide an interesting simulation in the experimental validation section. \n\nWeaknesses \n1. It is unclear if the NP hardness result necessarily implies that the one layer transformer can or cannot learn the task (see Questions).\n2. The experiments are fairly limited, including those in the Appendix. More explanation of the solution structure in the deep case and the GD implementation would be helpful to the reader.", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s encouraging comments as well as the suggestions for improving clarity and expanding the experimental section. We respond below to the specific questions about the NP-hardness and generalization to deeper layers.\n \n## 1. NP-hardness interpretation\n> **W1:** It is unclear if the NP hardness result necessarily implies that the one layer transformer can or cannot learn the task. \n> **Q2:** Can gradient descent find a solution close to the global minimum performance? \n\nThe NP-hardness result does not mean 1-layer transformer can’t learn Markov chains, but finding the best possible parameters is computationally hard. \n- **For 1-layer LSA:** As shown in Fig. 5 (Section B.1), as the sequence length d increases, the performance gap between the 1-layer LSA trained with gradient descent with 1000 epochs and the reparameterized optimum X* becomes larger. This aligns with our NP-hardness result, which suggests that recovering the global optimum becomes intractable as $d$ increases. \n- **For deeper LSA:** Fig. 3 in the main text shows that adding layers improves accuracy, suggesting that increased expressivity helps overcome the limitations we identify for shallow architectures. For the same task and number of training epochs, standard training can succeed in practice when the model class is sufficiently expressive to begin with.\n\n## 2. Generalization to multi-layers\n> **W2**: Limited experiments, including those in the appendix. Lack of explanation of deep case solution structure and GD implementation\n- **Explanation of deep-case solution**. \n    - **Parameter assumption.** The parameter assumption in Eq. 11 requires certain entries in the LSA weights $P, Q$ to be nonzero, with their locations matching the global minimum structure in Theorems 3.3 and 3.4, though their values remain unrestricted. We then interpret the forward computation of this family of LSA parameters as iterative steps of preconditioned gradient descent. The emergence of multiple objective types is due to the nonzero entries in $P,Q$. \n    - **Empirical verification (Fig. 2).** To verify the multi-layer interpretation, we fixed the LSA parameter obtained by gradient descent training. The objectives are evaluated for this fixed LSA at each layer.  We observed that initially, the model makes progress on both $R_1$ and $R_2$, but beyond a certain depth, optimization deteriorates, all losses begin increasing. The non-monotonic trend is likely due to deeper LSA preconditioners  negatively affecting the performance.\n\n> **Q1:** Would the improvements to the performance on this task continue if you keep increasing the depth? Have the authors tried training models deeper than 3 layers? Is the global minimizing loss / Bayes error / best possible accuracy computed?\n\nWe extended our experiments to deeper models and observed that performance continues to improve up to 6 layers. We fix the total number of prompts as 1000, number of in-context examples as 100 and consider binary Markov chains with prefix length $d=5$. We perform gradient descent for 5000 epochs for all layers. The accuracy and MSE w.r.t. # layers is shown below. These results support the intuition that additional layers can mitigate the expressivity limitation. We hypothesize that the decrease at 5 layers may be due to optimization instability. Nonetheless, the overall trend supports the benefit of increased depth.\n\n| # Layers | Accuracy   | MSE  |\n| --- | --- | --- |\n| 1 | 0.681 | 0.225 |\n| 2 | 0.687 | 0.216 |\n| 3 | 0.713 | 0.200 |\n| 4 | 0.747 | 0.188 |\n| 5 | 0.681 | 0.218 |\n| 6 | 0.784  | 0.166 |\n- **Baseline.** For the same dataset, the Bayes error is 0.315, computed by first estimating the transition kernel from in-context example and the first 5 tokens of the query sequence, then using this estimated kernel to make a Bayes-optimal prediction of the next token. The Bayes error is then computed as the fraction of prompts for which this prediction is incorrect."}]

Author rebuttal:

Meta-review: This paper considers an in-context learning problem to learn Markovian Dynamics instead of usual regression problems. The authors show the explicit form of the optimal solution of a linear self-attention network for the binary Markov chain, and prove NP-hardness of computing the optimal solution in general. In addition to that, the authors show that the multi-layer model corresponds to a pre-conditioned gradient descent. These theoretical findings are properly supported by numerical experiments.  
  
Although the model of the analysis is quite limited (LSA), this paper investigates an important problem which generalizes the conventional i.i.d. regression problem to next token prediction. Unlike some related works considering non-i.i.d. dynamics estimation, this paper analyzes a sequence dependent ICL problem. This is a novel contribution. Although it is still at a basic level, it will lead to more general analyses.   
  
The numerical experiments also well support the theoretical findings. The paper is well written and properly survey the existing literatures.   
  
For these reasons, I think this paper is a good paper. I recommend acceptance.
Decision: Accept (poster)