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

arXiv: https://arxiv.org/abs/2602.10346 · HTML: https://mlsanity.com/paper/2602.10346/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 40 references that are in the map.
- First posted on arXiv 10 Feb 2026.

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities](https://mlsanity.com/paper/2507.06261/) (2025)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [Hierarchical Neural Story Generation](https://mlsanity.com/paper/1805.04833/) (2018)
- [OpenAI GPT-5 System Card](https://mlsanity.com/paper/2601.03267/) (2025)
- [Length-Controlled AlpacaEval: A Simple Way to Debias Automatic Evaluators](https://mlsanity.com/paper/2404.04475/) (2024)
- [Plug and Play Language Models: A Simple Approach to Controlled Text Generation](https://mlsanity.com/paper/1912.02164/) (2019)
- [Contrastive Decoding: Open-ended Text Generation as Optimization](https://mlsanity.com/paper/2210.15097/) (2022)
- [GeDi: Generative Discriminator Guided Sequence Generation](https://mlsanity.com/paper/2009.06367/) (2020)
- [DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models](https://mlsanity.com/paper/2309.03883/) (2023)
- [Diverse Beam Search: Decoding Diverse Solutions from Neural Sequence Models](https://mlsanity.com/paper/1610.02424/) (2016)
- [A Contrastive Framework for Neural Text Generation](https://mlsanity.com/paper/2202.06417/) (2022)
- [Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs](https://mlsanity.com/paper/2407.01082/) (2024)
- [Truncation Sampling as Language Model Desmoothing](https://mlsanity.com/paper/2210.15191/) (2022)
- [Locally Typical Sampling](https://mlsanity.com/paper/2202.00666/) (2022)
- [Contrastive Search Is What You Need For Neural Text Generation](https://mlsanity.com/paper/2210.14140/) (2022)
- [Top-H Decoding: Adapting the Creativity and Coherence with Bounded Entropy in Text Generation](https://mlsanity.com/paper/2509.02510/) (2025)
- [Mirostat: A Neural Text Decoding Algorithm that Directly Controls Perplexity](https://mlsanity.com/paper/2007.14966/) (2020)

## Built on by (most cited first, 3 of 3)

- [Truthful AI Advisors: A Pre-Specified Benchmark for Large Language Model Honesty Under Preference Misalignment](https://mlsanity.com/paper/2606.01456/) (2026)
- [Same Game, Different Story: A Minimal Conservative Strategic Robustness Benchmark for Large Language Model Agents](https://mlsanity.com/paper/2607.19670/) (2026)
- [Truncation Samplers Are Nested: Comparing Decoding Rules Is a Measurement Problem](https://mlsanity.com/paper/2610.15570/) (2026)

## Abstract (arXiv)

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% improvement. Moreover, we find that Top-W not only improves accuracy-focused performance, but also boosts creativity under judge-based open-ended evaluation.