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

TL;DR

This work tackles the problem of in-context learning for dynamics-driven, Markovian functions using transformer models. By reformulating the 1-layer linear self-attention loss into a strictly convex, reparameterized objective, the authors derive a closed-form global minimizer $X^* = H^{-1} b$ and show that recovering the corresponding transformer parameters is NP-hard, highlighting a fundamental expressivity gap. They then reinterpret multilayer linear self-attention as preconditioned gradient descent solving a multi-objective problem, and provide empirical validation on simplified transformers and GPT-2-like models. The results illuminate why deeper architectures can better capture structured Markovian dynamics and why shallow, single-layer LSAs may fail to realize the reparameterized optimum. Overall, the paper clarifies optimization landscapes, computational limits, and depth-driven improvements for ICL with dynamics-based inputs, offering a principled lens on transformer expressivity in structured settings.

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

Figures (10)

• Figure 1: Comparison between the sequence-level in-context Markovian data based attention structures and the existing works. Top: Compared to ICL for structured sequences, the key difference is that the exiting studies of the attention mechanism DBLP:journals/corr/abs-2402-04161DBLP:conf/icml/SanderGSBP24rajaraman2024transformersDBLP:conf/icml/NichaniDL24 is 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 DBLP:conf/nips/AhnCDS23zhang2023trained, we consider input vectors generated through a Markovian transition kernel with parameters $p_{01},p_{10}$ from given initial distributions.
• Figure 2: Preconditioned multi-objective optimization performance of a trained LSA across layers. (a) Values for the six objectives in the first group ($R1$), (b) Values for the six objectives in the second group ($R1$), and (c) Generational Distance (GD) measuring deviation from the Pareto front. The model is a 10-layer LSA trained on 100 in-context Markov chains of length 5.
• Figure 3: Accuracy of LSA models trained on Markovian prompts w.r.t. number of layers. Each curve corresponds to a different state space size.
• Figure 4: 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.
• 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 5 random seeds. The growing performance gap highlights the increasing difficulty for LSA to recover the optimal solution as the dimension increases.

Theorems & Definitions (29)

• Lemma 3.1: Strict convexity
• Lemma 3.2: Global minimum for reparameterized objective
• Theorem 3.3: Global minima for i.i.d. in-context initial states
• Theorem 3.4: Global minima for generalized in-context initial states distribution
• Theorem 3.5: NP-hardness of transformer parameter reconstruction from reparameterization
• proof
• Theorem 4.1: Forward pass as minimizing multiple objectives
• Lemma C.1
• proof
• Theorem C.2: Theorem \ref{['thm:minimizer']} restated