【论文笔记】Top-nσ: Not All Logits Are You Need

小嗷犬2024-12-292024-12-29

基本信息

标题: Top- $n\sigma$ : Not All Logits Are You Need
作者: Chenxia Tang, Jianchun Liu, Hongli Xu, Liusheng Huang
arXiv: https://arxiv.org/abs/2411.07641

摘要

大型语言模型（LLMs）通常使用贪婪解码或低温采样进行推理任务，这反映了一种在多样性和准确性之间的权衡观念。

我们通过引入 top-nσ 这一新颖的采样方法挑战了这一传统方法。该方法直接基于 pre-softmax logits，利用统计阈值进行操作。

我们的核心见解是，logits 自然分为高斯分布的噪声区域和独立的信息区域，从而在不进行复杂概率操作的情况下，实现了高效的token过滤。

与现有方法（例如，top-p、min-p）在较高温度下无意中包含更多噪声token不同，top-nσ 能够在不同温度缩放情况下保持稳定的采样空间。

我们还从理论上分析了 top-nσ 的行为，以更好地理解其特性。

通过在四个专注于推理的数据集上的广泛实验结果表明，我们的方法不仅优于现有的采样方法，还在高温下保持了一致的性能，同时超越了贪婪解码的表现。

引言

LLAMA3-8B-Instruct 在 AQuA 样本上的 logits 分布及其按概率降序排列的分布。请注意，右图中概率较高的前导 token 对应于 logits 分布的右侧区域。最大 logit 大约比分布的平均值高出 10σ。

主要贡献

基于 Logit 的新视角：我们引入了一个专注于预软化 logit 分布的新分析框架，不仅为采样策略开发提供了基本见解，还为模型训练方法的潜在改进提供了可能性。
高效的 Top-nσ 算法：我们提出了一种概念上简单但功能强大的采样方法，该方法直接作用于 logits，在保持计算效率的同时，通过免排序和免 softmax 的实现，获得了更高的生成质量。将 top-nσ 集成到现有的 LLM 框架中非常容易。
适用于测试时扩展技术：我们的 top-nσ 算法能够更精细地探索解空间，在探索和利用之间实现更好的平衡。这一特性使其在测试时扩展技术中尤其有效。
理论分析：我们提供了对 top-nσ 的全面定量分析，研究了其累积概率质量特性，并证明了其温度不变性属性。这些理论基础为该方法的实现和理解奠定了坚实的基础。
广泛的实证验证：我们通过对四个不同数据集的严格实验，证明了我们方法的有效性，与现有方法相比，显著提升了生成质量。

洞察

从图1中，我们观察到模型的 logits 似乎遵循一个定义明确的统计分布。这一经验观察自然引出了几个基本问题：

我们如何从 logit 空间的角度解释基于概率的采样方法（例如，核采样）？
大型语言模型中 logit 分布的基本特征是什么？
我们如何利用这些分布来有效地区分嘈杂和信息丰富的区域？

在本节中，我们将系统地回答上述问题并展示我们的发现。

从概率到 logits

现代大型语言模型（LLMs）通过两步过程生成文本：首先生成 logits，然后通过 softmax 变换将它们转换为概率。

为了更好地理解 logit 空间中的采样方法，我们首先考察传统的基于概率的采样方法，特别是核采样，如何从 logit 的角度重新解释。

给定一个输入序列，LLM首先生成一个 logit 向量 $l = (l_1, \cdots, l_V) \in \mathbb{R}^V$ ，其中 $V$ 是词汇表大小。

这些logits首先按温度缩放： $l \leftarrow l / T$ ，然后通过softmax函数转换为概率 $p = (p_1, \cdots, p_V) \in \mathbb{R}^V$ ：

$p_i = \frac{e^{l_i}}{s}, \quad \text{where } s = \sum_{j=1}^V e^{l_j}, 1 \leq i \leq V$

从根本上说，所有采样方法都通过确定一个概率阈值 $p^{(t)} \in [0, 1]$ 来操作。概率高于此阈值的token形成采样核，其累积概率定义了核质量。

形式上，对于阈值 $p^{(t)}$ ，核 $\mathcal{N}$ 是：

$\mathcal{N} = \{i \mid p_i \geq p^{(t)}\} = \{i \mid l_i \geq t\}$

这种基于概率的选择可以在 logit 空间中等效地执行，通过确定相应的logit阈值 $t = \ln(s \cdot p^{(t)})$ 。

基于我们在图1中的经验观察，我们假设logits遵循某个分布 $f$ 。在这一假设下，我们可以推导出核采样参数和logit阈值 $t$ 之间的关系。

定理1 考虑 $V$ 个独立同分布的 logits $\{l_i \mid i = 1, \cdots, V\}$ ，它们根据 $f(x)$ 分布。对于任何阈值 $t$ ，我们有：

$\sum_{l_i > t} e^{l_i} \xrightarrow{P} V \int_t^{+\infty} e^x f(x) \, \mathrm{d}x$

完整的证明在附录 A.1 中提供。尽管该定理本身概念上简单，但它引发了一系列具有实际意义的强大引理：

引理2 设 $\mathcal{I}(t) = \int_t^{+\infty} e^x f(x) \, \mathrm{d}x$ ，因此 $s = V \cdot \mathcal{I}(-\infty)$ 。给定阈值 $t$ 的核质量为

$p_\mathcal{N} = \sum_{i \in \mathcal{N}} p_i = \frac{\mathcal{I}(t)}{\mathcal{I}(-\infty)}$

注释该引理的一个关键含义是，给定 logits 的概率分布，核采样的 logit 阈值可以解析地导出。我们通过推导两个基本案例的闭合形式表达式来展示这一点。给定累积概率阈值 $p$ ：

高斯分布： $l_i \sim N(\mu, \sigma^2)$ 。阈值为 $$t = \mu + \sqrt{2} \sigma (\operatorname{erf}^{-1}(1 - 2p)) + \sigma^2$$ 其中 $\operatorname{erf}(\cdot)$ 表示误差函数。
均匀分布： $l_i \sim U(M - a, M)$ ，阈值为 $$t = M - \ln \left[ \frac{1}{1 - p(1 - e^{-a})} \right]$$

这些表达式的详细推导见附录 A.2。这些解析表达式提供了关于 logit 分布与采样阈值之间关系的宝贵见解。

虽然语言模型中的实际 logit 分布可能更复杂，但这些基本案例作为重要的理论基准，并且可以指导设计更高效的采样算法。

理解 Logits 的真实分布

噪声区域

如图 1 所示，绝大多数的 token 表现出遵循高斯分布的 logits。由于这些 logits 对应的概率通常可以忽略不计，并且在之前的工作中被认为是噪声，我们将这个区域指定为噪声区域。

这种特征与统计直觉很好地吻合，在统计直觉中，高斯分布通常表明系统中存在随机噪声。

当噪声区域和信息区域之间的边界变窄时，经过 softmax 变换后的噪声衍生概率倾向于干扰模型的生成过程，潜在地降低输出质量。

这种现象在高温采样场景中尤为明显，所有当前的非确定性采样算法在这种情况下表现不佳。

我们将这种退化归因于在高温下两个区域之间减小的间隙，使得噪声分布主导了概率景观。

我们识别出几个导致这种噪声的关键因素：

训练数据噪声 训练数据中的固有噪声和变化自然会通过模型的学习过程传播，最终表现为 logit 空间中高斯分布的一部分。

正则化效应 在训练过程中采用的各种正则化技术，虽然对于防止过拟合至关重要，但也有副作用，即给词汇表中语义无关的 token 分配较小但非零的概率。

沉默噪声 模型架构约束将有限值分配给无关的 token（理想情况下应该是“沉默”的，具有 $-\infty$ logits），这导致了一种独特的噪声模式。这是 softmax 函数的一个内在缺陷。

尽管从根本上解决这些噪声源超出了本文的范围，但我们相信这些见解可能有助于未来工作中改进训练程序。

在本文中，我们将专注于如何在推理过程中消除这些噪声 token。

信息区域

如图 1 所示，一小部分 token 占据了概率质量的大部分。这种集中表明大型语言模型对这些 token 具有特定的知识，因此我们将此指定为信息区域。

由于该区域中的 token 数量有限，很难对该区域的基本分布做出明确的声明。然而，关于采样方法的近期经验观察提供了对该区域特征的有趣见解。

值得注意的是，我们观察到 min- $p$ 采样方法在生成质量方面取得了显著改进。这种方法建立了一个基准概率阈值 $p$ ，并消除了所有低于 $p_{max} \cdot p$ 的概率值，其中 $p_{max}$ 是最大概率。虽然 Nguyen 等人通过经验推导出这种方法，但我们做出了一个令人惊讶的理论发现：

定理3 对于遵循均匀分布的 logits，min- $p$ 采样等价于top- $(1 - p)$ 采样。

证明：我们首先将 min- $p$ 放入logits空间。对于 $l_i \sim U(-\infty, M)$ ，阈值为

$t = \ln(s \cdot p_{max} \cdot p) = \ln(e^M \cdot p) = M + \ln p$

因为 $l_i \sim U(-\infty, M)$ ， $a = -\infty$ ，所以 top- $(1 - p)$ 的阈值完全相同：

$t = M - \ln \left[ \frac{1}{1 - (1 - p)} \right] = M + \ln p$

这个分析揭示了尽管 min- $p$ 声称具有适应性，但它本质上在 logits 空间中执行静态截断。此外，min- $p$ 采样的有效性表明信息区域大致遵循均匀分布。

确定边界

区分信息区域和噪声区域的一种自然方法是将信息区域视为噪声分布的异常值。根据这一假设，可以应用传统的 μ + 3σ 规则来确定边界。然而，我们的经验观察表明，这种方法可能不是当前任务的最佳选择。

我们定义 $σ\text{-distance}$ 为最大概率与分布均值之间的标准差数量，即 $σ\text{-distance} = (M - μ) / σ$ ，其中 $M$ 是最大 logit，遵循均匀分布情况下的相同符号。

如图 2a 所示，最大概率与均值之间的距离始终超过 10σ，在生成过程中表现出显著波动。结合图 2b，出现了一个有趣的模式：较大的核大小与较低的 $σ\text{-distance}$ 相关。

这一观察结果与传统的异常值检测方法（如 μ + 3σ 准则）相矛盾，后者理论上认为较高的 $σ\text{-distance}$ 应包含更多的 token。这一发现强烈表明信息 token 不应被视为噪声 token 的异常值。事实上，较高的 $σ\text{-distance}$ 表明模型对其输出具有强烈的信心，从而导致较小的核大小。

这使我们从根本上改变了视角：与其将少数从高斯分布中视为异常值，我们应该认识到大多数是从均匀分布中的异常值。这种反直觉的转变挑战了传统异常值检测的常规假设，其中异常通常是罕见事件。在我们的案例中，噪声区域中的 token 基本上是无关的候选者，应该被排除在外，尽管它们的数量占多数。

为了实现这一视角，我们引入了 top- $nσ$ 算法。我们的算法从最大值开始向下延伸，使用分布的标准差动态调整边界。具体来说，我们捕捉一个向下延伸 $nσ$ 的最大值区域，其中 $n$ 是通过经验确定的，以平衡保留信息 token 和排除均匀噪声。在实践中，我们发现 $n = 1.0$ 实现了令人满意的效果。

算法描述

我们的方法引入了一个统计阈值来在采样前过滤候选 token。算法 1 概述了我们方法的主要步骤。

该算法计算效率高，因为它直接在 logits 上操作，而不需要额外的 softmax 变换或排序。核心运算（最大值和标准差）针对现代 GPU 实现进行了高度优化。执行时间通常保持在几十微秒内，其中大部分计算成本由内核启动开销占据。

核质量

尽管在 logits 空间中分布混合的确切机制尚不清楚，我们专注于分析两种边界情况：

当分布近似高斯分布时，这通常发生在高温场景中；
当分布近似均匀分布时，如在低温设置中观察到的情况，或者当模型表现出高置信度（通常伴随着较大的标准差）时。

这两种情况为我们分析方法在实际中的行为提供了理论界限。

在高斯情况下， $l \sim N(\mu, \sigma^2)$ ，我们得到以下方程：

$M - n\sigma = \mu + \sqrt{2}\sigma (\text{erf}^{-1}(1-2p)) + \sigma^2$

因此，

$p = \frac{1}{2} \left[ 1 - \text{erf} \left( \frac{M - \mu - n\sigma - \sigma^2}{\sqrt{2}\sigma} \right) \right]$

该公式的一个关键属性是，当 $\sigma$ 接近零时， $\text{erf}$ 内的项接近正无穷大，导致 $p$ 收敛到零。这有效地防止了来自噪声区域的token被包含进来，为采样过程中的潜在不稳定性提供了一种自然的保护。

在均匀分布的情况下， $l \sim U(M-a, M)$ ，我们有：

$M - n\sigma = M - \ln \left[ \frac{1}{1-p(1-e^{-a})} \right]$

因此，

$p = \frac{1-e^{-n\sigma}}{1-e^{-a}}$

为了确定 $a$ 的适当值，我们利用 logit 的总体标准差为 $\sigma$ 这一事实。考虑到均匀分布的方差为 $a^2/12$ ，我们有 $a^2/12 \leq \sigma^2$ 以保持与观察到的标准差的一致性。这为我们提供了上限 $a \leq 2\sqrt{3}\sigma$ 。

因此，

$p \geq \frac{1-e^{-n\sigma}}{1-e^{-2\sqrt{3}\sigma}}$

对于典型的参数值，如 $n=1$ 和 $\sigma=1.9$ ，我们得到下限约为 0.85，表明我们的 top- $n\sigma$ 算法有效地保留了信息丰富的token。

这种分析也为选择超参数 $n$ 提供了有价值的指导：它不仅应该是正数，而且应保持在 $2\sqrt{3} \approx 3.46$ 以下，因为 $p \leq 1$ 。超过这个上限可能会导致包括噪声 token，从而可能降低算法的性能。

温度不变性

我们采样方法的一个关键特性是其温度不变性。具体来说，候选 token 的集合在使用任何温度值时都保持不变。

定理4 对于任何温度 $T > 0$ ，top- $nσ$ 的核心保持不变。

证明：考虑任意 token $i$ 和温度 $T > 0$ 。设 $l_i$ 为原始 logit。经过温度缩放后，对于所有 token，我们有 $l'_i = l_i/T$ 。对于给定的 token $i$ ：

$M' = \max(l'_j) = \frac{M}{T}$

$\sigma' = \sqrt{\frac{1}{N}\sum_j (l'_j - \mu')^2} = \sqrt{\frac{1}{N}\sum_j (\frac{l_j}{T} - \frac{\mu_j}{T})^2} = \frac{\sigma}{T}$

token $i$ 被选中当且仅当 $l'_i ≥ M' − nσ'$ 。代入：

$l'_i \geq M' - n\sigma' \iff \frac{l_i}{T} \geq \frac{M}{T} - \frac{n\sigma}{T} \iff l_i \geq M - n\sigma$

这个最终条件与 $T$ 无关。因此，对于任何 token $i$ ，它是否被选中取决于相同的条件，无论温度如何。

这种温度不变性特性将我们的方法与其他常见的采样方法区分开来。对于 top- $p$ 和 min- $p$ 采样，选择的 token 集随温度变化。随着温度升高，logit 分布趋向于高斯噪声部分，导致 top- $p$ 和 min- $p$ 在其采样池中包含更多噪声 token。虽然 top- $k$ 采样确实保持了温度不变性，但它使用固定的 $k$ 值，这仅仅转移了问题：固定的 $k$ 无法在不同上下文中动态地区分有效 token 和噪声 token。

细微的探索控制 使用 top- $nσ$ ，探索被分解为两个不同的参数。参数 $n$ 控制核心的大小，而温度仅调节如何在核心内进行探索。这种机制作为一种软屏蔽方法，有效地消除了潜在的不相关 token，同时仍允许一定程度的受控探索。这种解耦提供了对采样过程更细微的控制：参数 $n$ 确定有效 token 和噪声 token 之间的边界，而温度参数则在验证的 token 空间内微调探索策略。