Research map / Learning theory and algorithms / Matrix and tensor methods
Foundations of Top-$k$ Decoding for Language Models
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.
Develops a theory of top-k decoding as l0-regularized separable Bregman divergence minimization (primal and dual), proves greedy selection and discrete convexity in k, and evaluates alpha-Bregman decoders against top-k on GPT-2 Large and Llama 3.1 8B.
Problem
Top-k decoding is widely used for LLM sampling but lacks a precise theoretical justification; choosing k adaptively per distribution is typically heuristic.
Approach
Views per-token decoding as sparse distribution recovery: minimize a separable Bregman divergence (primal or dual argument) plus an l0 penalty; shows the optimal support is top-k (greedy) and the cost is discretely convex in k, so the adaptive k* can be found in logarithmic time. Top-k is recovered as the alpha=1 primal case under KL.
Key results
- Primal alpha=2.0 decoding shows the smallest perplexity and repetition gaps versus human text among top-k and alpha in {1.5, 2.0} on GPT-2 Large and Llama 3.1 8B (partial evaluation, k in 5 to 50).
- On GSM8K with Llama 3.1 8B at temperature 0.3, primal decoding with adaptive k* reaches about 85.1% (alpha=1.5, lambda=0.01) versus 83.6% for matched top-k.
- At temperature 1.5 on GSM8K, top-k accuracy drops to 43.2% to 76.0% while primal alpha decoding retains 57.2% to 80.1%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use adaptive primal Bregman decoding instead of matched top-k decoding | Llama 3.1 8B · GSM8K, 8-shot CoT, temperature 1.5 | accuracy: 79.00 vs 75.97; 80.06 vs 75.97; 57.24 vs 43.21; 64.97 vs 58.53 | Helped |
| Use adaptive primal Bregman decoding instead of matched top-k decoding | Llama 3.1 8B · GSM8K, 8-shot CoT, temperatures 0.3–1.0 | accuracy: 85.14 vs 83.62; 84.38 vs 84.69; 83.24 vs 83.78; 81.73 vs 84.69; 84.69 vs 85.14; 84.46 vs 83.62; 82.03 vs 82.11; 82.03 vs | Mixed |
Surprising findings
- At high temperature (1.5), top-k degrades much more rapidly on GSM8K than primal alpha-decoding with adaptive k.
Limitations
- Experiments are small-scale and aimed only at showing competitiveness; no comparison to other established decoders such as top-p or min-p.
- Dual decoding is not evaluated empirically; the dual objective can be non-convex.
- Theory covers only per-token decoding with separable Bregman divergences and does not analyze top-p or min-p.
Experimental setup
Datasets: WebText test set (GPT-2 output dataset) GSM8K Chain-of-Thought TriviaQA
Benchmarks: WebText open-ended generation GSM8K
Metrics: perplexity difference repetition difference accuracy
Baselines: top-k decoding
Builds on: top-k sampling Bregman projection onto the simplex Tsallis alpha-entropy
Closest papers by meaning
- Geometry-Aware Decoding with Wasserstein-Regularized Truncation and Mass Penalties for Large Language Models (ICML 2026 · similarity 0.88)
- Confidence-Based Decoding is Provably Efficient for Diffusion Language Models (NeurIPS 2026 · similarity 0.88)
- Top-H Decoding: Adapting the Creativity and Coherence with Bounded Entropy in Text Generation (NeurIPS 2025 · similarity 0.87)
- Sampling-Efficient Test-Time Scaling: Self-Estimating the Best-of-N Sampling in Early Decoding (NeurIPS 2025 · similarity 0.87)
- LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding (ICML 2026 · similarity 0.86)
- Entropy-informed Decoding: Adaptive Information-Driven Branching (ICML 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.