Fast Best-of-N Decoding via Speculative Rejection
Hanshi Sun, Momin Haider, Ruiqi Zhang, Huitao Yang, Jiahao Qiu, Ming Yin, Mengdi Wang, Peter Bartlett, Andrea Zanette
2024-10-29
Summary
This paper introduces a new method called Speculative Rejection, which improves the efficiency of aligning large language models (LLMs) with human preferences during their response generation.
What's the problem?
When using large language models to generate text, it's important that their responses align with what humans want. Traditional methods for achieving this alignment often require complex adjustments to the model after it has been trained, which can be time-consuming and resource-intensive. One popular method, called Best-of-N, generates multiple responses and selects the best one but is very resource-heavy, making it impractical for real-time use.
What's the solution?
The authors propose Speculative Rejection as a new inference-time alignment algorithm that generates high-quality responses more efficiently. Instead of generating many responses like Best-of-N, Speculative Rejection quickly evaluates responses during generation and stops producing lower-quality ones early on. This approach makes the process 16 to 32 times more efficient in terms of computational resources while still providing high-quality results.
Why it matters?
This research is significant because it allows for faster and more efficient use of large language models in real-time applications, making them more practical for everyday use. By improving how these models align with human preferences without the heavy computational costs, Speculative Rejection can enhance user experience in various applications, such as chatbots and content generation.
Abstract
The safe and effective deployment of Large Language Models (LLMs) involves a critical step called alignment, which ensures that the model's responses are in accordance with human preferences. Prevalent alignment techniques, such as DPO, PPO and their variants, align LLMs by changing the pre-trained model weights during a phase called post-training. While predominant, these post-training methods add substantial complexity before LLMs can be deployed. Inference-time alignment methods avoid the complex post-training step and instead bias the generation towards responses that are aligned with human preferences. The best-known inference-time alignment method, called Best-of-N, is as effective as the state-of-the-art post-training procedures. Unfortunately, Best-of-N requires vastly more resources at inference time than standard decoding strategies, which makes it computationally not viable. In this work, we introduce Speculative Rejection, a computationally-viable inference-time alignment algorithm. It generates high-scoring responses according to a given reward model, like Best-of-N does, while being between 16 to 32 times more computationally efficient.