TLDR: OverFill is a novel method that optimizes large language model (LLM) inference by separating the prefill and decode stages. It uses a full, larger model for the initial, compute-bound prefill stage and then switches to a smaller, pruned model for the memory-bound decode stage. This two-stage approach significantly boosts accuracy and efficiency, outperforming standalone smaller models and matching larger models trained from scratch, all while minimizing latency overhead, especially for long output sequences.
Large language models (LLMs) have become incredibly powerful, handling tasks from answering questions to generating code. However, their sheer size makes them challenging and costly to deploy, especially when it comes to inference – the process of generating responses. This inference process typically involves two main stages: prefill and decode.
The prefill stage is where the model processes all your input at once, building a memory (known as a Key-Value or KV cache). This stage is usually limited by how much computation can be done. The decode stage, on the other hand, generates each output token one by one, and its main bottleneck is repeatedly loading the model’s large components into memory. Current LLM designs often treat these two stages uniformly, even though they have very different computational needs.
A new approach called OverFill aims to tackle this challenge by decoupling these two stages to achieve a better balance of accuracy and efficiency. OverFill starts by using a full, larger model to handle the prefill stage, processing the initial system and user inputs in parallel. Once this is done, it switches to a smaller, pruned version of the model for the decode stage, generating tokens sequentially. This strategy leverages more computational power during the initial prefill, which helps improve the quality of the generated output, all while adding minimal delay.
The core idea behind OverFill is to use the model’s capacity asymmetrically. The full set of model parameters is used for prefill, while a carefully selected subset of these parameters (the pruned version) is used for decoding. This pruning specifically targets the ‘width’ of the model, meaning it reduces the size of certain internal dimensions without changing the overall structure of the KV cache, ensuring compatibility between the full and pruned models. Once the pruned subset is determined, the full model is kept frozen, and only the smaller, pruned model is trained further on output tokens. This makes the training process more efficient and allows the pruned decoder to integrate smoothly with existing full models.
Experimental results demonstrate the effectiveness of OverFill. For instance, a configuration where a 3-billion parameter model is used for prefill and a 1-billion parameter model for decode (3B-to-1B OverFill) significantly outperforms standalone 1-billion parameter pruned models by 83.2% on average across various benchmarks. Similarly, an 8B-to-3B configuration shows a 79.2% improvement over 3B pruned models. Impressively, OverFill can even match the performance of models of the same size that were trained from scratch, but it achieves this with significantly less training data.
In terms of efficiency, OverFill introduces minimal overhead compared to a small standalone model. While there’s a slight increase in prefill cost, this is negligible when considering the dominant decoding latency, especially for longer generated sequences. The approach shows consistent speedups compared to using a full model throughout the entire process. OverFill also demonstrates Pareto optimality, offering a better balance between latency and accuracy compared to other fine-tuned or pruned models.
Compared to speculative decoding, another method for speeding up LLM inference, OverFill shows competitive performance, particularly in scenarios with multiple concurrent requests. Its advantage comes from accepting all tokens from the smaller ‘drafter’ model without rejection sampling, leading to higher throughput. The benefits of OverFill also extend to long generations, with the model consistently predicting tokens more accurately than a standalone pruned model, even as the output length increases.
Also Read:
- LessIsMore: A Training-Free Approach to Efficient AI Reasoning
- Smart LLM Adaptation: DP-LLM Adjusts Precision on the Fly
This innovative two-stage approach offers a promising direction for making LLMs more practical and cost-effective for deployment. For more technical details, you can refer to the full research paper here.
