Reply to Reviewer cxoC's questions
We thank the reviewer for the feedback and raise in score!
**Comment**: In Remark 1, it seems that the only assumption for data and initialization is that $B_0$ is full rank. Then, by over-parameterization, i.e. $D=\Omega(Nn)$, the loss landscape is nearly convex around initialization. This sounds very similar to NTK or a lazy learning regime. I wonder whether there are any feature learning or rich learning regime effects in your analysis.
**Response**: Yes, our optimization analysis is similar to the NTK regime with lazy update. Our current theory does not consider feature learning or the rich learning effect. However, we believe it is possible to extend our analysis towards the feature learning regime. For example, the high-dimensional embedding can be further modeled by a function that extracts data features, e.g, networks, allowing the attention head to be interpreted as computing the "correlation" of extracted features. We plan to investigate the feature learning regime in our future work.
**Comment**: Is there any idea about the generalization with some assumptions on data distribution?
**Response**: Yes, generalization analysis can be derived from our theory with certain assumptions on data distribution and model weights, although it is not analyzed in this paper as we primarily focus on the optimization landscape. As a concrete example, it is possible to extend our analysis by following the framework in [1], where the data is balanced and a binary classification problem is considered. We can further derive a similar generalization bound as [1] Theorem 1.
**Comment:** The formula of Gaussian kernel.
**Response**: Respectfully, we would like to clarify that the formula of Gaussian kernel in your comment is missing an additional activation function. With different activations included, the Gaussian kernels (defined in [2, 3] and equation (11)) and Softmax kernels differ significantly, as the Softmax activation considers the entire row of the attention matrix, whereas the Gaussian does not.
Specifically, the $k$-th row and $j$-th column of Gaussian attention is given by (see equation (11)):
$$S\left(W\_h^Q, W_h^K ; X_i\right)\_{k j}=\operatorname{\exp}\left(-\frac{1}{\sqrt{d}}\left(X\_{ik\cdot} W_h^Q-X\_{ij\cdot} W\_h^K\right)^2\right)$$.
The Softmax attention is given by:
$$S\_{ih}:=\text{Softmax}\left(\frac{X\_{i} W_h^Q\left(X\_{i} W\_h^K\right)^{\top}}{\sqrt{d}} \in \mathbb{R}^{n\times n}\right).$$
Notice that the entry in Gaussian kernel is only related to the $k$-th and $j$-th token in sample $X_i$. In contrast, in Softmax attention, each entry is also related to the entries in the same row, since the Softmax activation computes the regularization based on a row. This leads to a more complex optimization landscape and makes the convergence more challenging. These two different kernels also result in different performance on some tasks [2, 3].
[1] Li, Hongkang, et al. "A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity." arXiv preprint arXiv:2302.06015 (2023).
[2] Lu, Jiachen, et al. "Soft: Softmax-free transformer with linear complexity." Advances in Neural Information Processing Systems 34 (2021): 21297-21309.
[3] Chen, Yifan, et al. "Skyformer: Remodel self-attention with gaussian kernel and nystr" om method." Advances in Neural Information Processing Systems 34 (2021): 2122-2135.