Higgs Audio V2-语音大模型

Higgs Audio V2模型,不仅能处理文本,还能同时理解并生成语音。除了一些常规语音任务外,这个模型还具备一些较为罕见的能力,比如生成多种语言的自然多说话人对话、旁白过程中的自动韵律调整、使用克隆声音进行旋律哼唱以及同时生成语音和背景音乐。

整个过程堪称“大力出奇迹”,直接将1000万小时的语音数据整合到LLM的文本训练,

Higgs Audio v2 采用上图架构图中所示的“generation variant”。其强劲的性能源于三项关键技术创新:

  • 开发了一套自动化注释流程,该流程充分利用了多个 ASR 模型、声音事件分类模型以及我们内部的音频理解模型借助该流程,我们清理并注释了 1000 万小时的音频数据,并将其命名为 AudioVerse 。该内部理解模型在 Higgs Audio v1 Understanding 的基础上进行了微调,并采用了架构图中所示的“理解变体”。
  • 从零开始训练了一个统一的音频分词器,它可以同时捕捉语义和声学特征。
  • 提出了 DualFFN 架构,它增强了 LLM 以最小计算开销建模声学 token 的能力。

Higgs Audio V2 在音频 AI 能力上实现了重大飞跃:

  • 多说话人对话自然流畅:多说话人对话往往难以处理,尤其是在角色无法匹配彼此的情绪和语气时。而借助 Higgs Audio V2,这种对话轻松自然,仿佛现场交流,充满生命力。
  • 支持长音频生成:生成长音频时需要保持声音的一致性,同时兼顾真实感、吸引力和生动性。Higgs Audio 提供条件控制与提示机制,使长音频表现出色。
  • 高保真音质:为了在高品质扬声器和耳机上实现逼真音效,V2 将音频处理管线从 16kHz 升级至 24kHz,带来更佳音质。
  • 高效推理,资源友好:无论是个人项目还是商用部署,推理效率都很重要。我们最小的模型可以在 Jetson Orin Nano 上运行;最新的 3B Audio Generation V2 模型则至少需要 RTX 4090 才能高效推理。
  • 生成真实、有情感的语音表现领先:在 EmergentTTS-Eval 基准测试中,其胜率超过 75%,超越 ChatGPT 4o。
  • 开源:模型开源。
  • 训练数据超千万小时:为实现更高音质与更逼真的语音效果,模型在超过1000万小时的音频上训练,并依托精细的处理与标注流程自动生成训练数据

内容:

传统的语音和文本模型之间相互独立,李沐老师就想,欸,能不能将两者结合起来,直接让LLM用语音进行沟通。那么首先就要知道文本语言模型的本质是用给定的一段指令去生成预测结果,就是将任务先拆解为系统指令(system)用户输入(user)模型回复(assistant)三个部分。system告诉模型,需要做什么事情,例如回答该问题、写一段文字或者其他,user就是告知事情的详细内容,例如问题具体是什么、文字要什么风格。

所以如果要让模型支持语音,就需要为模型增加一个系统命令,在user里输入要转录为语音的文字,让模型从system里输出对应语音数据。这样语音任务就能转换成相同的处理格式,直接打通语音和文本之间的映射,通过追加更多的数据和算力,直接scaling law“大力出奇迹”。

这就引出了新的问题,语音信号本质是连续的,要如何才能在离散的文本token中表示呢?

现有的方法是将一秒的语音信号裁切成多段(如100毫秒一段),为每一段匹配最相似的预定义模板(如45个模板),然后将其表示为长度为10的编号序列,也就是一个个token。

但这样做,虽然可以将一小时的音频从60兆压缩到0.16兆,但质量相当糟糕,所以需要优先保留语音的语义信息而声学信号只保留少量部分,后续再通过其他手段还原

于是他们训练了一个统一的离散化音频分词器,以每秒25帧 [40ms/帧] 的速度运行,同时保持甚至提高音频质量,以捕获语义和声学特征。

新的离散化音频分词器运行速度仅为每秒25帧,同时在音质上保持甚至优于码率翻倍的分词器。该模型是首个在 24 kHz 数据上训练的统一系统覆盖语音、音乐与声音事件。同时,该模型采用简单的非扩散式编码器/解码器,实现快速批量推理

音频分词器性能:

然后要让模型很好地理解和生成声音,就需要利用模型的文本空间,将语音的语义尽量地映射回文本,当中需要大量的数据支持。

由于版权问题,沐导没有使用B站或YouTube这类公开视频网站数据,而是购买或从允许抓取的网站获取。这样得到的数据质量参差不齐,需要删除其中的90%才能满足1000万小时的训练数据需求。

其次,将语音对话表示为相应的system(场景描述、声学特征、人物特征等)、user(对话文本)、assistant(对应音频输出)的形式。由于OpenAI和谷歌一向禁止使用他们的模型输出再训练,且训练成本过高,为了实现这种标注,他们利用相同的模型架构额外训练出一个语音模型AudioVerse

该模型接收用户语音输入,分析并输出场景、人物、情绪、内容等信息,再将输出反过来作为生成模型的system提示和user输入,实现模型的共同进步。

举个例子就是,如果想要教一个徒弟同时会拳脚功夫,但师傅一次又教不了,那就同时教两个徒弟,一个学打拳,一个学踢腿,然后让他们俩天天互相打,打着打着两个就都会拳脚功夫了。

最终,这个多模态模型就完成了,不仅可以完成简单的文本转语音,还能实现更复杂的任务,比如让它写一首歌并唱出来,再加上配乐。

还能根据语音分析场景、人物(性别、年龄、情绪状态)、环境音(室内外),并进行复杂的理解和推理。

在实时语音聊天上,还可实现低延迟、理解情绪并表达情绪的自然语音交互,而不仅仅是机械的问答。

EmergentTTS-Eval基准上,相较于其他模型,性能可以说是遥遥领先,尤其是在“情绪”和“问题”类别中,相比GPT-4o-mini-tts高出了75.7%和55.7%的胜率。

此外,它在Seed-TTS Eval和情感语音数据集 (ESD) 等传统TTS基准测试中也取得了最佳性能。

Evaluation:

Seed-TTS Eval & ESD

我们使用参考文本、参考音频和目标文本对 Higgs Audio v2 进行零样本语音合成(TTS)测试。评估采用 Seed-TTS Eval 和 ESD 中的标准评估指标。【SIM 指标一般是指 Speaker Similarity

EmergentTTS-Eval (“Emotions” and “Questions”):根据 EmergentTTS-Eval 论文,我们报告了在使用 “alloy” 音色时,相较于 “gpt-4o-mini-tts” 的胜率。评判模型为 Gemini 2.5 Pro。

多说话人评估:我们还设计了一个多说话人评估基准,用于评估 Higgs Audio v2 在多说话人对话生成方面的能力。该基准集包含三个子集:

  • two-speaker-conversation:包含1000条双人合成对话。我们固定两段参考音频,用以评估模型在双人语音克隆方面的能力,对话轮数在4到10轮之间,角色随机选择。
  • small talk(无参考音频):包含250条合成对话,生成方式与上类似,但特点是发言简短、轮数较少(4–6轮)。本集合未提供参考音频,旨在评估模型自动为角色分配合适声音的能力。
  • small talk(有参考音频):同样包含250条合成对话,发言更短。该集合在上下文中包含参考音频片段,类似于 two-speaker-conversation,用于评估基于参考音频的表现。

我们在这三个子集上报告了词错误率(WER)和说话人内相似度与说话人间差异度的几何平均值。除 Higgs Audio v2 外,我们还评估了 MoonCast 和 nari-labs/Dia-1.6B-0626 这两个当前最受欢迎、支持多说话人对话生成的开源模型。结果总结在下表中。由于 nari-labs/Dia-1.6B-0626 对话语长度及输出音频的严格限制,我们未能在 “two-speaker-conversation” 子集上运行该模型。

Seed LiveInterpret 2.0 端到端同声传译大模型

!!!总结:必须认识到数据在模型训练的重要性。模型经过数十万小时语音数据的训练,数据质量中的任何瑕疵都可能在最终效果中被显著放大,这些潜在问题包括口音差异、准确读音、时间戳的准确预测,以及句子衔接的流畅度等关键要素。良好的性能正是建立在海量优质训练数据之上。

Seed LiveInterpret 2.0 是首个延迟&准确率接近人类水平的产品级中英语音同传系统,在中英同传翻译质量达到业界 SOTA 的同时,实现了极低的语音延迟水平。

它基于全双工端到端语音生成理解框架,支持中英互译,可实时处理多人语音输入,像人类同传译员一样以极低的延迟 “边听边说”,一边接收源语言语音输入,一边直接输出目标语言的翻译语音。同时,Seed LiveInterpret 2.0 还支持 0 样本声音复刻,让沟通更加流畅自然。

  • 接近真人同传的翻译准确率 精准的语音理解能力保障了翻译准确度,在多人会议等复杂场景中英双向翻译准确率超 70%,单人演讲翻译准确率超 80%,接近真人专业同传水平。
  • 极低延迟的 “边听边说” 能力 采用全双工语音理解生成框架,翻译延迟可低至 2-3 秒,较传统机器同传系统降低超 60%,实现了真正的 “边听边说” 翻译。
  • 零样本声音复刻,音色真实自然 只需采样实时语音信号,便能提取声音特征,用说话人的音色特质实时 “说出” 外语,提升交流的沉浸感和亲和力。
  • 智能平衡翻译质量、延迟和语音输出节奏 可根据语音清晰度、流畅度、复杂程度,调整输出节奏,并适配不同语言特性。面对超长信息,依然能保证传译语音节奏的自然流畅。
同声传译系统评估:左右两图比较了人工评估的翻译质量分数和响应效率  对于语音转文本 (S2T) 和语音转语音 (S2S) 模式,响应效率是相对于人工翻译延迟来衡量的。人工评估准确度反映了翻译输出对说话者原始意图的忠实程度

框架:

系统会克隆每位说话者的声音,并以相应的语调将其翻译为另一种语言
全双工流式端到端模型架构: Hibiki ,模型架构和数据相关可参考该论文

提出一种端到端的语音到语音同步翻译模型,在一个统一框架内无缝整合了同步语音翻译和语音克隆功能。

  • 语言模型预训练:使用 Seed LLM 系列的方法对初始语言模型进行预训练,建立基础的文本生成与理解能力。
  • 多模态扩展:集成一个预训练的音频编码器,使模型能够接受流式音频输入,扩展为具备音频处理能力的多模态 LLM。
  • 多任务持续学习训练:在大规模多任务数据上进行自回归训练,生成包括文本 token(可选)和音频 token 的输出,实现语音合成。
  • 高质量数据微调:使用人工标注的高质量数据进行微调,进一步优化模型在指令理解、多说话人识别、翻译策略关键能力上的表现。

问题:面临严格延迟约束下的同步翻译优化难题,需要在翻译质量时序控制之间权衡。

核心思路:优化两个互补目标

  • 片段内一致性确保每个翻译片段自身准确、流畅
  • 片段间连贯性确保不同翻译片段之间逻辑衔接自然

奖励机制设计

  • 多维单轮奖励(step-level):为每一步生成即时反馈,评估翻译准确性与时序控制,实现片段内部一致性优化
  • 统一多轮奖励(sequence-level):从全局角度评估整个翻译段落的连贯性,优化跨片段一致性

两阶段训练策略

  • 第一阶段:单轮奖励训练
    • 仅使用 step-level 奖励,学习人类翻译的先验知识,确保训练稳定
  • 第二阶段:联合优化训练
    • 引入 sequence-level 奖励,与 step-level 奖励联合优化,平衡过程指标(每步表现)与结果指标(整体输出质量)

主要贡献包括:统一的语音到语音架构、跨语言的语音克隆机制,以及接近人类水平的翻译性能。

Training

Continual Training and Supervised Fine-tuning

为实现文本与语音之间的有效模态对齐,并提升跨语言能力,我们采用了全面的多任务多模态持续训练(CT)策略。该策略有效促进了语音与文本模态之间的对齐,并强化了模型的跨模态与跨语言泛化能力

具体措施如下:

  1. 多模态多任务训练数据
    • CT 数据集涵盖约 1000 亿 tokens,任务类型包括:
      • 语音转文本(Audio-to-Text Transcription)
      • 文本转语音(Text-to-Audio Synthesis)
      • 纯文本处理(Text-Only Tasks)
  2. 数据质量控制
    • 为提升训练效率并确保数据质量,我们引入了基于语音质量指标的严格过滤流程,对语音数据进行筛选。

在持续训练之后,我们在高质量的人类标注数据上进行有监督微调,以激活同步语音传译所需的关键能力。该过程使模型能够建立以下数据驱动能力:

  1. 读-写策略(read-write policy)
  2. 多说话人区分能力
  3. 语音翻译能力
  4. 声音克隆能力

有监督微调显著提升了模型的指令跟随能力以及在核心传译任务上的整体表现。经过微调后的模型为后续的强化学习阶段提供了强大基础,使得后续优化更具针对性有效性

Reinforcement Learning

现代同声传译系统采用双工处理,将输入流分割成连续的音频块。形式上,我们将输入输出序列表示为:

每个音频片段(audioₜ)对应一个增量翻译 yₜ。我们将(audioₜ, yₜ)表示为序列中的第 t 个片段,并将 audio :=(audio₁, audio₂, …, audioₜ)表示为从 1 到 T 的聚合音频。在每个 t 片段中,我们有 yₜ :=(yₜ₁, yₜ₂, …, yₜₙ, …, yₜₙ),其中 N 是输出的长度。该模型利用当前音频片段(audioₜ)和之前的上下文 x<t,通过策略生成翻译 yₜ。

其中 πθ 是具有参数 θ 的策略 决定翻译策略。完整的轨迹概率定义为:

我们将 rtn 表示为 t 个块中第 n 个 token 的奖励。强化学习的目标是最大化每条轨迹上的累积奖励,即:

其中 𝒟 是训练数据集。以下部分详细说明了 rtn  的设计方式。

 奖励设计:平衡单轮反馈和多轮反馈

  • 单轮奖励(Single-turn rewards):在每个决策点提供即时反馈,用于评估中间的推理或生成步骤。
  • 多轮奖励(Multi-turn rewards):评估整个输出序列的质量,反映多个决策步骤的长期、累积效果。

同步翻译系统尤其具有独特挑战,因此需要精细化的奖励设计。该任务需同时优化两个互补目标:

片段内一致性(Intra-segment consistency):要求模型在逐步输出时保持语义与时间上的准确性和完整性,适合采用单轮奖励(single-turn reward)进行即时评估。

片段间连贯性(Inter-segment coherence):确保整个翻译序列在语义和时间上的连续性与一致性,适合采用多轮奖励(multi-turn reward),从全局角度评估累积的序列质量。

基于上述考量,我们提出了一种新颖的框架,将多维单轮奖励统一多轮奖励相结合。

Single-turn Reward:方法利用了每个增量步骤的细粒度反馈,我们通过实证研究发现,这些反馈与人类的评估指标高度相关。

给定一个音频序列 {audiot}1T 和相应的真实值 {yt}1T ,沿着五个派生维度定义段内奖励:

检测准确性奖励(rl​):该奖励旨在鼓励模型在翻译前进行充分“倾听”,避免过早输出,从而提升语义单元完整性。

I(⋅) 为指示函数,条件成立时取值为 1,否则为 0;∣yt​∣ 表示模型在第 t 步生成的 token 数量;∣yt∗​∣ 表示参考翻译在第 t 步应生成的 token 数量。当模型和参考翻译在当前步都没有输出(token 数为 0)时,奖励为 1,否则为 0该设计鼓励模型在语义信息尚不完整时保持“静默”,从而提升翻译的延迟-准确性权衡表现。

翻译主动奖励 ( rs ):通过奖励尽快生成已确认的语义单元来鼓励语音翻译:鼓励模型在语义单元一旦可用时立即翻译

翻译质量奖励(rq:衡量当前步生成内容与参考翻译的相似度(可通过 BLEU、BERTScore 等):

时序匹配奖励(rc​):鼓励模型生成的语音时长与参考时长一致,惩罚过长或过短:

格式一致性奖励(rf):保证输出结构正确,如标点、格式符号等符合预设正则表达式:

最终单轮奖励定义如下:

多轮奖励:单轮奖励机制提供了详细的、逐步的反馈,能够在每一步的递增中平衡延迟和翻译质量,但它未能完全捕捉同声传译中固有的长期依赖关系和累积效应。尤其是,当生成的目标音频逐渐落后于源音频时,会导致破坏性延迟,从而降低用户体验。为了解决这些全局序列级的动态问题,我们设计了一个互补的多轮奖励机制,可以整体评估整个输出序列。

延迟惩罚(rL​):惩罚翻译滞后,鼓励更及时的输出:

  • l:可接受的最大等待阈值
  • K:翻译片段数
  • dk​:第 k 个翻译片段前等待的片段数量

序列级翻译质量奖励(rQ​):衡量整个翻译序列与参考的匹配度(例如通过全局对齐算法):

多轮奖励定义为:

最终奖励融合与正则项

  • 每个子奖励在 batch 中进行标准化(均值为 0,方差为 1),提高数值可比性。
  • 总奖励为标准化后的单轮与多轮奖励之和,融合了局部细粒度指导全局一致性目标
  • 引入 KL 散度正则项:

用于鼓励当前策略 πθ​ 向参考策略靠拢,提升训练稳定性与可控性。

稳定强化学习训练:通过近端策略优化 (PPO)来优化定义的目标,该算法通过修剪的目标函数实现稳定高效的策略更新。训练目标公式如下:

 audio={audiot}1T 表示输入的音频序列, y={yt}1T 表示从旧策略 πθo⁢l⁢d 采样的翻译响应。优势估计 Atn 使用广义优势估计 (GAE)计算。由于这些奖励之间紧密耦合且差异化,调整它们各自的权重颇具挑战性,而且通常效果不佳。为了解决这些问题并稳定训练,我们采用了两种主要策略:自适应 KL 惩罚 和两阶段强化学习训练方案。

 Adaptive KL

对于包含音频和文本 token 的序列,由于其长度较长,控制 KL 散度会更加困难,这自然会导致更高的累积 KL 散度。因此,KL 惩罚系数 β 必须设置为高于传统 RLHF 的设置。

采用对数空间中的比例控制器来自适应地调整 β ,以确保 KL 散度始终接近预定目标。

两阶段强化学习训练方案:在第一阶段,通过仅优化多维单轮奖励来预热模型,使其内化人类先验知识并实现稳定的学习动态。在第二阶段,使用结合过程和结果成分的多轮奖励对模型进行进一步训练,使其能够有效地优化和平衡延迟与翻译质量。、

Experiments

 评估指标:

对于文本翻译质量评估,我们主要依赖于人工评估指标——有效信息比例 (VIP),该指标衡量翻译输出对每个语义片段传达说话者原始意图的准确程度,与人工翻译的判断高度一致。

在语音到语音评估中,我们提出了“语音有效信息比例”(SVIP)作为一种全面的人类评估指标。该指标建立在已有的“有效信息比例”(VIP)框架[6]之上,用于衡量完整语音会话中有效语义片段所占的比例。

当一个语音语义片段能够有效传达源语音的核心信息、准确表达说话者的原始意图、在可接受的延迟范围内完成传递、保持适合听众理解的语速,并达到清晰和易懂的声音质量标准时,即被视为有效。

在延迟评估方面,我们采用“首字母出现延迟”(FLAL)指标来衡量系统在段落级别输出第一个确定翻译所需的时间。在句子级别,我们使用广泛应用的“平均延迟”(AL)和“长度自适应平均延迟”(LAAL)指标,以比较不同方法之间的延迟表现。

在延迟表现上,Seed LiveInterpret 2.0 在语音到文本场景中,输出首字平均延迟仅 2.21 秒,在语音到语音场景中,输出延时仅 2.53 秒,做到了对翻译质量以及时延的均衡。

针对 Seed LiveInterpret 2.0 中译英和英译中两个方向的表现进行了客观评估,与其他翻译系统在翻译质量(BLEURT/ COMET)和延迟(AL/ LAAL/FLAL)等指标上进行对比。

结果显示,Seed LiveInterpret 2.0 在两个数据集上均表现出最高的翻译质量。在延迟方面,Seed LiveInterpret 2.0 在英到中方向上实现了语音到语音翻译的最低平均滞后(AL),在中到英方向上也表现出竞争力,展现了速度与准确度的良好平衡。

总体来看,Seed LiveInterpret 2.0 在句子级基准测试中,有效平衡了翻译质量与延迟。这不仅缓解了传统同传中 “译得准则慢,译得快则偏” 的痛点,配合音色复刻能力,让中英跨语言交流首次具备自然对话般的流畅感。

总结与展望

在本研究中,团队进一步认识到数据对模型训练的重要性。模型经过数十万小时语音数据的训练,数据质量中的任何瑕疵都可能在最终效果中被显著放大,这些潜在问题包括口音差异、准确读音、时间戳的准确预测,以及句子衔接的流畅度等关键要素。Seed LiveInterpret 2.0 良好的性能正是建立在海量优质训练数据之上。

Seed LiveInterpret 2.0 已初步展现出一定优势,其边界仍有拓展空间。比如,在语言覆盖方面,目前模型主要支持中英互译,其他语种尚未较好支持。此外,其声音复刻的稳定性、语音表现力、情绪复刻能力、极复杂情况下的翻译准确性等仍有进步空间。

在未来研究中,我们希望进一步挖掘模型潜力,通过优化算法、增强数据及改进训练策略等方式,逐步拓展同传模型的能力边界,提升其在复杂场景下的适应性和性能表现。

高效 LLM 训练方法:Packed samples和 sorted batching

 要让大型语言模型更有效地处理长文本上下文,需要在相似长度的输入序列上进行指令微调。LongAlign 方法,它可以帮助大型语言模型有效处理长达 64k 的长上下文,并展现出强大的长文本理解和生成能力。

LongAlign :

动机:

  • 目前缺乏用于有监督微调(SFT)的长文本指令跟随数据集,更缺乏构建此类数据的方法。
  • 长上下文数据的长度分布差异较大,在多GPU环境中严重降低了传统批处理方法的训练效率——处理较短输入的GPU必须等待处理较长输入的GPU完成任务后才能继续运行。
  • 亟需一个强健的基准评估体系,用于衡量大型语言模型在面对真实世界长文本查询时的处理能力。

贡献:

LongAlign 方法,分别从数据构建、高效训练和评估三个方面入手:

在数据方面,为构建一个多样化的长文本指令跟随数据集,从九个来源收集了长文本序列,并使用 Self-Instruct生成了 1 万条长度在 8k 到 64k 之间的指令数据。

在训练方面,为应对不均匀批处理导致的效率问题,采用了 packing 策略,即在将数据分发到 GPU 之前,将多个序列打包为接近最大长度的组合。但我们发现这种打包训练中的损失计算存在偏差:不同数量序列的打包在最终损失计算中被赋予相同权重。为缓解这一问题,我们提出了“损失加权策略”,对每条文本的损失进行加权平均,以平衡不同序列对整体损失的贡献。此外,我们还引入了“排序批处理”方法,将长度相近的序列分组,从而减少批内空闲时间

在评估方面,开发了 LongBench-Chat 基准测试,它包含长度为 10k-100k 的开放式问题,这些问题由博士生进行标注。评估内容涵盖推理、编程、摘要以及多语种翻译等多种长文本指令跟随能力。使用 GPT-4(OpenAI,2023b)结合人工标注结果和少量示例,对模型生成的回答进行评分。

结论:

数据量与多样性的影响:长文本指令数据的数量和多样性都会显著影响模型处理长上下文的能力,最终性能差异最高可达 30%。

长文本指令数据的益处:增加长文本指令数据有助于提升模型在长上下文任务中的表现,同时不会削弱其处理短上下文任务的能力。

训练策略的有效性采用的打包和排序批处理策略可将训练速度提升超过 100%,且不影响模型性能。此外,提出的损失加权技术还能将长文本任务的性能提升 10%。

数据集构建:

构建了一个包含10,000条长度在8k-64k之间的长文指令跟随数据集,这些数据来自于9个不同的数据源,包括学术论文、书籍、百科全书等,覆盖了多样化的任务类型。

高效训练方法:

为了确保模型在有监督微调(SFT)后依然具备处理长文本和短文本(即通用能力)的能力,将长文本指令数据与通用指令数据集混合用于训练。这种训练策略使得大量通用短文本数据与相对较少的长指令数据结合,从而形成了一个“长尾”式的数据长度分布。探索了两种训练方法:packingsorted batching

Packing(打包)

该方法通过将不同长度的数据拼接,直至达到最大长度,生成的打包数据整体长度接近最大限值。这些打包后的数据再进行批处理并在多 GPU 上处理,有效减少了每个批次中的空转时间。

此外,为防止同一 pack 中的不同序列在自注意力计算中发生“交叉污染”,我们传入了每个序列的起始与结束位置列表,并使用了 FlashAttention 2 中的 flash_attn_varlen_func 该方法支持高效的块对角注意力计算,计算量与 IO 时间均优于传统的二维注意力掩码。

Packing 策略存在的偏差

不过我们注意到,packing 会带来对长序列目标 token 较多的序列偏向。这是因为:不同的打包组(pack)在最终损失计算中被赋予相同权重,而每个打包组中包含的序列数量和每个序列的目标 token 数量却不同。

因此,在对每个批次求平均损失时,包含序列较少(通常是较长序列)或目标 token 较多的 pack,会对最终损失产生更大影响。

形式上,设将 M 个序列打包成 K 个 pack,第 i 个 pack 包含索引区间为 [Pi−1,Pi)的序列,其中 P0=1,PK=M+1。设 Li 为第 i个序列在其 Ni​ 个目标 token 上的总损失。如果我们希望对每个序列赋予相等的权重[ SFT中算loss ],则损失应当为:

而在 packing 情况下计算得到的损失为:

(3)与公式 (2) 相比,在 packing 情况下,相当于为第 j个序列分配了一个权重:

也就是说,损失更偏向于目标 token 数较多的序列,以及位于pack 较小的组中的序列。

为了解决这种不公平,我们提出对第 i 个序列的损失进行缩放,缩放因子为:K/(NiM),然后对每个 pack 中缩放后的损失求和,这样得到的总损失将与公式 (2)(即平均每个序列损失)保持一致,从而消除了不同序列在损失计算中所受到的偏倚。

损失加权策略在下游任务中带来了约 10% 的性能提升

Sorted Batching(排序批处理)

还提出了一种高效的 排序批处理策略。为确保每个 batch 中的序列长度相近,我们先按照序列长度对数据进行排序,然后在每轮训练中从中随机选取一段连续的数据组成一个 batch,且不重复使用。

不过,该方法不可避免地会引入 批次间数据分布的不均衡某些 batch 可能全部由长序列组成,另一些则全是短序列。这种偏差可能对 SGD(随机梯度下降)优化过程造成严重影响。

尽管如此,我们在实验中发现,排序批处理显著加快了训练速度,且几乎不会对模型性能产生负面影响。这可能得益于我们使用了较大的梯度累积步数(gradient accumulation steps)和优化器本身较强的适应能力。

训练方法细节

这里介绍 packing 策略与损失加权的具体实现方式。

Packing 策略实现

在打包训练过程中,每个数据批次会传入一个特殊的一维注意力掩码。在该掩码中,第 i个元素表示第 i 个序列在该批次中的起始位置。掩码的第一个元素为 0,最后一个元素等于 batch_size × seq_len

在注意力计算时,我们使用 FlashAttention 2 提供的 flash_attn_varlen_func 函数,并将该掩码传入其参数 cu_seqlens_qcu_seqlens_k。该函数会根据掩码中相邻元素表示的起始和结束位置,在每个序列内部进行注意力计算因此,每个序列的 Query 只能与自身的 Key 进行注意力操作,实现了“序列内独立注意”。

损失加权策略实现

在实现损失加权策略时,首先对训练数据进行预处理:为每个 pack 中的序列生成一个加权的一维掩码。该掩码中,对应目标 token 的位置权重为 1/N(其中 N 是当前序列的目标 token 数),其他位置为 0。

训练时,根据当前配置动态设置 M 和 K,表示即当前批次中序列的数量和 pack 的数量。然后,损失计算方法为:对每个 token 的交叉熵损失乘以比例系数 K/(MN),再求和得到最终损失值。

Packing 加权loss代码实现:

SFT中算loss通常来讲都是样本内作token-level mean,样本间作sequence-level mean,也就是等式(2)的计算方式。如果不同样本间作token-level mean,则会使target token数量多的样本更受重视(相当于被upsample),从而引入不同样本间的不平衡。

### Support loss weighting for packing ###
        loss = None
        if labels is not None:
            lm_logits = lm_logits.to(torch.float32)
            # Shift so that tokens < n predict n
            shift_logits = lm_logits[..., :-1, :].contiguous()
            if isinstance(labels, tuple) or isinstance(labels, list):
                labels, weights = labels
            shift_labels = labels[..., 1:].contiguous()
            if self.pack_loss:
                shift_weights = weights[..., 1:].contiguous()
                loss_fct = CrossEntropyLoss(ignore_index=-100, reduction='none')
                loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
                loss = (loss * shift_weights).sum()
            else:
                loss_fct = CrossEntropyLoss(ignore_index=-100)
                loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

            lm_logits = lm_logits.to(hidden_states.dtype)
            loss = loss.to(hidden_states.dtype)
        ### -------------------------------------- ###

图解OpenRLHF中基于Ray的分布式训练流程

摘自:图解OpenRLHF中基于Ray的分布式训练流程

本文着重分析OpenRLHF中的PPO-Ray训练架构设计,不了解Ray的朋友也可以通过本文快速上手,全文共分成五块:

  1. 为什么用Ray
  2. 使用图例抽象出整体训练流程
  3. Ray核心知识速过
  4. 使用图例进一步抽象出核心代码细节,包括:
  • 训练入口
  • 部署PPO-Actor/Ref/Critic/RM实例
  • 部署vllm_engines实例
  • PPO-Actor与vllm_engines之间的通讯
  • PPO-Actor/Critic训练

5. RLHF-PPO算法细节介绍

Ray介绍

https://github.com/ray-project/ray

Ray 是一个用于扩展 AI 和 Python 应用程序的统一框架。Ray 由一个核心分布式运行时和一组用于简化 ML 计算的 AI 库组成:

一站式平台
包含多个模块(Ray Core, Ray Data, Train, Tune, RLlib, Serve)支持从数据加载、并行任务调度,到分布式训练、超参调优、强化学习、模型部署的完整 AI 工作流。

Python 原生,门槛低
采用装饰器 @ray.remote 标记函数或 Actor,即可并行执行,无需了解底层分布式原理。

高性能 & 容错能力
利用分布式调度、共享内存对象存储、层级调度机制,实现低延迟、高吞吐、故障恢复。

一、为什么要使用Ray

对于通常的rlhf框架,在训练时会在单卡上同时部署actor/ref/reward/critic四类模型,这种单一的部署方式可能存在如下问题:

  • 难以突破单卡显存的限制。
  • 无法实现更多的并行计算。例如在收集exp阶段,拿到(prompt, responses)结果的四类模型其实可以做并行推理;在训练阶段,拿到exp的actor和critic也可以做并行训练。但受到单卡显存等因素影响,通常的rlhf框架中使用更多的是串行。
  • 无法独立优化训练和推理过程。诸如vllm之类的框架,是可以用来提升actor生成(prompt, responses)的速度的,而对于其它模型,我们也可能会视算法需要有不同的推理需求。因此我们期望能更加灵活地设计训练、推理过程


而解决以上问题,需要开发者能设计一套较为灵活的分布式计算框架,能够实现资源定制化分配、分布式调度、节点内外通信等目标,同时相关的代码不能太复杂,能够让使用者更专注于算法部分的研发。而Ray天然可以帮我们做这件事:我们只需提供自己的资源分配方案,告诉Ray我想怎么部署这些模型,不管是分开还是合并部署Ray都可以帮我们实现。而复杂的调度策略和通信等事项,就由Ray在后台去做,我们无需关心这个过程。

二、整体流程

本节我们将提供2个例子,帮助大家更好理解使用Ray可以做什么样的“定制化”部署。注意,这些例子只做讲解用,不代表它们一定是训练的最优配置。

2.1 非共同部署


这个例子展示如何完全独立部署各个模型。假设我们有3台node,每台node 8张卡。以下展示其中一种可行的部署方式:

(1)部署4类模型

在这个例子中,4类模型分开部署在node0和node1上。以Actor为例,它分布在“node0的gpu0/1 + node1的gpu0/1”上。这一点是由Ray实现的:我们自己定制化资源分配的方案,进而管控模型的分配方式

而当实际训练时,我们还可进一步引入Deepspeed zero做优化:以Actor为例,上图中的4个Actor构成zero中的数据并行组(world_size = 4),根据zero的配置,我们可以在这4张卡间做optimizer/gradients/weights的切片

(2)部署vllm_engines

前文说过,对于Actor模型,在收集exp阶段我们可以采用vllm之类的框架加速(prompt, responses)的生成。在这个例子中:

  • 1个vllm_engine维护着一个vllm实例,每个vllm实例下维护一个完整的Actor模型,这里我们还假设一个vllm实例按tp_size = 2的方法切割模型。
  • 在node2中,共有4个vllm_engines(也即4个vllm实例),这种分配方式是通过Ray实现的。而每个vllm实例内的分布式推理则是由vllm自己管控。


(3)Actor与vllm_engines之间的通讯

我们称:

  • vllm_engines中的actor为vllm_actor
  • node0/1中的actor为ds_actor

在整个训练过程中,vllm_actor需要和ds_actor保持权重一致。我们来看这个一致性是如何维护的:

  1. 初始化阶段

假设pretrain路径下存储着sft模型,当我们首次开始训练时,ds_actor和vllm_actor都直接从pretrain中加载权重,两者互不影响,独立加载。


2. 训练中

在1个step结束后,ds_actor需要把更新后的权重broadcast给vllm_actor,具体步骤如下:

  • 首先,对ds_rank0 + all_vllm_ranks创建一个通讯组。在本例中:
    • node0/gpu0上的actor是ds_rank0
    • node2中所有的gpu构成all_vllm_ranks。
    • 我们就是把这两者纳入一个通讯组内,这个通讯组的world_size = 9。如果我们多一台node3来做vllm_engines,那么这个通讯组的world_size = 17,以此类推。
  • 假设我们使用ds_zero1/2,则ds_rank0上维护的是完整的actor权重,我们把ds_rank0上的权重broadcast到每一个vllm_rank,如有设置tp,vllm会自动帮我们完整接下来的模型切割。
  • 假设我们使用ds_zero3,则ds_rank0上只维护部分actor权重,那么:
    • ds_rank0先从ds_actor组内all gather回完整的模型权重
    • 再将完整的模型权重brocast给每一个vllm_rank

3. 从检查点恢复训练(load_checkpoint)

当我们需要从检查点恢复训练时,ds_actor会负责把检查点权重broadcast给vllm_actor,方式同2。

(4)整体运作流程

结合2.1开头的图例,我们来简述一下整体运作流程。

  • 首先明确一些表达。例如,node0中的Actor0/1 + node1中的Actor0/1属于相同的数据并行组,所以接下来我们会用它们在dp组中的rank来描述它们,也就是分别改称Actor0/1/2/3。对于其余三类模型也是同理。
  • 接着进行分组:
    • Actor0 / Ref0 / RM0 / Critic0 / vllm_engine0为一组
    • Actor1 / Ref1 / RM1 / Critic1 / vllm_engine1为一组
    • Actor2 / Ref2 / RM2 / Critic2 / vllm_engine2为一组
    • Actor3 / Ref3 / RM3 / Critic3 / vllm_engine3为一组
    • 你可以把每一组想象成原来的一张单卡,那么它的作用就是负责一个micro_batch的训练,这样我们就能大致想象到它们之间是如何配合运作的了。需要注意的是,在我们的例子中,这些实例都是一一对应的(各自有4个实例),但在实际操作中,根据不同用户的资源配置,不一定存在这个一一对应的关系。例如你可能用4卡部署Actor,2卡部署Critic,8个vllm_engines…以此类推。不管怎样,我们应该尽量在处理micro_bathes的各个组间均匀分配负载,在代码里相关的操作如下:

为每个actor分配其对应的critic/reward/ref,并启动每个分组的训练:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/launcher.py#L278-L299
为每个actor分配对应的vllm_engine,并使用vllm_engine进行推理:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ppo_utils/experience_maker.py#L627

2.2 共同部署


同样,我们可以按照自己的需求,选择性地在单卡上部署不同种类的模型,例如下面的例子中,actor/ref共部署,critic/reward共部署,图例如下,运作流程和2.1相似,这里不赘述:

三、Ray的核心概念

在传统的编程中,我们经常使用到2个核心概念:function和class。而在分布式系统中,我们希望可以分布式并行执行这些function和class。Ray使用装饰器@ray.remote来将function包装成Ray task,将class包装成Ray actor,包装过后的结果可以在远端并行执行。接下来我们就来细看task/actor(注意,这里的actor是ray中的概念,不是rlhf-ppo中actor模型的概念)

请大家一定仔细看3.1和3.2节代码中的注释。

3.1 Ray Task

import ray
ray.init()

@ray.remote
def f(x):
    return x * x
# ===================================================================
# 创建driver进程,运行main
# ===================================================================
if __name__ == "__main__":
    # ===================================================================
    # 创建4个worker进程,可以在远端并行执行。
    # 每执行1次f.remote(i),会发生以下事情:
    # - 创建1个worker进程,它将在远端执行函数f(i)
    # - 在driver进程上立刻返回一个引用(feature),该引用指向f(i)远程计算的结果
    # ===================================================================
    futures = [f.remote(i) for i in range(4)]
    # ===================================================================
    # 阻塞/同步操作:等待4个worker进程全部计算完毕
    # ===================================================================
    results = ray.get(futures)) 
    # ===================================================================
    # 确保全部计算完毕后,在driver进程上print结果
    # ===================================================================
    print(f"The final result is: {results}") # [0, 1, 4, 9]



3.2 Ray Actor

import ray
ray.init()

@ray.remote
class Counter(object):
    def __init__(self):
        self.x = 0
    
    def inc(self):
        self.x += 1
    
    def get_value(self):
        return self.x

# ===================================================================
# 创建driver进程,运行main
# ===================================================================
if __name__ == "__main__":
    # ===================================================================
    # 创建1个worker进程,具体做了以下事情:
    # - 在远端创建Counter实例
    # - 在driver端即刻返回对该实例的引用c(称为actor handler)
    # - 我们可以在Ray集群的任何节点上传递和使用这个actor handler。即在任何地方,
    #   我们可以通过c来invoke它对应Counter实例下的各种方法
    # ===================================================================
    c = Counter.remote()

    # ===================================================================
    # 阻塞/同步:通过c来invoke远端Counter实例的get_value()方法,并确保方法执行完毕。
    # 执行完毕后才能接着执行driver进程上剩下的代码操作
    # ===================================================================
    print(ray.get(c.get_value.remote()))  # 0
    
    # ===================================================================
    # Increment the counter twice and check the value again.
    # 道理同上,不赘述
    # ===================================================================
    c.inc.remote()
    c.inc.remote()
    print(ray.get(c.get_value.remote()))  # 2

3.3 Ray cluster架构简图


现在我们已经通过以上例子对Ray运作原理有了一些基本感知,我们来进一步探索一个ray cluster的组成

  • 在一个ray cluster中,会有一台head node和若干worker node
  • Driver process是一种特殊的worker process,它一般负责执行top-level application(例如python中的__main__),它负责提交想要执行的任务,但却不负责实际执行它们。理论上driver process可以运行在任何一台node内,但默认创建在head node内。
  • Worker process负责实际任务的执行(执行Ray Task或Ray Actor中的方法)。
  • 每台node中还有一个Raylet process,它负责管控每台node的调度器和共享资源的分配。
  • Head node中的GCS将会负责维护整个ray cluster的相关服务。

四、代码细节


本章将解读更多代码实践上的重要细节。我们通过图例的方式抽象出代码运行的过程,而具体代码可参考文中给出的相关链接

4.1 训练入口

ppo_ray相关的训练入口在:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/cli/train_ppo_ray.py


在main中我们启动了driver进程,并执行训练函数train(args),这里主要做了如下几件事:

  • 在ray集群上部署Actor/Ref/Critic/RM实例
  • 在ray集群上部署vllm_engines实例
  • 配置Actor和vllm_engines之间的通讯,用于传递权重
  • 训练Actor和Critic模型

我们依次来解读这几个关键步骤。同时为了在表述上消除歧义,我们接下来谈到“Actor”时,会使用Ray-Actor和PPO-Actor来做区分,从之前的介绍中可知,Ray-Actor是指部署在Ray集群中的远端class,PPO-Actor/Ref/Critic/RM都属于Ray-Actor。

4.2 部署Actor/Ref/Critic/RM实例


(1)非共同部署

针对图2.1的情况,我们以PPO-Actor为例,看代码是如何将其部署到Ray集群上的。

  • PPORayActorGroup创建在driver进程上,可将它理解成一种部署方案,专门负责部署PPO中的4类模型
    • PPORayActorGroup中维护着self._actor_handlers,它是一个List[ray.actor.ActorHandle],列表中每个元素表示某个远端Ray-Actor的引用,而这个远端Ray-Actor可以是PPO-Actor/Ref/Critic/RM实例。如前文所说,我们可以在ray集群中的任何位置调用这个handler,来对相应的远端Ray-Actor执行操作。
    • 在本例中,我们创建了4个Ray-Actor(1个master-actor,3个worker_actor)。每个Ray-Actor都运行在一个worker进程中。在创建Ray-Actor的同时,我们也会去修改worker进程的环境变量。后续当我们在这些worker进程中启动ds_zero相关的分布式配置时,ds会读取这些环境变量信息,这样我们就知道哪些Ray-Actor同时又构成ds中的数据并行组。
    • 使用PPORayActorGroup部署模型实例的代码如下:
model = PPORayActorGroup(
        # 为部署该模型的全部实例,我们想用多少台node,例如本例中为2
        args.actor_num_nodes,
        # 为部署该模型的全部实例,我们每台node上想用多少gpu,例如本例中为2
        args.actor_num_gpus_per_node,
        # Actor/Critic/Reward/ReferenceRayActor
        ActorModelRayActor, 
        # pg可理解为,在ray cluster中锁定/预留一片资源,然后只在这片资源上部署该模型全部实例。
        # (pg维护在Head Node的GCS上,参见3.3)
        # 例如本例中,pg锁定的资源为node0 gpu0/1, node1 gpu0/1,
        # 我们只在上面部署ActorModelRayActor全部实例
        pg=pg,
        # 当我们在pg指向的预留资源中分配模型实例时,再进一步指定每个实例占据一张gpu的多少部分
        # 等于1说明每个实例占满一张gpu,即“非共同部署”
        # 小于1说明每个实例只占部分gpu,即“共同部署”,例如PPO-Actor/Ref共同部署在一张卡上
        num_gpus_per_actor=0.75 if pg else 1,
    )
  • ActorModelRayActor创建在远端worker进程上,是Ray-Actor。它包含了设置ds_zero分布式环境、加载模型权重、数据集准备、optimizer/scheduler准备、训练等一系列操作。

PPORayActorGroup代码参见:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/launcher.py#L143
根据这份代码,大家可自行去找Actor/Critic/Reward/ReferenceRayActor的相关实现。

(2)共同部署

针对图2.2的情况,我们以PPO-Actor为例,看代码是如何将其部署到Ray集群上的。

  • PPORayActorGroup:在driver进程上创建2个PPORayActorGroup,分别管理PPO-Actor,PPO-Ref的部署
  • 使用actor_model = PPORayActorGroup(..., pg = pg, num_gpus_per_actor=0.75)创建PPO-Actor部署方案实例;使用ref_model = PPORayActorGroup(..., pg = pg, num_gpus_per_actor=0.25)创建PPO-Ref部署方案实例
  • 这里,两个方案实例使用的pg都是同一个,即这个pg都指向“1台node,每台node 8张卡”这片预留好的资源。
  • num_gpus_per_actor = 0.75/0.25是一种创建trick,虽然我们的最终目的是为了让PPO-Actor和PPO-Ref对半分一张卡(对半=共享,不是指显存上对半分),但是:
    • 假设设置为0.5,当我们实际部署ActorModelRayActor时,Ray先在单卡上部署1个ActorModelRayActor实例,当它准备部署第二个ActorModelRayActor实例时,它发现由于每个实例只占0.5块卡,因此完全可以把第2个实例接着第1个实例在同一张卡上部署,这样就导致最终无法让PPO-Actor和PPO-Ref共享一张卡
    • 假设设置0.75,当我们在单卡上部署完1个ActorModelRayActor实例后,ray发现单卡剩下的空间不足以部署第2个ActorModelRayActor实例,所以就会把第二个实例部署到别的卡上,这样最终实现PPO-Actor和PPO-Ref共享一张卡
    • 所以,这个设置是为了达到不同类型模型的实例共享一张卡的目的,而并非真正指模型实际占据的单卡显存空间。
  • 最后,在这一步中,我们对全部ActorModelRayActor共创建8个worker进程,对全部RefenreceModelRayActor共创建8个worker进程,一共创建16个工作进程。

相关代码依然在:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/launcher.py#L143

4.3 部署vllm_engines实例

  • create_vllm_engines:在driver端,我们通过运行该函数来创建vllm_engines,过程相似于4.2节中的介绍,信息都在图中,这里不赘述。
  • LLMRayActor:worker端Ray-Actor,它主要是把vllm实例进行了一些包装,包装的目的是为了让ds_rank0和all vllm ranks间可以进行PPO-Actor的权重通讯(参见2.1(3))
  • 在上面的例子中,我们会创建4个worker进程(不占gpu资源,只占cpu资源),用于运行管理4个vllm_engine。在每个worker进程内,vllm实例还会创建属于自己的worker进程做分布式运行(这些worker进程会实际占据gpu资源)。

相关代码参见:
https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/vllm_engine.py
https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/vllm_worker_wrap.py



4.4 ds_rank0与vllm_ranks之间的通讯

在2.2中,我们说过,PPO-Actor的ds_rank0需要和all_vllm_ranks进行通讯,传递最新的PPO-Actor权重,例如以下ds_rank0要把完整的权重broadcast给16个vllm_ranks:


我们分成如下几步实现这个目标:

(1)创建通信组



Step1:

代码来自:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/ppo_actor.py#L58
这段代码执行在PPO-Actor0(ds_rank0)所在的worker进程中。这个worker进程将通过handler引用,触发远端每个vllm_engine上的init_process_group操作,并将ds_rank0纳入通讯组

        # Create torch group with deepspeed rank 0 and all vllm ranks
        # to update vllm engine's weights after each training stage.
        #
        # Say we have 3 vllm engines and eache of them has 4 GPUs,
        # then the torch group is:
        # [    0,      1, 2, 3, 4,  5, 6, 7, 8,  9, 10, 11, 12]
        # |ds rank 0 |  engine-0  |  engine-1  |   engine-2   |
        #
        # For ZeRO-1/2:
        #   1. Broadcast parameters from rank 0 to all vllm engines
        # For ZeRO-3:
        #   1. AllGather paramters to rank 0
        #   2. Broadcast parameters from rank 0 to all vllm engines
        if self.vllm_engines is not None and torch.distributed.get_rank() == 0:
            ...
            # world_size = num_of_all_vllm_ranks + 1 ds_rank0
            world_size = vllm_num_engines * vllm_tensor_parallel_size + 1
            ...
            # =====================================================================
            # 遍历每个vllm_engines,将其下的每个vllm_rank添加进通讯组中,这里又分成两步:
            # 1. engine.init_process_group.remote(...):
            #    首先,触发远程vllm_engine的init_process_group方法
            # 2. 远程vllm_engine是一个包装过的vllm实例,它的init_process_group
            #    方法将进一步触发这个vllm实例下的各个worker进程(见4.4图例),
            #    最终是在这些worker进程上执行“将每个vllm_rank"添加进ds_rank0通讯组的工作
            # =====================================================================
            refs = [
                engine.init_process_group.remote(
                    # ds_rank0所在node addr
                    master_address, 
                    # ds_rank0所在node port
                    master_port,
                    # 该vllm_engine的第一个rank在"ds_rank0 + all_vllm_ranks“中的global_rank,
                    # 该值将作为一个offset,以该值为起点,可以推算出该vllm_engine中其余vllm_rank的global_rank
                    i * vllm_tensor_parallel_size + 1, 
                    world_size,
                    "openrlhf",
                    backend=backend,
                )
                for i, engine in enumerate(self.vllm_engines)
            ]
            # =====================================================================
            # 将ds_rank0添加进通讯组中
            # =====================================================================
            self._model_update_group = init_process_group(
                backend=backend,
                init_method=f"tcp://{master_address}:{master_port}",
                world_size=world_size,
                rank=0,
                group_name="openrlhf",
            )
            # =====================================================================
            # 确保all_vllm_ranks都已添加进通讯组中
            # =====================================================================
            ray.get(refs)



Step2:

代码来自:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/vllm_worker_wrap.py#L11
这段代码实际运行在每个vllm_engine(即每个包装后的vllm实例)下的worker进程内。例如tp_size=2,那么每个vllm实例下就有2个worker进程,这两个worker进程都会运行这段代码。

class WorkerWrap(Worker):
    def init_process_group(self, master_address, master_port, rank_offset, world_size, group_name, backend="nccl"):
        """Init torch process group for model weights update"""
        assert torch.distributed.is_initialized(), f"default torch process group must be initialized"
        assert group_name != "", f"group name must not be empty"
        # =====================================================================
        # torch.distributed.get_rank(): 在当前vllm_engine内部的rank,
        #                               例如在tp_size = 2时,这个值要么是0,要么是1
        # rank_offset:当前vllm_engine中的第一个rank在“ds_rank0 + all_vllm_ranks"中的global_rank
        # 两者相加:最终得到当前rank在“ds_rank0 + all_vllm_ranks"中的global_rank
        # =====================================================================
        rank = torch.distributed.get_rank() + rank_offset
        self._model_update_group = init_process_group(
            backend=backend,
            init_method=f"tcp://{master_address}:{master_port}",
            world_size=world_size,
            rank=rank,
            group_name=group_name,
        )
        ...




(2)_broadcast_to_vllm

构建好通讯组,我们就可以从ds_rank0广播PPO-Actor权重到all_vllm_ranks上了,这里也分成两步。


Step1:PPO-Actor ds_rank0发送权重

代码在:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/ppo_actor.py#L146
这段代码运行在ds_rank0对应的worker进程中

      def _broadcast_to_vllm(self):
        # avoid OOM
        torch.cuda.empty_cache()
        model = self.actor.model.module
        count, num_params = 0, len(list(model.named_parameters()))
        for name, param in model.named_parameters():
            count += 1  # empty_cache at last param

            # Fire all vllm engines for broadcast
            if torch.distributed.get_rank() == 0:
                shape = param.shape if self.strategy.args.zero_stage != 3 else param.ds_shape
                refs = [
                    # 远端vllm_engine的每个rank上,初始化一个尺寸为shape的empty weight张量,
                    # 用于接收广播而来的权重
                    engine.update_weight.remote(name, dtype=param.dtype, shape=shape, empty_cache=count == num_params)
                    for engine in self.vllm_engines
                ]

            # For ZeRO-3, allgather sharded parameter and broadcast to all vllm engines by rank 0
            # ds_rank0发出权重(视是否使用zero3决定在发出前是否要做all-gather)
            with deepspeed.zero.GatheredParameters([param], enabled=self.strategy.args.zero_stage == 3):
                if torch.distributed.get_rank() == 0:
                    torch.distributed.broadcast(param.data, 0, group=self._model_update_group)
                    ray.get(refs) # 确保所有vllm_ranks接收权重完毕




Step2: 各个vllm_ranks接收权重

代码在:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/vllm_worker_wrap.py#L29
代码运行在每个vllm_engine(即每个包装后的vllm实例)下的各个worker进程中。例如tp_size = 2,那么每个vllm实例下有2个worker进程,这2个worker进程都会运行这段代码。

    def update_weight(self, name, dtype, shape, empty_cache=False):
        """Broadcast weight to all vllm workers from source rank 0 (actor model)"""
        if torch.distributed.get_rank() == 0:
            print(f"update weight: {name}, dtype: {dtype}, shape: {shape}")

        assert dtype == self.model_config.dtype, f"mismatch dtype: src {dtype}, dst {self.model_config.dtype}"
        # 创建同尺寸空张量用于接收ds_rank0广播来的权重
        weight = torch.empty(shape, dtype=dtype, device="cuda")
        # 接收权重
        torch.distributed.broadcast(weight, 0, group=self._model_update_group)
        # 使用接收到的权重进行更新
        self.model_runner.model.load_weights(weights=[(name, weight)])

        del weight




4.5 PPO-Actor/Critic Training


正如2.1(4)中所说,我们将部署在ray集群上的PPO-Actor/Ref/Critic/RM实例们进行分组,每组分别负责一份micro-batch的训练,上图刻画了某个组内的训练流程。一组内的训练流程发起自PPO-Actor实例(fit方法),注意不同颜色的worker0表示的是不同工作进程。共分成如下步骤执行。


Step1:发送prompts,并从vllm_engine上收集(prompt, response)。

代码参见:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ppo_utils/experience_maker.py#L627



Step2:从Ref/Reward/Critic上收集并处理exps

代码参见:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ppo_utils/experience_maker.py#L492



Step3: 确保将处理后的exps传送给Critic,并行执行Actor和Critic的训练

将exps传送给Critic:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ppo_utils/experience_maker.py#L470
Actor训练:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/ppo_actor.py#L125
Critic训练:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/ppo_actor.py#L122
我们在Actor实例所在的worker进程上出发Actor和Critic的训练。以上代码只给出了训练入口,更多细节需要顺着入口去阅读。



Step4:vllm_engine权重更新。

代码参见:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ray/ppo_actor.py#L130

五、RLHF-PPO算法细节

到目前为止,我们已将Ray分布式训练部分介绍完毕。有了对工程架构的认识,很多朋友可能想进一步探索算法上的细节:例如在4.5图例中,我们收集的是什么样的exps?又是如何对PPO-Actor/Critic进行更新的?

我们知道整个RLHF-PPO训练过程大致分成2步:

  • Stage1:收集exps
  • Stage2:使用收集到的exps计算actor_loss和critic_loss,用于训练actor和critic

在OpenRLHF中的核心代码为:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ppo_trainer.py#L19

下面我们分别解读这2个stage的过程

5.1 Stage1:collect exps


假设大家已读过上述说的【代码篇】和【理论篇】,那么对这张图的大部分细节应该不难理解。这里我们对优势(假设选择的是GAE)再做一些补充说明


在理论篇8.3(2)中,我们详细介绍过GAE定义和使用它的原因,第t步GAE的定义为:

由上式可知,如果我们想要计算出每个t时刻的GAE,我们需要从头到尾遍历2次序列:第1次计算 δt ;第2次计算 At 。我们期望尽可能减少计算复杂度,仅在1次遍历中同时计算出 δt,At ,那么我们可以按如下步骤改写

在OpenRLHF中,还设计到“advantage norm”,“reward clip & norm”等细节操作,篇幅限制不展开,大家可以自行阅读。

核心代码在:https://github.com/OpenRLHF/OpenRLHF/blob/bb46342711a203c457df2fbca5967fd0549557e0/openrlhf/trainer/ppo_utils/experience_maker.py

5.1 Stage2:Training

大家可以对照理论篇8.5和8.6节对actor_loss和critic_loss的介绍来理解上面的流程图~另外,ybq佬写过一篇对OpenRLHF中的loss分析,也可以对照阅读:ybq:OpenRLHF学习笔记-loss篇

六、参考

1、OpenRLHF
2、Ray official document
3、Ray official architecture whitepaper(要梯子,建议想看ray架构的朋友,直接看这个最新的官方白皮书,不要看2018年的那篇paper了,那个比较老了)
4、推荐一篇快速了解Ray 应用层核心概念的blog(要梯子)
5、Ray
6、vllm

ThinkSound:多模态大语言模型中的链式思维推理,用于音频生成与编辑

ThinkSound 是一个统一的 Any2Audio 生成框架,通过链式思维(Chain-of-Thought, CoT)推理进行流匹配指导

基于 PyTorch 的多模态音频生成与编辑实现:可基于视频、文本、音频及其组合,生成或编辑音频,底层由多模态大语言模型(MLLMs)逐步推理驱动。

主要特性

  • Any2Audio:支持任意模态(视频、文本、音频或其组合)生成音频。
  • 视频转音频 SOTA:在多个 V2A 基准上取得最新最优结果。
  • CoT 驱动推理:基于链式思维推理,实现可组合、可控的音频生成。
  • 交互式面向对象编辑:通过点击视觉对象或文本指令,细化或编辑特定声音事件。
  • 统一框架:单一基础模型,支持生成、编辑与交互式工作流。

Abstract

ThinkSound 将音频生成与编辑分为三个交互式阶段均由基于 MLLM链式思维(CoT)推理指导

  1. 拟音生成(Foley Generation): 从视频生成基础、语义与时序对齐的声景。
  2. 面向对象的细化: 通过点击或选择视频中的对象区域,对用户指定对象的声音进行细化或添加。
  3. 定向音频编辑: 使用高级自然语言指令对生成音频进行修改。

在每个阶段,一个多模态大语言模型都会生成与上下文相符的 CoT 推理内容,用以指导统一的音频基础模型。此外,我们还引入了 AudioCoT,一个包含结构化推理标注的综合数据集,用于建立视觉内容、文本描述与声音合成之间的联系。

带有链式思维(CoT)的 ThinkSound:(1) 由 CoT 驱动的拟音合成,捕捉语义与时间细节;(2) 面向对象的交互式精细化处理,实现用户控制;(3) 有针对性的音频编辑。

为视频生成真实的声音不仅仅是识别物体,它还需要对复杂的视觉动态和上下文进行推理,比如判断一只猫头鹰是在鸣叫还是在拍打翅膀,识别树枝轻微的摆动,并在一个场景中同步多个声音事件。

ThinkSound——在技术上,提出了三个关键创新:

  • a) 对 MLLM 进行 AudioCoT 微调,使其能生成结构化、面向音频的推理链,明确捕捉时间依赖关系、声学属性与复杂音频事件的分解过程;
  • b) 设计了一个基于 flow matching 的统一音频基础模型,支持所有三个阶段,能够从任意组合的输入模态(视频、文本、音频)中合成高保真音频。该模型直接受益于 MLLM 提供的细致 CoT 推理,将复杂音频场景分解为可控组件,在保证整体连贯性的同时实现重点声音事件的精确合成;
  • c) 引入了一个新颖的基于点击的交互界面,使用户能够选择特定视觉对象进行音频精修,CoT 推理机制将视觉关注转化为语境合理的声音合成过程。

AudioCoT Dataset for CoT-Guided Generation and Editing

Multimodal Data Sources

AudioCoT 数据集包含视频-音频和音频-文本对。对于视频-音频数据,我们利用 VGGSound 和 AudioSet中精选的非语音子集,以确保广泛覆盖现实世界的视听事件。对于音频-文本数据,我们聚合了来自 AudioSet-SL 、Freesound 、AudioCaps和 BBC Sound Effects数据对,从而构建了一个用于训练多模态模型的多样化且具有代表性的语料库。

首先移除静默的音频-视频片段,仅保留含有效内容的素材。针对AudioSet子集,根据标签信息进一步剔除了含有人声的片段,以专注于非语音音频。随后将所有音视频片段统一分割为9.1秒的固定时长,舍弃较短片段以保证数据 uniformity(统一性)。为实现数据平衡,保持音乐与音效样本的比例约为1:1,确保两类别的均衡表征。

自动化 CoT 生成流程:

 AudioCoT 数据集构建流程概览

第一阶段:基础拟音思维链生成

  • 视频-音频对处理
    1. 使用VideoLLaMA2通过差异化提示策略提取视频的时序与语义信息
    2. 结合Qwen2-Audio生成音频描述
    3. 将视频描述与音频描述通过GPT-4.1-nano整合为完整思维链
  • 纯音频-文本对处理
    采用简化流程(无需VideoLLA2),直接生成音频描述后与现有文本标注整合
    该阶段生成的思维链能捕捉内容条件与对应音频元素的复杂关联,确保两类数据共同促进音频生成推理的全面理解

第二阶段:交互式对象中心思维链生成
为实现对象聚焦的音频生成,开发基于Grounded SAM2的ROI提取框架:

  1. 对象定位:以音频描述为提示,生成潜在发声物体的边界框
  2. 时序追踪:跨视频帧持续跟踪坐标变化
  3. 语义增强:VideoLLA2为每个ROI片段提供详细语义描述
  4. 复杂操作处理
    • 建立分层推理结构,合并目标视频CoT 与参考视频的思维链CoT 以构建全局上下文
    • 结合ROI特定生成信息,通过GPT-4.1-nano生成连贯的操作逻辑

第三阶段:基于指令的音频编辑思维链生成
针对指令引导的音频编辑任务,基于四类核心(扩展、修复、添加和移除)分析并整合来自第一阶段的 CoT 信息。这些操作涵盖从扩展序列到删除不需要的片段的各种场景。GPT-4.1-nano 处理这些整合的信息,生成特定于指令的 CoT 推理链,同时执行相应的音频操作,创建(指令-CoT、输入音频、输出音频)三元组,用于模型训练和评估。

ThinkSound

 ThinkSound 架构概览。 左图: 我们的多模态 LLM 框架,其中经过微调的 VideoLLaMA 2 模型生成用于音频生成和编辑的 CoT 推理。 右图: 我们增强型多模态 Transformer 架构,该架构以 MM-DiT 为骨干,具有用于处理多模态输入的专用路径和 CoT 驱动的条件反射,从而实现高保真、基于情境的音频生成。

Overview

ThinkSound 引入了一个新颖的分步式交互式音频生成和编辑框架,该框架由 CoT 推理引导。我们的方法将复杂的 V2A 任务分解为三个直观的阶段:(1) 基础拟音生成,创建语义和时间匹配的音景;(2) 通过用户点击进行基于区域的交互式细化;以及 (3) 基于高级指令的定向音频编辑。在每个阶段,MLLM 都会生成 CoT 推理,引导统一的音频基础模型制作和细化音轨。

使用多模态 LLM 进行 CoT 推理

为了实现分步式、情境感知的音频生成,我们利用 VideoLLaMA2 作为核心多模态推理引擎。VideoLLaMA2 之所以被选中,是因为其在融合视频、文本和音频模态方面拥有领先的能力,而其先进的时空建模对于捕捉视觉事件与其对应听觉表现之间的微妙相互作用至关重要。

通过AudioCoT数据集对VideoLLA2进行微调,使其适配音频推理领域。该数据集包含专为视听任务定制的丰富标注推理链,通过微调过程使模型具备三大核心能力:

(1)音频中心化理解能力

  • 声学特性推断(如材料属性、空间混响等)
  • 声音传播建模
  • 视听对应关系推理(包括音频事件间的时序与因果关系,例如”脚步声先于开门声,随后出现对话声”)

(2)结构化思维链分解能力
将复杂音频生成/编辑任务拆解为明确可执行的步骤序列

(3)多模态指令跟随能力
可靠地解析并执行跨模态的多样化生成/编辑指令

如图3所示,微调目标采用标准的下一个token预测交叉熵损失。通过这种针对性适配,VideoLLA2被转化为专用音频推理模块,能够生成上下文精确的思维链指令,驱动ThinkSound流程的每个阶段。

CoT 引导的统一音频基础模型

ThinkSound 的核心是我们统一的音频基础模型,它能将 CoT 推理无缝转换为高质量音频,具体细节见图 3 右侧部分。使用预训练的 VAE将音频编码为潜在表示,并采用条件流匹配对模型进行训练,其中速度场预测以多模态上下文为条件,包括视觉内容、CoT 推理、文本描述和音频上下文。为了支持任意组合的输入模态,我们在训练过程中引入了无分类器引导的随机丢弃方法。通过以概率 p_drop 随机丢弃不同模态的组合,使模型在推理阶段能够处理任意输入配置——这对于我们的交互式框架至关重要。我们还结合了策略性音频上下文遮蔽,以支持音频修补和扩展等高级编辑操作

在文本处理方面,我们采用了双通道编码策略:MetaCLIP对视觉字幕进行编码,提供场景级上下文;而 T5-v1-xl则处理结构化的 CoT 推理,以捕捉详细的时间和因果关系。这两种互补的表示被有效融合,MetaCLIP 的特征作为全局条件信号,而 T5 的输出则支持基于推理的精细控制。

我们改进的 MM-DiT 架构基于多模态生成建模领域的最新进展,包含三大关键组件:(1)采用混合型 Transformer 主干网络,在模态专用与共享处理之间交替进行。多流 Transformer 块为每个模态维护独立参数,同时共享注意力机制,从而高效处理多样输入,同时兼顾跨模态学习。(2)设计了自适应融合模块,通过门控机制对视频特征进行上采样并与音频潜变量融合。这不仅能够突出显著的视觉线索、抑制无关信息,还确保视频信息直接参与后续的单流 Transformer 块。通过将视频整合到音频潜变量空间,模型可以更好地捕捉细微视觉细节及其对声景的微妙影响,实现比仅依赖音频潜变量更丰富的跨模态推理。(3)通过对字幕和视频的 CLIP 特征进行均值池化,实现全局条件控制,并借鉴 MMAudio,引入同步特征以提升音视频时间对齐效果。最终得到的全局条件被添加到时间步嵌入中,并通过自适应层归一化(AdaLN)注入多流与单流块。

 逐步 CoT 引导的音频生成和编辑

通过支持输入模式与 CoT 的灵活组合,ThinkSound 支持将音频生成分解为图 1 所示的三个直观阶段。该三阶段流程通过直观的交互式工作流程,实现了逐步精细化、高度定制化的音频生成,CoT 推理在每个步骤中将用户意图与音频合成连接起来。

阶段 1:基于 CoT 的拟音生成
在第一阶段,系统分析整段视频以识别声学要素及其关系。经过微调的 MLLM 生成详细的 CoT 推理,明确识别主要声事件、环境元素、声学属性以及它们的时间依赖关系——确定物体何时发声及声音间的相互作用。这种结构化推理指导音频基础模型生成高保真音频,精准匹配视觉场景的语义内容与时间动态。借助 CoT 推理将复杂音频场景拆解为显式声源,模型能够生成多样且连贯的声景,捕捉微妙视觉线索与运动动态,实现逼真的音频合成。

阶段 2:交互式对象聚焦音频生成
第二阶段引入交互框架,让用户通过关注特定视觉元素来优化初步声景。借助简单的点击界面,用户可以选择感兴趣的物体进行音频强化。不同于第一阶段的整体生成方式,此阶段采用基于目标区域(ROI)的局部细化,利用分割出的目标区域指导定向音频合成。经过微调的 MLLM 针对所选 ROI 生成专门的 CoT 推理,关注该物体在全局背景下的声学特性。模型在这些结构化推理引导下生成物体专属声音,与第一阶段生成的音轨自然融合。值得注意的是,此阶段的基础模型将已有音频上下文作为附加条件信号纳入考虑。

阶段 3:基于指令的音频编辑
在最后阶段,用户可通过高层次的编辑指令来优化音质或修改特定元素。MLLM 将自然语言指令转译为具体的音频处理操作,利用 CoT 推理综合视觉内容和当前音频状态。基础模型在此推理及现有音频上下文条件下执行定向修改,同时保持整体连贯性。通过自然语言控制,非专业用户也可以完成复杂的音频操作,包括添加声音、移除声音、音频修补以及音频延展。

Results

虽然目前的 MLLM 模型能够很好地理解和推理语义信息,但它们在理解视频的精确时间和空间信息方面仍然存在局限性。例如,在定位声音事件的精确时间戳时,MLLM 模型经常无法提供准确的结果或给出错误的结果。此外,目前用于音频生成的开源视音频数据集在多样性和覆盖范围方面存在局限性,可能缺少稀有或特定文化的声音事件。未来,我们将继续探索更加多样化和全面的数据集,以提升模型的性能。此外,我们还将探索更有效的方法来提升生成音频的时间和空间对齐效果。

Ke-Omni-R :通过思考实现高级音频推理

Github:https://github.com/shuaijiang/Ke-Omni-R 【开源训练和推理代码】

贡献:用于将GRPO/思考过程 加入到语音大模型的强化训练过程中。

  • [1] Xie, Zhifei, et al. “Audio-Reasoner: Improving Reasoning Capability in Large Audio Language Models.” arXiv preprint arXiv:2503.02318.
  • [2] Ma, Ziyang, et al. “Audio-CoT: Exploring Chain-of-Thought Reasoning in Large Audio Language Model.” arXiv preprint arXiv:2501.07246.
  • [3] Li, Gang, et al. “Reinforcement Learning Outperforms Supervised Fine-Tuning: A Case Study on Audio Question Answering.” arXiv preprint arXiv:2503.11197
  • [4] Xu, Jin, et al. “Qwen2.5-Omni Technical Report.” arXiv preprint arXiv:2503.20215

Ke-Omni-R 是基于 Qwen2.5-Omni 构建的高级音频推理模型。构建音频推理模型,通过强化学习引入深度思考过程,提升复杂任务的理解和推理能力。仅使用 10,000 个训练后样本,Ke-Omni-R 就在 MMAU Test-mini 和 Test 基准测试中取得了最佳性能。其开发过程中的关键洞察包括:

  • GRPO 算法 :GRPO 算法显著增强了已经很强大的基础模型(Qwen2.5-Omni-7B)的性能,即使在看不见的语音领域也表现出卓越的泛化能力。
  • 思考过程 :融入简洁的思考过程(少于 50 个字)对于提高推理能力起着至关重要的作用。
  • KL 散度 :通过利用 KL 散度,在 GRPO 训练期间观察到轻微的改进。
  • 领域比例 vs. 数据量 :领域多样性比数据量更重要。我们仅使用了 10,000 个样本,其中 5,000 个从 AVQA 中随机选取,另外 5,000 个从 MusicBench 中选取。

Performance: Accuracies (%)↑ on MMAU Test-mini and Test benchmark

ModelMethodSound (Test-mini)Sound (Test)Music (Test-mini)Music (Test)Speech (Test-mini)Speech (Test)Average (Test-mini)Average (Test)
Human*86.3178.2282.1782.23
Gemini Pro 2.0 FlashDirect Inference*56.4661.7358.6856.5351.6561.5355.6059.93
Audio Flamingo 2Direct Inference*61.5665.1073.9572.9030.9340.2655.4859.42
GPT4o + Strong Cap.Direct Inference*57.3555.8349.7051.7364.8668.6657.3058.74
Llama-3-8B-Instruct + Strong Cap.Direct Inference*50.7549.1048.9348.9355.2562.7052.1053.57
Qwen2-Audio-7B-InstructDirect Inference*54.9545.9050.9853.2642.0445.9049.2052.50
SALAMONNDirect Inference*41.0040.3034.8033.7625.5024.2433.7032.77
Audio-Reasoner(Qwen2-Audio-7B-Instruct)[1]60.0664.3060.7061.71
Audio-Cot(Qwen2-Audio-7B-Instruct)[2]61.8656.2955.2657.80
R1-AQA(Qwen2-Audio-7B-Instruct)[3]68.7769.7664.3761.4063.6662.7065.6064.36
Qwen2.5-Omni-7B[4]67.8769.1659.7665.60
Qwen2.5-Omni-3B[4]70.2760.4859.1663.30
Ke-Omni-R-3B(Qwen2.5-Omni-3B)GRPO w/ think (ours)72.3771.8765.5759.6064.2664.1767.4065.17
Ke-Omni-R(Qwen2.5-Omni-7B)GRPO w/o think (ours)69.6770.5767.6664.0066.3767.1767.9067.24
Ke-Omni-R(Qwen2.5-Omni-7B)GRPO w/ think (ours)69.3771.9069.4667.1367.8767.1068.9068.71

Performance: CER/WER (%)↓ on ASR benchmark

ModelMethodWenetSpeech test-netWenetSpeech test-meetingLibriSpeech test-cleanLibriSpeech test-other
Qwen2.5-Omni-3B[4]6.38.12.24.5
Qwen2.5-Omni-7B[4]5.97.71.83.4
Ke-Omni-3Bours11.716.11.83.8
Ke-Omni-7Bours7.59.81.63.1

DPO为什么会让大语言模型输出变长

摘自:https://zhuanlan.zhihu.com/p/5830338806

总的来说,DPO让模型输出变长主要可以分为以下几个原因:

  1. RM和模型评测的长度偏好。不管是Reward Model还是当前用与评测的模型(即便是GPT4)都会存在比较明显的长度偏好,即倾向于给更长的回答一个更高的分数。这一点已经有非常多工作给出过分析了。
  2. 训练数据本身长度分布不均衡。实战过程中往往就是用RM进行排序构造训练数据,RM的长度偏好就是会导致训练数据中容易出现chosen比rejected更长的情况。训练数据的长度差异(chosen比rejected长)就会导致训练后模型输出变长。
  3. 数据长度差异导致的reward被高估或低估。《Eliminating Biased Length Reliance of Direct Preference Optimization via Down-Sampled KL Divergence》中发现,DPO的算法本身也存在对response长度的依赖,chosen和rejected之间的长度差异可能会导致reward被高估/低估(overestimated or underestimated rewards)。即,当chosen过短时,reward会被低估,而当chosen过长时,reward会被高估
  4. DPO算法本身的长度敏感性。《Length Desensitization in Direct Preference Optimization》中提到,response长度会影响到似然概率的大小,并且进一步影响到训练优化方向:当chosen更长时,DPO会往chosen的方向进行优化(增大chosen概率),从而使输出变长;而rejected更长时,DPO会往远离rejected的方向优化(降低rejected概率),但却未必会让输出变短。

如何解决:

  1. RM的优化:前面讲的都是对DPO进行长度控制的工作,但对RM本身的长度偏好进行优化的工作没有看到太多,如果大家有看到相关的也可以在评论区提供一下。如果将RM本身的长度偏好问题解决的话,那就可以极大程度上解决训练数据的长度分布均衡问题了。
  2. 数据的优化:有些工作会在数据构造时对长度进行综合考虑,如对RM打分进行长度归一后再排序、采样多个答案进行排序时根据均值方差限制chosen的长度等,通过这些方式可以减少长度差距过大的情况。如果数据本身的长度分布均衡了,也能一定程度上减缓这种问题。
  3. 训练算法上的优化:如果从LD-DPO的分析上看,即便数据分布比较均衡,只要存在长度差异,DPO本身的长度敏感性就是会导致模型输出变长,因此可能还是需要一些算法层面的优化,比如在DPO阶段加入SFTloss就是一种简单有效的方法,在很多公开的大模型技术报告中也都有用到该方法。另外R-DPO、SamPO和LD-DPO的长度控制效果都算是比较好的方法。

DPO面临的一个问题(准确来讲是一种现象)就是会让大模型的输出变长,且多轮DPO的话会让模型输出越来越长。本篇文章我们将结合搜集到的一些相关工作,探讨一下业界对该现象的一些分析,探究这一现象产生的根本原因,以及如何有效地解决。

首先我们需要思考一个问题,模型输出变长到底是不是一件坏事?一般来说,输出变长可能会使内容更加详细,信息量更丰富,回复质量更高,用户体验更好。但如果过度长,输出了很多冗余信息,回复质量没有明显改善,反而带来了推理成本的增加,回复变得啰嗦,用户体验反而变差了。

因此,无论是从用户体验的角度还是多轮DPO能否run下去的角度,做好长度控制都是有必要的。

相关工作

先简要介绍一些相关工作,然后后面详细总结。

1.《Disentangling Length from Quality in Direct Preference Optimization》(简称R-DPO)

在这之前的一些RL的工作也有分析过长度爆炸问题,但该文章可能是第一个提出DPO的长度爆炸问题的。

文章中发现,无论是RL训练中使用的Reward Model还是用来评测模型效果的打分模型(如GPT-4)都表现出明显的长度偏好,即会给更长的答案一个更高的分数(如下图)。且在一些公开的DPO训练数据集中,chosen的长度往往会比rejected更长,而这可能就是DPO后的模型输出长度明显比SFT模型更长的原因

为了解决这个问题,该文章提出了一种长度正则化的策略,即在计算loss的时候添加一个长度正则项,从而避免模型对长度的过度拟合,公式如下:

其中 |yw| 表示chosen的长度, |yl| 表示rejected的长度,从公式中可以看出,当chosen与rejected的长度差距越大,正则项的值越大,从而实现对长度的“惩罚”效果。

从文章中的实验结果可以看出,该方法确实可以在尽可能减少性能损失的前提下有效解决长度增长问题。(有时还是会损失一定的性能。)

2.《SimPO: Simple Preference Optimization with a Reference-Free Reward》(简称SimPO)

陈丹琦团队的工作,直接去掉了reference model,用长度归一的方式实现长度控制。其loss如下:

文章中提到了很多输出长度相关的内容,但核心贡献并不是做长度控制,而是用一种更简单高效的方法实现偏好训练。从公式上看,和原始DPOloss相比主要有两处不同,一个是分母从reference model的logp替换成了长度,另外就是增加了一个 γ ,类似一个offset的作用。不过其中对chosen和rejected的reward做长度归一的部分,直觉上看起来应该是能起到一定的长度控制效果的。

不过从论文中的实验结果看,该方法的效果还是比较好的(当时声称训出最强8B大模型),但与标准DPO相比似乎并没有实现长度控制的效果。

3.《Eliminating Biased Length Reliance of Direct Preference Optimization via Down-Sampled KL Divergence》(简称SamPO

这篇论文对DPO后长度变长的问题进行了一定的分析,提出的一个核心观点是:DPO的算法本身也存在对response长度的依赖,chosen和rejected之间的长度差异可能会导致reward被高估/低估(overestimated or underestimated rewards)。即,当chosen过短时,reward会被低估,而当chosen过长时,reward会被高估。

这篇工作中提出的一种方式就是在token级别下采样的概率特征,以计算正则化的KL散度,从而减少因pair长度不同而导致的奖励偏差。其loss的计算如下:

从公式可以看出,该方法的核心就是在计算reward的时候不再是全部token的条件概率的累乘(取log后就是累加),而是随机采样公共数量的token进行累乘。这样即便chosen和rejected长度不同,参与reward计算的token数是一样的。也就是说,在SamPO训练过程中,魔都看到的chosen和rejected相当于是完全等长的。

从文章中的实验结果看,该方法确实能有效控制模型输出长度的增长,甚至在多轮DPO依然能有效控制长度。但是在性能上看依然做不到碾压标准DPO的效果。

但该方法有两个风险便是:

  1. 本身DPO就存在一定的波动,随机下采样可能会导致训练的稳定性不强;
  2. 随机采样必然会导致一些信息缺失,如果采样时舍弃掉了一些非常重要的token可能会影响到训练效果。

4.《Length Desensitization in Direct Preference Optimization》(简称LD-DPO)

该论文可能是第一个从理论层面分析DPO后模型输出变长的原因的,其核心分析主要包括两方面:

  1. DPO的梯度优化方向和chosen/rejected的似然概率成反比。
  2. Response长度对似然概率的影响极大,因此长度会直接影响reward的计算,并影响到DPO的优化方向。

上图是一个对训练数据的统计热力图,图中,横坐标为chosen的长度,纵坐标为rejected的长度,颜色深度表示 log⁡πθ(yl|x)−log⁡πθ(yw|x) 值的大小。第一张图(a)是标准DPO,可以看出长度差距越大时,颜色越深,也就说明长度差距可能会导致reward计算产生bias,且长度差距越大这种bias越大。而这种bias会进一步影响到DPO的优化方向,使其往输出更长的方向进行优化。

该文章提出的解决方案是在计算似然概率时对长度进行解耦,将更长的答案拆成“公共长度部分”和“额外部分”,并进一步将后者拆分为真实偏好和冗余偏好,并对其中的冗余部分进行降权操作,通过一系列推导后将 πθ(y|x) 转化为如下的形式(可近似理解为完整似然概率部分与公共长度部分似然概率的加权和):

从公式上看,这种方式可以让长度更长的那个response(不管是chosen还是rejected)实现一定的缩放(),减少长度带来的似然概率的断崖式下滑,使其与另一个短response(不受影响)之间更具可比性,同时又不会像SamPO那样完全舍弃掉额外部分的信息。

从论文中的实验结果看,这种方法能够实现比较好的长度控制,且模型性能还能有一定提升,并且可以通过调整参数 α 可以实现不同程度的控制效果。另外文章还提出一个比较有意思的发现,就是过度冗余的回答可能反而会损害模型的推理能力,他们通过这种方法控制长度后,模型的推理能力也有明显提升。

其他工作

除此之外,还有一些工作直接在数据上做文章,通过控制chosen和rejected的长度差距来实现长度控制,如《Following Length Constraints in Instructions》(简称LIFT-DPO)。以及在一些开源模型的技术报告中我们也能看到一些相关的长度控制方法,如在利用RM打分排序时就综合考虑长度问题等,这些数据工作就不再详细展开了。

如何实现有效的长度控制?

  1. RM的优化:前面讲的都是对DPO进行长度控制的工作,但对RM本身的长度偏好进行优化的工作没有看到太多,如果大家有看到相关的也可以在评论区提供一下。如果将RM本身的长度偏好问题解决的话,那就可以极大程度上解决训练数据的长度分布均衡问题了。
  2. 数据的优化:有些工作会在数据构造时对长度进行综合考虑,如对RM打分进行长度归一后再排序、采样多个答案进行排序时根据均值方差限制chosen的长度等,通过这些方式可以减少长度差距过大的情况。如果数据本身的长度分布均衡了,也能一定程度上减缓这种问题。
  3. 训练算法上的优化:如果从LD-DPO的分析上看,即便数据分布比较均衡,只要存在长度差异,DPO本身的长度敏感性就是会导致模型输出变长,因此可能还是需要一些算法层面的优化,比如在DPO阶段加入SFTloss就是一种简单有效的方法,在很多公开的大模型技术报告中也都有用到该方法。另外R-DPO、SamPO和LD-DPO的长度控制效果都算是比较好的方法。

最后结合我自己的一些尝试来直接对比一下上面的四种方法:

  1. R-DPO是通过加正则项的方式实现长度控制,说是正则项,但其实只是一个常数,其原理相当于是对每条数据加上一个权重(文章中也提到了这点),即当chosen和rejected长度差距大时降低该数据的权重。也就是说,该方法其实是让模型减少对长度差距大的数据的学习权重。这种方法确实可以实现一定的长度控制效果,但必然会减少一些数据的利用率,这可能也是训练效果会有一定损失的原因。我自己尝试了一下该方案,实验下来确实可以做到长度控制效果,但大部分情况下性能都会比标准DPO差一些。
  2. SimPO是用长度归一来替换Reference Model的KL约束,理论上和长度控制其实没有太大关系,更多的是简化训练和提升性能。实验结果确实也体现了并不会比标准DPO更短。(该方法热度很高,但网络上褒贬不一,很多人表示无法复现结果。)根据我自己实验经验来看,跑出好的结果需要仔细调参,论文推荐的超参不一定适合所有情况。
  3. SamPO是直接用下采样的方式,强行将模型视角下的长答案变得和短答案一样长,该方法给人的直观感受就是长度控制效果肯定很好,但是很可能会有性能损失。但我自己实验下来,长度控制效果和R-DPO差不多,但是性能也比较不稳定,更换随机种子就会导致性能产生波动。我也尝试过将随机下采样改为top-k采样,即保留概率最大的top-k个token,但效果并不会比随机更好(这么直觉的方法可能论文作者也尝试过了)。
  4. LD-DPO的方法是只对答案过长的部分做了解耦和降权处理,通过降低过长部分的权重来实现整个条件概率的缩放,看起来是四种方法中实现最优雅的一种,既降低了长度差异带来的reward bias问题,又不会丢弃信息,相当于是用极小的代价实现了概率缩放目的。从论文中贴出的结果看,确实也是性能最强的一个,长度控制效果也是最好的。但论文代码没有开源,所以没有实验验证。但从公式上看复现难度应该不是很大,有能力的可以尝试复现一下看看效果。

LLM训练-人工强化反馈对齐算法:RLHF, RLAIF, PPO, DPO and More

参考论文:A Comprehensive Survey of LLM Alignment Techniques: RLHF, RLAIF, PPO, DPO and More

参考代码:

端到端推理学习增强推理能力方法

a.推理时扩展(Inference-time scaling): 如链式思维(CoT)或自我一致性(Self-Consistency),以增强模型的推理能力;【cot:核心思想是将复杂问题分解为一系列可解释的中间步骤。通过明确的推理链条,模型能够逐步解决原本可能超出其直接推理能力的问题。思维链方法特别适用于涉及多步骤推理的任务,如数学题、多重逻辑推理问题等。 Self-Consistency 自我一致提示是在 CoT 基础上进一步优化,通过采样多条推理路径,找出最一致的答案。它适用于对结果准确性要求更高的场景,避免一次性推理路径的偶然性导致错误。】

b.纯强化学习(Pure Reinforcement Learning, RL): 通过强化学习训练模型,使其在没有监督数据的情况下,通过试错学习复杂任务; 【deepseek-R1-zero】

c.监督微调结合强化学习(SFT + RL): 首先对模型进行监督微调,然后使用强化学习进行进一步优化,以提高模型的推理能力。【deepseek-R1】

d.纯监督微调和蒸馏(Pure Supervised Fine-Tuning and Distillation)仅使用监督学习和模型蒸馏技术来增强模型的推理能力。【deepseek-R1-distill蒸馏模型

一个完整的LLM训练过程包含以下几步:

  • Model Initialization:加载模型和处理器
  • 数据准备:解析数据集并设置其格式
  • 模型推理:将数据输入到模型中并获取输出
  • 梯度更新:根据损失函数更新模型参数

对齐(alignment)其作用就是让 LLM 与人类的价值观保持一致。在对齐 LLM 方面,基于人类反馈的强化学习(RLHF)是一种突破性的技术。该方法催生了 GPT-4、Claude 和 Gemini 等强大模型。RLHF 之后,人们也探索了多种多样的对齐 LLM 的方法。但是,此前还没有人全面总结对齐 LLM 与人类偏好的方法。

Salesforce 决定填补这一空白,于近日发布了一份 37 页的综述报告,其中按类别总结了现有的研究文献,并详细分析了各篇论文。

Introduction

这篇论文分为四大主题:奖励模型、反馈、强化学习(RL)、优化。每个主题又包含进一步的子主题,如图 1 所示。

xPO LLM 与人类偏好保持一致的 13 个分类方向

奖励模型的子主题包括:1. 显式奖励模型与隐式奖励模型;2. 逐点奖励模型与偏好模型;3. 响应层面的奖励与 token 层面的奖励;4. 负偏好优化。

反馈的子主题包括:1. 偏好反馈与二元反馈;2. 成对反馈与列表反馈;3. 人类反馈与 AI 反馈。

强化学习的子主题包括:1. 基于参考的强化学习与无参考的强化学习;2. 长度控制式强化学习;3. 强化学习中的不同分支;4. 在线策略强化学习与离线策略强化学习

优化的子主题包括:1. 在线 / 迭代式偏好优化与离线 / 非迭代式偏好优化;2. 分离 SFT 和对齐与合并 SFT 和对齐。

Individual Paper Reviews in Detail

1. RLHF/PPO

LLM 的预训练要用到大量来自不同来源的语料库,而这本身就无法确保这些数据集的质量。此外,LLM 的主要目标是预测下一个 token,这个目标与「有用且安全地遵从用户指令」的目标并不一致。因此,LLM 可能会输出不真实、有害或对用户无用的内容。本质上讲,这些模型并未与用户意图对齐。RLHF/PPO 的主要目标是在各种任务上对齐语言模型与用户意图,其做法是使用人类反馈来微调模型。有关这个主题的研究有很多。

2. RLAIF

获取人类偏好数据集的成本不低,因此基于人工智能反馈的强化学习(RLAIF)诞生了。此外,随着 LLM 的能力不断进步,所能收集到的 AI 偏好数据集的质量也不断提高,由此可提升 LLM 的对齐效果。

3.直接人类偏好优化

传统 RLHF 方法通常涉及到优化源自人类偏好的奖励函数。该方法虽有效,但也可能带来一些难题,比如增大计算复杂度以及在估计和优化奖励时需要考虑偏置 – 方差权衡。参阅论文《High-dimensional continuous control using generalized advantage estimation》。

近期有研究探索了其它一些旨在根据人类偏好(无需依赖某个标量的奖励信号)来直接优化 LLM 策略的方法。

这些方法的目标是通过更直接地使用偏好数据来简化对齐流程、降低计算开销以及实现更稳健的优化。通过将该问题描述为一个偏好优化问题,而不是奖励估计和最大化问题,这些方法能提供一种将语言模型与人类判断对齐的不同视角

4.token 级 DPO

使用 DPO 时,奖励会被一起分配给 prompt 和响应。相反,使用 MDP 时,奖励会被分配给各个动作。后续的两篇论文在 token 层面阐述了 DPO 并将其应用扩展到了 token 级的分析。

5.迭代式 / 在线 DPO

使用 DPO 时,会使用所有可用的偏好数据集来对齐 LLM。为了持续提升 LLM,应当实现迭代式 / 在线 DPO。这就引出了一个有趣的问题:如何高效地收集新的偏好数据集。下面两篇论文深入探讨了这一主题。

  • 自我奖励式语言模型,参阅论文《Self-rewarding language models》。
  • CRINGE,参阅论文《The cringe loss: Learning what language not to model》。

6.二元反馈

事实证明,收集偏好反馈比收集二元反馈(比如点赞或点踩)的难度大,因此后者可促进对齐过程的扩展。KTO 和 DRO 这两项研究关注的便是使用二元反馈来对齐 LLM。

  • KTO,Kahneman-Tversky 优化,参阅论文《KTO: Model alignment as prospect theoretic optimization》。
  • DRO,直接奖励优化,参阅论文《Offline regularised reinforcement learning for large language models alignment》。

7.融合 SFT 和对齐

之前的研究主要还是按顺序执行 SFT 和对齐,但事实证明这种方法很费力,并会导致灾难性遗忘。后续的研究有两个方向:一是将这两个过程整合成单一步骤;二是并行地微调两个模型,最终再进行融合。

  • ORPO,比值比偏好优化,参阅论文《ORPO: Monolithic preference optimization without reference model》。
  • PAFT,并行微调,参阅论文《PAFT: A parallel training paradigm for effective llm fine-tuning》。

8.长度控制式 DPO 和无参考 DPO

之前有研究表明,LLM 的输出往往过于冗长。为了解决这个问题,R-DPO 和 SimPO 的关注重心是在不影响生成性能的前提下实现对响应长度的控制。

此外,DPO 必需参考策略来确保已对齐模型不会与参考模型有太大偏差。相较之下,SimPO 和 RLOO 提出了一些方法,可以在不影响 LLM 效果的情况下消除对参考模型的需求

9.逐列表的偏好优化

之前在 PPO 和 DPO 方面的研究关注的是成对偏好,而 RLHF 方面的研究则是收集逐列表的偏好来加速数据收集过程,之后再将它们转换成成对偏好。尽管如此,为了提升 LLM 的性能,直接使用逐列表的数据集来执行偏好优化是可行的。以下三篇论文专门讨论了这种方法。

  • LiPO,逐列表偏好优化,参阅论文《LIPO: Listwise preference optimization through learning-to-rank》。
  • RRHF,参阅论文《RRHF: Rank responses to align language models with human feedback without tears》。
  • PRO,偏好排名优化,参阅论文《Preference ranking optimization for human alignment》。

10.负偏好优化

这些研究有一个共同前提:当前这一代 LLM 已经在翻译和总结等任务上超越了人类性能。因此,可以将 LLM 的输出视为期望响应,而无需依靠将人类标注的数据视为偏好响应;这样做是有好处的。反过来,不期望得到的响应依然也可被用于对齐 LLM,这个过程就是所谓的负偏好优化(NPO)。

  • NN,否定负例方法,参阅论文《Negating negatives: Alignment without human positive samples via distributional dispreference optimization》。
  • NPO,负例偏好优化,参阅论文《Negative preference optimization: From catastrophic collapse to effective unlearning》。
  • CPO,对比偏好优化,参阅论文《Contrastive preference optimization: Pushing the boundaries of llm performance in machine translation》。

11.纳什学习

之前的研究通常是使用逐点奖励和 BT 模型来得到成对偏好。但是,这种方法比不上直接成对偏好建模并且无法解决成对偏好中的不一致问题。为了克服这些局限,一些研究提出了纳什学习方法。

  • 根据人类反馈的纳什学习,参阅论文《Nash learning from human feedback》。
  • SPPO,自博弈偏好优化,参阅论文《A minimaximalist approach to reinforcement learning from human feedback》。
  • DNO,直接纳什优化,参阅论文《Direct nash optimization: Teaching language models to self-improve with general preferences》。

LLM 对齐(Alignment)方法:SFT、PPO、DPO 、ORPOD、GRPO等方法详细介绍

LLM(大语言模型)的对齐(Alignment)方法旨在让 AI 的输出更加符合人类预期,减少错误信息、有害内容或不准确的回答。主要总结LLM训练中的基本的对齐算法, 监督微调 (SFT)、直接偏好优化 (DPO) 和近端策略优化 (PPO)等。

SFT(Supervised Fine-Tuning,监督微调)

监督微调(SFT)是 LLM 训练中的第一步,通过高质量的人工标注数据集对模型进行微调,使其具备基础的任务能力。SFT 是所有对齐方法的基础,如 RLHF、DPO 等都依赖于一个经过 SFT 训练的模型作为初始状态。

过程

  1. 数据准备:收集高质量的指令-响应(Instruction-Response)数据集,例如人类标注的数据或合成的数据。
  2. 模型微调:使用交叉熵损失(Cross-Entropy Loss)训练模型,使其学习提供与标注数据匹配的答案。
  3. 效果:使模型在常见任务(如问答、代码生成、对话等)中表现更好,提高其对指令的遵循能力。

给定输入 x(Prompt) 和目标输出 y(Response),模型的目标是最大化生成目标文本的概率:

其中:

  • Pθ​(yt​∣x,y<t​) 是 LLM 在给定上下文下预测下一个 token yt​ 的概率。
  • 训练时采用交叉熵损失(Cross Entropy Loss)来优化模型参数 θ。

SFT 仅依赖于人工标注数据,无法让模型学习偏好信息(比如不同回答的优劣)。无法动态调整:SFT 训练后,模型固定,难以针对用户反馈进行调整。缺乏探索性:模型只会学到训练数据中的模式,无法进行强化学习优化。

DPO(Direct Preference Optimization,直接偏好优化)

论文:https://arxiv.org/abs/2305.18290

参考代码:https://github.com/eric-mitchell/direct-preference-optimization

DPO(直接偏好优化)是一种比 RLHF 更简单、更高效的对齐方法。
它不需要训练奖励模型(RM)和使用强化学习(RL),而是直接优化 LLM,使其更符合人类偏好数据

偏好数据

  • 每个输入 Prompt 对应两个候选回答:一个优选(Preferred y+),一个劣选(Dispreferred y−)。
  • 例如:

Prompt: “如何写一封正式的电子邮件?”
Response 1 (优选): “在邮件中应保持正式语气,并包含称呼、正文和署名。”
Response 2 (劣选): “随便写就行了,不要太在意格式。”

优化 LLM 使其更倾向于优选回答

只需要加载2个相同的模型,其中一个推理[reference model:old策略模型],另外一个模型[policy model 策略模型]训练,直接在偏好数据上进行训练即可:

Reference Model(以下简称Ref模型)一般也用SFT阶段得到的SFT模型做初始化,在训练过程中,它的参数是冻结的。Ref模型的主要作用是防止Actor”训歪”

损失函数: DPO 直接优化模型输出的偏好分布:

其中:

  • σ :sigmoid函数
  • β :超参数,一般在0.1 – 0.5之间
  • yw :某条偏好数据中好的response,w就是win的意思
  • yl :某条偏好数据中差的response,l就是loss的意思,所以偏好数据也叫comparision data
  • πθ(yw|x) :给定输入x, 当前 策略policy model生成好的response的累积概率(每个tokne的概率求和,具体看代码)
  • πref(yl|x) :给定输入x, 原始模型(reference model)生成坏的response的累积概率

开始训练时,reference model和policy model都是同一个模型,只不过在训练过程中reference model不会更新权重。

为了方便分析,我们把log里的分式展开,然后β设为1,并且暂时不看前面的log_sigmoid,那么上面的loss可以简化为:

由于最初loss前面是有个负号的,所以优化目标是让本简化公式最大,即我们希望左半部分和右半部分的margin越大越好,左半部分的含义是good response相较于没训练之前的累积概率差值,右半部分代表bad response相较于没训练之前的累计概率差值,如果这个差值,即margin变大了,就意味着:

  • 1)左边变大,右边变小,理想情况,good response概率提升,bad response概率下降
  • 2)左边变小,右边更小,good response概率下降,但是bad response概率下降的更多,生成的时候还是倾向于good response
  • 3)左边变的更大,右边只大了一点点,和2)同理

所以这个loss颇有一种对比的感觉。

OPA-DPO:多模态大模型幻觉难题的高效解决方案

在视觉多模态大语言模型领域,生成与输入图像不一致甚至还有虚假内容的“幻觉”现象,是一个亟待攻克的核心难题。作为一种简单有效的解决方案,直接偏好优化 (DPO) [1] 正在引起越来越多的关注。研究者们通过比较模型在相同提示词和图像下的不同响应,根据幻觉程度直接构造偏好数据对,用于 DPO 训练。

然而,微软亚洲研究院的研究员们注意到,现有研究中不同的数据构造方法会导致显著的性能差异。因此,他们对“基于 DPO 解决多模态大模型幻觉问题”的算法进行了全面分析,总结了它们的表现及局限性,同时从理论角度揭示了各算法性能差异背后的根本原因,并指出决定模型性能的最关键因素是“用于构建偏好对的数据,相较于 DPO 开始前的策略(reference policy)是否为同策略(on-policy)”。

研究员们将此前的研究工作分为三类:

第一类是幻觉注入类,如 HALVA [2]和 POVID [3],通过在已有图像和提示的标准响应中人为注入幻觉片段来构建偏好对;

第二类是幻觉识别类,如 RLHF-V [4]、HA-DPO [5]和 HSA-DPO [6],先让模型根据图像和提示自行生成响应,然后利用专家反馈(人类或 GPT-4/4v)来识别和修改其中的幻觉,从而构建偏好对;

第三类是自我进化类,如 RLAIF-V[7],让模型针对同一图像和提示生成多个响应,并由一个在幻觉识别方面能力更强的导师模型对这些响应中的幻觉严重程度进行判断和排序,以此构建偏好对。

根据实验结果,这三类算法的性能总结为:自我进化类 > 幻觉识别类 > 幻觉注入类。

对于幻觉注入类,幻觉通常并不来自模型本身,因此通过 DPO 训练往往不能给模型带来很大增益。对于自我进化类,理论上由于维度灾难问题,让模型自行探索并找到完全正确的回复是十分困难的,所以那些存在于多个回复中的顽固幻觉通常无法通过这种方法消除。

直觉上,幻觉识别类的方法应该是最高效的解决幻觉的方案,那为什么在实践中这类方法却败下阵来?为了了解背后的原因,研究员们从 DPO 算法的细节入手进行研究。

与最常用的 RLHF 算法 PPO 的初始目标相同,DPO 的初始目标也是(π_θ 是模型的当前策略,π_ref 是模型的初始策略/参考策略,x 为提示词,m 为图像,y 为响应,r(x,y,m) 是通过 Bradley-Terry model 训练得到的奖励函数):

formular

即在最大化奖励的同时,约束模型当前策略与模型初始策略之间的 KL 散度。然而,研究员们重新审视 KL 散度的定义发现,给定任何一个提示词和图像 (x,m),若存在一个响应 (y) 使得 π_θ(y|x,m)>0,但 π_ref(y|x,m)→0,此时 KL 散度会趋于无穷大。这个性质说明——对于任何从目标函数 (1) 出发的算法,那些相对原始策略 (π_ref) 采样概率极低的响应(根据强化学习的命名规范,这种数据被称为异策略(off-policy)数据,相反则为同策略(on-policy)数据)将没有任何机会被模型学会。

formular

如果非要将这些异策略(off-policy)的优选响应(preferred response)拿来构建 DPO 偏好对,会导致梯度在下一次更新时几乎消失。

重温 DPO 训练的优化目标:

text

其中 y_w 是优选响应(preferred response),y_l 是被拒响应(rejected response),其梯度可表示为(σ(⋅) 是 sigmoid 函数):

diagram

训练开始前 πθ=π_ref,所以 sigmoid 函数内部的值应当为0,即当前策略会以 0.5β 为系数对 y_w 进行最大对数似然更新(max-loglikelihood update)。但是在这一步更新过后,logπ_ref(y_w∣x,m)πθ(y_w∣x,m) 将会趋近于极大值(因为分子 > 0,而分母趋近于0),从而导致 σ(−r_w+r_l)→0。因此,梯度会在下一次更新时几乎消失。

回顾幻觉识别类的方法,专家改动后的响应,大部分对于原模型来说都是异策略(off-policy)的,即使这些改动再微小也无济于事,所以根本无法指望这些专家反馈能被模型学会。相对应地,自我进化类方法即使存在学习效率不高的潜在问题,但是它构建的偏好对都来自模型本身,即全是同策略(on-policy)的,因此效果最好。

OPA-DPO:打破常规,重塑对齐策略

是否存在一种方法既能够利用专家的精确反馈,又能完全避免异策略(off-policy)导致的 KL 散度约束问题?

针对现有方法的局限性,微软亚洲研究院联合香港中文大学提出了一种简单而高效的算法 On-Policy Alignment(OPA)-DPO,将专家的精确反馈数据在 DPO 训练前与模型策略对齐。在仅使用4.8k数据的情况下,OPA-DPO 可以实现目前 SOTA 的性能,而之前的 SOTA 算法需要16k数据。该成果已获选计算机视觉领域顶会 CVPR 2025 的 Oral 论文。

Mitigating Hallucinations in Large Vision-Language Models via DPO: On-Policy Data Hold the Key

OPA-DPO 的具体实现方法如下:首先,给定图像和提示,让模型自行生成对应的响应;接着,利用专家反馈(如 GPT-4v)对生成内容进行细粒度修改,保留正确的响应部分,同时纠正其中存在的幻觉内容;然后,将数据集中的真实响应与专家修改后的响应进行 LoRA-SFT 微调,得到一个新的模型(研究员们将其称为 OPA 模型);最后,在 OPA 模型的基础上,进行后续的 DPO 训练,其中研究员们参考了 mDPO 的设定,在构建语言偏好对的同时也构建了图像偏好对以及锚点对,尽管这些元素都很重要,但对最终结果影响最大的还是 OPA 操作。

研究员们综合比较了基于 LLaVA-1.5-7B 和 13B 模型微调的各种 DPO-based 的算法,OPA-DPO 在使用 4.8k 数据的情况下可在多个指标上实现 SOTA 效果。

研究员们发现,使用 OPA-DPO 训练过的模型会呈现出一种略显保守的策略,尤其是在描述任务中,它通常只输出显著并且确定的观测,而忽略一些不重要的细节。

不仅如此,研究员们还观测到一个有趣的现象:基座模型往往默认 query 中的语言是准确无误的,即使这部分文字存在严重幻觉,模型也会顺着其描述图片,这或许可以理解为一种文字惯性现象。而通过 OPA-DPO 训练的模型则展现出了甄别 query 文字部分幻觉的能力。

PPO(Proximal Policy Optimization,近端策略优化)

论文:https://arxiv.org/abs/1707.06347

人人都能看懂的RL-PPO理论知识

是OpenAI在2017提出的一种强化学习算法,是基于策略优化的算法,用于训练能够最大化累积奖励的智能体。PPO算法通过在每次更新时限制新策略与旧策略之间的差异,从而更稳定地更新策略参数。这种方法有助于避免训练过程中出现的不稳定性和剧烈波动,使得算法更容易收敛并学习到更好的策略。

强化学习基本概念

  • 强化学习的两个实体:智能体(Agent)环境(Environment)
  • 强化学习中两个实体的交互:
    • 状态空间S:S即为State,指环境中所有可能状态的集合
    • 动作空间A:A即为Action,指智能体所有可能动作的集合
    • 奖励R:R即为Reward,指智能体在环境的某一状态下所获得的奖励。

以上图为例,智能体与环境的交互过程如下:

  • 在 t 时刻,环境的状态为 St ,达到这一状态所获得的奖励为 Rt
  • 智能体观测到 St 与 Rt ,采取相应动作 At
  • 智能体采取 At 后,环境状态变为 St+1 ,得到相应的奖励 Rt+1

奖励值 Rt ,它表示环境进入状态 St 下的即时奖励
但如果只考虑即时奖励,目光似乎太短浅了:当下的状态和动作会影响到未来的状态和动作,进而影响到未来的整体收益。
所以,一种更好的设计方式是:t时刻状态s的总收益 = 身处状态s能带来的即时收益 + 从状态s出发后能带来的未来收益写成表达式就是:Vt=Rt+γVt+1

其中:

  • Vt : t 时刻的总收益,注意这个收益蕴涵了“即时”和“未来”的概念
  • Rt : t 时刻的即时收益
  • Vt+1 : t+1 时刻的总收益,注意这个收益蕴涵了“即时”和“未来”的概念。而 Vt+1 对 Vt 来说就是“未来”。
  • γ :折扣因子。它决定了我们在多大程度上考虑将“未来收益”纳入“当下收益”。

关键概念:

策略函数是一个 概率密度函数(PDF),输入时当前状态s,输出为一个概率分布,表征每个 action 的概率:

动作值函数:评价在状态 st 的情况下做出动作 at的好坏程度。

状态值函数:

  • 消掉了动作 A ,这样 Vπ 只跟状态 s 与策略函数 π 有关了。
  • 给定 π,可以评价当前状态的好坏;给定状态st,可以评价策略 π的好坏。

优势函数:有些时候我们不需要描述一个行动的绝对好坏,而只需要知道它相对于平均水平的优势。也就是说,我们只想知道一个行动的相对 优势 。这就是优势函数的概念。

一个服从策略 \pi 的优势函数,描述的是它在状态 s 下采取行为 a 比随机选择一个行为好多少(假设之后一直服从策略 \pi )。数学角度上,优势函数的定义为:

长期价值可以表示为状态值函数(State Value Function)或动作值函数(Action Value Function)。

优化方法:

value-based:优化状态动作值函数Q(s) ,贪心选择(确定性策略) :Q-Learning

policy-based :直接优化策略函数π(s, a),按概率采样(随机性策略) :REINFORCE

Actor-Critic •融合上述方法,同时优化Q与π:TRPO、PPO

NLP中的强化学习

在第一部分介绍了通用强化学习的流程,那么我们要怎么把这个流程对应到NLP任务中呢?换句话说,NLP任务中的智能体、环境、状态、动作等等,都是指什么呢?

回想一下我们对NLP任务做强化学习(RLHF)的目的:我们希望给模型一个prompt,让模型能生成符合人类喜好的response。再回想一下gpt模型做推理的过程:每个时刻 t 只产生一个token,即token是一个一个蹦出来的,先有上一个token,再有下一个token。


复习了这两点,现在我们可以更好解读上面这张图了:

  • 我们先喂给模型一个prompt,期望它能产出符合人类喜好的response
  • 在 t 时刻,模型根据上文,产出一个token,这个token即对应着强化学习中的动作,我们记为At 。因此不难理解,在NLP语境下,强化学习任务的动作空间就对应着词表。
  • 在 t 时刻,模型产出token At对应着的即时收益为Rt,总收益为Vt复习一下, Vt 蕴含着“即时收益”与“未来收益”两个内容)。这个收益即可以理解为“对人类喜好的衡量”。此刻,模型的状态从St变为St+1,也就是从“上文”变成“上文 + 新产出的token”
  • 在NLP语境下,智能体是语言模型本身,环境则对应着它产出的语料

这样,我们就大致解释了NLP语境下的强化学习框架,不过针对上面这张图,你可能还有以下问题:


(1)问题1:图中的下标是不是写得不太对?例如根据第一部分的介绍, At 应该对应着 Rt+1  At+1 应该对应着 Rt+2 ,以此类推?
答:你说的对。但这里我们不用太纠结下标的问题,只需要记住在对应的response token位置,会产生相应的即时奖励和总收益即可。之所以用图中这样的下标,是更方便我们后续理解代码。


(2)问题2:我知道 At 肯定是由语言模型产生的,那么 ,Rt,Vt 是怎么来的呢,也是语言模型产生的吗?
答:先直接说结论, At 是由我们的语言模型产生的, ,Rt,Vt 则分别由另外两个模型来产生,在后文中我们会细说。


(3)问题3:语言模型的参数在什么时候更新?是观测到一个 Rt,Vt ,就更新一次参数,然后再去产生 At+1 吗?
答:当然不是。你只看到某个时刻的收益,就急着用它更新模型,这也太莽撞了。我们肯定是要等有足够的观测数据了(例如等模型把完整的response生成完),再去更新它的参数。这一点我们也放在后文细说。


(4)问题4:再谈谈 Rt,Vt 吧,在NLP的语境下我还是不太理解它们
答:

  • 首先,“收益”的含义是“对人类喜好的衡量”
  • Rt :即时收益,指语言模型当下产生token At 带来的收益
  • Vt : 实际期望总收益(即时+未来),指对语言模型“当下产生token At ,一直到整个response生产结束”后的期收益预估。因为当下语言模型还没产出 At 后的token,所以我们只是对它之后一系列动作的收益做了估计,因而称为“期望总收益”。

RLHF中的四个重要角色:

我们从第二部分中已经知道:生成token At 和对应收益 Rt,Vt 的并不是一个模型。那么在RLHF中到底有几个模型?他们是怎么配合做训练的?而我们最终要的是哪个模型?

如上图,在RLHF-PPO阶段,一共有四个主要模型,分别是:

  • Actor Model:演员模型,这就是我们想要训练的目标语言模型
  • Critic Model:评论家模型,它的作用是预估总收益 Vt
  • Reward Model:奖励模型,它的作用是计算即时收益 Rt
  • Reference Model:参考模型,它的作用是在RLHF阶段给语言模型增加一些“约束”,防止语言模型训歪(朝不受控制的方向更新,效果可能越来越差)

其中:

  • Actor/Critic Model在RLHF阶段是需要训练的(图中给这两个模型加了粗边,就是表示这个含义);而Reward/Reference Model参数冻结的。
  • Critic/Reward/Reference Model共同组成了一个“奖励-loss”计算体系(我自己命名的,为了方便理解),我们综合它们的结果计算loss,用于更新Actor和Critic Model

Actor Model (演员模型)

正如前文所说,Actor就是我们想要训练的目标语言模型。我们一般用SFT阶段产出的SFT模型来对它做初始化。

我们的最终目的是让Actor模型能产生符合人类喜好的response。所以我们的策略是,先喂给Actor一条prompt (这里假设batch_size = 1,所以是1条prompt),让它生成对应的response。然后,我们再将“prompt + response”送入我们的“奖励-loss”计算体系中去算得最后的loss,用于更新actor。

Reference Model(参考模型)

Reference Model(以下简称Ref模型)一般也用SFT阶段得到的SFT模型做初始化,在训练过程中,它的参数是冻结的。Ref模型的主要作用是防止Actor”训歪”,那么它具体是怎么做到这一点的呢?

“防止模型训歪”换一个更详细的解释是:我们希望训练出来的Actor模型既能达到符合人类喜好的目的,又尽量让它和SFT模型不要差异太大。简言之,我们希望两个模型的输出分布尽量相似。那什么指标能用来衡量输出分布的相似度呢?我们自然而然想到了KL散度

如图所示:

  • 对Actor模型,我们喂给它一个prompt,它正常输出对应的response。那么response中每一个token肯定有它对应的log_prob结果呀,我们把这样的结果记为log_probs
  • 对Ref模型,我们把Actor生成的”prompt + response”喂给它,那么它同样能给出每个token的log_prob结果,我们记其为ref_log_probs
  • 那么这两个模型的输出分布相似度就可以用ref_log_probs - log_probs来衡量,我们可以从两个方面来理解这个公式:
    • 从直觉上理解,两个分布的相似度越高,说明Ref模型对Actor模型输出的肯定性越大。即Ref模型也认为,对于某个 St ,输出某个 At 的概率也很高( P(At|St) )。这时可以认为Actor模型较Ref模型没有训歪。
    • 从KL散度上理解

(当然这里不是严格的等于,只是KL散度的近似),这个值越小意味着两个分布的相似性越高。

注:你可能已经注意到,按照KL散度的定义,这里写成log_probs - ref_log_probs更合适一些。但是如果你看过一些rlhf相关的论文的话,你可能记得在计算损失函数时,有一项Rt−KL散度 (对这个有疑惑不要紧,我们马上在后文细说),即KL散度前带了负号,所以这里我写成ref_log_probs - log_probs这样的形式,更方便大家从直觉上理解这个公式。

现在,我们已经知道怎么利用Ref模型和KL散度来防止Actor训歪了。KL散度将在后续被用于loss的计算,我们在后文中会详细解释。

Critic Model(评论家模型)

Critic Model用于预测期望总收益 Vt ,和Actor模型一样,它需要做参数更新。实践中,Critic Model的设计和初始化方式也有很多种,例如和Actor共享部分参数、从RW阶段的Reward Model初始化而来等等。我们讲解时,和deepspeed-chat的实现保持一致:从RW阶段的Reward Model初始化而来。


你可能想问:训练Actor模型我能理解,但我还是不明白,为什么要单独训练一个Critic模型用于预测收益呢?
这是因为,当我们在前文讨论总收益 Vt (即时 + 未来)时,我们是站在上帝视角的,也就是这个 Vt 就是客观存在的、真正的总收益。但是我们在训练模型时,就没有这个上帝视角加成了,也就是在 t 时刻,我们给不出客观存在的总收益 Vt ,我们只能训练一个模型去预测它。


所以总结来说,在RLHF中,我们不仅要训练模型生成符合人类喜好的内容的能力(Actor),也要提升模型对人类喜好量化判断的能力(Critic)。这就是Critic模型存在的意义。我们来看看它的大致架构:

deepspeed-chat采用了Reward模型作为它的初始化,所以这里我们也按Reward模型的架构来简单画画它。你可以简单理解成,Reward/Critic模型和Actor模型的架构是很相似的(毕竟输入都一样),同时,它在最后一层增加了一个Value Head层,该层是个简单的线形层,用于将原始输出结果映射成单一的 Vt 值。

在图中, Vt 表示Critic模型对 t 时刻及未来(response完成)的收益预估。

 Reward Model(奖励模型)

Reward Model用于计算生成token At 的即时收益,它就是RW阶段所训练的奖励模型,在RLHF过程中,它的参数是冻结的。


你可能想问:为什么Critic模型要参与训练,而同样是和收益相关的Reward模型的参数就可以冻结呢?
这是因为,Reward模型是站在上帝视角的。这个上帝视角有两层含义:

  • 第一点,Reward模型是经过和“估算收益”相关的训练的,因此在RLHF阶段它可以直接被当作一个能产生客观值的模型。
  • 第二点,Reward模型代表的含义就是“即时收益”,你的token At 已经产生,因此即时收益自然可以立刻算出。

你还可能想问:我已经用Critic预测出 Vt 了,而这个 Vt 包含了“即时”和“未来”的概念,那我还需要代表“即时”的 Rt 做什么呢?直接用 Vt 不就好了吗?


为了解答这个问题,我们先回顾下1.2部分中给出的价值函数: Vt=Rt+γVt+1
这个函数告诉我们,我们当前可以用两个结果来表示 t 时刻的总收益:

  • 结果1:Critic模型预测的 Vt
  • 结果2:Reward模型预测的 Rt 和critic模型预测的 Vt+1

那么哪一个结果更靠近上帝视角给出的客观值呢?当然是结果2,因为结果1全靠预测,而结果2中的 Rt 是事实数据。
我们知道Critic模型也是参与参数更新的,我们可以用MSE(上帝视角的客观收益-Critic模型预测的收益)来衡量它的loss。但是上帝视角的客观收益我们是不知道的,只能用已知事实数据去逼近它,所以我们就用 Rt+γ∗Vt+1 来做近似。这就是 Rt,Vt 同时存在的意义

Reward模型和critic模型非常相似,这里我们就只给出架构图,不再做过多的说明。

RLHF中的loss计算

到目前为止,我们已经基本了解了RLHF的训练框架,以及其中的四个重要角色(训练一个RLHF,有4个模型在硬件上跑,可想而知对存储的压力)。在本节中,我们一起来解读RLHF的loss计算方式。在解读中,我们会再一次理一遍RLHF的整体训练过程,填补相关细节。在这之后,我们就可以来看代码解析了。


在第三部分的讲解中,我们知道Actor和Critic模型都会做参数更新,所以我们的loss也分成2个:

  • Actor loss:用于评估Actor是否产生了符合人类喜好的结果,将作用于Actor的BWD上。
  • Critic loss:用于评估Critic是否正确预测了人类的喜好,将作用于Critic的BWD上。

我们详细来看这两者。

Actor loss

(1)直观设计

我们先来看一个直观的loss设计方式:

  • Actor接收到当前上文 St ,产出token At ( P(At|St) )
  • Critic根据 St,At ,产出对总收益的预测 Vt
  • 那么Actor loss可以设计为: 

求和符号表示我们只考虑response部分所有token的loss,为了表达简便,我们先把这个求和符号略去(下文也是同理),也就是说:

我们希望minimize这个actor_loss。


这个设计的直观解释是:

  • 当 Vt>0 时,意味着Critic对Actor当前采取的动作给了正向反馈,因此我们就需要在训练迭代中提高 P(At|St) ,这样就能达到减小loss的作用。
  • 当 Vt<0 时,意味着Critic对Actor当前采取的动作给了负向反馈,因此我们就需要在训练迭代中降低 P(At|St) ,这样就能到达到减小loss的作用。

一句话总结:这个loss设计的含义是,对上文 St 而言,如果token At 产生的收益较高,那就增大它出现的概率,否则降低它出现的概率。

(2)引入优势(Advantage)

在开始讲解之前,我们举个小例子:
假设在王者中,中路想支援发育路,这时中路有两种选择:1. 走自家野区。2. 走大龙路。
中路选择走大龙路,当她做出这个决定后,Critic告诉她可以收1个人头。结果,此刻对面打野正在自家采灵芝,对面也没有什么苟草英雄,中路一路直上,最终收割2个人头。
因为实际收割的人头比预期要多1个,中路尝到了甜头,所以她增大了“支援发育路走大龙路”的概率。
这个多出来的“甜头”,就叫做“优势”(Advantage)。


对NLP任务来说,如果Critic对 At 的总收益预测为 Vt ,但实际执行 At 后的总收益是 Rt+γ∗Vt+1 ,我们就定义优势为:

我们用 Advt 替换掉 Vt ,则此刻actor_loss变为:

(3)重新设计 Rt

总结一下,到目前为止,我们的actor_loss形式为:

同时注意,这个actor_loss应该是response的所有token loss的sum或者avg。这里为了表达方便,我们的公式略去了求和或求平均的符号。

按照这个理解, Rt 应该表示每个Actor产出token At 带来的即时收益,正如下图所示(其中 T 表示最后一个时刻):

但在deepspeed-chat的RLHF实践中,对 Rt 做了另一种设计:

基于这些,上面这个对 Rt 的设计可理解成:

  • t≠T时,我们更加关心Actor是否有在Ref的约束下生产token At
  • $ t=T时,我们不仅关心Actor是否遵从了Ref的约束,也关心真正的即时收益Rt


为什么只有最后一个时刻的 Rt 被纳入了考量呢?这是因为在Reward模型训练阶段,就是用这个位置的 Rt 来表示对完整的prompt + response的奖励预测(但不妨碍你理解成是执行完 AT 的即时奖励),然后用这个指标来做模型eval的(但是Reward训练阶段算loss时,还是考虑了response部分所有token输出的reward值)。所以到了RLHF的场景下,其余时刻的即时奖励,我们就用“Actor是否遵循了Ref的约束”来进行评价。


需要注意的是, Rt 的设计并不只有这一种。deepspeed在自己的代码注释中也有提过,可以尝试把最后一个时刻的 RT 替换成所有token的即时奖励的平均值。如果站在这个角度理解的话,我们同样也可以尝试在每一个位置的奖励衡量上引入 Rt 。

代码实践如下:

def compute_rewards(self, prompts, log_probs, ref_log_probs, reward_score,
                        action_mask):
        """
        reward_function:计算最终的reward分数
        复习一下几个相关参数的默认值:
        self.kl_ctl = 0.1
        self.clip_reward_value = 5
        
        对于batch中的某个prompt来说,它最终的reward分数为:
        (1) 先计算actor和ref_model的logit相似度: -self.kl_ctl * (log_probs - ref_log_probs)
            其实写成self.kl_ctl * (ref_log_probs - log_probs)更好理解些
            这个值越大,说明ref_model对actor生成的结果的认可度越高(即表明rlhf没有训歪),
            没有训歪的情况下我们也应该给模型一些奖励,这个奖励就是self.kl_ctl * (ref_log_probs - log_probs)
            
        (2)由于我们只取最后一个token对应位置的分数作为reward_score,因此我们只需要:
            self.kl_ctl * (ref_log_probs - log_probs)的最后一位 + reward_score
         
         (3) 同时我们对reward_score也做了大小限制,最大不超过self.clip_reward_value(超过统一给成self.clip_reward_value),
             最小不低于-self.clip_reward_value(低于统一给成-self.clip_reward_value)
        
         (4) 最后返回的rewards大小为:(batch_size, 各条数据的长度),对batch中的每条数据来说:
             - response的最后一位:self.kl_ctl * (ref_log_probs - log_probs)的最后一位 + reward_score
             - response的其余位置:self.kl_ctl * (ref_log_probs - log_probs)
        
        """

        kl_divergence_estimate = -self.kl_ctl * (log_probs - ref_log_probs)
        rewards = kl_divergence_estimate
        # ---------------------------------------------------------------------------------------------------
        # response开始的位置
        # (因为我们对prompt做过padding处理,因此batch中每个prompt长度一致,也就意味着每个response开始的位置一致)
        # (所以这里start是不加s的,只是一个int)
        # ---------------------------------------------------------------------------------------------------
        start = prompts.shape[1] - 1
        # ---------------------------------------------------------------------------------------------------
        # response结束的位置
        # (因为一个batch中,每个response的长度不一样,所以response的结束位置也不一样)
        # (所以这里end是加s的,ends的尺寸是(batch_size,)
        # ---------------------------------------------------------------------------------------------------
        ends = start + action_mask[:, start:].sum(1) + 1
        # ---------------------------------------------------------------------------------------------------
        # 对rewards_score做限制
        # ---------------------------------------------------------------------------------------------------
        reward_clip = torch.clamp(reward_score, -self.clip_reward_value,
                                  self.clip_reward_value)
        batch_size = log_probs.shape[0]
        for j in range(batch_size):
            rewards[j, start:ends[j]][-1] += reward_clip[j] # 

        return rewards

注意输入输出的维度,prompts 是一个 [batch size, padded prompt length] 的 matrix,ref_log_probs 和 log_probs 是 [batch size, padded prompt with response length] 大小的矩阵,然后只有从 prompt 结束到 response 结束这一块儿的 reward 才会实际有作用,prompt 的 reward 是不计算的。

prompt 有统一的 padding,所以 response 的 start 位置是唯一的,而 ends 则通过 action_mask 中的 1 元素的截止为止计算得到。最后,在这个 batch 中,每个 prompt 的 reward 的结尾那个 token 加上 reward_score 进过 clip 得到的 reward。

(4)重新设计优势

好,再总结一下,目前为止我们的actor_loss为:

同时,我们对 Rt 进行来改造,使其能够衡量Actor模型是否遵从了Ref模型的约束。
现在我们把改造焦点放在 Advt 上,回想一下,既然对于收益而言,分为即时和未来,那么对于优势而言,是不是也能引入对未来优势的考量呢?这样,我们就可以把 Advt 改写成如下形式:


(熟悉强化学习的朋友应该能一眼看出这是GAE,这里我们不打算做复杂的介绍,一切都站在直觉的角度理解)
其中,新引入的 λ 也是一个常量,可将其理解为权衡因子,直觉上看它控制了在计算当前优势时对未来优势的考量。(从强化学习的角度上,它控制了优势估计的方差和偏差)

看到这里,你可能想问:这个代表未来优势的 Advt+1 ,我要怎么算呢?
注意到,对于最后一个时刻 t ,它的未来收益( VT+1 )和未来优势( AdvT+1 )都是0,也就是 AdvT=RT−VT ,这是可以直接算出来的。而有了 AdvT ,我们不就能从后往前,通过动态规划的方法,把所有时刻的优势都依次算出来了吗?

代码实践如下(其中返回值中的returns表示实际收益,将被用于计算Critic模型的loss,可以参见4.2,其余细节都在代码注释中):注意这个函数一并返回了 returns,也即每个 token 的实际收益,这个收益之后会用于更新 critic model:

 def get_advantages_and_returns(self, values, rewards, start):
        """
        Adopted from https://github.com/CarperAI/trlx/blob/main/trlx/models/modeling_ppo.py#L134
        
        没有引入GAE前的t时刻的优势值:
        detal_t = r_t + gamma * V_t+1 - V_t
        其中:
            - r_t表示t时刻的即时收益
            - V_t+1表示未来时刻的预期收益
            - r_t + gamma * V_t+1可理解成t时刻的实际预期收益
            - V_t可理解成t时刻的预估预期收益(是模型,例如critic model自己估算出来的)
        
        引入GAE后的t时刻的优势值:
        A_t = delta_t + gamma * lambda * A_t+1
        粗暴理解为在t时刻时,不仅考虑当下优势,还考虑了未来的优势
        为了知道A_t, 我们得知道A_t+1,所以在本算法中采取了从后往前做动态规划求解的方法,也即:
        假设T是最后一个时刻,则有A_T+1 = 0, 所以有: A_T = delta_T
        知道了A_T, 就可以依次往前倒推,把A_t-1, A_t-2之类都算出来了
        
        引入GAE后t时刻的实际预期收益
        returns_t = A_t + V_t
                  = delta_t + gamma * lambda * A_t+1 + V_t
                  = r_t + gamma * V_t+1 - V_t + gamma * lambda * A_t+1 + V_t
                  = r_t + gamma * (V_t+1 + lambda * A_t+1)
        
        注意,这里不管是advantages还是returns,都只算response的部分
        """
        
        # Adopted from https://github.com/CarperAI/trlx/blob/main/trlx/models/modeling_ppo.py#L134
        lastgaelam = 0
        advantages_reversed = []
        length = rewards.size()[-1]
        # 注意这里用了reversed,是采取从后往前倒推计算的方式
        for t in reversed(range(start, length)):
            nextvalues = values[:, t + 1] if t < length - 1 else 0.0
            delta = rewards[:, t] + self.gamma * nextvalues - values[:, t]
            lastgaelam = delta + self.gamma * self.lam * lastgaelam
            advantages_reversed.append(lastgaelam)
        advantages = torch.stack(advantages_reversed[::-1], dim=1) # 优势
        returns = advantages + values[:, start:] # 实际收益
        # values: 预期收益
        return advantages.detach(), returns

(5)PPO-epoch: 引入新约束

总结一下,目前为止我们的actor_loss为:

同时

  • 我们已经对Rt进行来改造,使其能够衡量Actor模型是否遵从了Ref模型的约束。
  • 我们已经对Advt进行改造,使其不仅考虑了当前时刻的优势,还考虑了未来的优势

基于这些改造,我们重新理一遍RLHF-PPO的训练过程。

  • 第一步,我们准备一个batch的prompts
  • 第二步,我们将这个batch的prompts喂给Actor模型,让它生成对应的responses
  • 第三步,我们把prompt+responses喂给我们的Critic/Reward/Reference模型,让它生成用于计算actor/critic loss的数据,按照强化学习的术语,我们称这些数据为经验(experiences)。critic loss我们将在后文做详细讲解,目前我们只把目光聚焦到actor loss上
  • 第四步,我们根据这些经验,实际计算出actor/critic loss,然后更新Actor和Critic模型

这些步骤都很符合直觉,但是细心的你肯定发现了,文字描述中的第四步和图例中的第四步有差异:图中说,这一个batch的经验值将被用于n次模型更新,这是什么意思呢?

我们知道,在强化学习中,收集一个batch的经验是非常耗时的。对应到我们RLHF的例子中,收集一次经验,它要等四个模型做完推理才可以,正是因此,一个batch的经验,只用于计算1次loss,更新1次Actor和Critic模型,好像有点太浪费了

所以,我们自然而然想到,1个batch的经验,能不能用来计算ppo-epochs次loss,更新ppo-epochs次Actor和Critic模型?简单写一下伪代码,我们想要:

# --------------------------------------------------------------
# 初始化RLHF中的四个模型
# --------------------------------------------------------------
actor, critic, reward, ref = initialize_models()

# --------------------------------------------------------------
# 训练
# --------------------------------------------------------------
# 对于每一个batch的数据
for i in steps: 
    # 先收集经验值
    exps = generate_experience(prompts, actor, critic, reward, ref)
    # 一个batch的经验值将被用于计算ppo_epochs次loss,更新ppo_epochs次模型
    # 这也意味着,当你计算一次新loss时,你用的是更新后的模型
    for j in ppo_epochs:
        actor_loss = cal_actor_loss(exps, actor)
        critic_loss = cal_critic_loss(exps, critic)
        
        actor.backward(actor_loss)
        actor.step()
        
        critc.backward(critic_loss)
        critic.step()

而如果我们想让一个batch的经验值被重复使用ppo_epochs次,等价于我们想要Actor在这个过程中,模拟和环境交互ppo_epochs次。举个例子:

  • 如果1个batch的经验值只使用1次,那么在本次更新完后,Actor就吃新的batch,正常和环境交互,产出新的经验值
  • 但如果1个batch的经验值被使用ppo_epochs次,在这ppo_epochs中,Actor是不吃任何新数据,不做任何交互的,所以我们只能让Actor“模拟”一下和环境交互的过程,吐出一些新数据出来。

那怎么让Actor模拟呢?很简单,让它观察一下之前的数据长什么样,让它依葫芦画瓢,不就行了吗?我们假设最开始吃batch,吐出经验的actor叫 Actorold ,而在伪代码中,每次做完ppo_epochs而更新的actor叫 Actornew ,那么我们只要尽量保证每次更新后的 Actornew 能模仿最开始的那个 Actorold ,不就行了吗?

诶!是不是很眼熟!两个分布,通过什么方法让它们相近!那当然是KL散度!所以,再回到我们的actor_loss上来,它现在就可被改进成:

我们再稍作一些改动将log去掉(这个其实不是“稍作改动去掉log”的事,是涉及到PPO中重要性采样的相关内容,大家有兴趣可以参考这篇):

其中, Pold 表示真正吃了batch,产出经验值的Actor;P表示ppo_epochs中实时迭代更新的Actor,它在模仿 Pold 的行为。所以这个公式从直觉上也可以理解成:在Actor想通过模拟交互的方式,使用一个batch的经验值更新自己时,它需要收到真正吃到batch的那个时刻的Actor的约束,这样才能在有效利用batch,提升训练速度的基础上,保持训练的稳定。

在 PPO 强化学习中使用 KL 散度,是为了:

让当前策略(Actor)在更新时不要偏离旧策略太远,从而保证经验数据依然有效、训练过程更稳定。

换句话说:

  • P_old 是“真正经历过环境”的
  • P 是“后续更新后模拟交互的”
  • 所以你得让 P 尽量模仿 P_old,才能继续用旧数据去更新模型
  • KL 散度,就是这个“模仿程度”的衡量指标和约束手段

但是,谨慎的你可能此时又有新的担心了:虽然我们在更新Actor的过程中用 Actorold 做了约束,但如果 Actorold 的约束能力不够,比如说 P(At|St)/Pold(At|St) 还是超出了可接受的范围,那怎么办?

很简单,那就剪裁(clip)它吧!

我们给 P(At|St)/Pold(At|St) 设置一个范围,例如(0.8 ,1.2),也就是如果这个值一旦超过1.2,那就统一变成1.2;一旦小于0.8,那就统一变成0.8。这样就能保证 Actor 和 Actorold 的分布相似性在我们的掌控之内了。此时actor_loss变为:

这时要注意,如果超过变化范围,将 P(At|St)/Pold(At|St) 强制设定为一个常数后,就说明这一部分的loss和Actor模型无关了,而 Advt 这项本身也与Actor无关。所以相当于,在超过约束范围时,我们停止对Actor模型进行更新。

整体代码如下:

注意最后整个 response 每一处的 loss 取均值,就是这个 prompt + response 的 actor loss 了

    def actor_loss_fn(self, logprobs, old_logprobs, advantages, mask):
        """
        logprobs: 实时计算的,response部分的prob(只有这个是随着actor实时更新而改变的)
        old_logprobs:老策略中,response部分的prob (这个是固定的,不随actor实时更新而改变)
        advantages: 老策略中,response部分每个token对应的优势(这个是固定的,不随actor实时更新而改变)
        mask:老策略中,response部分对应的mask情况这个是固定的,不随actor实时更新而改变)
        
        之所以要引入logprobs计算actor_loss,是因为我们不希望策略每次更新的幅度太大,防止模型训歪
        
        self.cliprange: 默认值是0.2
        """
        ## policy gradient loss
        # -------------------------------------------------------------------------------------
        # 计算新旧策略间的KL散度
        # -------------------------------------------------------------------------------------
        log_ratio = (logprobs - old_logprobs) * mask
        ratio = torch.exp(log_ratio)
        # -------------------------------------------------------------------------------------
        # 计算原始loss和截断loss
        # -------------------------------------------------------------------------------------
        pg_loss1 = -advantages * ratio
        pg_loss2 = -advantages * torch.clamp(ratio, 1.0 - self.cliprange, 1.0 + self.cliprange)
        pg_loss = torch.sum(torch.max(pg_loss1, pg_loss2) * mask) / mask.sum() # 最后是取每个非mask的response token的平均loss作为最终loss
        return pg_loss

(6)Actor loss小结

(1)~(5)中我们一步步树立了actor_loss的改进过程,这里我们就做一个总结吧:

  • 我们已经对Rt进行来改造,使其能够衡量Actor模型是否遵从了Ref模型的约束
  • 我们已经对Advt进行改造,使其不仅考虑了当前时刻的优势,还考虑了未来的优势
  • 我们重复利用了1个batch的数据,使本来只能被用来做1次模型更新的它现在能被用来做ppo_epochs次模型更新。我们使用真正吃了batch,产出经验值的那个时刻的Actor分布来约束ppo_epochs中更新的Actor分布
  • 我们考虑了剪裁机制(clip),在ppo_epochs次更新中,一旦Actor的更新幅度超过我们的控制范围,则不对它进行参数更新。

Critic loss

我们知道,1个batch产出的经验值,不仅被用来更新Actor,还被用来更新Critic。对于Critic loss,我们不再像Actor loss一样给出一个“演变过程”的解读,我们直接来看它最后的设计。

首先,在之前的解说中,你可能有这样一个印象:

  • Vt :Critic对t时刻的总收益的预估,这个总收益包含即时和未来的概念(预估收益)
  • Rt+γ∗Vt+1 :Reward计算出的即时收益 Rt ,Critic预测出的 t+1 及之后时候的收益的折现,这是比 Vt 更接近t时刻真值总收益的一个值(实际收益)

所以,我们的第一想法是:

现在,我们对“实际收益”和“预估收益”都做一些优化。

(1)实际收益优化

我们原始的实际收益为 Rt+γ∗Vt+1 ,但是当我们在actor_loss中引入“优势”的概念时,“优势”中刻画了更为丰富的实时收益信息,所以,我们将实际收益优化为: Advt+Vt

(2)预估收益优化

我们原始的预估收益为 Vt 。
类比于Actor,Critic模型在ppo_epochs的过程中也是不断更新的。所以这个 Vt 可以理解成是 Criticold ,也就是真正吃了batch,参与产出经验的那个时候的Critic产出的收益预测结果。


我们同样想用旧模型去约束新模型,但对于Critic我们采用的约束策略就比较简单了,我们直接看代码,从中可以看出,我们用老 Vt 设计了了一个变动范围,然后用这个变动范围去约束新 Vt

# self.cliprange_value是一个常量
# old_values: 老critic的预测结果
# values:新critic的预测结果
values_clipped = torch.clamp(
            values,
            old_values - self.cliprange_value,
            old_values + self.cliprange_value,
        )

那么最终我们就取实际收益和预估收益的MSE做为loss就好,这里注意,计算实际收益时 Advt,Vt 都是老Critic(真正吃了batch的那个)产出的结果,而预估收益是随着ppo_epochs而变动的。


代码如下:

def critic_loss_fn(self, values, old_values, returns, mask):
        """
        values: 实时critic跑出来的预估预期收益(是变动的,随着ppo epoch迭代而改变)
        old_values:老critic跑出来的预估预期收益(是固定值)
        returns:实际预期收益
        mask:response部分的mask
        
        self.cliprange_value = 0.2
        """
        ## value loss
        # 用旧的value去约束新的value
        values_clipped = torch.clamp(
            values,
            old_values - self.cliprange_value,
            old_values + self.cliprange_value,
        )
        if self.compute_fp32_loss:
            values = values.float()
            values_clipped = values_clipped.float()
        
        # critic模型的loss定义为(预估预期收益-实际预期收益)**2
        vf_loss1 = (values - returns)**2
        vf_loss2 = (values_clipped - returns)**2
        vf_loss = 0.5 * torch.sum(
            torch.max(vf_loss1, vf_loss2) * mask) / mask.sum() # 同样,最后也是把critic loss平均到每个token上
        return vf_loss

总结:RLHF 的计算流

构造 Reward

给定一个 transformer 和任何一个 string,我都可以将整个 string 输入给 reward model 做一次 forward pass,得到每个位置的 token 的 logit。我们取出最后一个 token 的 logit,经过 logit processor 处理,再过一次 softmax 并取 log,得到此处的 log prob。此外,我们也可以对最后一个 token 的 logit 进行其他操作,譬如 pooling 和 projection 等等,拿到 embedding、reward 或者 value。由此可见,对于 string 里的每个 token,我们都可以得到前述所有计算值,但是在 RLHF 中,我们会用到 response 中每个 token 的 log prob 和 value,但是 reward 模型只会用最后一个 token 的 reward。这里直接给出 reward 的实际计算:

为什么只有最后一个时刻的 Rt被纳入了考量呢?这是因为在Reward模型训练阶段,就是用这个位置的 Rt来表示对完整的prompt + response的奖励预测(但不妨碍你理解成是执行完 AT的即时奖励),然后用这个指标来做模型eval的(但是Reward训练阶段算loss时,还是考虑了response部分所有token输出的reward值)。所以到了RLHF的场景下,其余时刻的即时奖励,我们就用“Actor是否遵循了Ref的约束”来进行评价。

需要注意的是,Rt的设计并不只有这一种。deepspeed在自己的代码注释中也有提过,可以尝试把最后一个时刻的 RT替换成所有token的即时奖励的平均值。如果站在这个角度理解的话,我们同样也可以尝试在每一个位置的奖励衡量上引入 Rt

对于第 t 个 response token,当 t 为最后一个 token T 时,才将 reward model 输出的对整个 response 的 reward 加到 Rt 上。换言之,实际上一个 prompt + response 只会让 reward model 推理一次,作为整个 response 的 reward。

至于其他部分,$kl _ ctl$ 是个常数,$ \log \frac{P(A_t|S_t)}{P_{ref}(A_t|S_t)} $ 是 reference model 和 actor model 生成 At 这个 token 的条件概率比值取对数,也即直接将 actor 的 log prob 和 reference 的 log prob 相减,体现到代码里就是 kl_ctl * (actor_log_probs - ref_log_probs)(KL 散度),这样就得到了每个 token 的 reward。注意这里的单复数,actor_log_probs 和 ref_log_probs 都是所有 response token 的 log prob 构成的 list。

得到 KL 散度后,再在这个 prompt + response 的最后一个 token 上加上此处的 reward(称为 reward score),整个 response 每一处的 reward 便构造完成了。当然,实际上的计算还需要考虑 reward score 的 clip 问题,也即不能让 reward 过大。在知乎里面给了非常好的伪代码。

构造 Advantage

Advanatage 可以某种程度理解为“意外之喜”,具体的描述参考知乎原文。这里直接给出 Advantage 的构造公式:

我们来拆解下,考虑到 Rt 是每个 token 的 reward,前面已经构造了。Vt 和 Vt+1 是当前 token 和下一个 token 的 value,而每个 token 的 value 在 value model 的 forward pass 中都可以得到,Adv_t 是当前 token 的 advantage,$\gamma, \lambda$ 都是常数。这种递归的构造方式,可以用尾递归来反推每个位置的 advantage。

构造 Actor Loss

这里还是直接给出 Actor Loss 的构造公式:

这个构造公式看着复杂,实际上一点也不简单。每个 response token 的 Advt 的构造已经在前文给出,而 P(At|St),Pold(At|St) 其实都是 actor model 的条件概率。之所以有个 old 是因为我们希望多利用每轮产生的 experiences,因此一组 experiences 会更新多轮。old 表示这一组 experiences 用于更新之前的 actor model,用这个 old actor model 对这几轮更新的大小做了约束。最后,考虑到某一轮更新里,当前 actor model 和 old actor model 的差距实在太大了,以至于条件概率的比值超出了人为预设的范围,此时 Advt 的系数(ratio)会取为约束边界。此时 actor model 的参数不再影响 ratio,换言之 actor model 的参数不再在 actor loss 的计算图中了,这个 loss 也就不会更新 actor 的参数了。 注意,advantage 的构造是由 old actor model 构造来的,计算结束就固定了,对于更新中的 actor model 没有梯度,所以整个 actor loss 的计算图中只有 ratio 对更新中的 actor model 有梯度。

构造 Critic Loss

注意到,在 advantage 的构造中,我们一并得到了 returns,将其视为每个 token 的实际收益。

而预估收益就是 Vt,然后我们构造 MSE loss 来最小化预估收益和实际收益的差距。

看上去似乎 Rett−Vt 就是 Advt但是实际使用的 values 是多轮更新中的 value model 的输出,也即 new value,而 returns 是多轮更新开始时就固定了的实际收益(old returns),所以 Rett−Vt 并不是 Advt

更新流程

  1. 准备一个 batch 的 prompts
  2. 将这个 batch 的 prompts 输入给 Actor,解码得到 responses
  3. 将 prompt + responses 输入给 Critic/Reward/Reference,分别计算得得到所有 token 的 values、最后一个 token 的 reward 和所有 token 的 log probs,按照强化学习的术语,称这些数据为经验(experiences)了;
  4. 根据 experiences 多轮计算 actor loss 和 critic loss 并更新 Actor 和 Critic 模型。

对于第 4 步,我们当然可以一轮 experiences 就更新一次 actor 和 critic,但是为了尽可能利用这个 batch 的 experiences,我们对 actor 和 critic 做多轮更新。我们将 experiences 中多轮更新开始前的 log probs 和 values 称为 old log probs 和 old values(reward 不会多轮计算)。在每一轮中,actor 和 critic 会生成 new log probs 和 new values,然后在 old 的基础上计算 actor loss 和 critic loss,然后更新参数。

整体流程:

PPO优化目标

(1)策略梯度算法:更新幅度大,不稳定

(2)TRPO(信任区域算法):加入KL散度约束条件,但需计算二阶导数,计算量大

(3)PPO(近端策略优化算法):

这里At为优势函数:Critic Model用于估计状态的价值函数 V(st),从而计算策略梯度中的优势值A(t),下面的 r(st,at) 函数就是 RM 模型的输出: 用于计算生成某个token的即时收益 。 下图转换参考:https://zhuanlan.zhihu.com/p/651780908

PPO训练流程

  • Actor Model:要训练的目标语言模型,策略网络
  • Critic Model:预估总收益
  • Reward Model:计算即时收益
  • Reference Model:在RLHF阶段给语言模型增加一些“约束”,防止语言模型训偏

ColossalChat RLFH过程也是非常接近ChatGPT的RLFH过程,RLFH过程主要涉及四个模型分别是Actor、Critic、RM、STF,损失函数也是由三个损失函数组成分别是策略损失、价值损失和 PTX 损失。

ColossalChat RLFH过程

策略损失函数计算:

策略损失计算过程

通过instruction dataset数据训练STF模型,通过计算sft model的logits和actor model(没有经过sft的model)的logits计算kl散度,然后加上reward model的打分变成 reward R奖励值,避免太过偏向reward model加入和sft model的kl散度,同时也避免强化学习将actor模型训歪。

这样做的目的就是避免模型训飞,让模型更新保持在一个小范围内。

价值损失函数计算:

上式R是reward model和sft model计算出来的反馈分数,V(s)是Critic Model输出的价值分数。主要是衡量reward分数和价值函数分数的均方误差。

ptx的损失计算:

计算Actor输出response和输入语料的回答部分的交叉熵损失函数,用来在PPO梯度中加入预训练梯度,以保持语言模型原有性能防止遗忘。这个就是instruct gpt论文中在强化学习中加入预训练梯度以防过度拟合ppo数据带来nlp通用任务能力的下降操作。

总的强化学习损失计算:

为什么RLHF中,PPO需要Critic模型而不是直接使用RewardModel

在强化学习中,PPO(Proximal Policy Optimization)算法是一种基于策略梯度的方法,用于训练强化学习智能体。PPO算法中引入Critic模型的主要目的是为了提供一个价值估计器,用于评估状态或状态动作对的价值,从而辅助策略的更新和优化。

虽然奖励模型(Reward Model)可以提供每个状态或状态动作对的即时奖励信号,但它并不能直接提供对应的价值估计。奖励信号只反映了当前动作的即时反馈,而并没有提供关于在长期时间尺度上的价值信息。

Critic模型的作用是估计状态或状态动作对的长期价值,也称为状态值函数或动作值函数。Critic模型能够学习和预测在当前状态下采取不同动作所获得的累积奖励,它提供了对策略改进的指导。PPO算法使用Critic模型的估计值来计算优势函数,从而调整策略的更新幅度,使得更有利于产生更高长期回报的动作被选择。

另外,Critic模型还可以用于评估不同策略的性能,为模型的评估和选择提供依据。PPO算法中的Actor-Critic架构允许智能体同时学习策略和价值函数,并通过协同训练来提高性能。

因此,在RLHF(Reinforcement Learning from Human Feedback)中,PPO算法需要Critic模型而不是直接使用奖励模型,是为了提供对状态或状态动作对的价值估计,并支持策略的改进和优化。Critic模型的引入可以提供更全面和准确的信息,从而增强算法的训练效果和学习能力。

即时奖励 与 状态动作对的长期价值 的差别是什么?

即时奖励(Immediate Reward)和状态动作对的长期价值(Long-Term Value)代表了强化学习中不同的概念和时间尺度。

即时奖励是指智能体在执行某个动作后立即获得的反馈信号。它通常由环境提供,用于表示当前动作的好坏程度。即时奖励是一种即时反馈,可以指示当前动作的立即结果是否符合智能体的目标。

而状态动作对的长期价值涉及更长时间尺度上的评估,它考虑了智能体在当前状态下选择不同动作所导致的未来回报的累积。长期价值可以表示为状态值函数(State Value Function)或动作值函数(Action Value Function)。

状态值函数(V-function)表示在给定状态下,智能体从该状态开始执行一系列动作,然后按照某个策略进行决策,从而获得的预期累积回报。状态值函数估计了智能体处于某个状态时所能获得的长期价值,反映了状态的优劣程度

动作值函数(Q-function)则表示在给定状态下,智能体选择某个动作后,按照某个策略进行决策,从该状态转移到下一个状态并获得预期累积回报的价值。动作值函数估计了在给定状态下采取不同动作的长期价值,可以帮助智能体选择在每个状态下最优的动作

长期价值考虑了智能体在未来的决策过程中所能获得的累积回报,相比之下,即时奖励只提供了当前动作的即时反馈。长期价值对智能体的决策具有更全面的影响,可以帮助智能体更好地评估当前状态和动作的长期效果,并指导智能体在长期时间尺度上作出更优的决策。

在强化学习中,长期价值的估计对于确定性策略选择和价值优化非常重要,而即时奖励则提供了对当前动作的直接反馈。这两者相互补充,结合起来可以帮助智能体实现更好的决策和学习效果。

PPO中优势函数指什么

在Proximal Policy Optimization(PPO)算法中,优势函数(Advantage Function)用于评估状态-动作对的相对优劣程度。它衡量了执行某个动作相对于平均水平的优劣,即在给定状态下采取某个动作相对于采取平均动作的效果。

优势函数可以用以下方式定义:Advantage(s, a) = Q(s, a) - V(s)

其中,Advantage(s, a)表示在状态 s 下采取动作 a 的优势函数值,Q(s, a) 表示状态动作对 (s, a) 的动作值函数(也称为动作优势函数),V(s) 表示状态值函数。

优势函数的作用在于帮助评估当前动作的相对价值,以便在策略更新过程中确定应采取的动作。通过比较不同动作的优势函数值,可以决定哪些动作是更好的选择。正的优势函数值表示执行的动作比平均水平更好,而负的优势函数值表示执行的动作比平均水平更差。

在PPO算法中,优势函数用于计算策略更新的目标,以便调整策略概率分布来提高优势函数为正的动作的概率,并降低优势函数为负的动作的概率,从而改进策略的性能。

总而言之,优势函数在PPO算法中用于评估状态-动作对的相对优劣,帮助确定应该采取的动作,并在策略更新过程中引导策略向更优的方向调整。

GRPO (Group Relative Policy Optimization)

传统的强化学习算法(如Proximal Policy Optimization,PPO)在应用于LLMs的推理任务时面临着重大挑战:

  1. 依赖批评者模型:
  • PPO需要一个独立的批评者模型来评估每个回答的价值,这使内存和计算需求增加了一倍。
  • 训练批评者模型非常复杂且容易出错,尤其是在需要对主观或细微差别进行评价的任务中。

2. 高昂的计算成本:

  • 强化学习流程通常需要大量计算资源来迭代评估和优化回答。
  • 将这些方法扩展到更大的LLMs会进一步加剧成本。

3. 可扩展性问题:

  • 绝对奖励评估难以应对多样化任务,使得跨推理领域的泛化变得困难。

GRPO如何应对这些挑战:

  • 无批评者优化: GRPO通过比较组内回答,消除了对批评者模型的需求,显著降低了计算开销。
  • 相对评估: GRPO不依赖外部评价者,而是利用组内动态来评估每个回答在同一批次中的相对表现。
  • 高效训练: 通过专注于组内优势,GRPO简化了奖励估计流程,使其对大型模型的训练更快且更具可扩展性。

GRPO的核心思想是通过组内相对奖励来估计基线(baseline),从而避免使用额外的价值函数模型(critic model)。传统的PPO算法需要训练一个价值函数来估计优势函数(advantage function),而GRPO通过从同一问题的多个输出中计算平均奖励来替代这一过程,显著减少了内存和计算资源的消耗。

Group Relative Policy Optimization (GRPO),不再需要像PPO那样加入额外的价值函数近似,而是直接使用多个采样输出的平均奖励作为Baseline,显著减少了训练资源的使用。

具体来说,对于每个问题 i,GRPO 从旧策略 πθold​​ 中采样一组输出 {i1​,i2​,…,iA​},然后通过最大化以下目标函数来优化策略模型:

其中,ϵ 和 β 是超参数,A^i,j​ 是基于组内奖励的相对优势估计。与 PPO 不同,GRPO 通过直接使用奖励模型的输出来估计基线,避免了训练一个复杂的值函数。此外,GRPO 通过直接在损失函数中加入策略模型和参考模型之间的 KL 散度来正则化,而不是在奖励中加入 KL 惩罚项,从而简化了训练过程。

此外,GRPO 通过直接在损失函数中加入策略模型和参考模型之间的 KL 散度来正则化,而不是在奖励中加入 KL 惩罚项,从而简化了训练过程。

GRPO的计算流程包括:

  1. 采样一组输出并计算每个输出的奖励。
  2. 对组内奖励进行归一化处理。
  3. 使用归一化后的奖励计算优势函数。
  4. 通过最大化目标函数更新策略模型。
  5. 迭代训练,逐步优化策略模型。

GRPO通过组内相对奖励估计基线,避免了传统PPO中价值函数的使用,显著减少了训练资源消耗,同时提升了模型在数学推理等复杂任务中的表现。

image-20240804164951112

GRPO 计算总结

GRPO的核心思想是相对评估:

  1. 对于每个输入查询,模型生成一组潜在回答。
  2. 根据每个回答在组中的相对表现进行评分,而不是孤立地评估单个回答。
  3. 一个回答的优势反映了其相对于组内平均表现的优劣程度。

这种方法消除了对独立批评者模型的需求,使GRPO既高效又稳健通过在组内引入竞争,GRPO推动模型不断提升其推理能力。正是这一创新使DeepSeek在推理任务中取得了卓越的成果。

以简单的方式理解GRPO目标函数

GRPO(Group Relative Policy Optimization,群体相对策略优化)的目标函数就像是一种“食谱”,通过比较模型的回答并逐步改进,教会模型生成更好的答案。让我们用一个易于理解的方式逐步解析它:

目标:假设你正在教一组学生解决一个数学问题。你不是单纯告诉他们谁对谁错,而是通过比较所有学生的答案,找出谁做得最好(以及原因)。然后,你通过奖励更好的方法并改进较弱的方法来帮助他们学习。这正是GRPO所做的——只不过它教的是AI模型,而不是学生。

逐步解析

第一步:从查询开始

从训练数据集 P(Q) 中选取一个查询 (q)。
例如:假设查询是“8 + 5 的和是多少?”

第二步:生成一组回答

模型生成一组 G(4) 个回答来应对查询。
例如:模型生成以下回答:
o₁: “答案是13。”
o₂: “十三。”
o₃: “是12。”
o₄: “和是13。”

第三步:为每个回答计算奖励

什么是奖励?奖励通过量化模型回答的质量来指导其学习。

GRPO中的奖励类型:

  • 准确性奖励:基于回答的正确性(例如,解决数学问题)。
  • 格式奖励:确保回答遵循结构化的指导(例如,用 标签包裹的推理过程)。
  • 语言一致性奖励:惩罚语言混杂或格式不连贯的情况。

基于其表现,为每个回答分配一个奖励 (rᵢ)。例如,奖励可能取决于:

  • 准确性:答案是否正确?
  • 格式:回答是否结构良好?

示例:
r₁ = 1.0(正确且格式良好)。
r₂ = 0.9(正确但不够正式)。
r₃ = 0.0(回答错误)。
r₄ = 1.0(正确且格式良好)。

第四步:比较回答(群体优势)

  • 计算每个回答相对于群体的优势 (Aᵢ):
  • 表现优于群体平均水平的回答会获得正分,而表现较差的回答则会获得负分。
  • 这种方法在组内引入了竞争机制,推动模型生成更好的回答。

第五步:使用截断技术更新策略

  • 示例:如果新策略开始为 o₁ 分配过多概率,截断技术会确保它不会过度强调这个回答。
  • 这种方法即使在复杂任务(如推理)中,也能实现稳定可靠的策略优化。

第六步:使用KL散度惩罚偏离

整体流程

GRPO 目标函数的执行过程如下:

  1. 为查询生成一组回答。
  2. 根据预定义标准(例如准确性、格式)为每个回答计算奖励。
  3. 在组内比较回答,计算它们的相对优势(A_i)。
  4. 更新策略,以倾向于具有更高优势的回答,并通过截断技术确保稳定性。
  5. 对更新进行正则化,防止模型偏离其基线太远。

为什么GRPO有效?

  1. 无需评判器:GRPO 通过依赖组内比较,避免了单独评估器的需求,降低了计算成本。
  2. 稳定学习:截断技术和KL正则化确保模型稳步改进,不会出现剧烈波动。
  3. 高效训练:通过关注相对表现,GRPO 特别适合像推理这样的任务,因为这些任务很难用绝对评分衡量。

RLOO(REINFORCE Leave-One-Out)

Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs

在RLHF过程中,PPO是最常用的对齐算法。PPO是一种包含了很多技巧的强化学习算法,例如GAE,importance weight,policy/value的clip等等。本文提出的观点是,直接采用更简单的policy gradient类强化学习算法也可以取得很好的效果,PPO中的value模型,clip操作等模块可能并不有效。本文提出的RLOO(REINFORCE Leave-One-Out)算法在多种大模型任务中都取得了优于PPO/DPO的结果,同时也对噪声和KL约束更robust。

核心理念:在经典 REINFORCE 中引入 per-prompt baseline,即“留一法”(Leave-One-Out)来估计 baseline,显著降低梯度方差,在线 RLHF,无需 critic 网络。

PPO算法是由policy gradient/actor-critic等强化学习算法发展而来的算法,已经在众多的强化学习经典任务中验证了效果。然而把PPO用在LLM中会有计算成本高(需要加载policy/ref-policy/value/reward/4个模型),众多模块耦合在一起难以判断错误出现的源头,以及算法表现不稳定等问题。

为了解决上述问题,本文使用更原始,更简单的RL算法去替代PPO。RL中policy-based最基础的reinforce算法如下所示:

b表示baseline,用来降低方差。RLOO使用一种蒙特卡洛的方式去计算b:

这样的方式能够避免使用value model和GAE,减少显存占用。PPO使用GAE的方式来平衡误差和方差。与PPO相比,reinforce算法的方差更大,但是由于预训练出的模型足够强大,方差不是主要问题,用RLOO的形式去进行梯度更新是可以接受的。

REINFORCE++: 比 GRPO 稳定比PPO快

An Efficient RLHF Algorithm with Robustness to Both Prompt and Reward Models

REINFORCE++: A Simple and Efficient Approach for Aligning Large Language Models

REINFORCE++-baseline 在 reasoning 等任务中非常好用,全局的标准差归一化避免了 std 太小带来的 advantage不稳定

REINFORCE++的核心思想是将PPO中的各种优化技巧整合到经典的强化学习算法REINFORCE中,以提升其性能和稳定性。这样REINFORCE++不需要 Critic 从而节省计算资源,又有加持了 PPO 相关的优化技巧实现高效训练。 REINFORCE++的特点是 比 GRPO 稳定比PPO快。

REINFORCE算法基于蒙特卡罗方法,通过以下步骤进行操作:

策略采样:智能体根据当前策略与环境交互,生成一条状态-动作-奖励序列(轨迹)。

回报计算:对每条轨迹进行回报计算,通常采用折扣累计奖励的形式,即:

其中, γ 是折扣因子, rk 是在时间步 k 获得的即时奖励。

 梯度估计:使用蒙特卡罗方法计算策略梯度,更新策略参数 θ 的公式为:

策略更新:通过梯度上升法更新策略参数:

其中, α 是学习率。

RLHF Implementation Tricks

在 REINFORCE 上集成下面的优化 Tricks 以稳定模型的训练。

Token Level KL-Penalty

Token Level KL-Penalty 是一种在序列生成任务中使用的正则化技术。其主要目的是控制生成的文本与训练数据之间的差异,以避免模型生成过于偏离训练分布的输出。具体方法如下:

这种 Token-level KL 的好处是可以无缝兼容 PRM 并且实现了KL reward的信用分配 (更新:最近也有网友提到用GRPO的外置kl 也可以)

Mini-batch Updates

Mini-batch Updates 是一种常用的优化策略,旨在提高训练效率和稳定性。其基本思想是:

  • – 小批量样本:将训练数据划分为多个小批量(mini-batch),而不是使用整个数据集进行更新。
  • – 频繁更新:通过在每个小批量上进行多次参数更新,可以更快地收敛,同时减少内存消耗。
  • – 随机性引入:小批量更新引入了随机性,有助于避免局部最优解,提高模型的泛化能力。

Reward Normalization and Clipping

Reward Normalization and Clipping 是处理奖励信号不稳定的一种方法。具体包括:

-奖励归一化:通过对奖励进行标准化(例如,减去均值并除以标准差),使得奖励信号更为平稳,从而提高训练过程的稳定性。

– 奖励裁剪:限制奖励值在某个范围内,以防止极端奖励对模型更新造成过大的影响。这有助于保持学习过程的稳定性,并防止梯度爆炸。

Advantage Normalization

Advantage Normalization 是一种用于处理优势函数(advantage function)估计方差的方法。REINFORCE++的优势函数定义为

其中 r 是Outcome奖励函数, KL 是per-token 的kl reward, t 是token位置。

优势归一化的步骤包括:

– **均值和方差计算**:对一个batch计算出的优势值进行均值和方差计算。

– **归一化处理**:将优势值减去均值并除以标准差,使得优势值具有更好的数值稳定性,进而提高学习效果。

PPO-Clip

PPO-Clip 是近端策略优化(Proximal Policy Optimization, PPO)算法中的一个关键技巧,用于限制策略更新幅度。其主要思想是:

剪切目标函数:通过引入一个剪切机制,限制新旧策略之间的比率变化,确保更新不会过大。这可以用以下公式表示:

提高稳定性和样本效率:这种剪切机制有效防止了策略更新过大导致的不稳定,提高了算法的收敛速度和样本效率。

ORPO偏好优化(Odds Ratio Preference Optimization)

ORPO: Monolithic Preference Optimization without Reference Model

核心: 最大化正样本的生成概率,最小化负样本的生成概率 。相比DPO 【加载2个模型,其中一个推理,另外一个训练,直接在偏好数据上进行训练】,只加载训练模型,直接在偏好数据上进行训练

本文提出的算法ORPO是对SFT的改进,通过修改SFT阶段的损失函数,将类似于DPO中偏好对齐的思想引入到SFT中,提出一种无需奖励模型和参考模型算法。同时,ORPO只有一阶段,不需要像DPO一样需要先SFT再DPO对齐。在众多大模型任务上的实验结果表明,与SFT,DPO等算法相比,ORPO更有优势。

本文提出的算法ORPO是对SFT的改进,修改了SFT阶段的损失函数。同时,与DPO/PPO相比,ORPO将原本分两步进行的过程(SFT+DPO/PPO)合并为一步,更加简洁高效。

现在有许多方法可以使大型语言模型(LLM)与人类偏好保持一致。以人类反馈为基础的强化学习(RLHF)是最早的方法之一,并促成了ChatGPT的诞生,但RLHF的成本非常高。与RLHF相比,DPO、IPO和KTO的成本明显更低,因为它们不需要奖励模型。

虽然DPO和IPO的成本较低,但它们仍需训练两个不同的模型。首先是监督微调(SFT)步骤,即训练模型按指令回答问题,然后使用SFT模型作为初始化和参考,以使模型与人类偏好一致。

ORPO是另一种新的LLM对齐方法,这种方法甚至不需要SFT模型。通过ORPO,LLM可以同时学习回答指令和满足人类偏好。

对于STF,它是在与选择的答案配对的提示上进行训练的。用于sft的数据集可以与偏好优化使用的相同,但不包括”被拒绝”的答案。所以可以直观地认为,应该能够微调一个基础LLM,使其在学习如何回答指令的同时,也学会惩罚和偏好某些答案。

SFT只用正样本更新策略,没有考虑到负样本,会把负样本生成的概率同时拉高,如下图所示:

由于SFT的损失函数对于rejected data没有惩罚项,SFT之后正样本和负样本的生成概率有可能同时上升。

odds定义:模型θ生成 输出序列y的可能性 比 不生成y序列的可能性 比值。

OR为正负样本的odds的比值:

ORPO算法要做的就是最大化OR,即最大化正样本的生成概率,最小化负样本的生成概率,LOR项用了和DPO类似的logsigmoid的形式:

ORPO就是在这个理论基础上建立的,ORPO简单地通过添加负对数似然损失与OR损失(OR代表奇异比)来修改训练损失:

OR损失对被拒绝的答案进行弱惩罚,而对选择的答案进行强有力的奖励。这里包含了一个超参数lambda用于加权OR损失。通过ORPO的损失,模型在学习了SFT期间的内容的同时,也学会了人类偏好。

ORPO需要数千个训练步骤来学习如何区分选择的响应和拒绝的响应。为了获得类似的结果,应该训练ORPO至少2000步,总批大小为64(如论文所述)。

ORPO 已经可以在Hugging Face库上使用了,并且它因为只修改了损失函数,所以可以很好的与现有的Lora方法集成

ORPO是一种单步微调和对准指令llm的新方法。它不需要任何奖励或SFT模型,并且ORPO比DPO和RLHF更简单。根据论文ORPO的性能与DPO相当或略好。但是ORPO需要几千个训练步骤来学习好的和坏的反应之间的区别。

SimPO 简单偏好优化:

算法的核心是将偏好优化目标中的奖励函数与生成指标对齐不需要ref参考模型

SimPO 包含两个主要组件:(1)在长度上归一化的奖励【/|y|】,其计算方式是使用策略模型的奖励中所有 token 的平均对数概率;(2)目标奖励差额 γ ,用以确保获胜和失败响应之间的奖励差超过这个差额 γ 。

DPO 是最常用的离线偏好优化方法之一。DPO 并不会学习一个显式的奖励模型,而是使用一个带最优策略的闭式表达式来对奖励函数 r 进行重新参数化:

其中 π_θ 是策略模型,π_ref 是参考策略(通常是 SFT 模型),Z (x) 是配分函数。通过将这种奖励构建方式整合进 Bradley-Terry (BT) 排名目标,DPO 可使用策略模型而非奖励模型来表示偏好数据的概率,从而得到以下目标:

DPO 的奖励与生成之间的差异。使用 (1) 式作为隐式的奖励表达式有以下缺点:(1) 训练阶段需要参考模型 π_ref,这会带来额外的内存和计算成本;(2) 训练阶段优化的奖励与推理所用的生成指标之间存在差异。具体来说,在生成阶段,会使用策略模型 π_θ 生成一个能近似最大化平均对数似然的序列,定义如下:

在解码过程中直接最大化该指标是非常困难的,为此可以使用多种解码策略,如贪婪解码、波束搜索、核采样和 top-k 采样。此外,该指标通常用于在语言模型执行多选任务时对选项进行排名。在 DPO 中,对于任意三元组 (x, y_w, y_l),满足奖励排名 r (x, y_w) > r (x, y_l) 并不一定意味着满足似然排名:

图片

事实上,在使用 DPO 训练时,留存集中大约只有 50% 的三元组满足这个条件。

构建在长度上归一化的奖励。很自然地,我们会考虑使用 (3) 式中的 p_θ 来替换 DPO 中的奖励构建,使其与引导生成的似然指标对齐。这会得到一个在长度上归一化的奖励:

其中 β 是控制奖励差异大小的常量。该团队发现,根据响应长度对奖励进行归一化非常关键;从奖励公式中移除长度归一化项会导致模型倾向于生成更长但质量更低的序列。这样一来,构建的奖励中就无需参考模型了,从而实现比依赖参考模型的算法更高的内存和计算效率。

SimPO 目标

目标奖励差额。另外,该团队还为 Bradley-Terry 目标引入了一个目标奖励差额项 γ > 0,以确保获胜响应的奖励 r (x, y_w) 超过失败响应的奖励 r (x, y_l) 至少 γ

两个类之间的差额已知会影响分类器的泛化能力。在使用随机模型初始化的标准训练设置中,增加目标差额通常能提升泛化性能。在偏好优化中,这两个类别是单个输入的获胜或失败响应。

在实践中,该团队观察到随着目标差额增大,生成质量一开始会提升,但当这个差额变得过大时,生成质量就会下降。DPO 的一种变体 IPO 也构建了与 SimPO 类似的目标奖励差额,但其整体目标的效果不及 SimPO。

目标。最后,通过将 (4) 式代入到 (5) 式中,可以得到 SimPO 目标:

总结起来,SimPO 采用了与生成指标直接对齐的隐式奖励形式,从而消除了对参考模型的需求。此外,其还引入了一个目标奖励差额 γ 来分离获胜和失败响应

KTO:Kahneman-Tversky Optimisation

特点:

KTO关注的是答案偏离平均水准的程度——比平均好还是坏。所以它的训练数据集是对单个问答的“好/差”标注,而不再是成对数据间谁好谁差(所以用户对LLM结果的点赞或踩就可以当做反馈使用了)。

KTO不需要偏好数据,可以直接利用二元信号标记的数据来训练算法,对于负样本更加敏感。 KTO并不需要一个数据对,只需要对生成的结果进行good/bad的二元标注即可。 【比如:ChatGPT UI 界面会输出两个答案,用户可以选择哪个更好,适用于从生产环境中运行的聊天模型的训练】

实验表明,KTO算法在一定参数范围内能够超过DPO算法,并且KTO可以处理数据正负样本不平衡的情况。同时,在跳过SFT阶段的情况下,直接使用KTO相比于直接使用 DPO,效果有很大提升。在数据正负样本比例失衡/偏好数据有非传递性/偏好数据有噪声/的情况下,使用KTO可能是更好的选择。

KTO 使用 Kahneman-Tversky 人类效用模型,论文提出直接最大化生成效用的 HALO, 而不是最大化偏好的对数可能性。

  • 在1B~30B尺度上与基于偏好的方法的性能相匹配或超过,尽管它只从二进制信号(0或者1)中学习输出是否可取。
  • 没有一个 HALO 普遍优越;
  • 最佳损失取决于最适合给定设置的归纳偏差,经常被忽视的考虑因素。

KTO算法的具体步骤如下:

  1. 定义效用函数:根据前景理论中的效用函数公式,定义一个效用函数,用于计算模型输出相对于参考点的效用。
  2. 计算参考点:根据概率分布Q(X’, Y’ | x, y),计算出一个参考点,用于衡量模型输出的效用。
  3. 计算模型输出的效用:对于每个输入,计算模型输出相对于参考点的收益或损失,然后使用效用函数计算这些收益或损失的效用。
  4. 优化模型参数:优化模型参数以最大化模型输出的总效用。

KTO 损失函数本质是把 pair-wise 公式变成 point-wise 方式,结合了HALOs以及二元信号数据的思想提出使用Kahneman-Tversky 优化的KTO算法:

其中 zo是KL散度项,参考点zo为最优策略下reward的期望值,最终可以推导成KL散度的形式,y’表示任意输出,在实际训练中,Z0表示batch平均水准的程度Z0当前batch里面的样本进行估计得到的,平均 reward,代表不好不坏的居中的结果。 LKTO 就是DPO中推导的reward函数形式。

按照上面的定义估计z0是不切实际的,因为从πθ采样很慢,人类无法感知πθ引起的完整分布。

这个估计是有偏差的,但这是可取的,因为它使我们更接近人类如何构建他们的主观参考点。

实际上KTO相对比DPO差异就两点

  • 对正负样本进行了加权:DPO里面是使用正负样本的reward差值进行sigmoid映射,但是KTO里面使用reward模型与KL散度之间的差异!(说是KL散度,但其实也是bad的log比值数值!不过不是同一个pair
  • 注意:在实践的时候,KL项并不参与反向传播,这其实就跟DPO更相似的。DPO使一个数据对,但是这里把DPO给拆分了,相当于对每一个样本单独进行最大化或最小化了,以及进行加权。另一个作用就是,如果 rKTO(x,y) 的差异与KL散度有足够区别的话,那对应的Loss也就比较小。因此,KTO会更加鼓励差异大的数据对。
  • 但其实我们可以从KTO的目标函数直接看到。由于KTO是分别针对单条数据,如果数据是正样本,那么一定要超过 zo 才会产生预测正确反馈;对于负样本,需要低于 zo才会产生预测正确反馈

KTO和DPO的选择

  • 数据比例:如果数据集是以good/bad形式进行标注,并且数据比例不平衡,那么选择KTO
  • 数据质量:如果你的偏好数据质量高,数据噪声小,那么DPO的效果更好。由于目前公开的数据集中存在的噪声较大,这就能解释为什么KTO的效果会超过DPO了。
  • 理论分析:KTO不会从负样本中学习到很高的反馈,也不会从正样本中学习到很低的反馈(所以对噪声比较鲁棒)

KTO 的工作原理:

  • 如果模型以直接(blunt manner)方式增加了理想示例的奖励,那么 KL 惩罚也会增加,并且不会取得任何进步。这迫使模型准确地了解是什么让输出变得理想,这样就可以增加奖励,同时保持 KL 项持平(甚至减少)。
  • 实际实现中,KL 项是通过当前batch里面的正负样本进行估计得到的可以认为是batch样本的平均水平】,详细 debug KTOTrainer 源代码

对成对偏好数据进行分配

  • 与大多数比对方法一样,DPO 需要一个成对偏好数据集(x, y_w, y_l),够根据一组标准(如有益性或有害性)来标记哪种模型响应更好。
  • 实践过程中,创建这些数据是一项耗时且成本高昂的工作。
  • ContextualAI 提出替代方案,称为 Kahneman-Taversky 优化(KTO),完全根据被标记为「好」或「坏」的样本(例如在聊天 UI 中看到的图标👍或👎)来定义损失函数。这些标签更容易获得, KTO 是一种很有前景的方法,不断更新在生产环境中运行的聊天模型。

与此同时,这些方法都有相应的超参数,其中最重要的是 β ,控制对使用模型的偏好程度的权重。这些方法已经在第三方库(如 huggingface TRL)中实现

KTO 数据集

KTO 不需要成对的偏好数据,实验时直接将 GPT-4 生成的响应归类为「好」标签,将 Llama Chat 13b 的响应视为「坏」标签。

KTO数据集与偏好数据集类似,但不同于给出一个更优的回答和一个更差的回答,KTO数据集对每一轮问答只给出一个 true/false 的 label。 除了 instruction 以及 input 组成的人类最终输入和模型回答 output ,KTO 数据集还需要额外添加一个 kto_tag 列(true/false)来表示人类的反馈。在一轮问答中其格式如下:

[
  {
    "instruction": "人类指令(必填)",
    "input": "人类输入(选填)",
    "output": "模型回答(必填)",
    "kto_tag": "人类反馈 [true/false](必填)"
  }
]

对于上述格式的数据, dataset_info.json 中的 数据集描述 应为:

"数据集名称": {
  "file_name": "data.json",
  "columns": {
    "prompt": "instruction",
    "query": "input",
    "response": "output",
    "kto_tag": "kto_tag"
  }
}

代码实现:

基于pytorch、deepspeed、transformers代码:

sft训练代码:

def loss(self, sft_batch: SupervisedBatch) -> dict[str, torch.Tensor]:
    """Loss function for supervised finetuning."""
    outputs = self.model(**self.infer_batch(sft_batch))
    return {'loss': outputs.loss}

def train_step(self, sft_batch: SupervisedBatch) -> dict[str, Any]:
    """Performs a single training step."""
    loss = self.loss(sft_batch)['loss']
    self.model.backward(loss)
    self.model.step()

    return {
        'train/loss': loss.item(),
        'train/lr': self.model.optimizer.param_groups[0]['lr'],
    }

dpo训练代码:

https://blog.csdn.net/weixin_43013480/article/details/141370269

# 从 logits(未归一化的概率分布)中,提取 labels 对应类别的对数概率(log probabilities)。
def gather_log_probabilities(
    logits: torch.Tensor,  # size = (B, L, V)
    labels: torch.LongTensor,  # size = (B, L)
) -> torch.Tensor:  # size = (B, L)
    """Gather log probabilities of the given labels from the logits."""
    log_probs = F.log_softmax(logits, dim=-1)  # size = (B, L, V)
    gathered_log_probs = torch.gather(  # size = (B, L, 1)
        log_probs,
        dim=-1,
        index=labels.unsqueeze(dim=-1).to(torch.int64),
    )
    return gathered_log_probs.squeeze(dim=-1)  # size = (B, L)
def compute_log_probs(
    self,
    model: AutoModelForCausalLM,
    batch: PreferenceBatch,
) -> torch.Tensor:
    """Compute log probabilities of given sequences."""
    # 获得所有可能输出的log概率,logits 表示每个 token 位置的 未归一化的概率分布
    logits = model(**self.infer_batch(batch)).logits
    device = logits.device
    input_ids = batch['input_ids']
    #取得每个样本的回复长度,用于截取模型输出
    batch_size = len(batch['meta_info']['response_lens'])
    logprob_list = []
    for idx in range(batch_size):
        response_length = batch['meta_info']['response_lens'][idx]
        # 去除填充 (PAD) token,避免计算无效 token 的概率。
        raw_input_id = strip_pad(input_ids[idx], self.tokenizer.pad_token_id)
        #只保留 回复部分的 logits,丢弃 prompt 部分。 
        logit = logits[idx][-response_length:].unsqueeze(0)
        input_id = raw_input_id[-response_length:].unsqueeze(0)
        #计算对应的better 和worse 序列token 对数概率
        log_p = gather_log_probabilities(logit[:, :-1], input_id[:, 1:])
        logprob_list.append(log_p.squeeze(0))
    # 不同样本的 log_probs 长度不同,使用 pad_sequence 进行 padding,补齐到相同长度。
    return torch.nn.utils.rnn.pad_sequence(
        logprob_list, batch_first=True, padding_value=0.0
    ).to(device)

def loss(  # pylint: disable=too-many-locals
    self,
    batch: PreferenceBatch,
) -> dict[str, torch.Tensor]:
    """Loss function for the DPO algorithm."""
   #计算当前模型 (self.model.module) 在 batch 上的 log 概率。
    sequence_log_probs = self.compute_log_probs(
        self.model.module,
        batch,
    )
# better_sequence_log_probs (用户偏好的回复)
# worse_sequence_log_probs (用户不喜欢的回复)
    (
        better_sequence_log_probs,  # size = (B, L - 1)
        worse_sequence_log_probs,  # size = (B, L - 1)
    ) = sequence_log_probs.chunk(chunks=2, dim=0)
# 计算参考模型 (self.reference_model.module) 的对数概率 (log_probs)。
# reference_model 通常是 原始未优化的模型,作为对比基准。
# torch.no_grad() 表示 不计算梯度,避免影响参考模型。
    with torch.no_grad():
        ref_sequence_log_probs = self.compute_log_probs(  # size = (2 * B, L - 1)
            self.reference_model.module,
            batch,
        )
        ref_better_sequence_log_probs, ref_worse_sequence_log_probs = (
            ref_sequence_log_probs.chunk(chunks=2, dim=0)
        )

    losses = []
    better_sample_rewards = []
    worse_sample_rewards = []

    batch_size = better_sequence_log_probs.size(0)
    for i in range(batch_size):
# 计算 更好/更差回复的总 log 概率(即累加 token 级别 log 概率)。
        better_log_prob = better_sequence_log_probs[i, :].sum(dim=-1)
        worse_log_prob = worse_sequence_log_probs[i, :].sum(dim=-1)
        ref_better_log_prob = ref_better_sequence_log_probs[i, :].sum(dim=-1)
        ref_worse_log_prob = ref_worse_sequence_log_probs[i, :].sum(dim=-1)
# 当前模型比参考模型更偏好 better 回复 的程度。
        better_log_ratio = better_log_prob - ref_better_log_prob
# 当前模型比参考模型更偏好 worse 回复 的程度。
        worse_log_ratio = worse_log_prob - ref_worse_log_prob
# 计算 better 和 worse 的 log 比值差
# 使用 -logsigmoid(x) 计算负对数 sigmoid 损失,优化模型使其更倾向 better 回复。
# logsigmoid 的性质:
# 如果 x 很大,logsigmoid(x) ≈ 0,意味着损失小,模型已经正确偏好 better response。
# 如果 x 很小或负,logsigmoid(x) ≈ x,意味着损失大,模型没有正确区分 better 和 worse,需要优化。
        losses.append(
            -F.logsigmoid(
                self.cfgs.train_cfgs.scale_coeff * (better_log_ratio - worse_log_ratio),
            ),
        )
        better_sample_rewards.append(
            self.cfgs.train_cfgs.scale_coeff * better_log_ratio.detach(),
        )
        worse_sample_rewards.append(self.cfgs.train_cfgs.scale_coeff * worse_log_ratio.detach())
    loss = torch.stack(losses).mean()  # size = ()
    better_sample_reward = torch.stack(better_sample_rewards)  # size = (B,)
    worse_sample_reward = torch.stack(worse_sample_rewards)  # size = (B,)
# 计算 奖励 (reward)、准确率 (accuracy) 和奖励间距 (margin)。
    reward = better_sample_reward + worse_sample_reward  # size = (B,)
    reward_accuracy = (better_sample_reward > worse_sample_reward).float().mean()  # size = ()
    reward_margin = better_sample_reward - worse_sample_reward  # size = (B,)

    return {
        'loss': loss,
        'reward': reward,
        'better_sample_reward': better_sample_reward,
        'worse_sample_reward': worse_sample_reward,
        'reward_accuracy': reward_accuracy,
        'reward_margin': reward_margin,
    }

def train_step(
    self,
    batch: PreferenceBatch,
) -> dict[str, Any]:
    """Perform a single training step for DPO."""
    loss_dict = self.loss(batch=batch)
    loss = loss_dict['loss']
    self.model.backward(loss)
    self.model.step()

    with torch.no_grad():
        reward = loss_dict['reward'].mean()
        better_sample_reward = loss_dict['better_sample_reward'].mean()
        worse_sample_reward = loss_dict['worse_sample_reward'].mean()
        reward_accuracy = loss_dict['reward_accuracy']
        reward_margin = loss_dict['reward_margin'].mean()

        loss = get_all_reduce_mean(loss)
        reward = get_all_reduce_mean(reward)
        better_sample_reward = get_all_reduce_mean(better_sample_reward)
        worse_sample_reward = get_all_reduce_mean(worse_sample_reward)
        reward_accuracy = get_all_reduce_mean(reward_accuracy)
        reward_margin = get_all_reduce_mean(reward_margin)

    return {
        'train/loss': loss.item(),
        'train/reward': reward.item(),
        'train/better_sample_reward': better_sample_reward.item(),
        'train/worse_sample_reward': worse_sample_reward.item(),
        'train/reward_accuracy': reward_accuracy.item(),
        'train/reward_margin': reward_margin.item(),
        'train/lr': self.model.optimizer.param_groups[0]['lr'],
    }

ppo训练代码:

#使用策略模型 (Actor Model) 生成文本,并返回其 input_ids 和 attention_mask。
def actor_step(self, mini_prompt_only_batch: PromptOnlyBatch) -> dict[str, Any]:
    infer_batch = self.infer_batch(mini_prompt_only_batch)
    actor_batch = copy.deepcopy(infer_batch)
    sequences = self.actor_model.module.generate(
        **infer_batch,
        generation_config=self.generation_config,
        synced_gpus=True,
        do_sample=True,
    )
    attention_mask = sequences.not_equal(self.tokenizer.pad_token_id)
    actor_batch['input_ids'] = sequences
    actor_batch['attention_mask'] = attention_mask

    return actor_batch
# 计算奖励值 (reward) 和对抗奖励值 (reward_values)。 
def reward_model_step(self, actor_batch: PromptOnlyBatch) -> dict[str, Any]:
        reward_batch = copy.deepcopy(actor_batch)
        if self.reward_tokenizer is not self.tokenizer:
            reward_tokenize_output = batch_retokenize(
                actor_batch['input_ids'],
                src_tokenizer=self.tokenizer,
                dest_tokenizer=self.reward_tokenizer,
                skip_special_tokens=True,
                device=self.args.device,
            )
            reward_batch['input_ids'] = reward_tokenize_output['input_ids']
            reward_batch['attention_mask'] = reward_tokenize_output['attention_mask']
        reward_infer_batch = self.reward_infer_batch(reward_batch)
        reward_batch['reward'] = self.reward_model(**reward_infer_batch).end_scores.squeeze(dim=-1)
        critic_infer_batch = self.reward_infer_batch(actor_batch)
        scores = self.reward_critic_model(**critic_infer_batch).scores
        reward_batch['reward_values'] = scores.squeeze(dim=-1)[:, :-1]

        return reward_batch
#冻结模型参数,避免影响训练,采样多个 mini-batch,生成文本,计算奖励,计算 log 概率 (log_probs),计算参考模型的 log 概率 (ref_log_probs)
# 经验回放:生成训练数据并计算指标
  @torch.no_grad()
    def rollout(self, prompt_only_batch: PromptOnlyBatch) -> list[dict[str, Any]]:
        """Rollout a batch of experiences."""
        # freeze the model for rolling out
        self.set_train(mode=False)

        total_batch_size = prompt_only_batch['input_ids'].size(0)
        micro_batch_size = int(self.cfgs.train_cfgs.per_device_train_batch_size)
        micro_inference_batches = []
        micro_training_batches = []
        mini_batch = {}
        for i in range(0, total_batch_size, micro_batch_size):

            mini_batch = {
                key: prompt_only_batch[key][i : i + micro_batch_size] for key in prompt_only_batch
            }

            # actor generation
            actor_batch = self.actor_step(mini_batch)
            # reward model and reward critic model scoring
            reward_batch = self.reward_model_step(actor_batch)
            # calculate the log probabilities
            logits = self.actor_model(**actor_batch).logits
            ref_logits = self.actor_reference_model(**actor_batch).logits
            log_probs = gather_log_probabilities(logits[:, :-1], actor_batch['input_ids'][:, 1:])
            ref_log_probs = gather_log_probabilities(
                ref_logits[:, :-1], actor_batch['input_ids'][:, 1:]
            )

            micro_training_batch = {}
            micro_training_batch['prompt_idx'] = mini_batch['input_ids'].size(-1) - 1
            micro_training_batch['log_probs'] = log_probs
            micro_training_batch['ref_log_probs'] = ref_log_probs
            micro_training_batch['reward'] = reward_batch['reward']
            micro_training_batch['reward_values'] = reward_batch['reward_values']

            mini_batch['input_ids'] = reward_batch['input_ids']
            mini_batch['attention_mask'] = actor_batch['attention_mask']
            # add rollout results to the batches
            micro_inference_batches.append(mini_batch)
            micro_training_batches.append(micro_training_batch)

        # unfreeze the model for training
        self.set_train()

        return micro_inference_batches, micro_training_batches

#计算策略梯度损失
# 计算 PPO 损失函数:
# ratios = exp(new_log_probs - old_log_probs)(新旧策略比)。
# 裁剪 ratios 避免策略剧烈变化(PPO 关键)。
# return -masked_mean(surrogate, mask):最大化优势 𝐴𝑡
   
def actor_loss_fn(
        self,
        log_probs: torch.Tensor,  # size = (B, L - S)
        old_log_probs: torch.Tensor,  # size = (B, L - S)
        advantages: torch.Tensor,  # size = (B, L - S)
        mask: torch.BoolTensor,  # size = (B, L - S)
    ) -> torch.Tensor:  # size = ()
        # size = (B, L - S)
        ratios = torch.exp(log_probs - old_log_probs)
        surrogate1 = advantages * ratios
        surrogate2 = advantages * torch.clamp(
            ratios,
            1.0 - self.clip_range_ratio,
            1.0 + self.clip_range_ratio,
        )
        surrogate = torch.minimum(surrogate1, surrogate2)
        return -masked_mean(surrogate, mask)  # size = ()
#  rl_step函数是训练过程中使用强化学习(RL)更新策略的一步。在PPo算法中,rl_step是用来更新策略网络(actor)和价值网络(critic)的一部分。具体来说,这个函数通过计算强化学习损失(actor loss和critic loss),并通过反向传播优化这两个网络。
# reward_critic_model 评估奖励函数的 价值估计,用于计算 优势函数 𝐴𝑡不是直接计算奖励,而是估算未来可能获得的奖励。主要用于时间差分(TD learning)更新策略,类似于 价值函数。

def rl_step(
        self, inference_batch: dict[str, torch.Tensor], training_batch: dict[str, torch.Tensor]
    ) -> dict[str, Any]:
        """Perform a single update step with RL loss."""
        old_log_probs = training_batch['log_probs']
        ref_log_probs = training_batch['ref_log_probs']
        reward = training_batch['reward']
        old_reward_values = training_batch['reward_values']
        start = training_batch['prompt_idx']

        input_ids = inference_batch['input_ids']
        attention_mask = inference_batch['attention_mask']

        sequence_mask = attention_mask[:, 1:]

        with torch.no_grad():
            old_rewards = self.add_kl_divergence_regularization(
                reward,
                old_log_probs,
                ref_log_probs,
                sequence_mask,
            )
            reward_advantages, reward_returns = self.get_advantages_and_returns(
                old_reward_values,
                old_rewards,
                sequence_mask,
                start,
            )

        logits = self.actor_model(**inference_batch, use_cache=False).logits
        log_probs = gather_log_probabilities(logits[:, :-1], input_ids[:, 1:])
        actor_loss = self.actor_loss_fn(
            log_probs[:, start:],
            old_log_probs[:, start:],
            reward_advantages,
            sequence_mask[:, start:],
        )
        self.actor_model.backward(actor_loss)
        self.actor_model.step()

        reward_values = self.reward_critic_model(**inference_batch).scores
        reward_values = reward_values.squeeze(dim=-1)[:, :-1]
        reward_critic_loss = self.critic_loss_fn(
            reward_values[:, start:],
            old_reward_values[:, start:],
            reward_returns,
            sequence_mask[:, start:],
        )
        self.reward_critic_model.backward(reward_critic_loss)
        self.reward_critic_model.step()

        with torch.no_grad():
            mask = sequence_mask[:, start:]
            kl_divergence = ((old_log_probs - ref_log_probs)[:, start:] * mask).sum(dim=-1).mean()
            mean_generated_length = mask.sum(dim=-1).float().mean()
            max_generated_length = mask.sum(dim=-1).float().max()

            reward = reward.mean()
            reward_with_kl_penalty = (old_rewards[:, start:] * mask).sum(dim=-1).mean()
            reward_advantage = masked_mean(reward_advantages, mask)
            reward_return = masked_mean(reward_returns, mask)
            reward_value = masked_mean(reward_values[:, start:], mask)

            actor_loss = get_all_reduce_mean(actor_loss)
            reward_critic_loss = get_all_reduce_mean(reward_critic_loss)
            reward = get_all_reduce_mean(reward)
            reward_with_kl_penalty = get_all_reduce_mean(reward_with_kl_penalty)
            reward_advantage = get_all_reduce_mean(reward_advantage)
            reward_return = get_all_reduce_mean(reward_return)
            reward_value = get_all_reduce_mean(reward_value)
            kl_divergence = get_all_reduce_mean(kl_divergence)
            mean_generated_length = get_all_reduce_mean(mean_generated_length)
            max_generated_length = get_all_reduce_max(max_generated_length)

        dist.barrier()

        return {
            'train/actor_loss': actor_loss.item(),
            'train/reward_critic_loss': reward_critic_loss.item(),
            'train/reward': reward.item(),
            'train/reward_with_kl_penalty': reward_with_kl_penalty.item(),
            'train/reward_advantage': reward_advantage.item(),
            'train/reward_return': reward_return.item(),
            'train/reward_value': reward_value.item(),
            'train/kl_divergence': kl_divergence.item(),
            'train/actor_lr': self.actor_model.optimizer.param_groups[0]['lr'],
            'train/reward_critic_lr': self.reward_critic_model.optimizer.param_groups[0]['lr'],
            'train/mean_generated_length': mean_generated_length.item(),
            'train/max_generated_length': max_generated_length.item(),
        }

    def ptx_step(self, ptx_batch: dict[str, torch.Tensor]) -> dict[str, Any]:
        """Perform a single update step with PTX loss."""
        ptx_loss = self.actor_model(**self.infer_batch(ptx_batch)).loss
        self.actor_model.backward(self.ptx_coeff * ptx_loss)
        self.actor_model.step()
        ptx_loss = get_all_reduce_mean(ptx_loss)
        return {
            'train/ptx_loss': ptx_loss.item(),
        }

    def train(self) -> None:
        """Train the model."""
        self.logger.print('***** Running training *****')

        progress_bar = tqdm(
            total=self.total_training_steps,
            desc=f'Training 1/{self.cfgs.train_cfgs.epochs} epoch',
            position=0,
            leave=True,
            disable=not is_main_process(),
        )

        if self.cfgs.data_cfgs.eval_datasets:
            self.logger.print('\n***** Evaluating at the beginning *****')
            self.eval()

        num_prompt_only_batches = len(self.prompt_only_dataloader)
        num_ptx_batches = len(self.ptx_dataloader)
        num_ptx_replicas = (num_prompt_only_batches + num_ptx_batches - 1) // num_ptx_batches
        for epoch in range(int(self.cfgs.train_cfgs.epochs)):
            for prompt_only_batch, ptx_batch in zip(
                self.prompt_only_dataloader,
                itertools.chain.from_iterable([self.ptx_dataloader] * num_ptx_replicas),
            ):
                inference_batches, training_batches = self.rollout(prompt_only_batch)

                if self.use_ptx:
                    ptx_batches = self.split_ptx_micro_batches(ptx_batch)
                else:
                    ptx_batches = [None for _ in range(len(inference_batches))]
                torch.cuda.empty_cache()

                for _ in range(self.cfgs.train_cfgs.update_iters):
                    for inference_batch, training_batch, ptx_batch in zip(
                        inference_batches, training_batches, ptx_batches
                    ):
                        rl_info = self.rl_step(inference_batch, training_batch)

                        torch.cuda.empty_cache()
                        self.logger.log(rl_info, step=self.global_step)
                        if self.use_ptx:
                            ptx_info = self.ptx_step(ptx_batch)
                            torch.cuda.empty_cache()
                            self.logger.log(ptx_info, step=self.global_step)

                        self.global_step += 1
                        progress_bar.set_description(
                            f'Training {epoch + 1}/{self.cfgs.train_cfgs.epochs} epoch '
                            f'(reward {rl_info["train/reward"]:.4f})',
                        )
                        progress_bar.update(1)

                        if self.global_step % self.cfgs.logger_cfgs.save_interval == 0:
                            self.logger.print(f'Saving checkpoint at step {self.global_step} ...')
                            self.save(tag=self.global_step)
                            self.logger.print('Checkpoint saved.')

                        if (
                            self.cfgs.data_cfgs.eval_datasets
                            and self.cfgs.train_cfgs.eval_strategy == 'steps'
                            and self.global_step % self.cfgs.train_cfgs.eval_interval == 0
                        ):
                            self.logger.print(
                                f'\n***** Evaluating at step {self.global_step} *****',
                            )
                            self.eval()

RM奖励模型训练代码:

    def loss(
        self,
        batch: PreferenceBatch,
    ) -> dict[str, torch.Tensor]:
        """Loss function for the reward model."""
        (
            better_input_ids,  # size = (B, L)
            worse_input_ids,  # size = (B, L)
        ) = batch[
            'input_ids'
        ].chunk(chunks=2, dim=0)
        assert better_input_ids.size(0) == worse_input_ids.size(0), 'batch size mismatch!'

# scores:一般来说,这代表模型在每个时间步骤(或输入分段)上的奖励得分,通常是一个形状为 (B, L, 1) 的张量,其中 B 是批量大小,L 是输入序列的长度,1 是奖励得分的维度。
#end_scores:通常表示输入序列的结束阶段的奖励得分,这可能是在整个序列处理完成后,模型计算出的最终奖励。
        output = self.model(**self.infer_batch(batch))
        scores = output.scores
        end_scores = output.end_scores
        higher_rewards, lower_rewards = scores.squeeze(dim=-1).chunk(chunks=2, dim=0)
        higher_end_reward, lower_end_reward = end_scores.squeeze(dim=-1).chunk(chunks=2, dim=0)

        loss = -F.logsigmoid(higher_end_reward - lower_end_reward).mean()

        if self.cfgs.train_cfgs.regularization > 0.0:
            loss = (
                loss
                + self.cfgs.train_cfgs.regularization
                * torch.stack([lower_end_reward, higher_end_reward]).square().mean()
            )

        accuracy = (higher_end_reward > lower_end_reward).float().mean()  # size = ()
        return {
            'loss': loss,  # size = ()
            'higher_end_reward': higher_end_reward,  # size = (B,)
            'lower_end_reward': lower_end_reward,  # size = (B,)
            'higher_rewards': higher_rewards,  # size = (B, L)
            'lower_rewards': lower_rewards,  # size = (B, L)
            'accuracy': accuracy,  # size = ()
        }

    def train_step(
        self,
        batch: PreferenceBatch,
    ) -> dict[str, Any]:
        """Perform a single training step."""
        loss_dict = self.loss(batch)
        loss = loss_dict['loss']
        self.model.backward(loss)
        self.model.step()

        accuracy = loss_dict['accuracy']

        loss = get_all_reduce_mean(loss)
        accuracy = get_all_reduce_mean(accuracy)

        return {
            'train/loss': loss.item(),
            'train/accuracy': accuracy.item(),
            'train/lr': self.model.optimizer.param_groups[0]['lr'],
        }

orpo 训练代码:

相关介绍:https://github.com/Paul33333/ORPO https://zhuanlan.zhihu.com/p/688583797

# 从 logits(未归一化的概率分布)中,提取 labels 对应类别的对数概率(log probabilities)。
def gather_log_probabilities(
    logits: torch.Tensor,  # size = (B, L, V)
    labels: torch.LongTensor,  # size = (B, L)
) -> torch.Tensor:  # size = (B, L)
    """Gather log probabilities of the given labels from the logits."""
    log_probs = F.log_softmax(logits, dim=-1)  # size = (B, L, V)
    gathered_log_probs = torch.gather(  # size = (B, L, 1)
        log_probs,
        dim=-1,
        index=labels.unsqueeze(dim=-1).to(torch.int64),
    )
    return gathered_log_probs.squeeze(dim=-1)  # size = (B, L)

# compute_log_probs 的作用是计算给定序列的 log 概率(对数概率),主要用于评估语言模型(LLM)的生成质量。
def compute_log_probs(
        self,
        model: AutoModelForCausalLM,
        batch: PreferenceBatch,
    ) -> torch.Tensor:
        """Compute log probabilities of given sequences."""
        logits = model(**self.infer_batch(batch)).logits
        device = logits.device
        input_ids = batch['input_ids']
        batch_size = len(batch['meta_info']['response_lens'])
        logprob_list = []
        for idx in range(batch_size):
            response_length = batch['meta_info']['response_lens'][idx]  # for the eos token
            logit = logits[idx][-response_length:].unsqueeze(0)
            input_id = input_ids[idx][-response_length:].unsqueeze(0)
# logit[:, :-1]取 response 部分的 logits,去掉最后一个 token(因为 logits 预测的是下一个 token)input_id[:, 1:]: 取 response 部分的 token IDs,从第二个 token 开始(因为 log_probs 计算的是下一个 token 概率)。
作用:计算 response 部分每个 token 的 log 概率(对 logit 的 softmax 取对数)。
            log_p = gather_log_probabilities(logit[:, :-1], input_id[:, 1:]) 
            logprob_list.append(log_p.squeeze(0))
#pad填充,返回张量形状 (B, max_L_resp)
        return torch.nn.utils.rnn.pad_sequence(
            logprob_list, batch_first=True, padding_value=0.0
        ).to(device)
class ORPOTrainer(DPOTrainer):

    def loss(  # pylint: disable=too-many-locals
        self,
        batch: PreferenceBatch, # size = (2*B, L)
    ) -> dict[str, torch.Tensor]:
        """Loss function for the ORPO algorithm."""
        sequence_log_probs = self.compute_log_probs(
            self.model.module,
            batch,
        )
        (
            better_sequence_log_probs,  # size = (B, L - 1)
            worse_sequence_log_probs,  # size = (B, L - 1)
        ) = sequence_log_probs.chunk(chunks=2, dim=0)

        losses = []
        better_sample_rewards = []
        worse_sample_rewards = []

        better_input_ids, worse_input_ids = batch['input_ids'].chunk(chunks=2, dim=0)
        better_attention_mask, worse_attention_mask = batch['attention_mask'].chunk(chunks=2, dim=0)

        batch_size = better_input_ids.size(0)
#diverge_index 代表 better 和 worse 输入序列开始不同的位置:diverge_index,即它之后的 token 是模型生成的部分。
        for i in range(batch_size):
            if torch.all(torch.eq(better_input_ids[i], worse_input_ids[i])).item():
                continue
            better_end_index = better_attention_mask[i].nonzero()[-1].squeeze().item()
            worse_end_index = worse_attention_mask[i].nonzero()[-1].squeeze().item()
            diverge_index = (
                (better_input_ids[i] != worse_input_ids[i]).nonzero()[0].squeeze().item()
            )
            assert 0 <= diverge_index <= better_end_index, 'diverge index is out of range!'
            assert 0 <= diverge_index <= worse_end_index, 'diverge index is out of range!'
# better_seq_slice 和 worse_seq_slice 取从 diverge_index 开始到序列结束的部分(即模型生成的 token)。
            better_seq_slice = slice(diverge_index, better_end_index + 1)
            worse_seq_slice = slice(diverge_index, worse_end_index + 1)
            better_seq_length = better_end_index + 1
            worse_seq_length = worse_end_index + 1

            # size = ()
# better_log_prob: 计算 better 部分的总 log 概率。
# worse_log_prob: 计算 worse 部分的总 log 概率。
# 计算 对数比率(log ratio):
            better_log_prob = better_sequence_log_probs[i, better_seq_slice].sum(dim=-1)
            worse_log_prob = worse_sequence_log_probs[i, worse_seq_slice].sum(dim=-1)
            better_log_ratio = better_log_prob / better_seq_length
            worse_log_ratio = worse_log_prob / worse_seq_length
# 计算 ORPO 的 odds ratio loss:
            log_odds = (better_log_ratio - worse_log_ratio) - (
                torch.log1p(-torch.exp(better_log_ratio)) - torch.log1p(-torch.exp(worse_log_ratio))
            )
#  better 的 log 概率明显高于 worse,从而优化生成策略。
            odds_ratio_loss = -F.logsigmoid(log_odds)
# 最终损失
            sft_loss = -better_log_ratio
            losses.append(
                sft_loss + self.cfgs.train_cfgs.scale_coeff * odds_ratio_loss,
            )
            better_sample_rewards.append(
                self.cfgs.train_cfgs.scale_coeff * better_log_ratio.detach(),
            )
            worse_sample_rewards.append(self.cfgs.train_cfgs.scale_coeff * worse_log_ratio.detach())

        loss = torch.stack(losses).mean()  # size = ()
        better_sample_reward = torch.stack(better_sample_rewards)  # size = (B,)
        worse_sample_reward = torch.stack(worse_sample_rewards)  # size = (B,)
        reward = better_sample_reward + worse_sample_reward  # size = (B,)
        reward_accuracy = (better_sample_reward > worse_sample_reward).float().mean()  # size = ()
        reward_margin = better_sample_reward - worse_sample_reward  # size = (B,)

        return {
            'loss': loss,
            'reward': reward,
            'better_sample_reward': better_sample_reward,
            'worse_sample_reward': worse_sample_reward,
            'reward_accuracy': reward_accuracy,
            'reward_margin': reward_margin,
        }


def main():
    # setup distribution training
    deepspeed.init_distributed()
    current_device = get_current_device()
    torch.cuda.set_device(current_device)

    # read default configs from the yaml file
    task = os.path.join('text_to_text', 'orpo')
    dict_cfgs, ds_cfgs = read_cfgs(mode='train', task=task)

    # get custom configs from command line
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    _, unparsed_args = parser.parse_known_args()
    keys = [k[2:] for k in unparsed_args[1::2]]
    values = list(unparsed_args[2::2])
    unparsed_args = dict(zip(keys, values))
    for k, v in unparsed_args.items():
        dict_cfgs = update_dict(dict_cfgs, custom_cfgs_to_dict(k, v))

    # setup training
    cfgs = dict_to_namedtuple(dict_cfgs)
    seed_everything(cfgs.train_cfgs.seed)

    # finetune the model
    trainer = ORPOTrainer(cfgs=cfgs, ds_cfgs=ds_cfgs)
    trainer.train()
    trainer.save()

SimPO训练代码:

https://blog.csdn.net/weixin_43013480/article/details/141370269

# compute_log_probs 的作用是计算给定序列的 log 概率(对数概率),主要用于评估语言模型(LLM)的生成质量。

def compute_log_probs(
        self,
        model: AutoModelForCausalLM,
        batch: PreferenceBatch,
    ) -> torch.Tensor:
        """Compute log probabilities of given sequences."""
        logits = model(**self.infer_batch(batch)).logits
        device = logits.device
        input_ids = batch['input_ids']
        batch_size = len(batch['meta_info']['response_lens'])
        logprob_list = []
        for idx in range(batch_size):
            response_length = batch['meta_info']['response_lens'][idx]
            raw_input_id = strip_pad(input_ids[idx], self.tokenizer.pad_token_id)
            logit = logits[idx][-response_length:].unsqueeze(0)
            input_id = raw_input_id[-response_length:].unsqueeze(0)
            log_p = gather_log_probabilities(logit[:, :-1], input_id[:, 1:])
            logprob_list.append(log_p.squeeze(0))
        return torch.nn.utils.rnn.pad_sequence(
            logprob_list, batch_first=True, padding_value=0.0
        ).to(device)
class SimPOTrainer(DPOTrainer):

    def loss(  # pylint: disable=too-many-locals
        self,
        batch: PreferenceBatch,
    ) -> dict[str, torch.Tensor]:
        """Loss function for the SimPO algorithm."""
        sequence_log_probs = self.compute_log_probs(
            self.model.module,
            batch,
        )
# 使用 chunk 将 sequence_log_probs 按照第0维(批次维度)进行切分。将批次数据分为两部分:一部分对应 "更好"(better_sequence_log_probs),另一部分对应 "更差"(worse_sequence_log_probs)。每部分的大小为 (B, L - 1),B 是批次大小,L 是序列长度。  L-1 是为了删除最后的 <eos>
        (
            better_sequence_log_probs,  # size = (B, L - 1)
            worse_sequence_log_probs,  # size = (B, L - 1)
        ) = sequence_log_probs.chunk(chunks=2, dim=0)

        losses = []
        better_sample_rewards = []
        worse_sample_rewards = []

        better_input_ids, worse_input_ids = batch['input_ids'].chunk(chunks=2, dim=0)
        better_attention_mask, worse_attention_mask = batch['attention_mask'].chunk(chunks=2, dim=0)

        batch_size = better_input_ids.size(0)
        for i in range(batch_size):
#检查当前样本的 "更好" 和 "更差" 部分的 input_ids 是否相同。如果相同,跳过这个样本,因为它们对比不出差异。
            if torch.all(torch.eq(better_input_ids[i], worse_input_ids[i])).item():
                continue

#分别计算 "更好" 和 "更差" 样本的结束位置(通过 attention_mask 中的非零元素位置来确定)。
            better_end_index = better_attention_mask[i].nonzero()[-1].squeeze().item()
            worse_end_index = worse_attention_mask[i].nonzero()[-1].squeeze().item()
            better_input_length = better_end_index + 1
            worse_input_length = worse_end_index + 1
# diverge_index 是 "更好" 和 "更差" 样本之间的第一个差异位置。
            diverge_index = (
                (better_input_ids[i] != worse_input_ids[i]).nonzero()[0].squeeze().item()
            )
            assert 0 <= diverge_index <= better_end_index, 'diverge index is out of range!'
            assert 0 <= diverge_index <= worse_end_index, 'diverge index is out of range!'
#根据 diverge_index 进行切片,获取差异区域的对数概率。
#better_log_prob 和 worse_log_prob 是对应于 "更好" 和 "更差" 样本的对数概率的总和。
            better_seq_slice = slice(diverge_index, better_end_index + 1)
            worse_seq_slice = slice(diverge_index, worse_end_index + 1)
# 计算损失和奖励
            better_log_prob = better_sequence_log_probs[i, better_seq_slice].sum(dim=-1)
            worse_log_prob = worse_sequence_log_probs[i, worse_seq_slice].sum(dim=-1)
#在长度上归一化的奖励【/|y|】,其计算方式是使用策略模型的奖励中所有 token 的平均对数概率;
            better_log_ratio = better_log_prob / better_input_length
            worse_log_ratio = worse_log_prob / worse_input_length
#目标奖励差额γ,用以确保获胜和失败响应之间的奖励差超过这个差额γ
            losses.append(
                -F.logsigmoid(
                    self.cfgs.train_cfgs.scale_coeff * (better_log_ratio - worse_log_ratio)
                    - self.cfgs.train_cfgs.gamma,
                ),
            )
            better_sample_rewards.append(
                self.cfgs.train_cfgs.scale_coeff * better_log_ratio.detach(),
            )
            worse_sample_rewards.append(self.cfgs.train_cfgs.scale_coeff * worse_log_ratio.detach())
        loss = torch.stack(losses).mean()  # size = ()
        better_sample_reward = torch.stack(better_sample_rewards)  # size = (B,)
        worse_sample_reward = torch.stack(worse_sample_rewards)  # size = (B,)
        reward = better_sample_reward + worse_sample_reward  # size = (B,)
        reward_accuracy = (better_sample_reward > worse_sample_reward).float().mean()  # size = ()
        reward_margin = better_sample_reward - worse_sample_reward  # size = (B,)

        return {
            'loss': loss,
            'reward': reward,
            'better_sample_reward': better_sample_reward,
            'worse_sample_reward': worse_sample_reward,
            'reward_accuracy': reward_accuracy,
            'reward_margin': reward_margin,
        }

KTO训练代码:

#  创建 不匹配的提示-回答对:错位传入批次(batch)中的 answer_input_ids 和 answer_attention_mask 数据,以创建不匹配的提示-回答对。获取当前索引前一个样本作为回应(response)。如果当前索引是 0,则取最后一个样本作为回应。这是为了创建“不匹配”的数据对,即提示和回应不一定是成对的。
class UnmatchedSupervisedDataset(SupervisedDataset):

    def preprocess(
        self, raw_sample_for_prompt: dict[str, Any], raw_sample_for_response: dict[str, Any]
    ) -> SupervisedSample:
        return_dict = {}
        formatted_text, _ = self.template.format_unmatched_supervised_sample(
            raw_sample_for_prompt, raw_sample_for_response
        )
        return_dict['input_ids'] = self.tokenize(formatted_text)

        return return_dict

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        """Get a tokenized data sample by index."""
        raw_sample_for_prompt = self.raw_data[index]
        if index == 0:
            raw_sample_for_response = self.raw_data[-1]
        else:
            raw_sample_for_response = self.raw_data[index - 1]
        data = self.preprocess(raw_sample_for_prompt, raw_sample_for_response)
        return data

    def get_collator(self) -> Callable[[list[dict[str, torch.Tensor]]], dict[str, torch.Tensor]]:
        return UnmatchedSupervisedCollator(self.tokenizer.pad_token_id)

class KTOTrainer(DPOTrainer):

# 计算kl散度:通过计算当前模型(self.model.module)和参考模型(self.reference_model.module)之间的 KL 散度来比较它们的概率分布
# 选择最后一个 batch 的 KL 值可能只是实现上的简化。实际中,计算所有 batch 的 KL 散度并取平均,或者采取其他更复杂的策略,可能会增加额外的计算负担,而选择最后一个 batch 的 KL 值是一种更直接、简便的实现方式。
def compute_kl(self):
    random_dataset = UnmatchedSupervisedDataset(
        path=self.cfgs.data_cfgs.train_datasets,
        template=self.train_template,
        tokenizer=self.tokenizer,
        processor=self.processor,
        name=self.cfgs.data_cfgs.train_name,
        size=self.cfgs.data_cfgs.train_size,
        split=self.cfgs.data_cfgs.train_split,
        data_files=self.cfgs.data_cfgs.train_data_files,
        optional_args=self.cfgs.data_cfgs.train_optional_args,
    )
    seed = torch.randint(0, 100000, (1,)).item()
    torch.manual_seed(seed)
    self.random_dataloader = DataLoader(
        random_dataset,
        collate_fn=random_dataset.get_collator(),
        sampler=DistributedSampler(random_dataset, shuffle=True),
        batch_size=self.cfgs.train_cfgs.per_device_kl_batch_size,
    )
    for batch in self.random_dataloader:
        log_probs = self.compute_log_probs(  # size = (2 * B, L - 1)
            self.model.module,
            batch=batch,
        )
        ref_log_probs = self.compute_log_probs(  # size = (2 * B, L - 1)
            self.reference_model.module,
            batch=batch,
        )
        kl = (log_probs - ref_log_probs).mean()

        self.kl = max(kl, 0)
# 此方法是 DPO (Direct Preference Optimization) 算法的核心部分。它计算了在当前模型和参考模型之间的对比损失
    def loss(  # pylint: disable=too-many-locals
        self,
        batch: PreferenceBatch,
    ) -> dict[str, torch.Tensor]:
        """Loss function for the DPO algorithm."""
        sequence_log_probs = self.compute_log_probs(
            self.model.module,
            batch,
        )
        (
            better_sequence_log_probs,  # size = (B, L - 1)
            worse_sequence_log_probs,  # size = (B, L - 1)
        ) = sequence_log_probs.chunk(chunks=2, dim=0)

        with torch.no_grad():
            ref_sequence_log_probs = self.compute_log_probs(  # size = (2 * B, L - 1)
                self.reference_model.module,
                batch,
            )
            ref_better_sequence_log_probs, ref_worse_sequence_log_probs = (
                ref_sequence_log_probs.chunk(chunks=2, dim=0)
            )

        losses = []
        better_sample_rewards = []
        worse_sample_rewards = []

        better_input_ids, worse_input_ids = batch['input_ids'].chunk(chunks=2, dim=0)
        better_attention_mask, worse_attention_mask = batch['attention_mask'].chunk(chunks=2, dim=0)

        batch_size = better_input_ids.size(0)
        for i in range(batch_size):
            if torch.all(torch.eq(better_input_ids[i], worse_input_ids[i])).item():
                continue
            better_end_index = better_attention_mask[i].nonzero()[-1].squeeze().item()
            worse_end_index = worse_attention_mask[i].nonzero()[-1].squeeze().item()
            diverge_index = (
                (better_input_ids[i] != worse_input_ids[i]).nonzero()[0].squeeze().item()
            )
            assert 0 <= diverge_index <= better_end_index, 'diverge index is out of range!'
            assert 0 <= diverge_index <= worse_end_index, 'diverge index is out of range!'

            better_seq_slice = slice(diverge_index, better_end_index + 1)
            worse_seq_slice = slice(diverge_index, worse_end_index + 1)

            better_log_prob = better_sequence_log_probs[i, better_seq_slice].sum(dim=-1)
            worse_log_prob = worse_sequence_log_probs[i, worse_seq_slice].sum(dim=-1)
            ref_better_log_prob = ref_better_sequence_log_probs[i, better_seq_slice].sum(dim=-1)
            ref_worse_log_prob = ref_worse_sequence_log_probs[i, worse_seq_slice].sum(dim=-1)
            better_log_ratio = better_log_prob - ref_better_log_prob
            worse_log_ratio = worse_log_prob - ref_worse_log_prob

# 计算loss,kl值作为基准
            losses.append(
                self.cfgs.train_cfgs.scale_better
                * (1 - F.sigmoid(self.cfgs.train_cfgs.scale_coeff * (better_log_ratio - self.kl)))
                - self.cfgs.train_cfgs.scale_worse
                * (1 - F.sigmoid(self.cfgs.train_cfgs.scale_coeff * (self.kl - worse_log_ratio))),
            )
            better_sample_rewards.append(
                self.cfgs.train_cfgs.scale_coeff * better_log_ratio.detach(),
            )
            worse_sample_rewards.append(self.cfgs.train_cfgs.scale_coeff * worse_log_ratio.detach())
        loss = torch.stack(losses).mean()  # size = ()
        better_sample_reward = torch.stack(better_sample_rewards)  # size = (B,)
        worse_sample_reward = torch.stack(worse_sample_rewards)  # size = (B,)
        reward = better_sample_reward + worse_sample_reward  # size = (B,)
        reward_accuracy = (better_sample_reward > worse_sample_reward).float().mean()  # size = ()
        reward_margin = better_sample_reward - worse_sample_reward  # size = (B,)

        return {
            'loss': loss,
            'reward': reward,
            'better_sample_reward': better_sample_reward,
            'worse_sample_reward': worse_sample_reward,
            'reward_accuracy': reward_accuracy,
            'reward_margin': reward_margin,
        }
#执行训练步骤:这个方法在每一个训练步中计算并反向传播损失。它更新模型参数并计算并返回训练信息。
#奖励计算:通过 reward、better_sample_reward 和 worse_sample_reward 等指标来衡量模型的性能。
#全局平均:get_all_reduce_mean() 用于分布式训练,确保在多个设备上计算的值被平均,以保证训练的一致性。
def train_step(self, batch: PreferenceBatch) -> dict[str, Any]:
    """Perform a single training step for KTO."""
    loss_dict = self.loss(batch=batch)
    loss = loss_dict['loss']
    self.model.backward(loss)
    self.model.step()

    with torch.no_grad():
        reward = loss_dict['reward'].mean()
        better_sample_reward = loss_dict['better_sample_reward'].mean()
        worse_sample_reward = loss_dict['worse_sample_reward'].mean()
        reward_accuracy = loss_dict['reward_accuracy']
        reward_margin = loss_dict['reward_margin'].mean()

    loss = get_all_reduce_mean(loss)
    reward = get_all_reduce_mean(reward)
    better_sample_reward = get_all_reduce_mean(better_sample_reward)
    worse_sample_reward = get_all_reduce_mean(worse_sample_reward)
    reward_accuracy = get_all_reduce_mean(reward_accuracy)
    reward_margin = get_all_reduce_mean(reward_margin)

    return {
        'train/loss': loss.item(),
        'train/reward': reward.item(),
        'train/better_sample_reward': better_sample_reward.item(),
        'train/worse_sample_reward': worse_sample_reward.item(),
        'train/reward_accuracy': reward_accuracy.item(),
        'train/reward_margin': reward_margin.item(),
        'train/lr': self.model.optimizer.param_groups[0]['lr'],
    }

思考:

来源:https://wqw547243068.github.io/rlhf#%E6%80%9D%E8%80%83-1

0、KL惩罚

KL是放在奖励函数里面,还是放在外面

PPO 中的Rt计算:

的做法都能解释的通,其实实质其实是一个贝叶斯推断[介绍文章]。不过加入奖励函数里面控制粒度更细训练应该更加稳定。但是皓天大佬用REINFORCE+++复现时候加入KL约束会限制模型探索空间。皓天大佬的文章很有启发性,指出在base模型变强以后,其实现有的RL算法在规则奖励上应该都能work。更加应该探索如何基于强base模型来优化RL算法,不应拘泥于原来RL训练不稳定,难训练这种传统观念。

下面还是略微来解释一下KL的作用实质到底是什么: 一开始我们通过语料训练了一个预训练模型 πPT,然后有得到一个 πSFT模型,接下来通过RLHF我们要得到一个 πRLHF的模型。这其实是什么?是不断调整分布的过程,或者说是语言模型不断调整信念的过程。

符号说明D={(xi,yi)},其中 xi表示指令, yi是预训练的语言模型的输出。 π(y∣x)是从指令到输出的概率分布。

1、首先一开始有一个在大规模语料上训练的语言模型 π0(y∣x), 目前它表现欠佳,它的世界和人类的世界差别有点大,说起话来前言不搭后语。

2、好了现在有一个对话语料 D={(xi,yi)},这个对话预料的特点就是真实反应了人类世界的情况,或者说基于此我们能生成一个评分函数 r(x,y)这个函数能给语言模型基于指令 x生成的 y打分。既然如此不如这样思考

给定 x,y对评分 r附上一个信念或者概率 q,这里的 β是一个信念可调整的超参数。比较是个比较主观的东西,加个可调整参数来调节,以便让大多数人满意。

3、现在的问题就变成了如何根据初始模型 π0(y∣x)和人类的评分信念 q(r∣y,x)来调整模型参数 θ得到一个新的模型 πθ(y∣x)

如果说 π0(y∣x)是先验分布,那么人类的评分信念 q(r∣y,x)就是似然函数,于是我们可以构造一个后验分布:

现在我们是无法直接得到 πKL-RL(y∣x,r),计算证据Z(y,x,r)计算是巨大的。但是我们可以让一个分布接进它,或者最好的方式就是就地取材微调 π0(y∣x,θ)得到 πθ(y∣x)使得它接近 我们的后验πKL-RL(y∣x,r),这样我们就得到了对齐后的模型。我们自然就使用到了计算分布相似度的 KL散度。于是问题就变为了:

放在奖励函数里面,还是放在优势函数外面。不过是评分信念的不同,在token层级似然函数就是奖励,在优势函数外面就是优势函数。对贝叶斯更新的提供的信息不同、粒度与层次的不同。

1、RL 有用吗?

 Does Reinforcement Learning Really Incentivize Reasoning Capacity in LLMs Beyond the Base Model?

RL给不了新知识,只是激发了 Base Model 能力, 强化学习的边界被基座模型“锁死”。RL 只是让 BaseModel 朝更能给到正确答案的方向结题, 实际上,Base Model 不会的,可能永远不会, 会的,偶尔能做对, RL能增加这个做对的概率。

RLVR(可验证奖励的强化学习)在数学、代码、视觉推理等任务中表现突出,被视为提升大语言模型(LLM)推理能力的关键手段。

然而,核心问题始终存在:强化学习真能让大模型获得超越基座模型的新推理能力?

在数学、代码、视觉推理三大领域的系统性实验发现:

  • 能力边界未突破:RLVR模型所有推理路径均已存在于基础模型中,强化学习并未赋予模型新的推理能力。
  • 采样效率与覆盖能力权衡:RLVR 小采样次数下表现优于基座模型,但随着采样次数增加,基座模型逐渐追平并反超,显示出更广泛的覆盖能力。
  • 答案同源性:RLVR模型正确答案来自基座模型的输出分布,强化学习只是通过调整概率分布筛选高奖励路径。

显示:

  • 数学推理任务中,基座模型在多次采样后的能力表现逐渐追平并反超RL模型。
  • 代码生成任务中,RL模型提升了单样本准确率,但在更高采样次数下,基座模型仍展现出更强的覆盖能力。
  • 视觉推理任务中,RL训练后的模型在单次回答准确率上提升显著,但基座模型在多次采样后仍表现出更广泛的问题覆盖能力

RLVR只是让模型更偏向高奖励解决方案,而非创造新的推理能力。

对奖励路径的聚焦,削弱了模型的探索能力,限制了大规模采样时对可解问题的覆盖范围。

Key Insights

  1. 尽管RL训练后的模型在 pass@k(k=1) 情况下超越 Base模型, 但是 BaseModel 在k值不做限制情况下, 可能比RL后的模型pass率还高;
  2. RL 只是优化了 Base Model 采样效率, 一方面增加了Base Model一次就能做对题的概率,但同时限制了模型的探索能力,导致了在增加pass@k的k时候, Base Model 做对题的概率反而增加了;
  3. CoT 方法对模型 Finetune 更能激发模型的做题能力

对比 CoT对模型进行Finetune

  1. 从R1蒸馏的数据对模型直接进行CoT Finetune,在同样多次Sample看pass结果上, CoT 确实是在 Base Model上足量提升,超越 Base Model,并比RL的结果更好。但这个图里面奇怪的是Instruct的模型甚至没有Base版本在AIME24的表现上好?
  2. 不同RL算法整体差异并不大。不同的RL算法,比如PPO,DAPO,GRPO等

思考

为什么 AlphaGO 和玩游戏, RL能发掘新的胜利模式, 而 LLM 中的RL不行?

  1. LLM 输出token概率空间比游戏概率空间大很多, 因此,RL优化LLM 更难, 并且Reasoning 经常是从Pretrain Model开始训练, 而Pretrain模型本身受限制于预训练的语料,训练游戏的一般都是随机初始化,导致可能Pretrain模型本身就不包含所有能解决问题的先验(比如一个问题永远答不对,Reward永远是0),而随机初始化的可能本身就存在可能为1的情况,RL才有可能找到正确答案。
  2. Pretrain 模型的先验知识限制太强, 导致模型探索说话空间时,会因为错误格式或者语句不通被干掉, 即使有可能导出正确答案,也会因为中间步骤产生问题,而永远失败;
  3. RL算法设计机制潜在限制了模型探索正确答案的可能性, 比如 PPO算法中的KL Divergence约束了模型前后概率分布不能差别过大。

很多人认为,强化学习(RL)能够泛化到不同任务中,监督微调(SFT)可以记忆知识点,另外,还有类似 R1-zero 的结论等。

而如今整体开源社区的探索已经深入了许多。

  • 解题方面,构建出色的基础模型(Base Model)。
  • 实际上,很多基础模型都已经过指令微调,只是没有经过复杂指令微调,所以很难简单地将其认定为一个单纯的预训练(Pretrain)模型,毕竟预训练和监督微调的学习模式基本相同。在这个基础模型之上进行强化学习(RL)操作,能够提升它解决某类问题的能力。
  • 思维链(CoT)本质是什么,为什么能提高答案的准确性?思维链本质上就是 “大声思考”(Thinkout loud)。
    • 对于有明确答案的问题,通过思维链来检查其解决问题的步骤是否错误,以及结果是否正确;
    • 对于开放性问题而言,思维链增加回答的可信度。
    • 思维链本质是结构化思考,言之有理即可,模型可解释性的另一种体现。

因此,思维链方向可以继续鉴定的走下去,同时,模型本身的限制应该更少些,比如乱码没关系,预留更多探索空间

2、RL 核心在于 奖励函数

【2025-5-5】忽略强化学习算法细节,在reward上做点手脚,简单又重要

(1) rl 与 reward

RLHF 精髓: 将人类偏好转化为可量化的奖励信号。

奖励函数告诉模型”什么是好的输出”,而rl算法只是将这种反馈训练到模型参数中去。reward与构建高质量数据,对于rl最终的结果来说同样重要。

deepseek-r1的grpo则是针对数学和代码任务设计了规则判别的奖励函数

(2) 奖励函数构建策略

reward 构造策略

  • 任务相关性:奖励信号与任务目标相关。
    • 数学问题关注正确性,写作注重多样性,销售助手需要情商
  • 可量化:可量化的指标才可以交给rl进行训练学习。
    • 答案对错由规则判断给0,1布尔值。这个回答很好由reward model转化为0~1.0之间的得分
  • 相对性:PPO中给的是某个答案的绝对奖励值(通过pairwise 方式训练 reward model),DPO中则是构建答案间的相对偏好关系,GRPO 计算一批样本的相对奖励优势
  • 推理过程:对整个过程给一个最终奖励,还是每个推理步骤评估,以及是否需要推理过程,都可以设置为奖励信号

3、方法选择

何时使用 DPO 与 PPO 与 GRPO ?

偏好对齐组合中加入 GRPO 后,有几项决策因素需要考虑:

  1. 数据可用性(是否有偏好数据):DPO 使用偏好数据(选择/拒绝的答案),而 PPO 则需要先用这种偏好数据训练一个奖励模型。GRPO 则更具灵活性,因为它可以使用偏好数据,但并非必须使用。
  2. 奖励模型:DPO 通过直接基于偏好进行优化,将问题构建成分类问题,从而消除了对单独奖励模型的需求。相比之下,PPO 则需要训练和维护一个单独的奖励模型,这增加了复杂性。GRPO 则处于两者之间,既支持使用显式的奖励模型(如 PPO),也支持直接使用奖励函数。
  3. 计算资源:DPO 最高效,因为无需添加奖励模型。PPO 计算需求最高,因为它需要多个模型。GRPO 由于采用了基于组的方法,所以所需的资源适中。
对比项DPOPPOGRPO
数据可用性偏好数据(选择/拒绝的答案)先用偏好数据训练奖励模型更灵活,可用偏好数据,但并非必须
奖励模型直接基于偏好进行优化,将问题构建成分类问题,消除对单独奖励模型的需求训练和维护单独的奖励模型,增加了复杂性既支持使用显式的奖励模型(如 PPO),也支持直接使用奖励函数
计算资源最高效,无需添加奖励模型计算需求最高,需要多个模型由于采用基于组的方法,所需资源适中

要点

  • 当拥有高质量的偏好数据且计算资源有限时,选择 DPO。
  • 当需要精细控制、拥有充足的计算资源并且能够投入精力进行仔细调整时,选择 PPO。
  • 当想要整合多个奖励信号,或者没有全面的偏好数据时,选择 GRPO。

loss

为什么不用 梯度下降 ?

RLHF 为什么不直接对 loss 进行梯度下降来求解?

核心原因:

  • loss 或优化目标不可微,看一下优化目标的红色框部分:

损失函数表达式中的 y 是采样出来的, Dy~pi(y|x) , 可能是 greedy,beam search 等,在词表上进行采样或选择,而不是产生连续的、可微分的输出。所以,没法直接使用梯度下降,而是用 PPO 等策略梯度来求解。

RLHF 问题

【2025-2-6】Andrej Karpathy 最新视频盛赞 DeepSeek:R1 正在发现人类思考的逻辑并进行复现

  • 视频链接:youtube
  • DeepSeek R1 在性能方面与 OpenAI 模型不相上下,推动了 RL 技术的发展

如果只是模仿人类玩家,就永远无法超越极限。

强化学习的优势

  • 不受人类表现的限制。
    • 围棋游戏中,强化学习会自己与自己对弈,通过试错来学习哪些走法能赢得比赛。最终使AlphaGo能够超越人类顶尖棋手,甚至发明了一些人类棋手从未想到过的创新走法
    • AlphaGo 对弈中,实际上下了一步人类专家通常不会下的棋。评估来看,这步棋被人类玩家下的概率大约是1/10,000。

所有问题都属于可验证领域。任何时候都可以很容易地与一个具体答案进行比较评分。

基本思路:

训练人类的模拟器,并通过强化学习对这些模拟器进行优化

人类反馈中进行强化学习的优势

  • 能在任意领域进行强化学习,包括无法验证的领域。
    • 例如,像摘要生成、写诗、编笑话或任何其他创意写作
  • RLHF 却绕过了这个问题,不直接生成,而是排序

判别器和生成器之间的差距有关:对于人类来说,判别比生成要容易得多

RLHF显著缺点

  • 强化学习不是基于实际的人类判断,而是基于人类的一个有损模拟,可能会产生误导
  • 强化学习擅长“欺骗”模型,误导其做出许多错误的决定。

奖励模型

ppo 中 RM 如何工作

PPO 为啥不直接用 Reward Model

RLHF中,为什么 PPO 需要 Critic模型 而不是直接使用 Reward Model ?

强化学习中,PPO(Proximal Policy Optimization)基于策略梯度训练强化学习智能体。

PPO算法中引入Critic模型的主要目的:提供价值估计器,用于评估状态状态动作对的价值,从而辅助策略的更新和优化。

虽然奖励模型(Reward Model)可以提供每个状态或状态动作对的即时奖励信号,但它并不能直接提供对应的价值估计

  • 奖励信号只反映了当前动作的即时反馈,而并没有提供关于在长期时间尺度上的价值信息。

Critic模型估计状态或状态动作对的长期价值,也称为状态值函数动作值函数。Critic模型能学习和预测在当前状态下采取不同动作所获得的累积奖励,它提供了对策略改进的指导。

PPO算法使用Critic模型的估计值来计算优势函数,从而调整策略的更新幅度,使得更有利于产生更高长期回报的动作被选择。

另外,Critic模型还可用于评估不同策略的性能,为模型的评估和选择提供依据。PPO算法中的Actor-Critic架构允许智能体同时学习策略价值函数,并通过协同训练来提高性能。

因此,在 RLHF(Reinforcement Learning from Human Feedback)中,PPO算法需要Critic模型而不是直接使用奖励模型,是为了提供对状态或状态动作对的价值估计,并支持策略的改进和优化。Critic模型的引入可以提供更全面和准确的信息,从而增强算法的训练效果和学习能力。

即时奖励和长期奖励

即时奖励 与 状态动作对的长期价值 的差别是什么?

即时奖励(Immediate Reward)和状态动作对的长期价值(Long-Term Value)代表了强化学习中不同的概念和时间尺度。

  • 即时奖励是指智能体在执行某个动作后立即获得的反馈信号。由环境提供,用于表示当前动作的好坏程度。即时奖励是一种即时反馈,可以指示当前动作的立即结果是否符合智能体的目标。
  • 而状态动作对的长期价值涉及更长时间尺度上的评估,考虑了智能体在当前状态下选择不同动作所导致的未来回报的累积。长期价值可以表示为状态值函数(State Value Function)或动作值函数(Action Value Function)。
    • 状态值函数(V-function)表示在给定状态下,智能体从该状态开始执行一系列动作,然后按照某个策略进行决策,从而获得的预期累积回报。状态值函数估计了智能体处于某个状态时所能获得的长期价值,反映了状态的优劣程度。
    • 动作值函数(Q-function)则表示在给定状态下,智能体选择某个动作后,按照某个策略进行决策,从该状态转移到下一个状态并获得预期累积回报的价值。动作值函数估计了在给定状态下采取不同动作的长期价值,可以帮助智能体选择在每个状态下最优的动作。

长期价值考虑了智能体在未来决策过程中所能获得的累积回报

  • 相比之下,即时奖励只提供了当前动作的即时反馈。
  • 长期价值对智能体的决策具有更全面的影响,可以帮助智能体更好地评估当前状态和动作的长期效果,并指导智能体在长期时间尺度上作出更优的决策。

在强化学习中,长期价值的估计对于确定性策略选择和价值优化非常重要,而即时奖励则提供了对当前动作的直接反馈。这两者相互补充,结合起来可以帮助智能体实现更好的决策和学习效果。

PPO 优势函数

PPO 中优势函数指什么

在 Proximal Policy Optimization(PPO)算法中,优势函数(Advantage Function)用于评估状态-动作对的相对优劣程度。它衡量了执行某个动作相对于平均水平的优劣,即在给定状态下采取某个动作相对于采取平均动作的效果。

优势函数定义:

Advantage(s, a) = Q(s, a) – V(s)

其中

  • Advantage(s, a) 表示在状态 s 下采取动作 a 的优势函数值
  • Q(s, a) 表示状态动作对 (s, a) 的动作值函数(也称为动作优势函数)
  • V(s) 表示状态值函数。

优势函数的作用在于帮助评估当前动作的相对价值,以便在策略更新过程中确定应采取的动作。通过比较不同动作的优势函数值,可以决定哪些动作是更好的选择。正的优势函数值表示执行的动作比平均水平更好,而负的优势函数值表示执行的动作比平均水平更差。

在PPO算法中,优势函数用于计算策略更新的目标,以便调整策略概率分布来提高优势函数为正的动作的概率,并降低优势函数为负的动作的概率,从而改进策略的性能。

相关论文:

  • KTO,Kahneman-Tversky 优化,参阅论文《KTO: Model alignment as prospect theoretic optimization》。
  • DRO,直接奖励优化,参阅论文《Offline regularised reinforcement learning for large language models alignment》。
  • SimPO,简单偏好优化,参阅论文《SimPO: Simple preference optimization with a reference-free reward》