SlowGuess/NeurIPS_2023_OCR huggingface dataset:
CALM. To reproduce CALM [51] in our experimental setup, we have fine-tuned the 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.,  $\mathcal{L} = \sum_{i=1}^{L} w_i \mathcal{L}i$ , where  $w_i = i / \sum_{j=1}^{L} j$ , which was introduced in the paper to ensure the layer consistency. We use Adafactor optimizer [53] with constant learning rate of  $\{0.5, 1, 2, 5\}\mathrm{e}-4$  for 500k training steps. To make a fair comparison, we match the BLEU score of the fine-tuned model to that of BiLD's models. Among the two training-free confidence measures introduced in the CALM paper, softmax-based and hidden-state saturation-based 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 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 mT5, which is  $\sim 10 \times$  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-offs between latency and generation quality, we sweep over  $\lambda$  in  $[0.7, 0.98]$  and  $t$  in  $\{0, 1, 2, 4, 8\}$  in the decaying threshold function.
7.2.2 Performance Comparison between BiLD and CALM
Figure 7 illustrates 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  $\sim 1.5\times$  speedup regime. This can be attributed to two factors. First, in BiLD, even if an early exited prediction (i.e., prediction made by the smaller model) is incorrect, it can be corrected and replaced using the rollback policy.
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 error propagation and improved decoding stability.
7.3 Comparison with Other Speculative Decoding Frameworks
Concurrently and independently of our work, [38, 4] also propose an algorithm to accelerate generative inference using a more powerful model to score and speculatively sample predictions from a less powerful model. While the rejection sampling-based approach in [38, 4] offers unbiased estimators that match the stronger model's probability distributions, our extensive empirical evaluation shows that our approach can deliver superior latency-performance trade-offs, due to its non-random rollback (i.e., rejection) policy as well as the dynamic fallback window size. Below, we 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 in methodology
While the idea of using two models with different sizes can be deemed similar to the speculative decoding frameworks in [38, 4], 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 in [38, 4]. While [38, 4] propose 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 Fallback Window Size: Additionally, we introduce the dynamic fallback window size in our fallback policy. In [38, 4], the window size remains a 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, 4], 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.
7.3.2 Quantitative Comparisons
In Table 3, we provide a comprehensive quantitative comparison between our method and [38, 4] 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 all experiments. We maintained the same evaluation setup and hyperparameter space that are outlined in Appendix 7.1.2.
Table 3 includes 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. Across all experiments, BiLD consistently outperforms speculative decoding. It achieves either (1) notably improved BLEU/ROUGE-L scores with equivalent latency gains, or (2) superior latency gains while retaining the same BLEU/ROUGE-L scores.
Table 3: Comparison of BiLD to other rejection sampling based speculative sampling methods proposed in [38, 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.
Dataset	IWSLT	XSUM
BLEU	Speedup	ROUGE-L	Speedup
Vanilla Inference	40.32	-	35.08	-
Rejection Sampling Based [38, 4]	39.93	1.28×	35.00	1.25×
BiLD (Match Latency)	40.54	1.23×	35.30	1.42×
BiLD (Match BLEU/ROUGE-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 in a significantly better latency speedup.
Task	Method	BLEU/ROUGE-L	Speedup	% Fallback	% Rollback (rejection)
IWSLT	Rejection Sampling Based [38, 6]	39.93	1.28×	23.24%	9.81%
BiLD (Better BLEU)	40.33	1.43×	21.09%	1.56%
XSUM	Rejection Sampling Based [38, 6]	35.00	1.25×	36.84%	24.24%
BiLD (Better ROUGE-L)	35.12	1.48×	32.33%	6.41%
7.3.3 Insights on Better Latency and Performance
Quantitative analysis reveals a consistent trend where our method, when compared to 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 to each other [27]. 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 [38, 4].
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 counterpart, based on cross entropy-based distance metric. This contrasts with speculative decoding, where the decision involves a stochastic rejection sampling process. We empirically observe that BiLD's hard rejection policy allows a better BLEU/ROUGE-L score with significantly fewer number of rollbacks (rejections) than the stochastic rejection policy of speculative decoding as described in Table 4. 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 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 in contrast with speculative decoding which adopts a static window size. The advantages of the dynamic window size are two-fold:
Table 5: BiLD with nucleus sampling (p=0.8) on IWSLT and XSUM. Similar to the greedy decoding case, our method achieves a  $\sim 1.5\times$  speedup without compromising performance and a  $\sim 1.8\times$  speedup with a modest 1-point BLEU/ROUGE score reduction with sampling.
Dataset	IWSLT	XSUM
BLEU	Speedup	ROUGE-L	Speedup
Vanilla Inference	39.24	-	34.00	-
BiLD	39.72 (+0.48)	1.51×	34.34 (+0.34)	1.22×
39.26 (+0.02)	1.63×	34.04 (+0.04)	1.45×
38.27 (-0.97)	1.80×	33.10 (-0.90)	1.85×
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 of backups ( $23.24\% \to 21.09\%$  and  $36.84\% \to 32.33\%$ ) than [38, 4] while achieving better performance.
Less rollbacks/rejections: 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\% \to 1.56\%$  and  $24.24\% \to 6.41\%)$  than [38, 4] while achieving better performance.
Minimizing both backends and rollbacks/rejections reduces unnecessary computation which directly translates to end-to-end latency improvement.
7.4 BiLD with Sampling
Our approach isn't restricted to greedy decoding, but it can seamlessly extend to 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 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 [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 BiLD on the IWSLT and XSUM benchmarks as well as their relative speedup. As can be seen in the table, our method exhibits a similar trend to the greedy decoding case. It achieves a  $\sim 1.5\times$  speedup without compromising performance and a  $\sim 1.8\times$  speedup with a modest 1-point BLEU/ROUGE score reduction.
7.5 Additional Analysis
7.5.1 Model Analysis of BiLD: FLOPs, MOPs, and Arithmetic Intensity
Figure 8 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, we use the model with roughly the same ROUGE-L score as the vanilla inference, and all the numbers are normalized by the numbers of the vanilla inference. The figure illustrates that BiLD exhibits slightly higher FLOPs compared to the vanilla inference. This is due to the fact that the autoregressive and non-autoregressive executions have the same amount of FLOPs, and BiLD involves additional overhead of running the small model alongside. However, in the case of MOPs, BiLD demonstrates a significant  $\sim 5\times$  reduction of memory operations. This can be attributed to the capability of BiLD to process multiple tokens with a 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 a single token. Consequently, BiLD achieves a significantly higher arithmetic intensity, which is approximately 5 times larger than the vanilla inference. Arithmetic intensity [72] 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
Figure 8: FLOPs, MOPs (memory operations), arithmetic intensity, and latency speedup comparison of vanilla inference and BiLD on the CNN/DailyMail benchmark. BiLD approach results in a remarkable reduction in MOPs due to the improved token-level parallelism, resulting in significantly higher arithmetic intensity.
Figure 9: Example text sequences that BiLD generates with the validation set of IWSLT 2017 De-En, compared to 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 a small engagement of the large model, BiLD can correct not only inaccurate vocabulary but also wrong semantics of the text that the small model would have otherwise generated.
Ground Truth	And Siftables are an example of a new ecosystem of tools for manipulating digital information.
Large	And the Siftables are an example of a new generation of manipulation tools for digital data.
Small	And the if you look at the ifleses are an example of a new generation of technologies for manipulation of digital data.
BiLD (ours)	And the Siftables are an example of a new generation of manipulation of digital data.
Ground Truth	Which is great, because the Romans did not actually think that a genius was a particularly clever individual.
Large	That's great. The Romans didn't really think that a genius was a particularly smart individual.
Small	That's great. The tube didn't really think that a genius was a particularly lonely individual.
BiLD (ours)	That's great. The Romans didn't really think that a genius was a particularly smart individual.
Ground Truth	The viral particles then were released from the cells and came back and killed the E. coli.
Large	The viral particles then were released by the cells and came back and killed E. coli.
Small	The viral particles were then released by the cells and came back and killed E. Coke.
BiLD (ours)	The viral particles then were released by the cells and came back and killed E. coli.
and increasing arithmetic intensity can effectively alleviate the inference bottleneck. This leads to an overall latency speedup of  $1.85 \times$  on actual hardware.
7.5.2 Examples of Generated Sequences
Figure 9 provides examples of text sequences generated by BiLD on the validation set of IWSLT 2017 De-En, along with the ground truths (i.e., 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 in BiLD decoding not only improves the prediction accuracy but also prevents incorrect predictions from impacting the future ones.
Figure 10: 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 trade-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 backups.
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, we present a detailed analysis of how these thresholds affect the performance using the aligned BiLD model on two different summarization tasks, XSUM and CNN/DailyMail, as illustrated in Figure 10. 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 more rollback, while a larger fallback threshold implies more fallback.
We observe a 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 can be potentially inaccurate by sacrificing the latency. Similarly, there is also a general trend that smaller fallback thresholds (i.e., fewer backalls) result in faster latency but a worse generation quality. However, we 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 are later rolled back, incurring an extra 'flush' cost for the tokens that follow.