Top-W: Geometry-Aware Decoding with Wasserstein-Regularized Truncation and Mass Penalties for LLMs¶

会议: ICML 2026
arXiv: 2602.10346
代码: https://github.com/arashgholami/top-w-decoding (有)
领域: LLM 解码 / 评测 / 推理时控制
关键词: 截断解码, Wasserstein 距离, token embedding 几何, 熵约束, 高温稳健性

一句话总结¶

Top-W 把 next-token 截断写成"考虑 token embedding 几何的 Wasserstein-熵-质量"三项最小化问题，理论证明最优解要么是单 token、要么是按 \(f(i)+\lambda\log p_i\) 排序的前缀，工程实现只是 \(O(n\log n)\) 的扫描；在 GSM8K、GPQA、AlpacaEval、MT-Bench 上 15 个 (T, model) 组合多数胜出，高温下 GSM8K 比 Top-H 最多再提 33.7%。

研究背景与动机¶

领域现状：LLM 解码的截断采样早已成为基础设施——Top-\(k\)、Top-\(p\) (nucleus)、Min-\(p\)、locally typical sampling 都从"概率排序"角度裁掉低概率尾部；近期 Top-\(H\) 把"裁出来子分布的熵不超过阈值"显式作为约束，是首批从"分布塑造"视角入手的工作。

现有痛点：所有这些规则都把 token 当作无结构类别——只看概率，看不到 token 在 embedding 空间里的语义距离。结果：(i) 高温度时（\(T\geq 1.5\)）Top-\(p\) / Min-\(p\) 频繁拉到几乎整词表，输出崩坏；(ii) 即便加了熵控制（Top-\(H\)），仍可能把概率集中在同义/近义的近邻 token 上，"伪多样性"而失去真创造力。

核心矛盾：解码器要平衡 (i) faithfulness（不能太远离原分布），(ii) creativity（要够 diverse），(iii) coherence（不能裁太多质量）。前两者本质上要在 token 几何空间里衡量，但所有现成 sampler 都跳过了几何信息。

本文目标：把 token embedding 几何"显式"塞进截断目标，给出一个既有理论闭式、又能用 logits-processor 接口部署、对温度稳健的几何感知 sampler。

切入角度：作者从最优传输 (OT) 视角看截断——把"截断 + 重归一"看作把原分布 \(p\) 运输到一个支撑在 \(S\) 上的 \(q_S\)，自然引入 Wasserstein-\(1\) 距离 \(W_1(p,q_S)\) 作为 faithfulness 项，且 \(W_1\) 可用 token embedding 上的 Mahalanobis 距离做地面成本。

核心 idea：用"\(W_1\)（几何） + \(\lambda H(q_S)\)（创造力） − \(\beta\log\Gamma_S\)（质量）"三项的最优化定义最佳截断集，并证明这个问题有"前缀 / 单 token"的结构性解。

方法详解¶

整体框架¶

给定下一 token 分布 \(p\in\Delta^{|V|}\) 与 token embedding 几何（whitened embedding 的 Mahalanobis 距离），Top-W 每步：(a) 计算每个候选集 \(S\) 的目标 \(F_{\lambda,\beta}(S)=W_1(p,q_S)+\lambda H(q_S)-\beta\log\Gamma_S\)；(b) 由于 \(W_1\) 在 vocab 规模上不可解，改用 KR 对偶替身：\(W_1=\sup_{f\in\mathcal{F}}(\mathbb{E}_p[f]-\mathbb{E}_{q_S}[f])\)，再选 anchored potential \(f_S(i)=-\mathrm{dist}(i,S)\) 作为最攻击性的可行势；(c) 交替更新 \(f\)（用当前 \(S\) 算 distance-to-set）和 \(S\)（按定理 3.4 用 1D 前缀扫描），3-4 轮收敛后把不在 \(S\) 内的 logits 置 \(-\infty\) 走标准采样。

关键设计¶

Wasserstein-熵-质量三项目标 + 精确分解:
- 功能：把"几何忠诚 + 创造力控制 + 质量保留"三个目标统一进单一可优化目标。
- 核心思路：定义 \(F_{\lambda,\beta}(S)=W_1(p,q_S)+\lambda H(q_S)-\beta\log\Gamma_S\)。论文证明 \(W_1(p,q_S)=(1-\Gamma_S)W_1(p(\cdot|S^c),p(\cdot|S))\)，把"删除的质量"和"删除部分与保留部分的几何距离"分离；从而高概率但与保留集语义远的 token 会被强烈惩罚（避免裁掉同义近邻反留下噪声），低概率但靠近保留集的 token 反而可能进 \(S\)。
- 设计动机：单纯概率截断忽略 token 之间的语义相似性，会把"近义同义词聚团"或"语义孤岛 outlier"一视同仁；引入 \(W_1\) 让截断"看得见"语义结构。
几何锚定势 + 闭式 S-step:
- 功能：把不可解的 KR 对偶 LP 换成几乎零成本的距离查询，并保证截断子问题有闭式最优。
- 核心思路：固定一个可行势 \(f\)，论文证明 \(\arg\min_S F\) 等价于 \(\arg\max_S G_f(S)=\frac{1}{\Gamma_S}\sum_{i\in S}p_i\phi_i(f)+(\beta-\lambda)\log\Gamma_S\)，其中 \(\phi_i(f)=f_i+\lambda\log p_i\)；进一步证明 (i) 若 \(\beta\geq\lambda\)，最优 \(S\) 必是按 \(\phi_i\) 排序的某个前缀；(ii) 若 \(\beta\leq\lambda\)，最优 \(S\) 退化为单 token。把势固定成 \(f_S(i)=-\mathrm{dist}(i,S)\) ——所有 anchored 1-Lipschitz 函数里"最吸引"的那个，combined score 就是 \(\phi_i=-\mathrm{dist}(i,S)+\lambda\log p_i\)。
- 设计动机：原始 \(W_1\) 在 vocab 规模 \(10^5\) 上做 LP 不现实；用 distance-to-set 作为势既保证 1-Lipschitz 可行性，又把"远离当前 set 的 token"明确加负分。S-step 闭式让整个 sampler 成本只多一次排序 + 一次前缀扫描。
交替 f-step / S-step + 候选池剪枝:
- 功能：在不显式解 OT 的前提下逼近联合最优解。
- 核心思路：alt 循环 (i) 用当前 \(S^{(t)}\) 算 \(f^{(t)}_i=-\mathrm{dist}(i,S^{(t)})\)；(ii) 按 \(\phi_i^{(t)}\) 降序排序，前缀扫描 \(J_k=\Phi_k/\Gamma_k+(\beta-\lambda)\log\Gamma_k\) 取 \(k^\star\)；(iii) 收敛即停。为了不在全 vocab 上算距离，先用 nucleus warm-start 限制到 top_m=1200 候选池，论文附录给出"剪枝下仍精确"的充分条件。
- 设计动机：3 轮迭代足以收敛；候选池让每 token 只需 ms 级开销，实测比 Top-\(H\)/Top-\(p\)/Min-\(p\) 平均只慢 5.4%。

损失函数 / 训练策略¶

本文是推理时方法，无训练；唯一超参 \((\lambda,\beta)\) 默认 \((2.2,2.8)\)；当 \(\beta>\lambda\) 进入 prefix 区间，可通过升降 \(\beta\) 在 sharpness ↔ creativity 之间滑动。

实验关键数据¶

主实验¶

3 个 LLM（Qwen2.5-3B、LLaMA-3.1-8B-Inst、Phi-3-Mini）× 5 个温度 \(T\in\{0.5,0.7,1.0,1.5,2.0\}\)，共 15 组合：

Benchmark	Top-W 胜场	最大相对提升 vs Top-H	备注
GSM8K	13/15	+33.7%（\(T=2.0\)）	高温下 baselines 几乎崩
GPQA	12/15	普遍 1-3 点	全 3 模型 \(T\in\{1.5,2.0\}\) 都赢
AlpacaEval	12/15	judge 评分稳赢	length-controlled win-rate
MT-Bench	8/15	多轮一致性更好	高温保不漂移

GSM8K 在 \(T=2.0\) 上：Top-W 75.13% / 73.09% / 84.63%，而 Top-\(p\) 跌到 9.10% / 2.65% / 7.73%。

消融实验¶

配置	GSM8K@T=2.0 (LLaMA)	说明
\(\beta>\lambda\)（前缀区间）	73.09	默认设置
\(\beta\leq\lambda\)（singleton）	显著下降	退化为单 token
\(\beta\) 过大	创造力 ↑ 但 GSM8K ↓	留太多 mass
Top-W (creative rubric \(\beta=2.8\))	27 设置赢 12	比 Top-\(p\)/Top-\(H\)/Min-\(p\) 平均高

关键发现¶

几何 + 熵 + 质量三项缺一不可：仅有质量项 → Top-\(k\)；仅有熵 → Top-\(H\)；加入几何后高温稳健性质变。
\(\beta\) 是"创造力 ↔ 准确率"调节器：rubric 评测（Diversity/Originality/Narrative/Emotion/Imagery 5 维）显示提高 \(\beta\) 创造力升、严格答案分降；为不同任务可分别调。
统一视角：论文证明在 0-1 uniform metric 下 Top-W 退化成 Top-\(k\)（加 \(\lambda=\beta=0\)）或 Top-\(H\)（\(\beta=0\) 的 Lagrangian 松弛），把现有 sampler 纳入同一框架。
运行开销可控：3 轮 alt × top_m=1200，每 token 约 ms 级，比 Top-\(p\) 慢 5.4%——几何感知不是吞吐杀手。

亮点与洞察¶

结构性最优解的发现：Theorem 3.4 把组合搜索 \(2^{|V|}\) 降到 1D 扫描，对任何"加权平均 + 凹凸 mass 项"的截断目标都可复用，是个一般性技巧。
OT 视角统一截断 sampler：把 Top-\(k\)/\(p\)/\(H\) 都看成 \(W_1\)+熵+mass 不同特例，给后续解码研究一个统一坐标系。
anchored potential 的"白名单"思路：用 1-Lipschitz 包络做替身，避开 LP 求解，是把 OT 落地工程的关键妙手；这种"用 distance-to-set 当势"的招式对其他 OT-on-discrete 问题也有借鉴价值。
以"温度稳健性"作为新评测维度：以前 sampler 论文很少在 \(T=2.0\) 报告，本文系统地展示了几何感知截断的高温抗崩盘能力，是评测范式上的进步。

局限与展望¶

\(W_1\) 用 token embedding Mahalanobis 距离作 ground cost，但 LLM 的 embedding 实际并不严格反映"语义距离"——多义词、复合 token、罕用 token 可能误导几何。
候选池 top_m=1200 是经验值，对超大词表（>200k）或代码 token 仍可能漏掉远但合理的 token。
\((\lambda,\beta)\) 需要事先按任务调；论文给了 sensitivity 但没自动化方案，工业部署仍需调参。
实验集中在 instruction-tuned 模型 + QA / chat 场景，对 code generation、long-context summarization 等任务的效果未验证。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次把 token embedding 几何 + OT 视角带进截断 sampler，并证明结构性最优解。
实验充分度: ⭐⭐⭐⭐ 4 个 benchmark × 3 模型 × 5 温度共 60 组合 + rubric 创作评测 + 运行开销分析，覆盖充分；但缺代码生成等场景。
写作质量: ⭐⭐⭐⭐ 证明排版清晰、算法伪代码完整；少数符号（\(\phi,c,\beta-\lambda\)）来回切换略费脑。
价值: ⭐⭐⭐⭐⭐ 解码侧免训练即插即用且对高温稳健，是直接可用的 deployment 级别提升。