快速分叉：高效估计文本生成中的不确定性动态

# 高效估计文本生成中的不确定性动态
来源：https://arxiv.org/html/2608.19611
Amir Zur、Satchel Grant、Tal Haklay、Can Rager、Owen Lewis
附属机构：Thomas McGrath、Jack Merullo、Ekdeep Singh Lubana、Atticus Geiger
附属机构：![[未命名图片]](https://arxiv.org/html/2608.19611v1/gf_color.png)

###### 摘要
大语言模型的推理是随机的，因此要理解一个模型，需要处理它对于给定问题可能产生的推理链分布，即其不确定性。基于重采样的分析可以描述这种分布，揭示出 rollout 的哪些步骤决定了模型如何得出答案。然而，这类方法的一个主要局限性在于，在推理链的每个 token 或句子处重采样文本序列的成本非常高。我们的工作旨在提高重采样分析的计算效率，同时阐明一个重要的科学问题：解释文本生成中不确定性动态的正确统计模型是什么？我们发现，在对大量推理链进行重采样时，不确定性动态会收敛到稳定的模式，噪声在很大程度上是采样的产物，而非 LLM 对每个单独 token 或推理步骤的敏感性所致。我们开发了一个统计模型，用于平滑带噪声的低样本 rollout 数据，以更好地逼近高样本数据，从而显著降低采样成本。

## 1 引言
††脚注文本：代码与交互式仪表盘：https://github.com/ericb-goodfire/forking-fast
当大型语言模型进行长篇文本生成（如推理）时，它们在每一步或每个生成的 token 处都可能采取许多不同的路径。先前的工作表明，推理过程中的特定步骤可能对模型的整体行为产生重大影响，而许多其他步骤则不那么重要(1 (https://arxiv.org/html/2608.19611#bib.bib1);3 (https://arxiv.org/html/2608.19611#bib.bib8);15 (https://arxiv.org/html/2608.19611#bib.bib15);11 (https://arxiv.org/html/2608.19611#bib.bib9);4 (https://arxiv.org/html/2608.19611#bib.bib14))。理解推理（以及更广泛意义上的文本生成）中哪些步骤最重要的一种方法是重采样：给定单个推理 rollout，在推理的每一步重采样替代 rollout，然后收集所有 rollout 的最终答案，并将它们聚合成一个分布。诸如分叉路径分析这样的重采样方法可以提供关于文本生成中不确定性动态的关键见解，其中不确定性衡量了 LLM 选择某个答案而非另一个答案的可能性。然而，这种方法成本非常高，分析单个推理链通常需要数百万个 token。在这项工作中，我们提出了一种更有效的方法来近似这些不确定性动态，方法是开发一个描述其分布的统计模型，并利用该模型平滑带噪声的低样本估计。我们通过在 tinyMMLU 问题上研究 Llama-3-8B-Instruct(7 (https://arxiv.org/html/2608.19611#bib.bib6)) 和原生推理模型 DeepSeek-R1-Distill-Llama-8B(6 (https://arxiv.org/html/2608.19611#bib.bib7)) 的思维链推理(10 (https://arxiv.org/html/2608.19611#bib.bib13))来评估我们的方法，这些问题来自 tinyMMLU(12 (https://arxiv.org/html/2608.19611#bib.bib4);8 (https://arxiv.org/html/2608.19611#bib.bib5))。在收集并分析了总计近二十亿 token 的推理数据后，我们发现三个主要结果：
1.  当重采样次数较少时，推理 rollout 的不确定性动态噪声较大(1 (https://arxiv.org/html/2608.19611#bib.bib1))。然而，当重采样数百次时，不确定性动态会变得越来越平滑，除了在具有急剧变化的关键分叉点。我们发现，推理 rollout 之间的变化可以很好地建模为多项式采样噪声（第 3.3 节 (https://arxiv.org/html/2608.19611#S3.SS3)），并且重构误差与收集的样本量的平方根成比例衰减。
2.  我们开发了一个统计模型，允许我们通过平滑低样本来近似高样本不确定性动态。该模型使用变化点检测来识别分叉点，并使用核池化来平滑这些点之间的估计。通过平滑低样本数据，我们可以有效地将有效样本量乘以 3.3 倍（S=30）到 5 倍（S=5）（第 3.1 节 (https://arxiv.org/html/2608.19611#S3.SS1)）。
3.  我们发现，每隔 N 个 token 或步骤（而非每一步）进行重采样，可以在不进行平滑的情况下提高数据效率。然而，这种方法在估计分叉点时损失了精度，并且与 N=1 相比，从平滑中获益较少。通过将此方法与我们的平滑模型结合，我们可以将总预算削减至 1/8，而误差仅略有增加（第 3.2 节 (https://arxiv.org/html/2608.19611#S3.SS2)）。

图 1：一个问题的不确定性动态 o_t（tinyMMLU 问题 39；Llama-3-8B-Instruct）。颜色代表在时间 t 处以不同最终答案结束的 rollout 的比例。
(a) 当在每个 token（N=1）处收集大量重采样 rollout（S=1000）时，结果分布变得越来越平滑，除了在分叉点。
(c) 使用参考采样成本的 1% 进行较低样本分析（S=20, N=2）的不确定性动态。
(b) 平滑较低样本数据 o_t 与高样本数据非常接近，并恢复了相同的分叉点。更多示例见附录 B (https://arxiv.org/html/2608.19611#A2) 和我们的交互式仪表盘 (https://github.com/ericb-goodfire/forking-fast)。

## 2 方法

#### 分叉路径分析
遵循(1 (https://arxiv.org/html/2608.19611#bib.bib1))，我们分析相对于固定推理链 x（由贪婪解码生成的基础路径）的不确定性动态。在每个 token 或句子位置 t，每间隔 N 个（例如，每个 token 或每个句子，N=1），我们从前缀 x_t 以温度 τ=1.0 重采样延续，并提取每个延续的结果，在我们的情况下是多项选择答案 A/B/C/D/其他。然后我们将这些最终答案聚合成一个关于结果的加权分布，使得每个时间步 t 索引一个关于最终答案的分布 o_t，其权重由 token 的概率 p(x_t | x_{≤t}) 决定，这可以可视化为时间序列（图 1 (https://arxiv.org/html/2608.19611#S1.F1)）。在这项工作中，我们考虑两种降低 FPA 采样成本的方法：减少每个 token 位置的样本数 S，以及增加采样间隔，每隔 N 个 token 或句子进行重采样，而不是每一步（即 N=1）。我们的目标是使用这些低样本来逼近具有高 S 和 N=1 的不确定性动态 o_t。

#### 估计不确定性动态
结果分布 o_t 中的急剧变化何时对应于推理路径中的决策点，何时对应于采样噪声？图 1 (https://arxiv.org/html/2608.19611#S1.F1)a 显示了具有高重采样率（S=1000 且 N=1）的参考 o_t。我们发现，拥有大量重采样数据时，结果分布 o_t 在几乎所有位置 t 处都变得越来越平滑，除了关键的分叉点，在这些点 o_t 会突然且显著地改变(1 (https://arxiv.org/html/2608.19611#bib.bib1))。与图 1 (https://arxiv.org/html/2608.19611#S1.F1)c 相比，后者显示了 S=20 且 N=2 的较低样本结果分布 o_t。很难从低样本数据中直观辨别 o_t 中的急剧变化是真实的不确定性动态，还是这些波动是采样噪声的产物。由于我们对每个可能的下一个 token x_t = w 独立抽取 S 个延续，答案计数形成一个多项分布 c_t^{(w)} ∼ Multinomial(S, o_t^{(w)})，并且分支权重 \tilde{p}_w 可以从下一个 token 分布 p(x_t = w | x_{≤t}) 中精确获知。TVD 允许我们量化两个分布之间的差异：TVD(o_t, o_{t+N}) > ε，即两个相邻步骤之间结果分布存在显著差异的点。我们在分叉阈值 ε ∈ {0.10, 0.15, 0.20} 下报告结果。我们观察到，随着观测间距变宽且 N 增加，分叉区域的性能比整体性能更早、更快地恶化。在 S=30 时，随着间距从 N=4 增加到 N=8 再到 N=16 个 token，分叉区域的 TVD 从 0.0669 增加到 0.0936 再到 0.1587。然后在 32-64 个 token 的间距附近饱和于约 0.16，表明分叉实际上已被错过。平坦区域性能退化更慢，并且对于两个模型，在规模化时都保持相同的排序。其次，完整模型相比原始数据在分叉附近的优势集中于中等分叉（阈值 0.10），对于较大的分叉则消失，甚至对于最大的分叉会逆转（第 3.2 节 (https://arxiv.org/html/2608.19611#S3.SS2)）。在阈值 0.15 的面板中，实线和虚线曲线在密集间距下几乎重合。这种一致性反映了预先注册检验的功效性阴性结果，而非绘图伪影。

图 7：对应于图 2 (https://arxiv.org/html/2608.19611#S2.F2) 的分叉区域。仅在分叉区域评估相对于 S=200 参考的 TVD，使用分叉阈值 0.10（左）和 0.15（右）。实线显示平滑数据，虚线显示原始数据；颜色与正文匹配。

图 8：对应于图 3 (https://arxiv.org/html/2608.19611#S3.F3) 的分叉区域。绘制分叉区域的 TVD 与每个问题生成的 token 数的关系，分叉阈值为 0.10 和 0.15。实线显示平滑数据，虚线显示原始数据。

## 附录 E 完整模型的组件消融实验
在 S=200 tinyMMLU 数据集上对完整模型的组件进行逐一消融，揭示了哪些组件有助于平坦区域和分叉区域的准确性。我们报告了相对于完整模型的配对逐问题 TVD 差异，使用留一复制参考计算，并给出 95% t 区间。

### E.1 估计量
我们比较了四个复杂度递增的估计量：
- 原始数据：使用观测位置之间的线性插值的经验频率。
- 核池化：使用高斯核加权的相邻计数狄利克雷池化。该估计量设计用于分叉之间的区域，其中增量较小且附近位置具有相似的 o_t 值。
- 分段+池化：使用精确多项式代价的 PELT 变点检测(9 (https://arxiv.org/html/2608.19611#bib.bib2);14 (https://arxiv.org/html/2608.19611#bib.bib3))，随后进行逐段狄利克雷池化。所得的分段常数拟合在检测到的分叉处保留边界。
- 完整模型：PELT 分段后进行截断于变化点边界的核池化。该估计量适应段内的逐渐漂移，同时在分叉处保留急剧变化。惩罚和带宽超参数通过在低样本运行上交叉验证选择。

### E.2 进一步的模型消融
分段保留了分叉。移除分段并仅使用核池化保留了 93-96% 的池化改进，但在分叉区域准确性下降。相对于完整模型的 TVD 差异对于 Llama 在分叉阈值 0.15 时为 +0.0113 [+0.0039, +0.0187]，对于 DeepSeek 在阈值 0.10 时为 +0.0091 [+0.0036, +0.0145]。图 10 (https://arxiv.org/html/2608.19611#A5.F10) 显示了两个核池化模糊了完整模型所保留的分叉的例子。在记录的 DeepSeek 示例中，相应的 TVD 分别为 0.2720 和 0.0088。
交叉验证调整在大分叉中很重要。用固定值替换交叉验证的超参数会将 TVD 增加 0.0126 [+0.0051, +0.0202]，针对 Llama 在阈值 0.15 时。
核池化提供了一小部分但一致的池化增益。用平坦的逐段池化（分段+池化）替换核池化会将 Llama 在 S=30 时的池化 TVD 增加 0.0021 [+0.0015, +0.0026]。
检测代价对端到端影响很小。将 PELT 固定为 L2 代价在整个消融过程中使池化 TVD 增加 ≤0.0010。在两种主要分叉条件下，其性能与完整模型在统计上无法区分；例如，在 S=30 且阈值 0.10 时 Llama 的差异为 +0.0003 [-0.0004, +0.0009]。该变体保留了完整模型相对于原始数据 97-100% 的改进。

图 9 总结了这些配对差异。对于池化准确性，每个平滑变体都在完整模型的几千分之一之内，而原始数据具有明显更高的 TVD。在分叉区域，无分段变体与零的差异是可靠的，而 L2 检测变体则不是。

图 9：DeepSeek-R1-Distill-Llama-8B 在 1 句子间距下，每个消融变体与完整模型的配对逐问题 TVD 差异。误差条显示 tinyMMLU 上的 95% t 区间，虚线零线表示完整模型。结果如图所示（池化准确性（左），阈值 0.10 时的分叉区域准确性（右））。移除分段会损害分叉附近的准确性，但不会损害池化准确性；将检测代价固定为 L2 在两个面板中均与完整模型无法区分。

图 10：两个示例问题的消融叠加。重建使用 S=30 在最密集的间距 N 上，并与留一复制 S=200 参考进行评估；虚线标记参考分叉。左列显示 Llama-3-8B-Instruct 问题 46 使用的消融变体、交叉验证以及记录的 S=200 存储中的操作点。右列显示记录的 DeepSeek 示例，其中分叉发生在延续停止解析为答案时。移除分段（核池化）模糊了分叉，而移除核池化（分段+池化）将逐渐漂移减少为离散步骤。L2 检测代价变体在视觉上与完整模型无法区分。

## 附录 F 1/√S 噪声定律
图 11 (https://arxiv.org/html/2608.19611#A6.F11) 展示了第 3.3 节 (https://arxiv.org/html/2608.19611#S3.SS3) 背后的复制噪声测量。对于 S=1000 开发集，不相交复制之间的池化 TVD 绘制在对数-对数坐标轴上与 S 的关系图。拟合斜率总体为 -0.4903，在尾部为 -0.5047。该图还显示了精确的独立同分布多项式零假设，通过使用相同统计量重采样每个分支的直方图计算得出；测量 TVD 与零假设 TVD 的比率范围从 0.98 到 1.01。最后，来自 S=200 tinyMMLU 数据集的每个模型的平均值的斜率分别为 -0.4767 和 -0.4577。它们的绝对 TVD 较低，因为这些问题平均变异性较小，但误差以相似的速率下降。

图 11：复制 TVD 与每个位置样本数 S 的关系图，使用对数-对数坐标轴。测量曲线在 S=200 之前始终保持在精确独立同分布多项式零假设的 2% 以内。三角形显示两个模型的全数据集平均值。

## 附录 G PELT 实现中的分段代价
在本研究中，