Title: Learning Linear Attention in Polynomial Time
Venue: NeurIPS 2025
Accepted: yes
Keywords: Transformers

Learning Theory

PAC learning
Abstract: Previous research has explored the expressivity of Transformer models in simulating Boolean circuits or Turing machines. However, the efficient learnability of Transformers from data has remained an open question.  Our study addresses this gap by providing the first polynomial-time learnability results (specifically strong, agnostic PAC learning) for single-layer Transformers with linear attention.  We show that learning the optimal multi head linear attention can be recast as finding the optimal kernel predictor in a suitably defined RKHS.  Moving to generalization, we construct an algorithm that, given a dataset, checks in polynomial time whether the set of best fit multi head linear attention networks on this data all perform an identical computation--a powerful notion for out of distribution generalization.  We empirically validate our theoretical findings on several canonical tasks: learning random linear attention networks, key--value associations, and learning to execute finite automata. Our findings bridge a critical gap between theoretical expressivity and learnability of Transformer models.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper establishes the strong, agnostic PAC-learnability of a single layer of multi-head linear attention, for regression tasks.  This is accomplished by first transforming the function into an elementwise product of (1) a matrix derived from the input and (2) a matrix derived from the parameters.  This means that optimizing over H heads equates to optimizing over the set of rank-H matrices.  This leads to a result that the training-set optimal matrix can be identified by effectively fitting a linear regression and applying singular value decomposition. The next result shows that the minimizer is identifiable (surprising in my subjective opinion) if the dataset is realizable by MLHA or there are enough heads in the model.  Together, these results give a poly-time procedure for learning any circuit family that MHLA can handle when the training data satisfy that condition.  The procedure recovers a Turing machine that the learned MHLA will simulate on any input within a size budget.\n\nExperiments explore overparameterization by adding heads -- which leads to a convex problem in the theoretical setup -- gives optimization benefits (more so than adding layers).  Further experiments empirically verify the identifiability certificate in various settings.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nThe paper presents incredibly important theoretical advances in what (some key parts of an important variant of) the transformer network can learn.  The ideas are clearly conveyed and bring together tools from a wide range of CS and machine learning.  This is creative work and an important step toward a deeper understanding of an incredibly influential class of architectures.\n\nThe paper also does a nice job of situating the results in the context of extensive related work.\n\nThe experiments are thoughtfully selected and well chosen, e.g., connecting the theoretically derived algorithm with behavior of SGD, on synthetic data.\n\nWeaknesses\n\nThe intro is mostly clear and a good sketch to the paper.  It could be improved by defining the symbols when they are used, especially those that deviate from conventional transformer notation or are ambiguous (Z, d, etc.).  It could also be improved by unpacking the Turing machine -related reasoning around lines 64-66 (I read it several times but couldn't make sense of it) and taking more care to spell out the experiments (e.g., the result at lines 69-70 seems to be comparing multi-head linear attention to algorithm 1, which doesn't make any sense).  Far more people will read the intro to understand a result than will read the rest of the paper, and it's worth making sure they will actually be able to make sense of the results.\n\nGraphs in figure 1 were very hard to read; the assignment of colors and line styles is not systematic, which really slows down a reader trying to make sense of the graph.  There's a ton of space here, consider using larger fonts for the legend and axes.\n\nQUESTIONS:\nLine 60:  polynomial time *in what*?  I suspected at first that there's a dependency on the number of training instances, and this should be presented honestly up front, not buried later in the paper.  Later, if I understood correctly, there was not (it seems it's about the dimensionality of the key/query vectors ...).\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-03"}

comments: SUMMARY:
The paper investigates learnability of a simplified transformer architecture, i.e., a multi-head linear attention (MHLA) module. 

- They show that MLHA is learnable is PAC learnable, i.e., it requires polynomially many samples in the input size and runs in polynomial time w.r.t the input size. There proof idea reduces MHLA learning to a kernel learning problem. 
- They also characterize an identifiability criterion on the dataset that certifies that every empirical risk minimizer is functionally equivalent. 
- They show that  MHLAs can (autoregressively) express universal Turing machines
with polynomially bounded computation histories, implying a strong learnability condition for UTMs.
- They verify their certificate empirically on an associative memory task.

STRENGTHS AND WEAKNESSES:
**Strengths**

- The MLHA's reduction to a kernel learning problem and the resulting PAC learnability results are solid theoretical contributions towards understanding learnability of transformers.

- The identifiability condition for equivalence of MHLA's learnt via ERM is potentially very relevant to learning theoretic aspects of transformer learning

- The experimental results provide sufficient evidence to support the relevance of their theory.

**Weakness**

- MLHAs are quite a toy-setting compared to real-world transformers.

QUESTIONS:
-  Is the condition Lemma 2.3 a necessary condition for identifiability?

LIMITATIONS:
The authors address the limitations quite clearly in the paper.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper shows polynomial-time learnability of single-layer transformers with multi-head linear attention that outputs the last position. The key idea is to cast the learning problem of the model as a kernel least-squares regression problem on a lifted feature space defined by a fixed cubic polynomial function of the data entries. This makes the SVD of the learned regressor coincide architecturally with multi-head linear attention. Specifically, each pair of left and right singular vectors correspond to flattened query and key projection matrices of each head, respectively. Under the setup, the authors show polynomial-samples learnability of the kernel regressor (and hence linear attention) by applying the classical learning theoretic results (Theorem 2.2). Then, the authors derive an identifiability criterion based on the second moment of data computed by a fixed feature map that, if it is full-rank, all empirical risk minimizers of the kernel regression are essentially the same function (Lemma 2.3), while their parameters can differ up to the equivalence classes defined by a polynomial (Corollary 2.4). The authors present an application to autoregressive learning of universal Turing machines with bounded-length history (Section 3). While the theory uses kernel regression learning algorithm, the authors empirically verify the theoretical predictions also under SGD learning. Section 4.1 focuses on the role of number of heads, which is predicted to play the role of rank of SVD and hence control the well-behavedness of the loss landscape. Then, in Section 4.2, the authors verify the validity of the identifiability criterion on an associative memory task and a controlled mixture of identifiable and non-identifiable dataset, using both the kernel regression algorithm and SGD. An additional result on sample complexity of learning a deterministic finite automata is also provided (Appendix G.4).\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nS1. The paper studies an important and challenging problem of learnability of transformers. I agree with the motivating observation that studying expressive power of these architectures is not enough.\n\nS2. The paper presents a convincing answer to the raised problem for a single-layer, multi-head linear attention transformer architecture, which can be applied in an autoregressive manner, possibly with position encodings (Appendix B). This seems like an adequate simplification of the large architectures used today.\n\nS3. The theoretical results are strong, as they include polynomial-time learnability and a guarantee for out of distribution generalization under identifiability criterion. It is also notable that the theory puts no strong assumptions on structures on the data generating distribution, unlike many existing work on expressive power of transformers (which is partially thanks to the proposed kernel regression formulation, as far as I understand). The idea that overparameterization can lead to better learning behavior is also found in mean field analysis of neural network training, and NTK, but the idea of this work seem distinct from these works.\n\nS4. The key ideas underlying the theoretical results, e.g., viewing the singular vectors of learned kernel regressor as flattened query and key projection matrices, are also interesting and new as far as I am aware.\n\nS5. The theoretical results are validated empirically on an important set of learning tasks, not only using the kernel regression algorithm, but also SGD and AdamW.\n\nWeaknesses\n\nW1. Due to the SVD formulation, for convexity and to avoid the realizability assumption, the number of heads have to be at least the square of the feature dimension. This seems like a major difference from the hyperparameter choices of transformers in practice. While the authors empirically verify the theoretical results with reasonable number of heads, some discussion regarding this point could be helpful for readers.\n\nW2. Please see Q1 and Q2.\n\nMinor comments\n- It seems legend is missing in Figure 3 in Appendix.\n\nQUESTIONS:\nQ1. In Figure 2(b), the distance to ground truth parameters for certifiably identifiable data seems to be anticorrelated with the number of heads up to 4, which is in contrary to what the theoretical results would predict. Can the authors provide a possible explanation?\n\nQ2. In Line 881-885, and in Figures 1 and 3, are the performances (accuracy and MSE) measured on the held-out test set? If not, does measuring on test set substantially change the results?\n\nLIMITATIONS:\nThe authors discussed limitations in Section 6.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-03"}