
Theorem \ref{thm: main_theory} characterizes under what condition of the data the neural network with self-attention in (\ref{eqn: attention}) trained with Algorithm \ref{alg: sgd} can achieve zero generalization error. To show that the self-attention layer can improve the generalization performance by reducing the required sample complexity to achieve zero generalization error, we also quantify the sample complexity when there is no self-attention layer in the following proposition.

\begin{prpst}[Generalization without self-attention]\label{prpst: ViT_CNN}
Suppose assumptions in Theorem \ref{thm: main_theory} hold. 
    When there is no self-attention layer, i.e.,
    ${\boldsymbol W}_K$ and ${\boldsymbol W}_Q$ are not updated during the training,  if $N$ satisfies
    \begin{equation}\label{eqn:no_attention}
    N\geq \Omega(\frac{1}{(\alpha_*(\alpha_*-\sigma-\tau))^2})
    \end{equation}
    then after $T$ iterations with $T$ in (\ref{eqn: iterations}), the returned model achieves zero generalization error as 
    \begin{equation}
    f({\boldsymbol A}^{(0)}, {\boldsymbol W}_O^{(T)},{\boldsymbol W}_V^{(T)},{\boldsymbol W}_K^{(0)},{\boldsymbol W}_Q^{(0)})=0
\end{equation}
    \end{prpst}

\begin{rmk} (Advantage of the self-attention layer) Because $m\gg m_a, m_b, d$, the number of trainable parameter   remains almost the same with or without updating the attention layer. Combining Theorem \ref{thm: main_theory} and Proposition \ref{prpst: ViT_CNN}, we can see that with the additional self-attention layer, the sample complexity\footnote{The sample complexity bounds in (\ref{eqn: sample_complexity}) and (\ref{eqn:no_attention}) are sufficient but not necessary. Thus, rigorously speaking, one can not compare two cases based on sufficient conditions only. In our analysis, however, these two bounds are derived with   exactly the same technique with the only difference in handling the self-attention layer. Therefore, we believe it is fair to compare these two bounds to show the advantage of ViT.} is reduced by a factor $1/\alpha_*^2$ with an approximately equal number of network parameters. 
\end{rmk}

\begin{rmk}\label{rmk: sparsification} (Generalization improvement by token sparsification). (\ref{eqn: sample_complexity})  and (\ref{eqn: iterations}) show that the sample complexity $N$ and the required number of iterations $T$ scale with $1/\alpha_*^2$  and $1/\alpha_*$, respectively. Then, increasing $\alpha_*$, the fraction of label-relevant tokens, can reduce the sample complexity and speed up the convergence. Similarly,  $N$ and $T$ scale with  $1/(\Theta(1)-\tau)^2$  and $1/(\Theta(1)-\tau)$. Then decreasing $\tau$, the noise in the tokens, can also improve the generalization. Note that a properly designed token sparsification method can both increase $\alpha_*$ by removing label-irrelevant tokens and decrease $\tau$ by removing noisy tokens, thus improving the generalization performance. 
\end{rmk}

\begin{rmk}\label{rmk: initial_model} (Impact of the initial model) The initial model  ${\boldsymbol W}_V^{(0)}$, ${\boldsymbol W}_K^{(0)}$, ${\boldsymbol W}_Q^{(0)}$ affects the learning performance through $\sigma$ and $\delta$, both of which decrease as the initial model is improved. Then from (\ref{eqn: sample_complexity}) and (\ref{eqn: iterations}), the sample complexity reduces and the convergence speeds up for a better initial model.  
\end{rmk}
 
\iffalse
    Token sparsification can improve the generalization performance by reducing the sample complexity. On one hand, removing tokens with label-irrelevant patterns can enlarge $\alpha^*$ in (\ref{eqn: sample_complexity}). On the other hand, one can eliminate tokens of large initialization errors to reduce the upper bound of the error term $\sigma$ of tokens of tolerable initialization errors.
\fi 

    
Proposition \ref{lemma: concentration} shows that the attention weights are increasingly concentrated on label-relevant tokens during the training.  Proposition  \ref{lemma: concentration} is a critical component in proving Theorem \ref{thm: main_theory} and is of independent interest. 

\begin{prpst}\label{lemma: concentration}
The attention weights for each token become increasingly concentrated on those correlated with tokens of the label-relevant pattern during the training, i.e.,
\begin{equation}\label{eqn:attention}
    \sum_{i\in\mathcal{S}_*^n}\text{softmax}({{\boldsymbol x}_i^n}^\top{{\boldsymbol W}_K^{(t)}}^\top{\boldsymbol W}_Q^{(t)}{\boldsymbol x}_l^n)=\sum_{i\in\mathcal{S}_*^n}\frac{\exp({{\boldsymbol x}_i^n}^\top{{\boldsymbol W}_K^{(t)}}^\top{\boldsymbol W}_Q^{(t)}{\boldsymbol x}_l^n)}{\sum_{r\in\mathcal{S}^n}\exp({{\boldsymbol x}_r^n}^\top{{\boldsymbol W}_K^{(t)}}^\top{\boldsymbol W}_Q^{(t)}{\boldsymbol x}_l^n)}\rightarrow 1-\eta^C
\end{equation}
at a sublinear rate of $O(1/t)$ as $t$ increases for a large $C>0$ and all $l\in\mathcal{S}^n$ and $n\in[N]$.
\end{prpst}

\iffalse
This result comes from the growth of magnitude of label-relevant features ${\boldsymbol q}_1$, ${\boldsymbol q}_2$, ${\boldsymbol r}_1$, and ${\boldsymbol r}_2$ during the training, which increases the attention weight in terms of correlations with label-relevant patterns for all tokens.
\fi
Proposition \ref{lemma: concentration} indicates that only label-relevant tokens are highlighted by the learned attention of ViTs,
while other tokens have less weight. This provides a theoretical justification of
magnitude-based token sparsification methods.





\textbf{Proof idea sketch}: The main proof idea is to show that
the SGD updates scale up value, query, and key vectors of discriminative patterns, while  keeping the magnitude of the projections of non-discriminative patterns and the initial model error almost unchanged. To be more specific, by Lemma \ref{lemma: initial_WU}, \ref{lemma: update_WU}, we can identify two groups of neurons in the hidden layer  ${\boldsymbol W}_{O}$,
where one group only learns the positive pattern, and the other group only learns the negative pattern. Claim \ref{clm: W_O} of Lemma \ref{lm: training} states that   during the SGD updates, the neuron weights in these two groups evolve in the direction of
projected discriminative patterns,  ${\boldsymbol p}_1$ and ${\boldsymbol p}_2$, respectively. Meanwhile,  Claim \ref{clm: W_QK} of Lemma \ref{lm: training} indicates that  ${\boldsymbol W}_K$ and ${\boldsymbol W}_Q$ update in the direction of increasing the magnitude of the query and key vectors of label-relevant tokens from $1$ to $\Theta(\log T)$, such that the attention weights correlated with label-relevant tokens gradually become dominant.
Moreover, by Claim \ref{clm: W_V} of Lemma \ref{lm: training}, the update of ${\boldsymbol W}_V$ increases the magnitude of the value vectors of label-relevant tokens, by adding partial neuron weights of ${\boldsymbol W}_O$ that are aligned with the value vectors
to these  vectors.
Due to the above properties during the training, one can simplify the training process to show that
the output of neural network (\ref{eqn: attention}) changes linearly in the iteration number $t$. From the above analysis, we can develop the sample complexity and the required number of iterations for the zero generalization guarantee.

\textbf{Technical novelty}: Our proof technique is inspired by the feature learning technique in analyzing fully connect networks and convolution neural networks \citep{SWL21, BG21}. Our paper makes new technical contributions from the following aspects. First, we provide a new framework of studying the nonconvex interactions of multiple weight matrices in a three-layer ViT
while other feature learning works \citep{SWL21, BG21, KWLS21, AL22, WL21, ZWCL23} only study one trainable 
weight matrix in the hidden layer of a two-layer network. Second, we
analyze the updates of the self-attention module with the softmax function during the training, while other papers either ignore this issue  without exploring convergence analysis \citep{EGKZ22} or oversimplify the analysis by applying the neural-tangent-kernel (NTK) method
that considers impractical over-parameterization and updates the weights only around initialization.  \citep{HBSN20, Y20, ALL19, ADHL19}.