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
Paper
Similar papers
Peer review
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. To 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. This 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. [1] DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference - ACL Anthology [2] [2004.02178] FastBERT: a Self-distilling BERT with Adaptive Inference Time [3] [2004.07453] The Right Tool for the Job: Matching Model and Instance Complexities [4] BERxiT: Early Exiting for BERT with Better Fine-Tuning and Extension to Regression - ACL Anthology [5] [2104.08803] Consistent Accelerated Inference via Confident Adaptive Transformers [6] [2207.07061] Confident Adaptive Language Modeling [7] [1910.10073] Depth-Adaptive Transformer [8] [2204.06271] TangoBERT: Reducing Inference Cost by using Cascaded Architecture [9] [2012.14682] CascadeBERT: Accelerating Inference of Pre-trained Language Models via Calibrated Complete Models Cascade [10] https://arxiv.org/abs/2306.02307
Strengths
- An interesting and important problem - A natural extension of a rich line of work (see discussion above) - Impressive solutions to the challenges raised by this setup (see a-c above) - Impressive results - Interesting analysis and ablation experiments, hinting that all components contribute to the success of the approach. - The paper is overall very clearly written, though see some comments below.
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. - 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?
Questions
- Fig2: (a) what is the baseline (large) model performance? (b) why are we seeing a drop around 10-20%? - Did the authors consider experiments with models from different families? Would this make the alignment problem harder? - What is the performance of the small model in figure 4? This seems like an important baseline
Rating
8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
4 excellent
Presentation
4 excellent
Contribution
4 excellent
Limitations
N/A
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.
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. - 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. - 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. - 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? [1]: https://arxiv.org/abs/2207.07061
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. 2. How large is the big model compared to the small model? Any restrictions?
Rating
4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
2 fair
Contribution
2 fair
Limitations
Not addressed.
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. Given 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). The 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.
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. Experiments show that this significantly improves the speedup compared to speculative sampling.
Weaknesses
It is unclear how this can achieve quality better than speculative sampling. Speculative 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. However, 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. If 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.
Questions
As mentioned in the weakness section, what are your thoughts on why the model can perform superior compared to Speculative sampling?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
No limitations that I can think of, and none discussed in the paper.
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.
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. - 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. - **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. [1] Fast Inference from Transformers via Speculative Decoding, Leviathan et al., 2023
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. - 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. - 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? - 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.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
2 fair
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.
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.
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. 2. The fallback and rollback policies are simple yet effective mechanisms to coordinate the models and empirically show better performance than speculative sampling. 3. 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. 2. 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]. 3. The policies require tuning fallback and rollback thresholds, adding hyperparameters. Optimal values likely vary across tasks and models. 4. 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]. 5. 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. 6. A minor problem is that there are limited insights into what types of predictions the large model tends to correct. [1] Accelerating Large Language Model Decoding with Speculative Sampling [2] Lossless Acceleration for Seq2seq Generation with Aggressive Decoding [3] Blockwise Parallel Decoding for Deep Autoregressive Models [4] Levenshtein Transformer [5] Deep Encoder, Shallow Decoder: Reevaluating Non-autoregressive Machine Translation [6] Self-Consistency Improves Chain of Thought Reasoning in Language Models
Questions
Questions: 1. In Figure 2, why does 20% achieve better performance than 40 or 80 percent? Typos: line 142: regressive ==> autoregressive
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
2 fair
Presentation
3 good
Contribution
2 fair
Limitations
N/A
Thanks for the author response
I appreciate the author's thoughtful response, which has helped clarify several points. Regarding the first issue, I agree that batch size would be less problematic when utilizing iteration-level batching. However, the paper would be strengthened substantially if the authors could include quantitative speedup comparisons between BiLD and standard decoding with this implementation. Nonetheless, the batch size=1 results adequately demonstrate the approach's efficacy. Increasing batch sizes greater than 1 appears to require primarily engineering efforts rather than fundamental enhancements. For W2, thanks for providing the additional results. Given that, could you provide an intuitive understanding as to why BiLD provides additional efficiency gains compared to speculative decoding? Based on the authors thoroughly addressing the majority of the questions and concerns raised in my review, I have made the decision to increase my score to 6.
Further Comments
Thank you for your valuable comments and feedback. Regarding the question of why BiLD outperforms speculative decoding, we offer insightful explanations in our response to [reviewer wHu9](https://openreview.net/forum?id=EfMyf9MC3t¬eId=NJBmWF2n8y). 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. Below is a more detailed discussion on each point. (1) **Better text generation quality:** (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 [3]. 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]. (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. (2) **Lower end-to-end latency:** 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: (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. (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. Minimizing both fallback and rollback/rejection reduces unnecessary computation which directly translates to end-to-end latency improvement. | Task | Method | BLEU/ROUGE | Latency | % Fallback | % Rollback (rejection) | | ----- | ------------------- | ---------- | ------- | ---------- | ---------------------- | | IWSLT | Speculative | 39.93 | 1.28x | 23.24% | 9.81% | | | BiLD (Better BLEU) | **40.33** | 1.43x | 21.09% | **1.56%** | | | | | | | | | XSum | Speculative | 35.00 | 1.25x | 36.84% | 24.24% | | | BiLD (Better ROUGE) | **35.12** | 1.48x | 32.33% | **6.41%** | Reference: [1] Chen et al. Accelerating Large Language Model Decoding with Speculative Sampling [2] Leviathan et al. Fast Inference from Transformers via Speculative Decoding [3] Jiang at el. LLM-Blender: Ensembling Large Language Models with Pairwise Ranking and Generative Fusion
I appreciate the authors' clarifications in the rebuttal. > W1. BiLD with Sampling and comparison to Speculative Decoding Thanks for the added clarification and results. This will surely strengthen the submission. > W2. Hyperparameters While I understand that the evaluation-time hyperparameter can make searches more efficient, it's my belief that as we aim for a general text-based agent like ChatGPT to handle a wide range of tasks, tuning hyperparameters for each task can complicate its general application. > W4. HumanEval I remain unconvinced by the author's argument regarding the challenges of machine translation and text summarization. However, I value the results presented, especially those compared to CALM and speculative decoding. > Overall All in all, some of my concerns about the paper have been addressed, leading me to raise my score to 5.
Thank you for the clarifications, I am happy with the responses.
Acknowledge the reponse
Thanks for the rebuttal. The reviewer still feels CALM and method in this paper are fundamentally similar. Token-level routing is not practical for serving when batch size is greater than 1 while sequence level routing is not illustrated effective yet. The reviewer does not agree that QA benchmarks are not relevant to this method, as QA benchmarks are more challenging to get good scores when model capacity is limited.
Additional comments on token-level batching and comparison with CALM
We appreciate the reviewer for the further comments. We would like to make further clarification to the points brought up in the comments. ### 1. Token-level Batching/Routing > Token-level routing is not practical for serving when batch size is greater than 1 while sequence level routing is not illustrated effective yet. We politely claim that token-level batching/routing (or iteration-level batching) is **a more efficient and pragmatic solution** compared to basic batching. This method is **already a standard in various LLM serving systems**, including open-source frameworks like Huggingface’s Text Generation Inference (TGI) [1], vLLM [2], Periflow [3], and FlexFlow [4], as well as in proprietary systems. Additionally, as highlighted by Reviewer mW2r, **integrating the BiLD framework into token-level batching systems is a straightforward extension**, primarily requiring engineering work rather than fundamental changes. For instance, FlexFlow, a recent framework, supports speculative decoding, another way of coordinating multiple-sized models in run-time, on top of their token-level batching system. This emphasizes that the BiLD framework can be **seamlessly supported on modern batching systems**. [1] TGI, github: huggingface/text-generation-inference [2] vLLM, github: vllm-project/vllm [3] Periflow, github: friendliai/periflow-client [4] FlexFlow, github: flexflow/FlexFlow ### 2. Comparison with CALM > The reviewer still feels CALM and method in this paper are fundamentally similar. We've incorporated a requested comparison table between CALM and BiLD. Note that since BiLD consists of a small and a large model, we have utilized the *large model* for CALM to ensure a fair comparison. We fine-tuned the mT5-large model on IWSLT using CALM loss objective (weighted per-layer loss average) with the same hyperparameters in the paper. Furthermore, we have conducted an exhaustive hyperparameter search to determine the optimal inference setup, following the methodology outlined in our paper. | Method | BLEU | Norm. Latency | CALM avg. layer | | -------------------- | ----- | ------------------ | ---------------------------- | | Baseline (mT5-large) | 40.32 | 1x | | | | | | | | BiLD | **40.54** | **1.23x** | | | | **39.87** | **1.49x** | | | | | | | | CALM | 35.03 | 1.05x | 8.9 | | | 32.67 | 1.18x | 5.2 | | | 31.33 | 1.25x | 3.6 | The table clearly shows **BiLD's substantial superiority over CALM**. We hypothesize two reasons for CALM's reduced performance using the 24-layer large model (note that the original CALM paper only assessed an 8-layer small model): 1. Unlike in shallower models, deeper models' early layers might lack the representative power of the full model, making it more challenging to match predictive performance within a few initial layers. This is clear in the table, where utilizing only 3-5 layers leads to a notable BLEU drop of 8-9 points. 2. CALM involves early exiting decisions in every layer, introducing notable latency overhead due to frequent dense multiplication with a huge vocabulary embedding. This overhead is particularly exaggerated when the vocabulary size is big. For instance, the table indicates that running ~9 layers yields similar latency as the full model, demonstrating a huge overhead regarding the early exiting policy. In contrast, BiLD's fallback policy infrequently runs at per-token granularity and can be more latency efficient.
Decision
Accept (poster)