Vision Transformers learn patch association
Anonymous Author(s)
AffiliationAddress email
Abstract
Vision Transformers (ViTs) have recently achieved comparable or superior performance to Convolutional neural networks (CNNs) in computer vision.
Thisempirical breakthrough is even more remarkable since ViTs discards spatial information by mixing patch embeddings and positional encodings and do not embedany visual inductive bias (e.g. spatial locality).
Yet, recent work showed that whileminimizing their training loss, ViTs specifically learn spatially localized patterns.
This raises a central question: how do ViTs learn this pattern by solely minimizingtheir training loss using gradient-based methods from random initialization ?
Wepropose a structured classification dataset and a simplified ViT model to providepreliminary theoretical justification of this phenomenon.
Our model relies on asimplified attention mechanism –the positional attention mechanism– where theattention matrix solely depends on the positional encodings.
While the problemadmits multiple solutions that generalize, we show that our model implicitly learnsthe spatial structure of the dataset while generalizing.
We finally prove that learning the structure helps to sample-efficiently transfer to downstream datasets thatshare the same structure as the pre-training one but with different features.
Weempirically verify that ViTs using only the positional attention mechanism performsimilarly to the original one on CIFAR-10/100, SVHN and ImageNet.
Introduction
Transformers are deep learning models built on self-attention [38], and in the past several years theyhave increasingly formed the backbone for state-of-the-art models in domains ranging from Natural
Language Processing (NLP) [38, 14] to computer vision [15], reinforcement learning [8, 24], programsynthesis [2] and symbolic tasks [27].
Beyond their remarkable performance, several works reportedthe ability of Transformers to simultaneously minimize their training loss and learn inductive biasestailored to specific datasets e.g. in computer vision [32], in NLP [6, 40] or in mathematical reasoning[43].
In this paper, we focus on computer vision where convolutions are considered to be an adequateand biologically plausible inductive bias since they capture local spatial information [18].
This seemsintuitively reasonable: nearby pixels encode the presence of small scale features, whose patternsin turn determine more abstract features at longer and longer length scales.
Several seminal works[9, 15, 32] empirically show that although randomly initialized, the positional encodings in Vision
Transformers (ViTs) [15] learn structured patterns where closer patches have more similar positionalencodings – see Figure 1a-(5,6,7,8).
We refer to this structure as spatially localized patterns.
A priori,learning such specific structure is surprising.
Indeed, in contrast to convolutional neural networks(CNNs), ViTs are not built with the inductive bias of spatial locality and weight sharing.
They startby replacing an image by a collection of D patches p X 1 , . . .
X D q P R d ˆ D , each of dimension d .
While each X i represents (an embedding of) a spatially localized portion of the original image, therelative positions of the patches X i in the image are disregarded.
Instead, relative spatial informationis supplied through image-independent positional encodings P “ p p 1 , . . .
, p D q P R d ˆ D .
Unlike
CNNs, each layer of a ViT then learns, via trainable self-attention, a non-local set of filters that
[Figure] [Figure]
[Figure]
non-linearly depend on both the values of all patches X j and their positional encodings p j .
Despitetheir spatial non-locality, ViTs surprisingly learn the convolution inductive bias in vision datasets asshown in Figure 1a.
This empirical observation sets a central question: from a theoretical perspective , how do ViTsmanage to learn these spatially localized patterns by simply minimizing their training loss usinggradient descent from random initialization?
While it is known that attention can express localoperations as convolution [9], it remains unclear how ViTs learn it.
In this paper, we present a simplespatially-structured classification dataset for which it is sufficient (but not necessary) to learn thestructure in order to generalize.
We also present a simplified ViT model which we prove implicitlylearns a spatially localized inductive bias when it minimizes its training loss via gradient descent(GD).
We name this inductive bias patch association .
We prove that our ViT model leverages this biasto generalize and sample-efficiently transfer to downstream datasets that share the same structure.
Toour knowledge, our work is the first showing that there exists simple (yet natural) data distributionswhere simple ViTs implicitly learn spatially localized patterns during training.
Our contributions aredivided as follows:
– In Section 2, we introduce our framework to formalize patch association.
– In Section 3, we introduce our structured classification dataset and a simplified ViT model totheoretically study how ViTs learn their inductive bias and generalize.
– In Section 4, we present our main theorems (Theorem 4.1 and Theorem 4.2) along with theintermediate lemmas.
We prove that a one-layer single-head ViT model trained with GD on oursynthetic dataset learns its structure and thus generalizes.
– In Section 5, we show (Theorem 5.1) that after pre-training in our synthetic dataset, our model canbe sample-efficiently fine-tuned to transfer to a downstream dataset that shares the same structureas the source dataset (and may have different features).
– In Section 6, we conduct a set of experiments to validate our theoretical model.
Related work
CNNs and ViTs.
Many computer vision architectures can be considered as a form of hybridizationof Transformers with CNNs.
For example, DeTR [7] uses a CNN to generate features that are fedto a Transformer.
[16] show that self-attention can be initialized or regularized to behave like aconvolution and [11, 20] add convolution operations to Transformers.
Conversely, [5, 33, 4] introduceself-attention or attention-like operations to supplement or replace convolution in ResNet-like models.
In constrast, our paper attempts to explain how ViTs learn convolution-like patterns using GD.
Theory for attention models.
Early theoretical works have focused on the expressivity of attention.[39, 17] addressed this question in the context of self-attention blocks and [12, 41, 22] for Transformers.
On the optimization side, [46] investigate the role of adaptive methods in attention models and[35] analyze the dynamics of a single-head attention head to approximate the learning of a Seq2Seqarchitecture.
Our work rather belongs to the implicit bias line of research: we exhibit a settingwhere a single-head ViT trained with gradient descent provably learns convolution-like patterns andgeneralizes.
To our knowledge, this work is the first that proves generalization of attention modelsthrough the lens of optimization.
Notation.
We use lower case letters for scalars, lower case bold for vectors and upper case boldfor matrices.
Given an integer D , we define r D s “ t 1 , . . .
, D u .
Any statement made "with highprobability" holds with probability at least 1 ´ 1 { poly p d q .
Given a vector a P R d and k ď d , wedefine Top k t a j u dj “ 1 “ t a i 1 , . . .
, a i k u where a i 1 , . . .
, a i k are the k -largest elements.
For a function
F that implicitly depend on parameters A and v , we often write F A , v to highlight its parameters.
Defining patch association
The goal of this section is to define the notion of patch association.
Before, we define the datastructure we use to model spatially localized patterns.
Definition 2.1 (Data structure) .
Let D be a distribution over R d ˆ D ˆ t´ 1 , 1 u where each patch
X “ p X 1 , . . .
X D q P R d ˆ D has label y P t´ 1 , 1 u .
We partition r D s into L disjoint subsets i.e.r D s “ Ť Lℓ “ 1 S ℓ such that S ℓ Ĺ D and | S ℓ | “ C.
Note that Definition 2.1 is not an assumption since for any matrix X , it is possible to define a partitionon its columns.
We thus consider distributions D where there exists an underlying partition of thepatches into disjoint sets S ℓ .
A particular case of sets S ℓ is the one of localized sets which we use tomodel the patterns from Figure 1a.
We define them as follows.
Definition 2.2 (Informal: spatially localized sets) .
For X “ p X 1 , . . .
, X D q , we say that the sets S ℓare localized sets if @ ℓ P r L s , @
i P S ℓ , the X
i ’s form a rectangle.
An example of spatially localized sets is S 1 “ t 1 , . . .
, C u , . . .
, S t D { C u “ t D ´ C, . . .
, D u that wedepict in Figure 1b-(1) in the case D “ 16 and C “ 4 .
We emphasize that Definition 2.1 is notlimited to spatially localized sets and also covers non-contiguous sets as Figure 1b-(2).
Since we aimto show that ViTs learn the sets S ℓ while fitting the labels y , we need to assume the existence of alabeling function f ˚ that respects the partition in Definition 2.1.
Assumption 1 (Data structure) .
Let D as in Definition 2.1.
We assume that there exists a labelingfunction f ˚ such that P r yf ˚ p
{+X q ą 
+}
s “ 1 ´ d ´ ω p 1 q
and,
[Equation]
where ϕ : R d ˆ C Ñ R is some function, (1)
Assumption 1 states that there exists a labeling function that preserves the underlying structure byapplying the same function ϕ to each S ℓ as in (1).
For instance, when the sets S ℓ ’s are spatiallylocalized, f ˚ can be a one-hidden layer convolutional network.
In this paper, we are interested inpatch association which refers to the ability of an algorithm to identify the sets S ℓ ’s. Following[9, 15, 32], we investigate whether ViTs can learn patch association using its positional encodings.
Definition 2.3 (Patch association for ViTs) .
Let Assumption 1 holds.
Let M : R d ˆ D Ñ t´ 1 , 1 u be atransformer and P p M q its positional encodings matrix.
We say that M performs patch association iffor all ℓ P r L s and i P S ℓ , we have Top C tx p p M q i , p p M q j yu Dj “ 1 “ S ℓ .
Definition 2.3 states that patch association is learned when for a given i P S ℓ , its positional encodingmainly attends those of j such that i, j P S ℓ .
In this way, the transformer groups the X i according to
S ℓ just like the true labeling function – It is thus natural to ask would ViTs really learn those S ℓ aftertraining to match the labeling function f ˚ ?
[Figure]
ViTs do not always learn patch association under
Assumption 1.
Figure 1a shows that ViTs learnspatially localized patterns on ImageNet.
This raisesa question: while minimizing their loss, do ViTs always learn the structure of S ℓ for any data distributionsatisfying Assumption 1?
We give a negative answerthrough the following synthetic experiment.
Considerthe case where all the patches X j are i.i.d.
standard
Gaussian and f ˚ being a one-hidden layer convolutional network with cubic activation.
As shownin Figure 2, one-layer ViT reaches small test error.
However, P J P does not match the convolution pattern encoded in f ˚ .
To prove that ViTs learn patchassociation, we need additional assumptions on D .
Setting to learn patch association
In this section, we introduce our theoretical setting to analyze how ViTs learn patch association.
Wefirst define our binary classification dataset and finally present the ViT model we use to learn it.
Data distribution.
The binary classification problem analyzed in this paper is defined as follows.
Definition 3.1 (Data distribution) .
Let D be a distribution satisfying Assumption 1 and w ˚ P R d bethe underlying feature.
For each data-point X , we define two types of sets:
– Signal set : uniformly sample an index ℓ p X q from r L s and for j P S ℓ p X q , X j “ y w ˚ ` ξ j , where
y w ˚ is the informative feature and ξ j i.i.d. „
N p 0 , σ 2 I D q .
– Random sets : for ℓ P r L szt ℓ p X qu and j P S ℓ , X j “ δ j w ˚ ` ξ j , where the spurious
feature δ j “ 1 with probability q { 2 , ´ 1 with same probability and 0 otherwise, and ξ j i.i.d. „
N p 0 , σ 2 I D q .
Besides, D admits the labeling function
f ˚ p
X
q
“ ř ℓ Pr L s Threshold 0 .
9 C p ř i P S ℓ x w ˚ , X i
yq
, where
Threshold C p
z q “ z if | z | ą C and 0 otherwise.
-    +1 +1 0 -  +1 +1 +   - -1 +   0 - - 0 0 -1
0 + + than "+1"’s, the label of X is
since there are only "+1"’s inside the signal set.
We sketch a data-point of D in Figure 3.
Our datadistribution is a model to describe the patch association phenomenon in Figure 1-(a). Indeed, we candivide each image in sets S ℓ ’s that contain adjacentpatches. For instance, the patches in yellow in Figure 1-(a) belong to the same S ℓ . Thus, in our setting,proving patch association amounts to show that thepositional encodings in the same set S ℓ are highlycorrelated.
We make the following assumptions onthe parameters of D .
Assumption 2.
We set d “ poly p D q , C “polylog p d q , q “ poly p C q{ D , } w ˚ } 2 “ 1 andσ 2 “ 1 { d .
This implies C !
D and q ! 1 .
Assumption 2 may be justified by considering a "ViTbase-patch16-224" model [15] on Imagenet.
In thiscase, d “ 384 , D “ 196 .
σ is set to have } ξ j } 2 «} w ˚ } 2 .
q is chosen so that there are more spuriousfeatures than informative ones (low signal-to-noiseregime) which makes the data non-linearly separable.
Our dataset is thus non-trivial to learn since generalized linear networks fail to generalize.
Theorem 3.1.
Let D be as defined in Definition 3.1.
Let g p X q “ ϕ ´ř
Dj “ 1 x w j , X j y ¯ be a generalized linear model.
Then, g does not fit the labeling function i.e. P r f ˚ p X q g p X q ď 0 s ě 1 { 8 .
Intuitively, g fails to generalize because it is does not have any knowledge on the underlying partitionand the number of random sets is much higher than the signal one.
Thus, a model must have aminimal (but not necessarily complete) knowledge about the S ℓ ’s in order to generalize.
Theorem 3.2.
Let D be defined as in Definition 3.1.
There exists a (one-layer) transformer M so that
P r f ˚ p X q M p
X q ď 0 s “ d ´ ω p 1 q but for all ℓ P r L s , i P S ℓ , Top C tx p p M q i , p p M q j yu Dj “ 1 X S ℓ “ H .
Theorem 3.2 states that there exist classifiers that generalize without being able to learn patchassociation (see Appendix for proof).
In Section 4, we show that there exist a ViT model that learnspatch association via minimizing its training objective.
171 Simplified ViT model.
We here define our simplified ViT model.
Before, we remind the definitionof attention [3, 38] in the single-head case with average pooling.
Definition 3.2 (Self-attention) .
Let X P R
d
ˆ D be a patch from D and P
P R d ˆ D the positionalencoding.
The attention mechanism computes
the sum of patches and positional encodings i.e. XXX “ X ` P P R d ˆp D ` 1 q .
the attention matrix
A “ QK J where Q “ XXX J W Q , K “ XXX J W K , W Q , W K P R d ˆ d .
the score matrix S P R D ˆ D with coefficients S i,j “ exp p
A i,j {?
d q{ ř Dr “ 1 exp p
A i,r {?
d q .4.
the matrix V “ XXX J W V , where W V P R d ˆ d .
It finally outputs SA pp X
P qq “ SV P R d
ˆ D .
Given its highly non-linear nature with respect to the input, directly analyzing SA is difficult.
Wenow introduce the positional attention mechanism that circumvents this issue.
Definition 3.3 (Positional attention) .
Let X P R
d
ˆ D be a datum from D and P P R d ˆ D the positional encoding.
The positional attention mechanism takes as input the pair p X ; P q and computes:
the attention matrix
A “ QK J where Q “ P J W Q , K “ P J W K and W Q ,
W K
P
R
d ˆ d .2.
the matrix V “ X J W V , where W V P R d ˆ d .
Similarly to self-attention, it computes the score matrix and outputs PA pp X ; P qq “ SV .
In Definition 3.3, one of the main differences with self-attention is having A “ P J W Q W K P .
Thus, the advantage of positional attention is to isolate the positional encoding from the data.
Usingthis framework, we can thus use A to track the dynamics of P in order to prove patch association asin Figure 1a.
On the other hand, the matrix V encodes the data-dependent part.
In summary, a modelthat fits the labeling function and learns patch association in our setting must have i)
A learns patchassociation ii) V learns w ˚ in order to generalize.
Lastly, we make the following simplification inthe parameters to ease our analysis.
Simplification 3.1.
In the positional attention mechanism, we set d “ D , W K “ I D and W Q “ I Dwhich implies A “ P J P .
Besides, we set W V “ r v , . .
. , v s P R D
ˆ D where v P R D .
In Simplification 3.1, we set W K and W Q so that A “ P J P .
This Gram matrix encodes theinductive bias of the ViT as previously seen in Figure 1a.
Besides, since fitting the labeling function requires to learn one feature w ˚ , it is sufficient to parametrize W V with a vector v .
Under Simplification 3.1, our simplified ViT model is thus one-hidden attention layer with a single head:
[Equation]
where σ is an activation function.
Since we aim to the simplest ViT model, we opt for a polynomialactivation i.e. σ p x q “ x p ` νx where p ě 3 is an odd integer and ν “ 1 { poly p d q .
Note that thischoice of polynomial activation is common in the deep learning theory literature – see e.g. [28, 1, 42]among others.
The degree p is odd to make the ViT model compatible with the labeling function andstrictly larger than 1 because the data is not linearly separable (Assumption 2).
We add a linear partin the activation function to ensure that the gradient is non-zero when v has tiny coefficients.
Withthese simplifications, we formally prove that F is able to learn patch association and generalize.
What aspects of self-attention does positional attention capture?
As SA, PA is also permutationinvariant and processes all tokens simultaneously. Besides, PA also computes a score matrix betweenthe different tokens. This similarity matrix is also normalized in a sparse manner with the Softmaxoperator. The only aspect that PA misses from SA is the fact that S does not depend on the input.
Nevertheless, we believe that our PA model encapsulates most of the needed inductive bias incomputer vision since in Section 6, we observe that PA numerically competes with self-attention.
Trainable parameters.
The trainable parameters of our model (T) are specified as follows.
Assumption 3.
In the model (T) , we set A and v as trainable parameters.
Besides, without loss ofgenerality, we train all A i,j for i ‰ j and leave the diagonals of A fixed.
Although A “ P J P and P is trainable, we choose for simplicity to only optimize over A .
Besides,we leave the A
i,i ’s fixed because the softmax is invariant under the uniform shift of the input.
Learning problems.
Given a dataset Z “ tp X r i s , y r
i squ Ni “ 1 sampled from D , we solve theempirical risk minimization problem for the logistic loss defined by:
[Equation]
Instead of directly analyzing (E), we introduce a proxy where we minimize the population risk
[Equation]
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 inthe case of (P) for t P r T s and i, j P r D s is
[Equation]
i
[Equation]
where η ą 0 is the learning rate.
A similar update may be written for (E).
We now detail how to setthe parameters in (GD).
Parametrization 3.1.
When running GD on (P) and (E) , the number of iterations is any T ěpoly p
d q{ η.
We set the learning rate as η P ` 0 , 1 poly p d q ˘ .
The diagonal coefficient of the attentionmatrix are set for i P r D s as A p 0 q i,i “ p A p 0 q i,i “ σ A I D where σ A “ polyloglog p d q .
The off-diagonalcoefficients of A and the value vector are initialized as:
Idealized case: v p 0 q “ α p 0 q w ˚ where α p 0 q “ ν 1 {p p ´ 1 q and A p 0 q i,j “ 0 for i ‰ j.2.
Realistic case: p v p 0 q „ N p 0 , ω
2 I d q and p A p 0 q i,j „ N p 0 , ω
2 q where i ‰ j and ω “ 1 { poly p d q .
We now explain our initialization choice in the idealized case.
We remind that in Simplification 3.1,we have A “ P J P .
{+If one initializes P „ N p 0 , σ A I D { D q , then with high probability, A p 0 q i,i “} p p 0 q i } 22 “ Θ p σ A q and A p 0 q i,j “ x p p 0 q i , p p 
+}
q j y “ Θ p σ A {?
D q for i ‰ j . Since D " 1 , it is thenreasonable to set A p 0 q i,j “ 0 . Besides, we initialize v p 0 q in span p w ˚ q even though this latter shouldbe unknown to the algorithm.
We remind that the idealized case is only an easy-to-study proxy toultimately characterize the realistic dynamics.
Learning patch association via matching the labeling function
As announced above, we show that our ViTs (T) implicitly learns patch association and fits thelabeling function by minimizing the training objective.
We first study the dynamics in (P).
Using theanalysis in the idealized case, we then characterize the solution found in the realistic problem (E).
4.1 Learning process in the idealized case
In this section, we analyze dynamics of (P).
Our main result is that the final iterates of (P) learn patchassociation while generalize using GD on infinitely many examples sampled from D .
Theorem 4.1.
Assume that we run GD on (P) for T iterations with parameters set as in Parametrization 3.1.
With high probability, the ViT model (T)
learns patch association i.e. for all ℓ P r L s and i P S ℓ , Top C t
A p T q i,j u Dj “ 1 “ S ℓ .2. learns the labeling function
f ˚ i.e. P D r f ˚ p
X q F A p T q , v p T q p X q ą
s ě 1 ´ o p 1 q .
We now sketch the main ideas to prove the theorem and refer to Appendix D for a complete proof.
Invariance and symmetries.
The main feature making (P) appealing is the fact that we take theexpectation over D .
Indeed, 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 thatshows that A can be reduced to three variables in (P).
Lemma 4.1.
There exist β “ σ A , γ p t q , ρ p t q P R such that for all t ě 0 :
for all i P r D s , A p t q i,i “ β.2.
for all i, j P r D s such that i, j P S ℓ for some ℓ P r L s , A p t q i,j “ γ p t q . for all i, j P r D s such that i P S ℓ and j P S m for some ℓ, m P r L s with ℓ ‰ m , A p t q i,j “ ρ p t q .
[Figure]
o r
[Figure] [Figure] [Figure]
s i n e
S i
Cosine
[Figure]
Besides, using the initialization in Parametrization 3.1, we can show that v always lies in span p w ˚ q .
Lemma 4.2.
For all t P r T s , there exists α p t q P R such that v p t q “ α p t q w ˚ .
In summary, Lemma 4.1 and Lemma 4.2 imply that instead of optimizing over A and v in (P), wecan instead consider the scalar variables α p t q , γ p t q and ρ p t q .
The remaining of this section consists inanalyzing the dynamics of these three quantities.
Learning patch association.
We first analyze the dynamics of γ p t q and ρ p t q .
To this end, weintroduce the following softmax terms:
[Equation]
Note that Λ p t q , Γ p t q and Ξ p t q respectively correspond to the coefficients on the diagonal, those forwhich i, j P S ℓ for some ℓ and all the other coefficients of the attention matrix S .
Using thesenotations, we first derive the GD updates of γ p t q and ρ p t q .
Lemma 4.3.
Let t ď T .
The attention weights γ p t q and ρ p t q satisfy:
[Equation]
Lemma 4.3 shows that the increment of γ p t q is larger than the one of ρ p t q .
Since γ p 0 q “ ρ p 0 q “ 0 ,this implies that γ p t q ě ρ p t q for all t ě 0 .
This observation proves the first item of Theorem 4.1.
Wenow show how learning patch association leads to v highly correlated with w ˚ .
How do ViTs learn patch association? The proof of Theorem 4.1 relies on this mechanism. – Event I :
At the beginning of the process, the gradient increment for v is larger than the one of A p t q i,j .
Because the A p t q i,j ’s stay very small and v p t q is initialized in span p w ˚ q , v updates in the direction of w ˚ .
At this point, the model is nothing else than a generalized linear model that would not
generalize because there are much more noisy tokens than signal ones (see Theorem 3.1). – Event II : Once x v p t q , w ˚ y becomes large, the A p t q i,j ’s start to update.
At this point, because the A i,j ’s are not negligible, v p t q updates in the direction of w ˚ and of ξ j . However, the gradient
increment for γ (which corresponds to A p t q i,j ’s in the same set S ℓ ) is much larger than the one for ρ(Lemma 4.3).
Thus, γ p t q increases until being large enough so that the gradient of L with respectto v is in the direction of w ˚ again.
– Event III : Since the gradient of L is in the direction of w ˚ , v p t q updates again in the direction ofw ˚ until the population risk becomes a Ω p 1 q and we fit the labelling function.
Our mechanism highlights a main advantage of attention-based models: their ability to perform patchassociation i.e. to set high A i,j for i, j in the same set S ℓ . By doing so, we show that they are able toalign their gradient in span p w ˚ q and fit the labelling function. Another by-product of this ability isto have sample-efficient transfer as highlighted in Section 5.
4.2 From the idealized to the realistic learning process
The real learning process differs from the idealized one in that i) we have a finite number of samplesii)
we initialize both p A and p v as Gaussian random variables.
Using a polynomial number of samples,we show that (T) is still able to learn patch association and generalize
[Figure] [Figure]
Theorem 4.2.
Assume that we run GD on (E) for T iterations with parameters set as in Parametrization 3.1.
Assume that the number of samples is N “ poly p d q .
With high probability, the model
learns patch association i e. for all
and
[Figure] [Figure]
2. fits the labeling function i
[Equation]
Figure 4 numerically illustrates Theorem 4.2.
Similarly to [29], the proof introduces a "semi-realistic"learning process that is a mid-point between the idealized and realistic processes.
We show that p A p T qand p v p T q are close to their semi-realistic counterparts – see Appendix E for a complete proof.
Patch association yields sample-efficient fine-tuning with ViTs
A fundamental byproduct of our theory is that after pre-training on a dataset sampled from D , ourmodel (T) sample-efficiently transfers to datasets that are structured as D but differ in their features.
Downstream dataset.
Let r D a downstream data distribution defined as in Definition 3.1 such thatits underlying feature is r w ˚ with } r w ˚ } 2 “ 1 and r w ˚ potentially different from w ˚ .
In other words,the downstream r D and source D distributions share the same structure but not necessarily the samefeature.
We sample a downstream dataset r Z “ tp Ă X r
i s , r y r
i squ Ă Ni “ 1 from r D .
Learning problem.
We consider the model (T) pre-trained as in subsection 4.2.
We assume that p Ais kept fixed from the pre-trained model and we only optimize the value vector r v to solve:
[Equation]
We run GD on ( r E) with parameters set as in Parametrization 3.1 except that the p
A i,j ’s are fixed andr v p 0 q „ N p 0 , ω
2 I d q with ω “ 1 { poly p d q .
Our main results states that this fine-tuning procedurerequires a few samples to achieve high test accuracy in r D .
In contrast, any algorithm without patchassociation needs a large number of samples to generalize.
Theorem 5.1.
Let p A be the attention matrix obtained after pre-training as in subsection 4.2.
Assumethat we run GD for T iterations on ( r E) to fine-tune the value vector.
Using r N ď polylog p D q samples,the model (T) transfers to r D i.e. P r D r f ˚ p
X q F p A , r v p T q p X q ą
s ě 1 ´ o p 1 q .
Theorem 5.2.
Let A : R d ˆ D Ñ t´ 1 , 1 u be a binary classification algorithm without patch association knowledge.
Then, it needs D Ω p 1 q training samples to get test error ď o p 1 q on r D .
The proofs of Theorem 5.1 and Theorem 5.2 are in Appendix F.
These theorems hightlight thatlearning patch association is required for efficient transfer.
We believe that they offer a new perspectiveon explaining why ViTs are widely used in transferring to downstream tasks.
While it is possible that
ViTs learn shared (with the downstream dataset) features during pretraining, our theory hints thatlearning the inductive bias of the labeling function is also central for transfer.
Numerical experiments
[Figure] In this section, we empirically verify that i) ViTs learn patch association while miniziming its trainingloss and ii) the positional attention mechanism competes with the vanilla one on small-scale datasetssuch as CIFAR-10/100 [26], SVHN [31] and large-scale ones such as ILSVRC-2012 ImageNet[13].
For the small datasets, we use a ViT with 7 layers, 12 heads and hidden/MLP dimension384.
For ImageNet, we train a "ViT-tiny-patch16-224" [15].
Both models are trained with standardaugmentations techniques [10] and using AdamW with a cosine learning rate scheduler.
We run allthe experiments for 300 epochs, with batch size 1024 for Imagenet and 128 otherwise and averageour results over 5 seeds.
We refer to Appendix A for the training details.
ViTs learn patch association.
We consider the CIFAR-100 dataset where we divide each image into grids of size s ˆ s pixels.
For a fixed s P t 2 , 4 , 8 , 16 , 32 u , we permute the gridsaccording to π s to create the permuted CIFAR-100 dataset.
We call s the grid shuffle size.
[Figure] Figure 5b-(1) shows a CIFAR-100 image and itscorresponding shuffling in the permuted CIFAR100 dataset Figure 5b-(2-3-4).
We train a ViT andconvolutional models Resnet18 [21] and VGG-19[34] on the permuted CIFAR-100 dataset.
For the ViT, we set the patch size to 2 , althoughthis is sub-optimal in terms of accuracy, becausethe patch size needs to stay smaller or equal tos .
Figure 5a reports the test accuracy of thesethree models for different values of s .
When s issmall, the image does not have a coherent structuree.g. Figure 5b-(2) and thus, convolutional modelsstruggle to generalize.
As s increases e.g. Figure 5b-(4), the information inside a patch is meaningfuland thus, the convolutional models well-perform.
In contrast, ViTs are insensitive to permutationsas their performance remains unchanged for all s – see Figure 5a.
We believe that this is due to theability of ViTs to learn patch association.
Indeed, intuitively, when we permute the grids in Figure 5b,we lose the spatially localized structure and create new sets S ℓ ’s and a new labeling function f ˚ .
Despite this change, we verify that the ViT is able to recover the new S ℓ ’s: we feed the ViT with theshuffled pear image ( Figure 5b-(2)) and see that P J P does not have a spatially localized structurein Figure 5c-(1,2).
We then apply π ´ 1 s to P J P and observe that we recover the spatially localizedpatterns Figure 5c-(1’,2’).
This experiment highlights that ViTs do not just group nearby pixelstogether.
They learn patch association even when the S ℓ ’s are spatially delocalized.
[Figure]
ViTs with positional attention are competitive.
We numerically verify that ViTs using positionalattention compete with those with vanilla attention.
In Section 3, we introduced positional attention todefine our theoretical learner model.
Figure 6 and
Figure 7 show that ViTs using positional attentioncompete with vanilla ViTs on a range of datasets.
These experiments strengthen our intuition thatfor images, having an attention matrix that onlydepends on the positional encodings is sufficientto have a good test accuracy.
Conclusion, limitations and future works
Our work is a first step towards understanding how Transformers learn tailored inductive biases whentrained on large datasets with gradient descent.
Our analysis heavily relies on the positional attentionmechanism that disentangles patches and positional encodings.
In practice, self-attention mixes thesetwo quantities.
An interesting direction is to understand the impact of patch embeddings on theinductive bias learned by ViTs.
Moreover, our experiment on the Gaussian data shows that ViTs donot always learn the correct inductive bias under Assumption 1: characterizing the distributions underwhich ViTs recover the structure of the function is an important question.
Lastly, this work also pavesthe way to many extensions beyond convolution.
For example, can ViTs learn other inductive bias?
What are the inductive biases learnt by Transformers in NLP?
Answering those questions is central tobetter understand the underlying mechanism of attention.
Checklist
For all authors...
[Table] (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes] See Section 4 and Section 5.
(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 atheory paper.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
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 as a URL)?
[Yes] See supplementary 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 (e.g., with respect to the random seed after running experi- ments multiple times)?
[Yes] See Section
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If 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 are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]