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

AI summary

Overview

Research area: Natural Language Processing — specifically inference-time decoding algorithms for large language models (LLM sampling/truncation rules).

Technical level: Advanced. The method is built on optimal transport (Wasserstein-1 distance), Kantorovich–Rubinstein duality, 1-Lipschitz potentials, and embedding-geometry metrics, though the authors work hard to reduce the final algorithm to a simple sort-and-prefix-scan procedure.

Scope in one sentence: The paper proposes Top-W, a decoding rule that chooses which next-token candidates to keep by minimizing a Wasserstein-plus-entropy-minus-mass objective over token-embedding geometry, proves the optimal crop has a prefix structure solvable in linear time, and benchmarks it against Min-p, Top-p, and Top-H across four benchmarks and three instruction-tuned LLMs.

Authors: Arash Gholami Davoodi (Carnegie Mellon University), Navid Rezazadeh (University of California, Irvine), Seyed Pouyan Mousavi Davoudi, and Pouya Pezeshkpour (Megagon Labs). Posted as arXiv:2602.10346v2 [cs.CL], 14 May 2026; labeled with keywords "Machine Learning, ICML." Code is released at https://github.com/arashgholami/top-w-decoding.

What This Paper Is About

When an LLM generates text, it must decide at every step which subset of the vocabulary's next-token probabilities to keep before sampling. Existing truncation rules — Top-k, nucleus (Top-p), locally typical sampling, Min-p, and the entropy-bounded Top-H — make this decision using probability mass and entropy alone, treating tokens as an unordered list of categories with no notion of how semantically similar they are to one another.

Top-W's goal is to make that truncation decision geometry-aware: the crop should be penalized when it forces probability mass to "travel" between tokens that are far apart in embedding space. The paper asks whether explicitly balancing (i) faithfulness to the original distribution under a token-embedding metric, (ii) entropy of the kept set, and (iii) retained probability mass produces a better accuracy/diversity trade-off than probability-only truncators.

Key Contributions

1. A geometry-aware truncation objective with an exact factorization. The paper defines a per-step objective F_{λ,β}(S) = W₁(p, q_S) + λ H(q_S) − β log Γ_S, where q_S is the renormalized crop, Γ_S is the retained mass, and W₁ is Wasserstein-1 under an embedding-induced ground metric. Lemma 3.1 proves an exact factorization of the transport term: W₁(p, q_S) = (1 − Γ_S) · W₁(p(·|S^c), p(·|S)), separating how much mass is removed from how far the removed mass sits from what remains.

2. An exact, linear-time solve of the subset step. For any fixed feasible 1-Lipschitz potential f, Lemma 3.2 shows the S-step reduces to maximizing a single set function G_f(S) that depends on S only through the retained mass. Theorem 3.4 then proves a sharp dichotomy: when β ≥ λ the optimizer is a prefix of tokens sorted by φᵢ = fᵢ + λ log pᵢ (a one-dimensional scan of O(n) instead of a 2^|V| subset search), and when β ≤ λ the optimum collapses to a singleton. Corollary 3.5 shows retained mass is nondecreasing in β in the prefix regime — the theoretical justification for the paper's practical tuning advice to choose β > λ.

3. A practical alternating decoder with cheap, provably feasible geometry. Rather than solving the optimal-transport dual at each step (which the authors note is infeasible at vocabulary scale), the paper uses the anchored, pointwise-minimal feasible potential f_S(i) = −dist(i, S) (Lemma 4.2), combined with a fixed number of alternating f-step/S-step refinements. The standard truncate-and-sample interface is unchanged — the method is implemented as a logits processor.

4. Empirical evaluation across four benchmarks and analysis of reductions to existing samplers. The paper reports a unifying view: under the uniform 0–1 metric, Top-W recovers Top-k (with a cardinality budget |S| ≤ k and λ = β = 0) and becomes the Lagrangian relaxation of the entropy-constrained mass-maximization problem underlying Top-H (with β = 0).

Main Findings

• GSM8K: Top-W wins (or mutually wins) 13/15 (T, model) settings. The comparison spans three models (Qwen2.5-3B, LLaMA-3.1-8B, Phi-3-Mini) and five temperatures T ∈ {0.5, 0.7, 1.0, 1.5, 2.0}. The improvement over Top-H reaches up to 33.74% at T = 2.0. Concretely, at T = 2.0 on LLaMA-3.1-8B-Instruct, accuracy is 73.09 (Top-W) versus 39.35 (Top-H), 2.65 (Top-p), and 13.72 (Min-p).

• GPQA: Top-W wins (or mutually wins) 12/15 (T, model) settings (N = 448, 4 runs). Top-W dominates for all three models at T ∈ {1.5, 2.0} and is top or tied at T = 1.0. Margins shrink at T ∈ {0.5, 0.7}, where Min-p (LLaMA-3.1-8B, T = 0.5) and Top-H (LLaMA-3.1-8B, T = 0.7) occasionally win.

• Probability-only truncators collapse as temperature rises. Min-p and Top-p degrade sharply with increasing T on both reasoning benchmarks, while Top-H and Top-W remain substantially more stable. On GPQA at T = 2.0 with LLaMA-3.1-8B: Min-p 26.12, Top-p 23.88, Top-H 28.79, Top-W 31.02.

• AlpacaEval and MT-Bench: Top-W wins the majority of tuples. In the reported aggregates, Top-W wins 12/15 (AlpacaEval) and 8/15 (MT-Bench) (T, model) tuples out of 5 temperatures × 3 models. The corresponding bar-plot counts are: AlpacaEval — Min-p 0, Top-p 1, Top-H 2, Top-W 12; MT-Bench — Min-p 1, Top-p 4, Top-H 2, Top-W 8. Gaps are larger on MT-Bench, which the authors attribute to multi-turn coherence and instruction retention being more brittle under increased stochasticity.

• Creative-writing rubric scores also favor Top-W, but the reported aggregates differ by configuration. Using GPT-4o as judge on three open-ended storytelling prompts with five rubric dimensions (diversity, originality, narrative flow, emotional impact, imagery), the Table 3 caption (β = 2.8) reports that across all 27 (LLM, T, prompt) triplets, Min-p, Top-p, and Top-H win 8, 5, and 5 cases respectively, while Top-W wins 9. The main text separately describes a higher-β variant where Min-p, Top-p, and Top-H win 6, 5, and 5 settings and Top-W wins 12, with increasing β further improving rubric quality. The paper states Top-W achieves the strongest average rubric score and remains robust at higher T.

• Runtime overhead is described as modest but no numeric value is given. The paper states Top-W adds "only modest runtime overhead (ms/token)" without reporting a specific milliseconds-per-token figure in the provided content.

• Efficient implementation settings. Top-W uses a nucleus-style warm start, a fixed top_m = 1200 candidate pool, alt_iters = 3 alternating refinement steps, and default λ = 2.2, β = 2.8, applied to the temperature-scaled distribution.

• The β sensitivity study (Section 6.2) is not fully available in the provided content. The section is truncated mid-sentence, so the completed ablation results on how Top-W transitions from coherence to diversity as β and λ vary are not reported here beyond Figure 3 (GSM8K accuracy sensitivity of Top-W to β for fixed λ on LLaMA-3.1-8B-Instruct at T ∈ {1.0, 1.5, 2.0}).

Methodology in Plain English

The starting point is a reframing of decoding as distribution shaping. At each token step, the model gives a probability distribution p over the whole vocabulary. A truncation rule picks a subset S, renormalizes to get q_S, and samples from it. The paper asks: what is the best subset to pick?

To answer this, the authors write down a single objective with three competing terms. The first term is a Wasserstein-1 distance between the original distribution and the cropped one, measured using a metric built from the model's own input embeddings — each token embedding is diagonally whitened (to reduce anisotropy) and ℓ2-normalized, and distance between tokens is the Euclidean distance between those whitened vectors. Intuitively, if the mass you throw away sits very close to the mass you keep, the transport cost is low and the crop is "faithful." The second term, λ·entropy, penalizes diffuse crops. The third term, −β·log(Γ_S), rewards keeping a large share of the original probability. Minimizing the total is the goal.

The obstacle is that computing Wasserstein distance exactly at every decoding step is far too expensive. The authors' workaround is a Kantorovich–Rubinstein dual surrogate: instead of the exact transport value, they use a 1-Lipschitz "potential" f that assigns a score to each token, with φᵢ = fᵢ + λ log pᵢ as the combined score. Two theoretical results make this practical. First, once f is fixed, the best subset is no longer a combinatorial search: it is either the top-k prefix after sorting by φᵢ (when β ≥ λ), or a single token (when β ≤ λ). Second, instead of solving for the optimal f, they choose a specific cheap and provably feasible potential, f_S(i) = −dist(i, S), the negative distance from each token to the current kept set — the pointwise-minimal anchored 1-Lipschitz extension. This means "close to what we already kept" translates into a score bonus.

The resulting decoder alternates a few times: score tokens using −dist(i, S) plus λ log pᵢ, sort, scan prefixes to find the best retained-mass cutoff, update S, repeat. The default is three alternations, over a candidate pool of the 1200 most probable tokens. The logits of everything outside the final S are set to −∞, so the surrounding generation pipeline is unchanged.

For evaluation, the authors compare against Top-p (nucleus), Min-p, and Top-H under matched temperature, prompts, max-token budgets, and stopping criteria, and use lm-eval-harness for GSM8K and GPQA. Min-p, Top-p, and Top-H accuracy numbers for GSM8K and GPQA are quoted from Baghaei Potraghloo et al. [2]. AlpacaEval is scored with length-controlled win rates using the standard length-debiasing protocol; MT-Bench uses the average judge score (1–10). GPT-4o serves as the judge, with candidate order randomized per item to mitigate positional bias.

Why This Matters

Impact on research. The paper extends the "decoding as constrained distribution shaping" line of work by showing that the constrained quantity need not be a scalar summary (entropy, probability mass) but can be a full geometric object — a Wasserstein transport cost over embedding space. It also contributes a genuinely useful theoretical result: a proof that a seemingly exponential subset-selection problem has a prefix or singleton solution, plus monotonicity of retained mass in β. The unification result (uniform metric → Top-k and Top-H) gives the family a shared theoretical home.

Real-world applications

• Reliable reasoning and tool use. The paper's largest reported gains are on GSM8K at high temperature (73.09 vs 39.35 for Top-H at T = 2.0 on LLaMA-3.1-8B). Systems that need many diverse samples at high temperature — self-consistency, best-of-n, chain-of-thought voting — benefit most when sampling stays accurate rather than degenerating.
• Creative and open-ended generation. Open-ended writing and ideation are exactly where high temperature is desired; the rubric evaluation on diversity, originality, narrative flow, emotional impact, and imagery targets this use case, and Top-W is reported to achieve the strongest average rubric score.
• Chat assistants and instruction-following products. AlpacaEval length-controlled win rates and MT-Bench multi-turn judge scores are directly relevant to deployed assistants, where drift and inconsistency across turns are the visible failure modes.
• Predictable high-temperature deployment. Because Min-p and Top-p collapse as T rises, practitioners often keep temperature low and lose diversity; a truncator that stays stable up to T = 2.0 widens the usable operating range without retuning the whole pipeline.

Industry relevance. The method is implemented as a logits processor with a fixed alternating-step count, a bounded candidate pool (top_m = 1200), and only distance-to-set queries — no linear program solve per token. That keeps it compatible with existing serving stacks and the standard truncate-and-sample interface, which is the practical bar for adoption of any new sampler.

Future Directions

• Resolving the creative-evaluation bookkeeping. The paper reports two different win counts across the 27 (LLM, T, prompt) creative-writing triplets (8/5/5/9 at β = 2.8 in the Table 3 caption versus 6/5/5/12 for the higher-β variant described in the text), and points to Appendix M and Table 8 for the full breakdown. Reconciling these and clarifying the effect of β on judged creativity is a natural follow-up.

• Completing and extending the (λ, β) sensitivity analysis. Section 6.2's parameter sweep is truncated in the available content, and Figure 3 covers only fixed λ values at three temperatures on one model. A fuller characterization of the accuracy–creativity trade-off surface, and whether the β > λ guidance holds uniformly, remains open.

• Cheaper or better potentials. The implemented decoder uses a deliberately simple anchored potential f_S(i) = −dist(i, S) rather than a Kantorovich–Rubinstein-optimal dual. The paper explicitly notes that an optimal dual is too expensive at each decoding step (Appendix E), so approximations that retain feasibility while capturing more of the true dual are an obvious direction.

• Scale and generality. The reported evaluation uses three instruction-tuned models (3B, 8B, and Phi-3-Mini class) and four benchmarks. Whether the prefix/singleton theory and empirical stability carry to much larger models, longer contexts, and non-English vocabularies — where embedding whitening statistics and candidate-pool sufficiency may shift — is untested in the provided content.

Target Audience

This paper is most valuable to researchers and engineers working on inference-time decoding for LLMs — particularly those who build or tune samplers and are already familiar with Top-k, Top-p, Min-p, and Top-H. It also speaks to theoretically minded readers in optimal transport and structured prediction, because the central results are a factorization lemma, a dual-surrogate bound, a prefix-optimality theorem, and a monotonicity corollary. Practitioners looking for a drop-in logits processor will find the implementation guidance (λ = 2.2, β = 2.8, top_m = 1200, alt_iters = 3, β > λ to avoid singleton collapse) directly actionable, provided they are comfortable working with token-embedding geometry. Readers with no background in Wasserstein distance or Lipschitz duality will find the middle sections demanding, though the objective's three-term interpretation and the alternating-decoder description are accessible on their own.

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