Learning Linear Attention in Polynomial Time

We thank the reviewer for their valuable time and insightful feedback, which has significantly helped us refine and improve the clarity of our work. We have addressed each point raised below and hope our clarifications resolve any concerns.

Q1: n 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?

A1: We appreciate the reviewer's astute observation regarding Figure 2(b). Our theoretical results primarily focus on the conditions under which all empirical risk minimizers (ERMs) compute the same function, thereby guaranteeing certified identifiability. The convex algorithm, by its design, is indeed guaranteed to find this ERM with negligible numerical error.

However, the empirical results shown in Figure 2(b) for linear attention are obtained using Stochastic Gradient Descent (SGD), which does not necessarily converge to the exact ERM, especially with finite training steps and potential local optima in the non-convex landscape for fewer heads.

A core theme of our work is to provide theoretical and empirical evidence that increasing the number of heads improves both optimization and generalization properties of the model. This is precisely why the convex algorithm and SGD on 8-head linear attention achieves a remarkably small distance to ground truth, nearly $10^{-6}$. In this context, it is expected that configurations with fewer heads (1, 2, and 4) might exhibit slightly larger distances to the ground truth. While the observed differences are indeed tiny (from $10^{−4}$ to $10^{−6}$), the precise trend for 1, 2, and 4 heads presents an intriguing direction for future detailed investigation into the non-convex optimization landscape at lower head counts.

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

A2: Yes, we confirm that all reported performances, including accuracy and MSE in Lines 881-885, and in Figures 1 and 3, are absolutely measured on a held-out test set.

Our experimental setup is designed not only to validate optimization convergence but also to demonstrate Probably Approximately Correct (PAC) learnability, which directly accounts for and controls the generalization error on unseen data. This rigorous evaluation on a test set ensures that our results reflect the true generalization capabilities of the models.

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

A1: We genuinely thank the reviewer for this insightful comment and for highlighting a crucial point for practical interpretability. We agree that the theoretical requirement of the number of heads being at least the square of the feature dimension ($d^2$) to guarantee convexity via SVD formulation is indeed a significant distinction from typical transformer hyperparameter choices in practice.

Our theoretical analysis establishes the extreme case where the loss landscape becomes provably "convex" (in a precise sense) at $d^2$ heads, ensuring that any empirical risk minimizer computes the same function, irrespective of the optimization algorithm used. The practical intuition, however, is that adding even a few heads can significantly ameliorate the non-convexity of the optimization loss landscape, making it more benign and easier for SGD to navigate. It is highly plausible that much of the practical advantage derived from increasing heads is achieved long before reaching the $d^2$ theoretical threshold.

Investigating this regime of a constant, smaller number of heads (e.g., 2, 4, 8) and characterizing the precise benefits gained from their addition is an extremely interesting and relevant direction for future theoretical analysis. One promising avenue could involve drawing connections between the optimization of Multi-Head Linear Attention (MHLA) and the rich literature on matrix sensing or low-rank matrix optimization, where SGD's behavior for low-rank solutions (here, the rank corresponds to the number of heads) has been extensively studied. This would provide a theoretical bridge between our current work and the practical realities of transformer architecture design.

We thank the reviewer for their valuable time and insightful feedback, which has significantly helped us refine and improve the clarity of our work. We have addressed each point raised below and hope our clarifications resolve any concerns.

Question: Is the condition Lemma 2.3 a necessary condition for identifiability?

A: We thank the reviewer for this insightful question, which probes a crucial aspect of our theoretical findings. For Multi-Head Linear Attention (MHLA) with more than $d^2$ heads, the condition articulated in Lemma 2.3 is indeed both necessary and sufficient for identifiability. This precision stems directly from the nature of the Ordinary Least Squares (OLS) estimator: it yields a unique solution if and only if the data covariance matrix is full rank. In this overparameterized regime (> $d^2$ heads), the problem's structure ensures this tight relationship.

However, for MHLA with fewer than $d^2$ heads, the condition in Lemma 2.3 functions as a sufficient condition for identifiability. In this critically parameterized regime, it is theoretically possible (though perhaps hard to concretely imagine in all scenarios) that spurious solutions residing in the null space of the data covariance might exist and could coincidentally possess a rank smaller than or equal to the number of heads, thereby violating strict identifiability. Formulating a necessary and sufficient condition for identifiability in this regime (fewer than $d^2$ heads) or extending it to the complexities of softmax attention remains an intriguing and important open research question.

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

Regarding the broader context, we fully acknowledge that linear attention, in its standalone form, can be considered a "toy model" when compared to the full complexity of "real-world" Transformers. Nevertheless, algorithms developed for learning optimal weights in linear attention models serve several critical purposes as a warm start for more complex models; as a theoretical test bed; and as a natural extension of the classical system identification/control theory to the setting of modern neural sequence models.

Firstly, these algorithms can provide a strong "warm start" for optimizing more competitive linear attention variants, such as Gated Linear Attention, DeltaNet, and mLSTM. The insights gained are directly transferable.

Secondly, Linear attention offers an ideal test bed for building a rigorous theoretical understanding of how optimization algorithms are likely to converge for more complex models like softmax attention and deeper networks incorporating residual connections. Its simplified structure allows for analytical tractability that is often elusive in full Transformers.

Toward Polynomial-Time Learnable Transformers: From a different perspective, our work contributes to the vision of a polynomial-time learnable variant of the modern Transformer. Just as the classical Ho-Kalman algorithm from control and system identification provides a foundational box of tools for linear dynamical systems, we view our research as developing analogous tools and insights for the learning and identification of modern neural sequence models.

We thank the reviewer for their valuable time and insightful feedback, which has significantly helped us refine and improve the clarity of our work. We have addressed each point raised below and hope our clarifications resolve any concerns.

Q: Line 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 ...).

A: The "polynomial time" refers to the process of verifying our identifiability condition. This condition is fundamentally tied to the data covariance matrix being full rank, which is essential for ensuring the uniqueness of the Ordinary Least Squares (OLS) estimator adapted for Multi-Head Linear Attention (MHLA).

Therefore, the runtime for verifying this condition involves two main steps:

Computing the data covariance: This step is linear in the number of samples (N) in the dataset. For MHLA, it also scales with the fourth power of the feature dimension ($d^4$).

Computing the minimum eigenvalue: This involves checking if the $d^2$ by $d^2$ data covariance matrix is full rank, which takes time polynomial in d (poly(d)), depending on the specific algorithm used for estimating the minimum eigenvalue.

We will ensure this explanation is clearly stated earlier in the manuscript.

Q: 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)

A: Thank you very much for pointing out this lack of clarity; we agree this section requires more precise explanation.

The core point we aimed to convey in lines 64-66 is about generalization guarantees for algorithmic tasks, particularly when a model achieves zero training loss but faces out-of-distribution inputs. We used the Universal Turing Machine (UTM) as a canonical example of a "uniform circuit family" (this concept also applies to arithmetic operations, etc.).

Even if a standard Transformer (or any sequence model) achieves zero validation loss on a dataset of Turing machine computations, there's no inherent guarantee that the learned model will correctly execute any arbitrary Turing machine program, especially with out-of-distribution inputs or machines. The model might have simply memorized the training patterns.

Our contribution for MHLA is that we can establish a condition on the data that ensures the model will compute the correct function, even out-of-distribution. So long as the target circuit family (be it UTM computations, arithmetic, etc.) can be expressed by some instantiation of MHLA parameters, then certifiable identifiability guarantees true generalization. For the UTM example, this means the learned MHLA will robustly output the correct computation history token by token for any Turing machine and any input word, not just those seen during training. This addresses the fundamental challenge of ensuring that models don't just interpolate but truly learn the underlying algorithm.