Learning Linear Attention in Polynomial Time
Authors: Morris Yau
Ekin Akyürek
Jiayuan Mao
Joshua B. Tenenbaum
Stefanie Jegelka
Jacob Andreas
Venue: NeurIPS 2025 oral ()
Keywords: Transformers
Learning Theory
PAC learning
Primary area: theory
TL;DR: We develop algorithms that are guaranteed to PAC learn transformers.

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.

Learning Linear Attention in Polynomial Time
Morris Yau MIT CSAIL morr i isymmitmtt... edu
Ekin Akyürek MIT CSAIL akyurek@mit..en edu
Jiayuan Mao MIT CSAIL jiayuann@mit. edu
Joshua B. Tenenbaum MIT Brain and Cognitive Sciences jbtêmit. edu
Stefanie Jegelka TUM Munich, MCML, MIT CSAIL jtefje@mit.cc edu
Jacob Andreas MIT CSAIL j jda@mit. edu
Abstract
Previous research has explored the expressivity of Transformer models in simu- lating Boolean circuits Turing machines. However, the efficient learnability of Transformers from data has remained an open question. Our study addresses this learnability results (specifically strong,
gap by providing the first polynomial-time agnostic PAC learning) for single-layer Transformers with linear attention. We show that learning the optimal multi head linear attention can be recast finding the optimal kernel predictor in a suitably defined RKHS. Moving to generalization, we construct an algorithm that, given 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 out of distribution generalization. We empirically validate our theoretical findings several canonical tasks: learning random linear attention networks, key-value associations, and learning to execute finite automata. Our findings bridge critical gap between theoretical expressivity and lcarnability of Transformer models.
Introduction
Transformers the dominant neural architecture used in language modeling. A growing body of work seeks to explain the behavior of trained Transformers and characterize their learnability (Pérez
al. 2019 Edelman al., 2022b Hahn 2020 Merrill and Sabharwal 2023 Merrill al. 2022
2021 Liu et al., 2022 Feng et al.] 2023 Edelman et al. 2022a Wei et al. 2021 Zhang et al. 2024
Trauger and Tewari 2023 Chen and Li 2024). While large body of work shows that Transformers
are expressive enough h to implement important models of computation, it remains an open question whether these constructions may be efficiently learned. Even verifying that a trained model has successfully learned a computational procedure (uniform circuit family) has remained challenging.
Existing work shows positive results on how Transformer-like architectures computations, including simulating universal Turing machines [Li al.) 2024), evaluating sentences of first-order logic [Barceló et al.) 2020), and recognizing various formal languages (Strobl et al.)
express diverse
2024 On the other hand, results on learnability in polynomial time and samples with provable guarantees tend rely strong data-generating assumptions, e.g., Gaussian data, etc. (Zhang et al.
2023 Jelassi et al. 2022 Tian et al. 2023 Oymak et al., 2023 Fu et al. 2023 Tarzanagh et al.) 2024 Deora et al.) 2023 This brings us our first motivating question.
there an efficient algorithm in time and samples that learns the optimal parameters of class of Transformer models for any dataset?
In this paper, we establish the strong, agnostic PAC-learabillity linear attention. Linear attention
variants (kernel, gated, flash, etc.) Yang et al. [025 2024), mLSTM in LLLSM//Bcck et al. (0024)
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

2024]. DeltaNet Schlag et al. (202]) have recently matched or outperformed softmax attention language and vision benchmarks, underscoring the practical value of their theory; Ahn et al. 2024 Katharopoulos et al.) 2020). Despite its name, linear attention is not linear and its loss landscape is nonconvex. We focus our analysis on multi-head linear attention networks, or MHLAs for regression tasks. An MHLA is parameterized by two matrices (Vn, Qn) for each of И heads as such = {(V,. Qa) A one layer MHLA computes Y =befe(t) Vi, (ZF Z). Here key and query matrices are fused into one, as they multiply one another directly.
Reinet (un et al. [2023] Mamba2 Dao and Gu
We first show that the computation performed by MHLAs can be reformulated as elementwise product between two larger matrices (W, (2)),, where W = >eeer flatten( Vi. (latten(()))* and (Z) is fixed cubic polynomial function of Z. Consequently, optimizing over the class of H H-head MHLA models equivalent to optimizing over the class of rank-H matrices W. Furthermore, in the full-rank space of d² x dl matrices, optimization of W can be performed via linear regression with time polynomial in the inverse target error and size of the dataset. Finally, decomposing an optimal W via SVD recovers an MHLA model with no more than d² heads that is then guaranteed to compete against the best MHLA parameters establishing our agnostic learning result (the learned model competes against the best choice of parameters in the hypothesis class).
Next, achieving zero training and validation loss does not by itself certify that model has learned target computation well enough to generalize out of distribution. Imagine learning arithmetic from input output pairs alone. Many distinct parameter settings can fit the same data, and fail for larger length inputs. We therefore ask:
Is there data-dependent, efficiently checkable condition that forces every empirical-rist risk minimiser to realise the same function?
For MHLAs the answer is yes. Define the second-moment matrix of the cubic feature map X' as
An [(2
If Ap is full rank-out our certifiable identifiability criterion- then all empirical-risk minimisers of MHLA coincide every input. The test runs in polynomial time and is unaffected by parameter redundancies such as rescaling V and Q.
Combining this certificate with our expressivity result yields polynomial-time procedure that (i) learns any circuit family implementable by MHLA whenever the training data satisfy the criterion, and (ii) provably recovers, for example, a bounded- history universal Turing machine from its input-output traces (Appendix C). Once learned, the MHLA simulates any such Turing machine on any input within the prescribed size budget.
In the experimental section, we validate our theoretical findings. In Section 4.1 we train multiple models using stochastic gradient descent on dataset generated by single linear attention network's s output. Our results demonstrate that multi-head linear attention outperforms both single-layer linear attention and multi-layed linear attention, achieving comparable results to our Algorithm [] In Section 4.2 we show that our proposed certificate directly correlates with generalization error even for models trained using stochastic gradient descent. In summary:
We provide a polynomial time algorithm that, given any dataset, finds the best parameters for multi head linear attention and generalizes with polynomial data, i.c., strong agnostic PAC learning (Section 2.I
We find an efficiently checkable condition (certifiable identifiability) on the training dataset that certifies every empirical risk minimizer of a MHLA is functionally equivalent, and therefore has the same behavior out of distribution (Appendix A see Lemma A.3
We study empirically the value of overparameterizztion with multiple heads vs. multiple layers in Section 4.1 We verify our certificates empirically on the associative memory task in Section 4.2

---

Algorithm MHLA Learning via Regression
1: Imput: Data D {(Z,, (1) kEIN) for Z, idxn and R2d
2: [Ki)eIN ExtractFeature D), generates
(21:: 21:) Zim, (22: 21.)21
(21: #2:) 2in Zins (22: 22.) zin,
(21:) za. adms
(22 Rd.) Fdns
N; ==
(I)
(=d: 21:) 21m.
(Ed: 22:) Rim,
(Za:s zd:) Zdm.
3: Create dataset (Xi,a Keia) wejd Let Xi,a rows and everywhere else:
Rdd xal
be matrix that is comprised of X; in the a' th block
4:
X
X
o]"
(2)
Let W E xdd be regressor:
W
arg min
Σ (E|M)
Σ
(W,
yi.a)
(3)
Wead x
where Vi,a is the a'th coordinate of yi- Take SVD W = ABT EieAr A, where A the rank W.
7: Vn (o)) and Qs Fold( where Fold Rdd Rdxd takes vector p [py for df and j [d] and reshapes into matrix / E gex such that P.g Pij-
8: Return: (Vh, Quhnej8)
2 Technical Overview
We start with basic definitions of multi-head linear attention (MHLA) module, an attention module without the softmax activation.
Definition 2.1 (Multi-Head Linear Attention). Let Z E Rdx= be a matrix of input data. Let {(V,, he(H) be set of parameters where each Vn Qu.h E Rddd denotes value and key- query matrices for all heads ordered tuples of d x d matrices. We define mulli-head linear attention (MHHLA) to be the function MHLAg Rdxx Rdxx.
[H]. We say & QH where Qn the space of sets of H
Ý MHLAA = -Z)
Ene(m VrZ(ZTQnha QnZ).
where y E Bdxn is the output of the one layer linear attention. We will primarily be interested in the rightmost column vector output by MHLAg (e.g., as auto-regressive language models), which is:
ý MHLAA
E hEiH Vn. (nZT"))
where Z|: n1] is the n th column of Z.
2.1 Polynomial-time learnability
Our main result is that MHLA is learnable in polynomial time. Colloquially, Algorithm D re- turns an MHLA that attains the global minimum of the training loss and requires as few as poly(d, e1, log(5-1)) samples to achieve generalization gorithmic guarantees do not require the data to be "realizable" (that is, the data need not be generated by an underlying MHLA).
with probability 1 6. Our al-
Theorem 2.2 (Leamability of Linear Attention). Let D be dataset D (Zi, Vs Vie drawn i.i.d. from a distribution where each Z, E Rdd.,, Y: E IRd, The embedding dimension d fixed across the dataset, whereas n; can be different for each datapoint. Let n'az = max be the maximum sequence length, and let S2y be the space of H pairs of value and key-query matrices {V,, Qn. 2n)bne(H) for any H E oo). Then there is an algorithm (Algorithm ! that runs in time O(Nd' Dinan and that, given
pairs {ZZ, y.) y.) JE[N), returns è [(V,, Qx) 1e)meeff E Ωg for H < de
3

---

such that with probability 1 1-6 6,
BiZ.)DDD [IMHIAA =
mineedy RIwWED [[|MHLA0(Z) - -11 <c (6)
with sample complexity N O (2 (at 1o)(
Below describe the high-level ideas behind the algorithm; a formal proof is given in Appendix D Note that if we are purely concerned with guaranteeing that we can find global minimum of the training loss, we may remove the i.id. assumption: Algorithm i always within error e of the optimal training loss. This is also detailed in Appendix(D) Specific issues related to generalization over autoregressive sequences rather than i.idd data are handled in the UTM learning result with standard union bound on the sample complexity; see Section F.2
The main idea behind Algorithm 1 is to construct feature mapping X' : RR + dx from the data covariates Z with entries 24y for the entry the i'th row and j'th column and rows Z1: 22: .. Ad: F'dh feature space of dimension d x d The map (22) defined as:
X(2) =
(11: 21: 2in
(21:1 22.) 2im
11 2d: Zdn
(22: 21:)21n
(22: 22:)
(2 Zdt Zdn
(Zd: 21:)21n
(Zdts 22.) Fin
(Ed: Zdi Zdn
Here, index the of ( by id] and the columns by all tuples (k, e) [d] such that x(()),.(.)) = (Zy:, 1x:)2nn. At a high level, Algorithm|] [ is a kernel method defined by the feature mapping X. The learned kernel predictor regressor) can be mapped back onto set of parameters (Vn, Qn ) for MHLA with more than d² heads via SVD. Hence, the relaxation translates into more heads. Interestingly, in our experiments in Section 4.1 d² heads also benefit learning with SGD.
Proof Idea: Much of the notation in this section is defined in Algorithm I First we write down the loss, and observe that a one-layer attention network is a quadratic polynomial in (Vi, Qs heEH with input features Xi.a:
(e()(Z).9)) =
1 Σ N
Σ RE[N] aE|d]
(Te,X,,6 v1,a))
(8)
with
To = Σ
flatten(V,) (iltemme))lat
AE[m]]
V1.00 Qn.00
V1,00 Vv.01 Qh.01
Vn Vn.0n Qn.aa
Vn.01 Qn.oo
bE[H]
Vn.ad
Vh
VnaaaQn
Now we relax this objective by replacing Te with an unconstrained matrix W E xd While Te is rank- H matrix, we allow W to be general matrix, this relaxation is guaranteed to have a smaller loss. Furthermore, the loss can be optimized via ordinary least squares. Finally, we apply SVD to W obtain a set of d? left and right singular vectors scaled by the square root the magnitude of the singular value. Here the scaled left singular vectors correspond to Vi, and the scaled right singular vectors correspond Qn for h E [.t]. Since the rank of W is no greater than d² the resulting MHLLA satisfies / d2 The sample complexity follows from classical results VC theory Kearns and Vazirani
1994 For full proof see Appendix D

---

2.2 Identifiability
direct implication of our algorithmic result is the construction of an efficiently checkable condition on the data that guarantees every empirical risk minimizer in a family of MHLAs computes the same function. Let Ap be the second moment of specific mapping HHZ)) of the data, defined in Lemma A.3
An = B[H(Z) = Σ
ZED
Then if Ao is full rank or equivalently its minimum eigenvalue is greater than zero, then it is guaranteed that MHLA is identifiable with respect to the data.
Let dataset D {(Z, y: heIN be realizable (see Definition (A2) by an H Hhad MHIL for any H 2 1. Let H be the uniform family of polynomials H.. RRdxn + Re for vb := (2) + d' defined as in Algorithm 2 For convenience we drop the
Lemma 2.3 (Certificate of adentifabilityy-noormme
subscript of n and write H(Z) to mean Hz( for Z E Rdxn, Finally, define AD E &R to be the second moment of the data features:
Ap = ED [M(2))(()*]]
(10)
Then if the eigenvalue Amin (Ap) > we say that MHILAe is certifiably identifiable with respect to D. That is, for every pair of empirical risk minimizers O O' E S2
MHLAg = MHLLee
(11)
L.e, the two models have the same outputs on all inputs.
Corollary 2.4. There is a polynomial p QH - RR such that for any pair of parameters O O' E S2 we have MHLA MIlLAg if and only if p(0) p(O").
The polynomial p defines the equivalence class of parameters that compute the same function. For formal statement of Lemma 2.3 see Lemma A.3 For handling of errors for approximate empirical risk minimization see Lemma A.7 Moreover, the certificate given by Algorithm 2 is not the only choice of feature mapping H that would certify identifiability; Lemma E. gives a general certificate for identifiability. One way to interpret Corollary 1sroflary2.4jjs that two MHLA models parameterized by θ and O' compute the same function if and only if they are the same linear function in a specific feature space (akin to matching coefficients in polynomial regression), which in turn is true if p(0) = p(O') for the polynomial given in Corollary A.4 Comparing distance between the coefficients in the range of p essentially the only meaningful metric of distance that is agnostic to the choice of dataset.
Finally, we answer a few natural questions related to identifiability which we briefly summarize here. Firstly, perfectly noisy input data is identifiable under weak assumptions on the moments of the noise (see Lemma A.5). Secondly, the model class of MHLA with at least d² heads is certifiably identifiable from the second moment condition alone, and does not require realizability of the data (see Lemma A.6). Finally, we empirically verify the min eigenvalue of AD predicts the generalization behavior of SGD for MHLA for the problem of learning key-value memories (see Figure 2
3 Application to learning Universal Turing Machines.
In Appendix B we demonstrate that MHLAs can (autoregressivelyy) express universal Turing machines with polynomially bounded computation histories. In this context, our identifiability results imply that, given a certifiably identifiable dataset of Turing machines and their computation histories on input words, empirical risk minimization and in particular Algorithm [ will Iearn the universal Turing machine in a strong sense (Lemma CC5) for learning, Lemma A.8 with identifiability) That means at test time the learned MHLA will simulate any Turing Machine on any input word up to a given size for a bounded number of steps. For more detail see|C]
Lemma 3.1 (Learning UTM from Certifiably Identifiable Data). Let D = {(Zi, yi) [E)] be dataset satisfying y; = .HLLAe for θ E 22 being the expressibility parameters of Lemma B. for the set of TM's/words (M,z) E Q E, ú, ②). If D is cerlifiably identifiable with Xmin (A n) > then there is a poly(d, N, Q, E, n, t - time algorithm that outputs set of parameters θ Qgee such that for all TM's M and input words x in (1 E, n, 4). we have
CH&(M, x) (t) -ko] x²
(12)
5

---

The c(t) step of the autoregressive computation history of is equal to the thth step of the computation history of M on x.
4 Experiments
In our experiments, we validate our theoretical predictions settings where Transformers are trained using stochastic gradient descent (SGD), as follows: Firstly, Theorem 2.2 adding a sufficient number of heads to an MHLA leads to a convex optimization problem after reparameterization. This suggests that over parameterizing by adding heads may provide optimization benefits. We investigate the role of over-vvrannen aaee-aaaetttiizttionnn in multi-head and multi-layer linear attention networks. For random data generated from línear attention networks, we observe that adding more heads achieves faster convergence of training loss than adding more layers. This suggests that while depth is important for expressiveness, the number of heads is important for optimization (Figure
exploits that
Secondly, we empirically verify the certificate of identifiability provided by Lemma A.3 on datasets for associative memory [Bietti et al. 2023 Cabannes et al. 2024] with different choices of embeddings, demonstrating convergence to the equivalence class of the true parameters when Amin (Ap) > and converging to spurious solutions when (Ap) 0 (Figure
4.1 Do extra heads help optimization with SGD?
To probe whether more heads facilitate learning in general, we train our convex relaxation and different types of over peeamaarmeteri models with SGD on data generated from single-layer linear attention network. For the data, we initialize a single-layer linear attention network with parameters
V IRd and Q Rddxd sampled from Gaussian distribution (O, +-a)- Input sequences
Z' E Rxxd sampled from (0, vT ), where = 1, ...... N, T 100 is the maximum number of time steps, and N is the dataset size. We generate outputs by running the ground-trute network auto-regressively: y = V. (Z' =|QZ"| [E, f)), creating our dataset 2 = {(Z', y y'))
In addition learning with Algorithm|| | train three types of models on this data using SGD: (1) multi-head linear attention as in Equation (4): (2) multi-layer linear attention with a single head: (3) an ordinary Transformer network (Vaswani et al. 2017 with softmax attention, multi-layer perceptron blocks, and layer normalization.
Figure I illustrates the results. For same experiment with d 4 and N 2048 see Figure 3a in the appendix. Detailed hyperparameters and optimization procedures are described in Appendix G.1
We observe that multi-head attention scales effectively with an increasing number of heads, resulting in improved performance. Notably, for d or 4 input dimensions, using dl heads yields the best performance and is empirically comparable to Algorithm [ approaching floating point error precision. Theoretically, d² is the maximum rank in the relaxation in Algorithm 1 In contrast, multi-layer attention models show diminishing returns and perform worse than single-layer attention. Interestingly, adding more layers can sometimes degrade performance. The full transformer model, which incorporates softmax attention, MI.P layers and layer normalization, does not significantly outperform the single-layer linear attention model on this task.
These findings suggest that the type of weee-paramtterization matters significantly in learning linear attention networks. Interestingly, multi-head architectures appear be particularly effective-aligned with the structure of Algorithm where the relaxation corresponds to adding more heads.
4.2 Does certifiable identifiability predict generalization?
we developed certificate that provides a sufficient condition for identifiability. To assess the practical relevance of this certificate, we conducted an empirical analysis of convergence in cases where the condition is not satisfied. The results of this analysis are presented in Figure? 2
In Lemma A.3
Associative Memory Associative Memory [Bettii et al. 2023 Cabannes et al. 2024 a task of looking up value in table with a query. Via a single head layer linear attention model it can be

---

10
10
m-1:ael, lineer
10
MST.y,
20-,, 3-1, lnnerr
linesr
10"14
12-16, n-1, linaer
liner
10"13
2r1. Iineer
m-1, a-1, 6:ll
10
mel a=2, Hill
100
200
300
400
500
epochs
N d= 2
10
10
10
MST:Qy
10
10"H
100
200
300
400 500
epochs
(b) N = d
Figure 1: Performance comparison of multi-head, multi-layer linear attention models and the original Transformer model (denoted as full). We trained using SGD synthetic data generated from single-layer linear attention model for varying training set sizes (N) and input dimensions (d), number of heads m, and number of layers n. Results demonstrate that multi-head architectures converge faster on different input dimensions and match the performance of our algorithm | (convex algorithm). Increasing the number layers or incorporating multilayer perceptrons (MLPs) and layer normalization did not yield consistent improvements. Shading indicates the standard error over three different runs.
represented with ground truth parameters
(V,Q) where V, E R2dx2d,
0 laxd
Taxd
0
Q
O
The data Z is drawn as follows: let kj, ka,..., ka E Rdd be random variables corresponding to keys in a lookup table, let 01, 12 ... Va E Rdd be random variables corresponding to values in lookup table, let E R be random variable corresponding to query to the lookup table, and Ç ~ -V I) be random noise, such that Z and the output vector y are defined as:
ki
K2
Kd q Va
=
U1
U2
5
(13)
0
y = MHLAg(Z) = Eyeten (q. kz)b
(14)
Mixture of distributions: We generate two datasets, one that has identifiable Amin (Ap) > 0 and one that nonidentifiable with Amin (Ap) 0. The identifiable dataset generated with {kz Jeid and (u, Jaet drawn i.idd (O. The query q is chosen to be one of the (ky Je(d uniformly at random. The non-identifiable dataset is drawn such that (k, Jeed forms random unitary matrix,
7

---

distunce( p()))
10
10
Cartifleate Full Raci
10
Low Rank
Convet
I-Road
2-hold
+ Aead
8- hehd
Eigetvvauu
(a) Certificate of identifiability form of the minimum Eigenvalue Amin (Ap)) vs. Eu- clidean distance in p feature space of param- eters learned by Algorithm T
(b) Distance to ground truth parameters in feature space for certifiably identifiable data (min eigenvalue 0.06) vs. nonidentifiable data (min eigenvalue = 0), Here the paramcters MHLA are learned via SGD. Error bars are standard error on three diffcrent runs.
Figure Impact of data distribution on the associative lookup task performance: We generated training data for an associative lookup task [Bietti et al. 2023 Cabannes et al. 2024) using mixtures of two distributions: (1) Gaussian key and value vectors, and (2) random unitary key and value vectors. By adjusting the mixture probability, of the data covariance matrix), as unitary key-value vectors give rank-defici "certificates" (a) Algorithm T the minimum eigenvalue increases, Algorithm converges parameters. (b) SGD: SGD learns parameters that are equivalent to the ground truth parameters in feature space for certifiably identifiable data, but for unidentifiable data, they are far apart in p feature space and therefore compute different functions.
manipulate the certificate value (minimum eigenvalue
closely the true
ie.. ||سل|| = for all j E [d] and (ki, ky) O for all ن j'. Similarly, heedh is also drawn from a randomly generated unitary matrix. We draw new random unitary matrices for each datapoint, where q is again chosen to be one of the fei] uniformly at random. We set d = 4 dimensions for both datasets, and draw N 214 samples for each dataset. We mix the datasets together with a mixing probability ranging from 95% unidentifiable to 100% unidentifiable. In this manner we generate spread of datasets with different values for Amin (Ad) that tend zero.
Certifiable Identifiability for Algorithm For each dataset, we run Algorithm| I which returns Ô. We compare Ô to the ground truth θ in p feature space via the distance
- -(ê)]]
(15)
Here, p is the polynomial given in Lemma Recall from Corollary A.4 that p defines the equivalence class parameters that compute the same function, i.e., MHL.A = MHLAg if and only if p(0) p(ê). On each dataset, measure the certificate value Amin Amin(Ap) the x-axis vs. d(O) Ô) the y-axis. In Figure 24. we see that the certificate value increases, d(O, Q) decreases, indicating that MHLAg and MHLAg compute the same function.
Our notion of certifiable identifiability in Lemma A.3 applies to any empirical risk minimizer. Therefore, it applies to popular optimizers like SGD and Adam they achieve the minimum of the loss, which is in our synthetic case equal to zero. In Figure 2b] we train MHLA models via SGD with 1, 2, 4, and 8 heads. For identifiable data with minimum eigenvalue 0.06, we see that the learned parameters and ground truth parameters are same in feature space. However, for unidentifiable data with minimum eigenvalue 0, learned parameters and ground truth parameters are far apart in p feature space and therefore compute different functions.
Certifiable Identifiability for MHIA:

---

5 Related Work
5.1 Formal Expressivity of Transformers
A large body of work has been trying tackle the problem of quantifying what algorithmic tasks can Transformer do, in terms of various kinds of circuit families [éérez et al. [219 Edelman ct al. 2022b
Hahn (020 Merrill and Sabharwal 2023 Merrill al. 2022 2021 Liu al. 2022 Feng al. 2023
In particular, researchers have studied how Transformers can realize specific DSLs Weiss et al.
2021] logic expressions (Dong al. 2019 Barceló et al. 2020 2024.. Turing machines |Dehghani
ct al. 2018 Giannou et al., [023] Pérez et al. 2021]. formal language recognition (Hao et al 2022
Chiang et 2023], as well as automata and universal Turing machines [[Liu et al.) 2022 et al.
2024]. However, while these works primarily focus on determining the types of problems whose
solutions Transformer can express, they often overlook the crucial question of how these solutions can be learned from data. Moreover, there is limited discussion on the sufficiency of the dataset wtself-whether the data available can identify the underlying "true" function or algorithm that we aim to capture.
5.2 Learning Transformers
We break down the literature on learning transformers. First, there is the literature on statistical learnability, where the focus is on the amount of data required learn without considering whether
there is a tractable algorithm for learning (Edelman et al. 2022a Wei et al. 2021 Zhhng et al. 2024
Trauger and Tewari 2023].
Second, there are learnability results for single head transformers for data distributions under a variety of assumptions. In particular, ZZhang al. [2023) provide learnability results for in-context linear regression: Jelassi et al. [2022] show that data with spatial structure can be learned; the work of Tian
al. 20233 analyzes SGD training dynamics for a toy model for data; and Oymak et al. [2023 study
the prompt attention model.
Third, the literature provable guarantees for learning multi head attention is rather sparse. Fu et al. [2023] give learnability results in a regime where attention matrices are fixed and only the projection matrices are trained. (Tarzanagh et al. [2024) show connections between single layer attention optimization and SVM learning. Under a good gradient initialization condition, overparameteriizatio condition, and a condition on the scores of optimal tokens the global convergence of gradient descent to a particular SVM problem can be established. [eora et al. (203) analyze a setting of learning multi head attention with gradient descent under their Assumption 2. In the words of the authors "these conditions are related to the realizability condition, which guarantees obtaining small training error near initialization". which they instantiate with the separability of the data in an NTK space and a proximity of initialization to realizable parameters. Interestingly, they find that multi head attention has benign optimization properties. Finally, Chen and Li| [2024) study learning for multi head attention for well structured data that is drawn independent Bernoulli or Gaussian. They provide an extensive discussion of lower bounds for learning multi head attention.
6 Conclusion and Limitations
In this work we tackle the fundamental problem of finding an efficient algorithm that provably learns the weights of a linear Transformer. Our key theoretical ingredient is to consider a model class that's sufficiently "wide" (scaling number of heads), and to find that the loss is convex under this scaling, with generalization guaranteed by the classical VC theory. This reinforces the empirical observation that scaling model size enables efficient optimization and can still result in successful generalization. Our theory extends trivially when arbitrary feature maps Φ(-) are applied to keys and queries providing natural avenue for extending our theory to models that can approximate softmax transformers with custom key-query kernels. Of course the model class we consider is far simpler than modern L.M's, s, but we consider our work an important step towards designing algorithms with provable guarantees for training neural sequence models.

---

Acknowledgments
We gratefully acknowledge support from NSF grants IIS-221417) IIS 22382402 CCF-21 12665 and DMS-21 11NS-1331108, 108; from AFOSR grant FA9550-22-1-08 2--299, from ONR MURI grant N00014-22-1-211
and from ARO grant W911NF-23: 1-0034; from the OpenPhilanthropy Foundation; from MIT Quest for Intelligence; from the MIT-IBM Watson Al Lab; from ONR Science of AI; from Simons Center for the Social Brain; and from an Alexander von Humboldt professorship. Any opinions, findings and conclusions recommendations expressed in this material are those of the authors and do not necessarily reflect the views of our sponsors.
References
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization), 2024. URL https
//arxiv.org/abat baagg/abs//00 2310. 2w1..000000000
Ekin Akyürek, Bailin Wang. Yoon Kim, and Jacob Andreas. In-context language learning: Architec- tures and algorithms, 2024. URL htttp::/aarav..org/abbb////
12973
Pablo Barceló, Egor V Kostylev, Mikael Monet, Jorge Pérez, Juan Reutter, and Juan-Pablo Silva. The logical expressiveness of graph neural networks. In /CL.R, 2020.
Pablo Barceló, Alexander Kozachinskiy, Anthony Widjaja Lin, and Vladimir Podolskii. Logical languages accepted by transformer encoders with hard attention. 2024.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xLSTM: Extended long short-term memory. Vancouver, Canada, December 2024.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformrr: A memory viewpoint, 2023. URL httpt: org/ /bs22306.00002
Vivien Cabannes, Berfin Simsek, and Alberto Bietti. Learning associative memories with gradient descent, 2024. URL https:ts:: arxiv. www.i.www//...............owo.
org/ abs absp/402..11122 18724
Sitan Chen and Yuanzhi Li. Provably learning a multi-head attention layer, 2024. URL https
7/arxiv. org/ abs habp:/240204003
David Chiang, Peter Cholak, and Anand Pillay. Tighter bounds on the expressivity of transformer encoders. arXiv preprint arXiv:2301.10743 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024. URL https://arxiv, weexxivvorrgaaae wttp/oo/////aaa/eeeoeeee 2405 21060
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Eukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.0388) 2018
Puneesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimization and generalization of multi-head attention, 2023. URL https://aaxiv.. wwt/abg./////eeo.eeo....o.aaeeee httg//bb//310..... 2310. 12680
Honghua Dong, Jiayuan Mao, Tian Lin, Chong Wang, Lihong Li, and Denny Zhou. Neural logic machines. In ICLR, 2019.
Benjamin I.. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms, 2022a. WtRLhttts:://ett wwwpiii..rg/abb//////eaa att.aa////...1///0.0 10090
Benjamin I. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pages 2022b.
5793-5831. PMLR, PMLRR
Guhao Feng, Yuntian Gu, Bohang Zhang, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. arXiv preprint arXiv:2305. 15408, 2023.
10

---

Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can single attention layer learn? study through the random features lens, 2023. URL hatp:://arxiv.oogg///aaa abw.///7/a//330./a.//// 11353
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Pa- pailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301. 13196, 2023.
Michael Hahn. Theoretical limitations of self-attention neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171, 2020.
Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard attention trans- formers: Perspectives from circuit complexity. Transactions of the Association for Computational Linguistics, 10:800-810, 2022.
Samy Jelassi, Michael F. Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure, 2022.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention, 2020. URL https:// //tps: org/ abs 16236
Michael J. Kearns and Umesh Vazirani. An Introduction to Computationall Learning Theory. The MIT Press, 08 1994. ISBN 9780262276863,
doi: 10.755//miippess/39...11.00 wt0.:/8887..0000011 URL
https://doi.. org /00gg//1..00 10. 7551 htt61/mmtpress////0
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 141.66:41..6980 2014.
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. arXiv preprint 128 2024.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts automata. arXiv preprint arXiv:2210.100: 10749, 2022.
Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam, 2018. URL https http:p/rwwreeeeeeeeeeoommmaaaddo
William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of rre-precision transformers. Transactions of the Association for Computational Linguistics, 11:531-545, 2023.
William Merrill, Yoav Goldberg, and Noah A Smith. On the power of saturated transformers: A view from circuit complexity. arXiv preprint arXiv:2 16213, 2021.
William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843-856, 856, 2022.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning, 2023. URL https://aarx org/ abs//3//..2.////////// 03435
Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. In ICLR, 2019.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research. 22(1):3463-3499 2021.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear Transformers are secretly fast weight programmers. Virtual only, July 2021.
Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. What formal languages can transformers express? a survey. Transactions of the Association for Computational Linguistics, 12:543-561, 2024.

---

Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: successor transformer for large language models. Preprint arXiv:2307.08621, 2023.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines, 2024. URL
wwtp:iiv..rg/b////////aaa /20// 2308 16898
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer, 2023. htRp:ttps: org/ abs/
2305. 16380
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arxiv:2307.009288 2023.
Jacob Trauger and Ambuj Tewari. Sequence length independent norm- norm-based generalization bounds for transformers, 2023. URL.hhtpsss//axxxt..or/////////2 13088
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximatio::: a case study on approximating turing machines with transformers. CoRR, abs/2107.11663,, 2021. URL https
bexxiv.org/a abs/ 2107 13163
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pages 11080-11090 PMLR, 2021.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. Vienna, Austria, July 2024.
Songlin Yang. Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length, 2025. URL URL httttttt///aretee org abs/a//s//400 06484
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context, 2023. Utt.tttss:///rrt/ http://awww....rrg/a///36//6
Yufeng Zhang, Boyi Liu, Qi Cai, Lingxiao Wang. and Zhaoran Wang. An analysis of attention via the lens of exchangeability and latent variable models, 2024. URL httpst /tapss:///rxxxx //tp://./iivviiia org abs/ 2212. 14852]
12

---

A Certificate for identifiability of linear attention
We begin by defining identifiability of a model class with respect to a dataset.
Definition A.1 (Identifiability). Let D {(Z,, 34) heN Let Ue denote model class which is uniform circuit family parameterized by parameters @ E Ω. Let C be loss function and 9ERM be the set of empirical risk minimizers:
Pe 16 e Ω Ô arg arg mingen L(ue, D)}.
(16)
We say model class Ue is identifiable with respect to the dataset D if for all Z E Rd and for all pairs of empirical risk minimizers Q O' E 9ERM we have Ue and Uey compute the same function, i.e, they agree on all inputs (are the same uniform circuit family):
(e Uer (Z).
(17)
In establishing conditions for identifiability, it will be useful to refer to another condition relating models to datasets.
Definition A.2 (Realizability). Let θ Qe be MHLA parameterization. We say dataset D = ((4, 34) ) JE(N) is realizable by a parameterization θ if yi MHLAA(Z)).
The definition of realizability can be modified to include independent noise the expense of adding some terms to our analyses. See Lemma A.7 for details.
Next, we prove that for the model class MHLA there is an efficiently checkable condition (certificate) of the data D that guarantees the model class is identifiable with respect to D. Our results follow by reinterpreting the results of Theorem 2.2 with focus data conditions that uniquely determine the optimal regressor. In this section we denote the mapping from data to feature space to be H and the mapping from parameters feature space be p which are analogous to the X and Ta of Equation Sp. We instantiate the feature mapping H. and parameter mapping polynomial as follows.
Lemma A.3 (Certificate of Identifiability). Let dataset D = {(Zs, y4 (4))s be a realizable dataset.
Let H {² be family of polynomials H. : dR - R for (c) + d² defined as follows. We index the entries of H. by taking the Kronecker product between all sets of pairs {j, k} (for all j, & e (d) with with all E [d]. We define H(()) (j,K)e as in Algorithm 2 to be
H H(Z) == (Zz: aNa) =tm4
(18)
Then if Amin (En (2)))( > 0, we have that MHLAg is identifiable with respect to D.
Next construct mapping p Ω d that partitions the parameter space into equivalence classes of parameters that compute the same function. This akin to matching coefficients in polynomial regression. This mapping defines meaningful notion of "distance" between different attention parameters by constructing a feature space in which equivalent models have the same representation. We denote the a' th row of p be Pa Ω + R and define it as follows.
Corollary A.4. Let (pa Jaaba be a collection of polynomials such that Pa(O) Z2 - R is defined as follows. Each pa(O) is indexed by pairs {j, for j, E [d] and & E [d] defined to be
Pa (a(O) = Σ Qke Vn,ak
(19)
hE[H]
Let the polynomial p S - Rdx be p == (Pr,P2, P2, P2..... .... Pa). Then for any pair of parameters O O' E Qt we have MHLA = MHLAgr if and only if p(0) = p(e)).
We give an overview of a few results building on our certifiable identifiability machinery:
First, data drawn from independent noise is certifiably identifiable. If the data matrices [Zihee() are drawn with each entry being standard normal noise, then MHLAA for QH is identifiable with respect to the data. The statement holds beyond standard normals to distributions satisfying weak moment conditions. The result stated with population risk instead of empirical risk to simplify the statement.
13

---

Algorithm 2 Constructing Features for Certificates of Identifiability
1: Input: Data D == (EI) for Z, Rdx=.
2: Output: fcature vectors H(,,) for [N]
3: for Z, D do
4: Let 21:, 22:;...d:: be the rows of Zi and lct Rab be the (a, b) entry of Z:
5: for sets [j.k] in Distinct Pairs of Indices in [d]² do
6: for E [d] do
H(Z.) = H(Z)) [z ax:)=ens]
8: end for
9:
end for
10:
for E [d] do
11: for f E [d] do
12: H(Z,) = H(Z,) [||2]]|
13: end for
14: end for
15:
end for Return: (H(2.)){aN)
Lemma A.5 (Independent input noise yields identifiability). Let (Z.y) ~ D be a realizable dataset. Let Z be drawn from distribution Z where the (a, b) entry of Z denoted by Zub is drawn i.i.d. from a distribution v over R. for all a E [d] and E [n]. Let the second and fourth moment of v be denoted m2 and m4 respectively. Let m2 > 0 and m4 > m². Then MHLA for Q E Qu is identifiable with respect to D. That is to say, for any population risk minimizers O O' E 4PPMM:
MHLAg MHLA
(20)
Second, when specialized the case of Multi Head Linear Attention MHLAg with more than d2 heads we can avoid the realizability assumption entirely. This is because the class of MHLA with an arbitrary number of heads linear in the feature space H given in Lemma
Lemma A.6 (Identifiability without realizability for MHLA with arbitrarily many heads). Let dataset D {Ze, (6)))))e VEEN) be any dataset drawn i.i.d from a distribution D. Let be defined as in Lemma A.3 Then if Xmin (ED(H(Z)H(Z)T) > then MHLA for O Qt for any E [d o0) is identifabble with respect to the data D. That is,
MHILAe MHHLAe
(21)
for all pairs of empirical risk minimizers O O' E ΩERM-
We also add quantitative version of identifiability with precise treatment of issues related to error. (For a corresponding statement of realizability with noise see Lemma (.22)
Lemma A.7 (Identifiability with Emror). Let -RM be the of e-approximate empirical risk minimizers,
Ω Qe-ERM
oo EQ EIZ.DED (MHLAg(Z)) - w) sc
Then we have for any O, 0' e Q.-ERM that for all inputs Rdxn
[MHLA -MHLA (Z)Il $
Amin (Ap) 12
(22)
We prove all the above statements in Appendix E
Application to learning Universal Turing Machines. Appendix we demonstrate that MHLAS can (autoregressively) express universal Turing machines with polynomially bounded computation histories. this context, our identifiability results imply that, given a certifiably identifiable dataset of Turing machines and their computation histories on input words, empirical risk minimization and in particular Al Algorithm] [ will learn the universal Turing machine
strong sense (Lemma C.5 for learning, Lemma A.8 with identifiability). That means at test time the learned MHLA will simulate any Turing Machine on any input word up to a given size for bounded number of steps. For more detail see C
14

---

Lemma A.8 (Learning UTM from Certifiably Identifiable Data). Let D = {(Zz, yo tie be dataset satisfying Yi M for Qn being the expressibility parameters of Lemma B. for the set of TM' sswordss (M.x) ) ( h. b). If D certifiably identifiable with Amsn (Ap) p) > n. then there is a poly(d, (d, N, Q, (2,,,, f. & n time algorithm that oupurs a set of parameters θ E 22 such that for all TM's M and input words x in 4 E, n, 4) we have
(M, - kio] x²
(23)
The cl step of the autoregressive computation history of is equal to the 1'th step of the computation history of M on x.
B Realizability of Universal Automata in MHLA
We also include an application of our theory on learnability and identifiability to the problem of learning universal Turing machine (UTMs) with polynomially bounded computation length. We prove such a UTM is expressible via MHLA in Lemma B.T and show that for certifiably identifiable data the leamed MHLA generalizes any TM M and input word x: in Lemma A.8
Lemma B.1 (UTM Expressibility). Let ( E. f. 6) be the set of Turing machines M {6, Σ, 2, Astart, laccepts (reject) and words x E Σ" with number of states, size of alphabet, size of input, and number of steps in computation history bounded by Q, E, â. & respectively. For any (M,x) E A, let { be the computation history of the UTM on (M,x). Let the autoregressive computation history (see Definition
C.2) of MHLAg on input (M, x) be denoted CHe(M, = { ...,,... )) Then there exists a set of parameters θ E QH for H = O(iσL) and embedding dimension d = O max(,, ②)), such that for all (M,x) E A. the TM computa- tion history at time step is equivalent to the autoregressive computation history at time step c(t) where c(t) S O((n + t)t) i.e ( -length(=") [)]] = x². Furthermore, this can be achieved with bits of precision.
Our construction bears similarities to Pérez et al. 2019 Hahn 2020 Merrill and Sabharwal 2023
Merrill et al.] 2022 102 Liu et al. 2022 Feng al. 2023 the high-level idea is write down every letter in the computation history of M on æ. If we use orthogonal vectors to encode every letter, state, and positional embedding we arrive at natural construction involving few basic primitives copy, lookup, and if-then-else. For details see discussion section] and Proof F.
C Application to Learning Universal Turing Machines
We apply our algorithmic and identifiability machinery to show that an important computational procedure is representable and learnable as an MHLA: namely, a restricted class of universal Turing machines (UTMs) with bounded computation history. We must first generalize our previous MHLA definition to enable multi-step computation:
Definition C.1 (Autoregressive MHLA), Let zd be an input matrix in dimension RRdxn, We define the iterative process of s-step autoregressive MHLA as follows: starting from t 0, let the next token *11 E Rd be:
+ MHLAa(Z').
(24)
and, for all t [+), let Zt+1 e 1 be the concatenation:
z+l Z ooy
(25)
Next we define the computation history of an autoregressive model analogously to the computation history of Turing machine.
Definition C.2 (Autoregressive Computation History). We refer to CHe(?) {Z')e as the computation history of the h-step autoregressive MHLA. We denote the 1-th step of the computation history C (7) Z'
We will often use the notation 7 -k] to denote the last k E 2 tokens of 7, Often, Z will be the embeddings corresponding word x in a language L. in which case we will use the notation
15

---

(z) and CHa (Z) interchangeably. For pedagogical discussion on how to map embeddings to letters in an alphabet, see Section G
Although the theory derived in this paper applies to all functions expressible by MHLAs, we are particularly interested in the task of learning universal Tiaring machines (UTMs). Let Σ be an alphabet. Let Q be set of states that includes {gstart, (oceept>9rejeet] a start, accept, and reject state respectively, Let 6 : 2 x _ + Q x Σ x {L/It) be a transition function that takes an alphabet and state symbol and maps to state transition, an output symbol, and a head movement left or right. Typically there is also a tape alphabet for which the input alphabet Σ is a subset.
Definition C.3 (Accept TM), Let M = {6, Σ, D, 2, Istart; Gaccept: Greject } be a TM. Let x E Σ" be all strings in the alphabet Σ. Then let ATM be the language ATM = (M) | M accepts z}.
The UTM constructed in Turing's 's 1936 paper recognizes ATM- In practice, we are most often interested in the behavior of TMs that run in polynomial time, and focus below on implementing a universal simulator for this restricted class:
Definition C.4. (Polynomially Bounded Universal Turing Machine) In general, a UTM is a recognizer for the language ATM- That is if x is in ATM. the UTM accepts, else, the UTM rejects or does not halt.
Let ATM n P be the language of input pairs (M, x) for TM M and word æ E E' such that M decides in polynomial time. Here, we consider UTM to be the polynomial time decider for ATM 0 P.
To define what it means for an autoregressive MHLA to perform the same computation as a TM, our main idea is to construct parameters for MHLA such that it executes the computation history of TM M on input æ. Let the UTM computation history at step t include the contents Ty...... Ikk on the tape after transition steps of the Turing machine M, the current state qt, and the current head position h4 Here k, is the number of tokens at timestep t. Then, there is a single-layer MHLA capable of simulating UTM:
Lemma B.1 (UTM Expressibility). Let AQ,,L, be the set of Turing machines M {6, I, Q, Istart, laccepts Oreject) J and words x E E" with number of states, size of alphabet, size of input. and number of steps in computation history bounded by Q, E, n. & respectively. For any (M, E A, let [ be the computation history of the UTM on (M,x). Let the autoregressive computation history (see Definition
C.2) of MHLAe on input (M, x) be denoted CHe(M, = 12,,22,,.... 2. Then there exists a set of parameters O E Qu for H = O and embedding dimension d = - max 2)), such that for all (M,x) E A, the TM computa- tion history at time step is equivalent to the autoregressive computation history at time step c(t) where c(t) S O(nn + t)t) i. ( length(>** (x²))] = x². Furthermore, this can be achieved with bits of precision.
We include the full proof for the existence of θ in the appendix. For simplicity, we adopt a naive embedding scheme that represents different letters in an alphabet as orthogonal unit vectors. This makes easy to contrive embedding schemes that incorporate arbitrary polynomial-sized circuits which could compute whether x E L(M). Moreover, we adopt positional encodings that are simply orthogonal unit vectors. Thus, in order to give each of T tokens unique ID,, we would require ((T) dimensional positional embeddings.
This can be combined with the learnability results above to yield a specialized result for UTMs:
Lemma C.5 (Learning a UTM). Let θ E Qu in dimension d be the MHLA parameters in Lemma B Let (M, Ii JeEN be pairs of T'M's M and words x of maximum length n drawn i.i.d. from a distribution D. Let Zi
Emmee(MM;,a x.). For each TM/word pair (M,, xx) let (2,) {2, D2,,,., 2 } be the 4-step autoregressive computation history of MIILA on Z,. Let D be the dataset D = (CH)(Z.)', "3,"""" hE(N)) where vi = MHLAa (Z)) Then
Algorithm 1 applied to imput D returns è E Ω4 for H S de such that with probability
E(Z.)ED (MIILA - v) Se
(26)
for sample complexity N poly(d, c 1og(6-1)). Then with probability 1 1-6 8 over the randomness in the data, the probability over D that the --step autoregressive computation history CH (M, x) and CH&(M, x) differ is upper bounded by
PE(Ma))D(CHH(AM... # CHo(M, & 5(ee).
(27)
16

---

Finally, if the dataset D is certifiably identifiable, then generalization holds out-of-distribution For proof see Appendix F.2
Lemma A.8 (Learning UTM from Certifiably Identifiable Data). Lei ((Zi, yi HElN be
dataset satisfying y: = e for è E S2 being the expressibility parameters of Lemma B. for
the set of 7M's/words (M,x) ( E, n. b). If D is certifiably identifiable with Xmin (Ap) then there is poly(d, N, Q, E, ft, 2 n time algorithm that output set of parameters θ 2aze such that for all TM's M and input words x in (12, E, n, }), we have
(H) (: - -oo =x'
(23)
The elt) step of the auloregressive computation history of is equal the l'th step of the computation history of M on x.
D Proof of the Main Theorem
Theorem 2.2 (Learnability of Linear Attention). Let D be dataset D = (Ze Vslieln drawn i.i.d. from a distribution D where each Zi E Rdxx., yi E RR. The embedding dimension d is fixed across the dataset, whereas n, can be different for each datapoint. Let Amaz = maxie[N] T4 be the maximum sequence length, and let St be the space of H pairs of value and key-query matrices {Vh,, (n)]nE(BH for any H E [1, 20). Then there is an algorithm (Algorithm [ that runs in time and that, given such that with probability 1 5,
pairs {(Zx, yi) bE(.). returns {(Vn, be)M E Ω for H S d
((Z.))DD [[|MHLAg =1l?]
minee RIZWED [[|MHIAO(Z)- - -y|13] Se (6)
with sample complexity N O (2 (d loo((o0)))) 1)).
Proof. First write down the loss:
Le
Σ
Σ NE(H]
(28)
(E[N]
2
Σ
Σ hE(A)
T V. (n2/z, - yi,a
(29)
N
(E|N]
aE[d]
Observe that the one layer attention network is a quadratic polynomial in (V,, Qu heM
Σ Σ
((To, Xi,)) yk.a)²
(30)
N
JE[N] RE(d]
Here
Vn 01,00
Vn,00Qu,01
Vn Vn. Qn.ad
Vn.01 Q1,00
Vh.01 Qh.o1
Te ==
Σ
flatten(V, (itenn/)))aatee(())) (Q,)"
kE[Ay
he(m]
Vh.aaQr.or Qh
Vn.aaQiada
(31)
Now we relax the objective where we replace Te with an unconstrained matrix W &x Another way to put it is that Te is rank-H but W can be a general matrix. Because the space of general rank matrices is larger, we have written down relaxation guaranteed have a smaller loss. Furthermore the loss can be optimized via ordinary least squares.
17

---

min Lw (i)ENN)) = N
Σ iE|N] aEid]
Σ
(W,
Vi,a)?
WER xa2
min Le Coll(Z) +e (32) Ben
Thus the optimum of the regression with respect to the data achieves optimum of the loss to error
in time d The sample complexity to achieve error is then ( log(8-1)) 1]) with probability 1 8 over the data distribution. Furthermore, if we take the SVD of W Eie( A where absorb the singular values into the left and right singular vectors we have for θ (Fo,
Let Vn Fold(A,) and Qn (o,d(
be(((Z)W)))(M))
Σ iE[N]
Σ ME(M]
(z Q y4
N
2
N
E
E NE[N RE[d]
Σ (1 (A((A)
Vi.a
Se (33)
desired.
E
Proofs from Identifiability Section
First, we start with general lemma (Lemma E. which states sufficient condition for identifiability of any model class that can be written as an inner product of polynomial of parameters θ with a poly- nomial feature mapping H. If the data is realizable by the model class and AD = ED (2))) full rank then the model class is identifiable with respect to D.
The following is the certificate of identifiability written in an abstract form involving polynomials to map parameters to feature space and polynomials to map data to feature space. The proof does not require the model to be an MHLA, but we state it in MHLA terms for the sake of concreteness.
Lemma E.1 (General Certificate of Identifiability) Let dataset D) = {(Z,, [:)]iE be a dataset realizable by θ E Qg. Letp:= == beaoe be a collection of polynomials Pa 9 + R mapping the
parameters O E Ω to a feature space of fixed dimension vb E zt Let H. H J be uniform Let p and H satisfy
family of polynomials such that H, Rddx + Rdb
MHIA [a] (pa (z))
(34)
for all Rddx for all n (1). Then if Ami (ED [H(Z)H(Z)TT] > 0, we have
MHIAg MHIAe
(35)
for all empirical risk minimizers O O' E ΩERM- That is, all empirical risk minimizers compute the same function.
Proof. We construct map p: Ω > IR such that MHLAe MHLAA if and only if p(0) ppe)). Then we show that any empirical risk minimizer θERM and the ground truth θ satisfy p (0ERM) O()).
In more detail, we construct some polynomials {potaetd and family of polynomials H such that
MHLAa = (9) H(Z))
(36)
We construct linear model class R that takes as parameters E RV and data H(Z) E R such that
R. (H(2)) (v. H()))
(37)
18

---

Let OERM be defined as
FRM := {O' E 22/6' = arg min [RE[N] (4(MHLAA((2))
(38)
BEΩ
Let UERM be defined as
VERM = fu' R 1 arg aargginn Eig[N] (C(R y:)l
(39)
VER"
Observe that for all @ E GERM. we have p()) FERM. Here we use the fact that y is realizable by the ground truth O. Therefore if we show that FERM is unique, i.e comprised of single element then PERM == (p(0)|@ e OERM) is also unique. Therefore, MHLAg is the same function for any θ E θERM
To show UERM is unique, all we need is that the second moment of the features AD Ep [H(Z)H(2)²] is positive definite (the covariance has a minimum eigenvalue bounded away from zero).
Next we prove the main certifiable identifiability lemma by instantiating the polynomials H and from Lemma E.
Lemma A.3 (Certificate of Identifiability). Let dataset D = {(Z,, y4 HiENI be a realizable dataset.
Let H ) be family of polynomials Hn. -Rdxn - IR for u/ (2) d + dl defined as follows. We index the entries of H. by taking the Kronecker product betveen all sets of pairs {j. k} (for all j, k [d) with with all E [d]. We define (5,3e as in Algorithm be
H(Z) G.Aee == (2g: Ata)
(18)
Then if Amin (ED [H(Z)H(Z)?]) > 0, we have that MHIAg is identifiable with respect to D.
Proof. that
First construct a polynomial p. Ω R and H. Rdxx + R for 2b (2) d dl such
MHLAa [2) ((),
(40)
We begin by rewriting (2)] We index the first (4) entries of pa (9) by all pairs {j. k} for j, [d] and all e [d].
Σ (VnaaQuake + anaak
(41)
hE(M]
We define the entries of Pa (0) from d. ( + d2) as follows.
Pa ())) = Σ
Vh.anQh.fe
(42)
BE(m]
Similarly, we define H(Z) be be the following (9) + de features. H H(Z) and H(Z)
H(Z))au((0) == (2jz Dkn)
(43)
and
H ((()))))1)) -112j:112
(44)
Thus rewrite (H)A)(Z)(a) as
(2)A = Σ Σ
H(2 21 + Σ (2(0))))) j.eE/d
(.A)ES tE|d
= (p)(A), H(()) (45)
Here introduce the notation S to denote the set of all pairs [j for j. E [d]. We have
constructed a polynomial Pa (O) such that for any O E Ω in the same equivalence class MHLAgr. Furthermore, if there exists b E [n]] such that
Pa (+) = P. (O'), have MHLAg Amin EED [H(Z))((Z)]] > 0 then OLS returns a unique solution for (O). Since the data is realizable, conclude Pa (+) (+) for all θ E ΩERM-
19

---

Next we present the proof that realizability is not necessary to identify the function learned by MHLA with more than d² heads.
Lemma A.6 (Idenifiiaiility without realizability for MHLA with arbitrarily many heads). Let dataset D {(Zz. FEIN] be any dataset drawn i.i.d from a distribution D. Let H be defined as in Lemma A.3 Then if Xmin (E))H
> then MHIAg for O E Qu for any E [dP, 20) is identiiable with respect to the data D. That is,
MHLAg = MHLAey
(21)
for all pairs of empirical risk minimizers G 0,0' E 4ERMM
Proof. We know from [lemma main algorithm] there exists a surjective map (Q) that takes θ E f into E RV. This implies that for all v E E there exists a right inverse function p"(u) O satisfying p(@) given by SVD. Therefore, p(DERM) E VERM i.e optimizing over u R does no better than optimizing over @ E Ω. To prove this consider the contrary that there exists v' E UERM and there is @ S2 that achieves the same empirical risk as v'. However, p" (v) S is such θ, and we have a contradiction. The key point is that we avoid the assumption of realizability and replace with surjectivity of the polynomials Pa-
Finally we prove that data drawn from independent noise is certifiably identifiable. A subtlety in the proof is that normalize our features, which still satisfies the assumptions of the general certificate Lemma E
somewhat different set of polynomials than Lemma AA.3] as we center and
Lemma A.5 (Independent input noise yields identifiability). Let (Z,y)~ D be realizable dataset. Let Z be drawn from a distribution where the (a, b)-th entry of Z denoted by Zus is drawn i.i.d. from distribution over IR for all a E [d] and b E [n]. Let the second and fourth moment of be denoted m2 and ma respectively. Let my > and m, > m². Then MHLAg for Q E QH is identifiable with respect to D. That is to say. for any population risk minimizers O O' E -ppRs:
MHLAg = MHLAgr
(20)
Proof. We give the entries of A()) the following naming convention. Let the terms (j, and pairs {j', k HL') Terms that involve [{)] and {(²)(") are referred to as 'singles'.
E [Ho( (LAME) Ho ()(2))) = _ KZgs (Frt,, (..))
(46)
We give entries of the following form the name "singles to singles"
E Ho (2)()) (e) (z,(2) = - 2(0 - nmz) (1 nm2)
(47)
For the case of drawn with each entry i.i.d we can proceed via case work.
Case 1: Pairs to Pairs, j * and j' + K'
1. Subcase 1: {j,k) # (i','' and l = 8::
E B2(s,
(48)
n
2. Subcase 2: [j,k) = (5', k1 and i - E':
EZzz 220a m2
(49)
n
Case 2: Singles to Singles, k and jr = k'
1. Subcase 1: j # j' and = f/1
E. [[Il=]:1|2 nmmz) (11zy:112 nm22) 2261 =
(50)
n
2. Subcase 2: j 3 and = fl:
1 n
1 ((n2 nmm3 + mmm - n22m3) m2 (mm
[11=))11p -nm2)) 220 n
m2 )mz
(51)
20

---

Case 3: Singles to Pairs, j = k and j' + kv
1. Subcase 1: - f":
[[|z]il| nm2) (2y1) =mr))² = 0
(52)
Finally for the feature HZ)) m220 we have on the main diagonal Emm]z2]] m2 and everywhere else.
Therefore we've concluded that AZ) is block diagonal matrix because the & + l' blocks are near zero. All that remains is to verify that the diagonal blocks are full rank.
1. Pairs to Pairs: m I is full rank with min eigenvalue mi
2. Singles to Singles: (m4 mm) is full rank with min eigenvalue (m4 - 1014-m3)mma )
Finally we provide simple error bound for approximate empirical risk minimizers to demonstrate the robustness of the conclusions in Lemma A.3
Lemma A.7 (Identifiability with Error). Let Ω be the set of e-approximate empirical risk minimizers,
Q Qu-ERM
0 { E QH E(zen)en (MHLA - wi) sc).
Then we have for any O, E Qle-ERM that for all inputs E Rdxn
' Amiu (Ap) 11 1
[H MHLAGr S
(22)
Proof.
MMLLAa(2) - MHLAg (
Σ (1pa P. H(2))
aEid]
Ilaa Pa (2))))))
a
[ea pa()))) ||
/E/d]
12
Amia (Ap)
(53)
Here the first equality follows from the linearization exhibited in Lemma E. The first inequality cauchy schwarz. the second inequality apply a crude upper bound that more than th degree polynomials that are products of three squares of entries in Z are involved in [H())
12 <
Σ
2o 22 S |
(54)
a.a', Ea'dd, 6,6 M.((()
The last inequality comes from the fact that O, O' are approximate empirical risk minimizers. Therefore we know
Aminn (Ap)
Σ (l.(()) - P. (0))
Σ (p)(()) - (0)) (())) <e
(55)
aE(d
aE|d]
21

---

which implies
(lpa(0) - Pa (0' aEid
Σ
(56)
Amin (Ap)
which concludes the proof.
Lemma E.2 (Identifiability with Error and Noise in Realizability). Let D = {(Zi, (<), KEIN be dataset such that yi = MIILL(ZZ) + So for Sa iidd and bounded. Let -e-ERRM be the set of e-approximate empirical risk minimizers.
{o E Qn BEDDDED (MILLAe(Zz) -v:)² sc)
D 2R-FRMM FRM
(57)
Let max4e[M] || / r S B. Then we have for any O, O' E - that for all inputs Z E Rdxx
MHI - MHIA (Z) S c DiE(N C + B² N
log(6-1
|
(58)
Amiu (Ap)
Proof. The proof follows directly from Lemma A.7 but we incorporate the Si terms as is standard in analyses linear regression.
22

---

F Programs Expressible as Fixed Depth Linear Transformer
In this section we build out examples of programs that can be expressed as fixed depth linear transformers. Expressibility results can be carried out in a variety of equivalent ways. The main takeaway, is that the computation history of TM M word æ. when written down "step by step" can captured by next token prediction of linear attention. This is because the key-query-value naturally implements table lookup sometimes referred as "associative memory" or "in context linear regression" in the linear case.
The notion of an Autoregressive MHLA Program is useful for condensing the proofs of expressibility, We write such programs an object oriented syntax with each token representing an object with multiple attributes. Attributes can be updated and looked from other objects using a generalized lookup akin to associative memory.
Algorithm 3 Autoregressive MHLA Program
1: Instantiate N instances OBJ = {obj(i) of Class with set of Attributes (Atrri, Attra,
Attri
2: Each Attribute takes on values in alphabet Exunbue
3: for iter [T] do
Let obj[r]
rightmost token
5: Let obj(r + I] he new token initialized with positional embedding obj(r + 1]-pos = +
6: for each [AttrSource, AttrDest] in PPairs of Attributes in Class) do
7:
#AttrKey and Value can be any pair of Attrihutes (and can be distinct from VarSource/VarDest LookupDict {{ob). Attree::t obj.Attr Value) for obj in OBJ}
9:
# if multiple objects have same AbrrKu then returns sum of obj.Attr Values which we aim to avoid
10:
Let Be be any function from MAmsSouree to Minnee
11:
Let By be any function from EamrValue Ammeett
12:
Let query [
13:
query in LookupDict. Keys then
14:
= By (LookupDict(qurry)
15:
end if
16: end for
Append next token OBJ {obj[i] he(r) (obj[r 1]}
18:
r = +
19:
end for
Lemma F.1. For any program P written in the form of algorithm there exists corresponding MHLA parameters @ Qn such that MHLAg PZZ).
Proof. We set some matrices to implement lookup tables. For any function of f + for sets A and B there is canonical representation of the input domain as orthogonal unit vector U1, D2 ..., l/ll E IR and output domain as another set of orthogonal unit vectors U1 U2,,.., U|a] E R Therefore, there is a matrix G2 that maps input vectors to output vectors satisfying Gyv; = uj for ((i) for all E [A] and j E [B].
For functions f AuuSouree + ΣAmkoy and f EAneValue + EAndDest we associate matrices Bq E TR. Exursoase |x [Esnkey and By E R RREEurvitel x EAnntea respectively.
Then we form (Vn, QhENE follows. Let V be the matrix that is all zeros with By in the rows associated with Ammrsorre and the columns associated with Enmrey. Let Q be the matrix that is all zeros with By in the rows associated with AauVatue and the columns associated with EmmDest-
In each layer we have multiple heads, each one performs the lookup operation for each pair of attributes in the class.
F11 Construction of UTM
Now we proceed with our construction of an Autoregressive MHLA-Prograe for UTM. The UTM requires small number of operations captured by Autoregressive MHLA-Program.
We define embedding function that takes as input a TM M and word such that
23

---

Definition F.2 (Embedding). Let M be TM over state space Q, alphabet A, transition function 8. Then
qo
q1
#
do 8(10, ao) 8(90: 01)
0g
#
Embedding(M)
((q), ao)
8(g,, an)
#
(59)
Q1
#
8(11, (1))a1)
8 (g. a1)
#
Let .1.p2,..,,5 be "positional encodings" that assign unique id's for every letter in the word x.
P1
P2
Pi+1
Ps
#
Embedding(z)) =
21
22
24
28
#
(60)
0
0
q
0
#1
Then we define (mbeddinggM.x) to be
Em))ddinggMMx) [Fmbedding(M
(61)
Embeddingg())
Henceforth we will write the construction in the syntax of an Autoregressive MHIA Proggam instead of matrices with blocks of zeros and token embeddings to save space.
Lemma B.1 (UTM Expressibility). Let A(ê, E, n, ②) be the set of Turing machines M {ō, Σ, 2, Istart, quccept, Oreject) I and words x E E" with number of states, size of alphabet, size of inpul, and number of steps in computation history bounded by Q, S2 â. & respectively. For any x) E A, let {xe] (Ee))e be the computation history of the UTM on (M,x). Let the autoregressive computation history (see Definition (C2) of MHLAA on input (M,x) be denoted CHe(M, x)) = (Z), ,...... Z.... Then there exists a set of parameters E Qu for = O and embedding dimension d = OOâ max(E. 2)), such that for all (M, x) E A, the TM computa- tion history at time step t is equivalent the autoregressive computation history at time step c(t) where e(l) S O(n + 1)L) i. [= -length( x x². Furthermore, this can be achieved with bits of precision.
The construction is given in the language of Autoregressive MHLA-Progran in algorithm 6 which provides the instruction set for writing the next letter in the computation history onto the output tape.
Proof. Proof Idea: A few elementary operations can be captured by MHLA MHLA-program program which can be composed to output the computation history of M on æ. We begin by introducing some notation for the "Lookup" operation which we build into copy, move, and if-then which are all the operations required to construct the UTM.
General Lookup: For each lookup there are three objects that are involved. Let Token= obj[r] be the "source" which is always the rightmost token. An attribute from the source ob- ject known as AttrSource is linearly transformed to form "query". Lookup involves a table which used to match an AttrKey look up an Attr- Value from an object obj[p] that we denote the "target". Note, that if the obj[i] has an AttrKey that is zero, it is the same as not being in the table. In the pseudocode algorithm o these zero attributes are denoted as "None".
T T= Attrey: Aurri).AuVallue)
Given a query, we copy the associated Vatue from the lookup table T and update AttrDest in an object NextToken= obj + 1] which denote the "destination". Multiple lookup operations be performed in parallel by many heads with each head responsible for single lookup.
To output each letter of the computation history, we increase the number of tokens by a constant c. We refer to the set of contiguous tokens [0, c), [c, 2c], etc. involved in the computation of a single letter "block". Here block[i] act on each token and enforce that the nonzero rows that each block of tokens occupy are disjoint. Furthermore, within block, the states of each token occupies disjoint set of rows except when they are used to construct a table. Tables are the only case where we want tokens to occupy the same rows. In this manner the following abstraction can be made.
We construct different set of heads to
At the beginning of each block starting with obj[r), we can lookup attributes from anywhere OBJ that we want to load into different attributes in objír]. Then apply any sequence of if-then
24

---

statements involving the attributes of obj[r] to update the attributes (or create new attributes). To run the UTM we need few simple primitives denoted Lookup and Ih-Then.
Construction of Primitives: We write down the construction by constructing a sufficient set of primitives Lookup and If-Then. We also include Copy which is special case of Lookup that is used frequently.
Lookup: When the transforms Bg and By are the identity we denote the lookup operation for table I where we query an attribute s' of obj(r) to update the attribute s of obj[r+1] as obj[r+1].ss Lookup(T.ojjrl..s)
Copy: A special case of lookup is copy, where we need to copy attributes from tokens that are at an offset - for & E [r]. This can be done by setting Be to permute the positional encoding by positions. Then the query matches the key that is the positional encoding of the target object. Let s, s' be target and destination attributes. We denote the copy operation of the attribute s' the obj at offset k from into the attribute of the destination object to be obj[r+1].s Copy(objlr-k)ss),
If-Then: We write down an If- Then Program algorithm 4 and a corresponding Autoregressive
MHLLA-Progran algorithm to implement If-Then. An If-Then program looks up whether an attribute a is equal to any of attributes @) .... Og then we set attribute x' to bı L2, .... by respectively. This achieved by copying the attributes a; and by into dummy attributes s0 and sl for all in for
...,
series of l consecutive tokens. This creates a table with key s0 and value sl. Then we use attribute x as the query, which looks up the corresponding value s1 which we use to update an attribute x'.
Algorithm 4 If-Then Program
1: # attribute equal any of a1, 02, ag then attribute x' to bı, be. . be respectively
2: if Token.x Token.o then
3: NextToken.x' Token.b;
4: end if
5: if Token.x Token.ag: then
6: NextToken.x' Token.bz
end
8:
9: if Tokenx == Token.ag: then
10: NextToken.x = Token.b.
11: end if
Algorithm 5 MHLA Ih-Then Program
I: # attribute x is equal any of di, d2; ar then set attribute x' to bj, be.... by respectively
2: == token(r].a]
loken[r+1].s1== to
4: NEXT r
5: =oken[r+]],00=0000 token]rr) t0
6: token[r+1].sl = token[r].bb
...
8: NEXTTOKEN r= r
9:
10: token[r+1]] [].s|++++Is_l_l1 token[r].bb
11: NEXTT r
Table T= [objli].s0: [0b1iji.....oobjjjllshhaea--a--
13: token|r+1]x* Lookup(T.tokentri.x
F.2 Proofs For Learning UTM
Lemma C.5 (Learning a UTM). ④ QH in dimension be the MHLA parameters in
Lemma B. Let (M,eN be pairs of TM's M and words of maximum length drawn
25

---

Algorithm 6 Simplified Instruction Set MHLA Program for UTM for a single block
1: # Initialize Lookup Tables for TM M and tape Ti
2: # ō(g. a) = [next-state, next-letter, next-move]
3: M {q [do. 8(g,
(o)) (1, 8(g, a1 di)l}ee
4: 71 = [tokenfi,PossEncodingg tokenfi).Letter)iie
# Begin Loading Information from and previous tokens on tape
6: # First copy letter/state from token -N-1 positions away
7: # Attribute s(-1) = |letter, state] | where state can be equal to None
8: NextTookenss--
9: # Sccond copy letter/state from token -N positions away
10: Attrihute s0 = [letter, state) where state can cqual to None
11: NexiToken.ste s0 Coopooooooooooaaa...
12: # Third copy letter/state from token -N+1 positions away
13: # Attribute s [letter, state) where state be equal to None (oaoeooooonnNt]sss
14: NextToken. sl
15: NEXT TOKEN r = +
16: #Split into three branches to handle left, head, and right positions relative to head
17: RUN BRANCH (Token is Left of Head Position) See algorithm 7
18: RUN BRANCH (Token is Head Position) See algorithm y
19:
RUN BRANCH (Token is Right of Head Position) See algorithm
i..dd from a distribution D. Let Z, = Embed Embed(M;, For each TM/word pair (M,, 2;) let CHe(Z.) {Z, Z2 .... Z } be the --step autoregressive computation history of MHLAg Z,. Let D be the dataset D == {(CHe(Z) (Z.)* v, ++ he(N) where yf+1 MHHLAa (Z,). Then
A applied to input D returns E QH for H S de such that with probability 1
E(D,)ED (MHIAA -
(26)
for sample complexity N poly(d, E (oood,,,, Then with probability 1 over the randomness in the data, the probability over 'D that the D-step autoregressive computation history CH& (M, and CHe (M, x)) differ is upper bounded by
(CM.
+ CHe(M, O(e)).
(27)
Corollary F.3. In particular, for sample complexity N = poly(d, c- log(6-"),n. n,,,,, t), by Lemmal B. we have with probability 5 over the randomness in the data that the probability that the c(t) step of the computation history of MHLAg is equal to D, is
PI(M,))~D (M,)(M,))<0) -ki] =x]==x 21-c,
(62)
where c(t) < <0(nn + t)t). That is, the computation history of the MHLA returned by algorithm ? is equal the computation history of M on æ.
Proof. We have from Theorem 2.2 that algorithm returns Ô such that
(() (MHLAA(Z) v)²
min E(Z,))D (MHLAs(Z)- - Se
(63)
eesh
Then to obtain an error bound on the & step computation history, which involves O(n tokens, we just observe that by union bound each step rounds to an incorrect set of tokens with probability less than e. Therefore, over ((t) steps the error probability is upper bounded by $b. Equivalently
Pr (Mar)>>
[CH (M, x) * (M, < O(e)).
(64)
Then proving Corollary H3 is simple exercise. For larger sample complexity N by Lemma B.I we have that the probability that every token of the autoregressive computation history of MHLA & is equal to I, is
poly(d, ( "ooii==",ooooo00 log(8-1), ), t).
Pr (Ma)>D
(H, -hin] xt 2 -
(65)
26

---

Algorithm/ 7 Branches to handle cases Left Head, Head, and Right of Head
I: #Split into three branches to handle head, right positions relative to head BRANCH (Token Left of Head Position)
3: # we have loaded a state q into s1 (if left of head) and next we load (ao, 8(g, ao), 01. 8(g. a1)] into s2
4: NexcToken.s2 = To.een .staae))
5: NEXT TOKEN
+
Token.s2.letter do
7: NextToken.s3 = 6(4, 0o) = [q w ./R]
8: end
9: if Token.s2.letter = then
10: NextToken.s3 8(g, 01) lq _LR]
end
12: NEXT TOKEN r
13: if TooeTook.mmomm = L then
14: NextToken. rerTokeenretum-lettter Token.s0.letter
NextToken.retur- state =
16: end if
17: Token.s3.move = I. then
18: NexTToken.retum-letee Token.sO.letter
19:
Texffoken.returm-stur None
20: end if
21: BRANCH (Token is Head Position)
22: # we have loaded a state into s0 and next we load [4o, 6(4, a1, 001,0(0, (1).)))) into s2
23: NextToken.s2 Lookup(M, Tokae.sO.statee)
24: NEXT TOKEN
25: T.ltn.s2.letter = On then
26: NextToken.s3 = (gg ao) [q" w _LR]
27: end if
28: T.lten.22.eetter an then
29: NextToken.s3 8(g, 01) lq .w L/RJ
30: end if
31: NEXT TOKEN r= r+3
32: if moken.s3.next-leee is not None then
33:
token.....ext-letter
34: NexfTooee..rttum-ssie
None
35: end if
36: BRANCH 3 (Token is Right of Head Position)
37: have loaded state q into s(-1) and
load [ao: 8(q, ao), a1, 01, a:)] into s2
38: NextToken.s2 x2 Lookup(M, Token.s(-)))eeeee.e
39: NEXT TOKEN r r+3
40: if s.Toke..s2.etter = 0o then
41: NextToken.s3 = 5(g,
[q" wq"'"" L/R]
42: end if
43: 2etter = a1 then
44: NextToken.s3 = 6(g, 01) = [q" w [/R]
45: end if
46: NEXT TOKEN r+3
47: T'Toeen.s.mmnn move
Token.s0.lett 49: NexfTokennrehum--stae = None
48:
50: end if
51: ToToookk..mom. = then
NexiToken..elum-leter Token.s0l letttrr
53: NexfToken..etummsstte = Token.s 2..nn...... next-stte
54: end if
27

---

Lemma A.8 (Learning UTM from Certifiably Identifiable Data). Let D = {(Zz, y4 ii be dataset satisfying Yi e for Q2 being the expressibility parameters of Lemma B.
the set of TM' s/worrss (M,z) ( Σ, h. b). If D certifiably identifiable with Amsn (Ap) p) > n. then there is a poly(d, N, Q, E, f. & n7 ') time algorithm that outprt a set of parameters θ 22
such that for all TM's M and input words x in 4 E, n, 4) we have
2x) (: - kia] x² (M,
(23)
The c(t) step of the autoregressive computation history of Q is equal to the t I'th step of the computation history of M on x.
Proof. The proof follows from the quantitative version of Lemma A.7
Using the given that RRdx
Amia (Ap) > 77, we conclude that for any Ô E QR-ERM that for all inputs Z
|MHLAg MHLAa - " 1
(66)
we select a sufficiently small c c= 1/ N. [Q), [-], n, 1, ) then we can ensure
Pr (M.))~ (HQ(M,
2z)()) (()
kin] x² 2 1
(67)
The runtime then scales with N, [2) 17 m, 1) as desired.
28

---

G
Additional Definitions
Definition G.1 (Orthogonal Embeddings). Let Embed be a function Embed: Σ - RR Let E be an alphabet and let e "....., é|z] E: R 1s) be a basis of orthogonal unit vectors. Then for each letter in an alphabet Σ, we define Embed(a) ca where we associate a different unit vector to each letter.
We adopt naive "rounding" scheme for converting vectors into tokens. This be done in variety of ways, and we choose to simply round the entries of the vector embeddings to the nearest token embedding.
Definition G.2 (Rounding). For any vector v = (V1, U2, va) E R& let Round(e) = ej for j = arg me (v, ez). Since use orthogonal unit vectors for token embeddings we will refer (ound(v) as a token, We will often refer to a matrix Z E IRdxn as being equivalent to a series of n tokens 01 -2: Og mean Round ([;; ag for all i [m]
Algorithm 8 Extract Features
1: Input: Data == (SINN for Z, E gdx=; and ys E R2d
2: for Zi E D do
3: Let 21, #2, ... be the rows of Zi and let #a.b be the (a, b) entry of Z.
for j E [d] do
for E [d] do
6:
for & E [d] do
7:
Let Xr Rdxdd be defined as follows
8:
Xi kd f] [(2y, ax:) Ren))
end for
end
11: end for
12: end for
13: Return: {XI such that
$1:s #1. 2ing
(#1 #2. Dine
(21 Rdl Sing
(#1:s 2d.) ) Rdm.
(#2:, 21:)2in
(#2:) 22. sime
(222 Ad. Zime
(#22 Ae. 2ddi
X; =
(68)
(Ea:; 21m)
(Za:s #2:) aine
(Edi Ad:)
(#et )
G.1 Training details of attention networks
use Adam Kingma and Ba (014
optimizer to train linear attention model Equation 4 and models.
full Transformer Vaswani et al. 2017
hyper parameter
search space
d input dimension
[2, 4, 16)
number of heads
[1, 2, 4, 8, 16]
n number of layers
[1, 2, 4]
learning rate batch size optimizer
[0.01, 0.001]
[32, 64]
AdamW Loshchilov and Hutter (2018)
G.2 Training details in DFA Execution
We use Llama variant of the Transformer arhitecture from Touvron et al. 2023 We run each setting with N number of training examples with the following different values N E (16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 6144, 8192, 12290, 16384, 20480, 32768, 65536). The other hyper parameters are given in the below table.
29

---

hyper parameter
search space
d input dimension
[2048]
number of heads
[16]
n number of layers
[4]
learning rate epochs optimizer
[0.00025]
100
AdamW Loshchilov and Hutter [2018)
G.3 Additional Experiments
MSRiy.)
10
10
10
100
spochs
(a) N 2048, d
Figure Performance comparison multi-head, multi-layer linear attention models and the original Transformer model (denoted as full). We trained using SGD on synthetic data generated from single-layer linear attention model for varying training set sizes (N) and input dimensions (d), number of heads m, and number of layers We present mean squared error of the predictions w.r.t number training epochs. Results demonstrate that multi-head architectures converge faster on different input dimensions and match the performance of our algorithm T (convex algorithm). Increasing the number of layers or incorporating multilayer perceptrons (MLPs) and layer normalization did not yield consistent improvements. Shading indicates the standard error three different runs.
G.4 Learning the Computation History of Deterministic Finite Automata
Universal automata (like the universal Turing machine discussed in Appendix F22 receive descriptions of other automata as input, and simulate them to produce an output. Here we empirically evaluate the ability of MHLA models to perform universal simulation of deterministic finite automata (DFAs). We limit our study DFAs with a maximum number of states (N), alphabet size (V), and input length (L). While recent work on in-context learning [Akyürek et al., 2024] bas focused on inferring DFA behavior from input-output examples, here, we aim to simulate DFAS given explicit descriptions of their state transitions as input-a a task somewhat analogous to instruction following in large scale language models.
The construction in Lemma C.5 shows that a linear attention layer can output the polynomially bounded computation history of any TM (and therefore any DFA). Our construction requires embed- ding size linear with maximum length of computation history, number of states and alphabet size. Therefore, we predict the data requirements are polynomial in each of N, and L.
Dataset Our dataset consists of strings containing three components: the input DFA's transition function d : Q x Σ + Q, the input word x E " and the computation history h E Q which is the sequence of states visited in the DFA it decides if æ in its language. The first components are the input to the model, while the computation history is the target output. We adopt the following schema for representing 5, x and h:
(Si, W. (7,8))..... Vs E Wowl....ug wL
(s "ps')) (s² (ewwsss s²) ....................... (st-1w wus
DFA transition lunction
word
computation history
30

---

24
Number of States (2)
Word Longth (L)
A.ptabet Size
Figure 4: Data requirement for universal DFA simulation: We train fixed sized Transformer (4-layers, 16 heads and 2048 hidden dimensions) simulate DFA given transition table and input word. The vertical axis shows the number of tokens (expressed as word length L times the number of examples Q) required to obtain 99% next token accuracy.
We encode each input-output relation in the transition function as a sequence of three tokens (8,, w,8). 85) where 8(s; w) 8j. We also include two parantheses to separate each triplet of tokens for total of five tokens for each input-output relation. The total description length of 8 is then 5QE. We encode word of length L as sequence of L tokens. Finally, encode the computation history as the sequence state transitions the DFA visits when deciding if
its language. Here we designate s0 as the start state, and let a² = 0(s-_1 1 Each state transition is again represented by triplet (s, w, 8(s, w)). We train an autoregressive Transformer model using cross-entropy loss to predict the computation history tokens given the transition function and word. Please refer to Appendix G.2 for hyperparameter details.
Results In Figure we vary each of the parameters Q, and V, while the other parameters are fixed to constant (in this case we fix them to be 4). Then, on the vertical axis, we display the minimum number of tokens (number of examples times the word length) required get 99% accuracy on the next token prediction. Plots are suggestive of a sub-exponential dependence on DFA complexity.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Both in the abstract and the introduction, we accurately describes the scope and objective of the work.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss our limitations.
Guidelines:
31

---

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
We include all proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We will release code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
32

---

If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/ htppllccguides//// http://guuiiee//eeubbbieieeeee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https //nips.cc cc/ http://lllgggg..eee/eeeeeeuueeee http:///oooosuuumssssooooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
33

---

Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We provide details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
We provide statistical significance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emorr bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report I-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: We have dedicated section in Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
34

---

9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We do not foresee any such issue in work.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Not applicable
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
35

---

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Yes, we make sure to fully credit them.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, pappessii thcode www.com//a/oo/oc datasat has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not produce any data asset.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
This is not applicable to our research.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
36

---

Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This is not applicable to our research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing. editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
This is not applicable to our research.
Guidelines:
The answer NA means that the method development in this research does involve LLMSs as any important, original, or non-standard components.
Please refer to our LLM policy (https /tap:r:///euris.... cttpiiisssc.o///oooo /Cco///rnnnee//25///////// for what should or should not described.
37


Reviews:
[{"review_id": "kOx5QvEmpy", "reviewer": "Reviewer_w6rd", "summary": "This 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).", "questions": "Q1. 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?", "limitations": "The authors discussed limitations in Section 6.", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "My original concerns regarding the experiments (Q1, Q2) and practicality (W1) has been convincingly address by the rebuttal: The authors provided clarifications on Q1/Q2, and provided a plausible direction for future work to improve on W1. I think this work provides an original and significant idea for studying multi-head linear attention mechanism which opens a new avenue of subsequent work and would like to recommend acceptance.", "rebuttal": "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.\n\n> 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?\n\nA1: 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.\n\nHowever, 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.\n\nA 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.\n \n> 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?\n\nA2: 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.\n\nOur 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.\n\n> 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.\n\nA1: 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.\n\nOur 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.\n\nInvestigating 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."}, {"review_id": "grjniihiI8", "reviewer": "Reviewer_bW3t", "summary": "The paper investigates learnability of a simplified transformer architecture, i.e., a multi-head linear attention (MHLA) module. \n\n- 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. \n- They also characterize an identifiability criterion on the dataset that certifies that every empirical risk minimizer is functionally equivalent. \n- They show that  MHLAs can (autoregressively) express universal Turing machines\nwith polynomially bounded computation histories, implying a strong learnability condition for UTMs.\n- They verify their certificate empirically on an associative memory task.", "questions": "-  Is the condition Lemma 2.3 a necessary condition for identifiability?", "limitations": "The authors address the limitations quite clearly in the paper.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- The MLHA's reduction to a kernel learning problem and the resulting PAC learnability results are solid theoretical contributions towards understanding learnability of transformers.\n\n- The identifiability condition for equivalence of MHLA's learnt via ERM is potentially very relevant to learning theoretic aspects of transformer learning\n\n- The experimental results provide sufficient evidence to support the relevance of their theory.\n\n**Weakness**\n\n- MLHAs are quite a toy-setting compared to real-world transformers.", "was_revised": true, "final_justification": "I thank the authors for their detailed rebuttal. \nI was already quite convinced about the paper, so I will keep my score. \n\nAll the best.", "rebuttal": "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.\n\n> Question: Is the condition Lemma 2.3 a necessary condition for identifiability?  \n\nA: 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$\nheads, 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.\n\nHowever, 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.\n\n>MLHAs are quite a toy-setting compared to real-world transformers.\n\nRegarding 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.  \n\nFirstly, 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.\n\nSecondly, 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.\n\nToward 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."}, {"review_id": "H0016oSBwF", "reviewer": "Reviewer_EkNE", "summary": "The 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.", "questions": "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 ...).", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "Thanks for the response.", "rebuttal": "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.\n\n>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 ...).         \n\nA: 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).\n\nTherefore, the runtime for verifying this condition involves two main steps:\n\nComputing 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$).\n\nComputing 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.\n\nWe will ensure this explanation is clearly stated earlier in the manuscript.\n\n>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)\n\nA: Thank you very much for pointing out this lack of clarity; we agree this section requires more precise explanation.\n\nThe 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.).\n\nEven 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.\n\nOur 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."}]

Author rebuttal:

Meta-review: There is an enormous need to reconcile the success of modern transformer-based LLMs with the constraints implied by VC-dimension and PAC-learnability.  This paper takes an important step, which is why I recommend oral.  It shows the high representation power of linear attention while also admitting its limitations, and it proves strong agnostic learnability results for this representation.  The proofs are well-written and appear correct to me and the reviewers.  Reviewers had a great back and forth with the authors, and we ended with one accept and two strong accept recommendations.
Decision: Accept (oral)