Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions

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.

1. Setup clarification

Q1: Why study Markov chains instead of general dynamical systems for non-i.i.d. ICL

Thank 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.

• 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.

2. Global minima to NP-hard results

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?

The 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.

• 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.
• 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.

3. Multi-objective interpretation

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?

• 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$.
• 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.

References

[1] L. Arnold, Random dynamical systems, Springer, 1998.

[2] R. Bhattacharya and M. Majumdar, Random dynamical systems: theory and applications, CUP, 2007.

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.

1. Restricted assumptions in problem setup

W1: Analysis assumes 1-layer linear attention and binary first-order Markov chains with limiting direct applicability to richer architectures or dynamics.

• 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.
• 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.

Q1: How do results generalize to softmax attention, multi-head architectures, or higher-order/non-binary dynamics?

• 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.
• 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.
• 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.

2. NP-hard result interpretation

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?

• 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.
• 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.
• 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).

3. Generality of multi-Layer analysis

W3: Limited generality of the multi-layer analysis due to structured parameterization. Q3: Do real-world trained transformers actually follow this structure? Can you justify that the multi-objective interpretation holds more generally?

• 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$.
• 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.

4. Experimental scope

W4: Limited experimental scope and missing main-text results for larger models 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?

Due 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.

• 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.
  • 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.
• 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.

• 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.

References

[1] Kwangjun Ahn, et al., Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.

[2] Ruiqi Zhang, et al., Trained Transformers Learn Linear Models In-Context. J. Mach. Learn. Res. 25: 49:1-49:55 2024.

[3] Jianhao Huang, et al., Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought. ICLR 2025.

[4] Catherine Olsson et al., In-context Learning and Induction Heads. 2022

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.

1. NP-hardness interpretation

W1: It is unclear if the NP hardness result necessarily implies that the one layer transformer can or cannot learn the task. Q2: Can gradient descent find a solution close to the global minimum performance?

The NP-hardness result does not mean 1-layer transformer can’t learn Markov chains, but finding the best possible parameters is computationally hard.

• 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.
• 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.

2. Generalization to multi-layers

W2: Limited experiments, including those in the appendix. Lack of explanation of deep case solution structure and GD implementation

• Explanation of deep-case solution.
  • 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$.
  • 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.

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?

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.

• 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.

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.

1. Setup clarification

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.

• Relationship with prior studies
  • 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.
  • 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.
  • 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.
• 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.

2. Global minima and reparameterization

W2: Unclear interpretation of global minimum and its relation to mean prediction. Q1: What the solutions in Theorem 3.3 and 3.4 represent?

• 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.
• 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.

W2: Introduced several simplifications (reparameterization) that make it harder to assess how applicable the findings are in practice. 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?

• 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.
  • 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.
• 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.

3. NP-Hardness interpretation

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])?

We refer the reviewer to our response to W1 for a full walkthrough of the reparameterization-to-hardness chain.

• 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.

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?

• Theoretical perspective:
  • 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}$.
• Empirical perspective:
  • 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.
  • 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.
  • 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.

4. Multi-Layer behavior

W3: Lack of intuitive explanation for the multi-layer forward pass and the emergence of multiple objectives in Theorem 4. Q5: What do the objectives in Theorem 4.1 represent intuitively, and how does the observed depth-dependent optimization behavior arise?

• 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$.
• 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.
• 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.

Minor typos. Thank you for pointing out these typos. We will correct the misplaced underbrace and add the missing bracket in the final version.

---

References

[1] Johannes von Oswald et al., Transformers Learn In-Context by Gradient Descent. ICML 2023

[2] Ruiqi Zhang, et al., Trained Transformers Learn Linear Models In-Context. J. Mach. Learn. Res. 25: 49:1-49:55 2024

[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

[4] Kwangjun Ahn et al., Linear attention is (maybe) all you need (to understand Transformer optimization). ICLR 2024

[5] Kwangjun Ahn et al., Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023

[6] Yu Bai et al., Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection. NeurIPS 2023

[7] Hongkang Li et al., How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? ICML 2024

[8] Yedi Zhang, et al., Training Dynamics of In-Context Learning in Linear Attention, ICML 2025.

[9] Catherine Olsson et al. In-context Learning and Induction Heads. 2022