SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(b) He initialization,  $\gamma = 1$
(d) NTK initialization,  $\gamma = 10$
Figure 5: Comparison of convergence curve for different initialization schemes. The convergence speed under Lecun/He initialization is generally faster than that of NTK initialization with the same step size.
E Additional details on experiments
E.1 Additional result in Section 5.1
In this section, we present the convergence curve for different initialization schemes following the setup in Section 5.1. The results in Figures 5a to 5c show that when using the same step size  $\gamma = 1$  for LeCun and He initialization exhibits similar behavior while NTK initialization leads to slower convergence. This is consistent with our theoretical finding. We also depict the result with larger step size  $\gamma = 10$  for NTK initialization in Figure 5d, following our analysis in Appendix C.5.
E.2 Additional result and set-up in Section 5.2
We have validated Assumption 3 on language data in the main paper. Here we will validate it on image dataset. Specifically, we choose MNIST dataset and consider the embedding in the ViT as  $\mathbf{X}$ , then we plot  $\mathbb{P}\left(|\langle \mathbf{X}_n^\top \mathbf{X}_n,\mathbf{X}_{n'}^\top \mathbf{X}_{n'}\rangle |\geq t\right)$  as  $t$  increases in Figure 6, where we could see exponential decay. The architecture of ViT is the same as in Section 5.2.
Additionally, in Figure 7, we can see that  $d_m^{-1/2}$  setting achieves faster convergence in training loss and higher test accuracy. In Figure 8, we conduct the experiment on MNIST with regression task by using the same architecture. We take two classes  $\{0,1\}$  in MNIST for binary classification, the output of our regression is taken by  $\mathrm{sign}(f)$ , and the training error is given by  $\sum_{n=1}^{N}[\mathrm{sign}(f(\mathbf{X}_n)) - y_n]^2$ . We can see that the convergence speed of  $d_m^{-1/2}$  is faster than that of  $d_m^{-1}$ . Overall, these separation results on both classification and regression tasks provide good justification for the  $d_m^{-1/2}$  scaling.
Figure 6: Verification of Assumption 3 on MNIST dataset.
Figure 7: Test accuracy of classification task on MNIST dataset.
Figure 8: Regression task (with MSE loss) on MNIST dataset.
(c) Kernel distance.
Figure 9: Experimental validation of the theoretical results on Transformers with  $\tau_0 = d_m^{-1}$  scaling trained on synthetic data. (a) Linear convergence. (b) Rate of change of the weights during training. Observe that the weights change very slowly after the  $50^{\mathrm{th}}$  epoch. (c) Evolution of the NTK during the training. The result mirrors the plot (b) and demonstrates how the kernel varies significantly at the beginning of the training and remains approximately constant later. As the width increases, the NTK becomes more stable.
Firstly, in this work, we prove the global convergence guarantee for Transformer under different initialization schemes. Nevertheless, we only consider the gradient descent training under squared loss.
Our result does not cover SGD training and other loss functions, e.g., hinge loss and cross-entropy loss. Secondly, the model architecture being analyzed in this work is the encoder of Transformer which is widely used for regression problems or image classification problems. We do not consider the entire Transformer including the decoder, which is designed for sequence-to-sequence modeling. Lastly, our model does not cover the deep Transformer. We believe our analytic framework paves a way for further analyzing large-scale Transformer.
G Societal impact
This work offers a convergence analysis for the Transformer, which is a core element within contemporary large language models. Our theoretical analysis with realistic initializations and scaling schemes lays a theoretical foundation for the interested practitioner in the ML community to further study other priorities of Transformer such as generalization and in-context learning. We do not expect any negative societal bias from this work.