Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

When to Ensemble: Identifying Token-Level Points for Stable and Fast LLM Ensembling

Heecheol Yun, Kwangmin Ki, Junghyun Lee, Eunho Yang

TL;DR

This work tackles the instability and inefficiency of probability-level LLM ensembling for long-form generation by identifying tokenization mismatch and consensus in next-token distributions as key factors. It introduces SAFE, a Generate–Verify–Ensemble framework that designates a drafter to generate lookahead tokens and verifiers to select safe ensemble points, supplemented by a probability sharpening step to combat dispersed ensemble mass across subword tokens. SAFE reduces the number of ensemble operations, maintains latency near single-model inference, and improves accuracy across benchmarks with heterogeneous tokenizers, outperforming baselines like GaC and UniTE in many settings. The approach offers a practical, plug-and-play solution for robust, fast long-form LLM ensembling and suggests avenues for extending to reasoning-model ensembles.

Abstract

Ensembling Large Language Models (LLMs) has gained attention as a promising approach to surpass the performance of individual models by leveraging their complementary strengths. In particular, aggregating models' next-token probability distributions to select the next token has been shown to be effective in various tasks. However, while successful for short-form answers, its application to long-form generation remains underexplored. In this paper, we show that using existing ensemble methods in long-form generation requires a careful choice of ensembling positions, since the standard practice of ensembling at every token often degrades performance. We identify two key factors for determining these positions: tokenization mismatch across models and consensus in their next-token probability distributions. Based on this, we propose SAFE, (Stable And Fast LLM Ensembling), a framework that selectively ensembles by jointly considering these factors. To further improve stability, we introduce a probability sharpening strategy that consolidates probabilities spread across multiple sub-word tokens representing the same word into a single representative token. Our experiments on diverse benchmarks, including MATH500 and BBH, demonstrate that SAFE outperforms existing methods in both accuracy and efficiency, with gains achieved even when ensembling fewer than 1% of tokens.

When to Ensemble: Identifying Token-Level Points for Stable and Fast LLM Ensembling

TL;DR

This work tackles the instability and inefficiency of probability-level LLM ensembling for long-form generation by identifying tokenization mismatch and consensus in next-token distributions as key factors. It introduces SAFE, a Generate–Verify–Ensemble framework that designates a drafter to generate lookahead tokens and verifiers to select safe ensemble points, supplemented by a probability sharpening step to combat dispersed ensemble mass across subword tokens. SAFE reduces the number of ensemble operations, maintains latency near single-model inference, and improves accuracy across benchmarks with heterogeneous tokenizers, outperforming baselines like GaC and UniTE in many settings. The approach offers a practical, plug-and-play solution for robust, fast long-form LLM ensembling and suggests avenues for extending to reasoning-model ensembles.

Abstract

Ensembling Large Language Models (LLMs) has gained attention as a promising approach to surpass the performance of individual models by leveraging their complementary strengths. In particular, aggregating models' next-token probability distributions to select the next token has been shown to be effective in various tasks. However, while successful for short-form answers, its application to long-form generation remains underexplored. In this paper, we show that using existing ensemble methods in long-form generation requires a careful choice of ensembling positions, since the standard practice of ensembling at every token often degrades performance. We identify two key factors for determining these positions: tokenization mismatch across models and consensus in their next-token probability distributions. Based on this, we propose SAFE, (Stable And Fast LLM Ensembling), a framework that selectively ensembles by jointly considering these factors. To further improve stability, we introduce a probability sharpening strategy that consolidates probabilities spread across multiple sub-word tokens representing the same word into a single representative token. Our experiments on diverse benchmarks, including MATH500 and BBH, demonstrate that SAFE outperforms existing methods in both accuracy and efficiency, with gains achieved even when ensembling fewer than 1% of tokens.
Paper Structure (25 sections, 1 theorem, 7 equations, 10 figures, 10 tables, 1 algorithm)

This paper contains 25 sections, 1 theorem, 7 equations, 10 figures, 10 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Ensemble after inference
  4. Ensemble during inference
  5. SAFE: Toward Stable and Fast LLM Ensembling
  6. Generate
  7. Verify
  8. Ensemble: Sharpening Ensemble Distribution
  9. Experiments
  10. Implementation
  11. Experimental Settings
  12. Main Analysis
  13. Ablation Study
  14. Conclusion
  15. Experimental details
  16. ...and 10 more sections

Key Result

Theorem 1

Let $\mathbf{t}_{<j}$ denote the drafter's prefix. For each verifier model $LLM_v \in M_\mathrm{ver}$, let $t_{v_j}^v$ be its next token, aligned such that $\mathbf{t^v}_{<v_j}$ is the $LLM_v$'s tokenization of $\mathbf{t}_{<j}$. Let $P_v$ be the next-token probability distribution of $LLM_v$. Suppo Then $t_j$ is the token selected by the ensemble, i.e., $t_j = \arg\max_t P_\mathrm{ens}(t \mid \ma

Figures (10)

  • Figure 1: Illustration of the OOV-like token problem. When So is fed into $LLM_2$, which tokenizes Sofia as a single token, the probability distribution of the next token becomes corrupted.
  • Figure 2: (Left) Failure cases of existing ensemble methods in long sequence generation. (Right) Feeding OOV-like tokens into a model often causes it to produce wrong tokens.
  • Figure 3: Overview of SAFE. The drafter generates a sequence of tokens, which the verifiers re-tokenize it using their own tokenization schemes and identify where ensembling is required. At this position, ensembling is performed with our probability sharpening strategy.
  • Figure 4: Tokenization agreement rates between each model pair on Oxford 5000 words.
  • Figure 5: Latency comparison on MATH500. Our method shows similar latency compared to individual models, even when generating long sequences. w/ KV indicates that our KV caching strategy is applied. Note that the time-axis scale in (c) differs from (a) and (b).
  • ...and 5 more figures

Theorems & Definitions (2)

  • Theorem 1
  • proof