Overfill: Two-Stage Models for Efficient Language Model Decoding

Thank you for the constructive feedback and careful reviews.

1. Scope beyond width pruning

Our main contribution lies in the two-stage inference pipeline, which uses a large model for prefill and a smaller model for decoding. While we use an existing width pruning technique (Sreenivas et al., 2024) to construct an effective decoder, OverFill is not tied to this specific approach.

The training method is compatible with a range of strategies for building lightweight decoders, such as depth pruning, layer stacking, and even a separate pretrained model. OverFill can serve as a general framework for decoupled inference, rather than being limited to a single pruning technique.

2. Comparison to parameter efficient tuning

Thanks for the suggestion. Parameter-efficient tuning methods such as prompt tuning and LoRA reduce memory usage during training, which allows for larger batch sizes. However, they do not reduce inference cost, as the forward pass still goes through the full model parameters.

Our primary goal in this work is to reduce inference latency. We focus on reducing the size of the decoder model itself, which directly lowers inference latency, particularly in scenarios where model weights are the main bottleneck.

That said, our method is compatible with parameter-efficient tuning techniques. For example, one could apply LoRA to the decode model in OverFill to further reduce training-time memory usage.

3. Pruning ratio and scaling behavior

This is a great suggestion. In Section 6.2 of the paper, we present experiments with various pruning ratios (i.e., decoder sizes) using a fixed prefill model, to study how the capacity gap between the full and pruned models affects OverFill’s performance.

We observe that OverFill consistently outperforms the pruned model with the same decoder architecture. The robustness of OverFill under high pruning ratios varies across tasks—some tasks remain stable even under aggressive pruning, while others show more degradation as more channels are pruned.

As you suggest, a broader analysis across two axes (prefill model size and pruning ratio) would be valuable and could potentially lead to scaling heuristics. Since each combination of prefill model and pruning ratio requires separate training, we were limited by time and resources, but we view this as an exciting direction for future work.

Thank you for the constructive feedback and careful reviews.

1. Analysis on memory

Thank you for the suggestion. Whereas prefill is mostly compute-bound, decoding is known to be largely memory-bound, meaning memory usage is a direct driver of latency in this stage [1, 2]. As such, memory consumption correlates closely with our latency results.

2. Analysis on throughput

Throughput is indeed another important metric, especially at higher batch sizes. However, there is often a trade-off between throughput and end-to-end latency [3]. Prior work on speculative decoding has primarily focused on the low-batch, latency-sensitive setting, which we also follow in our evaluation.

That said, we believe OverFill has the potential to achieve better throughput than speculative decoding. Since OverFill avoids the overhead of rollback and verification mechanisms, it is more straightforward to scale to larger batch sizes. This is an interesting direction for future analysis.

3. Comparison with speculative decoding

We have added a new Pareto curve that includes the speculative decoding. The original speculative decoding method [4] is theoretically lossless, meaning it can exactly match the target model’s output. We also evaluate a lossy variant with lenient rejection sampling [5], where more tokens proposed by the draft model are accepted to achieve greater speedup.

We plot end-to-end latency versus accuracy on GSM8K with Chain-of-Thought (CoT) prompting. In our experiments, the average prompt length is 613 tokens (with 4 demonstrations), and the average generation length is 120 tokens. Latency is measured using vLLM on a single NVIDIA H100 GPU. For speculative decoding, the target model is a finetuned 3B model, and the draft model is a finetuned 1B model.

• Temperature 1 plot: https://ibb.co/YFdQsryg
• Temperature 0.01 plot: https://ibb.co/FqJB2ZQ3

In the temperature 1 setting, OverFill achieves a better tradeoff. In particular, OverFill 3B -> 1B can achieve 1.06× speedup compared to lossy speculative decoding at the same level of accuracy.

However, we found that under near-greedy decoding (i.e., low-temperature settings), speculative decoding performs very strongly, since the outputs of the target and draft models are very similar. This explains why lenient rejection sampling provides limited benefit in this case, as the acceptance rate is already high.

We plan to conduct a more thorough analysis with different OverFill configurations across more diverse tasks.

4. Comparison with the pruned model trained from scratch

We agree that training the pruned model from scratch would provide a cleaner comparison, but this is cost-prohibitive under our current resource constraints.

[1] Kwon, Woosuk, et al. "Efficient memory management for large language model serving with pagedattention." Proceedings of the 29th Symposium on Operating Systems Principles. 2023.

[2] Zhong, Yinmin, et al. "DistServe: Disaggregating prefill and decoding for goodput-optimized large language model serving." 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24). 2024.

[3] Agrawal, Amey, et al. "Taming {Throughput-Latency} tradeoff in {LLM} inference with {Sarathi-Serve}." 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24). 2024.

[4] Leviathan, Yaniv, Matan Kalman, and Yossi Matias. "Fast inference from transformers via speculative decoding." International Conference on Machine Learning. PMLR, 2023.

[5] Zhou, Yongchao, et al. "Distillspec: Improving speculative decoding via knowledge distillation." arXiv preprint arXiv:2310.08461 (2023)

Thank you for the constructive feedback and and careful reviews.

1. Clarification on Figure 1

Thank you for the feedback. Figure 1 illustrates the disaggregation of the prefill and decode stages. The yellow blocks represent the full model used for prefill. The decoder is initialized by selecting important channels (green blocks) from the full model. The blocks on the decoder side are in darker shades because they are updated during Overfill training. We will clarify this in the final version.

2. Related work on prefill-decode disaggregation

Our motivation stems from the distinct computational profiles of the prefill and decode stages, which has led to a growing trend of disaggregating them. The first line of work [1] is system-oriented, focusing on stage-specific resource allocation and parallelism while still using a single model across both stages. The second line of work (algorithmic-oriented) [2, 3] shares a similar motivation with ours and explores using models of different sizes for each stage. In contrast, Overfill proposes a more effective architecture without introducing additional modules. We will make this clearer in the final version.

3. Different model families

We have added results on the Qwen 2.5 models. Overfill consistently outperforms the pruned model and surpasses similarly sized models on many benchmarks.

[1] Zhong, Yinmin, et al. "DistServe: Disaggregating prefill and decoding for goodput-optimized large language model serving." 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24). 2024.

[2] Nair, Pranav Ajit, et al. "Tandem transformers for inference efficient llms." arXiv preprint arXiv:2402.08644 (2024).

[3] Bergner, Benjamin, et al. "Think big, generate quick: Llm-to-slm for fast autoregressive decoding." arXiv preprint arXiv:2402.16844 (2024).

Thank you for the constructive feedbacks and and careful reviews.

1. Bigger models

We have added new results with larger models. Results are shown in the general response.

2. Analysis of training cost

Thanks for the suggestion. We provide an analysis of the training costs (FLOPs) for our three setups. Using a formula for Transformer forward FLOPs per token:

$\text{Forward-FLOPs} \approx L \cdot (4H^2 + 2SH + 2HI),$

where $L$ is the number of transformer blocks, $H$ is the hidden dimension, $I$ is the intermediate MLP dimension, and $S$ is the sequence length. This formula accounts for projection, attention, and MLP operations. For total training cost, this should be multiplied by 2 to account for both forward and backward passes.

We compare three models in the 8B → 3B setup, assuming $S = 1024$, and compute their forward FLOPs per token:

We consider three training setups: OverFill, Pruned, and Finetuned. For simplicity, assume each sequence is evenly split between prompt and generation. While this slightly affects the $S$-dependent term, the difference is minor and ignored for clarity.

• OverFill processes the first half of tokens with the full model (forward only), and the second half with the pruned model (forward + backward). The full model is frozen, so the total training FLOPs are:

$FLOPs_{overfill} = (S/2) \cdot F + 2 \cdot (S/2) \cdot f \approx 5.20B \cdot S$

• Pruned uses the pruned model for the entire sequence:

$FLOPs_{pruned} = (S/2) \cdot f + 2 \cdot (S/2) \cdot f = 3.17B \cdot S$

• Finetuned uses the finetuned model similarly:

$FLOPs_{finetuned} = (S/2) \cdot f' + 2 \cdot (S/2) \cdot f' = 3.96B \cdot S$

Thus, OverFill has approximately 1.64× and 1.31× higher theoretical training cost compared to the pruned and finetuned models, respectively. In practice, OverFill is roughly 2× slower in runtime during training.

However, we would like to emphasize that our primary focus is inference efficiency, as training is a one-time cost.

3. Compatibility with existing model compression methods (quantization, distillation)

This is also a good point. We believe OverFill is highly compatible with distillation, whether at the sequence level (e.g., using sampled outputs instead of teacher forcing) or the token level (e.g., learning from the full model’s output logits). Such techniques can benefit both OverFill and the baselines, which we think is orthogonal improvements to the core idea of OverFill. Also, post-training quantization techniques can be applied independently to the prefill and decode models. However, we have not yet evaluated the robustness of our models to quantization.

An interesting future direction would be to prefill with a full-precision model and decode with a quantized model, which may require quantization-aware training. We leave this as future work.

4. Elaboration on prefill refresh

To clarify the idea we mentioned in future work: this direction can be implemented in a way similar to speculative decoding, where a bigger model and a small model are used together during decoding.

Specifically, decoding would proceed as follows:

1. Prefill the KV cache using the full model with the initial prompt.
2. Generate N tokens using the small model (decode model).
3. Feed the newly generated N tokens into the full model to refresh its KV cache.
4. Repeat steps 2–3 until the sequence is complete.

This approach queries the full model every N tokens instead of only once at the beginning. We believe it can preserve the benefit of "smart prefill" throughout generation, without breaking autoregressive continuity.

The main difference from speculative decoding is that we do not reject the tokens generated by the small model. As a result, this approach avoids verification and rollback steps, which may lead to higher throughput.