SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Cenk Baykal*
Dylan Cutler
Nishanth Dikkala
Nikhil Ghosh†
Rina Panigrahy
Xin Wang
It has been well established that increasing scale in deep transformer networks leads to improved quality and performance. However, this increase in scale often comes with prohibitive increases in compute cost and inference latency. We introduce Alternating Updates (AltUp), a simple-to-implement method to increase a model's capacity without the computational burden. AltUp enables the widening of the learned representation, i.e., the token embedding, while only incurring a negligible increase in latency. AltUp achieves this by working on a subblock of the widened representation at each layer and using a predict-and-correct mechanism to update the inactivated blocks. We present extensions of AltUp, such as its applicability to the sequence dimension, and demonstrate how AltUp can be synergistically combined with existing approaches, such as Sparse Mixture-of-Experts models, to obtain efficient models with even higher capacity. Our experiments on benchmark transformer models and language tasks demonstrate the consistent effectiveness of AltUp on a diverse set of scenarios. Notably, on SuperGLUE and SQuAD benchmarks, AltUp enables up to  $87\%$  speedup relative to the dense baselines at the same accuracy.
Contemporary machine learning models have been remarkably successful in many domains, ranging from natural language [6, 20] to computer vision [56, 40]. Many of these successes have come in part through sheer scale. A vast amount of empirical studies justify the conventional wisdom that bigger (models and data sets) is better [19, 24]. Accordingly, state-of-the-art Transformer [48] models often contain billions of parameters and are trained for weeks on enormously large data sets using thousands of AI accelerators. Their immense size leads to prohibitive compute and energy costs [36] and prevents their deployment to resource-constrained applications [31].
To alleviate these costs and enable scalability of modern Transformers, a recent line of works have proposed techniques to increase the capacity of models without drastically increasing the computational costs via conditional computation. A notable paradigm is sparsely-activated networks, such as Mixture-of-Experts (MoE) models [11, 59, 35, 1, 28, 43, 45]. The main idea of MoE is to effectively widen each network layer by accessing dynamically invoked parameters, i.e., experts, where each expert corresponds to a small subset of disjoint parameters that can be acted on by the input. During training and inference, a given input to the network is routed to a small subset of experts (parameters) to compute the output. As a result, the computation cost remains small relative to the total number of parameters. This scheme enables models with higher capacity with only a relatively small increase in computation.
*Correspondence to baykalc@google.com.
$^{\dagger}$ Work done as an intern at Google Research
While prior approaches in conditional computation have predominantly focused on the processing power of transformers, there is a research gap in efficiently incorporating widened learned representations. Recent works have empirically and theoretically established that a wider token representation (i.e., a larger model dimension) helps in learning more complicated functions by enabling more information to be packed in the representation vectors [19, 24, 54]. This phenomenon is also evident in modern architectures of increasing capability. For instance, the representation dimension grows from 512 (small) to 768 (base) and 1024 (large, 3B, and 11B) in T5 models [37], and from 4096 (8B) to 8192 (64B) and 18432 (540B) in PaLM models [6]. A widened representation dimension also significantly improves performance for dual encoder retrieval models [33, 34]. However, naively widening the learned representation requires accordingly increasing the model dimension (see Fig. 1), which quadratically increases the amount of computation in the feedforward computation. In light of the above, a natural question arises: can we leverage the benefit of wider representations without incurring the additional cost of wider transformer layers?
(a) 2x Wider Model
(b) 2x Wider Model with AltUp
Figure 1: An illustration of widening the token representation without (left) and with Alternating Updates (right). This widening causes a near-quadratic increase in computation in a vanilla transformer due to the increased layer width. In contrast, Alternating Updates keeps the layer width constant and efficiently computes the output by operating on a sub-block of the representation at each layer.
In this paper, we address this research gap by introducing Alternating Updates (AltUp), a technique to incorporate wider representations in a simple and efficient way. AltUp operates by partitioning the widened representation vector into blocks, processing only a single block at each layer, and using an efficient prediction mechanism to infer the outputs of the other blocks (see Fig. 1). Processing a single block in each transformer layer enables AltUp to simultaneously keep the model dimension, hence the computation cost, constant and take advantage of using an increased token dimension. Unlike prior approaches, e.g., Sparse Mixture of Experts, AltUp is easy to implement, requires minimal hyperparameter tuning, and does not necessitate sharding. Moreover, since AltUp focuses on increasing the representation dimension, it can be applied synergistically with orthogonal techniques like MoE [60] to obtain complementary performance gains.
In particular, our contributions are:
1. We introduce Alternating Updates (AltUp) to bridge the research gap in efficiency techniques by enabling wider representations with little additional computation cost. AltUp is simple-to-implement, requires minimal hyperparameter tuning, and does not necessitate sharding.
2. We develop two notable extensions of AltUp: (i) Recycled-AltUp, a faster variant of AltUp that requires virtually no additional learnable parameters and (ii) Sequence-AltUp, an extension of the AltUp idea to the sequence dimension.
3. We present an extensive evaluation of AltUp on T5 models on various benchmark language tasks. Our experimental results show that AltUp and its variants uniformly lead to models with improved speed-accuracy trade-offs. Notably, on SuperGLUE and SQuAD benchmarks, AltUp enables up to  $87\%$  speedup relative to the dense baselines at the same accuracy.
Prior work is rich with a diverse set of techniques to increase the efficiency of contemporary transformer models. Here, we cover the most relevant subset of state-of-the-art techniques and refer the interested reader to [47] for a more comprehensive survey.
Recent works have introduced extremely large, yet scalable models with the use of conditional routing of inputs to a learnable subset of parameters. These sparsely-activated models have achieved state-of-the-art performance on various benchmarks [11] and exhibit favorable theoretical properties [4, 1]. Notably, the Sparse Mixture of Experts (SMoE) [45, 57, 22] family of models use a learned softmax probability distribution to conditionally direct the computation to experts, i.e., subsets of network parameters. By routing the computation to a small subset of parameters on an input-dependent basis, SMOE leads to higher capacity models with a relatively small and controllable increase in computation. Switch Transformers [12] show that routing to a single expert on an input-dependent basis reduces computation and outperforms prior SMOE approaches on language tasks. Deja Vu [32] centers around selecting subsets of the attention and MLP parameters to apply for each input (contextual sparsity). Our work is orthogonal to Deja Vu and synergistic with these approaches at large as it focuses on conditional computation.
Follow-up work on SMoE include those that improve the load balancing of experts [60, 29], use reinforcement learning to learn the routing function [7], and leverage smooth top- $k$  expert selection [18] (see [11] for a survey). Other choices for the routing function include non-learnable ones such as Locality Sensitivity Hashing (LSH) [35] which generally maps similar inputs to the same expert, Hash Layers that use token-based hashing [43], and language-specific deterministic routing [10]. Residual Mixture of Experts [51] separates the expert weights into input-independent and input-dependent components.
Conditionally accessing external memory is another related approach to vastly increase model capacity at the cost of a relatively small increase in computation [15, 14]. For examples, Memorizing Transformers [53], Memformer [52], and Product key memory [27] leverage dynamic memory to encode and retrieve relevant information. Additional works include those that use an immensely large untrainable corpus, such as Wikipedia, REALM [17], or a 2 trillion token database, RETRO [2]. These prior works that focus on routing(expert)-based mechanisms often necessitate complicated, sharded implementations due to the sheer number of additional parameters that they introduce — often on the order of billions. Our work, on the other hand, is simple-to-implement and requires virtually no hyperparameter tuning. Moreover, AltUp can be synergistically combined with sparsely-activated models like MoE to obtain complementary improvements in efficiency.
Additional relevant works in the realm of efficient transformers include Funnel transformers [8], Reformers [25], Performers [5], Big-Bird [58], and LongT5 [16], among others. These works notably present methods to reduce the quadratic cost of the attention mechanism of transformers. Another flavor of methods complementary to our work is that of adaptive computation, e.g., CALM [44], DynaBERT [21] CascadeBERT [30] and DeeCap [13], where different amounts of computational power is allotted on an example-specific basis via some type of early-exit strategy. AltUp achieves its efficiency via the orthogonal direction of conditionally leveraging wider token representations, and hence can be easily combined with these efficiency techniques and others (e.g., quantization [55], FlashAttention [9]).
3 Alternating Updates
In this section, we introduce the method of Alternating Updates (AltUp), an approach to enable increased token dimension with little additional computation cost.
3.1 Background
At a high level, a standard transformer with  $L$  layers generates a  $d$ -dimensional representation by applying a sequence of layers transformer layers  $\mathcal{L}_1, \ldots, \mathcal{L}_L$  as follows. For a particular input token within a sequence of length  $N$ , the initial token representation  $x_1 \in \mathbb{R}^d$  is computed by an embedding table lookup. Subsequently, this  $d$ -dimensional representation is refined across the transformer layers by iteratively computing  $x_{i+1} = \mathcal{L}_i(x_i)$  for each layer  $i \in [L]$ ; here and throughout  $[N]$  denotes the set  $\{1, \ldots, N\}$  for  $N \in \mathbb{N}$ . Each transformer layer  $\mathcal{L}_i: \mathbb{R}^{d_{\mathrm{model}}} \to \mathbb{R}^{d_{\mathrm{model}}}$  has width  $d_{\mathrm{model}}$  (with  $d_{\mathrm{model}} = d$  in the standard setting) and contains an attention block and a FeedForward (FFN) block. The width of the layer  $d_{\mathrm{model}}$  controls the dimensions of the matrices involved in the attention and FFN blocks. Consequently, the computation cost of attention and FFN scales with  $\mathcal{O}(N^2 d_{\mathrm{model}})$  and  $\mathcal{O}(Nd_{\mathrm{model}}^2)$ , respectively. The output,  $x_{L+1}$  is the output token representation generated by the transformer. This computation is usually followed by a linear layer operation that maps from the  $d$ -dimensional representation  $x_{L+1}$  to  $|\mathcal{V}|$ -dimensional logits (in  $\mathcal{O}(|\mathcal{V}|d)$  time), followed by a softmax non-linearity to generate the probabilities over the vocabulary  $\mathcal{V}$ .
Increasing the representation dimension  $d$  is a way to enhance the capacity of the transformer model, as a wider representation enables the transformer to store richer information about the input and helps in learning more complicated functions [19, 24, 54]. Naively widening the token representation  $d$  requires widening each layer as well, since  $d_{\mathrm{model}}$  must match  $d$  in a standard transformer model. However, the computation time of each transformer layer grows roughly quadratically with  $d_{\mathrm{model}}$ , notably for relatively short input sequences. This means that, growing the token dimension from  $d$  to  $2d$ , for example, leads to a model that is at least 2 times (and closer to 4 times for small  $N$ ) slower than the original model with a  $d$ -dimensional representation.
3.2 Alternating Updates
The core idea of Alternating Updates is to widen the representation vector, but perform computation with a  $d$ -dimensional sub-block, and estimate the updated representation using a Predict-Compute-Correct algorithm, as illustrated in Figure 1, right. More specifically, AltUp expands the representation width from  $d$  to  $Kd$ , for integers  $K > 1$ ,  $d > 0$  (for example,  $K = 2$  in Fig. 1), but uses layers of width  $d_{\mathrm{model}} = d$  (not  $d_{\mathrm{model}} = Kd$ ) to transform the representation vector. By keeping the width of each transformer layer constant, AltUp avoids incurring the quadratic increase in computation cost that would otherwise be present with a naive expansion of the representation.
Alg. 1 depicts the details of the per-layer computation involved in a transformer with AltUp with a  $K$  dimensional representation vector. The input to the AltUp layer is assumed to be the concatenation of  $d$ -dimensional contiguous subblocks  $x_{\mathrm{old}} = \mathrm{concat}(x_{old}^{1}, x_{old}^{2}, \dots, x_{old}^{K}) \in \mathbb{R}^{dK}$ . Inspired by predictor-corrector methods used to solve ordinary differential equations [3], AltUp first generates a prediction  $\hat{x}^i$  for each of the subblocks  $i \in [K]$  (Line 1). This prediction takes the form of a mixture of subblocks  $\hat{x}^i = \sum_{j=1}^{K} p_{i,j} x_{old}^j$ , where  $p_{i,j} \in \mathbb{R}$  for  $i, j \in [K]$  are learnable scalars. Subsequently, one of the  $K$  sub-blocks is chosen and the computation with the unexpanded transformer layer of width  $d_{\mathrm{model}} = d$  is performed on this sub-block (Line 2). Finally, the result of this computation is used in the correction step to generate the updated representation for each sub-block (Line 3).
Computation time We see from Alg. 1 that AltUp introduces negligible amount of additional computation per layer, as the prediction and correction steps involve only vector addition and scalar-vector multiplications  $(\mathcal{O}(d)$  operations). Thus, relative to the computation cost of a transformer layer with width  $d$  (which we incur on Line 2 in AltUp), the cost of AltUp is only an additional  $\mathcal{O}(dK^2)$  per token, where  $d$  is the original model dimension and  $K$  is the factor of increase in the representation dimension (typically  $K = 2$  or  $K = 4$ , see Sec. 5). This additional  $\mathcal{O}(dK^2)$  cost per token is a factor of  $d$  smaller than the  $\mathcal{O}(d^2 K^2)$  per token cost of the FFN block alone in a  $K$ -times wider transformer layer. In fact, an AltUp layer does not lead to an increased computation time relative to a  $d$ -width transformer layer asymptotically, since the  $\mathcal{O}(dK^2)$  additional cost per token per layer is dominated by the cost of the FFN block as  $K \ll d$  in practice. At a higher level, AltUp requires using an embedding table with width  $Kd$  and invoking the final linear operation with  $Kd$ -dimensional vectors. The initial embedding lookup using a wider table and the linear + softmax operation with  $Kd$  (instead of  $d$ ) dimensional vectors may lead to a perceptible increase in computation time. However, since we only incur this additional cost in the beginning and the end, these factors are often inconsequential,
Algorithm 1 Alternating Updates (AltUp) Layer
Input:  $x_{old} = \mathrm{concat}(x_{old}^{1}, x_{old}^{2}, \ldots, x_{old}^{K}) \in \mathbb{R}^{dK}$ :  $dK$ -dimensional input representation vector to the layer, where  $x_{old}^{j} \in \mathbb{R}^{d}, j = 1, 2, \ldots, K$  are contiguous sub-blocks of  $x_{old}$ .
Output:  $x_{new} \in \mathbb{R}^{dK}$ : The layer's  $dK$ -dimensional output representation.
1: Predict: for each  $i \in [K]$ , predict the updated representation with a trainable linear map:
$$
\hat {x} ^ {i} = \sum_ {j = 1} ^ {K} p _ {i, j} x _ {o l d} ^ {j},
$$
where  $p_{i,j} \in \mathbb{R}, i, j \in [K]$  are trainable scalars.
2: Compute: select a sub-block  $j^{*} \in [K]$  and update this block with  $\mathcal{L}$ :
$$
\tilde {x} ^ {j ^ {*}} = \mathcal {L} (x _ {o l d} ^ {j ^ {*}}).
$$
3: Correct: for each  $i \in [K]$ , correct the prediction with the computation result:
$$
x _ {n e w} ^ {i} = \hat {x} ^ {i} + g _ {i} (\tilde {x} ^ {j ^ {*}} - \hat {x} ^ {j ^ {*}}),
$$
where  $g_{i}\in \mathbb{R},i\in [K]$  are trainable scalars.
and increasingly so for deeper transformers. Nevertheless, we present an extension to AltUp in Sec. 4 that avoids this slowdown altogether for specialized applications.
Parameter count AltUp introduces  $K^2 + K$  additional learnable parameters per layer, where  $K^2$  is due to  $p_{i,j}, i, j \in [K]$  and  $K$  is a result of  $g_i, i \in [K]$ . Since  $K \ll d$ , this is an imperceptible amount of additional parameters per layer in practice. Zooming out, AltUp with an expansion factor of  $K$  requires a  $Kd$ -width embedding table, and consequently requires  $(K - 1)|\mathcal{V}|d$  additional parameters, where  $|\mathcal{V}|$  is the vocabulary size. In Sec. 4, we present a variant that requires no additional embedding parameters to be added to the model.
Memory footprint Due to the increase in the representation dimension, AltUp incurs a slightly higher activation memory footprint compared to the baseline model during training. In particular, a transformer model with  $L$  layers, model dimension  $h$ , batch size  $b$ , sequence length  $s$ , and number of attention heads  $a$  has an activation memory footprint of  $sbhL(34 + 5as / h)$  [26]. Adding AltUp (with  $K = 2$ ) to this model leads to an additional activation memory of  $(sbh + 2sbh)L = 3sbhL$ , which is less than  $10\%$  of the vanilla transformer's. Moreover, a significant portion of memory is used by the weight parameters which only increase marginally with AltUp. Model parallelism and techniques such as gradient checkpointing, sequence parallelism, or selective activation recomputation can mitigate the memory impact of a wider activation vector even further. During inference, the memory footprint of activations is dominated by the Key-Value cache, which is unaffected by the additional activations introduced by AltUp.
Selection of sub-blocks The selection of the sub-block  $j^*$  for the computation step in Algorithm 1 can be any user-specified technique. We consider two simple, deterministic selection methods in this paper and leave more sophisticated methods for future work: (i) same: choose the same sub-block for all the layers in a neural network and (ii) alternating (default method): for a sequence of layers, alternating through the sub-blocks, that is, if the sub-blocks are indexed with zero-based index, then sub-block  $\ell \mod K$  is selected for the computation step for layer  $\ell \in [L]$ . This alternating selection is the default for Algorithm 1 (hence the name Alternating Updates). We compare the two selection methods empirically in the supplementary material and find that using alternating blocks performs better empirically.
4 AltUp Extensions
In this section, we present extensions of the core AltUp idea introduced in the previous section.
4.1 Recycled-AltUp: Faster AltUp via embedding recycling
The AltUp formulation presented in Sec. 3 adds an insignificant amount of per-layer computation, however, it does require using a  $K$ -times wider embedding table. In certain scenarios where the vocabulary  $\mathcal{V}$  is very large, this may lead to a non-trivial amount of added computation for the initial embedding lookup and the final linear + softmax operation. A colossal vocabulary may also lead to an undesirable amount of added embedding parameters. Recycled-AltUp is an extension of AltUp that avoids these computational and parameter costs by keeping the embedding table's width  $d$ -dimensional.
Figure 2 depicts an example application of Recycled AltUp with  $K = 2$ . The general idea is to recycle the initial  $d$ -dimensional lookup by replicating the  $d$ -dimensional lookup  $K$  times. Hence, Recycled-AltUp virtually adds no addi
Figure 2: Recycled-AltUp with  $K = 2$ .
tional parameters relative to the baseline width  $d$  model. Subsequently, the regular AltUp layers (Alg. 1) are applied until the last linear  $^+$  softmax operation. To avoid the computational cost of this final operation, Recycled AltUp downprojects the  $Kd$ -dimensional representation vector  $x_{L + 1}\in \mathbb{R}^{dK}$  to a  $d$ -dimensional representation by simply elementwise-adding the  $d$ -dimensional contiguous subblocks in  $\mathcal{O}(Kd)$  time. Applying the linear  $^+$  softmax operation on this down-projected vector implies that the computation cost of this operation is now  $\mathcal{O}(|\mathcal{V}|d)$  rather than  $\mathcal{O}(K|\mathcal{V}|d)$ , effectively reducing the amount of computation by  $\mathcal{O}((K - 1)|\mathcal{V}|d)$ . Our results in Sec. 5 show that Recycle-AltUp's improved speed and reduced parameter count may make it more appealing for certain applications.
4.2 Sequence-AltUp: Extension of AltUp to the sequence dimension
Here, we introduce Sequence-AltUp, a natural extension of Alternating Updates to reduce the apparent sequence length. This extension is motivated by the computation cost associated with the cost of the attention mechanism for long input sequence lengths. Our approach is similar in its goal to that of prior techniques focused on designing efficient attention mechanisms to reduce the quadratic dependency of attention cost on the sequence length: Funnel transformers [8], Reformers [25], Performers [5], Big-Bird [58], and LongT5 [16].
Similar to the Funnel transformer [8], Sequence-AltUp uses a simple striding operation to reduce the sequence length. Only sampled tokens are processed by the transformer layer while the rest of the tokens require little computation, leading to a computation cost reduction by a factor of  $k$ , where  $k$  is the stride parameter. In this way, Sequence-AltUp is similar to AltUp in that it applies the predict-compute-correct algorithm (Algorithm 1) to the sequence dimension, but it is different from AltUp in that it does not increase the sequence dimension.
Figure 3 depicts the baseline stride-and-skip technique (left) and the proposed Sequence-AltUp method (right). Given an input sequence of vectors  $(x_0, x_1, \dots, x_{T-1})$  to a transformer layer  $\mathcal{L}$ , we propose the following extension of Algorithm 1 to reduce the effective sequence length by a factor of  $k$ . First, we apply a lightweight predictor on the full sequence to obtain a predicted output  $\hat{y} = (\hat{y}_0, \dots, \hat{y}_{T-1})$ . Next, we subsample the input with a fixed stride  $k$  and apply  $\mathcal{L}$  on the subsampled input and get computed output  $\tilde{y} = (\tilde{y}_0, \tilde{y}_k, \tilde{y}_{2k}, \dots, \tilde{y}_{\lfloor(T-1)/k\rfloor *k}) = \mathcal{L}(x_0, x_k, \dots, x_{\lfloor(T-1)/k\rfloor *k})$ . Finally, we use a lightweight corrector to combine  $\hat{y}$  and  $\tilde{y}$  to form the final output sequence. This design allows unsampled token vectors to obtain contextual information, even though they are not processed by the transformer layer directly—analogous to the inactivated sub-blocks in original AltUp. In contrast, a simple stride-and-skip approach (Figure 3, left) lacks the ability to bring contextual information to the skipped tokens. We present the full algorithm pseudocode and implementation details of Sequence-AltUp in the supplementary material.
(a) Stride-and-skip