SlowGuess/NeurIPS_2023_OCR huggingface dataset:
B.5 The impact of clipped softmax hyperparameters  $(\gamma$  and  $\zeta)$  on ViT
γ	ζ	FP32 acc.	Max inf norm	W8A8 acc.
0 (= Vanilla)	1	78.80±0.42	426±69	71.27±0.88
0	1.001	78.78±0.29	411±88	71.24±0.59
0	1.002	78.90±0.17	420±47	70.74±0.34
0	1.004	78.80±0.45	377±67	72.31±0.06
0	1.01	78.81±0.30	419±77	71.35±0.26
-0.00001	1	78.81±0.21	432±76	69.02±0.19
-0.0001	1	78.81±0.36	380±64	64.04±10.8
-0.001	1	78.42±0.63	282±105	68.43±6.50
-0.003	1	78.26±0.06	99±36	76.49±0.48
-0.01	1	78.10±0.14	391±21	75.83±1.12
-0.03	1	70.26±1.46	197±2	65.80±1.41
-0.001	1.001	78.45±0.53	283±82	65.03±8.54
-0.003	1.003	78.25±0.14	119±17	76.37±0.45
Table 8: The impact of clipped softmax hyperparameters on ViT-S/16.
We investigate the effect of different values of the clipped softmax stretch parameters applied to the vision transformer and present the results in Table 8. To speed up training, for this experiment we trained ViT for 150 epochs instead of the usual 300 epochs. For this experiment, we did not apply LayerNorm after the patch embeddings.
We found similar observations compared to BERT. Specifically, most of the improvement happens when we use  $\gamma < 0$  (clipping at zero) whereas using  $\zeta > 1$  (clipping at one) yields similar results to the vanilla softmax and combining both  $\gamma < 0$  and  $\zeta > 1$  yields similar results compared to just clipping at zero.
B.6 Fine-tuning experiment
Method	FP16 ppl.↓	Max inf norm	Avg. Kurtosis
Vanilla fine-tuning	29.46	79.3	2086
Fine-tuning w/ Gated attention	29.18	50.9	665
Table 9: OPT-1.3B fine-tuning results with vanilla softmax and gated attention. We report the causal language modeling perplexity (ppl. for short) on the English Wikipedia validation set. We also report the maximum  $\| \mathbf{x}\|_{\infty}$  averaged across the validation set, and kurtosis of  $\mathbf{x}$  averaged across all layers, where  $\mathbf{x}$  is the output of an attention layer.
One of the drawbacks of our proposed framework is that it requires training from scratch, which could be expensive when applied to very large models. To address this, we explored whether fine-tuning using gated attention can still lead to improved performance and decreased outliers for larger models.
We used OPT-1.3B pre-trained checkpoint from HuggingFace and fine-tuned it on Bookcorpus + Wikipedia for 4000 steps with batch size 256, maximum sequence length 512, maximum learning rate  $10^{-5}$ , and linear LR schedule with 400 warmup steps. We use the same LR for both model parameters and gating module parameters. The rest of hyper-parameters are the same as for our pre-training setup. We adapted our gating approach as follows. We initialized bias as  $b_{\mathrm{init}} = 0$ , which corresponds to the expected initial gating probability output of  $\pi_{\mathrm{init}} = 0.5$ . We multiply the gating probability by 2 so that the expected gate output is 1 and we approximate the attention output of
the vanilla softmax at the start of fine-tuning. We add a small activation regularization term at the output of each FFN to further encourage the reduction in the magnitude of activations, as unlike when training from scratch outliers are already present in the pre-trained model and need to be suppressed.
As we can see from Table 9, fine-tuning with our proposed gated attention results in a better perplexity and also reduced maximum infinity norm and the average kurtosis compared to fine-tuning with vanilla softmax.
B.7 Low-bit quantization results
Bitwidths	Weight range estimation	Vanilla	Clipped softmax	Gated attention
FP16	-	4.49±0.01	4.39±0.00	4.45±0.03
W8A8	min-max	1294±1046	4.52±0.01	4.65±0.04
W6A8	min-max	598±254	4.64±0.01	4.79±0.03
W6A8	MSE	6.49±0.38	4.56±0.01	4.71±0.03
W4A8	MSE	6.52±0.02	4.90±0.02	5.02±0.03
W6A6	MSE	42.8±11.7	6.64±0.14	5.90±0.11
Table 10: A summary of results for our proposed methods applied to BERT-base and quantized to different bitwidths for weights and activations (using the same PTQ setup as in all previous experiments). We report the masked language modeling perplexity on the English Wikipedia validation set.
Note that our proposed methods are not limited to 8-bit quantization only and in general can be combined with other more advanced quantization and weight compression methods, including [18, 35, 36, 45, 63, 67]. In Table 10, we show the results of our proposed methods applied to BERT-base and quantized to different bitwidths using our simple post-training quantization setup. Unless stated otherwise, for low-bit (<8-bit) weights and activations we use MSE range estimator as recommended by [2, 7] since it gives better results.
As we can see, in all cases both of our methods significantly improve the perplexity compared to the vanilla softmax pre-training. We also notice that generally the performance progressively degrades as we decrease the bitwidths, which is to be expected. Achieving good results with low-bit activation quantization in general is a challenging problem. Further, we notice that the perplexity of the vanilla model significantly improves whenever we consider a low-bit weight quantization with MSE ranges compared to the INT8 case. This can be explained by the fact that using MSE range estimation for weights leads to an implicit clipping of activations (in the same and all subsequent layers in the network), which happen to be of the right amount so that it doesn't hurt the perplexity. We found that by going from W8A8 to W6A8 the average kurtosis is reduced from  $3406^{\pm 547}$  to  $631^{\pm 94}$  and the maximum infinity norm is reduced from  $577^{\pm 80}$  to  $158^{\pm 40}$ . However, in all cases the resulting model still has significantly larger outliers and a worse performance than both of our proposed methods. Finally, as said before, if achieving good low-bit quantization performance is the goal, it is recommended to combine our methods with more advanced quantization techniques.
C.1 BERT
Fine-tuning on MNLI dataset We use pre-trained checkpoint BERT-base-uncased (109M parameters) from HuggingFace repository. We follow standard fine-tuning practices from [14] and [65]. Each data sequence is tokenized and truncated to the maximum sequence length of 128. Shorter sequences are padded to the same length of 128 using a special [PAD] token. We fine-tune for 3 epochs using Adam [29] with a batch size of 16 and no weight decay. The learning rate is initially set to its maximum value of  $2 \cdot 10^{-5}$  and is linearly decayed to zero by the end of fine-tuning.
Pre-training from scratch We follow closely the pre-training procedure from [14]. We concatenate, tokenize, and split the training set into sequences of length 128 (to speed up training and experimentation, we do not fine-tune on longer sequences of 512). We use the masked language modeling objective with the probability of masking  $p = 0.15$ . We train with a batch size of 256
sequences for  $10^{6}$  steps, using AdamW optimizer [39] with the maximum learning rate of  $10^{-4}$ , learning rate warm up over the first  $10^{4}$  steps, following by a linear decay to zero by the end of training. We use L2 weight decay of 0.01, L2 gradient norm clipping of 1.0, and dropout probability of 0.1 on all layers. We also use FP16 mixed-precision from HuggingFace Accelerate library [20].
C.2 OPT pre-training
To speed up experimentation, we train OPT-125m sized model on the concatenation of Wikipedia and BookCorpus (same as BERT pre-training). We train with a batch size of 48 and 4 gradient accumulation steps (which results in the effective batch size of 192), so that we can perform pretraining on a single A100 80GB GPU. We concatenate, tokenize, and split the training set into sequences of length 512 and train for 125000 steps (500000 forward passes).
We use the rest of the hyper-parameters and follow pre-training practices from [74] and [65]. We initialize weights using a normal distribution with zero mean and a standard deviation of 0.006. All bias terms are initialized to zero. We use AdamW optimizer with  $(\beta_{1},\beta_{2}) = (0.9,0.95)$ . We use the linear learning rate schedule, warming up from 0 to the maximum value of  $4\cdot 10^{-4}$  over the first 2000 steps, following by a linear decay to zero by the end of training. We use L2 weight decay of 0.1, L2 gradient norm clipping of 1.0, and dropout probability of 0.1 on all layers. We also use FP16 mixed-precision from HuggingFace Accelerate library [20].
C.3 ViT pre-training
We use the model definition for ViT-S/16 and the training pipeline from PyTorch Image models library [64]. All training is done on resolution  $224 \times 224$  and  $16 \times 16$  patches. For data augmentation, we use RandAugment [10], Mixup [73], CutMix [70], random image cropping [56], horizontal flip, label smoothing  $\varepsilon = 0.1$ , color jitter 0.4, and random (between bilinear and bicubic) interpolation during training.
We train with a batch size of 512 for 300 epochs, using AdamW optimizer and the L2 weight decay of 0.03. We use the cosine learning rate schedule, warming up from  $10^{-6}$  to the maximum value of  $10^{-3}$  over the first 20 epochs, followed by a LR decay by a factor of 10 every 30 epochs, until it reaches the minimum value of  $10^{-5}$ .
C.4 Quantization settings
Weights In all cases, we use symmetric uniform quantization of weights. We use min-max weight quantization for all models except the OPT model, for which we found the MSE estimator to perform better in all cases.
Activations We adopt static range estimation approach, which determines quantization parameters for the network by passing a few batches of calibration data through the model before inference. Specifically, we use a running min-max estimator [32], which uses an exponential moving average of the min and max over multiple batches. In all cases, we use running min-max with 0.9 momentum over 16 batches randomly sampled from respective training sets.
For OPT model, we also experiment with using  $99.99\%$  and  $99.999\%$  percentiles instead of actual min and max. We select the best configuration for each experiment (including baseline), based on the model performance. In almost all cases, we found that setting activation quantization ranges using  $99.999\%$  percentiles gives the lowest W8A8 perplexity.
D Compute cost
We compare the runtime of our proposed methods in Table 11. As we can see, the clipped softmax is only marginally more expensive compared to using the vanilla softmax attention. The gated attention using the linear  $G$  adds the compute overhead between  $3\%$  and  $8\%$ , depending on the model. We found that adding weight decay on LayerNorm  $\gamma$  for OPT and adding the LayerNorm after the patch embeddings for ViT had a negligible effect on the runtime.
In our experiments, we found this value to perform better compared to the value of  $6 \cdot 10^{-4}$  listed in the paper.
Model	Vanilla	Clipped softmax	Gated attention (Linear / MLP)
BERT	92.8±1.2	93.6±0.8	97.7 / 119.1
OPT	53.6±0.4	54.4±0.4	55.7 / 64.7
ViT	101.8±0.3	104.0±0.7	110.8 / 122.9
Table 11: An overview of the runtime of the proposed methods, compared to the vanilla pre-training, measured in hours on Nvidia-A100 GPUs.
We estimated that the compute cost of producing the main results in the paper is about 320 GPU days (on A100) and the total cost of the project (including preliminary experiments and ablation studies) to be about 1400 GPU days.
(a) Attention layer #10, data sequence #16
(b) Attention layer #11, data sequence #16
(c) Attention layer #10, data sequence #21
(d) Attention layer #11, data sequence #21
(e) Attention layer #10, data sequence #61
(f) Attention layer #11, data sequence #61
(g) Attention layer #10, data sequence #88
Figure 10: Visualization of the self-attention patterns (attention probabilities, values, and their product in left, middle and right columns, respectively) in attention head #3 ( $\leftrightarrow$  channel dim #180) for BERT-base trained with vanilla softmax, computed on several random data sequences from MNLI-m validation set.
(h) Attention layer #11, data sequence #88
Figure 11: Visualization of the self-attention patterns (attention probabilities, values, and their product in left, middle and right columns, respectively) in attention head #12 ( $\leftrightarrow$  channel dim #720) for BERT-base trained with vanilla softmax, computed on several random data sequences from MNLI-m validation set.
(a) Attention layer #10, attention head #1
(b) Attention layer #11, attention head #1
(c) Attention layer #10, attention head #7
(d) Attention layer #11, attention head #7
(e) Attention layer #10, attention head #8
(f) Attention layer #11, attention head #8
(g) Attention layer #10, attention head #10
Figure 12: Visualization of the self-attention patterns (attention probabilities, values, and their product in left, middle and right columns, respectively) in attention heads that are not associated with the strong outliers for BERT-base trained with vanilla softmax, computed on data sequences #16 from MNLI-m validation set.
(h) Attention layer #11, attention head #10
(a) Attention layer #1
(b) Attention layer #2
(c) Attention layer #3
(d) Attention layer #4
(e) Attention layer #5
(f) Attention layer #6
(g) Attention layer #7
Figure 13: Visualization of the self-attention patterns (attention probabilities, values, and their product in left, middle and right columns, respectively) in attention head #3 ( $\leftrightarrow$  channel dim #180) and the first eight layers of BERT-base trained with vanilla softmax, computed on data sequences #16 from MNLI-m validation set.
(h) Attention layer #8
(a) Attention layer #10, Attention head #3, data sequence #1
(b) Attention layer #11, Attention head #3, data sequence #1
(c) Attention layer #10, Attention head #3, data sequence #5
(d) Attention layer #11, Attention head #3, data sequence #5
(e) Attention layer #10, Attention head #3, data sequence #7
(f) Attention layer #11, Attention head #3, data sequence #7
(g) Attention layer #10, Attention head #12, data sequence #1
(h) Attention layer #11, Attention head #12, data sequence #1
Figure 14: Visualization of the self-attention patterns (attention probabilities, values, and their product in left, middle and right columns, respectively) for BERT-base trained with Clipped softmax, computed on several random data sequences from MNLI-m validation set.
Figure 15: Visualization of the self-attention patterns (from left to right: gating probabilities  $\pi =$  sigmoid  $(G(\mathbf{x}))$ , output of softmax, values, and their combined product) for BERT-base trained with gated attention, computed on several random data sequences from MNLI-m validation set.
(d)
Figure 16: A summary of our outlier analysis for ViT demonstrated on a random subset from ImageNet validation set. (a) An input image. (b) Outliers in the output of layer #10. (c) Cumulative attention weight spent on every patch (matrix of attention probabilities summed over rows) in the attention head #1, in the next layer #11. (d) A corresponding matrix of attention probabilities. (e) An average magnitude of values  $(\bar{V})$  for outlier and non-outlier patches.
(d)
Figure 17: A summary of our outlier analysis for ViT demonstrated on a random subset from ImageNet validation set. (a) An input image. (b) Outliers in the output of layer #11. (c) Cumulative attention weight spent on every patch (matrix of attention probabilities summed over rows) in the attention head #1, in the next layer #12. (d) A corresponding matrix of attention probabilities. (e) An average magnitude of values  $(V)$  for outlier and non-outlier patches.