SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(b) Sequence-AltUp
Figure 3: An illustration of Sequence-AltUp (right) and the baseline Stride-and-skip method (left). Sequence-AltUp has virtually the same computation cost as Stride-and-skip, but enables contextual information passing to the skipped tokens.
In this section, we apply AltUp and its variants to benchmark language models and tasks. We proceed by outlining the experimental setting below. In Secs. 5.1 and 5.2 we present the performance of AltUp on standard benchmarks with varying configurations and model sizes; in Secs. 5.3 and 5.4 we evaluate the performance of AltUp extensions and demonstrate their effectiveness. We present the full details of our evaluations and additional experimental results in the supplementary; namely, the supplementary contains additional evaluations that demonstrate the synergistic combination of AltUp with other conditional compute techniques, additional finetune results, and complementary ablation studies. Overall, our results consistently show that AltUp and its variants enable sizeable performance gains, e.g., up to  $87\%$  faster models, across all evaluations on standard benchmarks.
Setting We performed all of our experiments using T5-model architectures [37] of varying sizes (small, base, large, and 3B) which we pretrained on the C4 dataset for 500,000 steps with a batch size of 256. The pretrained models were then finetuned on either the GLUE [50], SuperGLUE (SG) [49], SQuAD [39] or Trivia-QA (closed-book) [23, 42] benchmark tasks for a further 50,000 steps with a batch-size of 256. The pretraining task is to predict corrupted text spans, and the finetuning tasks are re-cast into text generation tasks. We report both pretraining and finetuning metrics: for pretraining, we report span prediction accuracy on a hold-out validation set, and for finetuning, we follow the same recipe as the T5 models, see [37] for more details. The supplementary contains the full details of our evaluations and hyperparameters.
5.1 Alternating updates on benchmarks
First, we investigate whether incorporating AltUp on a baseline model leads to an unambiguously better model when we consider the predictive performance and actual observed latency (not theoretical FLOPS). To this end, we compare the dense T5-Base/Large/XL models to models augmented with AltUp with  $K = 2$  on GLUE, SuperGLUE, SQuAD, and TriviaQA (closed-book) finetuning tasks. Figure 4 plots the performance and normalized speed of the evaluated models.
As the figure depicts, models augmented with AltUp are uniformly faster than the extrapolated dense models at the same accuracy. For example, we observe that a T5 large model augmented with AltUp leads to a  $27\%$ ,  $39\%$ ,  $87\%$ , and  $29\%$  speedup on GLUE, SuperGLUE, SQuAD, and Trivia-QA benchmarks, respectively. Moreover, we see that AltUp's relative performance improves as we apply it to larger models (compare relative speedup of T5 Base + AltUp to that of T5 Large + AltUp). This demonstrates the scalability of AltUp to and its improved performance on even larger models. Overall, AltUp consistently leads to models with better predictive performance than the corresponding baseline models with the same speed on all model sizes and benchmarks.
(a) GLUE
(c) SQuAD
(b) SuperGLUE
(d) Trivia-QA
Figure 4: Evaluations of AltUp on T5 models of various sizes and popular benchmarks. AltUp consistently leads to sizeable speedups relative to baselines at the same accuracy. Latency is measured on TPUv3 with 8 cores. Relative speedup is defined as latency delta divided by AltUp latency.
5.2 AltUp with varying representation size
In the previous subsection, we had used a value of  $K = 2$  for the runs with AltUp. As discussed in Sec. 3,  $K$  controls the width of the representation vector, and is the only hyperparameter required by AltUp. Can we obtain even more performant models by using a larger expansion factor  $K$ ? Here, we compare the performance of AltUp with  $K = 2$  to AltUp with a larger expansion factor  $K = 4$ .
Table 1: Performance of AltUp with varying representation dimension scaling parameter  $K$  on T5.
Model	Pretrain Accuracy	GLUE	SG	SQuAD (EM/F1)	TriviaQA (EM/F1)
S	61.21	75.83	59.52	76.44/84.97	19.03/22.83
S + AltUp (K=2)	61.86	76.82	59.60	77.51/85.79	19.27/22.95
S + AltUp (K=4)	62.00	76.40	59.54	76.38/84.86	19.07/22.84
B	66.42	84.25	73.56	83.78/91.19	23.1/27.56
B + AltUp (K=2)	66.96	85.32	75.80	85.24/92.36	24.35/28.78
B + AltUp (K=4)	67.18	84.95	78.91	84.82/92.07	24.41/28.90
L	69.13	87.23	81.21	86.77/93.56	26.15/30.76
L + AltUp (K=2)	69.32	88.20	82.75	87.81/94.29	27.10/32.04
L + AltUp (K=4)	69.55	88.42	82.94	87.59/94.02	27.36/32.42
Table 1 summarizes the results with AltUp instantiated on T5 small, base, and large sized models with hyperparameter  $K = 2$  and  $K = 4$ . We observe that a larger value of  $K = 4$  leads to strict improvements in pretrain accuracy over AltUp with  $K = 2$  for all models (Table 1, column 2). This is perhaps intuitive, as a wider representation vector enables more information to be learned during the pretraining stage. Interestingly, however, a larger  $K$  does not always lead to better finetune performance, especially for smaller models. For example, despite having a worse pretrain accuracy, AltUp with  $K = 2$  is better than AltUp with  $K = 4$  on all finetune tasks GLUE, SuperGLUE, and SQuAD. We see a similar phenomenon occur for the Base model, but here  $K = 4$  is better on GLUE;
and on Large, the trend reverses:  $K = 4$  is better on every metric except for SQuAD. Our results indicate that a larger value of  $K$  has potential to increase the performance of models on pretrain and fine-tune metrics when AltUp is applied to larger models. We note that there is an inherent trade-off between a larger factor  $K$  and trainability, however, as a larger value of  $K$  leads to less frequent activation of each sub-block which may impair performance. We envision that practitioners can pick a value of  $K$  other than the default  $K = 2$  to optimize performance on an application-specific basis.
5.3 Recycled-AltUp
Next, we consider the performance of the lightweight extension of AltUp, Recycled-AltUp, introduced in Sec. 4. We apply Recycled-AltUp with  $K = 2$  to T5 base, large, and XL models and compare its pretrain accuracy and speed to those of baselines. We record both the training speed and inference speed of the resulting models. Since Recycled-AltUp does not require an expansion in the embedding table dimension (see Sec. 4), we remark that the models augmented with it have virtually the same number of trainable parameters as the baseline models.
Figure 5: Recycled-AltUp on T5-B/L/XL compared to baselines. Recycled-AltUp leads to strict improvements in pretrain performance without incurring any perceptible slowdown.
The results of our experiments are shown in Fig. 5. The figures for both the training and inference speed show that models with Recycled-AltUp clearly improve over baselines in pretrain accuracy, without any perceptible slowdown. While Recycled-AltUp's predictive strength generally falls below that of standard AltUp (cf., pretrain values for AltUp in Table 1), its improved speed and reduced parameter count may make it more suitable for certain applications - see Recycled-AltUp Fine-tune evaluations in the Appendix (Sec. G). We present additional fine-tuning results with Recycled-AltUp in the supplementary material; overall, our results demonstrate that Recycled-AltUp is similarly effective on fine-tuning tasks.
5.4 Sequence-AltUp
Here, we evaluate Sequence-AltUp (from Sec. 4) to reduce the apparent sequence length for the T5 base model. In particular, we apply average pooling, stride-and-skip, and Sequence-AltUp to the encoder layers to reduce the apparent input sequence length. We apply stride-and-skip and Sequence-AltUp to layers  $2, \ldots, L - 1$  of the encoder, rather than all the layers, with stride length 4 as we found that this configuration results in a better accuracy/speed trade-off for both techniques. For average pooling, the sequence length is immutably reduced from the start according to the method.
Table 2: Performance and pretrain speed of different methods for sequence length reduction on T5.
Model	Pretrain Accuracy	Finetune GLUE	Finetune SG	Speed
S	61.21	59.52	76.44/84.97	166.1
B (Baseline)	66.42	73.56	83.78/91.19	52.4
Average pooling	63.89	57.85	71.37/81.87	91.9
Stride-and-Skip	65.02	65.98	79.72/87.64	79.4
Sequence-AltUp	65.39	66.94	81.67/89.37	74.9
Table 2 presents the comparisons on pretrain and finetune metrics (GLUE and SuperGLUE) and pretrain speed (measured by the number of sequences per second per core). The table additionally lists the relevant metrics for T5 Base (which is the baseline model) and T5 Small as reference points in the table. We observe that average pooling gives a large speed-up, but suffers from severe quality degradation, especially on the finetune metrics where it performs even worse than T5 small. Stride-and-skip and Sequence-AltUp, on the other hand, offer an improved quality and speed trade-off relative to T5 Base. In particular, Sequence-AltUp is only slightly slower than stride-and-skip (yet, still  $\approx$ $40\%$  faster than the baseline), but is much closer to the baseline model's quality.
We propose the method of Alternating Updates (AltUp) to increase the capacity of modern transformer models without incurring a significant increase in latency. Our approach bridges the research gap in efficient transformers by enabling the use of wider token representations without widening the transformer layers. AltUp utilizes lightweight prediction and correction steps to update a wider representation vector without increasing the transformer layer's computation cost. As a result, we achieve strong performance improvements on language modeling and language understanding benchmarks. We present extensions of AltUp that enable additional gains in efficiency. Given its orthogonal scope, AltUp can be synergistically applied with existing techniques like MoE. On popular language understanding and QA benchmarks, AltUp enables up to  $87\%$  speedup relative to the dense baselines at the same accuracy.
Limitations and future work A current limitation of the technique we propose is the lack of a deep theoretical understanding of its properties due to the complicated nature of rigorously analyzing transformer models. An interesting open question is whether it would be possible to analyze AltUp by relating its performance to a block compressed layer, and transitively relating that to a wide layer without block compression. A deeper understanding of AltUp may also shed light on the optimal hyperparameter  $K$  on an application-specific basis. In future work, we plan to conduct a theoretical investigation of alternating updates to develop a deeper understanding of its effectiveness across differing applications. We also plan to experiment with the use of a very large expansion factor  $K$ .
Broader Impact Training and deploying modern neural network models consumes colossal amounts of resources. This leads to detrimental effects on the environment and hampers the widespread applicability and democratization of AI. We envision that AltUp can serve as a valuable component of efficient architectures of the future and help alleviate these negative impacts.
We would like to thank Ravi Kumar, Erik Vee, and Nikhil Vyas for constructive discussions and their feedback.
[1] Cenk Baykal, Nishanth Dikkala, Rina Panigrahy, Cyrus Rashtchian, and Xin Wang. A theoretical view on sparsely activated networks. arXiv preprint arXiv:2208.04461, 2022. 1, 3
[2] Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. In International Conference on Machine Learning, pages 2206-2240. PMLR, 2022. 3
[3] John Charles Butcher. Numerical methods for ordinary differential equations. John Wiley & Sons, 2016. 4
[4] Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding mixture of experts in deep learning. arXiv preprint arXiv:2208.02813, 2022. 3
[5] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020. 3, 6
[54] Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.14522, 2020. 2, 4
[55] Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. Advances in Neural Information Processing Systems, 35:27168-27183, 2022. 3
[56] Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui Wu. Coca: Contrastive captioners are image-text foundation models. arXiv preprint arXiv:2205.01917, 2022. 1
[57] Seniha Esen Yuksel, Joseph N Wilson, and Paul D Gader. Twenty years of mixture of experts. IEEE transactions on neural networks and learning systems, 23(8):1177-1193, 2012. 3
[58] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontonon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17297, 2020. 3, 6
[59] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai, Zhifeng Chen, Quoc Le, and James Laudon. Mixture-of-experts with expert choice routing. arXiv preprint arXiv:2202.09368, 2022. 1
[60] Barret Zoph, Irwan Bello, S ameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. Designing effective sparse expert models. arXiv preprint arXiv:2202.08906, 2022. 2, 3, 17
Supplementary Material for Alternating Updates for Efficient Transformers
In this supplementary, we present the full details and hyperparameters of our evaluations, provide details of our algorithmic contributions, and present complementary empirical results that support the effectiveness the presented work.
We experiment with various T5 baseline model sizes: small (S), base (B), large (L), and XL. The base (B), large (L), and XL models follow the same model configurations as in the T5 paper, while the small model is shallower than the T5 paper [37] to cover a larger range of model sizes (4 encoder/decoder layers instead of 8 encoder/decoder layers). In particular, we use the T5 version 1.1 models with gated GELU feedforward network and pre layer norm. The models are implemented on top of the T5X [41] codebase. During pretraining, we use 256 batch size, Adafactor optimizer [46] with base learning rate 1.0 and reciprocal square-root decay with 10000 warmup steps, and zero dropout. During finetuning, we use 256 batch size, Adafactor optimizer with constant learning rate of 0.001 and 0.1 dropout. Unless explicitly stated, we pretrain for  $500k$  steps and finetune for  $50k$  steps. Our experiments were implemented in Python and run on TPUv3 with 8 cores.
B Parameter Counts and Speed
Here, we present the number of additional parameters needed by adding AltUp, its speed, and its pretrain accuracy on T5 models of varying sizes. In the following tables, the embedding parameters include input embedding table parameters (shared between encoder and decoder) and output embedding table. Non-embedding parameters include all the transformer blocks. Train speed is measured by number of examples per second per core. Table 3 documents the parameter count and training speed comparison. Note that Alternating Updates increases the number of embedding parameters while leaving the non-embedding parameters roughly the same. Since the narrow transformer layer computation is not changed by alternating updates and since the predict and correct steps are lightweight (see Sec. 3), we incur a relatively small increase in the computation cost compared to a dense 2x width model.
Model	# emb params	# non-emb params	train speed
S	3.29E+07	3.78E+07	166.1
S + AltUp	6.58E+07	3.99E+07	119.4
B	4.93E+07	1.98E+08	52.4
B + AltUp	9.87E+07	2.12E+08	42.3
L	6.58E+07	7.17E+08	17.1
L + AltUp	1.32E+08	7.68E+08	14.4
Table 3: Model size and train speed comparisons on T5X models with AltUp instantiated with  $K = 2$
Table 4 documents the parameter count, training speed and pretrain accuracy comparison when the representation dimension is scaled up with AltUp or dense scaling. Note that Alternating Updates increases the number of embedding parameters while leaving the non-embedding parameters roughly the same, providing an efficient way to scale up the representation dimension relative to a  $K$ -times wider model.
Table 5 contains pretrain performances for T5 XL sized models. We note the AltUp technique continue to offer quality boost at the billion parameters scale (note that T5XL has roughly 3B parameters), suggesting that AltUp is a robust technique for increasing model capacity for modern large language models.
C Combination with MoE
Here, we investigate whether AltUp can be combined with orthogonal techniques, namely MoE, to obtain additive performance gains in pretrain accuracy for T5 small, base, and large models. In
Model	# emb params	# non-emb params	Train speed	Pretrain accuracy
T5 Base	4.93E+07	1.98E+08	52.4	65.29
T5 Base + AltUp2x	9.87E+07	2.12E+08	42.3	65.78
T5 Base + Dense2X	9.87E+07	3.97E+08	32.9	66.45
T5 Base + AltUp4x	1.97E+08	2.41E+08	28.1	66.00
T5 Base + Dense4X	1.97E+08	7.93E+08	12.6	67.01
Table 4: AltUp compared with dense scaling evaluated at  ${250k}$  pretrain steps.
Model	# emb params	# non-emb params	Train speed	Pretrain accuracy
T5 XL	1.32E+08	2.72E+09	3.6	70.01
T5 XL + AltUp2x	2.63E+08	2.92E+09	3.0	70.61
Table 5: Pretrain performances for T5 XL sized models. Pretrain accuracy is measured at  $400k$  steps. AltUp continues to offer a performance boost even on the scale of models with billions of parameters.
particular, we consider the partial experts setting similar to [38, 35], where at each layer, in addition to the layer's module, we route the input to a smaller expert module and combine the outputs of the main and auxiliary modules as the input to the subsequent layer (see Fig. 6).
(a) Mixture of Experts
(b) Mixture of Partial Experts
Figure 6: The partial experts setting in the context of the evaluations in Sec. C. The standard MoE model (left) routes the inputs to one or more of  $n$  experts based on a routing function. Mixture of Partial Experts (right) always routes the input to the main expert and additionally routes the input to one or more partial experts; the output is a function of the main expert's and partial experts' outputs.
The MoE layer routes an input token  $x$  to  $k$  of  $n$  experts where each expert is itself a parametrized subnetwork (e.g., a fully-connected layer). Following [11], we let  $\{E_i(\cdot)\}_{i\in [n]}$  and  $E_{i}(x)$  denote the set of experts and the output of lookup the input token  $x$  to expert  $i$ , respectively. For an input token  $x$ , a learnable weight matrix  $W$  is applied to obtain the logits  $h(x) = Wx$ . The lookup probabilities are computed by taking the softmax of  $h(x)$
$$
p _ {i} (x) = \frac {\exp \left(h _ {i} (x)\right)}{\sum_ {j \in [ n ]} \exp \left(h _ {j} (x)\right)} \quad \forall i \in [ n ].
$$
The token  $x$  is routed to the expert(s)  $\mathcal{T} \subset [n]$  with the top-  $k$  probabilities  $p(x)$ . Since this operation is not differentiable, the output  $y$  is computed as a probability weighted combination of the experts' outputs to enable gradients to propagate back to the router parameters [45], i.e.,
$$
y = \sum_ {i \in \mathcal {T}} p _ {i} (x) E _ {i} (x).
$$
For MoE, we used the simplified implementation of the top-1 softmax routing of [12]. We use 128 experts each per encoder and decoder layer, with each expert representing a 2 layer fully-connected neural network with hidden dimension 16. For sake of the synergistic demonstration even with the core MoE implementation, we did not incorporate a sophisticated mechanism for load balancing such as load balancing loss [12] or router z loss [60]. We use multiplicative jitter noise sampled from a uniform distribution over  $[1 - \varepsilon, 1 + \varepsilon]^{d_{\mathrm{in}}}$  with  $\varepsilon = 0.01$ . The router matrix  $W$  was initialized by drawing from a zero mean Normal distribution with standard deviation  $2 \times 10^{-2}$ .
Method	T5 Small	T5 Base	T5 Large
Baseline	59.10	63.35	65.58
MoE [60]	59.42	63.62	65.71
AltUp (K=2)	59.67	63.97	65.73
AltUp (K=2) + MoE	59.91	64.13	65.95
Table 6 synthesizes the pretraining performance on the C4 dataset at 100k training steps of AltUp and compared techniques on T5 Small, Base, and Large models. Perhaps most notably, we show that combining AltUp and MoE leads to even further sizable improvements in the pretrain performance (last row of Table 6). For example, the combination of MoE and AltUp improves over the baseline by 0.81, over AltUp alone by 0.24, and over MoE alone by 0.49. For all model sizes, the combination of AltUp and MoE is synergistic and leads to significant improvements compared to not only the baseline, but also to each approach in isolation.
D Alternating updates with varying block selection
In this section, we present empirical results on the alternating updates technique and comparison with other techniques that widen token representation vectors. We increase the token representation dimension by a factor of 2 (corresponding to  $K = 2$  in Algorithm 1) unless otherwise specified. The model capacity increase comes from a wider embedding table at the bottom layer of the model while the transformer layers remain the same, which results in minimal additional computation cost.
Table 6: Pretrain accuracy at 100k steps of T5 models augmented with Alternating Updates (see Sec. 3) and MoE. MoE synergistically combines with Alternating Updates and enables further increases in model capacity.
Model	Pretrain Accuracy	Finetune GLUE	Finetune SG	Finetune SQuAD (EM/F1)
S (baseline)	61.21	75.83	59.52	76.44/84.97
S + Sum	61.67	77.54	59.63	75.06/83.82
S + SameUp	61.91	77.75	60.81	76.85/85.51
S + AltUp	61.86	76.82	59.60	77.51/85.79
B (baseline)	66.42	84.25	73.56	83.78/91.19
B + Sum	66.82	84.85	75.2	84.36/91.36
B + SameUp	66.82	84.06	74.15	84.41/91.76
B + AltUp	66.96	85.32	75.80	85.24/92.36
L (baseline)	69.13	87.23	81.21	86.77/93.56
L + Sum	69.09	86.18	78.93	86.19/93.08
L + SameUp	69.45	87.95	82.72	87.65/94.13
L + AltUp	69.32	88.20	82.75	87.58/94.27
Table 7: Comparison of Algorithm 1 with various sub-block selection methods on T5-S/B/L.
In Table 7, we compare the summation method (Sum) in which additional embedding vectors are added to the token representation vector, Algorithm 1 with same block selection (SameUp), and Algorithm 1 with alternating block selection (AltUp), all on top of the T5 version 1.1 small (S), base (B), and large (L) models. We observe the Prediction-Compute-Correct scheme as described in Sec. 3 with same and alternating block selection methods outperforms the summation method. For the small models, same block selection method performs better in most tasks, while for the base and
large models, alternating block selection method performs better in most tasks. We note all three methods bring improvements in both pretraining and finetuning, and AltUp is generally the most effective one. While pretraining accuracies for all three methods are mostly similar, differences in finetuning metrics are large, and AltUp generally achieves roughly twice the gains of the other two variants. Moreover, we observe that the gains of AltUp in pretraining accuracies show diminishing returns when model sizes grows, but the gains in finetuning metrics do not.
E Additional Evaluations
In addition to the T5-based evaluations presented in the main body of the paper, we conducted a preliminary study on a lightweight BERT model. The model has 12 layers, 256 model dimension, and 4 attention heads. On a masked language pretraining task (with BERT pretraining data), lightweight-BERT achieves 54.7 MLM accuracy while lightweight-BERT + AltUp ( $K = 2$ ) achieves 56.2 MLM accuracy.