SlowGuess/NeurIPS_2023_OCR huggingface dataset:
1:  $y\gets [
$
2: while  $y[-1] \neq <\mathrm{eos}>$
3:  $p_S \gets \text{SmallModel}(y)$
4: if  $\max(p_S[-1]) > \alpha_{FB}$
5: # Use the small model's predicton
6:  $y\gets y + [\mathrm{sample}(p_S[-1])]$
7: else
8: #Fallback to the large model
9:  $p_L \gets \text{LargeModel}(y)$
10:  $m\gets \min .$  index such that  $d(p_L[m],p_S[m]) > \alpha_{RB}$
11: if  $m$  exists
12: #Rollback:use the large model's prediction
13:  $y\gets y[:m] + [\mathrm{sample}(p_L[m])]$
14: else
15: # Don't rollback: use the large model's prediction
16:  $y\gets y + [\mathrm{sample}(p_L[-1])]$
17: return  $y$
3.5.1 Model Prediction Alignment
BiLD is a general framework that imposes no restriction on the selection of small and large models as long as they use the same vocabulary. Therefore, as will be demonstrated in Section 4.2, two
independently trained models can compose BiLD to achieve a significant latency improvement. Nevertheless, when two models are trained separately, they may generate sequences with similar or identical semantic meanings but using different vocabularies. For instance, one model may produce the phrase "writing is hard" while the other may generate "writing is difficult". Because the BiLD policy relies on the degree of agreement between the large and small models, such a vocabulary-level discrepancy can result in unnecessary disagreements that roll back the small model's prediction without any improvement in generation quality.
In order to address this issue and further optimize the BiLD performance, we present a simple approach called model prediction alignment that aids in aligning the predictions produced by the small and large models. To achieve this, we leverage a calibration dataset  $\mathcal{X}_{\mathrm{cal}} = \{x^{(i)}\}$  that well represents the input sentence distribution. We then generate the corresponding output sequence for each input sequence using the large model, resulting in  $\mathcal{Y}_{\mathrm{cal}} = \{y^{(i)}\}$  where  $y^{(i)} = \arg \max p_L(y|x^{(i)})$ . Subsequently, we fine-tune the small model using the calibration examples  $(x_{\mathrm{cal}}, y_{\mathrm{cal}}) \in (\mathcal{X}_{\mathrm{cal}}, \mathcal{Y}_{\mathrm{cal}})$ .
The underlying rationale of this approach is to increase the likelihood of the small model generating sequences that would have been generated by the large model. This can minimize the distance between the small model and the large model's predictions per each token, i.e.,  $d(p_{S}(y|x,y_{1:m}),p_{L}(y|x,y_{1:m}))$ , throughout the decoding process, thereby avoiding unnecessary rollbacks. Despite its simplicity, our experiments in Section 4.2 demonstrate that this approach can be incorporated into the BiLD framework with minimal effort to significantly enhance the performance. We further emphasize that this method does not introduce any additional complexity or hyperparameters to the normal training pipeline. This is comparable to knowledge distillation [22], an alternative method for aligning model predictions, which requires modifications to the training pipeline, access to internal hidden states such as logits, and additional hyperparameter tuning.
Models and Datasets. To access the generalizability and validity of BiLD in various text generation settings, we have selected IWSLT 2017 De-En [3] and WMT 2014 De-En [1] for machine translation benchmarks and XSUM [43] and CNN/DailyMail [21] for summarization benchmarks. We used mT5-large and small [78] for machine translation and T5-large and small [47] for summarization as our target models, where the size of the models differ by approximately a factor of 20. Our framework is built on top of PyTorch [45] and the HuggingFace Transformers library [73] along with their pre-trained checkpoints.
Training. We fine-tune the pre-trained models on the target benchmarks for 500k steps to obtain the baseline small and large models. To train the aligned small models via the prediction alignment method (Section 3.5.1), we generate output sequences from the input sequences of the training datasets using the fully trained large models to create calibration datasets. We then fine-tune the pre-trained small models on the calibration dataset using the same training recipes and the number of steps as the baseline small models. More training details can be found in Appendix 7.1.1. Throughout the paper, we refer to BiLD with the baseline and aligned small models as unaligned and aligned BiLD, respectively.
Inference. All inference evaluations are conducted on a single NVIDIA T4 GPU of a GCP n1-standard-4 instance, using a batch size 1, which is a common use case for online serving [51]. For the distance metric  $d$  in Equation 3 for the rollback policy, we use the cross-entropy loss between the small model's hard label and the large model's soft label. For BiLD inference, we sweep over different fallback and rollback thresholds to explore different trade-offs between generation quality and latency. More evaluation details can be found in Appendix 7.1.2.
The main results are illustrated in Figure 4, which shows the trade-off between text generation quality and average end-to-end latency per example, normalized by the vanilla inference latency of the pure large baseline models. The trade-offs are obtained by controlling the fallback and rollback thresholds. Table 1 summarizes the results, with the second and third rows corresponding to unaligned
Table 1: The summary of Figure 4 which compares the generation quality and latency speedup of BiLD against vanilla inference with large baseline models. The first row reports the vanilla inference, and the second and third rows report unaligned BiLD. The fourth and fifth rows report aligned BiLD. In both cases of unaligned and aligned BiLD, we report the speedup with minimal BLEU/ROUGE-L score degradation (second and fourth rows), and within  $\sim 1$  point degradation (third and fifth rows).
Task (Model)	Machine Translation (mT5)	Summarization (T5)
Dataset	IWSLT	WMT	XSUM	CNN/DailyMail
BLEU	Speedup	BLEU	Speedup	ROUGE-L	Speedup	ROUGE-L	Speedup
Vanilla Inference	40.32	-	31.38	-	35.08	-	41.54	-
BiLD (Unaligned)	40.33	1.43×	31.28	1.34×	35.12	1.48×	41.44	1.71×
39.44	1.58×	30.47	1.43×	34.02	1.72×	40.57	2.05×
BiLD (Aligned)	40.24	1.62×	31.26	1.47×	35.05	1.50×	41.52	1.85×
39.13	1.78×	30.33	1.70×	33.95	1.80×	40.96	2.12×
Figure 4: Generation quality and average end-to-end latency of processing a single example on 4 different benchmarks. We report BLEU for machine translation and ROUGE-L for summarization as performance metrics. The green and blue lines are unaligned and aligned BiLD, respectively. The X marks are the vanilla inference with the baseline large models. For comparison, two horizontal lines are plotted to indicate the BLEU/ROUGE-L score of the vanilla inference and 1 point degradation from it. The latency on the x-axis is normalized by the baseline latency.
BiLD. When coupled with the normally fine-tuned baseline small models, BiLD achieves an average speedup of  $1.50 \times$  across all benchmarks, with up to  $1.71 \times$  speedup on CNN/DailyMail without any degradation in text generation quality (2nd row). By allowing  $\sim 1$  point degradation, BiLD achieves an average speedup of  $1.70 \times$ , with up to  $2.05 \times$  speedup (3rd row). Note that unaligned BiLD is a pure plug-and-play solution that does not require additional training effort or cost beyond preparing small and large models independently.
In addition, Figure 4 shows the efficacy of the prediction alignment method, leading to a consistent improvement of aligned BiLD over unaligned BiLD. As summarized in the forth and fifth rows of Table 1, aligned BiLD that incorporates the aligned small models yields an average speedup of  $1.61 \times$  with up to  $1.85 \times$  speedup (4th row). Within  $\sim 1$  point degradation, it achieves an average speedup of  $1.85 \times$ , with up to  $2.12 \times$  speedup (5th row). The results also demonstrate that both unaligned and aligned BiLD outperform the baseline BLEU/ROUGE-L scores in the high-latency regime, which can be attributed to the ensembling effect of using two different models, as also studied in prior work [40]. In Appendix 7.5.2, we provide examples of text sequences generated by BiLD, which demonstrate that the large model's engagement in BiLD decoding not only improves the prediction accuracy but also prevents incorrect predictions from impacting the future ones.
We have additionally conducted a performance comparison of our method with the speculative sampling method proposed in [4] on the IWSLT 2017 De-En and XSUM benchmarks. We implement and evaluate it in the same environment as our main BiLD experiments using the same baseline large and small models. We apply a fixed window size of [3, 10]. On the IWSLT benchmark, speculative sampling achieves a BLEU score of 39.93 with a  $1.28 \times$  speedup, while BiLD (unaligned) achieves a 0.61 higher BLEU score with similar speedup, or a  $0.21 \times$  more latency gain with a similar BLEU score. On the XSUM benchmark, speculative sampling achieves a ROUGE-L score of 35.00 with a  $1.25 \times$  speedup. In contrast, BiLD achieves up to a 0.30 ROUGE-L score gain with a faster latency, or up to  $0.22 \times$  more latency gain with a better ROUGE-L score. We provide more detailed comparisons in Appendix 7.3.
Figure 5: Ablation study results for BiLD on (Left) IWSLT 2017 De-En translation and (Right) XSUM summarization tasks without the rollback or fallback policy. Aligned small models were used in all cases. The result demonstrates that BiLD experiences significant performance degradation without either policy in both tasks. The horizontal lines indicate the vanilla inference score and 1 point degradation from it.
Figure 6: Application of the BiLD framework to the early exit problem using the mT5-small model as the large model and its first layer as the small model, evaluated on (Left) the IWSLT 2017 De-En and (Right) WMT 2014 De-En benchmarks. The  $\times$  marks indicate the latency and BLEU score of the mT5-small models. The horizontal lines indicate the vanilla inference score and 1 point degradation from it.
We have further conducted two ablation studies to validate the individual components of BiLD by (1) removing the rollback policy, and (2) removing the fallback policy. When removing the rollback policy, we use the same fallback thresholds as the main experiments to control the generation quality and latency trade-off. When removing the fallback policy, we use the same rollback thresholds as the main experiments. Additionally, we apply fallback after a fixed number of small model executions (swept over [3, 10]), similar to [4].
Figure 5 illustrates the results of these ablation studies on IWSLT 2017 De-En for machine translation and XSUM for summarization with aligned BiLD. The results show that the rollback policy consistently produces better generation quality across all latency regimes, particularly in the high-BLEU/ROUGE regime where the large model's engagement via rollback is crucial in correcting small model's wrong predictions. This demonstrates that, despite the additional latency overhead from the duplicated computation of reverted tokens, the improvement in text generation quality outweighs this cost. Similarly, removing the fallback policy and periodically handing over control to the large model after a fixed number of token generations leads to significant performance degradation. Taken together, these results highlight that both policies are critical components of BiLD.
4.4 Early Exiting Strategy in the BiLD Framework
So far, we have demonstrated how BiLD can be used as a general framework for accelerating the text generation process by incorporating a small model and a large model. However, having two separate models is not a limitation as they can be combined into a single model by using a subset of a larger model, such as a few of its early layers, as a smaller model. This approach resembles the early exit strategy, which is a popular method for accelerating the decoding process [51]. This section demonstrates how the early exiting strategy can be reframed within the BiLD framework.
To demonstrate the applicability of using the early exiting strategy within the BiLD framework, we use mT5-small model as the large model and the first (out of 8) layer as the small model, and evaluate it on two machine translation benchmarks: IWSLT 2017 De-En and WMT 2014 De-En. To ensure consistency between the prediction made after the first layer and the one made after the last layer, we train the model with the average loss of these two layers, similar to [11, 51]. The prediction head is shared for these two layers. More training and evaluation details can be found in Appendix 7.2.1.
Figure 6 illustrates the results, where for each benchmark, BiLD achieves up to  $1.60 \times$  and  $1.74 \times$  speedup within less than one point BLEU score drop, respectively. This demonstrates the extensibility of the BiLD framework to early exit problems. In Appendix 7.2.2, we further provide a detailed comparison of our results with CALM [51], another framework that incorporates early exiting for fast Transformer decoding. Compared to CALM, BiLD offers two advantages that contribute to better generation quality: (1) in BiLD, even if an early exited prediction (i.e., prediction made by the smaller model) is incorrect, it can be corrected and replaced using the rollback policy; (2) the key and value caches for skipped layers are filled with actual values instead of being computed from the
exiting layer's hidden states, leading to reduced error propagation and improved decoding stability. As a result, when tested on IWSLT 2017 De-En and WMT 2014 De-En using mT5-small, BiLD achieves a BLEU score improvement of up to  $\sim 2$  points over CALM in both datasets (Figure 7).
In this work, we have introduced Big Little Decoder (BiLD), a framework that reduces end-to-end inference latency for a wide variety of text generation tasks without the need for training or modifying the existing models. In essence, our framework couples a large and small decoder model together to generate text more efficiently. In particular, we start inference with a small model which runs autoregressively for the majority of the time to generate text with a low inference cost, while the large model is executed non-autoregressively to refine the small model's inaccurate predictions. BiLD incorporates two policies, the fallback policy, which hands control to the large model when the small model is uncertain, and the rollback policy, which allows the large model to revert the small model's inaccurate predictions. Our framework is evaluated across various text generation scenarios, including machine translation, summarization, and language modeling. Running on an NVIDIA Titan Xp GPU, with no performance drop BiLD achieved an average speedup of  $1.52 \times$ , with improvements of up to  $2.18 \times$  on some tasks. Furthermore, when a 1 point degradation in performance was allowed, BiLD achieved an average speedup of  $1.76 \times$  with speedups of up to  $2.38 \times$  on some tasks.
We acknowledge gracious support from Google Cloud, Google TRC team, and specifically Jonathan Caton, Prof. David Patterson, and Dr. Ed Chi. Prof. Keutzer's lab is sponsored by Intel corporation, Intel VLAB team, Intel One-API center of excellence, as well as funding through BDD and BAIR. Sehoon Kim and Suhong Moon would like to acknowledge the support from Korea Foundation for Advanced Studies (KFAS). Amir Gholami was supported through funding from Samsung SAIT. Michael W. Mahoney would also like to acknowledge a J. P. Morgan Chase Faculty Research Award as well as the DOE, NSF, and ONR. Our conclusions do not necessarily reflect the position or the policy of our sponsors, and no official endorsement should be inferred.
[1] Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, Radu Soricut, Lucia Specia, and Ale's Tamchyna. Findings of the 2014 workshop on statistical machine translation. In Proceedings of the Ninth Workshop on Statistical Machine Translation, pages 12-58, Baltimore, Maryland, USA, June 2014. Association for Computational Linguistics.
[2] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020.
[3] Mauro Cettolo, Marcello Federico, Luisa Bentivogli, Jan Niehues, Sebastian Stüker, Katsuhito Sudoh, Koichiro Yoshino, and Christian Federmann. Overview of the IWSLT 2017 evaluation campaign. In Proceedings of the 14th International Conference on Spoken Language Translation, pages 2–14, Tokyo, Japan, December 14–15 2017. International Workshop on Spoken Language Translation.
[4] Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023.
[5] Daoyuan Chen, Yaliang Li, Minghui Qiu, Zhen Wang, Bofang Li, Bolin Ding, Hongbo Deng, Jun Huang, Wei Lin, and Jingren Zhou. Adabert: Task-adaptive bert compression with differentiable neural architecture search. arXiv preprint arXiv:2001.04246, 2020.
[6] Patrick H Chen and Cho-jui Hsieh. A comparison of second-order methods for deep convolutional neural networks. openreview under ICLR 2018, 2018.
7.1 Experimental Details
Table 2: Model configurations of the large and small models for each evaluation task. For comparison, the number of layers, hidden dimension, FFN dimension, and the number of decoder parameters (without embeddings) for each model are provided.
Task	Model	# Layers	dim	FFN dim	# Params
Machine	mT5-large [78]	24	1024	2816	409M
Translation	mT5-small [78]	8	512	1024	25M
Summarization	T5-large [47]	24	1024	4096	402M
T5-small [47]	6	512	2048	25M
7.1.1 Training Details
For machine translation, we use IWSLT 2017 German-English [3] and WMT 2014 German-English [1] as target benchmarks, and mT5 [78] as a target model. We use the 8-layer mT5-small and the 24-layer mT5-large as the small and large models. For summarization, we use XSUM [43] and CNN/DailyMail [21] as target benchmarks, and T5 [47] as a target model. We use T5-small and T5-large with 6 and 24 layers, respectively, for the small and large models. Table 2 summarizes the size and configuration of each model. All the models are fine-tuned from the pre-trained checkpoints of the HuggingFace library [73] for 500k steps using a batch size of 16. We use Adafactor optimizer [53] with constant learning rate of  $\{0.5, 1, 2, 5\} e - 4$  for the small models and  $\{0.5, 1\} e - 4$  for the large models. We refer to the normally fine-tuned models on the validation datasets as the baseline small and large models.
When training aligned small models via the prediction alignment method described in Section 3.5.1, we first generate calibration datasets using the input sequences from the training datasets of each benchmark. We then use the fully trained large model to generate output sequences through greedy sampling with a beam size of 1. To ensure a fair comparison, we fine-tune pre-trained small models (rather than the baseline small models that are already fine-tuned on the training datasets) on the calibration datasets using the same training recipes and the number of training steps as described above. This decision is based on our observation that fine-tuning a baseline model using the calibration dataset tends to improve generation quality, likely due to the increased number of training examples and data augmentation effects, which makes it difficult to make a fair comparison between unaligned BiLD and aligned BiLD. However, in practice, one can obtain aligned models by applying the prediction alignment method directly to the fine-tuned baseline small models to achieve the best performance.
7.1.2 Evaluation Details
All inference evaluations including latency measurement are conducted on a single NVIDIA T4 GPU of a GCP n1-standard-4 instance with 4 vCPUs and 15GB memory. For inference, we use batch size 1, which is a common use case for online serving [51]. For the distance metric  $d$  in Equation 3 for the rollback policy, we use the cross-entropy loss between the small model's hard label and the large model's soft label. This measures the (negative log) likelihood of obtaining the small model's prediction from the large model's output. For BiLD inference, we sweep over different fallback and rollback thresholds to explore different trade-offs between generation quality and latency. For the machine translation tasks, we use fallback thresholds in [0.5, 0.9] and rollback thresholds in [1, 10]. For the summarization tasks, fallback thresholds in [0.2, 0.6] and rollback thresholds in [2, 6]. We keep the maximum generation length of the small model to 10 to avoid high rollback costs. In Appendix 7.5.3, we provide a detailed analysis of how varying the fallback and rollback thresholds impacts the trade-offs between generation quality and latency in the BiLD framework.
Figure 7: The trade-off curves between inference latency and BLEU score for BiLD and CALM in the early exiting setting for (Left) IWSLT 2017 De-En and (Right) WMT 2014 De-En. The  $\times$  marks indicate the vanilla inference latency and BLEU score of the mT5-small models. The horizontal lines indicate the vanilla inference score and 1 point degradation from it. BiLD outperforms CALM across all speedup regimes by up to  $2\sim 2.5$  points better BLEU score, demonstrating the effectiveness of our approach for the early exiting strategy.
7.2 Details of Early Exiting Strategy in the BiLD Framework
7.2.1 Training and Evaluation Details
The training and evaluation details for BiLD as well as for CALM are as follows.
BiLD. We use the mT5-small model as the large model and the first (out of 8) layer as the small model, and evaluate it on two machine translation benchmarks: IWSLT 2017 De-En and WMT 2014 De-En. To ensure consistency between the prediction made after the first layer and the one made after the last layer, we fine-tune the pre-trained mT5 model using the average loss of the first and the final layers, similar to [11, 51]. That is,  $\mathcal{L} = \frac{1}{2} (\mathcal{L}_1 + \mathcal{L}_{-1})$  where  $\mathcal{L}_1$  and  $\mathcal{L}_{-1}$  are the negative log-likelihood loss after the first layer and the final layer. The prediction head is shared for these two layers. We fine the pre-trained mT5-small model on each benchmark for 500k steps using a batch size of 16. Similar to the main experiments, we use Adafactor optimizer [53] with constant learning rate of  $\{0.5, 1, 2, 5\} \mathrm{e}^{-4}$ . For evaluation, we use fallback thresholds in  $[0.2, 0.8]$  and rollback thresholds in  $[0.5, 1.5]$ .