(Transformer) Decoder-Only 架构作为现今的主流因果语言模型架构, 其通过自回归的方式生成文本, 其中的每一步生成都是在从分词表大小的概率分布中采样下一 Token, 而如何采样下一 Token 则依赖 “解码策略” (Decoding Strategy), 不同的解码策略会显著影响生成序列的质量 多样性 以及连贯性, 因此选择合适而正确的解码策略至关重要. 本文深入探讨了 4 种常见的解码策略 (Top-K, Top-P, Temperature, Beam Search) 的算法原理.

Top-K 解码: 硬范围内的采样

Top-K 是一种简单而有效的采样策略, 最早由 Facebook AI Research 团队在故事生成研究中提出$^{[4]}$, 其核心思想是在每个时间步仅从概率最高的 K 个候选 Token 中进行随机采样, 而其他概率较低的 Token 则被抛弃, 这一策略通过减少候选 Token 的数量, 降低了生成文本的随机性, 提高了生成文本的连贯性和合理性.

抛弃低概率 Token 后, 剩余的 K 个 Token 的概率会被重新归一化以确保总和为 1 (概率域), 并按重新归一化后的概率进行采样.

  • 算法简述

    1. 对当前 Step 的模型输出向量 (logits 向量) 进行降序排序 (依据维度排序), 找出 logit 最大的 Top-K 个维度索引.

      模型的原始输出通常是一个与”分词表大小”相同维度的 logits 向量, 关于什么是 logits, 请移步此处.

      词汇表是一个线性表, 所以可以被索引.

      该步骤所找的维度索引, 就是分词表中对应 token id 的位置, 找到的 K 个 token id 就是我们需要的候选 token ids.

    2. 取到 Top-K 个 token 的 logit 和其对应的 token id 后, 我们需要对其进行 Softmax 归一化, 以确保 K 个 token 均被映射到概率域, 且和为 1.

    3. 现在我们得到了 Top-K 个 token 的概率分布, 在该分布上进行采样, 即可确定下一个 Token 的 token id.

    4. 使用分词器对 token id 进行解码 (tokenizer.decode()), 得到 token 字符串, 也就是我们真正需要的下一 Token.

  • 性质

    • 随机性可控: 通过调整 K 值, 可以灵活控制文本生成的多样性, 当 K = 1 时, Top-K 解码会退化为贪心搜索解码 (Greedy Search), 生成结果完全确定, 当 K 等于分词表大小时, Top-K 解码则退化为随机采样.

    • 计算高效: 由于只需处理 K 个 Token, Top-K 采样的复杂性远低于全词表随机采样.

    • 平衡多样性与连贯性: Top-K 通过排除低概率 Token, 减少了生成无意义或语法错误内容的可能性, 同时也保留了一定的随机性.

Top-P (核采样) 解码: 自适应的概率空间优化

Top-P 采样，也称为核采样 (Nucleus Sampling)，是一种更高级的采样策略，由 Holtzman 等人在开放域文本生成研究中提出$^{[1]}$, 最初用于解决 “神经文本退化问题”. 与 Top-K 采样不同，Top-P 采样不是固定选择前 K 个 token，而是动态选择累积概率达到阈值 P 的最小 token 集合进行采样​.

神经文本退化, 即重复/通用化/不连贯等文本现象.

  • 算法简述

    1. 对当前 Step 的模型输出向量 (logits 向量) 进行降序排序 (依据维度排序).

    2. 对 logits 向量应用 Softmax 归一化, 得到降序排序的概率分布.

    3. 从高到低对概率进行累加(同时将遍历到的 token id 及其概率添加到候选集中), 直到累积概率达到阈值 P 停止, 抛弃之后的所有 Token 索引.

    4. 对候选集重新进行 Softmax 归一化, 确保概率总和依然为 1.

    5. 在归一化后的候选集上依据新的 token 概率进行随机采样, 确定写一个 token 的 token id.

    6. 使用分词器对 token id 进行解码 (tokenizer.decode()), 得到 token 字符串, 也就是我们真正需要的下一 Token.

  • 性质

    • 自适应采样空间: Top-P 采样能够根据概率分布的具体情况动态地调整采样空间的大小. 在概率分布集中的情况下, 采样空间小; 而在概率分布分散的情况下, 采样空间则被扩大, 从而实现更加灵活的控制.

    • 更好的多样性和连贯性平衡: Top-P 采样不仅考虑了概率的高低, 还考虑了概率分布的整体结构, 能够在保持连贯性的同时引入多样性输出.

    • 更好的可解释性: P 值直观地表示了采样空间占总概率质量的比例，相比 K 值更具解释性​.

    • 计算复杂度：Top-P 采样的计算复杂度略高于 Top-K 采样，因为需要计算累积概率并动态确定采样空间.

Temperature (温度): 概率分布的”放大器”

Temperature 并不是一种独立的解码策略, 但也值得拿出来讨论. 温度采样最初由深度学习之父 Geoffrey Hinton 在模型蒸馏研究中提出$^{[2]}$, 不同于直接在词汇集中选取子集的解码策略, 温度并没有直接取子集, 而是通过调整概率分布 (可以理解为控制方差 $\sigma$) 实现采样概率的集中化或平滑化, 其计算式如下:

\[P(w) = \frac {e^{\frac {s_w} {T}}} {\sum_{w' \in V}e^{\frac {s'_w} {T}}}\]

其中, $s_w$ 是原始 logits, $T$ 是温度参数, $V$ 是分词表. 温度的工作原理是通过缩放 logits 来改变概率分布的形状:

  • 当 T = 1, 概率分布保持不变.

  • 当 T < 1, 概率分布的差异被放大, 分布变得更加陡峭, 高概率 token 的概率进一步提升, 而低概率 token 的概率进一步降低, 解码输出则更具确定性.

  • 当 T > 1, 概率分布的差异被缩小, 分布变得更加平坦, 高概率 token 的概率可能被拉到和低概率 token 相近的值, 各 token 之间的区分度更低, 解码输出则倾倾向于多样性, 甚至不稳定.

Temperature 通常和其他解码策略进行集成, 旨在实现更加精细化的输出控制.

在 Hinton 最初的研究中, temperature（温度）主要应用于知识蒸馏过程，是一种调整 Softmax 归一化分布的关键手段，目的是更有效地将复杂模型（如集成模型或大型正则化模型）的知识迁移到小型模型中.

Beam Search (束搜索 / 最大化解码): 动态规划的贪心近似, 同时探索多条路径的确定性方法

束搜索 (Beam Search) 是一种确定性解码策略, 最初由 Alex Graves 在他的音素识别研究中提出$^{[3]}$，该策略通过维护多个候选序列来平衡探索和利用​. 束搜索的核心思想是在每个 Step 保留概率最高的 $K$ 个候选序列 (称为束宽，Beam Width)，而不是仅保留一个最优序列, 其目标是找到一个接近全局最优的序列 ($K$ 越大, 越接近全局最优).

  • 算法简述

    1. 初始化, 从初始状态开始, 生成概率最高的 $K$ 个 token 作为初始候选序列.

    2. 迭代扩展, 对于每个序列, 生成概率最高的 $m$ 个 next token, 得到 $K \times m$ 个新序列.

    3. 计算每个新序列的累积概率, 保留 Top-K 个候选序列, 抛弃其他序列.

    4. 当序列达到最大长度或生成结束符后, 输出结束, 返回 $K$ 个最终候选序列.

  • 性质

    • 动态规划的贪心近似: 束搜索通过维护多个候选序列，能够在探索不同可能性的同时利用已知的高概率路径，相比贪心搜索能找到更优解.

    • 生成质量高: 束搜索通常能生成比贪心搜索和采样方法更连贯、质量更高的文本，特别是在需要长距离依赖的任务中.

    • 时间复杂度高: 束搜索的计算复杂度随着束宽 $K$ 的增加呈指数级增长，因为每个步骤都需要处理 $K \times m$ 个序列.

    • 空间复杂度高: 束搜索需要存储多个候选序列及其状态，因此内存需求显著高于贪心搜索或采样方法.

  • 典型应用

    由于束搜索的生成质量最高, 其已经成为机器翻译任务的标准解码策略, 通常采用中等束宽 (3 到 10) 可以显著提升翻译生成质量. 而在其他 NLP 任务如文本摘要, 问答系统, 代码生成等, 其也能够保证输出的连贯性和稳定性.

参考文献

[1] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, Yejin Choi. The Curious Case of Neural Text Degeneration. arXiv preprint, 2019.

[2] Geoffrey Hinton, Oriol Vinyals, Jeff Dean. Distilling the Knowledge in a Neural Network. arXiv preprint, 2015.

[3] Alex Graves. Sequence Transduction with Recurrent Neural Networks. arXiv preprint, 2012.

[4] Angela Fan, Mike Lewis, Yann Dauphin. Hierarchical Neural Story Generation. arXiv preprint, 2018.