💬 NLP

K-Forcing: Joint Next-K-Token Decoding via Push-Forward Language Modeling

K-Forcing 是一种新颖的推前（push-forward）语言建模范式，通过将教师模型蒸馏为能够通过单次前向传递联合解码多个未来标记的条件映射，从而加速自回归文本生成，在仅有轻微质量下降的情况下，实现了高负载批处理服务下的显著推理加速。

原作者： Zhiwei Tang, Yuanyu He, Yizheng Han, Wangbo Zhao, Jiasheng Tang, Fan Wang, Bohan Zhuang

发布于 2026-06-10

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Zhiwei Tang, Yuanyu He, Yizheng Han, Wangbo Zhao, Jiasheng Tang, Fan Wang, Bohan Zhuang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图写一个故事，但你有一位非常严格且才华横溢的编辑（AI 模型），她只允许你一次只能写一个词。

每当你写下一个词，你都必须停下来，把纸递给编辑，等待她读完整个句子，思考下一个词，然后再把纸递回来。然后你写下下一个词，停下，重复这个过程。

这就是目前的 AI（被称为自回归或“AR”模型）的工作方式。它非常聪明，但因为陷入了这种“走走停停”的节奏而变得很慢。如果你想写一个长篇故事，你必须进行成千上上的次往返编辑的桌前。

问题所在：“一次一个词”的瓶颈

论文指出，这种“一次一个词”的方法就像是用小勺子去填满一个游泳池。即使这把勺子很快，整个过程也会受到你往返次数的限制。用计算机术语来说，AI 在进行下一次计算（勺子）之前，必须等待内存（池子）准备就绪。这导致了效率低下，尤其是在许多人同时尝试生成文本时。

旧的解决方案：为什么它们并不完全奏效

科学家们尝试通过两种主要思路来解决这个问题，但两者都有缺陷：

“草拟与检查”法（投机采样/Speculative Decoding）： 想象一个学生试图猜测接下来的几个词，然后由老师进行检查。如果老师同意，那就太好了！如果不同意，学生就必须从头开始。
- 缺陷： 有时学生猜对了 5 个词，有时只猜对了 1 个。这打乱了进度表。当你有一群人在做这件事时，每个人都会失去同步，系统再次变慢。
“扩散”法（Diffusion Method）： 想象你正在通过从空白画布开始并逐渐揭示部分内容的方式来画一幅画，但试图同时猜测多个部分。
- 缺陷： 论文声称，独立地猜测多个部分（比如分别猜测天空和草地）往往会导致一幅混乱的画作，其中天空和草地无法匹配。为了得到完美的画作，你通常仍需一次揭示极小的一部分，这反而抵消了提速的初衷。

新的解决方案：K-Forcing（“神奇蓝图”）

作者引入了 K-Forcing。与其向编辑索要一个词，不如教 AI 查看一张神奇蓝图，并在单次运行中同时写出多个词（例如 4 个词）。

以下是他们实现这一目标的方法，使用了一个简单的类比：

1. “前推”地图（蓝图）
想象你有一个机器，它可以将一个随机数（比如掷骰子）瞬间转化为一个特定的句子。

旧方法： 你掷一次骰子，得到“3”，机器说“那个（The）”。然后你再掷一次，得到“5”，机器说“猫（cat）”。
K-Forcing 方法： 你同时掷四个骰子。机器查看蓝图并说：“好吧，这四个数字对应短语‘那只猫坐下了（The cat sat down）’。”它能瞬间输出所有四个词。

2. 如何获得蓝图？（渐进式自我强迫/Progressive Self-Forcing）
你不能只是瞎猜蓝图；它必须是完美的。因此，他们使用了一种“教师-学生”游戏：

步骤 1： 他们采用缓慢但完美的“教师”AI。他们给它一个随机数，并让它写一个词。他们记录下这对组合：“随机数 0.45” = “那个（The）”。
步骤 2： 他们训练一个“学生”AI 来学习这种联系。
步骤 3（神奇技巧）： 一旦学生学会了写 1 个词，他们就利用这个学生来教自己如何写 2 个词。接着，他们利用 2 个词的学习成果去教自己如何写 4 个词。
为什么这很重要： 与其试图一次性学习整个复杂的规则，不如循序渐进地构建它，就像在脱掉辅助轮之前先学习骑自行车一样。

3. 结果：“批处理”超能力
因为 K-Forcing 每次运行时总是产生固定数量的词（例如，每次恰好是 4 个），所以计算机不会感到困惑。它可以安排 100 个人排队，每个人都能在同一时刻获得他们的 4 个词。

速度： 论文显示，在处理大量请求时，这使得 AI 的速度提升了 2.4 到 3.5 倍。
质量： 生成的文本比缓慢的“教师”稍逊一筹（可能只有 5% 的质量下降），但仍然非常出色，而且速度增益巨大。

总结

可以将 K-Forcing 理解为将快递员从“一次只送一个包裹”升级为“一次性卸下一整托盘包裹的货运卡车”。它并没有改变交付的内容（单词），但它改变了交付的方式，使得整个系统在繁忙时期更加快速高效。

论文证明了，通过改变 AI 对未来进行“思考”的方式（预测一组词而不是单个词），我们可以在不需要新硬件的情况下获得巨大的速度提升。

技术摘要：K-Forcing

问题陈述

自回归（AR）语言模型是文本生成的主流范式，但面临着一个根本性的效率瓶颈：严格的逐 token 顺序解码。这一过程受限于内存带宽，生成长度为 $L$ 的序列需要进行 $L$ 次前向传递，导致硬件加速器利用率不足，尤其是在高负载的批处理（batch）服务场景下，延迟较高。

现有的加速方法在工业级部署中难以有效解决这一瓶颈：

投机采样（Speculative Decoding，即“草拟-验证”模式）： 虽然这些方法能有效降低单次请求的延迟，但它们产生的输出长度是可变的。这破坏了常规的批处理机制，导致同步问题（例如不规则张量/ragged tensors）以及随批大小增长的开销，在重负载下往往会抵消吞吐量的提升。
扩散语言模型（Diffusion Language Models, DLMs）： 这些模型通过迭代去噪来解码多个 token。然而，标准的公式化方法（如掩码扩散语言模型）训练的是逐位置的边缘预测器。理论分析表明，从这些独立的边缘分布而非联合条件分布中进行采样会引入分布误差。为了保持质量，它们通常需要许多次精炼步骤，导致其在减少前向评估次数（NFE）方面相比于 AR 并没有显著优势。

方法论：K-Forcing

作者提出了 K-Forcing，一种旨在实现联合下 $k$ 个 token 解码的推前（push-forward）语言建模范式。其核心目标是学习一个条件映射，将独立的均匀噪声变量转换为单次前向传递中的 $k$ 个未来 token 的联合样本，从而保持固定长度的输出和常规的批处理。

1. 推前公式化（Push-Forward Formulation）

K-Forcing 并非显式地对联合分布 $p(x_{t+1}, \dots, x_{t+k} | x_{\le t})$ 进行建模（因为对于大规模词表而言，这在计算上是不可行的），而是学习一个确定性映射 $G_\theta: \mathcal{V}^t \times [0, 1]^k \to \mathcal{V}^k$ 。

输入： 上下文 $x_{\le t}$ 和 $k$ 个独立同分布（i.i.d.）的噪声变量 $z \sim \text{Uniform}(0, 1)^k$ 。
输出： $k$ 个未来 token 的联合样本 $(\hat{x}_{t+1}, \dots, \hat{x}_{t+k})$ 。
理论基础： 论文证明，通过展开自回归（AR）教师模型的逆累积分布函数（inverse CDF）展开，存在一个闭式推前映射。该映射在理论上可以精确复现 AR 教师的联合条件分布，而无需显式参数化联合概率表。

2. 渐进式自我强制蒸馏（Progressive Self-Forcing Distillation）

为了从预训练的 AR 教师模型中学习该映射，作者引入了渐进式自我强制蒸馏，解决了以往噪声反转基准方法中的两个关键局限（训练-推理不匹配以及数值脆弱性）：

阶段 1（前向蒸馏）： 引导出一个 $k=1$ 的 PFLM。AR 教师通过采样噪声 $z$ 并应用其逆 CDF 采样器来生成目标。学生模型学习将 $(x_{\le t}, z)$ 映射到对应的 token。这确保了推理时的噪声分布（均匀分布）与训练一致，并避免了反转狭窄的 CDF 分箱问题。
阶段 2（自我强制蒸馏）： 为了扩展到更大的 $k$ （例如 $k=2 \to 4$ ），训练好的 PFLM( $k$ ) 作为 PFLM( $2k$ ) 学生的教师。教师使用 $z_{1:k}$ 生成前 $k$ 个 token，然后以生成的这些 token 为条件，使用 $z_{k+1:2k}$ 生成接下来的 $k$ 个 token。学生模型学习在给定完整噪声向量的情况下，一次性预测完整的 $2k$ 序列。这个过程递归地使预测窗口翻倍（ $1 \to 2 \to 4$ ）。

3. 架构设计

作者提出了一个全因果架构，该架构复用了 AR 主干网络：

每个噪声变量被视为序列中的一个独立 token。
一个共享的预测头负责解码未来的 token。
因果注意力机制确保第 $j$ 个未来 token 仅依赖于上下文和前缀噪声变量 $(z_1, \dots, z_j)$ ，这与理论上的逆 CDF 构建过程相一致。
该设计避免了将所有随机决策都通过单个共享隐变量进行路由的信息瓶颈问题（如标准的多 Token 预测架构）。

核心贡献

现有局限性分析： 论文指出，由于输出长度不一，草拟-验证类方法在批处理服务中表现不佳；同时，扩散模型由于是从边缘分布而非联合分布进行采样，导致无法有效减少 NFE。
推前公式化： K-Forcing 被形式化为一个隐式生成模型，它将均匀噪声映射为联合 token 块，满足了固定长度输出和联合采样的要求。
训练策略： 引入了渐进式自我强制蒸馏，消除了与噪声反转相关的训练-推理不匹配和数值脆弱性问题。
实证验证： 证明了仅通过统计建模层面的改变，即可获得相对于 AR 解码在批处理服务上的显著吞吐量增益。

实验结果

在 LM1B 和 OpenWebText 上使用标准的因果 Transformer 主干网络进行了评估：

吞吐量： 当配置为每步前向传递生成 $k=4$ 个 token 时，在不同批大小（4, 16, 128）下，K-Forcing 相比于 AR 实现了 2.4–3.5 倍的加速。
质量： 相对于 AR 教师，生成质量出现了轻微下降。
- 在 LM1B 上， $k=4$ 在成对 LLM-as-a-Judge 评估中的胜率约为 42.9%。
- 在 OpenWebText 上， $k=4$ 的胜率约为 39.4%。
- 较小的 $k$ 值（如 $k=2$ ）提供了更平滑的权衡： $k=2$ 的表现与 AR 几乎无异（胜率约 50%），同时提供约 1.6 倍的加速。
与 SOTA 对比： K-Forcing 在质量-NFE 权衡方面优于 MDLM（MDLM 在解开多个 token 时质量下降剧烈）和投机采样基准（如 Medusa, PTP）。在等效 NFE 下，K-Forcing 实现了更低的生成困惑度（Gen-PPL）和更高的胜率，并且在达到同等质量时比投机类方法使用的前向传递次数更少。

意义与主张

论文声称，K-Forcing 为在真实世界高负载部署下加速 AR 生成提供了一条极具前景的路径。

它解决了现有方法在高负载批处理服务中面临的特定瓶颈，因为这些方法往往会导致不规则的批处理。
它证明了通过推前映射进行联合多 Token 采样是顺序解码的一种可行替代方案，能够保持高效的 KV 缓存管理和连续批处理所需的固定长度结构。
作者将这项工作定位为迈向更广泛研究方向的第一步，即开发推前语言模型。他们承认，与 AR 教师相比，特别是在较大的 $k$ 值时，仍存在不容忽视的质量差距，并指出了未来的研究方向，例如开发定制的块稀疏注意力算子以降低训练成本，以及探索不依赖于预训练 AR 教师的训练范式。