TLDR: The research paper introduces SAFE (StableAndFast LLM Ensembling), a novel framework designed to improve the accuracy and efficiency of combining Large Language Models (LLMs) for long-form text generation. It addresses the performance degradation seen in existing ensemble methods by selectively ensembling tokens. SAFE identifies optimal ensembling points by considering tokenization mismatches across models and the consensus in their next-token probability distributions. The framework employs a Generate-Verify-Ensemble cycle, a probability sharpening strategy, and an efficient KV cache implementation. Experiments show SAFE significantly outperforms current methods, achieving comparable speeds to individual models while requiring minimal ensembling operations.
Large Language Models (LLMs) have become incredibly powerful across many fields, from solving complex math problems to generating creative text. However, no single LLM is perfect for every task, and each has its own unique strengths. This has led to a growing interest in ‘ensembling’ LLMs – combining multiple models to leverage their complementary abilities and achieve even better performance than any individual model.
One particularly effective method is ‘probability-level ensemble,’ where the next-token probability distributions from several LLMs are aggregated to select the most confident next word or sub-word unit. While this approach has shown great success for short answers, its application to longer, more complex text generation has been less explored.
A recent research paper, titled WHEN TO ENSEMBLE: IDENTIFYING TOKEN-LEVEL POINTS FOR STABLE AND FAST LLM ENSEMBLING by Heecheol Yun, Kwangmin Ki, Junghyun Lee, and Eunho Yang, delves into this challenge. The authors found that simply ensembling at every single token during long-form generation often leads to a decrease in performance. This is primarily due to two critical issues: ‘tokenization mismatch’ across different models and a lack of ‘consensus’ in their next-token probability distributions.
Tokenization mismatch occurs when an ensemble selects a token that doesn’t fit well with how a participating model breaks down words. This can create ‘OOV-like’ (Out-Of-Vocabulary-like) tokens, forcing a model into an unfamiliar state and leading to incorrect or repetitive outputs. Imagine trying to build a word like “Sofia” where one model sees it as a single unit, but the ensemble generates “So” first. If “So” isn’t a natural prefix for the second model’s tokenization of “Sofia,” it can corrupt the model’s predictions, leading to errors that accumulate over long sequences.
To address these problems, the researchers propose a new framework called SAFE (StableAndFast LLM Ensembling). SAFE is designed to identify the optimal moments for ensembling by jointly considering both tokenization mismatches and the level of agreement among models’ next-token predictions. It adopts a speculative strategy, similar to how some LLMs generate text quickly. In SAFE, one model, called the ‘drafter,’ generates a short sequence of tokens. The other models, called ‘verifiers,’ then quickly examine these tokens to determine if ensembling is both stable (no OOV-like tokens introduced) and necessary (insufficient agreement among verifiers).
The SAFE framework operates in a three-step cycle: Generate–Verify–Ensemble. First, the drafter generates a small chunk of tokens. Next, the verifiers check these tokens in a single, efficient pass. Ensembling is only triggered if an OOV-like token is not introduced and if the verifiers don’t show enough agreement on the token. Finally, if ensembling is needed, the token is replaced with a more confident one derived from the combined distributions of all models. SAFE also introduces a ‘probability sharpening’ strategy to consolidate probabilities that might be spread across multiple sub-word tokens for the same word, ensuring more precise token selection.
The advantages of SAFE are significant. It offers improved efficiency by limiting costly autoregressive generation to just the drafter and by reducing the number of expensive ensemble operations. This allows SAFE to achieve inference speeds comparable to individual models, even for long sequences. It also greatly enhances stability by preventing the introduction of OOV-like tokens, leading to more accurate outputs. Furthermore, SAFE is ‘plug-and-play,’ meaning it can be easily integrated with existing ensemble methods, consistently improving their performance across various model combinations.
Experiments on diverse benchmarks like MATH500, BBH, and MMLU-redux demonstrated that SAFE consistently outperforms existing ensemble methods in both accuracy and efficiency. Remarkably, these gains were achieved even when ensembling fewer than 1% of tokens in some cases. The research also found that math-related tasks required less ensembling due to the structured nature of their responses, leading to higher model agreement. The authors also implemented a novel KV cache management strategy to maintain consistency and further boost efficiency during the ensembling process.
Also Read:
- TokenTiming: Accelerating LLM Inference with Universal Speculative Decoding
- Enhancing Language Model Reasoning with Calibrated Sampling
In conclusion, SAFE represents a practical advancement in LLM ensembling, making it more robust and efficient for generating long-form content. By intelligently deciding when to ensemble, SAFE helps overcome critical challenges, paving the way for more stable and deployable LLM collaboration in real-world applications.
