SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We have to pay the extra cost for training for the parameter change. The parameter optimizer manages  $A_{o}, b_{o}$ , but we use their recentered version  $\hat{A}_{o}, \hat{b}_{o}$  during training. The induced cost is small for several reasons. (1) The size of parameters  $A_{o}, b_{o}$  is relatively small since they do not depend on the batch size and sequence length. For reference, the computation cost of LayerNorm in the original Pre-LN Transformer is proportional to the batch size and the sequence length. (2) Obtaining  $\hat{A}_{o}, \hat{b}_{o}$  can be done ahead of time since it does not depend on the input. We may leverage the idle time of accelerators to compute them. (3) In data-parallel distributed training, the parameter server [19] manages the parameters. Each worker will receive  $\hat{A}_{o}, \hat{b}_{o}$  from the server and then pass the gradients of loss to  $\hat{A}_{o}, \hat{b}_{o}$  to the server. Only the server needs to maintain and update the original  $A_{o}, b_{o}$ . In short, it is much easier to process the model parameters than the intermediate activations.
Recentering. It is possible to fuse the recentering with the preprocessing, which usually handles the sum of several kinds of embeddings. For example, the input embeddings of the BERT model [11] are the accumulation of the token embeddings, the segmentation embeddings, and the position embeddings. Since Recenter  $(x + y) = \operatorname{Recenter}(x) + \operatorname{Recenter}(y)$ , recentering the input is equivalent to recentering each embedding before the addition. We can recenter the accumulated embeddings or each embedding separately before the accumulation. Suppose an embedding is from a linear layer. In that case, we can modify the linear layer such that it generates the zero-mean output, similar to how we edit the output linear projection.
For inference, we can recenter the related embeddings or linear layers in advance such that no extra computation is induced. For training, recentering induces extra cost since it is on the fly.
Replacing LayerNorm with RMSNorm. Section 2.1 introduces that RMSNorm can achieve speedup compared with LayerNorm, as demonstrated by the previous models. This replacement can help us accelerate the training and inference of the Pre-LN Transformer.
3.4.2 Pre-CRMSNorm Transformer
CRMSNorm, an extension of RMSNorm, saves execution time as it is more computationally efficient than LayerNorm. Additionally, CRMSNorm further reduces the hidden dimension from  $d$  to  $d - 1$ , which can reduce the model size, computation, communication, and memory consumption by  $1 / d$  in theory. However, most accelerators can not efficiently handle the vectors in  $\mathbb{R}^{d - 1}$  when  $d$  is a large even number, especially a power of two (e.g., 1024, 4096). This limitation arises because these accelerators are typically optimized for arithmetic with even dimensions. In some cases, handling  $\mathbb{R}^{d - 1}$  vectors may take much more time than  $\mathbb{R}^d$  vectors. Thus, we need to examine if the accelerators can support  $\mathbb{R}^{d - 1}$  vectors efficiently. If not, we have the following alternatives. For inference, we may either (1) add zero embeddings, or (2) decompress the vectors and translate the model into the Pre-RMSNorm variant. For training, we suggest keeping the hidden dimension  $d$ , which is equivalent to a Pre-LN model with the hidden dimension  $d + 1$ . In this way, the CRMSNorm can help us increase the model representability and computation efficiency at the same time.
3.5 Training and Inference Equivalence
Taking Pre-LN and Pre-RMSNorm Transformers as examples, we further explain the arithmetic equivalence.
Inference equivalence. Let  $f$  be a Pre-LN Transformer with parameter  $\theta$ , and  $g$  be a Pre-RMSNorm Transformer with parameter  $\phi$ . We demonstrate that for any input  $\pmb{x} \in \mathbb{R}^d$ , we can always have  $f(\pmb{x},\theta) = g(\pmb{x},\phi = h(\theta))$  and we show how we conduct the conversion  $\phi = h(\theta)$ . Thus, we claim that  $f$  and  $g$  are equivalent in terms of arithmetic functionality. Namely, when calculating  $f(\pmb{x},\theta)$  during inference, we can always compute the equivalent counterpart  $g(\pmb{x},\phi = h(\theta))$ . Our
proposed method is a re-parameterization technique [12], which builds equivalent models with different parameters.
Training equivalence. Despite  $f(\pmb{x}, \theta) = g(\pmb{x}, \phi = h(\theta))$ , there may be a large difference in the convergence speed and stability when training  $f$  and  $g$  with SGD and its variants. We avoid this issue by applying the gradient updates on  $\theta$  instead of  $\phi$ . Specifically, during the training process, we maintain a Pre-LN Transformer model  $f$  and its parameter  $\theta$ . In the forward and backward computation, we convert the model and its parameter into  $g$  and  $\phi$  to save training time. We do not update  $\phi$  directly. Instead, we calculate the gradients  $\nabla \theta$  and call the SGD optimizer  $\theta = \theta - \eta \nabla \theta$  where  $\eta$  is the learning rate. From the perspective of optimization, the training is still on the Pre-LN Transformer model  $f$  and its parameter  $\theta$ . Consequently, the training performance and stability are preserved and equivalent to the original Pre-LN Transformers.
We claim that our major contributions are the unification and equivalence of the three Transformer variants. We do not report the task-related performance, such as the classification accuracy and perplexity, since the training and inference equivalence are guaranteed. We discuss how we verify the task-related performance in Appendix D.4.
The efficiency is a free lunch accompanying the equivalence. Now that the efficiency of RMSNorm over LayerNorm has been shown in the previous work [44, 26], we focus on analyzing the efficiency of each component of our method.
We conduct experiments on ViT [13, 36] and GPT-3 [5] since they represent two mainstream architectures of Transformers, encoder-only and casual decoder. Other Transformer variants can be treated as an extension of these two architectures, such as encoder-decoder [39], and non-causal decoder [43]. Also, ViT and GPT cover the areas of computer vision and natural language, where Transformers have been popular and achieved great success.
We abstain from utilizing pre-trained weights. We replicate the ViT and GPT-3 architectures as described in their respective papers, employing dummy data to assess the training and inference speed. We use PyTorch 2.0 [28] to build the training and inference pipeline. We run iterations at least 100 times and report the 25th, 50th (median), and 75th percentile since these quartiles are more robust than the mean. We use the automatic mixed precision [24] for both training and inference.
4.1 Experiments on ViT
The ViT takes images with 3 channels and a resolution of  $224 \times 224$  and generates a classification over 1000 classes, which is the standard setting for ImageNet training and inference.
In the training recipe of ViT [13], dropout [33] is added at the end of each residual branch, which breaks the zero-mean property of the residual output. Hence, in Pre-RMSNorm Transformers, we need to recenter the output vectors explicitly at the end of the residual branch. The Pre-CRMSNorm Transformers are compatible with dropout since it uses vectors in  $\mathbb{R}^{d-1}$  to represent the compressed zero-mean vectors in  $\mathbb{R}^d$ . On the contrary, DeiT [36] disables the dropout and can guarantee the zero-mean output, in spite that the stochastic depth [17] and LayerScale [37] are applied. We follow the settings in DeiT [36] in this paper.
Inference. Figure 3 illustrates the inference time with different batch sizes and model sizes on a single A100 GPU. Taking the Pre-LN Transformer as the baseline, our proposed Pre-RMSNorm Transformer can reduce the inference time by  $1\% - 9\%$ . This reduction takes effect for ViTs with different models and various mini-batch sizes. The left subfigure demonstrates the inference latency, which is the inference time when the batch size is 1.
The proportion of LayerNorm in the total computation is  $12\% - 18\%$  in these experiments. As discussed in Sections 2.1 and 3.4.1, replacing LayerNorm with RMSNorm can help us accelerate the inference. RMSNorm can reduce the inference time of LayerNorm by  $20\% - 60\%$ , thus helping us achieve a stale faster inference for Pre-RMSNorm models.
5Dropout only impacts training and has no impact on inference. Thus, the Pre-RMSNorm variant can also be used for inference.
Figure 3: Normalized inference time on ViT with different model sizes and batch sizes.
(a) Training time breakdown on ViT-Ti/16.
Figure 4: Training time comparison and breakdown on ViT variants
(b) Normalized training time
For Pre-CRMSNorm, the GPU cannot efficiently handle vectors in  $\mathbb{R}^{d - 1}$  in some cases, where we add zero padding to obtain vectors in  $\mathbb{R}^d$ . With zero padding, Pre-CRMSNorm models are less efficient than Pre-RMSNorm ones due to the extra decompression. However, for the cases where  $\mathbb{R}^{d - 1}$  vectors can be accelerated efficiently, we can achieve an even  $10\%$  time reduction.
We also conduct experiments (1) with other precisions, (2) on CPUs, (3) with JAX [4] and observe similar performance. For these ViT variants, we have achieved an average of  $3.0\%$  inference time reduction. Please see Appendix D for more details.
Training. We train ViTs on a single workstation with 4 A100s with data parallel training [20], following the DeiT training recipe. Each training iteration consists of forward and backward computation on all the workers, gradient all-reduce, parameters update with an optimizer, and broadcasting the new parameters. Figure 4a visualizes the breakdown of the related components. In the Pre-LN Transformer, the layer normalization accounts for  $12.6\%$  of total training time. For the Pre-RMSNorm variant, we have to modify the output projection and recenter the input, which induces the  $0.89\%$  and  $0.09\%$  extra computation time. Then LayerNorm is reduced to RMSNorm, whose computation cost is  $9.27\%$ . Overall, the Pre-RMSNorm reduces the training time by  $2.36\%$ .
We train the Pre-CRMSNorm Transformers with  $d$  as the hidden dimension since we do not obtain a speedup from the compressed dimension since the GPUs cannot handle  $\mathbb{R}^{d - 1}$  vectors efficiently. The CRMSNorm is more computationally expensive than the RMSNorm given the same input but takes less time than LayerNorm. The Pre-CRMSNorm Transformer does not need the recentering and special linear layers to generate zero-mean results at the end of residual branches. Above all, training the Pre-CRMSNorm variant is  $1.74\%$  faster than the Pre-LN model.
Figure 4b illustrates the training time of different ViTs. We have achieved a speedup of  $1\% - 2.5\%$ , which is smaller than the inference speedup. Considering only the forward and backward computation, the speedup is similar between training and inference. Nevertheless, the training needs extra time on gradient all-reduce, optimizer update, and parameters broadcast, which shrinks the percentage of
normalizations in the whole computation. For reference, the percentage is  $10\% - 15\%$  for training these ViTs.
4.2 Experiments on GPT
We measure the inference time on the GPT-3 [5] variants with batch size 1 and sequence length 512. The results are shown in Figure 5. In small GPT-3 models, layer normalization takes considerable time. Hence, replacing the LayerNorm with (C)RMSNorm can reduce the inference time by up to  $10\%$ . However, as the model grows, the attention and MLP take charge of the main part of the computation [18]. The normalization takes  $< 1\%$  of the inference time for GPT-3 XL and larger models, which is the upper bound of the speedup with our methods.
Applying quantization may mitigate this issue to some extent. By applying int8 matrix multiplication [10], the percentage of the layer nor
malization increases to  $10\%$  for GPT-3 XL and 2.7B. Our Pre-RMSNorm can reduce the inference time by  $4\%$ . Training performance is similar to the ViT one. We have achieved  $1.5\%$  and  $1.8\%$  time reduction for GPT-3 Small and Medium, respectively.
Figure 5: GPT-3 inference performance
In this paper, we propose two equivalent and efficient variants for the widely used Pre-LN Transformers. We point out the inherent redundancy in the Pre-LN Transformer. By maintaining zero-mean on the main branch vectors and thus simplifying LayerNorm, we obtain the Pre-RMSNorm architecture. We further apply a lossless compression on the zero-mean vectors to obtain the Pre-CRMSNorm model. For the first time, We unify these normalization variants within the Transformer model.
We enable the more efficient utilization of Pre-LN Transformers, allowing for seamless transitions between normalization techniques with minimal overhead. We can replace a Pre-LN Transformer with an equivalent Pre-(C)RMSNorm Transformer with better training and inference efficiency, which is a free lunch. We strictly push the performance-efficiency Pareto frontier of foundational Pre-LN Transformers. As a result, pre-trained Pre-LN Transformers (such as ViT and GPT) can be deployed more efficiently, and new equivalent or superior models can be trained with higher throughput.
Extensions. We believe that our proposed CRMSNorm technique has the potential for application in other neural architectures. Further exploration of its usage in different contexts would be beneficial. Additionally, while our focus has been on pre-normalization Transformers, it would be valuable to investigate the application of our methods to other related architectures, especially the foundation models. Finally, integrating our proposed method into machine learning compilers could directly enable the generation of equivalent and simplified computation graphs.
Limitations. Detailed implementations and specific optimizations will play a crucial role in achieving practical performance gains. It is necessary to focus on developing highly optimized implementations of (C)RMSNorm to bridge the gap between theoretical and practical efficiency improvements. By addressing these challenges, we can fully leverage the potential benefits of (C)RMSNorm and enable more efficient utilization of Pre-LN Transformers.
We acknowledge NVIDIA for donating its A100 GPU workstations and the support from TILOS, an NSF-funded National Artificial Intelligence Research Institute.
[35] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1-28, 2022.
[36] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pages 10347-10357. PMLR, 2021.
[37] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 32-42, 2021.
[38] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Roziere, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[40] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR, 2020.
[41] Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[42] Daochen Zha, Zaid Pervaiz Bhat, Kwei-Herng Lai, Fan Yang, Zhimeng Jiang, Shaochen Zhong, and Xia Hu. Data-centric artificial intelligence: A survey. arXiv preprint arXiv:2303.10158, 2023.
[43] Biao Zhang, Behrooz Ghorbani, Ankur Bapna, Yong Cheng, Xavier Garcia, Jonathan Shen, and Orhan First. Examining scaling and transfer of language model architectures for machine translation. In International Conference on Machine Learning, pages 26176-26192. PMLR, 2022.
[44] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[45] Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al. A survey of large language models. arXiv preprint arXiv:2303.18223, 2023.
A Proof for Lemma 1.
Given a linear transformation  $\pmb{y} = \pmb{A}\pmb{x} + \pmb{b},\pmb{x}\in \mathbb{R}^n,\pmb{A}\in \mathbb{R}^{m\times n},\pmb{b},\pmb{y}\in \mathbb{R}^m$  , we have
$$
\begin{array}{l} \boldsymbol {y} = \boldsymbol {A} \boldsymbol {x} + \boldsymbol {b} (7) \\ = (\boldsymbol {A} - k \mathbf {1 1} ^ {T} \boldsymbol {A}) \boldsymbol {x} + k \mathbf {1 1} ^ {T} \boldsymbol {A} \boldsymbol {x} + (\boldsymbol {b} - \mu (\boldsymbol {b}) \mathbf {1}) + \mu (\boldsymbol {b}) \mathbf {1} (8) \\ = (\boldsymbol {A} - k \mathbf {1 1} ^ {T} \boldsymbol {A}) \boldsymbol {x} + (\boldsymbol {b} - \mu (\boldsymbol {b}) \mathbf {1}) + k \left(\mathbf {1} ^ {T} \boldsymbol {A} \boldsymbol {x}\right) \mathbf {1} + \mu (\boldsymbol {b}) \mathbf {1} (9) \\ = (\boldsymbol {A} - k \mathbf {1 1} ^ {T} \boldsymbol {A}) \boldsymbol {x} + (\boldsymbol {b} - \mu (\boldsymbol {b}) \mathbf {1}) + (k \mathbf {1} ^ {T} \boldsymbol {A} \boldsymbol {x} + \mu (\boldsymbol {b})) \mathbf {1} (10) \\ = \hat {A} \boldsymbol {x} + \hat {\boldsymbol {b}} + f (\boldsymbol {x}, k) \mathbf {1} (11) \\ \end{array}
$$
where  $\hat{\pmb{A}} = \pmb {A} - k\mathbf{1}\pmb{T}\pmb {A},\hat{\pmb{b}} = \pmb {b} - \mu (\pmb {b})\mathbf{1},f(\pmb {x},k) = k\mathbf{1}^{T}\pmb {A}\pmb {x} + \mu (\pmb {b}).$
If  $k = 1 / m$ , then we obtain
$$
\begin{array}{l} \mu (\hat {\boldsymbol {A}} \boldsymbol {x}) = \frac {1}{m} \mathbf {1} ^ {T} (\boldsymbol {A} - \frac {1}{m} \mathbf {1 1} ^ {T} \boldsymbol {A}) \boldsymbol {x} = \frac {1}{m} \left(\mathbf {1} ^ {T} \boldsymbol {A} - \mathbf {1} ^ {T} \boldsymbol {A}\right) \boldsymbol {x} = 0 (12) \\ \mu (\boldsymbol {y}) = \mu (\hat {\boldsymbol {A}} \boldsymbol {x}) + \mu (\hat {\boldsymbol {b}}) + \mu (f (\boldsymbol {x}, k = 1 / m) \mathbf {1}) (13) \\ = 0 + 0 + f (\boldsymbol {x}, k = 1 / m) (14) \\ = \frac {1}{m} \mathbf {1} ^ {T} \mathbf {A} \mathbf {x} + \mu (\mathbf {b}) (15) \\ \end{array}
$$
The  $\hat{A}$  is the recentered matrix of  $A$ , and all its column vectors have zero-mean. We decompose the output into two parts.
- The first part  $\hat{\pmb{A}}\pmb{x} + \hat{\pmb{b}} = \pmb{y} - \mu(\pmb{y})\pmb{1}$ , with zero mean, is another linear transformation with  $\hat{\pmb{A}} = \pmb{A} - \frac{1}{m}\pmb{1}\pmb{1}^T\pmb{A}, \hat{\pmb{b}} = \pmb{b} - \mu(\pmb{b})\pmb{1}$ .
- The second part corresponds to the mean information  $\mu (\pmb {y})\mathbf{1} = (\frac{1}{m}\mathbf{1}^{T}\pmb {A}\pmb {x} + \mu (\pmb {b}))\mathbf{1}$
B Post-LN Transformers
Different from the Pre-LN Transformers, the Post-LN Transformers have the following blocks.
$$
\boldsymbol {x} _ {l + 1} = \mathrm {L N} \left(\boldsymbol {x} _ {l} + \mathcal {F} _ {l} \left(\boldsymbol {x} _ {l}\right)\right), l = 0, 1, \dots , L - 1, \tag {16}
$$
Layer normalization is on the main branch instead of the beginning of residual branches. We can keep a zero-mean branch on the main branch without impacting the functionality.
$$
\begin{array}{l} \boldsymbol {x} _ {l + 1} = \operatorname {L N} \left(\boldsymbol {x} _ {l} + \mathcal {F} _ {l} \left(\boldsymbol {x} _ {l}\right)\right) (17) \\ = \operatorname {L N} \left(\left(\boldsymbol {x} _ {l} - \mu (\boldsymbol {x} _ {l}) \mathbf {1}\right) + \left(\mathcal {F} _ {l} \left(\boldsymbol {x} _ {l}\right) - \mu \left(\mathcal {F} _ {l} \left(\boldsymbol {x} _ {l}\right)\right) \mathbf {1}\right)\right) (18) \\ = \operatorname {L N} \left(\hat {\boldsymbol {x}} _ {l} + \hat {\mathcal {F}} _ {l} (\boldsymbol {x} _ {l})\right) (19) \\ = \operatorname {R M S N o r m} \left(\hat {\boldsymbol {x}} _ {l} + \hat {\mathcal {F}} _ {l} \left(\boldsymbol {x} _ {l}\right)\right) (20) \\ \end{array}
$$
For the residual branch  $\mathcal{F}_l$ , we can apply the same method in the Pre-LN Transformer. We can modify the output linear projection to obtain  $\hat{\mathcal{F}}_l$ , which generates the zero-mean part of the original result.
The recentering operation  $\hat{x}_l = x_l - \mu(x_l)\mathbf{1}$  requires extra computation. If elementwise affine transformation is disabled in LayerNorm,  $x_l$  is the output of a normalization such that  $\mu(x_l) = 0$  and  $\hat{x}_l = x_l$ . If the transformation is enabled,  $x_l$  is not guaranteed zero-mean such that explicit recentering is necessary.
C Revertible Conversions
The conversions of Pre-LN  $\rightarrow$  Pre-RMSNorm and Pre-RMSNorm  $\rightarrow$  Pre-CRMSNorm are listed in Sections 3.1 and 3.2. These steps are fully reversible. We write the inverse steps explicitly below.
Coverting Pre-RMSNorm into Pre-LN Transformers.
1. Remove the recentering.
2.  $\mathbf{A}_o = \hat{\mathbf{A}}_o + \mathbf{1}c^T$ , where  $c$  can be a random vector,  $\mathbf{b}_o = \hat{\mathbf{b}}_o + d\mathbf{1}$  where  $d$  can be a random scalar.
3. Replace RMSNorm with LayerNorm.
Coverting Pre-CRMSNorm into Pre-RMSNorm Transformers.
1. Decompression vectors in  $\mathbb{R}^{d - 1}$  into zero-mean vectors in  $\mathbb{R}^d$ . The compression is lossless, so the decompression is its invertible operation.
2. Replace CRMSNorm with RMSNorm.
3.  $\mathbf{A}_i = \hat{\mathbf{A}}_i + \mathbf{a}_d\mathbf{1}^T$ , where  $\mathbf{a}_d$  is a random vector.
4. Recover the last row of  $\hat{A}_o$  and the last element of  $\hat{b}_o$  such that each column of  $\hat{A}_o$  and  $\hat{b}_o$  is zero-mean.
D.1 Implementation of Normalization
We have provided our implementation with JAX and PyTorch in the supplementary material. The reported results are based on the following implementations.
For JAX, we use the APIs of LayerNorm and RMSNorm in the flax library. For PyTorch, we use the implementations of LayerNorm and RMSNorm from NVIDIA's apex extension  $^{6}$ . For CRMSNorm, we use our own customized implementations. We also provide our customized implementation of LayerNorm and RMSNorm.
We notice that there are lots of APIs for the standard LayerNorm and RMSNorm. For example, PyTorch has provided the official LayerNorm API but lacks RMSNorm implementation. These different implementations are mixed. We do not find one implementation dominant over others for all the cases. For instance, torch(nn.layersNorm is usually faster than apex.normalization.FusedLayerNorm when the input vectors are small in inference, while it is slower than the apex when the input vectors are large. PyTorch's official implementation is also slower than the apex for training.
D.2 Extended Experiments in ViT
Name	Dimension	Depth	Heads	MLP Dimension
Tiny-16	192	12	3	192 × 4
Small-16	384	12	6	384 × 4
Base-16	768	12	12	768 × 4
Large-16	1024	24	16	1024 × 4
Huge-14	1280	32	16	1280 × 4
Giant-14	1664	48	16	8192
Table 2: ViTs with different sizes. The number in the model name is the patch size.
no norm	Pre-LN	Pre-RMS	Pre-CRMS
PyTorch, single A100, amp	0.8567	1.000	0.9699	0.9783
amp → float32	0.9353	1.000	0.9850	0.9951
single A100 → 16-thread CPU	0.8697	1.000	0.9012	0.8857
PyTorch → JAX	0.9610	1.000	0.9873	1.0005
Table 3: Normalized inference time of ViT.
Table 2 lists the architecture parameters of Vision Transformer. We first measure the inference time. We sweep these 6 ViTs with 6 batch sizes (1, 4, 16, 64, 256, 1024) and collect the medians of these
$^{6}$ https://github.com/NVIDIA/apex
Figure 6: The test accuracy of training a ViT-S/16 on ImageNet-1k from scratch.