SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Entropy Collapse for Large Learning rates. To understand the sources of training instability, we keep track of the entropy of the probability distribution induced by the Softmax, as it has been observed that the Transformer's training is unstable in the low-entropy regime [57]. The entropy is calculated for each row of the Softmax matrix, and it is averaged across rows and heads. The results are in Fig. 6. Notice how for the large learning rate regime observed in Fig. ??, the entropy collapses for the baseline model, but not for the proposed shaped Transformer. Entropy collapse indicates that the Softmax distribution degenerates to a point mass, which is itself caused by large logits. Remarkably, this phenomenon does not affect the recover setting, despite recovering the Transformer architecture (without layer normalization) after the warm-up period.
Dataset. We use a subset of the English Wikipedia 20220301.en and English bookcorpus datasets [69, 70]. The sentences are tokenized using by pre-training a tokenizer on the training set. We use a vocabulary size of 32000, and a maximum sequence length of 128 tokens.
Model parameters. We use an embedding size of  $n = 768$  and 8 multi-attention heads. The batch size is fixed to 32 sequences. All the initial weights are sampled from  $\mathcal{N}(0, n^{-1})$ , with the exception of the queries and keys' weights  $W^{K}, W^{Q}$  in the shaped attention case, that are sampled from  $\mathcal{N}(0, n^{-3/2})$  (as explained in Appendix D). The feedforward layer maps the  $n = 768$ -dimensional embedding to the larger dimension 3072, as in the Hugging face implementation of Bert [73].
Figure 6: Dynamics of the mean entropy across heads for selected layers (first and seventh) and learning rates. (Above): shaped attention in recover setting. (Below): Baseline transformers. Notice that at  $lr = 0.001$  the entropy collapses for the baseline model.
**Optimization.** We train using Adam [71] with betas parameters (0.9, 0.999) and learning rate chosen in the grid (0.0001, 0.0005, 0.001, 0.005). We do not use weight decay.
Computational Resources. The experiments are executed on Nvidia DGX-1 GPU nodes equipped with 4 20-core Xeon E5-2698v4 processors, 512 GB of memory and 8 Nvidia V100 GPUs.
E Additional Figures
Figure 7: Kernel density estimate and histogram of covariances from the covariance SDE in Theorem 4.2 and shaped attention NN (Eq. 9). Simulated with  $n = 200$ ,  $d = 150$ ,  $\gamma = 1 / \sqrt{8}$ ,  $\tau_0 = 1$ ,  $V_0^{\alpha \beta} = 0.2$ , SDE step size 0.01, and  $2^{12}$  samples.
Figure 8: Median of the stopping time capped at 1, of the shaped attention neural network with respect to its parameters  $\gamma$  and  $\tau_0$ . Stopping time is defined as  $t^* = d^* / n$  with  $d^*$  the maximum depth beyond which one of the eigenvalues of the covariance matrix exceeds  $10^4$  or drops below  $10^{-4}$ . Simulated with  $n = d = 200$ , and 100 samples used to estimate the median. To demonstrate the potential numerical instabilities, we had to choose an adversarial set of parameters: in particular, an unrealistically large norm (approx.  $10\sqrt{n}$ ) for the initial tokens  $X_0$ , which enlarges the eigenvalues of  $V_0$  to the order of 100.