Speculative Decoding with Big Little Decoder
Authors: Sehoon Kim
Karttikeya Mangalam
Suhong Moon
Jitendra Malik
Michael W. Mahoney
Amir Gholami
Kurt Keutzer
Venue: NeurIPS 2023 poster (2023)
Keywords: Transformer
efficient inference
efficient model
decoding
TL;DR: We propose a plug-and-play framework that uses small and large Transformer models together to speed up text generation by ~2x.

Abstract: 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.12x 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.

Speculative Decoding with Big Little Decoder
Sehoon Kim
Karttikeya Mangalam'
Suhong Moon'
Jitendra Malik
Michael W. Mahoney 123 Amir Gholami
Kurt Keutzer'
University of California, Berkeley PICSI LLBNL
(sehoonkim, mangalam, suhong.movt moon, malik, mahoneymu, amirgh, keutzer)@berkeley.
Abstract
The recent emergence Large Language Models based on the Transformer ar- chitecture 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 appli- cations. The inference latency further exacerbated by autoregressive generative tasks, models need to iteratively to generate tokens sequentially without leveraging token-level parallclization. To address this, we propose Big Little De- coder (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 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 correct the small model's inaccurate predictions. To evaluate our framework across different tasks and models, we apply BiLD various 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.12x speedup with minimal generation quality degradation. Furthermore, our framework fully plug-and-play and can be applied without any modifications in the training process or model architecture. Our code is open-sourced"
Introduction
In recent years, the Transformer 63 has become the de-facto model architecture for 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 their performance, efficiently running these models for inference is challenge due to their large model size and runtime complexity. This limits their use in many applications that require real-time responses.
53 83 62
Despite
These computational inefficiencies are particularly pronounced in autoregressive generative tasks such as machine translation (3] M. summarization 21.. and language modeling [I]. For these tasks, models need run iteratively to generate tokens sequentially, as each token 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
https://gituu...om//ssseeee htt.s//seeeeeeee/////gi/iiteeeee
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

① zeveral @ famous @ songs are ( composed
X composed
①
Several lamoux vonge
@
Model
Small Model (#parameter
Large Model
(éparameter =1ON)
(*parameter - (N)
songs
Several
famous
songs
Sievrnal
famous
songs
and
Autoregrssaive (sequemial)
Autoregressive (sequentiai)
Non-sutonegresshe (paralle)
Latenay:
@+@++.+- (
Latenoy: @
e++ @ @ (*)
Figure 1: Illustration of (Left) the normal autoregressive decoding procedure of 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 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 []. 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 LLMs, is critical to improve the inference latency and runtime efficiency of autoregressive decoding processes despite the potential sacrifice in generation quality.
To overcome this, autorerrreeees decoding [6
19 has been explored to maximize leven-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, con-autoregressivv 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 and knowledge distillation [84].
69
39
59
the decoding process 669 39
59
70
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 model architecture. As illustrated in Figure [ (Right), the BiL.D framework consists of two decoder models, 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-autoregresssve execution. This autoregressive small, non-autoregressive large scheme results in substantial improvement of up to ~2x 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 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
In summary, our main contributions are as follows:
We introduce BiLD, 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 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 '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.I.

---

We apply BiLD for different text generation scenarios including IWSLT 2017 De-En [3] and WMT 2014 De-En [ for machine translation, XSUM 44] and Maall [2T for summarization. Compared to the full autoregressive execution, BiLD achieved speedup of up to 1.85x without generation quality degradation and 2.12x allowing ~ point degradation on an NVIDIA T4 GPU (Section
2 Related Work
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-knows
approaches include efficient architecture design 26
36
60
67 [75,, quantization [31 54
81
82
35 and neural 80 While these methods are generally suitable for Transformer-
79 74, 91, pruning [14] 49
34
42
64
12
architecture search [ 56 57
66
76
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
70 52 to ensure more accurate parallel decoding, or by allowing multiple additional iterations refine any inaccurate predictions B7 15 19].
39
59
Such multi- iteration decoding scheme has also been proposed in [7]77588 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 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 (5T) 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 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 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 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 either infeasible or prohibitively expensive.
Section 4.4. we also show that the early exiting strategy 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 larger model, such as the attention map [60] 28 68]. In
contrast to knowledge distillation, which leverages the knowledge of large model solely during the training time to improve the training of a smaller model, our method is a run- cun-time solution applied during the decoding process. Therefore, our approach can be more adaptive 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 [4 25
40 24).
3

---

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] also propose an interesting algorithm to accelerate generative inference using more powerful model score and speculatively sample predictions from a less powerful model. While [3 offer unbiased estimators that match the stronger model's 's probability distributions, our extensive empirical evaluation shows that our approach can deliver superior
raadeoofss due to its non-random rollback (i.e., rejection)
policy as well as the dynamic fallback window size. See Section Section|4.2.ann and Appendix [.3 an in-depth comparison.
3 Methodology
3.1 Motivating Examples
Although large models tend 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 model that is an order of magni- tude smaller than a larger model can achieve com- parable generation quality to the larger model, pro- vided that few erroneous predictions are corrected. This implies that only small fraction of the small model's predictions deviate from those of the larger model. To validate this claim, we evaluate two differ- ent generative scenarios, machine translation with mT5 (78) on WMT 2014 De-En [I] and summa- rization with T5 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 like- lihood of the large model predicting the same token 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.
WNT 2017 De-En, mT55
CNN/DaillMMill T5
22
muzu (ligher better)
Whigher
19 15
% Lange Model
LaLgg Model
Figure 2: Quality of text generation for dif- ferent proportions of the large model's en- gagement on the small model's prediction, evaluated on the validation datasets of (Left) WMT 2014 De-En translation [ and (Right) CNN/Daily summarization 2T.. We see that the small models can achieve compa- rable or better generation quality to the large models if ~20% of their incorrect predictions were substituted.
47 on CNNiDailyMail (21] by
Figure plots the text generation quality on the validation dataset of cach benchmark for different proportions of the large model's engagement. The results exhibit clear trend across the tasks where the small models with ~ 100 smaller sizes can retain the large model's generation quality only 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.
3.2 Problem Formulation
At nthh decoding iteration, the small model and the large model each take input partially generated output text V1-m-1 (y) ... 1v-1), and then generate a probability distribution over entire vocabulary ( (1) and pr. (1) respectively. Then, the next token yn.s and yr.r. are sampled from the probability distributions,
Fn.S ~ ps (vly and Vn.L ~ PL. (vllinn-)))
(I)
Depending on whether to use the small model or the large model for the with decoding step, the with token In can be either yr or yn, 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

---

Fallback Threshold
Fallback
Several
tamous
composing
Several
famous
people
sincing
Large Model
Model ( peremeetrr N)
(1f (fpprrmmtter =1NN)
<S>
Several
famous
people
are
Several
Tamous
Rollback Threshold Sge
Rollback
distance: d'
Several
Several
femous
Pepple
shr
Small Model
(#parameter N)
Large Model
(#pparmmtter -1ON)
<SP
Several
famous
people
2
Several
famous
Generated Text: Several femous songs
Figure 3: (Top) The fallback policy. When the small model generates tokens autoregressivelly if the prediction probability of a specific token is below the predefined fallback threshold value OFB, 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 's prediction by a distance metric d exceeding the predetermined rollback threshold anz, the small model's prediction regarded as incorrect. In such a case, we roll back predictions made by the small model that follow the corresponding token.
large model only when the small model is likely make an inaccurate prediction based on a policy T (8):n-1) that returns boolean value {0, 1) indicating whether to use the large model:
ya.S if x (y):n-1) InnL
Yn
if x(f1n-1)) =
The objective, therefore, is to design lightweight policy 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 simple has generated tokens y/ through Yn autoregressively. If the large model takes over the control and now take multiple input tokens (уı through yn)) in parallel, thus predicts the next token, In++1, inference. It worth noting that this non-autoregressive approach allowing for non-outoreereesivve would require the same amount of FLOPs as regressive approach that predicts 311 through y sequentially; however, it is much faster on hardware due to its token-level parallelism and increased arithmetic intensity 172 In other words, processing multiple tokens 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 from running the large model an--attoregresivelyy outweights the additional cost of running the small model, there is net latency reduction. Therefore, the aim of this approach is not reduce the number of FL.OPs, 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 provides a high-level overview of how the small and the large models in BiL. coordinates for text generation.
where the small model
the latency saving
We now focus on constructing an appropriate policy π 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 the following subsections.
3.3 Fallback Policy: Small Model Knows When to Stop Predictions
The first principle of the policy that the small model should be able determine when to hand over control to the large model. Whenever the small model lacks confidence in its prediction, it is
5

---

better to allow the large model to take over. Confidence (or uncertainty, in reverse) quantification
has been an active research area [3 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., maxy, (s similar to the observations made in (20] If the maximum prediction probability is lower than certain threshold OFB. then the small model's '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 (Top) illustrates the fallback policy,
Fallback Policy: If max, Ps (v (im-1) < OFF, then fallback the large model and set yn, 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 [5 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 Vie.. it generates pr. (vlVV ) for all previous and current decoding steps m 1, ,,, which can be used to validate the small model's previous predictions.
Therefore, for some distance metric .((:-) that compares two probability distributions, we find the smallest decoding step m such that
d(ps( (1ls(ul)mm)) ), PL (ulyi.m)) CRB
(3)
for a predetermined threshold a an- such m exists, we regard the small model's previous prediction Um be inaccurate, and we rollback all predictions that follow, i.e., ym through y/n since they all dependent on the wrong prediction. We then replace ym with 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 s soft label (which measures the likelihood of obtaining the small model's prediction from the large model's 's output) 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 detailed illustration of the rollback policy.
Rollback Policy: If there exists a minimum E [1 1] such that d(ps (M)) Pr. (yy (1:m))
ORB, then rollback the predictions (ym. yn and set ym Vm
3.5 Big Little Decoder
Algorithm 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
1: y [<S>]
while y|-1] <cos> ps f -SmallMoee(y))
if max(ps(-1))> > ars
5:
# Use the small model's predicton
6:
y /++++ "=sp=smplepps(-----
comprises two components: the fallback
else
policy to fall back to the large model
Fallback to large model
when the small model's prediction is not confident enough; and rollback 10:
9:
Pr. + LargeModel(y)
m < min. index such that d(pz [m], ps [m]] >> aRB
policy to roll back the small model's
11:
if m exists
predictions if they deviate from the pre-
12:
# Rollback: use the large model's prediction
dictions of the large model. Algorithm]] provides summary of the end-to-end
[]
13:
y < w(s m] jsample(Pr. mm]]
14:
else
15:
# rollbackk use the large model's prediction y + y + [sample(pz|-11]
algorithm.
16:
17:
return y
3.5.1 Model Prediction Alignment
BiL.D a general framework that imposes no restriction on the selection of small and large models as long they use the same vocabulary. Therefore, as will be demonstrated in Section 4.2 two

---

independently trained models can compose BiLD to achieve 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 hard" while the other may generate "writing is difficult". Because the BiL.D 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 aligmment that aids in aligning the predictions produced by the small and large models. To achieve this, we leverage calibration dataset X = {x:(*) that well represents the input sentence distribution. We then generate the corresponding output sequence for each input sequence using the large model, resulting in Veat = {y where y( arg max PL (1 (z])z())).. Subsequently, we fine-tune the small model using the calibration examples (#eat, Bew)) E (Xeat, Vea).
The underlying rationale of this approach is to increase the likelihood of the small model gen- erating 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., (fps(y|z, 31:m). PL (y (1:m)), throughout the decoding process, thereby avoiding unnecessary rollbacks. Despite its simplicity, our experiments in Section 4.2 demonstratee 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 hyperparame- ters 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.
4 Evaluations
4.1 Experiment Setup
Models and Datasets. To access the generalizability and validity of BiLD in various text generation settings, we have selected IWSLT 2017 De-En B3 and WMT 2014 De-En [ for machine translation benchmarks and XSUM 4 and CNN/DailyMain [11] for summarization benchmarks. We used -arg and small [78 for machine translation and Ta-large and small [47] for summarization as our target models, where the size of the models differ by approximately 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.I) 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.I 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 nl- standard-4 instance, using batch size 1, which is common case for online serving [51]. For the distance metric di in Equation for the rollback policy, we use the cross-entropy loss between the small model's '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
4.2
Main Results
The main results are illustrated in Figure 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. [ the results, with the second and third rows corresponding to unaligned
7

---

Table I: 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 ~ point degradation (third and fifth rows).
Task (Model)
Machine Translation (miS5)
Summmarization (T5)
IWSLT
WMT
XSUM
CNNyDailyMail
Dataset
BLEU
Speedup
BLEU
Speeduu
ROUGE-L
Speelup
ROUGE-L
Speedup
Vanilla Inference
40.32
31,38
35.08
41,54
BiLD (Unaligned)
40.33
1.41x
31.28
1.34x
35.12
x
41.44
1.71x
39.44
1.58%
30.47
1.43%
34,02
1.72x
40.57
2,05x
BiLD (Aligned)
40.24
1
31.26
1.47x
35.05
1.50x
41.52
1.85x
1.78x
30.33
1.70x
33.95
2.12X
IWSLT 2017 De-En, MTS
WMT 2014 De-En, MTS
XSUM, TS
CNN/DailyMall T5
41.0
RIEU (higher better)
60.0
ALEU (ligher teten
31.50
$1.25
SASTIC
ROUCE- Phighter (neeter))
-5.0
1-5 feottrr
ROUCE Phigheen better)
41
43.
39.5
$1.00
1- taster
-6
2.8h Taster
#1
7.4 hoteer
20.75
32.0
Va nifrence
30.30
Vani la Inferenco
24.0
Vari ia Inference
42
Varilla Inference
BILD (A1g)od)
20.25
BLD (Algnod)
BLD (Aligned)
BLD (Algres)
BLD (Unaligne))
BLD (Uhuligne))
23.5
BLD (Uhua) ggres
40.6
BLD (Uhaligned)
0.6 07
0.6
0. 06
54
077
06
05
06
07
08 09
Norma lzed latency
Normakled latency excmmpe
Normalized ialency ecample
Normalized
example
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 normalized by the baseline latency.
BiLD. When coupled with the normally fine-tuned baseline small models, BiLD achieves an average speedup of 1.50x across all benchmarks, with up to 1.71 x speedup on CNN/DailyMail without any degradation text generation quality (2nd row). By allowing ~1 point degradation, BiLD achieves an average speedup of 1.70x, with up to 2.05 x speedup (3rd row). Note that unaligned BiLD is a pure plug-and-pll 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 aligmment method, leading to a consistent improvement of aligned BiLD over unaligned BiLD. As summarized in the forth and fifth rows of [ aligned BiLD that incorporates the aligned small models yields an average speedup of 1.61 x. with up to 1.85x speedup (4th row). Within ~-1 point degradation, it achieves an average speedup of 1.85 x, with up to 2.12x speedup (5th row). The results also demonstrate that both unaligned and aligned BiLD outperform the baseline BLEU/ROUGE-L scores in the high- 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 fixed window size of [3, 10]. On the IWSLT benchmark, speculative sampling achieves a BLEU score of 39.93 with a 1.28x speedup, while BiLD (unaligned) achieves a 0.61 higher BLEU score with similar speedup. or a 0.21 x more latency gain with similar BLEU score. On the XSUM benchmark, speculative sampling achieves a ROUGE-L score of 35.00 with a 1.25x speedup. In contrast, BiLD achieves up to 0.30 ROUGE-L score gain with a faster latency, or up to 0.22x more latency gain with a better ROUGE-L score. We provide more detailed comparisons in Appendix 7.3
8

---

IWSU 2017 DeEn, mT5
XSUM. T5
INSLT 2017 De-En, m15
WMT 2014 De-En, m75
BLEU (lgine seuenn
$2.0
ROLLEE (higher taet
34:22
16.4
hodge
cogner better
27.0
20.0
22.3
2.00
Wardia
25.0
35.0
1r Ispler
20.8
1.7x
BLEU
20
36. No
B11 the B1. the FR)
31.4
58. (Mo FS3
242
0.6
37
6.8
3.0
16
0.6
3.
3.0 1.
15.4
6.0
67 03 03
Normalized sug atency per example
205
argiotency per exampie
Normalized aug sttrry per example
Normalzed ovg atency per example
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 mod- els were used in all cases. The result demon- strates that BiL. experiences significant perfor- mance 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 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 x marks indicate the latency and BLEU score of the mm55-mmall models. The horizontal lines indicate the vanilla inference score and I point degradation from it.
4.3 Ablation Studies
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 jillusratts the results of these ablation studies on IWSLT 2017 De-En for machine trans- lation 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 is engagement via rollback is crucial in correcting small model's ss 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 fixed number of token generations leads to significant performance degradation. Taken together, these results highlight that both policies are critical components of BiL.D.
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 small model and 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 larger model, such 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.
few of its carly layers,
smaller model. This approach resembles the
To demonstrate the applicability of using the early exiting strategy within the BiLD framework, we use mm5ll model as the large model and the first (out of 8) layer as the small model, and evaluate 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 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 iillstrraes the results, where for each benchmark, BiLD achieves up to 1.60x and 1.74x 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..22 we further provide 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, 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 's hidden states, leading to reduced error propagation and improved decoding stability. As 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 ~2 points over CALM in both datasets (Figure 77
Conclusion
In this work, we have introduced Big Little Decoder (BiLD), a framework that reduces end-to aud-to-end inference latency for wide variety of text generation tasks without the need for training 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 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 von-auloregressively 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 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.52x, with improvements of up to 2.18x on some tasks. Furthermore, when I point degradation in performance was allowed, BiLD achieved an average speedup of 1.76x with speedups of up to 2.38 x on some tasks.
Acknowledgements
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 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.
References
[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 few-shot learners. arXiv preprint arXiv:2005. /4165, 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 Spoken Language Translation, pages 2-14, Tokyo, Japan, December 14-15 2017. International Workshop on Spoken Language Translation.
[4] Charlic Chen, Sebastian Borgcaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint 230.v22302.01311 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.00 04246, 2020.
[6] Patrick Chen and Cho-jui Hsieh. A comparison second-order methods for deep convolu- tional neural networks. openreview under ICLR 2018, 2018.
10

---

[7] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et Palm: Scaling language modeling with pathways. arXiv preprini
2022.
[8] Michiel de Jong. Yury Zemlyanskiy, Joshua Ainslie, Nicholas FitzGerald, Sumit Sanghai, Fei Sha, and William Cohen. Fido: Fusion-in- faster inference. arXiv preprint arXiv: 22.222222../111 1533, 2022.
optimized for stronger performance and
[9] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3, int8 (): 8-bit matrix multiplication for transformers at scale. In Advances in Neural Information Processing Systems, 2022.
[10] Nan Du Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, al. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference on Machine Learming, pages 5547-5569. PMLR, 2022.
[11] Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. arXiv preprint arXiv: 10073, 2019.
[12] Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: 1156,, 2019.
[13] Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050-1059. PMLR, 2016.
[14] Trevor Gale. Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint 17Xiv:1922.09544 2019.
[15] Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zeutlemyee. Mask-predict: Parallel decoding of conditional masked language models. arXiv preprint arXiv:1904.0932. 2019.
[16] Jiatao Gu, James Bradbury, Caiming Xiong, Victor OK Li, and Richard Socher. Non- autoregressive neural machine translation. arXiv preprint arXiv: 01117/111.222111 2017.
[17] Jiatao Gu, Changhan Wang, and Junbo Zhao Levenshtein transformer. Advances in Neural
Information Processing Systems, 32, 2019.
[18] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. International conference on machine learning, pages 1321-1330. PMLR, 2017.
[19] Junliang Guo, Linli Xu, and Enhong Chen. Jointly masked pouuence-to-sqquenee model for non non-autoregressivee neural machine translation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 376-385, 2020.
[20] Dan Hendrycks and Kevin Gimpel. baseline for detecting misclassified and out-of-distributtin examples in neural networks. arXiv preprint arXiv: 2016.
[21] Karl Moritz Hermann, Tomas Kocisky, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. Teaching machines to read and comprehend. In Advances in neural information processing systems, pages 1693-1701, 2015.
[22] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. Workshop paper in NIPS, 2014.
[23] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimet large language models. arXiv preprint arXiv:2203.
2022. [24] Chris Hokamp, Demian Gholipour Ghalandari, Nghia The Pham, and John Glover. Dyne: Dy- namic ensemble decoding for multi-document summarization. arXiv preprint 08X1v:2000.08748. 2020.
[25] Ngo Quang Huy, Tu Minh Phuong, and Ngo Xuan Bach. Autoencoding language model based ensemble learning for commonsense validation and explanation. arXiv preprint arXiv:2204.03324 2022.
[26] Forrest N landola, Albert E Shaw, Ravi Krishna, and Kurt W Keutzer. Squeezebert: What can computer vision teach nlp about efficient neural networks? arXiv preprint arXiv:2006. 1316, 2020.
11

---

[27] Dongfu Jiang, Xiang Ren, and Bill Yuchen Lin. LIm-blender: Ensembling large language models with pairwise ranking and generative fusion. arXiv preprint arXiv:2306.02550 02561, 2023,
[28]
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. arXiv preprint arXiv:1909. 10351, 2019.
[29] Jungo Kasai, Nikolaos Pappas. Hao Peng, James Cross, and Noah A Smith. Deep en- coder, shallow decoder: Reevaluating non con-autoregressive machine translation. arXiv preprint arXiv:2006. I0369, 2020.
[30] Alex Kendall and Yarin Gal. What uncertainties do need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
[31] Sehoon Kim, Amir Gholami, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. I-bert: Integer-only bert quantization. arXiv preprint arXiv:2101.00. 0.1121,, 2021.
[32] Sehoon Kim, Coleman Hooper, Thanakul Wattanawong, Minwoo Kang, Ruohan Yan, Hasan Genc, Grace Dinh, Qijing Huang, Kurt Keutzer, Michael W Mahoney, et al. Full stack optimization of transformer inference: survey. arXiv preprint arXiv:2302. 2023.
[33] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2019.
[34] Eldar Kurtic, Daniel Campos, Tuan Nguyen, Elias Frantar, Mark Kurtz, Benjamin Fineran, Michael Goin, and Dan Alistarh. The optimal bert surgeon: Scalable and accurate second-order pruning for large language models. arXiv preprint arXiv:2203.07255 2022.
[35] Woosuk Kwon, Sehoon Kim, Michael W Mahoney, Joseph Hassoun, Kurt Keutzer, and Amir Gholami. A fast post-training pruning framework for transformers. arXiv preprint arXiv:2204. 09656, 2022.
[36] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervise, learning of language representations. arXiv preprint arXiv: 1909. / 1942, 2019.
[37] Jason Lee. Elman Mansimov, and Kyunghyun Cho. Deterministic non auo-rutorrgressive neural sequence modeling by iterative refinement. arXiv preprint arXiv: vrXiiv:802.000 080200690, 2018.
[38] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. International Conference on Machine Learning, pages 19274-19286. PMLR, 2023.
[39] Zhuohan Li, Zi Lin. Di He, Fei Tian, Tao Qin, Liwei Wang, and Tie-Yan Liu. Hint- bist- training for oon-autoregressive machine translation. arXiv preprint arXiv:1909.0 1909.0670, 2019.
[40] Yoshitomo Matsubara, Luca Soldaini, Eric Lind, and Alessandro Moschitti. Ensemble trans- former for efficient and accurate ranking tasks: an application to question answering systems. arXiv preprint arXiv:2201. 2276677777 2022.
[41] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
[42] Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? arXiv preprint arXiv: 1905. 10650, 2019.
[43] Shashi Narayan, Shay B Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. arXiv 2018.
preprint
[44] Liu Pai. Qiaoning at semeval-2020 task 4: Commonsense validation and explanation system based on ensemble of language model. In Proceedings of the Fourteenth Workshop on Semantic Evaluation, pages 415-421, 2020.
[45] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vitooo, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
[46] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. arXiv preprint arXiv:2211.05100, 2022.
12

---

[47] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21(140):
2020.
[48] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01100 2019.
[49] Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine-tuning. Advances in Neural Information Processing Systems, 0--20388-20353 2020.
[50] Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilić, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 1 766-paramettr open- apen-access multilingual language model. arXiv preprint arXiv:2211.05100, 2022.
[51] Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Q Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. arXiv preprint 22iv:2207.070611 2022.
[52] Chenze Shao, Jinchao Zhang, Yang Feng, Fandong Meng. and Jie Zhou. Minimizing the bag-of-ngrams difference for autorrrorrereeevvee neural machine translation. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 198-205, 2020.
[53] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604, 2018.
[54] Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. Q-BERT: Hessian based ultra low precision quantization of bert. 2020.
In AAAL, pages
[55] Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deepspeed and megatron train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint 22X:::22011119900 2022.
[56] David So, Quoc Le, and Chen Liang. The evolved transformer. In International Conference on Machine Learning, pages 5877-5886. PMLR, 2019.
[57] David R So, Wojeiech Mańke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc V Le. Primer: Searching for efficient transformers for language modeling, arXiv preprint DeXii:2109.08660 2021.
[58) Mitchell Stern, William Chan, Jamie Kiros, and Jakob Uszkoreit, Insertion transformer: Flexible sequence generation via insertion operations. In International Conference on Machine Learning, pages 5976-5985. PMLR, 2019.
[59] Zhiqing Sun, Zhuohan Li, Haoqing Wang, Di He, Zi Lin, and Zhihong Deng. Fast structured decoding for sequence models. Advances in Neural Information Processing Systems, 32, 2019.
[60] Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobilebert: compact task-agnostic bert for resource-limited devices. arXiv preprini urXiv:2004.0298. 2020.
[61] Raphael Tang, Yao Lu, Linging Liu, Lili Mou, Olga Vechtomova, and Jimmy Lin. Distilling task- specific knowledge from bert into simple neural networks. arXiv preprint arXiv: 12136, 2019.
[62] Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.0822
2022.
[63] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008 2017.
[64] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:1905.09418, 2019.
[65] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. GLUE: multi-task benchmark and analysis platform for natural language understanding. arXiv preprint
2018.
13

---

[66]
Hanrui Wang, Zhanghao Wu, Zhijian Liu, Han Cai, Ligeng Zhu, Chuang Gan, and Song Han. Hat: Hardware-award transformers for efficient natural language processing. arXiv preprint arXiv:2005. 14187, 2020.
[67]
Sinong Wang. Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprin! urXiv:2006.04768. 2020.
[68]
Wenhui Wang. Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. arXiv
preprint Xr0: 10957, 2020.
[69] Yiren Wang. Fei Tian, Di He, Tao Qin, ChengXiang Zhai, and Tie- Yan Liu. Non-autoregressive machine translation with auxiliary regularization. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 5377-5384, 2019.
[70] Bingzhen Wei, Mingxuan Wang, Hao Zhou, Junyang Lin, Jun Xie, and Xu Sun. Imitation learning for non-autoregressive neural machine translation. arXiv preprint arXiv:1906.000 02041, 2019.
[71] Sean Welleck, Kianté Brantley, Hal Daumé lii, and Kyunghyun Cho. Non-monotonic sequential text generation. In International Conference on Machine Learning, pages 6716-6726 PMLR, 2019.
[72] Samuel Williams, Andrew Waterman, and David Patterson. Roofline: an insightful visual performance model for multicore architectures. Communications of the ACM, 52(4):65-76. 2009.
[73] Thomas Wolf, Julien Chaumond, Lysandre Debut, Victor Sanh, Clement Delangue, Anthony Moi, Pierric Cistac, Morgan Funtowicz, Joe Davison, Sam Shleifer, et al. Transformers: State- of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38 45, 2020.
[74] Xiaoxia Wu, Zhewei Yao, Minjia Zhang, Conglong Li, and Yuxiong He. Extreme compression for pre-trained transformers made simple and efficient. arXiv preprint arXiv:2206.0/88 2015229966.8, 2022.
[75] Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, and Song Han. Lite transformer with long-short range attention. arXiv preprint arXiv:2004. 1886, 2020.
[76] Jin Xu, Xu Tan, Renqian Luo, Kaitao Song, Jian Li, Tao Qin, and Tie- Yan Liu. Nas-bert: task- agnostic and adaptive-size bert compression with neural architecture search. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pages 1933-1943, 2021.
[77] Yige Xu, Xipeng Qiu, Ligao Zhou, and Xuanjing Huang. Improving bert fine-tuning via self-ensemble and self-distillation. arXiv preprint arXiv:2002. 10345, 2020.
[78] Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami Al-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mt5: massively multilingual pre-trained text-to-text trans- former. arXiv preprint arXiv:2010.1 1194. 2020.
[79] 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. arXiv preprint arXiv:2206. 20666666..66..... 2022.
[80] Yichun Yin, Cheng Chen, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. Autotinybert: Automatic hyper-parameter optimization for efficient pre-trained language models. arXiv preprint arXiv:2107. / 33686, 2021.
[81] Ali Hadi Zadeh, Isak Edo, Omar Mohamed Awad, and Andreas Moshovos. Gobo: Quantizing attention-based nlp models for low latency and energy efficient inference. In 2020 53rd Annual IEEE/ACM International Symposium on Microarchitecture (MICRO), pages 81 1 1-824. IEEE, 2020.
[82] Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8BERT: Quantized 8bit bert. arXiv preprint arXiv: 1910.0668, 2019.
[83] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01065 2022.
[84] Chunting Zhou, Graham Neubig, and Jiatao Gu. Understanding knowledge distillation in non-autoregressivee machine translation. arXiv preprint arXiv:1911.02777 2019.
14

---

7 Supplementary Material
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
FEN dim
# Params
Machine
mTS-large
78
24
1024
2816
409M
Translation
mTS small
78
8
512
1024
25M
Summarization
TS-lage
24
1024
4096
402M
T5 small
+Y
512
2048
25M
7.1.1 Training Details
For machine translation, 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 - as the small and large models. For summarization, we use XSUM 43 and CNai [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[]summan the size and configuration of each model. All the models are fine-tuned from the pre-trained check- points of the HuggingFace library (73 for 500k steps using 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, 1je- for the large models. We refer the normally innettned models 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 the fully trained large model to generate output sequences through greedy sampling with a beam size of I. To ensure a fair comparison, we fine-tune pre-trained small models (rather than the baseline small models that already finettuned on the training datasets) the calibration datasets using the same training recipes and the number of training steps as described observation that fine-tuning baseline model using the calibration
above. This decision is based dataset tends to improve generation quality, likely due to the increased number training examples and data augmentation effects, which makes it difficult to make 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 n -s-anndddd--4 instance with 4 vCPUs and 15GB memory. For inference, we use batch size which is a common use case for online serving (5I) For the distance metric in Equation 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 '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 [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.
15

---

IWSLT 2017 De- De-En, mT5
WMT 2014 De-En Be-En, mT5
36.5
BLEU thigher better)
BLEU (higher belter)
36.0
26.5
35.5
26.0
35.0
25.5
34.5
Vanilla Inference
BILD
25.0
X
Vanilla Inference
BILD
CALM
CALM
34.0
24.5
0.7
0.8
0.9
10
Normalized avg latency per example
Normalized avg lalency per example
Figure 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 marks indicate the vanilla inference latency and BLEU score of the mT5-smaal models. The horizontal lines indicate the vanilla inference score and point degradation from it. BiLD outperforms CALM across all speedup regimes by up to 2 ~ 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 the mT5-small model as the large model and the first (out of 8) layer as the small model, and evaluate 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 11 55 That is, L 1/40 (C1 + L-1) where Li and L- are the negative log- bogelikeldhood loss after the first layer and the final layer. The prediction head is shared for these two layers. We fine the pre-trained mmS-l model on each benchmark for 500k steps using batch size of 16. Similar to the main experiments, we use Adafactor optimizer (53) with constant learning rate of [0.5, 1, 2, 5)e- 4. For evaluation, we use fallback thresholds in [0.2, 0.8] and rollback thresholds in [0.5, 1.5].
CALM. To reproduce CALM [51] in our experimental setup, we have fine-tuned the pre- pre-trained mT5-small model on IWSLT 2017 De-En and WMT 2014 De-En datasets. We employ the averaged loss across all layers, i.e., L = Ete Li, where = i/ Σラ= j. which was introduced in the paper to ensure the layer consistency. We use Adafactor optimizer [53] with constant learning rate 4 for 500k training steps. To make a fair comparison, we match the BLEU
of (0.5, 1, 2, 4 score of the fine-tuned model to that of BiLD's models Among the two training- confidence measures introduced in the CALM paper, softmax-bust based and hidden-state saturation-bass measures, we have chosen to use the latter approach as an early exiting criterion. That said, if the cosine similarity between the current layer's hidden states and the previous layer's 's hidden states exceeds a certain threshold, we perform early exiting. We have found that the softmax- based alternative is not applicable in our evaluation scenario due to the large output vocabulary (more than 200k for mT55 which is 10x larger than T5), which significantly increases latency overhead. As described in the paper, when early exiting happens, the hidden states of the exited layer are propagated down to the remaining layers to compute the key and value caches. To achieve different trade-off's between latency and generation quality, we sweep over À in [0.7, 0.98] and / {0, 1, 2, 4, 8) in the decaying threshold function.
7.2.2 Performance Comparison between BiLD and CALM
Figure 7 jillssrrates the BLEU score and latency curves of BiLD compared to CALM in the early exiting setting. In both tasks, our method achieves significantly better BLEU scores with the same latency speedup, yielding up to around 2 point better BLEU score in the 1.5x speedup regime. This can be attributed to two factors. First, in BiLD even an early exited prediction (i.e., prediction made by the smaller model) is incorrect, it can be corrected and replaced using the rollback policy.
16

---

Therefore, an error in the early exited layer is propagated less drastically to the future prediction. Second, the key and value caches for skipped layers are filled with actual values instead of being computed from the exiting layer's hidden states. This also leads to reduced enror propagation and improved decoding stability.
7.3 Comparison with Other Speculative Decoding Frameworks
Concurrently and independently of our work, [38 also propose an algorithm to accelerate generative inference using more powerful model to score and speculatively sample predictions from a less powerful model. While the rejection sampling-bast approach in [3 4 offers unbiased estimators that match the stronger model's probability distributions, our extensive empirical evaluation shows that our approach can deliver superior lalency-performance trade-offs, due to its non-random rollback (i.e., rejection) policy well as the dynamic fallback window size. Below, provide distinctions in detailed methodologies and quantitative comparison, as well as our insights on better latency and performance of our approach.
7.3.1 Differences methodology
While the idea of using models with different sizes can be deemed similar the speculative decoding frameworks in 38
we have clear distinctions in detailed methodologies.
(1) Non-Random Prediction Rollback Approach: The primary difference lies in how we decide the rollback (e.g., rejection) of predictions from the small model. In our rollback policy, we propose to make the rejection decision based on the distance between the small and large model predictions, which differs from the rejection sampling policy outlined (38 4 While [38] 4 propose an unbiased estimator the large model's 's prediction. Figure demonstrates that combining predictions
from both models through our distance-based rejection approach can surpass the exclusive utilization of the large model's prediction probability. BiLD seeks to find and utilize this optimal performance point without introducing much runtime cost. We have a further discussion below about how our rejection policy benefits text-generation performance.
(2) Dynamic Fallback Window Size: Additionally, we introduce the dynamic fallback window size in our fallback policy. In (38 4), the window size remains fixed hyperparameter; however, it is also highlighted in 4 that the window size can have a noticeable impact on end-to-end latency. Our approach offers an efficient and robust solution: adjusting the window size at runtime based on the small model's confidence level in run-time. Our ablation study (Figure 5 demonstrates that omitting the fallback policy and periodically transitioning control to the large model, as proposed in [38 can result in notable latency degradation.
(3) Model Alignment Enhancement: Beyond the core framework, we introduce model alignment method align the small model's predictions with those of the large model. This enhances the framework by reducing unnecessary rejections and can be incorporated with minimal adjustments to the training pipeline.
7.3.2 Quantitative Comparisons
In Table we provide a comprehensive quantitative comparison between our method and [38 across two different datasets: IWSLT for machine translation and XSum for summarization. In order to ensure a fair comparison that isolates the impact the frameworks themselves, we employ the baseline (non-aligned) small model for all experiments. We maintained the same evaluation setup and hyperparameter space that are outlined in Appendix 7...22
Table incluuess two BiLD configurations: the one that matches latency and the other that matches BLEU/ROUGE-L scores as compared to the rejection sampling-base methods. Across all experi- ments, BiLD consistently outperforms speculative decoding. It achieves either (1) notably improved BLEU/ROUGE-L scores with equivalent latency gains, (2) superior latency gains while retaining the same BLEU/ROUGE-L scores.
17

---

Table Comparison of BiLD to other rejection sampling based speculative sampling methods proposed in 338 4 on IWSLT and XSUM. For BiLD, we include two BiLD configurations: the one that matches latency and the other that matches BLEU/ROUGE-L scores as compared to the rejection sampling based methods. Note that BiLD consistently outperforms other methods by achieving either (1) improved BLEU/ROUGE-L scores with equivalent latency gains, or (2) improved latency gains while retaining the same performance score.
IWSLT Speedup
XSUM
Dataset
ROUGE-I.
Speedup
Vanilla Inference
40.32
35.08
Rejection Sampling Based [
39.93
1.28
35.00
1.25x
(Match Latency)
1.23
35.30
Bil.D (Match BLEUROUGE-L)
39.87
1.49
34.96
1.50
Table 4: Comparison of the percentage of fallback and rollback (rejection) occurrences of BiLD and other rejection sampling based speculative sampling methods [38] 4] While achieving even better BLEU/ROUGE-L scores in IWSLT and XSUM, BiLD involves noticeably fewer number of fallbacks and rollbacks, resulting significantly better latency speedup.
Task
Method
BLEUROUGE-L
Speodup
% Fallback
% Rollback (rejection)
IWSLT
Rejection Sampling Based BiLD (Better BLEUU
38
1,28x
23.24%
40.33
1.43x
21.09%
1.56%
XSUM
Rejection Sampling Based BiLD (Better ROUGE-L)
38
35.00
1.25x
36.84% 32.33%
24.24%
35.12
1.48x
6.41 %
7.3.3 Insights Better Latency and Performance
Quantitative analysis reveals consistent trend where our method, when compared other speculative decoding frameworks, effectively enhances both text generation quality and latency. We provide insights and explanations into why our approach surpasses speculative decoding frameworks.
(1) Better text generation quality
Ensembling effect: The power of blending outputs from multiple models has been well-explored in various fields. This is also the case in open-source LLM models which exhibit diverse strengths and weaknesses due to variations in data, architectures, and hyperparameters, making different models complementary each other (27]. In fact, show such effects of blending multiple model outputs in Figure 2, where a combination of 20% of the large model's prediction with the small model's prediction outperforms the exact imitation of the large model's behavior. Our approach offers fallback and rollback policies that efficiently exploit optimal ensemble point, which produces superior output quality compared to the unbiased estimate of the large model in 38
Rollback policy that leads to higher performance: Our method adopts a rejection policy that completely discards the small model's prediction if it significantly deviates from the large model's s counterpart, based on cross entropy-based distance metric. This contrasts with speculative decoding, where the decision involves stochastic rejection sampling process. We empirically observe that BiL.D 's hard rejection policy allows a better BI B..FUOOOGGGLL score with significantly fewer number of rollbacks (rejections) than the stocastic rejection policy of speculative decoding described in Table
4 We hypothesize that this boost in predictive performance stems from our hard rollback policy, which prevents potentially erroncous predictions by ruling out stochasticity. We additionally hypothesize that such a strategy can address exposure bias, mitigating the impact of a single early- stage misprediction on subsequent predictions.
(2) Lower end-to-end latency Furthermore, our fallback policy introduces a dynamic fallback window size (i.e. number of small model's consecutive iterations) that is determined based on the run-time prediction confidence of the small model. This is contrast with speculative decoding which adopts a static window size. The advantages of the dynamic window size are two-fold:
18

---

Table 5: BiLD with nucleus sampling (p=0.8) on IWSLT and XSUM. Similar to the greedy decoding case, our method achieves a ~1.5x speedup without compromising performance and a ~ 11.8x speedup with a modest 1-point BLEU/ROUGE score reduction with sampling.
TWSLT
XSUM
Dataset
BLEU
Speedup
ROUGE-L
Speodup
Vanilla Inference
34.00
BiLD
39.72 (40.48)
1.51x
34.34 (+0.34)
1.22%
39.26 (+0.02)
1.63x
34,04 (+0.04)
1.45x
38.27 (.0.97)
1.80
33.10 (-0.90)
1
Less fallbacks: The dynamic window size enables the small model to persist in making predictions when it is confident, thus minimizing the unnecessary engagement of the large model. This is supported by Table 4 where BiLD involves fewer number fallbacks (23.24% + 21.09% and 36.84% - 32.33%) than [38 4
while achieving better performance.
Less rollbacks/rejeecionn:: The dynamic window size further enables preemption of the small model when it is uncertain, which avoids rollback of the small model's wrong predictions. This is also supported by Table 4 where BiLD involves significantly fewer number of rollbacks (9.81% % + 1.56% and 24.24% - 6.41%) than [3 4
while achieving better performance.
Minimizing both fallbacks and rollbacks/rejections reduces unnecessary computation which directly translates to end-to-end latency improvement.
7.4 BiLD with Sampling
Our approach isn't 't restricted greedy decoding, but it can seamlessly extend sampling methods. The only modification is to perform random sampling instead of greedy sampling when drawing a token from both the small model and the large model while using the same fallback and rollback policy. This is because both the fallback and rollback policies, based on the maximum prediction probability, serve an effective indicator of the small model's uncertainty in prediction and potential inaccuracies, regardless of the sampling method. The following table illustrates the latency versus performance trade-off of the sampling-based based approach, specifically using nucleus sampling with p=0.8, similar to (6) This evaluation follows the same environment as other experiments outlined in the paper, and both cases involve aligned small models.
Table 5] exhibits the BLEU/ROUGE-L score of BiL.D on the IWSLT and XSUM benchmarks as well their relative speedup. As can be seen in the table, our method exhibits similar trend to the greedy decoding case. It achieves a ~~ 11.5x speedup without compromising performance and a 1.88 speedup with modest 1-point BLEU/OUGE score reduction.
7.5
Additional Analysis
7.5.1 Model Analysis of BiLD: FLOPs, MOPs, and Arithmetic Intensity
Figure g compares average FLOPs, MOPs (memory operations), arithmetic intensity, and the latency speedup of the vanilla inference and BiLD on the CNN/DailyMail benchmarks. For BiLD, use the model with roughly the same ROUGE-L score as the vanilla inference, and all the numbers normalized by the numbers of the vanilla inference. The figure illustrates that BiLD exhibits slightly higher FL.PPs compared to the vanilla inference. This is due to the fact that the autoregressive and autorrrrseseeiv exccutions have the same amount of FLOPs, and BiLD involves additional overhead of rumning the small model alongside. However, in the case of MOPs, BiLD demonstrates a significant ~5x reduction of memory operations. This can be attributed to the capability of BiLD to process multiple tokens with single weight load, thereby enhancing token-level parallelism and maximizing data reuse. In contrast, this is not the case in the vanilla inference where a single weight load can only process single token. Consequently, BiLD achieves significantly higher arithmetic intensity, which approximately times larger than the vanilla inference. Arithmetic intensity 772) measures the number of arithmetic operations that can be performed per memory operation. Given that memory operations can contribute more to the overall inference latency than arithmetic operations in many Transformer decoding scenarios [32], decreasing memory operations
19

---

Model Analysis Vanilla Inference vs BiLD
BILD
ded Values
Normalizzed
3-
1.85
FLOPs(4)
MOPs(4)
Arithmetic
Speedupl
Intersilyy f)
Figure FLOPs, MOPs (memory operations), arithmetic intensity, and latency speedup comparison of vanilla inference and BiLD on the CNN/DailyMaii benchmark. BiLD approach results in remarkable reduction MOPs due the improved token-level parallelism, resulting in significantly higher arithmetic intensity.
Ground Truth
And Siftables example of ecosystem of tools for manipulating digital information.
Large
And the siftables are an example of new generation of manipulation tools for digital data.
And you look at the ifleses are example of new generation technologies manipulation digital data.
BiLD
And Siftables example
of manipulation of digital data.
Ground Truth
Which is great, because the Romans didd not actually think that genius was particularly clever individual.
Large
great. Romans really think genius particularly smart individual.
Small
That's The tube didn'tt really think that genius particularly lonely individual.
BiLD (ours)
That's great. The Romans didn't really think that genius was particularly smart individual.
Ground Truth
particles then released from the cells and
killed coll.
Large
The viral particles then were released by the cells and came back and killed coll.
Small
The viral particles were then released by the cells and came back and killed Coke.
BiLD (ours)
The viral particles then were released by the cells and came back and killed E. coli.
Figure 9: Example text sequences that BiLD generates with the validation set of IWSLT 2017 De-En, compared the ground truths and the outputs of the large and small baselines. For BiLD, tokens generated by the large model are highlighted in red, while all the other tokens are generated by the small model. This illustrates that with small engagement of the large model, BiLD correct not only inaccurate vocabulary but also wrong semantics the text that the small model would have otherwise generated.
and increasing arithmetic intensity can effectively alleviate the inference bottleneck. This leads to an overall latency speedup of 1.85x actual hardware.
7.5.2 Examples of Generated Sequences
Figure 9 provides examples of text sequences generated by BiLD on the validation of IWSLT 2017 De-En, along with the ground truths (i.c., labels) and outputs of the pure large and small baseline models. The tokens generated from the large model of BiLD are highlighted in green, while all the other tokens are generated by the small model. The results illustrate that the small model often produces low-quality texts, by predicting inaccurate tokens which can alter the meaning of the entire sentence. To contrast, it is observed from the examples that BiLD is able to improve the text generation quality by letting the large model interrupt when the small model generates incorrect tokens. Particularly, in the examples provided, BiLD tends to be as strong as the large model at predicting terminologies. Overall, the large model's engagement BiL.D decoding not only improves the prediction accuracy but also prevents incorrect predictions from impacting the future ones.
20

---

XSUM, T5
CNNiDaailyMail.
ROUGE: (higher better)
34.75
41.6
ROUGE- (higher
34.50
34.25
ballac Brasheic1.
33.75
becs Braihci22.
back 0rasnel33
41.0
tabask Sranse
33.50
becs treiniice..
0.35
0.60
0.65
0.70
0.475 0.500 0.525 0.550 0.575 0.600
Normallzed avg latency per example
Normalized avg latency per example
The trade-off between latency and generation quality (ROUGE-L) for the aligned BiLD model on two summarization tasks: (Left) XSUM and (Right) CNN/DailyMail. Each curve represents a different rollback threshold, with smaller thresholds indicating more rollbacks. The tradc-off can be further obtained within each curve with different fallback thresholds, where larger scatter sizes indicate larger fallback thresholds. A larger fallback threshold implies more fallbacks.
Figure 10:
7.5.3 Impact of Fallback and Rollback on Performance
We have explored how the BiLD framework can achieve different trade-offs between latency and generation quality by adjusting fallback and rollback thresholds. In this section, present detailed analysis of how these thresholds affect the performance using the aligned BiLD model on two different summarization tasks, XSUM and CNN/DailyMaill, as illustrated in Figure T0] Different curves in the plot represent different rollback thresholds, and each scatter point within the curve represents different fallback thresholds. Note that a small rollback threshold implies while a larger fallback threshold implies more fallback.
rollback,
We observe general trend where smaller rollback thresholds (i.e., more rollbacks) result in better generation quality but longer latency. This trend is expected because, with more rollback, we preempt more small model's predictions that be potentially inaccurate by sacrificing the latency. Similarly, there is also a general trend that smaller fallback thresholds (i.e., fewer fallbacks) result in faster latency but worse generation quality. However, observed that lowering the fallback rates beyond a certain point can actually hurt both the latency and generation quality. This is because inaccurate predictions that the small model should have fallen back later rolled back, incurring extra 'flush' cost for the tokens that follow.
21


Reviews:
[{"review_id": "QRvR1pgNTr", "reviewer": "Reviewer_NSPU", "summary": "The paper proposes Big Little Decoder (BiLD), a framework to accelerate autoregressive text generation by using a small decoder model to generate most tokens efficiently and invoking a large decoder model occasionally to refine inaccurate predictions. BiLD manually designs two policies to invoke refinement, 1) a fallback policy (based on the small model's uncertainty) to hand control from the small to large model when uncertain, and 2) a rollback policy (based on the distance of two models' probability) for the large model to correct inaccurate small model predictions. Experiments on machine translation and summarization tasks show BiLD can achieve up to 2x speedup on a GPU with minimal quality loss.", "questions": "Questions:\n\n1. In Figure 2, why does 20% achieve better performance than 40 or 80 percent?\n\nTypos:\n\nline 142: regressive ==> autoregressive", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The idea of combining small and large models is intuitive and aligns with the observation that small models can achieve comparable quality if a few inaccurate predictions are corrected.\n\n2. The fallback and rollback policies are simple yet effective mechanisms to coordinate the models and empirically show better performance than speculative sampling.\n\n3. The framework is general and does not require additional training, while the prediction alignment technique further improves BiLD's performance with minimal extra effort.", "weaknesses": "1. A few very similar prior works are not cited and discussed [2-3], i.e., the idea of using a fast-inference model to generate tokens quickly and refining with a large model.\n\n2. The paper criticizes [1] as \"suggests an unbiased estimator that recovers the exact probability distributions of the more powerful model, it can still have a large variance that may impact current and future predictions\". But the proposed BiLD approach itself even does not have any guarantee of unbiased or low-biased recovery of the large model's probability. It seems that BiLD can only be used for greedy decoding, rather than arbitrary sampling from the large model's distribution, which is quite important in a few downstream tasks [6].\n\n3.  The policies require tuning fallback and rollback thresholds, adding hyperparameters. Optimal values likely vary across tasks and models.\n\n4. While sharing several common techniques to the non-autoregressive transformer (NAT) models, the literature on NAT is not discussed. For example, the prediction alignment technique is very similar to knowledge distillation adopted to train NAT models [4].\n\n5. The proposed method is only evaluated on the machine translation and text summarization datasets, which have been shown to not that hard for small or non-autoregressive sequence models [4-5]. More challenging datasets, such as HumanEval used in [1], should be evaluated.\n\n6. A minor problem is that there are limited insights into what types of predictions the large model tends to correct.\n\n[1] Accelerating Large Language Model Decoding with Speculative Sampling\n\n[2] Lossless Acceleration for Seq2seq Generation with Aggressive Decoding\n\n[3] Blockwise Parallel Decoding for Deep Autoregressive Models\n\n[4] Levenshtein Transformer\n\n[5] Deep Encoder, Shallow Decoder: Reevaluating Non-autoregressive Machine Translation\n\n[6] Self-Consistency Improves Chain of Thought Reasoning in Language Models", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### W1. BiLD with Random Sampling\n\nOur approach **isn't restricted to greedy decoding, but it seamlessly extends to sampling methods.** Even though the fallback policy is based on the maximum prediction probability of a small model, we can still perform decoding based on random sampling. The only modification is to perform random sampling instead of greedy sampling when drawing a token from both the small model and the large model. We keep the same fallback and rollback policy. This is because both the fallback and rollback policies, based on the maximum prediction probability, serve as an effective indicator of the small model's uncertainty in prediction and potential inaccuracies, regardless of the sampling method. The following table illustrates the latency versus performance trade-off of the sampling-based approach, specifically using nucleus sampling with p=0.8, similar to [1]. This evaluation follows the same environment as other experiments outlined in the paper, and both cases involve aligned small models.\n\nWithin the sampling context, our method **exhibits a similar trend to the greedy decoding case**. It achieves a ~1.5x speedup without compromising performance and a 1.8x speedup with a modest 1-point BLEU/ROUGE score reduction. For a better comparison, we include the performance/latency results of speculative decoding under the same sampling scheme for comparative analysis. We will add this point in the revised version.\n\n**IWSLT**\n\n| Method | BLEU  | Norm. Latency |\n| -- | -- | -- |\n| Baseline | 39.24 | 1x  |\n| BiLD | 39.72 (+0.48) | 1.51x |\n|  | 39.26 (+0.02) | 1.63x  |\n|  | 38.27 (-0.97) | 1.80x |\n| Speculative Decoding | 39.48 (+0.24) | 1.30x |\n\n**XSum**\n\n| Method | ROUGE-L | Norm. Latency |\n| -- | -- | -- | \n| Baseline | 34.00 | 1x |\n| BiLD  | 34.34 (+0.34) | 1.22x |\n|   | 34.04 (+0.04) | 1.45x  |\n|  | 33.10 (-0.90) | 1.85x  |\n| Speculative Decoding | 32.65 (-1.35) | 1.23x  |\n\n\n### W2. Hyperparameters\n\nAlthough optimal values can differ among models and tasks, they generally lie within the restricted tange across various scenarios. Notably, the **hyperparameter search has a very small overhead, as it is an evaluation-time hyperparameter** rather than a training-time one. For instance, in IWSLT with <1k val examples, the full exploration across the hyperparameter space in the paper in under 5 GPU hours. This is similar to the larger dataset’s case like XSum, whose hyperparameters can be swiftly determined using a subset of the evaluation dataset (1k). Despite the small size, it still maintains a high correlation (R^2 > 0.92) in both latency and ROUGE to the full dataset counterpart.\n\n\n### W3. Comparison to NAT\n\nWe acknowledge the value of NAT techniques and the shared objective of BiLD in improving token-level parallelism. In Section 2.1., we have elaborated on the comparison between BiLD and well-known NAT models including Levenshtein Transformer as follows:\n\nBiLD is different from NAT frameworks in that it eliminates the need for (1) from-scratch training and (2) adjustments to the training and inference processes, which is the core design principle of BiLD. This contrasts with the typical requirements of most NAT frameworks. These two attributes are becoming increasingly significant as more powerful pre-trained models become open-sourced, while the costs associated with training models from scratch become restrictive across various applications. As a result, only methods that circumvent the (1) need for from-scratch training and (2) alterations to the training pipeline can fully harness the benefits of pre-trained checkpoints, which is a feature not inherently straightforward within NAT.\n\n### W4. HumanEval\n\nWe appreciate the request to see further evaluation of our work. However, we believe that the datasets already used in this work are sufficiently challenging and fair benchmarks for assessing the effectiveness and generalizability of our framework. In fact, the small model’s performance on IWSLT/XSum are 37.47/31.23, which is significantly lower than the large model’s counterpart that are 40.33/35.10. Qualitatively, Figure 9 (Appendix 6.3.2) also visually illustrates the small model's difficulties in solving machine translation problems. As a result, these datasets have been widely adopted in previous studies [1, 2, 3] for evaluating similar frameworks. \n\nWe do appreciate and agree with the opinion that HumanEval would be a even more challenging benchmark for the smaller models. Should we extend BiLD to target HumanEval, the small model will need to be enlarged in model capacity to achieve reasonable predictive performance on the new task, and in turn, the large model would be made even larger to serve as a fallback option. Such an extension is essentially a direct translation of the proposed framework to a new set of small and large models. In our work, we extensively evaluated BiLD against CALM (Figure 7, Appendix 6.2.2) and speculative decoding (Section 4.2, Global rebuttal) over four benchmarks, showing that the proposed framework can achieve substantial performance-latency improvements in the general context of collaborative use of differently sized models. Therefore, our results already demonstrate the efficacy of BiLD across a diverse set of NLP tasks, and future adapations of this work could consider coordination between larger “small” models and billion-sized “large” models to target emerging LLM tasks. \n\n### W5. Comparison to Speculative Decoding\n\nWe acknowledge the comment that our paper does not include a guarantee of an unbiased estimate. In the [global rebuttal](https://openreview.net/forum?id=EfMyf9MC3t&noteId=laUzvEbREz), however, we provided an empirical observation that BiLD can outperform [1, 2] in various scenarios, along with intuitive explanations on this. \n\n[1] Accelerating Large Language Model Decoding with Speculative Sampling\n\n[2] Fast Inference from Transformers via Speculative Decoding\n\n[3] Confident Adaptive Language Modeling"}, {"review_id": "m1lUVziC6Y", "reviewer": "Reviewer_mW2r", "summary": "This paper presents an efficient generation approach that utilizes a smaller model during decoding to reduce computational requirements. In essence, the large model is only employed when the small model's generation exhibits low confidence (fallback policy) or when there is a significant distribution difference between the previously generated tokens of the small and large models (rollback policy). The authors demonstrate that this approach achieves a speedup of up to 2X for downstream tasks using FLAN-T5 models. However, a potential limitation of this approach is its scalability issue when the batch size exceeds 1, which may impede its practical applicability. Additionally, there is a lack of discussion regarding previous works proposing similar ideas, which would provide valuable context and comparisons.", "questions": "- How does the method generalize to models of different scales? Given a large model, how to most effectively select a corresponding small model? A discussion of scales would be very helpful.\n- The main results presented in Table 1 and Figure 4 lack benchmarking for the small model's performance and speedup. It would be helpful to compare the performance of the small model against its larger counterpart to provide a baseline for understanding the achieved improvements.\n- The evaluation of the study focuses on in-domain tasks. Would the method provide advantages for out-of-domain generalizations? For examples, tasks that are not trained with FLAN, e.g., GSM8K?\n- The fallback policy involves regeneration for (n-m) tokens and appears to be computationally expensive. It would be insightful to provide a more detailed analysis, such as the percentage of occurrences of rollback and fallback events. Similarly, an analysis indicating the percentage of predictions made by the small model versus the large model would be insightful as well.", "limitations": "The authors did not discuss the limitations of the work explicitly. However, I believe that the work appears to have minimal negative societal impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well written, presenting a clear and easily understandable explanation of the proposed method, which is both straightforward and effective. As generation inference continues to gain importance, this paper offers a valuable solution to enhance performance in that regard.\n- The paper effectively illustrates the tradeoff between efficiency and performance through the figures presented. The visual representations provide a clear understanding of how the proposed method strikes a balance between the two aspects.", "weaknesses": "- **Batch size:** The biggest limitation of the approach is that it probably only works for a batch size of 1 (the authors only tested this setting as well). Extending the approach to handle multiple examples within a batch poses challenges due to the significant variation in low-confidence positions across the generations for different examples. Consequently, it becomes difficult to envision the practical usability of the approach, particularly considering the importance of batching for efficient hardware utilization. It is crucial for the authors to conduct experiments to ascertain whether the method retains its advantages when the batch size exceeds 1 and provide a more comprehensive discussion on this aspect.\n- **Novelty and comparison to previous work:** it is important to note that the core idea of this work bears significant resemblance to existing work [1]. Therefore, a quantitative comparison and discussion concerning the relationship between this work and the aforementioned existing work is necessary.\n\n[1] Fast Inference from Transformers via Speculative Decoding, Leviathan et al., 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### W1. Batch sizes\n\nWhile we only evaluate the use case of batch size 1, it is **not a limitation and it is straightforward in modern LLM serving systems to extend the BiLD framework to multi-batch applications**. This is because the modern LLM serving systems (especially for multi-batch applications) such as HuggingFace TGI are equipped with “iteration-level batching” [1, 2]. The main idea is to improve the throughput of LLM inference by removing “bubbles” when batching multiple examples with different lengths, and this is critical in many applications as different requests in a batch could have different numbers of prompt tokens and decoding iterations, resulting in some requests finishing earlier than the others. Iteration-level batching prevents this by storing the decoding jobs in a request pool in the iteration granularity and dynamically batching multiple examples on different decoding iterations together in every decoding iteration.\n\nTherefore, BiLD (as well as other speculative decoding frameworks) can be supported seamlessly on top of the existing inference system, with the only difference of having two separate request pools one for the large model, and the other for the small model. For example, even though one example A falls back at its 2nd step and the other B falls back at its 3rd step due to different low-confidence positions, the query will be handled as follows.\n\n- 1st and 2nd iteration: A and B are batched together from the small model’s request pool\n- 3rd iteration:\n    - A will be registered to the large model’s request pool, from which it will be batched with other jobs in the same request pool for its 3rd decoding iteration.\n    - B is still at the small model’s request pool and can be batched together with other jobs at the same request pool if there are any. \n    - Note that within the iteration-level batching system, neither A nor B needs to be at the same decoding step with the other examples that are batched together\n\n### W2. Quantitative Comparison to Speculative Decoding\n\nIn the table below, we present a comprehensive quantitative comparison between our method and speculative decoding [3, 4] across two different datasets: IWSLT for machine translation and XSum for summarization. We will include this comparison in the revised paper. In order to ensure a fair comparison that isolates the impact of the frameworks themselves, we employ the baseline (non-aligned) small model for both BiLD and speculative decoding. We maintained the same evaluation setup and hyperparameter space that are outlined in Appendix 6.1.\n\nThe table below shows two BiLD configurations across two different tasks and two different sampling methods of greedy sampling and nucleus sampling with p=0.8, and temperature of T=1.0 as in [4]. In all cases, the first row of BiLD matches latency, and the second row matches BLEU/ROUGE-L scores as speculative decoding. Across all experiments, BiLD consistently outperforms speculative decoding. It achieves either **(1) notably improved BLEU/ROUGE scores of up to an additional 2 points with equivalent latency gains, or (2) superior latency gains of up to +0.4x while retaining the same score.**\nWe will further elaborate on this point in the revised version.\n\n**IWSLT**\n\n| Sampling | Method  | BLEU  | Norm. Latency |\n| -- | -- | -- | -- |\n| Greedy   | Baseline  | 40.32 | 1x   |\n| | Speculative | 39.93 | 1.28x  |\n|  | BiLD (match latency) | 40.54 | 1.23x  |\n|  | BiLD (match BLEU) | 39.87 | 1.49x |\n| Sampling | Baseline | 39.24 | 1x |\n|  | Speculative | 38.79 | 1.36x |\n|  | BiLD (match latency) | 39.56 | 1.35x |\n|  | BiLD (match BLEU) | 38.88 | 1.49x |\n\n**XSum**\n\n| Sampling | Method  | ROUGE-L | Norm. Latency |\n| -- | -- | -- | -- |\n| Greedy | Baseline | 35.08 | 1x |\n|  | Speculative | 35.00 | 1.25x  |\n|  | BiLD (match latency) | 35.30  | 1.42x  |\n|  | BiLD (match ROGUE) | 34.96 | 1.50x  |\n| Sampling | Baseline | 34.02   | 1x  |\n|  | Speculative | 32.32 | 1.23x |\n|  | BiLD (match latency) | 34.32   | 1.22x |\n| | BiLD (match ROGUE) | 32.81 | 1.61x |\n\n### Q1: How to most effectively select a small model\n\nWe have provided a detailed discussion in [Q2 of Reviewer asbG](https://openreview.net/forum?id=EfMyf9MC3t&noteId=nrwJEv4oRT).\n\n### Q2. Performance of the small model\n\nWe have provided the small model's performance in [Q1 of Reviewer DxpA](https://openreview.net/forum?id=EfMyf9MC3t&noteId=8Hcp16yZeo).\n\n### Q3. More Statistics\n\nWe present an in-depth comparison of rollback and fallback event percentages in two different datasets, using the unaligned small model. Please note that % fallbacks indicates % large model.\n\n- With a fallback engagement ranging from 20% to 30% (i.e. 20-30% engagement of the large model), BiLD attains performance on par with the baseline, which resonates with our observations in Figure 2.\n- The occurrence of rollback in the BiLD framework is infrequent and is ~2% for IWSLT and ~6% for XSum. Consequently, the overhead incurred from regenerating (n-m) tokens remains small. While a similar regeneration mechanism exists in the speculative decoding framework [3, 4], our approach demonstrates notably lower rollback occurrences. In the speculative decoding, the % rollback/rejection is around 10% for IWSLT and 24% for XSum. Such a reduced rollback overhead in our framework directly translates to enhanced speedup.\n\n| Task  | Method  | BLEU/ROUGE | Latency | % Fallback | % Rollback |\n| -- | -- | -- | -- | -- | -- |\n| IWSLT | Speculative  | 39.93 | 1.28x | 23.24%  | 9.81%  |\n| | BiLD  | 40.33 | 1.43x | 21.09%  | 1.56% |\n| XSum  | Speculative | 35.00      | 1.25x | 36.84% | 24.24% |\n|  | BiLD  | 35.12  | 1.48x | 32.33% | 6.41% |\n\n[1] Yu et al. Orca: Progressive Learning from Complex Explanation Traces of GPT-4\n\n[2] Kwon et al. vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention\n\n[3] Leviathan et al. Fast Inference from Transformers via Speculative Decoding\n\n[4] Chen et al. Accelerating Large Language Model Decoding with Speculative Sampling\n"}, {"review_id": "kSxH2RTbXb", "reviewer": "Reviewer_wHu9", "summary": "The paper presents an approach to improve decoding using a smaller model to create drafts, and fallback to a larger model whenever the smaller model predicts any token with probability smaller than a fixed threshold. Additionally, during this fallback, if the larger model produces a token that is not in sync with the prediction from the smaller model, the predictions are rolled back to the point until the predictions are in sync and then prediction continues using the predicted token from the larger model.\n\nGiven that the larger model can use causal attention masking and predict tokens in parallel using the predictions from the smaller model as hints for teacher forcing, the whole process becomes latency efficiency (at the cost of flops).\n\nThe method is inspired from `Accelerating Large Language Model Decoding with Speculative Sampling`, which too uses two models. The small one for creating draft predictions using non-autoregressive decoding, and the second (larger) model for iterating over the predictions from the smaller model using the larger model.", "questions": "As mentioned in the weakness section, what are your thoughts on why the model can perform superior compared to Speculative sampling?", "limitations": "No limitations that I can think of, and none discussed in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The method is an improvement over speculative sampling in terms of latency by not just using the small model for the initial prediction, but using a procedure that keeps shifting control back and forth between the smaller and the larger model.\n\nExperiments show that this significantly improves the speedup compared to speculative sampling.", "weaknesses": "It is unclear how this can achieve quality better than speculative sampling.\n\nSpeculative sampling just uses the small model as draft. The draft is refined by the larger model, but the predictions are always going to be the same as that of the large model.\n\nHowever, that is not the case with this paper. In this paper, all the predictions may come from the small model if the model never produces a token with low probability.\n\nIf that is the case, I would expect the methodology in this paper to only outperform Speculative sampling in terms of latency. Quality should be at par or below the larger model, but seems like that is not the case.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n### W1. Insights on Better Latency and Performance\n\nQuantitative analysis reveals a consistent trend where our method, when compared to speculative decoding [1, 2], effectively enhances both text generation quality and latency. We provide insights and explanations into why our approach surpasses speculative decoding frameworks.\n\n\n- **(1) Better text generation quality:** \n    - (1-1) **Ensembling effect:** The power of blending outputs from multiple models has been well-explored in various fields. This is also the case in open-source LLM models which exhibit diverse strengths and weaknesses due to variations in data, architectures, and hyperparameters, making different models complementary to each other [Ref: LLM-Blender]. In fact, we show such effects of blending multiple model outputs in Figure 2, where a combination of 20% of the large model's prediction with the small model's prediction outperforms the exact imitation of the large model's behavior. Our approach offers fallback and rollback policies that efficiently exploit optimal ensemble point, which produces superior output quality compared to the unbiased estimate of the large model in [1, 2].\n    - (1-2) **Rejection policy that leads to higher predictive performance:** Our method adopts a rejection policy that *completely* discards the small model's prediction if it significantly deviates from the large model's counterpart, based on cross entropy-based distance metric. This contrasts with speculative decoding, where the decision involves a stocastic rejection sampling process. We empirically observe that BiLD’s “hard rejection policy” allows better BLEU/ROUGE score with significantly fewer number of rejections (rollbacks) than the “stocastic rejection policy” of speculative decoding as described in the table below. We hypothesize that this boost in predictive performance stems from our hard rollback policy, which prevents potentially erroneous predictions by ruling out stochasticity. We additionally hypothesize that such a strategy can addresses exposure bias, mitigating the impact of a single early-stage misprediction on subsequent predictions.\n- **(2) Lower end-to-end latency:**\n    - Furthermore, our fallback policy introduces a dynamic window size (i.e. number of small model's consecutive iterations) that is determined based on the run-time prediction confidence of the small model. This is in contrast with speculative decoding which adopts a static window size. The advantages of the dynamic window size is in two-fold:\n        - (2-1) **Less fallback:** The dynamic window size enables the small model to persist in making predictions when it is condifent, thus minimizing the unnecessary engagement of the large model. This is supported by the table below where BiLD involves fewer number of fallbacks (23.24 → 21.09 and 36.84 → 32.33) than speculative sampling while achieving better performance.\n        - (2-2) **Less rollback/rejection:** The dynamic window size further enables preemption of the small model when it is uncertain, which avoids rollback of the small model’s wrong predictions. This is also supported by the table below where BiLD involves significantly fewer number of rollbacks (9.81 → 1.56 and 24.24 → 6.41) than speculative sampling while achieving better performance.\n        \n        Minimizing both fallback and rollback/rejection reduces unnecessary computation which directly translates to end-to-end latency improvement.\n\n\n| Task  | Method              | BLEU/ROUGE | Latency | % Fallback | % Rollback (rejection) |\n| ----- | ------------------- | ---------- | ------- | ---------- | ---------------------- |\n| IWSLT | Speculative         | 39.93      | 1.28x   | 23.24%     | 9.81%                  |\n|       | BiLD (Better BLEU)  | **40.33**  | 1.43x   | 21.09%     | **1.56%**              |\n|       |                     |            |         |            |                        |\n| XSum  | Speculative         | 35.00      | 1.25x   | 36.84%     | 24.24%                 |\n|       | BiLD (Better ROUGE) | **35.12**  | 1.48x   | 32.33%     | **6.41%**              |\n\n\n### W2. Can all the predictions come from the small model?\n\nAlthough it is highly unlikely for the small model to generate all tokens with high probability, we prevent this occurrence by limiting the maximum consecutive predictions from the small model to 10 tokens. This is also stated in the Appendix for reference.\n\n\n[1] Chen et al. Accelerating Large Language Model Decoding with Speculative Sampling\n\n[2] Leviathan et al. Fast Inference from Transformers via Speculative Decoding"}, {"review_id": "R6FRj6Qz9v", "reviewer": "Reviewer_asbG", "summary": "The paper proposes BiLD, a framework 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. The proposed two policies allows the small model to hand control over to the large model if it is not confident enough and the large model to review and correct the small model's inaccurate predictions.\n\n", "questions": "1. How to set the thresholding to fallback? The paper briefly mentions that the small model can be over confident therefore, even with a high confidence, the prediction can be wrong. \n\n2. How large is the big model compared to the small model? Any restrictions?", "limitations": "Not addressed. ", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The proposed fallback method is novel in multiple ways: 1) use small model's generation as part of big model's input, thus transforming an autoregressive generation problem to semi-autoregressive. 2) Improves compute intensity as tokens can be generated in parallel in the big model, which could significantly reduce latency while not sacrificing model quality much. \n\n- Rollback policy mitigate the issue of small model being over confident. ", "weaknesses": "- Evaluation benchmarks are mostly on translation, not on QA tasks. The paper should evaluate more QA tasks including Triviaqa, Squad, webqs and Nqs. \n\n- The paper should compare with CALM [1] to demonstrate the effectiveness of big/small model with fallback and rollback rather than having early stopping in a per layer basis. \n\n- Inference batch size of 1 is quite limited. Will this approach work for batch size larger than one? If the batch size is larger than one, there will be a parameter loading overhead and very likely both models need to be loaded. Loading parameters can dominate the latency. Batch size of one makes serving easier for any conditional computation approaches, however, it can lead to under utilization of hardware. Could you provide some evidence of batch size larger than one?\n\n[1]: https://arxiv.org/abs/2207.07061 ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### W1. Evaluation on QA Benchmarks\n\nWe acknowledge the comment that our paper lacks evaluations on QA tasks. \nHowever, we highlight that **QnA tasks are not the ideal benchmarks for assessing our methodology's effectiveness**. This is because the main focus of our framework lies in enhancing *decoding processes* to efficiently generate lengthy sequences. In contrast, QnA benchmarks, often involving multiple-choice or short-answer questions (Triviaqa, Squad, Webqs and Nqs), mostly entail brief decoding lengths. \nTherefore, these benchmarks primarily evaluate the ability of a given language model to encode and understand the question prompt, rather than its ability to generate a response, making them unsuitable for evaluating decoding optimization methodologies like BiLD. \n\n\n### W2. Comparison to CALM\n\nFigure 7 (Appendix 6.2.2) contains the information this comment is requesting for. We directly compared CALM applied to the mT5-small model with BiLD applied to the same model. The first layer of the pre-trained BiLD model serves as our small model, but this choice is one of many design points and other choices are viable as well. Choosing the first layer of the large model as the small model showcases the framework's extensibility when only a single model is accessible while preserving the core idea of using both big and small models.\n\n### W3. Multi-batch Inference\n\nWhile we only evaluate the use case of batch size 1, it is **not a limitation and it is straightforward in modern LLM serving systems to extend the framework to multi-batch applications**. This is because the modern LLM serving systems (especially for multi-batch applications) such as HuggingFace TGI are equipped with “iteration-level batching” [1, 2]. The main idea is to improve the throughput of LLM inference by removing “bubbles” when batching multiple examples with different lengths, and this is critical in many applications as different requests in a batch could have different numbers of prompt tokens and decoding iterations, resulting in some requests finishing earlier than the others. Iteration-level batching prevents this by storing the decoding jobs in a request pool in the iteration granularity and dynamically batching multiple examples on different decoding iterations together in every decoding iteration.\n\nTherefore, BiLD (as well as other speculative decoding frameworks) can be supported seamlessly on top of the existing inference system, with the only difference of having two separate request pools one for the large model, and the other for the small model. For example, even though one example A falls back at its 2nd step and the other B falls back at its 3rd step due to different low-confidence positions, the query will be handled as follows.\n\n- 1st and 2nd iteration: A and B are batched together from the small model’s request pool\n- 3rd iteration:\n    - A will be registered to the large model’s request pool, from which it will be batched with other jobs in the same request pool for its 3rd decoding iteration.\n    - B is still at the small model’s request pool and can be batched together with other jobs at the same request pool if there are any. \n    - Note that within the iteration-level batching system, neither A nor B needs to be at the same decoding step with the other examples that are batched together\n\n\n\n### Q1. How to set the thresholding to fallback and Over-confidence of the small model\n\nThe fallback threshold is explored within the hyperparameter space outlined in Appendix 7.1.2. It is possible that the small model to be overconfident and the fallback threshold might not always detect such occurrences; however, this doesn't impact text generation quality as such occurrences are fixed by the large model during the rollback procedure.\n\n### Q2. How large is the big model compared to the small model?\n\nThis is indeed a very interesting question, and finding the right-scaled small model given a large model with the optimal performance tradeoff can be an open-ended research direction. \nThis is indeed a very interesting question, and identifying the appropriately scaled small model to a given large model for the optimal performance tradeoff can lead to an open-ended avenue of research.\n\nBelow, we present a comparison table that combines (non-aligned) mT5-base and mT5-small to mT5-large. While the base model enhances the lower threshold of attainable text generation performance, the considerable speedup it offers offsets this gain, leading to a less favorable performance-latency trade-off. Across our experiments, we recognize the significance of maintaining the small model's latency at a sufficiently low level to achieve a favorable balance. We have observed that setting the small model size to approximately 10 times smaller than the large model size serves as a reasonable guideline in this regard. However, delving into a systematic analysis of scaling rules and developing decision-making policies would be an interesting future research direction.\n\n\n|                      | BLEU  | Latency |\n| -------------------- | ----- | ------- |\n| Large                | 40.32 | 1x      |\n| Base                 | 38.13 | 1.9x    |\n| Small                | 36.61 | 2.7x    |\n|                      |       |         |\n| Base + Large         | 40.30 | 1.19x   |\n| **Small + Large (BiLD)** | **40.33** | **1.43x**   |\n\n\n\n[1] Yu et al. Orca: Progressive Learning from Complex Explanation Traces of GPT-4\n\n[2] Kwon et al. vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention"}, {"review_id": "oIUMEEpPD8", "reviewer": "Reviewer_DxpA", "summary": "This paper presents a new adaptive computation method to be used for text generation tasks. The method works by applying a small model, and if needed, revising its predictions using a larger, more expensive model. The resulting model obtains a far better speed-accuracy curve compared to the baseline model.\n\nTo put this work in a bit of context, this paper follows a rich line of work on adaptive computation in NLP. Most of these works explored the idea of early exiting. Many of them applied this approach to classification tasks (see [1-5] below for a few examples), but some also to generation tasks (e.g., [6] cited in this work, but also [7]). A much smaller number of works used a different implementation of this idea, by considering several models of different sizes. To the best of my knowledge, this has only been applied to classification tasks (e.g., [8-9]). A very recent work also compared these two approaches [10], also in a classification setup. \n\nThis work is the first work I’ve seen that applies the idea of combining multiple-models to generation tasks, and in this sense it is novel and interesting. The authors build on the standard adaptive computation framework (using confidence threshold to decide whether to use the larger model or not), and apply several other tricks to improve it: (a) using the large model in a non-autoregressive manner, to allow parallelization; (b) allowing the model to rollback “bad” generations; and (c) aligning the predictions of both models for better consistency between the two models. These ideas are interesting and a non-trivial contribution, but also hint that this approach requires some engineering to work (unlike the classification solutions, which generally work without much effort), and raise the question of how generalizable the proposed solution is. All-in-all, this is a strong paper that should be accepted to NeruIPS. \n\n[1] DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference - ACL Anthology\n[2] [2004.02178] FastBERT: a Self-distilling BERT with Adaptive Inference Time\n[3] [2004.07453] The Right Tool for the Job: Matching Model and Instance Complexities\n[4] BERxiT: Early Exiting for BERT with Better Fine-Tuning and Extension to Regression - ACL Anthology\n[5] [2104.08803] Consistent Accelerated Inference via Confident Adaptive Transformers\n[6] [2207.07061] Confident Adaptive Language Modeling\n[7] [1910.10073] Depth-Adaptive Transformer\n[8] [2204.06271] TangoBERT: Reducing Inference Cost by using Cascaded Architecture\n[9] [2012.14682] CascadeBERT: Accelerating Inference of Pre-trained Language Models via Calibrated Complete Models Cascade\n[10] https://arxiv.org/abs/2306.02307\n", "questions": "- Fig2: (a) what is the baseline (large) model performance? (b) why are we seeing a drop around 10-20%?\n- Did the authors consider experiments with models from different families? Would this make the alignment problem harder?\n- What is the performance of the small model in figure 4? This seems like an important baseline\n", "limitations": "N/A", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- An interesting and important problem\n- A natural extension of a rich line of work (see discussion above)\n- Impressive solutions to the challenges raised by this setup (see a-c above)\n- Impressive results\n- Interesting analysis and ablation experiments, hinting that all components contribute to the success of the approach.\n- The paper is overall very clearly written, though see some comments below.\n\n", "weaknesses": "- It seems much engineering is required to make this method work as well. This is not such a strong critique, as the authors show it works for four different tasks, but still somewhat of a limitation.\n- I am not entirely sure how the authors convert the decoder model to a non-autoregressive model. Do they make multiple decoding operations simultaneously? Or perhaps replace the decoder with an encoder? This feels like an important part of the construction, and requires some discussion and explanation. How much does it contribute to the performance gain?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### W1. Elaboration on non-autoregressive use of the model\n\nWe can use the same decoder model, but rather in a non-autoregressive (parallel) fashion by providing the model with multiple tokens and processing them simultaneously.\nThis mechanism is identical to how decoder models can simultaneously handle multiple tokens in a prompt, without necessitating architectural modifications.\nBy using an upper diagonal attention mask, tokens are restricted to attending only to preceding tokens; therefore, running multiple decoding operations simultaneously maintains the same functionality as the autoregressive counterparts when processing the same tokens (hence resulting in no performance degradation) but only with a improved latency.\n\n\n###  Q1. Performance of the small model\n\nBelow, we provide a comparison of the small model to the larger counterpart and BiLD.\nWe will elaborate on this point in the revised version of our paper.\n\n\n| Task   |                    | BLEU  | Latency |\n| ------ | ------------------ | ----- | ------- |\n| IWSLT  | Large              | 40.33 | 1x      |\n|        | Small              | 37.47 | 2.61x   |\n|        | BiLD (match BLEU)  | 40.33 | 1.43x   |\n|        | BiLD (BLEU - 1)    | 39.44 | 1.58x   |\n|        |                    |       |         |\n| WMT    | Large              | 31.38 | 1x      |\n|        | Small              | 28.74 | 2.50x   |\n|        | BiLD (match BLEU)  | 31.28 | 1.34x   |\n|        | BiLD (BLEU - 1)    | 30.47 | 1.43x   |\n|        |                    |       |         |\n| XSum   | Large              | 35.10 | 1x      |\n|        | Small              | 31.23 | 4.49x   |\n|        | BiLD (match ROUGE) | 35.12 | 1.48x   |\n|        | BiLD (ROUGE- 1)    | 34.02 | 1.72x   |\n|        |                    |       |         |\n| CNN/DM | Large              | 41.54 | 1x      |\n|        | Small              | 40.43 | 4.15x   |\n|        | BiLD (match ROUGE) | 41.44 | 1.71x   |\n|        | BiLD (ROUGE- 1)    | 40.57 | 2.05x   |\n\n\n### Q2. Elaboration on Figure 2\nThe large model performance is the right-most scatter (100% Large Model) on each plot, which is around 31.3 and 41.5 for WMT and CNN/DM, respectively.\nThe drop around 10-20% is due to the increased engagement of the small model in text generation, which converges to the small model's own performance on the left-most scatter (0% Large Model).\n\n\n### Q3. Experiments on different families of models\n\nAlthough we have yet to experiment with models beyond two model families, T5 and mT5, in our experiments, we believe that extending the proposed framework to different small and large model combinations is a straightforward translation given the strong generalizability across different models and tasks we presented in our paper. Investigating generalizability across varied models and scales, along with exploring the effects on model alignment, is indeed an intriguing avenue for future research."}]

Author rebuttal:
We appreciate all the reviewers for taking the time to review our work and providing us with their valuable feedback. We provided responses to the questions that each of the reviewers has commented on.
Additionally, we present a comprehensive comparison with previous works [1, 2] as a global rebuttal, which multiple reviewers had requested in common.

---

## Comparison of BiLD and Speculative Decoding

We thank the reviewers for bringing up the concurrent work [1, 2]. Given that the paper [2] is publicly known to be accepted to ICML (4/27) 2 weeks before the NeurIPS submission deadline (and their cam-ready version was only available after 5/31), we believe that the mentioned works however should be considered as concurrent and independent work as ours, and should not be the main reason for rejection:

"Note that authors are excused for not knowing about all non-refereed work (e.g. those appearing on ArXiv). Papers (whether refereed or not) appearing less than two months before the submission deadline are considered contemporaneous to NeurIPS submissions; authors are not obligated to make detailed comparisons to such papers."

Nonetheless, **we appreciate the contributions and find it useful to discuss as a concurrent work**. Following are the reviews provided to the comments provided by the reviewers. 

### 1. **Differences in methodology**

While the idea of using two models with different sizes can be deemed similar to the speculative decoding framework presented in [1, 2], we have clear distinctions in detailed methodologies. 


- (1) **Distinct Prediction Rejection Approach:** The primary difference lies in how we **decide the "rejection" of predictions from the small model**. In our rollback policy, we propose to make the rejection decision based on the ‘distance’ between the small and large model predictions, which differs from the rejection sampling policy outlined in [1, 2]. While [1, 2] proposes an unbiased estimator on the large model’s prediction, Figure 2 demonstrates that combining predictions from both models through our distance-based rejection approach can surpass the exclusive utilization of the large model’s prediction probability. BiLD seeks to find and utilize this optimal performance point without introducing much runtime cost. We have a further discussion below about how our rejection policy benefits text-generation performance.
- (2) **Dynamic Window Size:** Additionally, we introduce the dynamic window size in our fallback policy. In [1, 2], the window size remains a fixed hyperparameter; however, it is also highlighted in [2] that the window size can have a noticeable impact on end-to-end latency. Our approach offers an efficient and robust solution: adjusting the window size at runtime based on the small model's confidence level in run-time. Our ablation study (Section 4.3, Figure 5) demonstrates that omitting the fallback policy and periodically transitioning control to the large model,  as proposed in [1, 2], can result in notable latency degradation.
- (3) **Model Alignment Enhancement:** Beyond the core framework, we introduce a model alignment method to align the small model’s predictions with those of the large model. This enhances the framework by reducing unnecessary rejections and can be incorporated with minimal adjustments to the training pipeline.

### 2. **Quantitative Comparison**

Furthermore, we present in [W2 of Reviewer mW2r](https://openreview.net/forum?id=EfMyf9MC3t&noteId=sb2OQpyUG0) a comprehensive quantitative comparison between our method and speculative decoding [1, 2] across two different datasets: IWSLT for machine translation and XSum for summarization. In order to ensure a fair comparison that isolates the impact of the frameworks themselves, we employ the baseline (non-aligned) small model for both BiLD and speculative decoding. We maintained the same evaluation setup and hyperparameter space that are outlined in Appendix 6.1.

The table below shows two BiLD configurations across two different tasks and two different sampling methods of greedy sampling and nucleus sampling with p=0.8, and temperature of T=1.0 as in Chen et. al. [2]. In all cases, the first row of BiLD matches latency, and the second row matches BLEU/ROUGE-L scores as speculative decoding. Across all experiments, BiLD consistently outperforms speculative decoding. It achieves either (1) notably improved BLEU/ROUGE scores of up to an additional 2 points with equivalent latency gains, or (2) superior latency gains of up to +0.4x while retaining the same score. We will include this comparison in the revised paper.
Please refer to [W2 of Reviewer mW2r](https://openreview.net/forum?id=EfMyf9MC3t&noteId=sb2OQpyUG0) for more details.



### 3. **Insights on Better Latency and Performance**

Quantitative analysis reveals a consistent trend where our method, when compared to speculative decoding [1, 2], effectively enhances both text generation quality and latency. In [W1 of wHu9](https://openreview.net/forum?id=EfMyf9MC3t&noteId=NJBmWF2n8y), we provide insights and explanations into why our approach surpasses speculative decoding frameworks.
This can be summarized as (1) **better text generation quality** due to the ensembling effect and BiLD's rejection policy that leads to higher predictive performance as well as (2) **lower end-to-end latency** due to less fallback and rollback/rejections. 
Please refer to [W1 of wHu9](https://openreview.net/forum?id=EfMyf9MC3t&noteId=NJBmWF2n8y) for more details.

[1] Accelerating Large Language Model Decoding with Speculative Sampling

[2] Fast Inference from Transformers via Speculative Decoding

Meta-review: This paper presents a method for swapping between a large Transformer and a small Transformer to minimize decoding costs while maintaining performance. The algorithm involves a fallback policy (which essentially invokes the large model when the small model is not confident) and a rollback policy (which overwrites the small model's predictions when they disagree with the large model's predictions). Experiments were included on a variety of generation tasks, and the method was shown to decrease decoding costs while beating various baselines. Reviewers generally agreed that the paper should be accepted, with the exception of one reviewer who felt it should be compared more directly to CALM. The authors have sufficiently addressed this concern in the rebuttal and the paper should therefore be accepted.
Decision: Accept (poster)