Geometry-Aware Decoding with Wasserstein-Regularized Truncation and Mass Penalties for Large Language Models

Arash Gholami Davoodi, Navid Rezazadeh, Seyed Pouyan Mousavi Davoudi, Pouya Pezeshkpour
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:23159-23178, 2026.

Abstract

Large language models (LLMs) must balance diversity and creativity against logical coherence in open-ended generation. Existing truncation-based samplers are effective but largely heuristic, relying mainly on probability mass and entropy while ignoring semantic geometry of the token space. We present Top-$W$, a geometry-aware truncation rule that uses Wasserstein distance—defined over token-embedding geometry—to keep the cropped distribution close to the original, while explicitly balancing retained probability mass against the entropy of the kept set. Our theory yields a simple closed-form structure for the fixed-potential subset update: depending on the mass–entropy trade-off, the optimal crop either collapses to a single token or takes the form of a one-dimensional prefix that can be found efficiently with a linear scan. We implement Top-$W$ using efficient geometry-based potentials (nearest-set or $k$-NN) and pair it with an alternating decoding routine that keeps the standard truncation-and-sampling interface unchanged. Extensive experiments on four benchmarks (GSM8K, GPQA, AlpacaEval, and MT-Bench) across three instruction-tuned models show that Top-$W$ consistently outperforms prior state-of-the-art decoding approaches achieving up to 33.7 percentage improvement. Moreover, we find that Top-$W$ not only improves accuracy-focused performance, but also boosts creativity under judge-based open-ended evaluation.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-davoodi26a, title = {Geometry-Aware Decoding with {W}asserstein-Regularized Truncation and Mass Penalties for Large Language Models}, author = {Davoodi, Arash Gholami and Rezazadeh, Navid and Davoudi, Seyed Pouyan Mousavi and Pezeshkpour, Pouya}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {23159--23178}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/davoodi26a/davoodi26a.pdf}, url = {https://proceedings.mlr.press/v306/davoodi26a.html}, abstract = {Large language models (LLMs) must balance diversity and creativity against logical coherence in open-ended generation. Existing truncation-based samplers are effective but largely heuristic, relying mainly on probability mass and entropy while ignoring semantic geometry of the token space. We present Top-$W$, a geometry-aware truncation rule that uses Wasserstein distance—defined over token-embedding geometry—to keep the cropped distribution close to the original, while explicitly balancing retained probability mass against the entropy of the kept set. Our theory yields a simple closed-form structure for the fixed-potential subset update: depending on the mass–entropy trade-off, the optimal crop either collapses to a single token or takes the form of a one-dimensional prefix that can be found efficiently with a linear scan. We implement Top-$W$ using efficient geometry-based potentials (nearest-set or $k$-NN) and pair it with an alternating decoding routine that keeps the standard truncation-and-sampling interface unchanged. Extensive experiments on four benchmarks (GSM8K, GPQA, AlpacaEval, and MT-Bench) across three instruction-tuned models show that Top-$W$ consistently outperforms prior state-of-the-art decoding approaches achieving up to 33.7 percentage improvement. Moreover, we find that Top-$W$ not only improves accuracy-focused performance, but also boosts creativity under judge-based open-ended evaluation.} }
Endnote
%0 Conference Paper %T Geometry-Aware Decoding with Wasserstein-Regularized Truncation and Mass Penalties for Large Language Models %A Arash Gholami Davoodi %A Navid Rezazadeh %A Seyed Pouyan Mousavi Davoudi %A Pouya Pezeshkpour %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-davoodi26a %I PMLR %P 23159--23178 %U https://proceedings.mlr.press/v306/davoodi26a.html %V 306 %X Large language models (LLMs) must balance diversity and creativity against logical coherence in open-ended generation. Existing truncation-based samplers are effective but largely heuristic, relying mainly on probability mass and entropy while ignoring semantic geometry of the token space. We present Top-$W$, a geometry-aware truncation rule that uses Wasserstein distance—defined over token-embedding geometry—to keep the cropped distribution close to the original, while explicitly balancing retained probability mass against the entropy of the kept set. Our theory yields a simple closed-form structure for the fixed-potential subset update: depending on the mass–entropy trade-off, the optimal crop either collapses to a single token or takes the form of a one-dimensional prefix that can be found efficiently with a linear scan. We implement Top-$W$ using efficient geometry-based potentials (nearest-set or $k$-NN) and pair it with an alternating decoding routine that keeps the standard truncation-and-sampling interface unchanged. Extensive experiments on four benchmarks (GSM8K, GPQA, AlpacaEval, and MT-Bench) across three instruction-tuned models show that Top-$W$ consistently outperforms prior state-of-the-art decoding approaches achieving up to 33.7 percentage improvement. Moreover, we find that Top-$W$ not only improves accuracy-focused performance, but also boosts creativity under judge-based open-ended evaluation.
APA
Davoodi, A.G., Rezazadeh, N., Davoudi, S.P.M. & Pezeshkpour, P.. (2026). Geometry-Aware Decoding with Wasserstein-Regularized Truncation and Mass Penalties for Large Language Models. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:23159-23178 Available from https://proceedings.mlr.press/v306/davoodi26a.html.

Related Material