SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yongtao Wu
yongtao.wu@epfl.ch
Fanghui Liu*
fanghui.liu@warwick.ac.uk
Grigorios G Chrysos*
chrysos@wisc.edu
In this paper, we aim to build the global convergence theory of encoder-only shallow Transformers under a realistic setting from the perspective of architectures, initialization, and scaling under a finite width regime. The difficulty lies in how to tackle the softmax in self-attention mechanism, the core ingredient of Transformer. In particular, we diagnose the scaling scheme, carefully tackle the input/output of softmax, and prove that quadratic overparameterization is sufficient for global convergence of our shallow Transformers under commonly-used He/LeCun initialization in practice. Besides, neural tangent kernel (NTK) based analysis is also given, which facilitates a comprehensive comparison. Our theory demonstrates the separation on the importance of different scaling schemes and initialization. We believe our results can pave the way for a better understanding of modern Transformers, particularly on training dynamics.
Transformers [Vaswani et al., 2017] have demonstrated unparalleled success in influential applications [Devlin et al., 2019, Brown et al., 2020, Wang et al., 2018, Dosovitskiy et al., 2021, Liu et al., 2022b]. A fundamental theoretical topic concerns the global convergence, i.e., the training dynamics of Transformers, which would be helpful for further analysis, e.g., in-context learning [von Oswald et al., 2022, Akyurek et al., 2023], generalization [Li et al., 2023]. In fact, even within a simplified Transformer framework under certain specific regimes, the global convergence guarantees still remain an elusive challenge.
To theoretically understand this, let us first recall the exact format of the self-attention mechanism, the core ingredient of the Transformer. Given the input  $\mathbf{X} \in \mathbb{R}^{d_s \times d}$  ( $d_s$  is the number of tokens and  $d$  is the feature dimension of each token), a self-attention mechanism is defined as:
Self-attention  $(\pmb{X}) \triangleq \sigma_{s}\left(\tau_{0}(\pmb{X}\pmb{W}_{Q}^{\top})(\pmb{X}\pmb{W}_{K}^{\top})^{\top}\right)(\pmb{X}\pmb{W}_{V}^{\top}) = \sigma_{s}\left(\tau_{0}\pmb{X}\pmb{W}_{Q}^{\top}\pmb{W}_{K}\pmb{X}^{\top}\right)(\pmb{X}\pmb{W}_{V}^{\top})$
where  $\sigma_{s}$  is the row-wise softmax function,  $\tau_0\in \mathbb{R}^+$  is the scaling factor, and the learnable weights are  $W_{Q},W_{K},W_{V}\in \mathbb{R}^{d_{m}\times d}$  with the width  $d_m$ . Given  $\pmb{X}$ , the input of softmax depends on  $\tau_0\pmb{W}_Q^\top \pmb{W}_K$ , including the scaling factor  $\tau_0$  and initialization schemes for learnable parameters, and thus determines the output of softmax and then affects the performance of Transformers in both
*Work done at LIONS, EPFL.
theory and practice. There are several scaling schemes in previous literature. For instance, given  $W_{Q}$  and  $W_{K}$  initialized by standard Gaussian, the scaling factor  $\tau_0$  is chosen by
-  $\tau_0 = d_m^{-1/2}$  in the original Transformer [Vaswani et al., 2017]: each element in  $\tau_0\mathbf{W}_Q^\top \mathbf{W}_K$  is a random variable with mean 0 and variance 1. This scaling avoids the blow-up of value inside softmax as  $d_m$  increases [Hron et al., 2020].
-  $\tau_0 = d_m^{-1}$ : This scaling stems from the neural tangent kernel (NTK) analysis [Jacot et al., 2018], a commonly used technical tool for convergence analysis of fully-connected (or convolutional) networks under an infinite-width setting  $d_m \to \infty$ . However, for Transformer, if one uses this scaling under the infinite-width setting, then by the law of large numbers, we have  $\lim_{d_m \to \infty} \tau_0[\mathbf{W}_Q^{\top} \mathbf{W}_K]^{(ij)} = 0$ . As a result, the input of softmax is zero and the softmax degenerates to a pooling layer. That means, the non-linearity is missing, which motivates researchers to carefully rethink this setting.
For instance, under the  $\tau_0 = d_m^{-1}$  setting, Yang [2020] use the same query and key matrices to prevent the softmax from degenerating into a pooling layer. Besides, to avoid the analytic difficulty of softmax due to the fact that each element of the output depends on all inputs, Hron et al. [2020] substitute softmax with ReLU under the  $\tau_0 = d_m^{-1/2}$  and infinite width setting for simplicity.
Clearly, there exists a gap between theoretical analysis and practical architectures on the use of softmax, and accordingly, this leads to the following open question:
How can we ensure the global convergence of Transformers under a realistic setting?
The primary contribution of this work is to establish the convergence theory of shallow Transformer under a realistic setting. Despite its shallow and encoder-only architecture, our Transformer model captures all the fundamental components found on typical Transformers, including the self-attention mechanism with the softmax activation function, one feedforward ReLU layer, one average pooling layer, and a linear output layer, cf. Eq. (1.2). We adopt the  $\tau_0 = d_m^{-1/2}$  scaling under the finite-width setting and compare the results of LeCun/He initializations, which are commonly used in practical applications. Besides, the convergence result under the  $\tau_0 = d_m^{-1}$  setting (as well as the NTK based analysis) is also studied, which facilitates a comprehensive comparison. Our theoretical results demonstrate notable separations among scaling settings, initializations, and architectures as below:
- Scaling: The global convergence can be achieved under both  $\tau_0 = d_m^{-1/2}$  and  $\tau_0 = d_m^{-1}$ . Nevertheless, as suggested by our theory: for a small  $d_m$ , there is no significant difference for these two scaling schemes on the convergence; but for a large enough  $d_m$ , the  $\tau_0 = d_m^{-1/2}$  scaling admits a faster convergence rate of Transformers than that with  $\tau_0 = d_m^{-1}$ . Interestingly, under this  $\tau_0 = d_m^{-1}$  setting, our theory also demonstrates the separation on the convergence result, depending on whether the input is formed along sequence dimension ( $d = 1$ ) or embedding dimension ( $d_s = 1$ ).
- Initialization: Under LeCun and He initialization, our shallow Transformer admits a faster convergence rate than the NTK initialization. This could be an explanation for the seldom usage of NTK initialization for Transformer training in practice.
- Architecture: Quadratic over-parameterization is enough to ensure the global convergence of our shallow Transformer. As a comparison, if the self-attention mechanism is substituted by a feed-forward ReLU layer, our shallow Transformer is close to a three-layer fully-connected ReLU neural networks to some extent, requiring cubic over-parameterization for global convergence.
We firmly believe that our theoretical analysis takes a significant step towards unraveling the mysteries behind Transformers from the perspective of global convergence. We hope that our analytical framework and insights on various initialization and scaling techniques would be helpful in training modern, large-scale Transformer-based models [Radford et al., 2018, Brown et al., 2020].
Self-attention, Transformer: Regarding training dynamics, Snell et al. [2021] explain why singlehead attention focuses on salient words by analyzing the evolution throughout training. Hron et al. [2020] show that the output of Transformer converges to Gaussian process kernel and provide the
NTK formulation of Transformer. Recently, Li et al. [2023] provide sample complexity of shallow Transformer to study its generalization property under a good initialization from pretrained model. The separation between the Transformer and CNN is recently explored. Jelassi et al. [2022] provably demonstrate that Vision Transformer (ViT) has the ability to learn spatial structure without additional inductive bias such as the spatial locality in CNN. Chen et al. [2022] study the loss landscape of ViT and find that ViT converges at sharper local minima than ResNet. Park and Kim [2022] show that ViT is a low-pass filter while CNN is a high-pass filter, thus, these two models can be complementary.
NTK, lazy training, Hessian: The NTK was introduced by Jacot et al. [2018] to connect the infinite-width neural network trained by gradient descent and the kernel regression. The roles of such kernel include analysis of the training dynamics of the neural network in the over-parameterization regime [Allen-Zhu et al., 2019a, Chizat et al., 2019, Du et al., 2019a,b, Zou et al., 2020]. The global convergence, generalization bound, and memorization capacity largely depend on the minimum eigenvalue of the NTK [Cao and Gu, 2019, Zhu et al., 2022, Nguyen et al., 2021, Bombari et al., 2022]. Even though the NTK is extended from FCNN to several typical networks including Transformer [Tirer et al., 2020, Huang et al., 2020, Arora et al., 2019b, Alemohammad et al., 2021, Nguyen and Mondelli, 2020], it has not been used to analyze the global convergence of Transformer. On the other hand, the stability of the tangent kernel during training is required when connecting to kernel regression, but such stability can not be explained by the phenomenon of lazy training [Chizat et al., 2019], which indicates a small change of the parameters from initialization. The hessian spectral bound is the main reason for the stability of kernel, as mentioned in Liu et al. [2020].
Over-parameterization for convergence analysis: Due to over-parameterization, neural networks (NNs) can fit arbitrary labels with zero training loss when trained with (stochastic) gradient descent (SGD), both theoretically Li and Liang [2018], Du et al. [2019b] and empirically [Zhang et al., 2017]. This leads to an interesting question in theory: how much overparameterization is enough to ensure global convergence of NNs? A common recipe for the proof of global convergence relies on the variant of Polyak-Lojasiewicz condition [Polyak, 1963, Liu et al., 2022a], NTK [Du et al., 2019b,a, Zou and Gu, 2019, Allen-Zhu et al., 2019a], or the minimum eigenvalue of the gram matrix [Nguyen, 2021, Bombari et al., 2022]. In Appendix B.3, we provide a comprehensive overview of a recent line of work that improves the over-parameterization condition for ensuring the convergence of NNs. However, the over-parameterization condition for Transformer to achieve global convergence remains elusive from existing literature and we make an initial step towards this question.
This section includes the problem setting with notations and model formulation of the shallow Transformer that is studied in this paper.
Vectors (matrices) are symbolized by lowercase (uppercase) boldface letters, e.g.,  $\boldsymbol{w}$ ,  $\boldsymbol{W}$ . We use  $\| \cdot \|_{\mathrm{F}}$  and  $\| \cdot \|_{2}$  to represent the Frobenius norm and the spectral norm of a matrix, respectively. The Euclidean norm of a vector is symbolized by  $\| \cdot \|_{2}$ . The superscript with brackets is used to represent the element of a vector/matrix, e.g.,  $w^{(i)}$  is the  $i^{\text{th}}$  element of  $\boldsymbol{w}$ . The superscript without brackets symbolizes the parameters at different training step, e.g.,  $\theta^{t}$ . We denote by  $[N] = \{1, \dots, N\}$  for short. We use  $\sigma_{\min}(\cdot)$  and  $\lambda_{\min}(\cdot)$  to represent the minimum singular value and minimum eigenvalue of a matrix. The NTK matrix and hessian matrix of the network are denoted by  $K$  and  $H$ , respectively. The order notation, e.g.,  $\widetilde{\mathcal{O}}$ ,  $\widetilde{\Omega}$ , omits the logarithmic factor. More detailed notation can be found in Table 2 of the appendix.
Let  $X \subseteq \mathbb{R}^{d_s \times d}$  be a compact metric space and  $Y \subseteq \mathbb{R}$ , where  $d$  is the dimension of each token,  $d_s$  is the total sequence length of the input. The training set  $\{(X_n, y_n)\}_{n=1}^N$  is assumed to be iid sampled from an unknown probability measure on  $X \times Y$ . In this paper, we focus the regression task by employing the squared loss. The goal of our regression task is to find a hypothesis, i.e., a Transformer  $f: X \to Y$  in our work, such that  $f(\boldsymbol{X}; \boldsymbol{\theta})$  parameterized by  $\boldsymbol{\theta}$  is a good approximation of the label  $y \in Y$  corresponding to a new sample  $\boldsymbol{X} \in X$ . We use a vector  $\boldsymbol{\theta}$  to denote the collection of all learnable parameters.
Table 1: Common initialization methods with their variances of Gaussian distribution and scaling factors. The choice of  $\tau_{1} = d_{m}^{-1 / 2}$  is based on standard NTK initialization on prior literature [Du et al., 2019b].
Init.	ηO	ηV	ηQ	ηK	τ1
LeCun	dm-1	d-1	d-1	d-1	1
He	2dm-1	2d-1	2d-1	2d-1	1
NTK	1	1	1	1	dm-1/2
Algorithm 1: Gradient descent training
Input: data  $(X_{n},y_{n})_{n = 1}^{N}$  , step size  $\gamma$
Initialize weights as follows:
$$
\boldsymbol {\theta} ^ {0} := \{\bar {\boldsymbol {W}} _ {Q} ^ {0}, \boldsymbol {W} _ {K} ^ {0}, \boldsymbol {W} _ {V} ^ {0}, \boldsymbol {W} _ {O} ^ {0} \}.
$$
for  $t = 0$  to  $t' - 1$  do
$$
\boldsymbol {W} _ {Q} ^ {t + 1} = \boldsymbol {W} _ {Q} ^ {t} - \gamma \cdot \nabla_ {\boldsymbol {W} _ {Q}} \ell (\boldsymbol {\theta} ^ {t}),
$$
$$
\boldsymbol {W} _ {K} ^ {t + 1} = \boldsymbol {W} _ {K} ^ {t} - \gamma \cdot \nabla_ {\boldsymbol {W} _ {K}} \ell (\boldsymbol {\theta} ^ {t}),
$$
$$
\boldsymbol {W} _ {V} ^ {t + 1} = \boldsymbol {W} _ {V} ^ {t} - \gamma \cdot \nabla_ {\boldsymbol {W} _ {V}} \ell (\boldsymbol {\theta} ^ {t}),
$$
$$
\boldsymbol {W} _ {O} ^ {t + 1} = \boldsymbol {W} _ {O} ^ {t} - \gamma \cdot \nabla_ {\boldsymbol {W} _ {O}} \ell (\boldsymbol {\theta} ^ {t}).
$$
Output: the model based on  $\theta^{t'}$ .
3.2 Model formulation of shallow Transformer
Throughout this work, we consider the encoder of Transformer, which can be applied to both regression and classification tasks [Yüksel et al., 2019, Dosovitskiy et al., 2021]. Given an input  $\mathbf{X} \in \mathbb{R}^{d_s \times d}$ , the model is defined as below:
$$
\boldsymbol {A} _ {1} = \text {S e l f - a t t e n t i o n} (\boldsymbol {X}) \triangleq \sigma_ {s} \left(\tau_ {0} \left(\boldsymbol {X} \boldsymbol {W} _ {Q} ^ {\top}\right) \left(\boldsymbol {X} \boldsymbol {W} _ {K} ^ {\top}\right) ^ {\top}\right) \left(\boldsymbol {X} \boldsymbol {W} _ {V} ^ {\top}\right), \tag {1.1}
$$
$$
\boldsymbol {A} _ {2} = \tau_ {1} \sigma_ {r} (\boldsymbol {A} _ {1} \boldsymbol {W} _ {H}), \quad \boldsymbol {a} _ {3} = \varphi (\boldsymbol {A} _ {2}), \quad f (\boldsymbol {X}; \boldsymbol {\theta}) = \boldsymbol {a} _ {3} ^ {\top} \boldsymbol {w} _ {O}, \tag {1.2}
$$
where the output is  $f(\mathbf{X};\pmb {\theta})\in \mathbb{R}$ ,  $\tau_0$  and  $\tau_{1}$  are two scaling factors. The ingredients of a Transformer with width  $d_{m}$  are defined as follows:
- A self-attention mechanism (Eq. (1.1)):  $\sigma_s$  is the row-wise softmax function; the learnable parameters are  $W_{Q}, W_{K}, W_{V} \in \mathbb{R}^{d_{m} \times d}$ . We employ Gaussian initialization  $W_{Q}^{(ij)} \sim \mathcal{N}(0, \eta_{Q})$ ,  $W_{K}^{(ij)} \sim \mathcal{N}(0, \eta_{K})$ ,  $W_{V}^{(ij)} \sim \mathcal{N}(0, \eta_{V})$  with  $i \in [d_m]$  and  $j \in [d]$ . Refer to Table 1 for typical initialization examples.
- A feed-forward ReLU layer (in Eq. (1.2)):  $\sigma_r$  is the ReLU activation function; the learnable parameter is  $W_H \in \mathbb{R}^{d_m \times d_m}$ . Following Yang et al. [2022], we combine  $W_V$  and  $W_H$  together (by setting  $W_H = I$ ) for ease of the analysis. Note that it does not mean its training dynamics are the same as the joint-training of these two adjacent matrices.
- An average pooling layer (in Eq. (1.2)):  $\varphi$  indicates the column-wise average pooling. Note that the average pooling layer is applied along the sequence length dimension to ensure the final output is a scalar, which is commonly used in practical Vision Transformer or theoretical analysis [Dosovitskiy et al., 2021, Yang, 2020].
- An output layer (in Eq. (1.2)) with learnable parameter  $\pmb{w}_O \in \mathbb{R}^{d_m}$ , initialized by  $\pmb{w}_O^{(i)} \sim \mathcal{N}(0, \eta_O)$ .
Remarks: Proper initialization and scaling are required to ensure the convergence and learnability, as seen in previous work [Jacot et al., 2018, Tirer et al., 2022, Lee et al., 2019]. For our convergence analysis, we consider standard Gaussian initialization with different variances and different scaling factor that includes three typical initialization schemes in practice. In Table 1, we detail the formula of LeCun initialization, He initialization, and NTK initialization.
Given  $N$  input samples  $\{\pmb{X}_n\}_{n=1}^N$ , the corresponding ground truth label, the final output of network, and the output of the last hidden layer, are denoted by:
$$
\boldsymbol {y} \triangleq \left\{y _ {n} \right\} _ {n = 1} ^ {N} \in \mathbb {R} ^ {N}, \quad \boldsymbol {f} (\boldsymbol {\theta}) \triangleq \left\{f (\boldsymbol {X} _ {n}; \boldsymbol {\theta}) \right\} _ {n = 1} ^ {N} \in \mathbb {R} ^ {N}, \quad \boldsymbol {F} _ {\mathrm {p r e}} (\boldsymbol {\theta}) \triangleq \left\{\boldsymbol {a} _ {3} (\boldsymbol {X} _ {n}; \boldsymbol {\theta}) \right\} _ {n = 1} ^ {N} \in \mathbb {R} ^ {N \times d _ {m}}.
$$
We consider standard gradient descent (GD) training of Transformer, as illustrated in Algorithm 1. Here the squared loss is expressed as  $\ell (\pmb {\theta}) = \frac{1}{2}\| \pmb {f}(\pmb {\theta}) - \pmb {y}\| _2^2$
In this section, we study the convergence guarantee of Transformer training by GD under the squared loss. Firstly, we provide a general analytical framework in Section 4.1 covering different initialization
schemes, where we identify the condition for achieving global convergence. Next, we validate these conditions for several practical initialization schemes under  $\tau_0 = d_m^{-1/2}$  in Section 4.2 and  $\tau_0 = d_m^{-1}$  in Section 4.3, respectively. We include NTK-based results for self-completeness. Discussion on the convergence results with different scalings, initializations, and architectures is present in Section 4.4.
4.1 General framework for convergence analysis
Before presenting our result, we introduce a basic assumption.
Assumption 1. The input data is bounded  $\| X\|_{\mathrm{F}}\leq \sqrt{d_s} C_x$  with some positive constant  $C_x$
Remark: This assumption is standard as we can always scale the input. The embedding of token is usually assumed to have a unit norm [Li et al., 2023], which is unrelated to  $d$ .
Now we are ready to present our proposition, with the proof deferred to Appendix C.2. Notice that our proposition holds with high probability under some conditions, which depend on certain initialization schemes and scaling factors. Since our proposition is devoted to a unifying analysis framework under various initialization schemes, we do not include specific probabilities here. The validation of the required conditions and probability is deferred to Sections 4.2 and 4.3, respectively.
Proposition 1. Consider the Transformer with  $d_m \geq N$ . Let  $C_Q, C_K, C_V, C_O$  be some positive constants, and define the following quantities at initialization for simplification:
- The norm of the parameters:
$$
\bar {\lambda} _ {Q} \triangleq \left\| \boldsymbol {W} _ {Q} ^ {0} \right\| _ {2} + C _ {Q}, \bar {\lambda} _ {K} \triangleq \left\| \boldsymbol {W} _ {K} ^ {0} \right\| _ {2} + C _ {K}, \bar {\lambda} _ {V} \triangleq \left\| \boldsymbol {W} _ {V} ^ {0} \right\| _ {2} + C _ {V}, \bar {\lambda} _ {O} \triangleq \left\| \boldsymbol {w} _ {O} ^ {0} \right\| _ {2} + C _ {O}.
$$
- Two auxiliary terms:  $\rho \triangleq N^{1/2} d_s^{3/2} \tau_1 C_x$ ,  $z \triangleq \bar{\lambda}_O^2 \left(1 + 4 \tau_0^2 C_x^4 d_s^2 \bar{\lambda}_V^2 \left(\bar{\lambda}_Q^2 + \bar{\lambda}_K^2\right)\right)$ .
Under Assumption 1, we additionally assume that the minimum singular value of  $\mathbf{F}_{\mathrm{pre}}^{0}$ , i.e.,  $\alpha \triangleq \sigma_{\min}\left(\mathbf{F}_{\mathrm{pre}}^{0}\right)$  satisfies the following condition at initialization:
$$
\alpha^ {2} \geq 8 \rho \max  \left(\bar {\lambda} _ {V} C _ {O} ^ {- 1}, \bar {\lambda} _ {O} C _ {V} ^ {- 1}, 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} C _ {Q} ^ {- 1}, 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \bar {\lambda} _ {Q} \bar {\lambda} _ {V} \bar {\lambda} _ {O} C _ {K} ^ {- 1}\right) \sqrt {2 \ell (\boldsymbol {\theta} ^ {0})}, (2)
$$
$$
\alpha^ {3} \geq 3 2 \rho^ {2} z \sqrt {2 \ell (\boldsymbol {\theta} ^ {0})} / \bar {\lambda} _ {O}. \tag {3}
$$
If the step size satisfies  $\gamma \leq 1 / C$  with a constant  $C$  depending on  $(\bar{\lambda}_Q,\bar{\lambda}_K,\bar{\lambda}_V,\bar{\lambda}_O,\ell (\pmb {\theta}^0),\rho ,\tau_0)$  then  $GD$  converges to a global minimum as follows:
$$
\ell \left(\boldsymbol {\theta} ^ {t}\right) \leq \left(1 - \gamma \frac {\alpha^ {2}}{2}\right) ^ {t} \ell \left(\boldsymbol {\theta} ^ {0}\right), \quad \forall t \geq 0. \tag {4}
$$
Remark: The parameter  $\alpha$  in Eqs. (2) and (3) controls the convergence rate of global convergence, and the condition will be verified in the next subsection. The step-size  $\gamma$  is inversely proportional to  $N^{1/2}$  due to the construction of  $C$  in Appendix C.2.
Proof sketch: The main idea of our convergence analysis is based on the variant of Polyak-Lojasiewicz (PL) inequality [Polyak, 1963, Nguyen, 2021], i.e.,  $||\nabla \ell(\theta)||_2^2 \geq 2\lambda_{\min}(K)\ell(\theta) \geq 2\lambda_{\min}(F_{\mathrm{pre}}F_{\mathrm{pre}}^\top)\ell(\theta)$ . Thus, if the minimum singular value of  $F_{\mathrm{pre}}$  is strictly greater than 0, then minimizing the gradient on the LHS will drive the loss to zero. To this end, Proposition 1 can be split into two parts. First, by induction, at every time step, each parameter in the Transformer can be bounded w.h.p; the minimum singular value of  $F_{\mathrm{pre}}$  is bounded away for some positive quality at the initialization point. Secondly, we prove that the Lipschitzness of the network gradient, which means the loss function is almost smooth. Combining the above two results, the global convergence can be achieved. Furthermore, based on different initialization and scaling schemes, we are able to validate Eqs. (2) and (3) via the spectral norm estimation of  $\bar{\lambda}_Q$ ,  $\bar{\lambda}_K$ ,  $\bar{\lambda}_V$ ,  $\bar{\lambda}_O$  and a positive lower bound for  $F_{\mathrm{pre}}$  in the following section.
4.2 LeCun and He initialization under the  $d_m^{-1/2}$  setting
Here we aim to show that, under the LeCun/He initialization with the setting of  $d_m^{-1/2}$ , the conditions in Eqs. (2) and (3) will be satisfied with high probability and scaling schemes in Table 1 and hence
lead to global convergence. To derive our result, we need the following assumptions on the input data regarding the rank and dissimilarity of data.
Assumption 2. We assume that the input data  $X$  has full row rank.
Remark: This assumption requires that each row  $\mathbf{X}^{(i;)}$  is linearly independent for any  $i \in [d_s]$ , which is fair and attainable in practice. For example, for language tasks, even though there might be some repeated token in  $\mathbf{X}$ , each row in  $\mathbf{X}$  can be uncorrelated when added with positional embedding. Similarly, in image tasks with Visual Transformer, the raw image is grouped by patch and mapped via a linear layer to construct  $\mathbf{X}$ , thus each row in  $\mathbf{X}$  can be uncorrelated.
Assumption 3. For any data pair  $(X_{n}, X_{n'})$ , with  $n \neq n'$  and  $n, n' \in [N]$ , then we assume that  $\mathbb{P}\left(|\langle X_{n}^{\top} X_{n}, X_{n'}^{\top} X_{n'} \rangle| \geq t\right) \leq \exp(-t^{\hat{c}})$  with some constant  $\hat{c} > 0$ .
Remark: We discuss the rationale behind this assumption: