Research map / Large language models / Reasoning and chain of thought

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

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Test-time compute scaling · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Top-W, a geometry-aware truncation sampler that crops the next-token distribution by trading off Wasserstein-1 distance (under an embedding-derived token metric), entropy, and retained mass, solved by a prefix scan with an alternating decoder.

Problem

Probability- and entropy-based truncation samplers for LLMs ignore token-embedding geometry, so they can trade off coherence and diversity heuristically.

Approach

Define a crop objective W1(p,q_S) + lambda H(q_S) - beta log Gamma_S under an embedding-induced ground metric; use a dual lower bound with a distance-to-set potential so that, for fixed potential, the optimal crop is a prefix (or singleton) found by linear scan, and alternate potential and crop updates.

Key results

• On GSM8K, Top-W wins or mutually wins 13/15 (temperature, model) settings, with up to 33.74% improvement over Top-H at T=2.0.
• On GPQA, Top-W wins or mutually wins 12/15 settings, dominating all three models at T in {1.5, 2.0}.
• Runtime overhead is about 5.4% slower than Top-H, Top-p, and Min-p on average.

Surprising findings

• Min-p and Top-p collapse sharply at high temperature (e.g., Top-p on GSM8K with Qwen2.5-3B falls to 9.10% at T=2.0), while Top-W stays near 75%.

Limitations

• Relies on hyperparameters (lambda, beta) tuned on a 50-example AlpacaEval dev split; sensitivity to these is nontrivial.
• Exact KR/OT step is replaced by a distance-to-set surrogate with a fixed number of alternations, so the optimality guarantee does not cover the full W1 objective.
• Creative-writing rubric uses only three prompts and a single judge (GPT-4o).

Experimental setup

Datasets: GSM8K GPQA AlpacaEval MT-Bench

Benchmarks: GSM8K GPQA AlpacaEval MT-Bench creative writing rubric

Metrics: exact-match accuracy length-controlled win rate judge score (1-10) ms/token

Baselines: Top-p Min-p Top-H

Builds on: Top-H Min-p sampling nucleus sampling Kantorovich-Rubinstein duality

Closest papers by meaning

• Top-H Decoding: Adapting the Creativity and Coherence with Bounded Entropy in Text Generation (NeurIPS 2025 · similarity 0.89)
• Foundations of Top-$k$ Decoding for Language Models (NeurIPS 2025 · similarity 0.88)
• Entropy-Gated Latent Recursion (NeurIPS 2026 · similarity 0.87)
• Free Draft-and-Verification: Toward Lossless Parallel Decoding for Diffusion Large Language Models (NeurIPS 2026 · similarity 0.87)
• Entropy-informed Decoding: Adaptive Information-Driven Branching (ICML 2026 · similarity 0.87)
• Optimized Multi-Token Joint Decoding With Auxiliary Model for LLM Inference (ICLR 2025 · similarity 0.86)