Vision Transformers provably learn spatial structure
Authors: Samy Jelassi
Michael Eli Sander
Yuanzhi Li
Venue: NeurIPS 2022 Accept (2022)
Keywords: deep learning theory
vision transformers
implicit bias

Abstract: Vision Transformers (ViTs) have recently achieved comparable or superior performance to Convolutional neural networks (CNNs) in computer vision. This empirical breakthrough is even more remarkable since ViTs discards spatial information by mixing patch embeddings and positional encodings and do not embed any visual inductive bias (e.g.\ spatial locality). Yet, recent work showed that while minimizing their training loss, ViTs specifically learn spatially delocalized patterns. This raises a central question: how do ViTs learn this pattern by solely minimizing their training loss using gradient-based methods from \emph{random initialization}? We propose a structured classification dataset and a simplified ViT model to provide preliminary theoretical justification of this phenomenon. Our model relies on a simplified attention mechanism --the positional attention mechanism-- where the attention matrix solely depends on the positional encodings. While the problem admits multiple solutions that generalize, we show that our model implicitly learns the spatial structure of the dataset while generalizing. 
We finally prove that learning the structure helps to  sample-efficiently transfer to downstream datasets that share the same structure as the pre-training one but with different  features. We empirically verify that ViTs using only the positional attention mechanism perform similarly to the original one on CIFAR-10/100, SVHN and ImageNet.

Vision Transformers provably learn spatial structure
Samy Jelassi Princeton University edu
Michael E. Sander ENS, CNRS
Yuanzhi Li Carnegie Mellon University yuanzhi 1 @aandrew. cmu. edu
mi
Abstract
Vision Transformers (ViTs) have achieved comparable or superior performance than Convolutional Neural Networks (CNNs) in computer vision. This empirical breakthrough even more remarkable since, in contrast to CNNs, ViTs do not embed any visual inductive bias of spatial locality. Yet, recent works have shown that while minimizing their training loss, ViTs specifically learn spatially localized patterns. This raises central question: how do ViTs learn these patterns by solely minimizing their training loss using gradient-based methods from random initialization? In this paper, we provide some theoretical justification of this phenomenon. We propose spatially structured dataset and simplified VIT model. In this model, the attention matrix solely depends on the positional encodings. We call this mechanism the positional attention mechanism. On the theoretical side, we consider a binary classification task and show that while the learning problem admits multiple solutions that generalize, our model implicitly learns the spatial structure of the dataset while generalizing: we call this phenomenon patch association. We prove that patch association helps to sample-efficient downstream datasets that share the same structure as the pre-training one but differ in the features. Lastly, we empirically verify that ViT with positional attention performs similarly to the original one on CIFAR- 10/100, SVHN and ImageNet.
transfer to
Introduction
Transformers are deep learning models built self-attention [65], and in the past several years they have increasingly formed the backbone for
models in domains ranging from Natural Language Processing (NLP) [65, 23) computer vision [24], reinforcement learning [13,38], 38), program synthesis [5] and symbolic tasks [44]. Beyond their remarkable performance, several works reported the ability of transformers simultaneously minimize their training loss and learn inductive biases tailored to specific datasets c.g. in computer vision [55]. in NLP [10, 67] or in mathematical reasoning [73]. In this paper, we focus on computer vision where convolutions are considered be an adequate and biologically plausible inductive bias since they capture local spatial information [27] by imposing sparse local connectivity pattern. This seems intuitively reasonable: nearby pixels encode the presence of small scale features, whose patterns in turn determine more abstract features longer and longer length scales. Several seminal works [17, 24, 55] empirically show that although randomly initialized, the positional encodings in Vision transformers (ViTs) [24] actually learn this local connectivity: closer patches have more similar positional encodings, as shown Figure la. A priori, learning such spatial structure is surprising. Indeed, in contrast to convolutional neural networks (CNNs), VITs are not built with the inductive bias of local connectivity and weight sharing. They start by replacing an image by collection of D) patches (X....., Xp) dx each of dimension While each X; represents (an embedding of) a spatially localized portion of the original image, the relative positions of the patches X; in the image disregarded. Instead, relative spatial information is supplied through image-inded -ineeedddeeeneen positional encodings P = (Pı...., PD) E Rdx D Unlike CNNs, each layer of ViT then learns, via trainable self-attention, non-local set of filters that non- loneaiiy depend on both the values of all patches X, and their positional encodings p,.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

10
X1
X,
Xx
X
08
X5
X
Xg
X5
X
X9
X
06
Xg
X10
X1
X2
Xg
X10
X1
X22
0.4
X13
X14
X15
X1c
X1
X14
X15
X16
(1)
(2)
2
0.0
(a)
(b)
Figure 1: (a) Visualization of the positional encodings similarities P (pi,, Ps)) Piy)e. at initialization (1) and after training on Imagenel (2) using a VVT-ssmall-patch222222
[24]. We normalise the values P P between 1 and 1 and apply a threshold of 0.55. In contrast with the initial arrays that are random, the final ones show local connectivity patterns: nearby patches have similar positional encodings. (b) Partition of the patches into sets S;
Definition 2.1. Squares the same color belong to same set Se. We refer to (I) as a "spatially localized set" since all the elements a St are spatially contiguous. This the type of sets appearing in Figure la at the end of training. Definition 2.1 also covers sets with non-contiguous elements as (2).
Contributions. The empirical observation of Figure la sets central question: from a theoretical perspective, how do ViT's manage to learn these local connectivity patterns by simply minimizing their training loss using gradient descent from random initialization? While it known that attention can express local operations convolution [17]. it remains unclear how ViTs learn it. In this paper, we present a simple spatially--truucuree classification dataset for which is sufficient (but not necessary) to learn the structure in order to generalize. We also present simplified VIT model which we prove implicitly learns sparse spatial connectivity patterns when it minimizes its training loss via gradient descent (GD). We name this implicit bias patch association (defined in Definition 2.2). We prove that our ViT model leverages this bias to generalize. More precisely, we make the following contributions:
In Section 2, we formally define the concept of performing patch association, which refer to the ability of learning spatial connectivity patterns dataset.
Section 3, we introduce a structured classification dataset and simplified ViT model. This model is simplified in the sense that its attention matrix only depends on the positional encodings. We then present the learning problems we are interested in: empirical risk (realistic setting) and population risk (idealized setting) minimization for binary classification.
In Section 4, we prove that one-layer single-head ViT model trained with gradient descent synthetic dataset performs patch association and generalizes, in the idealized (Theorem 4. and realistic (Theorem 4.2) settings. We present a detailed proof, based on invariance and symmetries of coefficients in the attention matrix throughout the learning process.
In Section 5, we show (Theorem 5.1) that after pre-training in our synthetic dataset, our model can be sample-efficiently fine-tuned to transfer to a downstream datasct that shares the same structure as the source dataset (and may have different features).
On the experimental side, we validate in Section 6 that ViTs learn spatial structure in images from the CIFAR- 100 dataset, even when the pixels of the images are permuted. This result validates that, in contrast to CNNs, ViTs learn more general form of spatial structure that is not limited to local patterns (Figure 5). We finally show that our ViT model -where the attention matrix only depends the positional encodings- is competitive with the vanilla ViT on the ImageNet, CIFAR-10/100 and SVHNs datasets (Figure and Figure 7).
Notation. We use lower case letters for scalars, lower case bold for vectors and upper case bold for matrices. Given an integer D, we define [D] {1..... D). Any statement made "with high probability" holds with probability least 1/poly(d). Given vector I and k & d, define To F that implicitly depend on parameters and u, we often write FA.v to highlight its parameters. We use the asymptotic complexity notations when defining the different constants.
1 1a,
{a,,, ..... ,.,.) where ai ................................ aig are the k-largest elements. For a function
Related work
CNNs and ViTs. Many computer vision architectures can be considered as form of hybridization
between Transformers and CNNs. For example, DeTR [11] use a CNN to generate features that
are fed to Transformer. [25] show that self-attention can be initialized or regularized to behave like a convolution and [19, 30] add convolution operations to Transformers. Conversely, [8, 56, 7] introduce self-attention attention-like operations to supplement or replace convolution in ResNet-

---

Labelling function Definition 2.1 states that there exists labelling function that preserves the under- lying structure by applying the same function & to each St as in (1). For instance, when the sets Sy's like models. In contrast, our paper does not consider any form of hybridization with CNN, but rather a simplification of the original ViT to explain how ViTs learn spatially structured patterns using GD.
Empirical understanding of ViTs. A long line of work consists in analyzing the properties of ViTs, such as robustness [9, 54, 51] or the effect of self-supervision [12, 14). Closer to our work, some papers investigate why ViTs perform so well. [55] compare the representations of VIT' and CNNs and [50, 64] argue that the patch embeddings could explain the performance of ViTs. We empirically show in Section 6 that applying the attention matrices to the positional encodings which contains the structure of the dataset- approximately recovers the baselines. Hence, our work rather suggests that the structural learning performed by the attention matrices may explain the success of ViTs.
Theory for attention models. Early theoretical works have focused on the expressivity of attention. [66, 26] addressed this question the context of self-attention blocks and [21, 68, 34] for Transform- ers. On the optimization side, [76] investigate the role of adaptive methods in attention models and [59] analyze the dynamics of a single-head attention head to approximate the learning of a Seq2Seq architecture. In work, we also consider single-head ViT trained with gradient descent and exhibit a setting where it provably learns convolution- like patterns and generalizes.
Algorithmic regularization. The question address concerns algorithmic regularization which characterizes the generalization of an optimization algorithm when multiple global solutions exist in over wer-parametrrized models. This regularization arises in deep learning mainly due to the non- convexity of the objective function. Indeed, this latter potentially creates multiple global minima seattered in the space that vastly differ in terms of generalization. Algorithmic regularization appears in binary classification [60, 49, 16), matrix factorization [29, 3], convolutional neural networks 29,,77, generative adversarial networks [2], contrastive learning [70] and mixture of experts [15]. Algorithmic regularization is induced by and depends on many factors such as learning rate and batch size [28, 33, 41. 58, 47], initialization [1], momentum [39], adaptive step-size [42, 53, 20, 71, 78, 40], batch normalization [4, 32, 36] and dropout [61, 69]. However, all these works consider the case of ford-forward neural networks which does not apply to ViTs.
2 Defining patch association
The goal of this section is to formalize the way ViTs learn sparse spatial connectivity patterns. We thus introduce the concept of performing patch association for a spatially structured dataset.
Definition 2.1 (Data distribution with spatial structure). Let D be a distribution over RdxD D x (-1,1)
where each patch X = structured if
E RdxD has label y E {-1,1). We say that D is spatially
there exists a partition of [D] into L disjoint subsets i.e. [D] U Se with Se & D and |St] = C.
there exists a labeling function I satisfying Plys((X) 0] d1= and,
(x) Σ (X.)) ies,), where 6: R RE[T]]
IR is arbitrary function.
(1)
Test error
Examples. A partieular case for the sets S is the one of spatially localized sets as in Figure 1b-
Test errar
(I). In this case, we have D 16. C
and S = 21,2,5, 5, 6), S2 = (384.7.8). Sg
(9, 10, 13, 14): S4 (11, 15, 16). We empha- size that Definition 2.1 is not limited to spatially lo- calized sets and also covers non-contiguous sets as Figure lb-(2).
10% 5%
500
1000
Number of gradient descent steps
Figure 2: Left: Test error of the ViT on the convolution structured dataset. Upper Right: Grid displaying the input patches. Yellow squares represent spatially localized sets Se. Those sets are taken into account when com- puting the convolutional function f* Lower Right Learnt P. P looks random compared to upper one.
3

---

are spatially localized, f* can be one- me-hidden layer convolutional network. In this paper, we are interested in patch association which refers the ability of an algorithm identify the sets Sy's, ss and is formally defined as follow.
Definition 2.2 (Patch association for ViTs). Let DD be as in Definition 2.1. Let M: RdxD {-1,1} be a transformer and P(M) its positional encodings matrix. We say that M performs patch association D for all [L] and i E Se, have Tope (p!
(M) PP > St.
Definition 2.2 states that patch association is learned when for a given i E Se, its positional encoding mainly attends those of j such that i, E St. this way, the transformer groups the X, according to St just like the true labeling function. Definition 2.2 formally describes the empirical findings in Figure la-(2), where nearby patches have similar positional encodings. natural question is then: would ViTs really learn those Sy after training to match the labeling function f*? Without further assumptions on the data distribution, we next show that the answer is no.
ViT's do not always learn patch association under Assumption 1. We give negative answer through the following synthetic experiment. Consider the case where all the patches X, are i.i.d. standard Gaussian and f* is one-hidden layer CNN with cubic activation. The label y of any X is then given by y = sign( f (X)). As shown in Figure 2, one-layer ViT reaches small test error on the binary classification task. However, P P does not match the convolution pattern encoded in f*. This is not surprising, since the data distribution D is Gaussian, and thus lacks spatial structure. Thus, order to prove that ViTs learn patch association, we need additional assumptions on D, which we discuss in the next section.
3 Setting to learn patch association
In this section, we introduce our theoretical setting to analyze how ViTs learn patch association. We first define our binary classification dataset and finally present the VIT model
classify it. Assumption 1 (Data distribution with specific spatial structure). Let D be distribution as in Definition 2.1 and w* Rd be an underlying We suppose that each data data-point X is defined as follow
Uniformly sample an index ((X) from [L] and j S((x). X, -yu* + Ez, where ywu is the informative feature and Ej s.ddd. N(0,0"(1 w" T)) (signal set).
For é [L]]((X] and j St. X, = d + Es, where dz = 1 with probability q/2, -1 with same probability and O otherwise, and Ej isd. o20, - w w T)) (random sets).
To
keep the analysis simple, the noisy patches are sampled from the orthogonal complement of w" Note that D admits the labeling function
0
0
0
0
1
-
11
0
0
0
-1
I
f* (x) Ex((C) Thresholdo.oc (u* X,)), where Thresholde()) = 2 |2| C' and otherwise.
0
+1
0
O
0
+1
-1
0
We sketch a data-point of D in Figure Our dataset can be viewed extreme simplification of real-worke image datasets where there is a set of adjacent patches that contain a useful feature (e.g. the nose of dog) and many patches that have uninformative or spurious features e.g. the background of the image. We make the following assumption on the param- eters of the data distribution.
-1
-1
-1
Figure 3: Visualization of data-point X when the S, 's are spatially lo- calized. Each square depicts a patch X, and squares of the same color be- long to the same set St. "0" indicates that the patch does not have feature, "I" stands for feature 1 I-w* and "-1" for feature -1 w² The large red square depicts the signal set ((X). Although there are more "-1"'s than "+1" ""s,, the label X +1 since there are only "41"'s inside the signal set.
Assumption 2. We suppose that = poly(D), C = (olylog(d), q poly(C)) | 112 and o² = l/d. This implies C < D and & 1. Assumption 2 may
justified by considering a "ViT-base- patch Mncch1-222" model [24] on ImageNet. In this case, d = 384, D 196. a is set to have | | g is chosen so that there are more spurious features than informative ones

---

(low signal-to- noise regime) which makes the data non-linearty separable. Our dataset non- con-trivial to learn since generalized linear networks fail to generalize, as shown in the next theorem (see Appendix J for a proof).
Theorem 3.1. Let 'D be as in Assumption Let g(X) Φ - x,)) be a generalized linear model. Then, does not fit the labeling function i.e. (X)*( & 0] 1/8.
Intuitively, g fails to generalize because it does not have any knowledge on the underlying partition and the number of random sets is much higher than those with signal. Thus, a model must have a minimal knowledge about the Sy's in order to generalize. In addition, the following Theorem 3.2 states the existence of transformer that generalizes without learning spatial structure (see Appendix J for proof), thus showing that the learning process has priori no straightforward reason to lead to patch association.
Theorem 3.2. Let D be defined as in Assumption ([) X) & 0] d-w(1) but for all & E [T.] E St, Topo (p! .D _p ) S Ø.
1. There exists a (one-layer) transformer MM so that
(M)
Simplified VIT model. We now define our simplified VIT model for which show in Section 4 that it implicitly learns patch association via minimizing its training objective. We first remind the self-attention mechanism that ubiquitously used in transformers.
Definition 3.1 (Self-attention [6, 65)). The attention mechanism [6, 65] in the single-head case defined as follow. Let X Iddx data point and P Rd its positional encoding, The self-attee attention mechanism computes
1. the sum of patches and positional encodings ie.. X X the attention matrix A =QK where Q X We K = XW. Wq. WK E ladxd the score matrix Se E RxDxD with coefficients Sn exp(At../VO)) E. exp(Ar.r/Vd).
the matrix V = X Wv, where Wv E gdxd. finally outputs SA((X;P)) SV RdxD
In this paper, our ViT model relies on a different attention mechanism -the "positional attention"- that we define follows.
Definition 3.2 (Positional attention), Let X E RRuDD and PE RdxD the positional encoding. The positional attention mechanism takes input the pair (X; P) and computes:
1. the attention matrix A -KK where Q P PP K P and Wo, WK E Raxd
2. the score matrix S DxD with coefficients Sug exp(Ar./ E.- exp( h/t
3. the matrix V = =X Wv, where Wy E Rddxd
outputs (A P)) SV.
Positional attention isolates positional encoding P from data X: A encodes the dynamics of P and tracks whether patch association is learned. encodes the data-dependent part and monitors whether the feature is learned. Indeed, given its highly non-linear nature with respect to the input, directly analyzing self-attention is difficult. Yet, positional attention similar to self-attention. As this latter, positional attention is also permutation-inyarian and processes all tokens simultaneously. Besides, positional attention also computes a score matrix between the different tokens. This similarity matrix is also normalized in sparse manner with the Softmax operator. The only aspect that positional attention misses from self-attention is the fact that does not depend on the input. Nevertheless, we empirically show that our positional attention model competes with self-attention in Section 6. Lastly, we make the following simplification the parameters to ease our analysis.
Simplification 3.1. In the positional attention mechanism, ........ v] E RdxD where E Rdd Finally, we set A and e as trainable parameters. Besides, without loss of generality, we train all A for i # and leave the diagonals of A fixed.
set d = D, WK ID and Wq ID
which implies A PP P. We set Wv = [[o....,
In Simplification 3.1, we set WK k and Wq to the identity so that A PTP. This Gram matrix encodes the spatial patterns learned by the ViT shown in Figure la. Besides, since fitting the
5

---

labeling function requires to learn one feature w", it is sufficient to parameterize Wy with a vector v. Also, although A -PP and P trainable, we choose for simplicity to only optimize over A. Besides, we leave the Ai,'' fixed because Softmax is invariant under the uniform shift of the input. Under Simplification 3.1, our simplified ViT model then two attention layer with a single head:
D Σ
F(X)
D 2
Sty (u,
with S = -xp=:(AA,//V)) Σ exp(Arr/Vd), (T)
where a is an activation function. Since we aim to the simplest ViT model, we opt for a polynomial
activation i.e. (xx) = x² vz where p > 3 is odd integer and v 1/poly(d). Note that this choice of polynomial activation is common in the deep learning theory literature see e.g. [45, 1, 72] among others. The degree p is odd to make the ViT model compatible with the labeling function and strictly larger than because the data not linearly separable (Theorem 3.1). We add linear part in the activation function to ensure that the gradient is non-zero when v has small coefficients. With these simplifications, we formally prove that is able to learn patch association and generalize, in the two following settings.
Idealized and realistic learning problems. Given a dataset 2 = [ sampled from D, we solve the empirical risk minimization problem for the logistic loss defined by:
log (1 e ((IP((E))) O),
min A
Instead of directly analyzing (E), we introduce a proxy where minimize the population risk
min Ep log (1 + (==()) == L(A.v).
(P)
A,
We refer to (E) as the realistic problem while (P) as the idealized problem.
Algorithm. We solve (P) and (E) using gradient descent (GD) for T iterations. The update rule in the case of (P) for L [T] and i,j [D]
( 1 n x(A ())), (f+1)) (+) V/ (0)),
(GD)
where 77 > 0 the learning rate. A similar update may be written for (E). We now detail how to set the parameters in (GD).
Parametrization 3.1. When running GD on (P) and (E), the number of iterations is any T > poly(d)/n. We set the learning rate as 17 E (0, p The diagonal coefficient of the attention matrix are set for i E [D] A 1 GAID where a polyloglog(d). The off-diagonal coefficients of A and the value vector are initialized as:
1. Idealized case: (0) = a (0) x where (0) = (1-11 and A 0 for *
2. Realistic case: (0)) 100,w21a and A10) -N(0w where + j and w 1/ 1/poly(d).
We remind that in Simplification 3.1, we have A = PT If one initializes P ~ (100 (IDD))) then with high probability, ! pp
(0) 1 (( and ! (p.))
(0) O(oA/vD) for ()) p,
i A j. Since >> 1, it is then reasonable to set 1 Note that, also the idealized setting, we initialize (0) in span(w"), even though this latter should be unknown to the algorithm. We remind that the idealized case a proxy to ultimately characterize the realistic dynamics.
4 Learning spatial structure via matching the labeling function
As announced above, show that our ViT (T) implicitly learns patch association and fits the labeling function by minimizing the training objective. We first study the dynamics (P). Using the analysis in the idealized case, we then characterize the solution found in the realistic problem (E).

---

Test error
Cosine Similarity
Test srror
10%
0.5
6%
similarity
10
10'
10º
Number of gradient descent Steps
Number gradient descent steps
Figure 4: Illustration of Theorem 4.2. We consider the exact same setting (data generation, parameter settings...) as for the realistic case. From left to right, we first display in grey the tuples (i, such that (i,j) e St. We then plot the learned matrix A and that coefficients with high value exactly correspond to their grey scale counterpart the left plot. We also display test error and cosine similarity between w and v wart the number of training steps.
4.1 Learning process in the idealized case
In this section, we analyze the dynamics of (P). Our main result is that after minimizing (P), our model (T) performs patch association while generalizing.
Theorem 4.1, Assume that GD (P) for T iterations with parameters set in Parametriza- tion 3.1. With high probability, the ViT model (T)
1. learns patch association i.e. for all f [L] and Se. Tope (4 13-1 St.
2. learns the labeling function f i.e.
(X) > 0] > 1 o(1)
We now sketch the main ideas prove the theorem for which refer to Appendix D for complete proof.
Invariance and symmetries. In (P), take the expectation over D. Since (T) is permutation- invariant and the data distribution is symmetric, we can thus dramatically simplify the variables in (P). An illustration of this is the next lemma that shows that A can be reduced to three variables in (P).
Lemma 4.1. There exist B = AA. (0, p(()) E R such that for all t > 0:
1. for all i DD] A 3.
2. for all i,j E [D] such that i, E St. for some E [L]. A!9 7(0).
3. for all i,j [D] such that i St and j e Sm. for some [L]] with E * m, A!3 p(t).
Besides, using the initialization in Parametrization 3.1, we show that always lies in span w) Lemma 4.2. For all [T], there exists (1) such that w(*) w
In summary, Lemma 4.1 and Lemma 4.2 imply that instead of optimizing over A and in (P), can instead consider the scalar variables (t()) 7(()) and p(t), The remaining of this section consists in analyzing the dynamics of these three quantities.
Learning patch association. We first analyze the dynamics of 7/0) and pl)) To this end, we introduce the following terms:
e (t)
(1)
(+)
@ + (C - 1) + (D - C)eeees
Φ + (C - 1) +(D C)
en(t)
G = ((A + (C- ())())
@ (C 1)
+ (D -C)
Note that A0) (() and (+) respectively correspond to the coefficients the diagonal, those for which i.j Sp for some e [L] and all the other coefficients of the attention matrix S. Using these notations, we first derive the GD updates of (() and p(t).
Lemma 4.3. Let & T. The attention weights (0) and ((0) satisfy:
((+1) = (0))) + (oooyloog(d((a"))) ((0 (G(t))p1
| () (polylog(d)(o) (
1 D (
7

---

Lemma 4.3 shows that the increment of 7(() is larger than the one of pl)) Since 700)
p()) = 0,
this implies that ((t) > pl for all t > O. This observation proves the first item of Theorem 4.1. We now explain how learning patch association leads to v highly correlated with w*.
Event I: At the beginning of the process, the update of o(1) is larger than the one of A which implies that only v(4) updates during this first phase. We show that (x = (v(") w=) increases until time To 0 where it reaches some threshold (Lemma D.2). At this point, the model is nothing else than generalized linear model that would not generalize because there are much more noisy tokens than signal ones (see Theorem 3.1).
Event II: During this phase, the attention weights must update. Indeed, assume by contradiction
that the A stay around initialization and that v(t) is optimal i.e. v(t) = a where ( >> 1. Then, the predictor g we would have is
D D 9 g(X) Σ7 OO
D ΣΣ
D
ins (w* X,)
(2)
Such predictor g would yield high population loss because there many more data with random labels (q)) poly(C)) than with the exact label. Therefore, Al9's 's start to update. The gradient increment for 70() (which corresponds to and j in the same set St) is much larger than the one for p( (Lemma 4.3). Thus, (00) increases until a time Ti [To: 7] such that (Ti) max{=[T] l
Event III: Because we have y (T)) > maX (, we again have a (1+1) (+ as Phase (Lemma 1.I Thus, (t) increases again until the population risk becomes a o(1).
Main insights of our analysis. Our mechanism highlights two important aspects that are proper to attention models:
because of the initialization and the data structure, we have patch association for any time I (Lemma 4.3).
our ViT model uses patch association to minimize the population loss (Event III). Without patch association, the model would only be generalized linear model that does not minimize the loss.
4.2 From the idealized the realistic learning process
The real learning process differs from the idcalized one in that we have finite number of samples and initialize both Á and è as Gaussian random variables. Using polynomial number of samples, we show that (T) still learns patch association and generalizes.
Theorem 4.2. Assume that we run GD on (E) for T iterations with parameters set as in Parametriza- tion 3.1. Assume that the number of samples is N poly(d). With high probability, the model
1. learns patch association i.e. for all f [L.] and i E St. Tope
(1 )-1 = St.
Am,,am 2. fits the labeling function Le. Po[s (X) Am, soy (X) > o(1)
Similarly to [46], the proof introduces "semi-realistic" learning process that a mid-point between the idealized and realistic processes. We show that (TT) and @(T) are close to their semi-realistic counterparts Appendix E for a complete proof. Figure 4 numerically illustrates Theorem 4.2.
5 Patch association yields sample-efficient fine-tuning with ViTs
fundamental byproduct of our theory is that after pre-training on dataset sampled from DD our model (T) sample-efficiently transfers datasets that are structured D but differ in their features.
Downstream dataset. Let D downstream data distribution defined as in Assumption such that its underlying feature is t with || l and w* potentially different from w* other words, the downstream D and source D distributions share the same structure but not necessarily the same feature. We sample downstream dataset Ż
{XX M from D.

---

Test accuracy (Permuted CIFAR-100 100)
80-
Resnet
(1)
(9)
(b)
ViT
50-
Rasaline
16
32
Shuffle Grid Size
(1)
(a)
(c)
Figure 5: (a): Test accuracy obtained with ViT (patch size 2), ResNet-18 and VGG-19 permuted (in solid lines) on original (in dashed lines) CIFAR-I00. While convolutional models are very sensitive to permutations, the VIT performs equally whether the dataset is permuted or not. (b): (2) CIFAR-100 image (1) and Permuted CIFAR- 100 image when shuffle grid size is 2 (2). 4 (3) and 8 (4). (c): (1-2) Visualization of positional encoding similarities after training ViT (patch size
permuted CIFAR- 100 (shufflc grid size 2).
Here, we display P, where i is some fixed index and reshape such vector into matrix
16. We observe
that these similaritics (1-2) do not have any spatially localized structure. However, when applying the inverse of the permutation, we recover spatially localized patterns in i (1'-2").
Learning problem. We consider the model (T) pre-trained as in subsection 4.2. We assume that A kept fixed from the pre-trained model and we only optimize the value vector ò to solve:
min Σ log (1+ (1+ = Z()). N
1
(E)
We GD (E) with parameters set in Parametrization 3.1 except that the A.j 's are fixed and ((0) NN(0 (Ia) with 1/poly(d). Our main results states that this fine-tuning procedure requires a few samples to achieve high test accuracy in D. In contrast, any algorithm without patch association needs large number of samples to generalize.
Theorem 5.1. Let Å be the attention matrix obtained after pre-training as in subsection 4.2. Assume that we run GD for T iterations on (E) to fine-tune the value vector. Using N & polylog( D) samples, the model (T) transfers to D i.e. Se..TPpo (F))(()FF)))))))) A, (X) 11
Theorem 5.2. Let A: R2xDD 1-1,13 13 be a binary classification algorithm without patch associa- tion knowledge. Then, it needs DQ(1)) training samples to get test error & o(1) on D.
The proofs of Theorem 5.1 and Theorem 5.2 are Appendix F. These theorems hightlight that learming patch association is required for efficient transfer. We believe that they offer a new perspective on explaining why ViTs are widely used in transferring to downstream tasks. While it possible that ViTs learn shared (with the downstream dataset) features during pretraining, our theory hints that learning the inductive bias of the labeling function is also central for transfer.
Numerical experiments
In this section, we first empirically verify that ViTs learn patch association while miniziming their training loss. We then numerically show that the positional attention mechanism competes with the vanilla one on small-scale datasets such as CIFAR- ILSVRC-2012 ImageNet [22]. For the small datasets, we hidden/MLPP dimension 384. For ImageNet, we train a "ViT-tiny-pater 6-224" [24]. Both models are trained with standard augmentations techniques [18] and using AdamW with cosine learning rate scheduler. We run all the experiments for 300 epochs, with batch size 1024 for Imagenet and 128 otherwise and average our results over seeds. We refer to Appendix A for the training details.
[43], SVHN [52] and large-scale ones such ViT with 7 layers, 12 heads and
VITs learn patch association. We consider the CIFAR-100 dataset where we divide each image into grids size s x pixels. For a fixed s {2, 4, 8, 16, 32), we permute the grids according to n. to create the permuted CIFAR-100 dataset. We call the grid shuffle size. Figure 5b-(1) shows CIFAR- 100 image and its corresponding shuffling in the permuted CIFAR-101 dataset Figure 5b-(2-3-4). We train ViT and CNNs ResNet18 [31] and VGG-19 [57] on the permuted CIFAR- 100 dataset.

---

Training loss (Imagenet)
Accuracy (Imagenet)
For the VIT, we set the patch size to 2, although this is sub-optimal in terms of accuracy, because the patch size needs to stay smaller equal to 8. Indeed, intuitively, when we permute the grids in Figure 5b, we lose the local aspect of the spa- tial structure and create new sets S s and a new labeling function f*. Figure 5a reports the test accuracy of these three models for different values of s. When is small, the image does not have a coherent structure e.g. Figure 5b-(2) and thus,
66.3
VIT
6-
Durs
5-
30
3.0
3.5
100
200
300
100
200
300
Number spachs
Number apochs
Figure 6: Training loss (1) and test accuracy (2) obtained using a ViT-tiny-patch 16-2h 16-224 on Imagenet. VIT using positional attention (Ours) gets 68. 98.9% test accuracy while vanilla ViT (ViT) gets 71.9%.
CNNs struggle to generalize. As increases e.g. Figure 5b-(4), the information inside a patch is meaningful and thus, the CNNs well-perform. Unsurprisingly, since ViTs are permatation invariant, their performance remains unchanged for all 8 see Figure 5a. Despite this change, we verify that the ViT is able to recover the new Sy's: feed the ViT with the shuffled pear image Figure 5b-(2)) and consider for some i the similarity matrix P: P. We see that it does not exhibit a local spatial structure in Figure (c-(1,2). We then apply म. to p: P and observe that we recover the spatially localized patterns Figure 5c-(1' (")) This experiment highlights that ViTs do not just group nearby pixels together as convolutions. They learn a more general spatial structure, in accordance to our theoretical results.
ViTs with positional attention competitive. We numerically verify that ViTs using positional attention compete with those with vanilla attention. In Section 3, we introduced positional attention to define our theoretical learner model. Figure 6 and Figure 7 show that ViTs using positional attention compete with vanilla ViTs on range of datasets. These experiments strengthen our intuition that for images, having an attention matrix that only depends on the positional encodings is sufficient have a good test accuracy.
CIFAR-10
CIFAR-100
SVHN
99
91
68.5
98
50
67.5
97
VIT
Ours
VIT
Ours
VIT
Ours
Figure 7: Test accuracy obtained with a ViT using vanilla attention (ViT) and positional attention (Ours) CIFAR-10 (1), CIFAR- 100 (2) and SVHN (3). Our model competes with the vanilla VIT. Patch size 4 and average over 10 seeds for this experiment.
Conclusion, limitations and future works
Our work is a first step towards understanding how Transformers learn tailored inductive biases when trained with gradient descent. Our analysis heavily relies the positional attention mechanism that disentangles patches and positional encodings. In practice, self-attention mixes these two quantities. An interesting direction is to understand the impact of patch embeddings on the inductive bias learned by ViTs. Moreover, our experiment on the Gaussian data shows that ViTs do not always learn the correct inductive bias under Definition 2.1: characterizing the distributions under which ViTs recover the structure of the function is an important question. Lastly, this work also paves the way to many extensions beyond convolution. For example. can ViTs learn other inductive biases? What are the inductive biases Iearnt by Transformers in NLP? Answering those questions is central to better understand the underlying mechanism of attention.
Acknowledgments and Disclosure of Funding
The authors would like to thank Boris Hanin for helpful discussions and feedback on this work.
References
[1] Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:20/2.09816, 2020.
[2] Zeyuan Allen-Zhu and Yuanzhi Li. Forward super-resolution: How can gans learn hierarchical generative models for real-world distributions. arXiv preprint
2021.
10

---

[3] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization, arXiv preprint arXiv:1905. 2019.
[4] Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. arXiv preprint arXiv: 18Xiv:/8/2.0398// 2018.
[5] Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoe Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07733: 2021.
[6] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv: 1409.0473, 2014.
[7] Irwan Bello. Lambdanetworks: Modeling long-range interactions without attention. arXiv preprint arXiv:2102.08862, 2021.
[8] Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens, and Quoe V Le. Attention augmented convolutional networks. In Proceedings of the IEEE/CVF international conference on computer vision, pages 3286-3295, 2019.
[9] Srinadh Bhojanapalli, Ayan Chakrabarti, Daniel Glasner, Daliang Li, Thomas Unterthiner, and Andreas Veit. Understanding robustness of transformers for image classification. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10231-10241, 2021.
[10] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1900 2020.
[11] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision, pages 213-229. Springer, 2020.
[12] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in sulf-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9650 9660,,660,00 2021.
[13] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Adiyya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning vía sequence modeling. Advances in neural information processing systems, 34, 2021.
[14] Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages
2021.
[15] Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding mixture of experts in deep learning. arXiv preprint arXiv:2208. 2,200, 2022.
[16] Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory, pages 1305-1338. PMLR, 2020.
[17] Jean- lean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. On the relationship between self-attention and convolutional layers. arXiv preprint arXiv: 1.005584, 2019.
[18] Ekin D Cubuk, Barret Zoph. Dandelion Mane, Vijay Vasudevan, and Quoc Le. Autoaugment: Learning augmentation policies from data. arXiv preprint arXi: 2018.
[19] Zihang Dai. Hanxiao Liu, Quoc V Le, and Mingxing Tan. Coatnet: Marrying convolution and attention for all data sizes. Advances in Neural Information Processing $ystems,,4.3965555 34:3965-3977, 2021.
[20] Amit Daniely. Sgd learns the conjugate kernel class of the network. In Advances in Neural Information Processing Systems, pages 2422-2430, 2017.

---

[21] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Uni- versal transformers. arXiv preprint arXiv:1807.03819, 2018.
[22] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages 248 255.. leee, 2009.
[23] Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 2018.
[24] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11 1929, 2020.
[25] Stéphane d' Ascoli, Hugo Touvron, Matthew L Leavitt, Ari Morcos, Giulio Biroli, and Levent Sagun. Convit: Improving vision transformers with soft convolutional inductive biases. In International Conference on Machine Learning, pages 2286-2296. PMLR, 2021.
[26] Benjamin I. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. arXiv preprint (0.:://10.00900 2021.
[27] Kumihiko Fukushima. Neocognitron for handwritten digit recognition. Neurocomputing. 51:161- 180, 2003.
[28] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv: 1766000776 2017,
[29] Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in Neural Information Processing Systems, 31, 2018.
[30] Jianyuan Guo, Kai Han, Han Wu, Chang Xu, Yehui Tang, Chunjing Xu, and Yunhe Wang. Cmt: Convolutional neural networks meet vision transformers. arXiv preprint arXiv:2107.06663, 2021.
[31] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference computer vision, pages 630-645. Springer, 2016.
[32] Elad Hoffer, Ron Banner, Itay Golan, and Daniel Soudry. Norm matters: efficient and accurate normalization schemes in deep networks, 2019.
[33] Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. arXiv preprint 17Xiv:1705.08741 2017.
[34] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nogp and ntk for deep attention networks. In International Conference on Machine Learning, pages 4376-4386. PMLR, 2020.
[35] Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In European conference on computer vision, pages 646-661. Springer, 2016.
[36] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint 102.01150. 2015.
[37] Meena Jagadcesan, Ilya Razenshteyn, and Suriya Gunasckar. Inductive bias of multi-channel linear convolutional networks with bounded weight norm. In Conference on Learning Theory, pages 2276-2325. PMLR, 2022.
[38] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34, 2021.
12

---

[39] Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 9965-10040. PMLR, 17-23 Jul 2022.
[40] Samy Jelassi, Arthur Mensch, Gauthier Gidel, and Yuanzhi Li. Adam is no better than normal- ized SGD: Dissecting how adaptivity improves GAN performance, 2022.
[41] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large- training for deep learning: Generalization gap and sharp minima. arXiv preprint 1al9:1609.48800 2016.
[42] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 11X1::1412.6980, 2014.
[43] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[44] Guillaume Lample and François Charton. Deep learning for symbolic mathematics. arXiv preprint arXiv: 19120114// 2019.
[45] Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over- parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Thheory pages 2-47. PMLR, 2018.
[46] Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over wer-parametrized two-layer relu neural networks beyond ntk. arXiv preprint arXiv:2007.04596 2020.
[47] Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. arXiv preprint arXiv:1907.0455) 2019.
[48] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 1rX1v:1711.05101, 2017.
[49] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv: 1906.0589, 2019.
[50] Luke Melas- Kyriazi. Do you even need attention? a stack of ferw-forward layers does surpris- ingly well on imagenet. arXiv preprint arXiv:2105.02723. 2021.
[51] Muhammad Muzammal Naseer, Kanchana Ranasinghe, Salman H Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Yang. Intriguing properties of vision transformers. Advances in Neural Information Processing Systems, 34, 2021.
[52] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
[53] Behnam Neyshabur, Ruslan Salakhutdinov, and Nathan Srebro. Path-sgd: Path-normalized optimization in deep neural networks. arXiv preprint arXiv:1506.02617, 2015.
[54] Sayak Paul and Pin-Yu Chen. Vision transformers are robust learners. arXiv preprint arXiv:2105.0758/ 2(3), 2021.
[55] Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Dosovitskiy,
Do vision transformers see like convolutional neural networks? Advances in Neural Information Processing Systems, 34, 2021.
[56] Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jon Shlens. Stand-alone self-attention in vision models. Advances in Neural Information Processing Systems, 32, 2019.
[57] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556 2014.
13

---

[58] Samuel L. Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe V. Le. Don't 't decay the learning rate, increase the batch size, 2018.
[59] Charlic Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.0760. 2021.
[60] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1):2822-2878, 2018.
[61] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way prevent neural networks from overfitting. The Journal of Machine Learning Research, 15(1):1929-1958 2014.
[62] Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 281 2818-2826, 2016.
[63] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pages
PMLR, 2021.
[64] Asher Trockman and J Zico Kolter. Patches are all you need? arXiv preprint arXiv:2201.0977 097999 2022.
[65] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neurul information processing systems, 30, 2017.
[66] James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. A mathematical theory of attention. arXiv preprint WKXi:2007.008777 2020.
[67] Alex Warstadt and Samuel R Bowman. Can neural networks acquire structural bias from raw linguistic data? arXiv preprint arXiv:2007.06761 2020.
[68] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. arXiv preprint arXiv:2107.13163, 2021.
[69] Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout, 2020.
[70] Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th Inter- national Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11112-11122. PMLR, 18-24 Jul 2021.
[71] Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nathan Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. arXiv preprint 17X1v: 2017.
[72] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages
PMLR, 2020.
[73] Yuhuai Wu, Markus N Rabe, Wenda Li, Jimmy Ba, Roger B Grosse, and Christian Szegedy. Lime: Learning inductive bias for primitives of mathematical reasoning. In International Conference Machine Learning, pages 1251-11262. PMLR, 2021.
[74] Sangdoo Yun, Dongyoon Han, Scong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In Proceedings of the IEEE/CVF international conference computer vision, pages 6023-6032, 2019.
[75] Hongyi Zhang, Moustapha Cisse, Yann Dauphin, and David I Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint
2017.
14

---

[76] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383-15399 2020.
[77] Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmentation. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 13001-13008 2020.
[78] Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. Understanding the generalization of adam in learning neural networks with proper regularization. arXiv preprint arXiv:2108. 11371, 2021.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes] See Section and Section
(b) Did you describe the limitations of your work? [Yes] See Section 6.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] This is theory paper.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section 3.
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URLY? [Yes] See supplemen-
tary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix.
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes] See Section 6
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
15


Reviews:
[{"review_id": "zQDhA-qWVoS", "reviewer": "Reviewer_gkYb", "summary": "Vision Transformers\n\nThe paper focuses on trying to explain how Vision Transformers learn the spatial locality biases on images. To do so, authors mathematically analyze the learning of convolution-like structures on ViT and test their idea with some synthetic datasets based on CIFAR.\n\n", "questions": "- The test of L116-129 is not so clear to me. Where is the cubic activation applied? Why 1 layer? Would you have the same result with more layers? not clear to me\n- Why do we need to learn so well how positional encodings are learned?", "limitations": "I think the authors do not have major limitations or negative societal impact.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "PROS:\n- the paper seem to be very detailed and provides a mathematical analysis of ViT\n\nCONS:\n- very hard to read\n- limited/niche contribution\n\nDETAILS:\nGenerally speaking, this paper is VERY difficult to follow. My background is computer science but I found it hard to read this paper that is FULL of definitions, assumptions, definitions, theorems but, most importantly, the paper introduces many many letters in each section.\n\nThe math seems reasonable but I have to admit that around line 200 I was already a bit lost and I could not follow the paper so well. I might have missed something. This raises also a problem for me: I did not understand how ViT learns the spatial inductive bias after reading 9 pages only about that. Moreover, I am not even sure that the tests done with 1-layer ViTs are generalizable to more layers.\n\nMoreover, I think the paper misses one main point: explaining WHY it is important to learn how ViT learns the inductive biases. I know it sounds a stupid comment, but without this explanation a reader is not intrigued by the paper and the heavy math inside it. Then, once we know that “ViT do not just group nearby pixel together” why is it useful? I mean, I am very naive but self-attention is based on content, so I am not surprised that it does not group nearby pixel together.\n\nHere some details from my review:\n- L29: not clear. Are you just saying that near pixels are more similar to distant pixels? Because it not so clear that distant pixels encode “abstract\" features. \n- Figure 1 b should not be there. Better to move it near the text.\n- L80 this notation is not a related work, is it?\n- L90 at this point is not clear what C is\n- L91 why is a definition expected to be an assumption? L90-94 is not clear at all\n- The test of L116-129 is not so clear to me. Where is the cubic activation applied? Why 1 layer? Would you have the same result with more layers? not clear to me\n- L140 is Threshold_0.9C might be Threshold_0.9\n- I would modify the title in something more clear like “How do Vision Transformers learn…\"", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "sYm6uDkgi0q", "reviewer": "Reviewer_4CvW", "summary": "As observed by the authors, the advent of the vision transformer architecture is surprising given that their network structure has nothing specifically curated for vision problems. This raises the question of how it is that, in training a vision transformer, certain desirable properties for vision problems are obtained in the resulting models.  In the abstract authors refer to these as \"convolutional-like patterns\", by which they mean the patches of the image are embedded with high cosine similarity to adjacent patches. The question the authors seek to address is how it is that there resulting positional encodings are learnt.\n\nThe model considered is called by the authors \"positional attention\" which, if I understand correctly, consists of a single self-attention layer where only the positional embeddings are fed into the attention mechanism when computing the attention matrix. In this setting, when the query and key matrices are held fixed as identities (Assumption 3.1), the one-layer model learns to associate patches in a manner consistent with the assumed data generating process, (1), i.e. Theorem 4.1.", "questions": "* What is missing from this paper are two things, one which may be more easily addressed than the other: \n\t1. what the relationship is between this simplified structured partitioning problem and and the real phenomenon of local cosine similarity among positional embeddings in a vision transformer; and\n\t2. what these results tell us about the \"*how*\" portion of the question posed above. \n\n* Supposing (1) is answered, the results in the paper appear merely to prove that the phenomenon happens, rather than answering questions about how or why it happens.\n\n* The closest insight to (2) this appears to be Theorem 3.1 which says that GLM's cannot solve these structured problems. Okay, so we must need something more complicated than a GLM? Is it a ViT? is it a single layer positional attention model? Is it something else we haven't even thought of?\n\n* By removing the inputs from the attention mechanism, is it really fair to call the resulting architecture \"attention\" at all? \n\n* The conclusion reads *\"Our work is a first step towards understanding how Transformers learn tailored inductive biases when trained on large datasets with gradient descent.\"* What exactly the first that first step? What new understanding should I take away from this work about the behaviour of vision transformers. \n", "limitations": "* The work would be strengthened by an analysis of the limitation of the data generating process assumption. Since this is not the way labels yielded conventionally for vision transformers. I.e. in the case of the ViT the features were extracted from the `[pad]` token output. \n\n* Much of the discussion about \"competitiveness\" of \"positional attention\" seems like it would have benefitted by an analysis looking at how much of the error gradient flowed to the original input matrix $X$ in the sum $X+P$ in the pre-transformer layer. In the sense that if really it makes no difference to tweak the architecture in this fashion then $\\mathrm{SA}(X;P) \\approxeq \\mathrm{PA}(X; P)$. Really the point here should be to argue for the modifications to the architecture for positional attention, which I am not observing. \n\n* Still the more fundamental question about translation invariance of convolution seems to be have forgotten despite \"convolution\" appearing throughout.\n\n* Given the implications and assumptions made here, and the claim that this is a first step towards understanding something, what are the possible next steps? What is the next step get more of a handle on features of the transformer architecture that either make it unique, or ordinary?", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I will preface this section by observing that I believe the term \"convolutional pattern\" to be a misnomer and misleading. The defining quality of the convolutional layer is translation invariance. What the authors study instead is how transformers learn positional embeddings with high local cosine similarity; which is quite different from a more interesting question which is how transformer models adapt to domains where convolutional architectures have remained dominant. For my review I will use the terminology consistent with the paper. The confusion extends to the description of the partitions $(\\mathcal{S}_\\ell)_\\ell$ being \"convolutional\".\n\n#### Strengths\n* While I have not examined the supplementary material in detail, the authors have a technical work with nontrivial results which I applaud. \n\n* The paper is well written with good, clear notation. \n\n* In Section 3 that the authors make an effort to show the nontrivially of the data generating process in sense that it cannot be fit with low error by a GLM. (Theorem 3.1) . I.e. learning the positional embeddings is necessary generalisation given the data generating process. \n\n* The permutation experiments in Section 4 are interesting. \n\n#### Weaknesses\nThe biggest weakness for this work I feel is the overall style. The original question posed by the authors \"****how**** do ViTs learn [convolutional patterns] by solely minimizing their training loss?\" is an interesting one. However it is not clear to me how the (well written) theorems and results shed light on this question. \n\n\n\n#### Typographic:\n- [90] $\\mathcal{S}_\\ell  \\subsetneq  D$ should read $\\mathcal{S}_\\ell  \\subsetneq  [D]$ and a quantifier over $\\ell$ should be inserted: \"for all $\\ell \\in [L]$\"\n- [195] Observe your hyperlink to \"Assumption 3.1\" is linking to a theoremlike environment titled \"Simplification 3.1\"", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "M4zT-6TMqi", "reviewer": "Reviewer_1xfL", "summary": "The paper provides a theoretical analysis of gradient descent on logistic loss under simplified formulations of model and data (represented as a collection of $D$ tokens/patches). The assumptions and their relevance are detailed below, but briefly at high level the data assumption models presence of structured association between patches, and the model captures a (significantly) simplified attention mechanism that is loosely motivated from transformer models.  \nThe key result shows that gradient descent in a range of hyperparameters learn the patch associations under the model and data assumptions (more details below). \n", "questions": "Technical questions:\n\nIn proof of Lemma 4.2 in Appendix I.1, I believe the second line in eq. in 205 does not hold $\\xi$ also appears as part of $X$ in other terms. To elaborate, from Definition 3.1, X depends on samples of $\\xi_j, \\delta_j, l(X)$, so the expectation is of the form $E_{\\xi,l(X),\\delta_j) g(X)X$ cannot in general be split as $E[g(X)E_\\xi[X]]$. From briefly checking the options, I think we need $\\xi_j=0$ as assumption for he lemma to go through.  \n\nOther comments:\n\n(a)\tIn the context of the analysis and results in the paper, it’s unclear what the significance of the the Gaussian noise $\\xi_j$ in the data model in Defn. 3.1. From my understanding of the proof, although the noise technically captures a more general model, it does not provide any new insight into the analysis/results itself and merely leads to complicated bookkeeping. If I had missed their significance, I request the authors to clarify in their response. If not, having a simpler analysis without the noise as the main result would improve the readability of the paper significantly (the noisy general analysis could be deferred to separate appendix section). \n   \n(b)\tDefinition 2.3 might need slight modification of wording to ensure there are no-tie breaking required for top_C, for example for a model that learns all position encoding as the same, say $p_i=p$, whether it learns patches or not according to this definition depends on how we do the tie breaking. \n\n(c)\tDefinition 2.1: where each *patch* X -> *input* X \n\n(d)\tFig 2: what is the kernel size of f^*? Also, it is unclear what exactly is being plotted in UpperRight in the  DxD matrix – if it is the ideal convolutional structure, why is it not all 1 in the diagonally entries? For LowerRight, what is the normalization used (not that P^\\top P  is not intrinsically normalized)\n\n(e)\tOne way to rephrase the insight of Fig 2. Is that the patch learning may not happen when merely $f^*(X)=y|X$ has a patch structure, but rather we might require conditions on marginals of distribution over X to have patch structure – this might have been obvious to the authors and experts, but I believe some readers might benefit from explicitly stating this intuition. \n", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The main results to my knowledge are correct except for the point under “Technical questions” below – although I did not check all the proof details. This work is in the line of results that aim to identify interesting inductive biases of models learned using gradient descent. Modern neural networks are extremely complex to theoretically capture in full generality and inevitably any mathematical analysis would rely on significant simplifications. My evaluation of the paper relies on to what extent the paper provides new ideas/insights into the workings of attention based models and I discuss the strengths and weaknesses of the assumptions below.\n\n*Data assumption:* The data distribution in Definition 3.1 captures scenarios where learning patch/token associations are important, in a sense that the labels as well as marginals depend on certain grouping of tokens. In a nutshell, the $D$ data patches are grouped into L disjoined sets ${S_l:l=1,2,…,L}$ where each set has $C$ patches, with $C<<D$. The goal is to understand when a model/algorithm can learn the patch association/groupings $S_l$.  The dimension scaling in the data model captures high SNR scenarios with spurious features. \nThis is a good simple data model to theoretically understanding of learning structures present in data and is of interest beyond this paper. Secondary results in the paper establish that this model is *not* learnable by any generalized linear model, and moreover it is possible for a one-layer transformer model can represent the labeling distribution even without learning the patch association structure (Thm 3.1 & 3.2). \n\n*Model and hyperparameter assumption:* The model on inputs denoted as $X\\in\\mathbb{R}^{d\\times D}$ that is studied by the paper can be written as \n$$ F(X)=\\sigma(\\langle X, v\\cdot s_i^\\top\\rangle),\\text{ where } s_i=softmax(a_i)$$\nwhere $a_i\\in\\mathbb{R}^D$ for $i=1,2,\\ldots, D$ and $v\\in\\mathbb{R}^d$ are unconstrained parameters over which gradient descent is performed, and $\\sigma$ is polynomial activation. $S,A$ are $D\\times D$ matrices with ${s_i},{a_i}$ along the rows. \n\nIn my evaluation, the main shortcoming of the paper is that the simplifications in model and hyperparameter assumptions do not capture aspects of transformer/attention-based models to a sufficiently satisfactory extent. The authors acknowledge that in their paper in multiple places and I am not opposed extreme simplifications if they capture some new insight about the models the paper aims to address. Unfortunately, with the combination of above modeling choice, it is not clear what aspects of attention model are still being captured. \n\n1.\tThe motivation for the formulation is based loosely on transformers, where A and S correspond to attention and score matrices. However, the connection is very loose in that A does not depend on the inputs at all.  Even with the simplification that the attentions only depend on the positional encoding (which I think is reasonable for image tasks and is also justified with experiments), the standard form would involve parametrizing $A$ in factorized form as $P^\\top W_QW_K^\\top P$ -- even though the factorized form does not change representation, it is certainly known to change the gradient descent dynamics which is not captured in this form. As stated in the above equation, it is unclear how the model even distinguishes attention mechanism from any other simplification of fully connected networks.\n\n2.\tA more important concern for me arises from the choice of initialization hyperparameters. With initialization $a_{ii}=polyloglog(d)$ and $a_{ij}=1/poly(d) $. This provides a strong prior on initial score matrices as $s_ii\\approx 1-o(1)$ and $s_ij=O(1/d)$ -- i.e., the scores are explicitly chosen to be diagonally dominant. This in comparison to even the simplified parametrization as $A=P^\\top W_QW_K^\\top P$ would lead to a more balanced scores at initialization—thus the initialization is truly agnostic to structure and the mechanism if learns patch association at all will have to learn from scratch.\n\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "B_lPl9-hmhY", "reviewer": "Reviewer_7w7a", "summary": "This work aims to understand how Transformers learn inductive biases as convolution when trained on large datasets with gradient descent. They provide detailed analysis and find that ViT relies on the positional attention mechanism that disentangles patches and positional encodings. And it also provides experimentS on the Gaussian data that show that ViTs do not learn the correct inductive bias under the assumption that \"characterizing the distributions under which ViTs recover the structure of the function is an important question.\" Lastly, it also provides experiments on CIFAR and ImageNet.", "questions": "Please refer to Weaknesses.", "limitations": "Please refer to Weaknesses.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The analysis of this work is completed and somehow interesting.\n2. The definition and theorem on the patch association are good.\n\n\nWeaknesses:\n1. Even though the analysis is complete, the final results are not surprising and the improvement on downstream tasks is incremental. I think good work should be simple but effective or complex but effective. This work does not give me any new insight.\n2. In Figure 1(a), why does the author state that ViTs learn the convolution inductive bias in vision datasets? I do not see any convolution inductive bias in this figure.\n\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept