Large language model decoding improved by diversity aware token selection
Beyond Truncation: Rethinking LLM Decoding as Ensemble Pruning
Artificial IntelligenceMachine Learning
Summary
When large language models predict words, they often pick the most likely options without thinking about how similar those choices are. This can lead to repetitive or redundant predictions. The authors introduce a new method called ME-Decoding that views this word selection process like choosing a diverse team, focusing on picking varied but still good options. This approach uses math to measure how different the candidate words are and efficiently selects a diverse set, improving the quality without slowing down the model. They tested this method on multiple tasks and found it consistently performs better.
What this means in practice
- •For natural language processing engineers: Improve text generation quality in language model deployments by selecting more diverse yet probable tokens during decoding.
- •For chatbot developers: Enable chatbots to generate more varied and meaningful responses by reducing token redundancy without adding computational cost.
Authors
Dunyao Xue, Chengshuo Du, Zhengbo Wang, Wenlin Dai, Cheng Meng
Abstract
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.