SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sehoon Kim $^{1}$  Karttikeya Mangalam $^{1}$  Suhong Moon $^{1}$   
Jitendra Malik $^{1}$  Michael W. Mahoney $^{123}$  Amir Gholami $^{12}$  Kurt Keutzer $^{1}$
<sup>1</sup>University of California, Berkeley <sup>2</sup>ICSI <sup>3</sup>BNLN {sehoonkim, mangalam, suhong.moon, malik, mahoneymw, amirgh, keutzer}@berkeley.edu
The recent emergence of Large Language Models based on the Transformer architecture has enabled dramatic advancements in the field of Natural Language Processing. However, these models have long inference latency, which limits their deployment and makes them prohibitively expensive for various real-time applications. The inference latency is further exacerbated by autoregressive generative tasks, as models need to run iteratively to generate tokens sequentially without leveraging token-level parallelization. To address this, we propose Big Little Decoder (BiLD), a framework that can improve inference efficiency and latency for a wide range of text generation applications. The BiLD framework contains two models with different sizes that collaboratively generate text. The small model runs autoregressively to generate text with a low inference cost, and the large model is only invoked occasionally to refine the small model's inaccurate predictions in a non-autoregressive manner. To coordinate the small and large models, BiLD introduces two simple yet effective policies: (1) the fallback policy that determines when to hand control over to the large model; and (2) the rollback policy that determines when the large model needs to correct the small model's inaccurate predictions. To evaluate our framework across different tasks and models, we apply BiLD to various text generation scenarios encompassing machine translation on IWSLT 2017 De-En and WMT 2014 De-En, and summarization on XSUM and CNN/DailyMail. On an NVIDIA T4 GPU, our framework achieves a speedup of up to  $2.12 \times$  speedup with minimal generation quality degradation. Furthermore, our framework is fully plug-and-play and can be applied without any modifications in the training process or model architecture. Our code is open-sourced<sup>1</sup>.
In recent years, the Transformer [63] has become the de-facto model architecture for a wide range of Natural Language Processing tasks. The potential of the Transformer architecture has been further enhanced by the emergence of Large Language Models (LLMs) with up to hundreds of billions of parameters trained on massive text corpora [2, 47, 50, 10, 23, 7, 55, 83, 62]. Despite their performance, efficiently running these models for inference is a challenge due to their large model size and runtime complexity. This limits their use in many applications that require real-time responses.
These computational inefficiencies are particularly pronounced in autoregressive generative tasks such as machine translation [3, 1], summarization [21], and language modeling [41]. For these tasks, models need to run iteratively to generate tokens sequentially, as each token is dependent on the previously generated tokens. This requires the models to load weight matrices, as well as the cached keys and values of previously generated tokens [46], for each token generation, thus preventing
$^{1}$ https://github.com/kssteven418/BigLittleDecoder
Figure 1: Illustration of (Left) the normal autoregressive decoding procedure of a large model and (Right) BiLD that consists of a small model and a large model. In BiLD, the small model generates tokens autoregressively (i.e., sequentially) until it hands over control to the large model. The large model then takes as input the tokens generated by the small model in parallel, allowing for non-autoregressive (i.e., parallel) execution to generate the next token. This improves end-to-end latency by allowing for more efficient utilization of underlying hardware.
parallelization of the loaded values across multiple tokens. This makes autoregressive text generation memory bandwidth constrained during inference [8]. As a consequence, autoregressive generative tasks suffer from low hardware utilization as well as high inference latency [32]. In contrast, non-autoregressive tasks, such as text classification [65], can process the entire input sequence with a single weight load, which is then shared across all input tokens in parallel. Given the increasing popularity of text generation tasks, in light of advancements in LLMs, it is critical to improve the inference latency and runtime efficiency of autoregressive decoding processes despite the potential sacrifice in generation quality.
To overcome this, non-autoregressive decoding [16, 69, 39, 59, 70, 52, 37, 15, 19] has been explored to maximize token-level parallelization and reduce the inference latency of generative tasks by generating multiple tokens simultaneously. This approach can be more computationally efficient than the regular autoregressive process. However, non-autoregressive decoding suffers from text generation quality issues due to its assumption of conditional independence between output tokens [29]. In order to achieve comparable performance to that of autoregressive processes, it generally requires complex and often task-dependent training strategies, supplementary hinting information that guides the decoding process [69, 39, 59, 70, 52], and knowledge distillation [84].
In this paper, we introduce a novel framework named Big Little Decoder (BiLD) that can be applied to various text generation scenarios to reduce inference latency without additional training iterations or modifications to the existing training pipeline or model architecture. As illustrated in Figure 1 (Right), the BiLD framework consists of two decoder models, a large model and small model, that work collaboratively to generate text sequences. In particular, only the small model is executed autoregressively to generate the majority of the text, taking advantage of its small runtime overhead. The large model only engages occasionally to refine the small model's inaccurate predictions, thus allowing for efficient non-autoregressive execution. This autoregressive small, non-autoregressive large scheme results in a substantial improvement of up to  $\sim 2\times$  in end-to-end inference latency, compared to regular autoregressive execution, while maintaining similar or better generation quality. The effectiveness of our framework is also supported by our observation that the predictions made by small and large models only slightly disagree, and thus the small model can match the performance of the large model with a minimal refinement of its own predictions (Figure 2, Section 3.1).
- We introduce BiLD, a general framework that allows faster inference of various text generation applications. Our framework is designed to coordinate a large model and a small model such that the large model is only executed infrequently and efficiently in a non-autoregressive manner to refine the small model's inaccurate predictions.
- We propose two policies for BiLD: the fallback policy that allows the small model to hand control over to the large model if it is not confident enough (Section 3.3), and the rollback policy that allows the large model to review and correct the small model's inaccurate predictions (Section 3.4).
- We introduce a simple yet effective technique for aligning the predictions of the small model with those of the large model. By incorporating this prediction alignment technique into the BiLD framework, we can further enhance its performance with minimal additional effort (Section 3.5.1).
- We apply BiLD for 4 different text generation scenarios including IWSLT 2017 De-En [3] and WMT 2014 De-En [1] for machine translation, XSUM [43] and CNN/DailyMail [21] for summarization. Compared to the full autoregressive execution, BiLD achieved a speedup of up to  $1.85 \times$  without generation quality degradation and  $2.12 \times$  allowing  $\sim 1$  point degradation on an NVIDIA T4 GPU (Section 4.2).
2.1 Efficient Transformer Decoding Inference
A variety of approaches have been proposed to increase the speed and reduce the overall inference costs of Transformers. Well-known approaches include efficient architecture design [26, 33, 36, 60, 67, 75], quantization [31, 54, 81, 82, 79, 74, 9], pruning [14, 49, 34, 42, 64, 12, 35], and neural architecture search [5, 56, 57, 66, 76, 80]. While these methods are generally suitable for Transformer-based tasks, some of the works have been focused on efficient decoding mechanisms to reduce the cost of autoregressive tasks.
One popular line of research that shares similarity to our work is non-autoregressive decoding. Non-autoregressive decoding, also known as parallel decoding, was first introduced in [16] as a method to reduce inference latency by producing multiple output tokens in parallel, thus avoiding sequential text generation. Subsequent work has further improved the performance of non-autoregressive models by incorporating auxiliary or hinting information [69, 39, 59, 70, 52] to ensure more accurate parallel decoding, or by allowing multiple additional iterations to refine any inaccurate predictions [37, 15, 19]. Such a multi-iteration decoding scheme has also been proposed in [71, 17, 58], which generates texts with fewer steps than the autoregressive scheme by inserting or deleting multiple tokens per iteration. However, these works require complex and often task-dependent training strategies and/or auxiliary information to achieve comparable performance to that of autoregressive models. In contrast, our methodology aims for a plug-and-play solution that does not require any complex training pipeline.
Our work is also related to the approaches that reduce the decoding cost by making decoders shallow. [29] demonstrates that increasing the depth of encoders and decreasing the depth of decoders can reduce decoding latency while still preserving performance. CALM [51] recently introduces early exiting, which dynamically adjusts the depth of the decoder for each token generation by terminating the inference at a middle layer, rather than executing until the end layer. While our method shares the same goal of accelerating decoding, we take a different approach by improving decoding parallelism rather than by skipping unnecessary computation. In addition, our framework offers several advantages over CALM: (1) our method is a fully black box approach that does not involve any modifications to model structures, while CALM requires modifications such as state propagation for the skipped layers; (2) our approach does not require changes to the training pipeline, whereas CALM requires averaged loss across all layers to ensure layer consistency; (3) our approach can be also applied without any training which is critical in various LLM use cases where training is either infeasible or prohibitively expensive. In Section 4.4, we also show that the early exiting strategy can be implemented in our framework to yield significantly better generation quality, further demonstrating the generalizability of our method to a wider range of problems.
2.2 Use of Multiple Models
Coordinating the use of multiple models has also been explored in knowledge distillation and ensemble learning. Knowledge distillation is a widely adopted methodology for enhancing the performance of smaller models by training them to replicate the behavior of larger, more complex models [22]. When applied to the Transformer architecture, this approach involves distilling the final logits [48, 61] and/or hidden states of a larger model, such as the attention map [60, 28, 68]. In contrast to knowledge distillation, which leverages the knowledge of a large model solely during the training time to improve the training of a smaller model, our method is a run-time solution applied during the decoding process. Therefore, our approach can be more adaptive to run-time behaviors and does not add complexity to training.
Ensemble learning is another approach for coordinating multiple models, wherein multiple models are trained independently and their predictions are combined to improve overall performance. Ensemble learning has been found to yield promising results for Transformer inference [44, 25, 77, 40, 24],
particularly when the models aggregated are pre-trained on different datasets and use different techniques. However, ensemble learning generally requires running multiple models and combining their predictions at run-time, which can be computationally expensive and not optimized for latency. Our research aims to optimize both model performance and run-time latency.
Concurrently and independently of our work, [38, 4] also propose an interesting algorithm to accelerate generative inference using a more powerful model to score and speculatively sample predictions from a less powerful model. While [38, 4] offer unbiased estimators that match the stronger model's probability distributions, our extensive empirical evaluation shows that our approach can deliver superior latency-performance trade-offs, due to its non-random rollback (i.e., rejection) policy as well as the dynamic fallback window size. See Section 4.2 and Appendix 7.3 for an in-depth comparison.
3.1 Motivating Examples
Although large models tend to produce higher-quality text, they also result in longer end-to-end latencies, which can be further exacerbated by the regressive process of predicting one token at a time. However, in many text generation scenarios we demonstrate that a model that is an order of magnitude smaller than a larger model can achieve comparable generation quality to the larger model, provided that a few erroneous predictions are corrected. This implies that only a small fraction of the small model's predictions deviate from those of the larger model. To validate this claim, we evaluate two different generative scenarios, machine translation with mT5 [78] on WMT 2014 De-En [1] and summarization with T5 [47] on CNN/DailyMail [21] by running the large model along the small model for every decoding iteration. See Section 4.1 for more details on these models. Then, we measure the likelihood of the large model predicting the same token
Figure 2: Quality of text generation for different proportions of the large model's engagement on the small model's prediction, evaluated on the validation datasets of (Left) WMT 2014 De-En translation [1]; and (Right) CNN/DailyMail summarization [21]. We see that the small models can achieve a comparable or better generation quality to the large models if  $\sim 20\%$  of their incorrect predictions were substituted.
that the small model generates. If the likelihood is below a certain threshold, we assume that the small model's prediction is not accurate enough, and we replace it with the large model's prediction. By controlling the threshold, we can adjust the proportion of the large model's engagement.
Figure 2 plots the text generation quality on the validation dataset of each benchmark for different proportions of the large model's engagement. The results exhibit a clear trend across the tasks where the small models with  $\sim 10\times$  smaller sizes can retain the large model's generation quality only if approximately  $20\%$  of their inaccurate predictions were substituted by the large model. While this experiment assumes an ideal case where the predictions of the large model are available as ground truth in every iteration, it nonetheless demonstrates the feasibility of achieving the text generation quality of the large model while maintaining the low inference latency of the small model.
At  $n$ th decoding iteration, the small model and the large model each take as input a partially generated output text  $y_{1:n-1} = (y_1, \dots, y_{n-1})$ , and then generate a probability distribution over entire vocabulary  $p_S(y|y_{1:n-1})$  and  $p_L(y|y_{1:n-1})$ , respectively. Then, the next token  $y_{n,S}$  and  $y_{n,L}$  are sampled from the probability distributions,
$$
y _ {n, S} \sim p _ {S} (y | y _ {1: n - 1}) \quad \text {a n d} \quad y _ {n, L} \sim p _ {L} (y | y _ {1: n - 1}). \tag {1}
$$
Depending on whether to use the small model or the large model for the  $n$ th decoding step, the  $n$ th token  $y_{n}$  can be either  $y_{n,S}$  or  $y_{n,L}$ . When deciding which model to use, it is not feasible to run the large model along with the small model for every decoding step to verify the predictions of the small model, as in the experiments in Section 3.1. Thus, it is necessary to hand over the control to the
Figure 3: (Top) The fallback policy. When the small model generates tokens autoregressively, if the prediction probability of a specific token is below the predefined fallback threshold value  $\alpha_{FB}$ , the prediction is deemed to be not confident enough, and control is then shifted to the larger model to produce the corresponding token. (Bottom) The rollback policy. If the large model takes over the control, it produces its own predictions for all previous tokens, as well as the current token. If the prediction probability from the large model for a previously generated token deviates from the small model's prediction by a distance metric  $d$  exceeding the predetermined rollback threshold  $\alpha_{RB}$ , the small model's prediction is regarded as incorrect. In such a case, we roll back all predictions made by the small model that follow the corresponding token.
large model only when the small model is likely to make an inaccurate prediction based on a policy  $\pi(y_{1:n-1})$  that returns a boolean value  $\{0, 1\}$  indicating whether to use the large model:
$$
y _ {n} = \left\{ \begin{array}{l l} y _ {n, S} & \text {i f} \pi \left(y _ {1: n - 1}\right) = 0 \\ y _ {n, L} & \text {i f} \pi \left(y _ {1: n - 1}\right) = 1. \end{array} \right. \tag {2}
$$
The objective, therefore, is to design a lightweight policy  $\pi$  that leads to high text generation quality with the minimum end-to-end latency by invoking the large model only when necessary. In order to illustrate the mechanism by which latency is reduced, consider a simple case where the small model has generated tokens  $y_{1}$  through  $y_{n}$  autoregressively. If the large model takes over the control and predicts the next token,  $y_{n + 1}$ , it can now take multiple input tokens ( $y_{1}$  through  $y_{n}$ ) in parallel, thus allowing for non-autoregressive inference. It is worth noting that this non-autoregressive approach would require the same amount of FLOPs as a regressive approach that predicts  $y_{1}$  through  $y_{n + 1}$  sequentially; however, it is much faster on hardware due to its token-level parallelism and increased arithmetic intensity [72]. In other words, processing multiple tokens in a single memory operation is more efficient than individually processing tokens with separate memory operations, as memory accesses can be more costly than arithmetic operations in decoding tasks [32]. If the latency saving from running the large model non-autoregressively outweighs the additional cost of running the small model, there is a net latency reduction. Therefore, the aim of this approach is not to reduce the number of FLOPs, but rather to improve the hardware utilization and arithmetic intensity of the decoding process. More detailed analysis on this can be found in Appendix 7.5.1. Figure 1 provides a high-level overview of how the small and the large models in BiLD coordinates for text generation.
We now focus on constructing an appropriate policy  $\pi$  for our framework. Here, we introduce two simple policies, the fallback and rollback policies, which (despite their simplicity) result in high performance with significant latency reduction. We discuss the details in the following subsections.
3.3 Feedback Policy: Small Model Knows When to Stop Predictions
The first principle of the policy is that the small model should be able to determine when to hand over control to the large model. Whenever the small model lacks confidence in its prediction, it is
better to allow the large model to take over. Confidence (or uncertainty, in reverse) quantification has been an active research area [13, 30], and any lightweight confidence metric can serve as a potential candidate. Here, we find it sufficient with a simple policy based on the maximum prediction probability, i.e.,  $\max_y p_S(y|y_{1:n - 1})$ , similar to the observations made in [20]. If the maximum prediction probability is lower than a certain threshold  $\alpha_{FB}$ , then the small model's prediction is regarded to be not confident enough, and we fallback to the large model to generate the next token. Note that this does not entail a runtime overhead. Figure 3 (Top) illustrates the fallback policy.
**Fallback Policy:** If  $\max_y p_S(y|y_{1:n-1}) < \alpha_{FB}$ , then fallback to the large model and set  $y_n = y_{n,L}$ .
3.4 Rollback Policy: Large Model Knows When to Revert Predictions
While the fallback policy allows the large model to take over when the small model is not confident enough, it is still possible that the small model is over-confident in its incorrect predictions [18]. Moreover, a single incorrect prediction at an early decoding iteration can lead to a catastrophic effect [51], as it will affect all subsequent token predictions. To avoid such cases, it is desirable to have the large model review the small model's predictions and ensure the validity of each prediction. In our framework, this comes without any extra cost. When the large model is provided with the tokens generated by the small model for its non-autoregressive prediction of the next token, it also produces its own predictions for all the previous decoding steps. That said, given the partially generated text  $y_{1:n}$ , it generates  $p_L(y|y_{1:m})$  for all previous and current decoding steps  $m = 1,\dots ,n$ , which can be used to validate the small model's previous predictions.
Therefore, for some distance metric  $d(\cdot, \cdot)$  that compares two probability distributions, we find the smallest decoding step  $m$  such that
$$
d \left(p _ {S} \left(y \mid y _ {1: m}\right), p _ {L} \left(y \mid y _ {1: m}\right)\right) > \alpha_ {R B} \tag {3}
$$
for a predetermined threshold  $\alpha_{RB}$ . If such  $m$  exists, we regard the small model's previous prediction  $y_{m}$  to be inaccurate, and we rollback all predictions that follow, i.e.,  $y_{m}$  through  $y_{n}$ , since they are all dependent on the wrong prediction. We then replace  $y_{m}$  with  $y_{m,L}$  of the large model. We will discuss in Section 4.2 that the cross-entropy loss between the small model's hard label and the large model's soft label (which measures the likelihood of obtaining the small model's prediction from the large model's output) is a good choice for the metric  $d$ . Rollback may incur additional latency due to the need for duplicated computation for the reverted tokens. However, we demonstrate in Section 4.3 the net advantage of rollback as the improved text generation quality outweighs the additional latency. See Figure 3 (Bottom) for a detailed illustration of the rollback policy.
Rollback Policy: If there exists a minimum  $m \in [1, n - 1]$  such that  $d(p_S(y|y_{1:m}), p_L(y|y_{1:m})) > \alpha_{RB}$ , then rollback the predictions  $(y_m, \dots, y_n)$  and set  $y_m = y_{m,L}$ .
3.5 Big Little Decoder
Taken together, the Big Little Decoder (BiLD) framework consists of one small model, one large model, and a policy that determines which model to use for each decoding iteration. The policy comprises two components: the fallback policy to fall back to the large model when the small model's prediction is not confident enough; and a rollback policy to roll back the small model's predictions if they deviate from the predictions of the large model. Algorithm 1 provides a summary of the end-to-end algorithm.
Algorithm 1: Big Little Decoder