When can transformers reason with abstract symbols?
Authors: Enric Boix-Adserà
Omid Saremi
Emmanuel Abbe
Samy Bengio
Etai Littwin
Joshua M. Susskind
Venue: ICLR 2024 poster (2024)
Keywords: transformers
language models
reasoning
theoretical analysis
variable binding
Primary area: learning theory
TL;DR: Transformers learn to reason relationally given enough data, and we improve data efficiency with theory-inspired architecture modifications.

Abstract: We investigate the capabilities of transformer models on relational reasoning tasks. In these tasks, models are trained on a set of strings encoding abstract relations, and are then tested out-of-distribution on data that contains symbols that did not appear in the training dataset. We prove that for any relational reasoning task in a large family of tasks, transformers learn the abstract relations and generalize to the test set when trained by gradient descent on sufficiently large quantities of training data. This is in contrast to classical fully-connected networks, which we prove fail to learn to reason. Our results inspire modifications of the transformer architecture that add only two trainable parameters per head, and that we empirically demonstrate improve data efficiency for learning to reason.

Published conference paper at ICLR 2024
WHEN CAN TRANSFORMERS REASON WITH ABSTRACT SYMBOLS?
Enric Boix-Adsera Apple, MIT eboix@mit.ccm edu
Omid Saremi Apple
Emmanuel Abbe
Apple, EPFL
osaremi @aaniiaappeecc.mm
enmanuel wwwp::elll.abbeeee...... cch
Samy Bengio Apple bengio@apple.cco
Etai Liftwin
Joshua Susskind Apple jsusaki meaaakind@apple.com
Apple el elittwingapple..o com
ABSTRACT
We investigate the capabilities of transformer models on relational reasoning tasks. In these tasks, models are trained on a set of strings encoding abstract relations, and are then tested out-of-distribution did not appear in the training dataset. We prove that for any relational reasoning task in a large family of tasks, transformers learn the abstract relations and gener- alize to the test when trained by gradient descent sufficiently large quantities contrast classical tully-connected networks, which
data that contains symbols that
of training data. This we prove fail to learn to reason. Our results inspire modifications of the trans- former architecture that add only two trainable parameters per head, and that we empirically demonstrate improve data efficiency for learning reason.
INTRODUCTION
As large language models (LLMs) are trained with increasing quantitics of data, they begin exhibit the ability reason mathematically (Kaplan et al., 2020; Yuan al., 2023). Why does more data help an LLM learn to reason? And can we make L.LMs more data-efficient at learning to reason?
In this paper, we study relational reasoning with abstract symbols, which is basic capability that has been hypothesized underlie more complex abilities in human cognition (Fodor, 1975; Newell, 1980; Snow et al., 1984; Marcus, 1998; Holyoak, 2012; Kriete et al., 2013; Webb et al., 2020b), One example is in mathematics or computer science, where relational reasoning is necessary to parse proof program: variable names abstract symbols and the functionality the proof or program only depends on how they relate to each other and not on the variable names themselves.
Our contributions are threefold: (i) we formalize relational reasoning through "template tasks"; (ii) we conduct an analysis of when transformers can learn template tasks when trained by gradient descent and show separation with classical fully-connected neural network architectures; (iii) we propose modifications to transformers that improve data efficiency for learning to reason.
1.1 CAPTURING RELATIONAL REASONING WITH TEMPLATE TASKS
Building line of work in neuroscience (Marcus, 1998; Martinho III & Kacelnik, 2016: Kim al., 2018; Webb et al., 2020b; Kerg et al., 2022; Altabaa et al., 2023; Webb et al., 2023a; Geiger et al., 2023), we formalize framework of reasoning tasks called template tasks.
Figure 1:
Tasks from Raven (1938); Webh al. (2020h) which fall under our theory. Networks are truined with one alphahet of symhols and then tested on held-out symbols. Details in Appendix A.
(b)
(c)
o

---

Published conference paper at ICLR 2024
Regression setting In the regression setting, template task is specified by a collection of "tem- plate" strings labeled by real numbers, which are used to generate the train and test data. The simplest way to describe these is through an example. Consider, for instance, the templates
"a-1; 1o:1;:--1;prin (a) - label=+1 and "a-1;;0-1; prinn (,8) label=-1. (1)
These are used to generate the datasets in Figure 2. where every sample (æ;, V4) A x y is formed by picking template and replacing the placeholders a, 3 (which we call "wildcards") with variable names. Memorizing the training data is casy (Zhang et al., 2021a), , but we wish to measure reason- ing: will the model learn treat the variable names abstract symbols, enabling generalization beyond its training distribution? To evaluate this, we adopt an out-of-distribution setting, where the train and test data distributions differ (Marcus, 1998; Abbe et al., 2023). The test dataset consists of the same programs, but with new variable names never seen during training. By resting symbols unseen in the train set, we measure the ability of an LLM to learn logical rules on the relations between symbols. To succeed, the LLM must effectively infer the templates from training data, and test time match samples to the corresponding templates to derive their labels.
(a) Train data
(b) Test data
(c) Transformer performance
2.0
Transformar
a=1;b=-1 :print (a)
+1
fest
feet
Test loss
1.5
Thansformer KD identity
10:aa-11;piing (a)
-1
R-1; b-1;A=-1;prin (R)
+1
2.0
f=1;c=-1; 1:p-int (f)
+1 Q=1;V=-1;print
-1
h-1 1p::iittt
0.0
30
10
10a
1G
Number of training samples
Figure 2: (a.b) Variable names in the test data never appear in the train data (indicated by lowerfupper-case names). (c) Remarkably, as the training set size inereases, the LLM's ability reason outside is training data improves, as learns to use the relations between the variable names to classify, instead of simply memorizing the training data. Our theory motivates a modified transformer architecture (see Observation 1.2), which solves the reasoning task with less training data. Details in Appendix A.
Apart from programming tasks as in Figure 2, this framework captures several natural problems:
Sameldifferent task. The simplest relational reasoning task is when the templates "ax" and "cv39" labeled by +1 and 1. This encodes learning to classify two symbols as equal (e.g., AA, BB) or as distinct (c.g.. AB, BC), even when the symbols were unseen in the training data. This task has been studied empirically in animal behavior (Martinho III & Kacelnik, 2016) and in neural networks (Kim et al., 2018; Webb et al., 2020b).
Word problems. Word problems often have building blocks that follow simple templates. For example, the template "If a gives 3 y. how many y does have?" labeled by +5, could generate the data "If Alice gives Bob 5 oranges, how many oranges does Bob have?" the data "If Rob gives Ada 5 apples, how many apples does Ada have?"
Psychometric tests. Psychometric tests of relational reasoning, which have recently been used to probe LLMs (Raven, 1938; Webb et al., 2020b; Altabaa et al., 2023; Kerg et al., 2022; Webb et al., 2023a:b), are often template tasks. Figure I illustrates some examples.
Next-token-preddee prreiiicinnn setting In the next-token- prreecnpreeittion setting, there is one extra layer of complexity: each sample is labeled with a symbol. For the LLM to generalize to symbols unseen at train time, not only must it learn to track the value stored in a variable, but it also must learn to predict labels at test time that might not occur in its training data. For example, the train and test datasets in Figure 3 are generated by:
"a="""" ;==""""""""""""""""""""""""""""" :rrint:pintt.... (ox)" + label=y and "u=">"";==*"" sprinttppiinide (9)" - label=8, (2)
where a. B, 7,8 are wildcards. Other problems covered by these tasks include:
Programming. The template "print ("a""" labeled with a generates (print ("A"), A) or (print ("dog)) dog), and so an LLM that learns on the corresponding task can robustly evaluating print statements on symbols not seen in the training data.

---

Published conference paper at ICLR 2024
Mathematical functions. For example, the set of templates {aaa, a,9a, aag, 9ax) la- beled by a encode the task of outputting the majority token in a length-3 string with a vocabulary of two symbols. Similarly, for length-k strings, the task of outputting the ma- jority element can be encoded with 2 templates.
(a) Train data
(b) Test data
(c) Transformer performance
2.0
2; a- a-"d"; ipiint (a)
yi
Waasformer
x jest
M/
test
1.5
Tansfomer KQ.V identity
o= c="x" ja="""""=""""""""""""""""""""""" iprint "=""""""""""""p""""""""""" h
A="2" (R)
1.0
y
Q="B";V""A :vrint(V)
0.5
0.0
10-
10
10
10
10"
Num ber of tra training samples
Figure 3: (a.b) The labels symbols. (c) We propose modified that transformer learns the reasoning task with less data (see Observation 1.2 and Theorem 1.4). Details Appendix A.
1.2 MAIN RESULTS
The phenomenon from Figures 2 and that we seek understand is: why does the out-of- distribution performance of the transformer architecture improve as the number of training samples increases? We analyze the regression and next-token-prediction settings separately.
MLPSs fail to generalize to unseen symbols A classical criticism of connectionism by Marcus (1998) is that neural networks do not learn relational reasoning when trained. We support this crit- icism in Appendix by proving that classical MLP architectures (a.k.a. fully-connected networks) trained by SGD or Adam will not generalize in template tasks on symbols unseen during training, even in the regression setting. This failure to reason relationally occurs regardless of the training data size. The proof uses permutation equivariance property of MLP training (Ng, 2004; Shamir, 2018; Li et al., 2020; Abbe et al., 2022; Abbe & Boix-Adsera, 2022).
(2) Transformers generalize to unseen symbols, but require large data diversity Neverthe- less, prove that he criticism of Marcus (1998) is not valid for modern transformer architectures (Vaswani al., 2017). We analyze the training dynamics of a transformer model and establish that it can learn to reason relationally:
Theorem 1.1 (Informal Theorem 3.4), For any regression template task, a widee-nouuhh transformer architecture trained by gradient flow on sufficiently many samples generalizes on unseen symbols.
Here the key points are: (a) Universality. The transformer architecture generalizes on symbols un- seen in train data regardless of which and how many templates are used to define the reasoning task. (b) Large enough number samples. Our theoretical guarantees require the training dataset size be large, and even for very basic tasks like the two-template task Figure 2, good generalization begins to occur only at very large number of training samples considering the simplicity of the task. This raises the question of how the inductive bias of the transformer can be improved.
The proof of Theorem 1.1 inspires parametrization modification that empirically lowers the quan- tity of data needed by an order of magnitude. A standard transformer attention head that takes in an input X E R is given by
smax X K
where WK, K, Wq. WV, Wo are trainable parameters. Our modification makes it easier for the transformer access the incidence matrix x e x of the input, which is invariant to permu- tations of the symbol alphabet and can be used to solve the relational reasoning task:
Observation 1.2. Adding one trainable parameter to each attention head that W W W is replaced by W x T + al improves transformers' data-efficiency on template tasks.

---

Published a conference paper at ICLR 2024
(3) Transformers fail at copying unseen symbols The story slightly different for next-token- prediction tasks, because of the bottleneck of learning to output symbol that was never seen in the training dataset. Transformers' performance degrades as the model grows (an "inverse scaling" law (McKenzie et al., 2023)). Large transformers fail even for the task of copying the input.
Theorem 1.3 (Informal Theorem 4.1). Transformers with large embedding dimension fail to gener- alize on unseen symbols for the copy-tusk ouipaating label "a" " on template "a".
However, we propose adding an iitention-modulaae skip connection, which corrects this failure, making it easy for the transformer learn copy data between its residual streams:
Theorem 1.4 (Informal Theorem 4.2). Adding one trainable parameter b to each head that W is replaced by W + bl makes transformers generalize on the task of Theorem 1.3.
(4) Experiments We conclude with experimental validation of our architecture modifications, and find that they improve data efficiency on relational reasoning tasks by an order of magnitude, and improve sanguage-modeling performance when training the GPT-2 architecture Wikitext.
1.3 RELATED LITERATURE
spate recent work studies whether and how LLMM perform various reasoning tasks, each focus- ing on one component of reasoning: these include recognizing context-free grammars (Zhao et al., 2023; Allen-Zhu & Li, 2023), learning sparse functions (Edelman et al., 2022), learning compo- sitionally (Hupkes et al., 2020), generalizing out-of-distribution when learning Boolean functions (Abbe et al., 2023), performing arithmetic (Nanda et al., 2023), learning in context (Garg et al., 2022; Ahn et al., 2023; Zhang al., 2023), and evaluating indexing (Zhang et al., 2021b). Our setting closest to that of empirical work studying neural networks on relational reasoning tasks (Geiger et al., 2023; Webb ct al., 2023b). For example, the four tasks in Webb et al. (2020b), the matrix digits task in Webb et al. (2023a), the SET game task in Altabaa et al. (2023), and most of the tasks in Kerg et al. (2022) (with the exception of the relational games tasks), are examples of regres- sion template tasks that fall under our theory. Furthermore, Kim et al. (2018) shows experimentally that MLPs fail the same/different template task, and we provide proof for this in Appendix I. There is also a literature on modifying training to improve relational reasoning: (Webb et al., 2020a) proposes applying Temporal Context Normalization during training, and Santoro et al. (2017; 8); Palm et al. (2018); Shanahan et al. (2020); Webb et al. (2020b); Kerg al. (2022); Altabaa et al. (2023) propose new architectures. Finally, some recent works in mechanistic interpretability look for subnetworks within trained networks that are responsible for tasks such as variable binding (Ols- son et al., 2022; Davies et al., 2023). In contrast, our focus is on proving when the transformer architecture learns or fails to learn. and on applying this theoretical understanding improve its data efficiency for relational reasoning.
2 FORMAL DEFINITION OF TEMPLATE TASKS
We formally define regression template tasks. For next-token prediction, see Appendix J.
Definition 2.1. A template is a string (X' U W/),, where X' is alphabet of tokens, and W an alphabet "wildcards". A substitution map is an injective function & : - X. We write sub(2, x) E x for the string where each wildcard is substituted with the corresponding token: sub(z, s); 2; if 2; X, and sub(z, s); s(x) if 24 W. The string æ E X matches the template x sub(2,s) for some substitution map and also s(W) n (=hew = 0: i.e. the substituted tokens did not already appear in the template z.
Example Using Greek letters to denote the wildcards and Latin letters denote regular tokens, the template "0a:357" matches the string "QQRST", but not "QQQST" (because the substitution map is not injective) and not "QQSST" (because 3 is replaced by S which is already in the template).
A template task's training data distribution is generated by picking template randomly from distribution, and substituting its wildcards with a random substitution map.
Definition 2.2. A template data distribution D D(JAempt:: {Muub 1z: f., is given by

---

Published conference paper at ICLR 2024
a template distribution Pemplt supported on templates in (X' W)*
for each supp(/Amplee) distribution Msub,z over substitution maps W + X,
template labelling function f- supp((Mmppt) + R, and label-noise parameter a 2 0.
We draw a sample (x, y) (sub(z, s), (z) + E) D, by drawing a template Hemplt a substitution map x Maubi and label noise (2)
Finally, we define what it means for a model to solve the template task and generalize on unseen symbols; namely, the model should output the the correct label for any string 2: E 2 matching template, regardless of whether the string is in the support
training distribution.
Definition 2.3. A (random) estimator f 2 + IR generalizes unseen symbols with (e, -eeror if the following is true. For any æ E 2² that matches template = supp(Hemppt), we have
(f(x) f+(=))² <,,
with probability at least 1 - 8 over the randomness of the estimator f.
Example If the training data is generated from uniform distribution on templates "aa" with label and "a,"" for label -1, then it might consist of the data samples {(AA, (BB, 1), (1), (BA, -1)]. An estimator that generalizes to unseen symbols must correctly label string CC with +1 and string CD with -I, even though these strings consist of symbols that do not appear in the training set. This is nontrivial reasoning task since it requires learning to use the relations between the symbols to classify rather than the identities of the symbols.
3 ANALYSIS FOR TEMPLATE TASKS IN THE REGRESSION SETTING
We establish that lay-r transformers of large enough width generalize unseen symbols, when trained with enough data on regression template tasks, It important to note that this not true for all architectures, as we prove in Appendix that MLPs trained by SGD or Adam will not succeed.
3.1 TRANSFORMER RANDOM FEATURES KERNEL
The one-layer transformer architecture that we analyze consists of an embedding layer, a multihead attention mechanism, an MI. layer, and an unembedding layer wy. This is written mathematically in Appendix H. We analyze training only the final wy layer of the transformer, keeping the other weights fixed at their random Gaussian initialization. Surprisingly. even though we only train the final layer of the transformer, this is enough to guarantee generalization on unseen symbols. Taking the width and embedding and head dimensions to infinity, and the step size 0, the SGD training algorithm with weight decay converges to kernel gradient flow with the following kernel K in the infinitely-wide, infinitely-smin saattell-small-ttee step-sste--izze limit. Here and throughout the remainder of the paper, interchangeably denote input by a string E X matrix X E Bexm constructed by stacking the one-hot vectors X = [ex]..... ..., e of the string's tokens. : R R. is the MLP activation layer, 3,7 E R are hyperparameters controlling the temperature and magnitude of positional activations.
Kattn (X,X) Katto (X, )
Krran X, Y)= Eu.v(d(u))(u)) for u N(0,
(4)
K (Y.)) X) Koun (oonnYY)) (Y,Y)
where Kattn (X,Y) = Em(X)). (smax//mmmmmmmmmmmmmmmm (XY" + 222 I smax((3((())
Xx T + + YxT + YY +
XYT + +1
ml X). (Y))
(O
T
The function outputted by kernel gradient flow is known to have a closed-font solution in terms of the samples, the kernel, and the weight-decay parameter X, which we recall in Proposition 3.1.
Proposition 3.1 (How kernel gradient flow generalizes; see e.g., (Welling, 2013).). Let (X,, (X1.1))..... vI....... (X.. Mn) be training samples. With the square loss and ridge-regularicc of mag- nitude X. kernel gradient flow with kernel K converges to the following solution
j(x) = T(K AI) k(X),
(5)

---

Published a conference paper at ICLR 2024
where y = lyI...., Mrl IR" are the train labels, K Ruxx is the empirical kernel matrix and has entries Kig- K( K(X;:
s: x3), and xxx) E Rn has entries kyy ,,
3.2 TRANSFORMERS GENERALIZE ON UNSEEN SYMBOLS
We prove that trasformers will generalize out-of- distriiitiitttonn unseen symbols when trained on template tasks, We require the templates in the distribution Pemplt to be "disjoint", since otherwise the correct label for string æ is not uniquely defined, as a could match more than one template:
Definition 3.2. Two templates 2, =' E (XUW)& are disjoint if no E X matches both and z'.
Furthermore, in order to ensure that the samples are not all copies of each other (which would not help generalization), we have to impose a diversity condition on the data.
Definition 3.3. The data diversity measured by p = ming csupp()aon) m Purttyubl [E](W]]
When the data diversity P is large, then no token is much more likely than others to be substituted. If p is on the order of the number of samples n, then most pairs of data samples will not be equal.
Theorem 3.4 (Transformers generalize on unseen symbols), Let Pemplt be supported on a finite set of disjoit templates ending with [CLS] tokens. Then, for almost any 3, ", by. b2 parameters (except for
set), the transformer random features with ((t) cos cos(bbtt- +22 generalizes on unseen symbols.' Formally, there are constants C > O and ridge regularization parameter X > 0 that depend only 3,7,, b1, b2, Pampht Jw. such that for any matching template E supp((emmtt) the kernel ridge regression estimator f in (5) with kernel Krrans satisfies
1f(z) - (-(=)) S C Co/log(1/8)/ + CVL/p, with probability at least 1 - 5- exx(-on) over the random samples.
The first term is due the possible noise in the labels. The second term quantifies the amount of sample diversity in the data. Both the sample diversity and the number of samples must tend infinity for an arbitrarily small error guarantee.
Proof sketch (1) In Lemma 3.5 we establish with sufficient condition for kernel ridge regression generalize on unseen symbols. (2) We prove that Krrans satisfies it.
(1) Sufficient condition. Let Pemple be supported on templates 21. Zy. Let R. UEe 2 be the tokens that appear in the templates. Let [n]] = Ii Iz .J...IZA be the partition of the samples such that a I, then sample (æ., Me) is drawn by substituting the wildcards of template zj. Two samples Fa 2y that are drawn from the same template zj may be far apart as measured by the kernel: i.e., the kernel inner product K K(x,,, 2.)) may be small. However, these samples will have similar relationship to most other samples:
K(x, x.) = K K( 2.) for most i E [n].
(6)
Specifically, if the wildcards of C.. Hy and æ; substituted by disjoint sets of tokens that do not appear the templates, then (6) holds. Therefore, as the sample diversity p increases, the empirical kernel matrix K hecomes approximately block-structured with blocks I, Iy.. For most samples Va, Ty corresponding to template zj, and most ae', By corresponding to template zy' we have
K ((x. ) K ((ss, my K (sub(zz,,)) sub((zz Nj5',
(7)
where s, s': W + X are substitution maps satisfying
x xWW = o and
x((W) n 0.
(8)
One can check that (7) and (8) uniquely define matrix N E Rexr which gives the entries in the blocks of K, with one block for each pair of templates.? 2 See Figure 4.
We analyze the shifted and rescaled cosine activation function @(t) = cos(b, + b2) out of technical convenience, but conjecture that most non-polynomial activation functions should succeed. "This assumes ttoken-symmetr" property K that is satisficd by transformers; details in the full proof.

---

Published a conference paper at ICLR 2024
I
L I1 1
K(AA, BB) K(AA, BC) N
K-
R2x2
AAB ((AB,CD)
L2
Figure 4: Illustration of structure of K and for the same/different task, which has templates 21 = an and 22 = 0,9. As the sample diversity p increases and the number of samples n increases, the empirical kernel matrix K E R=x" becomes approximately (r x r)-bblock-structuredd and within each block most of the entries are given by N E R= exceptions where this is true, including the diagonals, are drawn in black. Furthermore, the spectrum K is increasingly determined by the spectrum of N, and if N is nonsingular then the top eigenspace increasingly aligns with the span of the indicator vectors on IT. Ir.
If the matrix N is nonsingular and the number of samples is large, then the span of the top eigen- vectors of K will align with the span of the indicator vectors on the sets Iy.... II.. Furthermore. when testing a string x fest that matches template 2,. but might not have appeared in the training set, it holds that for most a E Iy. we have
k( f = [K (atest Fest ...)..... K (xt xx)] Ro [K (1),)))))........ ... K (x (x)) Ko.--
words, the similarity relationship of zest to the training samples is approximately the same as the similarity relationship of x. to the training samples. So the kernel ridge regression solution (5) approximately equals the average of the labels of the samples corresponding to template 2,, which turn is approximately equal to the template label by a Chernoff bound,
y yJ(K + -k(x)
v R= fz())). AEI;
Σ
IT
Therefore, kernel ridge regression generalizes on x fest, It is important to note that the number of samples needed until (9) is good approximation depends on the nonsingularity of N. This yields the sufficient condition for kernel ridge regression to succeed (proof in Appendix C).
Lemma 3.5 (Informal Lemma C.3). If N is nonsingular, then (5) generalizes to unseen symbols.
(2) Krrans satisfies the sufficient condition. We now show that for any collection of disjoint templates $1,..., ky. the matrix N trans N E Texr defined with kernel K K is nonsingular. The challenging is that Ktrans does not have a closed-form solution because of the expectation over soft- max terms in its definition (4). Therefore, our amalysis of the transformer random feature kernel is, to the best of our knowledge, the first theoretical analysis showing that the transformer random features learn a nontrival class of functions of sequences. We proceed by analyzing the MLP layer and the attention layer separately, observing that a a"weak" condition on Katt can he lifted into the "strong" result that N trans nonsingular. The intuition is that as long Kattn is not very degenerate kernel, is unlikely that the MLP layer has the cancellations that to make N trans nonsingular.
Lemma 3.6 (Nonsingularity of N Wraans). Suppose for every non-identity permutation T E S, [id],
Katin (sub(zs s), ())()),6))) # Σ Katn(sub(e sub(())) 5'))) Re[r]
Σ
(10)
where 8, 8' are the substitution maps in the definition of N Nrrass in (8). Let the MLP layer's 's activation function be o(t) = cos(byt + b2)- Then for almost any choice of b1, b2 (except for Lebesgue- measure-zero set), the matrix N Ntrans is nonsingular.
This is proved in Appendix E, by evaluating a Gaussian integral and showing Nu has Vander- monde structure. Although we use the cosine activation function, we conjecture that this result holds for most non-polynomial activation functions. Next, we prove the condition on N attn-
Lemma 3.7 (Non- degenergyyy of Katn). The condition (10) holds for Lebesgue-almen almost any 3, %
The proof is in Appendix F. First, we prove the analyticity of the kernel Kattn in terms of the hyperparameters 3 and 7. Because of the identity theorem for analytic functions, it suffices show at least one choice of hyperparameters 3 and 7 satisfics (10) for all non-identity permutations т. Since Kann does not have a closed-form solution, we find such choice of 3 and by analyzing the Taylor-series expansion of Katn around 3 0 and 7
up to order- 10 derivatives.

---

Published a conference paper at ICLR 2024
3.3 IMPROVING TRANSFORMER DATA EFFICIENCY WITH WX W + al PARAMETRIZATION
Can we use these insights to improve transformers' data-efficiency in template tasks? In the proof. the nonsingularity of N in Lemma 3.5 drives the model's generalization unseen symbols. This suggests that an approach to improve data-efficiency is make N hetter-conditioned by modifying the transformer parametrization. We consider here the simplest task, with templates "aa" and "a,3" labeled with +1 and 11, respectively. For tokens A, B, C E X, the matrix N is
N= (((AA K(AA, BC) K(BC, AA) ((A
K is an inner-product kernel, x(xxx') = K(Eigg (z: = z()), as from an MLP, then ((AA, BB) = (AAA BC) K K(BC, AA) = ((AB,CD) = x(0), so is singular and generalization is not achieved. Intuitively, every sample x, has approximately the same "similarity profile to other data" K;;- = [x;;, z1)..... K( xr,)), ),, so the kernel method cannot identify the samples that come from the same template as x. In contrast, the transformer kernel (4) succeeds by using information about the incidence matrix Xx"" which differs between templates, and does not depend on the symbol substitution. We thus propose to emphasize the incidence matrix X Xx" by reparametrizing each head to W x W Q + ad,, where is trainable parameter. This adds scal- ing of x" in the attention, and can empirically improve data efficiency by an order of magnitude on several template tasks (sec Figures 2 and 3, as well as additional experiments in Appendix B).
4 ANALYSIS FOR TEMPLATE TASKS IN SEXT-TOXEE-PREDIICIOON SETTING
We switch gears to the next-token prediction setting with the cross-entropy loss, where the output label may be token in the example of Figure 3; formal definition is in Appendix J. The simplest task consists of template "a" labeled by "a". An example train set is {(A, A), (B, B), ()) where A, B, C' E X' are tokens, and then we test with [ ff = (D. D) which is not in the train set. This task captures the ability of a model to learn how to copy symbol, which is important for LLMs that solve problems with multi-stage intermediate computations and must copy these later parts of a solution (Csordás et al., 2021). From now on, we only consider this "copying" task.
We consider an architecture fatm 0) with just a multi-head attention layer, and we tic the em- bedding and unembedding weights as in practice (Brown al., 2020). Define the train loss and test loss as follows, where f is the cross-entropy loss and x fest is a token unseen in the training data: Levain (0) = E ([fatt([z; 0), U1) and Liest = é {f (xfestt, y, We prove this network does not generalize
unseen symbols when trained, we take the embedding dimension large. Our evidence is from analyzing the early time of training, and showing that the test loss on unseen symbols does not decrease.
Theorem 4.1 (Failure of transformers at copying). For any learning rates such that DCenstn O(1), we musl have that @Cr.tt
1==0
0t
1==0 0 as demb >> co.
dt
The proof idea is that since the input string has length A: = 1, the architecture simplifies: all soft- maxes in the attention heads output 1, and the network is a sum of attention heads of the form W E- At early times the evolution of the weights Wv W. O will roughly lie in the W&her Ehejj), which as the embedding dimension becomes large will be approxi-
XW E W v W O
span of (WT. E e. e,, W mately orthogonal to the direction W e W y that would lower the test loss. This suggests the following modification transformers allows them to copy symbols never seen at training:
Theorem 4.2 (Adding one parameter allows copying). After reparametrizing the attention (3) so that in each head W v W is replaced by w. + where b a trainable parameter, there are learning rates such that Bfaraan I==0= O(1) and &Lyust 1==0= (2(1) as demb -> co.
Figures 3 and illustrate the benefit of this additional per-head parameter on the copying task, It is not equivalent to adding trainable skip connection as in ResNet (He et al., 2016). Instead, the addition of by I encodes the transformer's ' streams. A related modification of adding a head with the hardcoded XX" as its attention matrix was proposed in Zhang al. (2022).
intention-mddllatee skip connection that allows copying tokens between

---

Published conference paper at ICLR 2024
(a) Vanilla transforee
(b) Transformer with Wy W bT
1.5
dim=64
dim-128
Test erier
1.0
dim-256
194t ener
1.0
dim-1024
dim=2041
dim=0096
0.0
10
10.
10"
0.0
10
10
10a
Number training samples
Number of training samples
Figure 5: (a) Transformers fail on the copying task as embedding dimension demh grows (Theorem 4.1);
(b) Success when reparametrizing Wy W. as w + (Theorem 4.2). Details in Appendix A.
EXPERIMENTS
Figures 2 and 3 (and additional experiments in Appendix B) show that our reparametrizations can give significant data-efficiency benefit template tasks. Figure shows they also give im- provements on real data. In Figure 7, we see that pretraining outperforms random initialization on a template task. This might be explained by several heads of the pretrained model with diagonals stronger from other weights (originally observed in (Trockman & Kolter, 2023)). These learned di- agonals resemble our proposed transformer modifications and so might be driving data-efficiency of fine-tuning pretrained model. Appendix B provides extensive experiments on the effect of hy- perparameters, inductive biases of different models, and varying levels of task difficulty.
Dataset
GPT-2
GPT-2 + trainable identity scalings (ours)
Wikitex12
64.00)
60.46
W1
16.83
16.40
Figure 6: Perplexity of GPT-2 trained from random initialization with Adam learning rate 3e-4 for 20 epochs Wikitext (smaller perplexity hetter). GPT-2 has 117M parameters, and add 288 parameters per head). Interestingly, even though the task is Wikipedia modeling, and therefore not a pure reasoning task, the transformer modifications still give improvement.
Effect of pretraining
Wx W Head 12, Layer 5
Wy W Head 12, Layer 11
GPT-2 from strath
Test loss
1.0
0.5
10
10'
'02
10a
Number training samples
50
72
15
Figure Left: Pretrained versus randomly-inttialized GPT-2 loss when fine-tuned aBa ag3 template task. Right: some GPT-2 pretrained heads have strong diagonals (zoomed to 100x top-left corner).
DISCUSSION
We show that transformers are universal architecture for template tasks in the regression setting: when trained with gradient descent with enough training data they learn reason relationally. How- ever, transformers are not optimal - empirically they require large amounts of data to learn hasic tasks, and in the next-token- prediction setting they fail at copying unseen symbols. Thus. we have proposed architectural modifications improve their inductive bias towards logical reasoning. It seems promising to explore other reasoning tasks (for example, reasoning with syllogisms, reason- ing by symmetry. and compositional reasoning). It may also be fruitful to study data augmentation approaches (c.g., concatenating the tensorization X x to the imput, so as to encourage use of rela- tional information). Additionally, tight quantitative upper and lower bounds on the data and width of the architecture needed, depending on the template task,
interesting open direction.

---

Published conference paper at ICLR 2024
REFERENCES
Emmanuel Abhe and Enric Boix-Adsert On the non-universality of deep learning: quantifying the cost of symmetry. Advances in Neural Information Processing Systems, 35:17188-177 2022.
Emmanuel Abbe, Elisabetta Cornacchia, Jan Hazla, and Christopher Marquis. An initial alignment between neural network and target is needed for gradient descent lear.. In International Con- ference on Machine Learning, pp. 33-52. PMLR, 2022.
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, and Kevin Rizk. Generalization on the unseen, logic reasoning and degree curriculum. arXiv preprint arXiv:2301.13105. 2023,
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306.0020 2023.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint 230. 2023.
Awni Altabaa, Taylor Webb, Jonathan Cohen, and John Lafferty. Abstractors: Transformer modules for symbolic message passing and relational reasoning. arXiv preprint arXiv:2304.00195, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot leamners Advances in neural information processing systems, 33:1877-1901, 2020.
Lenaic Chizat and Francis Bach. On the global convergence gradient descent for over- parameterized models using optimal transport. Advances in neural information processing sys- tems, 31, 2018.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The neural data router: Adaptive control flow transformers improves systematic generalization. arXiv preprint 1rXiv:2110.07732, 2021.
George Cybenko. Approximation by superpositions sigmoidal function, Mathematics of control, signals and systems, 2(4):303-314 1989.
Xander Davies, Max Nadeau, Nikhil Prakash, Tamar Rott Shaham, and David Bau. Discovering variable binding circuitry with desiderata. arXiv preprint urXiv:2307.03633 2023.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference Machine Learning, pp. 5793-5831. PMLR, 2022.
Jerry A Fodor. The language of thought, volume 5. Harvard university press, 1975.
Shivam Garg, Dimitris Tsipras, Perey S Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
Atticus Geiger, Alexandra Carstensen, Michael C Frank, and Christopher Potts. Relational reasoning and generalization using nonsymbolic neural networks. Psychological Review, 130(2):308, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Keith Holyoak. Analogy and relational reasoning. reasoingg The Oxford handbook thinking and reasoning. pp. 234-259, 2012.
Dieuwke Hupkes, Vern Dankers, Mathijs Mul, and Elia Bruni Compositionalliyy decomposed: How do neural networks generalise? Journal of Artificial Intelligence Research, 67:757-795. 2020.
10

---

Published conference paper at ICLR 2024
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08.30 2020.
Giancarlo Kerg, Sarthak Mittal, David Rolnick, Yoshua Bengio, Blake Richards, and Guillaume La- joie. On neural architecture inductive biases for relational tasks. arXiv preprint arXiv:2206.05000 2022.
Junkyung Kim, Matthew Ricci, and Thomas Serre. Not-so-clevr: learning same-different relations strains feedforward neural networks. Interface focus, 8(4))20180011, 2018.
Steven G Krantz and Harold R Parks. A primer of real analytic functions. Springer Science & Business Media, 2002.
Trenton Kriete, David C Noelle, Jonathan D Cohen, and Randall C O'Reilly. Indirection and symbol-like processing in the prefrontal cortex and basal ganglia. Proceedings of the National Academy of Sciences, 110(41):6390-1699) 2013.
Zhiyuan Li, Yi Zhang, and Sanjeev Arora. Why convolutional nets more sample-efficient than fully-connected nets? arXiv preprint arXiv:2010.08511 085155 2020.
Gary F Marcus. Rethinking eliminative connectionism. Cognitive psychology, 37(3):243-2822 1998.
Antone Martinho III and Alex Kacelnik. Ducklings imprint on the relational concept of "same or different". Science, (5666)))) 288, 2016.
lan R McKenzie, Alexander Lyzhov, Michael Pieler, Alicia Parrish, Aaron Mueller, Ameya Prabhu, Euan McLean, Aaron Kirtland, Alexis Ross, Alisa Liu, al. Inverse scaling: When bigger isn' better. arXiv preprint vrXiv:2306.09477 2023.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115(33) 2018.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pp. 2388- 2464. PMLR, 2019.
Boris Samuilovich Mityagin. The zero set of a real analytic function. Mathematical Notes, 107 (3-4):529-530, 2020.
Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXt:: 2023.
Allen Newell. Physical symbol systems. Cognitive science, 4(2): 135-1135, 1980.
Andrew Y Ng. Feature selection, 11 12 regularization, and rotational invariance. In Proceedings of the twenty-first international conference on Machine learning, pp. 78, 2004.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 1895, 2022.
Rasmus Palm, Uirich Paquet, and Ole Winther. Recurrent relational networks. Advances in neural information processing systems, 31, 2018.
John C Raven. Progressive matrices: A perceptual test of intelligence. London: HK Lewis, 19:20. 1938.
Grant Rotskoff and Eric Vanden- Vanden-Eijnden, Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. Advances in neural information processing systems, 31, 2018.

---

Published conference paper at ICLR 2024
Adam Santoro, David Raposo, David Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning. Ad- vances in neural information processing systems, 30, 2017.
Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski. Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy Lillicrap. Relational recurrent neural networks. Advances in neural information processing systems, 31, 2018.
Ohad Shamir. Distribution-spene speciffc hardness of learning neural networks. The Journal of Machine Learning Research, 1135:::113116:: 2018.
Murray Shanahan, Kyriacos Nikiforou, Antomia Creswell, Christos Kaplanis, David Barrett, and Marta Garnelo. An explicitly relational neural network architecture. In International Conference on Machine Learning, pp. 8593-8603. PMLR, 2020.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of deep neural networks. Math- ematics of Operations Research, 120-12 2022.
Richard Snow, Patrick C Kyllonen, and Brachia Marshalek. The topography of ability and learning correlations. Advances in the psychology of human intelligence, pp. 47-103, 1984.
Asher Trockman and J Zico Kolter. Mimetic initialization of self-attention layers. arXiv preprint arXiv:2305.0988 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Taylor Webb, Zachary Dulherg, Steven Frankland, Alexander Petrov, Randall O'Reilly, and Jonathan Cohen. Learning representations that support extrapolation. In International confer- ence on machine learning. pp. 10136-10146, PMLR, 2020a.
Taylor Webb, Keith Holyoak, and Hongjing Lu. Emergent analogical reasoning in large language models. Nature Human Behaviour, pp. 1-16, 2023a.
Taylor Webb, Ishan Sinha, and Jonathan D Cohen. Emergent symbols through binding in external memory, arXiv preprint arXiv:2012. 14601, 2020b.
Taylor W Webb, Steven M Frankland, Awni Altabaa, Kamesh Krishnamurthy, Declan Campbell, Jacob Russin, Randall O' O'Reilly, John Lafferty, and Jonathan D Cohen. The relational bottleneck as an inductive bias for efficient abstraction. arXiv preprint arXiv:2309.0662 2023b.
Max Welling. Kernel ridge regression. Max Welling's classnotes in machine learning, pp. 1-3, 2013.
Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. International Conference on Machine Learning, pp. 11727-11737. PMLR, 2021.
Zheng Yuan, Hongyi Yuan, Chengpeng Li, Guanting Dong, Chuanqi Tan, and Chang Zhou. Scal- ing relationship on learning mathemattcal reasoning with large language models. arXiv preprint arXiv:2308.00 01825, 2023.
Chiyuan Zhang, Samy Bengio. Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021a.
Chiyuan Zhang, Maithra Raghu, Jon Kleinberg, and Samy Bengio. Pointer value retrieval: A new benchmark for understanding the limits of neural network generalization. arXiv preprint arXiv:2107.125880 2021b.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint 2rXhv:2306.09927, 2023.
12

---

Published conference paper at ICLR 2024
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: synthetic reasoning task. arXiv preprint arXiv:220.044300 2022.
Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while pre- dicting the masked word? arXiv preprint arXiv:2303.08/17 2023.
13

---

Published conference paper at ICLR 2024
CONTENTS
1 Introduction
1.1 Capturing relational reasoning with template tasks
1.2 Main results
3
1.3 Related literature
2 Formal definition of template tasks
4
3 Analysis for template tasks in the regression sefting
5
3.1 Transformer random features kernel
5
3.2 Transformers generalize on unseen symbols
6
3.3 Improving transformer data-efficiency with Wx
W WxW3
al parametrization
8
4 Analysis for template tasks in next-token-prediction setting
8
5 Experiments
9
6 Discussion
9
A Details for figures in main text
15
Additional experiments
B. Effect of transformer hyperparameters
16
B.2 Effect of complexity of task
17
B.3 Effect of inductive bias of model
17
Proof of Theorem 3.4
25
C.1 Part 1. General sufficient condition for good test loss
25
C.2 Part 2. Analyzing the transformer random features kernel
26
C.3 Concluding the proof of Theorem 3.4
26
Sufficient condition for kernel method to generalize on unseen symbols (Proof of
Lemma C.3)
26
D.1 Deferred proofs of claims
29
D.2 Remark: explicit dependence on |N-1]
30
E Nonsingularity of random features after MLP layer (Proof of Lemma 3.6)
31
Analysis of attention layer features (Proof of Lemma 3.7)
32
E11 Low-order derivatives of attention kernel
33
F.2 Simplifying terms
34
F.2.1
Assuming 1 1xx==117 -1=T
34

---

Published conference paper at ICLR 2024
F.2.2
Assuming [[]|| = M|
34
F.2.3
Assuming 1"XX"1 1'yY"1
35
F.2.4
Assuming 1" XXX" 1"YY""
................................
35
E3 Proof Lemma F.I
35
G Analyticity of attention kernel (technical result)
37
G.1 Technical lemmas for quantifying power series convergence
38
G.2 Application of technical lemmas attention kernel
39
Derivation of transformer kernel
40
H.I Transformer architecture
40
H.2 Random features kernel
................................
41
H.3 Informal derivation
41
1 MLPs fail to generalize on unseen symbols
43
J Deferred details for next-token- next-ioken-rreiicioon template tasks
45
Definition of next-token predicctoo template tasks
45
Failure transformers to copy modification that succeeds
45
A DETAILS FOR FIGURES IN MAIN TEXT
Code available at http:/tw..cmmmeoooee/eeeataaaaaa
Psychometric tasks
We describe how the tasks in Figure fall under the template framework.
(a) Distribution of The task to complete the bottom row so that the set of elements is the same as in the top row (answer: 2). To input this task into language model, token is used represent each symbol. The example in the figure matches template "cv3- yal ea;9y" with label +2. There are other templates for this task, corresponding to different arrange- ments of the objects, such as "a:8 By0 aryc with label +1, and "a,3 780 <Baay" with label +3. In total there are 144 templates, since the first elements of the template are always a,8h, and then there are choices for the permutation in the next row, and finally 24 choices for the permutation in the final row.
(b) Relational match-to-saaee to-sample. The task is to match the first row one of two alternative patterns (answer: 1). Again, a token is used represent each symbol. The example in the figure matches "a333 768 eer" with label +1 simple combinatorial calculation gives a total of 40 templates (5 possible patterns in the first row, times 2 choices for whether the first option the second option is correct, times 4 choices for the pattern of alternative option).
(c) Raven's progressive matrices. A standard Raven's progressive matrices task (Raven, 1938) (answer: three dark circles). For each of the dimensions of shape, number, and color, we have "distribution of 3" task with symbolic lahel. For example, for the shapes in the figure, the task is "a,y S-ya 7/8?" with label a. Since another possibility is for each row to be constant (as in, e.g., the case of numbers), another possible template is "aaa 933 yy?? with label y, and there total of 36+1 37 possible templates per dimension. This discussion assumes that the only patterns in the progressive matrices are distribution 3, and constant. If progressions are also allowed as in Webb et al. (2023a), these can be incorporated by adding corresponding templates.
15

---

Published conference paper at ICLR 2024
Transformer performance In all experiments, standard transformer architectures are used. Figure 2, The architecture is a 2-layer transformer with 16 heads per layer, embedding dimension 128, head dimension 64, MLP dimension 256, trained with Adam with learning rate le-3 and batch- size 1024. The training samples are chosen by picking the variable names random from an alphabet of n tokens. The test set the same two programs but with disjoint variable names. The reported error bars are average over 5 trials. The learning rate for each curve is picked the one achieving best generalization in {10 10-4 except that the transformer is 4-layer transformer and has embedding dimension 512. In Figure the same hyperparameters as in Figure are used. In order to measure the generalization performance of the learned model on unseen symbols, we evaluate on test set and a validation set which each consist of 100 samples drawn in the same way as the training dataset, but each using a disjoint alphabet of size 100. Therefore, there is no overlap in the support of the train, test, and validation distributions. We use the validation loss select the best epoch of training out of 1000 epochs. We report the test loss on this saved model.
-,0440000 10 10 2). In Figure 3, the setting is the same
ADDITIONAL EXPERIMENTS
We report extensive additional experiments probing the template task framework. In each of these, the training dataset consists of random training samples. Each sample drawn according template distribution. The following are template tasks on which we test.
aB vs. a:33 task. Uniform on two templates aga and c,3, with labels 1, -1 respectively and and / are wildcards.
a,333 vs. 00.3 task. Same above, except with templates 0,3a, and ev,333
Length-k: majority task. Uniform on 2*-1 templates a x {a, -) where a and /3 are wildcards. A template has label if its first token occurs in the majority of the rest of the
1, Hi Mi::: = h > (k- + 1)/2
string, and lI otherwise. Namely, f=(z) (-1, otherwise
Random template task. A certain number of templates are drawn uniformly from (WU x²) conditioned being pairwise distinct. The task the uniform distribution over these r templates, with random Gaussian labels centered and scaled so that the trivial MSE is 1.
For any of these tasks, generate training samples follows. We substitute the wildcards for regular tokens using randomly chosen injective function s W X where X' is an alphabet of size n (which is the same size as the number of samples). For example, if given sample is generated from template arßer with substitution map mapping s(A) = 12, a(B) = 5, then the sample will be [12,5, 12). Error bars are over 5 trials, unless otherwise noted.
B.1 EFFECT OF TRANSFORMER HYPEERPARAMETERS
We test a standard transformer architecture on the a:a vs. aB3 task, varying some of the hyperpa- rameters of the transformer to isolate their effect while keeping all other hyperparameters fixed. The base hyperparameters are depth 2, embedding dimension 128, head dimension 64, number of heads per layer 16, trained with Adam with minibatch size 1024 for 1000 epochs. Our experiments are as follows:
Learning and In Figure 8 vary the learning and n.
Learning rate and depth. In Figure and Figure 10, we vary the learning rate and the depth, for 512 and 1024, respectively.
Learning rate and number of heads. In Figure 11 and 12, we vary the learning rate and number of heads, for n 512 and 1024. respectively.
Learning rate and embedding dimension. In Figure 13 we vary the learning rate and em- bedding dimension for 1024.
Learning rate and batch size. In Figure 14, we vary the learning rate and batch-size for n 512. In Figure 16 we vary the batch-size and n for learning rate 0.001.
16

---

Published conference paper at ICLR 2024
Training just the last layer. In Figure 13, we train just the last layer, and see that the network does learn to generalize out of distribution, as predicted by our theory. However, the number of samples and number of epochs needed is larger than when all parameters are trained. We train for 10000 epochs and have 64 heads per layer in this experiment.
B.2 EFFECT OF COMPLEXITY OF TASK
We test out-of-the thx-box transformer architecture with depth 2, embedding dimension 128, head dimension 64, number of heads 16, trained with Adam with batch-size 1024 for 1000 epochs, on various template tasks.
Comparing difficulty of various tasks. Figure 17 plot the performance various simple tasks.
Random tasks. In Figures 18, 19, 20, and 21, test on random template tasks, and in- vestigate effects of template length, wildcard alphabet size, regular token alphabet size, number of templates.
B.3 EFFECT INDUCTIVE BIAS OF MODEL
We provide experiments probing the effect of the inductive bias of the model:
Different architectures. In Figure 22, we plot the test loss for different architectures the 0,30 a:3, template task, including transformers with trainable identity perturbations
to Wo to Wy W.. to both WoWk and Wv W or to neither. Figure 23 illustrates on the beneficial effect of the transformer modification for the majority task with different lengths, lowering the amount of data needed by an order of magnitude.
Size of model. In Figure 24 we compare the test loss of fine-tuning small, medium and large pretrained GPT-2 networks the aBa a,3 template task.
MLP with XX" data augmentation vs. transformer. In Figure 25, we compare the test loss of transformer with the test loss of an MLP where the input data has been augmented by concatenating vec(XX""), which a data augmentation that improves performance under the NTK criterion similarly to the discussion in Section 3.3 and the discussion section.
Transformer test loss vs, learning rate and number of samples
Transformer train loss vs. leamning rate and number of samples
1.75
32
n.32
64
64
Oand train loos 1000 epoctts
1.25
6 1024
1024
Text loex
2040
2048
1.00
1.00
0.75
0.75
0.50
0.50
6.25
0.5
0.00
0.00
10
10
10
10
10
10
10
10
learning rabe
learning rate
Figure 8: Learning rate versus n number of samples training alphabet size. Taking too large too small a learning rate can hurt generalization even when the train loss is close to zero.
17

---

Published conference paper at ICLR 2024
Transformerr test loss leaining rate and depth, 512
Transformer train loss leaming rate and depth, at 512
2.00
2.00
depth
depth 1
1.75
depih
1.75
depth 2
depth
depth 4
1.50
depth
depth
apache
1.50
depth
depth
1.25
1.25
Text loxk
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
10
10
10
10
10"
10
Tearning rate
leamnng rañe
Figure 9: Learning rate vs. depth at n = 512. No clear relationship between depth and generaliza- tion. Too large or too small of a learning rate can hurt generalization.
Transforeer test loss vs. learning rate and depth, at 1024
Transformer train loss vs. leamnnng rate and depth, at n 1024
depth
depth
depth
depth
depth
16
Bestrran loss 2000 spo0
16
Iox
0.50
0.50
0.25
0.25
0.00
0.90
learning rate
kaminn rate
Figure 10: Learning depth 1024. Unlike depth typically performs better.
512 case, in previous figure, larger
Transformer leamnnn rate number at n.n.
Transformer leaming rate number heads, at
2.00
2,50
number of heads
number of heads 1
1.25
1.75
number heads
number heads
1.50
Beint train in eypocts
1.50
1.25
1.25
Text box
1.00
1.00
10
10
10
1
Insrning
leaming
Figure 1I: Learning rate vs. number of heads per layer at n = 512. More heads are better than one head.
18

---

Published conference paper at ICLR 2024
Transformmr test learning number heads, 1024
Transfo mer loss leamning and number heads, 1024 2.00
2,00
number heads
number of heads
1.75
number heads
1.75
number of hawds
heads
number
1.50
1.50
1.25
1000
1.25
Text loxs
1.00
1.00
0.75
0.75
0.50
0.50
0.25
10
10
10
10:4
10"
10"2
learning
learning
Figure 12: Learning rate vs. number of heads at 1024. More heads are better.
mer loss learning rate embedding dimension, 1024
Transformmr vs. learning rate embedding dimension,
2.00
embed dim 32
enbbed dim 32
dim
dim
1.25
Besttrain in 2000
1e01 bo0
1.00
10
10
10
leaming
learning
Figure 13: Learning rate vs. embedding dimension at 1024. Smaller embedding dimension is generally better.
former test loss learning and batch 512
Transformer train leaming and batch
512
2.00
2.0
batch size 32
batch_stee
1.75
batch 120
1.75
batch Fize 128
batch_sion
512
1.50
3.50
1.25
1.25
Text 1003
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
10
1
10
10:4
10"
10
Iearning
learning
Figure 14: Learning rate vs. batch-size n = 512, Smaller batch size better.
19

---

Published conference paper at ICLR 2024
Transformer test loss when training last layer vs. init, magnitude
Transtormer train loss when training last layer vs. init. magnitude
intalization mutiplier mitillzationn mullipiier miBalization scale muttiplar 64
Iribalizttin multiplir? istiaizztton scale multipler Initaliztion scale multipller 64
1006
Tea
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
- train siphabet size samples
n- train alphabet size zamples
Figure 15: Training just the final unembedding layer suffices for the transformer generalize distribution, as predicted by our theory. However, the number of samples and number of epochs needed is larger than when all parameters of the network are trained. Understanding why training all parameters gives better performance than training, just the last layer is interesting future direction. report results for different magnitudes of initializing the weights attention mechanism (1 times, 8 times, and 64 times the standard initialization), and find that larger initialization helps, which conjecture is due the softmax being in the saturated regime, which leads more weight on the relational features.
Transformer test batch and number samples
Transformer train learning and number of samples, batch 1024
2.00
2.00
n1
n6
1.75
n.32
1.75
1.32
1.50
Beat train loso 10000 spocle
1.50
n52
n 512
1.25
1014
1.25
1024
Text lox3
1.00
1.00
0.75
0.75
0.50
0.50
10
10a
102
10"
batch
learning
Figure 16: Batch size vs. n number training samples training alphabet size. Smaller batch size is generally better, which is most visible at n = 512.
20

---

Published conference paper at ICLR 2024
Transformer test loss various tasks
Transformer train loss on various tasks
2,00
2.00
ABA ABB, 0.001
ABA ABS, Ir 0.001
1.25
AAB9
0.001
1.75
AASS
0.000
longth-2 majority,
lunoth-2 majority.
longth majorit. 0.000
longth majority.
1,50
longth-6 majoity, 0.001
length-6 majority. .r 0010
longth majoniyy 0.000
longth 8 majority.
1.25
2000
Teat
10
10f 10'
107
trah siphabetsize zamples
alghabet samples
Figure 17: Test and train loss of transformer for various tasks. The aßa vs. aß9 task consists of templates a,aa and a,3 with labels +1, The aa33 vs. 3 task has templates +1, -l. For each k, the length-k majority task consists of all templates in {a} x {a. 3)==1, where each template has label if a occurs more times in the last I entries, and label +1 if occurs fewer times in the last entries. The trivial model that outputs always will achieve test loss of 1.
Transformer test loss random tasks with different lengths
Transformer train less on random tasks with different lengths
1.75
0.32
n2
128
120
8192
1.25
1000
Text loak
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
17.5
template length
template length
Figure 18: Performance tasks corresponding of two, distinct random templates with two wild- cards a, B. and with labels 1, 1,, respectively. Performance degrades as the template length in- creases.
21

---

Published conference paper at ICLR 2024
Transformer
alphabet sizes
Transformer train on random
alphabet
n8
1.75 6222
1.75
n25h
256
n512
1.50
n51
3024
1024
1.25
88192
Teet
wildcand alphabet
wildrard alphabet
Figure 19: Performance on tasks corresponding of two random templates of length 5, labeled with 1, 1, respectively. Each template is sampled randomly from W conditioned on the two templates being distinct. We vary the wildcard alphabet size [W]. Performance generally degrades the wildcard alphabet size increases.
Tansformer test loss random tasks with different regular alphabet sizes
Transfommer train loss fandom tasks with different regular alphabet sizes 2.00
2.00
1.75
n2
1.75
32
128
128
256
1.50
512
1.50
1024
1024
1.25
2048
Beed Eren IOSN 1000
1.25
2040
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
0.00
0.25
0.30
0.75
1.00
1.25
1.50
1.75
2.00
requiaa aiphabet size
regular alphabet size
Figure 20: Performance on tasks corresponding of random templates of length labeled with 1, -1 respectively. Each template is sampled randomly from (W U X)5 conditioned on the two templates being distinct. We keep [W] 2 and vary the regular token alphabet size | between 0 and Performance quickly improves the regular token alphabet size increases.
22

---

Published conference paper at ICLR 2024
Transfo mer test loss random tasks with different numbers of templates
Transformer train loss random tasks with different numbers of templall complates 2.00
7.00
1.75
n2
1.75
n.32
6.64
1.50
Best treen lovs 1000 apeEEe
1.50
128
512
n12
1.25
1024
1.25
1024
2348
2046
1.00
1.00
8119
0.75
00.75
0.50
0.50
0.25
0.25
0.00
0.00
12
12
number of templates
number of templates
Figure 21: Performance tasks corresponding of random templates of length labeled with 1, 1.-1, respectively. Each template is sampled randomly from (W X) conditioned on two templates being distinct. We keep [W] 2 and vary the regular token alphabet size |A between 0 and 2. Performance quickly improves the regular token alphabet size increases.
Comparison of architectures, test less
Comparison of architectures, train loss
Transformer,
Transformar KO identity. ir 8.001
Tranfommr K0 Identity. 0.001
Transformar vD idenitty. Ir 0.001
Transfrmer: vo idertity a
Transformer identity, 0.001
Vansformer KG, vo identity. 0,001
RNN, 0.001
1.25
LSTM, 0.001
1.25
USTM, 0.001
Train loss
1.00
1.00
0.75
0.75
0.50
0.50
6.25
0.25
0.00
0.00
10r
10
100
number of samples = train alphabet size
1 number samples train alphabet size
Figure 22: Different architectures on 0, vs. a33 task. Transformer outperforms the other archi- tectures, especially with the reparametrization that prioritizes identities in heads.
23

---

Published conference paper at ICLR 2024
Transformer majority
Tranffrmer identity, on majority
2,00
2.00
majority Jen=2, 0.001
majority jenn Ir 0.001
1.75
majorityleem Jenes, 0.00
175
magonty_Jeee Jetitty_emmmmeeee tro0,
majonty_Jenie 0.00
majorty_en=s6..rr
majority_enn- 0.000
mnxority_eenBBBrO.00
1.50
1.50
1.25
1.25
Text
1.00
Test
1.00
0.75
0.50
10
107
107 10f
10
10a
107
10
alphabet
train alphabet #samples
Figure 23: Comparison of test loss of architectures on length-k majority task with different k. Left: vanilla transformer architecture. Right: transformer architecture plus the trainable identity scal- ings on each attention head's Wx W and W'. matrices. Notice that again the transformer reparametrivation lowers the amount of data needed by at least an order of magnitude.
GPT-2 pretrained test loss
GPT-2 pretrained by size, train loss
2.00
2,00
9pt2-smal pretrained, 0.001
gpt2-small pretrained, 0.001
1.75
gpt2-medium pretrained, o 0001
1.75
gpt2-medium pretrsined, 0.0003
9p12-large pretrained, Ir 0.001
gpt2-large pretraited, Ir 0.001
1.50
1,50
1.25
1.25
Toat l008
Thein loos
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
103
107
102
10"
number of train samples - train alphabet size
number of train samples - train alphabet s:20
Figure 24: Pretrained GPT-2 of different sizes fine-tuned on a,a vs. 0,3,8 task.
MLP MLP with XX" data augmentation. test loss
MLP MLP with xx" date augmentation, train loss
2.00
2.00
MLP. 0.000
MLP. it 0.000
1.75
MIP Xx" dats nugmentation, 001
1.75
MLP xx" data w.grentation, 0.00
1.50
1.50
1.25
1.25
Text loss
Train 0oc
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
10a
10a
10
101
10:
10a
10
10t
number train samples - size train alphabet
rumber train samples dize train alphabet
Figure 25: Test loss of MLP with x x data augmentation, where it is concatenated to input, versus ML. without data augmentation, versus transformer.
24

---

Published conference paper at ICLR 2024
C PROOF OF THEOREM 3.4
There are two main parts to the proof. First, Section C.I we establish lemma with a sufficient condition for kernel method to have good test loss. Second, in Section C.2 we prove that the trans- former random features kernel Ktrans satisfies this condition for almost any B, 7, b1 b2 parameters. We conclude in Section C.3.
The reason that we state our result with mean-squared error loss is that we have the closed-form solution (5) for the function that the kernel method learns in terms of its kernll and the data. Such an expression is not known for the cross-entropy loss.
Remark C.I.
C.I PART 1. GENERAL SUFRICIENT CONDITION FOR GOOD TEST LOSS
We restrict ourselves to loken-symmetni kernels, which are kernels whose values are unchanged if the tokens are relabeled by a permutation.
Definition C.2 (Token-symmetric kernel., K is tymen-symmetrri if for any permutation X - X we have (x. y) K ([T(1))..... T<==)), (my)). ... [(y]]]]
Token-symmetry is a mild condition, as most network architectures used in practice (including trans- formers) have ookensyymmetric neural tangent kernels at initialization. We emphasize that token- symmetry is not sufficient for good test loss since MLPs are counterexample (see Appendix L.)
To state sufficient condition for good test loss, let {zi,..., .,, 2, = supp(Mempit) the template distribution support. Define also the set R = Usejkl,jerr| { of tokens that appear in the templates. Finally, define N R by
N₂ = K s), s'))
(11)
where 8/: W + X' are substitution maps satisfying
s(W)
0 and s(W) ((W)) 0.
(12)
One can check that because of the token-symmetry of the kernel K, the matrix N is uniquely- defined regardless of the substitution maps s, chosen, as long as they satisfy (12).
Lemma C.3 (It suffices for N be nonsingular). If K is token-symmetric kernel, and N is nonsingular, then kernel ridge regression achieves vanishing test loss.
Formally, there are constants c, C > 0 and ridge regularization parameter X > O depending only on Pample, o, [W), ||N~ and || I|x max K ((x,, v)) such that for any x matching template E supp((Mtmptt) the kernel ridge regression estimator f in (5) with kernel K satisfies
(f(x)) [a(=)) S CV [og((/)) + +VV
T
n
with probability least 1 exp(-cn) over the random samples.
The proof is in Appendix D, but we develop an intuition here on why the nonsingularity of the matrix N is important. Let [m] I, Li Iz LJ... In be the partition of the samples such that if i I then sample (r,; yi) is drawn by substituting the wildcards of template zj with substitution map 8; : W + X. We show that for any string æ matching template zj. the kernel ridge regression solution (5) is approximately equal to the average of the labels of the samples corresponding to template
[K -
Σ vs fx(x;). ET
(13)
In order to see why this is true, consider the regime in which the sample diversity is very high, i.c.,
o > 1. Since p is large, any particular token is highly unlikely to be substituted. This has the following implications:
For most sample pairs i o i' E [n], the maps and sy have disjoint range: ((W))s(WW)
For most samples i (n], the substitueed tokens are in the templates: s;(W) RR 0.
25

---

Published conference paper at ICLR 2024
These are the same conditions as (8). So by the token-symmetry of the kernel, for most pairs of samples the empirical kernel matrix is given by N:
Rs = ((,,, = N3.5 for most E I,, EIy
So N is nonsingular, then K has r large eigenvalues, and n much smaller eigenvalues. This turns out to be sufficient for (9) to hold. We refer the reader to Appendix D for more details.
C.2 PART 2. ANALYZING THE TRANSFORMER RANDOM FEATURES KERNEL
We show that the transformer random features kernel Irans satisfies the sufficient condition of Lemma C.3 for vanishing test loss. It is clear that the kernel is token-symmetric because the defi- nition is invariant to the permutation relabelings of the tokens. The difficult part is to show that the matrix trans == N defined with kernel K = Kaann in (11) is nonsingular. The main challenge is that the transformer kernel does not have known closed-form solution because of the softmax terms in its definition (4), Furthermore, the result is especially challenging prove because it must hold for any collection of disjoint templates :
We analyze the MLP layer and the attention layer the transformer separately. We observe that a "weak" condition on Kattn can be lifted into the "strong" result that N trans is nonsingular. Intu- itively, long as Kattn is not very degenerate kernel, is very unlikely that the MLP layer has the cancellations that would be needed to make trans nonsingular.
Lemma C.4 (Nonsingularity of Ntrans, restatement of Lemma 3.6). Suppose for every non-identity permutation S, {id},
Σ Kann(suble 8), subb,((2,,))) # EE(ri
Σ RE(r!
Kattn (sub((z,, sub((z ()))
(14)
where 8, are the substitation maps in the definition of N trans in (12). Let the MLP layer's ac- tivation function Φ0t) = cos(bit + b2). Then for almost any choice of bi, (except for a Lebesque-measure-znn set), the matrix N Vrans is nonsingular.
This lemma is proved in Appendix E, by explicitly evaluating the Gaussian integral, which is pos- sible since the activation function the cosine function. Although in proof we the cosine activation function, we conjecture that this result should morally hold for sufficiently generic non- polynomial activation functions. Next, we prove the condition on N Natm-
Lemma C.5 degenergyy of Kattn restatement of Lemma 3.7). The condition (14) holds for Lebesgue-almost any 3,7-
The proof is in Appendix F. First, we prove the analyticity of the kernel Kattn in terms the hyper- parameters 3 and which control the softmax inverse temperature and the positional embeddings. Because of the identity theorem for analytic functions, it suffices to show at least one choice of hy- perparameters 3 and yy satisfies (14) for all non-identity permutations 7. Since Kattn does not have a closed-form solution, we find such choice of 3 and y by analyzing the Taylor-series expansion of Kattn around / 0 and 7 0 up order-10 10 derivatives, which happens to suffice.
C.3 CONCLUDING THE PROOF OF THEOREM 3.4
By Lemma C.3, it suffices to prove the nonsingularity of the matrix Ntrans defined in (11) with kernel K Kirans- Lemma 3.6 gives condition for nonsingularity that holds for almost any by by. Lemma 3.7 proves this condition for almost any 3, 7- Therefore, Theorem 3.4 follows.
D SUFFICIENT CONDITION FOR KERNEL METHOD TO GENERALIZE ON UNSEEN SYMBOLS (PROOF OF LEMMA C.3)
restate and prove Lemma C.3. Let K be a ooken-symmetric kernel as Definition C.2. Let Pempl be a distribution supported on disjoint templates 2,...., 2, and define R. = [ Recall the definiton of the matrix N e Rexr with
N ((sub(zi, 8), sub(zv, 8'))-
26

---

Published conference paper at ICLR 2024
for substitution maps W 8' s'()) n n 0. Recall that this is well-defined by the token-symmetry of the kernel K.
X satisfying s(W) n s'(W) s(W) n
Lemma D.1 (Restatement of Lemma C.3). Suppose that K is -ymme-ryymmttic and is nonsingu- lar. Then there are constants O < c C and O < c < C' depending only on Pemple. 0, [W], |N~ and [ x| maxx (xx, 2) such that the following holds. Consider any regularization parame- ter E [c'n, C'n,, and any string a matching template z E suppl suup(()mmp)). Then with probability 2 1- - - exp(-cn), the kernel ridge regression estimator f achieves good accuracy on x:
(0/(1/8) + +C VV/2.
[f(x) - f-(2)] S CV
7
Proof. Note that some proofs helper claims are deferred to Section D.I. Let (#1. (r1.81)...... (I,, Vn) be the samples seen by the kernel method. We know from (5) that kemnll ridge regression outputs the estimator
f(z) - v" K u(x),
(Kerell ridge regression)
where the empirical kernel matrix K E Ruxa is
Rij K K( 25),
and y [3n ,, Mal. and v(x)= (x x)... ..... (a x)] ER"
Idealized estimator when sample diversity is high If the sample diversity is sufficiently high, then for most pairs of samples i * i' E [n]]. it will be the case that æ, and ay do not share any the wildcard substitution tokens. In other words, the wildcard substitution map used to form will that should
have disjoint range from the wildcard substitution map used form æy. This expect the estimator to perform similarly to the following idealized estimator:
fideal (x) = "f ideat
+ +AI)+ + (ddcat
(15)
R and y pideal (ar) Ra are idealized versions of K and w(x),, formed below. They correspond to the limit of infinitely-diverse samples, when all token substitution maps have disjoint range. For each j E [r]. let I₁ S [n] be the indices of samples æ; formed by substituting from template zj. For any E I, E Iyy,let
where K iideal
Reidaal = =N53'1
(16)
Also, similarly define videal (x) e R". For any i E Iy, let
(doall (x) (suub(z).)):2)
(17)
where W X is substitution map with s(W) n R s(W) (ziNee 0. i.e. it does not overlap with the templates or with æ in the tokens substituted for the wildeards. The expressions (16) and (17) are well-defined because of the token-symmetty of the kernel.
the sample diversity is high, then we show that the idealized estimator fideal is indeed close to the kernel ridge regression solution f.
Claim D.2 (Idealized estimator good approximation to true estimator). Suppose ||]]
max [[x,, æ)] < xc. Then there are constants C, c >> O depending only on [W), | |oo,,,, kr such that the following holds. For any x, with probability at least 1 exp(-cn).
C. Cn
fiduad (x) ( f(x))
X/P'
where p is defined in Definition 3.3 and measures the diversity of the substitution map distribution.

---

Published conference paper at ICLR 2024
Analyzing the idealized estimator using its block structure The matrix K ideal ture with blocks .,....,Iy Iy- Namely, it equals K,,"" = Nujy' for all E 2j x' E Iz- Similarly, gideat( (a)) also has block structure with blocks : Z,. This structure allows us to analyze esti- mator fideat and prove accuracy.
has block struc-
order to analyze the estimator, we prove the following technical claim. The interpretation of this ideal claim is that if matches template Za- then (x)) is equal any of the rows in K that correspond to template a. In other words, we should have (K (didaal)+ y (ideal / which is the indicator vector for samples that come from template The following technical claim is more robust version this observation.
Claim D.3. Let be a string that matches template Suppose that 0 A < ideal
minsely |///|| Then (Ked + XI) is invertible and the following are satisfied
II(K ideal
AT)
o idoal (x)) S
V
Iial
and, letting 1x, R be the indicator vector for set In.
1z. In
(K sdeat
A1) + b ideal (a)) S
1).
In
Using the above technical claim, we can prove that fideal is an accurate estimator. The insight is + - ) is approximately the indicator vector / for samples corresponding to template a, the output of the idealized estimator is the average of the labels for samples corresponding template
that since (k ideal
Claim D.4 (Idealized estimator gets vanishing test loss on unseen symbols). There are e, C > 0 depending only on [W], Pample: o, | such that the following holds for any < X < en/ Let x be any string that matches template z E supp( Then, for any > 0, with probability
- 5 exp(-cnn) over the random samples, the idealized estimator has error upper-bounded by
log(1/8)
fideal - fo(()] SCV
CV
n
Proof of Claim D.4. Let Eg be the event that [T,] 2 m1/tmpp (1ppmpt(2220211 (z,) for all jE [r], i.e. all templates well-represented in the dataset. By Hoeffding bound.
P[ 2 exp(-cn).
Suppose that matches template z By Claim D.3, under event Ej, there is constant C >> 0 such that
jideoa( - -s((za) = 10" ((K""K"
ideal
+ - (2)) - s2(zz))
1t.
I 10
f.(za)l +
V Ia T 1)
Iaa
T
1r.
f.(2)))
+ C
In
We conclude since PlyyT
17, I
fa(za) > C log(1/61
E1] S by a tail bound for Gaussians.
Putting the elements together to conclude the proof of the lemma Combined, Claims D.2 and D.4 imply the lemma if we take À = (n)), then we obtain error Oly 1og([//88/ v1/e) with probability least 8 exp(-S2(1))
28

---

Published conference paper at ICLR 2024
D.I DEFERRED PROOFS CLAIMS
Proof of Claim D.3. Let Wi....y We be orthogonal basis of eigenvectors for K ideai with eigen- ideal
values Vi...., Vn- tdeal
Notice that these are also cigenvectors of K XI. Because of the block structure of K eigenvectors and eigenvalues have simple form. Define
M diag( (Iv I.l........ v IT Ndiag(( [
(I.I)).
The nonzero eigenvalues of ideal correspond to the nonzero eigenvalues of M, hecause for any cigenvector u E R' M there is corresponding eigenvector of K idieal with the same eigenvalue by letting each of the blocks Zj consist of copies of the entry uz/ VIzl. Therefore, all nonzero eigenvalues of K have magnitude at least
|4 |u 2 1/1 /
min
1 =
JClr
So K ideal
+ AI is invertible, which is the first part of the claim. Write
Lza
the eigenbasis as [Za
1I.
Σ
cwwi
Iza
for some coefficients ez. By construction,
ideat xx)
idoal 1x
Σ
Ial
IK(K
idoal
A11 -11)-1 12)-1110 ideal iddeal (x:) 12
Σ
vs www,
Σ
14 1222
4s +
S max
Vi )2
< max
)
1/1 +
In
7
Similarly,
1r.
(K Ideal
(x)]1
v; X Je,w;11
14
Σ (1
X )2
[Ia
Us
max(1
Vi
)a 1
max
)2
14
In
Claim (Bound on difference between kernel regressions). Suppose that K p.s.e and that (K ideat + +XI) -lydee (x) is well-defined, Then, for any À >> 0,
fideat(s) f(xx) (I (ddaadde))))) (w) v(x)]] + JK ideal ||I| [K
ideal
A1) -1 I ideat (ar)l)
Proof of Claim D.5. By triangle inequality,
[f(z) fidteal (x)) 1 (K + AI) - v(x)) 3² "'K ideal
+ AI) -1) ideal
|lll [KK + AT) - v(x) (K + AT) poddenl (ax)
Term
|gy|
[K + x1) 21) ideeal (x) (K
ideal
+ A1) -1 iddidca (x:]
Term
29

---

Published conference paper at ICLR 2024
The first term can be upper-bounded because
IK + 111--11 - [1( 1/A,
|nyidoal (ar) x()))|
Term S
A
The second term can be upper- by
Term 2 = II(K + ((K + (T)(K ideal
+ +X1)-1 - (Kr ideal
+ (T)(K
ideat
+ 211)-1)
ided (ar)
= |I(K + +I) - k iddeal
)( + +XI) - (xideaa)))) (x))
ideal
$ K - K I1I(K
ideal
ideal
21) pldool (x)]].
Proof of Claim D.2. Let Ej be the event that [Z,] 2 MH (z,) for all j E [r). By Hoeffding, there is a constant >> 0 such that P[E]] 2 1 - exp(-cnn) -cn). By Claim D.3, under event E₁, there is a constant such that
I[K ideal
+ -1)) ideut (x)) S
C Vn
(18)
Next, recall the parameter p used to measure the spread of the substitution map distributions {Heub.s (=sspp(()oe))
defined in (3.3). For each [n]]. let s; : W + X' be the substitution map used to generate the sample 2;. Let P1 be the number of samples (i, f') such that their substi- tution maps overlap, or have range that overlaps with the regular tokens in the templates. Formally:
Pi 7 ((1 <i< 5 n (:V)) n s (WV) of 0 or s; (WW) n R. * 0 or (wW) n R * 03].
Similarly, let P2 be the number of samples that (i, (') such that their substitution maps overlap with that used to generate x, or they overlap with the regular tokens in the templates:
1/2 = |(1 S i S n: 8; (WV) n R. * (;W)) n (Z)b)E # 0 By the definition of we can upper-bound the expected number of "bad" pairs Pi and "bad" indices Ps by:
On
Cn
Cn²
E[Pi]
Σ
Σ P(si(w)= = (w
+ Σ
Σ SE[M] FR
P(t E s(W))
p
p
p
(.PE(M) W.w'EW
Cn
E[P2 Σ Σ ((e)( 5
p
1E(m] (EEEJR
By Hoeffding's incquality, the event Eg that P1 < Cu² and P2 s Cn occurs with probability exp(- (cn). Under event Ez.
K K ideat
S - Cn/Vp and |,(()) ((doal()))) < CC/ppp
(19)
By Claim D.5 and (18) and (19), under events E1, F2, and using that || S CCV/. n, we have
CV (CCn//++ (C Cn/vp)
C On
f (ax) f(x)) s A
vn X
A/p
D.2 REMARK: EXPLICIT DEPENDENCE ON |N-1|
In the case that p oo, let obtain explicit dependence |N || in the bound of Lemma D.I.
Lemma D.6. Suppose that K token-symmelric and N is nonsingular. Suppose also that p co. Then there are constants O < c < and 0 < d' < C' depending only on jempls. o, [W/, and ||]] = max, K(x, æ)) such that the following holds. Consider any regularization parameter N |
/ || []], and any string x matching template x E suup((Pump)) Then with probability 2 1 - 8 exp(-cn (cn), the kernel ridge regression estimator f achieves good accuracy
A E
on
log(1/8)
IN
l5(m) - [(()) 5 CV
7
30

---

Published conference paper at ICLR 2024
Proof. First, by Claim D.2, we have ffideul (x:) f(x)) $ X- Next, by Claim D.4, we have ((deul(x)) - fx(z) & CV 16)
NONSINGULARITY OF RANDOM FEATURES AFTER MLP LAYER (PROOF OF LEMMA 3.6)
Consider a kernel K₂ formed from a kernel K1 as follows:
(2 y) = (())(v)) Ei = Ki (a, x) (x, Ki
Ki(z,y)
Here Φ + IR is a nonlinear activation function, Such a random features kernel arises in a neural network architecture by appending an infinite-widhh MLP layer with Gaussian initialization to neural network with random features with kernel K1.
We wish to prove that a certain matrix N E Rexr given by
Ny = K2(2..y;))
(20)
nonsingular, where ......y Tr. activation function, then only weak condition on K1 is required for the matrix to be invertible. We provide general lemma that allows us to guarantee the invertibility if the activation function a shifted cosine, although we conjecture such result to be true for most non-polynomial activation functions This is generalization of Lemma 3.6, so implies Lemma 3.6.
.........3 y. are inputs. The intuition is that if is "generic"
Lemma E.1 (Criterion for invertibility of N). Consider the matrix N e Rxxx defined in (20) where
D. xr and D1.... y, are inputs. Suppose that for all nontrivial permutations T E S, \ {id} we have
Σ Ki EE
# Σ Ki(ai, (r(1)) VElr
(21)
Suppose also that the MLP activation function @(t) (kt e) for two hyperparameters Then, N nonsingular for all (k, c) B2 except for a Lebesgue-measure- subset of R?.
Proof. Let e(k,c) == dct(N). We wish to show that ((k, c) f(k,c) 0} is measure-zero set. By Claim E.2, is an analytic function of and k and by the identity theorem for analytic functions (Mityagin, 2020), suffices to show that $ O. Fixing c 7/4, by Claim E.2,
22 K2(0, y) = 2 exp( (Ki Ki (x, v))).
Therefore
Σ sgn(T) II Kg( (+())
FES.
= e
(gu(())xp((
Σ
K1 (2t, (.()))
7ES.
se(r
remains to prove that as a function of we have
Σ
egp(7)
K1(x,, Mr(0)) # 0,
r
ke(r)

---

Published conference paper at ICLR 2024
This holds because for any distinct 01, ....... &j the functions exp(e14).... exp(eed are linearly inde-
pendent functions of i, since their Wronskian is rescaled Vandermonde determinant
exp(ctt) exp(cit)
exp(ed)
1
d exp(ct)
expl
c,))
df explett)
a exp(e(t)
expl
cd
II
(0--)) #0
1siejS1
Below is the technical claim used in the proof of the lemma.
a
N(0, Claim E.2. U, NNO,
b)
Then for any k c R.
BKCos(KU + c) cov + c)] = ze 3A 12p cos((c)) 2c) + 2 PR))
Proof. By Mathematica, we have the following Gaussian integrals
B ikVV (a+bb+2)))
Elc Elcth0 1
Ele ihU +hU+
1= = (6
Since cos(kt + = (ethl+ie + + c
Elcos(kU + c) cos(kVV +c]]
El(eakUtic + e-kU- (e te e
(e (A (c + 22n + 2c =- (a+b-22))))
20 11% (0 cos(2c)
PP.
ANALYSIS OF ATTENTION LAYER FEATURES (PROOF OF LEMMA 3.7)
For any inputs X, Y, we write the kernel of the random features of the attention layer
Katto () () (smaxx X (XY
XX" + XY + Yx" T +
m(X), m (Y) ~ (o,
YY +
stated Section 3.1; see also Section for the derivation of this kernel the infinite-wice limit of the transformer architecture. For shorthand, we write (8, attention kernel's dependence on the hyperparameters 3 and which control the softmax's inverse temperature and the weight the positional embeddings, respectively.
7)= Kattn (X,,, Y) to emphasize the
We prove Lemma 3.7, which is Katn satisfies the property (10) required by Lemma 3.6 the transformer random features kernel to succeed at the template task.
Namely, consider any disjoint templates 2i..... Z, and two substitution maps 8, W + X
that have disjoint range: s(W) n s' s'(W) d.
and the substituted tokens do not overlap with any of the tokens in the templates: s(W) 0
s'WW) where (=
32

---

Published conference paper at ICLR 2024
Then we define ;, Y; E R be the strings (where we abuse notation slightly by viewing them as matrices with one-hot rows) after substituting 2; by s, respectively:
X; sub(z;:
= sub(z.;;''))
Lemma E.1 (Restatement of Lemma 3.7), Define 9r (3. = EvE(r , ris) 7). Then for all
but Lebesgue-meausur-zee set of (8,7) 7) E R²² we have (,a 7) # gr (3,371) 7) for all permutations * id.
No closed-form expression known for (3, 7), so our approach is to analyze its Taylor series expansion around 3 = = 0. Our proof proceeds in stages, where, in each stage, examine higher derivative and progressively narrow the of that might possibly have (+, (8, 7) Section F.I, we list certain low-order derivatives of xx,y (3, that will be sufficient for our analysis. In Section F.2, we analyze some of the terms in these expressions. In Section F.3 put the previous lemmas together prove Lemma F.1.
To avoid notational overload, in this section we will not use bolded notation refer to the matrices X, Y. but rather use the lowercase X, Y.
F. LOW-ORDER DERIVATIVES OF ATTENTION KERNEL
the following table we collect several relevant derivatives 3 3. Hix Kxx.y (0,0), for S
j < 4. For each j we use Ci2 ..... denote constants that depend only k. and the derivative k.j being computed. Certain constants that are important for the proof are provided explicitly. These derivatives were computed using Python script available in our code. The colors are explained in Section F:2.
Derivative
Expansion
(x,
tep I
0/3 (0,0)
+o, Xxx +o2tr(XY")
+e1"XY"] +o 1o21'111o11XXXXXX1 XY"1 Hel
Ør (0,00
1 +Gz(1= "=""""""""(((")"""""
XX==XXXXX"XXXX (<01x -
t06
x"1""""x""">>11+00 """""""X"">>>1
(LXX GEXXDOG MXXX ( Ti..
+G2(
XX ( Mo+
teis -
+ te20 (1xXX 1
() 1
+021 (F ( (X
toza (1 ()( (
SEAAAA 4ot XXX
tog 1 X +0g (XY ryt) ( 
033 17 Rx,r = + +q +q1=xy2] +x2 t0s T
+cs(tr(XX )1 XXX"1 +co (X " 11
(1 100000 1
tc (1'""""'(((X ()(((XXY))) +
ton "1 11
te2
te (1
te 1d xxx"XXXXXX +c1o
33

---

Published conference paper at ICLR 2024
8/99 8, RX,Y (0.0)
8t
+ql +ozt((XY)) +og"" +
to +0 +ex(tr(XX) X1"xX"1 +o((t(XY")
to,1" tall
tc10 "ia""""""")a"""""""
t012 +x221xxyxxy"1 (()((((X))(((1))))))) (0)66)0666)))
tap P PXx
ters te1s (axry)""++httx00""0000""
ta1s
tc19 +egg""XX"""""
fr (xXX +cal NY YY" +027 1
teg "'X 1X toey tom 1"yy1
+ +c23 (( te2 11"" "
+025
(
Furthermore,
the expression for 6x, (0, 0) we have 61 1/k1²= >
in the expression for
83 85 xx,y 0), we have 02 8/k² 0.
d² a²
the expression for xx.y y (0, 0), we have c2o = 24/k3 0,
8*
8.4
in the expression for ,,
8² 0²
x y (0, 0), we have cre - -88/k* > 0,
and in the expression for
o Ø²
Kx, (0, 0), we have Czs 17280 / > 0.
0.
F.2 SIMPLIFYING TERMS
Let X E kkmm and Y E IRkxm be matrices with one-hot rows (i.e., all entries are zero except for one).
For the submatrix corresponding rows S and columns T, we the notation [X]s sxx RRXXT v is a vector, then the subvector consisting of indices [u]r.
Let R S [m]] be a set containing the intersection of the column support of X and Y: i.e., for all i [m]] R. either [X]]]x Section F.1 below.
Y[JA) 0. We analyze the terms in the expressions of
F.2.1 ASSUMING X/x [l'
Suppose that XRR YI' Then any of the pink terms can be written as function of only X only? Y.
"xxF|=|||*x|a||
X ( (1" (1"Y) ||[1 x 1 (")dddi"g""""""'1""""""""""""""" (1" (1'Y))02 = X|="
X"xxXx"xy"1 = 1 X ("X))"1 |||""X
XXXTyx"xy"1 X (1 (F ||[I 1l
VYX'XX'Y'"1 = X)diagg (Y)Y" = ||11"|
X)Y") = Exe (1 ΣieWA EVER X (1" = 1" X (1)x |[[1' x|
traco(XXX' XY =
L1) || x|
F.2.2 ASSUMING [||X [|
Suppose that X[R|RR YKRRR (i.e., the restriction of and Y the rows is equal). Then any of the orange terms can be written as function of only X or only Y.
tr(XY)) Eveim) Eiety NioYio EVER E x 12 ITYlR

---

Published conference paper at ICLR 2024
1'XYXX'1
Ennoeyy 1
86)l(zb
ye)
(W XM) ((W))))
=1"X (X [H))))) (X(
FXXTXxTI= = Ease ( 280 = =xa) = Labe ( = 11) te E R)
-Earoo 1(z0 = 20 E R)1(3n
ER) = = 11 Ib E R)](22 2c E R)
I'XMMxx(XXp xx) x 1xx
PYXTRTTI =1² [A)R (X (M)TR)) ( (Xx))x)) ((xR)
trace(XXy* x Ea. ( yn 1 = va) = Las 1 = yb E R)1(2b = yo E
R)= = =a, 1 = 26 E R) = tracel 1x (X1)RR))
E.2.3 ASSUMING 1'X'11 = 1"YY"1
Suppose that 1'XX = 1. Then any of the blue terms can be written as function of only X only Y.
=
1TYYT 1=xX
E22.4 ASSUMING XXX = 1"YYT
Suppose that 1"xXT = (Tyy. Then any of the teal terms can written as function of only X or only Y.
1'xXTyX -
F.3 PROOF OF LEMMA F.I
We combine the above calculations to prove Lemma F.1.
Proof. By the technical Lemma G.1. know that (3,7) is analytic function for each Therefore, by the identity theorem for analytic functions (Mityagin, 2020), it suffices to show that for each T E S, \ {id} we have gia(3,7) 7) $ (r, r)-
Stage I. Matching regular token degree distributions.
Claim F2. If g0a(0,0) 0) = 9- (0, 0), then [1 X1}R = (1" for all i [r].
Proof. From the table in Section F.I, there is positive constant ei such that
gr (0,00)) = c Σ
xxXT 78
1=
Σ
- Y 1|
kEir
S
3V
[ X+}x
|||[1 [|[[*T+())]]]
sE(r]
I33
X.|x 12
Σ sE|r
[[ITY 1
sE(r]
Σ KE|
||| ||T2|111
where (a) is by Cauchy-Schwarz Schwarz and holds with equality if and only X.)R x [1"Yraal for all i. Similarly (b) is by Cauchy-Schwand and holds with equality and only if [|]1a, [[] ([1'Y=m) Y=l)al for all i. Notice that (a) and (b) hold with equality if = id, since [1' X, = Y
for all
35

---

Published conference paper at ICLR 2024
Stage 2. Matching regular token positions.
Claim F.3.
8² 8 B3 00.0)(000) 8 372 tha (0, and [1= (I for all [r], then
8 &
we have [Xi] [Y [ for all E [r].
Proof. For a constant C2 > 0,
0 8
9- (0.0)
Σ
X2Y ogtraco(Xx,
032 872
sEpl
||
x
tracol (x,oo(X((x70)3))
C2
vefri
vEjrl
by the calculation in Section E.2.1. sum does not depend on t, so we analyze second sum. Here.
trace
X,Y
c2
Σ
Σ
X Faylao
sE
DE
KilanDYstslla
sE VER oE
E (xlla aE|k)] VER
Σ (Vr1gla0)
SE|r
12X,
SE
where (a) is by Cauchy-Schwarz and holds with equality if and only if x (7(()) for constant c. We must have e = because of the CLS token, so (a) holds with equality if and only if [X.][[x (axR for all E [r]. Specifically (a) holds with equality if r = id.
Stage 3. Matching wildcard token degree histogram norm.
Claim F4. Suppose that [1 X,Iæ [1'Y~(s)R,, and that 84 8/ (0,00) IT XTX Y TTg) for all i [r].
yt 8,44 Ird (0, 0). Then
calculations Section F.2.1 the pink terms. Every term of gt (-, 0) can be written as depending only on one of X, or Y-()- with the exception of the C2p term. Namely, we have
Proof. Use x1=x = Y
81 23.0-(0.0)
Σ a((x)) ((Yr(g) sE|r]
+ (1 (1" O)(1 X,XX)((XY 1),
1),
for some functions a, b. Since is permutation, only the term with coefficient czo depends on Here, Cop This term corresponds to
C2o Σ (1 x
() X,X 1)(( (Y
||||1 ( (1 X.X 1) Yom
iE(r]
|[IXX]]]|P((XXXXX711111
| 12 1)2
Fe[r]
velr]
Σ
36

---

Published conference paper at ICLR 2024
where (a) is by Cauchy-Schwer Schwarz and holds with equality if and only if [[]"Xx x
((a [/"1" Yoto Fin for all i and some constant c. This constant c c= because the former
is a permutation of the latter over E [r]. Since [[[1 x1|x and since we have the CLS token, we know that (a) holds with equality if and only if i X,X" 1 (() Y70) (() for all i E [r]. This is the case for = id by construction of X, and Y,
[| > 1 by assumption
Stage 4. Matching wildcard degree distributions.
(6, Claim F5. Suppose that [xil|x>x [Y and i x, = (at) (() for all E [r]. 8 82 Suppose also that age 0n2 9- (0,0) = i E [r]
Then 1 XXX, = 1 (6) for all agt 0x2 (0, 0).
8t 8e
Proof. Similarly to the proof of the previous claim, because of the calculations in Sections F.2.1, 3 0² F.2.2 and F.2.3 for the pink, orange, and blue terms, respectively, we write
of terms that each depends on either X, or (-()) plus Even Cia1 (xx))Y=()) (()) This latter sum the only term that depends on т, and constant C1s satisfies cr6 > Similarly to the previous claim, by Cauchy-Schward
35 0-2 as
Σ ciel x,,,x,X Y 1() I K
Σ Cia |1
e X.X 110
+E(r]
with equality if and only if X = 1r
YT r(g)
for all i, since (X,X²) je a permutation
(Yr() High This condition holds for r = id.
Stage 5. Matching wildcard positions.
Claim F.6. Suppose that [Xil[s] xR.
[Y [Y and 1 X;XT = 1" Y7,) for all i e [r].
Suppose also that
05 क. Or (0, 0) 05 उ. gia (0, 0). Then x,X? Vi for all i e [r].
o gr
of ar
8a Proof. Write 8 3x4 9- (0, 0) as a sum terms each depending only on either X; or Y-() by using the calculations Sections F.2.1, F.2.3, F2.2, and F.2.4 handle the pink, orange, blue, and terms, plus (for C25 > 0),
Σ
cen trace X.X ((
Σ can X.X ||| XXX1)
K
FE(r)
with equality and only X,X! = Yr() YT concluding the claim.
for all e [r]. This equality holds r = id,
(0)
Combine the above four claims to conclude that if gr (3, = gia(3,)) 7), then we have X,XT Y-() Y7o and [XI]ı| [Y~o]uи for so id.
G ANALYTICITY OF ATTENTION KERNEL (TECHNICAL RESULT)
We prove the analyticity of x (x,x)),,) = Ka XX X) as function of 3 and y.
Lemma G.1 (Analyticity of Kaun), For any X, X, the function xx,x analytic in IR?.
Proof. Note that we write
m XÇ + PP,
m: m( xČ+
where -~ NoO, Im) and p NO0 1x)) independent Gaussians. So can rewrite Xx,
x x (8, B L
37

---

Published conference paper at ICLR 2024
where
1(B.7: = =5 (XXX" + +-21)3.
smax(3X( 3yp)"
(maxx(XX5 3rp)
The main obstacle is to prove the technical Lemma G.9, which states that for any k1, 22, we have
th aho m.c. C. pill & C(1
E E. 03h. 8 f(3 7i
hnl
So by smoothness of f and dominated convergence, we know that we can differentiate under the integral sign, and
t t d3 drd xx
ghr Ihz
(8, | 03% 8yh f(3, m x, X, C.C.
C + 1b [8] + kkt)).
Because of the bound the derivatives and its smoothness, x (B,
(8, 7) is real-analytic.
The proof the technical bound in Lemma G.9 is developed in the subsections below.
G.1 TECHNICAL LEMMAS FOR QUANTIFYING POWER SERIES CONVERGENCE
In order to show that the values of the attention kernel are real-analytic functions of in terms of 7. we will need to make quantitative certain facts about how real-analyticity of preserved under compositions, products, and sums. For this, we introduce the notion of the convergence-lype of a real-analytic function.
Definition G.2 (Quantifying power series convergence real-analytic functions). Let U R be an open set. We say that real-analytic function f U > R has (1I, r2)-tee for functions U + Rso and T2 U + Rz0 if the following holds. For any So. consider the power series of around So.
Eacon Su)
Then for any such that || Solloe S (So) this power series converges absolutely.
leco.ulls - Sold & r2(Co).
# s.t. 1A1
We provide rules for how convergence type affected by compositions, products, and sums.
Lemma G.3 (Composition rule for type: quantitative version of Proposition 2.2.8 of (Krantz Parks, 2002)). Let U C Rm and let V C IR be open. Let fi...., :U V be real-analytic with (T), Te)-type, and let : Va IR be real-analytic with (01, 0ty-eye. Then the composition h = g fn) is real-analytic with (min(r1, (01 f) T2 2,,02 o f--type.
Proof. Fix some En and let yo (f() ... fn and let a Cova be the coefficients of the power series expansion for fr around Ço- Define p min(1, on (No)/Ta(Ca)) (Co)). Then, for any Ç such that II Sa S pri(So) and E [n]] we have
Σ
(9 Go Sol" & Σ
(9)
(1())) s prz(Co) S (yo))
S.t. |A|21
S.t. |a|21
So, letting E e yu) be the series expansion of g around yo- we have the following absolute convergence
14
byo II
Σ
la ( Ic
Sol"
S 502(10)
S.L. 1|
P s.t. Jel
38

---

Published conference paper at ICLR 2024
So we may rearrange the terms of
Σ yyn
II
Σ
22 - So)A
Aa S.t. Ja|21
as we please, and we get an absolutely convergent series for around So-
Lemma G.4 (Sum and product rules for type). Let f R R and R R be real- analytic Junctions functtonn of (T, ra)-yep and (01, 02)-type respectively. Then h = f + y real-unalytic of (min((1, (1) 72 + 72) -type, and h fg is real-analytic of (inn(11,01) 7202 + Talgi + 12
Proof. Both of these are straightforward from the definition.
Lemma G.5 (Derivative bound based on type). Let f Rm - be real-analytic with (7), Ta)-type. Then, for any multi-index js,
Olml a
T2(Gn)
f(Co))
TI(So)lml
Proof. Let aCou be the coefficients of the power series f at So- Since f is of (Ti, -ty-e, we have
Σ Ilco=al s r2(Co)-
1221
Since terms the sum are nonnegative, for all j with |p| 2
|oc, r2(Cu) (1/T1 (Su))(l))
ginl The lemma follows by Remark 2.2.4 of Krantz & Parks (2002), which states aç" f(So)) = |
G.2 APPLICATION OF TECHNICAL LEMMAS ATTENTION KERNEL.
We now the above general technical lemmas specifically prove that the attention kernel is analytic in terms of and 7/ Lemma G.6. For any j E [m]] the fienctiionn f 12 - R given by J(C) = smax((), is reul-analytic of (1/(2e²), 1)-rype
Proof. Write = g h for g : Ry + R and h R E given by g(v) 1/y. and h(() = E ehi
The power expansion of g(y) around yo R is given by
(-1)++1
g(v)
(y vo)*
y0 k+1
so one can see that g is of (p1, Py)-type for Pi (yo) series expansion for h(() around So
2o/2 and p2(10) 1/20. Finally, write the
Go. (So.i Gi
(Si So.s)
A(C) = 1 +e &
Σ
So,s
e
21
k!
E ())
xE(m)( 1=0
k=0
Note that this expansion converges absolutely for all Ç. as the absolute series is
Go,s (Coss - Gil Coo Ise
Co.el*
1+ Σ
11
k!
sE(m]/b
Σ
1
sE/m]//
< 221l-6ol--h(6

---

Published conference paper at ICLR 2024
Specifically, is of (1, eyhp)-yyp. So by the composition rule Lemma G.3, must be that f real-
analytic of (21,72)-tyr ttype for TI = min(1, (p1 oh) (4)) = 1/(2²²) and 72 = o / = 1/h _ 1. D
Lemma G.7. For any j E [m]] and X, the function f IRa + given by f(8,7)
smax (5XC + 3yp), is real-analytic of (min(1, 1, 1/(e2 || + 2e 2|3( + [nDlpp]l>)) -type.
Proof. Write f = g h for 9 R + R and h R + Rm given by g(v) smax(v), and h(3. 7) = 3XÇ + Sryp. We know from Lemma G.6 that is real-analytic of (1/(2e"). 1) 1yypee. And easy to see that h is real-analytic of |
(15) + hnllppl (ppplpllp>))ty -type. Apply the
composition rule of Lemma G.3 to conclude.
Lemma G.8. For any X, X. S. C.P. the function f : given by f(3,7) = smax( (Tp)* + (a)mmax(xXXX Brp)) is real-analytic and of type
1
(minll, 26a |X (5) >D p 20 | + (18) + 1plp
). + +7)),
where C is constant depending the context length l..
Proef. Each entry of (XX* + +f) real-analytic y and (I, 71-yte.... So by combining with Lemma G.7 the product rule and sum rule (Lemma G.4), and the fact that each entry of the smax is most one.
As consequence, 9hp) (XX +m² + Byp),, which was what we needed to prove Lemma G.I.
bound the derivatives of S(3, 7; X, x,,;x,x,,,, p; (
Lemma G.9. For any k1. k2 2 0,
gki gh2 03h 02 [33, x, X, C, E,P)
$ cC(++>)) max ((22)) )II + || + (13) (nDlppl)ss)))) k) + )ky !k?!.
Proof. Direct consequence of Lemma G.5 and Lemma G.8.
H DERIVATION OF TRANSFORMER KERNEL
We state the transformer architecture and informally derive its random features kernel in the infinite- width limit.
H.I TRANSFORMER ARCUITECTURE
We consider depth-1 transformer architecture (without skip connections or layernorm, for sim- plicity). This architecture has heads, each with parameters W K,A. Wo. Wv, o. E Xedes, and embedding layer W.. E Ruxdemo, positional embeddings P Rexdems, an MLP layer with parameters W A. WB E Rmmmm Xdems, and a final unembedding layer with weights Wy E Rdemss The network takes in X E Rkxm and outputs
frans x;0)= = =
(Unembedding)
where
WHo 1 W (1z1)= E Rdi.m demt
22
(MLP layer)
dmll
21 VH bE[H]
Ex E Rdes
(Attention layer output at CLS token)
3Z. WT R.
W Q. 20
W W O.h V.A
Won
A Smax
Zo
E ixdddm
(Attention heads)
dembs v dhood
V /cheaddems
Za XW E 7pP xxdoo
(Embeeddng layer)
40

---

Published conference paper at ICLR 2024
Here B.Y 2 0 two hyperparameters that control the inverse temperature of the softmax and the strength of the positional embeddings. respectively. Note that only the output of the attention layer the final kth position CLS token used, since this is depth-1 network. The smax is softmax applied row-wise.
H.2 RANDOM FEATURES KERNEL
The derivation this kernel assumes that every string æ ends with special [CLS] classification token that does not appear elsewhere in the string. We choose that initialization so that each of the entries of the intermediate representations Zo. 21, is of order (11). In order to accomplish this, we initialize We g. P. W K,A, W Q W V.A, W o. W. A. W W with i..d. NN0, 1) entries.
also initialize wy = O, and only train uy while maintaining the rest of parameters at initializa- tion. The random features kernel corresponding to training wy is
Kerans (X, X,Y)
x2 ==(x("222)))) 22 (2))) /demb.
=
where view 2y function of the input (either X or Y), and depending the randomly- initialized parameters of the network.
the limit infinitely-many heads H, infinite embedding dimension demb and MLP dimension dmip and head dimension dheod, the kernel Kirans tends to deterministic limit Krans which can be recursively computed (see, e.g., Jacot et al. (2018)). Assuming that the final token of both X and Y is the same token (i.c., CLS token), the deterministic limiting kernel Krrans is given by:
K [X X) Kann (X, Y)
K Kerans (xrans(X,Y (X,Y) [u,(o(u)((()) for u, N(0,
Katn (Y, X)
Katm (Y, (Y,Y) YY)
(22)
where Katt (X,Y) Emm(X). |smax (yn/X).mmyy//mmx(((m
(XY" + 2 (mmax((3(()))
XX" + XY" + YX + YYT +721
m(X), m(X).m(Y) N(0, x2)
Notice that the covariance matrix in above definition the distribution of m mXX), mYY) rescaled compared to that in the main text in Section 1.1, but this is inessential, since we can simply reparametrize 3 as 8 + 3/v B/VI+?? x3 to recover the expression in the main text.
H.3 INFORMAL DERIVATION
We provide an informal derivation of (22) below. Informally, by law of large numbers we have the following almost sure convergence
22(X) zz(Y
a Themi
W (X)) W WEWTE W (x)(
Kerane X, Y)
demb
demidmmby
dembrtno of Vilemb W Wazi(X))
W aY))
Vilemb
dmii
Amip Bu,, (u))a(U)) for N(O,
Katta (X,
Katm (ant(X, (X,8)
Katm X)
Katn Katt(Y (Y,Y))
Kurane X, (YY),

---

Published conference paper at ICLR 2024
where Kattn is the kernel corresponding the attention layer the infinite-wiin -iithh limit, defined as:
2 (x))x{())
En (X)
K Kaun (x, Y)
demb
demmb
3Zo ( K.h W Q. (x())
e smax
)Zo War
dhead cmb b.
demb v Chood
BZo W.
K. W Q. Zo()),
W O WV. V.N (Y smax smaa
Jee
demb v dhoad
XZo(X)W K.A W Q. (x))
dhead Semb 400
ef smax
XXY +
demb V dhead
he(H]
3Zo (Y W K.A. W Q Zo(Y)"
smax
yee
demb V dhoad
3Zo X)W K.N W Q. q. )
33 smax
XXYT + +
demb v dhead
3Z
W Q. Zo(Y)²
y'er]]
Smax
demby dhead
3e Zo(X) W (x))
E|smax
1XY +1
demb dhead
SET Zo ZolY) K.A. W Wo Zo (Y
smax
)"1
demby dhead
damb +00 -oooehe +00
E ()) (smxxs/sammmm( (XYT + s)smax(3m(())
= Katta (X, S).
where
XX" +72 XY +P2 YX" + YY +121
m(X) ~ (1+
because due to the randomness in W K.b and Q.h we have that
Zu(X) Q.A W K.A ZT ex
demby dread
and
Zo(Y) Q.A W K.K. (T) denby dendyd dread
are jointly Gaussian with covariance:
Zol Z Q.b W Z, e of ZY(Y) R.h. W Q. Zo()7
E(X, Y)
1..
denb v dhead
demb' dhead
42

---

Published conference paper at ICLR 2024
Since this an expectation over products of jointly Gaussian variables, any i. [k] we can calculate:
( Y) = 18 18w, n:
Zo( 1r exe Zo(Y))
cmb
r,e (demm)
= EE
Zp Zo0
cmmh (e.ee()emm)
EEwg,
E
[ E + YW - T
Plap
de
cmb redte(demm)
Σ
Ew w.P. P YW Pl
d cmb r.seldami
EWE.P(XWE P (YW + O(1/demb)
tE|demb
(1 7 + ((1/demb)
demb
n.
r,eeddmms)
EW EXWE+12 + - + +
demb
rE (elemm))
[XY") Dis 28. (1 (1+ + ((1/dab))
where (a) we use that [XW8 & abb and (YW & + + Plab are independent of [XWs & P
and (YW E + 7 ed unless = d. So
E(X,Y)
demb (1+ (XY" +T²I).
I MLPS FAIL TO GENERALIZE ON UNSEEN SYMBOLS
A natural question is whether classical architectures such the MLP architecture (a.k.a,, fully- connected network) would exhibit the same emergent reasoning properties when trained with enough data. In this section, we prove a negative result: an SGD-trained or Adam-trained MI.P not reach good test performance on the template task. This is in sharp contrast to the positive result for transformers proved in the previous section.
The input the MLP concatenation of the token one-hot encodings. The MLP alternates linear transformations and nonlinear elementwise activations. Formally, the MLP has weights 0 = (WI..., Wz. w) and outputs
MLP architecture
(aLP( 0) = w (x; 0) where <x; 0) o( + (2: 0)) Rd for
(23)
=,( 0) =y(x>) (e.1 end Rm
We consider training the MLP with SGD.
Definition L.1 (One-pass SGD training). The learned weights 0' after t steps SGD training are the random weights given by initializing 0 that each of w ...... W w have i.i.d. Gausian entries, and then updating with 0 -'-11_ Vo(fMLP(z' y² le- for (x* (w'"")) ~ D some step size 7)
We show that SGD-tran trained MLPs fail at the template task since they do not generalize well in the case when the templates consist only of wildcard tokens. In words, if the template labels f. are a non-constant function, the MLP will
reach arbitrarily low error no matter how many training steps are taken. Let Xmns C X be the subset of tokens not seen in the train data. We assume that |xans 2 k which guarantees that for any template there is at least one string matching it where all the wildcards substituted by tokens in X Under this condition:

---

Published conference paper at ICLR 2024
Theorem 1.2 (Failure of MLPs at generalizing on unseen symbols). Suppose that the label function Ja is non-constant, and that all templates in the suppori of Pemple consist only of wildcards: z E W for all & E supp(pumpit). Then, for any SGD step there is string x (Rune)* ) that matches a template E supp(pumppt) such that
Eg ((fm 0¹) f.(2))
where is constant that depends only on pempht and f..
The proof relies on the key observation that SGD-training of MLPs satisfies a permutation invariance property (Ng, 2004). This property guarantees that MLP cannot consistently distinguish between the unseen tokens, and therefore, in expectation over the weights 0', ouiputs the same value for any sequence a E (Nuns)*. We make four remarks.
Remark I,3, MLPs are universal approximators (Cybenko, 1989), so there are choices of weights 0 such that [ 0) has good generalization unseen symbols. The theorem proves that these weights are not found by SGD.
Remark 1.4. The theorem does not assume that training is in the NTK regime, i.e., it holds even for nonlinear training dynamics.
Remark 1.5. The theorem also holds for training with Adam, gradient flow, and minibatch- since the permutation crrmutation-invariance invariance property of MLP training also holds for these.
Remark 1.6. As sanity check, we verify that MLP kernel does not meet the sufficient condition for generalizing unseen symbols from Lemma 3.5. The kernel for an MLP is an inner product kernel of the form K (1r, x²) x(Enn [xx; = x!)) for a function IR - IR. Therefore, the matrix N Rex" has all of its entries equal to No. x(0). it singular and the condition Lemma 3,5 is not met.
We now prove Theorem 1.2. We first show that trained MIPs cannot differentiate between tokens in the set Xuns- Let X = Xeeen Xane be the partition of tokens into those seen and seen in the train data. Here Xzeen is defined as the smallest set such that x E Xx almost surely for (ar, v) ~ D.
Lemma 1.7 (Trained MLPs cannot distinguish unseen tokens). For any number of SGD steps t, and any learning rate schedule m..... the learned MLP estimator cannot distinguish between sequences of unseen tokens. Formally, for any æ1 32 X we have
Eg [fMLP(æ1; 0')] = Eg [fMLP(22;; 0')
Proof of Lemma 1.7. The proof of this result is based a well-known permutation invariiancei prop- erty of ML.ss trained by SGD. This property has previously been used show sample complexity lower bounds for learning with SGD-trained MLPs (Ng. 2004; Li et al., 2020), as well as time- complexity lower bounds (Shamir, 2018; Abbe al., 2022: Abbe & Boix-Adsera, 2022). In this lemma, we use the permutation invariance property to show poor diitt--iistiibuttinn generalization of SGD-traiied MLPs.
First, construct a permutation II E Bgmxxkn such that ITz)(=1) = (x2)), but which also satisfies that for any (Xyeen)* ) we have ITz((²) zo(Ø). This permutation can be easily constructed since neither æi nor 22 contains tokens in Xgeen- Next, define the following network fMLP- analo- gously to (23) but with the first-layer inputs permuted by II
MM:(<; 0) = - wT= =F<<;;00 0) R
where
=x; = o S.W (x; 0)) Rd for
= (x; 0) = =!!' (m) = ../e......, epal km
Now let us couple the weights 8" ..... 0¹ from SGD training of fMLP on dataset D, with the weights
11 ...... ...., " from SGD training of fll dataset D. The coupling is performed inductively the time step, and we set 6" @ll For the inductive step, 7 2 1, we update the weights with the gradient from some sample (ar y" Since xt (Kaeenjk almost surely, we know that 2 20(=") z!f (x7) almost
maintain the property that 0- gll.r for all For the base case 0.

---

Published conference paper at ICLR 2024
surely, which means that 0" =0" almost surely. We conclude the equality in distribution of the weights
0' _ ou.
(24)
Next, let us inductively couple the weights 0 ......, O' with the weights @ ..... OT ....,0".
different
way, so as to guarantee that for any time S S L, we have
W W. II and W w... for all 5 513 1. and un w TTr
almost surely. The base case and is also invariant to permutations since it is Gaussian. For the inductive step, couple the sample updates so that SGD draws the same sample (æ" y²" D. One can see from the chain rule that the invariant is maintained. We conclude the equality in distribution of the weights
follows because the distribution of W and w. is equal
Ø (Wi...., W2. wi) - (W. W
W.. w..))
(25)
Combining (24) and (25), we get
0' i.... W L w'} & (WiIl, W W W w').
which,since [Z)((r)) z0())). immediately implies
fMLP(Z1: 0') = fMLP (x);; [Wi W W2, w')) $ fMLP(Z2;
which proves the lemma.
Theorem 1.2 follows as a consequence. Note that the key lemma proved above only relied on a permutation invariance property of SGD on MLPs that also holds for Adam training, gradient flow training, and SGD with minibatch (see Li et al. (2020)). Therefore, the result holds for training with those algorithms well, beyond just SGD.
Proof of Theorem 1.2. Pick any two templates 2, E supp((rmppt) such that f.(2) * f=(2") Recall that z, Wk by assumption. Since we assumed that (Xunsl 2 k, there are strings x, x E X matching templates z and z', respectively. Furthermone, by Lemma 1.7, if we define a Eg [MM 0')] = Eg [MMLP(x': 0')], we have
max [(MMLP(a; 0²) f=(=)) Egg (/MML(('''' fa(2())²)
2 max((a f.(z))² (a- f=(=)))²
- fo(2))² =c> 0.
J DEFERRED DETAILS FOR NEXT-TOKEN
TEMPLATE TASKS
J.I DEFINITION NEXT-TOOKEN TEMPLATE TASKS
eext-token-predeetine template tasks, the output is a token in X, with the cross-entropy loss for multiclass classification. The formal definition of these tasks is:
Definition J.1 (Multi-class prediction version of template). The data distribution D Pmatticloss
Dmssitictase (Humph:: {Meab. J. fx) is specified by: (i) a template distribution Humplt supported on (AU W) (ii) for each template z. distribution Paud, over substitution maps W + X: (ii) a labelling function f- supp((tmppi + XUW, A sample (x; y) x x X drawn from Dmutticlase drawn by taking sub(2,s) s) and y sub(f. (2),,) where z ~ Hemple and ~ Psub,z
J.2 FAILURE OF TRANSFORMERS TO COPY AND MODIFICATION THAT SUCCEEDS
We provide the deferred proofs for Section 4.
45

---

Published conference paper at ICLR 2024
Attention layer architecture For simplicity in this section we consider a transformer with the attention layer only, since the MLP layer does not play a role in the ability to copy unseen symbols. Our architecture has H heads with parameters W K.A. W Q,h. WV... W O.h E Rehendadamb, an cmbecding/umembecdinng layer We E E Bmxddmm positional embeddings P E R an MLP layer with parameters W A, W E Be Xdem,, a final unembedding layer and activation function Φb. The network takes in E B and outputs
fatto ( (x:0) WEZ R
(Unembedding layer)
where
21 E Ali ck bE[H]
A smax (32 WZ K.N W Q.A 2.) V.b Wos E Rixdems
(Attention heads)
Zo XWE E + 7 E R Rexdemt
(Embedding layer)
and we tie the embedding and unembedding weights, as often done in practice, for example in GPT- 2 (Brown et al., 2020). Here 3, y 2 0 are two hyperparameters that control the inverse temperature of the softmax and the strength of the positional embeddings, respectively.
Simplification in our case We consider here a next-token prediction setup, where there is no final [CLS] token appended to the string. Namely, given string æ X this is inputted to the network as a stacked matrix of one-hot vectors for the tokens of the string X = =ler..... Ezy]. We study very basic template task: template "a" labeled by a, where a is wildcard. An example dataset generated from this template could be (AA,A), (B, B), C)), where A, C E X are tokens. Because the template has length & = 1, X E RRxx is a one-hot vector encoding the input token. Furthermore, the softmax output is always a 1 x 1 matrix with the entry 1, so the architecture simplifies to
Sattn (X:0) = (E() Σ hE(m]]
W W WW + +TP").
(26)
We initialize the entries of P and WE be i.d. NO. 1/demb). the entries of W o. be (0, [d(emm)), and the entries of W v.a. be (0, 1/dneed). so that as demb -> de the variance of the output vanishes as ((1 /dedm)))) as in the mean-field scaling (Mei et al., 2018; 2019; Sirignano & Spiliopoulos, 2022; Chizat & Bach, 2018; Rotskoff & Sanden-Eijnde 2018; Yang & Hu, 2021).
Derivation of kernels driving dynamics at small times Despite the simplicity of the task, the architecture does not generalize well on unseen symbols. Our evidence for this will be by analyzing the early times of training. For these times, the dynamics are governed by the neural tangent kernel (NTK) of the network at initialization (Jacot et al., 2018; Chizat et al., 2019). Let us derive the neural langent kernel of this architecture. This is network with output of dimension m, so for each i.j e [m] we will derive (x,, Kg,v X'), K'y, (X, K',,
X') which give the dynamics at small times for training the {W O. (.A)))(B)) the { v.a valkeem the Wp, p. and the W g weights at small times, respectively. Writing We g = [wg,1.... W.m by the law of large numbers.
(alanm [[;ann(;;01) (x:01.
[Olant (X):0),
Xx,,(X,X'))
Σ
dWo..
aW o.
hE[m]
Σ (XW B + + V.A
W WW + +,p" Yu WEs
x H
hE[M]
dnead 00 000
8ig (6,),, + +72)
46

---

Published conference paper at ICLR 2024
(8laatn (bL)tt((:00)))
V (XIatt(X':00), VW V.A
xx)x(X, E
VWv. v.h
WE(H]
demb
Σ
WE, O,A W (W W
dhead
ke(m
drean
w W E.j
(XWE + PP
demb
011(52n11 +1²)
XX)) = (ilSattnl (Xlatt((X000 (x:0];
/8lSattol (xtt: x'0]]
= (/w] g dampstc 700
ap
OP
@flaatnn [0lattn(X:0)) (X:0)); T (X):
Ki,E(X,
OW E
OW E
(XW + P)(
Σ
WT Wi F.aWank (
Σ
W W WWw)(W +TP*)
hE]l
+ oz w W
W W WiNl
Σ
w WoN
hE(M
bE(M)
+ 01 WE,(
Σ
T..
W v.
Σ
W W W. +TP")
kE(m
+ OZ w
Σ
W.,.
W v.
Σ
W
W Wv,)) +7PT)
hE(M
heim
dpcad 00, demb - H4000
6zg (26 +)2),
since only the first two terms do not vanish the embedding dimension and number of heads go infinity.
be a training set data points drawn from this task, where due to the structure of the template task each of the context strings is length-1 and we have x; = yt. We will test the model on data point (gtest yyeet), which does not appear in the test set: i.c., fnat = rest & Enp).
Training loss and testing loss Let (Z1, yi)...,, (In: Vn)
The training loss is given by
Liroin (0) 1 -0 Σ (faunn 0).
where f is the cross-entropy loss, and the test loss is given by
Creat()) e(farn (eest)) gleat).
Theorem J.2. For any learning rates no. no, TP, TE such that DCuria|| O(1) demb. dhead- and H we have (Peuut o(1). In other words, the error for generalization unseen symbols does not decrease during training for infinite-widue width transformers.
Proof. Consider training with gradient flow with learning rates TJ, MI, Mp, Bi on the parameters [W o.a.ee m) [W V.v. bhe W p, and W g, respectively. In the limit as demb - Do we have facn (X; 0p) + 0, so
aLtaan
1 ex, yr Üftnt (X: 0)
00
10=00
00
0=00

---

Published conference paper at ICLR 2024
So time 1 0, the training loss decreases as
aLtrain
Σ Σ (1/m - 63..) (1.mm-6.a)))
11=0
0t
n
(nv Ky (x xx) x)) no x:)
+ ty (X xx) Th jr x))).
So we must take N1 mV O(demb/ ([dems//heed)) up O(1), and Me = (11) for us to have OCrrain O() be bounded by constant that does not grow with demb, dhead, and H.
Under these choices of learning rates, the test loss on token weest which is not in the training dataset {Z1,.... Ln ,....,,n.... evolves as
Ctrst
Σ (1/m (1/m (lco)))))
1=0
7h
ml j.j'E(mm)
(nv K Kyr,v v (X X X + no KoKK o0 (X,, X
+ K p (x,, X leed + ne ji (x x
Σ Σ (1/m
dhead
T ) (62, + +
demb
Np720 2H 11 j (0s + 72))
(1/m- 61,2,) (1/m )) dhead B/ + Hno + mp + +ne) denb
71
ke(n] 5[|m]
C'
Σ
10l1/mmjja - - Σ (1/m (1/m
kE(m]JE|m)
= -C/m C/m + C/m = C/m 20.
On the other hand, now we consider the sat architecture where each head we replace WT. V.J. Wo, with W.. Wo. b,I, where by is trainable parameter and I gelembxXann the identity matrix:
f (X:0) = W 8 E
(Unembedding layer)
where
21
Σ (AA))es ek HE[M]
4 smax( (3 Zo W. K.b W Q. 2 Zo W Wo,+b,l) E B kxdemb
(Attention heads)
Zo XW +P R kd...
(Embedding layer)
Again, for the case of k = 1 that we consider, the network simplifies considerably
facn 0) WE( Σ W W v () +*P").
(27)
hE[M]
We initialize by for all h, so that the neural tangent kernels Kij,,, Ki.., Kirr KiJE the
same as above. Now we also have a neural tangent kernel for training the parameters [bn heem)
O[fatml xx0)]; affa (X'; (X::0), 0),
Ki
abn
Obp.
NE[M
(
X + + TT (XW E + 7 wEs
E.
E
1 Mema 81 01 65,25
48

---

Published conference paper at ICLR 2024
We prove that under this parametrization the test loss does decrease with training, which shows that adding this trainable identity scaling allows transformers to succeed at this task.
Theorem J.3. There a choice of learning rates TJb VV, noo TE, TIP such that demby dheads 00 we have @Cersin 0t
-==0= (11) and Bbresst
öt 1==== Q(1).
Proof. Training just the parameters {bi, heezn with learning rate Thb (keeping the learning rates V N1, Ap, lE = O, so the training loss decreases as
ÖLtrain Of
Th n2
Σ Σ
(1/m-d KK X,),
(PE(N)
should take rib ()/HH) for the train loss have derivative on the order of (()). The test loss decreases as:
iLiest
1
Σ
Σ
(1/m - 1or--jjj. )1/m - a K K ( X
1=00
7
n
iE j.'e(m)
HT
Σ
Σ
(1/m 88.n (1/m - 8g stere gra
x [i] E.
Hij
(1/m 1)(1mm- 1)
71
11 1/m)
211),
for 1hb (HH), as demb, H + oc.
49


Reviews:
[{"review_id": "MCQ34LT6VP", "reviewer": "Reviewer_5tFF", "summary": "This work studies the ability / inability of transformers to learn certain symbolic reasoning tasks, specifically real valued (“a=+1, print(a)” -> +1) and symbolic valued (“a = `q’ , print(a)” => q) template tasks. To study whether models can learn a rule for an arbitrary new input symbols, they invoke a slightly different concept of generalization under distribution shift known as “generalization to the unseen.” This measures the intuitive notion of generalization of a rule to arbitrary inputs rather than inputs from the same distribution used for training (which is the usual notion of generalization in ML).  For both of these symbolic tasks, the authors argue that the permutation symmetry of multi-layer perceptrons (MLPs) trained with SGD prevents them from generalizing to unseen symbols (Appendix I). For transformer architectures, an analysis of the initial feature kernel under random Gaussian weights reveals that transformers can generalize on real-valued symbolic reasoning tasks by training the readouts only. The reason for this generalization follows from an approximation of the empirical kernel matrix as block diagonal with each block representing the inputs from a single template: K(x_a, x_{a’}) = K(x_b,x_{b’}) for x_a, x_b from template 1 and x_a’ , x_b’ from template 2. This approximation becomes increasingly valid as the “data diversity” metric introduced by the authors increases. This approximation allows the authors to work with a # template x # template gram matrix N, which they show to be non-singular. The importance of having a well-conditioned the N matrix suggests a possible improvement to transformers where an identity matrix is added to the weight innerproducts WK WQ -> WK Wq +  a * Identity, etc. The authors show this improves generalization in their task. For symbolic tasks, the authors show that the gradient descent updates are orthogonal to the directions which would reduce the loss when the embedding dimension is large. They provide experiments verifying that large embedding dimension worsens performance on copy tasks.", "questions": "1. Do the sample complexity results / rank of N matrix change if you allow the network to learn its internal weights rather than just the readout weights? \n2. Do the authors have any empirical evidence showing the emergence of this block diagonal structure in the NNGP kernel at initialization (I imagine this would be easy to compute)? The block-diagonal approximation seems reasonable, but some empirical support + visualization could be useful. A similar comparison of the predictor from kernel regression and the average of the labels over a template could also be useful. \n3. Do the authors have a sense of how wide the transformer must be compared to samples n, templates r, etc for the theory to be accurate (N non-singular etc)? Wouldn’t the kernels become singular at some finite width? What if every layer is trained at large width (NTK vs NNGP)? Do real networks (which learn features) have to be as wide as kernel analysis would suggest? \n4. The perplexity improvements on GPT + Wikitext seem rather small. Does this suggest that transformers trained in practice evolve to have well conditioned N matrices, removing the need to add an explicit identity matrix? Is this related to the pronounced diagonals in Wv Wo the authors report in the next section? Or is the idea that real data behaves differently than the proposed template tasks? \n5. In Theorem 1.1 it may be helpful to clarify that the gradient flow is only performed on the last layer. \n\n\nSmall typos\n1. Figure 3b caption should read W_V W_O^T  (transpose dropped in second writing)\n2. Page 28 d/dx in the Wronskian should be d/dt", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This paper studies an important problem of generalization to unseen symbols in transformer architectures. Much of the contribution is formalizing a set of symbolic reasoning tasks which can be analyzed in different architectures. The authors give a theoretical analysis that suggests implementation tweaks (adding the identity matrices) which improve performance on real-valued template tasks and preserve good performance on natural language modeling. They also give an important separation result between multi-layer perceptrons and transformer architectures and explain why the transformers can learn the real-valued regression tasks at large width. Lastly they observe that large embedding dimension can be harmful in transformers on copy tasks. I appreciate the effort the authors took in performing experiments (several in the Appendix) to evaluate their theoretical claims, which are often lacking in theory papers.", "weaknesses": "Though the paper provides some generic results about architectures for MLPs, the main result for transformers relies on operating the model in the kernel regime. It is unclear at the moment if any of the results would change (1) at finite width or (2) in the feature learning regime (see questions below). Further, the proposed theory bounds the generalization error in terms of the data diversity metric, which does not have an obvious scaling with samples n, making it hard to reason about the number of samples needed to attain good generalization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kgr1ZBgXff", "reviewer": "Reviewer_gjp1", "summary": "In this work, the authors aim to establish and understand the abilities of self-attention architecture to solve variable binding tasks. I find the targeted problem both valuable and exciting---indeed, variable binding is a core intelligence capability and yet we have not seen sufficient work in the community targeting it as a direct benchmark for understanding limitations of existing models. I think this work is a step in the right direction therefore.", "questions": "The proposed changes based on the theoretical models involve addition of a scaled identity matrix to parts of the attention matrices. This looks, at a surface, like addition of a residual connection in fact. Can the authors clarify this further? I found the discussion of this point to be quite unclear.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "As noted in the summary, I find the targeted problem valuable and exciting. The theoretical framework itself involves several assumptions that are arguably impractical (making this essentially a kernel regression take on LLMs paper), but nonetheless the tasks established earlier in the paper are very well described and the results are accordingly insightful.", "weaknesses": "1. I think the authors are over-claiming their results at several points. The simplified architecture that is studied in this work, i.e., a self-attention + MLP model, is not justified to be labeled a Transformer. Residual connections play a huge role in both optimization stability and, e.g., some of the seemingly emergent abilities of Transformers. In this sense, I would argue the paper is actually focused on understanding inductive biases and limitations of a feedforward, MLP model with self-attention. This is totally fine, but the over-claiming can yield, in my opinion, an incorrect takeaway that Transformers broadly cannot perform the variable binding tasks. It may also be worth mentioning that the results focus on a inference with a single pass through the model, i.e., tools like in-context learning or chain-of-though can impact the results. The latter is just a nitpick to better contextualize the paper's findings.\n\n2. Lack of experimental details: While the authors often state that their theoretical claims focus on a simplified architecture, it is unclear to me if the experiments involve this simplified architecture or an actual Transformer, i.e., with residual connections, is used. Generally, I think the experimental details can be improved.\n\n3. Relation with prior theoretical works on inductive biases of self-attention: A specific work that addresses a similar question as this paper is Edelman et al. [1]. Therein, the authors show a self-attention model can express sparse boolean circuits and performs implicit variable creation during forward inference. I think the work deserves thorough discussion in the current paper's related work. A few more empirical works also exist on this problem, e.g., Davies et al. [2], that will be worth discussing.\n\n4. Discussion with Systematic generalization: I would argue the tasks that the authors studied are extremely related to what NLP community has studied under the label of systematic generalization before. I think such papers warrant discussion as well (e.g., see [3]).\n \n\n\n[1] Inductive Biases and Variable Creation in Self-Attention Mechanisms, Edelman et al., 2022.\n\n[2] Discovering Variable Binding Circuitry with Desiderata, Davies et al., 2023\n\n[3] Compositionality Decomposed: How do Neural Networks Generalise?, Hupkes et al., 2020", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dgI4VD7mg4", "reviewer": "Reviewer_jCDP", "summary": "This manuscript studies the ability of transformers to reason with abstract binding --- or symbol binding. The authors contrast MLPs and transformers and argue how MLPs trained with gradient descent are incapable of abstract reasoning. They then show how transformers on the other hand are capable of abstract reasoning given lots of data. Finally, they propose an a transformer variant that generalizes faster to the symbol binding task and verify the same empirically.", "questions": "1. The proof for the failure of MLPs relies on taking an expectation over $\\theta_t$ obtained through gradient descent. Does this also hold for a single run of an MLP training. Is it possible for MLPs to converge some of the time (and not with probability at least $1 - \\delta$ )?\n2. The proof seems to rely on the kernel ridge regression estimator. Does this mean that representation learning is not critical for symbol binding?\n3. If the data diversity is small (data follows a Zipf distribution), then do these results hold? In practice we expect datasets to follows such distributions and I am unsure if I am misunderstanding the theorem incorrectly. \n4. Are there technical reasons for using the kernel regression with gradient flow in section 4, but the cross-entropy loss with next-token prediction for section 5? Is it possible to conduct an analysis similar to section 4 if we consider next token prediction?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper was a great read! The authors conduct a thorough theoretical analysis with practical insights suggesting changes to the transformer architecture. The work studies if transformers are capable of symbol binding, which is very important question and of broad interest to the community. The paper presents a diverse set of results, but makes them easy to understand to a broader audience too.\n\nThe template tasks are elegant and easy to understand. It grounds the problem in a concrete scenario that can be studied. The authors present concrete insights for real labels and symbolic labels and identify different sample complexity behaviors for both scenarios. The work identifies important factors like data diversity or transformer embedding size that change the convergence behavior.\n\nThe authors use their theory to develop practical insights for developing transformer architectures. In particular, they show experiments on a synthetic dataset and on Wikitext to validate their proposed modification.", "weaknesses": "The theory for the transformers seems to rely on training only the final layer of the transformer. Does that mean that there are other architectures that can also do symbol binding? It is unclear which aspect of transformers are important for symbol binding: is the attention, the MLP, the invariance properties or something else? It isn't entirely clear to me how \n\nA limitation of most theory that uses the PAC framework often has vacuous bounds. While the paper derives asymptotic trends, it is unclear if they will predictive of what happens in practice as a result. However, I also acknowledge that this is the case for lot of theorems about deep networks. It is unclear if this probabilistic framework is the right tool for analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LK6chKsVwT", "reviewer": "Reviewer_xzui", "summary": "The paper investigates the \"variable-binding\" capability of transformers from a theoretic perspective, which refers to the generalization ability to handle unseen symbols during training. To formulate the variable-binding ability, the authors propose a framework called template tasks. Then the authors prove that a simplified transformer block can generalize to unseen input symbols when the output is a binary label, but cannot generalize when the output label is also dependent on novel input symbols. Finally, the authors propose a simple mechanism to augment the attention layer in transformers which is empirically shown to be effective.", "questions": "As I mentioned in the weakness section, I am especially interested in authors' opinions about how to extend the main proof in the paper to more practical transformer architecture with residual connections and layer normalization.", "limitations": "", "rating": 6, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The author provides a theoretical formulation of variable-binding tasks called template tasks.\n- Proofs are provided for the capability and limitations of a simplified transformer architecture on two instantiations of template tasks. \n- Extensive empirical experiments are conducted on the template task framework.", "weaknesses": "My main concern lies in the assumption of the proof. Throughout the whole paper, the authors assume a simplified depth-1 transformer architecture without residual connections and layer normalization. However, residual connections and layer normalization are known to be crucial for the performance of transformers [1][2]. The simplification limits the applicability of the main theoretical conclusions to more practical settings. On the other hand, all empirical results use a vanilla depth-2 transformer without modification. This makes it hard to draw direct parallels to justify the practical implications of the theoretical results.\n\n[1] On Layer Normalization in the Transformer Architecture. ICML 2020\n[2] Transformer Feed-Forward Layers Are Key-Value Memories. Mor Geva, et al. EMNLP 2021", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "T80D4k3H6r", "reviewer": "Reviewer_Fsbi", "summary": "This work presents an analysis of the abstract reasoning capability of the transformer architecture. The analysis yields a novel modification that improves out-of-distribution reasoning in symbolic tasks. The proposed modification also yields improved performance in an LLM (GPT-2) on a more realistic task.", "questions": "I found it interesting that GPT-2 already contains heads that implement a strategy similar to the modifications proposed in this paper. A preliminary investigation of GPT-4 suggests that it has a very strong ability to perform the abstract reasoning tasks from this paper based on just a few in-context examples. Do the authors expect that larger models (e.g. GPT-3 and GPT-4) may be able to perform these tasks without the proposed modifications, and if so how does that relate to the theoretical results? Do the authors expect that this capability would similarly depend on the emergence of heads that implement a diagonal attention strategy? Does the analysis in this paper have anything to say about the ability to learn these kinds of abstract reasoning tasks through *in-context* learning (as opposed to direct training of the model parameters)?\n\nI am also curious how the proposal in this work might relate to the idea of 'induction heads' [1]. Intuitively these seem to be related, in that both exploit the inductive biases of transformers to perform abstract tasks involving copying.\n\n[1] Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., ... & Olah, C. (2022). In-context learning and induction heads. arXiv preprint arXiv:2209.11895.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper presents a formal analysis of OOD symbolic reasoning in both transformers and MLPs, finding that transformers, but not MLPs, are theoretically capable of a form of abstract symbolic reasoning, given sufficient training data diversity. This provides interesting insights into the source of reasoning capabilities in LLMs, and contrasts with the common view that their capabilities are due exclusively to scale. This sort of analysis is less common in reasoning papers, so it is a very welcome contribution.\n- The analysis yields a novel modification to transformers that further improves their abstract reasoning capabilities, improving sample efficiency on one class of tasks, and enabling generalization on another. \n- The proposed modification also yields improved performance in a pretrained LLM (GPT-2) on a realistic task (wikitext).", "weaknesses": "I think this paper provides a strong contribution. My comments are primarily about clarification and presentation:\n- The abstract and introduction refer to 'transformer large language models'. However, if I understand correctly, the experiments in Figures 1 and 2 are performed on transformers trained from scratch, not on LLMs (though an LLM is tested later on). If this is correct, I think it would be clearer to simply refer to transformers (rather than LLMs) in the abstract and introduction.\n- It would be helpful if an intuitive characterization of the proposed modifications ($aI$ and $bI$) could be provided earlier in the paper, along with some sense of why this modification helps.\n- Do the wikitext experiments involve OOD generalization in any sense (in the same way that the reasoning tasks do)? If not, I wonder whether this partly explains the smaller relative improvement from the modified models vs. GPT-2.\n- I was somewhat confused by the distinction between GPT-2 and 'GPT-2 pretrained' -- isn't GPT-2 by default pretrained (since this term refers to the trained model)? Just to clarify, do the results in Figure 4 reflect pretrained GPT-2, which is then fine-tuned on the wikitext datasets, or both of the models in this figure are trained from scratch only on wikitext?\n\nMinor comment:\n- The legend is cut off in figure 3a. In general the figure text is very small and difficult to read.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)