# Vision Transformers provably learn spatial structure

Samy Jelassi
Princeton University
<EMAIL_ADDRESS>
&Michael E. Sander
Ecole Normale Supérieure
<EMAIL_ADDRESS>
&Yuanzhi Li
Carnegie Mellon University
<EMAIL_ADDRESS>

###### Abstract

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

Transformers are deep learning models built on self-attention (Vaswani et al.,
2017), and in the past several years they have increasingly formed the
backbone for state-of-the-art models in domains ranging from Natural Language
Processing (NLP) (Vaswani et al., 2017; Devlin et al., 2018) to computer
vision (Dosovitskiy et al., 2020), reinforcement learning (Chen et al., 2021a;
Janner et al., 2021), program synthesis (Austin et al., 2021) and symbolic
tasks (Lample and Charton, 2019). Beyond their remarkable performance, several
works reported the ability of transformers to simultaneously minimize their
training loss and learn inductive biases tailored to specific datasets e.g. in
computer vision (Raghu et al., 2021), in NLP (Brown et al., 2020; Warstadt and
Bowman, 2020) or in mathematical reasoning (Wu et al., 2021). In this paper,
we focus on computer vision where convolutions are considered to be an
adequate and biologically plausible inductive bias since they capture local
spatial information (Fukushima, 2003) by imposing a 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 at longer and longer length scales. Several seminal works (Cordonnier
et al., 2019; Dosovitskiy et al., 2020; Raghu et al., 2021) empirically show
that although randomly initialized, the positional encodings in Vision
transformers (ViTs) Dosovitskiy et al. (2020) actually learn this local
connectivity: closer patches have more similar positional encodings, as shown
in 1(a). 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 a collection of $D$ patches
$(\bm{X}_{1},\dots,\bm{X}_{D})\in\mathbb{R}^{d\times D}$, each of dimension
$d$. While each $\bm{X}_{i}$ represents (an embedding of) a spatially
localized portion of the original image, the relative positions of the patches
$\bm{X}_{i}$ in the image are disregarded. Instead, relative spatial
information is supplied through image-independent positional encodings
$\bm{P}=(\bm{p}_{1},\dots,\bm{p}_{D})\in\mathbb{R}^{d\times D}$. Unlike CNNs,
each layer of a ViT then learns, via trainable self-attention, a non-local set
of filters that non-linearly depend on both the values of all patches
$\bm{X}_{j}$ and their positional encodings $\bm{p}_{j}$.

(a)

(b)

Figure 1: (a) Visualization of the positional encodings similarities
$\bm{P}^{\top}\bm{P}=(\langle\bm{p}_{i},\bm{p}_{j}\rangle)_{(i,j)\in[D]^{2}}$
at initialization (1) and after training on Imagenet (2) using a "ViT-small-
patch32-224" (Dosovitskiy et al., 2020). We normalise the values
$\bm{P}^{\top}\bm{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 $\mathcal{S}_{\ell}$ as in Definition 2.1.
Squares in the same color belong to the same set $\mathcal{S}_{\ell}$. We
refer to (1) as a "spatially localized set" since all the elements in a
$\mathcal{S}_{\ell}$ are spatially contiguous. This is the type of sets
appearing in 1(a) at the end of training. Definition 2.1 also covers sets with
non-contiguous elements as (2).

##### Contributions.

The empirical observation of 1(a) sets a central question: from a theoretical
perspective, how do ViTs manage to learn these local connectivity patterns by
simply minimizing their training loss using gradient descent from random
initialization? While it is known that attention can express local operations
as convolution (Cordonnier et al., 2019), it remains unclear how ViTs learn
it. In this paper, we present a simple spatially-structured classification
dataset for which it is sufficient (but not necessary) to learn the structure
in order to generalize. We also present a 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 on a
dataset.

  * –

In Section 3, we introduce a structured classification dataset and a
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 a one-layer single-head ViT model trained with
gradient descent on our synthetic dataset performs patch association and
generalizes, in the idealized (Theorem 4.1) 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 dataset 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 a more general form of spatial structure that is not limited to local
patterns (Figure 3). We finally show that our ViT model –where the attention
matrix only depends on the positional encodings– is competitive with the
vanilla ViT on the ImageNet, CIFAR-10/100 and SVHNs datasets (Section 6 and
Section 6).

##### 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,\dots,D\\}.$
Any statement made "with high probability" holds with probability at least
$1-1/\mathrm{poly}(d).$ Given a vector $\bm{a}\in\mathbb{R}^{d}$ and $k\leq
d$, we define
$\mathrm{Top}_{k}\\{a_{j}\\}_{j=1}^{d}=\\{a_{i_{1}},\dots,a_{i_{k}}\\}$ where
$a_{i_{1}},\dots,a_{i_{k}}$ are the $k$-largest elements. For a function $F$
that implicitly depend on parameters $\bm{A}$ and $\bm{v}$, we often write
$F_{\bm{A},\bm{v}}$ to highlight its parameters. We use the asymptotic
complexity notations when defining the different constants.

## Related work

##### CNNs and ViTs.

Many computer vision architectures can be considered as a form of
hybridization between Transformers and CNNs. For example, DeTR (Carion et al.,
2020) use a CNN to generate features that are fed to a Transformer. (d’Ascoli
et al., 2021) show that self-attention can be initialized or regularized to
behave like a convolution and (Dai et al., 2021; Guo et al., 2021) add
convolution operations to Transformers. Conversely, (Bello et al., 2019;
Ramachandran et al., 2019; Bello, 2021) introduce self-attention or attention-
like operations to supplement or replace convolution in ResNet-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 (Bhojanapalli et al., 2021; Paul and Chen, 2021; Naseer et al.,
2021) or the effect of self-supervision (Caron et al., 2021; Chen et al.,
2021b). Closer to our work, some papers investigate why ViTs perform so well.
Raghu et al. (2021) compare the representations of ViTs and CNNs and Melas-
Kyriazi (2021); Trockman and Kolter (2022) 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.
(Vuckovic et al., 2020; Edelman et al., 2021) addressed this question in the
context of self-attention blocks and (Dehghani et al., 2018; Wei et al., 2021;
Hron et al., 2020) for Transformers. On the optimization side, (Zhang et al.,
2020) investigate the role of adaptive methods in attention models and (Snell
et al., 2021) analyze the dynamics of a single-head attention head to
approximate the learning of a Seq2Seq architecture. In our work, we also
consider a single-head ViT trained with gradient descent and exhibit a setting
where it provably learns convolution-like patterns and generalizes.

##### Algorithmic regularization.

The question we address concerns algorithmic regularization which
characterizes the generalization of an optimization algorithm when multiple
global solutions exist in over-parametrized 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 scattered in
the space that vastly differ in terms of generalization. Algorithmic
regularization appears in binary classification (Soudry et al., 2018; Lyu and
Li, 2019; Chizat and Bach, 2020), matrix factorization (Gunasekar et al.,
2018; Arora et al., 2019), convolutional neural networks (Gunasekar et al.,
2018; Jagadeesan et al., 2022), generative adversarial networks (Allen-Zhu and
Li, 2021), contrastive learning (Wen and Li, 2021) and mixture of experts
(Chen et al., 2022). Algorithmic regularization is induced by and depends on
many factors such as learning rate and batch size (Goyal et al., 2017; Hoffer
et al., 2017; Keskar et al., 2016; Smith et al., 2018; Li et al., 2019),
initialization Allen-Zhu and Li (2020), momentum (Jelassi and Li, 2022),
adaptive step-size (Kingma and Ba, 2014; Neyshabur et al., 2015; Daniely,
2017; Wilson et al., 2017; Zou et al., 2021; Jelassi et al., 2022), batch
normalization (Arora et al., 2018; Hoffer et al., 2019; Ioffe and Szegedy,
2015) and dropout (Srivastava et al., 2014; Wei et al., 2020). However, all
these works consider the case of feed-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 $\mathcal{D}$ be a distribution over $\mathbb{R}^{d\times
D}\times\\{-1,1\\}$ where each patch
$\bm{X}=(\bm{X}_{1},\dots,\bm{X}_{D})\in\mathbb{R}^{d\times D}$ has label
$y\in\\{-1,1\\}$. We say that $\mathcal{D}$ is spatially structured if

  * –

there exists a partition of $[D]$ into $L$ disjoint subsets i.e.
$[D]=\bigcup_{\ell=1}^{L}\mathcal{S}_{\ell}$ with
$\mathcal{S}_{\ell}\subsetneq D$ and $|\mathcal{S}_{\ell}|=C$.

  * –

there exists a labeling function $f^{*}$ satisfying
$\mathbb{P}[yf^{*}(\bm{X})>0]=1-d^{-\omega(1)}$ and,

$\displaystyle
f^{*}(\bm{X}):=\sum_{\ell\in[L]}\phi(\big{(}\bm{X}_{i}\big{)}_{i\in\mathcal{S}_{\ell}}),\quad\text{where
}\phi\colon\mathbb{R}^{d\times C}\rightarrow\mathbb{R}\text{ is an arbitrary
function}.$ (1)

[13]r0.45

Left: Test error of the ViT on the convolution structured dataset. Upper
Right: Grid displaying the input patches. Yellow squares represent spatially
localized sets $\mathcal{S}_{\ell}$. Those sets are taken into account when
computing the convolutional function $f^{*}$. Lower Right: Learnt
$\bm{P}^{\top}\bm{P}$ looks random compared to upper one.

##### Examples.

A particular case for the sets $\mathcal{S}_{\ell}$’s is the one of spatially
localized sets as in 1(b)-(1). In this case, we have $D=16$, $C=4$ and
$\mathcal{S}_{1}=\\{1,2,5,6\\},\;\mathcal{S}_{2}=\\{3,4,7,8\\},\;\mathcal{S}_{3}=\\{9,10,13,14\\},\;\mathcal{S}_{4}=\\{11,12,15,16\\}.$
We emphasize that Definition 2.1 is not limited to spatially localized sets
and also covers non-contiguous sets as 1(b)-(2).

##### Labelling function

Definition 2.1 states that there exists a labelling function that preserves
the underlying structure by applying the same function $\phi$ to each
$\mathcal{S}_{\ell}$ as in (1). For instance, when the sets
$\mathcal{S}_{\ell}$’s are spatially localized, $f^{*}$ can be a one-hidden
layer convolutional network. In this paper, we are interested in patch
association which refers to the ability of an algorithm to identify the sets
$\mathcal{S}_{\ell}$’s, and is formally defined as follow.

######  Definition 2.2 (Patch association for ViTs).

Let $\mathcal{D}$ be as in Definition 2.1. Let
$\mathcal{M}\colon\mathbb{R}^{d\times D}\rightarrow\\{-1,1\\}$ be a
transformer and $\bm{P}^{(\mathcal{M})}$ its positional encodings matrix. We
say that $\mathcal{M}$ performs patch association on $\mathcal{D}$ if for all
$\ell\in[L]$ and $i\in\mathcal{S}_{\ell}$, we have
$\mathrm{Top}_{C}\;\\{\langle\bm{p}_{i}^{(\mathcal{M})},\bm{p}_{j}^{(\mathcal{M})}\rangle\\}_{j=1}^{D}=\mathcal{S}_{\ell}.$

Definition 2.2 states that patch association is learned when for a given
$i\in\mathcal{S}_{\ell},$ its positional encoding mainly attends those of $j$
such that $i,j\in\mathcal{S}_{\ell}$. In this way, the transformer groups the
$\bm{X}_{i}$ according to $\mathcal{S}_{\ell}$ just like the true labeling
function. Definition 2.2 formally describes the empirical findings in
1(a)-(2), where nearby patches have similar positional encodings. A natural
question is then: would ViTs really learn those $\mathcal{S}_{\ell}$ after
training to match the labeling function $f^{*}$? Without further assumptions
on the data distribution, we next show that the answer is no.

##### ViTs do not always learn patch association under Assumption 1.

We give a negative answer through the following synthetic experiment. Consider
the case where all the patches $\bm{X}_{j}$ are i.i.d. standard Gaussian and
$f^{*}$ is a one-hidden layer CNN with cubic activation. The label $y$ of any
$\bm{X}$ is then given by $y=\mathrm{sign}(f^{*}(\bm{X}))$. As shown in
Section 2, one-layer ViT reaches small test error on the binary classification
task. However, $\bm{P}^{\top}\bm{P}$ does not match the convolution pattern
encoded in $f^{*}$. This is not surprising, since the data distribution
$\mathcal{D}$ is Gaussian, and thus lacks spatial structure. Thus, in order to
prove that ViTs learn patch association, we need additional assumptions on
$\mathcal{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 we use to classify it.

######  Assumption 1 (Data distribution with specific spatial structure).

Let $\mathcal{D}$ be a distribution as in Definition 2.1 and
$\bm{w}^{*}\in\mathbb{R}^{d}$ be an underlying feature. We suppose that each
data-point $\bm{X}$ is defined as follow

  * –

Uniformly sample an index $\ell(\bm{X})$ from $[L]$ and for
$j\in\mathcal{S}_{\ell(\bm{X})}$, $\bm{X}_{j}=y\bm{w}^{*}+\bm{\xi}_{j}$, where
$y\bm{w}^{*}$ is the informative feature and
$\bm{\xi}_{j}\overset{i.i.d.}{\sim}\mathcal{N}(0,\sigma^{2}(\mathbf{I}_{D}-\bm{w}^{*}\bm{w}^{*\;\top}))$
(signal set).

  * –

For $\ell\in[L]\backslash\\{\ell(\bm{X})\\}$ and $j\in\mathcal{S}_{\ell}$,
$\bm{X}_{j}=\delta_{j}\bm{w}^{*}+\bm{\xi}_{j}$, where $\delta_{j}=1$ with
probability $q/2$, $-1$ with same probability and $0$ otherwise, and
$\bm{\xi}_{j}\overset{i.i.d.}{\sim}\mathcal{N}(0,\sigma^{2}(\mathbf{I}_{D}-\bm{w}^{*}\bm{w}^{*\;\top}))$
(random sets).

[20]r0.35  Visualization of a data-point $\bm{X}$ in $\mathcal{D}$ when the
$\mathcal{S}_{\ell}$’s are spatially localized. Each square depicts a patch
$\bm{X}_{j}$ and squares of the same color belong to the same set
$\mathcal{S}_{\ell}.$ "0" indicates that the patch does not have a feature,
"1" stands for feature $1\cdot\bm{w}^{*}$ and "-1" for feature
$-1\cdot\bm{w}^{*}$. The large red square depicts the signal set
$\ell(\bm{X}).$ Although there are more "-1"’s than "+1"’s, the label of
$\bm{X}$ is $+1$ since there are only "+1"’s inside the signal set.  To keep
the analysis simple, the noisy patches are sampled from the orthogonal
complement of $\bm{w}^{*}.$ Note that $\mathcal{D}$ admits the labeling
function
$f^{*}(\bm{X})=\sum_{\ell\in[L]}\mathrm{Threshold}_{0.9C}(\sum_{i\in\mathcal{S}_{\ell}}\langle\bm{w}^{*},\bm{X}_{i}\rangle)$,
where $\mathrm{Threshold}_{C}(z)=z$ if $|z|>{C}$ and $0$ otherwise.

We sketch a data-point of $\mathcal{D}$ in Section 3. Our dataset can be
viewed as an extreme simplification of real-world image datasets where there
is a set of adjacent patches that contain a useful feature (e.g. the nose of a
dog) and many patches that have uninformative or spurious features e.g. the
background of the image. We make the following assumption on the parameters of
the data distribution.

######  Assumption 2.

We suppose that $d=\mathrm{poly}(D)$, $C=\mathrm{polylog}(d)$,
$q=\mathrm{poly}(C)/D$, $\|\bm{w}^{*}\|_{2}=1$ and $\sigma^{2}=1/d$. This
implies $C\ll D$ and $q\ll 1.$

Assumption 2 may be justified by considering a "ViT-base-patch16-224" model
Dosovitskiy et al. (2020) on ImageNet. In this case, $d=384$, $D=196$.
$\sigma$ is set to have $\|\bm{\xi}_{j}\|_{2}\approx\|\bm{w}^{*}\|_{2}$. $q$
is chosen so that there are more spurious features than informative ones (low
signal-to-noise regime) which makes the data _non-linearly separable_. Our
dataset is non-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 $\mathcal{D}$ be as in Assumption 1. Let
$g(\bm{X})=\phi\left(\sum_{j=1}^{D}\langle\bm{w}_{j},\bm{X}_{j}\rangle\right)$
be a generalized linear model. Then, $g$ does not fit the labeling function
i.e. $\mathbb{P}[f^{*}(\bm{X})g(\bm{X})\leq 0]\geq 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
$\mathcal{S}_{\ell}$’s in order to generalize. In addition, the following
Theorem 3.2 states the existence of a transformer that generalizes without
learning spatial structure (see Appendix J for a proof), thus showing that the
learning process has a priori no straightforward reason to lead to patch
association.

######  Theorem 3.2.

Let $\mathcal{D}$ be defined as in Assumption 1. There exists a (one-layer)
transformer $\mathcal{M}$ so that
$\mathbb{P}[f^{*}(\bm{X})\mathcal{M}(\bm{X})\leq 0]=d^{-\omega(1)}$ but for
all $\ell\in[L]$, $i\in\mathcal{S}_{\ell}$,
$\mathrm{Top}_{C}\;\\{\langle\bm{p}_{i}^{(\mathcal{M})},\bm{p}_{j}^{(\mathcal{M})}\rangle\\}_{j=1}^{D}\cap\mathcal{S}_{\ell}=\emptyset$.

##### Simplified ViT model.

We now define our simplified ViT model for which we show in Section 4 that it
implicitly learns patch association via minimizing its training objective. We
first remind the self-attention mechanism that is ubiquitously used in
transformers.

######  Definition 3.1 (Self-attention (Bahdanau et al., 2014; Vaswani et al.,
2017)).

The attention mechanism (Bahdanau et al., 2014; Vaswani et al., 2017) in the
single-head case is defined as follow. Let $\bm{X}\in\mathbb{R}^{d\times D}$ a
data point and $\bm{P}\in\mathbb{R}^{d\times D}$ its positional encoding. The
self-attention mechanism computes

  1. 1.

the sum of patches and positional encodings i.e.
${\large\boldsymbol{\mathpzc{X}}}=\bm{X}+\bm{P}.$

  2. 2.

the attention matrix $\bm{A}=\bm{Q}\bm{K}^{\top}$ where
$\bm{Q}={\large\boldsymbol{\mathpzc{X}}}^{\top}\bm{W}_{\bm{Q}}$,
$\bm{K}={\large\boldsymbol{\mathpzc{X}}}^{\top}\bm{W}_{\bm{K}}$,
$\bm{W}_{\bm{Q}},\bm{W}_{\bm{K}}\in\mathbb{R}^{d\times d}$.

  3. 3.

the score matrix $\bm{S}\in\mathbb{R}^{D\times D}$ with coefficients
$S_{i,j}=\exp(A_{i,j}/\sqrt{d})/\sum_{r=1}^{D}\exp(A_{i,r}/\sqrt{d})$.

  4. 4.

the matrix $\bm{V}={\large\boldsymbol{\mathpzc{X}}}^{\top}\bm{W}_{\bm{V}}$,
where $\bm{W}_{\bm{V}}\in\mathbb{R}^{d\times d}.$

It finally outputs
$\mathrm{SA}((\bm{X};\bm{P}))=\bm{S}\bm{V}\in\mathbb{R}^{d\times D}.$

In this paper, our ViT model relies on a different attention mechanism –the
"positional attention"– that we define as follows.

######  Definition 3.2 (Positional attention).

Let $\bm{X}\in\mathbb{R}^{d\times D}$ and $\bm{P}\in\mathbb{R}^{d\times D}$
the positional encoding. The positional attention mechanism takes as input the
pair $(\bm{X};\bm{P})$ and computes:

  1. 1.

the attention matrix $\bm{A}=\bm{Q}\bm{K}^{\top}$ where
$\bm{Q}=\bm{P}^{\top}\bm{W}_{\bm{Q}}$, $\bm{K}=\bm{P}^{\top}\bm{W}_{\bm{K}}$
and $\bm{W}_{\bm{Q}},\bm{W}_{\bm{K}}\in\mathbb{R}^{d\times d}$.

  2. 2.

the score matrix $\bm{S}\in\mathbb{R}^{D\times D}$ with coefficients
$S_{i,j}=\exp(A_{i,j}/\sqrt{d})/\sum_{r=1}^{D}\exp(A_{i,r}/\sqrt{d})$.

  3. 3.

the matrix $\bm{V}=\bm{X}^{\top}\bm{W}_{\bm{V}}$, where
$\bm{W}_{\bm{V}}\in\mathbb{R}^{d\times d}.$

It outputs $\mathrm{PA}((\bm{X};\bm{P}))=\bm{S}\bm{V}.$

Positional attention isolates positional encoding $\bm{P}$ from data $\bm{X}$:
$\bm{A}$ encodes the dynamics of $\bm{P}$ and tracks whether patch association
is learned. $\bm{V}$ 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 is similar to self-attention. As this latter, positional attention
is also permutation-invariant and processes all tokens simultaneously.
Besides, positional attention also computes a score matrix between the
different tokens. This similarity matrix is also normalized in a sparse manner
with the Softmax operator. The only aspect that positional attention misses
from self-attention is the fact that $\bm{S}$ 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
in the parameters to ease our analysis.

######  Simplification 3.1.

In the positional attention mechanism, we set $d=D$,
$\bm{W}_{\bm{K}}=\bm{I}_{D}$ and $\bm{W}_{\bm{Q}}=\bm{I}_{D}$ which implies
$\bm{A}=\bm{P}^{\top}\bm{P}.$ We set
$\bm{W}_{\bm{V}}=[\bm{v},\dots,\bm{v}]\in\mathbb{R}^{d\times D}$ where
$\bm{v}\in\mathbb{R}^{d}.$ Finally, we set $\bm{A}$ and $\bm{v}$ as trainable
parameters. Besides, without loss of generality, we train all $A_{i,j}$ for
$i\neq j$ and leave the diagonals of $\bm{A}$ fixed.

In Simplification 3.1, we set $\bm{W}_{\bm{K}}$ and $\bm{W}_{\bm{Q}}$ to the
identity so that $\bm{A}=\bm{P}^{\top}\bm{P}.$ This Gram matrix encodes the
spatial patterns learned by the ViT as shown in 1(a). Besides, since fitting
the labeling function requires to learn one feature $\bm{w}^{*}$, it is
sufficient to parameterize $\bm{W}_{\bm{V}}$ with a vector $\bm{v}$. Also,
although $\bm{A}=\bm{P}^{\top}\bm{P}$ and $\bm{P}$ is trainable, we choose for
simplicity to only optimize over $\bm{A}.$ Besides, we leave the $A_{i,i}$’s
fixed because Softmax is invariant under the uniform shift of the input. Under
Simplification 3.1, our simplified ViT model is then a two attention layer
with a single head:

$F(\bm{X})=\sum_{i=1}^{D}\sigma\bigg{(}D\sum_{j=1}^{D}S_{i,j}\langle\bm{v},\bm{X}_{j}\rangle\bigg{)}\quad\mathrm{with}\quad
S_{i,j}=\exp(A_{i,j}/\sqrt{d})/\sum_{r=1}^{D}\exp(A_{i,r}/\sqrt{d}),$ (T)

where $\sigma$ is an activation function. Since we aim to the simplest ViT
model, we opt for a polynomial activation i.e. $\sigma(x)=x^{p}+\nu x$ where
$p\geq 3$ is an odd integer and $\nu=1/\mathrm{poly}(d)$. Note that this
choice of polynomial activation is common in the deep learning theory
literature – see e.g. (Li et al., 2018; Allen-Zhu and Li, 2020; Woodworth et
al., 2020) among others. The degree $p$ is odd to make the ViT model
compatible with the labeling function and strictly larger than 1 because the
data is not linearly separable (Theorem 3.1). We add a linear part in the
activation function to ensure that the gradient is non-zero when $\bm{v}$ has
small coefficients. With these simplifications, we formally prove that $F$ is
able to learn patch association and generalize, in the two following settings.

##### Idealized and realistic learning problems.

Given a dataset $\mathcal{Z}=\\{(\bm{X}[i],y[i])\\}_{i=1}^{N}$ sampled from
$\mathcal{D}$, we solve the empirical risk minimization problem for the
logistic loss defined by:

$\displaystyle\min_{\widehat{\bm{A}},\hat{\bm{v}}}\;\;\frac{1}{N}\sum_{i=1}^{N}\log\big{(}1+e^{-y[i]F(\bm{X}[i])}\big{)}:=\widehat{\mathcal{L}}(\widehat{\bm{A}},\widehat{\bm{v}}).$
(E)

Instead of directly analyzing (E), we introduce a proxy where we minimize the
population risk

$\displaystyle\min_{\bm{A},\bm{v}}\;\;\mathbb{E}_{\mathcal{D}}\big{[}\log\big{(}1+e^{-yF(\bm{X})}\big{)}\big{]}:=\mathcal{L}(\bm{A},\bm{v}).$
(P)

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 $t\in[T]$ and $i,j\in[D]$ is

$A_{i,j}^{(t+1)}=A_{i,j}^{(t)}-\eta\partial_{A_{i,j}}\mathcal{L}(\bm{A}^{(t)},\bm{v}^{(t)}),\quad\bm{v}^{(t+1)}=\bm{v}^{(t)}-\eta\nabla_{\bm{v}}\mathcal{L}(\bm{A}^{(t)},\bm{v}^{(t)}),$
(GD)

where $\eta>0$ is 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\geq\mathrm{poly}(d)/\eta.$ We set the learning rate as
$\eta\in\big{(}0,\frac{1}{\mathrm{poly}(d)}\big{)}$. The diagonal coefficient
of the attention matrix are set for $i\in[D]$ as
$A_{i,i}^{(0)}=\widehat{A}_{i,i}^{(0)}=\sigma_{\bm{A}}\mathbf{I}_{D}$ where
$\sigma_{\bm{A}}=\mathrm{polyloglog}(d).$ The off-diagonal coefficients of
$\bm{A}$ and the value vector are initialized as:

  1. 1.

Idealized case: $\bm{v}^{(0)}=\alpha^{(0)}\bm{w}^{*}$ where
$\alpha^{(0)}=\nu^{1/(p-1)}$ and $A_{i,j}^{(0)}=0$ for $i\neq j.$

  2. 2.

Realistic case:
$\widehat{\bm{v}}^{(0)}\sim\mathcal{N}(0,\omega^{2}\mathbf{I}_{d})$ and
$\widehat{A}_{i,j}^{(0)}\sim\mathcal{N}(0,\omega^{2})$ where $i\neq j$ and
$\omega=1/\mathrm{poly}(d)$.

We remind that in Simplification 3.1, we have $\bm{A}=\bm{P}^{\top}\bm{P}$. If
one initializes $\bm{P}\sim\mathcal{N}(0,\sigma_{A}\mathbf{I}_{D}/D)$, then
with high probability,
$A_{i,i}^{(0)}=\|\bm{p}_{i}^{(0)}\|_{2}^{2}=\Theta(\sigma_{A})$ and
$A_{i,j}^{(0)}=\langle\bm{p}_{i}^{(0)},\bm{p}_{j}^{(0)}\rangle=\Theta(\sigma_{A}/\sqrt{D})$
for $i\neq j$. Since $D\gg 1$, it is then reasonable to set $A_{i,j}^{(0)}=0$.
Note that, also in the idealized setting, we initialize $\bm{v}^{(0)}$ in
$\mathrm{span}(\bm{w}^{*})$, even though this latter should be unknown to the
algorithm. We remind that the idealized case is a proxy to ultimately
characterize the realistic dynamics.

##  4 Learning spatial structure via matching the labeling function

As announced above, we show that our ViT (T) implicitly learns patch
association and fits the labeling function by minimizing the training
objective. We first study the dynamics in (P). Using the analysis 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 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 we run GD on (P) for $T$ iterations with parameters set as in
Parametrization 3.1. With high probability, the ViT model (T)

  1. 1.

learns patch association i.e. for all $\ell\in[L]$ and
$i\in\mathcal{S}_{\ell}$,
$\mathrm{Top}_{C}\;\\{A_{i,j}^{(T)}\\}_{j=1}^{D}=\mathcal{S}_{\ell}.$

  2. 2.

learns the labeling function $f^{*}$ i.e.
$\mathbb{P}_{\mathcal{D}}[f^{*}(\bm{X})F_{\bm{A}^{(T)},\bm{v}^{(T)}}(\bm{X})>0]\geq
1-o(1).$

Figure 2: 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,j)$ such that
$(i,j)\in\mathcal{S}_{\ell}$. We then plot the learned matrix $\bm{A}$ and see
that coefficients with high value exactly correspond to their grey scale
counterpart in the left plot. We also display test error and cosine similarity
between $\bm{w}^{*}$ and $\bm{v}$ w.r.t the number of training steps.

We now sketch the main ideas to prove the theorem for which one can refer to
Appendix D for a complete proof.

##### Invariance and symmetries.

In (P), we take the expectation over $\mathcal{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 $\bm{A}$ can be reduced to three variables in (P).

######  Lemma 4.1.

There exist $\beta=\sigma_{\bm{A}}$, $\gamma^{(t)},\rho^{(t)}\in\mathbb{R}$
such that for all $t\geq 0$:

  1. 1.

for all $i\in[D]$, $A_{i,i}^{(t)}=\beta.$

  2. 2.

for all $i,j\in[D]$ such that $i,j\in\mathcal{S}_{\ell}$ for some
$\ell\in[L]$, $A_{i,j}^{(t)}=\gamma^{(t)}.$

  3. 3.

for all $i,j\in[D]$ such that $i\in\mathcal{S}_{\ell}$ and
$j\in\mathcal{S}_{m}$ for some $\ell,m\in[L]$ with $\ell\neq m$,
$A_{i,j}^{(t)}=\rho^{(t)}.$

Besides, using the initialization in Parametrization 3.1, we can show that
$\bm{v}$ always lies in $\mathrm{span}(\bm{w}^{*}).$

######  Lemma 4.2.

For all $t\in[T]$, there exists $\alpha^{(t)}\in\mathbb{R}$ such that
$\bm{v}^{(t)}=\alpha^{(t)}\bm{w}^{*}.$

In summary, Lemma 4.1 and Lemma 4.2 imply that instead of optimizing over
$\bm{A}$ and $\bm{v}$ in (P), we can instead consider the scalar variables
$\alpha^{(t)}$, $\gamma^{(t)}$ and $\rho^{(t)}$. The remaining of this section
consists in analyzing the dynamics of these three quantities.

##### Learning patch association.

We first analyze the dynamics of $\gamma^{(t)}$ and $\rho^{(t)}$. To this end,
we introduce the following terms:

$\displaystyle\Lambda^{(t)}$
$\displaystyle=\frac{e^{\beta}}{e^{\beta}+(C-1)e^{\gamma^{(t)}}+(D-C)e^{\rho^{(t)}}},$
$\displaystyle\Gamma^{(t)}=\frac{e^{\gamma^{(t)}}}{e^{\beta}+(C-1)e^{\gamma^{(t)}}+(D-C)e^{\rho^{(t)}}},$
$\displaystyle\Xi^{(t)}$
$\displaystyle=\frac{e^{\rho^{(t)}}}{e^{\beta}+(C-1)e^{\gamma^{(t)}}+(D-C)e^{\rho^{(t)}}},$
$\displaystyle G^{(t)}=D(\Lambda^{(t)}+(C-1)\Gamma^{(t)}).\hskip 56.9055pt$

Note that $\Lambda^{(t)}$, $\Gamma^{(t)}$ and $\Xi^{(t)}$ respectively
correspond to the coefficients on the diagonal, those for which
$i,j\in\mathcal{S}_{\ell}$ for some $\ell\in[L]$ and all the other
coefficients of the attention matrix $\bm{S}$. Using these notations, we first
derive the GD updates of $\gamma^{(t)}$ and $\rho^{(t)}.$

######  Lemma 4.3.

Let $t\leq T$. The attention weights $\gamma^{(t)}$ and $\rho^{(t)}$ satisfy:

$\displaystyle\gamma^{(t+1)}$
$\displaystyle=\gamma^{(t)}+\eta\mathrm{polylog}(d)(\alpha^{(t)})^{p}\cdot\Gamma^{(t)}(G^{(t)})^{p-1},$
$\displaystyle|\rho^{(t+1)}|$
$\displaystyle\leq|\rho^{(t)}|+\eta\mathrm{polylog}(d)(\alpha^{(t)})^{p}\Big{(}\frac{1}{D}+\frac{1}{D}\Gamma^{(t)}(G^{(t)})^{p-1}\Big{)}.$

Lemma 4.3 shows that the increment of $\gamma^{(t)}$ is larger than the one of
$\rho^{(t)}$. Since $\gamma^{(0)}=\rho^{(0)}=0$, this implies that
$\gamma^{(t)}\geq\rho^{(t)}$ for all $t\geq 0.$ This observation proves the
first item of Theorem 4.1. We now explain how learning patch association leads
to $\bm{v}$ highly correlated with $\bm{w}^{*}.$

  * –

Event I: At the beginning of the process, the update of $\bm{v}^{(t)}$ is
larger than the one of $A_{i,j}^{(t)}$ which implies that only $\bm{v}^{(t)}$
updates during this first phase. We show that
$\alpha^{(t)}=\langle\bm{v}^{(t)},\bm{w}^{*}\rangle$ increases until a time
$\mathcal{T}_{0}>0$ where it reaches some threshold (Lemma D.2). 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: During this phase, the attention weights must update. Indeed, assume
by contradiction that the $A_{i,j}^{(t)}$ stay around initialization and that
$\bm{v}^{(t)}$ is optimal i.e. $\bm{v}^{(t)}=a^{(t)}\bm{w}^{*}$ where
$a^{(t)}\gg 1.$ Then, the predictor $g$ we would have is

$\displaystyle
g(\bm{X})=\sum_{i=1}^{D}\sum_{j=1}^{D}S_{i,j}^{(0)}\langle\bm{v}^{(t)},\bm{X}_{j}\rangle\propto\sum_{i=1}^{D}\sum_{j=1}^{D}e^{A_{i,j}^{(0)}}\langle\bm{w}^{*},\bm{X}_{j}\rangle$
(2)

Such predictor $g$ would yield high population loss because there many more
data with random labels ($qD=\mathrm{poly}(C)$) than with the exact label.
Therefore, $A_{i,j}^{(t)}$’s start to update. The gradient increment for
$\gamma^{(t)}$ (which corresponds to $i$ and $j$ in the same set
$\mathcal{S}_{\ell}$) is much larger than the one for $\rho^{(t)}$ (Lemma
4.3). Thus, $\gamma^{(t)}$ increases until a time
$\mathcal{T}_{1}\in[\mathcal{T}_{0},T]$ such that
$\gamma^{(\mathcal{T}_{1})}>\max_{t\in[T]}|\rho^{(t)}|$.

  * –

Event III: Because we have
$\gamma^{(\mathcal{T}_{1})}>\max_{t\in[T]}|\rho^{(t)}|$, we again have
$\alpha^{(t+1)}>\alpha^{(t)}$ as in Phase I (Lemma D.11). Thus, $\alpha^{(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 $t$ (Lemma 4.3).

  * –

our ViT model uses patch association to minimize the population loss (Event
III). Without patch association, the model would only be a generalized linear
model that does not minimize the loss.

###  4.2 From the idealized to the realistic learning process

The real learning process differs from the idealized one in that we have a
finite number of samples and we initialize both $\widehat{\bm{A}}$ and
$\widehat{\bm{v}}$ as Gaussian random variables. Using a 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
Parametrization 3.1. Assume that the number of samples is
$N=\mathrm{poly}(d).$ With high probability, the model

  1. 1.

learns patch association i.e. for all $\ell\in[L]$ and
$i\in\mathcal{S}_{\ell}$,
$\mathrm{Top}_{C}\;\\{\widehat{A}_{i,j}^{(T)}\\}_{j=1}^{D}=\mathcal{S}_{\ell}.$

  2. 2.

fits the labeling function i.e.
$\mathbb{P}_{\mathcal{D}}[f^{*}(\bm{X})F_{\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)}}(\bm{X})>0]\geq
1-o(1).$

Similarly to Li et al. (2020), the proof introduces a "semi-realistic"
learning process that is a mid-point between the idealized and realistic
processes. We show that $\widehat{\bm{A}}^{(T)}$ and $\widehat{\bm{v}}^{(T)}$
are close to their semi-realistic counterparts – see Appendix E for a complete
proof. Figure 2 numerically illustrates Theorem 4.2.

##  5 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 $\mathcal{D}$, our model (T) sample-efficiently transfers to
datasets that are structured as $\mathcal{D}$ but differ in their features.

##### Downstream dataset.

Let $\widetilde{\mathcal{D}}$ a downstream data distribution defined as in
Assumption 1 such that its underlying feature is $\widetilde{\bm{w}}^{*}$ with
$\|\widetilde{\bm{w}}^{*}\|_{2}=1$ and $\widetilde{\bm{w}}^{*}$ potentially
different from $\bm{w}^{*}$. In other words, the downstream
$\widetilde{\mathcal{D}}$ and source $\mathcal{D}$ distributions share the
same structure but not necessarily the same feature. We sample a downstream
dataset
$\widetilde{\mathcal{Z}}=\\{(\widetilde{\bm{X}}[i],\widetilde{y}[i])\\}_{i=1}^{\widetilde{N}}$
from $\widetilde{\mathcal{D}}$.

##### Learning problem.

We consider the model (T) pre-trained as in subsection 4.2. We assume that
$\widehat{\bm{A}}$ is kept fixed from the pre-trained model and we only
optimize the value vector $\widetilde{\bm{v}}$ to solve:

$\displaystyle\min_{\widetilde{\bm{v}}}\;\;\frac{1}{\widetilde{N}}\sum_{i=1}^{\widetilde{N}}\log\big{(}1+e^{-\widetilde{y}[i]F(\widetilde{\bm{X}}[i])}\big{)}:=\widetilde{\mathcal{L}}(\widetilde{\bm{v}}).$
($\widetilde{\text{E}}$)

We run GD on ($\widetilde{\text{E}}$) with parameters set as in
Parametrization 3.1 except that the $\widehat{A}_{i,j}$’s are fixed and
$\widetilde{\bm{v}}^{(0)}\sim\mathcal{N}(0,\omega^{2}\mathbf{I}_{d})$ with
$\omega=1/\mathrm{poly}(d)$. Our main results states that this fine-tuning
procedure requires a few samples to achieve high test accuracy in
$\widetilde{\mathcal{D}}$. In contrast, any algorithm without patch
association needs a large number of samples to generalize.

######  Theorem 5.1.

Let $\widehat{\bm{A}}$ be the attention matrix obtained after pre-training as
in subsection 4.2. Assume that we run GD for $T$ iterations on
($\widetilde{\text{E}}$) to fine-tune the value vector. Using
$\widetilde{N}\leq\mathrm{polylog}(D)$ samples, the model (T) transfers to
$\widetilde{\mathcal{D}}$ i.e.
$\mathbb{P}_{\widetilde{\mathcal{D}}}[f^{*}(\bm{X})F_{\widehat{\bm{A}},\widetilde{\bm{v}}^{(T)}}(\bm{X})>0]\geq
1-o(1).$

######  Theorem 5.2.

Let $\mathcal{A}\colon\mathbb{R}^{d\times D}\rightarrow\\{-1,1\\}$ be a binary
classification algorithm without patch association knowledge. Then, it needs
$D^{\Omega(1)}$ training samples to get test error $\leq o(1)$ on
$\widetilde{\mathcal{D}}$.

The proofs of Theorem 5.1 and Theorem 5.2 are in Appendix F. These theorems
hightlight that learning 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 is 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.

(a)

(b)

(c)

Figure 3: (a): Test accuracy obtained with ViT (patch size 2), ResNet-18 and
VGG-19 on permuted (in solid lines) and on original (in dashed lines)
CIFAR-100. 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 a ViT (patch size 2) on permuted CIFAR-100
(shuffle grid size 2). Here, we display $\bm{p}_{i}^{\top}\bm{P}$ where $i$ is
some fixed index and reshape such vector into a matrix $16\times 16$. We
observe that these similarities (1-2) do not have any spatially localized
structure. However, when applying the inverse of the permutation, we recover
spatially localized patterns in (1’-2’).

##  6 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-10/100 (Krizhevsky et al., 2009), SVHN (Netzer et al.,
2011) and large-scale ones such as ILSVRC-2012 ImageNet (Deng et al., 2009).
For the small datasets, we use a ViT with 7 layers, 12 heads and hidden/MLP
dimension 384. For ImageNet, we train a "ViT-tiny-patch16-224" Dosovitskiy et
al. (2020). Both models are trained with standard augmentations techniques
(Cubuk et al., 2018) and using AdamW with a 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 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\times s$ pixels. For a fixed $s\in\\{2,4,8,16,32\\}$, we permute the
grids according to $\pi_{s}$ to create the permuted CIFAR-100 dataset. We call
$s$ the grid shuffle size. 3(b)-(1) shows a CIFAR-100 image and its
corresponding shuffling in the permuted CIFAR-100 dataset 3(b)-(2-3-4). We
train a ViT and CNNs ResNet18 (He et al., 2016) and VGG-19 (Simonyan and
Zisserman, 2014) on the permuted CIFAR-100 dataset. [12]r0.48  Training loss
(1) and test accuracy (2) obtained using a ViT-tiny-patch16-224 on Imagenet.
ViT using positional attention (Ours) gets $68.9\%$ test accuracy while
vanilla ViT (ViT) gets $71.9\%$.  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 or equal to $s$. Indeed, intuitively, when we permute
the grids in 3(b), we lose the local aspect of the spatial structure and
create new sets $\mathcal{S}_{\ell}$’s and a new labeling function $f^{*}.$
3(a) reports the test accuracy of these three models for different values of
$s$. When $s$ is small, the image does not have a coherent structure e.g.
3(b)-(2) and thus, CNNs struggle to generalize. As $s$ increases e.g.
3(b)-(4), the information inside a patch is meaningful and thus, the CNNs
well-perform. Unsurprisingly, since ViTs are permutation invariant, their
performance remains unchanged for all $s$ – see 3(a). Despite this change, we
verify that the ViT is able to recover the new $\mathcal{S}_{\ell}$’s: we feed
the ViT with the shuffled pear image ( 3(b)-(2)) and consider for some $i$ the
similarity matrix $\bm{p}_{i}^{\top}\bm{P}$. We see that it does not exhibit a
local spatial structure in 3(c)-(1,2). We then apply $\pi_{s}^{-1}$ to
$\bm{p}_{i}^{\top}\bm{P}$ and observe that we recover the spatially localized
patterns 3(c)-(1’,2’). 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 are competitive.

[10]r0.48

Test accuracy obtained with a ViT using vanilla attention (ViT) and positional
attention (Ours) on 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. 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. Section 6 and Section 6
show that ViTs using positional attention compete with vanilla ViTs on a range
of datasets. These experiments strengthen our intuition that for images,
having an attention matrix that only depends on the positional encodings is
sufficient to have a good test accuracy.

## 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 on 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 learnt 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

  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Dosovitskiy et al. (2020) 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.11929_ , 2020.
  * Chen et al. (2021a) Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  _Advances in neural information processing systems_ , 34, 2021a.
  * Janner et al. (2021) Michael Janner, Qiyang Li, and Sergey Levine.  Offline reinforcement learning as one big sequence modeling problem.  _Advances in neural information processing systems_ , 34, 2021.
  * Austin et al. (2021) Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al.  Program synthesis with large language models.  _arXiv preprint arXiv:2108.07732_ , 2021.
  * Lample and Charton (2019) Guillaume Lample and François Charton.  Deep learning for symbolic mathematics.  _arXiv preprint arXiv:1912.01412_ , 2019.
  * Raghu et al. (2021) 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.
  * Brown et al. (2020) 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–1901, 2020.
  * Warstadt and Bowman (2020) Alex Warstadt and Samuel R Bowman.  Can neural networks acquire a structural bias from raw linguistic data?  _arXiv preprint arXiv:2007.06761_ , 2020.
  * Wu et al. (2021) 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 on Machine Learning_ , pages 11251–11262. PMLR, 2021.
  * Fukushima (2003) Kunihiko Fukushima.  Neocognitron for handwritten digit recognition.  _Neurocomputing_ , 51:161–180, 2003.
  * Cordonnier et al. (2019) Jean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi.  On the relationship between self-attention and convolutional layers.  _arXiv preprint arXiv:1911.03584_ , 2019.
  * Carion et al. (2020) 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.
  * d’Ascoli et al. (2021) Stéphane d’Ascoli, Hugo Touvron, Matthew L Leavitt, Ari S 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.
  * Dai et al. (2021) Zihang Dai, Hanxiao Liu, Quoc V Le, and Mingxing Tan.  Coatnet: Marrying convolution and attention for all data sizes.  _Advances in Neural Information Processing Systems_ , 34:3965–3977, 2021.
  * Guo et al. (2021) 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.06263_ , 2021.
  * Bello et al. (2019) Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens, and Quoc V Le.  Attention augmented convolutional networks.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pages 3286–3295, 2019.
  * Ramachandran et al. (2019) 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.
  * Bello (2021) Irwan Bello.  Lambdanetworks: Modeling long-range interactions without attention.  _arXiv preprint arXiv:2102.08602_ , 2021.
  * Bhojanapalli et al. (2021) 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.
  * Paul and Chen (2021) Sayak Paul and Pin-Yu Chen.  Vision transformers are robust learners.  _arXiv preprint arXiv:2105.07581_ , 2(3), 2021.
  * Naseer et al. (2021) 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.
  * Caron et al. (2021) Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin.  Emerging properties in self-supervised vision transformers.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 9650–9660, 2021.
  * Chen et al. (2021b) 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 9640–9649, 2021b.
  * Melas-Kyriazi (2021) Luke Melas-Kyriazi.  Do you even need attention? a stack of feed-forward layers does surprisingly well on imagenet.  _arXiv preprint arXiv:2105.02723_ , 2021.
  * Trockman and Kolter (2022) Asher Trockman and J Zico Kolter.  Patches are all you need?  _arXiv preprint arXiv:2201.09792_ , 2022.
  * Vuckovic et al. (2020) James Vuckovic, Aristide Baratin, and Remi Tachet des Combes.  A mathematical theory of attention.  _arXiv preprint arXiv:2007.02876_ , 2020.
  * Edelman et al. (2021) Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang.  Inductive biases and variable creation in self-attention mechanisms.  _arXiv preprint arXiv:2110.10090_ , 2021.
  * Dehghani et al. (2018) Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser.  Universal transformers.  _arXiv preprint arXiv:1807.03819_ , 2018.
  * Wei et al. (2021) 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.
  * Hron et al. (2020) Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak.  Infinite attention: Nngp and ntk for deep attention networks.  In _International Conference on Machine Learning_ , pages 4376–4386. PMLR, 2020.
  * Zhang et al. (2020) 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–15393, 2020.
  * Snell et al. (2021) Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt.  Approximating how single head attention learns.  _arXiv preprint arXiv:2103.07601_ , 2021.
  * Soudry et al. (2018) 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.
  * Lyu and Li (2019) Kaifeng Lyu and Jian Li.  Gradient descent maximizes the margin of homogeneous neural networks.  _arXiv preprint arXiv:1906.05890_ , 2019.
  * Chizat and Bach (2020) 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\.
  * Gunasekar et al. (2018) 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.
  * Arora et al. (2019) Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo.  Implicit regularization in deep matrix factorization.  _arXiv preprint arXiv:1905.13655_ , 2019.
  * Jagadeesan et al. (2022) Meena Jagadeesan, Ilya Razenshteyn, and Suriya Gunasekar.  Inductive bias of multi-channel linear convolutional networks with bounded weight norm.  In _Conference on Learning Theory_ , pages 2276–2325. PMLR, 2022\.
  * Allen-Zhu and Li (2021) Zeyuan Allen-Zhu and Yuanzhi Li.  Forward super-resolution: How can gans learn hierarchical generative models for real-world distributions.  _arXiv preprint arXiv:2106.02619_ , 2021.
  * Wen and Li (2021) 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 International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 11112–11122. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/wen21c.html.
  * Chen et al. (2022) Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li.  Towards understanding mixture of experts in deep learning.  _arXiv preprint arXiv:2208.02813_ , 2022.
  * Goyal et al. (2017) 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:1706.02677_ , 2017.
  * Hoffer et al. (2017) Elad Hoffer, Itay Hubara, and Daniel Soudry.  Train longer, generalize better: closing the generalization gap in large batch training of neural networks.  _arXiv preprint arXiv:1705.08741_ , 2017.
  * Keskar et al. (2016) Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.  On large-batch training for deep learning: Generalization gap and sharp minima.  _arXiv preprint arXiv:1609.04836_ , 2016.
  * Smith et al. (2018) Samuel L. Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc V. Le.  Don’t decay the learning rate, increase the batch size, 2018.
  * Li et al. (2019) 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.04595_ , 2019.
  * Allen-Zhu and Li (2020) Zeyuan Allen-Zhu and Yuanzhi Li.  Towards understanding ensemble, knowledge distillation and self-distillation in deep learning.  _arXiv preprint arXiv:2012.09816_ , 2020.
  * Jelassi and Li (2022) 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.  URL https://proceedings.mlr.press/v162/jelassi22a.html.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Neyshabur et al. (2015) Behnam Neyshabur, Ruslan Salakhutdinov, and Nathan Srebro.  Path-sgd: Path-normalized optimization in deep neural networks.  _arXiv preprint arXiv:1506.02617_ , 2015.
  * Daniely (2017) Amit Daniely.  Sgd learns the conjugate kernel class of the network.  In _Advances in Neural Information Processing Systems_ , pages 2422–2430, 2017.
  * Wilson et al. (2017) Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nathan Srebro, and Benjamin Recht.  The marginal value of adaptive gradient methods in machine learning.  _arXiv preprint arXiv:1705.08292_ , 2017.
  * Zou et al. (2021) 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.
  * Jelassi et al. (2022) Samy Jelassi, Arthur Mensch, Gauthier Gidel, and Yuanzhi Li.  Adam is no better than normalized SGD: Dissecting how adaptivity improves GAN performance, 2022.  URL https://openreview.net/forum?id=D9SuLzhgK9.
  * Arora et al. (2018) Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu.  Theoretical analysis of auto rate-tuning by batch normalization.  _arXiv preprint arXiv:1812.03981_ , 2018.
  * Hoffer et al. (2019) Elad Hoffer, Ron Banner, Itay Golan, and Daniel Soudry.  Norm matters: efficient and accurate normalization schemes in deep networks, 2019.
  * Ioffe and Szegedy (2015) Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  _arXiv preprint arXiv:1502.03167_ , 2015.
  * Srivastava et al. (2014) Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov.  Dropout: a simple way to prevent neural networks from overfitting.  _The Journal of Machine Learning Research_ , 15(1):1929–1958, 2014.
  * Wei et al. (2020) Colin Wei, Sham Kakade, and Tengyu Ma.  The implicit and explicit regularization effects of dropout, 2020.
  * Bahdanau et al. (2014) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.  Neural machine translation by jointly learning to align and translate.  _arXiv preprint arXiv:1409.0473_ , 2014.
  * Li et al. (2018) Yuanzhi Li, Tengyu Ma, and Hongyang Zhang.  Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations.  In _Conference On Learning Theory_ , pages 2–47. PMLR, 2018.
  * Woodworth et al. (2020) 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 3635–3673. PMLR, 2020\.
  * Li et al. (2020) Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang.  Learning over-parametrized two-layer relu neural networks beyond ntk.  _arXiv preprint arXiv:2007.04596_ , 2020.
  * Krizhevsky et al. (2009) Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Netzer et al. (2011) Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng.  Reading digits in natural images with unsupervised feature learning.  2011\.
  * Deng et al. (2009) 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 on computer vision and pattern recognition_ , pages 248–255. Ieee, 2009.
  * Cubuk et al. (2018) Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le.  Autoaugment: Learning augmentation policies from data.  _arXiv preprint arXiv:1805.09501_ , 2018.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Identity mappings in deep residual networks.  In _European conference on computer vision_ , pages 630–645. Springer, 2016.
  * Simonyan and Zisserman (2014) Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ , 2014.
  * Szegedy et al. (2016) Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna.  Rethinking the inception architecture for computer vision.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 2818–2826, 2016.
  * Zhang et al. (2017) Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  _arXiv preprint arXiv:1710.09412_ , 2017.
  * Yun et al. (2019) Sangdoo Yun, Dongyoon Han, Seong 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 on computer vision_ , pages 6023–6032, 2019.
  * Touvron et al. (2021) 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 10347–10357. PMLR, 2021.
  * Huang et al. (2016) 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.
  * Zhong et al. (2020) 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.
  * Loshchilov and Hutter (2017) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.

## Checklist

  1. 1.

For all authors…

    1. (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.

    2. (b)

Did you describe the limitations of your work? [Yes] See Conclusion,
limitations and future works.

    3. (c)

Did you discuss any potential negative societal impacts of your work? [N/A]
This is a theory paper.

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [Yes]
See Section 3.

    2. (b)

Did you include complete proofs of all theoretical results? [Yes] See
Appendix.

  3. 3.

If you ran experiments…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]
See supplementary material.

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes] See Appendix.

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes] See Section 6

    4. (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. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes]

    2. (b)

Did you mention the license of the assets? [N/A]

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [N/A]

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [N/A]

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [N/A]

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]

##  Appendix A Additional experimental details

In this section, we provide additional details on our experiments and
additional plots.

###  A.1 Details on the implementation

We used Pytorch and Nvidia Tesla V100 GPUs. We conduct experiments on small-
scale (CIFAR-10/100 and SVHN) and large-scale datasets (ImageNet). The choice
of architecture and training parameters depend on the size of the dataset as
we detail below.

##### Small-scale experiments.

We use the code available at https://github.com/omihub777/ViT-CIFAR. The model
is made of 7 layers, 12 heads, hidden and MLP dimension 384, dropout 0. We use
"mean-pooling" and not the CLS pooling. We set the patch size to 2 in the
experiment Figure 3 and to 4 in the experiment Section 6. Indeed, we
empirically found that setting patch size 4 was the optimal choice. We apply
label smoothing [Szegedy et al., 2016] with coefficient 0.1 and do not apply
any cutmix [Zhang et al., 2017] nor mixup [Yun et al., 2019]. We use Adam
[Kingma and Ba, 2014] as optimizer and set the learning rate to $10^{-3}$,
minimum learning rate to $10^{-5}$, $\beta_{1}$ to $0.9$, $\beta_{2}$ to
$0.999$, batch size to $128$, weight decay to $5\cdot 10^{-5}$, number of
warmup epochs to 5 and number of total epochs to 200. The scheduler is a
cosine learning rate. We used the AutoAugment procedure [Cubuk et al., 2018]
as in the repository to generate data augmentations. The model has been
trained over a single GPU.

Regarding the convolutional models in the experiment Figure 3, we trained a
ResNet-18 and a VGG-19 with batch normalization. We trained the two
architectures using the same training procedure and hyperparameters as for the
ViT.

##### Large-scale experiments.

We use the code available at https://github.com/facebookresearch/deit. Due to
limited computational resources, we train a ViT-tiny-patch16-224 [Dosovitskiy
et al., 2020] where "CLS-pooling" is applied. A detailed table with the
hyperparameters used for the ImageNet experiment may be found in Table 9
(column "DeiT-B") in [Touvron et al., 2021]. We set no dropout but set
stochastic depth [Huang et al., 2016] 0.1. We used label smoothing 0.1.
Regarding the augmentations, we set RandAugment [Cubuk et al., 2018] 9/0.5,
mixup 0.8, cutmix 1, erasing probability [Zhong et al., 2020] 0.25. Lastly, we
trained the model using AdamW [Loshchilov and Hutter, 2017] and set the batch
size to 1024, learning rate to $5\cdot
10^{-4}\cdot\frac{\mathrm{batchsize}}{512}$ as in [Goyal et al., 2017], weight
decay to $0.05$, warmup epochs $5$ and number of total epochs to 300. The
total number of epochs is 300. The model has been trained over 16 GPUs (8
nodes and 2 GPUs per node) and batch size for each device is 64.

###  A.2 Additional plots

In Figure 3, we plot the positional encoding similarities for a few patches.
Figure 4 provides these plots for all the patches. One should think of Figure
3 as a Figure displaying just two of the arrays present in Figure 4. We
consistently verify that the ViT is always able to recover the convolution-
like patterns which shows that it is able to learn the right patch
association.

(a)

(b)

Figure 4: (a) Visualization of the positional encodings similarities when
feeding the ViT with Permuted CIFAR-100 data. Each array represents
$\\{\langle\bm{p}_{i},\bm{p}_{j}\rangle\\}_{j=1}^{D}$ for a fixed $i.$ (b)
displays the positional encodings similarities obtained after inverting the
permutation. We see that the ViT is able to recover the convolution-like
structure in all the cases.

##  Appendix B Induction hypothesis

In this section, we present the induction hypothesis that we use in the
analysis of the idealized case. This hypothesis is ultimately proved in
subsection D.7.

######  Induction hypothesis B.1.

During the idealized learning process, the following holds for $t\leq T$.

  * –

the sofmax denominator is large i.e.
$\Lambda^{(t)}+(C-1)\Gamma^{(t)}+(D-C)\Xi^{(t)}=\Theta(D).$

  * –

$\Xi^{(t)}$ is not too small i.e. $\Xi^{(t)}=\Theta(1/D).$

  * –

$\Gamma^{(t)}$ and $\Lambda^{(t)}$ are in a good range i.e.

$\displaystyle\Lambda^{(t)}=\frac{e^{\mathrm{polyloglog}(d)}}{D},\quad(C-1)\Gamma^{(t)}\in\left[\frac{\Omega(C)}{D},\frac{\lambda_{0}}{D}\right],$

where $\lambda_{0}=\Theta(D^{0.01}).$

##  Appendix C Notations

In this section, we introduce the different notations used in the proofs.

##### General purpose.

We first define notations that are used everywhere in the appendix.

  * –

Sigmoid function: Given $z\in\mathbb{R}$, $\mathfrak{S}(z)=1/(1+e^{-z}).$

  * –

Softmax function: Given $\bm{z}=(z_{1},\dots,z_{D})\in\mathbb{R}^{D}$,
$(\mathrm{softmax}(z_{1},\dots,z_{D}))_{m}=\frac{e^{z_{m}}}{\sum_{j=1}^{D}e^{z_{j}}}.$

  * –

Loss for a data-point $(\bm{X},y)$: $L(\bm{X})=\log(1+e^{-yF(\bm{X})}).$

##### Analysis in idealized case.

We now provide notations used in the analysis of the idealized case.

  * –

$\kappa(\bm{X})=\sum_{\ell\neq\ell(\bm{X})}\sum_{j\in\mathcal{S}_{\ell}}\delta_{j}+C.$

  * –

for $i\in[D]$, $\bm{O}_{i}^{(t)}=\sum_{j=1}^{D}S_{i,j}^{(t)}\bm{X}_{j}.$

##### Analysis in realistic case.

We now provide notations used in the analysis of the realistic case.

  * –

Score matrix: $\bm{S}\in\mathbb{R}^{D\times D}$ with coefficients
$\widehat{S}_{i,j}=\exp(\widehat{A}_{i,j}/\sqrt{d})/\sum_{r=1}^{D}\exp(\widehat{A}_{i,r}/\sqrt{D})$.

  * –

Given a data-point $(\bm{X}[i],y[i])$ and $j\in[D]$,
$\bm{O}_{j}^{(t)}[i]=\sum_{k=1}^{D}S_{j,k}^{(t)}\bm{X}_{k}[i]$

##  Appendix D Learning process in the idealized setting

###  D.1 Roadmap of the proof

From Lemma 4.2, we know that $\bm{v}^{(t)}\in\mathrm{span}(\bm{w}^{*})$ for
all $t\in[T]$. The main idea of the proof consists in analyzing the GD
dynamics of $\alpha^{(t)}=\langle\bm{v}^{(t)},\bm{w}^{*}\rangle$ that satisfy

$\displaystyle\alpha^{(t+1)}=\alpha^{(t)}+\eta\underbrace{\mathbb{E}\Bigg{[}y\mathfrak{S}(-yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X}))\sum_{i\in\mathcal{S}_{\ell(\bm{X})}}\sigma^{\prime}\Big{(}\sum_{j=1}^{D}S_{i,j}^{(t)}\langle\bm{v}^{(t)},\bm{X}_{j}\rangle\Big{)}\sum_{k=1}^{D}S_{i,k}^{(t)}\langle\bm{w}^{*},\bm{X}_{m}\rangle\Bigg{]}}_{{\Large\boldsymbol{\mathpzc{S}}}^{(t)}}$
$\displaystyle+\eta\underbrace{\mathbb{E}\Bigg{[}y\mathfrak{S}(-yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X}))\sum_{i\not\in\mathcal{S}_{\ell(\bm{X})}}\sigma^{\prime}\Big{(}\sum_{j=1}^{D}S_{i,j}^{(t)}\langle\bm{v}^{(t)},\bm{X}_{j}\rangle\Big{)}\sum_{m=1}^{D}S_{i,m}^{(t)}\langle\bm{w}^{*},\bm{X}_{m}\rangle\Bigg{]}}_{{\Large\boldsymbol{\mathpzc{N}}^{(t)}}}.$
(GD-$\alpha$)

We divide the idealized learning process as follows.

  * –

Event I ($t\in[0,\mathcal{T}_{0}]$, subsection D.2): at initialization,
$\alpha^{(0)}$ is small. Therefore, the sigmoid $\mathfrak{S}(-yF(\bm{X}))$ is
large. Besides, around $\alpha^{(0)}$, it stays constant i.e.
$\mathfrak{S}(-yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X}))\approx\mathfrak{S}(-yF_{\bm{A}^{(0)},\bm{v}^{(0)}}(\bm{X}))$
in (GD-$\alpha$). This implies that ${\large\boldsymbol{\mathpzc{N}}}^{(t)}=0$
which yields $\alpha^{(t)}$ to increase until reaching a specific value where
the sigmoid is not constant anymore.

  * –

Event II ($t\in[\mathcal{T}_{0},\mathcal{T}_{1}]$, subsection D.3): at time
$\mathcal{T}_{0}$, $\alpha^{(t)}$ is large. This fact along with Lemma 4.3
imply that $\gamma^{(t)}$ increases. Eventually, $\Gamma^{(\mathcal{T}_{1})}$
becomes large enough so that
${\large\boldsymbol{\mathpzc{S}}}^{(t)}\geq\max_{\tau\leq
T}{\large|\boldsymbol{\mathpzc{N}}}^{(\tau)}|$.

  * –

Event III ($t\in[\mathcal{T}_{1},T]$, subsection D.4): Since
${\large\boldsymbol{\mathpzc{S}}}^{(t)}\geq\max_{\tau\leq
T}{\large|\boldsymbol{\mathpzc{N}}}^{(\tau)}|$, $\alpha^{(t)}$ increases
again. It increases until the population risk is at most $o(1).$

After $T$ iterations, $\alpha^{(t)}$ is large and the population risk thus
converges (subsection D.5). Since the logistic loss is a surrogate for the 0-1
loss, we prove that the learner model fits the labeling function (subsection
D.6) which implies the first statement of Theorem 4.1.

Remark : Since we initialize $\alpha^{(0)}\geq\nu^{1/(p-1)}$, Lemma D.7
implies that we can overlook the linear part of the activation in this
section. Therefore, we only consider $\sigma(x)=x^{p}$ in the idealized
process.

###  D.2 Event I: $\alpha^{(t)}$ initially increases

A first question that arises is: starting from $\alpha^{(0)}$, what is the
value of $\alpha^{(t)}$ that makes the sigmoid non-constant? The following
lemma addresses this question.

######  Lemma D.1.

The value $\alpha^{(t)}$ at which the sigmoid $\mathfrak{S}(-yF(\bm{X}))$
becomes non-constant is:

$\displaystyle\tilde{\alpha}=\frac{\Theta(1)}{C^{2}\lambda_{0}}.$

###### Proof of Lemma D.1.

The update of $\alpha^{(t)}$ is

$\displaystyle\alpha^{(t+1)}$
$\displaystyle=\alpha^{(t)}+\eta\mathbb{E}[y\mathfrak{S}(-yF(\bm{X}))\mathscr{G}^{(t)}]$
(3)
$\displaystyle=\alpha^{(t)}+\eta\mathbb{E}[y\mathfrak{S}(-\alpha^{(0)})\mathscr{G}^{(t)}]+\eta\mathbb{E}\big{[}y\big{(}\mathfrak{S}(-yF(\bm{X}))-\mathfrak{S}(-\alpha^{(0)})\big{)}\mathscr{G}^{(t)}\big{]},$

where
$\mathscr{G}^{(t)}:=\sum_{j=1}^{D}\langle\bm{O}_{j}^{(t)},\bm{w}^{*}\rangle^{p}.$
Since $x\mapsto\mathfrak{S}(-x)$ is $1/4$-Lipschitz, we rewrite (3) as:

$\displaystyle\big{|}\alpha^{(t+1)}-\alpha^{(t)}-\eta\mathbb{E}[y\mathfrak{S}(-\alpha^{(0)})\mathscr{G}^{(t)}]\big{|}$
$\displaystyle\leq\frac{\eta}{4}\mathbb{E}\big{[}|F(\bm{X})-y\alpha^{(0)}|\cdot|y\mathscr{G}^{(t)}|\big{]}$
(4)
$\displaystyle\leq\Theta(C\eta)(\alpha^{(t)})^{p}\mathbb{E}\big{[}(C^{1+1/p}\lambda_{0})^{p}|y\mathscr{G}^{(t)}|\big{]}$
$\displaystyle\leq\Theta(\eta)(\alpha^{(t)})^{p}\mathbb{E}\big{[}(C^{2}\lambda_{0})^{p}|y\mathscr{G}^{(t)}|\big{]},$

where we applied Lemma D.3, Induction Hypothesis B.1 and the fact that
$\alpha^{(0)}$ is small in the penultimate inequality. Using Lemma D.4 and
Lemma D.5, we have
$\mathbb{E}[|y\mathscr{G}^{(t)}|]\leq\Theta(1)\mathbb{E}[y\mathscr{G}^{(t)}]$
which yields

$\displaystyle\big{|}\alpha^{(t+1)}-\alpha^{(t)}-\eta\mathbb{E}[y\mathfrak{S}(-\alpha^{(0)})\mathscr{G}^{(t)}]\big{|}$
$\displaystyle\leq\Theta(\eta)(\alpha^{(t)})^{p}\mathbb{E}\big{[}(C^{2}\lambda_{0})^{p}y\mathscr{G}^{(t)}\big{]}.$
(5)

(4) shows that when $\alpha^{(t)}$ is small, we have
$\alpha^{(t+1)}\approx\alpha^{(t)}+\eta\mathbb{E}[y\mathfrak{S}(-\alpha^{(0)})A]$.
Besides, we have $\mathbb{E}[y\mathscr{G}^{(t)}]\geq 0$ so $\alpha^{(t)}$
increases. However, during this increase, the right-hand side of (4) increases
and this approximation does not hold anymore. Therefore, the sigmoid is
approximately constant when $\alpha^{(t)}$ satisfies:

$\displaystyle\Theta(1)(\alpha^{(t)})^{p}\mathbb{E}[(C^{2}\lambda_{0})^{p}y\mathscr{G}^{(t)}]\leq\mathbb{E}[\mathfrak{S}(-\alpha^{(0)})y\mathscr{G}^{(t)}]\iff\alpha^{(t)}\leq\frac{\Theta(1)}{C^{2}\lambda_{0}}.$
(6)

∎

######  Lemma D.2.

Let $\mathcal{T}_{0}=\Theta\left(\frac{1}{\eta C(\alpha^{(0)})^{p-1}}\right)$.
For all $t\in[0,\mathcal{T}_{0}]$, we have
${\large\boldsymbol{\mathpzc{N}}^{(t)}}=0.$ Therefore, $\alpha^{(t)}$ is
updated as

$\displaystyle\alpha^{(t+1)}$ $\displaystyle=\alpha^{(t)}+\Theta(\eta
C)(G^{(t)})^{p}(\alpha^{(t)})^{p-1}.$

Consequently, $\alpha^{(t)}$ is non-decreasing and after $\mathcal{T}_{0}$
iterations, we have $\alpha^{(t)}\geq\frac{\Omega(1)}{C^{2}\lambda_{0}}$ for
$t\geq\mathcal{T}_{0}.$

###### Proof of Lemma D.2.

For $t\in[0,\mathcal{T}_{0}]$, we know that the sigmoid
$\mathfrak{S}(-yF(\bm{X}))$ is constant. We apply Lemma D.4 and Lemma D.6 to
respectively bound ${\large\boldsymbol{\mathpzc{S}}}^{(t)}$ and
${\large\boldsymbol{\mathpzc{N}}}^{(t)}$ in the update of $\alpha^{(t)}$.

$\displaystyle\alpha^{(t+1)}$ $\displaystyle=\alpha^{(t)}+\Theta(\eta
C)(\alpha^{(t)})^{p-1}(G^{(t)})^{p}.$ (7)

We apply Induction Hypothesis B.1 and $\Omega(C)\geq 0$ and
$e^{\mathrm{polyloglog}(d)}\leq\lambda_{0}$ in (7) and obtain:

$\displaystyle\begin{cases}\alpha^{(t+1)}\geq\alpha^{(t)}+\Theta(\eta
C)e^{\mathrm{polyloglog}(d)}(\alpha^{(t)})^{p-1}\\\
\alpha^{(t+1)}\leq\alpha^{(t)}+\Theta(\eta
C)\lambda_{0}^{p}(\alpha^{(t)})^{p-1}\end{cases}.$ (8)

(8) indicates that $\alpha^{(t)}$ is a non-decreasing sequence. Therefore,
there exists a time $\mathcal{T}_{0}$ such that
$\alpha^{(\mathcal{T}_{0})}=\frac{\Theta(1)}{C^{2}\lambda_{0}}$. Using Lemma
K.1, the time $\mathcal{T}_{0}$ is equal to:

$\displaystyle\mathcal{T}_{0}=\frac{1}{\eta
C(\alpha^{(0)})^{p-1}e^{\mathrm{polyloglog}(d)}}+(\lambda_{0})^{p}e^{-\mathrm{polyloglog}(d)}\Big{\lceil}-\log\big{(}C^{2}\lambda_{0}\alpha^{(0)}\big{)}\Big{\rceil}.$
(9)

∎

####  D.2.1 Auxiliary lemmas

In this section, we present the auxiliary lemmas needed to prove the main
results of subsection D.2. We first present a lemma that bounds the learner
model.

######  Lemma D.3.

Let $t\in[0,T]$. The learner model $F$ is bounded for all
$(\bm{X},y)\sim\mathcal{D}$ as:

$\displaystyle
yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X})\leq\Theta(1)(C^{1+1/p}\alpha^{(t)}\lambda_{0})^{p}.$

###### Proof of Lemma D.3.

By definition, the learner model is:

$\displaystyle
yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X})=y\sum_{i=1}^{D}\langle\bm{v}^{(t)},\bm{O}_{i}^{(t)}\rangle^{p}=\sum_{i\in\mathcal{S}_{\ell(X)}}\langle\bm{v}^{(t)},\bm{O}_{i}^{(t)}\rangle^{p}+\sum_{i\not\in\mathcal{S}_{\ell(X)}}\langle
v^{(t)},\bm{O}_{i}^{(t)}\rangle^{p}.$ (10)

We successively apply Lemma D.4, Induction Hypothesis B.1 and Lemma D.5 to
bound (10).

$\displaystyle
yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X})\leq(\alpha^{(t)})^{p}\Big{(}C\Theta(DC\Gamma^{(t)})^{p}+o(\lambda_{0})^{p}\Big{)}.$
(11)

Finally, we apply Induction Hypothesis B.1 in (11) to obtain the desired
result. ∎

We now present lemmas that bound ${\large\boldsymbol{\mathpzc{S}}}^{(t)}$ and
${\large\boldsymbol{\mathpzc{N}}}^{(t)}$.

######  Lemma D.4.

Let $t\in[T]$ and $i\in\mathcal{S}_{\ell(\bm{X})}.$ We have
$y\langle\bm{w}^{*},\bm{O}_{i}^{(t)}\rangle^{p}=\Theta(G^{(t)})^{p}.$ As long
as the population risk is not $o(1)$, we have
${\large\boldsymbol{\mathpzc{S}}}^{(t)}=C\Theta(G^{(t)})^{p}$ for all $t\leq
T.$

###### Proof of Lemma D.4.

We apply Lemma G.3 and obtain:

$\displaystyle y\langle\bm{w}^{*},\bm{O}_{i}^{(t)}\rangle^{p}$
$\displaystyle=D^{p}\bigg{(}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})+y\Xi^{(t)}\sum_{h\in[L]\backslash\\{\ell(\bm{X})\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\bigg{)}^{p}$
$\displaystyle=\Theta\Big{(}D\big{(}\Lambda^{(t)}+(C-1)\Gamma^{(t)}\big{)}\Big{)}^{p}$
$\displaystyle=\Theta(G^{(t)})^{p}.$

∎

######  Lemma D.5.

Let $t\in[T]$. We have
$\sum_{i\not\in\mathcal{S}_{\ell(\bm{X})}}y\langle\bm{w}^{*},\bm{O}_{i}^{(t)}\rangle^{p}\leq
o(\lambda_{0})^{p}.$ In particular, this implies
${\large\boldsymbol{\mathpzc{N}}}^{(t)}\leq o(\lambda_{0})^{p}$ for all $t\leq
T.$

###### Proof of Lemma D.5.

We have:

$\displaystyle\sum_{k\in[L]\backslash\\{\ell(X)\\}}\sum_{i\in\mathcal{S}_{k}}y\langle\bm{w}^{*},\bm{O}_{i}^{(t)}\rangle^{p}$
(12) $\displaystyle=$ $\displaystyle
D^{p}\sum_{k\in[L]\backslash\\{\ell(X)\\}}\sum_{m\in\mathcal{S}_{k}}y\bigg{(}\Lambda^{(t)}\delta_{m}+\Gamma^{(t)}\sum_{s\in\mathcal{S}_{k}\backslash\\{m\\}}\delta_{s}+\Xi^{(t)}\Big{[}Cy+\sum_{h\in[L]\backslash\\{\ell(X),k\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\Big{]}\bigg{)}^{p}.$

We distinguish two cases.

  * –

$\delta_{r}=0$ for all $r\in\mathcal{S}_{\ell}:$ we apply Lemma K.3 and
obtain:

$\displaystyle\sum_{k\neq\ell(\bm{X})}\sum_{i\in\mathcal{S}_{k}}y\langle\bm{w}^{*},\bm{O}_{i}^{(t)}\rangle^{p}$
$\displaystyle=\sum_{k\neq\ell(\bm{X})}\sum_{m\in\mathcal{S}_{k}}y\Big{(}D\Xi^{(t)}\Big{[}Cy+\sum_{h\neq\\{\ell(\bm{X}),k\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\Big{]}\Big{)}^{p}$
(13) $\displaystyle\leq D\cdot\Theta(D\Xi^{(t)}qD\log(d))^{p}.$

  * –

$\exists r\in\mathcal{S}_{\ell}$ such that $\delta_{r}\neq 0$: let
$i\in\mathcal{S}_{k}$. We apply Lemma K.4 and obtain:

$\displaystyle
y\langle\bm{w}^{*},\bm{O}_{i}^{(t)}\rangle^{p}\leq\Theta\big{(}D(\Lambda^{(t)}+O(1)\Gamma^{(t)})\big{)}^{p}\cdot\mathbf{1}_{\exists
r\in\mathcal{S}_{k},\delta_{r}\neq 0}.$ (14)

We now sum (LABEL:eq:jcneiencae) and apply Lemma K.3 to obtain:

$\displaystyle\sum_{k\neq\ell(\bm{X})}\sum_{i\in\mathcal{S}_{k}}y\langle\bm{w}^{*},\bm{O}_{i}^{(t)}\rangle^{p}$
$\displaystyle\leq\Theta\big{(}D(\Lambda^{(t)}+O(1)\Gamma^{(t)})\big{)}^{p}\sum_{k\neq\ell(\bm{X})}\sum_{i\in\mathcal{S}_{k}}|\delta_{r}|$
(15)
$\displaystyle\leq\Theta\big{(}D(\Lambda^{(t)}+O(1)\Gamma^{(t)})\big{)}^{p}qD\log(d).$

We finally apply Induction Hypothesis B.1 to have
$\Gamma^{(t)}\leq\lambda_{0}/C$ in (15) and get

$\displaystyle\sum_{k\neq\ell(\bm{X})}\sum_{i\in\mathcal{S}_{k}}y\langle\bm{w}^{*},\bm{O}_{i}^{(t)}\rangle^{p}$
$\displaystyle\leq\Theta(\lambda_{0}/C)^{p}qD\log(d).$ (16)

We finally plug (13) and (16) in (LABEL:eq:jwfejfw) and obtain:

$\displaystyle\sum_{k\neq\ell(\bm{X})}\sum_{i\in\mathcal{S}_{k}}y\langle\bm{w}^{*},\bm{O}_{i}^{(t)}\rangle^{p}$
$\displaystyle\leq
D\cdot\Theta(D\Xi^{(t)}qD\log(d))^{p}+\Theta(\lambda_{0}/C)^{p}qD\log(d)\leq
o(\lambda_{0})^{p}.$

∎

######  Lemma D.6.

Let $t\in[0,\mathcal{T}_{0}].$ We have
${\large\boldsymbol{\mathpzc{N}}}^{(t)}=0.$

###### Proof of Lemma D.6.

By definition of ${\large\boldsymbol{\mathpzc{N}}}^{(t)}$, we have:

$\displaystyle{\large\boldsymbol{\mathpzc{N}}}^{(t)}$ (17) $\displaystyle=$
$\displaystyle\frac{1}{2}\bigg{(}\mathbb{E}\bigg{[}\mathfrak{S}(-F(\bm{X}))\sum_{j\not\in\mathcal{S}_{\ell(\bm{X})}}\langle\bm{w}^{*},\bm{O}_{j}^{(t)}\rangle^{p}\Big{|}y=1\bigg{]}-\mathbb{E}\bigg{[}\mathfrak{S}(F(\bm{X}))\sum_{j\not\in\mathcal{S}_{\ell(\bm{X})}}\langle\bm{w}^{*},\bm{O}_{j}^{(t)}\rangle^{p}\Big{|}y=-1\bigg{]}\bigg{)}$
$\displaystyle=$
$\displaystyle\Theta(1)\bigg{(}\mathbb{E}\bigg{[}\sum_{j\not\in\mathcal{S}_{\ell(\bm{X})}}\langle\bm{w}^{*},\bm{O}_{j}^{(t)}\rangle^{p}\Big{|}y=1\bigg{]}-\mathbb{E}\bigg{[}\sum_{j\not\in\mathcal{S}_{\ell(\bm{X})}}\langle\bm{w}^{*},\bm{O}_{j}^{(t)}\rangle^{p}\Big{|}y=-1\bigg{]}\bigg{)},$

where we use
$\mathfrak{S}(-F(\bm{X}))\approx\mathfrak{S}(-F_{\bm{A}^{(0)},\bm{v}^{(0)}}(\bm{X}))$
for $t\in[0,\mathcal{T}_{0}]$ in the last equality of (LABEL:eq:foewjjewf). We
now show that each of the summands in (LABEL:eq:foewjjewf) is zero. Without
loss of generality, let’s focus on the first summand. The same reasoning holds
for the second one. In particular, for $j\in\mathcal{S}_{k}$ with $k\neq j$,
we now compute
$\mathbb{E}[\langle\bm{w}^{*},\bm{O}_{j}^{(t)}\rangle^{p}|y=1].$ Using the
binomial theorem and the independence of the $\delta_{r}$’s, we have:

$\displaystyle\mathbb{E}[\langle\bm{w}^{*},\bm{O}_{j}^{(t)}\rangle^{p}|y=1]$
(18) $\displaystyle=$
$\displaystyle\sum_{a=0}^{p}\binom{p}{a}(\Lambda^{(t)})^{p-a}\mathbb{E}[(\delta_{j})^{p-a}]\mathbb{E}\bigg{[}\bigg{(}\Gamma^{(t)}\sum_{s\in\mathcal{S}_{k}\backslash\\{j\\}}\delta_{s}+\Xi^{(t)}\bigg{[}C+\sum_{h\neq\\{\ell(\bm{X}),k\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\bigg{]}\bigg{)}^{a}\bigg{]}.$

For $a$ even, we have $p-a$ odd which implies
$\mathbb{E}[(\delta_{u})^{p-a}]=0$. Therefore, the summands with even $a$ are
zero. We now focus on the case $a$ odd. We again apply the binomial theorem
and the independence of the $\delta_{r}$’s to get:

$\displaystyle\mathbb{E}\bigg{[}\bigg{(}\Gamma^{(t)}\sum_{s\in\mathcal{S}_{k}\backslash\\{j\\}}\delta_{s}+\Xi^{(t)}\bigg{[}C+\sum_{h\neq\\{\ell(\bm{X}),k\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\bigg{]}\bigg{)}^{a}\bigg{]}$
(19) $\displaystyle=$
$\displaystyle\sum_{b=0}^{a}\binom{a}{b}(\Gamma^{(t)})^{a-b}\mathbb{E}\bigg{[}\bigg{(}\sum_{s\in\mathcal{S}_{k}\backslash\\{j\\}}\delta_{s}\bigg{)}^{a-b}\bigg{]}(\Xi^{(t)})^{b}\mathbb{E}\bigg{[}\bigg{(}Cy+\sum_{h\neq\\{\ell(\bm{X}),k\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\bigg{)}^{b}\bigg{]}.$

For $b$ even, we have $a-b$ odd. This implies that
$\mathbb{E}\big{[}\big{(}\sum_{s\in\mathcal{S}_{k}\backslash\\{j\\}}\delta_{s}\big{)}^{a-b}\big{]}=0$.
In the case $b$ odd, we exactly use the same argument and obtain:

$\displaystyle\mathbb{E}\bigg{[}\bigg{(}Cy+\sum_{h\neq\\{\ell(\bm{X}),k\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\bigg{)}^{b}\bigg{]}=0.$
(20)

(20) implies (LABEL:eq:jcenjnac) is zero and which lastly implies
(LABEL:eq:fidekowe) is zero. We conclude that
$\mathbb{E}[\langle\bm{w}^{*},\bm{O}_{j}^{(t)}\rangle^{p}|y=1]=\mathbb{E}[\langle\bm{w}^{*},\bm{O}_{j}^{(t)}\rangle^{p}|y=-1]=0$
and thus ${\large\boldsymbol{\mathpzc{N}}}^{(t)}=0.$

∎

######  Lemma D.7.

Let $(\bm{X},\cdot)\sim\mathcal{D}$ and $j\in[D]$. Assume that
$\alpha^{(t)}\geq\nu^{1/(p-1)}$. Then, we have:

$\displaystyle\sigma^{\prime}\Big{(}D\sum_{k=1}^{D}\bm{S}_{j,k}^{(t)}\langle\bm{v},\bm{X}_{k}\rangle\Big{)}$
$\displaystyle=\Theta(1)\Big{(}D\sum_{k=1}^{D}\bm{S}_{j,k}^{(t)}\langle\bm{v},\bm{X}_{k}\rangle\Big{)}^{p-1}.$

###### Proof of Lemma D.7.

We remind that the derivative of the activation function
$\sigma^{\prime}(x)=px^{p-1}+\nu$. We first remark that for all $x,$
$\sigma^{\prime}(x)\geq px^{p-1}$. Besides, we have:

$\displaystyle px^{p-1}+\nu\leq\frac{3px^{p-1}}{2}\iff
x\geq\Big{(}\frac{2\nu}{p}\Big{)}^{1/(p-1)}.$ (21)

In our case, we have
$x=D\alpha^{(t)}\sum_{k=1}^{D}\bm{S}_{j,k}^{(t)}\langle\bm{w}^{*},\bm{X}_{k}\rangle.$
Using Induction Hypothesis B.1, we have $x\leq C\lambda_{0}\alpha^{(t)}.$
Therefore, a sufficient condition for (21) to hold is
$\alpha^{(t)}\geq\nu^{1/(p-1)}/(C\lambda_{0}).$ Since
$\lambda_{0},C\ll\mathrm{poly}(d)$, we can simplify this condition as
$\alpha^{(t)}\geq\nu^{1/(p-1)}$.

∎

###  D.3 Event II: $\Gamma^{(t)}$ significantly increases

In this section, we show the increase of $\alpha^{(t)}$ for
$t\in[0,\mathcal{T}_{0}]$ leads to the increase of $\Gamma^{(t)}$. At time
$\mathcal{T}_{1}>\mathcal{T}_{0},$ $\Gamma^{(t)}$ is significantly large.

######  Lemma D.8.

Let $\mathcal{T}_{1}=\mathcal{T}_{0}+\Theta\Big{(}\frac{\lambda_{0}^{p}}{\eta
e^{\mathrm{polyloglog}(d)}}\Big{)}$. For all $t\in[\mathcal{T}_{1},T]$, we
have $\Gamma^{(t)}\geq\frac{\Omega(\lambda_{0})}{D}$. This implies
${\large\boldsymbol{\mathpzc{S}}}^{(t)}\geq\max_{\tau\leq
T}|{\large\boldsymbol{\mathpzc{N}}}^{(\tau)}|$.

###### Proof of Lemma D.8.

Let $t\in[\mathcal{T}_{0},T]$ and $\tau\in[\mathcal{T}_{0},t].$ Using
Corollary G.1 and Induction Hypothesis B.1, $\gamma^{(\tau)}$ satisfies:

$\displaystyle\gamma^{(\tau+1)}$ $\displaystyle\geq\gamma^{(\tau)}+\Omega(\eta
C)(\alpha^{(\tau)})^{p}e^{\mathrm{polyloglog}(d)}.$ (22)

Summing (22) for $\tau=\mathcal{T}_{0},\dots,t-1$ yields

$\displaystyle\gamma^{(t)}$
$\displaystyle\geq\gamma^{(\mathcal{T}_{0})}+e^{\mathrm{polyloglog}(d)}\Omega(\eta
C)\sum_{\tau=\mathcal{T}_{0}}^{t-1}(\alpha^{(\tau)})^{p}.$ (23)

We successively apply Lemma D.9 and $(a-b)^{p}\geq a^{p}-pb$ for $a\ll b$ to
lower bound (23) to obtain:

$\displaystyle\gamma^{(t)}$
$\displaystyle\geq\gamma^{(\mathcal{T}_{0})}+e^{\mathrm{polyloglog}(d)}\Omega(\eta
C)\sum_{\tau=\mathcal{T}_{0}}^{t-1}\Big{(}\frac{\Omega(1)}{C^{2p}\lambda_{0}^{p}}-o(\eta)\Big{)}.$
(24)

We apply Induction Hypothesis B.1 in (24) to obtain a bound on $\Gamma^{(t)}$.

$\displaystyle C\Gamma^{(t)}$ $\displaystyle\geq
C\Gamma^{(\mathcal{T}_{0})}\exp\Big{(}e^{\mathrm{polyloglog}(d)}\Omega(\eta
C)\sum_{\tau=\mathcal{T}_{0}}^{t-1}\big{(}\frac{\Omega(1)}{C^{2p}\lambda_{0}^{p}}-o(\eta)\big{)}\Big{)}$
(25)
$\displaystyle\geq\frac{\Omega(C^{2})}{D}\exp\Big{(}e^{\mathrm{polyloglog}(d)}\Omega(\eta
C)\sum_{\tau=\mathcal{T}_{0}}^{t-1}\big{(}\frac{\Omega(1)}{C^{2p}\lambda_{0}^{p}}-o(\eta)\big{)}\Big{)}.$

(25) shows that $\Gamma^{(t)}$ is an non-decreasing sequence. We thus deduce
the time $\mathcal{T}_{1}$ such that $C\Gamma^{(t)}\geq\Omega(\lambda_{0})/D.$

$\displaystyle\frac{C^{2}}{D}\exp\Big{(}e^{\mathrm{polyloglog}(d)}\Omega(\eta
C)\big{(}\frac{\Omega(1)}{C^{2p}\lambda_{0}^{p}}-o(\eta)\big{)}(\mathcal{T}_{1}-\mathcal{T}_{0})\Big{)}=\frac{\lambda_{0}}{D}$
(26)
$\displaystyle\implies\mathcal{T}_{1}=\mathcal{T}_{0}+\Theta\Bigg{(}\frac{C^{2p}\lambda_{0}^{p}}{\eta
e^{\mathrm{polyloglog}(d)}}\log\Big{(}\frac{\lambda_{0}}{C^{2}}\Big{)}\Bigg{)}.$

We now prove the second part of the lemma. We respectively apply Lemma D.4 and
Lemma D.5 to bound ${\large\boldsymbol{\mathpzc{S}}}^{(t)}$ and
$\max_{\tau\leq T}|{\large\boldsymbol{\mathpzc{N}}}^{(\tau)}|$.

$\displaystyle{\large\boldsymbol{\mathpzc{S}}}^{(t)}\geq\Omega(\lambda_{0})^{p}\quad\text{and}\quad\max_{\tau\leq[T]}|{\large\boldsymbol{\mathpzc{N}}}^{(\tau)}|\leq
o(\lambda_{0})^{p}.$ (27)

(27) implies for all $t\in[\mathcal{T}_{1},T]$,
${\large\boldsymbol{\mathpzc{S}}}^{(t)}\geq\max_{\tau\leq[T]}|{\large\boldsymbol{\mathpzc{N}}}^{(\tau)}|.$

∎

####  D.3.1 Auxiliary lemmas

In this section, we present the auxiliary lemmas needed to prove the main
results in subsection D.3.

######  Lemma D.9.

Let $t\geq\mathcal{T}_{0}.$ Then, we always have
$\alpha^{(t)}\geq\frac{\Omega(1)}{C^{2}\lambda_{0}}-o(\eta).$

###### Proof of Lemma D.9.

For $t\in[0,\mathcal{T}_{0}]$, $\alpha^{(t)}$ increases and eventually
satisfies $\alpha^{(t)}\geq\frac{\Omega(1)}{C^{2}\lambda_{0}}$ (Lemma D.2).
However, for $t\geq\mathcal{T}_{0}$, $\alpha^{(t)}$ may be non-increasing.
Here, we want to quantify the maximum amount of decrease for
$t>\mathcal{T}_{0}$. The worst-case scenario is when
$\alpha^{(t)}=\frac{\Omega(1)}{C^{2}\lambda_{0}}$. We bound $\alpha^{(t+1)}$
by using Lemma D.4, Lemma D.10 and Lemma D.5.

$\displaystyle\alpha^{(t+1)}$
$\displaystyle\geq\alpha^{(t)}+\Theta(C\eta)(\alpha^{(t)}G^{(t)})^{p}-\eta(\alpha^{(t)})^{p}o(\lambda_{0})^{p}.$
(28)

We now apply Induction Hypothesis B.1 in (28) and get:

$\displaystyle\alpha^{(t+1)}$
$\displaystyle\geq\frac{\Omega(1)}{C^{2}\lambda_{0}}+\Theta(\eta)\frac{e^{\mathrm{polyloglog}(d)}}{C^{2p-1}\lambda_{0}^{p}}-\frac{o(\eta)}{C^{2p}}$
(29)
$\displaystyle\geq\frac{\Omega(1)}{C^{2}\lambda_{0}}-\frac{o(\eta)}{C^{2p}}.$

At time $t+1,$ we potentially have
$\alpha^{(t+1)}<\frac{\Omega(1)}{C^{2}\lambda_{0}}$. In this case,
$\alpha^{(t+1)}$ starts to increase again because it is in the range of
$\alpha$’s that satisfies Event I (and therefore the update rule in Lemma D.2
holds). Thus, for all $t\geq\mathcal{T}_{0},$ we have
$\alpha^{(t)}\geq\frac{\Omega(1)}{C^{2}\lambda_{0}}-o(\eta).$ ∎

######  Lemma D.10.

When the population risk is $\Omega(1)$, we have
$\mathbb{E}[\mathfrak{S}(-yF(\bm{X}))]\geq\Omega(1).$

###### Proof of Lemma D.10.

Let $(\bm{X},y)$ be a data-point. We distinguish two cases:

  * –

$yF(\bm{X})>0$: we apply Lemma K.6 which implies
$\mathfrak{S}(-yF(\bm{X})\geq\log(1+e^{yF(\bm{X})}).$ Since the population
loss is $\Omega(1)$, this implies the aimed result.

  * –

$yF(\bm{X})\leq 0$: we have necessarily
$\mathfrak{S}(-yF(\bm{X}))\geq\Omega(1)$ since the sigmoid function is large
for non-positive values.

Therefore, we have $\mathbb{E}[\mathfrak{S}(-yF(\bm{X}))]\geq\Omega(1).$ ∎

###  D.4 Event III: $\alpha^{(t)}$ keeps increases again

For $t\in[\mathcal{T}_{0},\mathcal{T}_{1}]$, $\Gamma^{(t)}$ increases until
reaching $C\Gamma^{(t)}\geq\Omega(\lambda_{0})/D.$ In this section, we show
that this implies that $\alpha^{(t)}$ increases again.

######  Lemma D.11.

Let $\mathcal{T}_{1}=\mathcal{T}_{0}+\Theta\Big{(}\frac{\lambda_{0}^{p}}{\eta
e^{\mathrm{polyloglog}(d)}}\Big{)}$ and $t\in[\mathcal{T}_{1},T]$. Since
${\large\boldsymbol{\mathpzc{S}}}^{(t)}\geq\max_{\tau\leq
T}|{\large\boldsymbol{\mathpzc{N}}}^{(\tau)}|$, $\alpha^{(t)}$ updates as

$\displaystyle\alpha^{(t+1)}$ $\displaystyle=\alpha^{(t)}+\Theta(\eta
C)(G^{(t)})^{p}(\alpha^{(t)})^{p-1}.$

Consequently, $\alpha^{(t)}$ is non-decreasing until the population risk
satisfies $\mathcal{L}(\bm{A}^{(t)},\bm{v}^{(t)})\leq o(1).$ Eventually,
$\alpha^{(T)}$ is as large as
$\alpha^{(T)}=\frac{\mathrm{polylog}(d)}{C^{2}\lambda_{0}}.$

###### Proof of Lemma D.11.

Since ${\large\boldsymbol{\mathpzc{S}}}^{(t)}\geq\max_{\tau\leq
T}{\large|\boldsymbol{\mathpzc{N}}}^{(\tau)}|$ (Lemma D.8), the update of
$\alpha^{(t)}$ is:

$\displaystyle\alpha^{(t+1)}-\alpha^{(t)}=\Theta(\eta
C)(\alpha^{(t)})^{p-1}(G^{(t)})^{p}\mathbb{E}[\mathfrak{S}(-yF(\bm{X}))],$
(30)

Since the population loss is at least $\Omega(1)$ for
$t\in[\mathcal{T}_{1},T]$, Lemma D.10 implies that
$\mathbb{E}[\mathfrak{S}(-yF(\bm{X}))]\geq\Omega(1)$. Besides, we apply
Induction Hypothesis B.1 and $\Omega(C)\geq 0$ and
$e^{\mathrm{polyloglog}(d)}\leq\lambda_{0}$ in (LABEL:eq:fejdnjwjbbje) and
obtain:

$\displaystyle\begin{cases}\alpha^{(t+1)}\geq\alpha^{(t)}+\Theta(\eta
C)e^{\mathrm{polyloglog}(d)}(\alpha^{(t)})^{p-1}\\\
\alpha^{(t+1)}\leq\alpha^{(t)}+\Theta(\eta
C)\lambda_{0}^{p}(\alpha^{(t)})^{p-1}\end{cases}.$ (31)

(31) and Lemma D.12 show that $\alpha^{(t)}$ increases until reaching
$\alpha^{(t)}\leq\frac{\mathrm{polylog}(d)}{C^{2}\lambda_{0}}$.

∎

####  D.4.1 Auxiliary results

######  Lemma D.12.

The values of $\alpha$ such that
$\mathbb{E}[\mathfrak{S}(-yF(\bm{X}))]\geq\Omega(1)$ is

$\displaystyle\alpha\leq\frac{\mathrm{polylog}(d)}{C^{2}\lambda_{0}}.$

###### Proof of Lemma D.12.

We say that the sigmoid term is small for a constant $\kappa$ that satisfies

$\displaystyle\sum_{\tau=0}^{T}\frac{1}{1+\exp(\kappa)}$
$\displaystyle\leq\mathrm{polylog}(d)\implies\kappa\geq\log(T)\iff\kappa\geq\mathrm{polylog}(d).$
(32)

Intuitively, (32) means that the sum of the sigmoid terms for all time steps
is bounded (up to a logarithmic dependence). In our case, by using (32) and
Lemma D.3, the sigmoid $\mathfrak{S}(-yF(\bm{X}))$ is small when

$\displaystyle(C^{2}\alpha\lambda_{0})^{p}\geq\mathrm{polylog}(d)\implies\alpha\geq\frac{\mathrm{polylog}(d)}{C^{2}\lambda_{0}}.$
(33)

∎

###  D.5 Convergence rate of the population loss

######  Lemma D.13.

Let $t\in[\mathcal{T}_{1},T]$. Then, the population loss linearly converges to
zero i.e.

$\displaystyle\mathcal{L}(\bm{A}^{(t)},\bm{v}^{(t)})\leq\frac{\mathrm{polylog}(d)}{\eta\lambda_{0}^{2p}(t-\mathcal{T}_{1}+1)}$
(34)

###### Proof of Lemma D.13.

To ease the explanation in this proof, we use the $\tilde{\Omega}$,
$\tilde{\Theta}$, $\tilde{O}$ notations to hide the logarithmic dependence. We
hide for instance the constant $C$ under this notation. From Lemma D.11, we
know that $\alpha^{(t)}$ is lower bounded as:

$\displaystyle\alpha^{(t+1)}\geq\alpha^{(t)}+\Theta(\eta
C)(\alpha^{(t)})^{p-1}(G^{(t)})^{p}\mathbb{E}[\mathfrak{S}(-yF(\bm{X}))].$
(35)

Using Lemma D.3, we have
$\mathbb{E}[\mathfrak{S}(-yF(\bm{X}))]\geq\mathfrak{S}(-(\alpha^{(t)}G^{(t)})^{p}).$
Plugging this in (35) yields:

$\displaystyle\alpha^{(t+1)}\geq\alpha^{(t)}+\Theta(\eta)G^{(t)}\frac{(\alpha^{(t)})^{p-1}(G^{(t)})^{p-1}}{1+\exp((\alpha^{(t)}G^{(t)})^{p})}.$
(36)

Since $0<\alpha^{(t)}G^{(t)}\leq\tilde{O}(\lambda_{0}^{p-1})$,we apply Lemma
K.5 and get:

$\displaystyle\alpha^{(t+1)}\geq\alpha^{(t)}+\frac{\tilde{\Omega}(\eta)}{\lambda_{0}^{p}}G^{(t)}\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}}).$
(37)

Lemma D.8 implies that $G^{(t)}\geq\tilde{\Omega}(\lambda_{0}).$ Therefore, we
have:

$\displaystyle\alpha^{(t+1)}\geq\alpha^{(t)}+\frac{\tilde{\Omega}(\eta)}{\lambda_{0}^{p}}\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}}).$
(38)

Let’s now assume by contradiction that for $t\in[\mathcal{T}_{1},T]$, we have:

$\displaystyle\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}})>\frac{\tilde{\Omega}(1)}{\eta\lambda_{0}^{2p}(t-\mathcal{T}_{1}+1)}.$
(39)

For $t\in[\mathcal{T}_{1},T]$, we know that $\alpha^{(t)}G^{(t)}$ is non-
decreasing which implies that $(\alpha^{(t)}G^{(t)})^{p}$ is also non-
decreasing. Since $x\mapsto\log(1+\exp(-x))$ is non-increasing, this implies
for $s\leq t$ that

$\displaystyle\frac{\tilde{\Omega}(1)}{\eta\lambda_{0}^{2p}(t-\mathcal{T}_{1})}<\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}})\leq\log(1+e^{-(\alpha^{(s)}G^{(s)})^{p}}).$
(40)

Plugging (40) in the update (36) yields for $s\in[\mathcal{T}_{1},t]$:

$\displaystyle\alpha^{(s+1)}>\alpha^{(s)}+\frac{\tilde{\Omega}(1)}{\lambda_{0}^{p}(t-\mathcal{T}_{1}+1)}.$
(41)

Let $t\in[\mathcal{T}_{1},T]$. We now sum (41) for $s=\mathcal{T}_{1},\dots,t$
and obtain:

$\displaystyle\alpha^{(t+1)}>\alpha^{(\mathcal{T}_{1})}+\frac{\tilde{\Omega}(1)(t-\mathcal{T}_{1}+1)}{\lambda_{0}^{p}(t-\mathcal{T}_{1}+1)}>\frac{\tilde{\Omega}(1)}{\lambda_{0}^{p}},$
(42)

where we used the fact that
$\alpha^{(\mathcal{T}_{1})}\geq\alpha^{(\mathcal{T}_{0})}\geq\tilde{\Omega}(1)/(C\lambda_{0})>0$
(Lemma D.2) in the last inequality. Therefore, we have for
$t\in[\mathcal{T}_{1},T],$
$\alpha^{(t)}\geq\tilde{\Omega}(1/\lambda_{0}^{p})>0$. Let’s now show that
(42) implies a contradiction. Indeed, we have:

$\displaystyle\eta\lambda_{0}^{2p}(t-\mathcal{T}_{1}+1)\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}})$
$\displaystyle\leq\eta\lambda_{0}^{2p}T\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}})$
$\displaystyle\leq\eta\lambda_{0}^{2p}T\log(1+e^{-\tilde{\Omega}(1)}),$ (43)

where we used $G^{(t)}\geq\Omega(\lambda_{0})$ (Lemma D.8) and (42) in the
last inequality. We now apply Lemma K.6 and obtain:

$\displaystyle\eta\lambda_{0}^{2p}(t-\mathcal{T}_{1}+1)\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}})\leq\frac{\eta\lambda_{0}^{2p}T}{1+\exp(\tilde{\Omega}(1))}.$
(44)

Given the values of $T,\eta,\lambda_{0}$, we finally have:

$\displaystyle\eta\lambda_{0}^{2p}(t-\mathcal{T}_{1}+1)\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}})<\tilde{O}(1),$
(45)

which contradicts (39). Therefore, we obtain the convergence rate:

$\displaystyle\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}})\leq\frac{\tilde{O}(1)}{\eta\lambda_{0}^{2p}(t-\mathcal{T}_{1}+1)}$
(46)

We apply Lemma D.14 to bound the left-hand side of (46) and get the aimed
result. ∎

####  D.5.1 Auxiliary lemmas

######  Lemma D.14.

Let $t\in[\mathcal{T}_{1},T]$. We have:

$\displaystyle\Theta(1)\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}})\geq\mathcal{L}(\bm{A}^{(t)},\bm{v}^{(t)}).$

###### Proof of Lemma D.14.

The proof is similar to the one of Lemma D.3. We apply Lemma D.4, Lemma D.8
and Lemma D.5 and get:

$\displaystyle
yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X})\geq(\alpha^{(t)})^{p}\left(C(G^{(t)})^{p}-o(\lambda_{0})^{p}\right)\geq\Theta(1)(\alpha^{(t)}G^{(t)})^{p}.$
(47)

Using (47), we deduce:

$\displaystyle\mathbb{E}[\log(1+e^{-yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X})})]\leq\log(1+e^{-\Theta(1)(\alpha^{(t)}G^{(t)})^{p}})\leq\Theta(1)\log(1+e^{-(\alpha^{(t)}G^{(t)})^{p}}),$
(48)

where we applied Lemma K.7 in the last inequality. ∎

###  D.6 Fitting the labeling function

We now show that the learner model fits the labeling function.

######  Lemma D.15.

After $T$ iterations, the population risk converges i.e.
$\mathcal{L}(\bm{A}^{(T)},\bm{v}^{(T)})\leq O(1/\mathrm{poly}(d)).$ Therefore,
$\mathbb{P}_{\mathcal{D}}[f^{*}(\bm{X})F_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X})>0]\geq
1-o(1).$

###### Proof of Lemma D.15.

Since the logistic loss is a surrogate for the 0-1 loss, we have:

$\displaystyle\mathbb{E}_{\mathcal{D}}[\mathbf{1}_{yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X})<0}]\leq\mathbb{E}_{\mathcal{D}}[\log(1+e^{-yF_{\bm{A}^{(t)},\bm{v}^{(t)}}(\bm{X})})].$
(49)

We now apply Lemma D.13 to bound the right-hand side of (49). Given the value
of $T$, we have:

$\displaystyle\mathbb{P}_{\mathcal{D}}[yF_{\bm{A}^{(T)},\bm{v}^{(T)}}(\bm{X})]=\mathbb{E}_{\mathcal{D}}[\mathbf{1}_{yF_{\bm{A}^{(T)},\bm{v}^{(T)}}(\bm{X})<0}]\leq\frac{\mathrm{polylog}(d)}{\eta\lambda_{0}^{2p}T}\leq\frac{1}{\mathrm{poly}(d)}.$
(50)

We now use (50) and Definition 2.1 to show that the learner model fits the
labeling function. Indeed, we rewrite
$\mathbb{P}_{\mathcal{D}}[f^{*}(\bm{X})F_{\bm{A}^{(T)},\bm{v}^{(T)}}(\bm{X})]$
as

$\displaystyle\mathbb{P}_{\mathcal{D}}[f^{*}(\bm{X})F_{\bm{A}^{(T)},\bm{v}^{(T)}}(\bm{X})]$
$\displaystyle\geq\mathbb{P}_{\mathcal{D}}[yf^{*}(\bm{X})>0]\mathbb{P}_{\mathcal{D}}[yF_{\bm{A}^{(T)},\bm{v}^{(T)}}(\bm{X})>0]$
$\displaystyle\geq(1-d^{-\omega(1)})\Big{(}1-\frac{1}{\mathrm{poly}(d)}\Big{)}$
$\displaystyle=1-o(1).$ (51)

∎

###  D.7 Proof of the induction hypothesis

In this section, we prove Induction Hypothesis B.1.

###### Proof of Induction Hypothesis B.1.

We start by proving that $|\rho^{(t)}|=\Theta(1)$ for all $t\in[T].$
Let $t\in[\mathcal{T}_{1},T]$ and $\tau\in[t].$ Using Corollary G.2 and
Induction Hypothesis B.1, we upper bound $|\rho^{(\tau)}|$ as:

$\displaystyle|\rho^{(\tau+1)}|$
$\displaystyle\leq|\rho^{(\tau)}|+\frac{\eta\Theta(C)}{D}(\alpha^{(\tau)})^{p}\bigg{(}1+\frac{\lambda_{0}^{p+1}}{D}\bigg{)}.$
(52)

Summing (52) for $\tau=0,\dots,t-1$ and using $\rho^{(0)}=0$ lead to

$\displaystyle|\rho^{(t)}|$
$\displaystyle\leq\frac{\eta\Theta(C)}{D}\bigg{(}1+\frac{\lambda_{0}^{p+1}}{D}\bigg{)}\sum_{\tau=0}^{T}(\alpha^{(\tau)})^{p}$
(53)

We now apply Lemma D.16 to bound the sum of $\alpha^{(t)}$’s in (53).

$\displaystyle|\rho^{(t)}|$
$\displaystyle\leq\Theta\left(\frac{1}{D(\alpha^{(0)})^{p-1}(\lambda_{0})^{p}e^{\mathrm{polyloglog}(d)}}+\frac{1}{e^{\mathrm{polyloglog}(d)}}\right).$
(54)

Given the values of the different parameters, (54) implies that
$|\rho^{(t)}|\leq\Theta(1).$

We now prove $e^{\gamma^{(t)}}\in[\Omega(1),\lambda_{0}].$ Since
$e^{\gamma^{(t)}}$ is non-decreasing (Corollary G.1), we have
$e^{\gamma^{(t)}}\geq e^{\gamma^{(0)}}\geq\Omega(1)$ for all $t\geq 0.$ We now
prove the upper bound on $e^{\gamma^{(t)}}$. We assume that for all $\tau\leq
t$, $e^{\gamma^{(\tau)}}\leq\lambda_{0}.$ Let’s show this inequality for
$t+1$. Using Corollary G.1, we have:

$\displaystyle e^{\gamma^{(t+1)}}$ $\displaystyle\leq
e^{\gamma^{(t)}}\exp\Big{(}\Theta(C\eta)(\alpha^{(t)})^{p}\Gamma^{(t)}(G^{(t)})^{p-1}\Big{)}$
(55)
$\displaystyle=\exp\Big{(}\Theta(C\eta)\sum_{\tau=0}^{t}(\alpha^{(\tau)})^{p}\Gamma^{(\tau)}(G^{(\tau)})^{p-1}\Big{)}.$

We now apply the induction hypothesis in (55) and get:

$\displaystyle e^{\gamma^{(t+1)}}$ $\displaystyle\leq
e^{\gamma^{(t)}}\exp\big{(}\Theta(C\eta)(\alpha^{(t)})^{p}\Gamma^{(t)}(G^{(t)})^{p-1}\big{)}$
(56)
$\displaystyle=\exp\Big{(}\frac{\Theta(\lambda_{0}^{p}C\eta)}{D^{p}}\sum_{\tau=0}^{t}(\alpha^{(\tau)})^{p}\Big{)}$
$\displaystyle\leq\exp\Big{(}\frac{\Theta(\lambda_{0}^{p}C\eta)}{D^{p}}\sum_{\tau=0}^{T}(\alpha^{(\tau)})^{p}\Big{)}.$

We apply Lemma D.16 in (56) and obtain:

$\displaystyle e^{\gamma^{(t+1)}}$
$\displaystyle\leq\exp\bigg{(}\Theta\Big{(}\frac{1}{D^{p}(\alpha^{(0)})^{p-1}e^{\mathrm{polyloglog}(d)}}+\frac{\lambda_{0}^{p}}{D^{p-1}e^{\mathrm{polyloglog}(d)}}\Big{)}\bigg{)}$
$\displaystyle=\Theta(1)\exp\left(\frac{\lambda_{0}^{p}}{D^{p-1}e^{\mathrm{polyloglog}(d)}}\right)$
$\displaystyle\leq\Theta\bigg{(}1+\frac{\lambda_{0}^{p}}{D^{p-1}e^{\mathrm{polyloglog}(d)}}+\frac{\lambda_{0}^{2p}}{D^{2(p-1)}e^{\mathrm{polyloglog}(d)}}\bigg{)},$
(57)

where we used the inequality $e^{x}\leq 1+x+x^{2}$ for $x\leq 1$ in (57).
Given the values of the different parameters, we deduce that
$e^{\gamma^{(t+1)}}\leq\lambda_{0}.$

We now prove $e^{\beta^{(t)}}=e^{\mathrm{polyloglog}(d)}$ for $t\in[0,T].$
Since $\beta^{(t)}$ is not updated i.e. $\beta^{(t)}=\beta^{(0)}$ and
$\beta^{(0)}=\sigma_{\bm{M}}=\mathrm{polyloglog}(d)$, we therefore have the
aimed result.

Lastly, we prove that
$e^{\beta^{(t)}}+(C-1)e^{\gamma^{(t)}}+(D-C)e^{\rho^{(t)}}=\Theta(D)$ for
$t\in[T].$ Since $e^{\gamma^{(t)}}\geq\Theta(1)$, $e^{\rho^{(t)}}=\Theta(1)$
and $e^{\beta^{(t)}}\geq\Theta(1)$, we have:

$\displaystyle
e^{\beta^{(t)}}+(C-1)e^{\gamma^{(t)}}+(D-C)e^{\rho^{(t)}}\geq\Theta(D).$ (58)

On the other hand, we have $e^{\gamma^{(t)}}\leq\lambda_{0}$,
$e^{\beta^{(t)}}=e^{\mathrm{polyloglog}(d)}$ and $e^{\rho^{(t)}}=\Theta(1)$
which imply:

$\displaystyle e^{\beta^{(t)}}+(C-1)e^{\gamma^{(t)}}+(D-C)e^{\rho^{(t)}}$
$\displaystyle\leq e^{\mathrm{polyloglog}(d)}+(C-1)\lambda_{0}+(D-C)\Theta(1)$
(59) $\displaystyle\leq C\lambda_{0}+(D-C)\Theta(1)$
$\displaystyle\leq\Theta(D).$

∎

####  D.7.1 Auxiliary lemmas

######  Lemma D.16.

The sum of the $\alpha^{(t)}$’s is bounded as:

$\displaystyle\sum_{\tau=0}^{T}(\alpha^{(\tau)})^{p}$
$\displaystyle=\Theta\left(\frac{1}{\eta
C(\alpha^{(0)})^{p-1}(\lambda_{0})^{p}e^{\mathrm{polyloglog}(d)}}+\frac{D}{C\eta
e^{\mathrm{polyloglog}(d)}}\right).$

###### Proof of Lemma D.16.

We first decompose the sum of $\alpha^{(t)}$’s.

$\displaystyle\sum_{\tau=0}^{T}(\alpha^{(\tau)})^{p}$
$\displaystyle=\sum_{\tau=0}^{\mathcal{T}_{0}-1}(\alpha^{(\tau)})^{p}+\sum_{\tau=\mathcal{T}_{0}}^{\mathcal{T}_{1}-1}(\alpha^{(\tau)})^{p}+\sum_{\tau=\mathcal{T}_{1}}^{T}(\alpha^{(\tau)})^{p}.$
(60)

We apply Lemma D.2 and Lemma D.8 to rewrite (60).

$\displaystyle\sum_{\tau=0}^{T}(\alpha^{(\tau)})^{p}$
$\displaystyle=\frac{\Theta(1)}{\eta
C(\alpha^{(0)})^{p-1}(\lambda_{0})^{p}e^{\mathrm{polyloglog}(d)}}+\frac{\log\left(\lambda_{0}\right)}{\eta
e^{\mathrm{polyloglog}(d)}}+\sum_{\tau=\mathcal{T}_{1}}^{T}(\alpha^{(\tau)})^{p}.$
(61)

Now, we aim to obtain the value of the last summand in (61). Using Corollary
G.1, we have

$\displaystyle C\Gamma^{(T)}\geq
C\Gamma^{(\mathcal{T}_{1})}\exp\left(\frac{C\eta}{D}e^{\mathrm{polyloglog}(d)}\sum_{\tau=\mathcal{T}_{1}}^{T}(\alpha^{(\tau)})^{p}\right).$
(62)

We finally apply Lemma D.8 and Induction Hypothesis B.1 in (62) to get:

$\displaystyle\sum_{\tau=\mathcal{T}_{1}}^{T}(\alpha^{(\tau)})^{p}\leq\frac{\Theta(D)}{C\eta
e^{\mathrm{polyloglog}(d)}}.$ (63)

To obtain the aimed result, we plug (63) in (61) and use
$\frac{\log\left(\lambda_{0}\right)}{\eta
e^{\mathrm{polyloglog}(d)}}\leq\frac{\Theta(D)}{C\eta
e^{\mathrm{polyloglog}(d)}}$. ∎

##  Appendix E From idealized to real learning process

In Appendix D, we analyzed the ideal learning process. We now aim to bridge
the gap between the idealized and realistic cases. Given our initialization,
$\widehat{\bm{v}}^{(t)}$ has a component in
$\mathrm{span}(\bm{w}^{*})^{\perp}$ i.e.

$\displaystyle\widehat{\bm{v}}^{(t)}=\widehat{\alpha}^{(t)}\bm{w}^{*}+\varepsilon_{\bm{v}}^{(t)}\bm{u}^{(t)}$

where $\bm{u}^{(t)}\in\mathbb{R}^{d}$ such that $\bm{u}^{(t)}\perp\bm{w}^{*}$
and $\|\bm{u}^{(t)}\|_{2}=1.$ Thus, one main difference between the two cases
is that we initialize $\alpha^{(0)}\geq\nu^{1/(p-1)}$ in the idealized case
while $\widehat{\alpha}^{(0)}\leq\omega$. Thus, the proof strategy consists in
i) $t\in[0,\mathscr{T}]$, $\widehat{\alpha}^{(t)}$ increases until having
$\widehat{\alpha}^{(t)}\geq\nu^{1/(p-1)}$ (subsubsection E.1.1) while
$\varepsilon_{\bm{v}}^{(t)}$ (subsubsection E.1.2) and
$\widehat{A}_{i,j}^{(t)}$ (subsubsection E.1.3) stay tiny. ii)
$t\in[\mathscr{T},T]$, compare the realistic and idealized iterates. We remind
the GD update of $\widehat{\alpha}^{(t)}.$

$\displaystyle\widehat{\alpha}^{(t+1)}-\widehat{\alpha}^{(t)}$
(GD-$\widehat{\alpha}$)
$\displaystyle=\eta\underbrace{\frac{D}{N}\sum_{i=1}^{N}y[i]\mathfrak{S}\big{(}-y[i]F(\bm{X}[i])\big{)}\sum_{j\in\mathcal{S}_{\ell(\bm{X}[i])}}\sigma^{\prime}\Big{(}\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{k}[i]\rangle\Big{)}\sum_{r=1}^{D}\widehat{S}_{j,r}^{(t)}\langle\bm{w}^{*},\bm{X}_{r}[i]\rangle}_{{\Large\widehat{\boldsymbol{\mathpzc{S}}}^{(t)}}}$
$\displaystyle+\eta\underbrace{\frac{D}{N}\sum_{i=1}^{N}y[i]\mathfrak{S}\big{(}-y[i]F(\bm{X}[i])\big{)}\sigma^{\prime}\Big{(}\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{k}[i]\rangle\Big{)}\sum_{r=1}^{D}\widehat{S}_{j,r}^{(t)}\langle\bm{w}^{*},\bm{X}_{r}[i]\rangle}_{{\Large\widehat{\boldsymbol{\mathpzc{N}}}^{(t)}}}.$

###  E.1 Bound on the iterates during the initial steps
($t\in[0,\mathscr{T}]$)

Since we randomly initialize $\widehat{\bm{v}}^{(0)}$ with tiny variance, we
need to take into account the linear part of the activation function. Lemma
E.8 shows that we can overlook the power part of the activation and consider
$\sigma(x)=\nu x$ as long as $\widehat{\alpha}^{(t)}\geq\nu^{1/(p-1)}$.

####  E.1.1 $\widehat{\alpha}^{(t)}$ initially increases

######  Lemma E.1.

Let
$\mathscr{T}=\Theta\Big{(}\frac{1}{\eta\nu^{(p-2)/(p-1)}e^{\mathrm{polyloglog}(d)}}\Big{)}$.
For all $t\in[0,\mathscr{T}]$, $\widehat{\alpha}^{(t)}$ is updated as

$\displaystyle\widehat{\alpha}^{(t+1)}$
$\displaystyle=\widehat{\alpha}^{(t)}+\Theta(\eta\nu)e^{\beta}.$

Consequently, $\widehat{\alpha}^{(t)}$ is non-decreasing and after
$\mathscr{T}$ iterations, we have $\widehat{\alpha}^{(t)}\geq\nu^{1/(p-1)}$
for $t\geq\mathscr{T}.$

###### Proof of Lemma E.1.

Let $t\geq 0$. We apply Lemma E.5 and Lemma E.6 to respectively bound
${\large\widehat{\boldsymbol{\mathpzc{S}}}}^{(t)}$ and
$\widehat{\large\boldsymbol{\mathpzc{N}}}^{(t)}$ in the update of
$\widehat{\alpha}^{(t)}$.

$\displaystyle\widehat{\alpha}^{(t+1)}$
$\displaystyle=\widehat{\alpha}^{(t)}+\Theta(C\eta\nu)e^{\beta}.$ (64)

(64) indicates that $\widehat{\alpha}^{(t)}$ is a non-decreasing sequence.
Therefore, there exists a time $\mathscr{T}$ such that
$\widehat{\alpha}^{(\mathscr{T})}=\nu^{1/(p-1)}$. Summing (64) for
$t=0,\dots,\mathscr{T}-1$ yields
$\mathscr{T}=\Theta\Big{(}\frac{1}{\eta\nu^{(p-2)/(p-1)}e^{\beta}}\Big{)}$. ∎

####  E.1.2 Bound on $\varepsilon_{\bm{v}}$

We now show that for $t\in[0,\mathscr{T}]$, the orthogonal component
$\varepsilon_{\bm{v}}^{(t)}$ stays small.

######  Lemma E.2.

Assume that we run GD on the empirical risk (E) for $T$ iterations with
parameters set as in Parametrization 3.1. For $t\in[0,\mathscr{T}]$, the
orthogonal component $\varepsilon_{\bm{v}}$ satisfies

$\displaystyle\varepsilon_{\bm{v}}^{(t+1)}$
$\displaystyle\leq\Big{(}1+\eta\nu\frac{\mathrm{poly}(D)}{\sqrt{d}}\Big{)}\varepsilon_{\bm{v}}^{(t)}+\eta\zeta,$

where $\zeta=\frac{\mathrm{poly}(D)}{\sqrt{N}}.$

###### Proof of Lemma E.2.

Let $\mathbf{P}=(\mathbf{I}-\bm{w}^{*}\bm{w}^{*\top})$ and
$\accentset{\circ}{\bm{v}}^{(t)}=\widehat{\alpha}^{(t)}\bm{w}^{*}$. The
projected update of $\widehat{\bm{v}}$ satisfies:

$\displaystyle\|\mathbf{P}\widehat{\bm{v}}^{(t+1)}-\mathbf{P}\widehat{\bm{v}}^{(t)}\|_{2}$
$\displaystyle\leq$ $\displaystyle
D\eta\nu\biggr{\|}\frac{1}{N}\sum_{i=1}^{N}y[i]\mathfrak{S}(-y[i]F_{\widehat{\bm{v}}}(\bm{X}[i]))\sum_{m=1}^{D}\sum_{b=1}^{D}\widehat{\bm{S}}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}[i]$
(65)
$\displaystyle-\mathbb{E}\biggr{[}y\mathfrak{S}(-yF_{\widehat{\bm{v}}}(\bm{X}))\sum_{m=1}^{D}\sum_{b=1}^{D}\widehat{\bm{S}}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}\biggr{]}\biggr{\|}_{2}$
(66) $\displaystyle+$ $\displaystyle
D\eta\nu\biggr{\|}\mathbb{E}\biggr{[}y\mathfrak{S}(-yF_{\widehat{\bm{v}}}(\bm{X}))\sum_{m=1}^{D}\sum_{b=1}^{D}\widehat{\bm{S}}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}\biggr{]}$
(67)
$\displaystyle-\mathbb{E}\biggr{[}y\mathfrak{S}(-yF_{\accentset{\circ}{v}}(\bm{X}))\sum_{m=1}^{D}\sum_{b=1}^{D}\widehat{\bm{S}}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}\biggr{]}\biggr{\|}_{2}.$
(68)

##### Summand 1: $\|\eqref{eq:fewpwpkjkje}-\eqref{eq:fewefdfvdfewpe}\|_{2}$.

Using the matrix Hoeffding inequality, we have with high probability,
$\|\eqref{eq:fewpwpe}-\eqref{eq:fewefewpe}\|_{2}\leq
16\sqrt{\log(d)\sum_{i=1}^{N}M_{i}^{2}},$ where
$\Big{\|}\frac{\eta\nu
Dy[i]}{N}\mathfrak{S}(-y[i]F_{\widehat{\bm{v}}}(\bm{X}^{(i)}))\sum_{m=1}^{D}\sum_{b=1}^{D}\widehat{\bm{S}}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}^{(i)}\Big{\|}_{2}^{2}\leq
M_{i}^{2}.$ Induction Hypothesis B.1 and
$\|\mathbf{P}\bm{X}_{u}^{(i)}\|_{2}\leq\sigma^{2}d\log(d)$ imply
$M_{i}^{2}\leq\frac{\eta^{2}\nu^{2}\mathrm{poly}(D)}{N^{2}}.$ We deduce that
$\|\eqref{eq:fewpwpkjkje}-\eqref{eq:fewefdfvdfewpe}\|_{2}\leq\eta\nu\frac{\mathrm{poly}(D)}{\sqrt{N}}.$

##### Summand 2:
$\|\eqref{eq:fewkkjjpwpkjkje}-\eqref{eq:fewefdfvdfoiookewpe}\|_{2}$.

We use the 1-Lipschitzness of the sigmoid function and get:

$\displaystyle\|\eqref{eq:fewkkjjpwpkjkje}-\eqref{eq:fewefdfvdfoiookewpe}\|_{2}$
$\displaystyle\leq$
$\displaystyle\eta\nu\Biggr{\|}\mathbb{E}\Biggr{[}Dy\sum_{m=1}^{D}\sum_{b=1}^{D}\widehat{\bm{S}}_{m,b}^{(t)}\mathbf{P}\bm{X}_{u}\cdot$
$\displaystyle\biggr{[}\mathfrak{S}\biggr{(}-y\nu
D\sum_{a=1}^{D}\sum_{a^{\prime}=1}^{D}\widehat{\bm{S}}_{a,a^{\prime}}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{a^{\prime}}\rangle\biggr{)}-\mathfrak{S}\biggr{(}-y\nu
D\sum_{a=1}^{D}\sum_{a^{\prime}=1}^{D}\widehat{\bm{S}}_{a,a^{\prime}}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)},\bm{X}_{a^{\prime}}\rangle\biggr{)}\biggr{]}\Biggr{]}\Biggr{\|}_{2}$
$\displaystyle\leq$ $\displaystyle\Theta(\eta\nu
D)\mathbb{E}\biggr{[}\sum_{m=1}^{D}\sum_{r=1}^{D}\widehat{\bm{S}}_{m,r}^{(t)}\sum_{a=1}^{D}\biggr{|}D\sum_{a^{\prime}=1}^{D}\widehat{\bm{S}}_{a,a^{\prime}}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)}-\widehat{\bm{v}}^{(t)},\bm{X}_{a^{\prime}}\rangle\biggr{|}\sum_{b=1}^{D}\widehat{\bm{S}}_{m,u}^{(t)}\|\mathbf{P}\bm{X}_{b}\|_{2}\biggr{]}$
$\displaystyle\leq$ $\displaystyle\Theta(\eta\nu
D)\mathbb{E}\biggr{[}\sum_{m=1}^{D}\lambda_{0}\varepsilon_{\bm{v}}^{(t)}\cdot\biggr{|}\langle\bm{u}^{(t)},\sum_{a^{\prime}=1}^{D}\bm{\xi}_{a^{\prime}}\rangle\biggr{|}\cdot\sum_{b=1}^{D}\lambda_{0}\|\bm{\xi}_{b}\|_{2}\biggr{]}.$
(69)

where we applied Induction Hypothesis B.1 in (69). Since with high
probability, $\|\bm{\xi}_{b}\|_{2}\leq\sigma\sqrt{d\log(d)}$,
$\big{|}\langle\bm{u}^{(t)},\sum_{r=1}^{D}\bm{\xi}_{r}\rangle\big{|}\leq\sqrt{D\log(d)}\sigma$,
we finally have:

$\displaystyle\|\eqref{eq:fewkkjjpwpkjkje}-\eqref{eq:fewefdfvdfoiookewpe}\|_{2}\leq\eta\nu\cdot\mathrm{poly}(D)\varepsilon_{\bm{v}}^{(t)}\sigma=\eta\nu\frac{\mathrm{poly}(D)}{\sqrt{d}}\varepsilon_{\bm{v}}^{(t)}.$

Combining the bounds on Summands 1 and 2 yields the aimed result. ∎

We now use Lemma E.2 to show that $\varepsilon_{\bm{v}}$ stays small.

######  Lemma E.3.

For all $t\leq\mathscr{T}$,
$\varepsilon_{\bm{v}}^{(t)}\leq\omega\sqrt{d\log(d)}+\nu^{1/(p-1)}\frac{\mathrm{poly}(D)}{\sqrt{N}}.$
By setting $N=\mathrm{poly}(d)$, we have: $\varepsilon_{\bm{v}}^{(t)}\leq
1/\mathrm{poly}(d)$.

###### Proof of Lemma E.3.

Unraveling Lemma E.2 for $t=0,\dots,\mathscr{T}$ and using
$\varepsilon_{\bm{v}}^{(0)}\leq\omega\sqrt{d\log(d)}$ (with high probability)
leads to:

$\displaystyle\varepsilon_{\bm{v}}^{(\mathscr{T})}$
$\displaystyle\leq\omega\sqrt{d\log(d)}+\eta\nu\zeta\frac{\Big{(}1+\eta\nu\frac{\mathrm{poly}(D)}{\sqrt{d}}\Big{)}^{\mathscr{T}}-1}{\eta\nu\frac{\mathrm{poly}(D)}{\sqrt{d}}}\leq\omega\sqrt{d\log(d)}+2\mathscr{T}\eta\nu\zeta,$
(70)

where we used $(1+x)^{y}\leq 1+2yx$ for $x\ll 1$ and $y\geq 0.$ Plugging the
value of $\mathscr{T}$ in (70) yields the aimed result. ∎

####  E.1.3 $\widehat{A}_{a,b}^{(t)}$ stays small

We finally show that $\widehat{A}_{a,b}^{(t)}$ remains tiny for
$t\in[0,\mathscr{T}].$

######  Lemma E.4.

Let $a,b\in[D]$. We have
$|\widehat{A}_{a,b}^{(t)}|\leq\omega\sqrt{d\log(d)}+\frac{\Theta(\nu^{2/(p-1)})}{De^{\beta}}.$

###### Proof of Lemma E.4.

We remind that the GD update of $\widehat{A}_{a,b}^{(t)}$ is

$\displaystyle\widehat{A}_{a,b}^{(t+1)}$
$\displaystyle=\widehat{A}_{a,b}^{(t)}+\frac{\nu\eta}{N}\sum_{i=1}^{N}\mathfrak{S}\big{(}-y[i]F(\bm{X}[i])\big{)}\widehat{S}_{a,b}^{(t)}\sum_{m\neq
b}\widehat{S}_{a,m}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{m}[i]-\bm{X}_{b}[i]\rangle.$
(71)

The proof is by induction. We assume that
$\widehat{A}_{a,b}^{(t)}\leq\omega\sqrt{d\log(d)}+\frac{\Theta(\nu^{2/(p-1)})}{De^{\beta}}.$
We first apply Cauchy-Schwarz on (71) and get:

$\displaystyle|\widehat{A}_{a,b}^{(t+1)}|$
$\displaystyle\leq|\widehat{A}_{a,b}^{(t)}|+\nu\eta\widehat{S}_{a,b}^{(t)}\sum_{m\neq
b}\widehat{S}_{a,m}^{(t)}\|\widehat{\bm{v}}^{(t)}\|_{2}.$ (72)

Using the induction hypothesis, we have
$\widehat{S}_{a,b}^{(t)}\leq\Theta(1)/D.$ Thus, we have

$\displaystyle|\widehat{A}_{a,b}^{(t+1)}|$
$\displaystyle\leq|\widehat{A}_{a,b}^{(t)}|+\frac{\Theta(\nu\eta)}{D}\|\widehat{\bm{v}}^{(t)}\|_{2}\leq|\widehat{A}_{a,b}^{(t)}|+\frac{\Theta(\nu\eta)}{D}\sqrt{(\widehat{\alpha}^{(t)})^{2}+(\varepsilon_{\bm{v}}^{(t)})^{2}}.$
(73)

We sum (73) and get:

$\displaystyle|\widehat{A}_{a,b}^{(\mathscr{T})}|$
$\displaystyle\leq|\widehat{A}_{a,b}^{(0)}|+\frac{\Theta(\nu\eta)}{D}\sum_{t=0}^{\mathscr{T}-1}\sqrt{(\widehat{\alpha}^{(t)})^{2}+(\varepsilon_{\bm{v}}^{(t)})^{2}}.$
(74)

We now use Lemma E.1 and Lemma E.3 in (74) and get:

$\displaystyle|\widehat{A}_{a,b}^{(\mathscr{T})}|$
$\displaystyle\leq\omega\sqrt{d\log(d)}+\frac{\Theta(\nu^{1/(p-1)})}{De^{\beta}}\sqrt{\nu^{2/(p-1)}+\bigg{(}\omega\sqrt{d\log(d)}+\nu^{1/(p-1)}\frac{\mathrm{poly}(D)}{\sqrt{N}}\bigg{)}^{2}}$
$\displaystyle\leq\omega\sqrt{d\log(d)}+\frac{\Theta(\nu^{2/(p-1)})}{De^{\beta}}.$
(75)

∎

####  E.1.4 Auxiliary lemmas

######  Lemma E.5.

Let $\mathscr{T}$ be the time where $\widehat{\alpha}^{(t)}\geq\nu^{1/(p-1)}$.
Let $t\in[0,\mathscr{T}]$, $i\in[N]$ and $j\in\mathcal{S}_{\ell(\bm{X}[i])}.$
We have:

$\displaystyle y[i]\langle\bm{w}^{*},\bm{O}_{j}^{(t)}[i]\rangle$
$\displaystyle=e^{\beta}.$

This implies $\widehat{{\large\boldsymbol{\mathpzc{S}}}}^{(t)}=Ce^{\beta}$ for
all $t\in[0,\mathscr{T}].$

###### Proof of Lemma E.5.

We successively apply Lemma E.4 and Lemma K.3 to get:

$\displaystyle y[i]\langle\bm{w}^{*},\bm{O}_{j}^{(t)}[i]\rangle$
$\displaystyle=D\bigg{(}\widehat{S}_{j,j}^{(t)}+\sum_{k\in\mathcal{S}_{\ell(\bm{X}[i])}}\widehat{S}_{j,k}^{(t)}+y[i]\sum_{h\neq\ell(\bm{X}[i])}\sum_{r\in\mathcal{S}_{h}}\widehat{S}_{j,r}^{(t)}\delta_{r}\bigg{)}$
$\displaystyle\geq\Theta\bigg{(}e^{\mathrm{polyloglog}(d)}+\Theta(C-1)-\Theta(1)\sum_{h\neq\ell(\bm{X}[i])}\sum_{r\in\mathcal{S}_{h}}\delta_{h,r}\bigg{)}$
(76)
$\displaystyle\geq\Theta\bigg{(}e^{\mathrm{polyloglog}(d)}+\Theta(C-1)-\Theta(qD\log(d))\bigg{)}$
(77) $\displaystyle\geq e^{\mathrm{polyloglog}(d)}.$

Similarly, we also have $y[i]\langle\bm{w}^{*},\bm{O}_{j}^{(t)}[i]\rangle\leq
e^{\mathrm{polyloglog}(d)}.$ ∎

######  Lemma E.6.

Let $t\in[0,\mathscr{T}].$ Assume that $N=\mathrm{poly}(d).$ With high
probability, ${\large\widehat{\boldsymbol{\mathpzc{N}}}}^{(t)}\leq
1/\mathrm{poly}(d).$

###### Proof of Lemma E.6.

During this time phase, the sigmoid stays constant. Therefore, we have
$\mathbb{E}[{\large\widehat{\boldsymbol{\mathpzc{N}}}}^{(t)}]=0.$ Therefore,
we apply Hoeffding inequality and Lemma D.6 to get:

$\displaystyle{\large\widehat{\boldsymbol{\mathpzc{N}}}}^{(t)}\leq\sqrt{8N\log(d)\max_{i\in[N]}\;M_{i}^{2}},$
(78)

where $A_{i}$ is a constant such that
$\Big{|}\frac{1}{N}\sum_{j\not\in\mathcal{S}_{\ell(\bm{X}[i])}}y[i]\langle\bm{w}^{*},\bm{O}_{j}^{(t)}[i]\rangle\Big{|}\leq
M_{i}\leq\frac{\mathrm{poly}(D)}{N^{2}}.$ Since $N=\mathrm{poly}(d)$, we
finally proved ${\large\widehat{\boldsymbol{\mathpzc{N}}}}^{(t)}\leq
1/\mathrm{poly}(d).$ ∎

######  Lemma E.7.

Let $t\leq\mathscr{T}$. The sum of $\widehat{\alpha}^{(t)}$’s is bounded as:

$\displaystyle\sum_{\tau=0}^{t}\widehat{\alpha}^{(\tau)}$
$\displaystyle=t\widehat{\alpha}^{(0)}+\Theta(\eta\nu)e^{\beta}t^{2}.$ (79)

###### Proof of Lemma E.7.

Let $\tau\in[0,\mathscr{T}].$ We sum the update rule of
$\widehat{\alpha}^{(t)}$ (Lemma E.1) and obtain:
$\widehat{\alpha}^{(\tau)}=\widehat{\alpha}^{(0)}+\Theta(C\eta\nu)e^{\beta}\tau.$
Summing again this update yields the aimed result.

∎

######  Lemma E.8.

Let $(\bm{X},\cdot)\sim\mathcal{D}$ and $j\in[D]$. Assume that
$\widehat{\alpha}^{(t)}\leq\nu^{1/(p-1)}$. Then, we have:

$\displaystyle\sigma^{\prime}\Big{(}D\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widehat{\bm{v}},\bm{X}_{k}\rangle\Big{)}$
$\displaystyle=\begin{cases}\Theta(\nu)&\text{if
}\widehat{\alpha}^{(t)}\leq\nu^{1/(p-1)}\\\
\Theta(p)\Big{(}D\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widehat{\bm{v}},\bm{X}_{k}\rangle\Big{)}^{p-1}&\text{otherwise}\end{cases}.$

###### Proof of Lemma E.8.

We remind that the derivative of the activation function
$\sigma^{\prime}(x)=px^{p-1}+\nu$. We first remark that for all $x,$
$\sigma^{\prime}(x)\geq\nu$. Besides, we have since $p-1$ is even,

$\displaystyle px^{p-1}+\nu\leq
3\nu\iff|x|\leq\Big{(}\frac{2\nu}{p}\Big{)}^{1/(p-1)}.$ (80)

In our case, we have
$x=D\widehat{\alpha}^{(t)}\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\bm{w}^{*},\bm{X}_{k}\rangle.$
Using Induction Hypothesis B.1, we have $|x|\geq
e^{\mathrm{polyloglog}(d)}\alpha^{(t)}.$ Therefore, a sufficient condition for
(80) to hold is
$\widehat{\alpha}^{(t)}\leq\nu^{1/(p-1)}e^{-\mathrm{polyloglog}(d)}.$ Since
$e^{\mathrm{polyloglog}(d)}\ll\mathrm{poly}(d)$, we can simplify this
condition as $\widehat{\alpha}^{(t)}\leq\nu^{1/(p-1)}$. Proving the second
part of the lemma can be done as in the proof of Lemma D.7. ∎

###  E.2 Coupling between the semi-idealized and realistic processes
($t\in[\mathscr{T},T]$)

In this section, we aim to bound the realistic iterates
$\widehat{A}_{i,j}^{(t)}$ and $\widehat{\bm{v}}^{(t)}$ for
$t\in[\mathscr{T},T].$ For this reason, we introduce a "semi-idealized"
learning process (subsubsection E.2.1) which may be viewed as a mid-point
between the idealized and realistic process. We first bound the iterates in
this process. Then, using this process, we show that
$\varepsilon_{\bm{v}}^{(t)}$ (subsubsection E.2.2) and
$\Delta_{\bm{A}}^{(t)}:=\max_{i\neq
j}|\widehat{A}_{i,j}^{(t)}-\widecheck{A}_{i,j}^{(t)}|$ (subsubsection E.2.4)
stay small. Here, $\widecheck{A}_{i,j}^{(t)}$ is the semi-idealized attention
matrix coefficient. Finally, since $\varepsilon_{\bm{v}}^{(T)}$ and
$\Delta_{\bm{A}}^{(T)}$ are small, the final iterates $\widehat{\alpha}^{(T)}$
and $\alpha^{(T)}$ are equal (subsubsection E.2.6) and thus, the model fits
the labeling function (subsubsection E.2.7).

####  E.2.1 Defining the semi-idealized process

We define an intermediate learning process that we refer to as the "semi-
idealized" process. This process starts at time $t=\mathscr{T}$ involves two
parameters: the semi-idealized value vector $\widecheck{\bm{v}}$ and semi-
idealized attention matrix $\widecheck{\bm{A}}$ defined as

  * –

the value vector $\widecheck{\bm{v}}$ is fixed and satisfies
$\widecheck{\bm{v}}^{(t-\mathscr{T})}=\widehat{\alpha}^{(t)}\bm{w}^{*}$ for
$t\in[\mathscr{T},T].$

  * –

$\widecheck{A}_{i,j}^{(t-\mathscr{T})}$ is a trainable parameter and is
initialized as $\widecheck{A}_{i,j}^{(0)}=0$ for $i\neq j$.

Therefore, the only trainable parameter in this process is
$\widecheck{\bm{A}}$. In the semi-idealized process, we minimize the
population risk

$\displaystyle\min_{\widecheck{\bm{A}}}\;\;\mathbb{E}_{\mathcal{D}}\big{[}\log\big{(}1+e^{-yF(\bm{X})}\big{)}\big{]}:=\widecheck{\mathcal{L}}(\widecheck{\bm{v}},\widecheck{\bm{A}}).$
($\tilde{\text{P}}$)

We remark that such process present similarities to the idealized case. In
particular, it satisfies all the invariance and symmetry properties from Lemma
4.1. We thus define

  * –

$\widecheck{A}_{i,j}^{(t)}=\widecheck{\gamma}^{(t)}$ for all $\ell\in[L]$ and
$i,j\in\mathcal{S}_{\ell}.$

  * –

$\widecheck{A}_{i,j}^{(t)}=\widecheck{\rho}^{(t)}$ for all $\ell,m\in[L]$ such
that $m\neq\ell$ and $i\in\mathcal{S}_{\ell}$ and $j\in\mathcal{S}_{m}.$

Therefore, $\widecheck{\gamma}^{(t)}$ and $\widecheck{\rho}^{(t)}$ are
respectively updated as in Lemma G.1 and Lemma G.2. We define also the softmax
terms

$\displaystyle\widecheck{\Lambda}^{(t)}$
$\displaystyle=\frac{e^{\beta}}{e^{\beta}+(C-1)e^{\widecheck{\gamma}^{(t)}}+(D-C)e^{\widecheck{\rho}^{(t)}}},$
$\displaystyle\widecheck{\Gamma}^{(t)}=\frac{e^{\widecheck{\gamma}^{(t)}}}{e^{\beta}+(C-1)e^{\widecheck{\gamma}^{(t)}}+(D-C)e^{\widecheck{\rho}^{(t)}}},$
$\displaystyle\widecheck{\Xi}^{(t)}$
$\displaystyle=\frac{e^{\widecheck{\rho}^{(t)}}}{e^{\beta}+(C-1)e^{\widecheck{\gamma}^{(t)}}+(D-C)e^{\widecheck{\rho}^{(t)}}},$
$\displaystyle\widecheck{G}^{(t)}=D(\widecheck{\Lambda}^{(t)}+(C-1)\widecheck{\Gamma}^{(t)}).\hskip
56.9055pt$

We finally assume Induction Hypothesis B.1 for this process. This latter can
be proved using the same arguments as in subsection D.7.

####  E.2.2 Realistic dynamics are mainly on $\mathrm{span}(\bm{w}^{*})$

We previously showed in Lemma E.3 that $\varepsilon_{\bm{v}}^{(t)}$ is small
in the initial steps. We now show that it stays small during the whole
process.

######  Lemma E.9.

Assume that we run GD on the empirical risk (E) for $T$ iterations with
parameters set as in Parametrization 3.1 and the number of samples is
$N=\mathrm{poly}(d).$ Then, $\widehat{\bm{v}}^{(t)}$ mainly lies in
$\mathrm{span}(\bm{w}^{*})$ i.e. for $t\leq T$,
$\varepsilon_{\bm{v}}^{(t)}\leq 1/\mathrm{poly}(d).$

We now proceed to the proof of Lemma E.9. We first characterize the recursion
satisfied by $\varepsilon_{\bm{v}}^{(t)}.$

######  Lemma E.10.

Assume that we run GD on the empirical risk (E) for $T$ iterations with
parameters set as in Parametrization 3.1. Then, $\varepsilon_{\bm{v}}^{(t)}$
satisfies for $t\in[\mathscr{T},T]$

$\displaystyle\varepsilon_{\bm{v}}^{(t+1)}$
$\displaystyle\leq\big{(}1+\eta(\widehat{\alpha}^{(t)})^{p-1}\frac{\mathrm{poly}(D)}{\sqrt{d}}\big{)}\varepsilon_{\bm{v}}^{(t)}+\eta\zeta,\qquad\text{where
}\zeta=\frac{\mathrm{poly}(D)}{\sqrt{N}}.$

###### Proof of Lemma E.10.

Let $\mathbf{P}:=(\mathbf{I}-\bm{w}^{*}\bm{w}^{*\top})$,
$\accentset{\circ}{\bm{v}}^{(t)}:=\widehat{\alpha}^{(t)}\bm{w}^{*}$ and
$t\in[\mathscr{T},T].$ The projected update of $\widehat{\bm{v}}$ satisfies:

$\displaystyle\|\mathbf{P}\widehat{\bm{v}}^{(t+1)}-\mathbf{P}\widehat{\bm{v}}^{(t)}\|_{2}$
$\displaystyle\leq$
$\displaystyle\eta\biggr{\|}\frac{1}{N}\sum_{i=1}^{N}Dy[i]\mathfrak{S}(-y[i]F_{\widehat{\bm{v}}}(\bm{X}[i]))\sum_{m=1}^{D}\biggr{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{r}^{(i)}\rangle\biggr{)}^{p-1}\sum_{b=1}^{D}\widehat{S}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}[i]$
(81) $\displaystyle-$
$\displaystyle\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widehat{\bm{v}}}(\bm{X}))\sum_{m=1}^{D}\biggr{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{r}\rangle\biggr{)}^{p-1}\sum_{b=1}^{D}\widehat{S}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}\biggr{]}\biggr{\|}_{2}$
(82) $\displaystyle+$