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

AI summary

Overview

• Research area: Machine learning theory — in-context learning (ICL) in transformer architectures, specifically linear self-attention (LSA), applied to Markovian dynamical functions rather than the usual i.i.d. Gaussian linear regression tasks.
• Technical level: Advanced. The paper is built from convex reparameterization arguments, closed-form matrix expressions, and an NP-hardness reduction from bilinear separability.
• Scope: The paper characterizes the global minimum of the 1-layer LSA in-context loss for short Markov chains, proves that recovering the transformer parameters achieving that minimum is NP-hard in general, and shows that a multilayer LSA forward pass is equivalent to preconditioned multi-objective optimization.

What This Paper Is About

Transformers can solve new tasks from examples given in a prompt, a phenomenon called in-context learning. Most theory on why this works assumes simple i.i.d. inputs and linear regression, so it says little about tasks driven by temporal dynamics. This paper asks what happens when each in-context example is a short trajectory from a binary Markov chain sharing a common transition kernel, and whether a one-layer linear self-attention model can actually realize the optimal solution for such structured, stochastic prediction problems.

Key Contributions

1. Closed-form global minimum for 1-layer LSA. The authors characterize the global minimum of the in-context loss for length-2 binary Markov chains, under both independent initial conditions (Theorem 3.3) and correlated/generalized initial conditions (Theorem 3.4), showing the optimal transformer parameters have a denser structure than in the i.i.d. linear regression case.
2. NP-hardness of parameter recovery. For chains of arbitrary length, even though the global optimum is well defined and analytically characterized in an enlarged reparameterized space, recovering the 1-layer LSA parameters that realize it is NP-hard with respect to the in-context data dimension d (Theorem 3.5).
3. A realizability gap. Because the reparameterization maps onto the enlarged space, the optimum in that space gives a lower bound on what the transformer can achieve, exposing a structural limitation of 1-layer LSA for Markovian dynamics.
4. Multilayer LSA as multi-objective optimization. Within a specific parameter subspace, the forward pass of an L-layer linear transformer is equivalent to gradient descent on two groups of objectives, R1 and R2, preconditioned by the layer parameters (Theorem 4.1), and this holds for any prompt instance rather than only in expectation.

Main Findings

• Strict convexity in an enlarged space: Under initial probability π1 = [1 − p, p] with p ∈ (0, 1) and transition probabilities drawn from U(0, 1), the reparameterized objective f̃ over X ∈ ℝ^{dm}, with m = (d+2)(d+1)/2, is strictly convex, so a unique global minimizer X* exists (Lemma 3.1).
• Closed-form minimizer: X* = H⁻¹b, where H and b are expectations over powers of the transition kernel P, reflecting the first-order dependence of the in-context sequences (Lemma 3.2).
• Nontrivial optimal transformer parameters: For length-2 chains with Bernoulli(p) initial states and p01, p11 ~ U(0, 1), the global minimizer (P, Q) is given in closed form and has nonzero key entries, unlike the i.i.d. zero-mean Gaussian linear regression setting where the first d entries of b and the last row of A are zero.
• Robustness to correlated initial states: Relaxing the independence assumption on initial states (Theorem 3.4) leaves the structure of the global minimizer of f(P, Q) unchanged; the optimizer is now governed by c1 = Σᵢ 𝔼[xᵢ x_{n+1}] and c2 = Σᵢ Σ_{j ≠ i} 𝔼[xᵢ xⱼ x_{n+1}], i.e., joint expectations of the query and in-context samples.
• The sparsity that made linear regression easy disappears: In the prior i.i.d. linear setting, X* is sparse — each group of reparameterized variables tied to a column of A contains only one nonzero entry — reducing the system to d independent constraints. With Markovian inputs, statistical dependencies and nonzero mean patterns persist across in-context samples, producing dense coupling between variables.
• NP-hardness: Solving for (b, A) satisfying the reparameterization equation is NP-hard with respect to d. The problem is recast as bilinear constraints bᵀD^{(r)}A_{:,j} = X_r, and hardness is established by reduction from the bilinear separability problem of Theorem 3.1, system (13), in reference [48], using variable splitting and variable fixing.
• A lower bound, not an attainable value: Since f̃* = min_X f̃(X) satisfies f̃* ≤ min f(P, Q), the optimal value in the enlarged space need not correspond to any realizable transformer configuration. Example C.16 in Appendix C.3 illustrates the absence of a corresponding LSA parameter for X* in longer Markov chains, while the optimal (P, Q) keeps the same structure as in the length-2 case.
• Multilayer forward pass as optimization: With P = [0; b_lᵀ] and Q = −[Ā_l 0; a_l 0], the layer output satisfies y_{n+1}^{(l)} = ⟨w_l, x_{n+1}⟩ with w₀ = 0 and w_{l+1}ᵀ = w_lᵀ − b_lᵀ(∇R1(w_l)Ā_l + ∇R2(w_l)[0; a_lᵀ]). R1 combines d terms of the form x_{j,k}wᵀxⱼ with a squared alignment term (wᵀxⱼ − yⱼ)². R2 emphasizes the last coordinate w_d and ends in a cubic penalty (1/(3x_{j,d}))(yⱼ − wᵀxⱼ)³ when x_{j,d} ≠ 0, switching to a quadratic penalty when x_{j,d} = 0.
• Numerical illustration: Figure 2 reports a 10-layer LSA trained on 100 in-context Markov chains of length 5, showing values of the six objectives in the first group (R1), six objectives in the second group, and Generational Distance (GD) measuring deviation from the Pareto front. The abstract states the theory is validated with simplified transformers; no further benchmark or accuracy numbers are reported.

Methodology in Plain English

The authors set up an in-context learning problem where each prompt contains n input-output pairs plus a query. Each pair is a short binary Markov chain: the first d states form the input x, and the final state forms the label y. These are stacked into an embedding matrix Z₀, and the transformer output is read off as a single entry.

To make the loss tractable, they strip the softmax out of the attention layer, merge the key and query matrices, and then reparameterize: instead of optimizing the transformer matrices P and Q directly, they optimize a new variable X built from products of entries of the transformer parameters. In this new space the loss becomes an ordinary squared loss of a linear model, which they prove is strictly convex, so the global optimum is unique and can be written in closed form by solving a linear system. They then ask the reverse question — given that optimum, can one find transformer parameters that produce it? Answering that means solving a system of bilinear equations, which they show is computationally intractable by reducing a known NP-complete geometric separation problem to it. Finally, they analyze deeper stacks of the same linear attention layers and show algebraically that the sequence of layer updates matches preconditioned gradient steps on a pair of composite objectives.

Why This Matters

Impact on research. The paper moves ICL theory past i.i.d. Gaussian linear regression into temporally dependent data, showing that the clean, sparse, analytically invertible structure that made the linear case tractable does not survive. It also draws a sharp line between two notions of "optimal": the loss landscape can have a perfectly characterized global optimum while the transformer architecture may be unable to represent it at all. That distinction is useful framing for anyone reasoning about expressivity limits of shallow attention models.

Real-world applications (grounded in the paper's framing of Markov chains as a classical language-modeling abstraction):

• Next-token prediction over text-like sequences, where the state space resembles a vocabulary and transition probabilities resemble the conditional probability of the next word given previous text.
• Modeling statistical structure in sequential data drawn from discrete-state processes with shared transition dynamics across examples.
• Diagnosing when a shallow attention model will silently fall back on approximations rather than the true dynamics, which matters for sequence-modeling systems of moderate depth.
• Understanding multi-objective trade-offs in attention-based predictors, since the multilayer analysis reveals the forward pass balancing several simultaneously optimized objectives.

Industry relevance. The NP-hardness result warns that for structured sequential tasks, achieving the true optimum with a single attention layer is not merely an optimization difficulty but a representational one — relevant to architecture and depth choices. The multi-objective interpretation of multilayer linear attention also offers a lens for interpreting what trained attention layers are optimizing, including objective values and Pareto-front deviation as measured by Generational Distance.

Future Directions

• Extend beyond binary, first-order chains. The analysis deliberately starts with binary Markov chains with first-order memory; higher-order and larger state spaces are left open.
• Close the realizability gap. Identifying which parts of the enlarged parameter space are actually reachable by valid (b, A) pairs, and characterizing the true achievable loss min f(P, Q) relative to the lower bound f̃*, remains unresolved beyond the length-2 case.
• Understand the role of depth and initialization. The paper notes prior work showing shallow models can converge to oversimplified solutions and that convergence is sensitive to initialization, but does not study training dynamics here.
• Move past the simplified architecture. The results rely on linear self-attention with no softmax; whether the optimality and hardness conclusions extend to softmax attention or to models with feedforward layers is not established.
• Broader empirical validation. The theory is checked with simplified transformers and a 10-layer LSA experiment; how these findings scale to realistic transformer training setups is not reported.

Target Audience

Theoretical machine learning researchers working on in-context learning, transformer expressivity, and learning theory; optimization researchers interested in loss landscape analysis and multi-objective formulations; and practitioners weighing architecture depth against representational limits when applying attention models to temporally structured or stochastic sequential data.

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