PaperlayerView exact revision on arXiv

Source-linked AI summary

Vision Transformers provably learn spatial structure

Samy Jelassi, Michael E. Sander, Yuanzhi Li

arXiv:2210.09221v1cs.CVcs.LG

TL;DR

ViTs learn spatially localized patterns despite lacking CNNs’ built-in local connectivity, but how gradient descent learns those patterns remains unclear. The paper analyzes a simplified ViT with positional attention on a structured dataset, proving patch association and generalization and reporting competitive performance and sample-efficient transfer within the stated scope.

  • Problem

    Although ViTs lack CNNs’ local-connectivity bias, they learn spatially localized patterns, raising the question of how gradient descent learns these patterns from random initialization.

    Source: §1 · pp. 1–2

  • Method

    The paper introduces a spatially structured classification dataset and a simplified ViT whose attention matrix depends only on positional encodings.

    Source: §1 · p. 2, §3 · pp. 5–6

  • Results

    The simplified ViT learns patch association and generalizes, transfers sample-efficiently across datasets sharing structure, and is competitive with vanilla ViT on ImageNet, CIFAR-10/100, and SVHN.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support studying positional attention as a way to analyze how ViTs learn spatial structure and transfer across data with shared structure but different features.

    Source: §1 · p. 2, §3 · p. 6

    The analysis relies on positional attention, which disentangles patches from positional encodings, whereas practical self-attention mixes them; Gaussian data also shows structure recovery is not guaranteed.

    Source: Conclusion, limitations and future works · pp. 10–11

Abstract

from arXiv · show

Vision Transformers (ViTs) have achieved comparable or superior performance than Convolutional Neural Networks (CNNs) in computer vision. This empirical breakthrough is 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 a 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 a spatially structured dataset and a 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-efficiently transfer to downstream datasets that share the same structure as the pre-training one but differ in the features. Lastly, we empirically verify that a ViT with positional attention performs similarly to the original one on CIFAR-10/100, SVHN and ImageNet.

1 Introduction

ViTs can learn spatially localized patterns despite lacking CNNs’ built-in local connectivity, raising the question of how gradient descent selects such structure. The paper studies this through a structured classification dataset and simplified ViT, proving patch association and generalization and examining transfer and empirical performance.

1 Introduction

  • The paper asks how ViTs learn local connectivity from random initialization and proves that a simplified ViT can implicitly learn patch association while generalizing.The analysis studies a spatially structured binary-classification dataset and uses invariance and attention-matrix symmetries to analyze learning.Source: §1 · p. 2
  • The proposed dataset makes spatial structure sufficient, but not necessary, for generalization; the simplified model’s attention matrix depends only on positional encodings.The paper studies empirical-risk and population-risk minimization for binary classification.Source: §1 · p. 2
  • After pre-training, the model can be sample-efficiently fine-tuned on downstream data that shares the source structure but may have different features.This transfer result is stated as Theorem 5.1.Source: §1 · p. 2
  • The positional-attention ViT is competitive with the vanilla ViT on ImageNet, CIFAR-10/100, and SVHN; it also validates broader spatial structure on pixel-permuted CIFAR-100.The pixel-permutation result is presented as evidence that ViTs learn spatial structure beyond local patterns.Source: §1 · p. 2

Related work

  • The paper excludes CNN–Transformer hybrids, instead simplifying ViT attention to study gradient-descent learning of spatial patterns.The related-work discussion also contrasts this focus with prior analyses of ViT representations, patch embeddings, and attention matrices.Source: Related work · p. 3
  • Prior attention theory studies expressivity and optimization; this work analyzes a single-head ViT that provably learns convolution-like patterns and generalizes.Source: Related work · p. 3
  • The paper frames its question as algorithmic regularization in ViTs, noting that prior cited work on this topic concerns feed-forward networks.Source: Related work · p. 3

2 Defining patch association

This section defines patch association as learning the spatial groups used by the dataset’s labeling structure. A Gaussian-patch example shows that low test error alone does not ensure the model learns those groups.

2 Defining patch association

  • Spatial structure is defined using a partition of patches and a labeling function that applies the same function φ to each set Sℓ.When the sets are spatially localized, the labeling function can be a one-hidden-layer convolutional network; the definition also allows non-contiguous sets.Source: §2 · pp. 3–4
  • Figure 2 juxtaposes test error, the convolution-defined localized sets, and learned positional similarities that look random relative to those sets.Source: Figure 2 · p. 4, §2 · p. 4
  • Patch association means identifying the dataset’s sets Sℓ, so patches are grouped according to the structure used by the true labeling function.For a given patch i in Sℓ, its positional encoding mainly attends to encodings of patches j in that same set.Source: §2 · p. 4
  • On Gaussian patches, a one-layer ViT reaches small test error without matching the convolution pattern, showing that generalization alone does not ensure patch association.The example motivates additional assumptions on the data distribution to prove patch association.Source: §2 · p. 4

3 Setting to learn patch association

The section defines a spatially structured binary-classification task and a simplified ViT whose positional attention separates spatial structure from data features. It sets up idealized and finite-sample learning problems to study whether training recovers that structure.

Spatially structured dataset

  • Under Assumption 2, with C << D and q << 1 (a low signal-to-noise regime), the dataset has useful adjacent patches and uninformative or spurious others; generalized linear networks fail to generalize.Source: §3 · p. 5

Why study patch association

  • Generalization does not uniquely imply patch association: a one-layer transformer can generalize without learning the underlying partition.This existence result motivates studying which solution gradient descent selects.Source: §3 · p. 5

Simplified ViT model

  • The simplified ViT uses positional attention: its attention matrix depends on positional encodings, while its value matrix carries the data-dependent features.The model parameterizes the value projection with a vector and optimizes the off-diagonal attention coefficients while keeping the diagonal fixed.Source: §3 · pp. 5–6

Learning problems

  • The analysis considers both empirical risk on finite samples and population risk as an idealized proxy, optimized by gradient descent.The idealized initialization places the value vector in the span of the underlying feature, whereas the realistic setting uses Gaussian initialization.Source: §3 · pp. 6–7

4 Learning spatial structure via matching the labeling function

The analysis shows that gradient descent on the simplified ViT learns patch association while fitting the labeling function, first in an idealized population-risk setting and then with finite samples. The realistic result uses a polynomial number of samples.

4.1 Learning process in the idealized case

  • In the idealized setting, gradient descent learns patch association and fits the labeling function with high probability.The analysis states that patch association holds throughout training; its learning dynamics progress from value-vector updates to stronger within-set attention updates and renewed feature learning.Source: §4.1 · pp. 7–8
  • Permutation invariance and data symmetry reduce the idealized attention dynamics to a small set of coefficients.The attention matrix is summarized by diagonal, within-set, and between-set values, while the value vector remains aligned with the underlying feature.Source: §4.1 · p. 7

4.2 From the idealized to the realistic learning process

  • With N = poly(d) samples, gradient descent learns patch association and fits the labeling function with high probability in the realistic setting.The theorem reports label agreement with probability at least 1 − o(1).Source: §4.2 · pp. 8–9
  • The realistic-case proof compares the learned parameters with those of an intermediate semi-realistic process.The paper states that the realistic and semi-realistic parameters are close and uses a polynomial sample count.Source: §4.2 · pp. 8–9

5 Patch association yields sample-efficient fine-tuning with ViTs

The paper studies fine-tuning after pre-training when source and downstream datasets share structure but may have different features. It shows that fixing the pre-trained attention matrix permits transfer with few downstream samples, unlike algorithms without patch association knowledge.

Fine-tuning setup

  • The transfer setup keeps the pre-trained attention matrix fixed and optimizes only the value vector on downstream data.The downstream feature may differ from the source feature, while the structure is shared.Source: §5 · p. 9

Transfer result

  • With at most polylog(D) downstream samples, fine-tuning the value vector with pre-trained attention transfers to a dataset sharing the source structure.The theorem states that the transferred model achieves label agreement with probability at least 1 − o(1).Source: §5 · p. 9

Comparison without patch association

  • Algorithms without patch association knowledge need D^Ω(1) samples to achieve test error at most o(1) on the downstream distribution.The comparison concerns downstream data that share the source structure but can differ in their underlying feature.Source: §5 · p. 9

Implication

  • The theory identifies learning patch association as necessary for efficient transfer in this setting.The paper connects the result to learning the inductive bias of the labeling function, even if shared features may also be learned during pre-training.Source: §5 · p. 9

6 Numerical experiments

Experiments show positional attention is competitive with vanilla ViT, while ViTs preserve accuracy under grid permutations and recover spatial structure. The discussion notes that the analysis relies on positional attention and leaves patch-embedding effects open.

6 Numerical experiments

  • ViT accuracy remains unchanged across grid-shuffle sizes, whereas CNNs struggle with small shuffled grids and perform well when patches retain meaningful content.The ViT uses patch size 2 so patches are no larger than the shuffled grids.Source: §6 · p. 10, Figure 5 · p. 10
  • 68.9% test accuracy for positional attention versus 71.9% for vanilla ViT on ImageNet; it also competes on CIFAR-10/100 and SVHN.The CIFAR and SVHN comparison uses patch size 4 and averages over 10 seeds.Source: Figure 6 · p. 10, Figure 7 · p. 10, §6 · p. 10
  • After undoing the permutation, positional-encoding similarities recover spatially localized patterns, indicating structure beyond grouping nearby pixels.Before inversion, the displayed similarities on shuffled CIFAR-100 lack localized structure.Source: §6 · p. 10

Conclusion, limitations and future works

  • The analysis relies on positional attention, which separates patches from positional encodings; the effect of patch embeddings on learned inductive bias remains open.The discussion also notes that ViTs do not always learn the correct inductive bias on the Gaussian data.Source: Conclusion, limitations and future works · pp. 10–11

Checklist

The checklist reports that the paper states its theoretical assumptions and provides proofs, while also documenting experimental reproducibility and asset-related disclosures.

Checklist

  • The checklist reports that assumptions for theoretical results are stated and complete proofs are included in the appendix.It points to Section 3 for assumptions and the appendix for proofs.Source: Checklist · p. 15
  • The checklist reports that code, data, instructions, training details, error bars, and compute information are provided or referenced.It points to supplementary material, the appendix, and Section 6 for these items.Source: Checklist · p. 15
  • The checklist marks creator citation as addressed, while licensing, new assets, and human-subject disclosures are marked not applicable.The consent and participant-risk questions are also marked not applicable.Source: Checklist · p. 16

A.1 Details on the implementation

The appendix details small- and large-scale experiment configurations, reports additional positional-similarity plots, and introduces a softmax-denominator induction hypothesis and proof notation.

A.1 Details on the implementation

  • Small-scale experiments use a seven-layer, 12-head ViT with hidden and MLP dimensions 384; patch size is 2 for Figure 5 and 4 for Figure 7.Training uses Adam, cosine scheduling, 200 epochs, and batch size 128.Source: §A.1 · p. 17
  • ImageNet experiments use ViT-tiny-patch16-224 with CLS pooling, batch size 1024, and 300 epochs across 16 GPUs.The setup uses AdamW and training details based on the cited DeiT-B configuration.Source: §A.1 · p. 17

A.2 Additional plots

  • Across the plotted patches, positional-encoding similarities recover convolution-like patterns after the permutation is inverted.The appendix describes Figure 5 as showing two of the arrays displayed in Figure 8.Source: §A.2 · p. 17, Figure 8 · p. 18

B Induction hypothesis

  • The idealized-learning induction hypothesis assumes the softmax denominator is Θ(D) throughout the stated training interval.The denominator is expressed using Λ(t), Γ(t), Ξ(t), C, and D.Source: B Induction hypothesis · p. 17

C Notations

  • Appendix notation defines sigmoid, softmax, and per-example loss, and introduces a score matrix for the realistic-case analysis.The supplied score-matrix passage is partial.Source: C Notations · p. 18

D.1 Roadmap of the proof

The proof roadmap bounds gradient-descent dynamics across learning phases, including growth of the signal and control of attention and softmax-related quantities. It uses comparisons with a semi-idealized process and auxiliary bounds to establish low population risk and agreement with the labeling function.

D.1 Roadmap of the proof

  • After T iterations, α is large and population risk converges; the logistic-loss guarantee implies the learner fits the labeling function and proves Theorem 4.1’s first statement.The roadmap divides training into three events and treats the activation as σ(x) = x^p after initialization makes its linear part negligible.Source: §D.1 · p. 19

D.2 Event I: αptq initially increases

  • During Event I, α increases from a small initialization while the sigmoid is approximately constant, reaching Θ(1) by time T0.The proof identifies a value of α where the sigmoid ceases to be approximately constant; the supplied passages do not expose that threshold’s formula.Source: §D.2 · pp. 19–20

D.2.1 Auxiliary lemmas

  • Lemma D.3 states that learner model F is bounded for all (X, y) drawn from D.Source: §D.2.1 · p. 20

D.3 Event II: Γptq significantly increases

  • In Event II, α’s earlier increase drives Γ upward until S(t) is at least the maximum |N(τ)| over the training horizon.The proof characterizes Γ as non-decreasing and identifies T1 as the time when its bound reaches the stated threshold.Source: §D.3 · pp. 22–23

D.3.1 Auxiliary lemmas

  • For t ≥ T0, α(t) remains at least Ω(1), while population risk Ω(1) implies an expected sigmoid loss of at least Ω(1).The auxiliary argument bounds possible decreases in α after T0.Source: §D.3.1 · pp. 23–24

D.4 Event III: αptq keeps increases again

  • Once Γ reaches its Event II threshold, α is non-decreasing until population risk is o(1), and α(T) reaches polylog(d).The stated growth phase applies while the population loss is at least Ω(1).Source: §D.4 · p. 24

D.4.1 Auxiliary results

  • The auxiliary result interprets a small sigmoid term as making the sum of sigmoid terms over time bounded up to logarithmic dependence.The supplied passage does not show the full condition under which the sigmoid is small.Source: §D.4.1 · p. 24

D.5 Convergence rate of the population loss

  • From T1 onward, the population loss converges linearly to zero.The provided passages do not expose the full rate expression.Source: §D.5 · p. 25

D.7 Proof of the induction hypothesis

  • The induction proof bounds |ρ(t)| by Θ(1), keeps e^γ(t) between Ω(1) and λ₀, and gives e^β(t) = polyloglog(d).The proof also establishes a Θ(D) scale for the weighted combination of these parameters.Source: §D.7 · p. 27

E From idealized to real learning process

  • The proof bridges idealized and realistic learning by first increasing the realistic signal component to the idealized initialization threshold while keeping ε_v and attention components small, then comparing iterates.The threshold is ν^(1/(p−1)); the passage identifies separate subsubsections for the signal increase, the ε_v bound, and the attention-component bound.Source: E From idealized to real learning… · p. 28

E.1 Bound on the iterates during the initial steps (t P r0, T s)

  • For sufficiently small pα(t), Lemma E.8 treats the activation as linear, σ(x) = νx.This condition is used to account for the linear part of the activation under the small-variance initialization.Source: §E.1 · p. 28, §E.1.4 · p. 31

E.1.1 pαptq initially increases ´

  • The realistic signal parameter pα(t) is non-decreasing and reaches ν^(1/(p−1)) after T iterations.The proof sums its update rule to establish the hitting time T.Source: §E.1.1 · pp. 28–29

E.1.2 Bound on εv

  • During the initial phase, the orthogonal component ε_v stays small, and the attention entries pA_{a,b}(t) remain tiny.The supplied passages state these bounds but do not provide the complete formulas for their magnitudes.Source: §E.1.2 · pp. 29–30

E.1.4 Auxiliary lemmas

  • The initial-phase auxiliary lemmas bound the accumulated pα updates by summing the update rule.The cited passages also state that the per-step update has a Θ(Cην)e^βτ term.Source: §E.1.4 · p. 31
  • Auxiliary bounds use the activation derivative lower bound and an induction hypothesis to reduce a sufficient condition to a bound on α(t).The supplied proof fragment states the simplified condition as pα(t) ≤ ν^(1/(p−1)).Source: §E.1.4 · p. 31

E.2 Coupling between the semi-idealized and realistic processes (t P rT , Ts)

  • When the value-vector and attention discrepancies are small, the final realistic and semi-idealized iterates are equal, so the model fits the labeling function.The coupling argument bounds the realistic iterates through the semi-idealized process before establishing this conclusion.Source: §E.2 · p. 32

E.2.1 Defining the semi-idealized process

  • The semi-idealized process fixes its value vector along w* at scale α(t), trains only its attention matrix from zero off-diagonal initialization, and retains the idealized process's invariance and symmetry properties.It minimizes population risk with the fixed value vector and trainable attention matrix.Source: §E.2.1 · p. 32

E.2.2 Realistic dynamics are mainly on spanpw˚q

  • With polynomially many samples, the realistic value vector remains mainly in span(w*) throughout training, and the proof controls its orthogonal error through a recursion.The section states a high-probability polynomial-sample bound on that error, though the supplied passage does not preserve its full expression.Source: §E.2.2 · pp. 32, 34, 35

E.2.3 Auxiliary lemmas

  • For odd p ≥ 3 under the stated lower-bound condition on the coordinates, the auxiliary sigmoid-like function ψ is p-Lipschitz.The passage gives the condition as x_m^p ≥ −1/D and bounds |ψ(x) − ψ(y)| by p times the L1 distance.Source: §E.2.3 · p. 35

A stays small during the learning process

  • The proof controls the attention-matrix gap between realistic and semi-idealized training, showing its coefficients stay within an inverse-polynomial bound over the training interval.This is the stated conclusion of Lemma E.13 under gradient descent with the specified parametrization and polynomially many samples.Source: A stays small during the learning… · p. 35
  • For N = poly(d), Lemma E.15 bounds Δ^A(t) by 1/poly(d) for every t ≤ T.The proof also handles the learning interval in two phases, using an induction hypothesis and a product-term bound in the first phase.Source: A stays small during the learning… · p. 39

E.2.5 Auxiliary lemmas

  • For fixed vectors and odd p ≥ 3, Lemma E.16 bounds a softmax-weighted expression’s change using the maximum coordinate difference between its inputs.The expression uses softmax weights and fixed vectors X_m and v.Source: §E.2.5 · p. 39
  • Lemma E.17 bounds each coordinate’s softmax change by an exponential function of the maximum coordinate-wise input difference.This is stated for vectors a and b in R^D.Source: §E.2.5 · p. 40

E.2.6 Dynamics of pvptq

The analysis shows how the realistic training dynamics track an idealized process and culminate in generalization. In the realistic case, the learned model fits the labeling function with high probability.

E.2.6 Dynamics of pvptq

  • Under the stated dynamics, α is non-decreasing and eventually reaches polylog(d).The updates for α follow the stated analogues of Events I and III.Source: §E.2.6 · p. 40
  • By time T, the realistic iterates are close to the ideal ones, yielding nearby test loss and generalization.The analysis relates the value-vector gradient in the realistic process to that in the semi-idealized process.Source: §E.2.6 · pp. 40–41

E.2.7 The realistic model fits the labeling function

  • With high probability, the realistic model fits the labeling function, with population 0-1 error at most 1/poly(d).The passage also gives a high-probability correct-classification statement under the data distribution.Source: §E.2.7 · p. 42

F Transfer Learning

Pre-training on the structured dataset supports transfer to a downstream dataset with the same structure, even when its features differ. The section also states a sample lower bound for algorithms without patch-association knowledge and introduces gradient updates for the idealized process.

F Transfer Learning

  • Using at most polylog(D) samples to fine-tune the value vector, the pre-trained model transfers to a downstream dataset sharing the structure.The downstream dataset may have different features from the source dataset.Source: F Transfer Learning · p. 43
  • Without patch-association knowledge, a binary classification algorithm needs D^Ω(1) training samples to achieve test error at most o(1) on the stated distribution.Source: F Transfer Learning · p. 44

G Gradient descent updates in the idealized process

G.1 Indices in the same set: i, j P Sℓ

The section derives idealized-process updates for attention parameters and supplies auxiliary bounds and gradient expressions used in that analysis.

G.1 Indices in the same set: i, j P Sℓ

  • For indices i,j in the same set, the section derives gradient-based updates for γ(t) and ρ(t) during the idealized process.The stated results apply while the population loss is at most o(1).Source: §G.1 · pp. 44, 46, 50

G.3 Auxiliary lemmas

  • Two auxiliary lemmas provide high-probability bounds for the idealized learning process.Source: §G.3 · p. 50

H Gradients

  • The section gives expressions for the loss gradients with respect to the value vector v and attention entries A_i,j.Source: H Gradients · p. 51

I.1 Invariance of the parameters

The parameter invariants restrict the attention matrix to shared values within and across the dataset's sets, while the value vector remains aligned with the target direction. The paper also uses permutation invariance and examines why its data distribution is relevant.

I.1 Invariance of the parameters

  • At every stated iteration, the value vector has the form v(t)=α(t)w* for some scalar α(t).Source: §I.1 · p. 51
  • Throughout training, diagonal attention entries share a fixed value, within-set entries share γ(t), and cross-set entries share ρ(t).Source: §I.1 · pp. 51–52

I.2 Invariance by permutation

  • The data distribution and model are invariant under the specified permutations of set and within-set indices.Source: §I.2 · p. 53

J Justification of our data distribution

J.1 Generalized linear models poorly generalize

The appendix shows that generalized linear models can fail to fit the labeling function, while a one-layer transformer can fit it without patch association.

J.1 Generalized linear models poorly generalize

  • For a generalized linear model, the stated result is that g does not fit the labeling function, with Pr[f*(X)g(X) ≤ 0] ≥ 1/8.Source: §J.1 · p. 54
  • The proof compares opposite-label examples with specified values of Δ and uses their approximately matching data densities.The examples have Δ(X)=p−2C and Δ(X′)=p; their densities differ by a 1±o(1) factor.Source: §J.1 · p. 54

J.2 Classifiers fitting the labelling function without patch association

  • A one-layer transformer achieves error d^-ω(1) despite fitting the labeling function without patch association.The construction sets attention weights by patch-set membership.Source: §J.2 · p. 55

K.1 Tensor Power Method

The technical lemmas bound threshold-reaching times for recursively growing sequences, apply a concentration bound to a random count, and relate a loss to its derivative.

K.1 Tensor Power Method

  • The sequence lemmas bound the time to reach a threshold using lower and upper bounds on its recursive growth.The recursions use positive initial value, integer k>1, and positive growth parameters.Source: §K.1 · pp. 55–56

K.2 Probabilistic lemmas

  • A Chernoff bound is applied because the count of nonzero variables over Sℓ is binomial.The variables are independent and take values ±1 with probability q, and zero otherwise.Source: §K.2 · p. 57

K.3 Logarithmic inequalities

  • For x>0, S(−x) lies between 0.1 and 10 times log(1+exp(−x)).Source: §K.3 · p. 58
Loading 2210.09221v1…