Source-linked AI summary
Beyond Truncation: Rethinking LLM Decoding as Ensemble Pruning
Dunyao Xue, Chengshuo Du, Zhengbo Wang, Wenlin Dai, Cheng Meng
TL;DR
Probability-based decoding can overlook semantic relationships among candidate tokens, while some geometry-aware approaches require complex optimization or reweight token probabilities. ME-Decoding frames token selection as ensemble pruning with a Mahalanobis objective, and experiments report improvements over representative baselines across reasoning and open-ended generation tasks.
Problem
Probability-centric decoding overlooks semantic relationships among candidate tokens, while existing geometry-aware approaches can involve complex optimization or direct probability reweighting.
Source: §1 · p. 1, §2.1 · p. 2Method
ME-Decoding selects a compact token subset by maximizing a Mahalanobis-Ensemble Score that combines token probabilities with embedding-based similarities to discount redundancy.
Source: §3.1 · p. 4, §5 · p. 9Results
Takeaways & Limitations
The ensemble-pruning formulation offers a way to reduce redundancy in token selection while maintaining reasoning performance with negligible inference overhead.
Source: §1 · p. 2The static token-embedding kernel may miss context-dependent semantic distinctions, and the MES criterion may not be optimal for all token distributions.
Source: Limitations · pp. 9–10
Abstract
from arXiv · showhide
We introduce Mahalanobis-Ensemble Decoding (ME-Decoding), a novel Large Language Model (LLM) decoding framework that frames candidate token selection as ensemble pruning. Existing selection strategies rely predominantly on scalar probabilities, ignoring geometric semantic relationships and causing candidate redundancy. Meanwhile, current geometry-aware methods often require complex optimization or directly reweighting the original token probabilities, leading to significant computational overhead or inference instability. To address this, we formulate decoding as a subset optimization problem using a Mahalanobis distance-driven objective to enhance semantic diversity while preserving high probabilities. Specifically, we dynamically discount redundant generation paths using a token similarity matrix, constructed via an adaptive-bandwidth kernel over token embeddings. We further devise an efficient greedy selection algorithm with near-linear complexity in the candidate size under early stopping, while establishing its theoretical approximation guarantees. This renders ME-Decoding a robust, plug-and-play module with negligible inference overhead. Extensive experiments across diverse reasoning and generation tasks demonstrate that our method consistently achieves strong performance.
1 Introduction
LLM decoding must balance high-quality continuations with diversity, but probability-based selection overlooks semantic relationships among candidate tokens. ME-Decoding treats selection as ensemble pruning, using token geometry to reduce redundancy while preserving probabilities.
1 Introduction
- Deterministic decoding can be repetitive or conservative, whereas sampling may give low-quality tokens non-negligible probability and produce illogical continuations or hallucinations.Source: §1 · p. 1
- Probability-based truncation overlooks semantic relationships among candidates, potentially retaining redundant tokens and limiting the selected sampling support.Top-k, nucleus sampling, Min-p, entropy-aware sampling, and p-less mainly use probabilities or distributional statistics.Source: §1 · p. 1
- Existing geometry-aware methods can require complex optimization or hyperparameter tuning, or depend on downstream post-processing.Top-W does not explicitly optimize redundancy in the selected support, while CraEG relies on an auxiliary decoding method.Source: §1 · p. 1
- ME-Decoding defines a Mahalanobis distance-driven objective that selects a compact, informative candidate set while preserving token probabilities.Source: §1 · p. 2
- Its greedy selection algorithm has candidate-linear complexity under early stopping, with trajectory unimodality and theoretical approximation guarantees.Source: §1 · p. 2
- ME-Decoding consistently outperforms strong existing baselines across diverse reasoning and generation tasks.It reformulates token selection as ensemble pruning to reduce redundancy while maintaining reasoning performance with negligible inference overhead.Source: §1 · p. 2
2 Motivation
Probability-based decoding selects a subset from a candidate pool but treats tokens as independent, leaving semantic dependencies unmodeled. The paper connects token selection to ensemble pruning and uses Mahalanobis geometry to penalize redundancy while retaining high-likelihood candidates.
2.1 Background of LLM Decoding
- Probability-based decoding forms a candidate pool and selects a subset under a probability criterion, but prevailing methods overlook semantic dependencies among tokens.The omitted relationships motivate incorporating inter-token semantics into the subset-selection objective.Source: §2.1 · p. 2
2.2 Connection to Ensemble Pruning
- Ensemble pruning balances individual accuracy with diversity; analogously, token probabilities represent confidence and embedding similarities represent semantic redundancy.An effective decoding subset should retain high-likelihood tokens while suppressing redundant candidates.Source: §2.2 · p. 3
2.3 Generalizing Ensemble Pruning to LLM Decoding via Mahalanobis Distance
- The Mahalanobis norm, formed using the token similarity matrix, downweights correlated directions and separates information along non-redundant axes.This geometric property supports suppressing repetitive generation paths and enhancing token diversity.Source: §2.3 · p. 3
3 Method: ME-Decoding
ME-Decoding selects a compact token subset by balancing token probabilities with embedding-based redundancy. It uses an adaptive similarity kernel, greedy optimization, and theoretical guarantees for the search path and its approximation to the global optimum.
3.1 Mahalanobis-Ensemble Score.
- MES balances high-probability, non-redundant token contributions against an entropy-dependent penalty on subset size.The optimal subset size is unknown in advance, so the penalty discourages expansion for only marginal gains.Source: §3.1 · p. 4
3.2 Construction of Similarity Matrix
- The similarity matrix uses a Gaussian kernel over normalized token embeddings, with adaptive bandwidth set by probability-weighted semantic dispersion.Concentrated candidates yield a more localized, probability-driven kernel; dispersed candidates preserve semantic relations over a broader range.Source: §3.2 · p. 4
3.3 Optimization Strategy
- Because direct subset maximization is NP-hard, greedy forward selection builds the subset while maintaining an inverse Cholesky factor to avoid repeated matrix inversions.The procedure starts with the highest-scoring token.Source: §3.3 · pp. 4–5
3.4 Theoretical Properties
- Under a conditioning assumption, the greedy MES path cannot rise again after its score drops, enabling early stopping; the stopped value approximates the global optimum.The approximation factor depends on the restricted minimum eigenvalue of the token-similarity matrix.Source: §3.4 · p. 5
4 Experiments
The experiments evaluate ME-Decoding on reasoning and open-ended generation tasks, then analyze its performance, efficiency, and accuracy–diversity trade-offs. Across the reported benchmarks, it performs strongly while retaining stable accuracy as temperature increases.
4.1 Experimental Details
- The evaluation compares three instruction-tuned models with representative decoding baselines under shared prompts, temperature settings, generation limits, and stopping criteria.ME-Decoding uses λ = 0.9 by default; Top-W and ME-Decoding use the same top-N pool with N = 512.Source: §4.1 · pp. 6–7
- Reasoning benchmarks use final-answer accuracy, while AlpacaEval reports candidate win rate and MT-Bench reports an average judge score.Open-ended responses are evaluated using DeepSeek-V4-Pro under a fixed protocol.Source: §4.1 · p. 6
4.2 Main Results
- ME-Decoding achieves the best average performance across GSM8K and GPQA and maintains a more stable accuracy profile as temperature increases.The reported comparison covers three models and three temperatures.Source: §4.2 · p. 7, Figure 2 · p. 8
- ME-Decoding achieves the best averaged AlpacaEval win rate and MT-Bench score and the best overall average rank across models, temperatures, and benchmarks.The reported open-ended results are averaged over three models; detailed numerical results are in the Appendix.Source: §4.2 · p. 7, Figure 3 · p. 8
4.3 Analysis
- With early stopping, the practical cost scales nearly linearly with candidate-pool size N and approaches O(Nd) when the number selected τ is a small constant.The analysis gives kernel cost O(Nτd) and greedy evaluation cost O(Nτ^3).Source: §4.3 · p. 7
- ME-Decoding is 2.7× faster than Top-W in the Medium setting and 7.4× faster in the Large setting, with overhead comparable to Top-p and Top-H in the Large setting.The CPU-only synthetic-logits benchmark reports average sampling overhead per token for the two settings.Source: §4.3 · p. 7, Table 4 · p. 9
- On GSM8K, ME-Decoding achieves the highest accuracy and stronger reasoning performance at comparable diversity, while Top-p reaches higher diversity but loses accuracy.The analysis uses 500 sampled questions and 10 responses per question at each temperature.Source: §4.3 · p. 8, Figure 4 · p. 8
- Against probability-only controls matched on support size, entropy, or pairwise cosine, ME-Decoding achieves the highest accuracy and Mahalanobis-Ensemble Score across both datasets.The controls are evaluated on Qwen2.5-1.5B at T = 1.0 over three generation seeds.Source: §4.3 · pp. 8–9, Table 5 · p. 9
5 Conclusion
The paper frames LLM token selection as ensemble pruning and introduces ME-Decoding to select compact, informative subsets. Experiments report gains over representative baselines with low inference overhead, while identifying broader settings and alternative objectives as future work.
5 Conclusion
- ME-Decoding reformulates LLM token selection as ensemble pruning and selects compact token subsets by combining token probabilities with embedding-based similarities.The objective maximizes the Mahalanobis-Ensemble Score to discount semantic redundancy.Source: §5 · p. 9
- The greedy algorithm comes with theoretical guarantees, including greedy-trajectory unimodality and an approximation bound.Source: §5 · p. 9
- Experiments on reasoning and open-ended generation tasks show improvements over representative probability-based and geometry-aware baselines with low inference overhead.The authors present ensemble pruning as a principled perspective for compact, diverse, and reliable token candidate sets.Source: §5 · p. 9
- Future work includes simpler and more effective objectives and extending ME-Decoding to long-form generation, multi-sample reasoning, and verification-augmented decoding.Source: §5 · p. 9
Limitations
The authors identify limits in the similarity kernel, output diversity, and subset-size selection. These constrain how well the method may represent context-specific distinctions and fit different token distributions.
Limitations
- Performance depends on the similarity matrix, and the static token-embedding kernel may miss context-dependent semantic distinctions.The authors identify contextualized or hidden-state-conditioned kernels as a direction for future work.Source: Limitations · p. 9
- ME-Decoding does not always produce the highest output diversity, leaving room to promote diversity while preserving its accuracy gains.Source: Limitations · p. 9
- Choosing the optimal subset size remains challenging, and MES may not be optimal for all token distributions.Adaptive subset-size rules and alternative objectives are identified as promising directions.Source: Limitations · pp. 9–10
Ethical Considerations
The paper presents ME-Decoding as a plug-and-play inference-time method intended to improve generation without changing model parameters or requiring additional training. It also notes that deployment context shapes downstream risks and calls for responsible evaluation and oversight.
Ethical Considerations
- ME-Decoding is an inference-time method that aims to improve generation without modifying model parameters or requiring additional training.The paper characterizes such methods as potentially improving usability with limited computational overhead.Source: Ethical Considerations · p. 10
- Improved decoding may amplify risks including misinformation, plausible hallucinations, biased or harmful content, and scaled production of low-quality or malicious text.The authors say broader impact depends strongly on deployment context, access control, and downstream use cases.Source: Ethical Considerations · p. 10
- The authors recommend safety, bias, toxicity, and hallucination evaluation, usage policies, and rate limits; high-stakes outputs should be verified by qualified human experts.They also call for data governance, documentation, and licensing practices when models involve copyrighted, private, or sensitive data.Source: Ethical Considerations · p. 10
- The experiments use publicly available datasets, models, and evaluation tools under their respective licenses and standard research-evaluation settings.The authors report collecting no new user data or personally identifying information.Source: Ethical Considerations · p. 10
A Proofs of Theoretical Results
The proofs characterize the greedy MES trajectory as unimodal, justifying early stopping at its saturation point. They also bound the greedy solution against the best subset using spectral properties of the kernel.
A Proofs of Theoretical Results
- The greedy MES sequence increases strictly before τ and is nonincreasing afterward, so its maximum occurs at the saturation point.This trajectory shape underlies the early-stopping rule.Source: A Proofs of Theoretical Results · p. 14
- Across subset sizes, the greedy trajectory’s best MES is at least α_N times the best size-matched subset MES.The bound follows from the greedy MEE approximation guarantee and a lower bound on its submodularity ratio using restricted minimum eigenvalues of the kernel.Source: A Proofs of Theoretical Results · pp. 15, 17
B.1 Additional Results for Reasoning Benchmarks
- Across nine model–temperature settings, ME-Decoding achieves the highest average accuracy on both GSM8K and GPQA.It leads in eight of nine settings on each benchmark.Source: §B.1 · p. 19
B.2 Ablation Studies
The ablations compare ME-Decoding with greedy token selection and test the semantic kernel and its adaptive bandwidth. The results favor the full method and show that removing either component reduces accuracy.
B.2 Ablation Studies
- At T = 1.0 on Qwen2.5-1.5B, ME-Decoding improves GSM8K accuracy from 61.41% to 63.20 ± 0.57% and GPQA accuracy from 32.37% to 32.66±1.05% versus Greedy decoding.Greedy is deterministic; ME-Decoding results average three generation seeds.Source: §B.2 · p. 19
- The correctly aligned semantic kernel achieves the highest average accuracy on both datasets, while removing or permuting its geometry reduces performance.The permuted control preserves the kernel’s eigenvalues, condition number, entry multiset, and numerical scale while disrupting alignment with token identities.Source: §B.2 · p. 19
- Removing the adaptive bandwidth degrades performance, especially at higher temperatures, while replacing the semantic kernel with the identity matrix also reduces accuracy.Table 9 reports absolute accuracy for ME-Decoding and changes relative to it for the other columns.Source: §B.2 · p. 19, Table 9 · p. 20
B.3 Empirical Validation of Theorem 3.1
ME-Decoding adds modest latency relative to the fastest baseline, and its early-stopped implementation avoids explicitly constructing the full similarity matrix. The appendix also probes trajectory unimodality and reports accuracy stability across tested λ values.
B.3 Empirical Validation of Theorem 3.1
- The 512-trajectory probe for each benchmark examines complete greedy MES trajectories with early stopping disabled; the observation supports the practical relevance of the stopping behavior, while the theorem’s condition is sufficient rather than necessary.The approximation result has a separate scope and depends on restricted minimum eigenvalues of the selected kernel.Source: §B.3 · p. 20, Table 10 · p. 20
B.4 Efficiency and Complexity
- Total latency rises by approximately 5.8% on Qwen2.5-1.5B and 1.5% on Qwen3-4B relative to the fastest baseline.Model inference remains the dominant cost.Source: Table 11 · p. 21, §B.4 · p. 21
- With τ selected tokens, ME-Decoding avoids the full kernel’s O(N^2d) time and O(N^2) memory, using O(Nτ) cached-kernel memory and O(τ^2) memory for R instead.The early stopping rule usually yields τ ≪ N.Source: §B.4 · p. 21
B.5 Hyperparameter Analysis
- Accuracy is relatively stable across tested λ values, and λ = 0.9 achieves the best average accuracy on both GSM8K and GPQA.The paper uses λ = 0.9 as the default setting in the main experiments.Source: §B.5 · p. 21
B.6 Diversity Analysis
The diversity analysis compares accuracy and diversity across methods and examines whether probability-only controls reproduce ME-Decoding’s results. It reports a stronger accuracy–diversity trade-off for reasoning-oriented generation.
B.6 Diversity Analysis
- ME-Decoding achieves the highest accuracy and MES on both datasets while retaining a compact support and comparable probability mass; approximately matched probability-only controls do not reproduce both results.The reported pattern is not explained solely by stronger truncation or one support statistic.Source: §B.6 · pp. 21–22
- Top-p often has the highest diversity score, especially at high temperature, but its accuracy drops substantially; ME-Decoding maintains the strongest accuracy across temperatures.The diversity average combines min-max-normalized Distinct-1, Distinct-2, and 1 −Self-BLEU within each temperature.Source: §B.6 · p. 22
B.7 Additional Results for Open-Ended Generation
Across three instruction-tuned models and temperatures, ME-Decoding achieves the highest overall averages on MT-Bench and AlpacaEval, while detailed evaluations report gains under two judges and protocols.
B.7 Additional Results for Open-Ended Generation
- ME-Decoding achieves the highest result under both automatic judges and evaluation protocols, with positive paired 95% confidence intervals against Top-W, p-less, and Top-H.The gains are not explained by verbosity or conservative refusal behavior.Source: §B.7 · p. 22, Table 15 · p. 24, Table 16 · p. 24
- Figure 5 presents MT-Bench judge scores and AlpacaEval candidate win-rates across temperatures for three instruction-tuned models, aggregated over three runs.Source: §B.7 · p. 22, Figure 5 · p. 24
- 7.17 on MT-Bench and 14.80% on AlpacaEval are ME-Decoding’s highest overall averages across the reported model-temperature combinations.Individual baselines lead in several settings.Source: §B.7 · p. 22, Table 17 · p. 25, Table 18 · p. 25
- Figure 6 reports average ranks across model-temperature settings, where ME-Decoding has the best overall average rank and competitive ranks in multiple settings.Lower ranks indicate better performance; the average is over AlpacaEval and MT-Bench.Source: §B.7 · p. 22, Figure 6 · p. 25