Jiahong 的个人博客

凡事预则立,不预则废


  • Home

  • Tags

  • Archives

  • Navigation

  • Search

NLP——LLM-API调用示例


Qwen

  • Qwen API 申请:获取API Key
  • Qwen API 调用文档:Qwen-API Doc
  • 吐槽:Qwen 的文档和申请链接写的很差,阿里云东西太多,需要翻来翻去找
  • Qwen API 调用示例:
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    def qwen_api():
    import requests

    url = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
    headers = {
    "Authorization": "Bearer $API_KEY",
    "Content-Type": "application/json"
    }

    data = {
    "model": "qwen-plus",
    "messages": [
    {"role": "user", "content": "你好,请介绍一下你自己"}
    ],
    "max_tokens": 50,
    "temperature": 0.0, # 贪心采样示例
    "top": 0.2, # 贪心采样示例
    "logprobs": True, # 可以打开 logprobs 看每个 token 的 logprobs,使用 e^logprob 即可得到最终概率
    "top_logprobs": 2,
    }

    response = requests.post(url, headers=headers, json=data)
    print(response.json())

    if __name__ == "__main__":
    qwen_api()

LongCat

  • LongCat 文档:LongCat API开放平台快速开始

  • 文档写的清晰明了,Qwen 应该学习一下

  • LongCat API 调用示例:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    def longcat_api():
    import requests

    url = "https://api.longcat.chat/openai/v1/chat/completions"
    headers = {
    "Authorization": "Bearer $API_KEY",
    "Content-Type": "application/json"
    }

    data = {
    "model": "LongCat-Flash-Chat",
    "messages": [
    {"role": "user", "content": "你好,请介绍一下自己"}
    ],
    "max_tokens": 1000,
    "temperature": 0.7,
    # "logprobs": True, # 打开这个参数会报错
    }

    response = requests.post(url, headers=headers, json=data)
    print(response.json())

    if __name__ == "__main__":
    longcat_api()
    • 特别强调:目前 LongCat 不支持返回 logprobs 信息

NLP——LoRA-Without-Regret

注:本文包含 AI 辅助创作

  • 参考链接:
    • 原始博客:LoRA Without Regret, 20250929, Thinking Machines

Blog Summary

  • 本文是 Thinking Machines 发布的一篇博客,核心是研究 LoRA 在 LLM 中的作用以及的各超参的详细影响
  • 本文包含非常详细的实验细节,讨论了各种 LoRA 相关超参数(学习率和 Rank 等)设置技巧,是一篇非常实在的技术博客,值得深读

Introduction

  • 如今主流的语言模型(Language Model)包含超过一万亿(1T)个参数,这些模型在数十万亿个 Token 上进行预训练
  • Base Model 的性能随着规模的扩大而不断提升,因为这万亿级别的参数对于学习和表征人类书面知识中的所有模式至关重要
  • In contrast, Post-Traning 使用的数据集规模更小,且通常聚焦于更狭窄的知识领域和行为范围
    • 用万亿比特(Terabit)级别的权重来表征来自十亿比特(Gigabit)或百万比特(Megabit)级训练数据的更新,似乎存在资源浪费
    • 这种直觉推动了参数高效微调(Parameter-Efficient Fine-Tuning, PEFT)技术的发展,该技术通过更新一组规模小得多的参数来调整大型网络
  • 目前主流的参数高效微调方法是 LoRA(Low-Rank Adaptation)
    • LoRA 将原始模型中的每个权重矩阵 \( W \) 替换为经过修改的版本 \( W’ = W + \gamma BA \),其中矩阵 \( B \) 和 \( A \) 加起来的参数数量远少于 \( W \),而 \( \gamma \) 是一个常数缩放因子
    • 实际上,LoRA 为微调所带来的更新创建了一个低维表征
  • LoRA 在 Post-Traning 的成本和速度方面可能具有优势,此外,从操作角度来看,相比全参数微调(Full Fine-Tuning,以下简称 FullFT),选择 LoRA 还有几个理由:
    • 多租户服务(Multi-Tenant Serving) :
      • 由于 LoRA 仅训练一个 Adapter (即矩阵 \( A \) 和 \( B \)),同时保持原始权重不变,因此单个推理服务器可以在内存中存储多个 Adapter (不同的模型版本),并以批量方式同时从中采样
      • 现代推理引擎(如 vLLM 和 SGLang)已实现此功能(2023)
    • 训练 Layout 大小(Layout Size for Training) :
      • 对整个模型进行微调时,优化器状态需要与原始权重一起存储,且通常需要更高的精度
      • FullFT 通常比对同一模型进行采样所需的加速器数量多一个数量级,进而需要不同的 Layout (理解:这里的 Layout 应该是指显存?)
        • 注:在训练过程中,除了存储权重外,通常还需要存储所有权重的梯度和优化器动量;而且,这些变量的存储精度(float32)通常高于推理时权重的存储精度(bfloat16 或更低)
      • 由于 LoRA 训练的权重数量少得多,占用的内存也少得多,因此其训练所需的 Layout 仅比采样所需的 Layout 略大
      • 这使得训练更容易实现,且通常效率更高
    • 加载与迁移便捷性(Ease of Loading and Transfer) :
      • 由于需要存储的权重更少,LoRA Adapter 的设置速度更快、更便捷,也更容易在不同机器之间迁移
  • 上述原因足以解释自 2021 年 LoRA 原始论文发表以来,该技术日益普及的现象(2021)
    • 但现有文献中关于 LoRA 与 FullFT 性能对比的结论尚不明确
  • 研究界普遍认为,在类似预训练的场景下,LoRA 性能表现较差(2024)
    • 即当数据集规模极大,超出 LoRA 参数的存储限制时,LoRA 会处于劣势
  • 但对于 Post-Traning 中常见的数据集规模,LoRA 具备足够的容量来存储关键信息
    • 不过,这一事实并不能保证 LoRA 在样本效率和计算效率上的表现
  • 核心问题在于:LoRA 是否能达到全参数微调的性能水平?如果可以,需要满足哪些条件?
  • 论文实验发现,当正确把握几个关键细节时,LoRA 的样本效率确实与 FullFT 相同,并且能达到相同的最终性能

What Matters for LoRA(什么最关键?)

  • 论文介绍了论文开展的一系列 SFT 和 RL 实验,旨在确定 LoRA 与 FullFT 效率匹配的条件
  • 为此,论文的实验方法与以往关于 LoRA 的实验存在以下几点不同:
    • 1)论文研究了训练集大小与 LoRA 参数数量之间的一般关系 ,而非聚焦于特定的数据集和任务
    • 2)在监督学习中,论文测量的是对数损失(Log Loss) ,而非采用基于采样的评估方法(Sampling-Based Evals),这样做同样是为了保证结论的通用性
      • 对数损失测量能在不同训练步骤和训练参数范围内,提供清晰的结果和缩放定律(Scaling Laws)
  • 论文的研究发现如下:
    • 在中小型指令微调(Instruction-Tuning)和推理数据集(Reasoning Dataset)上进行 SFT 时,LoRA 的性能与全参数微调相同
    • 当数据集规模超出 LoRA 的容量时,LoRA 性能会低于 FullFT
      • 此时,LoRA 并非会达到一个无法突破的损失下限(Distinct Floor)
      • 而是其训练效率会下降 ,且下降程度取决于模型容量与数据集大小之间的关系
      • 理解:这里的意思是训练慢,但是多训练一段时间,也能达到不错的效果
    • 在某些场景下,LoRA 对大批量大小(Large Batch Size)的容忍度低于 FullFT(当批量大小超过某个阈值后,LoRA 的损失 penalty 会更大)
      * 增加 LoRA 的 Rank 无法缓解这种 penalty ,这是矩阵乘积参数化(Product-of-Matrices Parametrization)的固有属性,该参数化方式的 Training Dynamics 与优化原始权重矩阵的 Training Dynamics 不同
      • 注:Training Dynamics 指的是训练过程中模型的各种属性随时间或迭代轮数的变化情况,包括但不限于损失函数值、准确率、模型参数的更新、梯度的分布等属性
      • 问题:这里的意思是 Batch Size 太大,不适合 LoRA?
    • 即使在小数据场景下,将 LoRA 应用于所有权重矩阵(尤其是 MLP 和 MoE 层)时,其性能也会更优
      • 仅将 LoRA 应用于注意力层(Attention-Only LoRA)的性能较差,即便通过提高 Rank 来匹配可训练参数的数量,结果依然如此
    • 在强化学习中 ,即使使用较小的 Rank ,LoRA 的性能也能与 FullFT 相当
      • 论文发现强化学习所需的容量非常低 ,这一结果与论文基于信息论的推断一致
  • 论文还研究了 LoRA 超参数对其学习率(相对于 FullFT 学习率)的影响
    • 论文分析了初始缩放(Init Scales)、乘数(Multipliers)等超参数的一些不变性,并解释了为何 \( \frac{1}{r} \) 预因子能使最优学习率(LR)大致独立于 Rank
    • 论文还通过实验展示了 LoRA 的最优学习率与 FullFT 最优学习率之间的关系
  • 论文的实验结果明确了“低遗憾区间(Low-Regret Regime)”的特征
    • 在该区间内,LoRA 在数据集大小和参数数量方面的性能与 FullFT 相近
    • 论文发现该区间涵盖了大多数 Post-Traning 场景,这为在众多应用中使用高效微调技术开辟了道路

Methods and Results

  • 论文设计实验的目的是,在不同条件下详细测量 LoRA 相对于 FullFT 的性能。以下是实验设置的一些细节:
    • 论文将 LoRA 的 Rank 在三个数量级范围内变化( Rank 的取值为 1 到 512),并将这些不同 Rank 的 LoRA 与全参数微调进行对比
    • 为了排除因使用非最优学习率可能带来的干扰,论文针对每个实验条件都进行了学习率扫描(LR Sweep)
      • 论文采用恒定学习率调度(Constant Learning Rate Schedule),不进行预热(Warmup)或冷却(Cooldown)
    • 实验使用的模型为 Llama 3 系列模型(2024)和 Qwen3 模型(2025),其中包括一个 MoE 模型
    • 主要的监督学习实验使用了 Tulu3(2024)和 OpenThoughts3(2025)数据集,分别聚焦于Instruction Following和推理任务
      • 这两个数据集在范围、结构和应用场景上存在显著差异,这有助于验证论文结果的通用性
    • 强化学习实验使用数学推理任务,以答案的正确性作为奖励(Reward)

LoRA Rank

  • 论文在 Tulu3 数据集和 OpenThoughts3 数据集的一个子集上进行了单轮训练(Single Epoch)
  • 针对每个数据集和模型大小,论文都对 LoRA 的 Rank 和学习率进行了扫描
  • 在下图中,每条彩色线条代表一个特定的 Rank ,线条是通过在每个训练步骤中取所有学习率对应的最小点得到的:
  • 论文观察到,FullFT 和高 Rank LoRA 具有相似的训练曲线,损失随训练步骤对数的增加而线性下降
    • 当训练步骤达到某个与 Rank 相关的阈值时,中低 Rank LoRA 会偏离最小损失训练曲线
    • 从直觉上看,当 Adapter 达到容量上限时,学习速度会减慢,而容量上限由 Rank 决定
  • 接下来,论文绘制了损失随学习率变化的曲线,以验证论文的扫描是否覆盖了每个 Rank 对应的最优学习率
    • 理解:LoRA 比 FullFT 需要更大的最优学习率
  • 论文发现,FullFT 的最优学习率比高 Rank LoRA 的最优学习率低一个数量级(即 1/10)(2024)
    • 论文将在后面讨论 LoRA 超参数时再次提及这一点
  • 所有不同 Rank 的 LoRA 运行对应的最优学习率似乎相近(论文将在下文从理论角度解释这一发现)
    • 不过,最优学习率确实存在一定的 Rank 依赖性(Rank 为 1 时的最优学习率低于更高 Rank LoRA 的最优学习率)
    • 在 Rank 为 4 到 512 的范围内,最优学习率的变化因子小于 2

Batch Size Effects

  • 论文发现,在某些设置下,LoRA 对大批量大小的容忍度低于 FullFT
  • 随着批量大小的增加,两者的性能差距会扩大,且这种差距与 Rank 无关
  • 在接下来的实验中,论文使用了 OpenThoughts3 数据集中一个包含 10,000 个样本的小子集
  • 图 3 中的左图显示
    • 在大批量大小时,LoRA(虚线)与 FullFT(实线)的训练曲线之间存在持续的差距
    • 而在批量大小为 32(较小批量)时,这种差距更小,且会随着训练的进行而缩小
    • 问题:左图是不是少了些其他的曲线?
  • 右图展示了最终损失随批量大小变化的情况
    • 随着批量大小的增加,LoRA 与 FullFT 之间的损失差距越来越大
  • 大批量下的性能差距似乎与 Rank 无关,而是 LoRA 的固有属性
    • 其可能原因是,在该数据集上,矩阵乘积参数化(\( BA \))的优化 Dynamics 不如全矩阵(\( W \))的优化 Dynamics 有利
    • 不过,LoRA 和 FullFT 在较小批量大小时都能达到最佳损失,因此在实际应用中,这种差距可能不会产生太大影响
  • 问题:为什么 Batch Size 对最终结果的影响这么大,是因为 Batch Size 和 学习率没有对应调优吧!
    • 理解:理论上,Batch Size 对于训练不应该有那么大的差异才是

Layers Where LoRA Is Applied(LoRA 应用的层)

  • 论文研究了将 LoRA 应用于网络中不同层所产生的影响
  • Hu 等人的原始论文建议仅将 LoRA 应用于注意力矩阵(Attention Matrix),此后许多论文也遵循了这一做法
    • 但最近的趋势是将 LoRA 应用于所有层
  • 与论文的结果类似,QLoRA 论文也发现,仅将 LoRA 应用于注意力层的性能低于应用于 MLP 层或同时应用于 MLP 层与注意力层的性能
    • 不过该论文发现“MLP 层+注意力层 > MLP 层 > 注意力层”,而论文发现 MLP 层+注意力层 和 MLP 层 的性能大致相当
    • 事实上,当论文将 LoRA 应用于所有层(尤其是 MLP 层,包括 MoE 层)时,取得了好得多的结果
    • 实际上,将 LoRA 应用于注意力矩阵,相比仅将其应用于 MLP 层,并没有带来额外的收益(2024)
      • 理解:虽然加入注意力层没有带来额外收益,博客最终还是使用了所有层都微调,其实有些矛盾
  • 仅应用于注意力层的 LoRA 性能较差,并非因为参数数量少
    • 在特定案例中, Rank 为 256 的仅注意力层 LoRA,其参数数量与 Rank 为 128 的仅 MLP 层 LoRA 大致相同,但前者的性能却低于后者(可对比下表中的粗体数字)
      LoRA Configuration Params
      mlp, rank=256 0.49B
      attn, rank=256 0.25B
      all, rank=256 0.70B
      mlp, rank=128 0.24B
  • 在 MoE 实验中,论文为每个专家(Expert)训练了一个单独的 LoRA,每个 LoRA 的 Rank 等于总 Rank 除以活跃专家(Active Expert)的数量(Qwen3 MoE 模型的活跃专家数量为 8)
    • 这种缩放方式使得 MoE 层中 LoRA 参数与 FullFT 参数的比例,与其他层中两者的比例保持一致
    • 理解:这可以保证最终激活的 Rank 数是一致的
  • 论文还在另外两个场景中进行了类似实验,比较不同 LoRA 应用层的效果:
    • (1)在 OpenThoughts3 数据集的一个小子集上进行 Rank 为 256 的监督学习;
    • (2)在 MATH 数据集上进行强化学习
    • 论文将在下一节描述实验设置
  • 在这些场景中,仅注意力层的 LoRA 性能依然低于仅 MLP 层的 LoRA(后者与“MLP 层+注意力层”的 LoRA 性能相近)

RL

  • 论文的实验得出一个关键发现:在使用策略梯度算法(Policy Gradient Algorithm)进行强化学习时,即使 LoRA 的 Rank 低至 1,其学习性能也能完全匹配 FullFT
  • 在这些实验中,论文使用了一种带有重要性采样校正(Importance Sampling Correction)的基础策略梯度算法,其目标函数为(2024):
    $$ \text{objective} = \sum_t \frac{p_{\text{learner} } }{p_{\text{sampler} } } Adv_t $$
  • 论文采用了类似 GRPO 的中心化方案(2024),即针对每个问题采样多个补全结果(Completion),并减去每组的平均奖励
  • 下图(图 6)展示了在 MATH(2021)和 GSM(2021)数据集上的学习率扫描结果,每个数据集都使用了典型的超参数
    • 论文使用 Llama-3.1-8B 基础模型进行实验,因为正如 Qwen 技术报告(2024)所描述的,Qwen2.5 和 Qwen3 模型的预训练数据已提升了其数学性能,这会使得论文难以衡量仅在强化学习阶段学到的内容
    • LoRA 具有更广泛的有效学习率范围,并且能达到与 FullFT(黑线)相同的峰值性能,至少在强化学习的噪声所允许的精度范围内是如此
  • 这一结果与信息论的推断一致(注意:下面是重点,需要重点理解)
    • 监督学习每个样本(Episode)可提供约 \( O(\text{ Token 数量}) \) 比特的信息
    • 在策略梯度方法中,学习由优势函数(Advantage Function)驱动 ,每个样本仅能提供 \( O(1) \) 比特的信息
    • 当每个样本包含数千个 Token 时,强化学习在训练过程中每 Token 吸收的信息约为监督学习的 1/1000
  • 我们可以根据实验数据给出更精确的数值
    • 在 MATH 数据集的实验中,论文在约 10,000 个问题上进行训练,每个问题采样 32 次
    • 假设每个补全结果提供 1 比特信息,那么整个训练过程仅需吸收 320,000 比特信息
    • Llama-3.1-8B 模型的 Rank 为 1 的 LoRA 已有 300 万个参数(论文通过对模型中所有权重矩阵求和 \( \text{rank}\cdot d_{\text{in} } \)(矩阵 \( A \) 的参数数量)和 \( \text{rank} \cdot d_{\text{out} } \)(矩阵 \( B \) 的参数数量)计算得出),几乎是所需信息比特数的 10 倍
      • 问题:信息和信息比特数可以直接比较吗?
    • 即便 Rank 为 1,LoRA 也拥有足够的容量来吸收训练过程中提供的所有信息
  • 作为另一个对比案例,DeepSeek-R1-Zero 模型在 5.3M 个样本上进行了训练(训练共进行 10,400 步,每步包含 32 个独特问题,每个问题采样 16 次),这对应着 5.3M 比特的信息
    • 这一数量少于低 Rank LoRA 的参数数量,因此论文预测,使用 LoRA 也能复现该模型的训练结果
  • 为了进一步验证 LoRA 在推理强化学习中有效性的发现,论文使用 Qwen3-8b-base 模型在 DeepMath 数据集(2025)上进行了更大规模的实验
    • DeepMath 数据集比 MATH 数据集大得多,且通常包含更难的问题
    • 为了加快实验速度,论文将训练和评估的样本长度限制为 8192 个 Token
    • 这个长度足以支持回溯(Backtracking)和推理过程,但与更长的思维链(Chain-of-Thought)相比,会限制模型性能
  • 论文观察到,当为每种设置选择最优学习率后,不同大小的 LoRA 和 FullFT 的训练过程几乎完全一致
    • 此外,在预留的 AIME 2024 和 AIME 2025 测试集上评估模型时,论文也得到了类似的结果
    • 而且,论文发现 LoRA 和 FullFT 的训练过程表现出相似的定性行为:
      • 两者都发展出了先进的推理能力,如回溯、自我验证(Self-Verification)和上下文内探索(In-Context Exploration),这一点可以从模型思维链长度的增加中看出

Setting LoRA Hyperparameters

  • LoRA 应用的一个障碍是需要选择最优超参数,而这些超参数与为 FullFT 优化的超参数不同
  • 在本节中,论文将说明这个问题并不像最初看起来那么棘手,并讨论论文在超参数选择方面的发现

Optimal Learning Rate and Rank

  • 参照 Hu 等人的研究,论文采用以下参数化方式表示 LoRA:
    $$ W’ = W + \frac{\alpha}{r} BA $$
    • 其中,\( r \) 是 LoRA 的 Rank
    • \( \alpha \) 是 LoRA 的缩放因子(论文遵循其他实现的标准做法,设置 \( \alpha = 32 \))
    • \( A \) 和 \( B \) 是 LoRA 的权重矩阵( Rank 为 \( r \))
  • \( \frac{1}{r} \) 缩放因子使得最优学习率大致独立于 Rank
    • 事实上,存在一个更强的条件:训练初期的学习曲线完全相同,与 Rank 无关
    • 这一现象非常显著,在实验中,不同 Rank 的学习曲线如此接近,以至于论文曾担心是程序漏洞导致 Rank 参数未被正确使用
    • 由此可见,在短期训练中,最优学习率也与 Rank 无关
    • 然而,正如论文在之前的学习率-损失图(图 2)中所展示的,在长期训练中,最优学习率会表现出一定的 Rank 依赖性
  • 我们可以通过分析训练首次更新后 LoRA 矩阵的预期更新来部分解释这一结果(注:即解释 \( \frac{1}{r} \) 缩放因子使得最优学习率大致独立于 Rank 这个事实)
    • 我们可以将 LoRA 矩阵乘积 \( BA \) 表示为 \( r \) 个 Rank-1 外积的和:
      $$ BA = \sum_{i=1}^r b_i a_i^T = \sum_{i=1}^r \Delta_i $$
      • 其中论文定义 \( \Delta_i = b_i a_i^T \)
      • 这里,所有 \( i \) 对应的 \( \frac{\partial \text{Loss}}{\partial \Delta_i} \) 都是相同的;
      • 但梯度 \( \frac{\partial \text{Loss}}{\partial b_i} \) 和 \( \frac{\partial \text{Loss}}{\partial a_i} \) 会依赖于初始化(例如,\( \frac{\partial \text{Loss}}{\partial b_i} \) 依赖于 \( a_i \))
      • 由于 \( a_i \) 和 \( b_i \) 的初始化与 Rank 无关,因此所有 \( i \) 对应的 \( \mathbb{E}[\Delta_i] \) 都是相同的,且与 Rank 无关
      • 在训练的第一步,每个外积项的预期更新是相等的,且与 Rank 无关
      • 由此可知,\( (\frac{1}{r}) \sum_{i=1}^r \Delta_i \) 本质上是 \( r \) 个具有相同期望的项的样本均值,因此该均值的期望(即 Adapter \( (\frac{1}{r})BA \) 的变化量)与 Rank 无关

Parametrization Invariances(参数化不变性)

  • LoRA 可能涉及四个超参数:
    • 1)缩放因子 \( \alpha \)(出现在 \( \alpha/r \) 中)
    • 2)下投影矩阵 \( A \) 的学习率 \( LR_A \)
    • 3)上投影矩阵 \( B \) 的学习率 \( LR_B \)
    • 4)矩阵 \( A \) 的初始化缩放 \( \text{init}_A \)
      • 对于随机初始化,这是 \( A \) 初始元素的标准差
      • 矩阵 \( B \) 初始化为零,因此无需定义 \( \text{init}_B \)
  • 调整这四个参数似乎难度较大,但 Training Dynamics 的不变性意味着其中两个参数是冗余的,学习行为仅由两个参数决定
    • 问题:如何理解这里的 Training Dynamics 的不变性
  • 论文通过以下分析说明这种不变性:
  • 在使用 Adam 优化器且 \( \varepsilon = 0 \) 时(我们可以将此结果扩展到 \( \varepsilon > 0 \) 的情况;此时需要将 \( \varepsilon \) 按 \( 1/q \) 缩放,因为梯度会按该因子缩放),优化过程在以下两参数变换下保持不变
    • 对于 \( p, q > 0 \),满足:
      • \( \alpha \rightarrow \frac{1}{pq} \cdot \alpha \)
      • \( \text{init}_A \rightarrow p \cdot \text{init}_A \)
      • \( LR_A \rightarrow p \cdot LR_A \)
      • \( LR_B \rightarrow q \cdot LR_B \)
  • 由于四个参数中的两个自由度不影响学习过程,因此实际需要调整的参数空间为二维
  • 我们可以为这个二维空间选择不同的基,例如以下具有直观解释的基:
    • 1)\( \alpha \cdot \text{init}_A \cdot LR_B \):
      • 该参数决定初始更新的规模,或者说学习曲线的初始斜率
      • 由于 \( B \) 初始化为零,\( LR_A \) 和 \( A \) 的初始更新在此阶段无关紧要
    • 2)\( \text{init}_A / LR_A \):
      • 由于 Adam 优化器每步对 \( A \) 元素的更新量约为 \( LR_A \),因此这个时间尺度参数决定了 \( A \) 从初始状态发生显著变化所需的步数
  • 我们可以用这个基来重新解释以往关于 LoRA 的一些研究提案:
    • LoRA+(2024)提出对 \( A \) 和 \( B \) 使用不同的学习率,且 \( B \) 的学习率更高。用上述基表示,提高 \( LR_B \) 等价于提高 \( \text{init}_A / LR_A \),从而使 \( A \) 的变化时间尺度更长
    • Unsloth 的 LoRA 超参数指南建议对高 Rank LoRA 使用更大的 \( \alpha \) 值(例如,避免 \( \frac{1}{r} \) 缩放)
      • 这也等价于提高 \( \text{init}_A / LR_A \)
      • 当论文提高 \( \alpha \) 时,需要相应降低 \( LR_A \) 和 \( LR_B \) 以保持更新规模不变,这反过来会使 \( LR_A \) 相对于 \( \text{init}_A \) 更小
  • 在论文的实验中,论文采用了 Hugging Face PEFT 库(2022)中 Hu 等人提出的标准参数化方式:
    • \( A \) 采用均匀分布初始化(缩放因子为 \( 1/\sqrt{d_{\text{in} } } \))
    • \( B \) 初始化为零,\( A \) 和 \( B \) 使用相同的学习率,且 \( \alpha = 32 \)
    • 在实验中,论文未能找到比这些超参数更优的设置

Optimal Learning Rates for LoRA vs. FullFT(最优学习率对比)

  • 论文的实验表明,无论是在监督学习还是强化学习中,LoRA 的最优学习率始终是同一应用场景下 FullFT 最优学习率的 10 倍
    • 这一规律在所有以性能(损失或奖励)为纵轴、学习率为横轴的 U 型图中都能体现
    • 这一发现应能简化从 FullFT 到 LoRA 的学习率超参数迁移过程
  • 目前,论文尚未对这一现象做出充分的理论解释
    • 论文尝试从“LoRA 最优学习率与 Rank 无关”和“满 Rank LoRA 可直接与 FullFT 对比”这两个事实出发推导该结果,但分析得出的学习率比例为“模型隐藏层大小除以 \( 2 \cdot \alpha \)”,这与“最优比例固定为 10(与基础模型无关)”的实验结果不符
  • 在实证分析中,论文在 Tulu3 数据集上对 14 个不同的 Llama 和 Qwen 模型进行了 LoRA 和 FullFT 的学习率扫描
    • 基于这些扫描结果,论文拟合了一个函数,该函数可根据模型的隐藏层大小和模型类型(Llama 或 Qwen)预测最优学习率。所用函数形式如下:
      $$ LR = M_{\text{LoRA} } \cdot \left( \frac{2000}{\text{hidden size} } \right)^{\text{model pow} + \text{LoRA pow} } $$
      • 其中:
        • \( M_{\text{LoRA} } \) 是使用 LoRA 时的乘数(使用 FullFT 时为 1)
        • \( \text{model pow} \) 是针对不同模型来源(Llama 和 Qwen)分别计算的指数调整项
        • \( \text{LoRA pow} \) 是针对 LoRA 的额外指数调整项
        • \( \text{hidden size} \) 是模型残差流(Residual Stream)的维度
  • 论文通过线性插值基于扫描数据预测损失,以此对预测的学习率进行评分,并通过对 14 个模型的预测损失求和来评估参数
  • 优化结果显示:
    • LoRA 相对于 FullFT 的乘数为 9.8;
    • Qwen3 和 Llama 模型对隐藏层大小的依赖不同,但 LoRA 学习率对隐藏层大小的依赖与 FullFT 相同(即优化结果显示 \( \text{LoRA pow} = 0 \))

Learning Rates in Short and Long Runs

  • LoRA 的典型初始化会导致有效学习率产生一个隐含的变化过程,这使得短期训练和长期训练存在差异,且学习曲线的形状与 FullFT 也有所不同
  • 在训练初期,\( B \) 初始化为零
    • 当 \( B \) 非常小时,\( A \) 的变化对 Adapter \( BA \)(将被添加到原始网络权重中)的影响可忽略不计
    • 随着 \( B \) 逐渐增大,\( A \) 的更新对网络输出的影响开始变大
      • 当 \( B \) 的规模接近 \( A \) 时,有效学习率会在训练过程中逐渐提高
    • 论文发现,在 Tulu3 和 OpenThoughts 数据集上完成完整训练后,\( B \) 矩阵的谱范数(Spectral Norm)最终会大于 \( A \) 矩阵的谱范数
  • 这意味着,对于短期训练,应设置更高的最优学习率
    • 初步证据表明,短期训练(根据经验,约 100 步以内)的最优乘数约为 FullFT 的 15 倍 ,而对于长期训练,该乘数会收敛到前文提到的 10 倍

Discussion

  • 除了实证结果外,论文还希望探讨与 LoRA 性能和适用性相关的更广泛问题,这些问题对研究人员和工程实践者都具有参考价值
  • 首先,论文深入分析核心结论,LoRA 与 FullFT 性能相近的两个条件:
    • 1)LoRA 应用于网络的所有层,尤其是参数占比最高的 MLP/MoE 层
    • 2)LoRA 未受容量限制,即可训练参数数量超过待学习信息的总量(可根据数据集大小估算)
  • 当条件(1)满足时,训练初期 LoRA 的学习 Dynamics 与 FullFT 相似;
  • 然后,在条件(2)的保障下,LoRA 会继续保持与 FullFT 相近的性能,直到达到容量限制

Why LoRA Might Be Needed on All Layers(Why LoRA 需要应用在所有层上)

  • 正如论文之前所展示的,若仅将 LoRA 应用于注意力层,即使在小数据场景下,学习速度也会变慢
  • 一种可能的解释来自对经验神经正切核(Empirical Neural Tangent Kernel, eNTK)的分析(该核可用于近似少量微调时的模型行为(2022))
    • eNTK 基于梯度的点积,具体而言,梯度:
      $$ g_i = \partial/\partial \theta \log p(\text{token}_i \mid \text{prefix}_i) $$
    • 且 \( K(i,j) = g_i \cdot g_j \)
    • 因此,参数数量最多的层通常对核的影响最大
  • 该论文还指出,当对所有层进行训练时,LoRA 的 eNTK 与 FullFT 的 eNTK 大致相同,因此:
    $$\text{LoRA training} \approx \text{eNTK(LoRA)}\approx \text{eNTK(FullFT)} \approx \text{FullFT} $$
    • 而 “\(\text{eNTK(LoRA)}\approx \text{eNTK(FullFT)}\)” 这一近似关系仅在将 LoRA 应用于“构成梯度点积的主要参数所在层”时成立

How Much Capacity Is Needed by Supervised and Reinforcement Learning?

  • 监督学习与强化学习需要多少容量?
    • 注:这里指参数容量
  • 以往的研究(2024)表明,神经网络每个参数可存储 2 比特信息
    • 但该结果针对的是长期训练极限下模型可吸收的最大信息量,而非计算效率或学习速度
  • “每个参数 2 比特”的结论基于精心构建的合成数据集(这些数据集包含精确的信息量)
    • 对于现实学习任务,估算所需的信息量则更为复杂
  • 一个经典结论是:
    • 在最小化对数损失时,首轮训练的总对数损失可用于衡量数据集的描述长度,即记忆该数据集所需比特数的上限
    • LLM 数据集的损失通常约为每 Token 1 比特(0.69 纳特),具体数值因数据集和模型大小而异

      One classic observation is that when minimizing log-loss, the total log-loss measured during the first epoch of training provides a measurement of the dataset’s description length. That is, an upper bound for the number of bits required to memorize the dataset

  • 该估算值衡量的是“完美记忆数据集”所需的容量,而这一数值高估了“降低测试集对数损失”的泛化学习所需的实际容量
  • 目前,监督学习的容量需求及其与可训练参数数量的关系仍是有待未来研究的开放问题
  • 对于强化学习,作者认为,由于每个样本的末尾仅有一个奖励值,策略梯度算法每个样本大致仅学习 1 比特信息
    • 但这并非强化学习的固有属性,其他算法可能从每个样本中学习更多信息
    • 例如,基于模型的强化学习(Model-Based RL)算法会训练智能体预测观测结果并构建世界模型(World Model),从而可能从每个样本中提取更多信息
    • “每个样本 1 比特”的结论可能仅适用于策略梯度算法
    • 注:这里每个样本仅 1 比特的理论不敢苟同,因为强化学习的信号也可以不止是 0-1 信号,可以是 Reward 啊,此时就不再是 1 比特能够表达的了吧?
  • 我们可以从信息论角度更精确地阐述比特计数论证
    • 考虑一个样本(包含轨迹 \( \tau \) 和最终奖励)作为一条消息(即一个噪声信道),该消息包含关于未知奖励函数 \( R \) 的部分信息
    • 论文基于当前策略和训练历史,分析策略梯度估计量与 \( R \) 之间的互信息
    • REINFORCE 算法的更新公式为
      $$ G = S \cdot \text{Adv} $$
      • 其中 \( S = \nabla \log p_\theta(\tau) \)
    • 在给定历史的情况下,\( S \) 与 \( R \) 无关,因此唯一依赖 \( R \) 的部分是标量优势函数(Advantage)
  • 根据数据处理不等式:
    $$ I(G; R \mid \text{history}) \leq I((S, \text{Adv}); R \mid \text{history}) = I(\text{Adv}; R \mid S, \text{history}) \leq H(\text{Adv}) $$
  • 若将优势函数量化为 \( B \) 个区间,则 \( H(\text{Adv}) \lesssim \log(B) \)
    • 即每个样本可获取的有效信息比特数为 \( O(1) \),与模型大小无关
    • 这些比特信息告诉论文,当前面对的是离散奖励函数集合(或等效的最优策略类别集合)中的哪一个
    • 这种互信息分析与一些优化算法的理论分析方法一致(2009)
  • 需要注意的是,该估算值是训练可吸收信息量的上限,实际学习到的信息量会依赖于策略初始化和其他细节
    • 例如,若初始化的策略无法获得任何奖励,则优势函数的熵为 0(而非 \( \log(B) \)),模型无法学习到任何内容

Compute Efficiency Advantage of LoRA

  • 上述实验通过训练步数衡量学习进度,但论文也关注不同方法的计算效率
  • 论文计算得出,LoRA 每轮前向-反向传播所需的浮点运算数(FLOPs)略多于 FullFT 的 2/3
    • 因此,总体而言,LoRA 的计算效率通常会优于 FullFT
  • 论文通过分析给定权重矩阵上“前向-反向传播”操作的 FLOPs 来推导这一 2/3 比例,这些操作占据了神经网络模型中绝大部分的 FLOPs。论文使用以下符号:
    • \( W \in \mathbb{R}^{N \times N} \):权重矩阵
    • \( x \in \mathbb{R}^N \):输入向量
    • \( y = Wx \in \mathbb{R}^N \):输出向量
    • \( \bar{x}, \bar{y} \in \mathbb{R}^N \):损失分别对 \( x \) 和 \( y \) 的梯度(在反向传播中计算)
    • \( \bar{W} \in \mathbb{R}^{N \times N} \):损失对 \( W \) 的梯度
  • 全参数微调(FullFT)执行以下操作:
    • 1)前向传播(Forward)
      • \( y = Wx \)(需 \( N^2 \) 次乘加运算)
    • 2)反向传播(Backward)
      • \( \bar{x} = W^T \bar{y} \)(需 \( N^2 \) 次乘加运算)
      • \( \bar{W} += x \bar{y}^T \)(需 \( N^2 \) 次乘加运算)
    • 前向传播需 \( N^2 \) 次乘加运算,反向传播需额外 \( 2 \cdot N^2 \) 次乘加运算,总计 \( 3N^2 \) 次
      • 因此,包含前向和反向传播的训练过程,其 FLOPs 是仅前向推理的 3 倍
  • 在 LoRA 中,论文将 \( W \) 替换为 \( W + BA \),其中 \( B \in \mathbb{R}^{N \times R} \)、\( A \in \mathbb{R}^{R \times N} \),且 \( R \ll N \)
    • 由于论文仅更新 \( \bar{A} \) 和 \( \bar{B} \),因此无需执行 FullFT 中“更新 \( \bar{W} \)”的第三步,而是替换为代价低得多的操作
    • \( A \) 和 \( B \) 均为 \( N \cdot R \) 规模的矩阵,因此各自的“前向-反向传播”完整计算需 \( 3NR \) 次乘加运算,而非 FullFT 中 \( W \) 所需的 \( 3N^2 \) 次;两者合计需 \( 6NR \) 次乘加运算
    • 同时,论文仍需对 \( Wx \) 和 \( \bar{x} \) 执行前向-反向传播(等效于 FullFT 的前两步)
    • 因此,LoRA 总乘加运算次数为 \( 2N^2 + 6NR \)
    • 由于 \( R \ll N \),该数值略多于 FullFT 总运算量(\( 3N^2 \))的 2/3
  • 若论文以 FLOPs(而非训练步数)为横轴绘制 LoRA 性能曲线(本分析未包含注意力机制的 FLOPs,在长上下文场景中,注意力 FLOPs 可能占比显著),LoRA 相对于 FullFT 的优势将更为明显

Open Questions

  • 关于论文的研究结果,仍有几个问题有待未来探索:
    • 1)细化 LoRA 性能预测 :论文已大致描述了 LoRA 与 FullFT 性能相当的区间,并能通过 Token 或样本数量估算所需容量,但尚未能做出精确预测
      • 未来需进一步明确 LoRA 性能的预测方法及与 FullFT 匹配的具体条件
    • 2)LoRA 学习率与 Training Dynamics 的理论解释 :目前论文对 LoRA 学习率和 Training Dynamics 的理论理解仍有限
      • 若能建立完整理论解释 LoRA 与 FullFT 学习率的比例关系,将具有重要价值
    • 3)LoRA 变体的性能评估 :如 PiSSA(2024)等 LoRA 变体,若采用论文方法评估,其性能表现如何?
    • 4)MoE 层的 LoRA 应用方案 :将 LoRA 应用于 MoE 层有多种可选方案
      • 未来需研究不同方案的性能,以及每种方案与“张量并行”“专家并行”等大型 MoE 模型关键技术的兼容性

Closing Thoughts

  • 在 Thinking Machines,作者相信微调技术能提升人工智能在多个专业领域的实用性
  • 论文对 LoRA 的关注,源于让这种技术“广泛可及”并“易于定制以满足特定需求”的目标
  • 除实际应用价值外,LoRA 研究还促使论文深入探索模型容量、数据集复杂度和样本效率等基础问题
  • 通过分析“学习速度与性能如何依赖容量”,论文获得了研究机器学习基础问题的新视角
  • 论文期待未来能进一步推进这一领域的研究

NLP——样本packing与权重讨论


整体总结

  • 样本 packing 训练可以大幅提升模型训练速度
  • 由此引发的训练样本的组织方式和 loss 权重有多种组合形式
  • 本文对 样本 packing 和 packing 后的样本权重(packing weight)相关的各种情况进行总结

packing 模式的一些结论

  • 下面两个数据处理方式在算法理论层面不等价(特别是在 tool 等 prompt 很长,response 很短的场景):
    • 数据处理方案一:将多轮拆成多个单轮样本训练
    • 数据处理方案二:保留多轮,即多轮本身是一个样本,但训练计算损失时每个轮次视作一个样本(即按照轮次处理 Loss 归一化)
  • 上述两种训练方式不等价的原因包括:
    • 样本训练时的分布不等价
    • 样本训练时的权重不等价
  • 多轮拆单轮/保留多轮 的数据处理方式暂无法简单从算法层面实现等价
    • 除非 GBS 无穷大,达到所有样本一次过完

非 packing 下的训练模式(与 packing 无关的朴素模式)

数据集内所有 token 权重相等训练模式

  • 具体实现方式为:
    $$
    \begin{align}
    \text{Loss}(\text{batch}) &= \sum_{\text{token}_i \in \text{batch}} \text{loss}_{\text{token}_i} \\
    \text{Loss}(\text{batch}) &= \frac{1}{\text{GBS}} \sum_{\text{token}_i \in \text{batch}} \text{loss}_{\text{token}_i}
    \end{align}
    $$
    • 每个 token 的 loss 直接累加,不做任何归一化 或 仅除以 Global Batch 内样本数 GBS(一般对同一次训练来说是固定值)
  • 这种方式可以确保训练的整个过程中,所有 token 的权重都是相同的
  • 可能会导致部分短序列得到的关注度不够高
  • 预训练中就应该使用这种方式,每个 Token 是等权的

Global Batch 内部每个 token 权重相等训练模式(Global Batch 间的 token 权重可能不相等)

  • 具体实现方式:
    $$
    \text{Loss}(\text{batch}) = \frac{1}{\text{GBS_Token_Num}} \sum_{\text{token}_i\in \text{batch}} \text{loss}_{\text{token}_i}
    $$
    • 每个 token 的 loss 累加,除以 Global Batch 内的 token 总数
  • 这种方式可以确保同一个 Global Batch 内部的 token 权重相等
  • 不同 Global Batch 内部的权重可能是不相同的,跟 Global Batch 内部的 token 总数有关
  • 学习率相同的情况下,不同 Global Batch 对模型训练的贡献相等
  • 当 Global Batch Size 非常大时,这个实现方式几乎等价于 数据集所有 token 权重相等训练模式

样本内部每个 token 权重相等训练模式(样本间的 token 权重可能不相等)

  • 具体实现方式:
    $$
    \begin{align}
    \text{Loss}(\text{batch}) &= \sum_{\text{sample} \in \text{batch}} \frac{1}{|\text{sample}|} \sum_{\text{token}_i \in \text{sample}}\text{loss}_{\text{token}_i} \\
    \text{Loss}(\text{batch}) &= \frac{1}{\text{GBS}} \sum_{\text{sample} \in \text{batch}} \frac{1}{|\text{sample}|} \sum_{\text{token}_i \in \text{sample}} \text{loss}_{\text{token}_i}
    \end{align}
    $$
    • 每个 token 的 loss 累加,除以 样本 内的 token 总数,同时样本间 不做任何归一化 或 除以 Global Batch 内样本数 GBS
  • 这种情况下可以防止长序列样本权重过高,让不同样本对训练的影响权重相同
  • 长序列样本中的 Token 可能因为平均后权重过低而学不好
  • 很多框架实现中的目标就是实现这样的方式,想要保证样本粒度的 Loss 贡献相同

样本 packing 下的训练模式

数据集内所有 token 权重相等训练模式

  • 具体实现方式:
    $$
    \begin{align}
    \text{Loss}(\text{batch}) &= \sum_{\text{token}_i \in \text{batch}} \text{loss}_{\text{token}_i} \\
    \text{Loss}(\text{batch}) &= \frac{1}{\text{GBS}} \sum_{\text{token}_i \in \text{batch}} \text{loss}_{\text{token}_i}
    \end{align}
    $$
    • 每个 token 的 loss 直接累加,不做任何归一化 或 除以 Global Batch 内样本数 GBS(一般对同一次训练来说是固定值)
  • 这种方式下,与非 packing 的训练方式相同
  • 注:如果想要与非 packing 情况下的 数据集内所有 token 权重相等训练模式 完全等价,自回归下 Attention 时需要对前面的样本做 mask 才可以
  • 预训练中就应该使用这种方式

Global Batch 内部每个 token 权重相等训练模式(Global Batch 间的 token 权重可能不相等)

  • 具体实现方式:
    $$\text{Loss}(\text{batch}) = \frac{1}{\text{GBS_Token_Num}} \sum_{\text{token}_i\in \text{batch}} \text{loss}_{\text{token}_i}$$
    • 每个 token 的 loss 累加,除以 Global Batch 内的 token 总数
  • 此时已经是 packing 以后的样本,但理论上仍然与 非 packing 下的训练方式相同
  • 注:如果想要完全等价于 非 packing 模式下的 Global Batch 内部每个 token 权重相等训练模式 ,自回归下 Attention 时需要对前面的样本做 mask 才可以

packing 样本内部每个 token 权重相等训练模式(样本间的 token 权重可能不相等)

  • 具体实现方式:
    $$
    \begin{align}
    \text{Loss}(\text{batch}) &= \sum_{\text{sample} \in \text{batch}} \frac{1}{|\text{sample}|} \sum_{\text{token}_i \in \text{sample}}\text{loss}_{\text{token}_i} \\
    \text{Loss}(\text{batch}) &= \frac{1}{\text{GBS}} \sum_{\text{sample} \in \text{batch}} \frac{1}{|\text{sample}|} \sum_{\text{token}_i \in \text{sample}} \text{loss}_{\text{token}_i}
    \end{align}
    $$
    • 每个 token 的 loss 累加,除以 样本 内的 token 总数,同时样本间 不做任何归一化 或 除以 Global Batch 内样本数 GBS
  • 这种情况下由于短序列样本被 packing 为固定长度的样本,所以与 非 packing 下的 样本内部平均完全不同 ,无法做到每个样本贡献度一致了(长样本占优)

真实样本内部每个 token 权重相等训练模式

  • 每个 token 的 loss 累加,除以 真实样本(需要特殊手段去识别) 内的 token 总数,同时有归一化有三种实现方式:
  • 方式一:样本间 不做任何归一化
    $$
    \text{Loss}(\text{batch}) = \sum_{\color{red}{\text{sample}_{true}} \in \text{batch}} \frac{1}{|\color{red}{\text{sample}_{true}}|} \sum_{\text{token}_i \in \color{red}{\text{sample}_{true}}}\text{loss}_{\text{token}_i}
    $$
  • 方式二:除以 Global Batch 内样本数 GBS
    $$
    \text{Loss}(\text{batch}) = \frac{1}{\text{GBS}} \sum_{\color{red}{\text{sample}_{true}} \in \text{batch}} \frac{1}{|\color{red}{\text{sample}_{true}}|} \sum_{\text{token}_i \in \color{red}{\text{sample}_{true}}} \text{loss}_{\text{token}_i}
    $$
  • 方式三:除以 Global Batch 内真实样本数 GBS_true
    $$
    \text{Loss}(\text{batch}) = \frac{1}{\color{red}{\text{GBS_true}}} \sum_{\color{red}{\text{sample}_{true}} \in \text{batch}} \frac{1}{|\color{red}{\text{sample}_{true}}|} \sum_{\text{token}_i \in \color{red}{\text{sample}_{true}}} \text{loss}_{\text{token}_i}
    $$
  • 特别注意,此时除以 GBS 和 除以 GBS_true 是不同的,为了保证真实样本间的权重一致,应该除以 GBS_true

多轮场景的训练思考

  • 有时候想要将多轮中的每一轮单独看做一个样本(即一次回复一个样本)
    • 这种思路的基础逻辑是认为一次回复就应该是一个样本
    • 实际上,评估指标中,可能涉及多轮回复只对应一个分数,即多轮回复对整体评估的贡献跟单轮回复权重一样,这时候要注意训练指标和评估指标的一致性问题
  • 多轮场景下,即使不做样本 packing,也相当于隐含的存在 样本 packing 了,如果想要做到轮次间权重一致,可以参考 真实样本内部每个 token 权重相等训练模式 的实现,这种实现方式可以分为两种:
  • 方式一:将多轮拆成多个样本,labels 仅保留最后一个回复,再使用用样本平均的方式实现
    • 注:这种拆开的方式会引入一些特别的变化,且是不可恢复的
      • 这种拆开会导致重复计算一些 prompt ,且会将同一个多轮样本分散到不同的 Batch 上学习,是不可恢复的
      • 这种拆开还会导致 拆单轮 情况 vs 不拆单轮 情况的单个 Batch 内数量和分布不一致,导致一些差异,也是不可恢复的
        • 这里的影响应该还好,主要强调的是不可恢复性
    • 总结:拆开的方式虽然方便,但与合并训练的方式相比,无论如何也无法恢复到真实水平
  • 方式二:不管是否 packing,都将单个轮次识别为当个真实样本做归一化和学习
    • 此时不管是否 packing,都可以做到 轮次即样本,两者基本等价

在 DP 不为 1 的情况思考

  • 在 DP 不为 1 的情况下,每个 DP 内部是独立计算梯度的,所以需要先做归一化

NLP——GShard

注:本文包含 AI 辅助创作

  • 参考链接:
    • 专家并行论文:GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, 2020, Google

Paper Summary

  • 评价:
    • 本文是 2020 年的文章,是非常重要的 Scaling 方向的作品,对模型规模的扩展意义重大
    • 核心思路:
      • 条件计算(Conditional Computation):采用了稀疏门控混合专家(Sparsely-Gated MoE)架构
      • 自动分片(Automatic Sharding):通过轻量级标注 API(replicate/split/shard)和 XLA 编译器扩展,实现张量的自动分布式划分
      • SPMD 编译优化:提出单程序多数据(SPMD)编译范式(替代传统多程序多数据(MPMD)模式),为所有设备生成单一程序,编译时间与设备数量无关
    • Google 出品,必属精品
  • GShard 是一个由一组轻量级标注接口(API)和 XLA 编译器扩展组成的模块
    • Gshard 提供了一种简洁的方式,只需对现有模型代码进行少量修改,就能实现多种并行计算模式
    • Gshard 解决了神经网络扩展过程中的 计算成本、编程便捷性以及在并行设备上的高效实现等问题
  • 作者使用 GShard 将带有稀疏门控混合专家层的多语言神经机器翻译 Transformer 模型扩展到了 600B 参数以上(应该是当时第一次有这么大的模型)
  • 论文的实验表明:
    • 这样一个巨型模型能够在 2048 个 TPU v3 加速器上高效训练 4 天,在 100 种语言到英语的翻译任务中,取得了远超现有技术水平的翻译质量

Introduction and Discussion

  • 在各类机器学习问题中,神经网络扩展都带来了显著的质量提升
    • 在计算机视觉领域,提升模型容量使得多种计算机视觉架构在图像分类和检测精度上有了更好的表现
    • 在自然语言处理领域,Transformer 模型的扩展在语言理解任务(2018, 2019)、跨语言下游迁移任务(2018, 2019)以及(大规模)多语言神经机器翻译任务(2019)中都实现了稳定的性能提升
  • 这种普遍趋势促使近期研究开始深入探究影响扩展成功的关键因素(2017, 2019, 2020),其中包括过往研究发现的训练数据量、模型大小和计算资源利用率等
  • 尽管研究发现最终模型质量与数据量、计算资源和模型大小呈幂律关系(2017, 2020),但更大模型带来的显著质量提升也伴随着各种实际挑战
  • 训练效率便是其中至关重要的挑战之一,论文将其定义为:为获得优于现有最佳系统的模型质量所消耗的计算资源和训练时间 ,而这一指标往往被忽视
  • 图1:
    • 多语言翻译质量(与双语基线模型相比的平均 \(\delta\) BLEU 值)随着 MoE 模型规模增长到 600B 参数而不断提升,而端到端训练成本(以 TPU v3 核心年为单位)仅呈亚线性增长
    • 将模型规模从 37.5B 参数扩大到 600B 参数(16 倍),计算成本仅从 6 个核心年增加到 22 个核心年(3.6倍)
    • 达到最佳翻译质量的 600B 参数模型,使用 2048 个 TPU v3 核心训练了 4 天,总成本为 22 个 TPU v3 核心年
    • 相比之下,训练所有 100 个双语基线模型则需要 29 个 TPU v3 核心年
    • 论文训练的最佳质量密集型单一 Transformer 模型(2.3B 参数), \(\delta\) BLEU 值为 6.1,该模型使用 GPipe(2019)在 2048 个 TPU v3 核心上训练了 6 周,总成本达 235.5 个 TPU v3 核心年

Practical Challenges for Scaling(扩展面临的实际挑战)

  • 本节将列举在训练超大规模模型时面临的主要实际挑战,这类模型的规模远超单个加速器内存(如 GPU 或 TPU)的容量限制
  • 特定架构的模型并行支持问题 :在 TensorFlow(2016)和 PyTorch(2017)等常用深度学习框架中,缺乏对高效模型并行算法的支持
    • 框架虽支持基于图分割的简单模型并行,但由于网络的顺序依赖性和基于梯度的优化方式,这种并行方式会导致设备利用率严重不足
    • 为了高效扩展现有模型,用户通常需要投入大量工程工作,例如将模型代码迁移到专门的框架中(2018, 2019)
  • 计算成本与模型规模的超线性扩展问题 :通过增加网络深度或宽度来直接扩展模型规模(2018, 2019),通常会导致训练步长时间至少呈线性增长
    • 为解决这一问题,通常需要通过在多个设备间分割层权重和计算过程来实现模型并行,但这会带来网络通信开销和设备利用率不足的问题
    • 设备利用率不足源于神经网络底层计算任务分配不均衡以及存在顺序依赖关系
    • 这种计算成本与模型规模之间的超线性关系,无法通过简单增加设备数量来解决,使得训练超大规模模型变得不切实际
  • 巨型模型表示的基础设施扩展性问题 :对于分布在数千个设备上的超大规模模型,简单的图表示可能会成为深度学习框架及其优化编译器的瓶颈
    • 例如,通过操作间分割(inter-op partitioning)增加 D 倍的层数,或通过操作内分割(intra-op partitioning)在 D 个设备上增加模型维度,都可能导致图节点数量大幅增加
    • 设备间的通信通道可能会进一步使图大小增加(例如,分割聚合或转置操作时)
    • 对于超大规模模型而言,图大小的这种增长会导致图构建和编译时间达到无法实现的程度
  • 实现分片策略的复杂工作 :将模型高效地分割到多个设备上运行具有挑战性,因为这需要协调设备间的通信
    • 在图级分割方面,需要复杂的算法(2019, 2018)来减少因不同设备上分配的图分割部分之间存在顺序依赖而引入的开销
    • 在算子级并行方面,不同分割后的算子会有不同的通信模式,这取决于算子的语义,例如是否需要累积部分结果或重新排列数据分片
    • 根据论文的经验,由于 TensorFlow 等框架包含大量具有特殊语义的算子,在模型中手动处理这些问题需要耗费大量精力
    • 在所有情况下,实现模型分片对研究人员和工程师来说都是一项负担,因为改变模型架构就需要修改底层的设备通信逻辑,从而产生连锁反应

Design Principles for Efficient Training at Scale

  • 论文展示了如何通过构建一个拥有 600B 参数、带有稀疏门控混合专家层(Sparsely-Gated Mixture-of-Experts layers)的 sequence-to-sequence Transformer 模型,来克服上述挑战
    • 该模型的计算成本呈亚线性增长,编译时间为 \(O(1)\)
    • 论文在 2048 个 TPU v3 设备上,针对多语言机器翻译任务训练该模型 4 天,最终实现了单个非集成模型在 100 种语言到英语翻译任务中远超现有技术的翻译质量
    • 论文对不同规模的模型进行了实验,结果发现,随着模型规模增大,翻译质量不断提升,而训练总耗时(wall-time)相对于模型规模仅呈亚线性增长,如图1所示
  • 为构建如此庞大的模型,论文做出了以下关键设计选择
    • 亚线性扩展(Sub-linear Scaling) :首先,模型架构的设计应确保计算和通信需求相对于模型容量呈亚线性增长
      • 条件计算(2015, 2019, 2020, 2020)通过在每个输入样本的基础上激活一个子网络,使论文能够兼顾训练和推理效率
      • 通过在基于循环神经网络(RNN)的机器翻译和语言模型中添加位置感知稀疏门控混合专家层(Position-wise Sparsely Gated Mixture-of-Experts, MoE)(2019),可以在实现最先进性能的同时,使计算成本呈亚线性增长
      • 因此,论文将在第2节中详细介绍如何用混合专家层扩展 Transformer 架构
    • 抽象的力量(The Power of Abstraction) :其次,模型描述应与分片实现和优化相分离
      • 这种关注点分离使模型开发人员能够专注于网络架构,并灵活地改变分片策略,而底层系统则负责执行语义保持转换并实现高效的并行执行
      • 为此,论文提出了 GShard 模块,用户只需在模型中对少数关键张量标注分片策略即可
      • 该模块包含一组简单的标注接口,以及一个用于自动并行化的 XLA(2019)编译器扩展
      • 模型开发人员可以像在一个拥有超大内存和计算能力的单一设备上编写模型,编译器会根据标注信息和自身的启发式算法,自动为目标设备分割计算任务
      • 论文将在3.2节中提供更多标注示例
    • 可扩展编译器(Scalable Compilers) :第三,包括计算表示和编译在内的系统基础设施,必须能够支持数千个设备的并行执行
      • 例如,图2 展示了在 4 个设备上分割点积运算(用颜色编码)的两种不同方式
      • 需要注意的是,图2a 中常用的 MPMD(多程序多数据,Multiple Program Multiple Data)方法在扩展性方面面临挑战
        • 因为图中的节点数量会随着设备数量的增加而线性增长
      • 相反,论文开发了一种用于 SPMD(单程序多数据,Single Program Multiple Data)转换的编译器技术
        • 该技术生成一个可在所有设备上运行的单一程序,使得编译时间与设备数量无关,保持恒定,如图2b所示
        • 论文将在3.3节中详细讨论 SPMD 框架
  • 论文其余部分的结构如下:
    • 第2节详细描述带有稀疏门控混合专家层的 Transformer 架构;
    • 第3节介绍论文开发的 GShard 模块;
    • 第4节展示混合专家模型在100个语言对的多语言机器翻译任务中的应用;
    • 第5节对实现的性能和内存使用情况进行评估;
    • 第6节讨论相关工作
  • 图2:
    • 在4个设备上对点积算子( \([M, K] \times [K, N] = [M, N]\) )进行 MPMD 分片与论文提出的 SPMD 分片的对比
    • 在该示例中,两个操作数均沿收缩维度 K 进行分片,每个设备计算本地结果后,通过 AllReduce 操作进行全局合并
    • MPMD 分片会为每个设备生成独立的算子,限制了其扩展性;
    • 而 SPMD 分片则生成一个可在所有设备上运行的程序
    • 需要注意的是,使用论文提出的 SPMD 分片时,编译时间与所使用的设备数量无关

Model

Sparse scaling of the Transformer architecture

  • Transformer架构(2017)已被广泛应用于自然语言处理领域,成为许多 sequence-to-sequence 任务(如机器翻译)的事实标准
  • Transformer 包含两个计算模块,即 Encoder 和 Decoder,两者均通过堆叠多个 Transformer 层实现
  • Transformer Encoder 层由两个连续的层组成,即自注意力层(self-attention layer)之后紧跟一个位置感知前馈层(position-wise feed-forward layer)
  • Decoder 则额外增加了第三个交叉注意力层(cross-attention layer),该层会关注 Encoder 的输出
  • 论文通过条件计算(conditional computation)对 Transformer 进行稀疏扩展,在 Encoder 和 Decoder 中,每隔一个前馈层就用一个位置感知混合专家层(Position-wise Mixture of Experts, MoE)(2019)替代,该混合专家层采用了一种改进的 top-2 门控机制(图3)
  • 论文通过改变 Transformer 层的数量和每个混合专家层中专家的数量来调整模型容量
  • 每个训练样本由一对子词 Token(subword tokens)序列组成
  • 在训练和推理过程中,每个 Token 都会激活混合专家 Transformer 的一个子网络
  • 子网络的大小大致与每个混合专家层中的专家数量无关,这使得计算成本能够像前一节所述的那样呈亚线性增长
  • 3.1节将进一步分析计算复杂度,5节将分析训练性能

Position-wise Mixture-of-Experts Layer

  • 论文模型中使用的 MoE 层基于文献(2019)的设计,并在稀疏门控函数和辅助损失函数方面进行了改进
  • Transformer 的混合专家层由 E 个前馈网络(FFN)组成,其计算过程如下:

$$ g_{s,e} = \text{GATE}(x_s) \quad \tag{1} $$

$$ \text{FFN}_e(x_s) = w o_e \cdot \text{ReLU}(w i_e \cdot x_s) \quad \tag{2} $$

$$ y_s = \sum_{e=1}^{E} g_{s,e} \cdot \text{FFN}_e(x_s) $$

  • 其中:
    • \(x_s\) 是混合专家层的输入 Token
    • \(w i_e\) 和 \(w o_e\) 分别是前馈层(即一个专家)的输入投影矩阵和输出投影矩阵
    • 向量 \(g_{s,e}\) 由门控网络计算得到
    • \(g_{s,e}\) 中每个专家对应一个非负值,其中大部分值为零,这意味着该 Token 不会被分配给对应的专家
      • 每个 Token 仅被分配给极少数专家,论文设定每个 Token 最多被分配给两个专家
    • \(g_{s,e}\) 中对应的值非零,表示该专家对最终网络输出的贡献程度
    • 每个专家 \(\text{FFN}_e\) 都采用含 ReLU 激活函数(2010)的两层全连接网络对 \(x_s\) 进行处理
    • 混合专家层的输出 \(y_s\) 是所有被选中专家输出的加权平均值
  • 门控函数 \(\text{GATE}(\cdot)\) 是混合专家层的核心,它采用 softmax 激活函数来表示每个专家在处理输入 Token 时的权重,即表示某个专家处理该输入 Token 的适合程度
  • 此外,门控函数必须满足以下两个目标:
    • 负载均衡(Balanced load) :对于给定的 Token ,混合专家层最好能稀疏地激活专家
      • 一种简单的方案是根据 softmax 概率分布选择前k个专家,但研究表明这种方法会导致训练中的负载不均衡问题(2019):训练过程中看到的大多数 Token 会被分配给少数几个专家,导致这几个(繁忙的)专家的输入缓存变得极大,而其他专家则处于未充分训练的状态,从而减慢训练速度
        • 问题:分配给少数专家也不一定会导致训练速度变慢吧,除非有 EP (专家并行)?
      • 同时,许多其他专家根本无法得到充分训练
      • 因此,门控函数需要设计得更合理,以实现所有专家间处理负载的更均匀分配
    • 大规模下的效率(Efficiency at scale) :如果门控函数采用顺序执行方式,实现负载均衡相对容易
      • 但对于输入批次中所有 N 个 Token 和 E 个专家 ,仅门控函数的计算成本就至少为 \(O(N \cdot E)\)
      • 在论文的研究中,N 达到数百万量级,E达到数千量级,若门控函数采用顺序实现,会导致大部分计算资源在大部分时间处于空闲状态
      • 因此,论文需要一种高效的并行门控函数实现方式,以充分利用多个设备的计算能力
  • 为满足上述要求,论文在门控函数 \(\text{GATE}(\cdot)\) 中设计了以下机制(细节如算法1所示):
    • 专家容量限制(Expert capacity) :为确保负载均衡,论文强制每个专家处理的 Token 数量不超过某个统一的阈值,论文将该阈值定义为专家容量
      • 假设一个训练批次中的 Token 总数为 N,每个 Token 最多被分配给两个专家,那么专家容量设置为 \(\frac{2N}{E}\)
      • 门控函数 \(\text{GATE}(\cdot)\) 会为每个专家维护一个计数器 \(c_e\) ,用于记录分配给该专家的 Token 数量
      • 当一个 Token 选中的两个专家都已超过其容量限制时,该 Token 被视为溢出 Token ,此时 \(g_{s,e}\) 退化为零向量
      • 这类 Token 的表示会通过残差连接传递到下一层
        • 问题:这里是说专家容量限制也传递到下一层?
    • 本地组分配(Local group dispatching) :
      • 门控函数 \(\text{GATE}(\cdot)\) 将训练批次中的所有 Token 均匀划分为 G 个组,即每个组包含 \(S = \frac{N}{G}\) 个 Token ,所有组独立并行处理
      • 每个组被分配到每个专家的部分容量为 \(\frac{2N}{G \cdot E}\)
      • 每个组会确保分配给每个专家的 Token 数量不超过该部分容量
      • 通过这种方式,既能保证专家容量限制得到遵守,又能实现整体负载均衡
    • 辅助损失(Auxiliary loss) :门控函数不应总是选择相同的几个专家,否则会导致仅少数专家出现容量溢出,而其余专家利用率不足(这一点至关重要)
      • 借鉴文献(2019)的方法,论文定义了一个辅助损失项来强制执行这一约束
      • 该辅助损失项以一个常数系数 k 添加到模型的总损失函数中
      • 算法1第13行中辅助损失项的具体形式基于以下考虑:
        • \(\frac{c_e}{S}\) 表示分配给每个专家的输入 Token 比例,作者希望最小化 \(\frac{c_e}{S}\) 的均方值
        • 但由于 \(c_e\) 是通过 top-2 操作得到的,不具有可微性,因此论文使用每个专家的平均门控值作为可微近似,并用 \(m_e \cdot \frac{c_e}{S}\) 替代 \((\frac{c_e}{S})^2\) ,这样就可以通过梯度下降法对其进行优化
    • 随机路由(Random routing) :
      • 直观上,由于 \(y_s\) 是所选专家输出的加权平均值,如果第二个专家的权重非常小,我们可以直接忽略该专家 ,以节省整体专家容量
      • 因此,除了遵守专家容量约束外,门控函数 \(\text{GATE}(\cdot)\) 还会以与第二个最佳专家权重 \(g_2\) 成比例的概率,将 Token 分配给该专家

Highly Parallel Implementation using GShard(高度并行化)

  • 本节将介绍第2节中模型在张量处理单元(TPU)设备集群上的高效运行实现方案
  • 实现的第一步是将模型用线性代数运算表示,因为论文的软件栈(TensorFlow(2016))和硬件平台(TPU)在这类运算上经过了高度定制和优化
    • 与原始 Transformer 模型类似,将模型的大部分部分用线性代数运算编写非常容易
    • 但由于混合专家层(MoE Layer),尤其是算法1 中提出的 \(\text{GATE}(\cdot)\) 函数具有顺序执行特性,要用线性代数运算表示该层则需要额外付出努力,论文将在3.1节详细介绍相关细节
  • 接下来,论文对线性代数计算进行标注,以体现并行性
  • 通过3.2节中的分片接口(sharding APIs),可以对计算中的每个张量进行标注,指定其在设备集群中是采用复制(replication)还是分布(distribution)方式存储
  • 分片标注能够实现模型描述与高效并行实现的关注点分离,让用户可以灵活地表达各种并行化策略;例如:
    • (1)注意力层(attention layer)通过沿批次维度(batch dimension)分割并将权重复制到所有设备上来实现并行化;
    • (2)由于混合专家层中专家(expert)的规模极大,无法在所有设备上进行复制,因此唯一可行的策略是将专家分片存储到多个设备中
    • 此外,整个模型会在这两种模式(1)-(2)之间切换
    • 通过标注,模型开发人员无需关注系统优化工作,也不必将并行实现细节和底层细节融入模型代码中
  • 最后,编译器基础设施会接收(部分)标注后的线性代数计算,并生成可在数千个设备上高效运行的并行程序
  • 如3.3节所述,编译器会应用单程序多数据(SPMD,Single Program Multiple Data)分片转换来表示每个设备的计算任务,插入必要的跨设备通信操作,处理非规则模式(如非均匀分片),最终生成一个可在所有设备上启动并执行并行计算的单一程序

Positions-wise Mixture-of-Expert Layer Expressed in Linear Algebra(线性代数)

  • 论文的模型实现(算法2)将整个加速器集群视为一个单一设备,并用少量与集群具体配置无关的张量运算来表示其核心数学算法

    • 爱因斯坦求和符号(Einstein summation notation)(1923)(即tf.einsum)是一种能够简洁表示模型的强大工具,论文在实现中大量使用了该符号
    • softmax 门控计算可以通过一个爱因斯坦求和运算(einsum)后紧跟 softmax 函数来轻松表示;
    • 将输入分配给选定专家的操作,可以通过分配掩码(dispatching mask)与输入之间的单次爱因斯坦求和运算来表示;
    • 所有 \(\text{FFN}_e\) 的权重被组合成单个三维张量 \(wi\) 和 \(wo\)
    • \(\text{FFN}_1 \dots \text{FFN}_E\) 的计算则通过三个算子(两个爱因斯坦求和运算和一个relu运算)来表示;
    • 最后,对所有专家输出进行加权平均以得到最终输出的操作,也可以通过另一个爱因斯坦求和运算来表示
  • 算法2中的 Top2Gating 计算了算法1中所有组本地(group-local)门控决策的并集

    • combine_weights是一个四维张量,形状为[G, S, E, C]
    • 其中,combine_weights[g, s, e, c] 非零时,表示组 g 中的输入 Token s被发送到专家 e 的输入缓存中,且位于缓存位置 c
    • 对于特定的 g 和 s,combine_weight 切片中最多包含两个非零值
    • 通过将所有非零值设为 1,可由 combine_weights 生成二进制分配掩码(binary dispatch_mask)
  • 论文需要合理选择组数(G)和专家数量(E),以确保该算法能够在包含D个设备的集群上实现扩展

    • 有必要分析在一个包含 N 个 Token 的训练批次中,该算法在一个训练步骤内的整体计算复杂度(浮点运算总数)
  • 算法2 :位置感知混合专家层的前向传播(Forward pass of the Positions-wise MoE layer),下划线字母(如G和E)表示张量将沿该维度进行分片

    1
    2
    3
    4
    5
    6
    7
    1: gates = softmax(einsum("GSM, ME -> GSE", inputs, wg))
    2: combine_weights, dispatch_mask = Top2Gating(gates)
    3: dispatched_expert_inputs = einsum("GSEC, GSM -> EGCM", dispatch_mask, reshaped_inputs)
    4: h = einsum("EGCM, EMH -> EGCH", dispatched_expert_inputs, wi)
    5: h = relu(h)
    6: expert_outputs = einsum("EGCH, EHM -> GECM", h, wo)
    7: outputs = einsum("GSEC, GECM -> GSM", combine_weights, expert_outputs)
    • 在分析算法2的计算复杂度随设备数量D的扩展情况时,论文做出以下假设:
      • a)每个设备上的 Token 数量 \(\frac{N}{D}=O(1)\) (在实际应用中,为避免设备内存溢出,这一假设通常是必要的),且保持恒定;
      • b) \(G=O(D)\) 、 \(S=O(1)\) ;
      • c) \(H=O(1)\) ;
      • d) \(E=O(D)\) ;
      • e) \(C=O(\frac{2S}{E})=O(\frac{1}{D})\) ,且 \(D< S\) ,D 为正整数
  • 算法2中的浮点运算总数计算如下:
    $$
    \begin{align}
    \text{FLOPS}_{\text{Softmax} } + \text{FLOPS}_{\text{Top2Gating} } + \text{FLOPS}_{\text{Dispatch/Combine} } + \text{FLOPS}_{\text{FFN} } &= \\
    O(G S M E) + O(G S E) + O(G S M E C) + O(G E C H) &= \\
    O(D \cdot 1 \cdot 1 \cdot D) + O(D \cdot 1 \cdot D) + O(D \cdot 1 \cdot 1 \cdot D \cdot \frac{1}{D}) + O(D \cdot D \cdot \frac{1}{D} \cdot 1) &= \\
    O(D^2) + O(D^2) + O(D) + O(D)
    \end{align
    }
    $$

    • 因此,每个设备的浮点运算量为 \(O(D)\)
    • 每个设备上 softmax 的计算复杂度 \(\text{FLOPS}_{\text{Softmax} } / D = O(D)\) 与设备数量呈线性关系,但实际上,由于 \(D< S\) ,该复杂度会被其他项主导,因此可将整体复杂度视为 \(O(1)\) ,满足亚线性扩展的设计要求
    • 5节通过实验验证了这一分析结果
  • 除计算成本外,跨设备通信成本并非恒定,但如5节所述,当设备数量增加时,通信成本仅以 \(O(\sqrt{D})\) 的温和速率增长

GShard Annotation API for Parallel Execution

  • 由于算法1中张量的规模和计算需求极大,论文必须在多个设备上对该算法进行并行化处理
  • 算法2中带下划线的字母展示了对每个张量进行分片的一种直接方案
  • GShard 中的分片接口允许论文对程序中的张量进行标注,选择性地指定其分片方式
  • 这些信息会传递给编译器,以便编译器自动应用转换以实现并行执行
  • 在论文的研究中,使用了 TensorFlow/Lingvo(2019)中的以下接口
    • replicate(tensor):对张量进行标注,使其在各个分片(partition)中复制,并返回标注后的张量。该接口通常用于模型中的非混合专家层(non-MoE layers),以实现权重复制
    • split(tensor, split_dimension, num_partitions):对张量进行标注,使其沿split_dimension维度进行分片,并返回标注后的张量。第i个分片会被放置在第i个设备上,且num_partitions(分片数量)不得超过系统中的设备数量
    • shard(tensor, device_assignment):是split()接口的泛化形式,支持对多个维度进行分片,并指定每个分片的放置位置。附录A.3对该接口进行了更详细的描述
  • 需要注意的是,调用split或shard接口仅会添加标注,不会改变用户程序中张量的逻辑形状。用户仍可使用完整形状的张量进行操作,无需担心非均匀分片等问题
  • GShard 具有良好的通用性,其简单接口可同样应用于所有维度
  • 根据具体应用场景,分片维度可以包括批次维度(数据并行)、特征维度、专家维度,甚至图像模型中的空间维度
  • 此外,由于分片标注是基于每个张量单独进行的,模型的不同部分可以采用不同的分片方式
  • 这种灵活性使论文能够对巨型混合专家层权重进行分片,并在混合专家层和非混合专家层之间切换分片模式,同时也支持论文未涉及的其他应用场景,例如对大型图像进行空间分片(2019)(附录A.4)
  • 通过上述分片接口,可以将算法2中所示的分片策略表示如下:
    • 输入张量沿第一个维度(组维度G)进行分片,门控权重张量(wg)采用复制方式存储
    • 计算得到分配后的专家输入(dispatched expert inputs)后,应用split接口将分片维度从组维度(G)切换为专家维度(E)。其中,D为设备数量
      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      1: # 沿组维度(G)对输入进行分片
      2: inputs = split(inputs, 0, D)
      3: # 复制门控权重
      4: wg = replicate(wg)
      5: gates = softmax(einsum("GSM, ME -> GSE", inputs, wg))
      6: combine_weights, dispatch_mask = Top2Gating(gating_logits)
      7: dispatched_expert_inputs = einsum("GSEC, GSM -> EGCM", dispatch_mask, reshaped_inputs)
      8: # 沿专家维度(E)对分配后的输入进行分片
      9: dispatched_expert_inputs = split(dispatched_expert_inputs, 0, D)
      10: h = einsum("EGCM, EMH -> EGCH", dispatched_expert_inputs, wi)
Per-tensor sharding assignment
  • 如上述示例所示,用户无需对程序中的每个张量都进行标注
  • 通常只需对模型中的少数重要算子(如爱因斯坦求和算子)进行标注,编译器会通过自身的启发式算法推断出其余张量的分片方式
    • 注:由于反向传播计算通常由前端框架自动生成,用户无法访问这些张量,因此编译器推断缺失的分片信息也非常重要
    • 例如,由于输入张量沿组维度(G)进行分片,而权重张量采用复制方式存储,编译器会选择沿相同维度(组维度G)对爱因斯坦求和运算的输出进行分片(第5行)
    • 类似地,由于输入分配的爱因斯坦求和运算(第7行)的两个输入均沿组维度(G)进行分片,因此输出的分片方式会被推断为沿组维度(G)进行分片,之后论文通过添加split标注将输出的分片维度切换为专家维度(E)
    • 上述示例中的某些标注(如replicate(wg))也可由编译器自动确定,但建议对计算的初始输入张量和最终输出张量进行标注
  • 目前,编译器采用迭代数据流分析(iterative data-flow analysis),从用户标注的算子开始,将分片信息传播到其相邻算子(操作数和使用者)
    • 该分析通过对齐相邻算子的分片决策,尽量减少重分片(resharding)的需求
    • 虽然还可以采用整数规划或机器学习等其他方法,但改进自动分片分配并非论文的重点,论文将其留作未来的研究工作
Mixing manual and automatic sharding(手动分片与自动分片结合)
  • 在常见情况下,通过分片标注实现自动分片通常已足够,但GShard也支持灵活地将手动分片算子与自动分片算子结合使用
  • 这让用户能够更好地控制算子的分片方式,例如当用户掌握算子语义之外的运行时知识时
  • 例如,XLA 和 TensorFlow 的 Gather 算子定义均未包含输入中不同范围的索引边界信息,但用户可能知道某个特定的 Gather 算子仅在每个分片内对数据进行重排
  • 在这种情况下,用户只需缩小维度大小并执行本地 Gather 操作,即可轻松对该算子进行分片;
  • 否则,编译器需要对索引范围采取保守处理,并添加不必要的通信开销
  • 例如,算法2中使用独热矩阵(one-hot matrix)分配输入的爱因斯坦求和算子(第3行),也可以通过手动分片的 Gather 算子来实现,而模型的其余部分仍采用自动分片方式
  • 以下伪代码展示了该应用场景
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    1: # 输入形状为[G, S, M],split()不改变逻辑形状
    2: input = split(input, 0, num_devices)
    3: # s_indices形状为[E, G, C, 1],值为输入中S维度的索引
    4: s_indices = split(s_indices, 1, num_devices)
    5:
    6: # 开始手动分片
    7: # partitioned_input形状为[G/num_devices, S, M]
    8: partitioned_input = auto_to_manual_spmd_partition(input)
    9: # partitioned_s_indices形状为[E, G/num_devices, C, 1]
    10: partitioned_s_indices = auto_to_manual_spmd_partition(s_indices)
    11: # 在partitioned_input中拼接G维度的索引:在G维度上生成Iota张量
    12: partitioned_gs_indices = concat(
    13: iota([E, G/num_devices, C, 1], 1),
    14: partitioned_s_indices,
    15: 3
    16: )
    17: # partitioned_data形状为[E, G/num_devices, C, M]
    18: partitioned_data = gather(partitioned_input, partitioned_gs_indices)
    19:
    20: # 切换回自动分片
    21: # data形状为[E, G, C, M]
    22: data = manual_to_auto_spmd_partition(partitioned_data)
    23: ...

The XLA SPMD Partitioner for GShard

  • 本节将介绍基于分片标注自动对计算图进行分片的编译器基础设施
  • 分片标注告知编译器每个张量应如何在设备间分布
  • SPMD(单程序多数据,Single Program Multiple Data)分片器(简称“分片器”)是编译器的一个组件,它将计算图转换为可在所有设备上并行执行的单一程序
    • 注:另一种方案是MPMD(多程序多数据,Multiple Program Multiple Data),如图2所示,该方案的扩展性较差
  • 这使得编译时间几乎与分片数量无关,从而支持论文将模型扩展到数千个分片⁴
  • 论文在 XLA 编译器(2019)中实现了该分片器
  • TensorFlow、JAX、PyTorch 和 Julia 等多个前端框架已具备将其图表示转换为 XLA HLO 图的降级逻辑(lowering logic)
  • 与 TensorFlow 等流行前端框架相比,XLA 的算子数量要少得多,这在不影响通用性的前提下降低了实现分片器的负担,因为现有前端的降级过程已实现了强大的表达能力
    • 尽管论文在 XLA 中开发了该基础设施,但论文介绍的技术也可应用于其他机器学习框架的中间表示(如 ONNX(2019)、TVM Relay(2018)、Glow IR(2018))
  • XLA 将计算表示为数据流图(dataflow graph),其中节点表示算子,边表示在算子之间流动的张量
  • 分片器的核心是对每个算子进行处理,根据输入和输出指定的分片方式,将全尺寸算子转换为分片尺寸算子
  • 对计算进行分片时,会引入各种跨设备数据传输模式
  • 为了在大规模场景下最大化性能,定义一组核心通信原语(communication primitive)并针对目标平台对其进行优化至关重要
Communication Primitives,通信原语
  • 由于分片器强制所有设备运行相同的程序,通信模式也具有规律性,XLA 定义了一组集合算子(collective operator)来执行类似 MPI(2009)的通信操作
  • 下面列出了论文在 SPMD 分片器中使用的常见通信原语
    • CollectivePermute :该算子指定一组源-目标设备对(source-destination pairs),源设备的输入数据会发送到对应的目标设备
      • 它主要用于两个场景:改变分片张量在各分片间的设备顺序,以及本节后续讨论的 Halo Exchange
    • AllGather :该算子按照指定顺序将所有参与设备的张量连接起来,用于将分片张量转换为复制张量
    • AllReduce :该算子对所有参与设备的输入执行按元素归约(如求和)操作,用于合并来自不同分片的部分归约中间张量
      • 在 TPU 设备网络中,当分片数量增加时,AllReduce 的成本保持恒定(5.2节)
      • 它也是一种在其他类型网络拓扑中具有高效实现的常用原语(2019)
    • AllToAll :该算子从逻辑上沿某个维度对每个参与设备的输入进行分割,然后将每个数据片段发送到不同的参与设备
      • 每个设备在接收来自其他设备的数据片段后,会将这些片段连接起来生成结果
      • AllToAll 用于将分片张量从一个维度重分片到另一个维度
      • 在 TPU 设备网络中,AllToAll 是实现此类重分片的高效方式,其成本随分片数量的增加呈亚线性增长(5.2节)
Per-Operator SPMD Partitioning
  • 分片器的核心是根据指定的分片方式,对每个算子进行从全尺寸到分片尺寸的转换
  • 虽然某些算子(如按元素算子)的分片支持非常简单,但论文将讨论需要跨分片通信的几种常见情况
  • 一般情况下,存在一些重要的技术挑战,论文将在3.3.3节中介绍
  • 为了使讨论与混合专家模型更相关,本节将重点关注爱因斯坦求和算子(Einsum)的分片,以展示几种通信模式
  • 为简化讨论,此处假设所有张量均采用均匀分片方式,即待分片维度的大小是分片数量的整数倍
Einsum Case Study
  • 爱因斯坦求和算子是实现混合专家模型的最关键算子
  • 在 XLA HLO 中,它们被表示为点积(Dot)运算,其中每个操作数(左操作数 LHS 或右操作数 RHS)包含三种类型的维度:
    • 批次维度(Batch dimensions) :是易并行化维度(embarrassingly parallel dimensions)。所有左操作数、右操作数和输出都必须包含相同的批次维度集合,且输出中的每个元素仅依赖于左操作数和右操作数中对应的批次元素
    • 收缩维度(Contracting dimensions) :仅存在于操作数中。左操作数和右操作数必须包含相同的收缩维度集合,这些维度在输出中会被求和并压缩
    • 非收缩维度(Non-contracting dimensions) :也是并行维度,存在于某个操作数和输出中。左操作数和右操作数分别具有自己的非收缩维度集合,这些维度会被输出继承
  • 分片传播(Sharding propagation)优先选择在左操作数、右操作数和输出的批次维度上采用相同的分片方式,因为这样可以避免任何跨分片通信
  • 但在实际情况中,这并不总是可行的,以下三种情况需要跨分片通信:
    • 1)重分片(Resharding) :在论文构建的混合专家模型中,专家分配逻辑(算法2第3行)需要在爱因斯坦求和运算后切换分片维度。由于使用AllToAll可以高效实现重分片(5.2节),论文首先在本地执行爱因斯坦求和运算,然后将结果重分片到目标维度,如图4a所示
    • 2)累积部分结果(Accumulating partial results) :如果输入沿收缩维度进行分片,本地计算得到的结果只是部分结果,需要使用AllReduce将这些部分结果合并以生成最终结果,如图4b所示
    • 3)循环切片(Slicing in a loop) :在某些场景下,论文还实现了一种类似于Cannon算法(1969)的方法,以限制每个分片上张量的大小。例如,如果两个操作数均沿非收缩维度进行分片,由于操作数的非收缩维度不同,无法直接在本地执行爱因斯坦求和运算。复制其中一个操作数不会导致冗余计算,但需要确保被复制的操作数能够放入设备内存
      • 因此,如果操作数的规模过大,论文会保持两个操作数的分片状态,通过循环迭代计算结果的每个切片,并使用 CollectivePermute 来传输输入切片(图4c)
  • 图4:带有跨设备通信的爱因斯坦求和算子分片示例:
    • (a)分片爱因斯坦求和算子
      • 彩色字母(G和E)表示每个张量的分片维度
      • 分片器决定首先沿G维度执行批次并行爱因斯坦求和运算,然后将结果重分片到E维度。(注:图中省略了S和M维度)
    • (b)在收缩维度上分片的简单爱因斯坦求和算子(矩阵乘法)
      • 每个分片计算部分结果,然后通过AllReduce合并得到完整结果
    • (c)在循环中使用CollectivePermute的爱因斯坦求和算子(矩阵乘法)
      • 通过循环每次计算一个切片,整个过程中不会出现全尺寸张量
Supporting a Complete Set of Operators
  • 为了使 SPMD 分片器能够支持完整的算子集合,且不对张量形状或算子配置施加额外限制,论文解决了若干额外挑战
  • 这些挑战通常涉及分片间的非对称计算或通信模式,而由于单一程序需要对所有分片都具有通用性,这些模式在 SPMD 中尤其难以表达
  • 论文不能简单地根据运行时设备ID在单一程序中创建多个分支,因为这会导致程序规模激增

Static shapes and uneven partitioning

  • XLA要求张量形状是静态的(注:中间表示中有限的动态性对于高效适配加速器通常是必要的)
  • 然而,对计算进行分片时,由于维度大小可能无法被分片数量整除,并非所有分片的输入/输出形状都相同
  • 在这种情况下,会将形状大小向上取整到分片数量的下一个整数倍,填充区域(padded region)中的数据可以是任意值
  • 在计算算子时,为了保证正确性,可能需要在填充区域中填充已知值
    • 例如,对 Reduce-Add 算子进行分片时,需要使用单位元0作为填充值
    • 考虑一个示例:待分片维度大小为 15,无法被 2(分片数量)整除,因此分片1比所需大小多一列
    • 论文生成一个范围为 [0, 8) 的 Iota 算子,加上分片偏移量(由分片 ID×8 计算得到),并与全尺寸偏移量(15)进行比较。根据比较得到的谓词值(predicate value),选择从操作数中取值或从 0 中取值,最终得到掩码操作数(masked operand)

Static operator configurations

  • XLA算子具有静态配置,例如卷积(Convolution)中定义的填充(padding)、步幅(stride)和扩张(dilation)
  • 但不同分片可能不会使用相同的算子配置
    • 例如,对于卷积算子,最左侧的分片会在其左侧应用填充,而最右侧的分片会在其右侧应用填充
    • 在这种情况下,分片器可能会选择让某些分片生成略多于所需的数据,然后切分出无关部分
  • 附录A.4讨论了卷积和类似算子的示例

Halo Exchange

  • 某些算子具有一种通信模式,需要与相邻分片交换部分数据,论文称之为 Halo Exchange
  • 论文使用 CollectivePermute 算子在分片间交换 Halo 数据
  • Halo Exchange 最典型的应用场景是对基于窗口的算子(如卷积、ReduceWindow)进行分片,因为相邻分片可能需要重叠的输入数据(图5a)
  • 在实际应用中,由于窗口配置(扩张、步幅和填充)的复杂使用以及非均匀 Halo 大小,这些算子的 Halo Exchange 通常需要结合适当的填充、切片和掩码操作
    • 论文在附录A.4中描述了各种场景
  • Halo Exchange 的另一个应用场景是改变形状大小的数据格式化算子(data formatting operator)
    • 例如,经过 Slice 或 Pad 算子后,张量的形状会发生变化,分片间的边界也会随之改变
    • 这需要论文重新对齐不同分片上的数据,而这可以通过 Halo Exchange 的形式来处理(图5b)
  • 其他数据格式化算子(尽管在逻辑上不改变形状大小)也可能需要 Halo Exchange ,这主要是由于静态形状约束和非均匀分片
    • 例如,Reverse 算子会反转张量中元素的顺序,但如果对其进行非均匀分片,需要在分片间移动数据,以确保填充在逻辑上位于结果张量的右侧
    • 另一个示例是Reshape算子:
      • 考虑将形状为 [3, 2] 的张量重塑为 [6],其中输入在第一个维度上采用2种非均匀分片方式(分片形状为 [2, 2]),输出也采用2种分片方式(分片形状为[3])
      • 由于非均匀分片,输入中存在填充,但重塑后输出张量不再有填充;因此,需要采用与Slice类似的 Halo Exchange 方式(图5c)

Compiler optimizations

  • SPMD 分片器会创建各种数据格式化算子,以执行切片、填充、连接、掩码和 Halo Exchange 操作
  • 为解决这一问题,论文利用了 XLA 在 TPU 上的融合能力(fusion capabilities),以及针对切片和填充的代码移动优化(code motion optimizations),在很大程度上隐藏了数据格式化的开销
  • 因此,即使在大量使用掩码和填充的卷积网络中,运行时开销通常也可以忽略不计

Massively Multilingual, Massive Machine Translation, M4

  • 待补充

Performance and Memory Consumption

  • 待补充

NLP——Megatron-LM原始论文解读-第一篇

注:本文包含 AI 辅助创作

  • 参考链接:
    • 官网主页:Megatron Core User Guide
    • GitHub开源地址:github.com/NVIDIA/Megatron-LM
    • Megatron 系列目前公认的三篇核心论文如下,它们分别对应了张量并行、3D 并行 与 序列并行/激活重计算优化 三大阶段:
    • 第一篇:Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, arXiv 2019, NVIDIA
      • 核心贡献:提出 张量并行(Tensor Parallelism) ,将 Transformer 的 Attention 头与 FFN 权重按列/按行切分,实现层内模型并行,首次在 GPU 集群上训练出 8.3 B 参数的 GPT-2 模型
    • 第二篇:Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021, NVIDIA
      • 核心贡献:提出 3D 并行(数据 + 张量 + 流水线),并给出 interleaved 1F1B 流水线调度,显著降低流水线气泡;在 3072 块 A100 上训练出 530 B 参数的 GPT-3 级模型,GPU 利用率达到 76 %
    • 第三篇:Reducing Activation Recomputation in Large Transformer Models, MLSys 2023, NVIDIA
      • 核心贡献:提出 序列并行(Sequence Parallelism) 与 选择性激活重计算 ,将激活显存占用再降 3–5 倍,支持更长序列与更大批训练;与 Flash-Attention 思想互补,现已成为 Megatron-Core 默认配置
    • 以上三篇即社区常说的“Megatron 三篇论文”,覆盖了从单层切分、多层多维并行到显存优化的完整技术路线,本文介绍第一篇论文

Paper Summary

  • 整体总结:
    • 本文是 Megatron-LM原始论文解读-第一篇
    • 论文的核心贡献:提出 张量并行(Tensor Parallelism) ,将 Transformer 的 Attention 头与 FFN 权重按列/按行切分,实现层内模型并行,首次在 GPU 集群上训练出 8.3 B 参数的 GPT-2 模型
    • 论文通过对现有 PyTorch Transformer 实现仅进行少量修改来实现模型并行,成功突破了传统的单 GPU 单模型训练的限制
    • 论文在 512 块 NVIDIA V100 GPU 上使用 8 路模型并行高效训练了高达 8.3B 参数的基于 Transformer 的模型,并在整个应用中实现了高达 15.1 PetaFLOPs 的持续性能
    • 论文表明,对于 BERT 模型,随着模型规模的增长,类 BERT 模型中 Layer Normalization 的位置对准确率的提升至关重要
    • 论文研究了模型规模对下游任务准确率的影响,在下游任务上取得了远超现有水平的结果,并在 WikiText103、LAMBADA 和 RACE 数据集上建立了新的最佳结果(SOTA)
    • 非常值得称赞的是:论文开源了论文的代码,以支持未来基于模型并行 Transformer 的研究工作
  • 背景 & 问题提出:
    • Transformer 模型推动了 NLP 技术发展,但内存限制,导致超大型模型的训练难度极大
  • 方法内容:
    • 论文提出了训练超大型 Transformer 模型的技术,并实现了一种简单、高效的层内模型并行(intra-layer model parallel)方法,该方法能够训练具有数十亿参数的 Transformer 模型
    • 本方法不需要新的编译器或库修改,与流水线模型并行(pipeline model parallelism)相互独立且互补,完全可以通过在原生 PyTorch 中插入少量通信操作来实现
  • 论文通过使用 512 个 GPU 训练出高达 8.3B 参数的 Transformer 模型验证了该方法的有效性
    • 与单个GPU基准(可持续 39 TeraFLOPs,达到峰值浮点运算的 30%)相比,论文在整个应用中实现了 15.1 PetaFLOPs 的性能,扩展效率为76%
    • \({39 \times 512 / 1000}{39} \approx 76%\)
  • 为了证明 LLM 能进一步推动技术发展,论文训练了一个与 GPT-2 类似的 8.3B 参数 Transformer 语言模型和一个与 BERT 类似的 3.9B 参数模型
    • 论文发现,对于类 BERT 模型,随着模型规模的增长, Layer Normalization的位置对性能提升至关重要
    • 使用 GPT-2 模型,论文在 WikiText103 数据集和 LAMBADA 数据集上取得了最佳结果
      • WikiText103 数据集:困惑度为 10.8(之前最佳困惑度为 15.8)
      • LAMBADA 数据集:准确率为 66.5%(之前最佳准确率为63.2%)
    • 使用 BERT 模型,论文在在 RACE 数据集上取得了最佳结果
      • RACE 数据集:准确率为 90.9%(之前最佳准确率为 89.4%)

Introduction and Discussion

  • NLP 正在快速发展,部分原因是可用计算资源和数据集规模的增加
    • 丰富的计算资源和数据使得通过无监督预训练(unsupervised pretraining)训练越来越大的语言模型成为可能(2018;2019)
  • 实证证据表明,更大的语言模型在文章补全、问答和自然语言推理等 NLP 任务中更有用(2019;2019)
    • 近期研究显示,在下游自然语言任务上对这些预训练语言模型进行微调(finetuning),可以取得最佳结果(2018;2017;2016;2019b;)
  • 模型规模的增长超出了现代处理器的内存限制,需要额外的内存管理技术,如激活检查点(activation checkpointing)(2016)
    • ADAM 等广泛使用的优化算法(optimization algorithms)需要为每个参数额外分配内存来存储动量(momentum)和其他优化器状态,这限制了可有效训练的模型规模
    • 几种模型并行(model parallelism)方法通过对模型进行分区来克服这一限制,使得权重及其相关的优化器状态无需同时驻留在处理器上,比如:
      • GPipe(2018)和 Mesh-Tensorflow(2018)提供了不同类型的模型并行框架
      • 但它们需要重写模型,并依赖仍在开发中的自定义编译器和框架
  • 论文使用层内模型并行(intra-layer model-parallelism)实现了一种简单高效的模型并行方法
  • 论文利用基于 Transformer 的语言模型的固有结构,实现了一种简单的模型并行方案,该方案能在 PyTorch 中高效训练,无需自定义 C++ 代码或编译器
    • 这种方法与 GPipe(2018)等方法所倡导的基于流水线的模型并行相互独立
  • 为了证明论文方法的可扩展性,论文建立了一个基准:在单个 NVIDIA V100 32GB GPU 上训练 1.2B 参数的模型,可持续 39 TeraFLOPs
    • 这相当于 DGX-2H 服务器中配置的单个 GPU 理论峰值浮点运算的 30%,因此是一个很强的基准
    • 通过 8路 模型并行在 512 个 GPU 上训练 8.3B 参数的模型,论文在整个应用中实现了高达 15.1 PetaFLOPs 的持续性能
      • 与单GPU情况相比,这一扩展效率为76%
    • 图1显示了更详细的扩展结果
  • 为了分析模型规模扩展对准确率的影响:
    • 论文训练了从左到右的 GPT-2(2019)语言模型和 BERT(2018)双向 Transformer,并在多个下游任务上对它们进行了评估
      • 论文发现,现有 BERT 架构会随着规模的增加而导致模型性能下降
      • 论文通过重新排列 Transformer 层中的 Layer Normalization 和残差连接(residual connection)克服了这一挑战,并表明通过这种修改,下游任务在开发集上的结果会随着模型规模的增加而单调提升
    • 论文的模型在 WikiText103 测试集、LAMBADA 的完形填空预测准确率(cloze-style prediction accuracy)和 阅读理解 RACE 数据集上取得了最佳结果
  • 总结来看:论文的贡献如下:
    • 通过对现有 PyTorch Transformer 实现仅进行少量有针对性的修改,实现了一种简单高效的模型并行方法
    • 对模型和数据并行技术进行了深入的实证分析,并证明了在 512 个 GPU 上的扩展效率高达 76%
    • 论文实验表名:对于类 BERT 模型,随着模型规模的增长, Layer Normalization 的位置对准确率的提升至关重要
    • 论文证明:模型规模的扩展会提高 GPT-2(研究到 8.3B 参数)和 BERT(研究到 3.9B 参数)模型的准确率
    • 展示了论文的模型在测试集上取得了最佳结果:WikiText103的困惑度(10.8 ppl)、LAMBADA 的准确率(66.5%)和 RACE 的准确率(90.9%)
    • 开源了论文的代码以及训练和评估流水线,地址为 github.com/NVIDIA/Megatron-LM

Background and Challenges

Neural Language Model Pretraining

  • 预训练语言模型已成为 NLP 研究人员工具包中不可或缺的一部分
    • 利用大型语料库预训练来学习稳健的语言神经表征是过去十年中一个活跃的研究领域
  • 预训练和迁移语言神经表征的早期例子表明,预训练的词嵌入表(word embedding tables)比从头学习的词嵌入表能更好地提升下游任务结果(2013;2014;2010)
  • 后来的研究通过学习和迁移能捕捉单词上下文表征(contextual representations)的神经模型推进了这一领域(2016;2018;2017;2019)
  • 近期的相关研究(2016;2018;2019b;)进一步基于这些思想,不仅迁移语言模型来提取上下文单词表征,还在下游任务上对语言模型进行端到端微调
  • 通过这些工作,技术发展已经从仅迁移词嵌入表发展到迁移整个数十亿参数的语言模型
  • 这种方法的发展使得需要能够高效大规模运行并满足日益增长的计算需求的硬件、系统技术和框架
  • 论文的工作旨在提供必要的工具,以推动这一趋势向前迈出又一步

Transformer Language Models and Multi-Head Attention

  • 当前 NLP 研究趋势是使用 Transformer 模型(2017),因为它们具有更高的准确率和计算效率
  • 原始 Transformer 结构被设计为机器翻译架构,它使用 Encoder 和 Decoder 两部分将输入序列转换为另一个输出序列
  • 近期利用 Transformer 进行语言建模的研究(如BERT(2018)和GPT-2(2019))根据其需求仅使用 Encoder 或 Decoder
  • 本研究探讨了 Decoder 架构 GPT-2 和 Encoder 架构 BERT
  • 图2显示了论文使用的模型示意图
    • 有关模型架构的详细描述,论文建议读者参考先前的研究(2017;2018;2019)
    • 特别地,GPT-2 和 BERT都使用 GeLU(2016)非线性激活函数和 Layer Normalization (2016),并 将其应用于多头注意力(multi-head attention)和前馈层(feed forward layers)的输入,而原始 Transformer (2017)使用 ReLU 非线性激活函数,并将 Layer Normalization 应用于输出

Data and Model Parallelism in Deep Learning

  • 将深度神经网络训练扩展到多个硬件加速器有两种核心范式:数据并行(data parallelism)(1990)和模型并行(model parallelism)
    • 数据并行是将训练小批量(minibatch)分配到多个工作节点(workers)
    • 模型并行是将模型的内存使用和计算分配到多个工作节点
  • 通过与可用工作节点数量成比例地增加批量大小(即弱扩展),可以观察到训练数据吞吐量接近线性扩展
    • 但大批量训练会给优化过程带来复杂性,可能导致准确率下降或收敛时间延长,从而抵消训练吞吐量增加带来的好处(2017)
    • 进一步的研究(2017;2017;2019)开发了缓解这些影响的技术,并缩短了大型神经网络的训练时间
    • 为了进一步扩展训练,相关研究(2016)将数据并行与激活检查点相结合 :
      • 在反向传播(backward pass)中重新计算激活(activations),而不在前向传播(forward pass)中存储它们,以减少内存需求
  • 以上这些技术在处理问题规模上有一个根本限制:模型必须完全适合单个工作节点
    • 随着 BERT 和 GPT-2 等语言模型规模和复杂性的增加,神经网络已经接近现代硬件加速器的内存容量
    • 解决这一问题的一种方法是使用参数共享(parameter sharing)来减少模型的内存占用(2019),但这会限制模型的整体容量
    • 论文的方法是利用模型并行将模型分配到多个加速器上
      • 这不仅减轻了内存压力,还独立于微批量(microbatch)大小增加了并行性
  • 在模型并行中,还有两种进一步的范式:分层流水线并行(layer-wise pipeline parallelism)和更通用的分布式张量计算(distributed tensor computation)
    • 在流水线模型并行中,一组操作在一个设备上执行,然后将输出传递到流水线中的下一个设备,在那里执行另一组操作
      • 一些方法(2018;2018)将参数服务器(parameter server)(2014)与流水线并行结合使用,但这些方法存在不一致问题
      • TensorFlow 的 GPipe 框架(2018)通过使用同步梯度下降(synchronous gradient decent)克服了这一不一致问题
        • GPipe 需要额外的逻辑来处理这些通信和计算操作的高效流水线,并且会受到流水线气泡(pipeline bubbles)的影响而降低效率,或者需要修改优化器本身,从而影响准确率
    • 分布式张量计算是一种独立且更通用的方法,它将张量操作分配到多个设备上,以加速计算或增加模型规模
      • FlexFlow(2018)是一个协调这种并行计算的深度学习框架,它提供了一种选择最佳并行策略的方法
      • Mesh-TensorFlow(2018)引入了一种在 TensorFlow(2015)中指定通用分布式张量计算类别的语言
        • 用户在该语言中指定并行维度,然后使用适当的集合原语(collective primitives)编译生成的图
  • 论文利用与 Mesh-TensorFlow 中类似的见解,并利用 Transformer 注意力头(attention heads)计算中的并行性来并行化论文的 Transformer 模型
    • 论文没有为模型并行实现框架和编译器,而是仅对现有 PyTorch Transformer 实现进行了少量有针对性的修改
    • 论文的方法简单,不需要任何新的编译器或代码重写,并且可以通过插入一些简单的原语来完全实现,如下一节所述

Model Parallel Transformer s

  • 论文利用 Transformer 网络的结构,通过添加一些同步原语(synchronization primitives),实现了一种简单的模型并行方案

  • 如图2所示:一个 Transformer 层由一个自注意力块(self attention block)和一个两层的多层感知器(multi-layer perceptron, MLP)组成

    • 论文分别在这两个块中引入模型并行
  • 首先详细介绍 MLP 块

    • MLP 块的第一部分是一个通用矩阵乘法(general matrix multiplication, GEMM),后跟一个 GeLU 非线性激活函数:
      $$Y=GeLU(X A) \tag{1}$$
    • 并行化 GEMM 的一种方法是将权重矩阵 \(A\) 按行拆分,将输入 \(X\) 按列拆分,如下所示:
      $$X=\left[X_{1}, X_{2}\right], A=\left[\begin{array}{l}A_{1} \\ A_{2}\end{array}\right] \tag{2}$$
    • 这种划分将导致
      $$ Y=GeLU(X_{1} A_{1}+X_{2} A_{2}) $$
    • 由于 GeLU 是一个非线性函数
      $$ GeLU(X_{1} A_{1}+X_{2} A_{2}) \neq GeLU(X_{1} A_{1})+GeLU(X_{2} A_{2})$$
    • 因此这种方法需要在 GeLU 函数之前设置一个同步点
  • 另一种方法是将 \(A\) 按列拆分 \(A=[A_{1}, A_{2}]\)

    • 这种划分允许 GeLU 非线性激活函数独立应用于每个划分的 GEMM 输出:
      $$\left[Y_{1}, Y_{2}\right]=\left[GeLU\left(X A_{1}\right), GeLU\left(X A_{2}\right)\right]$$

    • 这一方法的优势在于消除了同步点

    • 论文以这种列并行的方式划分第一个 GEMM,并将第二个 GEMM 按行拆分,使其可以直接接收 GeLU 层的输出,而无需任何通信(如图3a所示)

    • 然后,第二个 GEMM 的输出在多个 GPU 之间进行归约(reduce),之后再传递到 dropout 层

    • 这种方法将 MLP 块中的两个 GEMM 都拆分到多个 GPU 上,并且在正向传播中只需要一个 \(g\) 算子,在反向传播中只需要一个全归约(all-reduce)( \(f\) 算子)

      • 这两个算子是共轭的,可以用 PyTorch 中的几行代码实现
    • 例如, \(f\) 算子的实现如下:

      1
      2
      3
      4
      5
      6
      class f(torch.autograd.Function):
      def forward(ctx, x):
      return x
      def backward(ctx, gradient):
      all_reduce(gradient)
      return gradient
    • \(g\) 与 \(f\) 类似,只是在正向传播中执行全归约,在反向传播中执行恒等操作

  • 如图3b所示

    • 对于自注意力块,利用多头注意力操作中固有的并行性 ,将与 \(K\)、 \(Q\) 和 \(V\) 相关的 GEMM 以列并行的方式划分 ,这样每个注意力头对应的矩阵乘法都在单个 GPU 上本地完成
    • 这使得我们可以将每个注意力头的参数和计算负载分配到多个 GPU 上,并且不需要立即通信就能完成自注意力计算
    • 自注意力之后的输出线性层(output linear layer)的 GEMM 按行并行化,可以直接接收并行注意力层的输出,无需 GPU 之间的通信
    • 这种针对 MLP 层和自注意力层的方法融合了两组 GEMM,消除了它们之间的同步点,从而实现了更好的扩展性
    • 这使得论文能够在一个简单的 Transformer 层中,仅在正向传播和反向传播中各使用两次全归约来执行所有 GEMM(见图4)
  • Transformer 语言模型的 Output Embedding 维度为隐藏层大小( \(H\) )乘以词汇表大小( \(v\) )

    • 现代语言模型的词汇表大小通常在数万个 token 量级(例如,GPT-2 使用的词汇表大小为 50257),因此并行化 Output Embedding 的 GEMM 是有益的
    • 但在 Transformer 语言模型中, Output Embedding 层与 Input Embedding 层共享权重,这需要对两者都进行修改
    • 对于 Input Embedding 权重矩阵 \(E_{H \times v}\) 沿词汇表维度按列拆分:
      $$ E=[E_{1}, E_{2}]$$
      • 由于每个分区现在只包含一部分嵌入,因此在 Input Embedding 之后需要一个 \(g\) 算子
    • 对于 Output Embedding :
      • 一种方法是执行下面的并行 GEMM 以获得 logits
        $$ [Y_{1}, Y_{2}]=[X E_{1}, X E_{2}] $$
        • 然后添加一个全收集(all-gather)操作 $$ Y = \text{all-gather}([Y_{1}, Y_{2}])$$
        • 再将结果传递给交叉熵损失函数(cross-entropy loss function)
        • 但这种情况下,全收集操作将传递 \(b \times s \times v\) 个元素( \(b\) 是批大小, \(s\) 是序列长度),由于词汇表大小很大,这一数据量非常庞大
      • 为了减少通信量,论文将并行 GEMM 的输出 \([Y_{1}, Y_{2}]\) 与交叉熵损失融合,这将维度减少到 \(b \times s\)
      • 传递标量损失而不是 logits 大大减少了通信量,提高了论文模型并行方法的效率
  • 论文的模型并行方法在很大程度上可以概括为旨在减少通信并使 GPU 保持计算受限(compute bound)的技术

  • 论文选择在多个 GPU 上复制计算,而不是让一个 GPU 计算 dropout、 Layer Normalization 或残差连接的一部分,然后将结果广播到其他 GPU

  • 具体来说:

    • 论文在每个 GPU 上维护 Layer Normalization 参数的副本
    • 并对模型并行区域的输出执行 dropout 和 残差连接
    • 然后将其作为输入馈送到下一个模型并行区域
    • 问题:如何理解这里的每个 GPU 独立维持 Layer Normalization 参数的副本?
  • 为了优化模型,论文允许每个模型并行工作节点优化自己的一组参数

    • 由于所有值要么在 GPU 上本地存储,要么在 GPU 上复制,因此在这种形式中不需要通信更新后的参数值
  • 关于混合模型并行与数据并行以及随机数生成的处理,论文在附录B中提供了更多细节

  • 总之,如上所述,论文的方法易于实现,只需要在正向传播和反向传播中添加少量额外的全归约操作

    • 它不需要编译器,并且与 GPipe(2018)等方法所倡导的流水线模型并行相互独立且互补

Setup

  • 预训练语言理解模型是自然语言处理和语言理解中的核心任务
  • 语言建模有多种形式,论文重点关注 GPT-2(2019)(一种基于 Transformer 的从左到右生成式语言模型)和 BERT(2018)(一种基于掩码语言模型的双向 Transformer 模型)
  • 论文在以下部分解释这些模型的配置,并建议读者参考原始论文以获取更多细节

Training Dataset

  • 为了收集具有长期依赖关系的大型多样化训练集,论文聚合了几个最大的语言建模数据集
  • 论文创建的聚合数据集包括 Wikipedia(2018)、CC-Stories(2018)、RealNews(2019)和 OpenWebtext(2019)
  • 为了避免训练集数据泄露到下游任务中
    • 移除了 WikiText103 测试集(2016)中存在的 Wikipedia 文章
    • 还移除了 CC-Stories 语料库中因预处理 artifacts 引入的不必要换行
  • 对于 BERT 模型,论文在训练数据集中包含了 BooksCorpus(2015),但由于该数据集与 LAMBADA 任务重叠,因此在 GPT-2 训练中排除了该数据集
  • 论文合并了所有数据集,然后做了以下操作:
    • 从聚合数据集中过滤掉内容长度小于 128 个 token 的所有文档
    • (聚合数据集中可能存在重复的相似内容),论文使用局部敏感哈希(locality-sensitive hashing, LSH)对 Jaccard 相似度大于 0.7 的内容进行去重
  • 最终的聚合语料库包含 174 GB 的去重文本

Training Optimization and Hyperparameters

  • 为了高效训练论文的模型,论文利用混合精度训练(mixed precision training)和动态损失缩放(dynamic loss scaling),以充分利用 V100 的Tensor Core(2017;2018)
  • 使用简单的正态分布初始化权重 \(W\) :
    $$ W \sim N(0,0.02)$$
  • 在残差层(residual layers)之前立即按下面的比例缩放权重:
    $$ \frac{1}{\sqrt{2 N} } $$
    • 其中 \(N\) 是由自注意力和 MLP 块组成的 Transformer 层数
  • 使用带权重衰减(weight decay)(2019)的 Adam 优化器(2014),其中 \(\lambda=0.01\)
  • 使用全局梯度范数裁剪(global gradient norm clipping)为 1.0,以提高大型模型训练的稳定性
  • dropout 率均设置为 0.1
  • 在每个 Transformer 层之后使用激活检查点(2016)
  • 对于 GPT-2 模型
    • 所有训练都使用 1024 个子词单元(subword units)的序列
    • 批大小为 512
    • 共训练 300K 次迭代
    • 初始学习率为 1.5e-4(3k 次迭代的预热期(warmup period))
    • 剩余的 297K 次迭代中遵循单周期余弦衰减(single cycle cosine decay)至最小学习率 1e-5
  • 对于 BERT 模型
    • 在很大程度上遵循(2019)中描述的训练过程
    • 论文使用原始 BERT 词典,词汇表大小为 30522
    • 按照(2019)的建议,用句子顺序预测(sentence order prediction)替换了下一句预测(next sentence prediction)头,并使用(2019)的全词 n-gram 掩码(whole word n-gram masking)
    • 批大小设置为 1024
    • 学习率为 1.0e-4
    • 在 10k 次迭代中预热,然后在 2M 次迭代中线性衰减
    • 其他训练参数与(2018)保持一致

Experiments

  • 论文所有的实验都使用了多达 32 台 DGX-2H 服务器(总共 512 块 Tesla V100 SXM3 32GB GPU)
  • 论文的基础设施针对多节点深度学习应用进行了优化,服务器内部的 GPU 之间通过 NVSwitch 实现 300 GB/秒 的带宽,服务器之间通过每台服务器配备的 8 个 InfiniBand 适配器实现 100 GB/秒 的互联带宽

Scaling Analysis

  • 为了测试论文实现方案的可扩展性,论文考虑了四组参数的 GPT-2 模型,详情见表1
  • 为了使自注意力层中的通用矩阵乘法(GEMM)大小保持一致,每个注意力头的隐藏层大小固定为 96,通过调整注意力头数和层数来获得参数规模从 1B 到 8B 的配置
    • 1.2B 参数的配置可以在单个 GPU 上运行,而 8B 参数的模型则需要 8 路模型并行(8 块 GPU)
    • 将原始词汇表大小从 50257 填充为 51200,原因如下:
      • 为了使 logit 层的通用矩阵乘法(GEMM)更高效,每个 GPU 上的词汇表大小最好是 128 的倍数
      • 论文的模型并行最多为 8 路,因此论文将词汇表填充为可被 \(128 \times 8 = 1024\) 整除的大小,最终填充后的词汇表大小为 51200
    • 论文研究了模型并行和模型+数据并行两种扩展性
      • 对于模型并行扩展性,所有配置都使用固定的批大小 8
      • 数据并行扩展对于训练许多 SOTA 模型是必要的,这些模型通常使用更大的全局批大小
      • 在模型+数据并行的情况下,论文所有实验的全局批大小都固定为 512,这对应于 64 路数据并行
Model and Data Parallelism
  • 本节展示模型并行和模型+数据并行情况下,相对于模型参数的弱扩展性
  • 弱扩展性通常通过扩展批大小来实现,但这种方法无法解决训练无法在单个 GPU 上运行的大型模型的问题,而且会导致大批量训练的收敛性下降
  • 论文不同,在这里使用弱扩展性来训练原本无法实现的更大模型
  • 所有扩展性数据的基准是表1中的第一种配置(1.2B 参数)在单个 GPU 上的运行情况
    • 这是一个很强的基准,因为它在整个训练过程中达到了 39 TeraFLOPs,相当于 DGX-2H 服务器中单个 GPU 理论峰值浮点运算的 30%
  • 图5显示了模型并行和模型+数据并行的扩展性数值:论文在两种设置下都观察到了优异的扩展性。例如,
    • 8.3B 参数的模型在 8 路(8 块 GPU)模型并行下实现了 77% 的线性扩展效率
    • 模型+数据并行由于需要额外的梯度通信,扩展性数值略有下降
    • 但即使是在 512 块 GPU 上运行的最大配置(8.3B 参数),相对于 1.2B 参数在单个 GPU 上的强基准配置,论文也实现了 74% 的线性扩展效率
    • 更多扩展性分析见附录D

Language Modeling Results Using GPT-2

  • 为了证明大型语言模型能进一步推动技术发展,论文考虑训练表2中列出的不同规模和配置的 GPT-2 模型

    • 355M 参数的模型与 BERT-Large 模型(2018)的规模和配置相当
    • 2.5B 参数的模型比之前最大的 GPT-2 模型更大,而据论文所知,8.3B 参数的模型比任何已训练的从左到右的 Transformer 语言模型都要大
  • 论文使用第4节中描述的流程来训练和评估论文的语言模型

  • 表2还列出了完成一个 epoch 所需的时间,一个 epoch 相当于 68507 次迭代

    • 例如, 8.3B 参数的模型在 512 块 GPU 上,每个 epoch 大约需要两天
    • 与表1中用于扩展性研究的配置相比:
      • 2.5B 参数的模型是相同的
      • 8.3B 参数的模型有 24 个注意力头(而不是 32 个)
      • 355M 参数的模型比之前的任何模型都小,但仍使用 64 块 GPU 进行训练,因此每个 epoch 的时间短得多
  • 图6显示了验证集困惑度(perplexity)随迭代次数的变化

    • 随着模型规模的增加,验证集困惑度降低, 8.3B 参数的模型达到 9.27 的验证集困惑度
  • 表3报告了训练后的模型在 LAMBADA 和 WikiText103 数据集上的零样本评估结果

    • 有关评估方法的更多细节,请参见附录E
  • 论文观察到,随着模型规模的增加,WikiText103 上的困惑度降低,LAMBADA 上的完形填空准确率(cloze accuracy)提高

    • 论文的 8.3B 参数模型在 WikiText103 测试集上实现了 10.81 的调整后困惑度,达到了最佳水平
    • 8.3B 参数的模型在 LAMBADA 任务上的准确率为 66.51%,同样超过了之前的完形填空准确率结果
    • 附录C 中包含了 8.3B 参数模型生成的样本
  • 最近,微软的研究人员与 NVIDIA 合作,使用 Megatron 训练了一个 17B 参数的 GPT-2 模型,称为 Turing-NLG(微软,2020),并表明随着模型规模的扩大,准确率进一步提高,凸显了更大模型的价值

  • 为了确保论文的训练数据中不包含任何测试集数据

    • 论文计算了测试集中的 8-grams 在训练集中出现的百分比,就像之前的工作(2019)所做的那样
    • WikiText103 测试集的重叠率最多为 10.8%,LAMBADA 测试集(2016)的重叠率最多为 1.4%
    • WikiText103 测试集与 WikiText103 训练集的重叠率已经达到 9.09%(2019)
    • 由于这些结果与之前的工作一致,论文有信心训练数据中没有无意中包含任何测试数据的文档

Bi-directional Transformer Results Using BERT

  • 在本节中,论文将论文的方法应用于类 BERT(BERT-style)的 Transformer 模型,并研究模型扩展对多个下游任务的影响
  • 先前的工作(2019)发现,当模型规模超过 336M 参数的 BERT-large 时,会出现意想不到的模型性能下降
    • 为了解决这种下降,该工作的作者(2019)引入了参数共享,并表明他们的模型比原始 BERT 模型的扩展性好得多
  • 论文进一步研究了这种现象,并通过实证证明,如图7所示,重新排列 Layer Normalization 和残差连接(residual connections)的顺序,对于使类 BERT 模型能够扩展到超过 BERT-Large 的规模至关重要
  • 图7中的架构(b)消除了使用原始BERT架构(a)时观察到的不稳定性,并且具有更低的训练损失
  • 据论文所知,论文是第一个报告这种修改能够训练更大 BERT 模型的研究
  • 使用图7(b)中的架构修改,论文考虑了表4中详细列出的三种情况
    • 336M 参数的模型与 BERT-large 的规模相同
    • 1.3B 参数的模型与之前被证明比 336M 参数的 BERT-large 模型结果更差的 BERT-xlarge 配置相同(2019)
    • 论文通过同时增加隐藏层大小和层数,将 BERT 模型进一步扩展到 3.9B 参数的情况
    • 在所有情况下,每个注意力头的隐藏层大小都固定为 64
    • 336M 和 1.3B 参数的模型训练了 2M 次迭代,而 3.9B 参数的模型训练了 1.5M 次迭代,目前仍在训练中
  • 在 3% 的保留集上,336M、1.3B和 3.9B 参数的模型分别实现了 1.58、1.30 和 1.16 的验证集困惑度 ,随着模型规模的增加而单调下降
  • 论文在多个下游任务上对训练后的模型进行微调,包括 GLUE 基准(2019)中的 MNLI 和 QQP,斯坦福问答数据集(2016;2018)中的 SQuAD 1.1 和 SQuAD 2.0,以及阅读理解 RACE 数据集(2017)
    • 微调时,论文遵循(2019b)中的相同流程
    • 论文首先对批大小和学习率进行超参数调优
    • 获得最佳值后,论文报告 5 个不同随机种子初始化的开发集结果的中值
    • 每个模型和任务使用的超参数见附录A
  • 表5显示了 MNLI、QQP、SQuAD 1.1 和 SQuAD 2.0 的开发集结果,以及 RACE 的测试集结果
    • 对于 RACE 的测试集结果,论文首先使用开发集找到在 5 个随机种子上给出中值分数的检查点,然后报告该检查点在测试集上的结果
    • 论文还报告了 SQuAD 开发集和 RACE 测试集的 5 路集成(ensemble)结果
  • 从表5中论文观察到:
    • (a)随着模型规模的增加,所有情况下的下游任务性能都有所提高;
    • (b)论文的 3.9B 参数模型在开发集上建立了比其他基于 BERT 的模型更好的最佳结果;
    • (c)论文的 3.9B 参数模型在 RACE 测试集上同时实现了单模型和集成模型的最佳结果

Future Work

  • 未来的工作有几个方向
    • 继续扩大预训练规模是一个很有前景的研究方向,这将进一步考验现有的深度学习硬件和软件
      • 为了实现这一点,需要提高优化器的效率和减少内存占用
      • 此外,训练超过16B 参数的模型将需要比 DGX-2H 服务器的16块GPU更多的内存
      • 对于此类模型,混合层内和层间模型并行以及节点间模型并行将更合适
    • 其他三个研究方向包括:
      • (a)预训练不同的模型系列(XLNet、T5);
      • (b)评估大型模型在更多更困难和更多样化的下游任务(如生成式问答、摘要和对话)上的性能;
      • (c)使用知识蒸馏(knowledge distillation)从这些大型预训练教师模型(teacher models)中训练小型学生模型(student models)

A. BERT Finetuning Hyperparameters

  • 表6展示了每个模型和任务在微调过程中使用的超参数

B. Model Parallel Supplementary Material

在本节中,论文将详细介绍混合模型并行与数据并行以及随机数生成的处理方法

B.1 Hybrid Model and Data Parallelism

  • 模型并行与数据并行是相互独立的,因此我们可以同时使用这两种方法,在合理的时间内训练大型模型
  • 图8展示了用于混合模型并行与数据并行的 GPU 分组方式
  • 同一服务器内的两块或多块 GPU 组成模型并行组(例如图8中的 GPU 1 至 8),这些 GPU 共同分布运行一个模型实例
  • 其余的 GPU(可能位于同一服务器内,但更常见的是位于其他服务器中)运行额外的模型并行组
  • 每个模型并行组中处于相同位置的 GPU(例如图8 中的 GPU 1、9、……、505)组成数据并行组,因此数据并行组内的所有GPU都持有相同的模型参数
  • 在反向传播过程中,论文并行运行多个梯度全归约(gradient all-reduce)操作,以在每个不同的数据并行组内归约权重梯度
  • 所需的 GPU 总数是模型并行组数与数据并行组数的乘积
  • 例如,对于 8.3B 参数的模型,论文每个模型并行组使用 8 块 GPU,并采用 64 路数据并行,因此总共需要 512 块 GPU
  • 所有通信都是在 PyTorch 中通过 Python 调用 NCCL 实现的
  • 每个模型并行组内的 GPU 在组内所有 GPU 之间执行全归约操作
  • 对于数据并行,每个全归约操作都在每个模型并行组中的一块 GPU 之间进行

B.2 Model Parallel Random Number Generation

  • 诸如 dropout 等利用随机数生成的技术是现代深度学习训练中的重要组成部分
  • Transformer 在模型并行区域之外的残差连接之前以及模型并行区域内的自注意力块中都有 dropout 层
  • 由于一些 dropout 层位于模型并行区域内,而另一些则不在,因此论文需要谨慎处理随机数生成,以确保 dropout 正常工作
  • 为了在模型并行工作节点之间同步残差连接 dropout,论文在训练开始时使用相同的种子初始化随机数生成器
  • 这使得所有模型并行工作节点的 dropout 模式完全相同
  • 然而,模型并行区域内的 dropout 应该为每个工作节点生成不同的随机模式,以实现整个操作的随机性
    • 为了实现这一点,论文为模型并行区域内的 dropout 维护一个单独的随机数生成器
    • 每个模型并行工作节点的这个随机数生成器都有唯一的种子

C. Text Samples

  • 本节展示了 Megatron-LM 使用上下文提示生成的一些文本样本

D. Further Scaling Analysis

  • 在本节中,论文研究注意力头数对扩展性结果的影响
  • 论文还展示了 1.2B 参数模型的强扩展性结果

D.1 Attention Heads and Scaling

  • 本节研究注意力头数对模型并行扩展性的影响
  • 为此,论文考虑 8.3B 参数配置(采用8路模型并行),并将注意力头数从 16 调整到 32
  • 结果如表7所示
    • 随着注意力头数的增加 ,自注意力层内的一些通用矩阵乘法(GEMMs)会变小 ,自注意力 softmax 中的元素数量也会增加
      • 这导致扩展性效率略有下降
    • 未来的研究在设计平衡模型速度和准确率的大型 Transformer 模型时,应注意这个超参数

D.2 Strong Scaling

  • 论文的模型并行主要是为了能够训练规模超过单个 GPU 内存的模型,但它也可以在不增加批大小的情况下加速较小模型的训练
  • 为了衡量这种加速效果,论文训练了一个固定的 1.2B 参数模型
  • 论文使用每次迭代 8 个样本的固定批大小,并通过模型并行增加 GPU 数量
  • 结果如表8所示
    • 使用两块GPU可使训练速度提高 64%
    • 超过这个数量后,论文看到收益递减,因为每个 GPU 的计算量减少,内存带宽和通信开销开始占据主导地位

E. 使用 WikiText103 和 LAMBADA 评估语言模型(Evaluating Language Models Using WikiText103 and LAMBADA)

  • 在本节中,论文详细介绍了 WikiText103 数据集(2016)和 LAMBADA 数据集(2016)上完形填空式预测准确率的评估方法

E.1 Wikitext103 困惑度(Wikitext103 Perplexity)

  • WikiText103 困惑度是过去几年中被广泛研究的评估标准,困惑度是语料库平均交叉熵的指数(2011)。是语言模型的自然评估指标,因为语言模型表示整个句子或文本的概率分布:
    $$\text{PPL}=\exp \left(-\frac{1}{T_{o} } \sum_{t}^{T} \log P(t | 0: t-1)\right) \tag{4}$$
  • 为了计算式(4)中的困惑度
    • 论文根据子词词汇表对 WikiText103 测试语料库进行分词,并对每个 token \([0, T]\) 的交叉熵损失求和
    • 然后,论文通过原始分词方案中的 token 数量 \(T_{0}\) 对交叉熵损失进行归一化
  • WikiText103 测试语料库已经过预处理,采用单词级分词,先前的研究使用这种分词来计算困惑度
  • 为了在公平的环境下与先前的研究比较论文模型的困惑度,论文必须通过原始 token 数量 \(T_{0}\) 进行归一化,而不是输入到论文模型中的实际分词数据的 token 数量 \(T\)
  • 这种预分词还会在文本中引入训练数据中不存在的 artifacts
  • 为了缓解这种分布不匹配,论文首先使用可逆的去分词器对 WikiText103 测试数据集进行预处理,以去除与标点符号和空格相关的各种 artifacts
  • \(T_{0}\) 的值是在这种预处理之前计算的
  • 对于WikiText103的测试集, \(T_{o}=245566\) , \(T=270329\)
  • 论文还必须对困惑度计算进行一项特定于 Transformer 的修改
    • 与基于 RNN 的语言模型不同, Transformer 在固定的窗口输入大小上运行
    • 因此,它们无法完全计算 \(P(t | 0: t-1)\) ,只能计算 \(P(t | t-w: t-1)\) ,其中 \(w\) 是论文的上下文大小:1024 个 token
    • 但为数据集中的每个 token 计算这个值的成本过高,因为论文必须对 \(w\) 大小的上下文进行大约 \(T\) 次评估
    • 为了高效地评估论文的模型,论文采用了一种称为重叠评估(overlapping evaluation)的折中方法,每次滑动窗口前进一定的重叠量 \(o\) ,只计算窗口最后 \(o\) 个 token 的交叉熵损失
    • 在论文的实验中,论文使用 32 的重叠量 \(o\) ,并以这种方式对所有滑动窗口计算损失

E.2 LAMBADA Cloze Accuracy

  • 处理长期上下文的能力对于 SOTA 语言模型至关重要,也是长文本生成和基于文档的问答等问题的必要前提
  • 像 LAMBADA 这样的完形填空式数据集旨在衡量模型在这些类型的长期上下文中运行和推理的能力
  • 完形填空式阅读理解使用单词 token 的上下文 \(x=x_{1: t}\) ,其中一个 token \(x_{j}\) 被掩码;
    • 模型的目标是正确预测缺失的第 \(j\) 个 token
  • 为了准确预测缺失的 token,模型需要深入理解周围的上下文以及在这种上下文中应如何使用语言
  • LAMBADA 使用完形填空式阅读理解来测试生成式从左到右语言模型,它构建了 45 个句子的示例,其中上下文 \(x_{t}\) 的最后一个单词被掩码
  • 论文的模型使用子词单元,因此在 LAMBADA 评估中,论文使用原始的、未处理的 LAMBADA 数据集,并要求论文的模型预测组成单词 token 的多个子词token
  • 论文使用教师强制(teacher forcing),只有当所有输出预测都正确时,才认为答案正确
    • 这种形式与原始的单词 token 预测任务等效

NLP——Megatron-LM原始论文解读-第三篇

注:本文包含 AI 辅助创作

  • 参考链接:
    • Megatron 系列目前公认的有三篇核心论文如下,它们分别对应了 张量并行、3D 并行 与 序列并行/激活重计算优化 三大阶段,本文是第三篇的解读
    • 第三篇:(Megatron-LM-3)Reducing Activation Recomputation in Large Transformer Models, MLSys 2023, NVIDIA
      • 核心贡献:提出 序列并行(Sequence Parallelism) 与 选择性激活重计算 ,将激活显存占用再降 3–5 倍,支持更长序列与更大批训练;与 Flash-Attention 思想互补,现已成为 Megatron-Core 默认配置
    • 论文补充材料:Supplementary Material of Megatron-LM-3

Paper Summary

  • 整体内容总结:
    • 论文提出了 序列并行(Sequence Parallelism) 与 减少激活重计算(Reducing Activation Recomputation)
    • 论文证明了将 SP (Sequence Parallelism) 与 TP (Tensor Parallelism) 结合使用可以显著减少所需的激活内存
    • 结合选择性激活重计算 (Selective Activation Recomputation),论文证明了可以实现内存占用减少 \(5\times\),并且恢复了使用完全激活重计算 (Full Activation Recomputation) 所引入的计算开销的 90% 以上
      • 传统上实现:通过不存储激活值,而是在反向传播时按需要重新计算它们,节省了内存但增加了冗余计算
      • 论文表明大部分这种冗余计算是不必要的,因为我们可以有选择的保存部分激活值
  • 补充:结合 张量并行 (tensor parallelism, TP),这些技术几乎消除了重计算激活的需要
  • 论文在规模 高达一万亿参数 的语言模型上评估了论文的方法,结果表明:
    • 论文的方法将激活内存减少了 \(5\times\),同时将激活重计算带来的执行时间开销降低了 90% 以上
  • 在 2240 个 NVIDIA A100 GPU 上训练一个 530B 参数的 GPT-3 风格模型 (2022) 时,论文实现了 54.2% 的模型浮点运算利用率 (Model Flops Utilization, MFU),这比使用重计算实现的 42.1% 快了 29%

Introduction and Discussion

  • 随着 Transformer 模型规模向万亿参数扩展,需要模型并行 (model parallelism) 来将模型参数、激活值和优化器状态分布到多个设备上,以便它们能够放入设备内存并在现实的时间量内可训练
  • 虽然模型并行线性地减少了每个设备的参数数量(例如,当模型并行规模加倍时,每个设备的参数数量减半),但模型并行的扩展是有限度的
    • 张量级模型并行 (Tensor-level model parallelism, TP) 增加了通信需求并引入了更小且性能较差的矩阵乘法,使得将模型分割到大量设备上效率低下
      • 因此,TP 通常仅限于通过高速带宽连接(例如 DGX 服务器内通过 NVLink 连接的 GPU)的相对较小的 GPU 组
    • 流水线并行 (Pipeline parallelism, PP) 需要存储几个微批次 (microbatches) 的激活值以减少流水线气泡 (pipeline bubble) (2021)
      • 因此, PP 只能帮助减少存储模型参数和优化器状态所需的内存,无法在保持高设备利用率的同时减少激活值所需的内存
    • 因此,激活值的存储迅速成为扩展大型 Transformer 模型的关键问题
      • 问题:实际上,对于不适用 Activation Recomputation 的场景,PP 也减少激活存储了;但是对于 Activation Recomputation 的场景,每一层存储的激活基本都还在,所以相当于此时 PP 对激活存储是没有节省的
  • 为了量化这一点,图 1 显示了从 22B 参数到 1 万亿(1T)参数的四种模型配置所需的内存(模型配置的详细信息在表 3 中提供)
    • 可以看出,对于所有这些情况,基线情况所需的内存都超过了 NVIDIA A100 GPU 提供的 80GB 内存
    • 缓解这种内存压力的标准方法是简单地不存储大部分激活值,并在反向传播 (backward pass) 期间根据需要重新计算它们以计算梯度 (2016)
      • 这种方法通常称为“梯度检查点 (gradient checkpointing)”或“激活重计算 (activation recomputation)”,会带来严重降低训练效率的代价
    • 对于 Transformer 架构,大多数先前的工作在 Transformer 层边界处检查点(或存储)激活值,并在反向传播中重新计算其余必要的激活值
    • 在论文中,论文将这种方法称为“完全激活重计算 (full activation recomputation)”
    • 在论文的训练运行中,论文观察到当使用完全激活重计算时 ,会产生 \(30-40\%\) 的执行时间开销
  • 论文提出了新技术,帮助缓解存储激活值的内存压力,从而减少重计算激活值的需求
    • 这些技术特定于 Transformer 架构,并且易于实现,对计算效率没有影响或影响非常小
  • 正如论文在第 2 节中详细说明的那样,还有其他几种技术可以减少训练大型模型的内存需求,例如:
    • 跨数据并行 Rank (data parallel ranks) 分区各种数据
    • 将数据卸载到 CPU 内存 (2020; 2021)
  • 以上这些技术与论文提出的技术是互补的,可以额外采用以获得更大的内存节省;但通常这些其他技术比论文提出的技术具有更高的实现成本和对计算效率的更大影响
  • 将这些技术与论文的技术进行比较的分析超出了论文的范围,留待未来工作
  • 论文工作内容包括:
    • 首先论文简要回顾 Transformer 架构,然后建立一个近似公式,用于存储 Single Stack Transformer 模型前向传播 (forward pass) 中激活值所需的内存
      • 使用这个公式,我们可以研究不同形式的模型并行如何影响激活内存需求
    • 然后论文引入序列并行 (sequence parallelism) 与 TP (tensor parallelism) 结合,以防止在不适合标准 TP 的区域中冗余存储激活值
      • 论文表明,通过有选择地保存哪些激活值以及重新计算哪些激活值,我们可以消除大部分重计算成本,同时仅使用不进行重计算时所需内存的一小部分
    • 最后论文提出了几个实验,测量这些技术对训练各个组成部分以及完整训练吞吐量的改进

Related Work

  • 模型并行 (model parallelism) 使得能够跨多个 GPU 训练非常大的模型
    • 模型参数以及这些模型相关的优化器状态需要大量内存,无法容纳在单个 GPU 上
    • 即使能够将模型放入单个 GPU(例如,通过在主存和设备内存之间交换参数 (2021)),所需的高计算操作数量可能导致不切实际的长训练时间
    • 结论:需要并行化
  • 两种常用的模型并行形式用于将模型参数分布到 GPU 上:
    • 1) TP (tensor parallelism),其中每层的参数分布在许多设备上 (2018; 2019; 2021);
    • 2) PP (pipeline parallelism),其中模型沿网络的层维度分割 (2019; 2021a; 2019)
  • 一些最近的方法结合了两种类型的模型并行,以支持训练高达 1T 参数的大型模型 (2021)
    • 模型并行的替代方案是结合一些训练技术与数据并行 (data parallelism),以实现大规模模型训练 (2020; 2021;)
      • 这种方法基于将优化器状态、梯度和参数跨数据并行 Rank 进行分片 (sharding)
    • 最近的一个扩展 (2021) 使用 CPU 卸载技术,使得能够在少量 GPU 上训练数万亿参数的模型
      • 与模型并行相比,这些基于数据并行的技术效率较低,并且不能很好地扩展到大量 GPU (2021),因此更适合在资源受限的环境中进行模型微调 (finetuning)
    • 论文仅关注模型并行优化,将这些技术与论文的技术进行比较的分析超出了论文的范围
  • 此外,Megatron-LM (2019) 中引入的 TP 在一定程度上帮助减少了激活内存
    • 在这种方法中,Transformer 的某些部分激活值没有在 TP Rank 之间分割,增加了激活内存开销
  • Li 等 (2021a) 提出的序列并行 (sequence parallelism),其中激活值沿序列维度分区贯穿整个网络可以缓解这个问题
    • 但他们的方法类似于数据并行,要求在所有设备上复制参数和优化器状态,这使得它不适合大型模型训练
  • Sagemaker (2021) 和 GSPMD (2021) 提出了内存高效的 TP 版本,这些版本在整个网络中沿隐藏维度 (hidden dimension) 在设备之间分割激活值
    • 这些方法的主要缺点是它们包含多设备层归一化 (multi-device layer normalization),这在计算/通信上非常低效
      • LayerNorm 操作需要沿隐藏维度顺序计算均值和方差
      • 如果 LayerNorm 沿隐藏维度分割(即使用 TP),它将增加两个额外的全规约 (all-reduce) 操作
      • 而且 LayerNorm 反向传播需要两个额外的 sequential reductions(即另外两个 all-reduce)
    • 与 LayerNorm 计算时间相比,每个 Transformer 层的这 4 次通信将引入显著的开销
    • LayerNorm 操作的这种通信开销可以通过 SP 来避免
  • 在论文中,论文提出了一种新技术,它利用了 TP 和 SP 的优点,而没有先前方法的任何缺点
    • 换句话说,论文的技术结合了 TP 和 SP ,显著减少了激活内存,而没有任何额外的计算、通信或内存开销
  • 论文的论文特别侧重于研究 Transformer 架构,因为作者相信它是一个重要且广泛使用的架构,足以证明手动检查各个层以了解其内存和计算贡献并推导最佳并行化方案是合理的
  • 通过分析建模,论文表明内存(包括激活值和参数)在设备之间均匀分区
    • 换句话说,论文的并行化策略在内存需求方面是最优的
    • Transformer 网络被广泛使用,足以高度利用这种手动搜索和验证最佳并行化策略的方法
    • 自动搜索 (2019; 2019) 有可能找到这种最优策略,但据论文所知,论文尚未看到这些方法应用于 Transformer 网络的任何已发表结果
  • 与并行化方案类似,论文提出的选择性激活重计算 (selective activation recomputation) 是基于手动搜索 Transformer 模型在激活内存需求和重计算之间的最佳权衡
  • 探索自动方法 (2021) 是否可以补充并可能改进这种权衡是未来工作的一个有趣方向

Transformer Architecture

  • 论文考虑一个具有 \(L\) 层的 Single Stack Transformer 编码器或解码器,如图 2 所示
  • 符号定义:\(s\) 是序列长度 (sequence length),\(h\) 是隐藏维度大小 (hidden dimension size),\(v\) 是词汇表大小 (vocabulary size)
  • 在网络的开始,输入 Tokens 被送入一个词嵌入表(word embedding table):
    $$v\times h$$
  • Token 嵌入与学习到的位置嵌入(positional embeddings) 相结合:
    $$s\times h$$
  • 嵌入层的输出,即 Transformer 块的输入,是一个 3 维张量:
    $$s\times b\times h$$
    • 其中 \(b\) 是微批次大小 (microbatch size)
  • 每个 Transformer 层由一个具有 \(a\) 个注意力头 (attention heads) 的自注意力块 (self-attention block) 和一个具有两层的多层感知器 (MLP) 组成,该 MLP 将隐藏大小增加到 \(4h\),然后将其减小回 \(h\)
  • 每个 Transformer 层的输入和输出具有相同的大小
    $$ s\times b\times h $$
  • 最后一个 Transformer 层的输出被投影回词汇表维度以计算交叉熵损失 (cross-entropy loss)
    • 论文假设词嵌入和输出层权重是共享的
  • 论文相关变量名称列于表 1 中以供参考:

Activation Memory

  • 论文推导出一个近似公式,用于计算如图 2 所示的 Single Stack Transformer 模型在前向传播中存储激活所需的内存
  • 请注意,论文中的 “激活(activations)” 指的是 在前向传播中创建并在反向传播期间梯度计算所必需的任何张量
    • 注:不包括模型的主要参数和优化器状态,但包括 Dropout 操作使用的掩码
  • 论文只考虑内存的主要贡献者,而忽略小的缓冲区
    • 小的缓冲区包括 LayerNorm 输入的均值和方差(\(2sb\))以及 GEMM 操作的偏置(\(O(h)\))
    • 总的来说,这些缓冲区占激活内存的比例远小于 1%,因为隐藏维度(\(h\))和序列维度(\(s\))的数量级都是数千,因此与 GEMM 和 LayerNorm 激活的 \(O(sbh)\) 大小相比,\(2sb\) 和 \(O(h)\) 都缺少了其中一个因子
  • 论文还假设网络和激活以 16-bit 浮点格式存储,因此每个元素需要 2 字节的存储空间
    • 唯一的例外是 Dropout 掩码,每个元素只需要 1 个字节
  • 请注意,除非明确提及,本节中报告的所有大小均以字节(Bytes)为单位,而不是元素数量

Activations Memory Per Transformer Layer

  • 如图 2 所示,每个 Transformer 层包含一个注意力块和一个 MLP 块,通过两个 LayerNorm 连接
  • 下面,论文推导存储这些元素激活所需的内存:
  • 注意力块 (Attention block):
    • 包括自注意力操作,后跟一个线性投影和注意力 Dropout
    • 线性投影存储其输入激活,大小为 \(2sbh\),注意力 Dropout 需要一个大小为 \(sbh\) 的掩码
    • 图 3 所示的自注意力包含几个元素:
      • Query (\(Q\)), Key (\(K\)), and Value (\(V\))矩阵乘法 (matrix multiplies): 论文只需要存储它们共享的输入,大小为 \(2sbh\)
      • \(QK^{T}\) 矩阵乘法 (matrix multiply): 需要存储 \(Q\) 和 \(K\),总大小为 \(4sbh\)
      • Softmax: 反向传播需要存储 Softmax 输出,大小为 \(2as^{2}b\)
      • Softmax Dropout: 只需要一个大小为 \(as^{2}b\) 的掩码
        • 问题:Softmax 也有 Dropout 吗?
        • 理解:在 Softmax 之后施加 Dropout,随机将部分注意力权重置为 0,目的是防止模型过度依赖某些特定的注意力权重,减少过拟合风险,提高模型的泛化能力,但原始 Transformer 论文中,没看到这部分的 dropout,只有 Residual Dropout(\(P_\text{drop} = 0.1\))
      • 注意力作用于值(\(V\))(Attention over Values (\(V\))): 论文需要存储 Dropout 输出(\(2as^{2}b\))和值(\(2sbh\)),因此需要 \(2as^{2}b+2sbh\) 的存储空间
    • 将上述值相加,注意力块总共需要 \(11sbh+5as^{2}b\) 字节的存储空间
  • MLP:
    • 两个线性层存储它们的输入,大小分别为 \(2sbh\) 和 \(8sbh\)
    • GeLU 非线性操作也需要其输入(大小为 \(8sbh\))用于反向传播
    • 最后,Dropout 存储其掩码,大小为 \(sbh\)
    • MLP 块总共需要 \(19sbh\) 字节的存储空间
  • LayerNorm:
    • 每个 LayerNorm 存储其输入,大小为 \(2sbh\),因此论文总共需要 \(4sbh\) 的存储空间
  • 最终,将注意力、MLP 和 LayerNorm 所需的内存相加,存储 Transformer 网络单层激活所需的内存为:
    $$\text{Activations memory per layer}=sbh\left(34+5\frac{as}{h}\right) \tag{1}$$
  • 上述方程适用于未应用任何形式模型并行的情况

Model Parallelism

  • 论文先量化 TP 对每层所需激活内存的影响,然后介绍一种将 SP 与 TP 相结合的新方法,以进一步减少每层激活所需的内存
  • 在本小节的最后,论文还讨论了 PP 对激活内存的影响,并推导了激活所需总内存的公式
TP (Tensor Parallelism)
  • 论文使用 Shoeybi 等人 (2019) 开发的 TP ,并按照图 4 所示并行化注意力块和 MLP 块
    • 这种形式的并行性引入了两个额外的通信操作 \(f\) 和 \(\bar{f}\)
    • 更多细节,请参阅该论文 (2019)
  • TP 不仅并行化了注意力块和 MLP 块内的模型参数和优化器状态,还并行化了这些块内部的激活
    • 这些块的输入激活(例如,输入到 \(Q\)、\(K\) 和 \(V\) 矩阵乘法的输入,或输入到 \(h \to 4h\) 线性层的输入)没有被并行化,只有每个块内的激活在 TP 组内被划分
    • 假设采用 \(t\) 路 TP ,存储激活所需的每层内存从公式 1 减少到:
      $$\text{Activations memory per layer}=sbh\left(10+\frac{24}{t}+5\frac{as}{ht}\right) \tag{2}$$
SP (Sequence Parallelism)
  • 如图 4 所示, TP 并行化了 Transformer 层中训练过程中最耗时的部分,因此它在计算上是高效的
  • 但它保留了 LayerNorm 以及注意力块和 MLP 块之后的 Dropout 操作不变,因此 LayerNorm 和 Dropout 在 TP 组内是复制的(单卡上需要完整保留所有参数和激活)
  • 这些元素不需要大量计算,但需要相当数量的激活内存
  • 定量地说,公式 2 中的 \(10sbh\) 部分就是由于这些复制操作造成的,因此它们没有除以 TP 大小 \(t\)
  • 论文注意到,在 Transformer 层的非 TP 区域中,操作在序列维度上是独立的
    • 这一特性允许论文沿序列维度 \(s\) 划分这些区域
    • 沿序列维度划分减少了激活所需的内存
  • 这种额外的并行性级别在 \(f\) 之前和 \(\bar{f}\) 之后引入了新的通信集合,这些操作将充当 SP 区域和 TP 区域之间的转换器
    • 例如,在前向传播中,论文需要在图 4 中的操作符 \(f\) 之前进行一次额外的 All-Gather
    • 这些额外的通信会引入开销并减慢训练速度
  • 为了避免这些额外的通信,论文将这些操作与 \(f\) 和 \(\bar{f}\) 操作符结合起来,并引入新的操作 \(g\) 和 \(\bar{g}\),如图 5 所示
    • 可以看出,\(g\) 和 \(\bar{g}\) 是 SP 区域和 TP 区域之间的转换器
    • 论文在本小节的剩余部分推导这些操作
  • 论文使用 MLP 块详细说明 \(g\) 和 \(\bar{g}\) 的推导。在非并行形式中,如图 2 所示,LayerNorm 后接 MLP 块可以表述为:
    $$
    \begin{align}
    Y &=\text{LayerNorm}(X), \\
    Z &=\text{GeLU}(YA), \\
    W &=ZB, \\
    V &=\text{Dropout}(W),
    \end{align}
    $$
    • 其中:
      • \(X\) 是 LayerNorm 的输入,大小为 \(s\times b\times h\)
      • \(A\) 和 \(B\) 是线性层的权重矩阵,大小分别为 \(h\times 4h\) 和 \(4h\times h\)
  • 上述操作的张量和 SP 组合形式如图 6 所示
  • 下标表示在加速器间的划分(注:理解为 Rank),上标表示划分所沿的维度 ,例如:
    • \(X_{1}^{s}\) 是 \(X\) 在第一个加速器上的部分,沿 \(s\) 维度(序列维度)划分
    • \(Z_{2}^{h}\) 是 \(Z\) 在第二个加速器上的部分,沿 \(h\) 维度(隐藏维度)划分
  • LayerNorm 的输入沿序列维度并行化
    $$ X=[X_{1}^{s},X_{2}^{s}]$$
  • 因此,LayerNorm 的输出也将沿序列维度并行
    $$Y=[Y_{1}^{s},Y_{2}^{s}]$$
  • 带有 GeLU 非线性的线性层需要完整的输入 \(Y\),因此论文需要执行一次 All-Gather
    • 这意味着 \(g\) 在前向传播中是一个沿序列维度的 All-Gather 操作
  • 通过将 \(A\) 沿其列划分(\(A_{1}^{c}\) 和 \(A_{2}^{c}\)),将 \(B\) 沿其行划分(\(B_{1}^{r}\) 和 \(B_{2}^{r}\)),论文避免了通信(更多细节请参见 (2019))并得到 \(W_{1}\) 和 \(W_{2}\)
    • 这两个张量不再是并行的,需要在输入到 Dropout 层之前求和为 \(W=W_{1}+W_{2}\) (对应 Reduce 操作)
    • 但 Dropout 需要其输入在序列维度 \(s\) 上是并行的 (对应 Scatter 操作)
  • 论文不是先求和再在序列维度上并行化,而是将这两个操作合并为一个 Reduce-Scatter 操作
    • 因此,\(\bar{g}\) 在前向传播中可以是一个单一的 Reduce-Scatter 操作
  • 综上所述,论文得到:
    $$
    \begin{align}
    [Y_{1}^{s},Y_{2}^{s}] &=\text{LayerNorm}([X_{1}^{s},X_{2}^{s}]) \\
    Y &=g(Y_{1}^{s},Y_{2}^{s}), \\
    [Z_{1}^{h},Z_{2}^{h}] &=[\text{GeLU}(YA_{1}^{c}),\ \text{GeLU}(YA_{2}^{c})],\\
    W_{1} &=Z_{1}^{h}B_{1}^{r}\ \text{ and }\ W_{2}=Z_{2}^{h}B_{2}^{r},\\
    [W_{1}^{s},W_{2}^{s}] &=\bar{g}(W_{1},W_{2}),\\
    [V_{1}^{s},V_{2}^{s}] &=[\text{Dropout}(W_{1}^{s}),\ \text{Dropout}(W_{2}^{s})]
    \end{align}
    $$
  • 如果论文对反向传播进行类似的分解,论文会发现 \(g\) 和 \(\bar{g}\) 是彼此共轭的(Conjugate)
    • \(g\) 在前向传播中是 All-Gather,在反向传播中是 Reduce-Scatter;
    • \(\bar{g}\) 在前向传播中是 Reduce-Scatter,在反向传播中是 All-Gather
  • 对 Transformer 层的 LayerNorm 后接注意力部分进行类似的分解,得到图 5
  • TP 在单次前向和反向传播中需要四次 All-Reduce,而 TP 与 SP 结合在单次前向和反向传播中需要四次 All-Gather 和四次 Reduce-Scatter
    • 乍一看,似乎张量与 SP 比 TP 需要更多的通信
    • 但论文注意到,Ring All-Reduce 由两个步骤组成:
      • 一个 Reduce-Scatter 后接一个 All-Gather
    • 因此,TP 与张量加 SP 所使用的通信带宽是相同的
      • 问题:Reduce-Scatter 和 All-Gather 只有在无缝连着的时候才能说等价于 Ring All-Reduce 吧
      • 理解:是的,但是无论是否无缝连着,Reduce-Scatter + All-Gather 的通信都等于 Ring All-Reduce,因为都是 Reduce,然后 Scatter,继而 Gather
        • 补充:论文后面的 6.2 章节也提到了:reduce-scatter 和 all-gather 组合的执行速度比单独的 all-reduce 慢
    • 所以, SP 不会引入任何通信开销!
  • 根据公式 3, SP 与 TP 结合,将反向传播所需的所有激活沿并行维度划分
    • 注:第一个线性操作所需的张量 \(Y\))不会被划分
    • 为了缓解这个问题,论文不存储完整的张量 \(Y\) 用于反向传播,只在第 \(i\) 个 TP Rank 上存储 \(Y_{i}^{s}\) 部分,并在反向传播中执行一次额外的 All-Gather
    • 为了消除这个额外 All-Gather 引入的延迟,论文将此通信与计算 \(Y\) 梯度所需的计算重叠,从而减少了开销
  • 使用 SP 与 TP 结合,存储每个 Transformer 层激活所需的内存从公式 2 减少到:
    $$
    \begin{align}
    \text{Activations memory per layer} &= sbh\left(\frac{10}{t}+\frac{24}{t}+5\frac{as}{ht}\right) \\
    &=\frac{sbh}{t}\left(34+5\frac{as}{h}\right) \tag{4}
    \end{align}
    $$
  • 上述方程现在是公式 1 除以 TP 大小
    • 这意味着使用张量和 SP ,我们可以在 TP 组内分布激活,并将所需内存减少 TP 大小 \(t\) 倍
PP (Pipeline Parallelism)
  • PP 简单地将 Transformer 的 \(L\) 层划分为 \(\frac{L}{p}\) 组层,其中 \(p\) 是 PP 大小
    • 但 PP 并不会将激活所需的总内存均匀地除以 \(p\)
    • 这是由于 PP 调度引入了重叠(overlapping)以减少流水线气泡 (2021)
      • 注:overlap 意味着对应的激活需要同时在内存中
  • 为了量化这一点,论文考虑 PipeDream (2020) 中开发的 1F1B 流水线调度
    • 具有最小化流水线气泡的调度对流水线的第一个阶段施加最大的内存压力
      • 流水线的第一阶段指的是第一个 \(\frac{L}{p}\) 层组,也包括输入嵌入
    • 激活内存随流水线阶段变化的可视化图见附录 B
    • 为了保持流水线加压并避免额外的空闲时间,第一阶段必须存储 \(p\) 个微批次的激活(更多细节参见 (2021) 的图 4-top)
    • 每个阶段包含 \(\frac{L}{p}\) 层,因此第一阶段必须存储激活层为:
      $$p \times \frac{L}{p} = L$$
      • 注:(与 PP 大小 \(p\) 无关)
    • 因此,第一阶段存储激活所需的总内存为:
      $$\text{Total activations memory of 1st stage}=\frac{sbhL}{t}\left(34+5\frac{as}{h}\right) \tag{5}$$
  • 对于其他流水线调度,所需的总内存会略有不同
    • 例如,Megatron-LM (2021) 中开发的交错调度需要存储 \(L(1+\frac{p-1}{pm})\) 层的激活,其中 \(m\) 是交错阶段的数量
      • 因此,如果使用交错调度,则总激活内存应乘以 \((1+\frac{p-1}{pm})\)

Total Activations Memory

  • 总的激活内存的大部分由公式 5 提供,但该公式未包括输入嵌入、最后一个 LayerNorm 和输出层所需的激活内存,如图 2 所示
  • 位置和词嵌入不需要存储任何大量的激活用于反向传播,但 Dropout 需要存储
    • 嵌入层中的 Dropout 也沿序列维度并行化
    • 因此,它将需要 \(\frac{sbhp}{t}\) 的存储空间
    • 注意,因子 \(p\) 来自 PP 以及论文需要存储 \(p\) 个微批次的事实(见第 4.2.3 节)
  • 输出层之前的 LayerNorm 也使用 SP ,因此需要 \(\frac{2sbh}{t}\) 的存储空间
  • 输出层投影到词汇维度需要存储其输入,大小为 \(\frac{2sbh}{t}\)
  • 交叉熵损失需要存储以 32 位浮点数计算的 logits,因此需要 \(\frac{4sbv}{t}\) 的存储空间
  • 请注意,由于论文只考虑流水线第一阶段的激活,上述激活(即总共 \(\frac{4sbh}{t} \cdot (1+\frac{v}{h})\))仅在没有 PP (\(p=1\))的情况下被计入
  • 加上上述内存,由于输入嵌入、最后一个 LayerNorm 和输出层产生的额外内存为:
    $$\frac{sbhL}{t}\left(\frac{p}{L}+\delta_{p=1}\frac{4}{L}\left(1+\frac{v}{h}\right)\right)$$
    • 其中 \(\delta_{p=1}\) 是一个 0-1 函数,在 \(p=1\) 时为 1,否则为 0
  • 论文注意到,与公式 5 中的项 \(34+5\frac{as}{h}\) 相比,\(\frac{p}{L}\) 和 \(\frac{4}{L}\cdot (1+\frac{v}{h})\) 都是可以忽略的
    • 例如,对于一个具有 22B 参数的模型,这些额外项占总激活内存需求的比例小于 \(0.01\%\)
    • 问题:\(\frac{v}{h}\) 对不同配置也不算小吧,也有几十了
  • 因此,公式 5 是总所需激活内存的一个良好近似,论文将在论文的其余部分使用它

选择性激活重计算,Selective Activation Recomputation

  • 公式 (5) 给出的总激活内存需求对于大型模型而言仍然相当可观
  • 激活重计算 (2016) 通过存储(或 checkpoint)一组层的输入激活,并在反向传播期间使用额外的前向传递重新计算其他所需的激活,来克服这一内存限制(论文中称之为完全激活重计算)
  • 假设检查组仅包含单个层,并忽略 Transformer 层之外的激活,该方法将激活所需的总内存减少到
    $$ 2sbhL $$
    • 论文也注意到,如果论文只在每个 TP Rank 上存储部分激活,这个所需内存可以进一步减少到
      $$ \frac{2sbhL}{t} $$
      • 但这种方法需要每层额外的 all-gather 操作 ,会增加通信开销,因此论文不考虑这种方法
  • 与存储所有激活(公式 (5))相比,检查(checkpointing)所有 Transformer 层显著减少了训练模型所需的内存量
    • 这种减少确实是以重计算(一次额外的前向传递)为代价的,这可能会带来高达 \( 30-40\% \) 的计算时间开销
  • 为了平衡内存节省和计算开销,理想的做法是:仅仅 Checkpointing 足够的激活,使得给定的模型并行配置能够在设备内存的限制下进行训练
  • SP 提供的内存节省使得比以往更多的配置可以在不进行重计算的情况下进行训练,但大型模型的最佳模型并行配置通常仍然需要存储和重计算部分激活
  • 选择存储与重计算的激活量的一种简单方法是仅 checkpoint 部分 Transformer 层,并存储其他所有层的激活
    • 这种方法对于大型模型的扩展性并不好;
      • 例如,在训练 MT-NLG 时,每个设备只有三层,限制了平衡内存与计算的粒度
    • 此外,论文注意到并非所有激活都需要相同数量的操作来重新计算,因此在选择存储哪些激活和重计算哪些激活时更智能是有益的
  • 论文不是 checkpoint 和重计算完整的 Transformer 层,而是建议仅 checkpoint 和重计算每个 Transformer 层中占用大量内存但重新计算计算成本不高的部分 ,即选择性激活重计算(Selective Activation Recomputation)
  • 为此,论文注意到公式 (5) 中的项 \( \frac{5as}{h} \) 是由于在通过计算 Q、K 和 V 值的线性层增加了网络宽度之后的注意力操作(after the width of the network is increased by the linear layer calculating the Q, K, and V values);
    • 即如图 3 所示的 \( QK^{T} \) 矩阵乘法、softmax、softmax dropout 和注意力作用于 \( V \)
    • 这些操作通常具有大的输入尺寸,因此激活也很大,然而,每个输入元素的浮点操作数非常低
    • Transformer 层的其余部分对应公式 (5) 中的常数项 34
    • 因此,对于 \( \frac{5as}{h}>34 \) 的大型模型,如果论文 checkpoint 并重计算 Transformer 层的这一部分,论文存储的激活不到一半,并且重计算那些未存储的激活只需付出适度的成本
  • 为了量化这一点,让论文考虑 GPT-3 (2020) 和 MT-NLG (2022) 模型,这些是迄今为止训练过的一些最大模型
    • 对于 GPT-3,\( a=96 \),\( s=2048 \),\( h=12288 \),因此 \( \frac{5as}{h}=80 \)
    • 对于 MT-NLG,\( a=128 \),\( s=2048 \),\( h=20480 \),所以 \( \frac{5as}{h}=64 \)
  • 将这些数字与常数项 34(即该层其余部分的因子)进行比较,我们可以看到这些激活占总激活的很大一部分
    • 因此,通过使用选择性激活重计算,我们可以分别为 GPT-3 和 MT-NLG 模型节省 \( 70\% \) 和 \( 65\% \) 的激活内存需求
    • 这些激活的重计算仅为这两个模型引入了 \( 2.7\% \) 和 \( 1.6\% \) 的 FLOPs 开销
    • 有关 FLOPs 计算的更多细节,请参见附录 A
  • 随着序列长度的增加,序列长度与隐藏大小的比率 \( (\frac{s}{h}) \) 将增加,执行选择性部分重计算的相对成本也会增加,从而使选择性激活重计算更加昂贵
  • 作者认为,对于任何实际的 \( \frac{s}{h} \) 比率,选择性激活重计算仍然是平衡内存节省和计算的最佳重计算策略
  • 使用这种形式的选择性激活重计算,存储激活所需的内存从公式 (5) 减少到:
    $$
    \text{Total required memory}=34\frac{sbhL}{t}.
    $$
    • 上述公式表明,使用选择性激活重计算使得所需的激活内存随序列长度线性缩放,并且与注意力头的数量无关
    • 如第 4.2.3 节所述,在使用交错流水线调度的情况下,上述公式需要乘以 \( (1+\frac{p-1}{pm}) \)
  • 当使用 PP 时,如第 4.2.3 节所述,尽管给定设备只有 \( \frac{L}{p} \) 层,但第一阶段仍然必须存储相当于 \( L \) 层的激活,因为它必须存储 \( p \) 个微批次的激活以保持流水线压力
    • 在这种情况下,可以采用一种额外技术来降低重计算成本,即根据可用设备内存尽可能多地存储所有微批次的激活,并对其余部分进行完全或选择性重计算
    • 实际上,论文发现应用 SP 和选择性激活重计算后,重计算开销已经足够小,以至于这种额外技术带来的改进非常有限
    • 该技术在附录 C 中有更详细的描述和分析

Evaluations

  • 在本节中,论文评估了论文提出的方法对内存使用和训练执行速度的影响
  • 表 3 列出了评估中使用的模型配置
  • 论文考虑了参数规模高达一万亿的模型,对于所有这些模型,TP 大小设置为 8
  • 对于 175B 和 530B 模型,论文使用具有三个交错阶段 (\( m=3 \)) 的交错调度
  • 在所有情况下,序列长度设置为 \( s=2048 \),词汇表大小设置为 \( v=51200 \)
  • 论文还注意到,在这些评估中没有考虑数据并行,因为论文的方法与数据并行无关
  • 因此,论文分析中使用的批大小远低于端到端训练中使用的批大小
  • 论文的所有结果都是在 Selene 超级计算机上使用混合精度运行的
  • 每个集群节点有 8 个 NVIDIA 80GB A100 GPU,通过 NVLink 和 NVSwitch 相互连接
  • 每个节点有八个 NVIDIA Mellanox 200Gbps HDR Infiniband HCA 用于应用程序通信

Memory Usage

  • 论文通过借助前向和后向钩子在运行的不同时间点跟踪 GPU 活跃内存和 GPU 总内存(在 PyTorch 中可用)来验证内存消耗模型
  • 论文验证了测量的内存使用情况与论文的分析模型密切匹配
  • 表 2 的第一列总结了论文讨论的不同技术所需的内存
  • 为了量化这一点,图 7 显示了不同技术使用的激活内存,表示为在所有 TP Rank 上保持所有激活所需内存的百分比,即公式 (2)
  • 单独来看,两种技术都将内存需求减少了近一半,结合使用则提供了 5 倍的减少,将内存需求降至 20% 以下。这仅是完全激活重计算(基线水平的 10%)的约 \( \sim 2\times \)
  • 如果没有 SP 和选择性重计算共同提供的内存节省,这些模型都无法装入内存
    • 请注意,所有这些结果都包括了附录 B 中描述的内存优化

Execution Time per Layer

  • 表 2 的第二列和第三列总结了论文讨论的不同技术下每个 Transformer 层的计算量(FLOPs)和通信字节数
  • 为了量化这一点,表 4 显示了针对 22B 模型的各种实验下,执行一个 Transformer 层前向和后向传递的时间
    • 前两行显示, SP 对完成一个 Transformer 层所需的时间提供了适度的改进,将前向时间从 7.7ms 减少到 7.2ms,加速了 6%
      • 这种改进来自于 LayerNorm 和 Dropout 层在 \( \frac{1}{t} \) 的数据上执行
    • 论文还发现,尽管移动的数据量是相同的 ,但 reduce-scatter 和 all-gather 组合的执行速度比单独的 all-reduce 慢 ,这降低了 SP 带来的改进
      • 请注意,这种加速是使用 SP 主要优势之外的额外好处,其主要优势是节省内存,从而减少激活的重计算
  • 表 4 的后两行显示,如果论文在重计算哪些操作上具有选择性(由于 SP ,我们可以在更多配置中做到这一点),我们可以显著减少后向传递中重计算的开销
    • 选择性重计算的开销是 \( 1.3 \)ms,占 11.9ms 基线的 \( 11\% \),而重计算整个层的开销是 \( 7.6 \)ms 或 \( 64\% \)
    • 对于前向和后向总时间,开销是 \( 7\% \) 对比 \( 39\% \)
    • 请注意,重计算整个层的开销为 \( 39\% \)(而非预期的 \( 33\% \))是由于后向传递中的一项优化,论文将 all-reduce 通信与线性层权重的梯度计算重叠
    • 正如论文后面将看到的,这种好处随着模型大小的增加而增加
    • 表 4 的最后一行显示了选择性重计算和 SP 结合的好处:当两种技术一起使用时,开销降至仅 \( 4\% \)
  • 图 8 显示了所有测试案例的相同细分
    • 论文看到,随着模型规模的增长,开销的减少也在增加
    • 对于 \( 530 \)B 和 \( 1 \)T 的情况,开销仅为 \( 2\% \),而完全重计算的开销为 \( 36\% \)

End-to-End Iteration Time

  • 表 5 列出了表 3 中列出的四种配置中每一种的完整端到端迭代时间
    • 论文发现,对于所有测试的配置,论文提出的技术相比不使用 SP 的完全重计算,在吞吐量上提供了介于 \( 29.0\% \) 和 \( 32.1\% \) 之间的改进。这些节省将直接转化为更短的训练时间
  • 论文定义了模型 FLOPs 利用率(MFU)和硬件 FLOPs 利用率(HFU) ,定义方式与 Chowdhery 等 (2022) 类似
    • 模型 FLOPs 是执行单次前向和后向传递(单次迭代)所需的浮点操作数(不包括激活值),与实现方式和硬件限制无关
      • 因此,模型 FLOPs 是硬件和实现无关的,仅取决于底层模型
      • 注:但 模型 FLOPs 利用率(MFU)的分母是理论最大峰值 FLOPs,是与硬件有关的
    • 硬件 FLOPs 表示硬件上每次迭代实际执行的浮点操作
      • 因此,如果实现需要激活重计算(例如论文的实现),那么硬件 FLOPs 将大于模型 FLOPs
      • 理解:因为增加了一些激活重计算的操作,这里的硬件主要强调的是实际计算量(包括激活重计算的操作量)
      • 注:硬件 FLOPs 利用率(HFU)的分母是也是理论最大峰值 FLOPs
    • 论文在附录 A 中提供了模型和硬件 FLOPs 的严格下界公式
    • 对于论文的方法,硬件 FLOPs 与模型 FLOPs 的比率约为 \( 1+\frac{s}{18h} \)
  • 随后,论文分别将模型和硬件每秒 FLOPs 定义为模型和硬件 FLOPs 除以迭代时间
    • 使用这些定义,MFU 和 HFU 分别定义为模型和硬件每秒 FLOPs 除以加速器的理论峰值每秒 FLOPs
      $$
      \begin{align}
      MFU &= \frac{\text{observed throughput (tokens-per-second) }}{\text{theoretical peak FLOPs of device}} \\
      HFU &= \frac{\text{observed FLOPs}}{\text{theoretical peak FLOPs of device}}
      \end{align}
      $$
    • 以上定义参考自:PaLM: Scaling Language Modeling with Pathways, 2022, Google
  • 表 5 提供了所有四种配置的 MFU 和 HFU
    • 随着模型规模的增加,论文实现了更好的 GPU 利用率,对于一万亿参数模型,论文分别达到了 \( 56.3\% \) 和 \( 57.0\% \) 的 MFU 和 HFU
  • 虽然论文的分析中没有考虑初始化、评估、检查点等时间,但这些时间与迭代时间乘以迭代次数(这主导了端到端训练时间)相比可以忽略不计
    • 用于训练大型语言模型的迭代次数根据训练目标而有很大差异
    • 因此,作者认为论文报告的迭代时间(即吞吐量),包括了所有必要的操作,如数据加载和优化器步骤,是大型语言模型端到端训练时间的一个良好代理
    • 论文还发现迭代时间在整个训练过程中是一致的,因此对吞吐量进行采样可以准确衡量预期结果
  • 特别说明:表 5 中的结果没有使用任何数据并行
    • 数据并行会由于数据并行组之间所需的梯度 all-reduce 而引入一些开销
    • 但对于大型 Transformer 模型,这种开销并不大
      • 例如,如果论文将 530B 模型扩展到 8 路数据并行(2240 个 GPU),同时保持每个模型实例的批大小不变(即批大小也乘以数据并行大小)每次迭代的时间从 \( 37.83 \) 秒略微增加到 \( 39.15 \) 秒
      • 这导致 MFU 从 \( 56.0\% \) 下降到 \( 54.2\% \),下降幅度不大
    • 论文注意到论文没有使用任何梯度 all-reduce 与反向传播的重叠,而高效的重叠几乎可以完全消除迭代时间的增加

Future Work

  • 作者未来计划通过解决由大批次 (Large Microbatches) 导致的内存碎片化问题以及由 PP (Pipeline Parallelism) 导致的非均匀内存分配问题,来进一步减少激活内存
  • 作者准备研究能够减轻流水线第一级 (First Stage of the Pipeline) 内存压力的方法

附录:论文补充材料

  • 详情见论文补充材料:Supplementary Material of Megatron-LM-3

NLP——Megatron-LM原始论文解读-第二篇

注:本文包含 AI 辅助创作

  • 参考链接:
    • Megatron 系列目前公认的有三篇核心论文如下,它们分别对应了张量并行、3D 并行 与 序列并行/激活重计算优化 三大阶段,本文是第二篇的解读
    • 第二篇:(Megatron-LM-2)Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021, NVIDIA
      • 核心贡献:提出 3D 并行(数据 + 张量 + 流水线),并给出 interleaved 1F1B 流水线调度,显著降低流水线气泡;在 3072 块 A100 上训练出 530 B 参数的 GPT-3 级模型,GPU 利用率达到 76 %

Paper Summary

  • 整体说明:
    • 本文是 Megatron-LM原始论文解读-第二篇
    • 论文展示了如何将 PTD-P(节点间 PP、节点内 TP 和 DP )组合起来,以在训练具有万亿参数的大模型时实现高聚合吞吐量
    • 论文首次实现了端到端训练能够在合理的时间内完成(万亿参数模型的估计时间约为 3 个月)
    • 论文通过分析讨论了与每种并行类型相关的各种权衡,以及在组合使用时需要仔细考虑它们之间的相互作用
    • 论文中的很多思想是与加速器无关:
      • 1)智能地划分模型训练计算图以在保持设备活跃的同时最小化通信量的思想
      • 2)通过算子融合和仔细的数据布局来最小化内存受限的内核数量
      • 3)其他领域特定的优化(例如,散射-聚集优化)
  • 背景 & 问题:
    • 高效地训练 LLM 模型具有挑战性,原因在于:
      • 1)GPU 内存容量有限,使得即使是在多 GPU 服务器上也无法容纳大型模型;
      • 2)所需的计算操作数量可能导致不切实际的长训练时间
  • 之前的工作:
    • 之前已经提出了 TP (Tensor Parallelism, TP)和 PP (Pipeline Parallelism, PP),但这些方法的简单使用会在数千个 GPU 上导致扩展性问题
  • 本文的核心方法贡献:
    • 论文展示了如何组合使用 TP 、PP 和 DP (Data Parallelism, DP),以扩展到数千个 GPU
    • 论文提出了一种新颖的交错流水线调度方法,在内存占用与现有方法相当的情况下,可以将吞吐量提高 10% 以上
  • 特别亮眼的表现:论文的方法使得论文能够在 3072 个 GPU 上以 502 petaFLOP/s 的速度对具有 1 万亿参数的模型执行训练迭代(每 GPU 吞吐量达到理论峰值的 52%)
    • 注意:这是在 2021 年实现的,而本文的方法已经是家喻户晓

Introduction and Discussion

  • NLP 中基于 Transformer 的语言模型近年来推动了快速进展,因为大规模计算变得更加可用且数据集变得更大
  • 最近的工作 (2020) 表明,大语言模型是有效的零样本或少样本学习者,在许多 NLP 任务和数据集上具有高准确率
  • 这些大语言模型有许多令人兴奋的下游应用,例如客户反馈摘要、自动对话生成、语义搜索和代码自动补全 (2021;)
  • SOTA NLP 模型中的参数数量呈指数级增长(图 1)
  • 训练此类模型具有挑战性,原因有二:
    • (a) 即使使用最大的 GPU(NVIDIA 最近发布了 80GB A100 卡),也无法将这些模型的参数容纳在其主内存中;
    • (b) 即使论文能够将模型装入单个 GPU(例如,通过在主机和设备内存之间交换参数 (2021)),所需的大量计算操作可能导致不切实际的长训练时间(例如,使用单个 NVIDIA V100 GPU 训练具有 175B 参数的 GPT-3 (2020) 将需要大约 288 年)
  • 所以需要并行化, DP 扩展通常效果良好,但受到两个限制:
    • a)超过某个点后,每 GPU 批大小变得太小,降低了 GPU 利用率并增加了通信成本;
    • b)可使用的最大设备数量等于批大小,限制了可用于训练的加速器数量
  • 为了应对这两个挑战,已经提出了各种模型并行技术
    • 一些工作 (2019; 2020) 展示了如何使用张量(层内)模型并行(即将每个 Transformer 层内的矩阵乘法拆分到多个 GPU 上)来克服这些限制
    • 尽管这种方法对于在 NVIDIA DGX A100 服务器(配备 8 个 80GB A100 GPU)上训练大小达 20B 参数的模型效果良好,但对于更大的模型则会失效
    • 更大的模型需要拆分到多个多 GPU 服务器上,这会导致两个问题:
      • (a) TP 所需的全局归约(All-Reduce)通信需要通过服务器间链路进行,这些链路比多 GPU 服务器内可用的高带宽 NVLink (2020) 慢;
      • (b) 高度的模型并行可能会产生小的矩阵乘法(GEMM),可能降低 GPU 利用率
  • PP (2019; 2020; 2021) 是另一种支持训练大模型的技术,它将模型的层分布到多个 GPU 上
  • 一个批次被分割成更小的微批次(microbatch),执行过程在这些微批次之间进行流水线化
    • 层可以以各种方式分配(assigned)给工作节点,并且可以使用各种输入的前向传播和后向传播调度策略
    • 层分配(layer assignment)和调度策略会导致不同的性能权衡
  • 无论采用何种调度,为了保持严格的优化器语义,优化器步骤需要在设备间同步,这导致在每个批次结束时进行一次 流水线刷新(pipeline flush) ,此时允许微批次完成执行(并且不注入新的微批次)
    • flushing the pipeline 可能会花费多达 50% 的时间(取决于注入流水线的微批次数量)
    • 微批次数量与流水线规模(阶段数)的比率越大,花费在流水线刷新上的时间就越少
    • 为了实现高效率,通常需要更大的批大小
    • 在这项工作中,论文还引入了一种新的流水线调度方案,该方案提高了小批大小下的效率
  • 这些技术可以组合使用,但组合这些技术会导致 non-trivial 的相互作用,需要仔细推理才能获得良好的性能
  • 在论文中,论文解决了以下问题:
    • 在给定批大小并保持严格优化器语义的前提下,应如何组合并行技术以最大化大模型的训练吞吐量?
  • 具体来说,论文展示了如何组合 PP 、TP 和 DP,论文称之为 PTD-P(pipeline, tensor, and data parallelism) 的技术,以在数千个 GPU 上以良好的计算性能(达到设备峰值吞吐量的 52%)训练大语言模型
  • 论文的方法利用跨多 GPU 服务器的 PP 、多 GPU 服务器内的 TP 以及 DP 的组合,在实际中训练具有万亿参数的模型,并在具有服务器内和服务器间 GPU 高速链路的优化集群环境中实现优雅的扩展
  • 给定更多的训练资源,也可以使用类似的思想来训练更大的模型
  • 在论文的实验中,论文在一个具有万亿参数的 GPT 模型 (2020) 上,使用混合精度,展示了接近线性的扩展到 3072 个 A100 GPU,实现了每 GPU 163 teraFLOP/s 的端到端训练吞吐量(包括通信、数据处理和优化),以及 502 petaFLOP/s 的总吞吐量
    • 这个吞吐量使得实际的训练时间成为可能:论文估计该模型的端到端训练大约需要 3 个月
  • 作者相信这是针对该规模模型实现的最快训练吞吐量:过去的系统 (2019; 2020) 无法训练如此大的模型,因为它们没有结合 PP 和 TP
  • 论文还将论文的方法与 ZeRO (2019) 进行了比较,发现由于跨节点通信更少,论文的方法对于 175B 和 530B 参数的模型比 ZeRO-3 性能高出 70%
    • 这些模型太大,无法容纳在单个多 GPU 服务器上
  • 在大规模下实现这种吞吐量需要在多个方面进行创新和精心的工程:
    • 高效的内核实现使得大部分计算是计算受限而非内存受限的
    • 在设备间智能划分计算图以减少通过网络链路发送的字节数同时限制设备空闲时间
    • 领域特定的通信优化以及快速的硬件( SOTA GPU 以及服务器内和服务器间的高速链路)
  • 作者希望论文开源的软件(可在 github.com/nvidia/megatron-lm 获取)将使其他团队能够高效地大规模训练大型 NLP 模型
  • 论文通过实验和尽可能的分析,研究了影响吞吐量的各个组件之间的相互作用
  • 基于这些研究,论文提供以下关于如何配置分布式训练的指导原则 :
    • 不同形式的并行化以 non-trivial 的方式相互作用:
      • 并行化策略会影响通信量 和 内核执行的计算效率,还会影响工作节点上的气泡时间(因流水线刷新(气泡)所花费的空闲时间)
        • 例如,在论文的实验中,论文发现,即使服务器间有高带宽网络链路,次优的 TP 和 PP 组合也可能导致吞吐量降低多达 2 倍;
        • TP 在多 GPU 服务器内是有效的 ,但对于更大的模型必须使用 PP
    • 用于 PP 的调度方案会影响通信量、流水线气泡大小以及用于存储激活值的内存
      • 论文提出了一种新颖的交错调度方案,与先前提出的调度方案 (2019; 2020) 相比,在内存占用相当的情况下,可以将吞吐量提高多达 10%
    • 微批次大小等超参数的值会影响内存占用、工作节点上执行的内核的算术效率以及流水线气泡大小
      • 在论文的实验中,微批次大小的最优值是问题相关的,并且可以将吞吐量提高 15%
    • 在大规模下,分布式训练是通信密集型的
      • 当在 3072 个 GPU 上训练一个万亿参数模型时,论文的实现用于 PP 通信的有效二分带宽为 892 GB/s,用于 DP 通信的有效二分带宽为 13 TB/s
      • 使用较慢的节点间互连或通信更密集的划分会阻碍扩展性能
  • 论文不会自动探索并行策略的搜索空间(例如 FlexFlow (2018), PipeDream (2019), Tarnawski 等 (2020), 和 DAPPLE (2021)),而是提出论文在实践中发现效果很好的启发式方法(在 章节3 中)

Models of Parallelism

  • 论文关注单个 GPU 无法装下的大型模型进行高效训练的并行技术
  • 论文将 PP (pipeline model parallelism) 和 TP (tensor model parallelism)(组合如图 2 所示)与 DP (data parallelism) 结合起来(论文将其简称为 PTD-P)

DP (Data Parallelism)

  • 使用 DP (2021) 时,每个工作节点 (worker) 都拥有完整模型的一个副本,输入数据集被分片,工作节点定期聚合它们的梯度,以确保所有权重版本一致
  • 对于无法放在单个工作节点上的大型模型,可以在较小的模型分片上使用 DP

PP (Pipeline Model Parallelism)

  • 使用 PP 时,模型的层被分片到多个设备上
  • 当用于具有重复相同 Transformer 块的模型时,可以为每个设备分配相同数量的 Transformer 层
  • 论文不考虑更不对称的模型架构 ,因为将层分配给流水线阶段更为困难;论文将这个问题留给相关工作 (2021;) 来解决
  • 一个批次 (batch) 被分割成更小的微批次 (microbatch);然后跨微批次进行流水线执行
    • 流水线方案需要确保输入在前向传播和后向传播中看到一致的权重版本,以获得明确定义的同步权重更新语义
  • 具体来说,简单的流水线可能导致一个输入在后向传播中看到在前向传播中未见的权重更新
    • 为了精确保持严格的优化器语义,论文引入了周期性的流水线刷新 (pipeline flush),以便跨设备同步优化器步骤
    • 在每个批次的开始和结束时,设备处于空闲状态,论文将此空闲时间称为流水线气泡 (pipeline bubble),并希望使其尽可能小
  • 诸如 PipeMare、PipeDream 和 PipeDream-2BW (2019; 2020; 2021) 之类的异步和有界陈旧度 (bounded-staleness) 方法完全取消了刷新,但放宽了权重更新语义
    • 论文将对此类方案的考虑留待未来工作
  • 有多种可能的方式来跨设备调度前向和后向微批次;每种方法在流水线气泡大小、通信量和内存占用之间提供了不同的权衡
  • 论文将在本节讨论两种这样的方法
Default Schedule
  • GPipe 提出了一种调度方案,即首先执行一个批次中所有微批次的前向传播(如图 3 所示),然后是所有微批次的后向传播
  • 我们可以量化 GPipe 的流水线气泡大小 (\(t_{pb}\))
    • 用一个批次中的微批次数量记为 \(m\),流水线阶段数(用于 PP 的设备数量)记为 \(p\),每次迭代的理想时间为 \(t_{id}\)(假设完美或理想缩放),执行单个微批次的前向和后向传播的时间记为 \(t_{f}\) 和 \(t_{b}\)
    • 在此调度中,流水线气泡包括批次开始时的 \(p-1\) 次前向传播和批次结束时的 \(p-1\) 次后向传播
      • 理解:图 3 中,对于 PP=4 的场景,即每个设备 都要有 3 个完整的前向和后向传播气泡
    • 花费在流水线气泡中的总时间为:
      $$t_{pb}=(p-1)\cdot(t_{f}+t_{b})$$
    • 该批次的理想处理时间为:
      $$t_{id}=m\cdot(t_{f}+t_{b})$$
    • 因此,理想计算时间中花费在流水线气泡中的比例为:
      $$\text{Bubble time fraction (pipeline bubble size)} =\frac{t_{pb} }{t_{id} }=\frac{p-1}{m}.$$
      • 理解:这里分母上算少了,分母上应该再加上分子才对
  • 为了使气泡时间比例变小,一般需要 \(m\gg p\),但对于如此大的 \(m\),这种方法具有很高的内存占用,因为它需要在一次训练迭代的整个生命周期内,将所有 \(m\) 个微批次的暂存中间激活值(或者在使用激活重计算时,仅为每个流水线阶段的输入激活值)保存在内存中
  • 论文使用 PipeDream-Flush 调度 (2021)
    • 论文首先进入一个预热阶段,工作节点执行不同数量的前向传播,如图 4(顶部)所示
      • 此调度将进行中的微批次(其后向传播尚未完成且需要维护激活值的微批次)的数量限制为流水线的深度,而不是一个批次中的微批次数量
    • 预热阶段之后,每个工作节点进入稳定状态,每个工作节点执行一次前向传播,然后执行一次后向传播(简称 1F1B)
    • 最后,在批次结束时,论文完成所有剩余进行中微批次的后向传播
    • 这种新调度花费在气泡中的时间是相同的,但对于 PipeDream-Flush 调度,未完成的前向传播数量最多为流水线阶段数
    • 因此,此调度需要为 \(p\) 个或更少的微批次暂存激活值(相比之下,GPipe 调度需要为 \(m\) 个微批次暂存)
    • 因此,当 \(m\gg p\) 时,PipeDream-Flush 比 GPipe 内存效率高得多
Schedule with Interleaved Stages
  • 为了减小流水线气泡的大小,每个设备可以执行多个层子集(称为模型块, model chunk)的计算,而不是单个连续的层集
  • 例如,如果之前每个设备有 4 层(即设备 1 有第 1-4 层,设备 2 有第 5-8 层,依此类推),我们可以让每个设备执行两个模型块(每个有 2 层)的计算,即设备 1 有第 1, 2, 9, 10 层;设备 2 有第 3, 4, 11, 12 层;依此类推
    • 使用这种方案,流水线中的每个设备被分配多个流水线阶段(每个流水线阶段的计算量比之前少)
    • 像之前一样,我们可以使用此调度方案的“全部前向,全部后向”版本,但这具有很高的内存占用(与 \(m\) 成正比)
  • 论文开发了一种交错调度 (interleaved schedule) ,它适配了之前的内存高效 1F1B 调度
    • 这种新调度如图 4 所示,并要求一个批次中的微批次数量是 PP 度(流水线中的设备数量)的整数倍
    • 例如,对于 4 个设备,一个批次中的微批次数量必须是 4 的倍数
  • 如图 4 所示,相同批次大小的流水线刷新在新调度中发生得更早
    • 如果每个设备有 \(v\) 个阶段(或模型块),那么每个阶段或块的微批次的前向和后向时间现在将是 \(t_{f}/v\) 和 \(t_{b}/v\)
  • 因此,流水线气泡时间减少到
    $$t_{pb}^{\text{int.} }=\frac{(p-1)\cdot(t_{f}+t_{b})}{v}$$
    • 气泡时间比例则为:
      $$\text{Bubble time fraction (pipeline bubble size)} =\frac{t_{pb}^{\text{int.} } }{t_{id} }=\frac{1}{v}\cdot\frac{p-1}{m}.$$
  • 这意味着新调度将气泡时间减少了 \(v\) 倍
    • 但这种减小的流水线气泡大小并非没有代价:此调度需要额外的通信
    • 定量地说,通信量也增加了 \(v\) 倍
  • 在下一节中,论文将讨论如何利用多 GPU 服务器(例如 DGX A100 节点)中的 8 个 InfiniBand 网络卡来减少这种额外通信的影响

TP (Tensor Model Parallelism)

  • 使用 TP 时,模型的各个层在多个设备上进行划分
  • 在论文中,论文使用 Megatron (2019) 针对语言模型基础 Transformer 层所使用的特定划分策略
  • 也可以将类似的思想应用于其他类型的模型,如 CNN
  • 论文简要概述此策略,如图 5 所示
  • 一个 Transformer 层由一个自注意力块 (self-attention block) 和一个两层的多层感知机 (MLP) 组成
  • MLP 块:MLP 块包含两个 GEMM(通用矩阵乘法,General Matrix Multiply)和一个 GeLU 非线性:
    $$Y=\text{GeLU}(XA).\quad Z=\text{Dropout}(YB).$$
    • 第一个权重矩阵 \(A\) 按列划分,即 \(A=[A_{1}, A_{2}]\)
      • 这种划分允许将 GeLU 非线性独立应用于每个划分后 GEMM 的输出:
        $$[Y_{1},Y_{2}]=\left[\text{GeLU}(XA_{1}),\text{GeLU}(XA_{2})\right].$$
      • 这样做的好处是它消除了同步的需要(如果 \(A\) 沿其行划分则需要同步,因为 GeLU 是非线性的)
    • 第二个权重矩阵 \(B\) 的按行划分,以消除 GEMM 之间的任何通信需要(如图 4(a) 所示),如下所示:
      $$B=\begin{bmatrix}B_{1}\ B_{2}\end{bmatrix}, Y=[Y_{1},Y_{2}].$$
    • 第二个 GEMM 的输出然后在 Dropout 层之前跨 GPU 进行归约 (reduce)
  • 自注意力块:利用多头注意力操作中固有的并行来划分自注意力块(如图 4(b) 所示)
    • 键 (\(K\))、查询 (\(Q\)) 和值 (\(V\)) 矩阵可以以列并行 (column-parallel) 的方式进行划分(每个设备上负责不同的注意力头,并行计算)
    • 输出线性层然后可以直接在注意力操作的划分输出上运行(权重矩阵跨行划分,不需要 GPU 同步通信)
  • TP 将 MLP 和自注意力块中的 GEMM 拆分到多个 GPU 上,同时在前向传播中仅需要两次全归约 (All-Reduce) 操作(\(g\) 运算符),在后向传播中需要两次全归约操作(\(f\) 运算符)
  • 论文用几行代码实现了 \(f\) 和 \(g\)

Performance Analysis of Parallelization Configurations

  • 考虑将 PP (pipeline parallelism) 和 TP (tensor model parallelism) 与 DP (data parallelism) 结合使用的性能影响
  • 给定固定的 GPU 预算和批次大小 (batch size),可以使用 PTD-P 中不同维度的并行类型来训练模型;
  • 每个维度都在内存占用 (memory footprint)、设备利用率 (device utilization) 和通信量之间进行权衡
  • 论文将在本节剩余部分讨论这些权衡,并在 5.4 节展示实证结果
    • 论文在分析流水线气泡 (pipeline bubble) 大小时会提供分析模型
    • 论文定性地描述通信时间的行为,并提供通信量的成本模型;但是,论文不提供通信时间的直接成本模型,因为对于层次化网络拓扑(同一服务器上 GPU 之间的互连带宽高于服务器之间的互连带宽)来说,通信时间更难建模
  • 据论文所知,这是首个分析这些并行化维度性能交互作用的工作

Notation

  • 论文在本节中使用以下符号表示:
    • \( (p, t, d) \):并行化维度
      • \( p \) 表示 PP 大小 (pipeline-model-parallel size),\( t \) 表示 TP 大小 (tensor-model-parallel size),\( d \) 表示 DP 大小 (data-parallel size)
    • \( n \):GPU 数量,论文要求 \( p \cdot t \cdot d = n \)
    • \( B \):全局批次大小 (global batch size)(作为输入提供)
    • \( b \):微批次大小 (microbatch size)
    • \( m = \frac{1}{b} \cdot \frac{B}{d} \):每个流水线 (per pipeline) 中一个批次内的微批次数量

TP 和 PP(Tensor and Pipeline Model Parallelism)

  • TP 和 PP 均可以用于在多个 GPU 对模型的参数进行分区
  • 正如前文所述,使用带有周期性刷新(periodic flushes)的 PP 会产生下面大小的流水线气泡(pipeline bubble):
    $$\frac{p-1}{m}$$
  • 假设 DP 规模 \(d=1\)( DP 大小,data-parallel size),那么此时总 GPU 数量满足
    $$ t \cdot p = n; \quad \text{s.t.} \ d=1 $$
  • 基于 TP 规模 \(t\) 的流水线气泡大小可表示为:
    $$
    \frac{p-1}{m} = \frac{n/t - 1}{m}
    $$
    • 说明:在固定批次大小 \(B\)、微批次大小 \(b\) 以及固定 \(d\)(此时 \(m = \frac{B}{b \cdot d}\) 也保持固定)的前提下,随着 TP 规模 \(t\) 的增大,流水线气泡会逐渐减小
  • 不同 GPU 之间的通信量同样会受到 PP 规模 \(p\) 和 TP 规模 \(t\) 的影响
    • PP 的特点是点到点通信(point-to-point communication)成本更低;TP 则需要使用 All-Reduce communication
      • TP:正向传播(forward pass)需执行两次归约操作;反向传播(backward pass)也需执行两次归约操作(详见第2.3节)
    • 在 PP 中,对于每个微批次,每对 PP 相邻设备之间需要执行的PP 通信总量为(一次 点对点通信):
      $$ bsh $$
      • 其中 \(s\) 代表序列长度(sequence length),\(h\) 代表隐藏层大小(hidden size)
      • 正向传播和反向传播各需要一次点对点通信,所以总共是 $$ 2bsh $$
    • 在 TP 中,对于每个层,总大小为 \(bsh\) 的张量需在 \(t\) 个模型副本(model replicas)间,正向传播和反向传播各两次通信,这使得每个设备、每个微批次、每个层的 TP 通信总量为
      $$8bsh \cdot \frac{t-1}{t}$$
      • 理解:TP 参数通信量估计详情
        • MLP 正向和反向各需要一次 TP 通信(All-Reduce),一次 All-Reduce 通信是 \(2\Phi_\text{TP} = 2bsh\)(使用 Ring All-Reduce,一次通信使用 \(2\Phi\frac{t-1}{t}\) 的通信量),单层模型上, MLP TP 累计通讯量是:
          $$ 4\Phi\frac{t-1}{t} = 4bsh\frac{t-1}{t} $$
        • Attention 部分正向和反向也各需要一次 TP 通信(All-Reduce),单层模型上,Attention TP 累计通讯量是:
          $$ 4\Phi\frac{t-1}{t} = 4bsh\frac{t-1}{t} $$
        • 注:上述通信量评估详情见:图解大模型训练之:张量模型并行(TP),Megatron-LM - 猛猿的文章 - 知乎
    • 通常每个设备会负责多个层的计算(TP 中,每个层都需要分别与其他设备交互),因此对于每个设备、每个微批次, TP 的总通信量可表示为
      $$l^{stage} \cdot (8bsh \cdot \frac{t-1}{t})$$
      • 其中 \(l^{stage}\) 代表一个流水线阶段(pipeline stage)中包含的层数,即一个设备负责的层数
  • 由此可见,TP 会增加设备间的通信量
    • 因此,当 TP 规模 \(t\) 大于单个节点(node)中的 GPU 数量时,跨节点链路(inter-node links)的传输速度较慢,此时在这类链路上执行 TP 的开销会变得难以承受
    • 这一结果已在第5.4节的实验中得到验证
  • Takeaway #1 :TP 不跨 节点
    • 对于使用每个节点 \(g\) 个 GPU 的服务器时, TP 的规模通常应不超过 \(g\);若要在多台服务器间扩展以训练更大的模型,则可使用 PP
    • 理解:这里指的是 单台机器的 GPU 数为 \(g\),TP 规模 \(t\) 大于 \(g\) 时会导致跨节点的 TP 通信,开销难以接受

DP Data and Model Parallelism

  • 注:这里的模型并行(Model Parallelism)包括了 PP(Pipeline Model Parallelism) 和 TP(Tensor Model Parallelism)
  • 论文还需要考虑 DP (data parallelism)与两种模型并行( TP 和 PP )之间的相互作用
    • 为简化分析,本节将分别对这些相互作用进行讨论
PP (Pipeline Model Parallelism)
  • 假设 TP 规模 \(t=1\)(tensor-model-parallel size),则每个流水线的微批次数量为
    $$ m = \frac{B}{d \cdot b} = \frac{b’}{d}$$
    • 其中:
      • \( b \) 是微批次大小 (microbatch size)
      • \(b’ = \frac{B}{b}\) 表示每个 Global Step 的 微批次数量;
    • 注:每个流水线的微批次数量 \(m\) 与 DP 大小 \(d\) 成反比
  • 在总 GPU 数量为 \(n\) 的情况下,流水线阶段的数量为
    $$ p = \frac{n}{(t \cdot d)} = \frac{n}{d} $$
  • 此时流水线气泡大小可表示为:
    $$
    \frac{p-1}{m} = \frac{n/d - 1}{b’/d} = \frac{n - d}{b’}
    $$
    • 随着 DP 规模 \(d\) 的增大,\(n - d\) 会逐渐减小,因此流水线气泡也会随之变小
  • 图6展示了在不同 \(d\)、\(n\) 和 \(b’\)(批次大小与微批次大小的比值,\(b’ = \frac{B}{b}\))下,流水线气泡大小的变化趋势
    • 需要注意的是,并非所有模型都能将 \(d\) 增大到 \(n\),因为部分模型的完整训练内存占用(full training memory footprint)可能超过单个加速器(accelerator)的内存容量
  • 若 DP 所需的归约通信量未随 \(d\) 的增大而大幅增加,那么整体吞吐量(throughput)会随之提升
    • 这一假设是成立的,因为基于环形实现(ring-based implementation, 即 Ring All-Reduce)的通信时间与 \(\frac{d-1}{d} = 1 - \frac{1}{d}\) 呈正相关
  • 论文还可以分析批次大小 \(B\) 增大带来的影响
    • 在特定的并行配置下,随着批次大小 \(B\) 的增加,\(b’ = \frac{B}{b}\) 会增大,进而导致流水线气泡大小 \(\frac{n - d}{b’}\) 减小,最终使吞吐量提升
    • 此外,DP 所需的归约通信频率会随 \(B\) 的增大而降低,这也会进一步提升吞吐量
      • 理解:这里是指相同的数据量和 epoch 下,\(B\) 越大,需要更新的总步数越少?
DP Data and Tensor Model Parallelism
  • TP 中,每个微批次都需要执行归约通信,而跨多 GPU 服务器执行此类通信的成本较高
  • DP 仅需在每个批次执行一次高成本的归约通信
  • 而且,在 TP 中,每个模型并行进程(model-parallel rank)仅负责模型每层计算中的一部分;
    • 若层的规模不够大,现代 GPU 执行这些子矩阵(sub-matrix)计算时可能无法达到峰值效率(peak efficiency)
  • Takeaway #2 :
    • 在同时使用 DP 和模型并行时,应将模型并行的总规模设为
      $$M = t \cdot p$$
      • 其中 \(t\) 为 TP 规模,\(p\) 为 PP 规模,以确保模型的参数和中间元数据(intermediate metadata)能够放入 GPU 内存;
      • 问题:模型并行的总规模本来就是:\(M = t \cdot p\) 吧,还需要什么特殊处理吗?
    • 而 DP 则可用于扩展训练规模,以适配更多的 GPU

Microbatch Size

  • 微批次大小 \( b \) 的选择也会影响模型训练吞吐量
  • 例如,论文在图 7 中看到,在单个 GPU 上,使用较大的微批次大小,每 GPU 吞吐量最多可提高 1.3 倍
  • 本节回答:确定给定并行配置 \( (p, t, d) \) 和批次大小 \( B \) 时,最优微批次大小 \( b \) 是多少呢?
  • 无论微批次大小如何, DP 通信量将是相同的
  • 给定函数 \( t_{f}(b) \) 和 \( t_{b}(b) \) 它们将微批次大小映射到单个微批次的前向和后向计算时间,则处理一个微批次的总计算时间(忽略通信成本)为(如前所述,定义 \( b’ \) 为 \( \frac{B}{d} \)):
    $$ \left(\frac{b’}{b} + p - 1\right) \cdot \left(t_{f}(b) + t_{b}(b)\right). $$
    • 因此,微批次大小既影响操作的算术强度 (arithmetic intensity),也影响流水线气泡大小(通过影响 \( m \))
      • 理解:随着 microbatch size \(b\) 增大,气泡占比变小,处理一个微批次的总时间减少,但总微批次数量增加
    • 注:进一步分析总的时间为:
      $$ b \cdot \left(\frac{b’}{b} + p - 1\right) \cdot \left(t_{f}(b) + t_{b}(b)\right). $$
      • 显然,上述式子存在一个处于最大值和最小值中间的最优点(类似二次函数)
  • 图 8 显示了一个具有十亿参数且 \( (p, t) = (8, 8) \) 的 GPT 模型的估计吞吐量(使用方程 (1) 估计处理时间)
    • 对于两种批次大小,最优的 \( b \) 都是 4
  • Takeaway #3:
    • 最优微批次大小 \( b \) 取决于模型的吞吐量和内存占用特性,以及流水线深度 \( p \)、 DP 大小 \( d \) 和批次大小 \( B \)

激活重计算,Activation Recomputation

  • 激活重计算 (activation recomputation) (2016; 2000; 2019; 2020) 是一种可选技术,通过 仅存储给定流水线阶段的输入激活,而不是存储整个中间激活集,后者要大得多,来权衡增加执行的计算操作数量以换取更少的内存占用
    • 注:成本是后向传递之前需要按需再次运行前向传递
  • 为了在可接受的低内存占用下训练合理的大模型,需要使用激活重计算
  • 之前的工作如 PipeDream-2BW (2021) 已经研究了激活重计算的性能影响
    • 激活检查点 (activation checkpoint) 的数量不影响吞吐量,但影响内存占用
    • 设 \( A^{\text{input} } \) 为一层的输入激活大小,\( A^{\text{intermediate} } \) 为每层的中间激活大小
    • 如果一个模型阶段有 \( l \) 层,并且有 \( c \) 个检查点,则总内存占用将为
      $$ c \cdot A^{\text{input} } + \frac{l}{c} \cdot A^{\text{intermediate} } $$
      • \(\frac{l}{c} \cdot A^{\text{intermediate} }\) 为每个时刻在使用的激活大小
    • 当 \( c = \sqrt{l \cdot (A^{\text{intermediate} }/A^{\text{input} })} \) 时,该函数取得最小值
    • 在实践中,论文通过经验测量 \( A^{\text{intermediate} } \)
    • 在大多数情况下,每 1 或 2 个 Transformer 层设置一个检查点是最优的
  • 其他技术,如激活分区 (activation partitioning) (2020),也可以与 TP 结合使用,以进一步减少由激活引起的内存占用

Implementation

  • 论文将 PTD-P 作为 Megatron-LM 代码库的一个扩展来实现
  • 论文的实现基于 PyTorch (2017)
  • 论文使用 NCCL (2018) 进行设备间的通信
  • 为了获得良好的性能,论文实施了针对通信和计算的优化,下面将概述这些优化

Communication Optimizations

  • 当使用 PP 时,作者希望并行地发送和接收前向和后向传播的张量
  • 每个 DGX A100 节点配备了 8 个 InfiniBand (IB) 网络卡
    • 不幸的是,发送和接收是点对点的,并且只发生在两个服务器上的一对 GPU 之间,这使得很难在流水线内的单个通信调用中利用所有 8 张卡
  • 然而,我们可以利用同时使用 TP 和 PP 这一事实来降低跨节点通信的开销
  • 特别地,论文注意到每个 Transformer 层的输出在 TP Rank 之间是复制的(在 MLP 块中的 \(g\) 操作之后,见图 5(a))
    • 因此,在执行 TP 的两个连续流水线阶段中的 Rank ,会发送和接收完全相同的张量集合(图 5(a))
  • 对于足够大的模型,论文使用大小为 8 的 TP
    • 这意味着论文在相邻的多 GPU 服务器上的对应 GPU 之间发送相同的张量集合 8 次
    • 为了减少这种冗余,我们可以在发送端将张量分割成大小相等的块,然后只将一个块发送到下一个节点上的对应 Rank ,使用该 Rank 自己的 InfiniBand 卡(例如,在图 9 中, Rank 1 发送给 Rank 3, Rank 2 发送给 Rank 4)
  • 对于 8 个 TP Rank ,每个块的大小将减小为原来的八分之一
  • 然后,在接收端,我们可以通过 NVLink 执行一个 all-gather 操作(这比 InfiniBand 互连快得多)来重新构建完整的张量(如图 5(b) 所示)
  • 论文称之为 分散/聚集通信优化 (scatter/gather communication optimization)
    • 这种优化有助于更好地利用 DGX A100 服务器上的多个 IB 卡,并使诸如交错调度这样通信密集的调度变得可行
  • 量化来看,通过分散-聚集通信优化,每对连续阶段之间需要执行的总通信量减少到 \(\frac{bsh}{t}\),其中 \(t\) 是 TP 大小,\(s\) 是序列长度,\(h\) 是隐藏层大小(在论文的实验中 \(t=8\))

Computation Optimizations

  • 论文对计算图实施了三个模型特定的优化以获得高性能
  • 第一,论文改变了 Transformer 层中的数据布局,以避免内存密集的转置操作,并启用跨步批处理 GEMM 内核
    • 具体来说,论文将数据布局从 \([b, s, a, h]\) 改为 \([s, b, a, h]\),其中 \(b\)、\(s\)、\(a\) 和 \(h\) 分别是批处理大小、序列长度、注意力头数和隐藏层大小的维度
  • 第二,论文使用 PyTorch JIT (2016) 为一序列逐元素操作(偏置 + GeLU 以及偏置 + Dropout + 加法)生成了融合内核
  • 第三,论文创建了两个自定义内核来实现缩放、掩码和 Softmax(归约)操作的融合:
    • 一个支持通用掩码(用于如 BERT 的模型),另一个支持隐式因果掩码(用于如 GPT 的自回归模型)
  • 论文将在下一节量化这些优化的效果

Evaluation

  • 在本节中,论文试图回答以下问题:
    • PTD-P 的性能如何?它是否能够满足现实的端到端训练时间?
    • 对于给定的模型和批处理大小, PP 的扩展性如何?交错调度对性能有多大影响?
    • 不同的并行化维度如何相互作用?微批处理大小等超参数的影响是什么?
    • 分散-聚集通信优化的影响是什么?在规模上运行训练迭代时,论文对硬件施加了哪些类型的限制?
  • 论文所有的结果都是在 Selene 超级计算机 (Selene, 2000) 上使用混合精度运行的
    • 每个集群节点有 8 个 NVIDIA 80-GB A100 GPU (2018),通过 NVLink 和 NVSwitch (2018) 相互连接
    • 每个节点有八个 NVIDIA Mellanox 200Gbps HDR InfiniBand HCA 用于应用通信,另外每个节点还有两个 HCA 用于专用存储
    • 节点通过具有 850 个交换机的三级(叶子、脊柱、核心)胖树拓扑连接
    • 这种拓扑支持高效的 All-Reduce 通信(深度学习训练中的主导通信模式)
    • 集群使用全 NVME 共享并行文件系统进行高性能数据访问和存储
  • 一个具有 16 位精度的 A100 GPU 的峰值设备吞吐量为 312 teraFLOP/s
  • 对于论文的大部分结果,论文报告每个 GPU 的吞吐量
  • 总吞吐量可以通过乘以使用的 GPU 数量来计算
  • 在论文的实验中,论文使用适当大小的 GPT 模型
    • 特别地,对于任何给定的微基准测试,模型需要能够适应实验中使用的模型并行 GPU 的数量
    • 适当的时候,论文使用标准的模型架构,如 GPT-3 (2016)

End-to-End Performance

  • 论文考虑了论文的系统在参数量从十亿到一万亿的 GPT 模型上的端到端性能,使用了张量、流水线和 DP (使用第 3 节中描述的启发式方法选择维度)
  • 特别地,论文使用了启用了分散/聚集优化的交错流水线调度
  • 所有模型使用词汇表大小(表示为 \(V\))为 51,200(1024 的倍数)和序列长度(表示为 \(s\))为 2048
    • 论文改变隐藏层大小(\(h\))、注意力头数和层数(\(l\))
    • 模型中的参数量 \(P\) 可以计算为:
      $$P=12lh^{2}\left(1+\frac{13}{12h}+\frac{V+s}{12lh}\right). \tag{2}$$
  • 随着模型大小的增加,论文也增加批处理大小(\(B\))和 GPU 的数量(\(n\))
  • 模型中大部分的浮点运算是在 Transformer 层和 logit 层中的矩阵乘法(GEMMs)中执行的
    • 仅考虑这些 GEMMs,每次迭代的 FLOPs 数量为(更多细节见附录):
      $$F=96Bslh^{2}\left(1+\frac{s}{6h}+\frac{V}{16lh}\right). \tag{3}$$
    • 这是真实 FLOP 计数的下限,但应接近实际值
    • 论文将 FLOP 计为浮点运算,无论精度如何
    • 方程 (3) 假设了激活重计算,并考虑了与额外前向传播相关的浮点运算
  • 表 1 显示了模型配置以及实现的 FLOP/s(包括每个 GPU 的和所有 GPU 的总和)
    • 论文看到在 3072 个 A100 GPU(384 个 DGX A100 节点)上实现了超线性扩展,因为随着模型变大(更大的矩阵乘法),GPU 利用率提高,而通信时间相对于计算时间没有显著增加
    • 吞吐量是针对端到端训练测量的,即包括所有操作,包括数据加载、优化器步骤、通信和日志记录
    • 对于最大的模型,论文达到了峰值设备吞吐量的 52%,对于最小的模型,达到了峰值设备吞吐量的 44%
训练时间估算
  • 给定这些吞吐量,论文还可以估算在 \(T\) 个 Token 上进行端到端训练所需的总时间
    • 训练需要 \(I=\frac{T}{B\cdot s}\) 次迭代(\(s\) 为序列长度)
    • 使用方程 (3) 中的 \(F\) 值和表 1 中的经验端到端吞吐量(表示为 X),我们可以估算总训练时间
    • 对于表 1 中的配置,论文有 \(12lh \gg (V+s)\) 和 \(16lh \gg V\)
  • 将这些观察结果与方程 (2) 和 (3) 结合,论文得到:
    $$\text{End-to-end training time}\approx\frac{8TP}{nX}.$$
    • 其中:
      • \(n\) 为 GPU 数量;
      • \(X\) 为每个 GPU 的吞吐量;
      • \(T\) 为 Token 数量
      • \(P\) 为模型参数量
  • 让论文以具有 \(P=\) 175B 参数的 GPT-3 模型为例
    • 该模型在 \(T=3000\) 亿个 Token 上进行了训练
      • 在 \(n=1024\) 个 A100 GPU 上使用批处理大小 1536,论文实现了每个 GPU \(X=140\) teraFLOP/s 的吞吐量
      • 因此,训练该模型所需的时间为 34 天
    • 对于 1 万亿参数模型,论文假设端到端训练需要 450B 个 Token
    • 使用 3072 个 A100 GPU,我们可以实现每个 GPU 163 teraFLOP/s 的吞吐量,端到端训练时间为 84 天
    • 作者认为这些训练时间(使用合理数量的 GPU)是可行的

Comparison to ZeRO-3

  • 论文在表 2 和图 10 中将 PTD-P 与 ZeRO-3 (2020; 2021) 进行了比较(针对标准的 GPT-3 模型架构以及表 1 中的 530B 参数模型)
    • 这些结果提供了一个与不使用模型并行的方法的对比点
    • 论文使用 DeepSpeed Python 库 (2020) 将 ZeRO 集成到论文的代码库中
    • 当论文增加 GPU 数量时,论文保持全局批处理大小不变
    • 在 GPU 数量较少且微批处理大小为 4 的情况下,PTD-P 对于 175B 和 530B 参数模型的吞吐量分别高出 6% 和 24%
    • 随着论文增加 GPU 数量,PTD-P 比单独使用 ZeRO-3 扩展得更优雅(见图 10)
      • 例如,通过将 GPU 数量加倍(保持批处理大小不变),由于跨节点通信更少,PTD-P 对两个模型的性能均优于 ZeRO-3 70%
    • 论文只考虑了不使用 TP 的 ZeRO-3
    • ZeRO-3 可以与模型并行结合,以潜在地改善其扩展行为

PP (Pipeline Parallelism)

  • 论文现在单独评估 PP 的弱扩展性能,并将非交错调度与交错调度的性能进行比较
Weak Scaling
  • 论文使用弱扩展设置、一个具有 128 个注意力头和隐藏层大小为 20480 的 GPT 模型以及微批处理大小为 1 来评估默认非交错 PP 调度的扩展性
  • 随着论文增加流水线阶段的数量,论文也按比例增加模型中的层数来增加模型的大小
    • 例如,当 PP 大小为 1 时,论文使用具有 3 个 Transformer 层和 15B 参数的模型,当 PP 大小为 8 时,论文使用具有 24 个 Transformer 层和 121B 参数的模型
  • 论文对所有配置使用 TP 大小为 8,并改变使用的 A100 GPU 总数,从 8 到 64
  • 图 11 显示了两种不同批处理大小下每个 GPU 的吞吐量,以说明流水线气泡的影响,其行为符合 \(\frac{p-1}{m}\)(第 2.2.1 节)
    • 正如预期的那样,较高的批处理大小扩展性更好,因为流水线气泡被分摊到更多的微批处理上
Interleaved versus Non-Interleaved Schedule
  • 图 12 显示了在具有 175B 参数(96 层,96 个注意力头,隐藏层大小为 12288)的 GPT-3 (2020) 模型上,交错和非交错调度每个 GPU 的吞吐量
    • 带有分散/聚集通信优化的交错调度比非交错(默认)调度具有更高的计算性能
    • 随着批处理大小的增加,这个差距会缩小,原因有二:
      • (a) 随着批处理大小的增加,默认调度中的气泡大小减小
      • (b) 流水线内的点对点通信量与批处理大小成正比,因此随着通信量的增加,非交错调度会赶上(交错调度每个样本的通信量更多)
    • 在没有分散/聚集优化的情况下,默认调度在较大批处理大小下表现优于交错调度(未显示)

Comparison of Parallel Configurations

  • 在本小节中,论文展示了结合不同并行化维度所带来的各种权衡
  • 特别地,论文展示了对于给定模型和多个批处理大小,使用相同数量 GPU 的并行配置的性能
TP 与 PP (Tensor versus Pipeline Parallelism)
  • 论文评估了对于给定模型和批处理大小,流水线和 TP 对性能的影响
  • 图 13 中的实证结果显示了结合使用张量和 PP 来训练一个 161B 参数 GPT 模型(32 个 Transformer 层以支持 PP 大小为 32,128 个注意力头,隐藏层大小为 20480)的重要性,以实现低通信开销和高计算资源利用率
  • 论文观察到
    • TP 在一个节点(DGX A100 服务器)内效果最好,因为其 All-Reduce 通信开销大
    • PP 使用成本低得多的点对点通信,可以在节点之间进行,而不会成为整个计算的瓶颈
    • 使用 PP 时,可能会在流水线气泡中花费大量时间:
      • 因此,流水线阶段的总数应受到限制,使得流水线中的微批处理数量是流水线阶段数量的合理倍数
    • 当 TP 大小等于单个节点中的 GPU 数量(对于 DGX A100 节点为 \(8\))时 ,达到了峰值性能
  • 这一结果表明,无论是单独使用 TP (由 Megatron (2019) 使用)还是单独使用 PP (由 PipeDream (2019) 和其他人使用),都无法与结合使用这两种技术的性能相媲美
PP 与 DP (Pipeline versus Data Parallelism)
  • 论文在图 14 中评估了数据和 PP 对一个具有 59 亿参数(32 个 Transformer 层,32 个注意力头,隐藏层大小为 3840)的 GPT 模型性能的影响
  • 论文使用比之前更小的模型,因为论文想展示当模型并行大小仅为 \(2\) 时模型能够容纳的性能
  • 为简单起见,论文在这些实验中保持微批处理大小等于 \(1\)
  • 论文看到,对于每个批处理大小,吞吐量随着 PP 大小的增加而降低,这与论文在第 3.3 节中的分析模型相符
  • PP 应主要用于支持训练无法容纳在单个工作器上的大模型,而 DP 应用于扩展训练规模
TP 与 DP (Tensor versus Data Parallelism)
  • 论文在图 15 中评估了数据和 TP 对同一个 59 亿参数 GPT 模型性能的影响(使用较小模型的原因同上)
  • 如前所述,论文最初保持微批处理大小等于 \(1\)
  • 在较大的批处理大小和微批处理大小为 1 的情况下, DP 通信不频繁;
  • TP 中所需的 all-to-all 通信需要对批次中的每个微批处理执行
  • 这种 TP 的 all-to-all 通信主导了端到端训练时间,特别是当通信需要在多 GPU 节点之间执行时
  • 随着 TP 大小的增加,论文在每个 GPU 上执行更小的矩阵乘法,降低了每个 GPU 的利用率
  • 尽管 DP 可以导致高效的扩展,但不能单独使用 DP 来处理具有有限训练批处理大小的非常大的模型,原因是
    • a) 内存容量不足
    • b) DP 的扩展限制
      • 例如,GPT-3 是以 1536 的批处理大小训练到收敛的
      • 因此, DP 仅支持扩展到 1536 个 GPU;然而,大约使用了 \(10,000\) 个 GPU 来在合理的时间内训练该模型

Microbatch Size

  • 论文在图 16 中评估了微批处理大小对结合流水线和 TP 的并行配置性能的影响(针对一个具有 91B 参数(\((t,p)=(8,8)\))的模型)
  • 论文看到:
    • 对于这个模型,最佳微批处理大小是 \(2\);
    • 对于其他模型,最佳微批处理大小是不同的(图中未显示)并且是模型依赖的
  • 对于给定的批处理大小,
    • 增加微批处理大小会减少流水线中的微批处理数量(\(m\)),导致更大的流水线气泡;
    • 但增加微批处理大小也可以通过增加执行内核的算术强度来提高 GPU 利用率
    • 以上这两个因素是相互矛盾的,这使得选择最佳微批处理大小具有挑战性
  • 论文来自第 3.3 节的分析模型合理地近似了真实性能,并且可以作为代理来确定如何为各种训练配置和模型选择此超参数值

Activation Recomputation

  • 图 17 显示了对于具有 145B 参数(80 个 Transformer 层,96 个注意力头,隐藏层大小为 12288)的 GPT 模型,在使用 128 个 A100 GPU、\((t,p)=(8,16)\) 以及一系列批处理大小的情况下,使用和不使用激活重计算的吞吐量
    • 对于小批处理大小,由于在反向传播期间需要执行额外的前向传播,激活重计算导致吞吐量(以每秒序列数计)降低高达 33%
    • 但激活重计算是支持更大批处理大小所必需的
  • 由于流水线气泡更小,使用激活重计算的大批处理大小的吞吐量比不使用激活重计算(对于较小批处理大小)实现的最佳吞吐量高出 up to 2\(\times\)

Scatter-Gather Optimization

  • 图 18 显示了对于具有 175B 参数的 GPT-3 模型,使用和不使用(未优化)分散/聚集通信优化时每个 GPU 的吞吐量
  • 论文看到,通过减少跨节点链路上的通信量,对于通信密集的调度(大批处理大小带交错),吞吐量提高了 up to 11%

Fused Operators(融合操作)

  • 论文还评估了第 4.2 节中描述的操作符融合对性能的影响
  • 对于 GPT-3 模型(175B 参数),通过融合,吞吐量提高了 19%(从每个 GPU 113 teraFLOP/s 到每个 GPU 135 teraFLOP/s)
  • 对于更大的 530B 参数 GPT 模型(图 1 中的模型配置),吞吐量提高了 11%(从每个 GPU 133 teraFLOP/s 到每个 GPU 148 teraFLOP/s)

Inter-Node Communication Bandwidth

  • 论文强劲的结果是优化软件和硬件栈共同使用的副产品
  • 特别地,论文利用了同一服务器内和跨服务器的 GPU 之间的高带宽通信链路
  • 在具有 3072 个 GPU 的万亿参数模型上,论文观察到流水线阶段之间点对点通信的有效二分带宽为 892 GB/s,而 DP 副本之间 All-Reduce 操作的有效二分带宽为 12.9 TB/s
  • 跨设备的操作符分区若优化不足,会导致更多的节点间通信,从而阻碍扩展性能

Checkpoint Loading and Saving

  • 训练大模型的一个重要实际考虑是加载和保存模型检查点,对于论文中考虑的模型,检查点尤其大
  • 例如,万亿参数模型的检查点大小为 13.8 TB
    • 问题:如果仅考虑参数,换算后不太对,1000000000000*4/1000/1000/1000/1000 = 4 TB,若考虑激活值才会到 12 TB
  • 所有 384 个节点(3072 个 GPU)对万亿参数模型的检查点初始加载达到了 1TB/s 的峰值读取带宽,这是并行文件系统可能的最大读取吞吐量
  • 检查点保存达到了峰值写入带宽的 40%(273 GB/s)

Related Work

Parallelism for Large Models

  • PP(Pipeline model parallelism)是一种用于训练大模型的常用技术
  • PP 有几种不同的模式:
    • 论文讨论的模式使用流水线刷新(flush)来确保严格(strict)的优化器语义
    • TeraPipe (2021) 为像 GPT 这样的自回归模型(auto-regressive models)在单个训练序列的 token 之间暴露了细粒度的 PP
    • PipeTransformer (2021) 通过冻结具有“稳定”权重的层,弹性地调整 PP 和 DP 的程度,并将资源专用于训练剩余的“活跃”层
    • HetPipe (2020) 在一组异构加速器上结合使用了流水线和 DP
    • PP 也可以使用松弛的语义来实现:
      • PipeDream-2BW (2021) 维护两个权重版本,并保证权重更新延迟为 1,而无需昂贵的流水线刷新;
      • PipeMare (2021) 和 Kosson 等人 (2021) 使用异步 PP
    • 与论文考虑的带有流水线刷新的技术相比,这些技术提高了吞吐量,但可能以收敛速度或最终精度为代价
    • 此外,单独的 PP 仍然只能扩展到与模型中层数相等数量的设备,这对于某些模型架构来说是有限的
  • PipeDream (2019) 以一种原则性的方式结合了 PP 和 DP ,以减少跨设备通信
    • DeepSpeed (2021) 将 PP 与 TP 和 DP 结合起来,以训练高达万亿参数规模的模型,但吞吐量低于论文所示的结果(峰值利用率的 52% 对比 36%),原因有几个:
      • 通过算子融合使大部分算子图保持计算受限(compute-bound)、使用更高效的 PP 调度以最小化流水线气泡(pipeline bubble)大小、快速的硬件(A100 与 V100 GPU 以及同一服务器内和不同服务器间 GPU 的高带宽链路)以及扩展到更多 GPU 的能力
    • 作者希望强调,这种更高的吞吐量使得估计的训练时间更加实用(约 3 个月);37.6 petaFLOP/s 的总吞吐量将需要大约 40 个月来训练一个同等规模的模型
      • 也可以扩展到更大的模型,但需要更多的 GPU 来保持训练时间的实用性
  • Mesh-TensorFlow (2018) 提出了一种语言,用于轻松指定结合数据和模型并行的并行化策略
    • Switch Transformers (2021) 使用 Mesh-Tensorflow 训练了一个具有 1.6 万亿参数的稀疏激活(sparsely activated)的基于专家(expert-based)的模型,其预训练速度比 T5-11B 模型 (2019) 有所提高

Sharded Data Parallelism(分片 DP)

  • 作为 MLPerf 0.6 (2019) 性能优化的一部分,引入了分片 DP (sharded data parallelism)(2019, 2020),其中优化器状态(optimizer state)在 DP 工作节点(data-parallel workers)上进行分片
  • 这种方法有两个优点:
    • (a) 它不会在原始 DP (vanilla data parallelism)的基础上引入额外的通信;
    • (b) 它将优化器的计算和内存成本分摊到 DP 分区中
  • ZeRO (2019, 2021) 扩展了这个思想:
    • 权重参数和梯度也在 DP 工作节点上进行分片,工作节点在执行计算前从其“拥有”相应状态的工作节点获取相关状态
      • 这增加了额外的通信,但可以通过仔细重叠计算和通信来部分隐藏
      • 但如果不使用 TP 或者批大小不够大以隐藏额外的通信开销,这会变得困难(图 10)
  • ZeRO-Infinity (2021) 使用 NVMe 来高效地交换参数,使得能够在少量 GPU 上训练非常大的模型
  • 论文注意到,使用少量 GPU 训练非常大的模型会导致不切实际训练时间(例如,收敛需要数千年)

Automatic Partitioning(自动分区)

  • FlexFlow (2018)、PipeDream (2019)、DAPPLE (2021) 和 Tarnawski 等人 (2020) 都借助成本模型(cost models)在多个设备上自动分区模型训练图
  • 然而,这些方法都没有考虑论文考虑的所有并行维度:流水线和 TP、DP、微批大小(microbatch size)、以及像激活重计算(activation recomputation)这样的内存节省优化对训练大于加速器内存容量的模型的影响
  • 这些增加的维度扩大了需要探索的搜索空间
  • Gholami 等人 (2018) 展示了如何对数据和模型并行组合的通信成本进行建模

高性能计算(HPC for Model Training)

  • Goyal 等人 (2017) 和 You 等人 (2018) 都展示了使用高性能计算(High Performance Computing, HPC)技术在几分钟内训练高精度 ImageNet 模型
  • 但所考虑的图像分类模型可以轻松地放在单个加速器上,使得模型并行变得不必要;
  • 它们支持非常大的批大小(\(>32k\)),允许将 DP 扩展到大量工作节点,且通信不频繁;
    • 并且它们由紧凑的卷积层组成,这些层本身就很适合 DP 通信

NLP——LLM-as-a-judge

注:本文包含 AI 辅助创作

  • 参考链接:
    • 原始论文:From Generation to Judgment: Opportunities and Challenges of LLM-as-a-judge, arXiv 20250206, UC Berkeley
    • 论文主页:llm-as-a-judge.github.io/
    • 作者给出的论文集合(作者会持续更新):github.com/llm-as-a-judge/Awesome-LLM-as-a-judge

Paper Summary

  • 核心内容:
    • 本综述很长,核心内容探讨了 LLM 作为评判的复杂性
    • 作者根据输入格式( Pointwise 、 Pairwise 和 Listwise )和输出格式(评分、排名和选择)对现有 LLM-based 评判方法进行分类
    • 然后,论文为 LLM 作为评判提出了一个全面的分类法,涵盖评判属性、方法论和应用
    • 在此之后,论文介绍了 LLM 作为评判的基准的详细集合,并伴随着对当前挑战和未来方向的深思熟虑的分析,旨在为这一新兴领域的未来工作提供更多资源和见解
  • 问题提出:
    • 评估(Assessment)与评价(Evaluation) 长期以来一直是 AI 和 NLP 领域的关键挑战
    • 传统方法(基于匹配 or 基于 Embedding),往往在评判(Judging)细微属性和提供令人满意的结果方面存在不足
  • LLM 的最新进展启发了“LLM-as-a-judge”范式,即利用 LLM 在各种任务和应用中执行评分、排序或选择
  • 论文对 LLM-based 评判和评估进行了全面综述
    • 首先从输入和输出的角度给出详细定义
    • 然后引入一个全面的分类法(taxonomy),从 what to judge, how to judge and where to judge(评判什么,如何评判,在哪里评判) 三个维度探讨 LLM-as-a-judge 的问题

Introduction and Discussion

  • 评估与评价长期以来一直是机器学习和自然语言处理(NLP)中必不可少但颇具挑战性的任务,尤其是对给定候选列表的各种属性(如质量、相关性和有用性)进行评分和比较(2022;2024)
  • 传统评估方法依赖于静态指标,如 BLEU(双语评估辅助工具)(2002)和 ROUGE(摘要评价指标)(2004)
    • 这些指标通过计算输出与参考文本之间的词重叠来衡量质量
    • 这些自动指标计算效率高,并在许多生成应用中得到使用(2022, 2023a, 2024c)
    • 但是,它们对 n-gram 匹配(n-gram matching)和基于参考设计(reference-based designs)的依赖严重限制了它们在动态和开放式场景中的适用性(2016;Reiter, 2018)
  • 随着深度学习模型(2019;2019)的兴起,许多基于 Embedding 的评估方法(如 BERTScore(2020)和 BARTScore(2021))也应运而生
    • 尽管这些基于小模型的指标从词级表示转向了 Embedding 级表示,并提供了更大的灵活性,但它们仍然难以捕捉相关性之外的细微属性(2024),如有用性和无害性
  • 最近先进的 LLM,如 GPT-4(2023)和 o1,在指令遵循、查询理解和响应生成方面表现出了惊人的性能
    • 这一进展促使研究人员提出了“LLM-as-a-judge ”的概念(2023),即利用强大的 LLM 对一组候选进行评分、排序和选择(图 1)
    • LLM 的强大性能(2020)与精心设计的评估 Pipeline (2023a;2024b;2023a)相结合,能够对各种评估应用进行细粒度和详细的评判,显著解决了传统评估方法的局限性,为 NLP 评估设定了新的标准
  • 除了评估之外,LLM-as-a-judge 还被广泛应用于整个 LLM 生命周期,包括对齐(2022;2023)、检索(2023;2024c)和推理(2023;2024b)等任务
    • 它赋予 LLM 一系列高级能力,如自我进化(2024a)、主动检索(2024c)和决策(2023),推动它们从传统模型向智能代理(2024)转变
  • 然而,随着 LLM-as-a-judge 的快速发展,评判偏差(judging bias)和脆弱性(vulnerability)等挑战(2023a;2024a;2024)也逐渐显现
    • 因此,对当前技术和未来挑战进行系统回顾对于推进 LLM-based 评判方法具有重要价值
  • 在本综述中,论文深入探讨了 LLM-as-a-judge 的细节,旨在全面概述 LLM-based 评判
    • 论文首先通过讨论其各种输入和输出格式(第 2 节)对 LLM-as-a-judge 进行正式定义
    • 接下来,论文提出了一个深入而全面的分类法,以解决三个关键问题(第 3-5 节):
      • 属性(Attribute) :评判什么? 论文深入研究 Judge LLM 评估的具体属性,包括帮助性、无害性、可靠性、相关性、可行性和整体质量
      • Methodology :如何评判? 论文探索了 LLM-as-a-judge 系统的各种调优和提示技术,包括手动标记数据、合成反馈、监督微调、偏好学习、交换操作、规则增强、多智能体协作、演示、多轮交互和比较加速
      • 应用(Application) :在哪里评判? 论文调查了 LLM-as-a-judge 已被应用的领域,包括评估、对齐、检索和推理
  • 此外,论文在第 6 节中收集了从不同角度评估 LLM-as-a-judge 的现有基准
  • 最后,论文在第 7 节中提出了当前的挑战和未来研究的有前途方向,包括偏差与脆弱性、动态与复杂评判、自我评判和人机共判

Differences from Other LLM-related Surveys

  • 近年来,LLM 已成为热门研究领域,涌现出诸多相关综述文献(2023b;2024;2024a)
  • 尽管已有若干综述专注于 LLM-based 自然语言生成(NLG)评估(2024a;2024n),但本研究旨在系统梳理”LLM-as-a-judge”这一方法论
  • 如前所述,”LLM-as-a-judge”的应用范畴已远超评估领域,因此有必要建立全局视角进行归纳与分类
  • 此外,现有综述多聚焦 LLM 赋能的具体应用场景
    • 如 LLM-based 数据标注(2024b)、数据增强(2024e)以及自我修正(2024)
    • 但目前仍缺乏针对”LLM-as-a-judge”范式的系统性与全面性综述研究

Preliminary

  • 在本节中,论文旨在提供 LLM-as-a-judge 的详细定义,分别在 2.1 节和 2.2 节讨论各种输入和输出格式

Input

  • 给定一个 Judge LLM \(J\),评估过程可以表示为:
    $$R = J\left(C_{1}, \ldots, C_{n}\right). \tag{1}$$
    • 这里 \(C_{i}\) 是第 \(i\) 个待评判的候选,\(R\) 是评判结果
  • 在本节中,论文根据候选数量 \(n\) 将输入格式分为两种类型:
    • ** Pointwise (Point-Wise)** :当 \(n = 1\) 时,成为逐点评判,此时 Judge LLM 将只关注一个候选样本(2023b)
    • 成对/ Listwise (Pair/List-Wise) :当 \(n \geq 2\) 时,成为成对(\(n = 2\))或 Listwise (\(n > 2\))评判,此时提供多个候选样本供 Judge LLM 比较并进行综合评估(2023;2024)

Output

  • 在本节中,论文根据 \(R\) 的不同格式讨论评判的三种输出类型:
    • 分数(Score) :当每个候选样本被分配一个连续或离散的分数时,\(R = \{C_{1}: S_{1}, \ldots, C_{n}: S_{n}\}\),这成为基于分数的评判
      • 这是最常见和广泛使用的协议,利用 LLM 进行评分以进行定量比较(2024a)或属性检测(2024a)
    • 排名(Ranking) :在基于排名的评判中,输出是每个候选样本的排名,表示为 \(R = \{C_{i} > \ldots > C_{j}\}\)
      • 这种比较方法在需要建立候选之间排名顺序的场景中很有用(2023a)
    • 选择(Selection) :在基于选择的评判中,输出涉及选择一个或多个最优候选,表示为 \(R = \{C_{i}, \ldots, C_{j}\} > \{C_{1}, \ldots, C_{n}\}\)
      • 这种方法在决策(2023a)或内容过滤(2024c)场景中特别有用

Attribute

  • 在本节中,论文从属性角度对当前 LLM-as-a-judge 的研究进行分类。图 3 概述了 Judge LLM 可以评估的方面

Helpfulness

  • 现代 SOTA LLM 经过指令调整和对齐过程,能够遵循用户指令并有效响应
  • 这一对齐步骤依赖于大量有用和无害的数据,通常作为人类偏好数据收集,然后用于对齐训练的强化学习
  • 鉴于获取此类对齐数据的成本很高,最近的工作探索了使用 LLM 来标记帮助性,以及生成或评估对齐数据(2022)
  • (2024)的作者以在线方式使用 LLM 来获取偏好 ,以直接对齐另一个 LLM
    • 最近的一些工作表明,来自 AI 的帮助性反馈,即 LLM 的反馈与人类反馈相当(2023)
    • 也有一些成功的 LLM(2024a)通过 AI 反馈数据进行了微调,从而证明了这种方法的可行性和有用性
  • 除了这些对齐方面的工作,利用 LLM 作为评估器的通用框架在评估候选响应的帮助性方面也至关重要(2023;2023;2024e)

Harmlessness

  • 评估文本数据的无害性对于内容审核和创建或管理合成数据集都很重要
  • 鉴于人类标注工作既昂贵又耗时,并且受先前研究中 Embedding 在 LLM 中的道德信念的启发(2023),许多最近的工作研究了使用 LLM 来评估无害性
  • SOTA LLM 能够有效地用于内容审核,要么在遵循一些政策指导方针的情况下现成使用,要么在安全/不安全数据上进行微调(2023;2024g)
  • Ye 等人 (2024b) 探索了使用 LLM 以细粒度方式评估无害性等属性的可行性,并发现专有模型的表现比开源模型好得多
  • Wang 等人 (2024l) 使用 OpenAI 的 GPT-4 来评估无害性,并进一步将其性能与针对该特定任务微调的较小预训练语言模型进行比较。此外,Bai 等人 (2022) 使用原则来指导 LLM 进行无害性评估以用于对齐目的,他们将这种范式称为“宪法 AI”(Constitutional AI)
  • (2023)使用相同的 LLM 来评估其响应是否有害,并为 GPT-3.5 和 Llama-2 提供了见解和最佳实践
  • Xie 等人 (2024a) 在 LLM 安全拒绝基准上对几个 LLM 进行了全面比较,发现小型 LLM 在微调设置下是有效的安全评判
  • 在推理时,(2023b) 提出了可回退自回归推理(RAIN),允许 LLM 进行自我评估和回退以确保 AI 安全

Reliability

  • 可靠性是指,LLM 能够生成事实性和忠实的内容 ,同时也能表达不确定性或承认对某些主题的知识空白(2024)
    • 关于事实性,Wang 等人 (2024a) 引入了 HALU-J,这是一个基于批判的幻觉评判(critique-based hallucination judge),通过选择相关证据并提供详细批判来增强事实性评估
  • Cheng 等人 (2023) 设计了一种使用 GPT-4 的自动评估方法,用于评判模型的输出是否是幻觉
  • 此外,一些工作采用 Judge LLM 进行长文本事实性评估
  • 在对话上下文中,Luo 等人 (2024) 收集了一个大规模的基准,用于自动对话级幻觉评估
    • 基于此数据集,他们引入了 HalluJudge,一个专门用于评估对话级幻觉的评判语言模型
  • Min 等人 (2023) 开发了 FactScore,一种细粒度的方法,用于评估长文本生成的事实性,首先将内容拆分为原子级句子,然后从维基百科检索相关语料库以评估其事实性
    • 在此基础上,Wei 等人 (2024b) 提出为 Judge LLM 配备 Google 搜索 API,以实现更灵活和高效的事实性评估
  • Jing 等人 (2024) 将这种细粒度的可靠性评估扩展到多模态领域,并引入了 FaithScore
    • Zhang 等人 (2024h) 在 LLM 对齐中采用了类似的策略,创建了一个合成对齐数据集,其中涉及使用声明提取和自我评判技术对每个生成的样本进行评估和过滤
  • Xie 等人 (2024b) 训练了一个外部基于批判的 LLM-as-a-judge,在生成阶段提供声明级的事实性反馈,以提高响应的事实性
    • 对于不确定性评判,Xu 等人 (2024c) 提出了 SaySelf,这是一个新的训练框架,通过自我一致性提示和基于组的校准训练,教 LLM 表达更细粒度的置信度估计(1-46, 1-50)

Relevance

  • 相关性衡量生成或检索的内容与原始查询的一致程度
  • 传统的相关性评估方法通常依赖于关键字匹配(2009)或语义相似度(2021),这些方法难以捕捉上下文中的细微差异或细微差别
  • 使用 Judge LLM 进行相关性评估已被探索并验证为在各种应用中更细粒度和有效的方式(2023;2024;2024b)
  • 在对话评估中,Lin 和 Chen (2023a) 首先提出用 LLM 评判来取代昂贵且耗时的人类注释,用于相关性评估,为 Judge LLM 提供对话上下文和生成的响应进行评估
  • 类似地,Abbasiantaeb 等人 (2024) 将 LLM-as-a-judge 应用于对话搜索,与人类注释者合作解决与不完整相关性评判相关的问题。在检索增强生成(RAG)场景中,Li 和 Qiu (2023) 利用 LLM 来确定哪些历史记忆最适合解决当前问题
  • 遵循这一概念,Li 等人 (2024c) 还提议采用 LLM 作为重新排序器,以评判和过滤子知识图中的噪声和不相关知识
    • 最近,LLM-as-a-judge 也被用于多模态应用中的相关性评判(2025;2024i;2024b;2024;2024b;2024;2024m;2024)
  • Yang 和 Lin (2024) 以及 Chen 等人 (2024a) 都采用多模态 LLM 来构建自动评估基准
    • 相比之下,Chen 等人 (2024f) 使用多模态奖励模型来评估文本到图像生成中的相关性
    • 在 SQL 生成评估中,(2023a) 提出利用 LLM 来确定 SQL 等价性
  • 此外,LLM-as-a-judge 还在许多传统检索应用中探索了相关性评估,如搜索(2024)、检索(2024a)和推荐(2024)

Feasibility

  • 复杂且设计良好的推理 Pipeline 可以进一步释放 LLM 的潜力
  • 在这些代理式 LLM 中,评估候选动作或步骤的可行性对于规划、推理和决策的成功至关重要
  • 虽然一些工作利用指标或外部工具进行这种可行性评估(2023a;Yuan 等人),但许多其他工作利用 LLM 本身来选择最适当和合理的动作来执行
  • Hao 等人 (2023) 首先提出提示 LLM 进行自我评估,并生成可行性评判作为奖励信号来执行蒙特卡洛树搜索(MCTS)
    • 类似地,Yao 等人 (2023a) 建议在其提出的“思维树”(ToT)框架中采用 LLM 作为状态评估器,用于潜在步骤搜索
  • Besta 等人 (2024) 将先前研究中使用的树结构替换为图结构,并使用 LLM 根据每个思维的可行性或正确性为其分配分数
  • 在多智能体协作系统中,Liang 等人 (2023)、Li 等人 (2024b) 和 Wang 等人 (2024d) 提议利用 Judge LLM 在多个候选响应中选择最可行和合理的解决方案
  • 此外,还有一些工作采用 Judge LLM 在 API 选择(2024b)、工具使用(2023)和 LLM 路由(2024)中进行可行性评估

Overall Quality

  • 如前所述,LLM-as-a-judge 可用于在各种任务和应用中执行多方面和细粒度的评估
    • 然而,在许多情况下,仍然需要一个综合评估来表示候选的整体质量,以进行比较或排名
    • 获得此整体分数的一种直接方法是计算特定方面分数的平均值或加权平均值(2023;2023b;Saad-2024)
  • 此外,其他一些研究呈现每个属性的评估结果,并提示 LLM 评判生成整体质量评判(2024d;2024;2024a)
    • 除了从多个属性总结整体质量外,其他研究专注于直接产生整体评判
  • 例如,在摘要(2023b;2023a;2024a;2024d,h)和机器翻译(2024;2023;2024)等传统 NLP 任务中,与更开放式的长文本生成任务相比,评估维度较少
    • 因此,在这些情况下,LLM-as-a-judge 通常用于直接提示 LLM 生成整体评判

Methodology

  • 在本节中,论文介绍 LLM-as-a-judge 常用的方法和技巧,将其分为调优方法(4.1 节)和提示策略(4.2 节)

Tuning

  • 为了增强通用 LLM 的评判能力,不同研究采用了各种调优技术
  • 在本节中,论文从两个角度讨论 LLM-as-a-judge 的调优方法:
    • 数据来源(4.1.1 节)
    • 训练方法(4.1.2 节)
  • 表 1 收集了所有专注于调优 Judge LLM 的研究论文
Data Source
  • 手动标注数据(Manually-labeled Data) :为了训练具有类似人类标准的 Judge LLM,一种直观的方法是收集手动标注的样本和相应的评判
    • 许多先前的工作已经利用并整合现有资源来构建全面的数据集,用于调优 Judge LLM
    • Vu 等人(2024)构建了一个大型且多样化的集合,包含超过 100 个质量评估任务,涵盖超过 500 万个人类评判,这些评判通过整理和标准化先前研究中公开的人类评估得到
    • 类似地,Wang 等人(2024k)提出了 PandaLM,并收集了多样化的人类标注测试数据,其中所有上下文均为人类生成,标签与人类偏好一致
    • 为了增强策略在对齐数据合成中的评判能力,Lee 等人(2024)通过成对评判任务扩充了 SFT数据集 ,该任务的指令是从一组选项中选择所选响应
  • 还有一些工作收集数据集用于细粒度的评判反馈(2024)
    • Xu 等人(2023)引入了 InstructScore,一种可解释的文本生成评估指标,并策划了 MetricInstruct 数据集,该数据集涵盖六个文本生成任务和 23 个数据集
    • Liu 等人(2024a)收集了 ASPECTINSTRUCT,这是第一个专为多方面 NLG 评估设计的指令调优数据集,跨越 65 个任务的 27 个不同评估方面
    • Yue 等人(2023)首先提出了属性评估,并使用来自相关任务(如问答、事实核查、自然语言推理和摘要)的数据对 Judge LLM 进行微调
    • Ke 等人(2024)采用了一种独特的方法,首先提示 GPT-4 生成反馈,并手动检查其为每个用户查询生成的文本,必要时进行修订以提高质量
  • 合成反馈(Synthetic Feedback) :尽管手动标注的反馈质量高且能准确反映人类评判偏好,但其在数量和覆盖范围上存在局限性
    • 一些研究人员将合成反馈作为调优 Judge LLM 的数据源。这一方向的一种方法依赖于 Judge LLM 自身生成合成反馈
      • Wu 等人(2024a)通过提示策略 LLM 评估自己的评判来构建成对反馈,以增强评判能力
      • Wang 等人(2024i)提示 LLM 生成原始指令的“noisy”版本 ,并将对该损坏指令的相应响应作为较差响应
      • Wang 等人(2024a)提示 GPT-4-Turbo 为每个实例基于原始证据生成多条证据,将其分类为完全不相关证据、部分不相关证据和高度相关证据,以训练幻觉评判 LLM
  • Park 等人(2024a)构建了 OFFSETBIAS,这是一个成对偏好数据集 ,利用 GPT-4 生成不良、离题和错误的响应 ,并进行难度过滤
  • 对于安全评判,Xie 等人(2024a)采用 GPT-4 作为分类器,将每个数据点映射到预定义的安全类别,以训练自动化评估器
  • 与先前工作不同,Li 等人(2024e)采用 GPT-4 合成成对和逐点数据,以训练生成式 Judge LLM
  • 对于逐点数据,他们采用“分而治之”的策略,从 GPT-4 收集单个响应的两条批评意见,将其合并为更全面的批评,并提供最终评分
  • 紧随其后,Kim 等人(2024b)使用 GPT-4 用详细的人类评估标准和口头反馈扩充偏好学习数据集
  • 在多模态领域,Xiong 等人(2024b)提出了 LLaVA-Critic,并采用 GPT-4o 生成给定分数或偏好评判背后的原因,用于训练数据构建
  • 此外,通过利用人类显式指令和 GPT-4 的隐式知识,Xu 等人(2023)基于 LLaMA 微调了一个 Judge LLM,为生成的文本同时生成分数和人类可读的诊断报告
  • Zhu 等人(2023)引入了 JudgeLM,并提出了一个全面、大规模、高质量的数据集,包含任务种子、LLM 生成的答案和 GPT-4 生成的评判,用于微调高性能评判
Tuning Techniques
  • SFT : SFT 是最常用的方法,用于帮助 Judge LLM 从成对(2024k;2024e;2023b;2023)或逐点(2024a;2023b;2023)评判数据中学习
    • 在采用 SFT 的众多工作中,Vu 等人(2024)提出了一种监督多任务训练方法,用于在各种任务的多个混合数据集上微调其基础大型自动评分模型(FLAMe)
    • 为了使 Judge LLM 同时具备成对和逐点评判能力,Kim 等人(2024b)在调优阶段新颖地提出了联合训练和权重合并方法,并发现后者在大多数情况下并未提高评估性能
    • 为了获得不仅能生成响应还能比较成对偏好的评判模型,Lee 等人(2024)设计了 Judge 增强监督微调( Judge-augmented Supervised Fine-tuning,JSFT),使用扩充的偏好学习数据集
    • 在训练阶段,Ke 等人(2024)通过添加简化提示来区分输入的不同部分,并通过交换两个生成文本的顺序和交换批评中的相应内容来扩充成对训练数据,从而增强他们的模型
    • Xu 等人(2023)进一步在自我生成的输出上微调其 INSTRUCTSCORE 模型,以优化反馈分数,生成与人类评判更好对齐的诊断报告
    • Liu 等人(2024a)还提出了两阶段监督微调方法,首先应用普通指令调优,使模型具备遵循指令进行各种评估的能力,然后使用辅助方面进行进一步调优,以丰富训练过程,纳入额外的指令调优阶段,利用与目标评估方面的潜在连接
  • 偏好学习(Preference Learning) :偏好学习与评判和评估任务密切相关,尤其是比较和排名评判
    • 除了直接采用或扩充偏好学习数据集用于监督微调 Judge LLM 的工作外,一些研究还应用偏好学习技术来增强 LLM 的评判能力
    • 为了提高 HALU-J 提供的评判质量,Wang 等人(2024a)在多证据设置下的 SFT 阶段后,使用定向偏好优化(DPO)(2023)进一步对其进行调优
    • 类似地,Park 等人(2024a)将 DPO 与包含关键错误但具有评判模型偏好的风格质量的合成“bad”响应一起应用,帮助减轻 Judge LLM 中的偏差
    • Wu 等人(2024a)新颖地提出了元奖励,利用策略 LLM 来评判自己评判的质量,并生成成对信号以增强 LLM 的评判能力
    • 这一概念也被 Wang 等人(2024i)采用,他们提出了自教评估器,使用损坏的指令生成次优响应作为偏好学习的劣质示例
    • 最近,Hu 等人(2024)提出了 Themis,一种专注于 NLG 评估的 LLM,其训练采用了设计的多视角一致性验证和面向评分的偏好对齐方法
    • Li 等人(2024o)提出了 PORTIA,一种基于对齐的方法,旨在以有效方式模拟人类比较行为,以校准位置偏差

Prompting

  • 在推理阶段设计适当的提示策略和 Pipeline 可以提高评判准确性并减轻偏差
  • 在本节中,论文总结和分类 LLM-as-a-judge 的现有提示策略(图 4)
Swapping Operation
  • 先前的研究表明, LLM-based 评判对候选的位置敏感,候选响应的质量排名很容易通过仅改变其在上下文中的顺序来操纵(2023c;2024;2023)
    • 为了减轻这种位置偏差并建立更公平的 LLM 评判系统,交换操作(2023)已被引入并广泛采用
    • 该技术涉及调用 Judge LLM 两次,在每个实例中交换两个候选的顺序
    • 在评估中,如果交换后的结果不一致,则标记为“平局(tie)” ,表明 LLM 无法自信地区分候选的质量(2023)
  • 几项研究还在自我对齐(2023;2024a;2024)中纳入了交换操作,以从 Judge LLM 获得更准确的成对反馈
  • Zhu 等人(2024a)提出了一种类似思维链(CoT-like)的提示技术,通过要求模型首先提供所有成对排名,然后用排名列表进行总结来减轻位置偏差
Rule Augmentation
  • 规则增强提示(Rule-augmented prompting)是指将一系列原则、参考依据和评估准则直接嵌入到 judge LLM 的提示指令中的方法
  • 这种方法通常用于 LLM-based 评估,其中引导 Judge LLM 评估特定方面(2024e;2023a;2024d;2024),并为其提供详细的标准(2023b;2024g;2024;2024l,h;2024c)以确保公平比较
  • Liu 等人(2024d)采用了一种独特的方法,提示 Judge LLM 通过对一组少样本示例的上下文学习来生成自己的评分标准。为了与 LLM-as-a-judge 保持一致,Bai 等人(2022)首先提出引入一系列原则(如帮助性、无害性、诚实性),使 Judge LLM 能够更精确和有方向地比较两个候选
  • 紧随其后,后续工作(2023, 2024;2024;2024a;2024a)通过为原则或标准的每个方面纳入更详细的解释来增强这种基于原则的提示
  • 此外,Li 和 Qiu(2023)以及 Li 等人(2024c)都提出提示 LLM 根据候选在解决特定问题中的帮助性来检索适当的演示/知识三元组
  • 为了从 LLMs 获得多样化的响应,Lahoti 等人(2023)提示多个 LLM 评判每个候选的多样性,并选择最多样化的一个进行进一步优化
  • Zhang 等人(2024g)提出了 RevisEval,其利用 LLM 的自我纠正能力自适应地修改响应,然后将修改后的文本作为后续评估的原则
  • 最近,一些工作研究了 LLM 作为个性化评判的可靠性(2024;2024;2024),将 persona 作为原则的一部分提供给 LLMs,以进行个性化评判
Multi-agent Collaboration
  • 由于 LLM 固有的各种偏差(2023c;2024;2023;2023a;2023c),访问单个 LLM 评判的结果可能不可靠
  • 为了解决这一限制,Li 等人(2023a)引入了 Peer Rank(PR)算法,该算法考虑每个对等 LLM 对所有答案对的成对偏好,并生成模型的最终排名
  • 在此基础上,出现了几种多智能体 LLM 的架构和技术,包括智能体混合(2023b)、角色扮演(2023)、辩论(2023;2024f;2024b)和投票(2024c;2024)。Jung 等人(2024)提出了级联选择性评估,其中成本较低的模型作为初始评判,仅在必要时升级到更强的模型(2024a)
  • 此外,一些工作将多智能体协作应用于对齐数据合成,利用多个 LLM 评判来完善响应(2024)或提供更准确的成对反馈(2024i)
  • 最近,(2024m)提出了 MATEval,其中所有智能体均由 GPT-4 等 LLM 扮演
    • MATEval 框架模仿人类协作讨论方法,整合多个智能体的交互以评估开放式文本
Demonstration
  • 上下文中的样本或演示(2020;2023;Agarwal 等人)为 LLM 提供了可遵循的具体示例,并已被证明是 LLM 上下文学习成功的关键因素
  • 几项研究引入了人类评估结果作为 LLM-as-a-judge 的演示,旨在引导 LLM 从几个具体的上下文中学习评估标准
  • Jain 等人(2023b)是第一个探索大型语言模型作为多维评估器在上下文学习中的功效的研究,无需大量训练数据集
  • Kotonya 等人(2023)对各种提示技术进行了系统实验,包括标准提示、基于注释器指令的提示和思维链提示,将这些方法与零样本和一样本学习相结合,以最大化评估效果
  • 为了提高 LLM 评估的鲁棒性,Hasanbeig 等人(2023)提出了 ALLURE,一种通过迭代纳入显著偏差的演示来增强评估器鲁棒性的方法
  • 此外,Song 等人(2024)引入并研究了两种多样本上下文学习(ICL)提示,使用两种版本的多样本 ICL 模板来帮助减轻 LLM 中的潜在偏差
Multi-turn Interaction
  • 在评估中,单个响应可能无法为 LLM 评判提供足够的信息来全面和公平地评估每个候选的表现
  • 为了解决这一限制,多轮交互通常被采用以提供更全面的评估
  • 通常,该过程从初始查询或主题开始,随后在 Judge LLM 和候选模型之间进行动态交互
  • Bai 等人(2023b)提出了一种多轮设置,其中评估器扮演面试官的角色,根据模型的先前答案提出越来越复杂的后续问题
  • 类似地,Yu 等人(2024d)引入了 KIEval,一种基于知识的交互式评估框架,其新颖地纳入了一个 LLM 驱动的交互器,以实现动态、抗污染的评估
  • 此外,一些方法促进候选之间的多轮辩论
    • 例如,Zhao 等人(2024c)设计了一个框架,其中两个 LLM 围绕查询进行多轮 peer battle ,使它们的真实性能差异得以显现
    • Moniri 等人(2024)提出了一个自动化基准系统,其中 LLM 进行辩论,最终评估由另一个 LLM 评判执行
Comparison Acceleration
  • 在 LLM-as-a-judge 的各种比较格式中(如 Pointwise 和 Listwise ),成对比较是直接比较两个模型或生成成对反馈的最常用方法
  • 然而,当需要对多个候选进行排名时(注:需要凉凉进行排列),这种方法可能非常耗时(2024)
  • 为了减轻计算开销,Zhai 等人(2024)提出了一种排名配对方法,所有候选首先与空白基线响应进行比较 ,然后根据每个候选与基线的比较表现确定其排名
  • Zhu 等人(2024a)提出了一种类似 CoT 的提示技术,通过强制模型首先提供所有成对排名,然后用列表总结这些成对排名来减轻位置偏差
  • 此外,Lee 等人(2024)在推理期间利用基于 tournament 的方法(2023a;2023c)进行拒绝采样,以加速成对比较
  • 他们构建了一个 tournament 树,其中叶节点表示采样的响应,非叶节点根据子节点之间的评判结果进行选择

Application

  • 尽管“LLM-as-a-judge ”最初是为评估应用而提出的,但其使用范围已大幅扩展到许多其他场景,如对齐、检索和推理
  • 因此,如图 5 所示,论文将全面介绍 LLM-as-a-judge 如何应用于各种领域

Evaluation

  • 自然语言处理中的传统评估依赖于预定义标准,通常通过指标来评估机器生成文本的质量
  • 一些著名的指标,如 BLEU、ROUGH 和 BERTScore 已在该领域广泛使用
  • 然而,基于指标的评估过度强调词汇重叠和相似性,当需要考虑许多有效响应和更细微的语义属性时,可能会显得不足(Post, 2018;2022)
  • 为解决这些局限性,LLM-as-a-judge 已被用于充当自动化评判,以增强许多任务的评估(2023b;2024)
  • LLM-as-a-judge 能够进行类似人类的定性评估,而不仅仅是对机器生成输出与 ground truth 的匹配程度进行简单的定量比较
  • 本节将讨论 LLM-as-a-judge 如何被用于评估开放式生成、推理和更多新兴 NLP 任务
Open-ended Generation Tasks
  • 开放式生成指的是生成内容应安全、准确且上下文相关的任务,尽管没有单一的“正确”答案
  • 此类任务包括对话响应生成、摘要、故事生成和创意写作(2024;2023a;2024;2024;Bermejo, 2024)。与传统的基于指标的评估方法不同,LLM-as-a-judge 提供了更细致、适应性更强的定制化评估
  • 正如 Zheng 等人 (2023) 所指出的,像 GPT-4 这样的 LLM 在评判开放式文本生成方面表现得可与人类媲美
  • 在实践中,LLM-as-a-judge 已被应用于评估单个模型的输出,以及在竞争环境中比较多个模型的输出
  • 例如,Gao 等人 (2023b) 使用 ChatGPT 进行类似人类的摘要评估
  • 同样,Wu 等人 (2023) 提出了一个基于比较的框架,让 LLMs 充当具有多种角色扮演的评判,以特定维度评估摘要质量并生成评估结果
  • 现代 LLM 擅长生成详细的长文本响应。然而,随着输出长度的增加,产生幻觉的可能性也随之增加
    • 为了更好地理解这一现象,Cheng 等人 (2023) 和 Zhang 等人 (2024d) 引入了一种评估方法,使用 GPT-4 来评判生成的输出是否包含逻辑结构合理但无意义的陈述
    • Wang 等人 (2024a) 提出了一个基于批判的评判系统,通过选择相关证据并提供深入批判来评估幻觉
  • 除了幻觉之外,LLM 生成有害(如鼓励自杀)和不安全(如指导非法活动)响应的问题也备受关注
    • 针对这一问题,Li 等人 (2024g) 引入了 MD-Judge 和 MCQ-Judge,用于评估与安全相关的问答对,尤其关注旨在引发不安全响应的查询。这种方法支持无缝且可靠的评估
  • 然而,对不安全查询过于谨慎的态度可能导致过度拒绝响应,从而阻碍正常功能并对用户体验产生负面影响
    • 为了探讨这一问题,Xie 等人 (2024a) 对各种 LLM-as-a-judge 的框架进行了元评估,评估了当前 LLMs 对潜在不安全查询的拒绝倾向
    • 此外,(2024a) 引入了一个 LLM-based 答案提取器,以准确确定开放式生成中答案的关键部分
  • 另外,An 等人 (2023) 提出了 L-Eval,这是一个 LLM-as-a-judge 的框架,旨在为长上下文语言模型建立更标准化的评估
    • 这一概念得到了 Bai 等人 (2024) 的跟进,他们提议利用 LLM-as-a-judge 来过滤长上下文 LLMs 的评估数据
  • 最近的研究还利用 LLM-as-a-judge 来评估生成模型的通用能力
    • 这种方法通常采用基于辩论的框架,其中多个 LLMs 生成响应,随后由单独的 Judge LLM 进行评估
    • 例如,Chan 等人 (2023) 引入了一个多智能体辩论框架,旨在促进自主讨论并评估开放式文本生成任务中不同 LLMs 生成响应的质量
    • 类似地,Moniri 等人 (2024) 提出了一个自动化辩论框架,该框架不仅根据领域知识评估 LLMs,还评估它们在问题定义和不一致识别方面的能力
Reasoning Tasks
  • LLMs 的推理能力可以通过它们在特定推理任务上的中间思维过程和最终答案来评估(2024;2024;2024c)
    • 最近,LLM-as-a-judge 已被用于评估模型中间推理路径的逻辑进展、深度和连贯性
  • 对于数学推理任务,Xia 等人 (2024) 引入了一个自动评估框架,使用专门设计的 Judge LLM 来评估问题解决过程中推理步骤的质量
  • LLM-as-a-judge 还可以应用于更复杂的推理任务,如时间推理,其中模型需要理解不同事件在时间上的关系
  • Fatemi 等人 (2024) 构建了合成数据集,专门用于评估 LLMs 在各种场景下的时间推理能力,测试它们在时间有序事件的序列、因果关系和依赖关系方面的推理熟练程度
  • 大量的训练数据带来了一个挑战,即如何确定模型是通过深度逻辑理解进行推理,还是仅仅利用记忆模式(2024)
    • Wang 等人 (2023a) 设计了一个辩论式框架来评估 LLMs 的推理能力。给定一个特定问题,LLM 和用户采取对立立场并讨论该主题,以达成正确的决策
  • Nan 等人 (2024) 开发了一个多智能体评估框架,模拟学术同行评审过程
  • 该框架让 LLM-as-a-judge 参与协作评审,为数据驱动任务中 LLMs 的推理能力提供了更细致的理解
Emerging Tasks
  • 随着 LLMs 能力的快速发展,机器越来越多地被用于以前被认为是人类专属的任务,尤其是在特定上下文领域
  • 一个突出的任务是社交智能,其中模型面临复杂的社交场景,需要理解文化价值观、伦理原则和潜在的社会影响
  • 例如,Xu 等人 (2024a) 评估了 LLMs 的社交智能,强调尽管这些模型在学术问题解决能力方面取得了进步,但在社交智能方面仍明显落后
  • 同样,Zhou 等人 (2023) 引入了 SOTOPIA 和 SOTOPIA-EVAL,以模拟 LLM 智能体之间的复杂社交互动并评估它们的社交智能
  • 在他们的工作中,GPT-4 被用作人类评判的代理,以评估模拟互动中的目标完成度、财务管理和关系维护
  • 另一项研究致力于评估大型多模态模型(LMMs)和大型视觉语言模型(LVLMs)
    • 例如,Xiong 等人 (2024b) 探索了 LMM 作为评判来评估多模态模型的性能,提供最终分数和评估的基本原理,以提高透明度和一致性
  • Chen 等人 (2024d) 提出了第一个用于 LVLMs 自动评估的基准,专门针对自动驾驶的边缘情况
    • 他们发现,由 LLMs 作为评判进行的评估比由 LVLMs 作为评判进行的评估更符合人类偏好
  • 最近,论文看到 LLM-as-a-judge 的更定制化应用,用于评估新兴任务,如
    • 代码理解(2024a;2025;2024c;2024a;2024a;2024)
    • 法律知识(2023)
    • 游戏开发(Isaza-2024)
    • 海洋科学(2023)
    • 医疗对话(2024n)
    • 辩论评判(2024a)
    • 检索增强生成(2024)
    • 跨能力(cross ability,2024)
    • 人机交互(HCI)(2024j;2025;2024)
    • 角色扮演(2024c)
    • RAG(2024)
    • 语音合成(2024b)
    • 反语音生成(counterspeech generation,2025b,a)等
    • 这一趋势反映了 LLM-as-a-judge 在评估多样化和专业领域方面的适应性日益增强

Alignment

  • 对齐调优(2022a;2022)是使 LLMs 与人类偏好和价值观保持一致的重要技术
  • 这一过程的一个关键组成部分是收集高质量的成对反馈,这对于奖励建模(2017)或直接偏好学习(2023)至关重要
  • 最近,越来越多的研究兴趣集中在通过在对齐调优中采用 LLM-as-a-judge 来自动化这种成对反馈机制
Larger Models as Judges
  • 在对齐调优中采用 LLM-as-a-judge 的一个直观想法是利用更大、更强的 LLMs 的反馈来指导较小、能力较弱的模型
  • (2022) 首先提出利用 AI 的反馈来构建无害的 AI 助手
    • 他们使用基于预训练语言模型偏好的合成偏好数据来训练奖励模型
  • 在此基础上,Lee 等人 (2023) 发现,即使 LLM 评判不够强大,RLAIF 方法也能通过 RLHF 取得可比的性能
    • 他们还引入了 DIRECT-RLAIF,直接使用现成的 LLM-as-a-judge 模型,以减轻奖励模型中的奖励陈旧问题
  • 为了避免对齐中的奖励欺骗,Sun 等人 (2024a) 设计了一个可指示的奖励模型,该模型在合成偏好数据上训练
    • 它使人类能够在 RL 时间进行干预,以更好地使目标策略与人类价值观保持一致
    • 除了上述研究之外,Guo 等人 (2024) 引入了在线 AI 反馈(OAIF),直接利用注释模型的偏好信号来训练目标模型
  • 还有一些工作利用多智能体合作来在对齐调优中获得更好的评判
    • Arif 等人 (2024) 和 (2024) 使用多智能体工作流构建合成偏好优化数据集,并采用具有各种提示策略和 Pipeline 的 LLMs 作为评判
    • 类似地,(2024i) 利用多个 LLMs 相互辩论,迭代提高响应质量,同时创建一个 Judge LLM 来选择偏好的响应,以增强指令调优
  • 为了使生成的代码与人类偏好保持一致,Weyssow 等人 (2024) 引入了 CodeUltraFeedback,这是一个使用 LLM-as-a-judge 方法构建的偏好编码数据集
    • 这个合成数据集后来被用于使用 SFT 和 DPO 对小型代码 LLMs 进行微调和对齐
  • 最近,Wang 等人 (2024f) 提出了 BPO,使用 GPT-4 作为评判,并在对齐过程中构建合成成对反馈,以实现知识深度和广度的平衡
Self-Judging
  • 另一类工作旨在利用同一 LLM 的偏好信号来自我改进
  • Yuan 等人 (2024e) 首先提出了自我奖励 LLM 的概念,其中通过让 LLM 自身充当评判来构建成对数据
  • 紧随其后,Wu 等人 (2024a) 引入了元奖励,对 LLMs 的评判进行评判,并使用反馈来改进它们的评判技能
    • 他们的 LLM 作为元评判的方法显著增强了模型评估和遵循指令的能力
  • 为了提高合成数据质量:
    • Pace 等人 (2024) 结合了 Best-of-N 和 Worst-of-N 采样策略,并引入了 West-of-N 方法
    • Lee 等人 (2024) 设计了 Judge 增强监督微调(Judge augmented Supervised Fine-Tuning,JSFT),以训练单个模型同时充当策略和评判
      • 为了充分利用这个评判模型,他们还提出了通过 tournament 进行自我拒绝的方法,以在推理时选择最佳响应
  • 与上述使用 LLM-as-a-judge 来构建成对数据的方法不同,Tong 等人 (2024) 将 LLM-as-a-judge 应用于自我过滤方法,以确保对齐任务中合成数据对的质量,用于推理
  • 为了减少成对评判中的计算开销,Zhai 等人 (2024) 提出了一种用于自我偏好语言模型的排名配对方法,通过测量每个响应相对于基线的强度来加速比较过程
  • Liu 等人 (2024e) 引入了元排名(meta-ranking),使较弱的 LLMs(2024b)能够充当可靠的评判并提供可信赖的反馈。他们还将元排名方法应用于 SFT 后训练,将其与 Kahneman-Tversky 优化(KTO)结合使用,以改进对齐
  • 为了提高合成指令调优数据的质量,Liang 等人 (2024c) 引入了迭代自我增强范式(I-SHEEP)。在训练期间,他们采用 LLM-as-a-judge 来为合成响应评分,并设置阈值来收集高质量的查询-响应对,用于后续训练迭代
  • 最近,Yasunaga 等人 (2024) 提出将 LLM-as-a-judge 与数据合成相结合,并成功构建了只需少量注释即可与人类对齐的 LLMs
  • 几项工作还在特定领域或针对特定属性采用了 LLM-as-a-judge
    • Zhang 等人 (2024h) 提出了一种自我评估机制,通过生成问答对来评判响应的事实性。然后,他们利用这些自我注释的响应通过 DPO 算法对模型进行微调,以提高事实性
    • 在机器人技术中,Zeng 等人 (2024) 利用 LLMs 的自我排名响应来迭代更新奖励函数,从而在没有人类监督的情况下提高学习效率
    • 在多模态领域,Ahn 等人 (2024) 提出了迭代自我回顾评判(iterative self-retrospective judgment, i-SRT),该方法采用自我反思来改进响应生成和偏好建模

Retrieval

  • LLM-as-a-judge 在检索中的作用包括传统文档排名和更动态的、上下文自适应的检索增强生成(Retrieval-Augmented Generation,RAG)方法
  • 在传统检索中,LLMs 通过先进的提示技术提高排名准确性,使它们能够在几乎没有标记数据的情况下按相关性对文档进行排序
  • 作为补充,RAG 框架利用 LLMs 生成内容的能力,这些内容由检索到的信息引导,支持需要复杂或不断发展的知识整合的应用
    • 这些技术共同强调了 LLMs 作为检索任务评判的适应性,从基础排名到特定领域的知识增强应用
Traditional Retrieval
  • 最近的研究探索了 LLMs 作为评判在信息检索中对文档进行排名的作用,旨在提高排名精度并减少对大量训练数据的依赖
    • Sun 等人 (2023) 探索了像 GPT-4 这样的生成式 LLMs 在信息检索中进行相关性排名的潜力
      • 他们提出了一种基于排列的方法来按相关性对段落进行排名,指示 LLMs 输出段落的有序排列,从而提高排名精度
    • 作为补充,Zhuang 等人 (2024a) 引入了一种方法,将细粒度的相关性标签 Embedding 到 LLM 提示中,使模型能够区分细微的相关性变化并产生更精细的文档排序
    • Listwise 排名的进一步创新由 Ma 等人 (2023) 展示,他们提出了使用大型语言模型的 Listwise 重新排序器(LRL),这是一种直接对文档标识符进行重新排序的工具,无需依赖特定任务的训练数据。此外,Zhuang 等人 (2024b) 提出了一种适用于零样本排名的集合式提示策略,通过减少 LLM 推理频率和标记使用来简化排名操作,在不牺牲性能的情况下提高了效率
  • 为了解决位置偏差(这是 Listwise 排名任务中的常见挑战),Tang 等人 (2024b) 引入了排列自我一致性技术,该技术对多个列表顺序进行平均,以产生与顺序无关的排名。这种方法有效减少了位置偏差,这在 LLM 驱动的 Listwise 排名中是一个特别成问题的问题
  • 最后,Qin 等人 (2024) 批评了现有方法中点式和 Listwise 排名提示的局限性,指出典型的 LLMs 往往缺乏理解复杂排名任务的深度
    • 为了缓解这一问题,他们提出了使用中等规模开源 LLMs 的成对排名提示(PRP),作为更昂贵的大型专有模型的有效且经济的替代方案
  • 除了一般检索任务之外,LLMs 还证明了它们作为专门应用评判的实用性。例如,Ma 等人 (2024a) 概述了一个少样本工作流程,该流程使用通用 LLM 进行法律信息检索中的相关性评判。该模型通过将任务分解为多个阶段,实现了与专家注释的高度一致性,促进了专家推理的整合,以提高法律环境中相关性评估的准确性
  • 在推荐系统中,Hou 等人 (2024) 研究了 LLMs 将项目排名视为条件排名任务的潜力。该框架考虑了用户交互历史以及候选项目,解决了 LLMs 已知的偏差,例如倾向于偏爱流行或排名靠前的项目。采用了专门的提示和引导技术来纠正这些偏差并提高解释准确性
  • 最后,在搜索系统领域,Thomas 等人 (2023) 发现 LLMs 在预测搜索者偏好方面的表现可与人类标注者相媲美,使它们成为识别高性能系统和标记具有挑战性查询的有价值工具。这项研究强调了 LLMs 作为复杂检索任务评判的有效性,能够在各种应用中实现更细致和准确的相关性评估
Retrieval-Augmented Generation, RAG
  • 检索增强生成(RAG)的最新发展探索了 LLMs 自我评估和自我改进的能力,无需注释数据集或参数调整(2024e)
    • Li 和 Qiu (2023) 引入了思维记忆(MoT)框架,这是一个两阶段自我反思模型,自主增强 LLM 的推理能力
  • 在第一阶段,模型在未标记数据集上生成高置信度推理,将其存储为记忆
    • 在测试阶段,模型通过评判每个记忆与当前问题的相关性来回忆这些记忆,并选择最相关的记忆作为演示
  • 类似地,Tang 等人 (2024a) 提出了自我检索,这是一种创新架构,通过自然语言索引将信息检索(IR)能力整合到单个 LLM 中,以将语料库内化。这种方法将检索转变为文档生成和自我评估过程,在单个模型中实现了完全端到端的 IR 工作流程
  • 此外,Asai 等人 (2024) 提出了 SELF-RAG(Self-Reflective Retrieval-Augmented Generation,自我反思检索增强生成),该模型通过检索和自我反思循环提高 LLM 响应的质量和事实性。通过使用“反思标记”来指导适应性响应,SELF-RAG 使模型能够根据特定任务要求动态评判和调整其响应
  • 在问答领域,LLMs 越来越多地被用作评估代理,以实时评估答案的相关性、质量和实用性
    • Rackauckas 等人 (2024) 引入了一个 LLM-based 评估框架,该框架从实际用户交互和特定领域文档生成合成查询
      • 在这个框架中,LLMs 充当评判,评估检索到的文档并通过 RAGElo(一种基于 Elo 的自动竞争)对 RAG 智能体变体进行排名
      • 这种结构为 QA 系统中的质量控制提供了可扩展的解决方案
  • 此外,Zhang 等人 (2024b) 对 LLMs 评估开放域 QA 中相关性与实用性的能力进行了广泛研究。他们的发现表明,LLMs 可以有效区分两者,并且在呈现反事实段落时具有高度适应性
    • 这种实用性评估能力使 LLMs 能够在评估过程中提供更细致和上下文相关的响应
  • 针对特定领域的 RAG 系统揭示了 LLMs 通过整合专业知识结构来导航复杂查询的潜力
    • 在特定领域检索中,Wang 等人 (2024b) 提出了 BIORAG,这是一种先进的 RAG 框架,通过分层知识结构增强向量检索
    • BIORAG 采用自我意识评估检索器来持续评判其收集信息的充分性和相关性,从而提高检索文档的准确性
  • 对于生物医学研究,Li 等人 (2024c) 引入了 DALK(LLMs 和知识图谱的动态协同增强),这是一种新颖的系统,将 LLM 与从科学文献中派生的不断发展的阿尔茨海默病(AD)知识图谱相结合
    • 使用新颖的自我意识知识检索方法,DALK 利用 LLMs 的评判能力进行噪声过滤,增强 LLM 在 AD 相关查询中的推理性能
  • 类似地,Jeong 等人 (2024) 提出了 SelfBioRAG,这是一种将 RAG 原理应用于生物医学应用的框架
    • Self-BioRAG 采用 LLM 来选择最佳证据并基于所选证据和编码知识生成答案
    • 最近,(Déjean, 2024) 提炼了一个 LLM-as-a-judge ,从 LLMs 的参数记忆中提取检索结果

Reasoning

  • 释放 LLMs 的推理能力提供了一种缓解缩放定律局限性的方法,仅靠缩放定律可能无法充分揭示模型的潜力
  • 推理是 LLMs 的一个关键方面,因为它直接影响它们解决复杂问题、做出决策和提供准确的上下文感知响应的能力
  • 最近,许多关于 LLMs 推理能力的研究都集中在如何利用 LLM-as-a-judge 来选择推理路径(5.4.1 节)和利用外部工具(5.4.2 节)
Reasoning Path Selection
  • Wei 等人 (2022b) 引入了思维链(CoT)提示的概念,以鼓励模型生成逐步推理过程。虽然已经提出了其他更复杂的认知结构(2023a;2023)来增强 LLMs 的推理能力,但一个关键挑战是如何为 LLMs 选择合理且可靠的推理路径或轨迹以遵循
    • 为了解决这个问题,许多工作采用了 LLM-as-a-judge
  • 一些工作专注于推理过程中的样本级选择
    • (2023a) 引入了策略评估器,以在验证集上进一步评估候选策略
    • (2024) 引入了 REPS(Rationale Enhancement through Pairwise Selection, 通过成对选择增强理由),通过使用 LLMs 进行成对自我评估来评判和选择有效理由,并基于这些数据训练验证器(verifier)
  • LLMs 的另一个推理路径选择问题是多样性,Lahoti 等人 (2023) 发现 LLMs 掌握了多样性的概念,并且可以识别响应缺乏多样性的方面。通过选择和聚合多个批评意见,LLMs 可以取得类似的收益,与多次批评和修订迭代相比
  • 在多智能体协作框架中,Liang 等人 (2023) 提出了多智能体辩论(multi-agent debating,MAD),这是一种促进多个智能体之间辩论和讨论的新范式。他们利用 Judge LLM 在辩论过程结束时选择最合理的响应作为最终输出
  • 类似地,Li 等人 (2024b) 在基于层的多智能体协作中提出了新角色,采用 Judge LLM 来选择高质量和合理的响应,从而显著提高整个系统的标记利用效率
  • 此外,还有许多工作专注于步骤级推理路径选择,利用 Judge LLM 作为过程奖励模型(process reward model,PRM)来评估状态分数。Creswell 等人 (2023) 将推理过程分解为选择和推理。在选择步骤中,他们利用 LLM 本身来评判和评估每个潜在的推理轨迹,选择合适的轨迹用于后续推理步骤
  • Xie 等人 (2024b) 提出了 Kwai-STaR 框架,该框架将 LLMs 转变为状态转换推理器,以在数学推理中为自己评判和选择最佳推理状态。Lightman 等人 (2023) 训练 LLM 作为 PRM 来进行推理时监督,并在推理阶段执行 Best-of-N 采样策略
  • 紧随其后,Setlur 等人 (2024) 进一步提出了过程优势验证器(process advantage verifiers,PAVs),基于未来产生正确响应的可能性变化生成奖励。其他工作模拟高级认知结构作为推理过程
  • Hao 等人 (2023) 采用 LLMs 作为世界模型来模拟环境状态,并执行蒙特卡洛树搜索(Monte Carlo Tree Search,MCTS)以提高需要谨慎路径选择的任务的性能
  • Besta 等人 (2024) 将 LLMs 生成的输出视为任意图(arbitrary graph)
    • LLM 思维被建模为顶点,而边是思维之间的依赖关系
    • 该框架能够对每个推理状态的连贯性和逻辑推理进行系统评判
  • Yao 等人 (2023a) 提出了思维树(ToT),其中每个思维都作为解决问题的中间步骤
    • 它将推理分解为多个步骤,在每个状态进行自我评估和评判进展,并使用带有 LMs 的搜索算法通过前瞻和回溯来评判思维路径
  • 此外,还有一些研究训练基于批评的 LLM 评判(2024;2024b;2024e;2024b;2024;2024),这些评判提供细粒度的口头反馈以促进推理过程
Reasoning with External Tools
  • Yao 等人 (2023b) 首先提出以交错方式使用 LLMs 来生成推理轨迹和特定任务的动作
    • 推理轨迹帮助模型评判和更新动作计划,而动作使其能够与外部源交互
  • 随后,Auto-GPT 由 (2023) 引入,通过将 LLM-as-a-judge 用于工具使用,提供更准确的信息
    • 通过配备一系列外部复杂工具,LLMs 变得更加通用和有能力,通过评判和推理使用哪些工具来提高规划性能
  • Sha 等人 (2023) 探索了 LLMs 在推理和评判方面的潜力,将它们用作需要人类常识理解的复杂自动驾驶场景的决策组件
    • Zhou 等人 (2024d) 利用自我发现过程,其中 LLMs 根据给定的查询进行评判,并选择最可行的推理结构用于后续推理阶段
  • 尽管 LLMs 在各种工具的评判能力方面表现出色,但选择使用哪个模型或 API 通常涉及性能和成本之间的权衡
    • 更强大的模型虽然有效,但成本也更高,而能力较弱的模型则更具成本效益
    • 为了解决这个难题,(2024) 的作者提出了一个路由模型,该模型可以在评判过程中动态选择较强或较弱的 LLM ,旨在平衡成本和响应质量
  • 出于效率考虑,类似地,Zhao 等人 (2024b) 引入了 DiffAgent,作为一种旨在根据用户特定提示评判和选择不同文本到图像 API 的智能体
    • DiffAgent 的评判与人类偏好更一致,优于传统的 API 选择方法

Benchmark: Judging LLM-as-a-judge

  • 对 LLM 作为评判的评估需要强大且目标明确的基准,以捕捉这一任务的多方面性质
  • 论文将现有基准按以下维度分类:通用性能、偏差量化、领域特定性能、多模态评估、多语言能力、评估指令遵循、脆弱性评估和挑战性任务性能
    • 这一分类法展示了LLM作为评判评估框架的多样化目标,为分析其设计、范围和影响提供了结构化视角
    • 表 2 展示了LLM作为评判的各种基准和数据集的集合

General Performance

  • 专注于通用性能的基准旨在评估LLM在各种任务中的整体能力。这些基准通常测量与人类评判的一致性、准确性和相关性
  • 值得注意的例子包括:
    • MT-Bench和Chatbot Arena(2023),使用一致性、偏差和错误等指标评估对话场景
      • 这些基准进一步探索特定偏差,包括位置偏差、冗长偏差和自我增强倾向
    • JUDGE-BENCH(2024a)、DHP(2024j)、RewardBench(2024)、SOS-BENCH(2024)和JuStRank(2024),它们在更大规模上运行,使用Cohen’s kappa、辨别分数和标准化准确性等指标来基准化通用LLM性能
    • LLM-judge-eval(2024a),评估摘要和对齐等任务,并使用额外指标如翻转噪声和长度偏差

Bias Quantification

  • 减轻LLM评判中的偏差对于确保公平性和可靠性至关重要(Schroeder和Wood-Doughty, 2024)。典型基准包括EvalBiasBench 和CALM(2024a),它们明确专注于量化偏差,包括从对齐中出现的偏差和对抗条件下的鲁棒性偏差。此外,Shi等人 (2024a) 在问答任务中评估位置偏差和一致性百分比等指标

Challenging Task Performance

  • 为困难任务设计的基准推动了LLM评估的边界。例如
    • Arena-Hard Auto(2024k)、JudgeBench(2024a)和Yuan等人 (2024b) 分别针对对话式QA和各种推理任务,根据LLM的表现选择更难的问题
    • CALM(2024a)探索对齐和挑战性场景,使用可分离性、一致性和破解准确性等指标,评估在手动识别的硬数据集上的性能

Domain-Specific Performance

  • 领域特定基准提供了针对特定上下文的任务评估,以评估LLM的有效性
  • 具体而言,Raju等人 (2024) 测量跨任务的可分离性和一致性,利用 Brier 分数等指标,深入了解编码、医疗、金融、法律和数学等特定领域
  • CodeJudge-Eval(2024a)专门评估 LLM 对代码生成的评判,使用以执行为中心的指标如准确性和 F1 分数
  • 这一想法也被后续代码摘要和生成评估的多项工作所采用(2024b;2024;2024;2024;2024;2024e;2024)(1-172, 1-173)

Other Evaluation Dimensions

  • 除了通用性能和偏差量化,一些基准还解决了使用 LLM 作为评判的其他关键评估维度:
    • 多模态(Multimodal) :MLLM-as-a-judge(2024a)将评估框架扩展到涉及多种数据模态的任务,专注于与人类评判的一致性、分析评分和幻觉检测
    • 多语言(Multilingual) :MM-EVAL(2024b)和KUDGE(2024a)等基准评估多语言和非英语性能,测量准确性和相关性等指标,尤其在挑战性场景中
    • Instruction Following :Murugadoss等人 (2024) 检查LLM遵循特定评估指令的程度,使用相关性指标量化性能

Challenges & Future Works

  • 在本节中,论文概述了LLM作为评判的当前挑战和值得探索的未来方向,包括偏差与脆弱性、动态与复杂评判、自我评判以及人机协同评判

Bias & Vulnerability

  • LLM 作为评判本质上将评估构建为生成任务,引入了与偏差和脆弱性相关的重大挑战
    • 这些偏差通常源于模型的训练数据,其中常常嵌入(embeds)了与种族、性别、宗教、文化和意识形态等人口身份相关的社会刻板印象(2021)
    • 当LLM被部署于多样化的评判任务时,此类偏差会显著损害公平性和可靠性
  • 除了这些一般性偏差,当LLM充当评判时还会出现特定的评估偏差
    • 位置偏差(Order Bias) 是一个突出问题,候选的顺序会影响偏好(2023;2023c;2023a;2024b)
      • 这种偏差会扭曲评估结果,尤其是在成对比较中,当竞争响应之间的质量差距较小时更为明显(2024b;2023c)
    • 自我中心偏差(Egocentric Bias) 出现时,LLM会偏爱同一模型生成的输出,损害客观性(2023c;2023a;2024;2024)
      • 当评估指标使用同一模型设计时,这一问题尤为突出,会导致对源自该模型的输出评分虚高(2023c)
    • 长度偏差(Length Bias) 是另一个普遍挑战,评估会不成比例地偏爱更长或更短的响应,而不论质量如何(2023;2023a)
    • 其他偏差进一步复杂化了 LLM 评估,如:
      • 错误信息忽视偏差(Misinformation Oversight Bias) :比如,错误信息忽视偏差反映了忽略事实性错误的倾向
      • 权威偏差(Authority Bias) :比如,权威偏差偏爱来自所谓权威来源的陈述
      • 美观偏差(Beauty Bias) :比如,美观偏差优先考虑视觉吸引力内容而非实质质量(2024b,e;2024)
      • 冗长偏差(Verbosity Bias) 表现为偏爱更长的解释,通常将冗长等同于质量,这可能误导评判过程(2024c)
      • 情感偏差(Sentiment Bias) 会使评估因情感基调而倾斜,偏爱带有积极表述的响应(2024a)
  • LLM评判还极易受到对抗性操纵
    • 诸如 JudgeDeceiver 等技术凸显了基于优化的提示注入攻击的风险,精心设计的对抗序列可操纵 LLM 评判以偏爱特定响应(2024a;2024;2024a;2024b)
    • 同样,通用对抗短语可在绝对评分范式中大幅抬高分数,揭示了零样本评估设置的脆弱性(2023;2024;2024)
    • 这些操纵引发了对 LLM 评判在排行榜、学术或法律评估等高风险场景中的可靠性的担忧(2024a;2024;2024e;2024)
  • 为解决这些偏差和脆弱性,诸如 CALM(2024a)和 BWRS(2024b)等框架提供了系统的偏差量化和缓解方法
    • 多重证据校准(Multiple Evidence Calibration, MEC)、平衡位置校准(Balanced Position Calibration, BPC)和人机协同校准(Human-in-the-Loop Calibration, HITLC)等技术已被证明在使模型评判与人类评估一致的同时减少位置偏差和其他偏差方面有效(2023c)
    • 此外,认知偏差基准如 COBBLER 已识别出需要在 LLM 评估中系统缓解的六种关键偏差,包括显著性偏差和从众效应(2023b)
  • 未来方向(Future Direction) :
    • 未来研究的一个有希望的方向是将检索增强生成(RAG)框架集成到 LLM 评估过程中(2024e)
      • 通过结合生成和检索能力,这些框架可通过将评估基于外部可验证的数据源来减少自我偏好和事实性问题等偏差
    • 另一个有前景的途径是使用偏差感知数据集 ,如 OFFSETBIAS,以系统地解决 LLM 作为评判系统中的固有偏差(2024a)
      • 将此类数据集纳入训练 Pipeline 可使 LLM 更好地区分表面特质与实质正确性,从而增强公平性和可靠性
    • 探索微调 LLM 作为可扩展评判,如 JudgeLM 框架所示,代表了另一个有趣的方向(2023)
      • 该框架中 Embedding 的交换增强和参考支持等技术可系统地减轻偏差,提高评估一致性,并将基于LLM的评判的适用性扩展到开放式任务
    • 此外,推进零样本(zero-shot)比较评估框架具有重要前景(2023)
      • 这些框架可完善成对比较技术并实施去偏策略,提高跨不同评估领域的公平性和可靠性,而无需广泛的提示工程或微调
    • 最后,需要进一步探索抗 JudgeDeceiver 校准和对抗短语检测策略,以保护 LLM 作为评判框架免受攻击(2024a;2024;2024b;2024d;2024c)

Dynamic & Complex Judgment

  • 早期关于 LLM 作为评判的工作通常采用静态和直接的方法,直接提示评判 LLM 执行评估(2023)
    • 最近,已提出更多动态和复杂的评判 Pipeline 来解决各种限制,提高LLM作为评判的鲁棒性和有效性
    • 一个方向上的方法遵循“LLM-as-a-examiner”的概念,系统根据候选 LLM 的表现动态交互地生成问题和评判(2024d;2023a;2024a)
    • 其他工作专注于基于两个或更多候选 LLM 的对抗和辩论结果进行评判(2024;2024c)
      • 这些动态评判方法极大地提高了评判 LLM 对每个候选的理解,并可能防止 LLM 评估中的数据污染问题(2024)
      • 此外,构建复杂和精密的评判 Pipeline 或智能体是另一个流行的研究领域(2023a;2023;2024;2024)
      • 这些方法通常涉及多智能体协作,以及精心设计的规划和记忆系统,使评判 LLM 能够处理更复杂和多样化的评判场景
  • 未来方向(Future Direction) :
    • 未来研究的一个有希望的方向是赋予 LLM 类似人类的评判能力(2024c;2024b;2024)
    • 这些设计可借鉴人类评判时的行为,如锚定与比较、后见之明与反思,以及元评判(meta-judgment)
    • 另一个有趣的途径是开发使用 LLM 的自适应难度评估系统(adaptive difficulty assessment system, 2024)
      • 该系统将根据候选的当前表现调整问题难度
      • 这种自适应和动态系统可解决 LLM 评估中的一个重大限制,因为静态基准通常无法准确评估具有不同能力的LLM

Self-Judging

  • LLM-based 评估器,如 GPT-4,广泛用于评估输出,但面临重大挑战:
    • 尤其是自我中心偏差(Egocentric Bias) ,即模型偏爱自己的响应而非外部系统的响应(2023b;2023a;2023;2025)
      • 这种自我偏好破坏了公正性,造成了“先有鸡还是先有蛋”的困境:强大的评估器对开发强大的 LLM 至关重要,而推进 LLM 又依赖于无偏的评估器
    • 其他问题包括自我增强偏差(Self-Enhancement Bias) ,模型会高估自己的输出(2023a),以及奖励欺骗(Reward Hacking),对特定信号的过度优化导致评估的泛化能力降低(2024a)
      • 此外,对静态奖励模型的依赖(Static Reward Models)限制了适应性,而位置(Positional)和冗长(Verbosity)等偏差通过偏爱响应顺序或长度而非质量来扭曲评判(2024e;2024i)
      • 人类注释的高成本和有限可扩展性进一步复杂化了动态和可靠评估系统的创建(2022;2022)
  • 未来方向(Future Direction) :
    • 未来研究的一个有希望的方向是开发如 Peer Rank and Discussion(PRD)(2023a)这样的协作评估框架
      • 这些框架利用多个 LLM 集体评估输出,使用加权成对评判和多轮对话来减少自我增强偏差,并使评估更接近人类标准
    • 另一个有趣的途径是采用自教评估器框架,生成合成偏好对和推理轨迹,以迭代改进模型评估能力(2024i)
      • 这种方法消除了对昂贵人类注释的依赖,同时确保评估标准适应不断发展的任务和模型
    • 集成自我奖励语言模型(Self-Rewarding Language Models, SRLM)提供了另一条有前景的路径(2024e)
      • 通过采用如直接偏好优化(DPO)等迭代机制,这些模型持续改进其指令遵循和奖励建模能力,缓解奖励欺骗和过拟合问题
    • 在 SRLM 的基础上,使用元奖励机制引入元评判角色,以评估和改进评判质量(2024a)
      • 这一迭代过程解决了冗长和位置等偏差,增强了对齐和评估复杂任务的能力
      • 最后,利用合成数据创建生成对比响应为训练评估器提供了可扩展的解决方案(2024i)
      • 通过在合成偏好对上迭代改进评估,模型可逐步提高其鲁棒性和适应性
      • 将这些方法与多样化基准(2022;2022)、多方面评估标准(2020)和人类反馈(2023;2022)相结合,可确保评估在各个领域公平、可靠且与人类期望一致

Human-LLMs Co-judgement

  • 如前所述,LLM 作为评判中的偏差和脆弱性可通过人类参与评判过程进行进一步干预和校对来解决。然而,仅有少数研究关注这一方法
    • Wang 等人 (2023c) 引入了人机协同校准,采用平衡位置多样性熵来衡量每个示例的难度,并在必要时寻求人类协助
    • 在相关性评判背景下,Faggioli等人 (2023) 提出了人机协作光谱,根据人类依赖机器的程度对不同相关性评判策略进行分类
  • 未来方向(Future Direction)*
    • 随着数据选择(2023;2024)成为提高 LLM 训练和推理效率的日益流行的研究领域,它也有望提升 LLM 评估
    • LLM 作为评判可从数据选择中汲取灵感,使评判 LLM 能够作为关键样本选择器,根据特定标准(如代表性或难度)选择一小部分样本供人类注释者评估
    • 此外,人机协同评判的发展可受益于其他领域成熟的人机交互解决方案,如数据标注(2024b)和主动学习(2023)

NLP——DeepSeek-GRM

注:本文包含 AI 辅助创作

  • 参考链接:
    • 原始论文:(DeepSeek-GRM)Inference-Time Scaling for Generalist Reward Modeling, DeepSeek & THU, 20250403-20250925
      • 本工作是 THU 学生在 DeepSeek 实习期间完成
    • 开源模型: huggingface.co/collections/BBQGOD/deepseek-grm 和 modelscope.cn/profile/BBQGOD
      • 开源包含三个模型:
        • BBQGOD/DeepSeek-GRM-16B
        • BBQGOD/DeepSeek-GRM-27B
        • BBQGOD/DeepSeek-GRM-27B-MetaRM

Paper Summary

  • 核心内容:
    • 论文提出了自 Principle Critique 调优 (Self-Principled Critique Tuning, SPCT) 方法
    • SPCT 是一种增强通用奖励建模推理时扩展性的方法
    • 通过 Rule-based Online RL,SPCT 实现了 Principle 和 Critique 的自适应生成,显著提升了 GRM 在多样领域中的奖励质量和推理时扩展性
    • 在实验中,DeepSeek-GRM 超越了基线方法和一些强大的公开 RM,并通过推理时扩展,尤其是在 Mata RM 的引导下,展现出显著的改进
  • 背景:
    • RL 在 LLMs 的后训练中已被广泛采用
    • 在 LLMs 中通过 RL 激励推理能力表明: 适当的学习方法能够实现有效的推理时扩展性(proper learning methods could enable effective inference-time scalability)
  • 问题提出:
    • RL 的一个关键挑战是,在可验证问题或人工规则之外的各个领域为 LLMs 获取准确的奖励信号
  • 论文研究了如何通过增加推理计算来改进通用 Query 的奖励建模(Reward Modeling, RM),即 通用奖励建模的推理时扩展性(inference-time scalability of generalist RM)
    • 对于 RM(即奖励建模)方法,论文采用 Pointwise 生成奖励建模(pointwise generative reward modeling, GRM)以实现对不同输入类型的灵活性和推理时扩展的潜力
    • 对于学习方法,论文提出了 Self-Principled Critique 调优(Self-Principled Critique Tuning, SPCT) ,通过 Online RL 在 GRMs 中培养(foster)可扩展的奖励生成行为,以自适应地生成 Principle 并准确地给出 Critique,从而产生了 DeepSeek-GRM 模型
    • Furthermore,为了实现有效的推理时扩展,论文使用并行采样来扩展计算使用,并引入一个元奖励模型(meta RM)来指导投票过程以获得更好的扩展性能
  • 实验表明,SPCT 显著提高了 GRMs 的质量和可扩展性,在各种 RM 基准测试中优于现有方法和模型,且没有严重的偏差,并且与训练时扩展(training-time scaling)相比,它能获得更好的性能
  • DeepSeek-GRM 在某些任务上仍面临挑战,作者相信未来在通用奖励系统上的努力可以解决这些问题
  • 注:相关模型均已开源
  • 图 1: 不同 RMs 在所有测试的 RM 基准上的推理时扩展性能
    • 结果显示为每种方法最多 8 个样本,论文的方法进一步扩展到 32 个样本
    • 非斜体字体表示基于 Gemma-2-27B 的模型
  • 特别说明:
    • DeepSeek-GRM 的用法看看附录 G 中 DeepSeek-GRM (Default) Prompt 细节更容易理解

Introduction and Discussion

  • LLMs (2023; 2024) 的显著进步推动了人工智能研究的重大转变,使模型能够执行需要理解、生成和细微决策能力的任务
  • Recently,RL 作为 LLMs 的后训练方法已被大规模采用,并在人类价值观对齐(human value alignment)(2024; 2025)、长期推理(long-term reasoning)(2023; 2024) 和 LLMs 的环境适应(environment adaptation)(2024) 方面带来了显著改进。奖励建模(RM)(2024) 作为 RL 中的一个关键组件,对于为 LLM Response 生成准确的奖励信号至关重要
    • 当前研究(2024; 2025)也表明,无论是在训练时还是推理时拥有高质量且稳健的奖励,LLMs 都能在特定领域取得强劲性能
  • However,这种特定领域的高质量奖励主要来源于具有明确条件的人工设计环境(2022; 2024)或为可验证问题(例如数学问题(2021; 2023)和编码任务(2024; 2025))手工制定的规则
    • 在通用领域,奖励生成更具挑战性,因为奖励标准更加多样化和复杂,并且通常没有明确的参考或真实答案
    • 因此,通用奖励建模(generalist reward modeling)对于从后训练(例如大规模 RL)或推理(例如 RM 引导的搜索)角度提高 LLMs 在更广泛应用中的性能至关重要
    • Furthermore,RM 性能应通过增加训练计算(2023)和推理计算(inference compute)来提升
  • 图 2:奖励生成的不同范式,包括 (a) Scalar,(b) Semi-scalar 和 (c) 生成式方法,以及不同的评分模式,包括 (i) Pointwise 和 (ii) Pairwise 方法
    • 论文列出了每种方法的代表性方法,以及相应的推理时扩展性(是否可以从多次采样中获得更好的奖励)和输入灵活性(是否支持对单个和多个 Response 的评分)
  • 在实践中,要使 RMs 既通用又在推理时有效可扩展(effectively scalable)存在挑战
    • 通用性 要求 RM:
      • (1)对不同输入类型的灵活性
      • (2)在各个领域生成准确的奖励
      • 论文将此范式称为 通用奖励建模(generalist reward modeling)
    • Moreover,有效的 推理时扩展性(inference-time scalability) 要求 RM
      • (3)能够通过增加推理计算生成更高质量的奖励信号
      • (4)学习可扩展的行为以实现更好的性能-计算缩放(performance-compute scaling)
    • 现有的奖励建模研究展示了奖励生成的几种范式,包括 Scalar(2021; 2024; 2024)、 Semi-scalar(2025; 2025; 2025) 和生成式(generative)(2024; 2024; 2024; 2024; 2025; 2025; 2025; 2025; 2025) 方法,以及各种评分模式,例如 Pointwise(1940; 2023; 2024; 2025; 2025) 和 Pairwise (2024; 2023; 2023; 2024; 2025)
      • 这些方法本质上决定了 RMs 的输入灵活性(flexibility)和推理时扩展性((1)和(3)),如图 2 所示
      • For Instance, Pairwise RMs 仅考虑 Pairwise Response 的相对偏好,缺乏接受单个或多个 Response 作为输入的灵活性; Scalar RMs 很难为同一 Response 生成多样化的奖励信号,这阻碍了通过基于采样的推理时扩展方法(2025)获得更好的奖励
      • Also,不同的学习方法(2024; 2024; 2024; 2024)用于提高奖励的质量,但其中很少关注推理时扩展性,并研究学习到的奖励生成行为与 RMs 推理时扩展有效性之间的相互联系,导致性能提升有限((2)和(4))
      • 当前研究(2025)表明,有效的推理时扩展性可以通过适当的学习方法实现,这引出了问题:论文能否设计一种旨在实现通用奖励建模有效推理时扩展的学习方法?

        Can we design a learning method aiming to enable effective inference-time scaling for generalist reward modeling?

  • 论文研究了不同的 RM 方法,发现 Pointwise 生成奖励建模(GRM)可以在纯语言表示中统一对单个、 Pairwise 和多个 Response 的评分,克服挑战(1)
  • 论文探索了某些 Principle 可以在适当的准则范围内指导 GRMs 的奖励生成,从而提高奖励质量,这表明 RM 的推理时扩展性可能通过扩展高质量 Principle 和准确 Critique 的生成来实现
  • 基于此初步发现,论文提出了一种新颖的学习方法 Self-Principled Critique 调优(Self-Principled Critique Tuning, SPCT) ,以在 GRMs 中培养有效的推理时可扩展行为
    • 通过利用 Rule-based Online RL,SPCT 使 GRMs 能够学习根据输入 Query 和 Response 自适应地提出 Principle 和 Critique ,从而在通用领域获得更好的结果奖励(挑战(2))
    • 然后论文推出了 DeepSeek-GRM-27B ,它是基于 Gemma-2-27B (2024) 使用 SPCT 进行后训练的
    • 对于推理时扩展,论文通过多次采样来扩展计算使用
      • 通过并行采样,DeepSeek-GRM 可以生成不同的 Principle 集和相应的 Critique ,然后投票决定最终奖励
      • 通过更大规模的采样,DeepSeek-GRM 可以基于更多样化的 Principle 进行更准确的判断,并以更细的粒度输出奖励 ,这解决了挑战(3)和(4)
    • Furthermore,除了投票,论文还训练了一个 Mata RM 以获得更好的扩展性能
  • 实验表明,SPCT 显著提高了 GRMs 的质量和可扩展性,在多个综合 RM 基准测试中优于现有方法和模型,且没有严重的领域偏差
    • 论文还将 DeepSeek-GRM-27B 的推理时扩展性能与参数高达 671B 的更大模型进行了比较,发现与模型大小的训练时扩展相比,它能获得更好的性能
  • 尽管当前方法在效率和特定任务上面临挑战,但作者相信,通过 SPCT 之外的努力,具有增强可扩展性和效率的 GRMs 可以作为通用奖励系统的多功能接口,推进 LLM 后训练和推理的前沿
  • In general,论文的主要贡献如下
    • 1)论文提出了一种新颖的方法 Self-Principled Critique 调优(Self-Principled Critique Tuning, SPCT) ,以培养通用奖励建模的有效推理时扩展性,从而产生了(resulting in) DeepSeek-GRM 模型
      • 论文还进一步引入了一个 Mata RM 来有效提升 DeepSeek-GRM 在投票之外的推理时扩展性能
    • 2)论文通过实验证明,与现有方法和几个强大的公共模型相比,SPCT 显著提高了 GRMs 的质量和推理时扩展性
    • 3)论文还将 SPCT 训练方案应用于更大尺寸的 LLMs,发现推理时扩展可以超越模型大小的训练时扩展

Preliminaries

Comparisons of Different RM approaches

  • 如图 2 所示,RM 方法主要由奖励生成范式和评分模式决定,这本质上影响了 RM 的推理时扩展性和输入灵活性
    • 对于 奖励生成范式(reward generation paradigms) ,论文区分了三种主要方法: Scalar 、 Semi-scalar 和生成式
    • 对于 评分模式(scoring patterns) ,论文区分了两种主要方法: Pointwise 和 Pairwise
  • 为了在推理时扩展计算使用,论文专注于基于采样的方法,这些方法为相同的 Query 和 Response 生成多组奖励,然后聚合最终奖励
    • RMs 的 推理时扩展性(inference-time scalability) 取决于是否可以从多次采样中获得不同的奖励 ,其中 Scalar RMs 在大多数情况下会因奖励的恒定生成而失败;
    • RMs 的 输入灵活性(input flexibility) 由 RM 是否支持对单个、 Pairwise 和多个 Response 的评分来定义,其中 Pairwise RMs 很难对单个 Response 评分,通常需要额外的技术(2023; 2025)来处理多个 Response
Reward Generation Paradigms
  • 经典的 RMs 采用 (a) Scalar 方法(scalar approach) 生成奖励(\(\mathcal{R}\))
    • 为给定的 Query 和 Response 分配 Scalar 值
  • Scalar 方法进一步扩展到 (b) Semi-scalar 方法(semi-scalar approach)
    • 除了 Scalar 值外还生成文本
  • 而 (c) 生成式方法(generative approach) 仅生成文本奖励
    $$
    \mathcal{R}=
    \begin{cases}
    S & \text{(Scalar)} \\
    (S, \boldsymbol{C}) & \text{(Semi-Scalar)} \quad \sim r_{\theta}\left(x,\{y_i\}_{i=1}^{n}\right)\\
    \boldsymbol{C} & \text{(Generative)}
    \end{cases} \\
    \tag{1}
    $$
    • 上面的公式表示如下含义:
      $$ \mathcal{R} \sim r_{\theta}\left(x,\{y_i\}_{i=1}^{n}\right) $$
    • \(x\) 是 Query
    • \(y_i\) 是第 \(i\) 个 Response
    • \(r_{\theta}\) 是由 \(\theta\) 参数化的奖励函数
    • \(S \in \mathbb{R}^{m}, m \leq n\) 是 Scalar 奖励
    • \(\boldsymbol{C}\) 是 Critique
Scoring Patterns
  • 论文区分了奖励的两种主要评分方法: Pointwise 和 Pairwise
  • (i) Pointwise 方法(pointwise approach) 为每个 Response 分配一个单独的分数:
    $$
    \{S_i\}_{i=1}^{n}=f_{\text{point} }\left(\mathcal{R},\{y_i\}_{i=1}^{n}\right), \quad \mathcal{R} \sim r_{\theta}\left(x,\{y_i\}_{i=1}^{n}\right), S_i \in \mathbb{R},
    \tag{2}
    $$
    • 其中 \(f_{\text{point} }(\cdot,\cdot)\) 是一个分割函数(Spliting Function)
    • 问题:这里 Pointwise 打分的情况下,输入的 \(y_i\) 仅一个就可以了吧?此时应该有 \(n=1\)? 还是说输入可以是多个,但是一个个分别打分?
  • (ii) Pairwise 方法(pairwise approach) 可以看作是一种最佳选择方法(best-of-\(n\) method),从所有候选中选择一个最佳 Response :
    $$
    \hat{y}=f_{\text{pair} }(\mathcal{R},\{y_i\}_{i=1}^{n}), \quad \mathcal{R} \sim r_{\theta}\left(x,\{y_i\}_{i=1}^{n}\right), \hat{y} \in \{y_i\}_{i=1}^{n},
    \tag{3}
    $$
    • 其中 \(f_{\text{pair} }(\cdot,\cdot)\) 是一个选择函数,在大多数情况下 \(n=2\)
    • 虽然 Pairwise 方法可以扩展到 \(n>2\),但不能应用于单个 Response 评分(\(n=1\))
    • 理解:这里 Pairwise 方法和 Pointwise 方法的最本质区别是:
      • Pairwise 方法在挑选最佳 Response,而 Pointwise 在给每个 Response 打分,但似乎并没有限制输入的 Response 数量
    • 问题:这里是不是也可以理解为 listwise?如何定义 listwise、pairwise 和 pointwise 更合适?
Representative Methods
  • 图 2 说明了三种奖励生成范式( Scalar 、 Semi-scalar 、生成式)如何与两种评分模式( Pointwise 、 Pairwise )结合
  • Specifically
    • Bradley-Terry 模型(1940)(Scalar + Pointwise)使用 Pairwise 偏好数据进行训练,并以 Pointwise 方式输出 Scalar 奖励:
      $$
      \{S_i\}_{i=1}^{n}=f_{\text{point} }\left(\mathcal{R},\{y_i\}_{i=1}^{n}\right)=\mathbf{S} \in \mathbb{R}^{n}.
      \tag{4}
      $$
      • 理解:上述公式的意思是,\(\{S_i\}_{i=1}^{n}\) 是一个 \(n\) 维的分数向量 \(\mathbf{S}\),且对应向量 \(\mathbf{S}\in \mathbb{R}^{n}\)
    • PairRM(2023)(Scalar + Pairwise)通过 Scalar 奖励的符号比较一对 Response :
      $$
      \hat{y}=f_{\text{pair} }\left(\mathcal{R},\{y_i\}_{i=1}^{n}\right)=y_{\left|\frac{1}{2}(3-\text{sgn}(S))\right|}, \quad n=2,S \in \mathbb{R}.
      \tag{5}
      $$
      • 上述两个 Scalar 方法由于奖励生成缺乏多样性,几乎无法进行推理时扩展
    • Cloud(2024)(Semi-Scalar + Pointwise)基于预先生成的 Critique 为每个 Response 生成 Scalar 奖励,类似于公式 4
      • 理解:先生成 Critique,然后通过 Critique 生成一个 Scalar 分数
    • LLM-as-a-Judge(2023; 2024)(Generative + Pairwise)以文本方式判断 Pairwise Response 之间的偏好顺序:
      $$
      \hat{y}=f_{\text{pair} }\left(\mathcal{R},\{y_i\}_{i=1}^{n}\right)=y_{f_{\text{extract} }(\boldsymbol{C})}, \quad n=2,
      \tag{6}
      $$
      • 其中 \(f_{\text{extract} }(\cdot)\) 从语言表示中提取最佳 Response 的索引
      • However,这种方法默认忽略了 Pairwise Response 的平局情况
      • 理解:LLM-as-a-Judge 其实是一种很广泛的泛指吧?只要是生成式的是不是基本上都属于 LLM-as-a-Judge 的范畴?
    • 根据 Zhang 等人(2025b),指示偏好顺序的 token 的生成概率可以用作 Scalar 奖励(Semi-Scalar + Pairwise):
      $$ \mathcal{S}=\text{TokenProb}(\hat{\boldsymbol{C} })=r_{\theta}(\hat{\boldsymbol{C} }|x,\{y_i\}_{i=1}^{n})$$
      • 其中 \(\hat{\boldsymbol{C} }\) 是与偏好顺序相关的预定义 token(pre-defined token related to the preference order)
      • 理解:这里的含义是:Token 的概率本身就可以作为一个偏好 Scalar 奖励
  • (论文的方法)在没有额外约束的情况下,GRMs 能够在纯语言表示中为多个 Response 生成 Pointwise 奖励(Generative + Pointwise):
    $$
    \{S_i\}_{i=1}^{n}=f_{\text{point} }\left(\mathcal{R},\{y_i\}_{i=1}^{n}\right)=f_{\text{extract} }(\boldsymbol{C}),
    \tag{7}
    $$
    • 其中 \(f_{\text{extract} }(\cdot)\) 从生成结果中提取分配给每个 Response 的奖励
    • Usually,奖励是离散的,在本工作中,我们设定为自然数 默认分配 \(S_i \in \mathbb{N},1 \leq S_i \leq 10\)
    • 这种方法同时实现了推理时扩展性和输入灵活性

Boosting Reward Quality with Principles

  • 通用 RM 需要在特定领域之外生成高质量的奖励(2021; 2024),在这些领域中奖励标准更加多样化和复杂,并且通常没有明确的参考或真实答案
    • 为此,对于通用领域,论文采用 Principle 来指导奖励生成,以代替人工规则
  • LLMs 的 Principle 首先在宪法式人工智能(Constitutional AI)中引入(2022b; 2025),这些是手工制定的准则(hand-crafted criteria),指导 LLMs 或精选的分类器构建安全的数据 Pipeline
  • 有了 Principle ,GRMs 的奖励生成变为:
    $$
    \mathcal{R}=\mathbf{C} \sim r_{\theta}\left(x,\{y_i\}_{i=1}^{n},\{p_i\}_{i=1}^{m}\right),
    \tag{8}
    $$
    • 其中 \(\{p_i\}_{i=1}^{m}\) 表示 Principle
    • 理解:这里就是指同时给出 Query、Response,评估指标(Principle),然后使用生成式模型评估奖励(即是否满足 Principle)
    • 论文进行了一项初步实验来检验适当 Principle 对奖励质量的影响,使用了 Reward Bench(2024)的 Chat Hard 子集和 PPE 基准(2025)的 IFEval 子集
  • 在实验中,数据样本包含一个 Query 和两个 Response ,真实标签表示更好的 Response
    • 论文使用 GPT-4o-2024-08-06 生成 Principle ,然后为每个样本生成四次 Pointwise 奖励
    • 论文从正确的奖励生成过程中筛选 Principle (正确的定义:即更大的奖励值被分配给 Token 为更好的 Response )
    • 理解:这里是在筛选对应 Chosen 分数高于 Rejected 分数的数据,Rubrics-based RL 筛选数据也可以这样来筛选
  • 论文用它们自己生成的 Principle 和筛选过的 Principle 测试不同的 LLMs,并将它们与无 Principle 指导的默认设置进行比较,结果如表 1 所示
    • 论文发现, 自生成的 Principle 几乎没有显著提升奖励质量 (注:经过过滤的 Principle 是可以提升奖励质量的 )
    • 这一结果并非微不足道(non-trivial),可以得出两个主要结论:
      • (a) 当前的 LLMs 可以生成多样化的 Principle ,但并非所有 Principle 都适合用于奖励生成
      • (b) 生成的 Principle 的一个子集可以在正确的准则下更好地指导奖励生成,这表明了自我引导(self-bootstrapping)的潜力
    • 这些发现是利用 Online RL 优化 GRMs 的基础,它们可以从自己生成的 Principle 中学习,并有一个清晰的信号来判断 Principle 是否合适
  • 其他细节在附录 D 中描述

Self-Principled Critique Tuning, SPCT

  • 受到初步结果的启发,论文为 Pointwise GRMs 开发了一种新颖的方法,学习生成能够有效指导 Critique 生成的自适应高质量 Principle ,称为 Self-Principled Critique 调优(Self-Principled Critique Tuning, SPCT)
  • 如图 3 所示,SPCT 包括两个阶段:
    • 第一阶段:作为冷启动的拒绝式微调(rejective fine-tuning)
    • 第二阶段:Rule-based Online RL,通过改进生成的 Principle 和 Critique 来强化通用奖励生成
      • SPCT 也在 GRMs 中培养了这些行为以实现推理时扩展
  • 图 3:SPCT 的图示,包括拒绝式微调、 Rule-based RL 以及推理期间相应的可扩展行为
    • 推理时扩展通过朴素投票或由大规模生成的 Principle 指导的 Mata RM 投票实现,从而在扩展的值空间内产生更细粒度的结果奖励

Unpinning Principles from Understanding to Generation(将核心原理从 “理解任务” 抽离并迁移到 “生成任务” 中)

  • 根据第 2.2 节的初步实验,论文发现适当的 Principle 可以在特定准则内指导奖励生成,这对于高质量奖励至关重要
    • However,大规模生成通用 RM 的有效 Principle 仍然具有挑战性
  • 为了解决这一挑战,论文提出 Unpinning Principles from Understanding to Generation,即将 Principle 视为奖励生成的一部分,而不是预处理步骤
  • 形式上, Principle 根据公式 8 指导奖励的生成,当 Principle 是预定义的时
    • GRMs 可以自己生成 Principle ,然后基于这些 Principle 生成 Critique ,形式化为:
      $$
      \{p_i\}_{i=1}^{m} \sim p_{\theta}\left(x,\{y_i\}_{i=1}^{n}\right), \quad \mathcal{R}=\mathbf{C} \sim r_{\theta}\left(x,\{y_i\}_{i=1}^{n},\{p_i\}_{i=1}^{m}\right),
      \tag{9}
      $$
    • 其中 \(p_{\theta}\) 是由 \(\theta\) 参数化的 Principle 生成函数,与奖励生成 \(r_{\theta}\) 共享同一模型
      • 在实践中,它们使用 LLMs 中的同一语言头实现
    • 这种转变使得 Principle 能够基于输入 Query 和 Response 生成,自适应地对齐奖励生成过程,并且 Principle 和相应 Critique 的质量和粒度可以通过 GRMs 的后训练进一步改进
      • 通过 大规模生成的 Principle ,GRMs 可能以更细的粒度和更广泛的考虑输出奖励,从而实现更好的推理时扩展性

Rule-Based Reinforcement Learning

  • 为了同时优化 GRMs 中的 Principle 和 Critique 生成,论文提出了 SPCT,它整合了拒绝式微调(rejective fine-tuning, RFT)和 Rule-based RL
    • RFT 的作用是冷启动
Rejective Fine-Tuning, Cold Start
  • 拒绝式微调阶段的核心思想是训练 GRM 以正确的格式为各种输入类型生成 Principle 和 Critique
  • 与之前混合不同格式的单个、 Pairwise 和多个(single, paired, and multiple) Response 的 RM 数据的工作(2024; 2024; 2025)不同,论文采用第 2.1 节介绍的 Pointwise GRM,灵活地以相同格式为任意数量的 Response 生成奖励
  • 对于数据构建,除了通用指令数据外,论文还通过给出 Query 和相应 Response 用预训练的 GRM 采样轨迹
    • 每个 RM 数据点(Data Point)包含一个 Query 和一个或多个对该 Query 的 Response ,以及表示最佳 Response 的真实标签
    • 对于每个 RM 数据点, Principle 和 Critique 的采样执行 \(N_{\text{RFT} }\) 次
      • 拒绝策略也是统一的,即拒绝以下两种情况:
        • 第一:预测奖励不正确的轨迹
        • 第二:所有 \(N_{\text{RFT} }\) 条轨迹都正确(太容易)的 Query 和 Response
    • 形式上,令 \(r_i\) 表示对 Query \(x\) 的第 \(i\) 个 Response \(y_i\) 的真实奖励,如果满足下面的条件,则认为预测的 Pointwise 奖励 \(\{S_i\}_{i=1}^{n}\) 是正确的
      $$
      \begin{cases}
      \forall i \neq j, \quad S_j > S_i, \quad j = \arg \max_{l}\{r_l\}_{i=l}^{n}, & \text{if } n \geq 2, \\
      S_1 = r_1, & \text{if } n = 1.
      \end{cases}
      \tag{10}
      $$
      • 只有一个 Response 时,当且仅当真实分数 \(r_1\) 和 预测分数 \(S_1\) 完全相等才算正确
      • 有多个 Response 时,当且仅当真实奖励中最大的 Response 对应的分数高于所有其他 Response(类似 Best-of-N)
    • 并保证真实奖励只包含一个最大值
  • However,与之前的工作类似(2025b),论文发现预训练的 GRMs 在有限的采样配额内很难为一部分 Query 和相应的 Response 生成正确的奖励
    • 理解:这里的问题是有限的采样次数可能是无法生成准确的评估轨迹的(即无法找到最佳的 Response)
  • Thus,论文可选择地将 \(\arg \max_{j}\{r_l\}_{l=1}^{n}\) 附加到 GRM 的提示中,称为 提示采样(hinted sampling) ,期望预测的奖励与真实情况一致,此外还有 非提示采样(non-hinted sampling)
    • 具体来说,将在输入中附加一个额外的片段 “The best response is: Response \(\arg \max_{i}\{r_l\}_{l=1}^{n}\)”
    • 对于提示采样,每个 Query 和相应的 Response 只采样一次 ,只有当轨迹不正确时才拒绝
      • 问题:这样的话,相当于先给答案(告诉模型最佳 Response 是谁),再让模型生成推理过程
    • 除了之前的研究(2024; 2024),论文观察到提示采样的轨迹有时在生成的 Critique 中走捷径,特别是对于推理任务,这表明了 Online RL 对 GRM 的必要性和潜在好处
Rule-Based RL
  • GRM 使用 Rule-based Online RL 进一步微调,论文使用 GRPO(2024)的原始设置以及 Rule-based 结果奖励
  • 在 rollout 期间,GRM 根据输入 Query 和 Response 生成 Principle 和 Critique ,然后提取预测的奖励并与真实值通过准确度规则进行比较
    • 与 DeepSeek-AI(2025)不同,论文不使用格式奖励
    • Instead,论文应用了更大的 KL 惩罚系数以确保格式并避免严重偏差
  • 形式上,对于给定 Query \(x\) 和 Response \(\{y_i\}_{i=1}^{n}\) 的第 \(i\) 个输出 \(o_i\) 的奖励是:
    $$
    \hat{r}_i=
    \begin{cases}
    1, & \text{if } n \geq 2 \text{ and } \forall i’ \neq j’, \quad S_{j’} > S_{i’}, \quad j’ = \arg \max_{l}\{r_l\}_{l=1}^{n}, \\
    1, & \text{if } n = 1 \text{ and } S_1 = r_1, \\
    -1, & \text{otherwise},
    \end{cases}
    \tag{11}
    $$
    • 其中 Pointwise 奖励 \(\{S_i\}_{i=1}^{n}\) 是从 \(o_i\) 中提取的
      • 问题:一个 \(o_i\) 中包含了所有的 Pointwise 奖励 \(\{S_i\}_{i=1}^{n}\) 吗?
    • 该奖励函数鼓励 GRMs 通过在线优化的 Principle 和 Critique 来区分最佳 Response ,有利于有效的推理时扩展
      • 奖励信号可以从任何偏好数据集和带标签的 LLM Response 中无缝获得
      • 理解:即有 Chosen/Rejected 或 Best-of-N 数据的样本都可以用来训练
    • 理解:上述奖励跟前面的 RFT 类似:
      • 只有一个 Response 时,当且仅当真实分数 \(r_1\) 和 预测分数 \(S_1\) 完全相等才算正确
      • 有多个 Response 时,当且仅当真实奖励中最大的 Response 对应的分数高于所有其他 Response(类似 Best-of-N)
        • 理解:在当前的设计下,有多个 Response 时,所有 Response 的分数是同时为 1(预测正确)或 -1(预测错误)的

Inference-Time Scaling with SPCT

  • 为了进一步利用更多推理计算资源来提升 DeepSeek-GRM 在通用奖励生成上的性能,论文探索了基于采样的策略,以实现有效的推理时扩展性

Voting with Generated Rewards

  • 投票是 RM 中广泛采用的实现推理时扩展的方法
  • 回顾第 2.1 节的方法,论文展示了 Semi-scalar RM 和生成式 RM 对于 \( k \) 个样本的投票结果
  • 对于 Semi-scalar RM (2024; ),投票以平均方式进行:
    $$
    S^* = \frac{1}{k} \sum_{i=1}^{k} S_{i}, \quad \{\mathcal{R}_{i}=(S_{i}, C_{i})\}_{i=1}^{k} \sim r_{\theta} \left( x, \{y_{i}\}_{i=1}^{n} \right),
    $$
    • 其中 \( S^* \) 是最终奖励
    • 在实践中, Scalar 值方差有限,这可能会阻碍扩展性
  • 对于 Pairwise GRM (2024; ),投票以多数表决方式选择被识别为最佳的 Response :
    $$
    \hat{y}^* = \arg \max_{y} \sum_{i=1}^{k} \mathbb{I}(y = \hat{y}_{i}), \quad \{\mathcal{R}_{i}=C_{i}\}_{i=1}^{k} \sim r_{\theta} \left( x, \{y_{i}\}_{i=1}^{n} \right),
    $$
    • 其中 \( \hat{y}^* \) 是最终预测的最佳 Response
    • \( f_{\text{pair} }(\cdot, \cdot) \) 是一个选择函数
    • \( \hat{y}_{i} = f_{\text{pair} }(\mathbf{C}_{i}, \{y_{i}\}_{i=1}^{n}) \) 是每个样本单独选择的最佳 Response
    • \( \mathbb{I}(\cdot) \) 是指示函数
    • 虽然投票过程是可扩展的,但由于每个样本中不允许出现平局,多数投票结果可能存在偏差,并且由于缺乏量化分数,可能无法区分 Response 之间的细微差别
  • Pointwise GRM 的投票过程定义为奖励求和:
    $$
    S_{i}^{*} = \sum_{j=1}^{k} S_{ij}, \quad \{p_{ij}\}_{i=1}^{m_{j} } \sim p_{\theta} \left( x, \{y_{i}\}_{i=1}^{n} \right), \mathcal{R}_{j} = C_{j} \sim r_{\theta} \left( x, \{y_{i}\}_{i=1}^{n}, \{p_{ij}\}_{i=1}^{m_{j} } \right), j=1,…,k,
    $$
    • 其中 \( S_{i}^{*} \) 是第 \( i \) 个 Response (\( i=1,…,n \)) 的最终奖励,且 \( \{S_{ij}\}_{i=1}^{n} = f_{\text{point} }(C_{j}, \{y_{i}\}_{i=1}^{n}) \) 是第 \( j \) 组 Pointwise 奖励
    • 由于 \( S_{ij} \) 通常被设定在一个较小的离散范围内, 例如 \{1,…,10\},投票过程实际上将奖励空间扩展了 \( k \) 倍,并使 GRM 能够生成大量 Principle ,这有益于最终奖励的质量和粒度
      • 一个直观的解释是,如果每个 Principle 可以被视为判断视角的代理,那么更多的 Principle 可能更准确地反映真实分布,从而产生扩展效果。值得注意的是,为了避免位置偏差并增加多样性,在采样前会对 Response 进行打乱

Meta Reward Modeling Guided Voting

  • DeepSeek-GRM 的投票过程需要多次采样,并且由于随机性或模型限制,少量生成的 Principle 和评判 (Critique) 可能存在偏差或质量低下
    • 因此,论文训练了一个 Mata RM 来引导投票过程(guide the voting process)
  • Mata RM 是一个 Pointwise Scalar RM,旨在识别 DeepSeek-GRM 生成的 Principle 和 Critique 的正确性,使用二元交叉熵损失,其中标签根据公式 10 确定
    • 提示模板见附录 G,整合了 Query 、候选 Response 、相应 Principle 和 Critique
  • 数据集包括 RFT 阶段来自非提示采样的轨迹,以及来自待引导的 DeepSeek-GRM 的采样轨迹,这既能提供足够的正负奖励,又能缓解训练与推理策略之间的差距,正如 Chow 等 (2025) 所建议
  • 引导的投票过程很简单(即使用 Meta RM 的方式):
    • Mata RM 输出 \( k \) 个采样奖励的元奖励,最终结果由元奖励最高的前 \( k_{\text{meta} } \leq k \) 个奖励进行投票得出,从而过滤掉低质量样本

Results on Reward Modeling Benchmarks

Experiment Settings

Benchmarks and Evaluation Metrics
  • 论文在不同领域的多个 RM 基准测试上评估不同方法的性能:
    • Reward Bench (RB) (2024),PPE(偏好和正确性子集)(2025),RMB (2025),Real.Mistake (2024)
  • 论文对每个基准测试使用标准评估指标:Reward Bench、PPE 和 RMB 中从一组 Response 中选取最佳 Response 的准确率,以及 Real.Mistake 的 ROC-AUC
  • 为了处理多个 Response 预测奖励出现平局的情况,论文打乱 Response 顺序,并通过 \( \arg \max_i S_i \) 确定最佳 Response ,其中 \( S_i \) 是打乱后第 \( i \) 个 Response 的预测奖励
  • 细节见附录 D
Method Implementation
  • 对于基线方法,论文基于 Gemma-2-27B (2024) 并采用与 DeepSeek-GRM 兼容的所有训练数据和设置,重新实现了 LLM-as-a-Judge (2023),DeepSeek-BTRM-27B(Bradley-Terry 模型)(1940),CLoud-Gemma-2-27B (2024) 和 DeepSeek-PairRM-27B (2023)
  • 对于论文的方法,论文基于 Gemma-2-27B 实现了 DeepSeek-GRM-27B-RFT,并在不同规模的 LLM 上实现了 DeepSeek-GRM,包括 DeepSeek-V2-Lite (16B MoE) (2024a),Gemma-2-27B,DeepSeek-V2.5 (236B MoE) 和 DeepSeek-V3 (671B MoE) (2024b)
  • Mata RM 在 Gemma-2-27B 上训练
  • 默认结果使用 贪婪解码(greedy decoding) 报告,推理时扩展(inference-time scaling)使用温度 = 0.5
  • 其他细节见附录 C

Results and Analysis

Performance on RM Benchmarks
  • 不同方法和模型在 RM 基准测试上的总体结果如表 2 所示
  • 论文将 DeepSeek-GRM-27B 的性能与公开模型的报告结果(Reported Results of Public Models)以及基线方法的复现结果(Reproduced Result of Baseline Methods)进行比较
    • 论文发现 DeepSeek-GRM-27B 在整体性能上优于基线方法,并且与强大的公开 RM(如 Nemotron-4-340B-Reward 和 GPT-4o)相比取得了有竞争力的性能;通过推理时扩展,DeepSeek-GRM-27B 可以进一步改进并获得最佳的整体结果
  • 详细比较来看, Scalar(DeepSeek-BTRM-27B)和 Semi-scalar(CLoud-Gemma-2-27B)RM 在不同基准测试上表现出有偏差的结果,在可验证任务(PPE Correctness)上的性能明显优于所有生成式 RM,但在其他不同基准测试上分别失败
    • 问题:如何理解 Skywork-Reward-Gemma-2-27B 反而是在 Reward Bench 上分数很高,其他任务上分数很低呢?是过拟合吗?
  • 尽管如此,大多数公开的 Scalar RM 也表现出严重的领域偏差
    • PairRM 方法可以缓解这个问题
  • LLM-as-a-Judge 与 DeepSeek-GRM-27B 显示出相似的趋势但性能较低,可能是由于缺乏对单 Response 评分的训练
  • 总之,SPCT 提升了 GRM 的通用奖励生成能力,与 Scalar 和 Semi-scalar RM 相比,偏差显著减少
Inference-Time Scalability
  • 不同方法的推理时扩展结果如表 3 所示,总体趋势如图 1 所示
  • 细节见附录 D.3
  • 在最多 8 个样本的情况下(即 Voting@8),论文发现 DeepSeek-GRM-27B 相对于贪婪解码和采样结果的性能提升最高
    • DeepSeek-GRM-27B 进一步显示出使用更多推理计算资源(最多 32 个样本)提升性能的强大潜力
    • 论文将这种有效性归因于细化的 Principle 生成,它以结构化的方式扩展了输出长度,并引导结果奖励更接近真实分布
  • Mata RM 也显示出其在每个基准测试上为 DeepSeek-GRM 过滤低质量轨迹的有效性
  • 使用 Token 概率进行投票的 LLM-as-a-Judge 也显示出显著的性能提升,这表明 作为量化权重(Quantitative Weights)的 Token 概率可以帮助提高仅基于离散索引进行多数投票的可靠性
  • 对于 CLoud-Gemma-2-27B,性能提升有限
    • 主要是因为 Scalar 奖励生成缺乏方差,即使 Critique 发生了很大变化
  • In Summary,SPCT 提升了 GRM 的推理时扩展性,而 Mata RM 进一步提升了通用场景下的扩展性能
Ablation Study
  • 表 4 展示了所提出的 SPCT 不同组件的消融研究结果,详细结果列在附录 D.3
    • 令人惊讶的是(Surprisingly),即使没有经过拒绝采样 Critique 数据的冷启动,经过通用指令微调的 GRM 在经历 Online RL 后性能仍有显著提升(66.1 → 68.7)
      • 理解:这里是指使用 Online RL 去训练 GRM,这里对比的是第8行模型(仅包含通用指令微调)和第3行模型(在通用指令微调模型上经过了 Online RL 的模型)
    • Also,非提示采样似乎比提示采样更重要
      • 可能是因为提示采样轨迹中出现了走捷径的现象
    • 以上这些都表明了 GRM 在线训练的重要性
  • 与先前工作一致,论文确认通用指令数据对于 GRM 的性能至关重要
    • 论文发现 Principle 生成对于 DeepSeek-GRM-27B 的贪婪解码和推理时扩展性能都至关重要
  • 对于推理时扩展, Mata RM 引导的投票在不同的 \( k_{\text{meta} } \) 下表现出鲁棒性
  • 关于通用 RM 性能的进一步分析,包括输入灵活性、训练数据的领域泛化等,在附录 E 中讨论
Scaling Inference and Training Costs
  • 论文通过在不同规模的 LLM 上进行后训练,进一步研究了 DeepSeek-GRM-27B 的推理时和训练时扩展性能
  • 模型在 Reward Bench 上进行测试,结果如图 4 所示
  • 论文发现,使用 32 个样本直接投票的 DeepSeek-GRM-27B 可以达到与 671B MoE 模型相当的性能,而 Mata RM 引导的投票仅用 8 个样本即可获得最佳结果
    • 这证明了 DeepSeek-GRM-27B 的推理时扩展相较于扩展模型规模有更高的有效性
  • Moreover,论文在包含 300 个样本的下采样测试集上测试了 DeepSeek-R1-0120,发现其性能甚至低于 236B MoE RFT 模型
    • 这表明扩展长思维链进行推理任务并不能显著提升通用 RM 的性能

Related Work

Generative Reward Models

  • GRM 代表了从 Scalar RM (2022) 的范式转变,将奖励建模为文本反馈或分数
  • (2024a; 2024; 2025a; 2024; 2024; 2025),实现了更丰富的奖励表示和更灵活的单个及多个 Response Critique
  • 此前,LLM-as-a-judge 方法 (2023; 2024c) 支持基于参考或无参考的 Pairwise Critique 来评估 LLM
  • 最近的研究使用离线和 Online RL 来训练 GRM (2024; 2024; 2025b; 2025b; 2025),将工具和外部知识与 GRM 结合 (2024b; 2025),甚至训练 GRM 作为调整环境奖励的接口 (2025)
  • 尽管这些方法在效率上面临挑战,但它们展示了大规模改进奖励的潜力,朝着更通用的奖励系统发展

Inference-Time Scaling for LLMs

  • LLM 的推理时扩展一直是一个与训练时扩展并行的重要研究方向
  • 研究集中于采样和 RM 引导的聚合 (2024; 2024; 2025; 2025)
  • 最近,从 LLM 中激励产生的长思维链 (2022) 显著提升了模型在解决 (OpenAI, 2024; DeepSeek-AI, 2025; OpenAI, 2025a) 和 Critique (2025; 2025) 困难可验证问题时的推理能力,这是推理时扩展的另一种形式
  • 然而,论文没有找到像 DeepSeek-AI (2025) 那样有效激励长范围奖励生成以实现通用奖励建模的方法,论文将推理与 Principle 引导的奖励生成的结合留待未来的工程努力
  • 也有研究使用可扩展的奖励或验证器来提升策略模型在编码 (2023)、推理 (2025) 等领域的性能
  • 因此,本工作中推理时可扩展的通用 RM 的发展,也可能通过推理时协同扩展,为策略模型的通用性能做出贡献

Ethics Statement

  • 论文提出的方法,自 Principle Critique 调优 (SPCT),旨在增强生成式奖励模型在通用领域的推理时扩展性
    • 尽管这一进展促进了奖励建模的准确性和一致性,但有几个伦理影响可能需要明确考虑
  • 首先,尽管通过论文的实证分析表明 DeepSeek-GRM 在不同领域表现出较少的偏差,但当训练数据存在毒性时,自动生成的 Principle 和 Critique 可能会无意中延续或放大偏差
    • 作者认为应该优先研究 Mata RM 和其他偏见缓解策略,以确保公平的结果
    • 此外,论文的方法并非旨在削弱人类监督
      • 相反,论文主张维护人在环路框架,并开发可靠的代理方法(如 SPCT)来更高效、更有效地扩展人类监督
  • 其次,推理时可扩展 GRM 在多样化领域的适用性扩大,可能会引发关于透明度、问责制等方面的担忧
    • 由于奖励生成行为很大程度上源于自我引导,不忠实的 Principle 和 Critique 的可能性是不可忽视的
    • 论文在附录 F.1 中展示了案例研究,在附录 B 中说明了局限性,并在公开监督下开源了模型,这对于维护信任和确保工件的负责任部署至关重要
  • 最后,在不同 RM 基准测试和实际场景中进行稳健的验证和持续的警惕仍然至关重要
    • 负责任地使用 DeepSeek-GRM 需要主动管理风险并持续评估偏见,这需要在 RM 评估研究方面付出努力

附录 A:Additional Related Work

Constitutional AI

  • Constitutional AI 已成为传统 RLHF (2022) 的一个可扩展替代方案,旨在通过一套指导 Principle 或“宪法”使语言模型与人类价值观对齐 (2022b; 2023, 2024),用基于这些人工制定 Principle 的AI生成反馈 (2024) 或分类器 (2025) 替代人类 Critique
  • 类似地, Rule-based 方法如 Sparrow (2022) 和 Rule-Based Rewards (RBR) (2024) 将明确的自然语言规则纳入特定领域(如安全性)的训练循环中
  • 尽管这些方法有效,但它们依赖于静态的、人工编写的宪法,这些宪法在范围上有限、可能存在偏见且不够灵活
    • 这激发了人们对自动化生成或改进 Principle 的兴趣,这也与论文本工作的目标相一致

Scalar Reward Models

  • Scalar 奖励模型最初是为 LLMs 提出,作为人类反馈的代理模型 (2020; 2023)
  • 近期的研究侧重于 Bradley-Terry 建模 (1940) 和其他回归方法,以提高 Scalar 奖励模型在通用偏好上的表达能力 (2024; 2024e, 2024b; 2024; 2025b)
  • 与这些结果奖励模型相比,过程奖励模型被提出作为推理问题(如数学等)的步骤验证器 (2021; 2024b; 2025b),展示了 Scalar RM 在具有广泛推理和知识的正式领域中的可行性
  • Scalar RM 的优点是简单且计算高效,但表达能力有限,并且难以跨不同输入类型进行泛化或在推理时细化奖励信号

Semi-Scalar Reward Models

  • Semi-scalar 奖励模型旨在通过文本中间表示来丰富 Scalar 奖励信号 (2025a; 2024)
    • (2025b) 提出通过提高生成的 critiques 的质量来最终改进奖励生成
  • 一些研究使用 token 概率来替代 Scalar 头部进行奖励提取 (2024; 2025a)
  • 以上这些工作表明
    • Semi-scalar RM 在基于采样和投票的推理时扩展方面面临挑战,导致性能提升有限
    • Semi-scalar 方法在效率和效果之间权衡了 Scalar RM 和 GRM

附录 B:Limitations and Future Directions

Limitation

  • 尽管 SPCT 显著提升了 GRM 的性能和推理时扩展性,并在通用领域超越了(Semi)Scalar RM,但它仍面临一些局限性
  • (1) 生成式 RM 的效率本质上远远落后于同等规模的 Scalar RM,这抑制了其在 Online RL Pipeline 中的大规模使用
    • 然而,由于论文采用并行采样进行推理时扩展,使用合理数量的采样(例如8次)进行奖励生成的延迟不会显著增加
    • 围绕 LLM 高效生成和 RM 应用创新的进一步研究可能缓解此问题
  • (2) 在特定领域(如可验证任务)中,DeepSeek-GRM 仍然落后于 Scalar 模型
    • 这可能是因为 Scalar RM 捕获了推理 Query 和 Response 的隐藏特征 ,而 GRM 需要更强的推理能力来彻底检查 Response
      • 然而, Scalar RM 存在严重的偏见和扩展性问题
    • 对于 GRM,论文发现基于参考的奖励生成(附录 E.1.3)和长链条推理(附录 D.3)可以缓解这一局限
  • (3) 由于 Pointwise GRM 方法的普适性,DeepSeek-GRM 除了作为结果 RM 外,还可能作为过程 RM
    • 尽管论文在论文中没有深入探索这个方向,但在 Reward Bench 的 Reasoning 子集(主要包含 MATH-prm 数据 (2024))上的性能部分支持了这种应用的潜力

Future Direction

  • 基于 SPCT 或 DeepSeek-GRM 模型,未来研究有几个有希望的方向
  • (1) 先前工作研究了 RM 的工具集成 (2024b),也可用于 DeepSeek-GRM 增强
    • 使用诸如代码解释器和搜索引擎接口等工具 ,生成的 critiques 对于需要严格流程或广泛知识的任务可能更准确,并且可以避免 GRM 在遵循与数值计算、模式匹配等相关 Principle 时失败的情况
  • (2) Principle 和 critiques 的生成范式可以分解 为不同阶段,即 Principle 可以为每个待评分的 Query 和 Response 预先生成并存储,然后使用 GRM、规则或其他智能体方法生成 critiques
    • Principle 生成作为后续 critiques 的接口
    • 这可能会提高当前 GRM 集成到 RL Pipeline 中的效率
  • (3) DeepSeek-GRM 可能用于 LLM 离线评估
    • 由于每个 Principle 反映了一个标准,我们可以从特定 LLM 劣于另一个 LLM 的所有数据点中获取标准,作为解释该特定 LLM 弱点的可解释协议
    • 问题:实践发现,如果 Principle 是 Query-Specific 的,此时使用 Chosen 和 Rejected 来作为 Rubrics 生成参考容易出现过拟合,是否在通用的 Rubrics 中使用更合适?
  • (4) DeepSeek-GRM 可能受益于长链条推理
    • 然而,这会进一步影响其效率
    • 这些方向应在未来工作中进行研究

附录 C:Implementation Details

C.1 Model Training

  • 对于 Rule-based Online RL,论文使用标准的 GRPO 设置 (2024),总体目标函数为:
    $$
    \begin{align}
    \mathcal{J}_{\text{GRPO} }(\theta)=\mathbb{E}_{[q\sim P(Q),\{o_{i}\}_{i=1}^{G}\sim\pi_{\theta_{old} }(O|q)]} &\frac{1}{G} \sum_{i=1}^{G} \frac{1}{|o_{i}|} \sum_{t=1}^{|o|} \\
    &\left\{\min\left[\frac{\pi_{\theta}(o_{i,t}|q.o_{i<t})}{\pi_{\theta_{old} }(o_{i,t}|q.o_{i<t})}\hat{A}_{i,t}, \text{clip}\left(\frac{\pi_{\theta}(o_{i,t}|q.o_{i<t})}{\pi_{\theta_{old} }(o_{i,t}|q.o_{i<t})}, 1-\epsilon, 1+\epsilon\right)\hat{A}_{i,t}\right]-\beta\mathbb{D}_{KL}\left[\pi_{\theta}||\pi_{ref}\right]\right\},
    \end{align}
    $$
    • 其中 \(\hat{A}_{i,t}=\frac{\hat{r}_{t}-\text{mean}(\hat{t})}{\text{std}(\hat{t})}\),\(G\) 是组大小,\(\beta\) 是 KL 惩罚系数,\(q=(x,\{y_{i}\}_{i=1}^{n})\) 包含 prompts
    • 论文对超参数 \(\beta\in\{0.00,0.01,0.02,0.08\}\) 进行了网格搜索,发现 \(\beta=0.08\) 是 DeepSeek-GRM-27B 最稳定的配置
      • 当 KL 系数太小时,DeepSeek-GRM-27B 倾向于在基准测试的几个子集上崩溃,例如 Reward Bench 中的 Chat 子集和 RMB 中的 Harmlessness 子集,并对其他一些领域表现出偏见
      • 对于较小的 DeepSeek-GRM-16B,论文使用 \(\beta=0.002\),因为它对 KL 损失系数不那么敏感
    • 论文设置 \(G=4\) 以在效率和性能之间取得更好的平衡
  • 训练集包含 1256K RFT 数据,包括 \(1070\)K 通用指令数据和 \(186\)K 拒绝采样数据,以及 \(237\)K RL 数据
    • 通用指令数据来自内部数据集
    • 拒绝采样数据和 RL 数据来自相同的 RM 数据集,包含对单个、 Pairwise 和多个 Response 的偏好,这些数据由内部数据和开源数据集构建,包括来自 MATH (2021)、UltraFeedback (2024)、OffsetBias (2024)、Skywork-Reward-Preference-80K-v0.2 (2024) 和 HelpSteer2-Preference (2025b) 的训练集
      • Specifically,由于 UltraFeedback 的部分数据存在质量问题,论文重新标记了其偏好标签;
      • 论文根据 Rule-based ground-truth 匹配对 MATH 进行采样和过滤轨迹,生成 Pairwise 偏好数据;
    • 对于评分单个 Response ,论文将正确 Response 的 ground-truth 奖励设置为 1,错误 Response 的奖励设置为 0,仅纳入可验证的问题
    • 对于拒绝采样,论文使用 DeepSeek-v2.5-0905 生成带有 Principle 和 critiques 的轨迹
    • 采样次数 \(N_{\text{RFT} }\) 设置为 3
    • 在 HelpSteer2 上进行 hinted sampling 时,论文添加原始数据集中标注的偏好强度作为提示
    • 论文还从 RL 数据中移除了对 DeepSeek-V2-Lite-Chat 来说过于简单的样本 ,即根据公式 (10),所有生成的奖励在三次生成中都是正确的
  • DeepSeek-GRM 模型和 meta RM 的派生关系(derivation)如图 5 所示
    • 所有 DeepSeek-GRM 模型都是从 LLM 的预训练版本开始训练的
    • 对于 meta RM 的训练,论文复用了 RFT 阶段的拒绝采样数据,并使用 DeepSeek-GRM-27B 进行拒绝采样,\(N_{\text{RFT} }=3\),以避免 meta RM 引导投票中的潜在偏见 (2025)
    • meta RM 训练的学习率为 \(1\times 10^{-5}\),批次大小为 512
  • DeepSeek-GRM-27B 的 RFT 和 RL 训练时间如表 5 所示,基于 Gemma-2-27B 的模型在 Fire-Flyer 平台 (2024) 上使用 128 个 A100 GPU 进行训练
    • RFT 阶段的学习率为 \(5\times 10^{-6}\),RL 阶段的学习率为 \(4\times 10^{-7}\),RFT 阶段的批次大小为 1024,RL 阶段为 512
      • 两个阶段均训练 900 步
    • 由于资源限制,大于 27B 的 DeepSeek-GRM 模型未经过 Rule-based RL,仅使用 50K 拒绝采样数据进行训练

C.2 Baseline Implementation

  • 对于基线方法,论文基于 Gemma-2-27B (Team, 2024) 并采用与 DeepSeek-GRM 兼容的所有训练数据和设置,重新实现了 LLM-as-a-Judge (2023)、DeepSeek-BTRM-27B (Kendall & Smith, 1940)、CLoud-Gemma-2-27B (2024) 和 DeepSeek-PairRM-27B (2023)
  • 对于 LLM-as-a-Judge ,
    • 论文使用与 DeepSeek-GRM-27B 完全相同的训练配置,包括使用 DeepSeek-v2.5-0905 的拒绝采样数据进行 RFT 和 Rule-based Online RL
    • 由于其评分模式,RL 阶段只能使用 Pairwise 数据
    • 对于 CLoud-Gemma-2-27B ,论文也使用相同的 prompt 模板从 DeepSeek-v2.5-0905 生成 pointwise critiques
      • 然而,由于没有训练好的价值头就无法提取奖励,执行拒绝采样是不可行的
      • 论文使用 DeepSeek-GRM-27B 的相同通用指令数据以及采样的 critique 对 Gemma-2-27B 进行微调,得到一个 critique 生成模型
      • 具体来说,论文微调了另一个带有价值头的 Gemma-2-27B 模型用于奖励生成,而不是在 critique 模型上进行事后的价值头训练
    • CLoud-Gemma-2-27B 的价值头、DeepSeek-BTRM-27B 和 DeepSeek-PairRM-27B (2023) 的训练使用与 DeepSeek-GRM-27B 的 RL 阶段相同的数据集,但排除了单个 Response 评分数据

附录 D:Experiment Details

D.1 Hyper-Parameters

  • 对于 DeepSeek-GRM-27B、DeepSeek-GRM-16B、LLM-as-a-Judge 和 CLoud-Gemma-2-27B 的推理时扩展结果,每个模型的温度 (temperature) 设置为 0.5
    • 对于其他实验,所有模型的温度设置为 0
  • 在没有特定说明的情况下,DeepSeek-GRM-27B 的 meta RM 引导投票中默认 \(k_{\text{meta} }=\frac{1}{2}k\)
  • 对于 DeepSeek-R1-0120 的推理,温度设置为 0.6
  • 请注意,论文让 DeepSeek-GRM 在 Real.Mistake 基准测试中为单个 Response 评分时,输出的奖励范围与其他基准测试相同

D.2 Benchmarks

  • 论文在不同领域的各种 RM 基准上评估不同方法的性能:
    • (1) Reward Bench (RB) (2024),一个常用的 RM 评估基准,包含半自动收集的聊天 (2023; 2023; 2024)、推理 (2024; 2024) 和安全性 (Rö2024; 2024d) 偏好数据,其中每个 Query 需要对两个 Response 进行排序;
    • (2) PPE (2025),一个包含众包偏好数据和可验证任务正确性数据的大规模基准,每个 Query 有两个 Response ;
    • (3) RMB (2025),一个更全面的基准,包含各种类型的偏好数据,侧重于帮助性和无害性,每个 Query 有两个或更多 Response ,分别在 Pairwise 和 best-of-N (BoN) 子集中;
    • (4) Real.Mistake (2024),一个用于诊断单个 Response 中错误的基准
    • 论文在总体分数计算中不包括 Reward Bench 基准测试的 prior sets (2022a; 2021; 2022; 2020)
    • 对于报告的公开模型结果,论文使用每个基准发布的分数
      • gpt-4o 的版本略有不同,因为论文报告的是 gpt-4o-2024-08-06 在 Reward Bench 和 PPE(Correctness 子集使用 AlpacaEval prompt 模板复现)上的结果,以及 gpt-4o-2024-05-13 在 RMB 上的结果
  • 论文为每个基准使用标准评估指标:在 Reward Bench、PPE 和 RMB 中从一组 Response 中挑选最佳 Response 的准确度,在 Real.Mistake 中使用 ROC-AUC
    • RMB 基准测试的 BoN 子集每个 Query 包含多个 Response ,只有当最佳 Response 被识别时,每个数据点才被视为正确
    • 评估模型在 RMB BoN 子集上的默认设置是,如果总共有 \(n\) 个 Response,则 Pairwise 评估 \((n-1)\) 对(每对包含最佳 Response 和另一个不同的 Response)
      • 对于基线方法,论文采用这种方法进行评估
      • 而对于论文的模型 (DeepSeek-GRM),论文直接将所有 Response 输入模型,并通过 \(\arg\max_{i} S_i\) 识别最佳 Response ,其中 \(S_i\) 是第 \(i\) 个 Response 的预测奖励
        • 这是一种更直接但也更困难的方式,并且几乎不影响性能
        • 请参阅附录 E.1.1 的经验分析
        • 问题:这里再次强调了论文是同时将所有 Response 输入模型的
  • 对于 DeepSeek-R1-0120,由于推理成本和延迟巨大,论文从 Reward Bench 基准测试中均匀下采样了 300 个数据点,并在该子集上测试 DeepSeek-R1-0120
    • 结果如图 4(b) 所示

D.3 Detailed Results

  • 论文在图 6 中提供了图 1 的详细结果,并提供了更多公开模型的性能作为参考
  • 论文在表 6 中提供了表 3 的详细结果
  • 在表 7 中提供了表 4 的详细结果,并给出了每个 RM 基准测试的分数
  • 此外,论文列出了所有测试方法在每个 RM 基准测试上的详细结果,Reward Bench 基准测试的结果在表 8 中,PPE Correctness 基准测试在表 9 中,RMB 基准测试在表 10 中
    • 论文发现,DeepSeek-R1 在 Reward Bench 的 Reasoning 子集中取得了最高结果,表明长链条推理可以提升 GRM 在广泛推理场景中的表现

附录 E:Additional Experiments

E.1 Input Flexibility of the Pointwise GRM Approach

  • 在章节 2.1 中,论文从理论上论证了 pointwise GRM 方法的输入灵活性
    • 在本节中,论文提供了各种输入类型的经验证据来支持这一观点
E.1.1 Generating Rewards for Many Responses
  • 在表 11 中,论文展示了 DeepSeek-GRM-27B 在 RMB 基准测试 BoN 子集上的实验结果,其中每个 Query 有多个 Response
  • 如果总共有一个 Query 有 \(n, (n>2)\) 个 Response , Pairwise 输入设置是评估 \((n-1)\) 对,每对包含最佳 Response 和其他 Response ,只有当最佳 Response 从所有 \((n-1)\) 对中被正确识别时,该数据点才被视为正确
    • 这也是原始基准测试的默认设置
  • 论文比较了 DeepSeek-GRM-27B 在 Pairwise 输入和列表输入设置下的性能,列表输入设置是输入所有 \(n\) 个 Response 来识别最佳 Response
    • 结果表明,DeepSeek-GRM-27B 几乎不受输入类型的影响,在帮助性和无害性子集上的性能差异都小于 1%
    • 这表明 pointwise GRM 可以灵活地输入多个 Response ,并且性能对输入类型不敏感
E.1.2 Generating Rewards for Single Responses
  • 在表 13 中,论文展示了 DeepSeek-GRM-16B 和 DeepSeek-GRM-27B 在 Real.Mistake 基准测试上的实验结果,其中每个 Query 只有一个 Response
    • 论文与公开模型(如 DeepSeek-V2.5-0905、GPT-4o-2024-08-06、DeepSeek-V2-Lite 和 Gemma-2-27B-it)以及 DeepSeek-BTRM-27B 进行了比较
    • 结果显示,DeepSeek-GRM 在同等规模的模型中取得了最佳性能,并且通过推理时扩展,性能与最佳公开模型相当
    • 这表明 pointwise GRM 可以有效地对单个 Response 进行评分
E.1.3 Generating Rewards with Reference
  • 在章节 5.2 中,论文展示了 Scalar 和 Semi-scalar RM 可能存在显著的领域偏见,并且通常在可验证问题上表现更好
  • 为了缓解这个问题,论文测试了 DeepSeek-GRM-27B 在这些任务中使用参考(即每个 Query 的 ground truth)生成奖励的能力
  • 结果如表 12 所示
    • 论文发现,在提供参考的情况下,DeepSeek-GRM-27B 可以达到超过 90% 的准确率
    • 这表明 pointwise GRM 可以有效地根据参考判断 Response ,从而缓解了在可验证任务上的性能问题

E.2 Transferability of Generated Principles

  • 论文用 DeepSeek-GRM-27B 生成的 Principle 扩展了章节 2.2 中的初步实验
  • 论文测试了 GPT-4o-2024-08-06 和 DeepSeek-GRM-27B 使用与表 1 完全相同的手动过滤 Principle 以及上述 DeepSeek-GRM-27B 生成的 Principle
  • 结果如表 14 所示
    • 论文发现 DeepSeek-GRM-27B 生成的 Principle 可以迁移到其他模型,甚至比手动从 GPT-4o 过滤的 Principle 略好
    • 这表明 DeepSeek-GRM-27B 生成的 Principle 是稳健的且可迁移到其他模型

E.3 eneralization beyond Training Data

  • 论文对 DeepSeek-GRM-27B 训练数据的泛化能力进行了消融研究
  • 论文移除了 MATH 训练集中的所有数据,并重新实施了训练方案
  • 在 Reward Bench 基准测试上的结果如表 15 所示
    • 论文发现,仅添加与数学相关的偏好数据也可以提升通用 RM 在不同领域上的性能,尤其是在 Chat Hard 子集上
    • 结果表明 DeepSeek-GRM-27B 可以泛化到训练数据覆盖范围之外的领域

E.4 Response Length Analysis for Rule-Based RL

  • 论文在图 7 中计算了 DeepSeek-GRM-27B 在进行 Rule-based Online RL 前后在 Reward Bench 基准测试各子集上的 Response 长度
    • DeepSeek-GRM-27B 的 token 计数基于 Gemma-2-27B 的 tokenizer 计算,而 DeepSeek-R1-0120 的结果使用其对应的 tokenizer
    • 论文发现,Chat 子集的 Response 长度在 RL 后几乎没有增加,而 Safety 子集的 Response 长度甚至略有下降
    • Response 长度增加最大的是在 Reasoning 子集,根据表 8,DeepSeek-GRM-27B 在该子集上的性能相比 DeepSeek-GRM-27B-RFT 也提升最大
    • 这可能表明 DeepSeek-GRM-27B 学会了在广泛推理任务上自适应地使用更多推理计算资源,并且在模型学会准确生成 Principle 后,可以节省一些其他领域(如安全性)的计算
    • 然而,DeepSeek-R1-0120 使用了多得多的 tokens 却取得了更低的结果(Reasoning 除外),这表明长链条推理也有助于与广泛推理相关的 RM 任务

附录 F:Qualitative Analysis

F.1 Case Study

  • 论文在表 16、17 和 18 中提供了 DeepSeek-GRM-27B 的案例研究

  • 第一个案例表明, Scalar RM DeepSeek-BTRM-27B 在特定情况下可能被攻击或存在偏见,而 DeepSeek-GRM-27B 生成文本 Principle 和 critiques,表现出更好的鲁棒性

  • 第二个案例展示了 DeepSeek-GRM-27B 的可扩展行为,通过在多个样本上投票生成准确的奖励

    • 相应的 meta RM 分数也显示了 meta RM 在引导投票过程中的有效性
  • 第三个案例展示了 DeepSeek-GRM-27B 的潜在失败,这是由于模型无法根据某些 Principle(例如评估实时数据)准确判断 Response ,并且每个 Principle 的权重可能不平衡

    • 包含权重分配和分数收集的 critique 过程主要来自 RFT 和 Online RL 训练,这在不同的采样和测试案例中可能有很大差异
  • 示例核心:

    • 在输入 <Query, Response1, Response2> 后, DeepSeek-GRM-27B-Result 的输入类似:

      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      12
      13
      14
      15
      16
      17
      18
      19
      20
      21
      22
      23
      24
      Specific Criteria:
      1. **Depth of Analogies and Metaphors** (Weight: 30%): The extent
      to which the analogies and metaphors used are deep, creative, and
      effectively convey complex concepts.
      2. **Comprehensiveness of Core Information** (Weight: 30%): The
      thoroughness with which the core information essential to behavioral
      neuroscience is covered.
      3. **Clarity and Organization** (Weight: 20%): The clarity of explanations
      and the logical organization of the information.
      4. **Engagement and Appeal** (Weight: 20%): The level of engagement
      and appeal the response has, making the information interesting and
      memorable.
      Analysis:
      - **Response 1**: This response excels in comprehensiveness and clarity,
      covering all major aspects of behavioral neuroscience with detailed
      and well-organized explanations. The analogies and metaphors are
      well-crafted and effectively enhance understanding. However, some
      metaphors could be more creative to enhance engagement further.
      - **Response 2**: This response is also comprehensive and clear, but
      it stands out more in the depth and creativity of its analogies and
      metaphors. The organization is logical, and the explanations are clear.
      The engagement level is high, making the information appealing and
      memorable.
      Scores: \boxed{8, 9}
      • 特点:先给出具体的 Principle 和 权重,再进行打分分析,最终给出得分
    • 对于使用 Meta RM 的场景,DeepSeek-GRM-27B-Result 模型会回复多个上面的类似打分,每个都是 \boxed{8, 9} 的形式,然后会根据 Meta RM 的评估分数汇总 DeepSeek-GRM-27B-Result 的打分

      • 注意:在使用 Meta RM 时,在输入 DeepSeek-GRM-27B-Result 时,可能会随机调换 Response1 和 Response2 的顺序

Failure Mode Analysis

  • 论文从 DeepSeek-GRM-27B 在每个基准测试上的测试结果中随机抽样了 10 个错误数据点 (只有 10 个数据,诚意不够),并在图 8 中总结了失败模式
  • 对失败案例的分析表明,挑战主要在于模型无法判断过于复杂或属于特定领域(如模式匹配、计数等)的 Response ,以及缺乏专家知识,从而导致错误的 critiques
    • 尽管在大多数情况下 Principle 生成正确,但模型为每个 Principle 分配的权重会影响奖励的生成,有时会导致错误的结果
  • 然而,论文也发现
    • RM 基准测试中少数数据点的 ground truth 与人类标注者的偏好不一致
    • 可能是由于小规模人工标注研究的偏见或 ground truth 标注中的潜在错误

附录 G:Prompt Templates

  • 论文在下面展示了用于 DeepSeek-GRM、训练期间用于单个 Response 的 DeepSeek-GRM、meta-RM 以及 LLM-as-a-Judge 的 prompt 模板
  • 对于 prompt 工程,论文设计了一些示例 Principle ,用于上下文学习和基本的 critique 指导
  • 论文对 meta RM 使用更简洁的模板,以确保 Query 、 Response 以及生成的 Principle 和 critiques 能够适应上下文窗口
  • 在组装 meta RM 的模板后,论文进一步将内容封装在专为 DeepSeek-V3-1226 (DeepSeek-AI, 2024b) 设计的聊天模板中,然后再进行输入

DeepSeek-GRM (Default)

  • DeepSeek-GRM (Default) Prompt

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    You are a skilled little expert at scoring responses. You should evaluate given responses based
    on the given judging criteria.\n Given the context of the conversation (the last round is the
    User’s query) and multiple responses from the Assistant, you need to refer to the [General
    Evaluation Criteria] to score the responses. Based on the general evaluation criteria, state
    potential other specific criteria to the query, the weights of different criteria, and then provide
    an overall comprehensive score upon them.\n Each score is an integer between 1 and 10,
    with a higher score indicating that the response meets the relevant criteria more closely. For
    example, a score of 1 means the response does not meet the criteria at all, a score of 6 means
    the response meets only some parts, and a score of 10 means the response perfectly meets the
    evaluation criteria.\n Before scoring, please analyze step by step. Your scoring needs to be as
    strict as possible.
    #### Evaluation Criteria ####
    1. Instruction Adherence:\n - Fully Adhered (9-10 points): The response fully complies with
    all instructions and requirements of the question.\n - Partially Adhered (6-8 points): The
    response meets most of the instructions but has some omissions or misunderstandings.\n -
    Basically Adhered (3-5 points): The response meets some instructions, but the main
    requirements are not fulfilled.\n - Not Adhered (1-2 points): The response does not meet any
    instructions.\n Example: If the question requires three examples and the response provides
    only one, it falls under “Partially Adhered.”
    2. Usefulness:\n - Highly Useful (9-10 points): The response provides comprehensive and
    accurate information, fully addressing the issue.\n - Useful but Incomplete (6-8 points):
    The response provides some useful information, but lacks details or accuracy.\n - Limited
    Usefulness (3-5 points): The response offers little useful information, with most content
    being irrelevant or incorrect.\n - Useless or Incorrect (1-2 points): The response is completely
    irrelevant or incorrect.\n Example: If there are factual errors in the response but the overall
    direction is correct, it falls under “Useful but Incomplete.”
    3. Level of Detail:\n - Very Detailed (9-10 points): The response includes ample details
    covering all aspects of the issue.\n - Detailed but Slightly Lacking (6-8 points): The response
    is fairly detailed but misses some important details.\n - Basically Detailed (3-5 points): The
    response provides some details but is not thorough enough overall.\n - Not Detailed (1-2
    points): The response is very brief and lacks necessary details.\n Example: If the response
    provides only a simple conclusion without an explanation, it falls under “Not Detailed.”
    4. Relevance:\n - Highly Relevant (9-10 points): The response is highly relevant to the
    question, with information closely aligned with the topic.\n - Generally Relevant (6-8 points):
    The response is generally relevant but includes some unnecessary information.\n - Partially
    Relevant (3-5 points): The response has a lot of content that deviates from the topic.\n - Not
    Relevant (1-2 points): The response is completely irrelevant.\n Example: If the response strays
    from the topic but still provides some relevant information, it falls under “Partially Relevant.”
    #### Conversation Context ####\n{conversation context & query}\n
    #### Responses to be Scored ####
    [The Begin of Response i]\n{the i-th response}\n[The End of Response i]\n
    #### Output Format Requirements ####
    Output with three lines
    Specific Criteria: <Other potential criteria specific to the query and the context, and the
    weights of each criteria>.
    Analysis: <Compare different responses based on given Criteria>.
    Scores: <the overall comprehensive score of all responses in order, separate by comma in the
    boxed, e.g., \boxed{x, x} if there exists 2 responeses>.
    • DeepSeek-GRM Prompt 中,输入的 Reponse 可以多个

Meta RM

  • Meta RM Prompt
    1
    2
    3
    4
    5
    6
    7
    8
    **Prompt:**
    Please score the responses.
    #### Conversation Context ####\n{conversation context & query}\n
    #### Responses to be Scored ####
    [The Begin of Response i]\n{the i-th response}\n[The End of Response i]\n
    -----
    **Response:**
    {principle & critique}

LLM-as-a-Judge

  • LLM-as-a-Judge Prompt

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    You are a skilled little expert at scoring responses. You should evaluate given responses based
    on the given judging criteria.\nGiven the context of the conversation (the last round is the
    User’s query) and multiple responses from the Assistant, you need to refer to the [General
    Evaluation Criteria] to score the responses. Based on the general evaluation criteria, state
    potential other specific criteria to the query, the weights of different criteria, and then select
    the best response among all candidates.\nBefore judging, please analyze step by step. Your
    judgement needs to be as strict as possible.
    #### Evaluation Criteria ####
    1. Instruction Adherence:\n - Fully Adhered: The response fully complies with all instructions
    and requirements of the question.\n - Partially Adhered: The response meets most of the
    instructions but has some omissions or misunderstandings.\n - Basically Adhered: The
    response meets some instructions, but the main requirements are not fulfilled.\n - Not
    Adhered: The response does not meet any instructions.\n Example: If the question requires
    three examples and the response provides only one, it falls under “Partially Adhered.”
    2. Usefulness:\n - Highly Useful: The response provides comprehensive and accurate
    information, fully addressing the issue.\n - Useful but Incomplete: The response provides
    some useful information, but lacks details or accuracy.\n - Limited Usefulness: The response
    offers little useful information, with most content being irrelevant or incorrect.\n - Useless or
    Incorrect: The response is completely irrelevant or incorrect.\n Example: If there are factual
    errors in the response but the overall direction is correct, it falls under “Useful but Incomplete.”
    3. Level of Detail:\n - Very Detailed: The response includes ample details covering all aspects
    of the issue.\n - Detailed but Slightly Lacking: The response is fairly detailed but misses
    some important details.\n - Basically Detailed: The response provides some details but is not
    thorough enough overall.\n - Not Detailed: The response is very brief and lacks necessary
    details.\n Example: If the response provides only a simple conclusion without an explanation,
    it falls under “Not Detailed.”
    4. Relevance:\n - Highly Relevant: The response is highly relevant to the question, with
    information closely aligned with the topic.\n - Generally Relevant: The response is generally
    relevant but includes some unnecessary information.\n - Partially Relevant: The response has
    a lot of content that deviates from the topic.\n - Not Relevant: The response is completely
    irrelevant.\n Example: If the response strays from the topic but still provides some relevant
    information, it falls under “Partially Relevant.”
    #### Conversation Context ####\n{conversation context & query}\n
    #### Responses to be Scored ####
    [The Begin of Response]\n{the response}\n[The End of Response]\n
    #### Output Format Requirements ####
    Output with three lines
    Specific Criteria: <Other potential criteria specific to the query and the context, and the
    weights of each criteria>.
    Analysis: <Compare different responses based on given Criteria>.
    Scores: <the index of the best response based on the judgement, in the format of \boxed{x}>.
    • 除了输入 Response 只有一个以外,其他评估指标等好像和 DeepSeek-GRM(输入的 Reponse 可以多个) 的 Prompt 差不多

NLP——LLM对齐微调-DPO

注:本文包含 AI 辅助创作

  • 参考链接:
    • 原始论文:Direct Preference Optimization: Your Language Model is Secretly a Reward Model, NeurIPS 2023, Stanford University
    • 论文阅读笔记:DPO——RLHF 的替代之《Direct Preference Optimization: Your Language Model is Secretly a Reward Model》论文阅读
    • 一些较为高阶的讨论:RLHF的替代之DPO原理解析:从RLHF、Claude的RAILF到DPO、Zephyr
    • DPO 的应用:利用直接偏好优化算法(DPO)微调语言模型, 消除幻觉

Paper Summary

  • 核心总结:
    • 写在前面:本文的理论推导非常漂亮,值得细看
    • DPO 是一个简单的训练范式,用于从偏好中训练语言模型而无需强化学习(24 年底补充:DPO 已经成为了某些领域在 RL 前必不可少的一个基线)
    • DPO 识别了语言模型策略和奖励函数之间的映射(区别于传统 RL 的偏好学习方法)
      • 使得能够直接训练语言模型以满足人类偏好,使用简单的交叉熵损失,无需强化学习或损失一般性
    • 在作者的实验上,在几乎不需要调整超参数的情况下,DPO 的表现与现有的 RLHF 算法(包括基于 PPO 的算法)相当或更好;
      • DPO 有降低了从人类偏好训练更多语言模型的门槛
      • 实际上比 PPO 好是需要打问号的,PPO 的理论上限是更高的
  • 背景:RLHF 是一个复杂且通常不稳定的过程,需要拟合一个反映人类偏好的奖励模型,然后使用 RL 来微调以最大化这个估计的奖励,同时不能偏离原始模型太远
  • 论文介绍了 RLHF 中奖励模型的一个新参数化(parameterization)方法(DPO),能够 以闭式(closed form)解提取相应的最优策略
    • 这使得我们仅使用一个简单的分类损失来解决标准的 RLHF 问题

Introduction and Discussion

  • 论文展示了如何直接优化语言模型以遵循人类偏好,而无需显式的奖励建模或强化学习
  • 论文提出了直接偏好优化 (DPO) 算法
    • DPO 隐式地优化了与现有 RLHF 算法相同的目标(带有 KL 散度约束的奖励最大化),但实现简单且训练直接
  • 直观地说
    • DPO 更新增加了优选响应相对于非优选响应的相对对数概率,但它引入了一个动态的、每个样本的重要性权重,防止了论文发现的在朴素概率比目标下发生的模型退化
  • 与现有算法类似
    • DPO 依赖于一个理论偏好模型(例如 Bradley-Terry 模型 (1952)),该模型衡量给定奖励函数与经验偏好数据的对齐程度
  • 现有方法使用偏好模型来定义训练奖励模型的偏好损失,然后训练一个优化所学奖励模型的策略
    • DPO 使用变量变换将偏好损失直接定义为策略的函数
    • 给定一个关于模型响应的人类偏好的数据集,DPO 可以使用简单的二元交叉熵目标来优化策略,产生拟合到偏好数据的隐式奖励函数的最优策略
  • 论文的主要贡献就是 DPO,一种简单的、无 RL 的从偏好中训练语言模型的算法
  • 实验表明,在使用多达 6B 参数的语言模型进行情感调节、摘要和对话等任务的偏好学习中,DPO 至少与现有方法(包括基于 PPO 的RLHF)一样有效

Related Work

  • 规模不断增长的自监督语言模型学会零样本 (2019) 或少样本提示 (2020; 2021; 2022) 完成某些任务
    • 但它们在下游任务上的表现和与用户意图的对齐可以通过在指令和人类编写的补全数据集上进行微调来显著提高 (2022; 2022; 2022; 2022)
  • 这种“指令微调”过程使得大型语言模型能够泛化到指令微调集之外的指令,并通常提高其可用性 (2022)
    • 但响应质量的相对人类判断通常比专家演示更容易收集,因此后续工作使用人类偏好数据集微调了大型语言模型,提高了在翻译 (2018)、摘要 (2020; 2022)、故事讲述 (2022) 和指令遵循 (2022; 2023) 方面的熟练度
  • 这些方法的大致流程是:
    • 首先在诸如 Bradley-Terry 模型 (1952) 等偏好模型下,优化一个神经网络奖励函数以与偏好数据集兼容
    • 然后通常使用强化学习算法,如 REINFORCE (1992)、PPO 或其变体 (2023),来微调语言模型以最大化给定的奖励
  • 这些方法代表了两类工作的融合:
    • 一类是关于使用强化学习训练语言模型以实现各种目标的工作 (2015; 2018; 2018)
    • 另一类是关于从人类偏好中学习的通用方法的工作 (2017; 2018)
  • 尽管使用相对人类偏好具有吸引力,但使用强化学习微调大型语言模型仍然是一个重大的实际挑战;
    • DPO 提供了一种理论上有依据的方法来优化相对偏好而无需 RL

在语言领域之外

  • 从偏好中学习策略已在赌博机和强化学习设置中都得到了研究,并且已经提出了几种方法
  • 使用偏好或行动排序而非奖励的情境赌博机学习被称为情境对决赌博机 (Contextual Dueling Bandit, CDB; (2012; 2015))
    • 在没有绝对奖励的情况下,CDB 的理论分析用冯·诺依曼赢家 (von Neumann winner) 的概念替代了最优策略,即一个对任何其他策略的期望获胜率至少为 50% 的策略 (2015)
    • 在 CDB 设置中,偏好标签是在线给出的
      • 在从人类偏好学习中,通常从固定的离线偏好标注行动对批次中学习 (2022)
  • 基于偏好的 RL (Preference-based RL, PbRL) 从由 未知“评分” 函数生成的二元偏好中学习,而不是从奖励中学习 (2014; 2023)
    • 存在各种 PbRL 算法,包括可以重用离策略偏好数据的方法,但通常涉及首先显式估计潜在评分函数(即奖励模型),然后对其进行优化 (2013; 2014; 2017; 2018; 2018)
  • 论文提出了一种单阶段策略学习方法,直接优化策略以满足偏好

Preliminaries

  • 论文回顾一下 Ziegler 等人(以及后来的 (2020; 2022; 2022))中的 RLHF 流程,它通常包括三个阶段:
    • 1)监督微调 (Supervised Fine-Tuning, SFT);
    • 2)偏好采样和奖励学习;以及
    • 3)RL 优化

SF

  • RLHF 通常首先在下游感兴趣任务(对话、摘要等)的高质量数据上,使用监督学习对预训练的语言模型进行微调,以获得一个模型 \(\pi^{\text{SFT} }\)

奖励建模阶段

  • 使用提示 \(x\) 来提示 SFT 模型,产生成对的答案
    $$(y_{1},y_{2})\sim\pi^{\text{SFT} }(y \mid x)$$
  • 然将这些答案呈现给人类标注者,他们表达对其中一个答案的偏好,记为
    $$ y_{w}\succ y_{l} \mid x$$
    • 其中 \(y_{w}\) 和 \(y_{l}\) 分别表示 \((y_{1},y_{2})\) 中优选和非优选的补全
    • 假设偏好是由某个论文无法访问的潜在奖励模型 \(r^{*}(y,x)\) 生成的
  • 有多种方法用于建模偏好,Bradley-Terry (BT) 模型 (1952) 是一个流行的选择(尽管如果论文能访问多个排名的答案,更通用的 Plackett-Luce 排序模型 (1975; 2012) 也与该框架兼容)
  • BT 模型规定人类偏好分布 \(p^{*}\) 可以写成:
    $$p^{*}(y_{1}\succ y_{2} \mid x)=\frac{\exp\left(r^{*}(x,y_{1})\right)}{\exp\left(r^{*}(x,y_{1})\right)+\exp\left(r^{*}(x,y_{2})\right)} \tag{1}$$
  • 假设我们可以访问一个从 \(p^{*}\) 中采样的静态比较数据集
    $$ \mathcal{D}=\left\{x^{(i)},y^{(i)}_{w},y^{(i)}_{l}\right\}^{N}_{i=1}$$
  • 我们可以参数化一个奖励模型 \(r_{\phi}(x,y)\) 并通过最大似然估计参数
  • 将问题构建为二元分类,论文得到负对数似然损失:
    $$\mathcal{L}_{R}(r_{\phi},\mathcal{D})=-\mathbb{E}_{(x,y_{w},y_{l})\sim\mathcal{D} }\big{[}\log\sigma(r_{\phi}(x,y_{w})-r_{\phi}(x,y_{l}))\big{]} \tag{2}$$
    • 其中 \(\sigma\) 是 logistic 函数
    • 在语言模型的上下文中,网络 \(r_{\phi}(x,y)\) 通常从 SFT 模型 \(\pi^{\text{SFT} }(y \mid x)\) 初始化,并在最终 transformer 层之上添加一个线性层,该层为奖励值产生单个标量预测 (2022)
  • 为了确保奖励函数具有较低的方差,先前的工作对奖励进行了归一化,使得对于所有 \(x\),有
    $$\mathbb{E}_{x,y\sim\mathcal{D} }\left[r_{\phi}(x,y)\right]=0$$

RL 微调阶段

  • 在 RL 阶段,学习到的奖励函数用于向语言模型提供反馈
  • 遵循先前的工作 (2017; 2020),优化问题表述为
    $$\max_{\pi_{\theta} }\mathbb{E}_{x\sim\mathcal{D},y\sim\pi_{\theta}(y|x)}\big{[}r_{\phi}(x,y)\big{]}-\beta\mathbb{D}_{\text{KL} }\big{[}\pi_{\theta}(y \mid x) \ | \ |\ \pi_{\text{ref} }(y \mid x)\big{]} \tag{3}$$
    • 其中 \(\beta\) 是一个控制与基础参考策略 \(\pi_{\text{ref} }\)(即初始 SFT 模型 \(\pi^{\text{SFT} }\))偏离程度的参数
  • 在实践中,语言模型策略 \(\pi_{\theta}\) 也被初始化为 \(\pi^{\text{SFT} }\)
  • 添加的约束很重要,因为它防止模型偏离奖励模型准确的分布太远,同时保持生成多样性并防止模式崩溃到单一高奖励答案
  • 由于语言生成的离散性质,该目标不可微,通常使用强化学习进行优化
    • 标准方法 (2022; 2020; 2022; 2022) 是构建奖励函数
      $$r(x,y)=r_{\phi}(x,y)-\beta(\log\pi_{\theta}(y \mid x)-\log\pi_{\text{ref} }(y \mid x))$$
      • 并使用 PPO (2017) 进行最大化

Direct Preference Optimization

  • 目标是推导出一种使用偏好直接进行策略优化的简单方法
  • 与先前学习奖励然后通过 RL 优化它的 RLHF 方法不同,论文的方法利用了一种特定的奖励模型参数化选择,使得无需 RL 训练循环即可闭式解提取其最优策略
  • 论文的关键见解是利用从奖励函数到最优策略的解析映射 ,这使论文能够将关于奖励函数的损失函数转换为关于策略的损失函数
    • 这种变量变换方法避免了拟合显式的、独立的奖励模型,同时仍然在现有的人类偏好模型(如 Bradley-Terry 模型)下进行优化
    • 本质上,策略网络同时代表了语言模型和(隐式的)奖励

推导 DPO 目标

  • 论文从与先前工作相同的 RL 目标开始,即公式 3,在一个通用的奖励函数 \(r\) 下
  • 遵循先前的工作 (2007; 2019; 2022; 2023),很容易证明公式 3 中 KL 约束的奖励最大化目标的最优解形式为:
    $$\pi_{r}(y \mid x)=\frac{1}{Z(x)}\pi_{\text{ref} }(y \mid x)\exp\left(\frac{1}{\beta}r(x,y)\right) \tag{4}$$
    • 其中 \(Z(x)=\sum_{y}\pi_{\text{ref} }(y \mid x)\exp\left(\frac{1}{\beta}r(x,y)\right)\) 是配分函数
    • 完整的推导见附录 A.1
  • 即使论文使用真实奖励函数 \(r^{*}\) 的 MLE 估计 \(r_{\phi}\),估计配分函数 \(Z(x)\) 仍然很昂贵 (2022; 2023),这使得这种表示在实践中难以利用
  • 我们可以重新排列公式 4,用其对应的最优策略 \(\pi_{r}\)、参考策略 \(\pi_{\text{ref} }\) 和未知的配分函数 \(Z(\cdot)\) 来表示奖励函数
    • 论文首先对公式 4 的两边取对数,然后进行一些代数运算得到:
      $$r(x,y)=\beta\log\frac{\pi_{r}(y \mid x)}{\pi_{\text{ref} }(y \mid x)}+\beta\log Z(x). \tag{5}$$
  • 我们可以将这个重新参数化应用于真实奖励 \(r^{*}\) 和相应的最优模型 \(\pi^{*}\),注意:Bradley-Terry 模型仅依赖于两个补全之间奖励的差异,即
    $$ p^{*}(y_{1}\succ y_{2} \mid x)=\sigma(r^{*}(x,y_{1})-r^{*}(x,y_{2}))$$
  • 将公式 5 中 \(r^{*}(x,y)\) 的重新参数化代入偏好模型公式 1,配分函数会抵消掉,我们可以仅用最优策略 \(\pi^{*}\) 和参考策略 \(\pi_{\text{ref} }\) 来表示人类偏好概率
  • 因此,在 Bradley-Terry 模型下,最优的 RLHF 策略 \(\pi^{*}\) 满足偏好模型:
    $$p^{*}(y_{1}\succ y_{2} \mid x)=\frac{1}{1+\exp\left(\beta\log\frac{\pi^{*}(y_{2} \mid x)}{\pi_{\text{ref} }(y_{2} \mid x)}-\beta\log\frac{\pi^{*}(y_{1} \mid x)}{\pi_{\text{ref} }(y_{1} \mid x)}\right)} \tag{6}$$
    • 推导过程在附录 A.2 中
  • 虽然公式 6 使用了 Bradley-Terry 模型,但我们可以在更一般的 Plackett-Luce 模型 (1975; 2012) 下类似地推导出表达式,如附录 A.3 所示
  • 现在论文有了用最优策略而非奖励模型表示的人类偏好数据的概率,我们可以为参数化策略 \(\pi_{\theta}\) 制定一个最大似然目标;类似于奖励建模方法(即公式 2),论文的策略目标变为:
    $$\mathcal{L}_{\text{DPO} }(\pi_{\theta};\pi_{\text{ref} })=-\mathbb{E}_{(x,y_{w},y_ {l})\sim\mathcal{D} }\bigg{[}\log\sigma\left(\beta\log\frac{\pi_{\theta}(y_{w} \mid x)}{\pi_{\text{ref} }(y_{w} \mid x)}-\beta\log\frac{\pi_{\theta}(y_{l} \mid x)}{\pi_{\text{ref} }(y_{l} \mid x)}\right)\bigg{]}. \tag{7}$$
  • 这样,论文使用一种替代参数化来拟合一个隐式奖励,其最优策略简单地就是 \(\pi_{\theta}\)
  • 此外,由于论文的过程等同于拟合一个重新参数化的 Bradley-Terry 模型,它在偏好数据分布的适当假设下享有某些理论性质,例如一致性 (2022)
    • 在第 5 节中,论文将进一步讨论 DPO 相对于其他工作的理论性质

DPO 更新做了什么?

  • 为了从机制上理解 DPO,分析损失函数 \(\mathcal{L}_{\text{DPO} }\) 的梯度是很有用的
  • 关于参数 \(\theta\) 的梯度可以写成:
    $$\nabla_{\theta}\mathcal{L}_{\text{DPO} }(\pi_{\theta};\pi_{\text{ref } })=\ -\beta\mathbb{E}_{(x,y_{w},y_{l})\sim\mathcal{D} }\bigg{[}\underbrace{\sigma(\hat{r}_{\theta}(x,y_{l})-\hat{r}_{\theta}(x,y_{w}))}_{\text{higher weight when reward estimate is wrong} }\quad\bigg{[}\underbrace{\nabla_{\theta}\log\pi(y_{w} \mid x)}_{\text{increase likelihood of } y_{w}}-\underbrace{\nabla_{\theta}\log\pi (y_{l} \mid x)}_{\text{decrease likelihood of } y_{l} }\bigg{]}\bigg{]}$$
  • 其中 \(\hat{r}_{\theta}(x,y)=\beta\log\frac{\pi_{\theta}(y|x)}{\pi_{\text{ref} }(y|x)}\) 是由语言模型 \(\pi_{\theta}\) 和参考模型 \(\pi_{\text{ref} }\) 隐式定义的奖励(更多内容在第 5 节)
  • 直观地说,损失函数 \(\mathcal{L}_{\text{DPO} }\) 的梯度增加了优选补全 \(y_{w}\) 的似然,并降低了非优选补全 \(y_{l}\) 的似然
  • 样本的权重由隐式奖励模型 \(\hat{r}_{\theta}\) 对非优选补全评分高出多少来衡量,按 \(\beta\) 缩放,即隐式奖励模型对补全排序的错误程度,同时考虑了 KL 约束的强度
  • 论文的实验表明了这种加权的重要性,因为没有加权系数的朴素版本的方法会导致语言模型退化(附录表 3)

DPO outline

  • 一般的 DPO 流程如下:
    • 1)为每个提示 \(x\) 采样补全 \(y_{1},y_{2}\sim\pi_{\text{ref} }(\cdot \mid x)\),用人类偏好进行标注,以构建离线偏好数据集 \(\mathcal{D}=\{x^{(i)},y^{(i)}_{w},y_{l})^{(i)}\}_{i=1}^{N}\);
    • 2)优化语言模型 \(\pi_{\theta}\) 以最小化给定 \(\pi_{\text{ref} }\)、\(\mathcal{D}\) 和期望 \(\beta\) 的 \(\mathcal{L}_{\text{DPO} }\)
  • 在实践中,人们可能希望重用公开可用的偏好数据集,而不是生成样本并收集人类偏好
    • 由于偏好数据集是使用 \(\pi^{\text{SFT} }\) 采样的,只要可用,论文就初始化 \(\pi_{\text{ref} }=\pi^{\text{SFT} }\)
    • 但当 \(\pi^{\text{SFT} }\) 不可用时,论文通过最大化优选补全 \((x,y_{w})\) 的似然来初始化 \(\pi_{\text{ref} }\),即
      $$\pi_{\text{ref} }=\arg\max_{\pi}\mathbb{E}_{x,y_{w}\sim\mathcal{D} }\left[\log\pi(y_ {w} \mid x)\right]$$
    • 理解:
      • 当生成数据的模型不可访问时,可以考虑将损失函数中的 \(\pi_\text{ref}\) 替换为这个?
      • 上面的公式本质是找到一个策略 \(\pi\) 使得已知 \(x\) 时 \(y_ {w}\) 出现的对数概率最大化的策略 \(\pi_\text{ref}\)
  • 这个过程有助于缓解真实参考分布(不可用)与 DPO 使用的 \(\pi_{\text{ref} }\) 之间的分布偏移
  • 与实现和超参数相关的更多细节可以在附录 B 中找到

Theoretical Analysis of DPO

  • 本节进一步阐释 DPO 方法,提供理论支持,并将 DPO 的优势与用于 RLHF 的 Actor-Critic 算法(如 PPO (2017))存在的问题联系起来

Your Language Model Is Secretly a Reward Model

  • DPO 能够绕过显式拟合奖励函数和执行强化学习这两个步骤,仅使用一个最大似然目标来学习策略
  • 优化目标方程 5 等价于一个具有奖励参数化 的 Bradley-Terry 模型:
    $$ r^{*}(x,y)=\beta\log\frac{\pi_{\theta}(y|x)}{\pi_{\text{ref} }(y|x)} $$
    • 并且在变量变换下,论文优化参数化模型 \( \pi_{\theta} \) 等价于方程 2 中的奖励模型优化
  • 在本节中,论文将建立这种重新参数化背后的理论,证明它不会约束所学奖励模型的类别,并且允许精确恢复最优策略
  • 论文首先定义奖励函数之间的等价关系
    • 定义 1.
      • 两个奖励函数 \( r(x,y) \) 和 \( r’(x,y) \) 是等价的,当且仅当 \( r(x,y)-r’(x,y)=f(x) \) 对于某个函数 \( f \) 成立
      • 很容易看出这确实是一个等价关系,它将奖励函数集合划分为不同的类
      • 理解:因为 \(f(x) \) 与 \(y\) 没有关系!
    • 我们可以陈述以下两个引理:
      • 引理 1.
        • 在 Plackett-Luce(特别是 Bradley-Terry)偏好框架下,来自同一等价类的两个奖励函数诱导出相同的偏好分布
      • 引理 2.
        • 来自同一等价类的两个奖励函数在约束强化学习问题下诱导出相同的最优策略
  • 证明是直接了当的,论文将它们推迟到附录 A.5
    • 第一个引理是 Plackett-Luce 模型族 (1975) 中存在的一个众所周知的欠定(under-specification)问题
      • 由于这种欠定性,论文通常必须施加额外的可识别性约束,以保证从方程 2 得到的 MLE 估计具有某些性质 (2022)
    • 第二个引理指出,来自同一类的所有奖励函数产生相同的最优策略,因此对于论文的最终目标,论文只关心恢复最优类中的任意一个奖励函数
  • 论文在附录 A.6 中证明了以下定理:
    • 定理 1.
      • 在温和的假设下,与 Plackett-Luce(特别是 Bradley-Terry)模型一致的所有奖励类都可以用重新参数化 \( r(x,y)=\beta\log\frac{\pi(y|x)}{\pi_{\text{ref} }(y|x)} \) 来表示,对于某个模型 \( \pi(y \mid x) \) 和给定的参考模型 \( \pi_{\text{ref} }(y \mid x) \)
    • 证明概要.
      • 考虑任何奖励函数 \( r(x,y) \),它诱导出一个相应的最优模型 \( \pi_{r}(y \mid x) \),由方程 4 指定
      • 论文将证明,\( r \) 的等价类中的一个奖励函数可以使用上面给出的重新参数化来表示
      • 论文定义投影 \( f \) 为:
        $$
        f(r;\pi_{\text{ref} },\beta)(x,y)=r(x,y)-\beta\log\sum_{y}\pi_{\text{ref} }(y \mid x)\exp\left(\frac{1}{\beta}r(x,y)\right)
        $$
        • 算子 \( f \) 只是用 \( \pi_{r} \) 的配分函数的对数对奖励函数进行归一化
        • 由于添加的归一化项仅是前缀 \( x \) 的函数,所以 \( f(r;\pi_{\text{ref} },\beta)(x,y) \) 是 \( r(x,y) \) 等价类中的一个奖励函数
      • 最后,将 \( r \) 替换为方程 5 的右边(这对任何奖励函数都成立),论文有 \( f(r;\pi_{\text{ref} },\beta)(x,y)=\beta\log\frac{\pi_{r}(y|x)}{\pi_{\text{ref} }(y|x)} \)
      • 也就是说,投影 \( f \) 产生了 \( r \) 等价类中具有所需形式的一个成员,并且论文提出的重新参数化没有损失奖励模型的任何一般性
  • 也可以将定理 1 视为精确指定了 DPO 重新参数化在每个等价类中选择哪个奖励函数,即满足以下条件的奖励函数:
    $$
    \sum_{y}\underbrace{\pi_{\text{ref} }(y \mid x)}_{\equiv\pi(y|x), \text{ using Thm. 1 reparam.} }\exp\left(\frac{1}{\beta}r(x,y)\right)=1,
    $$
    • 也就是说,\( \pi(y \mid x) \) 是一个有效的分布(概率为正且和为 1)
    • 然而,遵循方程 4,我们可以看到方程 9 是由奖励函数 \( r(x,y) \) 诱导的最优策略的配分函数
  • DPO 算法的关键见解是,我们可以对欠定的 Plackett-Luce(特别是 Bradley-Terry)偏好模型族施加某些约束,使得论文能够保留可表示的奖励模型的类别,但同时显式地使得方程 4 中的最优策略对于所有提示 \( x \) 都是解析可解的

Instability of Actor-Critic Algorithms(AC 算法的不稳定性)

  • 也可以使用论文的框架来诊断用于 RLHF 的标准 Actor-Critic 算法(如 PPO)的不稳定性
  • 论文遵循 RLHF 流程,并专注于第 3 节概述的 RL 微调步骤
  • 我们可以与约束强化学习问题的控制即推断框架 (2018) 建立联系
  • 论文假设一个参数化模型 \( \pi_{\theta}(y \mid x) \),并最小化
    $$ \mathbb{D}_{\text{KL} }[\pi_{\theta}(y|x) \mid \mid \pi^{*}(y \mid x)] $$
    • 其中 \( \pi^{*} \) 是由奖励函数 \( r_{\phi}(y,x) \) 诱导的方程 7 中的最优策略
  • 经过一些代数运算,这可以推出优化目标:
    $$
    \max_{\pi_{\theta} }\mathbb{E}_{\pi_{\theta}(y|x)}\bigg{[}\underbrace{r_{\phi}(x ,y)-\beta\log\sum_{y}\pi_{\text{ref} }(y \mid x)\exp\left(\frac{1}{\beta}r_{\phi}(x,y)\right)}_{f(r_{\phi},\pi_{\text{ref} },\beta)}-\underbrace{\beta\log\frac{\pi_{\theta}(y \mid x)}{\pi_{\text{ref} }(y \mid x)} }_{\text{KL} }\bigg{]}
    $$
    • 这是先前工作 (2022; 2020; 2022; 2022) 使用与 \( r_{\phi} \) 的奖励类等价的 DPO 奖励所优化的相同目标
  • 在这种设置下,我们可以将 \( f(r_{\phi},\pi_{\text{ref} },\beta) \) 中的归一化项解释为参考策略 \( \pi_{\text{ref} } \) 的软价值函数
    • 虽然这个项不影响最优解,但没有它,目标的策略梯度可能具有高方差,使得学习不稳定
  • 我们可以使用一个学习的价值函数来适应这个归一化项,但这可能也难以优化
    • 或者,先前的工作使用人类完成基线来归一化奖励,本质上是归一化项的单样本蒙特卡洛估计
    • 相比之下,DPO 重新参数化产生了一个不需要任何基线的奖励函数

Experiments

  • 在本节中,论文实证评估 DPO 直接从偏好中训练策略的能力
  • 在一个受控良好的文本生成环境中,论文提出:
    • 与常见的偏好学习算法(如 PPO)相比,DPO 在最大化奖励和最小化与参考策略的 KL 散度之间进行权衡的效率如何?
  • 论文在更大的模型和更困难的 RLHF 任务上评估 DPO 的性能,包括摘要和对话
  • 在几乎不需要调整超参数的情况下,DPO 的表现往往与像基于 PPO 的 RLHF 这样的强基线一样好甚至更好,并且在学习的奖励函数下也优于从 \( N \) 个采样轨迹中返回最佳结果的方法

Experiment Setup

  • 注:更多细节在附录 C 中
Tasks
  • 论文的实验探索了三种不同的开放式文本生成任务
  • 对于所有实验,算法从下面的偏好数据集中学习一个策略
    $$ \mathcal{D}=\{x^{(i)},y_{w}^{(i)},y_{l}^{(i)} \}_{i=1}^{N} $$
  • 在受控情感生成(controlled sentiment generation)中
    • \( x \) 是来自 IMDb 数据集 (2011) 的电影评论前缀,策略必须生成具有积极情感的 \( y \)
    • 为了进行受控评估,在这个实验中,论文使用预训练的情感分类器生成关于生成的偏好对,其中
      $$ p(\text{positive} \mid x,y_{w})>p(\text{positive} \mid x,y_{l}) $$
    • 对于 SFT,论文在 IMDB 数据集的训练分割中的评论上微调 GPT-2-large 直到收敛(更多细节见附录 C.1)
  • 在摘要(summarization)任务中
    • \( x \) 是来自 Reddit 的论坛帖子;策略必须生成一个总结帖子要点的摘要 \( y \)
    • 遵循先前的工作,论文使用 Reddit TL;DR 摘要数据集 (2017) 以及 Stiennon 等人收集的人类偏好
    • 论文使用在人类撰写的论坛帖子摘要上微调的 SFT 模型,并使用 TRLX (2023) 框架进行 RLHF
    • 人类偏好数据集是由 Stiennon 等人在一个不同但训练方式相似的 SFT 模型的样本上收集的
  • 在单轮对话(single-turn dialogue)中
    • \( x \) 是一个人类查询,可能涉及从天体物理学问题到关系建议的任何内容;策略必须对用户的查询产生一个引人入胜且有用的响应 \( y \);
    • 论文使用 Anthropic Helpful and Harmless 对话数据集 (2022),包含 17 万个人类与自动助手之间的对话
    • 每个记录都以一个大型(尽管未知)语言模型生成的一对响应结束,并带有一个表示人类偏好响应的人类偏好标签
    • 在这种设置下,没有预训练的 SFT 模型可用;因此,论文仅在偏好完成上微调一个现成的语言模型以形成 SFT 模型
Evaluation
  • 论文的实验使用两种不同的评估方法
  • 为了分析每种算法在优化约束奖励最大化目标方面的有效性,在受控情感生成设置中,论文通过其实现的奖励和与参考策略的 KL 散度的边界来评估每种算法;
    • 这个边界是可计算的,因为我们可以访问真实奖励函数(一个情感分类器)
  • 但在现实世界中,真实奖励函数是未知的;
    • 论文使用算法与基线策略的胜率来评估算法,使用 GPT-4 作为摘要质量和响应帮助度的代理评估,分别用于摘要和单轮对话设置
  • 对于摘要,论文使用测试集中的参考摘要作为基线;对于对话,论文使用测试数据集中的偏好响应作为基线。虽然现有研究表明 LM 可以比现有指标更好的自动评估器 (2023),但论文在第 6.4 节进行了一项人类研究来证明论文使用 GPT-4 进行评估的合理性。论文发现 GPT-4 的判断与人类高度相关,人类与 GPT-4 的一致性通常与人类注释者之间的一致性相似或更高
Methods
  • 除了 DPO,论文还评估了几种现有的训练语言模型以符合人类偏好的方法
  • 论文在摘要任务中探索了使用 GPT-J (2021) 的零样本提示,在对话任务中探索了使用 Pythia-2.8B (2023) 的 2-shot 提示
  • 论文还评估了 SFT 模型以及 Preferred-FT
    • 这是一个在选自 SFT 模型(在受控情感和摘要中)或通用 LM(在单轮对话中)的所选完成 \( y_{w} \) 上使用监督学习进行微调的模型
  • 另一种伪监督方法是 Unlikelihood (2019)
    • 它简单地优化策略以最大化分配给 \( y_{w} \) 的概率,并最小化分配给 \( y_{l} \) 的概率;
    • 论文在“非似然”项上使用一个可选的系数 \( \alpha\in[0,1] \)
  • 论文还考虑了使用从偏好数据学习的奖励函数的 PPO (2017) 和 PPO-GT ,后者是在受控情感设置中可访问真实奖励函数的 Oracle
    • 在论文的情感实验中,论文使用了 PPO-GT 的两种实现,一个现成的版本 (2023) 以及一个修改版本,该版本归一化奖励并进一步调整超参数以提高性能(论文在运行使用学习奖励的“普通”PPO 时也使用这些修改)
  • 最后,论文考虑了 Best of \( N \) 基线,从 SFT 模型(或对话中的 Preferred-FT)采样 \( N \) 个响应,并返回根据从偏好数据集学习的奖励函数得分最高的响应
  • 这种高性能方法将奖励模型的质量与 PPO 优化解耦,但在计算上是不切实际的,即使对于中等的 \( N \),因为它在测试时需要为每个查询采样 \( N \) 个完成
  • 图 2:
    • 左图:期望奖励与相对于参考策略的 KL 散度的边界。DPO 在所有 KL 值下提供了最高的期望奖励,证明了优化的质量
    • 右图:TL;DR 摘要相对于人工撰写摘要的胜率,使用 GPT-4 作为评估器
    • DPO 在摘要任务上超过了 PPO 的最佳性能,同时对采样温度的变化更加鲁棒

How well can DPO optimize the RLHF objective?

  • 典型 RLHF 算法中使用的 KL 约束奖励最大化目标在利用奖励的同时限制策略偏离参考策略太远之间取得平衡
  • 因此,在比较算法时,论文必须同时考虑实现的奖励以及 KL 差异;
    • 实现略高的奖励但 KL 高得多并不一定是可取的
  • 图 2 显示了在情感设置中各种算法的奖励-KL 边界
    • 论文对每种算法执行多次训练运行,每次运行使用不同的策略保守性超参数(PPO 的目标 KL \( \in\{3,6,9,12\} \),DPO 的 \( \beta\in\{0.05,0.1,1,5\} \),Unlikelihood 的 \( \alpha\in\{0.05,0.1,0.5,1\} \),Preferred-FT 的随机种子)
    • 扫描总共包括 22 次运行
    • 在每 100 个训练步骤直到收敛后,论文在一组测试提示上评估每个策略,计算在真实奖励函数下的平均奖励以及
  • 与参考策略的平均序列级 KL3 \( \text{KL} (\pi \mid \mid \pi_{\text{ref} }) \)
    • 论文发现 DPO 产生了迄今为止最有效的边界,实现了最高的奖励,同时仍然实现了低 KL
    • 这个结果尤其引人注目,原因有几个
      • 首先,DPO 和 PPO 优化相同的目标,但 DPO 明显更有效;
        • DPO 的奖励/KL 权衡严格优于 PPO
      • 其次,DPO 实现了比 PPO 更好的边界,即使 PPO 可以访问真实奖励(PPO-GT)
  • 图 3:
    • 左图:由 GPT-4 计算的 Anthropic-HH 单轮对话胜率;DPO 是唯一一个在 Anthropic-HH 测试集上优于所选摘要的方法
    • 右图:在训练过程中不同采样温度的胜率
    • 对于不同的采样温度,DPO 相对于数据集标签的改进在训练过程中相当稳定

Can DPO scale to real preference datasets?(将 DPO 扩展到真实的偏好数据集中)

  • 接下来,论文评估 DPO 在摘要和单轮对话上的微调性能
  • 对于摘要,自动评估指标如 ROUGE 可能与人类偏好相关性很差 (2020),并且先前的工作发现使用 PPO 在人类偏好上微调 LM 可以提供更有效的摘要
    • 论文通过在 TL;DR 摘要数据集的测试分割上采样完成来评估不同方法,并计算相对于测试集中参考完成的平均胜率
    • 所有方法的完成都是在从 0.0 到 1.0 变化的温度下采样的,胜率如图 2(右)所示
    • DPO、PPO 和 Preferred-FT 都微调了相同的 GPT-J SFT 模型4
      • 论文发现 DPO 在温度为 0.0 时胜率约为 61%,超过了 PPO 在其最佳采样温度 0.0 时的 57%
      • 与 Best of \( N \) 基线相比,DPO 也实现了更高的最大胜率
      • 论文没有刻意调整 DPO 的 \( \beta \) 超参数,所以这些结果可能低估了 DPO 的潜力
    • 论文还发现 DPO 对采样温度的鲁棒性远高于 PPO,后者的性能在高温度下可能退化到基础 GPT-J 模型的水平
      • Preferred-FT 相对于 SFT 模型没有显著改进
      • 论文还在第 6.4 节中在人类评估中比较了 DPO 和 PPO,其中温度为 0.25 的 DPO 样本在 58% 的情况下比温度为 0 的 PPO 样本更受偏好
  • 在单轮对话上,论文在 Anthropic HH 数据集 (2022) 测试分割的子集上评估不同方法,该子集包含一步人机交互
    • GPT-4 评估使用测试集上的偏好完成作为参考,计算不同方法的胜率
    • 由于此任务没有标准的 SFT 模型,论文从一个预训练的 Pythia-2.8B 开始,使用 Preferred-FT 在所选完成上训练一个参考模型,使得完成在模型的分布内,然后使用 DPO 进行训练
    • 论文还与从 Preferred-FT 完成中选出的最佳 128 个完成进行比较(论文发现 Best of \( N \) 基线在此任务上在 128 个完成时达到稳定;见附录图 4)以及一个 2-shot 提示的 Pythia-2.8B 基础模型版本,发现 DPO 在每种方法的最佳性能温度下表现相当或更好
    • 论文还评估了一个在 Anthropic HH 数据集5 上使用 PPO 训练的 RLHF 模型,该模型来自一个知名来源6,但未能找到能够提供优于基础 Pythia-2.8B 模型性能的提示或采样温度
    • 基于论文从 TL;DR 得到的结果以及两种方法优化相同奖励函数的事实,论文将 Best of 128 视为 PPO 级别性能的粗略代理
    • 总体而言,DPO 是唯一一个在计算上高效且改进了 Anthropic HH 数据集中偏好完成的方法,并且提供了与计算量大的 Best of 128 基线相似或更好的性能
    • 最后,图 3 显示 DPO 相对较快地收敛到其最佳性能

Generalization to a new input distribution

  • 为了进一步比较 PPO 和 DPO 在分布偏移下的性能,论文在不同的分布上评估了来自论文 Reddit TL;DR 摘要实验的 PPO 和 DPO 策略,即 CNN/DailyMail 数据集 (2016) 测试分割中的新闻文章,使用来自 TL;DR 的最佳采样温度(0 和 0.25)
  • 结果呈现在表 1 中
    • 论文计算了相对于数据集中真实摘要的 GPT-4 胜率,使用了论文用于 Reddit TL;DR 的相同 GPT-4 (C) 提示,但将“论坛帖子”一词替换为“新闻文章”
    • 对于这个新的分布,DPO 继续以显著优势优于 PPO 策略。这个实验提供了初步证据,表明 DPO 策略可以像 PPO 策略一样很好地泛化,尽管 DPO 没有使用 PPO 使用的额外的未标记 Reddit TL;DR 提示
  • 表 1:对于分布外的 CNN/DailyMail 输入文章,相对于真实摘要的 GPT-4 胜率
    • 理解:因为是跟 Ground Truth 的对比,所以他们的省略都不足 50%

Validating GPT-4 judgments with human judgments(用 Human Judgment 验证 GPT-4 judgment 是否准确)

  • 论文进行了一项人类研究来验证 GPT-4 判断的可靠性 ,使用了 TL;DR 摘要实验的结果和两种不同的 GPT-4 提示
    • GPT-4 (S)(简单)提示简单地询问哪个摘要更好地总结了帖子中的重要信息
    • GPT-4 (C)(简洁)提示还询问哪个摘要更简洁;
    • 论文评估这个提示是因为论文发现使用 GPT-4 (S) 提示时,GPT-4 比人类更喜欢更长、更重复的摘要
    • 完整的提示见附录 C.2
  • 论文进行了三次比较,使用了性能最高(DPO,温度 0.25)、最低(PPO,温度 1.0)和中等性能(SFT,温度 0.25)的方法,旨在覆盖多样性的样本质量;
    • 所有三种方法都与贪婪采样的 PPO(其最佳性能温度)进行比较
    • 论文发现,使用两种提示,GPT-4 与人类一致的程度通常与人类彼此一致的程度相似,这表明 GPT-4 是人类评估的合理代理(由于人类评分者有限,论文只收集了 DPO 和 PPO-1 比较的多人判断)
  • 总体而言,GPT-4 (C) 提示通常提供更能代表人类的胜率;因此论文在第 6.2 节的主要结果中使用此提示
  • 关于人类研究的更多细节,包括呈现给评分者的网络界面和人类志愿者列表,见附录 D.3
  • 表 2:
    • 在 TL;DR 摘要样本上比较人类和 GPT-4 的胜率以及每次判断的一致性
    • 人类与 GPT-4 的一致程度与人类彼此之间的一致程度大致相同
    • 每个实验将所述方法的摘要与温度为 0 的 PPO 摘要进行比较

Limitations & Future Work

  • 论文的结果提出了几个未来工作的重要问题
  • 第一:与从显式奖励函数学习相比,DPO 策略在分布外的泛化能力如何?
    • 论文的初步结果表明 DPO 策略可以类似于基于 PPO 的模型一样泛化,但需要更全面的研究
    • 例如,使用来自 DPO 策略的自标注进行训练是否同样能有效利用未标记的提示?
  • 第二:在直接偏好优化设置中,奖励过度优化(over-optimization)如何表现
    • 图 3 右侧性能的轻微下降是否是它的一个实例?
  • 第三:虽然论文评估了最多 6B 参数的模型,但将 DPO 扩展到数量级更大的最先进模型是未来工作的一个令人兴奋的方向
  • 关于评估
    • 论文发现 GPT-4 计算的胜率受到提示的影响;未来的工作可能会研究从自动化系统中引出高质量判断的最佳方式
  • 最后,DPO 的许多可能应用存在于从人类偏好训练语言模型之外,包括在其他模态中训练生成模型

附录 B:DPO Implementation Details and Hyperparameters

  • DPO 的实现相对简单;下面提供了 DPO 损失的 PyTorch 代码:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    import torch.nn.functional as F

    def dpo_loss(pi_logps, ref_logps, yw_idxs, yl_idxs, beta):
    """
    pi_logps: policy logprobs, shape (B,)
    ref_logps: reference model logprobs, shape (B,)
    yw_idxs: preferred completion indices in [0, B-1], shape (T,)
    yl_idxs: dispreferred completion indices in [0, B-1], shape (T,)
    beta: temperature controlling strength of KL penalty
    Each pair of (yw_idxs[i], yl_idxs[i]) represents the indices of a single preference pair.
    """

    pi_yw_logps, pi_yl_logps = pi_logps[yw_idxs], pi_logps[yl_idxs]
    ref_yw_logps, ref_yl_logps = ref_logps[yw_idxs], ref_logps[yl_idxs]

    pi_logratios = pi_yw_logps - pi_yl_logps
    ref_logratios = ref_yw_logps - ref_yl_logps

    losses = -F.logsigmoid(beta * (pi_logratios - ref_logratios))
    rewards = beta * (pi_logps - ref_logps).detach()

    return losses, rewards
  • 论文默认使用 \(\beta = 0.1\),批量大小为 64,以及学习率为 1e-6 的 RMSprop 优化器

  • 论文在前 150 步中将学习率从 0 线性预热到 1e-6

  • 对于 TL;DR 摘要任务,论文使用 \(\beta = 0.5\),其余参数保持不变


附录 C:Further Details on the Experimental Set-Up

  • 在本节中,论文包含了与实验设计相关的额外细节

C.1 IMDb 情感实验和基线细节 (C.1 IMDb Sentiment Experiment and Baseline Details)

  • 提示是来自 IMDB 数据集中长度为 2-8 个词符的前缀
  • 论文使用预训练的情感分类器 siebert/sentiment-roberta-large-english 作为真实奖励模型,并使用 gpt2-large 作为基础模型
  • 论文发现默认模型生成的文本质量较低且奖励有些不准确,因此使用了这些更大的模型
    • 论文首先在 IMDB 数据的一个子集上进行了 1 个周期的监督微调
    • 然后,论文使用该模型为 25000 个前缀生成了 4 个补全,并使用 ground-truth 奖励模型为每个前缀创建了 6 个偏好对
      • 问题:这里的 ground-truth 奖励模型是什么?
    • RLHF 奖励模型从 gpt2-large 模型初始化,并在偏好数据集上训练了 3 个周期,论文选择了在验证集上准确率最高的检查点
    • “TRL” 运行使用了 TRL 库中的超参数
    • 论文的实现在每个 PPO 步骤中使用更大的批量样本,为 1024

C.2 用于计算摘要和对话胜率的 GPT-4 提示 (C.2 GPT-4 prompts for computing summarization and dialogue win rates)

  • 论文实验设置的一个关键组成部分是 GPT-4 的胜率判断
  • 在本节中,论文包含了用于生成摘要和对话实验胜率的提示
  • 论文所有的实验都使用 gpt-4-0314
  • 每次评估时,摘要或响应的顺序是随机选择的
摘要 GPT-4 胜率提示 (S) (Summarization GPT-4 win rate prompt (S))
  • 提示词:
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    Which of the following summaries does a better job of summarizing the most \
    important points in the given forum post?
    Post:
    <post>
    Summary A:
    <Summary A>
    Summary B:
    <Summary B>
    FIRST provide a one-sentence comparison of the two summaries, explaining which \
    you prefer and why. SECOND, on a new line, state only "A" or "B" to indicate your \
    choice. Your response should use the format:
    Comparison: <one-sentence comparison and explanation>
    Preferred: <"A" or "B">
摘要 GPT-4 胜率提示 (C) (Summarization GPT-4 win rate prompt (C))
  • 提示词:
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    Which of the following summaries does a better job of summarizing the most \
    important points in the given forum post, without including unimportant or \
    irrelevant details? A good summary is both precise and concise.
    Post: <post>
    Summary A:
    <Summary A>
    Summary B:
    <Summary B>
    FIRST provide a one-sentence comparison of the two summaries, explaining which \
    you prefer and why. SECOND, on a new line, state only "A" or "B" to indicate your \
    choice. Your response should use the format:
    Comparison: <one-sentence comparison and explanation>
    Preferred: <"A" or "B">
Dialogue GPT-4 win rate prompt
  • 提示词:
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    For the following query to a chatbot, which response is more helpful?
    Query: <the user query>
    Response A:
    <either the test method or baseline>
    Response B:
    <the other response>
    FIRST provide a one-sentence comparison of the two responses and explain \
    which you feel is more helpful. SECOND, on a new line, state only "A" or \
    "B" to indicate which response is more helpful. Your response should use \
    the format:
    Comparison: <one-sentence comparison and explanation>
    More helpful: <"A" or "B">

C.3 非似然基线 (C.3 Unlikelihood baseline)

  • 虽然论文在情感实验中包含了非似然基线 (2019)(简单地最大化偏好响应的对数概率 \(\log p(y_{w}|x)\),同时最小化非偏好响应的对数概率 \(\log p(y_{l}|x)\)),但论文没有在摘要或对话实验中将之作为基线,因为它通常会产生无意义的响应,作者认为这是无约束似然最小化的结果

附录 D:Additional Empirical Results

D.1 不同 N 值的 Best of N 基线性能 (D.1 Performance of Best of N baseline for Various N)

  • 论文发现 Best of N 基线在论文的实验中是一个强大(尽管计算成本高,需要多次采样)的基线
  • 论文评估了 Best of N 基线在 Anthropic-HH 对话和 TL;DR 摘要任务中不同 N 值的性能;结果如图 4 所示

D.2 样本回复和 GPT-4 判断 (D.2 Sample Responses and GPT-4 Judgments)

  • 在本节中,论文展示了 DPO 与基线(摘要任务为 PPO temp 0,对话任务为数据集中选择的真实响应)之间比较的示例
  • 摘要示例见表 4-6,对话示例见表 7-10
  • 论文未列出,详情见博客

D.3 人类研究细节 (D.3 Human study details)

  • 为了验证使用 GPT-4 计算胜率的做法,论文的人类研究在 TL;DR 摘要设置中收集了几组对比的人类偏好数据
  • 论文选择了三种不同的算法对比,评估了 DPO (temp. 0.25)、SFT (temp. 0.25) 和 PPO (temp 1.0) 与参考算法 PPO (temp 0.) 的比较
  • 通过选择三种独特算法以及相对于参考算法胜率范围广泛的算法,论文捕捉了人类和 GPT-4 胜率在不同响应质量谱上的相似性
  • 论文对 DPO vs PPO-0 进行了 150 次随机比较采样,对 PPO-1 vs PPO-0 进行了 100 次随机比较采样,每次比较分配两个人进行标注,为 DPO-PPO’10 产生了 275 个判断,为 PPO-PPO 产生了 200 个判断
  • 论文对 SFT 进行了 125 次比较采样,每次分配一个人进行标注
  • 论文忽略了人类标记为平局的判断(仅占所有判断的约 1%),并测量了人类 A 和人类 B 之间(对于论文有两个人标注的比较,即不包括 SFT)以及每个人与 GPT-4 之间的原始一致百分比
  • 论文总共有 25 名志愿者人类评估者,每人比较了 25 个摘要(一名志愿者较晚完成调查,未纳入最终分析,但在此列出)
  • 评估者是斯坦福大学的学生(从本科到博士),或近期的斯坦福毕业生或访问学者,主修 STEM(主要是 CS)领域
  • 调查界面截图见图 5

附录:DPO 推导过程

  • 根据RLHF中的定义,策略 \(\pi_\theta\) 的训练目标是:在不太偏离 Reference Model 的情况下,最大化 reward,其定义可以写为如下形式
    $$
    \max_{\pi_\theta} \mathbb{E}_{x\sim D, y\sim \pi_\theta(y|x)} [r_{RM}(x,y)]-\beta \mathbb{D}_{KL}[\pi_\theta(y|x)||\pi_\text{ref}(y|x)]
    $$
    • \(r_{RM}(x,y)\) 是 Reward Model 的返回值
    • \(\beta\) 是控制目标策略和 Reference Model 差异的参数
    • 在强化学习中,可通过修改使用 reward(Reward Engineering)来实现训练目标:
      $$
      r(x,y) = r_{RM}(x,y) - \beta(\log\pi_\theta(y|x) - \log\pi_\text{ref}(y|x) )
      $$
  • DPO推导:
    • DPO-最优策略推导 :将原始目标进行转换有
      • 理解:(论文中未明确给出说明)上述证明中,转换 KL 散度形式为对形式时,要求 \(y \sim \pi\),即回答 \(y\) 是从当前策略采样的(实际中我们无法严格满足,只能近似做到)
      • 公式(12)最后一步推导为:将左边式子中括号内容按照对数法则全部展开,即可发现变化前后的式子相等
      • 实际上,推导公式(12)时, \(\log Z(x)\) 可以是任意函数,为了使得 \(\pi^{*}(y|x)\) 是一个分布,所以特意将 \(Z(x)\) 取值为分配函数的形式(物理学中的分配函数是一个统计量)
      • 注:可以证明这样构造的 \(\pi^*\) 正是原始 PPO 约束优化问题的最优解(这也是后续推导可以看到后续的目标就是让目标策略和 \(\pi^*\) 的 KL 散度变小),详细证明参考 RL——CQL 的附录部分
    • DPO-最优奖励形式推导 :进一步推导得到最优策略对应的奖励形式有
      • 这里由于 \(Z(x)\) 并不依赖于策略 \(\pi\),所以最小化公式14 的结果就是最小化其第一项的结果(注:第二项可以单独分出来),且 \(\pi\) 满足下面的形式时取得最小值,即在 KL 约束下,奖励最大化目标的 \(\pi(y|x)\) 的最优解形式为:
        $$
        \pi^{*}(y|x) = \frac{1}{Z(x)}\pi_\text{ref}(y|x)exp(\frac{1}{\beta}r(x,y))
        $$
      • 两边同时取对数可得:
        $$
        \log \pi^{*}(y|x) = -\log Z(x) + \log \pi_\text{ref}(y|x) + \frac{1}{\beta}r(x,y)
        $$
      • 进一步有:
        $$
        \frac{1}{\beta}r(x,y) = \log \pi^{*}(y|x) + \log Z(x) - \log \pi_\text{ref}(y|x)
        $$
      • 即最优策略对应的奖励形式为:
        $$
        r(x,y) = \beta \log \frac{\pi^{*}(y|x)}{\pi_\text{ref}(y|x)} + \beta\log Z(x)
        $$
        • 对上述奖励形式的理解:如果已知 公式11 问题中的 最优解 ,即 策略 \(\pi\) 的最优策略 \(\pi^{*}\) ,那么 最优的奖励值(最大奖励值)\(r^*(x,y)\) 的形式可以通过上述的形式来表示
    • DPO-损失函数推导 :在使用 Bradley-Terry (BT)方法建模人类偏好时 ,有 DPO 的训练损失函数如下:
      • 推导精髓:上述推导中消除了 \(Z(x)\) ,是通过引入 Bradley-Terry 偏好模型来实现的
      • 理解:上述基本推导思路是
        • 奖励和策略的关系:由上面已经已知,即指定任意策略为最优策略,则其对应的最优奖励就已知;反过来也一样
        • 奖励和概率的关系:借助 BT 偏好建模,可以建模 正样本 优于 负样本的概率 \(p(y_1 \succ y_2|x)\) 和奖励 \(r(x,y_1),r(x,y_2)\) 的关系
        • 由此,可以得到概率和策略的关系(注意:这样的好处是,不再需要预估奖励了,直接从优化概率即可优化策略)
        • 下面再进一步,最小化负对数概率即可得到目标,当负对数概率最小时,也就是策略最优时
      • 目标是让上面的概率 \(p^*(y_1 \succ y_2 )\) 越大越好,最终的损失函数可以定义为负对数概率在数据集上的期望:
        $$
        Loss_{\text{DPO}}(\pi_\theta;\pi_\text{ref}) = - \mathbb{E}_{(x,y_w,y_l) \sim D}\left [ \log \sigma \left( \beta\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} - \beta\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} \right)\right ]
        $$
      • 通过最小化上面的损失函数,即可求得最优参数 \(\theta^*\)

附录:DPO 训练流程

  • 一句话总结:用 \(\pi_\text{ref}\) 收集数据 + 人类离线标记偏好 + 按照 DPO 损失函数训练
  • 【与 RLHF Reward Model 训练相同】使用SFT模型对同一个Prompt重复生成多个不同Response,记为 \(y_a, y_b, y_c, y_d\)
  • 【与 RLHF Reward Model 训练相同】让人类对回答进行排序打分: \(y_a \succ y_b \succ y_c \succ y_d\) ,最终可获得数据集 \(D = \{x^{(i)}, y^{(i)}_w, y^{(i)}_l\}_{i=1}^N\)
    • 实际上,对于一个 \(x\) 对应K个 \(y\) 的情况,可以拆开成 \(C_K^2\) 个样本,然后按照Bradley-Terry(BT)Model来建模,在OPENAI原始论文:Training language models to follow instructions with human feedback中训练RM时是这样表述的,论文均按照BT模型来写
    • 实际上,DPO原始论文:Direct Preference Optimization: Your Language Model is Secretly a Reward Model中还给出了一种不做样本拆分的等价表达形式:Plackett-Luce Model,在论文最后我们会进行讨论
  • 【与 RLHF 不同】对于给定的 \(\pi_\text{ref}\) 和数据集 \(D\) ,优化语言模型 \(\pi_\theta\) 以最小化loss:
    $$
    Loss_{\text{DPO}}(\pi_\theta;\pi_\text{ref}) = - \mathbb{E}_{(x,y_w,y_l) \sim D}\left [ \log \sigma \left( \beta\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} - \beta\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} \right)\right ]
    $$
    • 注: \(\sigma(\cdot)\) 表示Sigmod函数 \(\sigma(x) = \frac{1}{1+e^{-x}}\)
    • 策略 \(\pi_\text{ref}\) 和数据集 \(D\) 必须是对应的,即数据集 \(D\) 应该是按照策略 \(\pi_\text{ref}\) 中采样得到的;
      • 一般来说:数据集 \(D\) 都是策略 \(\pi^{SFT}\) 采样得到的,此时 \(\pi_\text{ref}\) 就是 \(\pi^{SFT}\)
      • 当数据集 \(D\) 为公开数据集时,可以使用下面的方法先训练一个 \(\pi_\text{ref}\) (原始论文中描述 \(\pi_\text{ref}\) 的训练过程为initialize \(\pi_\text{ref}\) ,但实际上, \(\pi_\text{ref}\) 在DPO训练 \(\pi_\theta\) 过程中是不会被更新的):
        $$
        \pi_\text{ref} = \arg\max_{\pi} \mathbb{E}_{(x,y_w) \sim D} log \pi(y_w|x)
        $$
        • 理解:这种情况下 \(\pi_\text{ref}\) 的训练就是在数据集 \((x, y_w) \sim D\) 上的SFT
        • 这种设定可以减少未知的真实行为策略与 DPO 使用的 \(\pi_\text{ref}\) 之间的分布偏移(mitigate the distribution shift between the true reference distribution which is unavailable, and \(\pi_\text{ref}\) used by DPO)
    • 问题 :如何理解收集数据的行为策略必须和参考策略一致?
    • 回答(待补充):在推导过程中并未看到必须要求数据和参考策略一致,但是 DPO-最优策略推导过程中变换 KL 散度形式时,要求 \(y\sim\pi\),即回复 \(y\) 必须从当前策略采样,考虑到参考模型常常作为目标策略的初始值,且更新过程中数据无法变更,则保证数据从参考策略(目前策略初始值)采样是最合适的
      • 其他角度的进一步理解:如果能针对当前策略生成的样本进行反馈,则可以最大程度上对当前模型进行高效优化,重点优化模型可能会生成的样本,给与或正或负的高效反馈
  • 问题:为什么不能直接用 SFT 的结果作为 \(\pi_\theta\) ?DPO 相对普通 SFT 有什么优点?
  • 回答:因为 SFT 没有使用到 \(y_l\) 样本,没有体现出来人类偏好,而 DPO 可以

附录:DPO 更新梯度分析

  • DPO-loss 关于参数 \( \theta \) 的梯度可以写成:
    $$\begin{aligned}
    \nabla_{\theta} \mathcal{L}_{\mathrm{DPO} }\left(\pi_{\theta} ; \pi_{\text {ref } }\right) = -\beta \mathbb{E}_{\left(x, y_{w}, y_{l}\right) \sim \mathcal{D} }[\underbrace{\sigma\left(\hat{r}_{\theta}\left(x, y_{l}\right)-\hat{r}_{\theta}\left(x, y_{w}\right)\right)}_{\text {奖励估计错误时权重更高 } }[\underbrace{\nabla_{\theta} \log \pi\left(y_{w} | x\right)}_{\text {增加 } y_{w} \text { 的概率 } }-\underbrace{\nabla_{\theta} \log \pi\left(y_{l} | x\right)}_{\text {降低 } y_{l} \text { 的概率 } }]]
    \end{aligned}$$
    • 其中 \(\hat{r}_\theta(x,y) = \beta \log\frac{\pi_\theta(y|x)}{\pi_\text{ref}(y|x)}\) 在原始论文中被称为隐式奖励模型,\( \hat{r}_{\theta}(x, y)=\beta \log \frac{\pi_{\theta}(y | x)}{\pi_{\text {ref } }(y | x)} \) 是由语言模型 \( \pi_{\theta} \) 和参考模型 \( \pi_{\text {ref } } \) 隐式定义的奖励
    • 从梯度上可以看出:当 \(y_l\) 的奖励大于 \(y_w\) 的奖励时,梯度越大,而损失函数的梯度会增加生成 \(y_w\) 的概率,降低 \(y_l\) 的概率(对应负号 \(-\))
    • 直观地说,损失函数 \( \mathcal{L}_{\text {DPO} } \) 的梯度增加了偏好完成项 \( y_{w} \) 的可能性,并降低了非偏好完成项 \( y_{l} \) 的可能性
    • 由隐式奖励模型 \( \hat{r}_{\theta} \) 对非偏好完成项的评分高出多少来加权,由 \( \beta \) 缩放(详细讨论见附录),即隐式奖励模型对完成项排序的错误程度,考虑了 KL 约束的强度
      • 注:作者在论文中给出了实验来验证隐式奖励模型 \( \hat{r}_{\theta} \) 加权的重要性(注意不是 \(\beta\) 的实验,是整体加权系数的实验),没有加权系数的朴素版本会导致语言模型退化(原文附录表3)
  • 补充:梯度的推导过程

附录:DPO 与 RLHF 的关系是什么?

  • 对于给定的 \(\pi_\text{ref}\) 和数据集 \(D\)
  • DPO是在优化语言模型 \(\pi_\theta\) 以最小化loss:
    $$
    Loss_{\text{DPO}}(\pi_\theta;\pi_\text{ref}) = - \mathbb{E}_{(x,y_w,y_l) \sim D}\left [ \log \sigma \left( \beta\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} - \beta\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} \right)\right ]
    $$
  • RLHF的Reward Model则在优化Reward Model打分模型 \(r_\theta\) 以最小化loss:
    $$
    Loss_{\text{RM}}(\theta) = -\frac{1}{C_K^2} \mathbb{E}_{(x,y_w,y_l)\sim D}[\log \sigma(r_\theta(x, y_w) - r_\theta(x, y_l))]
    $$
  • 可以看到 DPO 的损失函数与 RLHF 中 Reward Model 的损失函数很相似,可以理解为 DPO 隐式建模了 Reward Model(在原始论文中 \(\hat{r}(x,y) = \beta \log\frac{\pi_\theta(y|x)}{\pi_\text{ref}(y|x)}\) 被称为隐式奖励模型,与 \(r_\theta(x,y)\) 替换后即和RM中的损失函数形式完全一致)
  • DPO 可以看做是一个 off-policy 且 offline 的 RL 方法(实际上不是 RL 方法),而 RLHF 中的 PPO 是 online 且 on-policy 的方法

附录:对于大于2个回答偏好的情况如何处理?

  • 处理方法一:Bradley-Terry(BT)Model,如论文所述,对多个偏好进行两两匹配生成多个样本,然后转换为只有两个回答偏好的情况
  • 处理方法二:Plackett-Luce Model,使用如下损失函数来训练:
    • 可以看到Plackett-Luce Model的损失函数展开以后就是 \(C_K^2\) 个偏好样本对,与先两两匹配生成样本,然后再利用Bradley-Terry Model建模的方法完全一致
    • 原论文附录中有关于 Plackett-Luce Model 的详细介绍和推导

附录:DPO 损失函数中 KL 散度和超参数 \(\beta\) 的理解

  • 从原始问题定义中理解 :在原始约束优化问题的定义中:
    $$
    \max_{\pi_\theta} \mathbb{E}_{x\sim D, y\sim \pi_\theta(y|x)} [r_{RM}(x,y)]-\color{red}{\beta} \mathbb{D}_{KL}[\pi_\theta(y|x)||\pi_\text{ref}(y|x)]
    $$
    • \(\beta\) 是作为控制 KL 散度的约束严格程度存在的: \(\beta\) 越大,KL 散度约束越强; \(\beta\) 越小,KL 散度约束越弱
  • 在 DPO 最终的损失函数中理解 :从损失函数的数学形式看
    $$
    Loss_{\text{DPO}}(\pi_\theta;\pi_\text{ref}) = - \mathbb{E}_{(x,y_w,y_l) \sim D}\left [ \log \sigma \left( \color{red}{\beta}\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} - \color{red}{\beta}\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} \right)\right ]
    $$
    • 从损失函数中不容易理解 \(\beta\) 的作用,需求导以后才好理解,参见下面梯度中的理解
  • 在 DPO 更新的梯度中理解 :从梯度更新的式子看:
    $$\begin{aligned}
    \nabla_{\theta} \mathcal{L}_{\mathrm{DPO} }\left(\pi_{\theta} ; \pi_{\text {ref } }\right) = -\color{red}{\beta} \mathbb{E}_{\left(x, y_{w}, y_{l}\right) \sim \mathcal{D} }[\underbrace{\sigma\left(\hat{r}_{\theta}\left(x, y_{l}\right)-\hat{r}_{\theta}\left(x, y_{w}\right)\right)}_{\text {奖励估计错误时权重更高 } }[\underbrace{\nabla_{\theta} \log \pi\left(y_{w} | x\right)}_{\text {增加 } y_{w} \text { 的概率 } }-\underbrace{\nabla_{\theta} \log \pi\left(y_{l} | x\right)}_{\text {降低 } y_{l} \text { 的概率 } }]]
    \end{aligned}$$
    • 其中 \(\hat{r}_\theta(x,y) = \color{red}{\beta} \log\frac{\pi_\theta(y|x)}{\pi_\text{ref}(y|x)}\)
    • 理解:从梯度公式看
      • \(\color{red}{\beta \rightarrow 0}\) 时, \(\sigma(\cdot) \rightarrow \sigma(0) = 0.5\),相当于是个固定值,此时模型专注于优化以增大正样本概率,减小负样本概率,几乎没有任何约束
      • \(\color{red}{\beta \rightarrow +\infty}\) 时:
        • 当 \(\left(\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} -\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} \right) > 0\),则 \(\sigma(\cdot) \rightarrow \sigma(+\infty) = 1\),梯度对模型的影响正常,模型会优化这种样本,使得 \(y_w\) 出现的概率变大, \(y_l\) 出现的概率变小
        • 当 \(\left(\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} -\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} \right) < 0\),则 \(\sigma(\cdot) \rightarrow \sigma(-\infty) = 0\),梯度对模型影响为 0,模型不再学习这种样本(即使 \(y_w\) 出现的概率比参考策略只大一点点,或\(y_l\) 出现的概率比参考策略仅小一点点)
        • 综上,\(\color{red}{\beta \rightarrow +\infty}\) 时,会导致模型学习时,刚刚好学到一个 \(y_w\) 出现的概率比参考策略只大一点点,或\(y_l\) 出现的概率比参考策略仅小一点点的策略即收敛,这会导致策略不再继续偏离参考策略,也就实现了对 KL 散度的极大约束
  • \(\beta\) 取值最佳实践 :原始论文中提到 \(\beta\) 默认取值为 \(\beta=0.1\),在 TL;DR summarization 任务(一个摘要生成任务)中取值为 \(\beta=0.5\)

附录:DPO Loss 的代码实现

  • 来自原始 DPO 论文的代码实现:
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    def dpo_loss(pi_logps, ref_logps, yw_idxs, yl_idxs, beta):
    """
    pi_logps: policy logprobs, shape (B,)
    ref_logps: reference model logprobs, shape (B,)
    yw_idxs: preferred completion indices in [0, B-1], shape (T,)
    yl_idxs: dispreferred completion indices in [0, B-1], shape (T,)
    beta: temperature controlling strength of KL penalty
    Each pair of (yw_idxs[i], yl_idxs[i]) represents the
    indices of a single preference pair.
    """
    pi_yw_logps, pi_yl_logps = pi_logps[yw_idxs], pi_logps[yl_idxs]
    ref_yw_logps, ref_yl_logps = ref_logps[yw_idxs], ref_logps[yl_idxs]
    pi_logratios = pi_yw_logps - pi_yl_logps
    ref_logratios = ref_yw_logps - ref_yl_logps
    losses = -F.logsigmoid(beta * (pi_logratios - ref_logratios))
    rewards = beta * (pi_logps - ref_logps).detach()
    return losses, rewards

附录:针对 DPO 的改进

  • DPO 中只保证目标策略生成整个句子 \(y_w\) 或 \(y_l\) 的概率与参考策略一致,并不限制中间每个 Token 生成的过程中概率一致
  • DPO 中的损失函数要求的是 \(\left(\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} -\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} \right) \) 变大;实际上,模型可能学习到的是,让 \(\pi_\theta(y_w|x)\) 和 \(\pi_\theta(y_l|x)\) 同时变小,只要正样本变小的幅度较小即可

附录:DPO 训练时的 \(\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)}\) 应该保证始终大于 1 吗?

  • 副标题:DPO 训练时的 \(\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)}\) 应该保证始终小于 1 吗?
  • 现象:在训练时发现 \(\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} < 1\) 且持续降低(此时 \(\beta=0.1\))
    • 设置 \(\beta=0.5\) 后 \(\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} > 1\) 波动
    • 其他指标均正常:
      • \(\left(\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} -\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} \right) \) 逐步变大
      • \(\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} < 1\) 且持续降低
  • 问题:让 \(\pi_\theta(y_w|x)\) 和 \(\pi_\theta(y_l|x)\) 同时变小的结果不是我们想要的吧,我们至少不想要 \(\pi_\theta(y_w|x)\) 降低
  • 如何理解 通过调整 \(\beta\) 值可以缓解这个问题呢?
    • 个人理解:DPO 的本质是让 \(\left(\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} -\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)} \right) \) 越大越好
    • 进一步分析有:
      • 当 \(\beta\) 过小时 ,KL 散度约束小,可能导致 \(\pi_\theta(y_w|x)\) 和 \(\pi_\theta(y_l|x)\) 同时变小(持续变小)
      • 当 \(\beta\) 过大时 ,KL 散度约束大,可能导致 \(\pi_\theta(y_w|x)\) 和 \(\pi_\theta(y_l|x)\) 都波动很小
        • 从梯度上体现就是,当 \(\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)}\) 比 \(\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)}\) 大一点点就不学了
      • 当 \(\beta\) 适中时 ,KL 散度约束适中,可约束 \(\pi_\theta(y_w|x)\) 和 \(\pi_\theta(y_l|x)\) 不能偏离当前策略太多,让 \(\pi_\theta(y_w|x)\) 变大, \(\pi_\theta(y_l|x)\) 变小,从而实现

附录:对齐人类偏好的其他方法

  • RLHF: 基于人类反馈的强化学习方法,基于人类偏好训练 RM 再 RL,ChatGPT 的训练方式
  • RLAIF: 基于AI反馈的强化学习方法,基于 AI 偏好模型训练 RM 再 RL,Claude 的训练方式
  • 其他参考RLHF的替代之DPO原理解析:从RLHF、Claude的RAILF到DPO、Zephyr:

    整个从 SFT 到 RLHF 或 RLAIF 的训练中,ChatGPT 只有最后 1/3 摆脱了人工参与的训练,而 Claude 则希望超过 1/2 摆脱人工参与的训练

1…131415…64
Joe Zhou

Joe Zhou

Stay Hungry. Stay Foolish.

638 posts
53 tags
GitHub E-Mail
© 2026 Joe Zhou
Powered by Hexo
|
Theme — NexT.Gemini v5.1.4