Composable Decoding on the Probability Simplex: Theory and Implementation
Xiaotong Ji, Ahmed Khaled Khamis, Rasul Tutunov, Matthieu Zimmer, Haitham Bou-Ammar
cs.LG, cs.AI
2026-09-28
华为Noah把下一token解码写成单纯形上的正则优化，用镜像上升混搭KL、覆盖与Diversity。Qwen2.5-7B在MATH500上KL+Diversity保住64.4% pass@1、pass@16达90%，纯Diversity则掉到56%；迭代求解约1.13到2.88倍延迟。
LLM 管线的最后一步几乎总是解码，但这一步仍是一堆互不沟通的技巧。工程上在 greedy、temperature、Top-k、Top-p、Min-p、typical、η-sampling 之间靠手感调参。前人已经看到这些变换共享一些性质，也测过质量与多样性的权衡，缺的是把这些偏好写成可配置、可相加、可对照的显式目标。
序列级优化走另一条路。DAEMON 对准文本指标，power sampling 锐化整段分布，Controlled Decoding 依赖学好的 prefix value，共同成本是外部奖励、critic 或额外 rollout。华为 Noah's Ark 与 UCL 把问题收回到每一步的下一 token 分布：先在概率单纯形上选出一个 q，再从 q 里采样或取众数。解码器差在隐式优化什么，不在采样器叫什么名字。
每一步模型给出分数向量 st。解码器要最大化 q 与 st 的内积，减去 λ 倍正则 Ω(q)，并要求 q 落在可行集 Ct 里。Ct 通常是 support 约束：Top-k、Top-p、Min-p 等先圈出候选集合 St，集合外概率钉死为 0。Ω 负责在这块可行面上怎么摊概率。
无正则时最优解把质量全堆在最高分 token 上，即 greedy。负 Shannon 熵给出 softmax，温度等于 λ。同一熵正则限制在 St 上，就是带温度的 Top-k 或 Top-p。二次惩罚给出 Sparsemax，零概率来自边界条件，不必预先指定 support。
组合是把多个 Ωi 按权重 αi 加总（α 非负且和为 1），仍解同一个问题。KKT 条件里，每个正则的梯度对「正则化分数」做一次加法。KL 与熵有闭式。组合后导数叠在一起，坐标往往解不开，于是在单纯形上做镜像上升：势函数取负熵，Bregman 散度变成 KL，更新变成乘法规则 q ← q ⊙ exp(ρ∇f) 再归一化。概率保持非负且和为 1。投影梯度上升用欧氏距离衡量分布，几何不对。
Best-of-K 是一个具体组合。KL 把 q 锚在参考分布 pt（St 上的 softmax）附近。覆盖项来自占有模型：token v 在 K 次独立抽样中至少出现一次的概率是 1-(1-q(v))^K。Coverage 把正权重均匀分给 pt 的 top-r；Diversity 按与最大 logit 的间隙加权，偏好「差一点但还在场」的替代 token。K>1 时，已经很可能会被抽到的 token 再加概率，边际收益递减。温度缩放按同一因子拉分数差，做不到这种非均匀奖励。
CompoSimplex 把 support、正则原语、求解器做成配置，接到 Transformers 与 vLLM 的 logits processor。不额外前向。主实验里 KL 与熵走解析解，其余镜像上升 10 步、步长 0.1。
四组设定：Qwen2.5-7B 配 MATH500，Qwen3-4B-Base 配 GPQA Diamond，LFM2.5-1.2B-Base 配 IFEval，Gemma-4-26B-A4B-IT 配 LiveCodeBench v6。默认每题 16 条样本、温度 0.5、λ=1、组合等权。
MATH500、Top-k（k=200）：
| 方法 | pass@1 | pass@16 | SC@16 | SemDiv |
| Top-k | 59.8 | 90.0 | 78.0 | 0.150 |
| KL | 64.4 | 88.4 | 76.8 | 0.132 |
| Diversity | 56.0 | 91.0 | 76.2 | 0.159 |
| KL+Diversity | 64.4 | 90.0 | 77.2 | 0.149 |
KL+Diversity 对齐 KL 的单样本，pass@16 从 88.4 收到 90.0，贴近 Diversity 的 91.0。Min-p 上 JS+Coverage 的 pass@1 为 67.2，高于 Min-p 基线 64.4，也高于 JS 的 66.6 和 Coverage 的 65.4。
相对对应基线采样器，四组里 pass@1 最大增益是 +10.6 个百分点，来自 GPQA Top-k 上 JS+Diversity 的 26.3 对 15.7；pass@16 最大 +5.1。正则更容易把质量堆到正确补全上，少样本时更明显。
组合并不处处占优。GPQA Top-k 上同一组 JS+Diversity 的 pass@16 是 79.8，低于 Diversity 的 85.4 和 JS 的 80.3。IFEval 的 η-sampling 上 KL+Coverage 把 pass@16 从 80.6 抬到 83.4，all-pass@16 从 17.2 掉到 16.1。LiveCodeBench Top-p 上 JS+Coverage 的 pass@4 到 66.9，高于两边的 65.7 和 64.6，pass@16 落到 67.4，两边都是 69.1。
λ 从 0.5 加到 2，各组合的效用上升、KL 或 JS 下降。纯 Coverage 与 Diversity 在 λ=2 时 pass@1 掉到 41.0 和 39.8；四个等权组合仍停在 57.0 到 62.6。加大效用权重 α 会单调抬效用、压低 KL，对 pass@1 和 pass@16 没有稳定方向。
代价方面，闭式 KL 与基线几乎同速，MATH500 上 2.238 对 2.258 ms/token。无闭式的单正则与组合大约 1.13 倍到 2.88 倍。主实验 10 步、步长 0.1，相对参考最优的平均 L1 在 0.033 到 0.062，是有限步近似。
这是解码期配置。模型权重不动，也不训 critic，逐步求解不依赖未来轨迹。已经在用 logits processor 的栈，可以把 support 和 Ω 换成配置项。适合既要单条质量、又要 K 条里有覆盖的推理场景，Best-of-K 的覆盖项就是按这个预算写的。
多数格子上是 1 到 3 个百分点的挪移，有的指标会低于单正则。能用的是一个可组合的目标空间，外加把 greedy 收到 Sparsemax 的同一套 KKT 推导。延迟按设置加约一到两倍；闭式 KL 和熵几乎不加点。
论文没有单独的 Limitations 节。正文已经写明：正则强度不好选；加大 α 并不稳定抬任务分；强效用正则会打穿准确率；有的指标低于组成成分；主实验用 10 步近似。
覆盖效用是前缀上的 token 占有，不是序列级 Best-of-K，和 self-consistency 或 rerank 的候选质量只有经验联系。评测集中在数学、科学问答、指令遵循和代码，没有开放生成。主表没有和 power sampling、Controlled Decoding 做同设置对照。GPQA 198 题，LiveCodeBench v6 增量 175 题，种子标准差只在 MATH500 Top-k 上用三个种子报过。vLLM 接了接口，主评测走 Transformers。