Jiahong 的个人博客

凡事预则立,不预则废


  • Home

  • Tags

  • Archives

  • Navigation

  • Search

RS——OneRec

OneRec:Unifying Retrieve and Rank with Generative Recommender and Preference Alignment,本文主要是对论文OneRec的翻译,其中包含了一些总结和理解

  • 参考链接:
    • OneRec: Unifying Retrieve and Rank with Generative Recommender and Preference Alignment, 202502, KuaiShou:快手周国睿团队作品

整体思路说明

  • 基于生成式检索的推荐系统介绍 :基于生成式检索(召回)的推荐系统(generative retrieval-based recommendation systems,GRs),通过自回归方式直接生成候选视频(注:论文针对短视频推荐场景,所以使用的是视频作为例子,实际上可以是任意 item),展现出巨大潜力
  • 当前问题 :然而,现代推荐系统大多采用“检索-排序”(retrieve-and-rank)策略,生成模型仅在检索阶段作为选择器发挥作用
  • 论文方案OneRec :用统一的生成模型取代传统的级联学习框架。具体而言,OneRec 包含以下核心组件:
    • 1)Encoder-Decoder 结构 :编码用户历史行为序列,逐步解码用户可能感兴趣的视频。采用稀疏专家混合(MoE)技术,在不显著增加计算量的情况下扩展模型容量;
    • 2)session 式生成方法 :与传统 next-item 预测不同,提出 session-wise 生成方式,相比依赖手工规则拼接生成结果的 point-by-point 生成更优雅且上下文连贯;
    • 3)迭代偏好对齐模块 :结合直接偏好优化(Direct Preference Optimization,DPO)提升生成结果质量。与 NLP 中的DPO不同,推荐系统通常只有一次展示机会,无法同时获取正负样本。为此,论文设计奖励模型模拟用户生成行为,并根据在线学习特性定制采样策略,大量实验表明,少量DPO样本即可对齐用户兴趣偏好,显著提升生成质量
  • OneRec已在快手主场景(日活数亿的短视频推荐平台)部署 ,观看时长提升1.6% ,实现了实质性突破
  • 注:据论文所知,这是首个在真实场景中显著超越当前复杂且精心设计的推荐系统的端到端生成模型

更多相关讨论

  • 推荐系统介绍 :为了平衡效率和有效性,大多数现代推荐系统采用级联排序策略。如图1(b)所示,一个典型的级联排序系统采用三阶段流程:召回、预排序和排序。每个阶段负责从接收到的 item 中选择前k个 item,并将结果传递到下一个阶段,共同平衡系统响应时间和排序准确性之间的权衡。尽管在实践中效率较高,但现有方法通常独立对待每个排序器,每个独立阶段的有效性成为后续排序阶段的上限,从而限制了整体排序系统的性能。尽管有各种努力通过实现排序器之间的交互来提高整体推荐性能,但它们仍然保持传统的级联排序范式
  • 最近,基于生成式检索的推荐系统(GRs)[36, 45, 51]作为一种有前途的范式出现,它通过自回归序列生成方式直接生成候选 item 的标识符。通过使用对 item 语义进行编码的量化语义ID对 item 进行索引[24],推荐系统可以利用 item 内丰富的语义信息。GRs的生成特性使其适合通过 beam search 解码直接选择候选 item,并产生更多样化的推荐结果。然而,目前的生成模型仅在检索阶段充当选择器,因为其推荐准确性仍无法与精心设计的多级联排序器相媲美
  • 为了应对上述挑战,论文提出了一种用于single-stage 推荐的统一端到端生成框架OneRec
    • 首先 ,论文提出了一种 Encoder-Decoder 架构。从训练 LLM 时观察到的缩放定律中获得灵感,论文发现扩展推荐模型的容量也能持续提高性能。因此,论文基于MoE 结构扩展模型参数,这显著提高了模型表征用户兴趣的能力
    • 其次 ,与传统的 point-by-point 预测 next-item 不同,论文提出了一种 session-wise 的列表生成方法,该方法考虑了每个 session 中 item 的相对内容和顺序。逐点生成方法需要 hand-craft 策略来确保生成结果的连贯性和多样性(coherence and diversity)。相比之下,session-wise 的学习过程使模型能够通过输入偏好数据自主学习最优的 session 结构
    • 最后 ,论文探索使用直接偏好优化(DPO)进行偏好学习 ,以进一步提高生成结果的质量。为了构建偏好对,论文从 hard negative sampling[37]中获得灵感,从 beam search 结果中创建 self-hard rejected samples(而不是随机采样)。论文提出了一种迭代偏好对齐(IPA)策略 ,根据预训练奖励模型(RM)提供的分数对采样响应进行排序,确定最佳选择和最差拒绝(best-chosen and worst-rejected)的样本
  • 论文在大规模工业数据集上的实验展示了所提方法的优越性。论文还进行了一系列消融实验,详细证明了每个模块的有效性
  • 这项工作的主要贡献总结如下:
    • OneRec框架提出 :为了克服级联排序的局限性,论文引入了 OneRec,这是一种 single-stage 生成式推荐框架。据论文所知,这是首批能够使用统一生成模型处理 item 推荐的工业解决方案之一,显著超越了传统的 multi-stage 排序流程
    • session-wise 的生成方式 :论文通过 session-wise 的生成方式强调了模型容量和目标 item 上下文信息的必要性,这使得预测更加准确,并提高了生成 item 的多样性
    • 迭代偏好对齐(IPA)策略 :论文提出了一种基于个性化奖励模型的新型 self-hard negative samples 选择策略。通过直接偏好优化,论文增强了 OneRec 在更广泛用户偏好范围内的泛化能力
    • 大量的离线实验和在线A/B测试证明了它们的有效性和效率

OneRec框架整体方法

  • 整体框架如下:

一些基本定义

  • 用户特征 :输入为用户正向历史行为序列 \(\mathbf{\mathcal{H}}_u=\{\boldsymbol{v}^h_1,\boldsymbol{v}^h_2,…,\boldsymbol{v}^h_n\}\),其中 \(\boldsymbol{v}\) 表示用户实际观看或交互(likes/follows/shares,点赞/关注/分享)的视频,\(n\)为序列长度
  • 输出 :视频列表构成的 session \(\mathcal{S}=\{\boldsymbol{v}_1,\boldsymbol{v}_2,…,\boldsymbol{v}_m\}\),\(m\) 为 session 内视频数
  • 视频信息 :每个视频 \(\boldsymbol{v}_i\) 通过多模态嵌入 \(\boldsymbol{e}_i\in\mathbb{R}^d\) 描述,其分布与真实 user-item 行为对齐。基于预训练多模态表示,现有生成推荐框架[25,36]使用RQ-VAE[49]将嵌入编码为语义标记。但该方法因 code 分布不平衡(hourglass phenomenon,“沙漏现象”[23],表现为第二层编码过于集中)存在缺陷。论文采用多层级平衡量化机制(multi-level balanced quantitative mechanism),通过残差K-Means量化算法[27]转换 \(\boldsymbol{e}_i\):
    • 第一层(\(l=1\))初始残差\(r^1_i=\boldsymbol{e}_i\)
    • 每层 \(l\) 维护 codebook \(C_l=\{c^l_1,…,c^l_K\}\)(\(K\) 为 codebook 大小)
    • 通过 \(\boldsymbol{s}^l_i=\arg\min_k|r^l_i-c^l_k|^2_2\) 生成最近中心节点索引,下一层残差 \(r^{l+1}_i=r^l_i-c^l_{s^l_i}\),即通过层级索引生成对应的 codebook tokens:
      $$
      \boldsymbol{s}^1_i = \arg\min_k|r^1_i-c^1_k|^2_2, \quad \boldsymbol{r}^2_i=r^1_i-c^1_{s^1_i} \\
      \boldsymbol{s}^2_i = \arg\min_k|\boldsymbol{r}^2_i-c^2_k|^2_2, \quad \boldsymbol{r}^3_i=r^2_i-c^2_{s^2_i} \\
      \vdots \\
      \boldsymbol{s}^L_i = \arg\min_k|\boldsymbol{r}^L_i-c^L_k|^2_2
      $$
      • 其中 \(L\) 为语义ID总层数
  • 为构建平衡 codebook \(C_l=\{c^l_1,…,c^l_K\}\),采用算法1的平衡K-means进行 item 集划分。给定视频集 \(\mathcal{V}\),该算法将其划分为 \(K\) 个簇,每簇包含恰好 \(w=|\mathcal{V}|/K\) 个视频。迭代计算中,每个中心点依次分配 \(\mathbf{w}\) 个最近未分配视频(基于欧氏距离),随后用分配视频的均值向量重新校准中心点,直至簇分配收敛
  • 算法1思路:
    • 算法1中仅展示了第 \(l\) 层 codebook 的迭代过程
    • 首先随机初始化 \(K\) 个中心点
    • 内部循环 \(K\) 次,每一次中,使用最接近当前 \(c_k^l\) 的 \(w\) 个当前层的残差 \(\mathbf{r}^l\) 的均值来编码当前中心点 \(c_k^l\),同时移除已经参与均值计算的残差(确保每个中心点都刚好被K个残差)
    • 整体循环以上步骤,直到算法收敛

session-wise 的列表生成

  • 与传统仅预测 next video 的 point-wise 推荐方法不同,session-wise 的生成旨在根据用户的历史交互序列生成高价值的 session 列表,这使得推荐模型能够捕捉推荐列表中视频之间的依赖关系。具体而言,session 是指响应用户请求返回的一批短视频,通常包含5到10个视频。session 中的视频通常会综合考虑用户兴趣、连贯性和多样性等因素。论文设计了几项标准来识别高质量 session,包括:
    • 用户在一个 session 中实际观看的短视频数量大于或等于5;
    • 用户观看该 session 的总时长超过某个阈值;
    • 用户表现出互动行为,例如点赞、收藏或分享视频;
  • 这种方法确保论文的 session 模型能够从真实的用户参与模式中学习,并更准确地捕捉 session 列表中的上下文信息。因此,论文的 session 模型\(\mathcal{M}\)的目标可以形式化为:
    $$
    \mathcal{S} := \mathcal{M}(\mathcal{H}_u)
    $$
    • 模型输入 :\(\mathcal{H}_u\) 通过语义ID表示( 长度为 \(n\) ):
      $$
      \mathcal{H}_u = \{(s^1_1, s^2_1, \cdots, s^L_1), (s^1_2, s^2_2, \cdots, s^L_2), \cdots, (s^1_n, s^2_n, \cdots, s^L_n)\}
      $$
    • 模型输出 :\(\mathcal{S}\) 也通过语义ID表示( 长度为 \(m\) ):
      $$\mathcal{S} = \{(s^1_1, s^2_1, \cdots, s^L_1), (s^1_2, s^2_2, \cdots, s^L_2), \cdots, (s^1_m, s^2_m, \cdots, s^L_m)\}$$
  • 如 图2(a) 所示,与T5架构一致,论文的模型采用基于Transformer的框架,主要由两部分组成:用于建模用户历史交互的 encoder 和用于 session 列表生成的 decoder。具体来说,encoder 通过堆叠的(stacked)多头自注意力机制(multi-head self-attention,MHA)和前馈层(feed-forward layers)处理输入序列\(\mathcal{H}_u\)。论文将编码后的历史交互特征表示为\(\boldsymbol{H} = \text{Encoder}(\mathcal{H}_u)\)
  • decoder 以目标 session 的语义ID为输入,并以自回归方式生成目标。为了以合理的经济成本训练更大的模型,论文在 decoder 的前馈神经网络(feed-forward neural networks,FNNs)中采用了Transformer语言模型中常用的MoE 架构 ,并将第\(l\)层的 FNN 替换为:
    $$
    \begin{split}
    \mathbf{H}^{l+1}_t &= \sum_{i=1}^{N_{\text{MoE} } } \left(g_{i,t} \text{FFN}_i \left(\color{red}{\mathbf{H}^l_t} \right)\right) + \color{red}{\mathbf{H}^l_t}, \\
    g_{i,t} &= \begin{cases}
    \mathcal{S}_{i,t}, & \mathcal{S}_{i,t} \in \text{TopK}(\{\mathcal{S}_{j,t}|1 \leq j \leq N\}, K_{\text{MoE} }), \\
    0, & \text{otherwise},
    \end{cases} \\
    \mathcal{S}_{i,t} &= \text{Softmax}_i \left({\mathbf{H}^{l}}^{\color{red}{T}}_t e^l_i \right),
    \end{split}
    $$
    • 其中,\(N_{\text{MoE} }\)表示专家总数,\(\text{FFN}_i(\cdot)\)是第\(i\)个专家的 FNN,\(g_{i,t}\)表示第\(i\)个专家的门控值。门控值\(g_{i,t}\)是稀疏的,即只有\(K_{\text{MoE} }\)个门控值为非零。这种稀疏性确保了MoE层的计算效率,每个 token 仅被分配到\(K_{\text{MoE} }\)个专家中进行计算
      • 问题:和常规的MoE一样,由于先做了 softmax 再取的 TopK 系数,所有系数和小于1,这里不需要保证系数和为1吗?
      • 回答:不需要,因为在每一个Transformer层中,FNN的结果和上一层的隐向量叠加后,都有LayerNorm存在(将每个token的隐向量分别归一化为均值为0,方差为1的向量),所以隐向量的值不会越来越小,在DeepSeek V3版本的公式中,会考虑在选择了TopK FNN后,再进行一次归一化(虽然不是softmax)
    • 注意 :\({\mathbf{H}^{l}}^{\color{red}{T}}_t e^l_i\) 中的 \(\color{red}{T}\) 是矩阵转置的含义
    • 修正 :这里的公式中表达有误,红色的 \(\mathbf{H}^l_t\) 在进入下一个 MoE(FNN)时,是需要整体经过 Self-Attention 编码的,即在更新前加一句 \(\color{blue}{\mathbf{H}^l_{1:T} = \text{Self-Att}(\mathbf{H}^l_{1:T}) + \mathbf{H}^l_{1:T}}\),详情见DeepSeekMoE的技术报告第2节(P4):DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
  • 在训练过程中,论文在编码(codes)的开头添加起始标记\(\boldsymbol{s}_{\text{[BOS]} }\)来构建 decoder 输入:
    $$
    \begin{split}
    \bar{\boldsymbol{\mathcal{S}} } &= \{\boldsymbol{s}_{\text{[BOS]} }, \boldsymbol{s}^1_1, \boldsymbol{s}^2_1, \cdots, \boldsymbol{s}^L_1, \boldsymbol{s}_{\text{[BOS]} }, \boldsymbol{s}^1_2, \boldsymbol{s}^2_2, \cdots, \boldsymbol{s}^L_2, \cdots, \boldsymbol{s}_{\text{[BOS]} }, \boldsymbol{s}^1_m, \boldsymbol{s}^2_m, \cdots, \boldsymbol{s}^L_m\}
    \end{split}
    $$
  • 论文使用交叉熵损失对目标 session 的语义ID进行 next token 预测。NTP(Next Token Prediction)损失\(\mathcal{L}_{\text{NTP} }\)定义为:
    $$
    \mathcal{L}_{\text{NTP} } = -\sum_{i=1}^m \sum_{j=1}^L \log P(\boldsymbol{s}^{j+1}_i \mid \left[\boldsymbol{s}_{\text{[BOS]} }, \boldsymbol{s}^1_1, \boldsymbol{s}^2_1, \cdots, \boldsymbol{s}^L_1, \cdots, \boldsymbol{s}_{\text{[BOS]} }, \boldsymbol{s}^1_i, \cdots, \boldsymbol{s}^j_i\right]; \Theta).
    $$
  • 在对 session 列表生成任务进行一定量的训练后,论文得到 seed 模型\(\mathcal{M}_t\)

基于奖励模型的迭代偏好对齐

  • 上一小节总结 :上一小节中定义的高质量 session 提供了有价值的训练数据,使模型能够了解什么构成一个好的 session,从而确保生成视频的质量
  • 本节内容 :在此基础上,论文旨在通过直接偏好优化(DPO)进一步提升模型的能力。在传统的 NLP 场景中,偏好数据是由人类明确标注的。然而,由于 user-item 交互数据的稀疏性,推荐系统中的偏好学习面临着独特的挑战,这就需要一个奖励模型(RM)。因此,论文引入了一个 session-wise 的奖励模型
  • 其他内容 :此外,论文通过提出一种迭代直接偏好优化方法来改进传统的DPO,使模型能够自我提升,这将在后续章节中进行描述
奖励模型训练
  • 论文用 \(R(\mathbf{u}, \mathcal{S})\) 表示奖励模型,它为不同的用户选择偏好数据。这里,输出 \(r\) 表示用户 \(\mathbf{u}\)(通常由用户行为表示)对 session \(\mathcal{S} = \{\mathbf{v}_1, \mathbf{v}_2, …, \mathbf{v}_m\}\) 的偏好所对应的奖励。为了使奖励模型具备对 session 进行排序的能力,论文首先提取 session \(\mathcal{S}\) 中每个 item \(\mathbf{v}_i\) 的目标感知表示 \(\mathbf{e}_i = \mathbf{v}_i \odot \mathbf{u}\),其中 \(\odot\) 表示目标感知操作(例如针对用户行为的目标注意力)。这样论文就得到了 session \(\mathcal{S}\) 的目标感知表示 \(\mathbf{h} = \{\mathbf{e}_1, \mathbf{e}_2, \cdots, \mathbf{e}_m\}\)。然后,session 中的 item 通过自注意力层相互作用,以融合不同 item 之间的必要信息:
    $$\mathbf{h}_f = \text{SelfAttention}(\mathbf{hW}_{\mathcal{s}}^Q, \mathbf{hW}_{\mathcal{s}}^K, \mathbf{hW}_{\mathcal{s}}^V)$$
    • 问题:\(\mathbf{W}_{\mathcal{s}}^Q\) 中的下标 \(\mathcal{s}\) 是表达什么?为不同的 Session 做了不同的 Attention 权重吗?或是在强调这里是 Session 粒度的 Attention?
  • 接下来,论文利用不同的塔对多目标奖励进行预测,并且使用丰富的推荐数据对奖励模型进行预训练:
    $$\begin{aligned}
    \hat{r}^{swt} = \text{Tower}^{swt}(\text{Sum}(\mathbf{h}_f)), &\hat{r}^{vtr} = \text{Tower}^{vtr}(\text{Sum}(\mathbf{h}_f)), \\
    \hat{r}^{wtr} = \text{Tower}^{wtr}(\text{Sum}(\mathbf{h}_f)), &\hat{r}^{ltr} = \text{Tower}^{ltr}(\text{Sum}(\mathbf{h}_f)), \\
    \text{where } \quad \text{Tower}(\cdot) &= \text{Sigmoid}(MLP(\cdot))
    \end{aligned}$$
    • 问题:swt,wtr,vtr,ltr等分别表示什么?
    • 回答:论文计算不同目标指标的平均奖励,包括 session 观看时间(session watch time,swt) 、观看概率(view probability,vtr)、关注概率(follow probability,wtr)和点赞概率(like probability,ltr),swt和vtr是观看时间指标 ,而wtr和ltr是交互指标
  • 在得到每个 session 的所有估计奖励 \(\hat{r}^{swt}, …\) 和真实标签 \(y^{swt}, …\) 之后,论文直接最小化二元交叉熵损失来训练奖励模型。损失函数 \(\mathcal{L}_{RM}\) 定义如下:
    $$\mathcal{L}_{RM} = -\sum_{swt, …}^{xtr} (y^{xtr} \log(\hat{r}^{xtr}) + (1 - y^{xtr})\log(1 - \hat{r}^{xtr}))$$
迭代偏好对齐(Iterative Preference Alignment,IPA)
  • 基于预训练的奖励模型 \(R(\mathbf{u}, \mathcal{S})\) 和当前的OneRec模型 \(\mathcal{M}_t\),论文通过 beam search 为每个用户生成 \(N\) 个不同的响应:
    $$\mathcal{S}_u^n \sim \mathcal{M}_t(\mathcal{H}_u) \quad \text{for all } u \in \mathcal{U} \text{ and } n \in [N]$$

    • 其中,论文用 \([N]\) 表示 \(\{1, 2, …, N\}\)
  • 然后,论文基于奖励模型 \(R(\mathbf{u}, \mathcal{S})\) 为这些响应中的每一个计算奖励 \(r_u^n\):
    $$r_u^n = R(\mathbf{u}, \mathcal{S}_u^n)$$

  • 接下来,论文通过选择奖励值最高的获胜响应 \((\mathcal{S}_u^w, \mathcal{H}_u)\) 和奖励值最低的失败响应 \((\mathcal{S}_u^l, \mathcal{H}_u)\) 来构建偏好对 \(D_t^{pairs} = (\mathcal{S}_u^w, \mathcal{S}_u^l, \mathcal{H}_u)\)。有了这些偏好对,论文现在可以训练一个新的模型 \(\mathcal{M}_{t+1}\),它从模型 \(\mathcal{M}_t\) 初始化,并使用结合了偏好对学习的DPO损失的损失函数进行更新。每个偏好对对应的损失如下:
    $$\begin{aligned}
    \mathcal{L}_{DPO} &= \mathcal{L}_{DPO}(\mathcal{S}_u^w, \mathcal{S}_u^l | \mathcal{H}_u) \\
    &= -\log \sigma \left(\beta \log \frac{\mathcal{M}_{t+1}(\mathcal{S}_u^w | \mathcal{H}_u)}{\mathcal{M}_t(\mathcal{S}_u^w | \mathcal{H}_u)} - \beta \log \frac{\mathcal{M}_{t+1}(\mathcal{S}_u^l | \mathcal{H}_u)}{\mathcal{M}_t(\mathcal{S}_u^l | \mathcal{H}_u)}\right)
    \end{aligned}$$

    • 注:上式和DPO的损失函数基本相同,对于给定的 \(\pi_{ref}\) 和数据集 \(D\) ,DPO是在优化语言模型 \(\pi_\theta\) 以最小化loss:
      $$
      Loss_{\text{DPO}}(\pi_\theta;\pi_{ref}) = - \mathbb{E}_{(x,y_w,y_l) \sim D}\left [ \log \sigma \left( \beta\log\frac{\pi_\theta(y_w|x)}{\pi_{ref}(y_w|x)} - \beta\log\frac{\pi_\theta(y_l|x)}{\pi_{ref}(y_l|x)} \right)\right ]
      $$
  • 如算法2和图2(b)所示,整个过程涉及训练一系列模型 \(\mathcal{M}_t, …, \mathcal{M}_t\)。为了减轻 beam search 推理过程中的计算负担 ,论文仅随机采样 \(r_{DPO} = 1\%\) 的数据用于偏好对齐。对于每个后续模型 \(\mathcal{M}_{t+1}\),它从先前的模型 \(\mathcal{M}_t\) 初始化,并利用由 \(\mathcal{M}_t\) 生成的偏好数据 \(D_t^{pairs}\) 进行训练

  • 对 随机采样 \(r_{DPO} = 1\%\) 的数据用于偏好对齐 的理解:

    • DPO数据大小是固定的,总数据量一定,每轮迭代时均考虑采样部分数据用于DPO
    • 随机采样的目的是,减少每轮迭代(\(\mathcal{M}_t \rightarrow \mathcal{M}_{t+1}\) 算是一轮迭代)的计算复杂度
    • 问题:为什么不直接采样固定的比例,然后固定下来,后续每轮DPO迭代都用相同数据呢?
      • 回答:为了数据的多样性和模型泛化性,类似于随机森林中每次采样部分样本训练一样

系统部署

  • OneRec已在现实世界的工业场景中成功实现。为了平衡稳定性和性能,论文部署了 OneRec-1B 用于在线服务
  • 如图3所示,论文的部署架构由三个核心组件组成:
    • 1)训练系统,
    • 2)在线服务系统,
    • 3)DPO样本服务器
  • 种子模型训练 :该系统将收集到的交互日志作为训练数据,最初采用 next token 预测目标 \(\mathcal{L}_{NTP}\) 来训练种子模型
  • DPO偏好对齐 :在收敛后,论文添加DPO损失 \(\mathcal{L}_{DPO}\) 进行偏好对齐,利用XLA和bfloat16混合精度训练来优化计算效率和内存利用率
    • 注:XLA(Accelerated Linear Algebra)是一种专为优化线性代数计算而设计的编译器,主要用于深度学习框架(如TensorFlow、JAX等)中,以提升计算效率和资源利用率
  • 参数部署 :训练好的参数被同步到在线推理模块和DPO采样服务器,用于实时服务和基于偏好的数据选择
  • 为了提高推理效果,论文实施了两项关键优化 :
    • 量化&KV缓存 :结合float16量化的键值缓存解码机制,以减少GPU内存开销;
    • 量化&KV缓存 :以及 beam size 为 128 的 beam search 配置,以平衡生成质量和延迟
    • MoE架构 :得益于MoE架构,在推理过程中只有13% 的参数被激活

相关实验

  • 首先,论文在离线设置下将OneRec与逐点方法和几种DPO变体进行比较
  • 然后,论文对提出的模块进行了一些消融实验,以验证OneRec的有效性
  • 最后,论文将OneRec部署到线上并进行A/B测试,以进一步在快手平台上验证其性能

Experiment Setup

实现细节
  • 论文使用Adam优化器训练模型,初始学习率为 \(2×10^{-4}\)
  • 论文使用NVIDIA A800 GPU对OneRec进行优化
  • 在整个训练过程中,DPO样本比率 \(r_{DPO}\) 设置为1%,并且论文通过 beam search 为每个用户生成 \(N = 128\) 个不同的响应;
  • 语义标识符聚类过程为每个 codebook 层采用 \(K = 8192\) 个聚类,codebook 层数设置为 \(L = 3\);
  • MoE 架构包含 \(N_{MoE} = 24\) 个专家,每次前向传递通过 top-\(k\) 选择激活 \(K_{MoE} = 2\) 个专家;
  • 对于 session 建模,论文考虑 \(m = 5\) 个目标 session item,并采用 \(n = 256\) 个历史行为作为上下文
基线方法
  • 论文采用以下具有代表性的推荐模型、DPO及其变体作为额外的比较基线,基线方法包括:
    • SASRec[22]:采用单向Transformer架构来捕获 user-item 交互中的顺序依赖关系,用于 next item 预测
    • BERT4Rec[40]:利用带有掩码语言建模的双向Transformer,通过序列重建来学习上下文 item 表示
    • FDSA[50]:实现了双自注意力路径,以在异构推荐场景中联合建模 item 级转换和特征级转换模式
    • TIGER[36]:利用分层语义标识符和生成式检索技术,通过自回归序列生成进行顺序推荐
    • DPO[35]:通过隐式奖励建模,使用从人类反馈数据中导出的封闭形式奖励函数来形式化偏好优化
    • IPO[1]:提出了一个有理论基础的偏好优化框架,绕过了标准DPO中固有的近似
    • cDPO[30]:引入了一个考虑噪声偏好注释的鲁棒性感知变体,纳入了标签翻转率参数 \(\epsilon\)
    • rDPO[5]:使用重要性采样开发了一个无偏损失估计器,以减少偏好优化中的方差
    • CPO[47]:通过联合训练序列似然奖励和监督微调目标,将对比学习与偏好优化统一起来
    • simPO[29]:通过采用序列级奖励裕度进行偏好优化,同时通过归一化概率平均消除参考模型依赖性
    • S-DPO[4]:通过硬负采样(hard negative sampling)和多 item 对比学习,使DPO适用于推荐系统,以提高排名准确性
评估指标
  • 论文使用几个关键指标来评估模型的性能。每个指标在评估模型输出的不同方面都有独特的作用,并且论文在每次迭代中对随机采样的测试用例集进行评估。为了估计每个特定 user-session 对的各种交互的概率,论文使用预训练的奖励模型来评估推荐 session 的价值。论文计算不同目标指标的平均奖励,包括 session 指标:
    • 观看时间(session watch time,swt)
    • 观看概率(view probability,vtr)
    • 关注概率(follow probability,wtr)
    • 点赞概率(like probability,ltr)
    • swt和vtr是观看时间指标 ,而wtr和ltr是交互指标

离线性能

  • 表1展示了OneRec与各种基线之间的全面比较。对于观看时间指标,论文主要关注 session 观看时间(swt),在交互指标中关注点赞概率(ltr)。论文的结果揭示了三个关键发现:
    • 首先,所提出的 session-wise 的生成方法显著优于传统的基于点积的方法和像TIGER这样的逐点生成方法。与 TIGER-1B 相比,OneRec-1B在最大swt上提高了1.78%,在最大ltr上提高了3.36%。这证明了 session-wise 的建模在保持推荐之间的上下文连贯性方面的优势,而逐点方法在生成输出中难以平衡连贯性和多样性
    • 其次,少量的DPO训练就能带来显著的提升。仅使用1% 的DPO训练比率(\(r_{DPO}\)),OneRec-1B + IPA在最大swt上比基础的 OneRec-1B 提高了4.04%,在最大ltr上提高了5.43%。这表明有限的DPO训练可以有效地使模型与期望的生成模式保持一致
    • 第三,所提出的IPA策略优于各种现有的DPO变体。如表1所示,IPA比其他DPO实现具有更优的性能。值得注意的是,一些DPO基线甚至比未进行偏好对齐的 OneRec-1B 模型表现更差,这表明通过迭代挖掘自生成输出来进行偏好选择比其他方法更有效
  • 从表1可以观察到一些补充结论:
    • DPO其实是一个不错的基线,且效果基本上优于其他DPO变体(包括IPO、cDPO、rDPO和CPO等),只有simPO和s-DPO在某些指标上能超过DPO
    • 实验评估是离线的,基于奖励模型的,评估准确性取决于奖励模型

Ablation Study

DPO样本比率消融
  • 为了研究DPO训练中样本比率 \(r_{DPO}\) 的影响,论文将DPO样本比率从 1% 变化到 5%
  • 如图4所示,消融结果表明,增加样本比率在多个评估目标上带来的性能提升微乎其微
    • 效果 :尽管计算成本增加,但性能提升不显著(超过1% 基线)
    • 算力 :在DPO样本服务器推理期间,样本比率与GPU资源利用率之间存在线性关系:5% 的样本比率比1% 的基线需要多5倍的GPU资源。这种缩放特性在计算效率和模型性能之间建立了明确的权衡
  • 因此,在平衡计算效率和性能的最佳权衡后,论文在训练中应用 1% 的DPO样本比率,它在仅需要更高样本比率所需计算资源的 20% 的情况下,实现了观察到的最大性能的平均95%
模型缩放消融
  • 当模型规模增加时OneRec的性能表现(如图6所示):
    • 将OneRec从0.05B扩展到1B实现了持续的精度提升,展示了一致的缩放特性
    • 与 OneRec-0.05B 相比,OneRec-0.1B 在精度上实现了显著的最大 14.45% 的提升
    • 在扩展到0.2B、0.5B和1B时,还可以分别实现 5.09%、5.70% 和5.69% 的额外精度提升

Prediction Dynamics of OneRec

  • 如图5所示,论文展示了不同层上 8192 个 code 的预测概率分布,其中红星表示奖励值最高的 item 的语义ID
  • 与OneRec基线相比,OneRec+IPA 在预测分布上表现出显著的置信度转移,这表明论文提出的偏好对齐策略有效地鼓励基础模型产生更符合偏好的生成模式
    • 理解 :OneRec+IPA 上模型输出更自信,更集中,熵更小
  • 此外,论文观察到第一层的概率分布(熵 = 6.00)比后续层(第二层平均熵 = 3.71,第三层熵 = 0.048)具有更大的分散性,后续层的分布逐渐集中。这种分层不确定性的降低可以归因于自回归解码机制:初始层的预测继承了前序解码步骤中更高的不确定性,而后续层受益于积累的上下文信息,这些信息限制了决策空间
    • 现象 :不论开始层输入分布如何,后续层都倾向于分布更集中
    • 理解(TODO):后续层是属于第一层的残差,残差分布倾向于一致,所以后续层对应的SID都差不多?这里跟RQ-VAE中的沙漏现象是否相似?

在线A/B测试

  • 为了评估OneRec的在线性能,论文在快手主页面的视频推荐场景中进行了严格的在线A/B测试
  • 线上流量划分 :论文将OneRec的性能与当前的 multi-stage 推荐系统进行比较,使用 1% 的主要流量进行实验
  • 评估指标 :论文使用总观看时间来衡量用户观看视频的总时长,平均观看时长计算用户在推荐系统展示请求的 session 时每视频的平均观看时间
  • 结论 :在线评估表明,OneRec在总观看时间上实现了1.68% 的提升,在平均观看时长上实现了6.56% 的提升,这表明OneRec实现了更好的推荐结果,并为平台带来了可观的收益增长

OneRec总结与讨论

  • 在论文中,论文重点介绍了一种用于 single-stage 生成式推荐的工业解决方案。论文的解决方案有三个关键贡献:
    • 第一,论文通过应用MoE架构,论文以高计算效率有效地扩展了模型参数,为大规模工业推荐提供了可扩展的蓝图
    • 第二,论文发现以 session-wise 的生成方式对目标 item 的上下文信息进行建模的必要性,证明了上下文序列建模本质上比孤立的逐点方式更能捕捉用户偏好动态
    • 第三,论文提出了一种迭代偏好对齐(IPA)策略,以提高OneRec在不同用户偏好模式下的泛化能力。广泛的离线实验和在线A/B测试验证了OneRec的有效性和效率
  • 在线实验 :此外,论文对在线结果的分析表明,除了用户观看时间外,论文的模型在点赞等交互指标方面存在局限性
  • 未来规划 :在未来的研究中,论文旨在增强端到端生成式推荐在多目标建模方面的能力,以提供更好的用户体验

补充:相关工作(直译,保留引用)

生成式推荐

  • 近年来,随着生成模型的显著进展,生成式推荐受到越来越多的关注
    • 传统基于嵌入的检索方法主要依赖双塔模型计算候选项排序分数,并利用高效的MIPS或ANN(2015; 2016)搜索系统检索Top-K相关项
    • 生成式检索(GR)(2016)将数据库相关文档检索问题转化为序列生成任务,按序生成文档标记(如标题、ID或预训练语义ID)。GENRE(2019)首次采用Transformer架构进行实体检索,基于上下文自回归生成实体名称。DSI(2016)提出为文档分配结构化语义ID的概念,并训练 Encoder-Decoder 模型实现生成式文档检索。TIGER(2019)将该范式扩展至推荐系统,提出生成式 item 检索模型
  • 除生成框架外,item 索引方法也备受关注。近期研究聚焦于语义索引技术(2017; 2019),旨在基于内容信息索引 item。例如,TIGER(2019)和LC-Rec(2019)对 item 标题和描述的文本嵌入应用残差量化(RQ-VAE)进行标记化;Recforest(2017)则通过层次化K均值聚类获取簇索引作为标记。EAGER(2019)等研究进一步探索将语义与协同信息共同融入标记化过程

语言模型的偏好对齐

  • 在 LLM 的 post-training 阶段(2018), RLHF(2017; 2019)是主流方法,其通过奖励模型表征人类反馈,指导模型对齐人类价值观
  • RLHF 存在不稳定和低效问题。直接偏好优化(DPO)(2024)通过闭式解推导最优策略,直接利用偏好数据优化模型。此外,多种改进方法被提出:
    • IPO(2024)绕过DPO的两步近似,提出通用目标;
    • cDPO(2024)引入超参数ε缓解噪声标签影响;
    • rDPO(2024)设计原始二元交叉熵损失的无偏估计
    • CPO(2024)、simDPO(2024)等变体也从不同角度优化DPO
  • 与NLP中人类显式标注偏好数据不同,推荐系统的偏好学习面临 user-item 交互数据稀疏的独特挑战,相关研究仍处于探索阶段。不同于S-DPO(2024)为基于语言的推荐器整合多负例数据,论文训练奖励模型并基于其分数为用户定制个性化偏好数据

附录:OneRec 的一些讨论

OneRec 冷启动

  • 用户冷启动:新用户没有历史行为序列怎么办?OneRec的冷起不好做吧?
  • 商家冷启动:由于语义空间的存在,只要ID语义接近,理论上就能做到冷启动?

附录:OneRec Technical Report

  • 参考链接:
    • OneRec Technical Report, 20250616, Kuaishou:论文是一个技术报告,相对原始论文补充了很多细节,快40页
    • 效果&成本双突破!快手提出端到端生成式推荐系统OneRec!,快手技术公众号

OneRec Technical Report 的整体说明

  • 整体架构和 OneRec 相似,多了些细节和改进
  • 论文暂时主要简单介绍一下与原始论文相比的重点改进
    • 架构上,GPU利用率(MFU,模型浮点运算利用率)从 10% 提升至 训练23.7% 和 推理28.6% (快追上业内主流 AI 模型效能)
    • Beam Search 从之前的 128 提升至 512,说明多采样才是硬道理
    • 对模型生成的合法性给与格式奖励,可以缓解模型生成的高概率样本不符合规范的情况,这里的不符合规范指 RQ 编码后输出的 Token 空间可能是无意义的
    • 使用了 RSFT + RL 的范式做训练,原始 OneRec 论文是抽取最好和最坏的样本做 DPO
    • 提出了一种创新的 RL 方法,ECPO(Early Clipped GRPO),是对 GRPO 方法的一种改进
  • 收益披露:
    • 在快手主流量上已经承接了 25% 的QPS
    • 在各种指标上均有涨幅,其中一些指标涨幅超过 5%
    • 本地生活场景收益尤其明显(在本地生活场景已经 100% 流量推全)

      除了短视频推荐的消费场景之外,OneRec在快手本地生活服务场景同样表现惊艳:AB对比实验表明该方案推动GMV暴涨21.01%、订单量提升17.89%、购买用户数增长18.58%,其中新客获取效率更实现23.02%的显著提升

  • 题外话:未披露作者信息,是怕被挖人吧

对齐微调升级为 ECPO(Early Clipped GRPO)

  • OneRec Technical Report 中使用的是 ECPO,原始 OneRec 论文用的是 DPO
  • 具体来说,对于用户 \( u \),论文使用旧策略模型生成 \( G \) 个物品。每个物品与用户一起输入偏好奖励模型,得到 P-Score 作为奖励 \( r_i \)。优化目标如下:
    $$
    \mathcal{J}_{\text{ECPO}}(\theta) = \mathbb{E}_{u \sim P(U), \{o_i\}_{i=1}^G \sim \pi_{\theta_{old} } } \left[ \frac{1}{G} \sum_{i=1}^G \min \left( \frac{\pi_\theta(o_i|u)}{\color{red}{\pi’_{\theta_{old} }}(o_i|u)} A_i, \text{clip} \left( \frac{\pi_\theta(o_i|u)}{\color{red}{\pi’_{\theta_{old}} }(o_i|u)}, 1 - \epsilon, 1 + \epsilon \right) A_i \right) \right], \\
    A_i = \frac{r_i - \text{mean}(\{r_1, r_2, \ldots, r_G\})}{\text{std}(\{r_1, r_2, \ldots, r_G\})},\\
    \color{red}{\pi’_{\theta_{old} }(o_i|u) = \max \left( \frac{\text{sg}(\pi_\theta(o_i|u))}{1 + \epsilon + \delta}, \pi_{\theta_{old} }(o_i|u) \right), \quad \delta > 0,}
    $$
    • \(\text{sg}\) 表示停止梯度操作(stop gradient operation)
    • \(\delta\) 是一个大于 0 的超参数
  • 论文对 GRPO(Group Policy Relative Optimization)(2024) 进行了修改,使其训练过程更加稳定
    • 如图 6 所示,在原始 GRPO 中,允许负优势(negative advantages)的策略比率(\(\pi_\theta / \pi_{\theta_{old} }\))较大,这容易导致梯度爆炸
    • 因此,论文预先对具有较大比率的策略进行截断,以确保训练稳定性,同时仍允许相应的负优势生效
    • \(\delta\) 越大,可容忍的策略比率越大,意味着可容忍的梯度越大,这可以根据实际需求确定
  • 在 OneRec 中,论文将 \(\delta\) 设为 0.1,表示允许负优势的策略比率略微超过 \(1 + \epsilon\)
    • 论文移除了 KL 散度损失(KL divergence loss),因为在 OneRec 中, RL 和 SFT 是一起训练的,SFT 损失确保了模型的稳定性
补充:对 ECPO 和 GRPO 的理解
  • 当 advantage 为正时,ECPO 和 GRPO 等价(可以证明,都会在 \(1+\epsilon\) 处截断梯度)
  • 当 advantage 为负时,ECPO 和 GRPO 不同:
    • 若 \(\delta \to +\infty\),那么 ECPO 和 GRPO 依然等价
    • 否则,相当于允许 \(1+\epsilon+\delta\) 样本的梯度更新,但是样本的更新比值是基于 \(1+\epsilon+\delta\) 的,而不是 \(\frac{\pi_\theta}{\pi_{\theta_\text{old}}}\),这相当于对梯度进行了一定限制
      • 当 \(A\) 为负时,此时若 \(\frac{\pi_\theta}{\pi_{\theta_\text{old}}}\) 很大,则 \(\pi_{\theta_\text{old}}\) 很小
      • 于是有取 \(\max\) 操作后 \(\pi’_{\theta_\text{old}} = \frac{\text{sg}(\pi_\theta)}{1+\epsilon+\delta}\),从而 Clip 结果是 \(1+\epsilon\)
      • 经过 \(\min\) 的结果是 \(\frac{\pi_\theta \cdot (1+\epsilon + \delta)}{\text{sg}(\pi_{\theta})}\)
      • 即梯度权重从原始的 \(\frac{1}{\pi_{\theta_\text{old}}}\) 变成了 \(\frac{\text{sg}(\pi_{\theta})}{1+\epsilon+\delta}\)
      • 结论:不论 \(\pi_{\theta_\text{old}}\) 多小,都能稳稳控制梯度权重为 \(\frac{\text{sg}(\pi_{\theta})}{1+\epsilon+\delta}\),这个值是有限的,而 \(\frac{1}{\pi_{\theta_\text{old}}}\) 这个值可能是无限大
    • 理解1:因为 \(\frac{\pi_\theta}{\pi_{\theta_\text{old}}}\) 的值可能是无穷大,这样可以缓解梯度爆炸的问题
    • 理解2:为什么 advantage 为正时不需要过小的 \(\frac{1}{\pi_{\theta_\text{old}}}\) 限制呢?因为此时的 \(\frac{\pi_\theta}{\pi_{\theta_\text{old}}}\) 值最小也就是 0,不会造成梯度爆炸,最多造成梯度消失(其实这部分也可以关注一下)

Post-Training 中加入了 RSFT

  • 在后训练中,除了使用了 ECPO 外,还引入了 RSFT(Reject Sampling Fine-Tuning)
  • 具体来说:拒绝后 50% 的样本

简单证明了存在 Scaling Law

  • 仅仅在损失函数上证明了存在 Scaling Law

披露了更多 RM 细节

  • Reward Model 细节披露

收益细节

  • 线上 AB 测试收益披露

RS——RARE

RARE: Real-time Ad REtrieval framework

  • 参考链接:
    • (RARE)Real-time Ad retrieval via LLM-generative Commercial Intention for Sponsored Search Advertising, 202504, Tencent
      • 相关博客:腾讯搜索广告生成式召回

整体思路

  • 现有的 LLM-based 的召回(retrieval,一些地方也称为检索)方法通过生成 numeric 或 content-based DocIDs 来检索 docs/ads
    • 问题 :1)DocIDs 与文档之间的 one-to-few 映射关系;2)耗时的内容提取过程
    • 两个问题共同导致语义效率低下,限制了其在大规模语料库中的扩展性
    • 注:在本文的写作中,会使用 文档(docs)或者 广告(ads)替代检索目标,而不是普通推荐系统中的 item
  • 论文提出了一种实时广告召回(Real-time Ad REtrieval,RARE)框架,该框架利用 LLM 生成的文本作为中间语义表示,直接从 Query 中实时召回广告
    • 论文将这些中间语义表示称为商业意图 ,Commercial Intentions (CIs),本质上就是 LLM 输出的剪短的文本
    • CIs 是由注入商业知识的定制化(Customized)LLM 生成的(有领域相关性)
    • 每个 CI 对应多个广告,从而形成一个轻量级且可扩展的 CIs 集合
  • 离线和在线实验验证了 RARE
    • 离线:在 4 个主要场景上超过 10 个基线
    • 在线 :Consumption+5.04%, GMV+6.37%, CTR+1.28%,Shallow Conversions+5.29%

一些讨论

  • 传统召回 :广告系统分召回和排序两个模块,其中召回召回模型常常采用两阶段流程,1)从 Query 中提取关键词;2)利用这些关键词获取广告
    • 问题 :现有的两阶段召回方法放大了用户 Query 与人工选择关键词之间的差异,导致大量漏检问题(missed retrieval issues)
    • 解决方案 :query-ad 单阶段(直接从Query 到 广告)方法通过直接召回广告解决了漏检问题,但由于推理能力和领域知识的限制,仍难以深入理解商业意图(Commercial Intentions)
  • LLM-based 的召回 :大多数 LLM-based 的召回方法(Lin等,2025)首先通过训练模型将文档与其标识符(DocIDs)关联,构建文档索引。在召回阶段,模型处理 Query 并生成相应的 DocIDs(Li等,2023a)
    • DSI(Tay等,2022)使用 numeric IDs 表示文档,建立用户 Query 与 numeric IDs 之间的连接;
    • LTRGR(Li等,2023b)则提取文档内容(如标题和正文)表示文档,实现从用户 Query 到文档的检索
  • 使用复杂的 DocIDs 存在以下缺陷 :
    • 第一,由于 DocIDs 与候选之间的 one-to-few 映射关系,推理效率较低(Wang等,2023b),难以在大规模场景中实现实时生成;
    • 第二,仅用复杂 DocIDs 表示文档或广告无法充分发挥 LLM 在商业意图挖掘和高级文本生成方面的能力,从而阻碍了对广告商意图的有效探索;
    • 第三,其泛化能力较差。当新候选出现时,通常需要重新训练模型或更新 FM-index (Ferragina and Manzini, 2000) 以适应其 DocIDs,难以及时更新或删除候选
  • 关键问题提出 :由于广告召回需要实时获取与用户商业意图匹配的大规模广告集合,现有语义效率低下的 DocIDs 在实际任务中既不实用也不适用。因此,利用 LLM 强大的语义能力设计更有效的语义标记索引,并开发更全面的端到端架构,已成为一项关键挑战
  • 论文开发了一种名为 RARE 的实时LLM生成广告召回框架。该框架利用 LLM 生成的商业意图(CIs)作为中间语义表示,直接连接 Query 与广告,而非依赖人工选择的关键词或复杂文档标识符,具体方法为:
    • 离线为广告生成 CIs,构建索引 :RARE 首先使用注入知识的LLM(离线)为语料库中的广告生成 CIs ,随后筛选出一组有限但全面的 CIs,并构建一个动态索引 ,将这些 CIs 以 one-to-many 的关系映射到对应的广告
      • 理解:相对之前的区别是:为候选广告生成 CIs 而不是 DocIDs
    • 在线为 Query 生成 CIs,实时召回 :接收到 Query 时,RARE 利用定制化 LLM(在线)实时生成 CIs,并从预构建的索引中召回对应的广告
  • 关键创新 :在于利用定制化 LLM 生成的 CIs 作为中间语义 DocIDs,链接 Query 与广告
    • 定制化 LLM 通过对基础 LLM 进行知识注入(knowledge injection)和格式微调(format fine-tuning)开发而成
      • 知识注入可融入领域特定信息以增强广告领域的专业性;
      • 格式微调则确保LLM仅输出 CIs 并提升解码效率
    • CIs 定义为关键词的聚合,由定制化LLM基于广告相关材料生成。与现有精心设计的 DocIDs 相比,CIs 充分利用了LLM的文本生成能力。CIs 与广告之间的 one-to-many 对应关系使得解码过程极为高效
      • 问题:为什么 one-to-many 能带来高效的解码?
      • 回答:这里是跟 one-to-few 而言的(TIGER 中就属于 one-to-few),他们要实时生成隐空间语义 ID,且多个语义 ID 才对应一个实际的广告,但 one-to-many 相当于生成的 CI 可直接检索很多个广告(离线会构建索引:让一个 CI 对应很多广告)
    • 对于新广告 ,RARE 可通过 constrained beam search 生成 CIs,无需重新训练模型
    • 传统 query-keyword-ads 范式中的关键词竞价可能引发 index manipulation 问题。与关键词不同,CIs 由具备世界知识和商业专业知识的LLM生成,能够更好地挖掘广告和 Query 背后的商业意图
  • 论文的主要贡献如下:
    • (1)提出了一种名为 RARE 的新型端到端生成召回框架,实现实时召回
      • 注:文章中说是首个在百万级数据库上实现实时召回的 LLM 生成架构研究,但忽略了 TIGER?
    • (2)提出了一种知识注入与格式微调方法,使基础LLM能够挖掘广告商和用户的深层商业意图,并以 CIs 形式表达
    • (3)实现了在线部署 LLM-based 的召回服务
    • (4)离线+在线实验

相关工作(直译)

  • 广告召回 :传统广告召回(Zhao和Liu,2024;Wang等,2023c)通常遵循 query-keyword-ads 架构,即 Query 召回关键词,再通过关键词拉取广告。该方法包括基于词(word-based)和基于语义(semantic-based)的方法
    • Word-based Methods(Ramos等,2003;Robertson等,2009)解析用户 Query 获取关键词,并使用倒排索引召回候选广告;
    • Semantic-based Methods(Ramos等,2003;Yates等,2021)则利用双编码器在共享语义空间中获取 Query 和关键词的嵌入,实现基于语义相似性的召回
    • 这些方法依赖人工选择关键词,导致大量漏检问题。相比之下,Generative LLM 召回方法(Sun等,2024;Lin等,2024;Tang等,2023b)使用 DocIDs 表示广告,LLM 在接收 Query 时直接生成候选广告对应的 DocID
  • Generative LLMs Retriever :生成式召回利用 LLM 的生成能力构建端到端召回模型
    • 部分方法(如DSI(Tay等,2022)、NCI(Wang等,2022)、TIGER(Rajput等,2024)以文档ID为生成目标 ,实现 Query 对文档/广告的召回。这些方法利用LLM学习文档/广告与其ID的对应关系,直接生成相关文档/广告的 ID 以完成召回
    • 其他方法(如SEAL(Bevilacqua等,2022)和LTRGR(Li等,2023b)则以文档内容为中介实现文档召回,通过FM索引生成文档中出现的片段,辅助 Query 到文档的召回。MINDER(Li等,2023)采用伪 Query 和文档内容进行召回,但显著增加了索引量,不适用于候选集较大的场景
  • 语义 DocIDs :LLM生成召回通常使用 DocIDs 执行 Query 到文档的召回任务
    • 现有 DocIDs 主要包括 numeric IDs 和文档内容。例如,TIGER中的 numeric IDs 表示为离散语义标记的元组;LTRGR中的文档内容则由文档内预定义序列组成。然而,这些方法使用的语义标记为类ID特征,解码效率较低 ,因为每个DocID仅对应少量候选。对于新候选文档或广告,需重新训练模型或重建FM索引以获取其 DocIDs,难以及时更新或删除广告

RARE 方法整体介绍

  • 如图1所示,论文提出了一种新颖的端到端生成式召回架构,专为在线召回设计,命名为实时广告召回(RARE)
  • RARE 有效缩短了链路结构,使广告能够突破关键词竞价的限制,帮助广告主获取更精准的流量
    • 问题:如何理解“突破关键词竞价的限制”?
  • 图2给出了 RARE 与传统召回方法的一个对比:

端到端生成式架构

  • 当接收到用户 Query 时,RARE 首先生成对应的具有特定语义的文本,即商业意图(CIs),随后利用这些 CIs 召回最终的广告
  • CIs 到广告的索引构建过程以及召回流程如下:
    • 索引构建(Indexing) :RARE 首先生成整个广告库的 CIs 并确定商业意图集合,随后构建 CIs 到广告的倒排索引
      • 对于后续新增的广告,基于当前商业意图集合进行约束推理 ,确保每个新候选广告能够准确更新至索引中
      • 值得注意的是,CIs 是由定制化 LLM 生成的具有特定语义的文本,用于挖掘广告的商业意图。具体实现细节将在 3.3节 讨论
    • 召回流程(Retrieval) :基于离线缓存和在线推理,为 Query 实时生成 CIs
      • 高频 Query 的 CIs 会被存储在缓存中,当 Query 到达时,RARE 首先检查当前 Query 是否匹配缓存条目
        • 若匹配,则直接使用对应的 CIs 获取广告;
        • 否则,RARE 会使用定制化 LLM 结合 constrained beam search 进行实时推理(具体实现细节将在原文 3.4节 介绍)

定制化 LLM

  • 为了挖掘广告和用户 Query 背后的深层意图,论文通过知识注入和格式微调两个阶段对基础 LLM 进行定制化
  • 阶段1:知识注入 :此阶段旨在将广告领域的专业知识注入基础 LLM 中,使其能够理解广告和用户 Query 的商业意图。知识数据包括 Query 意图挖掘、广告意图挖掘以及广告词购买等任务。知识注入过程可形式化表示为:
    $$
    \theta^{\prime}=h(\theta,K),
    $$
    • 函数 \( h \) 以 LLM 模型参数 \( \theta \) 和广告知识数据 \( K \) 为输入 ,输出更新后的模型参数 \( \theta^{\prime} \)
    • 新参数 \( \theta^{\prime} \) 用于生成预测结果,即:
      $$
      y=P(y|x;\theta^{\prime}).
      $$
  • 阶段2:格式微调 :在具备商业知识的 LLM 基础上,此阶段专注于优化生成 CIs 的格式并提升其多样性。格式微调的训练数据来自真实线上数据,并进行了必要的格式调整。格式微调的生成损失如下:
    $$
    L(\theta)=\frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{T_{i} }logp(y_{i,t}|y_{i<t},x_{i};\theta),
    $$
    • 微调数据集为 \( D=(x_{i},y_{i})_{i=1}^{N} \),\( x_{i} \) 为输入序列,\( y_{i} \) 为目标输出序列
    • 概率 \( p(\cdot) \) 是由参数为 \( \theta \) 的模型基于 \( x_{i} \) 和已生成词 \( y_{i < t} \) 预测的概率
    • 问题:\(x_i\) 是历史序列吗?\(y_i\) 和 \(x_i\) 都是广告吗?
  • 定制化LLM显著增强了其理解和提取广告与用户 Query 意图的能力。它将广告压缩并总结至商业意图空间、聚类相似的广告,从而减少同质化召回,提升在线和离线的召回性能

索引构建

  • 论文使用定制化 LLM 生成广告的商业意图(CIs),随后构建 CIs 到广告的倒排索引
  • 商业意图(CIs) :CIs 是由定制化 LLM 生成的简短文本 ,用于描述用户或广告的商业意图。给定包含广告信息(如标题、描述等)的提示,定制化 LLM 会生成对应的 CIs。例如,对于广告标题“新款智能手机限时优惠”,生成的 CIs 可能包括“智能手机促销”、“限时折扣”和“电子产品优惠”
  • 动态索引 :CIs 与广告之间是 one-to-many 的关系 ,即一个 CI 可能对应多个广告。论文基于这种关系构建动态索引,支持高效召回。对于新增广告 ,RARE 通过约束推理生成其 CIs,无需重新训练模型或更新索引结构

高效推理

  • 高效推理对于从数百万候选集中实时召回至关重要,因为搜索广告对召回时间有严格要求
  • constrained beam search :本工作采用 constrained beam search 算法生成商业意图(CIs),确保模型输出限定在预定义的 CIs 集合内
    • 论文开发了基于 CUDA 的 constrained beam search 实现,并将其与 LLM 推理过程集成,支持并行生成 Beam size 的 CIs,从而提升解码效率
    • 此外,在 constrained beam search 框架中引入了截断功能,允许丢弃低分单标记以提升模型输出的准确性。具体约束过程如 图3 所示(新的广告加入时,为其生成对应的 CIs,然后将其加入到这些 CIs 的检索列表中,注意,从图1中看还可以加入新的 CIs)
  • 缓存技术 :搜索系统呈现显著的长尾效应 ,其中 5% 的 Query 占据了总 Query 请求的 60%
    • 为了提升推理效率,论文对这些高频 Query 进行离线推理和存储 ,当用户提交 Query 时,系统首先检查离线缓存
      • 若命中,则立即返回结果;
      • 否则,由推理服务处理请求
  • 离线处理对时间要求较低,因此论文使用 13B 规模的 LLM 处理这些 Query。在线推理有严格的时延要求(通常在毫秒级完成),因此使用 1B 规模的小模型
    • 通过缓存数百万头部 Query 的离线结果,论文能够减少 70% 的在线机器消耗,不仅降低推理时间,还提升了头部 Query CIs 的质量

相关实验和讨论

Experiment Setup

  • 训练数据集 :为了使基础LLM具备商业知识,论文使用了商业知识和合成数据。原始数据来自真实在线日志,通过开源LLM执行 Query 意图挖掘和广告意图挖掘等任务生成最终合成数据。格式微调主要涉及 Query 和广告的商业意图(CIs),这些数据来自真实在线交互,并按固定规则组合。详情参见附录A的表4
  • 评估数据集 :为了评估模型效果,论文从真实场景中收集了一天内头部 Query 及其对应点击广告的数据。清洗后,得到5,000条 Query 和150,000条广告作为基准数据,每条 Query 最多包含 1,000 个候选广告
  • 基线方法 :论文将 RARE 与 10 种竞争基线方法进行比较,涵盖四大类别,包括基于关键词的 BM25、基于语义的 BERT、基于生成的 T5 以及 LLM-based 的 Qwen 等
    • BM25 :将 Query 分割为词项并计算相关性得分,最终累加得到文本相似度结果
    • BERT-small :使用 4 层 Transformer 网络,隐藏层大小为 768,参数约 52.14M
    • BERT-base :使用 12 层 Transformer 网络,隐藏层大小为 768,共 12 个头,参数约 110M
    • SimBert-v2-R :结合生成与召回能力的模型,用于计算句子向量
    • T5 :基于生成的强基线,通过大量在线点击 Query 和关键词微调
    • DSI :使用语义ID进行召回的典型方法
    • Qwen-1.8B 和 Hunyuan-2B :参数规模与 RARE 相近的模型,通过格式微调确保仅生成 CIs
  • 评估指标 :
    • 广告覆盖率(ACR, Ad Coverage Rate) :在广告检索中,ACR表示覆盖率,即召回广告的请求占总请求的比例。如公式5所示,AdPV表示召回广告的请求数量,PV表示总请求数量
      $$ ACR = AdPV / PV \tag{5} $$
    • 命中率(HR@K, Hit Ratio) :如公式6所示,GT表示候选广告的真实集合,Hits@K表示在检索结果的前K个候选中属于真实集合的相关广告数量
      $$ HR@K = \frac{Hits@K}{|GT|} \tag{6} $$
    • 平均精度均值(MAP, Mean Average Precision) :是所有查询的平均精度(AP)的均值,如公式7所示
      $$ MAP = \frac{\sum_{q \in Q} AP_q}{|Q|} \tag{7} $$
    • 平均精度(AP, Average Precision) :如公式8所示
      $$ AP_q = \frac{1}{|\Omega_q|} \sum_{i \in \Omega_q} \frac{\sum_{j \in \Omega_q} h(p_{qj} < p_{qi}) + 1}{p_{qi} } \tag{8} $$
      • \(\Omega_q\) 表示 ground-truth results
      • \(p_{qj}\) 表示广告 \(ad_j\) 在生成列表中的位置
      • \(p_{qj} < p_{qi}\) 表示广告 \(ad_j\) 在列表中排在广告 \(ad_i\) 之前
      • 理解:对 ground-truth results 中的任意一个位置 \(p_{qi}\),希望出现在它前面的广告中,包含在 ground-truth results 中的数量越多越好;
        • 举例1:最完美的情况是所有 ground-truth results 都排在最前面,则该值为 1
        • 举例2:最极端的情况,如果有些 ground-truth results 的广告排在了最后,则该值非常小
  • 实现细节
    • 使用 Hunyuan 作为主干模型,参数包括 1B-Dense-SFT 和 13B-Dense-SFT
    • 离线缓存使用13B模型,Beam Search 大小为 256,温度系数为 0.8,最大输出长度为 6
    • 在线推理使用1B模型,Beam Search 大小为 50,温度系数为 0.7,最大输出长度为 4,确保延迟在 60ms 内
    • 新广告的 CIs 每小时更新索引 ,整个 CIs 集合每月更新一次 ,并定期注入新商业知识

实验结果

离线评估
  • 表1展示了 RARE 与 10 种基线方法在工业评估数据集上的对比结果。RARE 在 HR@500 和 MAP 上表现优异,同时保持高 ACR,表明其能够理解用户搜索意图并优化广告投放。其 ACR 超过90%,HR@500的高值验证了其召回高商业价值广告的能力
消融研究
  • 论文通过两种消融研究分析各组件贡献:
  • 表3展示了 RARE 在不同设置下的广告召回结果:
    • w/o. KI(无知识注入):召回率仅为 59.51%,显著低于 RARE 的95.05%
    • w/o. CBS(无 constrained beam search ):平均 CIs 数量仅为 4.84,显著低于 RARE 的 74.49
    • w/o. CBS & KI :HR@500、MAP 和 ACR 均为最低
  • 注:附录C 的表5通过案例定性分析了各组件的作用
在线A/B测试
  • 论文将 RARE 应用于腾讯三个在线召回场景(每日请求数十亿):
    • 微信搜索(WTS) :消费提升5.04%,GMV提升6.37%,CTR提升1.28%,浅层转化提升5.29%,深层转化显著提升24.77%
    • 需求方平台(DSP)和QQ浏览器搜索(QBS) :也观察到显著收益(表2),图5进一步展示了 RARE 在八大行业的有效性
在线推理支持
  • 图4显示了实时在线推理中不同输出长度的时间消耗。CIs 平均词数为3,确保满足安全阈值。论文开发了专用GPU集群(数百台L40),量化模型至FP8精度,每台L40支持约30 QPS,缓存命中率达65%,显著降低了计算成本

附录A 微调数据

  • 本节主要介绍微调数据的细节。定制化 LLM 的微调包括两个阶段:知识注入和格式微调
  • 知识注入阶段的微调数据主要包括 Query 意图挖掘、广告意图挖掘和广告词购买。我们将包含广告和用户信息的提示输入开源LLM(如ChatGPT),获取包含丰富推理过程与指导信息的输出,随后将这些数据作为知识注入Hunyuan模型,在微调阶段注入大量数据(规模达数十万或数百万)可能导致LLM丢失通用知识与推理能力,因此本阶段仅针对每项任务精选2,000条实例。表4详细展示了这两个阶段的微调数据

附录B 相关工作

  • 束搜索(Beam Search) :作为一种启发式搜索的解码策略, Beam Search 已被广泛应用于多项研究。例如,DSI使用 Beam Search 生成排序后的候选文档列表,TIGER通过 Beam Search 一次性生成多个候选商品ID。早在数年前,seq2seq 与 constrained Beam Search 的结合已在实体链接和文档检索任务中实现了效果与效率的双赢。例如GENRE[3]将 constrained Beam Search 应用于文档检索任务并达到SOTA性能
  • Query-关键词-广告架构(Query-kwds-ads Architecture) :传统 Query -关键词-广告方法存在两大缺陷:
    • (1) 关键词由广告主手动选择,质量参差不齐,可能导致匹配范围过宽或过窄,造成流量匹配效率低下;
    • (2) 广告主常购买大量关键词,导致关键词反转后的广告检索效率受损,给系统带来沉重负担
    • 相比之下,商业意图(CIs)由注入领域知识的LLM生成,能更精准体现广告主意图并与相关流量匹配,既带来经济效益,也保障系统长期健康运行
  • 基于编码器的LLM检索器(Encoder-based LLMs Retriever) :基于编码器的检索器利用LLM的语义能力获取文本嵌入[9]。例如:
    • cpt-text[21]通过对比学习从头训练GPT-3[5],生成高质量文本嵌入;
    • GTR[22]基于 T5 模型[23]微调获得文本向量表示;
    • NoteLLM[26]则通过新增训练任务和修改 LLM 结构实现文本-图像嵌入

附录C 定性分析

  • 表5 通过实例直观分析了 RARE 各组件的作用
  • 观察 表5 可知:
    • (1) 零样本 LLM 缺乏对提示的推理过程,仅依赖对 Query 的表层理解,导致大量低质案例;
    • (2) 知识注入阶段教会 LLM 如何推理,使其能分步分析 Query 、提供商业意图并判断相关性,但多样性不足;
    • (3) 格式微调阶段指导模型遵循规则的同时生成更多样化结果,由于训练数据多来自线上,商业意图较丰富,但自由生成过程限制了意图数量;
    • (4) 引入受限解码(constrained decoding)能在保证相关性和多样性的同时显著增加商业意图数量

RS——TIGER

TIGER:全称 Transformer Index for GEnerative Recommenders,本文主要是对论文TIGER的翻译,其中包含了一些总结和理解

  • 参考链接:
    • (TIGER)Recommender Systems with Generative Retrieval, NeurIPS 2023, Google

整体思路说明

  • 现代推荐系统常常采用retrieve-and-rank策略:
    • 其中 retrieve(检索/召回) 部分:通过将 query 和候选 item 嵌入到同一统一空间中进行大规模检索(large-scale retrieval),随后通过近似最近邻搜索(Approximate Nearest Neighbor search,ANN)选择给定 query 嵌入的候选 item (论文研究重点)
    • rank(排序) 部分:负责对检索阶段输出的 item 集合进行排序(不是论文的研究重点)
  • 论文提出了一种新颖的生成式检索方法 ,该检索模型通过自回归解码目标候选 item 的标识符
    • 为每个 item 创建具有语义意义的码元组(tuple of codewords),称为“*语义 ID *”(Semantic ID,一些文章中也简称为S ID )
    • 给定用户会话中 item 的语义 ID ** ,论文训练一个基于 Transformer** 的 sequence-to-sequence 模型,以预测用户可能交互的 next item 的语义 ID
  • 实验表明,采用该范式的推荐系统在多个数据集上显著优于当前的最先进模型。此外,将语义 ID 整合到 sequence-to-sequence 模型中增强了模型的泛化能力,表现为对无交互历史 item 的检索性能提升

一些讨论

  • 现代推荐系统采用“retrieve-and-rank”策略
    • 在 retrieve(检索/召回) 阶段 :筛选出一组候选内容(要求检索阶段返回高度相关的候选内容)
    • 在 rank(排序) 阶段:通过rank 模型进行排序(仅对检索阶段输出的候选内容进行处理)
  • 构建检索模型常用的标准方法是矩阵分解或双塔架构(dual-encoder) :
    • 矩阵分解 :将 query 和候选内容嵌入到同一空间中
    • 双塔架构(dual-encoder) :即一个塔用于 query,另一个塔用于候选内容,通过内积将 query 和候选内容嵌入到同一空间
      • 双塔架构可更好地捕捉数据中的非线性特征
      • 在推理阶段 ,这些模型会使用候选塔为所有内容创建嵌入索引。对于给定的 query,通过 query 塔获取其嵌入,并采用近似最近邻(ANN)算法进行检索候选内容
    • 近年来,双塔架构还被扩展用于序列推荐(这里有许多参考文献),显式地考虑用户-内容交互的顺序
  • 论文提出了一种构建生成式检索模型的新范式,用于序列推荐
    • 与传统 query-candidate matching 方法不同,论文的方法使用端到端的生成模型直接预测候选内容的 ID
    • 论文利用 Transformer 的记忆(参数)作为推荐系统中端到端检索的索引(与Tay等人使用 Transformer 记忆进行文档检索的思路类似),并将该方法称为生成式推荐器的 Transformer 索引(TIGER)
    • 图1展示了TIGER的方法概览。TIGER的独特之处在于提出了一种称为“语义 ID ”的新颖语义内容表示方法——即从每个内容的信息中提取的 token 序列:
  • 具体来说,给定内容的文本特征 ,论文使用预训练的文本 Encoder(例如SentenceT5[27])生成密集的内容嵌入。随后,对内容的嵌入应用量化方案 ,生成一组有序的 tokens/codewords ,称为该内容的语义 ID 。最终,这些语义 ID 被用于训练 Transformer 模型完成序列推荐任务
  • 将内容表示为语义 token 序列具有诸多优势
    • 在语义数据上训练 Transformer 记忆可以实现相似内容间的知识共享 ,从而摒弃以往推荐模型中使用的原子化(atomic) 且 随机的内容 ID
    • 通过语义 token 表示内容,模型不易陷入推荐系统中固有的反馈循环 ,从而能够泛化到新添加到语料库中的内容
      • 问题:这里的反馈循环是指马太效应吗?
    • 使用 token 序列表示内容有助于缓解内容语料库规模带来的挑战;
      • 通过 token 序列可表示的内容数量是序列中每个 token 基数(cardinality)的乘积
        • 理解:这里用一个包含多个 token 的序列表示一个 item,所以不需要太多的 token 就足够了
      • 通常,内容语料库的规模可能达到数十亿,为每个内容学习唯一的嵌入会占用大量内存。虽然可以采用基于随机哈希的技术来减少内容表示空间,但论文表明,使用语义 token 表示内容是一种更具吸引力的替代方案
  • 论文的主要贡献如下:
    • 论文提出了TIGER ,一种基于生成式检索的推荐框架,为每个内容分配语义 ID,并训练检索模型预测用户可能交互的内容的语义 ID
    • 论文证明TIGER在多个数据集上的召回率和NDCG指标均优于现有的SOTA推荐系统
    • 论文发现这种生成式检索的新范式为序列推荐系统带来了两项新能力:
      • (1) 冷启动推荐 :推荐新内容和低频内容的能力,从而改善冷启动推荐;
      • (2) 推荐多样性 :通过可调参数生成多样化推荐的能力

相关工作

  • 序列推荐系统 :基于深度序列模型的推荐系统已发展为一个丰富的研究领域
    • GRU4REC[11]首次将基于GRU的RNN用于序列推荐
    • Li等人[24]提出了神经注意力会话推荐(NARM),通过注意力机制和GRU层追踪用户的长期意图
    • Zhang等人的AttRec[41]利用自注意力机制建模当前会话的用户意图,并通过度量学习确保个性化
    • Kang等人提出的SASRec[17]使用类似仅 Decoder Transformer 的自注意力机制
    • 受掩码语言建模在语言任务中的成功启发,BERT4Rec[32]和 Transformer s4Rec[6]采用掩码策略的 Transformer 模型处理序列推荐任务
    • S³-Rec[44]进一步通过四项自监督任务预训练以提升数据表示
  • 上述模型通过学习高维 item 嵌入,并在最大内积搜索(MIPS)空间中进行近似最近邻搜索(ANN)来预测 next item。相比之下,论文提出的TIGER采用生成式检索直接预测 next item 的语义 ID
  • P5[8]通过微调预训练大语言模型实现多任务推荐系统。P5依赖LLM分词器(SentencePiece[29])从随机分配的 item ID *生成 *token ,而论文则基于 item 内容信息学习语义 ID 表示。实验(表2)表明,基于语义 ID 的推荐系统性能显著优于随机编码
  • *语义 ID *相关研究:
    • Hou等人提出的 VQ-Rec[12]利用内容信息生成“编码”(类似于语义 ID)表示 item,但其重点在于构建可迁移推荐系统,并未将编码用于生成式检索。虽然他们也使用乘积量化[15]生成编码,但论文采用RQ-VAE生成语义 ID,从而获得 item 的分层表示(第4.2节)
    • Singh等人[31]的并行研究表明,分层语义 ID 可替代大规模推荐系统中排名模型的 item ID,从而提升模型泛化能力
  • 生成式检索相关研究:
    • 尽管过去已提出学习搜索索引的技术[20],生成式检索是最近发展的一种文档检索方法,其任务是从数据库中返回相关文档集。例如GENRE[5]、DSI[34]、NCI[37]和CGR[22]。更详细的相关工作见附录A
  • 据论文所知,论文是首个提出基于语义 ID **的生成式检索**推荐系统的研究

TIGER整体框架

  • 论文的框架包含两个阶段:
    • 基于内容特征的语义 ID 生成 :将 item 内容特征编码为嵌入向量,并通过量化生成语义码元组,称为 item 的语义 ID 。
    • 基于语义 ID 训练生成式推荐系统 :使用语义 ID 序列训练 Transformer 模型完成序列推荐任务

语义 ID 生成

  • 如 图2(a) 所示,先将与每个 item 相关的内容特征通过预训练的 encoder 生成语义 embedding ,然后对语义 embedding 进行量化(quantized)以生成*语义 ID *
    • 这里假设每个 item 具有关联的内容特征(如标题、描述或图像),可通过预训练的内容 Encoder生成语义 embedding \(\boldsymbol{x} \in \mathbb{R}^d\)(比如,这里用通用预训练文本 Encoder (如Sentence-T5或BERT)可将文本特征转换为语义 embedding)
    • 理解:不同物品的描述应该尽量不同,且文本内容比较丰富的场景才可以吧?
  • 语义 ID 定义为长度为 \(m\) 的码元组(a tuple of codewords),每个 codeword 来自不同的 codebook(码本)
    • 语义 ID s 可唯一表示的 item 数量为各 codebook 大小的乘积
      • 理解:比如语义序列长度是3,codebook 的大小是 \(K\),则可以编码的 item 的数量是 \(K\times K \times K\)
    • 不同生成技术赋予语义 ID 不同的语义属性,但至少需满足以下性质:*相似 item (内容特征相似或语义 embedding 接近)应具有重叠的语义 ID *
      • 比如:语义 ID 为 \((10,21,35)\) 的 item 应比 \((10,23,32)\) 更接近 \((10,21,40)\)
基于RQ-VAE的语义 ID
  • 残差量化编码器(Residual-Quantized Variational AutoEncoder,RQ-VAE)是一种多级向量量化器,通过对残差进行量化生成码元组(即语义 ID)。通过联合训练更新量化 codebook 和DNN Encoder-Decoder 参数。图3展示了通过残差量化生成语义 ID 的过程
  • RQ-VAE首先通过 Encoder \(\mathcal{E}\) 将输入 \(\boldsymbol{x}\) 编码为潜在表示 \(\boldsymbol{z} := \mathcal{E}(\boldsymbol{x})\)
    • 在第零级(\(d=0\)),初始残差定义为 \(\boldsymbol{r}_0 := \boldsymbol{z}\)
    • 每一级 \(d\) 对应一个 codebook \(\mathcal{C}_d := \{\boldsymbol{e}_k\}_{k=1}^K\),其中 \(K\) 为 codebook 大小
    • \(\boldsymbol{r}_0\) 通过映射到该级 codebook 的最近嵌入进行量化
  • 具体残差量化的逻辑为:
    • 第零级最近嵌入 \(\boldsymbol{e}_{c_0}\) 的索引 \(c_0 = \arg\min_i |\boldsymbol{r}_0 - \boldsymbol{e}_k|\) 表示第零个码
    • 下一级 \(d=1\) 的残差定义为 \(\boldsymbol{r}_1 := \boldsymbol{r}_0 - \boldsymbol{e}_{c_0}\),并通过相同过程计算第一级的码
    • 该过程递归重复 \(m\) 次,生成表示语义 ID 的 \(m\) 元组
    • 这种递归方法实现了从粗到细的输入近似,语义 ID \((c_0, \ldots, c_{m-1})\) 的量化表示为 \(\widehat{\boldsymbol{z} } := \sum_{d=0}^{m-1} \boldsymbol{e}_{c_i}\),随后通过 Decoder 重建输入 \(\boldsymbol{x}\)
  • RQ-VAE的损失函数定义为:
    $$\mathcal{L}(\boldsymbol{x}) := \mathcal{L}_{\text{recon} } + \mathcal{L}_{\text{vq} }$$
    • 其中:
      $$\mathcal{L}_{\text{recon} } := |\boldsymbol{x} - \widehat{\boldsymbol{x} }|^2\\
      \mathcal{L}_{\text{vq} } := \sum_{d=0}^{m-1} |\text{sg}[r_i] - \boldsymbol{e}_{c_i}|^2 + \beta|r_i - \text{sg}[\boldsymbol{e}_{c_i}]|^2$$
      • 这里 \(\widehat{\boldsymbol{x} }\) 为 Decoder 输出,\(\text{sg}\) 为 stop gradient 操作。该损失函数联合训练 Encoder、 Decoder 和 codebook
  • 缓解Codebook Collapse 问题 :为防止 Codebook Collapse(码本坍塌)问题(即大多数输入映射到少数 codebook 向量),论文采用基于k-means聚类的 codebook 初始化方法。具体而言,对首训练批次应用k-means算法,并将质心作为初始化
其他量化方法
  • 局部敏感哈希(Locality Sensitive Hashing,LSH):生成语义 ID 的简单替代方法是局部敏感哈希(LSH),第4.2节的消融研究表明,RQ-VAE 性能优于 LSH
  • 分层k-means聚类 :另一种方法是 分层 k-means 聚类,但其会丢失不同聚类间的语义关联[37]
  • VQ-VAE :论文还尝试了 VQ-VAE,尽管在检索候选生成中性能与RQ-VAE相近,但丢失了语义 ID 的分层特性(第4.3节讨论)
处理冲突(Handling Collisions)
  • 冲突说明 :根据语义 embedding 分布、 codebook 大小和码长度,可能发生语义冲突(即多个 item 映射到同一语义 ID)
  • 解决方案 :为消除冲突,论文在有序语义码后附加额外 token 以确保唯一性
    • 例如,若两 item 共享语义 ID \((12,24,52)\),则通过附加 token 区分,表示为 \((12,24,52,0)\) 和 \((12,24,52,1)\)
    • 注意:冲突检测和修复仅在 RQ-VAE 训练完成后进行一次
  • 因为语义 ID s 是整数的元组,所以其查表(lookup table)效率比高维 embedding高

基于语义 ID 的生成式检索

  • 论文按时间顺序为用户构造 item 序列。给定序列 \((\text{item}_1, \ldots, \text{item}_n)\),推荐系统的任务是预测 next item( \(\text{item}_{n+1}\))。论文提出一种直接预测 next item 语义 ID 的生成式方法
  • 设 \((c_{i,0}, \ldots, c_{i,m-1})\) 为 \(\text{item}_i\) 的语义 ID 序列,将 item 序列转换为 \((c_{1,0}, \ldots, c_{1,m-1}, c_{2,0}, \ldots, c_{2,m-1}, \ldots, c_{n,0}, \ldots, c_{n,m-1})\),sequence-to-sequence 模型训练目标为预测 \(\text{item}_{n+1}\) 的语义 ID \((c_{n+1,0}, \ldots, c_{n+1,m-1})\)。由于生成性质,Decoder 生成的语义 ID 可能不匹配推荐语料库中的任何 item
    • 如附录图6所示,此类事件发生概率较低,附录E进一步讨论了处理方法

Experiments

  • 数据集 :论文在来自亚马逊产品评论数据集[10]的三个公开真实世界基准上评估所提出的框架,该数据集包含从1996年5月至2014年7月的用户评论和物品元数据。具体而言,论文使用亚马逊产品评论数据集中的三个类别进行序列推荐任务:“Beauty(美容)”、“Sports and Outdoors(运动与户外)”以及““Toys and Games(玩具与游戏)”
    • 论文在附录C中讨论了数据集的统计信息和预处理方法
  • 评估指标 :论文使用Top-K召回率(Recall@K)和归一化折损累积增益(NDCG@K),其中 \(K = 5, 10\) 来评估推荐性能
  • RQ-VAE实现细节 :如第3.1节所述,RQ-VAE 用于量化物品的语义 embedding。论文使用预训练的Sentence-T5[27]模型获取数据集中每个物品的语义 embedding。具体而言,论文利用物品的内容特征(如标题、价格、品牌和类别)构建句子,然后将其输入预训练的Sentence-T5模型,得到物品的768维语义 embedding
  • RQ-VAE模型包含三个组件:一个DNN Encoder ,将输入的语义 embedding 编码为潜在表示;残差量化器 ,输出量化表示;以及一个DNN Decoder ,将量化表示解码回语义输入嵌入空间,实验设置如下:
    • Encoder 具有三个中间层,大小分别为512、256和128,使用ReLU激活函数
    • 最终的潜在表示维度为32,为了量化这一表示,进行了三级残差量化。每一级维护一个基数为256的 codebook,其中 codebook 中的每个向量维度为32。计算总损失时,论文使用 \(\beta = 0.25\)
    • RQ-VAE 模型训练了 20k 个周期,以确保 codebook 使用率高(\(\geq 80%\))
    • 论文使用 Adagrad 优化器,学习率为 0.4,批量大小为 1024
    • 训练完成后,论文使用学习到的 Encoder 和量化组件为每个物品生成一个 3 元组语义 ID 。为避免多个物品映射到相同的语义 ID,论文为共享前三个码字的物品添加一个唯一的第4个码字,例如,两个与元组(7,1,4)关联的物品分别被分配(7,1,4,0)和(7,1,4,1)(如果没有冲突,论文仍然将第4个码字分配为0)。这为推荐语料库中的每个物品生成了一个长度为4的唯一语义 ID
  • sequence-to-sequence 模型实现细节 :
    • 论文使用开源的TSX框架[28]实现基于 Transformer 的 Encoder-Decoder 架构
    • 为了使模型能够处理序列推荐任务的输入,sequence-to-sequence 模型的词汇表包含每个语义码字的 token。具体而言,词汇表包含1024(\(256 \times 4\))个 token,用于表示语料库中的物品
    • 除了物品的语义码字外,论文还向词汇表中添加了用户特定的 token。为了限制词汇表大小,论文仅为用户 ID 添加了2000个 token。论文使用哈希技巧将原始用户 ID 映射到2000个用户 ID token 之一。论文构建输入序列为用户 ID token,后跟对应用户物品交互历史的语义 ID token 序列。论文发现,在输入中添加用户 ID **可以让模型个性化检索到的物品**
      • 问题:用户 token 如何生成?
    • 论文为基于 Transformer 的 Encoder 和 Decoder 模型各使用4层,每层有6个自注意力头,维度为64
    • 所有层均使用ReLU激活函数
    • MLP和输入维度分别设置为1024和128
    • 论文使用了0.1的dropout
    • 整个模型约有1300万个参数
    • 论文在“美容”和“运动与户外”数据集上训练该模型20万步。由于“玩具与游戏”数据集较小,仅训练了10万步
    • Batch Size 为 256
    • 学习率在前1万步为0.01,之后遵循逆平方根衰减规则(decay schedule)

序列推荐性能

  • 本节中,论文将提出的生成式检索框架与以下序列推荐方法进行比较(附录B中简要描述了这些方法):GRU4Rec[11]、Caser[33]、HGN[25]、SASRec[17]、BERT4Rec[32]、FDSA[42]、S\(^3\)-Rec[44]和P5[8]
    • 所有基线方法(除P5外)均使用双 Encoder 学习高维向量空间,其中用户的过去物品交互和候选物品被编码为高维表示,并使用最大内积搜索(MIPS)检索用户可能交互的下一个候选物品
    • 论文的新型生成式检索框架直接通过 sequence-to-sequence 模型逐 token 预测物品的语义 ID
  • 推荐性能 :论文对提出的TIGER在序列推荐任务上进行了广泛分析,并与上述基线进行比较。除P5外,所有基线的结果均来自Zhou等人[44]公开可用的结果(详情见github.com/aHuiWang/CIKM2020-S3Rec)。对于P5,论文使用了作者提供的源代码。但为了公平比较,论文更新了数据预处理方法以与其他基线和论文的方法保持一致。附录D中提供了更多相关细节
  • 结果如表1所示。论文观察到,TIGER在所有三个基准上均显著优于现有基线(如表9所示,TIGER在不同数据集上的标准误差也不大(3个种子计算得到))。特别是在“美容”基准上,TIGER表现尤为突出,NDCG@5比SASRec提高了29%,Recall@5比S\(^3\)-Rec提高了17.3%。在“玩具与游戏”数据集上,TIGER的NDCG@5和NDCG@10分别提高了21%和15%

物品表示

  • 本节中,论文分析了RQ-VAE语义 ID 的几个重要特性
    • 首先进行定性分析,观察语义 ID 的层次结构
    • 接着,通过对比另一种基于哈希的量化方法,评估RQ-VAE在框架中的重要性
    • 最后,论文通过消融实验研究语义 ID 的重要性,将TIGER与使用随机 ID 表示物品的 sequence-to-sequence 模型进行比较
  • 定性分析 :
    • 论文在图4中分析了为亚马逊“美容”数据集学习的RQ-VAE语义 ID 。为了说明,论文将RQ-VAE的级别数设为3,codebook 大小分别为4、16和256,即对于物品的语义 ID \((c_1, c_2, c_3)\),有 \(0 \leq c_1 \leq 3\),\(0 \leq c_2 \leq 15\) 和 \(0 \leq c_3 \leq 255\)。在图4a中,论文使用 \(c_1\) 标注每个物品的类别,以可视化数据集中整体类别分布中 \(c_1\) 特定的类别。如图4a所示,\(c_1\) 捕捉了物品的高级类别。例如,\(c_1 = 3\) 包含大部分与“头发”相关的产品,而 \(c_1 = 1\) 的物品主要是面部、唇部和眼部的“化妆品”和“护肤”产品
    • 论文通过固定 \(c_1\) 并可视化所有可能的 \(c_2\) 值的类别分布,进一步展示RQ-VAE语义 ID 的层次结构(图4b)。论文发现第二个码字 \(c_2\) 进一步将 \(c_1\) 捕获的高级语义细化为更精细的类别。RQ-VAE学习的语义 ID 的层次结构开启了广泛的新能力,这些能力将在第4.3节讨论。与现有基于随机原子 ID 学习物品嵌入的推荐系统不同,TIGER使用语义 ID,其中语义相似的物品具有重叠的码字,这使得模型能够有效共享数据集中语义相似物品的知识
  • *哈希与RQ-VAE语义 ID * :
    • 论文通过将RQ-VAE与局部敏感哈希(LSH)[14,13,2]在语义 ID 生成上的表现进行对比,研究RQ-VAE在框架中的重要性。LSH是一种流行的哈希技术,可以轻松适应论文的设置。为了生成LSH语义 ID,论文使用 \(h\) 个随机超平面 \(\boldsymbol{w}_1, \ldots, \boldsymbol{w}_h\) 对嵌入向量 \(\boldsymbol{x}\) 进行随机投影,并计算以下二进制向量:\((1_{\boldsymbol{w}_1^\top \boldsymbol{x} > 0}, \ldots, 1_{\boldsymbol{w}_h^\top \boldsymbol{x} > 0})\)。将此向量转换为整数码 \(c_0 = \sum_{i=1}^h 2^{i-1} 1_{\boldsymbol{w}_i^\top \boldsymbol{x} > 0}\)。使用独立的随机超平面集合重复此过程 \(m\) 次,得到 \(m\) 个码字 \((c_0, c_1, \ldots, c_{m-1})\),论文称之为LSH语义 ID
    • 在表2中,论文将LSH语义 ID 与提出的RQ-VAE语义 ID 的性能进行了比较。在此实验中,对于LSH语义 ID,论文使用 \(h = 8\) 个随机超平面,并设置 \(m = 4\) 以确保与RQ-VAE的基数相当。超平面的参数从标准正态分布中随机采样,以确保超平面是球对称的。论文的结果表明,RQ-VAE始终优于LSH。这说明在给定相同的内容语义 embedding 下,通过非线性深度神经网络(DNN)架构学习语义 ID 比使用随机投影产生更好的量化效果
  • *随机 ID 与语义 ID * :论文还比较了语义 ID 在生成式检索推荐系统中的重要性。具体而言,论文将随机生成的 ID 与语义 ID 进行比较。为了生成随机 ID 基线,论文为每个物品分配 \(m\) 个随机码字。物品的长度为 \(m\) 的随机 ID 为 \((c_1, \ldots, c_m)\),其中 \(c_i\) 从 \(\{1, 2, \ldots, K\}\) 中均匀随机采样。论文为随机 ID 基线设置 \(m = 4\) 和 \(K = 255\),以使基数与RQ-VAE语义 ID 相似。表2展示了随机 ID 与RQ-VAE和LSH语义 ID 的比较。论文发现语义 ID 始终优于随机 ID 基线,突出了利用基于内容的语义信息的重要性

新能力

  • 论文描述了两种直接源自提出的生成式检索框架的新能力,即冷启动推荐和推荐多样性。论文将这些能力称为“新”能力,因为现有的序列推荐模型(参见第4.1节中的基线)无法直接满足这些实际用例。这些能力源于RQ-VAE语义 ID 与框架生成式检索方法的协同作用。论文将在以下部分讨论TIGER在这些场景中的应用
  • 冷启动推荐 :
    • 本节中,论文研究了提出框架的冷启动推荐能力。由于现实世界推荐语料库的快速变化,新物品不断被引入。由于新增物品在训练语料库中缺乏用户印象,使用随机原子 ID 表示物品的现有推荐模型无法将新物品作为潜在候选检索出来。相比之下,TIGER框架可以轻松执行冷启动推荐,因为它在预测下一个物品时利用了物品的语义信息
    • 在此分析中,论文考虑了亚马逊评论中的“美容”数据集。为了模拟新增物品,论文从训练数据分割中删除了5%的测试物品。论文将这些被删除的物品称为“未见物品”。从训练分割中删除物品确保了未见物品没有数据泄漏。如前所述,论文使用长度为4的语义 ID 表示物品,其中前3个 token 由RQ-VAE生成,第4个 token 用于确保所有已见物品具有唯一 ID 。论文在训练分割上训练RQ-VAE量化器和 sequence-to-sequence 模型。训练完成后,论文使用RQ-VAE模型为数据集中的所有物品生成语义 ID,包括物品语料库中的任何未见物品
    • 给定模型预测的语义 ID \((c_1, c_2, c_3, c_4)\),论文检索具有相同 ID 的已见物品。根据定义,模型预测的每个语义 ID 最多匹配训练数据集中的一个物品。此外,将具有相同前三个语义 token 的未见物品(即 \((c_1, c_2, c_3)\))添加到检索候选列表中。最后,在检索Top-K候选集时,论文引入了一个超参数 \(\epsilon\),它指定了框架选择的未见物品的最大比例
    • 论文在图5中将TIGER与k近邻(KNN)方法在冷启动推荐设置下的性能进行了比较。对于KNN,论文使用语义表示空间执行最近邻搜索。论文将基于KNN的基线称为Semantic_KNN。图5a显示,论文的框架在 \(\epsilon = 0.1\) 时在所有Recall@K指标上始终优于Semantic_KNN。在图5b中,论文比较了不同 \(\epsilon\) 值下论文的方法与Semantic_KNN的表现。对于所有 \(\epsilon \geq 0.1\) 的设置,论文的方法均优于基线
  • 推荐多样性 :
    • 虽然召回率和NDCG是评估推荐系统的主要指标,但预测的多样性是另一个关键目标。多样性差的推荐系统可能对用户的长期参与产生不利影响。在此,论文讨论了如何利用生成式检索框架预测多样化的物品。论文展示了在解码过程中基于温度的采样可以有效地控制模型预测的多样性。虽然基于温度的采样可以应用于任何现有的推荐模型,但由于RQ-VAE语义 ID 的特性,TIGER允许在不同层次上进行采样。例如,采样语义 ID 的第一个 token 允许从粗粒度类别中检索物品,而采样第二个或第三个 token 则允许在类别内采样物品
    • 论文使用Entropy@K指标定量测量预测的多样性,其中熵是根据模型预测的Top-K物品的真实类别分布计算的。论文在表3中报告了不同温度值下的Entropy@K。论文观察到,在解码阶段使用温度采样可以有效增加物品真实类别的多样性。论文还在表4中进行了定性分析

消融研究

  • 论文在表5中测量了 sequence-to-sequence 模型中层数变化的影响。论文发现随着网络变大,指标略有提升
  • 论文还测量了提供用户信息的效果,结果见附录中的表8

无效 ID

  • 由于模型自回归地解码目标语义 ID 的码字,可能会预测无效 ID (即未映射到推荐数据集中任何物品的 ID)

  • 在论文的实验中,论文使用了长度为4的语义 ID,每个码字的基数为256(即每级 codebook 大小为256),这个组合可能的 ID 数量为 \(256^4\),约为4万亿。而论文考虑的数据集中的物品数量为10K-20K(见表6)。尽管有效 ID 数量仅占完整 ID 空间的一小部分,但论文观察到模型几乎总是预测有效 ID 。论文在图6中可视化了TIGER生成的无效 ID 比例随检索物品数量 \(K\) 的变化。对于Top-10预测,三个数据集的无效 ID 比例在约0.1%至1.6%之间变化。为了抵消无效 ID 的影响并始终获得Top-10有效 ID,我们可以增加束大小并过滤无效 ID (注:尽管生成无效 ID,TIGER在仍优于已有序列推荐方法)

  • 前缀匹配(未来工作) :处理无效 token 的一种扩展方法是在模型生成无效 token 时进行前缀匹配。语义 ID 的前缀匹配允许检索与模型生成的 token 具有相似语义含义的物品。鉴于RQ-VAE token 的层次结构,前缀匹配可以被视为模型预测物品类别而非物品索引。这种扩展可能会进一步提高召回率/NDCG指标。论文将此类扩展留作未来工作


结论

  • 论文提出了一种新颖的范式——TIGER,通过生成式模型在推荐系统中检索候选 item。该方法的核心是一种新颖的 item 语义 ID 表示,它利用分层量化器(RQ-VAE)对内容嵌入进行编码,生成构成语义 ID 的 token。论文的框架训练出的模型无需创建索引即可直接用于服务—— Transformer 的记忆模块充当了 item 的语义索引。值得注意的是,论文的嵌入表基数不会随 item 空间的基数线性增长,这与那些需要在训练时创建大型嵌入表或为每个 item 生成索引的系统相比具有优势。通过在三个数据集上的实验,论文证明了该模型能够实现 SOTA 检索性能,同时能够泛化到新的未见过的 item

附录A 相关工作(续)

  • 生成式检索 :

    • 传统的文档检索通常需要训练一个双塔模型,将 query 和文档映射到同一个高维向量空间,然后通过近似最近邻搜索(ANN)或最大内积搜索(MIPS)返回与 query 最接近的文档。这种方法存在一些缺点,例如需要存储庞大的嵌入表[22, 23]。生成式检索是最近提出的一种技术,旨在通过逐词生成文档标题、名称或文档 ID 字符串来解决传统方法的部分问题
    • Cao等人[5]提出了GENRE用于实体检索,该方法基于 Transformer 架构,逐词返回 query 中引用的实体名称
    • Tay等人[34]提出了DSI用于文档检索,这是首个为每个文档分配结构化语义DocID 的系统。给定 query 时,模型自回归地逐词返回文档的DocID 。DSI的工作标志着信息检索领域向生成式检索方法的范式转变,也是 Transformer 首次成功应用于端到端检索任务。随后,
    • Lee等人[23]证明生成式文档检索在多跳场景中同样有效,即当复杂 query 无法通过单一文档直接回答时,模型可以通过链式思考生成中间 query,最终输出复杂 query 的结果
    • Wang等人[37]在Tay等人[34]基于层次化\(k\)-均值聚类的语义DocID 基础上,提出了一种新的 Decoder 架构,专门考虑了语义DocID 中的前缀
    • 在CGR[22]中,作者提出了一种结合双 Encoder 技术和生成式检索技术的方法,通过让 Encoder-Decoder 模型的 Decoder 学习独立的上下文嵌入,从而内在地存储文档信息
  • 向量量化 :

    • 向量量化是指将高维向量转换为低维码字元组的过程
    • 最直接的技术之一是层次化聚类,例如[34]中使用的方法,即在每次迭代中进一步划分子簇
    • 另一种流行的方法是向量量化编码器(VQ-VAE),由[35]提出,用于将自然图像编码为一系列码字。该技术首先通过 Encoder 降低输入向量(或图像)的维度,然后将低维向量分区,每个分区分别量化,从而生成一系列码字:每个分区对应一个码字。这些码字随后被 Decoder 用于重建原始向量(或图像)
    • RQ-VAE[40, 21]对VQ-VAE Encoder 的输出应用残差量化,以实现更低的重建误差
    • 局部敏感哈希(Locality Sensitive Hashing,LSH)[14, 13]是一种流行的聚类和近似最近邻搜索技术。论文中用于聚类的特定版本是SimHash[2],它使用随机超平面生成二进制向量作为项的哈希值。由于其计算复杂度低且可扩展[13],论文将其作为向量量化的基线方法

附录B 基线方法

  • 基线方法包括以下几个模型:
    • GRU4Rec[11]:首个基于RNN的方法,使用定制的GRU处理序列推荐任务
    • Caser[33]:使用CNN架构捕获高阶马尔可夫链,通过水平和垂直卷积操作实现序列推荐
    • HGN[25]:层次化门控网络通过新的门控架构捕捉用户的长期和短期兴趣
    • SASRec[17]:自注意力序列推荐使用因果掩码 Transformer 建模用户的序列交互
    • BERT4Rec[32]:通过双向自注意力 Transformer 解决单向架构的局限性,用于推荐任务
    • FDSA[42]:特征级深度自注意力网络将物品特征与物品嵌入一起作为 Transformer 输入序列的一部分
    • S\(^3\)-Rec[44]:序列推荐的自监督学习通过自监督任务预训练双向 Transformer 以改进序列推荐
    • P5[8]:P5是一种新方法,使用预训练的 LLM 将不同推荐任务统一到单一模型中

附录C 数据集统计

  • 论文使用Amazon产品评论数据集[10]中的三个公开基准,包含1996年5月至2014年7月的用户评论和物品元数据。针对序列推荐任务,论文选择了三个类别:“美妆”、“运动与户外”和“玩具与游戏”。表6总结了数据集的统计信息。论文按时间戳对用户的评论历史进行排序并过滤出评论少于5条的用户。遵循标准评估协议[17, 8],论文采用留一法进行评估。对于每个物品序列,最后一个物品用于测试,倒数第二个用于验证,其余用于训练。训练时,论文将用户历史中的物品数量限制为20

附录D 对P5数据预处理的修改

  • P5源代码5对Amazon数据集进行预处理,首先为每个用户创建会话,包含按时间排序的用户评论物品列表。生成会话后,原始物品 ID 被重新映射为整数 \(1, 2, 3,\cdots\)。因此,第一个会话的第一个物品 ID 为“1”,第二个未出现过的物品 ID 为“2”,以此类推。值得注意的是,此预处理方案在划分训练和测试集之前应用。这导致生成的序列数据中许多序列形式为\(a, a+1, a+2, \ldots\)。由于P5使用Sentence Piece分词器[30](参见[8]第4.1节),用户会话中的测试和训练物品可能共享子词,从而导致推理时信息泄露
  • 为解决泄露问题,论文改为为物品分配随机整数 ID,然后划分训练和评估集。P5的其余代码与论文提供的源代码保持一致。该数据集的结果在表7中以“P5”行显示。论文还从头实现了P5的一个版本,并仅针对序列推荐任务训练模型。论文的实现结果 token 为“P5-ours”。论文还在实现中验证,使用连续整数序列作为物品 ID 可以帮助论文获得与P5报告相当的或更好的指标

附录E 讨论

  • 语义 ID 长度和 codebook 大小的影响 :论文尝试调整语义 ID 长度和 codebook 大小,例如使用6个码字,每个 codebook 大小为64的 ID 。论文发现TIGER的推荐指标对这些变化具有鲁棒性。但需要注意的是,更长的 ID (即每个物品 ID 包含更多码字)会增加输入序列长度,从而增加基于 Transformer 的 sequence-to-sequence 模型的计算开销
  • 可扩展性 :为测试语义 ID 的可扩展性,论文进行了以下实验:合并三个数据集,并为所有物品生成语义 ID 。然后,论文在美妆数据集上使用这些语义 ID 进行推荐任务。将实验结果与仅基于美妆数据集生成语义 ID 的原始实验进行对比,结果如表10所示。论文发现性能仅有小幅下降
  • 推理成本 :尽管论文的模型在序列推荐任务中表现优异,但由于使用束搜索(Beam Search)进行自回归解码,其推理计算成本可能高于基于ANN的模型。论文强调,优化TIGER的计算效率并非本工作的主要目标。相反,论文的工作开辟了一个新的研究领域:基于生成式检索的推荐系统。未来工作中,论文将探索缩小模型规模或提升推理效率的其他方法
  • lookup table 的内存成本 :TIGER维护两个哈希 lookup table :物品 ID 到语义 ID 的表和语义 ID 到物品 ID 的表
    • 这两个表仅在生成后冻结一次:它们在RQ-VAE语义 ID 生成模型训练完成后生成,随后在 sequence-to-sequence Transformer 模型训练期间冻结
    • 每个语义 ID **包含4个整数** ,每个整数存储为 8 bits(1 Byte) ,总计32 bits(4 Byte整数)。每个物品用32 bits整数 ID 表示。因此,每个 lookup table 的大小约为\(64N\)位,其中\(N\)为数据集中物品数量
      • 理解:对每个 lookup table 都需要包含 <key, value> 两项,故每个 lookup table 的大小为 \( 64N = 2 \times 32N \)
  • 嵌入表的内存成本 :与传统推荐系统相比,TIGER使用的嵌入表更小。传统推荐系统为每个物品存储一个嵌入,而TIGER仅为每个语义码字存储嵌入。实验中,论文使用4个码字(每个 codebook 大小为256)表示语义 ID,共1024(256×4)个嵌入。传统推荐系统的嵌入数量为\(N\)(数据集中物品数量)。实验中,\(N\)范围为10K至20K。因此,TIGER嵌入表的内存成本为\(1024d\)(\(d\)为嵌入维度),而传统推荐系统的嵌入 lookup table 内存成本为\(Nd\)

附录:一些个人思考

  • 用户历史行为类型利用 :是否可以对 item 按照曝光、点击和下单等动作进行分类,不同的动作增加一个标签来表示?
  • 数据既有偏差 :用户之前的行为与用户之前的排序机制有关,曝光数量不多,如何缓解这种偏差?增加随机流量?
    • 理解:这其实是一个长期以来的问题,生成式检索理论上应该能缓解这个问题
  • 用户 token 生成 :目前是直接做聚类映射吗?后续是否可以考虑也对用户进行一些描述,再使用 RQ-VAE 来生成?

RS——Meta-GRs-HSTU

  • 参考链接:
    • (HSTU)Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations, ICML 2024, Meta

整体说明

  • 当前推荐系统的特点 :
    • 大规模推荐系统的特点在于其依赖高基数(high cardinality,注:也称为高维)、异构特征(heterogeneous features),并且需要每天处理数百亿用户行为
    • 尽管大多数工业界的深度学习推荐模型(Deep Learning Recommendation Models,DLRMs)利用海量数据和数千个特征进行训练,却未能实现与计算资源的有效扩展(即增加了资源,效果得不到对应的提升)
  • 论文方案一句话说明 :(受 Transformer 在语言和视觉领域成功的启发)
    • 论文将推荐问题重新定义为生成式建模框架下的 Squential Transduction 任务(Generative Recommenders(GRs),生成式推荐器),并提出了一种新架构HSTU ,专为高基数、非平稳的流式推荐数据设计
    • 注:论文的 Squential Transduction 和 Transduction Learning(直推式学习)没有直接关系(虽然二者名字很像);
      • Transduction Learning(直推式学习) :利用训练数据和测试数据的整体信息,直接对见过的测试数据做出预测
      • Inductive Learning(归纳式学习) :利用训练数据的信息,训练一个可以泛化到未见过的测试数据的通用的模型
      • 从定义看,论文本质还是一个Inductive Learning(归纳式学习) ,学到的模型是要使用在未知数据场景的(包括未知用户、未知时间和未知商品)
  • 效果和效率 :
    • 离线效果 :在合成和公开数据集上,HSTU相比基线模型在 NDCG 指标上最高提升 65.8%
    • 效率 :并且在处理 8192 长度序列时比基于 FlashAttention2 的 Transformer 快 5.3 倍至 15.2 倍
    • 在线效果 :基于 HSTU 的生成式推荐器拥有1.5万亿参数,在线A/B测试中指标提升12.4%,已在拥有数十亿用户的互联网平台的多个场景中部署
  • scaling law发现 :更重要的是,生成式推荐器的模型质量随训练计算量呈幂律增长(power-law),扩展范围达到GPT-3/LLaMa-2规模,这减少了未来模型开发所需的 carbon footprint(碳排放总和),并为推荐系统中的首个基础模型铺平了道路
    • 理解:power-law说明模型的算力投入可以很高效的得到回报,所以不需要盲目增加很多计算资源,也就减少了碳排放
  • 注意:论文的方法依然保留了推荐系统中的 Retrieval-Ranking 两阶段(Retrieval也称为召回或 Retrieval 阶段),只是分别将两阶段都使用生成式模型来建模了,对于召回阶段,甚至可以将生成式模型作为一个新的召回通道使用,作为对原始 DLRM 的一种补足

一些背景讨论

  • 传统推荐系统 : SOTA 推荐方法主要基于深度学习推荐模型 (DLRM),其特点在于使用异构特征,包括数值特征(如 counters 和 ratios)、嵌入以及分类特征(如商品ID、用户ID等)
    • 由于每分钟都有新内容和商品加入,特征空间具有极高的基数,通常达到十亿级别
    • 为了利用这些数以万计的特征,DLRM 采用各种神经网络组合特征、转换中间表示并生成最终输出
    • 尽管 DLRM 利用了大量人工设计的特征和海量数据进行训练,但工业界大多数 DLRM 的计算扩展性较差。这一显著限制至今仍未得到解决
  • 论文的思路 :受 Transformer 在语言和视觉领域成功的启发,论文重新审视了现代推荐系统中的基础设计选择,但作者观察到:在十亿用户规模下,这需要克服三个挑战 :
    • 推荐系统中的特征缺乏明确结构 :虽然在小规模场景中已探索过序列化方法(附录B详细讨论),但工业级 DLRM 中异构特征(包括高基数ID、交叉特征、计数器、比率等)起着关键作用
    • 持续变化的十亿级词汇表 :与语言模型中十万级静态词汇表相比,十亿级动态词汇表带来了训练挑战,并由于需要考虑数万个目标感知(target-aware)候选而推高了推理成本
      • 目标感知(target-aware)指生成用户表示或进行预测时,能够动态地结合当前候选内容(target item)的信息,从而更精准地建模用户与特定候选内容之间的交互
      • 注:传统推荐系统(如DLRMs)中,target-aware通常指在特征交互阶段显式引入候选内容(target item)的信息
    • 计算成本大 :(这是大规模序列模型落地的主要瓶颈),GPT-3使用数千个GPU在1-2个月内处理了总计3000亿 token。这一规模看似惊人,但与用户行为规模相比则相形见绌。最大规模的互联网平台每天服务数十亿日活用户,用户每天与数十亿帖子、图片和视频互动。用户序列长度可达 \(10^5\) 。因此,推荐系统每天需要处理的 token 数量比语言模型1-2个月处理的量高出几个数量级
  • 一个特别重要的创新点 :在本工作中,论文将用户行为视为生成式建模中的新模态(在图片、文本和视频等多模态的基础上,增加一个模态叫做用户行为)。论文的 key insights 是:
    • a)给定适当的新特征空间,工业级推荐系统中的核心排序和召回任务可被转化为生成式建模问题;
    • b)这一范式使论文能够系统性地利用特征、训练和推理中的冗余(redundancies)来提高效率,基于这一新范式,论文部署的模型计算复杂度比之前最先进技术高出三个数量级(如图1所示),同时核心业务指标(英文中称为topline metrics)提升了12.4%
      • 问题:如何理解”特征、训练和推理中的冗余(redundancies)“这句话?
  • 论文的贡献可总结如下:
    • 第一 :提出生成式推荐器(Generative Recommenders,GRs) ,这一新范式将取代 DLRM
      • 论文将 DLRM 中的异构特征空间序列化并统一(注:用户行为视为特定的一个新模态),当序列长度趋近无穷时,新方法可近似完整的 DLRM 特征空间。使得:
        • 论文能将主要推荐问题(排序和召回)重新定义为GR中的纯 Squential Transduction 任务
        • 论文模型训练能以序列化、生成式的方式进行,从而允许论文在相同计算量下处理多几个数量级的数据
    • 第二 :解决训练和推理过程中的计算成本挑战 :
      • 提出新结构 :提出新的 Squential Transduction 架构——分层 Squential Transduction 单元(Hierarchical Sequential Transduction Units,HSTU)。HSTU针对大规模非平稳词汇表修改了注意力机制 ,并利用推荐数据集特性实现比基于 FlashAttention2 的 Transformer 在 8192 长度序列上快5.3倍至15.2倍
      • 提出新算法M-FALCON :通过新算法M-FALCON完全平摊计算成本,论文能在传统 DLRM 使用的相同推理预算下,服务复杂度高出285倍的GR模型,同时实现1.50-2.99倍的加速
    • 第三 :验证所提技术在合成数据集、公开数据集以及在拥有数十亿日活用户的大型互联网平台多个场景中的部署效果
      • 据论文所知,论文的工作首次展示了纯 Squential Transduction 架构(如HSTU)在生成式设置(GRs)中显著优于工业级大规模 DLRM。值得注意的是,论文不仅克服了传统 DLRM 中已知的扩展瓶颈,还进一步证明了扩展定律(scaling law)适用于推荐系统 ,这可能是推荐系统的”ChatGPT时刻“

推荐作为 Squential Transduction 任务:从 DLRM 到GR

统一 DLRM 中的异构特征空间

  • 现代 DLRM 模型通常使用大量分类(sparse)和数值(dense)特征进行训练。在GR中,论文将这些特征整合并编码为统一的时间序列 ,如图2所示
  • 分类(sparse)特征 :这类特征的例子包括用户喜欢的商品、用户关注的某类别(如户外)创作者、用户语言、用户加入的社区、请求发起的城市等。论文按以下方式将这些特征序列化:
    • 选择最长的时间序列(通常通过合并代表用户互动商品的特征)作为主时间序列
    • 其余特征通常是随时间缓慢变化的序列(如人口统计信息或关注的创作者),论文通过保留每个连续段的最早条目(entry)来压缩这些时间序列,然后将结果合并到主时间序列中。由于这些时间序列变化非常缓慢,这种方法不会显著增加总体序列长度
      • 理解:如图2所示,auxiliary time series 1 和 auxiliary time series 2 中,分别只有每个连续段的第一个entry被插入到 main time series 中(注:auxiliary time series 1 包含多个连续段,每个连续段的第一个 entry 都会保留并插入 main time series 中)
    • 这部分写得不是很清晰,有点晦涩,原文介绍如下:

      Categorical (‘sparse’) features. Examples of such features include items that user liked, creators in a category (e.g., Outdoors) that user is following, user languages, communities that user joined, cities from which requests were initiated, etc. We sequentialize these features as follows. We first select the longest time series, typically by merging the features that represent items user engaged with, as the main time series. The remaining features are generally time series that slowly change over time, such as demographics or followed creators. We compress these time series by keeping the earliest entry per consecutive segment and then merge the results into the main time series. Given these time series change very slowly, this approach does not significantly increase the overall sequence length.

  • 数值(dense)特征 :这类特征的例子包括加权和衰减计数器、比率等(例如,某个特征可能代表用户过去对匹配特定主题商品的点击率(CTR))
    • 与分类特征相比,这些特征变化更频繁 ,可能随每个(用户,商品)互动而变化,所以从计算和存储角度看,完全序列化这些特征并不可行
    • 然而一个重要观察是:论文执行这些聚合所基于的分类特征(如商品主题、位置)在GR中已被序列化和编码。因此 ,给定足够强大的 Squential Transduction 架构与 target-aware 公式相结合,当增加GR中的总体序列长度和计算量时,论文可以移除数值特征 ,因为它们能被有效地捕捉

将排序和召回重新定义为 Squential Transduction 任务

  • 给定按时间顺序排列的 \(n\) 个 token 列表 \(x_{0},x_{1},\ldots,x_{n-1}\) (\(x_{i}\in\mathbb{X}\)),观察到这些 token 的时间 \(t_{0},t_{1},\ldots,t_{n-1}\) ,Squential Transduction 任务将此输入序列映射到输出 token \(y_{0},y_{1},\ldots,y_{n-1}\) (\(y_{i}\in\mathbb{X}\cup\{\varnothing\}\)),其中 \(y_{i}=\varnothing\) 表示 \(y_{i}\) 未定义
  • 论文用 \(\Phi_{i}\in\mathbb{X}_{c}\) (\(\mathbb{X}_{c}\subseteq\mathbb{X}\))表示系统向用户提供的内容(即历史上展示给用户看过的内容 ,如图片或视频)。由于不断有新内容产生 , \(\mathbb{X}_{c}\) 和 \(\mathbb{X}\) 是非平稳的。用户可以用某个行为 \(a_{i}\) (如点赞、跳过、视频完播+分享) \(a_{i}\in\mathbb{X}\) 来回应 \(\Phi_{i}\) 。论文用 \(n_{c}\) 表示用户互动过的内容总数
    • 对 Retrieval 输入的理解:输入是 \(\{(\Phi_{0},a_0),(\Phi_{1},a_1),\cdots,(\Phi_{n_c-1},a_{n_c-1})\}\) ,表示历史上依次给用户分别展示了内容 \(\Phi_{i}\) 以后用户的行为 \(a_{i}\) 的元组序列
    • 对 Retrieval 输出的理解:每个元组 \((\Phi_{i},a_i)\) 对应输出一个目标值 \(\Phi_{i}^{\prime}\) :
      $$\Phi_{i}^{\prime}=
      \begin{cases}
      \Phi_{i}& \ a_i \ \text{ is positive}\\
      \varnothing& \ \text{ otherwise}
      \end{cases}$$
    • 对 Ranking 输入的理解:如表1所示,对输入 \(x_i\) 的建模是, \(\{\Phi_{0},a_0,\Phi_{1},a_1,\cdots,\Phi_{n_c-1},a_{n_c-1}\}\) 是历史行为序列,表示历史上依次给用户分别展示了内容 \(\Phi_{i}\) 以后用户的行为 \(a_{i}\)
    • 对 Ranking 输出的理解:如表1所示,输出和输入依次对应,输入为 \(\Phi_{i}\) 时预估目标是用户的行为 \(a_{i}\)
  • 在因果自回归设置中,标准排序和 Retrieval 任务可定义为 Squential Transduction 任务(表1)。论文得出以下观察:
    • 召回(Retrieval) :在推荐系统的召回阶段,论文学习 \(\Phi_{i+1}\in\mathbb{X}_{c}\) 上的分布 \(p(\Phi_{i+1}|u_{i})\) ,其中 \(u_{i}\) 是用户 \(i\) 的 token 表示。典型目标是选择 \(\arg\max_{\Phi\in\mathbb{X}_{c} }p(\Phi|u_{i})\) 以最大化某些奖励。这与标准自回归设置有两个不同:首先, \(x_{i}\) 的监督信号 \(y_{i}\) 不一定是 \(\Phi_{i+1}\) ,因为用户可能对 \(\Phi_{i+1}\) 做出负面回应。其次,当 \(x_{i+1}\) 代表非互动相关的分类特征(如人口统计信息)时, \(y_{i}\) 未定义
    • 排序(Ranking) :GR中的 Ranking 任务带来独特挑战,因为工业推荐系统通常需要“目标感知”(target-aware)公式。在此设置中,目标 \(\Phi_{i+1}\) 与历史特征的”交互”(interaction)需要尽早发生,这在标准自回归设置中不可行(理解:不建模动作时,”交互”发生较晚)。论文通过表1中的交错排列商品和行为来解决这一问题,使得 Ranking 任务可被公式化为 \(p(a_{i+1}|\Phi_{0},a_{0},\Phi_{1},a_{1},\ldots,\Phi_{i+1})\) (在分类特征之前)。实践中论文使用小型神经网络将在 \(\Phi_{i+1}\) 的输出转换为多任务预测。重要的是,这使论文能在单次传递中对所有 \(n_{c}\) 次互动应用 target-aware 交叉注意力
  • 补充:对论文 target-aware 实现方式的理解 :其实 Transformer 的 Attention 也有交叉功能,只要输入端有目标 item 和用户历史交互 item 即可,但论文所说的传统的自回归模型是指输入侧不包含目标 item 的情况,论文中,将动作也建模进去,则在输出对目标 item 的动作 token 时,自然就需要将目标 item 作为输入,从而也就实现了 target-aware 交叉注意力:
    • 传统自回归模型预估目标为 : \(p(\Phi_{i+1}|\Phi_{0},\Phi_{1},\ldots,\Phi_{i})\),目标物品 \(\Phi_{i+1}\) 与其他历史序列无交叉
      • 注:此时动作 \(a_i\) 信息可能会作为额外表征加入到 \(\Phi_{i}\)中,所以 表述为 \(p(\Phi_{i}|(\Phi_{0}, a_{0}), …, (\Phi_{i - 1}, a_{i - 1}))\) 也可以
    • 论文预估目标为 : \(p(a_{i+1}|\Phi_{0},a_{0},\Phi_{1},a_{1},\ldots,\Phi_{i+1})\),目标物品 \(\Phi_{i+1}\) 与其他历史序列有交叉

生成式训练

  • 工业推荐系统通常在流式设置中训练,其中每个样本在可用时被顺序处理。在此设置中,基于自注意力的 Squential Transduction 架构(如 Transformer)的总计算需求量级为 \(\sum_{i}n_{i}(n_{i}^{2}d+n_{i}d_{ff}d)\) ,其中 \(n_{i}\) 是用户 \(i\) 的 token 数量, \(d\) 是嵌入维度,设 \(N=\max_{i}n_{i}\) ,总体时间复杂度降至 \(O(N^{3}d+N^{2}d^{2})\) ,这对推荐场景来说成本过高
    • 注:更多对公示的理解看笔者补充的附录
  • 为应对训练 Squential Transduction 模型处理长序列的挑战,论文从传统曝光级(impression-level)训练转向生成式训练,将计算复杂度降低 \(O(N)\) 因子,如图2顶部所示。这样做可将编码器成本平摊到多个目标上。具体来说,当论文以 \(s_{u}(n_{i})\) 的速率采样第 \(i\) 个用户时,总训练成本现在按 \(\sum_{i}s_{u}(n_{i})n_{i}(n_{i}^{2}d+n_{i}d^{2})\) 缩放,通过设 \(s_{u}(n_{i})\) 为 \(1/n_{i}\) 可降至 \(O(N^{2}d+Nd^{2})\) 。在工业级系统中实现此采样的一种方法是在用户请求或会话结束时发出训练样本,使得 \(\hat{s_{u} }(n_{i})\propto 1/n_{i}\)
    • 理解:从「曝光级(impression-level)训练转向生成式训练」主要是强调从以前每个曝光都是一个样本,现在一个用户的所有行为组成一个样本
    • 问题:这里对用户进行采样不会带来效果损失吗?模型看不到部分用户是OK的吗?

面向生成式推荐的高性能自注意力编码器(A High Performance Self-Attention Encoder for Generative Recommendations)

  • 为了将生成式推荐系统(GRs)扩展至具有大规模非静态词表的工业级推荐场景,论文提出了一种新型编码器设计——分层序列转导单元(Hierarchical Sequential Transduction Unit, HSTU)

  • HSTU由多个通过残差连接堆叠的相同层构成,每层包含三个子层:Pointwise Projection(公式1)、Spatial Aggregation(公式2)和 Pointwise Transformation(公式3):

$$ U(X), V(X), Q(X), K(X) = \text{Split}(\phi_1(f_1(X))) \tag{1} $$

$$ A(X)V(X) = \phi_2\left(Q(X)K(X)^T + \text{rab}^{p,t}\right)V(X) \tag{2} $$

$$ Y(X) = f_2\left(\text{Norm}\left(A(X)V(X)\right) \odot U(X)\right) \tag{3} $$

  • 其中:
    • \( f_i(X) \) 表示多层感知机(MLP);为降低计算复杂度, \( f_1 \) 和 \( f_2 \) 使用单线性层 \( f_i(X) = W_i(X) + b_i \)
      • 注:还通过融合内核批量处理查询 \( Q(X) \) 、键 \( K(X) \) 、值 \( V(X) \) 和门控权重 \( U(X) \) (即一次性将Q,K,V和Q映射都做完)
    • \( \phi_1 \) 和 \( \phi_2 \) 为非线性激活函数,均采用SiLU(即Swish);
    • Norm 表示 LayerNorm
    • \( \text{rab}^{p,t} \) 为结合位置(\( p \))和时间(\( t \))信息的相对注意力偏置
  • HSTU的编码器设计允许用单一模块化块替代DLRMs中的异构模块。论文观察到DLRMs实际包含三个阶段:特征提取、特征交互和表征变换
    • 1:特征提取(Feature Extraction)阶段通过池化操作获取类别特征的嵌入表示,其高级形式可泛化为成对注意力和 target-aware 池化
      • 如 HSTU 层所实现,能够做到上述能力
    • 2:特征交互(Feature Interaction)是DLRMs的核心,常用方法包括因子分解机及其神经网络变体、高阶特征交互等
      • HSTU通过注意力池化特征直接与其他特征交互(即 \(\text{Norm}(A(X)V(X)) \odot U(X)\))替代传统特征交互模块
      • 该设计源于学习型MLP难以近似点积的挑战(Rendle等, 2020; Zhai等, 2023a),由于SILU作用于 \(U(X)\),此结构也可视为SwiGLU(Shazeer, 2020)的变体
    • 3:表征变换(Transformations of Representations)阶段通常采用 MoE 和路由机制处理异构用户群体,其核心思想是通过子网络 specialization 实现条件计算
      • HSTU中的逐点乘积操作(Element-wise dot products)本质上能在归一化因子范围内模拟MoE的门控机制
  • 以上实现中可以看出,HSTU Layer 实现了类似 Transformer Layer 的功能
    • \(X\) 是 HSTU 的输入, \(Y(X)\) 是 HSTU 的输出
    • \(A(X)\) 为输入 \(X\) 的注意力向量(注:这里没有做归一化),是带位置偏置 \(\text{rab}^{p,t}\) 的 \(Q,K\) 内积
  • HSTU的一层 与传统的 Transformer 相比,HSTU有以下不同:
    • HSTU 增加了一个 \(U(X)\),跟Q,K,V的 projection 在同一层处理,\(U(X)\) 会作为一个 point-wise的向量与 HSTU Attention 的结果相乘
    • HSTU 没有使用传统 Transformer 中的 Softmax Attention机制,放弃了 Attention 权重和为1的约束,文中提到,这能保留用户对 item 的动作强度信息(Softmax会使得这部分信息失真)
    • HSTU 没有使用 FFN 层,这一层被 \(f_2\left(\text{Norm}\left(A(X)V(X)\right) \odot U(X)\right)\) 替换了,实现了将 MLP 替换为 point-wise 乘法,速度更快
      • 理解1:这里更像是使用 \(U(X)\) 保留了原始特征,和 Attention 后的高阶交叉特征 point-wise 相乘,得到高阶交叉特征同时也保留低阶交叉;
      • 理解2:这里的操作可以看做是 对 \(U(X)\) 进行加权,如前文所述,像是门网络/个性化加权机制实现对原始 Embedding \(U(X)\) 的加权(可类比于MoE的思想)

逐点聚合注意力(Pointwise aggregated attention)

  • HSTU采用了一种新型的Pointwise aggregated attention(区别于传统的Softmax注意力)。其动机有二:
    • 第一:与目标相关的历史数据点数量是用户偏好强度的关键特征,而Softmax归一化会削弱这一信息(Softmax的归一化会将强度信息也归一化为比例,这会使得绝对值失真)
    • 第二:Softmax对噪声鲁棒,但在流式非静态词表场景中表现欠佳(频繁增加新的 item 会影响 Softmax)
  • 实验表明,在合成数据(基于 Dirichlet 过程生成的非静态词表流式数据)中,逐点聚合注意力机制相比Softmax注意力可将Hit Rate@10提升高达44.7%(见表2)

利用算法增加稀疏性

  • 在推荐系统中,用户历史序列的长度通常呈现偏态分布(skewed distribution,这里主要指大部分用户交互序列短,少部分用户交互序列长),导致输入序列具有稀疏性,尤其是在处理超长序列时。这种稀疏性可以被有效利用以显著提升编码器的效率。为此,论文开发了一种高效的GPU注意力核函数,其设计类似于[Dao等, 2022; Zhai等, 2023b],但实现了完全分组的注意力计算。这本质上将注意力计算转化为不同尺寸的分组GEMM(通用矩阵乘法)运算(详见附录G)。因此,HSTU中的自注意力变为内存受限(memory-bound)操作,其内存访问复杂度为 \(\Theta(\sum_{i}n_{i}^{2}d_{qk}^{2}R^{-1})\) ,其中 \(n_i\) 为样本 \(i\) 的序列长度, \(d_{qk}\) 为注意力维度, \(R\) 为寄存器大小。仅此一项改进即可带来2-5倍的吞吐量提升(详见第4.2节讨论)
  • 论文进一步通过随机长度(Stochastic Length, SL)算法增加用户历史序列的稀疏性。推荐场景中用户历史序列的一个关键特征是用户行为在时间上具有重复性,这种重复性体现在用户交互历史的多尺度行为模式中。这为论文提供了在不损失模型质量的前提下人为增加稀疏性的机会,从而显著降低编码器计算成本(其复杂度为 \(\Theta(\sum_{i}n_{i}^{2})\))
    • 问题:这种降低稀疏性的方式是有损的吧?此外,相同算力消耗下,随机长度是不是不如仅保留最近多个 item 的策略?
  • 假设用户 \(j\) 的历史序列为 \((x_i)_{i=0}^{n_{c,j} }\) ,其中 \(n_{c,j}\) 为用户交互过的内容数量。令 \(N_c = \max_j n_{c,j}\) , \((x_{i_k})_{k=0}^L\) 为从1原始序列 \((x_i)_{i=0}^{n_{c,j} }\) 中选取的长度为 \(L\) 的子序列。SL算法按以下方式选择输入序列:
    $$
    \begin{align}
    (x_i)_{i=0}^{n_{c,j} } \text{ if } n_{c,j} &\leq N_c^{\alpha/2}\\
    (x_{i_k})_{k=0}^{N_c^{\alpha/2} } \text{ if } n_{c,j} &> N_c^{\alpha/2}, \text{w/ probability } 1 - N_c^{\alpha}/n_{c,j}^2\\
    (x_i)_{i=0}^{n_{c,j} } \text{ if } n_{c,j} &> N_c^{\alpha/2}, \text{w/ probability } N_c^{\alpha}/n_{c,j}^2 \tag{4}
    \end{align}
    $$
    • 其中 “w/ probability xx” 是 “with probability”的简写,表示以 “xx” 概率采样
    • \(N_c^{\alpha/2}\) 表示 \(N_c\) 的 \(\alpha/2\) 次方, \(\alpha\) 是一个调节采样比例的数字
  • 公式4描述了Stochastic Length (SL)算法如何根据用户历史序列的长度(\(n_{c,j}\))动态调整子序列的采样策略:
    • 如果序列长度 \(n_{c,j} \leq N_c^{\alpha/2}\) ,则:
      • 直接使用完整序列
    • 如果序列长度 \(n_{c,j} > N_c^{\alpha/2}\) ,则
      • 以概率 \(1 - N_c^\alpha / n_{c,j}^2\) 选择长度为 \(N_c^{\alpha/2}\) 的子序列
      • 以概率 \(N_c^\alpha / n_{c,j}^2\) 仍使用完整序列
  • 该算法将注意力相关复杂度降低至 \(O(N_c^{\alpha}d) = O(N^{\alpha}d)\) ,其中 \(\alpha \in (1,2]\) 。关于子序列选择的更详细讨论见附录F.1。值得注意的是,将SL应用于训练能够实现高性价比的系统设计,因为训练的计算成本通常远高于推理
  • 表3展示了在具有30天用户历史的典型工业规模配置下,不同序列长度和 \(\alpha\) 值对应的稀疏性(定义见附录F)。其中对模型质量影响可忽略的设置用下划线和蓝色标出。标为“ \(\alpha=2.0\) ”的行表示未应用SL时的基准稀疏性。更低的 \(\alpha\) 值适用于更长的序列,论文测试的最长序列长度为8,192

最小化激活内存使用

  • 在推荐系统中,使用大的 Batch Size 训练对提升训练吞吐量和模型质量至关重要(理解:大的 Batch Size 更容易得到准确的梯度),因此,激活内存的使用成为扩展的主要瓶颈,这与通常以小的 Batch Size 训练且以参数内存为主的大型语言模型形成鲜明对比
  • 与 Transformer 相比,HSTU采用了一种简化且完全融合的设计,显著降低了激活内存的使用
    • 首先,HSTU将注意力外的线性层数量从 6 个减少到 2 个,这与近期通过逐元素门控减少MLP计算的工作[Hua等, 2022; Zhai等, 2023b]一致
    • 其次,HSTU将计算过程激进地融合为单一算子,包括公式(1)中的 \(\phi_1(f_1(\cdot))\) ,以及公式(3)中的层归一化、可选Dropout和输出MLP。这种简化设计将每层的激活内存使用量降低至 \(2d + 2d + 4hd_{qk} + 4hd_v + 2hd_v = 14d\) (以bfloat16格式存储)
    • 理解:这里的减少线性层数量主要是通过移除传统的FFN来实现(传统的一个FFN包含两个线性层),HSTU总体仅保留两个线性层,即输入层和输出层
    • 注:文章这里说的将线性层从 6 层降低到 2 层的表达应该是跳过了一些设计细节,比如网络结构中FFN的数量未明确
  • 作为对比, Transformer 在注意力后使用前馈层和Dropout(中间状态为 \(3hd_v\)),随后是一个逐点前馈块,包含层归一化、线性层、激活函数、线性层和Dropout,中间状态为 \(2d + 4d_{ff} + 2d + 1d = 4d + 4d_{ff}\) 。这里论文假设 \(hd_v \geq d\) 且 \(d_{ff} = 4d\) [Vaswani等, 2017; Touvron等, 2023a]。因此,在考虑输入和输入层归一化(\(4d\))以及QKV投影后,总激活状态为 \(33d\) 。HSTU的设计使其能够扩展到比 Transformer 深2倍以上的网络
  • 此外,用于表示词汇表的大规模原子ID也需要大量内存。对于包含100亿词汇、512维嵌入和Adam优化器的配置,仅存储嵌入和优化器状态(以fp32格式)就需要60TB内存。为缓解内存压力,论文采用了 row-wise AdamW 优化器并将优化器状态放置在DRAM中,从而将每个浮点数的HBM使用量从12字节减少到2字节
    • 注:论文给的引用中没有详细介绍 row-wise AdamW 优化器
    • HBM(High Bandwidth Memory,高带宽内存)和 DRAM(Dynamic Random Access Memory,动态随机存取存储器)都是计算机内存技术,HBM 将多个 DRAM 芯片堆叠在一起(HBM可以理解为更为昂贵的告诉内存),拥有比 DRAM 更高的数据传输性能
  • 总结来看 :HSTU通过以下设计显著降低内存占用:
    • 第一:将注意力外的线性层从6层减至2层,采用逐点门控减少MLP计算;
    • 第二:将计算融合为单一算子(如公式1中的 \( \phi_1(f_1(\cdot)) \) 和公式3中的层归一化、Dropout及输出MLP);
    • 第三:使用row-wise AdamW优化器,将优化器状态存储于DRAM,使每个浮点数在HBM的占用从12字节降至2字节
    • 最终,与 Transformer 相比,HSTU的激活内存占用从每层33d降至14d(bfloat16),支持构建2倍更深的网络

通过成本分摊扩展推理规模

  • 论文需要解决的最后一个挑战是推荐系统在服务时需要处理大量候选内容。论文主要关注 Ranking 任务 :
  • 在Retrieval 任务(无需论文关注)中,编码器成本完全可以分摊,且已有高效的算法支持基于量化(quantization)、哈希(hashing)或分区(partitioning)的最大内积搜索(MIPS),以及基于束搜索或分层检索的 non-MIPS 场景
    • MIPS 是 Maximum Inner Product Search(最大内积搜索)的缩写,主要用于高效地找到与给定查询向量内积最大的候选向量 ,在召回阶段,MIPS可以与量化、哈希或分区等技术结合使用,以提升效率
  • 对于Ranking 任务 ,候选内容数量可能高达数万。论文提出了一种算法M-FALCON(Microbatched-Fast Attention Leveraging Cacheable Operations)(可译为微批次-快速注意力利用可缓存操作),用于在输入序列长度为 \(n\) 时对 \(m\) 个候选内容进行推理
    • 批量推理 :在一次前向传播中,M-FALCON通过修改注意力掩码和 \(\text{rab}^{p,t}\) 偏置,使得 \(b_m\) 个候选内容的注意力操作完全一致,从而并行处理 \(b_m\) 个候选内容,此时原本的 \(b_m\) 次长度为 \(n\) 的推理转换为 1 次 \(b_m+n\) 的推理(注:这里有MTP的味道)
      • 这将交叉注意力的计算成本从 \(O(b_m n^2 d)\) 降低到 \(O((n + b_m)^2 d) = O(n^2 d)\) (当 \(b_m\) 相对于 \(n\) 可视为小常数时)
    • KV Caching机制 :论文还可以将全部 \(m\) 个候选内容划分为 \(\lceil m/b_m \rceil\) 个大小为 \(b_m\) 的微批次,以利用编码器级的KV缓存 ,从而在跨前向传播时降低成本,或在跨请求时最小化尾部延迟
    • 关于M-FALCON方法的更多细节见附录H
  • 总体而言,M-FALCON使得模型复杂度能够与传统DLRM排序阶段的候选数量线性扩展。论文成功地在典型的排序配置中(见第4.3节)应用了复杂度高28倍的目标感知交叉注意力模型 ,同时在恒定推理预算下实现了1.5倍至3倍的吞吐量提升
  • M-FALCON方法的整体结构图如下(From 附录):

Experiments

验证HSTU编码器的归纳假设

传统序列化设置
  • 论文首先在两个流行的推荐数据集(MovieLens和Amazon Reviews)上评估HSTU的性能。实验遵循文献中的序列化推荐设置,包括完全打乱和多轮次训练。基线采用 SOTA Transformer 实现SASRec(Kang & McAuley, 2018)。与近期工作(Dallmann等, 2021; Zhai等, 2023a)一致,论文报告整个语料库的Hit Rate@K和NDCG@K
  • 结果如表4所示。“SASRec (2023)”表示(Zhai等, 2023a)中报告的最佳SASRec配置。“HSTU”行使用与SASRec相同的配置(层数、头数等)。“HSTU-large”表示更大的HSTU编码器(层数4倍,头数2倍)。结果表明:
    • a)HSTU因其针对推荐优化的设计,在相同配置下显著优于基线;
    • b)HSTU在扩展后性能进一步提升
  • 需要注意的是,此处使用的评估方法与工业级设置差异显著,因为完全打乱和多轮次训练在工业级流式设置中通常不实用(Liu等, 2022)
工业级流式设置
  • 接下来,论文在工业级数据集的流式设置中比较HSTU、消融版HSTU和 Transformer 的性能。本节剩余部分中,论文报告排名的归一化熵(Normalized Entropy,NE)(详细定义见附录)。模型训练超过1000亿样本(DLRM等效),每个任务使用64-256块H100 GPU(!!!)。由于排名是多任务设置,论文报告主要互动事件(“E-Task”)和主要消费事件(“C-Task”)。在论文的上下文中,NE降低0.001通常意味着数十亿用户的顶层指标提升0.5%。对于 Retrieval 任务,由于设置与语言建模类似,论文报告对数困惑度。在小规模设置中固定编码器参数(排名任务: \(l=3\) , \(n=2048\) , \(d=512\) ; Retrieval 任务: \(l=6\) , \(n=512\) , \(d=256\)),并因资源限制对其他超参数进行网格搜索
  • 结果如表5所示:
    • 第一 :HSTU显著优于 Transformer,尤其是在排名任务中,这可能是由于点注意力机制和改进的相对注意力偏差
    • 第二 :消融版HSTU与完整HSTU之间的差距验证了论文的设计有效性。Softmax版HSTU和 Transformer 的最佳学习率比其他模型低约10倍,这是为了训练稳定性。即使采用较低学习率和预归一化残差连接(Xiong等, 2020),标准 Transformer 在排名任务中仍频繁出现损失爆炸
    • 第三 :HSTU优于LLM中流行的 Transformer 变体 Transformer ++(Touvron等, 2023b), Transformer ++ 使用了RoPE、SwiGLU等技术。总体而言,在此小规模设置中,HSTU在质量上表现更优,同时训练速度提升1.5-2倍,HBM使用减少50%

编码器效率

  • 随机长度(Stochastic Length, SL)。图4和图5(a)展示了随机长度(SL)对模型指标的影响
    • 当 \(\alpha=1.6\) 时,长度为 4096 的序列在大多数情况下被压缩为 776,移除了超过 80% 的 token。即使 sparse 率增加到 64%-84%,主要任务的NE下降不超过 0.002(0.2%)
    • 以上这一证据表明,对于合适的 \(\alpha\) 值,SL 不会对模型质量产生负面影响,同时可以通过高 sparse 性降低训练成本。论文还在附录F.3中验证了 SL 显著优于现有的长度外推技术
  • 编码器效率。图5比较了HSTU和 Transformer 编码器在训练和推理设置中的效率
    • 对于 Transformer,论文使用 SOTA FlashAttention-2(Dao, 2023)实现
    • 考虑序列长度从 1,024 到 8,192,并在训练中应用随机长度(SL)
    • 评估时,HSTU和 Transformer 使用相同配置(\(d=512\) , \(h=8\) , \(d_{qk}=64\)),并消融相对注意力偏差(如第4.1.2节所示,HSTU在不使用rab \(^{p,t}\) 时仍优于 Transformer)。在NVIDIA H100 GPU上以bfloat16比较编码器级性能。总体而言,HSTU在训练和推理中的效率分别比 Transformer 高15.2倍和5.6倍
  • 此外,如第3.3节讨论的激活内存使用减少,使论文能够构建比 Transformer 深2倍以上的HSTU网络

工业级流式设置中生成式推荐器(GR)与 DLRM 的比较

  • 最后,论文比较了GR与 SOTA DLRM 基线在工业级流式设置中的端到端性能。论文的GR实现反映了生产中使用的典型配置,而 DLRM 设置则是数百人多年迭代的结果
    • 由于推荐系统的召回阶段使用了多个生成器(理解:这里应该是指多路召回),论文报告了将 GR 作为新增召回通道(“new source”)和替换现有主要 DLRM 通道(“replace”)的在线结果。表6和表7显示,GR不仅在离线测试中显著优于 DLRM,还在A/B测试中带来了12.4%的提升
    • 从表6中可知,作为一个召回通道加入能得到的更好的效果
  • 如第2节所述,GR基于原始分类互动特征构建,而 DLRM 通常使用更多手工特征(handcrafted features)。如果论文将GR使用的相同特征集提供给 DLRM (“DLRM(abl. features)”,消融特征),DLRM 的性能会显著下降,这表明GR通过其架构和统一特征空间可以有效地捕捉这些特征
  • 论文进一步通过与传统序列推荐器设置的比较验证了第2.2节中的GR公式。传统设置仅考虑用户交互过的物品(Kang & McAuley, 2018),结果显示其性能显著下降(表6和表7中的“GR(interactions only)”行)。论文还包含了一个仅使用内容特征的GR基线(“GR(content-based)”)。内容基线与 DLRM /GR之间的巨大性能差距凸显了高基数用户行为的重要性
  • 图6比较了GR与生产 DLRM 的效率。尽管GR模型的复杂度是 DLRM 的285倍(其中 \(285\times \text{FLOPs}\) 表示其复杂度是DLRM的285倍),但由于HSTU和第3.4节中的M-FALCON算法,论文在评分1024/16384个候选时实现了1.50倍/2.99倍的更高QPS
推荐系统的 scaling law
  • 众所周知,在大规模工业设置中,DLRM 在特定计算和参数范围内会达到质量饱和(Zhao等, 2023)。论文比较了GR和 DLRM 的可扩展性以更好地理解这一现象
  • 由于特征交互层对 DLRM 性能至关重要(Mudigere等, 2022),论文尝试了 Transformer (Vaswani等, 2017)、DHEN(Zhang等, 2022)以及论文生产设置中使用的带残差连接的DCN变体(He等, 2015)来扩展 DLRM 基线
    • 对于召回基线,论文通过调整隐藏层大小、嵌入维度和层数来扩展模型
    • 对于基于HSTU的生成式推荐器(GR),论文通过调整HSTU的超参数(如残差层数、序列长度、嵌入维度、注意力头数等)扩展模型,并调整 Retrieval 任务的负样本数量
  • 结果如图7所示:
    • 在低计算量区域,由于手工特征(handcrafted features)的存在,DLRM 可能优于GR,这验证了传统 DLRM 中特征工程的重要性
    • 然而,GR在FLOPs方面表现出显著更好的可扩展性(scalability),而 DLRM 性能趋于饱和,这与先前的研究结果一致
    • 论文还观察到,GR在嵌入参数和非嵌入参数方面均具有更好的可扩展性,其模型参数达到 1,500B ,而 DLRM 在约 200B 参数时性能饱和(这个参数量的训练,Meta真够有钱!)
  • 最后,论文的所有主要指标(包括 Retrieval 任务的 Hit Rate@100 和 Hit Rate@500,以及排名任务的NE)在适当的超参数下 ,随计算量的增加呈现幂律扩展(power law of compute used)
    • 这一现象在三个数量级范围内均成立,直至论文测试的最大模型(序列长度8,192,嵌入维度1,024,24层HSTU)
    • 此时,论文使用的总计算量(按365天标准化,因为论文采用标准流式训练设置)接近GPT-3(Brown等, 2020)和LLaMa-2(Touvron等, 2023b)的总训练计算量,如图1所示
    • 与语言建模(Kaplan等, 2020)相比,序列长度在GR中扮演更重要的角色,因此需要同步扩展序列长度和其他参数
    • 这或许是论文提出的方法最重要的优势,因为论文首次证明了 LLM 的 scaling law 可能也适用于大规模推荐系统

相关工作

传统序列推荐方法

  • 先前关于序列推荐的研究将用户交互行为简化为单一同质的物品序列(Hidasi等,2016;Kang & McAuley,2018)。工业级应用中,序列方法主要作为深度推荐模型(DLRMs)的一部分,包括成对注意力(Zhou等,2018)或序列编码器(Chen等,2019;Xia等,2023)。为了提高效率,已有研究探索了多阶段注意力替代自注意力机制(Chang等,2023)。在 Retrieval 任务中,将ID表示为 token 序列的生成方法也有所探索(Zhuo等,2020)。附录B.1中对相关工作进行了更详细的讨论

高效注意力机制

  • 由于自注意力机制的 \(O(n^2)\) 复杂度,高效注意力一直是研究热点,主要工作包括因子分解注意力(Child等,2019)、低秩近似(Katharopoulos等,2020)等。近期,针对序列转导任务的新架构也被提出(Gu等,2022;Hua等,2022)。HSTU的逐元素门控设计尤其受到FLASH(Hua等,2022)的启发。硬件感知的优化方法显著降低了内存使用(Rabe & Staats,2021;Korthikanti等,2022;Zhai等,2023b),并大幅缩短了实际运行时间(Dao等,2022)

长度外推技术

  • 长度外推技术使模型能够在较短的序列上训练后泛化到更长的序列,但多数研究聚焦于微调或改进偏置机制(Press等,2022)。论文则通过在长度维度引入随机性,受深度维度随机性研究(Huang等,2016)的启发

大语言模型(LLMs)的应用

  • 大语言模型的兴起促使研究者将推荐任务视为上下文学习(Silco等,2022)、指令微调(Bao等,2023)或基于预训练LLMs的迁移学习(Li等,2023)。LLMs中嵌入的世界知识可以迁移到下游任务(Cui等,2022),并在零样本或少样本场景中提升推荐效果。用户行为序列的文本表示也在中等规模数据集上展现出良好的扩展性(Shin等,2023)。然而,多数关于LLMs用于推荐的研究集中在低数据量场景;在大规模场景中,它们尚未在MovieLens等数据集上超越协同过滤方法(Hou等,2024)

整体总结

  • 新范式提出 :论文提出了一种新范式:生成式推荐器(Generative Recommenders, GRs),将排序和 Retrieval 任务重新定义为序列直推式任务(Sequential Transduction Tasks),使其能够以生成式方式进行训练
    • 注:最核心的创新点可能是将用户行为也视为一个像是图片、视频和文本类似的新模态
  • 性能优化 :得益于新颖的HSTU编码器设计,其在8192长度的序列上比当前 SOTA Transformer 快5.3至15.2倍,同时结合了M-FALCON等新型训练和推理算法
  • 实验 :通过GRs,论文部署了复杂度提升285倍的模型,同时减少了推理计算量。GRs和HSTU在生产环境中实现了12.4%的指标提升,并展现出比传统 DLRMs 更优的扩展性能
    • 论文的结果验证了用户行为是生成式建模中一个尚未充分探索的模态——正如标题所言,Actions speak louder than words(行动胜于雄辩)
  • 论文通过统一特征空间的简化设计,为推荐、搜索和广告领域的首个基础模型铺平了道路。GRs的完全序列化设置还支持端到端的生成式推荐框架。这两点使得推荐系统能够更全面地辅助用户
  • 未来规划 :生成式推荐器的完全序列化特性有望进一步推动端到端推荐系统的发展,例如通过直接生成推荐序列而非传统的列表排序。作者相信,这一方向将为推荐系统带来更广阔的应用前景

附录(个人): Transformer 模型计算量评估

  • 在原论文2.3节的公式 \(\sum_{i} n_{i}(n_{i}^{2}d + n_{i}d_{ff}d)\) 的来源解释如下

  • 注: \(d_{ff}\) 表示 前馈神经网络(Feed-Forward Network, FFN)的隐藏层维度(即中间层的神经元数量),在原始 Transformer 中通常设为 \(d_{ff} = 4d\)

  • 在 Transformer 或类似的自注意力架构(如HSTU)中,每个编码器层通常包含两个核心子层:

    • 自注意力机制(Self-Attention):复杂度为 \(O(n_i^2 d)\) (与序列长度平方和嵌入维度相关)
      • 自注意力的计算涉及三个主要步骤:
        • 计算查询(Q)、键(K)、值(V)矩阵:复杂度为 \(O(n_i d^2)\) ,其中 \(n_i\) 是序列长度, \(d\) 是嵌入维度
        • 计算注意力分数 \(QK^T\) :复杂度为 \(O(n_i^2 d)\)
        • 加权求和 \(AV\) :复杂度为 \(O(n_i^2 d)\)
      • 因此,自注意力的总复杂度为 \(O(n_i^2 d + n_i d^2)\) 。对于长序列(\(n_i \gg d\)),主导项为 \(O(n_i^2 d)\)
    • 前馈神经网络(FFN):复杂度为 \(O(n_i d_{ff} d)\) (与序列长度、嵌入维度和FFN隐藏层维度相关,)
    • FFN的典型结构是:
      $$
      \text{FFN}(x) = W_2 \cdot \text{ReLU}(W_1 x + b_1) + b_2
      $$
      • 其中:
        • \(W_1 \in \mathbb{R}^{d \times d_{ff} }\) :将输入从 \(d\) 维映射到 \(d_{ff}\) 维(通常 \(d_{ff} = O(d)\) ,注意不是相等,一般是通常 \(d_{ff} = 4d\) 等)
        • \(W_2 \in \mathbb{R}^{d_{ff} \times d}\) :将结果映射回 \(d\) 维
    • 在 Transformer 中, \(d_{ff}\) 通常远大于 \(d\) (例如 \(d_{ff}=4d\)),这使得FFN成为计算瓶颈之一。HSTU通过点乘门控(如 \(U(X)\))替代了传统FFN(见式3),因此可能不再需要显式的 \(d_{ff}\) 参数。但在对比 Transformer 基线时,公式中仍保留了 \(d_{ff}\) 以反映传统架构的成本
  • 总计算复杂度

    • 对于每个用户的序列,计算量为自注意力和前馈网络复杂度的总和:
      $$
      O(n_i^2 d + n_i d^2)
      $$
    • 若对所有用户求和,总计算量为:
      $$
      \sum_i n_i (n_i^2 d + n_i d^2)
      $$
    • 假设最大序列长度为 \(N = \max_i n_i\) ,则复杂度可简化为:
      $$
      O(N^3 d + N^2 d^2)
      $$
  • 论文中的其他优化:论文提出通过生成式训练(generative training)对计算量进行优化:

    • 通过调整用户采样率 \(s_u(n_i) \propto 1/n_i\) ,将总计算量降低为:
      $$
      \sum_i s_u(n_i) n_i (n_i^2 d + n_i d^2) \approx O(N^2 d + N d^2)
      $$
    • 这一优化将复杂度从三次方(\(N^3\))降至二次方(\(N^2\))

附录(个人):关于NE的完整定义

  • 参考原始论文:Practical Lessons from Predicting Clicks on Ads at Facebook
  • NE(Normalized Cross Entropy,归一化交叉熵),也可更准确地称为Normalized Logarithmic Loss(归一化对数损失),其定义为平均每个展示的对数损失除以模型预测每个展示的 CTR 时的平均对数损失
  • 假设给定的训练数据集有 \(N\) 个样本,则计算公式为:
    $$NE=\frac{-\frac{1}{N} \sum_{i=1}^{n}\left(\frac{1+y_{i} }{2} log \left(p_{i}\right)+\frac{1-y_{i} }{2} log \left(1-p_{i}\right)\right)}{-(p * log (p)+(1-p) * log (1-p))} $$
    • \(y_{i} \in{-1,+1}\) 为真实标签
    • \(p_{i}\) 为第 \(i\) 个样本的CTR预估值,其中 \(i = 1,2,\cdots,N\)
    • \(p\) 平均经验点击率(基于统计的真实点击率)
  • 该指标用于评估模型预测的好坏,其值越低,模型的预测效果越好
  • 进行这种归一化处理的原因是,真实 CTR 预估值越接近0或1,就越容易获得更好的对数损失,而除以 真实 CTR 的熵可使 NE 对真实 CTR 不敏感
  • 其他相关指标:Relative Information Gain (RIG)
    $$ RIG = 1 − NE $$

附录A:符号说明

  • 论文在下表中总结了论文使用的关键符号(对应原始论文的表8和表9)
    符号 描述
    \(\Psi_{k}(t_{j})\) Feature Logging 系统在 \(t_{j}\) 时刻发出的第 \(k\) 个 Training Example(\(k\) 是全局排序的)
    在典型的深度推荐模型(DLRM)推荐系统中,用户消费某些内容 \(\Phi_{i}\) (通过诸如跳过、视频观看完成和分享等动作 \(a_{i}\) 做出响应)后,特征记录系统将元组 \((\Phi_{i}, a_{i})\) 与用于对 \(\Phi_{i}\) 进行排名的特征相结合,并发出 \((\Phi_{i}, a_{i}\) ,以及 \(\Phi_{i}\) 的特征) 作为 Training Example \(\Psi_{k}(t_{j})\)
    如2.3节所述,DLRM和生成式推荐器(GRs)处理的 Training Example 数量不同,GRs中的 Example 数量通常少1 - 2个数量级
    \(n_{c}(n_{c,i})\) 与用户/样本 \(i\) 交互的内容数量
    \(\Phi_{0}, …, \Phi_{n_{c}-1}\) 在推荐系统的上下文中,与用户交互的内容列表
    \(a_{0}, …, a_{n_{c}-1}\) 与内容 \(\Phi_{i}\) 对应的用户动作列表
    当所有预测事件都是 Binary 时,每个动作可以被视为一个 multi-hot vector 事件(如点赞、分享、评论、图片浏览、视频初始化、视频观看完成、隐藏等)
    \(E, F\) 图2中DLRM中的分类特征
    \(E_{0}, E_{1}, …, E_{7}, E_{8}\) 和 \(F_{0}, F_{1}, …, F_{7}\) 表示在不同时间点通过特征提取(例如,最近喜欢的10张图片、用户过去点击的与当前候选内容最相似的50个网址等)从 \((\Phi_{0}, a_{0}, t_{0}), …, (\Phi_{n_{c}-1}, a_{n_{c}-1}, t_{n_{c}-1})\) 获得的转换结果
    “merge & sequentialize” 表示获取原始参与系列 \((\Phi_{0}, a_{0}, t_{0}), …, (\Phi_{n_{c}-1}, a_{n_{c}-1}, t_{n_{c}-1})\) 的(虚拟)反向过程
    \(G, H\) 图2中DLRM中与 user-item 参与无关的分类特征。这些特征(例如人口统计信息或关注的创作者)被合并到主时间序列(用户参与的内容列表,例如 \(\Phi_{0}, a_{0}, …, \Phi_{n_{c}-1}, a_{n_{c}-1}\))中,如2.1节所述并在图2中说明
    \(n(n_{i})\) Squential Transduction 任务中的 token 数量(对于用户或样本 \(i\))。虽然 \(O(n) = O(n_{c})\) ,但即使没有任何与非交互相关的分类特征, \(n\) 也可能与 \(n_{c}\) 不同;例如,见表1
    \(x_{0}, …, x_{n - 1}\) Squential Transduction 任务中的输入 token 列表
    \(y_{0}, …, y_{n - 1}\) Squential Transduction 任务中的输出 token 列表
    \(t_{0}, …, t_{n - 1}\) 与观察到 \(x_{0}, …, x_{n - 1}\) 的时间对应的时间戳列表
    \(\mathbb{X}, \mathbb{X}_{c}\) 所有输入/输出 token 的词汇表(\(\mathbb{X}\))及其内容子集(\(\mathbb{X}_{c}\))
    \(N, N_{c}\) \(\max_{i}n_{i}\) , \(\max_{i}n_{c,i}\)
    \(u_{t}\) 在时间 \(t\) 的用户表示
    \(s_{u}(n_{i}), \hat{s}_{u}(n_{i})\) 生成式训练(2.3节)中用于用户 \(i\) 的采样率
    \(d\) 模型维度(嵌入维度)
    \(d_{qk}\) HSTU和 Transformer 中注意力维度的大小。这适用于公式(1)中的 \(Q(X)\) 和 \(K(X)\)
    \(d_{v}\) HSTU中 Value 维度的大小。对于 Transformer,论文通常有 \(d_{qk} = d_{v}\)
    \(d_{ff}\) Transformer pointwise 前馈层中(feedforward)的隐藏维度大小。HSTU不使用前馈层;见下面的 \(U(X)\)
    \(h\) 注意力头的数量
    \(l\) HSTU中的层数。对于 Transformer,注意力层和逐点前馈层共同构成一层
    \(Q(X), K(X), V(X)\) 根据公式(1)为给定输入 \(X\) 在HSTU中获得的查询、键、值(Query/Key/Value)。其定义与标准 Transformer 中的 \(Q, K, V\) 类似。 \(Q(X), K(X) \in \mathbb{R}^{h \times N \times d_{qk} }\) ,并且HSTU使用 \(U(X)\) (与 \(f_{2}(·)\) 一起)在公式(3)中 “门控” 注意力池化值(\(V(X)\)),这使得HSTU完全避免了前馈层。 \(U(X) \in \mathbb{R}^{h \times N \times d_{v} }\)
    \(A(X)\) 为输入 \(X\) 获得的注意力张量。 \(A(X) \in \mathbb{R}^{h \times N \times N}\)
    \(Y(X)\) HSTU层针对输入 \(X\) 的输出。 \(Y(X) \in \mathbb{R}^{d}\)
    \(Split(·)\) 将张量分割成块的操作。 \(Split(\phi_{1}(f_{1}(X))) \in \mathbb{R}^{N \times (2hd_{qk} + 2hd_{v})}\) ;论文通过分割较大的张量(并置换维度)获得 \(U(X), V(X)\) (两者形状均为 \(h \times N \times d_{v}\))、 \(Q(X), K(X)\) (两者形状均为 \(h \times N \times d_{qk}\))
    \(rab^{p,t}\) 结合了位置(Raffel等,2020)和时间(基于观察到 token 的时间 \(t_{0}, …, t_{n - 1}\) ;一种可能的实现方式是对 \((t_{j} - t_{i})\) 应用某种分桶函数得到 \((i, j)\))信息的相对注意力偏差。在实践中,论文在一层内的不同注意力头之间共享 \(rab^{p,t}\) ,因此 \(rab^{p,t} \in \mathbb{R}^{1 \times N \times N}\)
    \(\alpha\) 控制HSTU中随机长度算法(3.2节)稀疏性的参数
    \(R\) GPU上的寄存器大小,在3.2节讨论的HSTU算法的上下文中
    \(m\) 推荐系统Ranking 阶段考虑的候选数量
    \(b_{m}\) M-FALCON 算法(3.4节)中的微批次大小

附录B:生成式推荐器:背景与公式

  • 许多读者可能对经典的深度学习推荐模型(DLRM)更为熟悉(Mudigere等,2022),因为自YouTube DNN时代起它就颇受欢迎(Covington等,2016),并且在每个大型在线内容和电子商务平台上都得到了广泛应用(Cheng等,2016;Zhou等,2018;Wang等,2021;Chang等,2023;Xia等,2023;Zhai等,2023a)。DLRM在异构特征空间上运行,使用各种神经网络,包括特征交互模块(Guo等,2017;Xiao等,2017;Wang等,2021)、顺序池化或目标感知成对注意力模块(Hidasi等,2016;Zhou等,2018;Chang等,2023)以及先进的多专家多任务模块(Ma等,2018;Tang等,2020)。因此,论文在第2节和第3节中通过将生成式推荐器(GRs)与经典DLRM进行明确对比,概述了GRs。在本节中,论文从经典的顺序推荐文献出发,为读者提供另一种视角

附录B.1:背景:学术界和工业界的顺序推荐

附录B.1.1:学术研究(传统顺序推荐设置)
  • 循环神经网络(RNNs)最早在GRU4Rec(Hidasi等,2016)中应用于推荐场景。Hidasi等(2016)考虑了门控循环单元(GRUs),并将其应用于两个数据集,即RecSys Challenge 2015和VIDEO(一个专有数据集)。在这两种情况下,只有 positive 事件(点击的电子商务商品或用户观看至少一定时间的视频)被保留作为输入序列的一部分。论文进一步观察到,在由检索和Ranking 阶段组成的经典工业规模两阶段推荐系统设置中(Covington等,2016),Hidasi等(2016)解决的任务主要对应于检索任务
  • 后来, Squential Transduction 架构的进步,特别是 Transformer (Vaswani等,2017),推动了推荐系统的类似进展。SASRec(Kang和McAuley,2018)首次在自回归设置中应用 Transformer。他们将评论或评分的存在视为 positive 反馈 ,从而将像亚马逊评论和MovieLens这样的经典数据集转换为 positive item 的序列,类似于GRU4Rec。采用二元交叉熵损失,其中正目标定义为下一个 “positive” item (回想一下,这本质上只是评论或评分的存在),负目标从 item 语料库 \(\mathbb{X}=\mathbb{X}_{c}\) 中随机采样
  • 随后的大多数研究都基于与上述GRU4Rec(Hidasi等,2016)和SASRec(Kang和McAuley,2018)类似的设置,例如BERT4Rec(Sun等,2019)应用来自BERT(Devlin等,2019)的双向编码器设置,S3Rec(Zhou等,2020)引入明确的预训练阶段,等等
附录B.1.2:作为深度学习推荐模型(DLRM)一部分的工业应用
  • 顺序方法,包括顺序编码器和成对注意力模块,由于能够作为DLRM的一部分增强用户表示,已在工业环境中得到广泛应用。DLRM通常使用相对较短的序列长度 ,例如BST(Chen等,2019)中为20,DIN(Zhou等,2018)中为1000,TransAct(Xia等,2023)中为100。论文观察到,论文中的8192比传统DLRM序列长度大 1-3 个数量级
  • 尽管使用短序列长度,大多数DLRM仍能成功捕捉长期用户偏好。这可归因于两个关键方面:
    • 第一 :现代DLRM中通常使用预计算的用户配置文件/嵌入或外部向量存储(Chang等,2023),这两者都有效地扩展了回顾窗口
      • 理解,用户 Embedding 其实与用于历史长期用户偏好有关?
    • 第二 :通常会采用大量的上下文、用户和 item 侧特征,并且使用各种异构网络,如因子分解机(FMs)、深度交叉网络(DCNs)、混合专家(MoEs)等,来转换表示并组合输出
      • 理解:特征交叉能泛化到用户长期偏好?
  • 与附录B.1.1中讨论的顺序设置相比,所有主要的工业界工作都在(用户/请求,候选 item)对上定义损失。在排名设置中,通常使用多任务二元交叉熵损失。在检索设置中,双塔设置(Covington等,2016)仍然是主导方法。最近的工作研究了将下一个推荐 item 表示为(子) token 序列上的概率分布,如OTM(Zhuo等,2020)和DR(Gao等,2021)(注意,在其他近期工作中,相同的设置有时被称为 “生成式召回”)。它们通常利用 beam search 从子 token 中解码 item。随着现代加速器(如GPU、定制ASIC和TPU)的普及,还提出并部署了先进的学习相似性函数,如混合逻辑(Zhai等,2023a),作为双塔设置和 beam search 的替代方案
  • 从问题公式化的角度来看,考虑到模型架构、使用的特征和损失与附录B.1.1中讨论的学术顺序推荐研究有显著差异,作者认为上述所有工作都属于DLRM(Mudigere等,2022)的一部分。值得注意的是,在这项工作之前,工业界尚未成功应用完全顺序排名设置,特别是在每日活跃用户(DAU)达到数十亿规模的情况下

附录B.2:公式化表述:生成式推荐器(GRs)中作为序列转换任务的排序与检索

  • 接下来,论文讨论传统序列推荐器设置和深度推荐模型(DLRM)设置中的三个局限性,以及生成式推荐器(GRs)如何从问题公式化的角度解决这些问题
  • 问题1:忽略用户交互物品之外的特征 :以往的序列公式化表述仅考虑用户明确交互过的内容(物品) ,而在GRs出现之前,传统推荐系统会基于大量特征进行训练 ,以增强用户和内容的表征 ,GRs通过以下方式解决这一局限性:
    • a)压缩其他分类特征并将其与主时间序列合并;
    • b)如2.1节和图2所述,利用目标感知公式通过交叉注意力交互来捕捉数值特征。论文通过实验验证了这一点,结果表明,忽略这些特征的传统 “interactions only” 公式化表述会显著降低模型质量;实验结果可在表7和表6中标记为 “GR(interactions only)” 的行中找到,论文发现仅利用交互历史会导致检索的 HitRate@100 下降1.3%,排序的归一化熵(NE)下降2.6%(回想一下,如4.1.2节和4.3.1节所述,NE变化0.1% 即为显著变化)
  • 问题2:用户表征在与目标无关的设置中计算 :大多数传统序列推荐器,包括GRU4Rec(2016)、SASRec(2018)、BERT4Rec(2019)、S3Rec(2020)等,都是以与目标无关的方式构建的。在这种方式中,对于目标物品 \(\Phi_{i}\) , \(\Phi_{0}, \Phi_{1}, …, \Phi_{i - 1}\) 被用作编码器输入来计算用户表征,然后用于进行预测。相比之下,工业环境中使用的大多数主要DLRM方法在构建所使用的序列模块时考虑了目标感知,能够将 “目标”(排序候选)信息整合到用户表征中。这些方法包括DIN(2018)(阿里巴巴)、BST(2019)(阿里巴巴)、TWIN(2023)(快手)和TransAct(2023)(Pinterest)
    • 生成式推荐器(GRs)通过交错内容和动作序列(2.2节),在因果自回归设置中实现目标感知注意力机制 ,结合了两者的优点(序列推荐器和DLRM的优点)。论文在表10中对先前的工作和本工作进行了分类和对比
    • 理解:(详情见前文对论文 target-aware 实现方式的理解),其实 Transformer 的 Attention 也有交叉功能,只要输入端有目标 item 和用户历史交互 item 即可,但论文所说的传统的自回归模型是指输入侧不包含目标 item 的情况,论文中,将动作也建模进去,则在输出对目标 item 的动作 token 时,自然就需要将目标 item 作为输入,从而也就实现了 target-aware 交叉注意力:
      • 传统自回归模型预估目标为 : \(p(\Phi_{i+1}|\Phi_{0},\Phi_{1},\ldots,\Phi_{i})\),目标物品 \(\Phi_{i+1}\) 与其他历史序列无交叉
        • 注:此时动作 \(a_i\) 信息可能会作为额外表征加入到 \(\Phi_{i}\)中,所以 表述为 \(p(\Phi_{i}|(\Phi_{0}, a_{0}), …, (\Phi_{i - 1}, a_{i - 1}))\) 也可以
      • 论文预估目标为 : \(p(a_{i+1}|\Phi_{0},a_{0},\Phi_{1},a_{1},\ldots,\Phi_{i+1})\),目标物品 \(\Phi_{i+1}\) 与其他历史序列有交叉
  • 问题3:判别式公式限制了先前序列推荐器工作的适用性 :传统的序列推荐器本质上是判别式的。现有的序列推荐文献,包括GRU4Rec和SASRec等开创性工作,对 \(p(\Phi_{i}|\Phi_{0}, a_{0}, …, \Phi_{i - 1}, a_{i - 1})\) 进行建模(注:虽然这里传统这里输入中包含动作信息,但是实际上动作是作为额外的信息加入到 item 表征中的,不是单独将 action 作为一个 token,所以写成 \(p(\Phi_{i}|(\Phi_{0}, a_{0}), …, (\Phi_{i - 1}, a_{i - 1}))\) 更合适),即根据用户当前状态推荐下一个物品的条件分布
    • 推荐系统中实际上存在两个概率过程 :
      • 1)推荐系统向用户推荐内容 \(\Phi_{i}\) 的过程;
      • 2)用户通过某些动作 \(a_{i}\) 对推荐内容 \(\Phi_{i}\) 做出反应的过程
    • 生成式方法对推荐内容和用户动作序列的联合分布进行建模 ,如2.2节所述,即联合概率分布 \(p(\Phi_{0}, a_{0}, \Phi_{1}, a_{1}, …, \Phi_{n_{c}-1}, a_{n_{c}-1})\)(如表11(图8)所示,论文提出的生成式推荐器能够对这种分布进行建模):
      • Next action token (\(a_{i}\)) prediction 任务正是 Ranking 任务(即表1中讨论的 GR Ranking 设定)
      • Next content token (\(\Phi_{i}\)) prediction 任务对应 Retrieval 任务,目标是学习 next item
  • 重要的是,这种公式化表述不仅能够对数据分布进行适当建模,还能够通过例如 beam search 直接采样要推荐给用户的物品序列。论文假设这将产生一种比传统列表式设置(例如DPP(2014)和RL(2018))更优越的方法,论文将此类系统的完整公式化表述和评估(在6节中简要讨论)留作未来的工作

附录C:评估:合成数据

  • 如3.1节先前所述,标准的softmax注意力机制因其归一化因子,难以捕捉用户偏好的强度,而这对于用户表示学习至关重要。在推荐场景中,这一点很关键,因为系统可能不仅需要预测 item 的相对排序,还需要预测用户参与的强度(例如,未来对特定主题的 positive 行为数量)
  • 为理解这种行为,论文构建了遵循狄利克雷过程(Dirichlet Process)的合成数据,该过程在动态词汇集上生成流式数据。狄利克雷过程捕捉了用户参与历史中 “富者更富”(rich gets richer) 的行为。论文设置合成实验如下:
    • 论文将 20,000 个 item ID中的每一个随机分配到100个类别中的某一个
    • 论文生成 1,000,000 条长度为 128 的记录,其中前90%用于训练,最后10%用于测试。为模拟流式训练设置,论文最初提供40%的 item ID,其余的以相等的间隔逐步提供;即在记录 500,000 时,可以采样的最大ID是 \((40% + 60% * 0.5) * 20,000 = 14,000\)
    • 论文为每条记录从100个类别中随机选择最多5个类别,并为这5个类别随机采样一个先验 \(H_{c}\) 。按照狄利克雷过程,为每个位置顺序采样类别,具体如下:
      • 对于 \(n>1\) :
        • 以概率 \(\alpha / (\alpha + n - 1)\) ,从 \(H_{c}\) 中抽取类别 \(c\)
        • 以概率 \(n_{c} / (\alpha + n - 1)\) ,抽取类别 \(c\) ,其中 \(n_{c}\) 是先前具有类别 \(c\) 的 item 数量
        • 随机采样一个符合类别 \(c\) 且受流式约束的分配 item
      • 其中 \(\alpha\) 从(1.0, 500.0)中均匀随机采样
  • 结果见表2。由于此数据集没有时间戳,论文在HSTU中去除 \(rab^{p, t}\)。论文观察到,相对于标准 Transformer,HSTU的 HitRate@10 提高了100%以上。重要的是,将HSTU的逐点聚合注意力机制(Pointwise aggregated attention)替换为 Softmax(“HSTU w/ Softmax”)也会导致 HitRate 显著降低,这验证了类似逐点聚合注意力机制的重要性

附录D:评估:传统顺序推荐器设置

  • 论文在4.1.1节的评估重点是将HSTU与 SOTA Transformer 基线SASRec进行比较,使用最新的训练方法。在本节中,论文进一步考虑另外两种替代方法
  • 循环神经网络(RNNs)。论文考虑顺序推荐器的经典工作GRU4Rec(Hidasi等,2016),以帮助读者理解包括 Transformer 和HSTU在内的自注意力模型,在充分融入最新的建模和训练改进后,与传统RNNs相比如何
  • 自监督顺序方法。论文考虑最受欢迎的工作BERT4Rec(Sun等,2019),以了解双向自监督(BERT4Rec通过完形填空目标利用)与单向因果自回归设置(如SASRec和HSTU)相比如何
  • 结果见表12。论文重用Klenitskiy和Vasilev(2023)报告的BERT4Rec和GRU4Rec在ML-1M和ML-20M上的结果。由于使用了采样softmax损失,论文保持负样本数量不变(ML-1M、ML-20M为128,亚马逊图书为512),以确保方法之间的公平比较
  • 结果证实,在使用采样softmax损失的传统顺序推荐设置中,SASRec仍然是最具竞争力的方法之一(Zhai等,2023a;Klenitskiy和Vasilev,2023),而HSTU显著优于所评估的 Transformer 、RNN和自监督双向 Transformer

附录E:评估:传统DLRM基线

  • 第4节中使用的DLRM基线配置反映了数百名研究人员和工程师多年来的持续迭代,并且在部署HSTU/GR之前,是对拥有数十亿日活跃用户的大型互联网平台上生产配置的近似。本节对所使用的模型进行简要描述

排名设置

  • 如(Mudigere等,2022)所述,基线排名模型采用了大约一千个密集特征和五十个稀疏特征。论文结合了各种建模技术,如 Mixture of Experts(Ma等,2018)、Deep & Cross Network 的变体(Wang等,2021)、各种顺序推荐模块,包括 target-aware pairwise 注意力(在工业设置中常用的一种变体可参见(Zhou等,2018)),以及特殊交互层上的残差连接(He等,2015;Zhang等,2022)。在 scaling law 部分(4.3.1节)的低FLOP区域,一些计算成本高的模块被简化、替换为其他 state-of-the-art 的变体,如DCN,以达到所需的FLOP
  • 由于保密考虑,论文无法透露确切设置 ,但据论文所知,在充分纳入最新研究成果后,论文的基线代表了最优秀的DLRM方法之一。为验证这一说法并帮助读者理解,论文在表7中报告了基于相同特征,但仅利用主要已发表成果(包括DIN(Zhou等,2018)、DCN(Wang等,2021)和MMoE(Ma等,2018))的典型设置(“DLRM (DIN + DCN)”),并在图9中展示了组合架构。该设置在主要E任务的NE上比论文的生产DLRM设置低0.71%,在主要C任务的NE上低0.57%(0.1%的NE变化即具有显著性)

检索设置

  • 基线检索模型采用标准的双塔神经检索设置(Covington等,2016),并结合了批内(in-batch)和批外(out-of-batch)采样。输入特征集包括高基数稀疏特征(如 item ID、用户ID)和低基数稀疏特征(如语言、主题、兴趣实体)。使用带有残差连接的前馈层堆栈(He等,2015)将输入特征压缩为用户和 item 嵌入

特征和序列长度

  • DLRM基线中使用的特征,包括各种顺序编码器/成对注意力模块所利用的主要用户交互历史,是所有 GR 候选模型使用特征的严格超集(strict supersets)。这适用于论文中进行的所有研究,包括缩放研究(4.3.1节)中使用的特征

附录F:随机长度

附录F.1:子序列选择

  • 在公式(4)中,论文从完整的用户历史中选择长度为 \(L\) 的子序列以增加稀疏性。论文的实证结果表明,精心设计子序列选择技术可以提高模型质量。论文计算一个指标 \(f_{i}=t_{n}-t_{i}\) ,它对应于用户与 item \(x_{i}\) 交互后经过的时间量。论文使用以下子序列选择方法进行离线实验:
    • 贪心选择(Greedy Selection) - 从 \(S\) 中选择 \(L\) 个 \(f_{i}\) 值最小的 item,即保留最近交互的 item
    • 随机选择(Random Selection) - 从 \(S\) 中随机选择 \(L\) 个 item
    • 特征加权选择(Feature-Weighted Selection) - 根据加权分布 \(1 - f_{n, i} / (\sum_{j = 1}^{L}f_{j, i})\) 从 \(S\) 中选择 \(L\) 个 item(理解:对越近交互的样本,采样权重越大)
  • 在离线实验中,特征加权子序列选择方法产生了最佳的模型质量,如表13所示

附录F.2:随机长度对序列稀疏性的影响

  • 在表3中,论文展示了随机长度对具有30天用户参与历史的代表性工业规模配置中序列稀疏性的影响。序列稀疏性定义为1减去所有样本的平均序列长度与最大序列长度的比值。为了更好地描述稀疏注意力的计算成本,论文还定义了 \(s2\) ,它被定义为:1减去注意力矩阵的稀疏性(which is defined as one minus the sparsity of the attention matrix)。作为参考,论文在表14和表15中分别给出了60天和90天用户参与历史的结果

附录F.3:与序列长度外推技术的比较

  • 论文进行了额外的研究,以验证随机长度与语言建模中使用的现有序列长度外推技术相比具有竞争力。许多现有方法通过修改旋转位置嵌入(RoPE)(Su等,2023)来进行序列长度外推。为了与现有方法进行比较,论文训练了一个没有相对注意力偏差和旋转嵌入的HSTU变体(HSTU - RoPE)
  • 论文在HSTU-RoPE上评估以下序列长度外推方法:
    • 零样本(Zero-Shot) - 应用NTK感知的RoPE(Peng等,2024),然后直接评估模型,不进行 Fine-tune
    • 微调(Fine-tune) - 应用逐部分NTK(Peng等,2024)后,对模型进行1000步 Fine-tune
  • 论文在HSTU(包括相对注意力偏差,无旋转嵌入)上评估以下序列长度外推方法:
    • 零样本(Zero-Shot) - 根据最大训练序列长度夹紧相对位置偏差,直接评估模型(Raffel等,2020;Press等,2022)
    • 微调(Fine-tune) - 根据最大训练序列长度夹紧相对位置偏差,在评估模型之前对模型进行1000步 Fine-tune
  • 在表16中,论文报告了训练期间引入数据稀疏性的模型(随机长度、Zero-Shot、Fine-tune)与在完整数据上训练的模型之间的NE差异。论文将 Zero-Shot 和 Fine-tune 技术的稀疏性定义为训练期间的平均序列长度与评估期间的最大序列长度之比。所有 Zero-Shot 和 Fine-tune 模型都在1024序列长度的数据上进行训练,并在2048和4096序列长度的数据上进行评估。为了为这些技术找到合适的随机长度基线,论文选择了导致相同数据稀疏性指标的随机长度设置
  • 作者认为,Zero-Shot 和 Fine-tune 的序列长度外推方法不太适合处理高基数ID的推荐场景。从经验上看,论文观察到随机长度明显优于 Fine-tune 和 Zero-Shot 方法。作者认为这可能是由于论文的词表较大 ,Zero-Shot 和 Fine-tune 方法无法为较旧的ID学习良好的表示,这可能会损害它们充分利用较长序列中包含的信息的能力

附录G:稀疏分组通用矩阵乘法(GEMMs)和融合相对注意力偏差

  • 论文提供3.2节中介绍的高效HSTU注意力内核的更多信息。论文的方法基于内存高效注意力(Rabe和Staats,2021)和FlashAttention(Dao等,2022),是一种内存高效的自注意力机制,它将输入划分为块,并避免在反向传播中具体化大的 \(h×N×N\) 中间注意力张量。通过利用输入序列的稀疏性,我们可以将注意力计算重新表述为一组具有不同形状的连续GEMM运算。论文实现了高效的GPU内核来加速此计算。相对注意力偏差的构建也因内存访问而成为瓶颈。为解决此问题,论文将相对偏差构建和分组GEMM运算融合到单个GPU内核中,并在反向传播中使用GPU的快速共享内存来累积梯度。尽管论文的算法在反向传播中需要重新计算注意力和相对偏差,但它比 Transformer 中使用的标准方法明显更快且内存使用更少

附录H:Microbatched-Fast Attention Leveraging Cacheable OperatioNs (M-FALCON)

  • 在本节中,论文详细描述3.4节中讨论的 M-FALCON 算法
  • M-FALCON 引入了三个关键思想
  • 批量推理可应用于因果自回归设置 :GR中的排名任务以目标感知的方式制定,如2.2节所述。通常认为,在目标感知设置中,论文需要一次对一个 item 进行推理,对于 \(m\) 个候选 item 和长度为 \(n\) 的序列,成本为 \(O(mn^{2}d)\) 。但这里论文表明这不是最优解决方案;即使使用普通 Transformer,也可以修改自注意力中使用的注意力掩码以进行批量操作(“批量推理”),并将成本降低到 \(O((n + m)^{2}d)=O(n^{2}d)\)
    • 图11给出了说明。这里,图11(a)和(b)都涉及因果自回归设置的注意力掩码矩阵。关键区别在于,图11(a)在因果训练中使用大小为 \(2n_{c}\) 的标准下三角矩阵,而图11(b)通过将 \(i, j≥2n_{c}\) 且 \(i≠j\) 的条目设置为 False 或 \(-\infty\) 来修改大小为 \(2n_{c}+b_{m}\) 的下三角矩阵,以防止目标位置 \(\Phi_{0}’, …, \Phi_{b_{m}-1}’\) 相互关注。很容易看出,通过这样做,自注意力块对 \(\Phi_{i}’, a_{i}’\) 的输出仅取决于 \(\Phi_{0}, a_{0}, …, \Phi_{n_{c}-1}, a_{n_{c}-1}\) ,而不取决于 \(\Phi_{j}’\) (\(i≠j\))。换句话说,通过使用修改后的注意力掩码对 \((2n_{c}+b_{m})\) 个 token 进行前向传递,论文现在可以获得与对 \((2n_{c}+1)\) 个 token 进行 \(b_{m}\) 次单独前向传递相同的最后 \(b_{m}\) 个 token 的结果,在第 \(i\) 次前向传递中, \(\Phi_{i}’\) 位于第 \(2n_{c}\) (基于0)位置,使用标准因果注意力掩码
  • 微批次将批量推理扩展到大型候选集 :Ranking 阶段可能需要处理大量的排名候选 item ,多达数万个(Wang等,2020)。我们可以将总共 \(m\) 个候选 item 划分为 \(\lceil m / b_{m}\rceil\) 个大小为 \(b_{m}\) 的微批次,使得 \(O(b_{m}) = O(n)\) ,这在大多数实际推荐设置中,对于多达数万个候选 item ,保持了前面讨论的 \(O((n + m)^{2}d)=O(n^{2}d)\) 的运行时间
  • 编码器级缓存可在请求内和请求间实现计算共享 :最后,键值缓存(Pope等,2022)可在请求内和请求间应用。例如,对于论文中介绍的HSTU模型(3节), \(K(X)\) 和 \(V(X)\) 在微批次内和/或请求间完全可缓存。对于缓存的前向传递,论文只需要为最后 \(b_{m}\) 个 token 计算 \(U(X), Q(X), K(X)\) 和 \(V(X)\) ,同时为包含 \(n\) 个 token 的序列化用户历史重用缓存的 \(K(X)\) 和 \(V(X)\) 。同样, \(f_{2}(Norm(A(X)V(X))\odot U(X))\) 只需要为 \(b_{m}\) 个候选 item 重新计算。这将缓存前向传递的计算复杂度降低到 \(O(b_{m}d^{2}+b_{m}nd)\) ,即使 \(b_{m}=n\) ,也比 \(O((n + b_{m})d^{2}+(n + b_{m})^{2}d)\) 提高了2 - 4倍
  • 算法1说明了 M-FALCON 算法,有助于理解。论文注意到, M-FALCON 不仅适用于HSTU和GR,还广泛适用于其他基于自注意力架构的目标感知因果自回归设置的推理优化算法

附录H.1:推理吞吐量评估:使用 M-FALCON 的生成式推荐器(GRs)与DLRM的比较

  • 如3.4节所述, M-FALCON 在推理时并行处理 \(b_{m}\) 个候选 item ,以在所有 \(m\) 个候选 item 之间分摊计算成本。为理解论文的设计,论文在相同硬件设置下比较了GR和DLRM的吞吐量(即每秒评分的候选 item 数,QPS)
  • 如图12和图13所示,由于批量推理实现了成本分摊,GR的吞吐量在一定区域内(在论文的案例研究中 \(m = 2048\))随Ranking 阶段候选 item 数 \((m)\) 呈次线性增长。这证实了批量推理在因果自回归设置中的关键性。由于注意力复杂度按 \(O((n + b_{m})^{2})\) 缩放,利用多个微批次本身就提高了吞吐量。缓存进一步消除了微批次之上的冗余线性和注意力计算。两者结合,相对于使用单个微批次的 \(b_{m}=m = 1024\) 基线,实现了高达1.99倍的额外加速,如图13所示。总体而言,凭借高效的HSTU编码器设计和利用 M-FALCON ,基于HSTU的生成式推荐器在大规模生产设置中的吞吐量比DLRM高出2.99倍,尽管GR在FLOP方面复杂285倍

NLP——LLM-Rubric-RL相关总结

注:本文包含 AI 辅助创作


整体说明

  • 在 2025 年的 LLM RL 研究中,利用 Rubrics(评分细则/量规) 来构建更精细、可解释的奖励函数已经成为主流趋势之一
  • 这种方法主要解决传统标量奖励(Scalar Reward)无法提供细粒度指导的问题
  • 本文总结相关领域具有代表性的论文,论文详细阅读见其他内容

RaR

  • 原始论文:(RaR)Rubrics as Rewards: Reinforcement Learning Beyond Verifiable Domains, Scale AI, 20251003
  • 该论文提出了“可验证奖励强化学习”(RLVR)的概念,证明了在复杂推理任务中,使用明确的 Rubrics 作为奖励信号比单纯的人类偏好更有效,能够显著提升模型推理的正确性

DR Tulu: Reinforcement Learning with Evolving Rubrics for Deep Research

  • 原始博客:DR Tulu: An open, end-to-end training recipe for long-form deep research, 20251118, AI2
  • 原始论文:DR Tulu: Reinforcement Learning with Evolving Rubrics for Deep Research, 20251124 & 20251126, AI2
  • 针对深度研究(Deep Research)场景,提出了一种“演化细则”(Evolving Rubrics)机制
  • 随着模型能力的提升,评分标准也会动态调整,从而引导模型适应性地使用工具并完成更复杂的科学问答任务

RubricRL: Simple Generalizable Rewards for Text-to-Text Generation

  • 原始论文:RubricRL: Simple Generalizable Rewards for Text-to-Image Generation, 20251125, Microsoft CoreAI
  • 文生图 Rubric,提出了一个名为 RubricRL 的通用框架,旨在为文本生成任务设计简单且可泛化的基于 Rubric 的奖励
  • 该方法强调了奖励设计的可解释性和可组合性,使用户能更灵活地定制模型行为

AdvancedIF: Rubric-Based Benchmarking and Reinforcement Learning for Advancing LLM Instruction

  • 原始论文:AdvancedIF: Rubric-Based Benchmarking and Reinforcement Learning for Advancing LLM Instruction Following, 20251113 & 20251126, Meta Superintelligence Labs & CMU
  • 该工作发布了 AdvancedIF 框架,利用基于 Rubric 的流水线来提升大模型的Instruction Following能力
  • 不仅将 Rubric 用于评估(Benchmarking),还将其直接用于 RL 训练环节

(Rubicon) Reinforcement Learning with Rubric Anchors

  • 原始论文:(Rubicon) Reinforcement Learning with Rubric Anchors, 20250818, Inclusion AI & Ant Group & Zhejiang University
  • 探讨了在 RLVR(可验证奖励 RL)范式下,如何利用“Rubric Anchors”(评分锚点)来增强大模型
  • 通过锚点机制,模型能够更稳定地对齐到预期的细粒度标准上

(RuscaRL) Breaking the Exploration Bottleneck: Rubric-Scaffolded Reinforcement Learning for General LLM Reasoning

  • 原始论文:(RuscaRL) Breaking the Exploration Bottleneck: Rubric-Scaffolded Reinforcement Learning for General LLM Reasoning, 20250823-20251022, ZJU
  • 提出了 RuscaRL 框架,将 Rubric 作为一种教学脚手架(Instructional Scaffolding)
  • 该方法旨在帮助模型突破复杂任务中的“探索瓶颈”,通过结构化的细则引导模型逐步探索出正确的策略

Self-Rewarding Rubric-Based Reinforcement Learning for Open-Ended Reasoning

  • 原始论文:Self-Rewarding Rubric-Based Reinforcement Learning for Open-Ended Reasoning, 20250919, Ant Group
  • 针对开放式推理任务,提出了一种自我奖励机制
  • 模型能够根据预设的 Rubric 对自己的输出进行评分和反馈,从而在缺乏外部大规模标注的情况下实现自我迭代和提升

RLAC: Reinforcement Learning with Adversarial Critic for Dynamic Rubric Generation

  • 原始论文:RLAC: Reinforcement Learning with Adversarial Critic for Free-Form Generation Tasks, 20251103, SJTU & UC Berkeley
  • 提出了一种结合对抗性 Critic 的强化学习方法(RLAC),通过动态生成的 Rubric 来应对训练过程中的挑战,属于Post-training阶段的优化策略

PaTaRM: Bridging Pairwise and Pointwise Signals via Preference-Aware Task-Adaptive Reward Modeling

  • 原始论文:PaTaRM: Bridging Pairwise and Pointwise Signals via Preference-Aware Task-Adaptive Reward Modeling, BUPT & Meituan, 20251028

Chasing the Tail: Effective Rubric-based Reward Modeling for Large Language Model Post-Training

  • 原始论文:Chasing the Tail: Effective Rubric-based Reward Modeling for Large Language Model Post-Training, Scale AI, 20250925

Auto-Rubric: Learning to Extract Generalizable Criteria for Reward Modeling

  • 原始论文:Auto-Rubric: Learning to Extract Generalizable Criteria for Reward Modeling, 20251020
  • 构造静态 Rubrics

(Self-Rewarding Rubrics) Self-Rewarding Rubric-Based Reinforcement Learning for Open-Ended Reasoning

  • 原始论文:Self-Rewarding Rubric-Based Reinforcement Learning for Open-Ended Reasoning, 20250919
  • 将策略自己用作 Rubrics 生成器

QA-LIGN: Aligning LLMs through Constitutionally Decomposed QA

  • 原始论文:QA-LIGN: Aligning LLMs through Constitutionally Decomposed QA, 20250609-20251204, Arizona State University

AutoRubric-R1V: Rubric-Based Generative Rewards for Faithful Multimodal Reasoning

  • 原始论文:AutoRubric-R1V: Rubric-Based Generative Rewards for Faithful Multimodal Reasoning, 20251016, University of Notre Dame
  • 多模态 Rubric

(DeepSeek-GRM)Inference-Time Scaling for Generalist Reward Modeling

  • 原始论文:(DeepSeek-GRM)Inference-Time Scaling for Generalist Reward Modeling, DeepSeek & THU, 20250403-20250925
  • 推出了 DeepSeek-GRM 模型,是 Pointwise GRM,模型地址:huggingface.co/collections/BBQGOD/deepseek-grm

Rubric-ARM

  • 原始论文:(Rubric-ARM)Alternating Reinforcement Learning for Rubric-Based Reward Modeling in Non-Verifiable LLM Post-Training, 20260202, Emory University & Purdue University
  • 交替训练

Rationale Consistency

  • 原始论文:(Rationale Consistency)Outcome Accuracy is Not Enough: Aligning the Reasoning Process of Reward Models, 20260204, Qwen Team & Fudan & THU
  • 作者提出,目前的 RLHF 中,训练和评估时都仅考虑结果准确性,这本身是不够的

MaMs(Multi-agent Markov-state)

  • 原始论文:(MaMs)Learning Query-Specific Rubrics from Human Preferences for DeepResearch Report Generation, 20260203, Tencent & Fudan
  • for DeepResearch

NLP——LLM-RL训练指标记录

注:本文包含 AI 辅助创作


整体说明

  • 本文汇总 LLM-RL 训练过程中需要关注的各种指标
  • RL 训练指标重要性:
    • 宗旨:能在训练指标上体现的问题,都不该留到下游测试集上再体现
    • Insight:全面的指标监控有助于快速发现各种可能存在的隐藏 Bug 或 找到改进点

训练日志落盘

  • 训练时在每一步都落表下面的内容
    • Prompt
    • Response/Rollout(每个 Prompt 可能多个)
    • 真实 Rewards 打分情况(每个 Response/Rollout 至少一个,若多个 Reward Model/Rule 则有多个)
    • final_rewards 打分情况(包含长度惩罚等后的最终奖励)

KL 散度(当前策略 vs Reference 策略)

  • KL 散度是最常见的观测指标,有助于判断当前策略偏离 Reference 策略多少了
  • KL 散度的监控上报值往往是直接估计值
    • 注:KL 散度完整值需要对全词表积分,除了用于偶尔对齐指标或训推不一致问题排查外,一般不需要使用
    • KL 散度的近似估计方式有 k1,k2,k3 三种,建议监控 k1 和 k3,其中 k1 的值会更稳定些,k3 可能存在较多 spike
  • 经验:
    • 好的 KL 散度曲线应该是缓慢的逐步增长(若设定 Target KL 来动态调整 KL 惩罚系数,会逐步收敛到 Target KL 区间内)
    • 模型的崩溃往往伴随着 KL 散度的突然崩溃
      • 遇到 KL 突然崩溃时,一般其他指标也会随着崩溃,包括熵、Rewards 曲线等
      • 若 KL 突然崩溃,即使训练时的 Rewards 等暂时没有崩溃,大概率也出现问题了

Response 长度

  • Response 长度代表着模型对长度偏好的变化,是非常重要的关注指标
  • Response 相关的指标还有 8K 截断率、16K 截断率、32K 截断率等
  • 经验:
    • 回复长度可能会持平、缓慢上涨、缓慢降低,看具体策略而言可能不同

Entropy

  • 当前策略的熵代表当前模型的探索能力(熵越大,不确定性越高,探索能力越强)
  • 经验:
    • 好的熵没有特别固定的形状,但有一些参考:
      • 可能逐步降低然后收敛(收敛过快往往不是好事,说明模型过快丢失探索能力)
      • 可能先降低,再提升,再收敛(观察到有时候提升时往往伴随着长度的增加)

真实奖励 rewards/score(不包含 kl_rewards 等)

  • ORM 场景中,真实奖励仅出现在最后一个 Token,这个奖励可以看做是 Sequence 粒度的真实奖励(必须监控)
    • 注:这个真实奖励有时也叫做 scores 等
  • 若使用了 PRM,也建议使用详细的上报保证可监控(建议分 Step 进行上报和监控)
  • 针对不同领域同时训练的场景,应该分领域上报该值
  • 对于 Rewards 是 0-1 的场景,应该同步上报每个 Prompt 的正确率等,否则需上报 Rewards 的分布情况
  • 经验:
    • 好的 rewards/score 曲线应该是逐步上涨然后逐步收敛的
    • 若训练过程中突然下降,大概率是模型崩溃了(此时往往伴随着 KL 散度崩溃)

loss 类

Actor Loss

  • Actor Loss 本身不太具备特殊含义,比如 REINFORCE 中的目标函数/损失函数更多只是为了凑出梯度来
  • 经验:
    • Actor Loss 这个值可能是波动的,看不到收敛迹象,但不太会逐步提升
    • 若 Actor Loss 先稳定,然后突然提升, 那么大概率是在这里发生模型崩溃了

Critic Loss

  • Critic Loss 代表了 Critic 模型拟合状态价值函数的能力
  • 经验:
    • 这个值在 Pretrain 阶段是大幅下降的,然后逐步收敛,在 Actor 和 Critic 交叉训练过程中保持收敛状态
    • 若在 Actor 和 Critic 交叉训练过程中, Critic Loss 逐步升高或明显高于 Pretrain 阶段的收敛值,说明当前 Critic 学习速度跟不上 Actor,需要 降低 Actor 学习率 or 提升 Critic 学习率 or 增加 Critic 训练步数(一步 Actor 梯度更新对应多步 Critic 梯度更新)

KL Loss

  • 与 kl 散度本质一致,更多是计算 k2 或 k3

Entropy Loss

  • 与 Entropy 指标一致
  • Rewards 曲线一般直接上报样在进行 Advantages 计算前的奖励情况,一般会上报多个相关值

NLL Loss

  • 仅对 reward > 0 的 token 计算 -mean(action_logprobs * nll_loss_mask),再乘以 rewards_mean,相当于仅更新 reward > 0 的 token,且按照 rewards_mean 加权的交叉熵损失(或负对数似然)更新

MTP Loss

  • MTP Loss

PPO PG IS Clip 相关

  • 重要性采样权重 Clip 相关的各种情况和比例

训推不一致监控和修正各种指标

  • 训推不一致的监控指标
    • training_ppl(包括 log_training_ppl)
    • rollout_ppl(包括 log_rollout_ppl)
    • training_rollout_ppl_diff(包括 training_rollout_log_ppl_diff)
      • training_rollout_log_ppl_diff:一般是统计 Sequence 内部 Token 的 log_probs 均值作为 log_ppl
      • 所以 training_rollout_log_ppl_diff 和下面的 training_rollout_kl(k1) 可能完全相同(若 training_rollout_kl 也是先计算,公式上可看出是同一个值)
        • 注:与 training_rollout_kl(k3) 不同,但趋势几乎一致
    • training_rollout_kl(包含 k1,k3 都可以同时上报,这个值在一些论文中也称为类似 vllm_kl 的命名)
      • 经验:training_rollout_kl 指标在不同模型/显卡上表现不同:
        • 在 GPU H200 卡,Dense 20B 左右的小模型上几乎没有 spike(但训练过程中从 1.2e-3 提升到 2.1e-3 水平,说明这个指标除了与硬件、参数量、模型架构、引擎实现有关外,还与模型自身的输出分布有关)
        • 在 NPU 910B 卡,MoE 500B-A30B 左右的大模型上有许多 spike(训练过程中从 1.5e-3 级别会 spike 到 1.5e-2 级别,甚至偶尔到 5e-2 级别,但不会逐步提升,与 Dense 20B 不同)
  • 极端 Token(Rollout 或 Training 引擎上的低概率 Token 比例,一般如 < 1e-3 等)
  • IS(Importance Sampling)修复系数的统计(均值、标准差等)
    • Sequence-level 和 Token-level 的
  • RS(Rejection Sampling)过滤系数的统计(均值、标准差等)
    • 注:被拒绝等价于被 Mask
    • 按照配置上报,也可以同时上报下面三种:
      • Token-level:按照 Sequence-level 作为过滤粒度,Token 粒度的 IS 比例作为判断依据
      • Sequence-level:按照 Sequence-level 作为过滤粒度,序列内所有 Token 系数的乘积(联合概率分布)作为判断依据
      • Geometric(Sequence-level):按照 Sequence-level 作为过滤粒度,序列内所有 Token 系数的集合平均作为判断依据

Advantages

  • Advantages 表示了 Token 的打压 or 鼓励情况,是最关键的指标之一
  • PPO 的 Advantages 分为 Batch Norm 归一化前和归一化后的,要分别监控
  • 经验:
    • Advantages 的上报(比如均值,最大值等)应该完整符合推导结果(比如均值为 0, 最大值不能超过理论上限等)
    • PPO 中未经过归一化的 Advantages 均值也应该几乎贴近于 0,若大于 0 则大概率是 Critic 还没收敛

附录:遇到过的各种指标异常现象总结

PPO 训练过程中 \((\pi_\theta - \pi_{old})\) 逐步增大意味着什么?

  • 首先:对于每个 Rollout Step,第一个梯度 Step 的 \((\pi_\theta - \pi_{old})\) 始终为 0
    • 因为在每次 Rollout 结束后,当前策略网络 \(\pi_\theta\) 的权重与用于采样的旧策略网络 \(\pi_{old}\) 完全一致
  • 若观察到随着 Rollout 步数(或训练 Epoch)的增加,第二个梯度 Step 的 \((\pi_\theta - \pi_{old})\) 的方差越来越大,这是一个需要警惕的信号
    • 通常意味着单次梯度更新对策略产生了越来越剧烈且不均匀的扰动
  • 注:此时 PPO 训练是使用了 Advantages Batch Normalization 的
    • 即:无论 Reward Model 给出的分数多么离谱,无论原始的 TD Error 方差有多大,输入到 Loss 函数中的 Advantage 的均值严格为 0,方差严格为 1
可能问题1:优势函数(Advantage)的方差在急剧扩大(若没有进行 Advantages Batch Normalization )
  • 在 PPO 中,策略更新的梯度方向和步长主要由优势函数 \(A(s,a)\) 决定
  • 若没有进行 Advantages Batch Normalization,且 \((\pi_\theta - \pi_{old})\) 的方差变大,说明在同一个 Batch 内,某些 Token 的概率被剧烈推高,而另一些被剧烈压低
  • 问题发生的链路:
    • 随着 Rollout 的进行,策略模型逐渐偏离初始的 SFT 模型,开始探索到奖励模型未曾见过的状态(OOD 数据)
      • 这会导致 RM 输出极端的、充满噪声的奖励(即已经发生了 Reward Hacking)
    • 进一步导致Critic 无法准确拟合这些极端的奖励,导致计算出的 Advantage 出现极大的方差
    • 高方差的 Advantage 直接导致了高方差的梯度更新
  • 可能导致这一点的原因包含:
    • KL Penalty 不足,导致模型偏离 Reference Model 太多,从而进入 RM 的 OOD 区间,打分开始变得激进且误差大
      • 实验验证:完全关闭 KL Penalty 后,Advantages 的方差 和 \((\pi_\theta - \pi_{old})\) 方差
  • 解法:增加 KL Penalty 约束,防止模型发生 Reward Hacking
可能的问题2:Critic 跟不上 Actor 的更新
  • Critic 的任务是预测当前状态的价值 \(V(s)\)
  • 如果 Actor 的策略变化过快,或者 Critic 的学习率/拟合能力不足,Critic 的预测就会产生滞后
  • 不准确的 \(V(s)\) 会导致 Advantage 估计错误
    • 随着 Rollout 步数增加,这种误差会累积,导致 Actor 在第二步更新时接收到极其嘈杂的梯度信号(有偏的,有误差的信号),从而表现为更新幅度的方差增大
  • 解法:提升 Critic 学习率、降低 Actor 学习率,或者让 Critic 多更新几步

GRPO 训练时 Advantage 均值不为 0

  • Advantage 的均值上报逻辑为:经过组内归一化,再按照序列平均(Token 相等,平均前后值不变),最后在 Sequence 间做归一化
    • 理论上,这个值肯定为 0,因为组内均值为 0,多个组均值也应该为 0
  • 发生问题的原因:
    • 当同一个组的 rewards 完全相同且是无理数时,可能因为精度问题导致组内归一化 Bug
      • 自身减去组内均值不为 0,是一个相同的 1e-8 量级的正数或者负数(具体正负和原始 rewards 值有关)
      • 1e-8 量级的数再除以 (std + 1e-8) 得到在 1 左右的 Advantage(且同一组内所有 Token 值相等)
  • 可选的解法(亲测:std 相关的解法本质都是避免针对过小的 std 组做过大的更新,这类解法可以大幅缩小训练过程中的 kl spike):
    • mask std 过小的组
    • 提升 std 的兜底值
    • 不使用 std 作为分母(如 Dr.GRPO)
    • (精确修复)提升归一化计算的精度(默认 torch 使用 float32,可以切换为 float64 计算)

附录:VeRL 框架上报指标汇总(20260423)

  • 本次梳理 commit_id 为:6daa6e0c1626567a90454a9a79076b30d34d4d8b
  • VeRL 所有指标最终通过下面代码统一上报到 wandb/swanlab 等实验跟踪后端
    1
    logger.log(data=metrics, step=self.global_steps)

Actor 相关指标

actor/pg_loss,策略梯度损失
  • 文件:verl/workers/actor/dp_actor.py:477

  • 由compute_policy_loss_* 系列函数计算后返回,以标准 PPO(vanilla)为例:

    1
    2
    3
    4
    5
    # 令 ratio = exp(log_prob - old_log_prob)(当前策略与旧策略的概率比)
    pg_losses1 = -advantages * ratio
    pg_losses2 = -advantages * clip(ratio, 1 - ε_low, 1 + ε_high)
    pg_losses= max(pg_losses1, pg_losses2) #(带 dual-clip 时取 min 与 c·A)
    pg_loss = agg_loss(pg_losses, response_mask, loss_agg_mode)
  • 其中 agg_loss 支持四种聚合模式(由配置决定):

    • token-mean:所有有效 token 的均值
    • seq-mean-token-sum:先对序列内的 token 求和,再对序列取均值
    • seq-mean-token-mean:先对序列内的 token 取均值,再对序列取均值
    • seq-mean-token-sum-norm:先对序列内的 token 求和,再除以最大响应长度(然后对序列取均值)
actor/pg_clipfrac,策略梯度上裁剪比例
  • 文件:verl/workers/actor/dp_actor.py:478

    1
    pg_clipfrac = masked_mean(pg_losses2> pg_losses1, response_mask)
    • 即 ratio > 1 + ε_high 时被裁剪的token 比例(正优势方向被clip 的占比)
actor/pg_clipfrac_lower,策略梯度下裁剪比例(dual-clip)
  • 文件:verl/workers/actor/dp_actor.py:480

    1
    2
    3
    4
    pg_clipfrac_lower = masked_mean(
    (clip_pg_losses1 > pg_losses3) * (advantages< 0).float(), response_mask
    )
    # 其中 pg_losses3 = -advantages * clip_ratio_c(默认 c=3.0)
  • 仅在 advantages < 0(负优势)时,因 dual-clip 机制触发的 token 比例

    • GSPO/GPG/KL-Cov 等方法此项固定为 0
actor/ppo_kl,近似 KL 散度
  • 文件:verl/workers/actor/dp_actor.py:479

    1
    2
    3
    negative_approx_kl = clamp(log_prob - old_log_prob, -20, 20)
    ppo_kl = masked_mean(-negative_approx_kl, response_mask)
    = masked_mean(old_log_prob - log_prob, response_mask)
  • 这是 KL(π_old ‖ π) 的 token 级均值估计(一阶近似)

    • 在 kl_cov 变体中,上报的是绝对值均值 masked_mean(|log_prob - old_log_prob|, response_mask)
actor/entropy,策略熵
  • 文件:verl/trainer/ppo/ray_trainer.py:1189

  • 在 recompute old_log_probs 阶段从 Actor 的 forward 中获取:

    1
    2
    token_entropy = logsumexp(logits) - sum(softmax(logits) * logits)  (每个 token 的熵)
    entropy_agg = agg_loss(entropys, response_mask, loss_agg_mode)
    • 是对 response token 的熵按指定 loss 聚合模式计算的标量值
  • 补充:上述 token_entropy 的公式推导:
    $$
    \begin{align}
    \text{Entropy}
    &= -\sum_i p_i \log p_i \\
    &= -\sum_i p_i \log\left( \frac{e^{z_i}}{\sum_j e^{z_j}} \right) \\
    &= -\sum_i p_i \left( z_i - \log\sum_j e^{z_j} \right) \\
    &= -\sum_i p_i z_i + \log\sum_j e^{z_j} \cdot \sum_i p_i \\
    &= -\sum_i p_i z_i + \log\sum_j e^{z_j} \\
    &= \text{logsumexp}(z) - \sum_i \left( \text{softmax}(z)_i \cdot z_i \right)
    \end{align}
    $$

actor/kl_loss,KL 损失(有条件上报)
  • 文件:verl/workers/actor/dp_actor.py:465

  • 仅当 config.use_kl_loss = True 时上报:

    1
    2
    kld = kl_penalty(log_prob, ref_log_prob, kl_penalty_type)
    kl_loss = agg_loss(kld, response_mask, loss_agg_mode) * loss_scale_factor
    • 其中 kl_penalty_type 支持多种估计方式(详见 KL 惩罚计算)
actor/kl_coef,KL 损失系数(有条件上报)
  • 文件:verl/workers/actor/dp_actor.py:466

    1
    actor/kl_coef = config.kl_loss_coef
    • 仅当 use_kl_loss = True 时上报,为超参数常量
actor/reward_kl_penalty,奖励中的 KL 惩罚均值
  • 文件:verl/trainer/ppo/ray_trainer.py:191

  • 仅当 algorithm.use_kl_in_reward = True 时,在 apply_kl_penalty 函数中计算:

    1
    2
    3
    4
    kld = kl_penalty(old_log_probs, ref_log_prob, kl_penalty_type)# token level
    kld = kld * response_mask
    current_kl = mean(masked_mean(kld, response_mask, axis=-1)) # 先对序列内平均,再对 batch 平均
    actor/reward_kl_penalty = current_kl
  • 最终 token 级奖励 = token_level_scores - beta * kld

actor/reward_kl_penalty_coeff, KL 惩罚自适应系数
  • 文件:verl/trainer/ppo/ray_trainer.py:191

    1
    actor/reward_kl_penalty_coeff = kl_ctrl.value #(beta,自适应调整)
  • 实际实现时是设定一个 kl 目标,若未达到目标,则通过提升或降低 kl_coef 来调整模型朝达到目标的方向走

actor/grad_norm,Actor 梯度范数
  • 文件:verl/workers/actor/dp_actor.py:486

    1
    actor/grad_norm = clip_grad_norm_(actor_module, max_grad_norm).item()
  • 由 _optimizer_step() 在梯度裁剪时返回

actor/lr,Actor 学习率
  • 文件:verl/workers/fsdp_workers.py:727 / verl/workers/megatron_workers.py:610
    1
    actor/lr = actor_lr_scheduler.get_last_lr()[0]
perf/mfu/actor,Actor MFU(模型浮点利用率)
  • 文件:verl/workers/fsdp_workers.py:719
    1
    2
    estimated_flops = flops_counter.estimate_flops(global_num_tokens, delta_time)
    perf/mfu/actor = estimated_flops * ppo_epochs / promised_flops / world_size

Critic 相关指标

critic/vf_loss,价值函数损失
  • 文件:verl/workers/critic/dp_critic.py:254
    1
    2
    3
    4
    5
    vpredclipped = clip(vpreds, values - ε_v, values + ε_v)
    vf_losses1 = (vpreds - returns)**2
    vf_losses2 = (vpredclipped - returns)**2
    clipped_vf_losses = max(vf_losses1, vf_losses2)
    vf_loss = 0.5 * agg_loss(clipped_vf_losses, response_mask, loss_agg_mode) * loss_scale
critic/vf_clipfrac,价值函数裁剪比例
  • 文件:verl/workers/critic/dp_critic.py:255

    1
    vf_clipfrac = masked_mean(vf_losses2 > vf_losses1, response_mask)
  • 即被 clip 的价值预测所占 token 比例

critic/vpred_mean,价值预测均值
  • 文件:verl/workers/critic/dp_critic.py:256

    1
    critic/vpred_mean = masked_mean(vpreds, response_mask)
  • 当前 mini-batch 中,Critic 对 response token 的价值预测均值

critic/grad_norm,Critic 梯度范数
  • 文件:verl/workers/critic/dp_critic.py:263
    1
    critic/grad_norm = clip_grad_norm_(critic_module, max_grad_norm).item()
critic/lr,Critic 学习率
  • 文件:verl/workers/fsdp_workers.py:1307 / verl/workers/roles/critic.py:178
    1
    critic/lr = critic_lr_scheduler.get_last_lr()[0]
perf/mfu/critic,Critic MFU
  • 文件:verl/workers/fsdp_workers.py:1304 / verl/workers/roles/critic.py:176
    1
    perf/mfu/critic = estimated_flops * ppo_epochs / promised_flops / world_size

奖励/优势/回报统计指标(compute_data_metrics)

  • 这些指标在 verl/trainer/ppo/metric_utils.py:compute_data_metrics 中统一计算
得分(Score)指标
  • 具体指标:
    指标 计算方式
    critic/score/mean mean(token_level_scores.sum(-1)[non_aborted]),非中止样本的序列级得分均值,注意是对 Sequence 内部的 Token 做 sum
    critic/score/max max(token_level_scores.sum(-1)[non_aborted])
    critic/score/min min(token_level_scores.sum(-1)[non_aborted])
  • 注:token_level_scores 是 reward function 输出的原始得分(未减去 KL 惩罚)
奖励(Reward)指标
  • 具体指标
    指标 计算方式
    critic/rewards/mean mean(token_level_rewards.sum(-1)[non_aborted]),含KL 惩罚后的实际奖励均值
    critic/rewards/max max(token_level_rewards.sum(-1)[non_aborted])
    critic/rewards/min min(token_level_rewards.sum(-1)[non_aborted])
  • 注:token_level_rewards = token_level_scores - beta * kld(若开启 KL penalty)
优势值(Advantages)指标
  • 具体指标
    指标 计算方式
    critic/advantages/mean mean(advantages[response_mask]),所有有效 response token 的优势值均值
    critic/advantages/max max(advantages[response_mask])
    critic/advantages/min min(advantages[response_mask])
回报(Returns)指标
  • 具体指标:
    指标 计算方式
    critic/returns/mean mean(returns[response_mask]),有效 token 的回报均值(GAE 计算结果)
    critic/returns/max max(returns[response_mask])
    critic/returns/min min(returns[response_mask])
价值预测(Values)指标(仅 use_critic=True 时)
  • 具体指标
    指标 计算方式
    critic/values/mean mean(values[response_mask]),Critic 旧价值预测均值
    critic/values/max max(values[response_mask])
    critic/values/min min(values[response_mask])
    critic/vf_explained_var 1 - var(returns - values) / (var(returns) + 1e-5),Critic 价值函数解释方差,越接近 1 表示价值估计越准确

序列长度相关指标

  • 均在 verl/trainer/ppo/metric_utils.py:compute_data_metrics 中计算
    指标 计算方式
    response_length/mean mean(response_mask.sum(-1).float()),所有样本 response token 数均值
    response_length/max max(response_length) ,response_length=response_mask.sum(-1).float()
    response_length/min min(response_length)
    response_length/clip_ratio mean(response_length == max_response_length),达到最大长度被截断的样本比例
    response_length_non_aborted/mean 仅统计 response_length > 0 的样本的response 长度均值
    response_length_non_aborted/max 同上,取最大值
    response_length_non_aborted/min 同上,取最小值
    response_length_non_aborted/clip_ratio 非中止样本中达到最大长度的比例
    response/aborted_ratio mean(response_length == 0),响应长度为 0(中止/异常)的样本比例
    prompt_length/mean mean(prompt_mask.sum(-1).float()),prompt token 数均值
    prompt_length/max max(prompt_length)
    prompt_length/min min(prompt_length)
    prompt_length/clip_ratio mean(prompt_length == max_prompt_length),prompt被截断的比例

多轮对话相关指标(按条件展示)

  • 当 batch 中存在 __num_turns__ 字段时上报(多轮对话场景):
    指标 计算方式
    num_turns/mean num_turns.mean(),每条样本的对话轮数均值
    num_turns/max num_turns.max()
    num_turns/min num_turns.min()
  • 当 batch 中存在 tool_call_counts 字段时上报(工具调用场景):
    指标 计算方式
    tool_call_counts/mean tool_call_counts.mean()
    tool_call_counts/max tool_call_counts.max()
    tool_call_counts/min tool_call_counts.min()

时序性能指标

  • 在 verl/trainer/ppo/metric_utils.py 的 compute_timing_metrics 和 compute_throughout_metrics 中计算
各阶段耗时(timing_s/*)
  • 具体逻辑:

    1
    timing_s/{stage_name} = 该阶段实际执行时间(秒)
  • 阶段名称包括:

    • gen(生成)
    • ref(参考策略前向)
    • values(价值估计)
    • adv(优势计算)
    • update_critic(Critic 更新)
    • update_actor(Actor 更新)
    • reward(奖励计算)
    • old_log_prob(log prob 重计算)
    • testing(验证)
    • save_checkpoint(保存检查点)
    • step(总步时间)等
每 token 耗时(timing_per_token_ms/*)
  • 具体计算逻辑:

    1
    2
    timing_per_token_ms/gen = timing_s["gen"] * 1000 / num_response_tokens
    timing_per_token_ms/{其他阶段} = timing_s[stage] * 1000 / (num_prompt_tokens + num_response_tokens)
  • 注:gen 阶段只用 response token 数作分母,其他阶段用全序列 token 数

吞吐量与性能指标

  • 吞吐量相关指标:
    指标 计算方式
    perf/total_num_tokens sum(batch.meta_info["global_token_num"]),本step 处理的 token 总数
    perf/time_per_step timing_raw["step"],每个训练步的总耗时(秒)
    perf/throughput total_num_tokens / (time * n_gpus),每 GPU 每秒处理的 token 数
    perf/max_memory_allocated_gb torch.device.max_memory_allocated() / 1024³,GPU峰值显存分配(GB)
    perf/max_memory_reserved_gb torch.device.max_memory_reserved() / 1024³,GPU 峰值显存预留(GB)
    perf/cpu_memory_used_gb psutil.virtual_memory().used / 1024³,CPU 内存占用(GB)

训练状态指标

  • 在 verl/trainer/ppo/ray_trainer.py:1343 处直接赋值:
    指标 计算方式
    training/global_step 当前全局训练步数 self.global_steps
    training/epoch 当前 epoch 编号

批次负载均衡指标(按条件展示)

  • 当 trainer.balance_batch = True 时,在 _balance_batch 方法中调用 log_seqlen_unbalance(verl/utils/seqlen_balancing.py:194):
    指标 计算方式
    global_seqlen/min 各DP rank 的序列长度之和中的最小值(均衡前)
    global_seqlen/max 各 DP rank 的序列长度之和中的最大值(均衡前)
    global_seqlen/minmax_diff max - min(不均衡程度,越小越好)
    global_seqlen/balanced_min 均衡后各 rank 的序列长度最小值
    global_seqlen/balanced_max 均衡后各 rank 的序列长度最大值
    global_seqlen/mean 各 rank 序列长度总和的均值

Debug 调试指标(按照条件展示)

  • 当 batch 中存在 rollout_log_probs 字段时(开启 rollout vs actor prob 对比),在 verl/utils/debug/metrics.py:calculate_debug_metrics 中计算:
    指标 计算方式
    training/rollout_probs_diff_valid 固定为 1,表示输入有效
    training/rollout_probs_diff_max `max(
    training/rollout_probs_diff_mean `mean(
    training/rollout_probs_diff_std `std(
    training/rollout_actor_probs_pearson_corr corrcoef(actor_probs, rollout_probs)[response_mask],Pearson 相关系数,参考 arXiv:2506.13585

验证集指标(按条件展示)

  • 在 _validate() 方法中,通过 process_validation_metrics 函数处理后以多种统计形式上报:

    1
    2
    3
    4
    5
    6
    7
    8
    val-core/{data_source}/{var_name}/mean@N 		# N 个 rollout 的均值
    val-core/{data_source}/{var_name}/std@N # N 个 rollout 的标准差
    val-core/{data_source}/{var_name}/best@N/mean # bootstrap 采样最优值的均值(pass@N变体)
    val-core/{data_source}/{var_name}/best@N/std # bootstrap 采样最优值的标准差
    val-core/{data_source}/{var_name}/worst@N/mean # bootstrap 采样最差值的均值
    val-core/{data_source}/{var_name}/worst@N/std # bootstrap 采样最差值的标准差
    val-core/{data_source}/{var_name}/maj@N/mean # 多数投票结果的均值(如有 pred字段)
    val-core/{data_source}/{var_name}/maj@N/std # 多数投票结果的标准差
  • 其中 N 以 2 的幂次递增(2, 4, 8, …, n_resps),bootstrap 采样 1000 次。val-aux/ 前缀同理,还额外包含:

    • val-aux/num_turns/mean|max|min(多轮验证时)

补充:KL 散度计算方式汇总

  • kl_penalty 函数(core_algos.py:1272)支持多种 KL 估计变体:
    类型 公式
    kl / k1 logprob - ref_logprob(一阶线性近似)
    abs `
    mse / k2 0.5 * (logprob - ref_logprob)²
    low_var_kl / k3 exp(clamp(ref - log, -20, 20)) - (ref - log) -1(低方差估计,来自 Schulman 2020)

NLP——Megatron源码阅读笔记

  • 参考链接:
    • 各种并行的通信量:图解大模型训练之:张量模型并行(TP),Megatron-LM - 猛猿的文章 - 知乎
    • 模型分组及方式理解:Megatron-LM训练的大模型如何分组? - wx1997的文章 - 知乎
      • DP,TP,PP:TP > DP > PP,优先在机器内进行TP,其次是DP,最后是PP,因为通信量上是TP最多,DP其次,PP最后
    • 图解大模型系列之:Megatron源码解读1,分布式环境初始化 - 猛猿的文章 - 知乎
    • 图解大模型训练之:Megatron源码解读2,模型并行 - 猛猿的文章 - 知乎
    • 图解大模型训练系列之:Megatron源码解读3,分布式混合精度训练 - 猛猿的文章 - 知乎
    • Megatron-LM 中使用 DeepSpeed 加速:(DeepSpeed 官方文档)在 Megatron-LM 中加入 DeepSpeed
    • 专家并行论文:GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, 2020, Google
    • [张量/序列并行]📚图解 DeepSpeed-Ulysses & Megatron-LM TP/SP - DefTruth的文章 - 知乎:写的非常详细
    • [转]Megatron-LM源码系列(八): Context Parallel并行 - 李睿的文章 - 知乎:作者还有一些博客内容可供参考
      • Megatron-LM源码系列(一):模型并行初始化
      • Megatron-LM源码系列(二):Tensor模型并行和Sequence模型并行训练
      • Megatron-LM源码系列(三):详解Pipeline模型并行训练实现
      • Megatron-LM源码系列(四):重计算(recompute)
      • Megatron-LM源码系列(五): FP16使用
      • Megatron-LM源码系列(六):Distributed-Optimizer分布式优化器实现Part1
      • Megatron-LM源码系列(七):Distributed-Optimizer分布式优化器实现Part2
      • Megatron-LM源码系列(八): Context Parallel并行
    • Megatron 新版 MoE 源码阅读 - Fizzmy的文章 - 知乎
    • 知乎专栏:跟着执行流程阅读源码系列:
      • 跟代码执行流程,读Megatron源码(一)从目录结构开始 - AIer的文章 - 知乎
      • 跟代码执行流程,读Megatron源码(二)训练入口pretrain_gpt.py - AIer的文章 - 知乎
      • 跟代码执行流程,读Megatron源码(三)megatron训练脚本training.py之pretrain() - AIer的文章 - 知乎
      • 跟代码执行流程,读Megatron源码(四)megatron训练脚本initialize.py之initialize_megatron()分布式环境初始化 - AIer的文章 - 知乎

整体说明

  • 本文的代码以 Megatron-LM 20250904 的版本 e000263e21ac89571123303c4043ec9ea7261513 为主,还有部分更早的版本的代码(之前写的,没有修改)

Megatron 数据处理

  • 数据预处理负责将 .jsonl 的文本数据 tokenize 并处理成 Megatron 可以直接读取的数据格式(.bin 和 .idx 类型的文件),减少训练时的数据处理时间

  • 数据处理的使用方式详情参考:github.com/NVIDIA/Megatron-LM

    • 准备 .jsonl 文件,文件格式如下:

      1
      2
      {"text": "Your training text here..."}
      {"text": "Another training sample..."}
    • 数据预处理:

      1
      2
      3
      4
      5
      6
      7
      python tools/preprocess_data.py \
      --input data.jsonl \
      --output-prefix processed_data \
      --tokenizer-type HuggingFaceTokenizer \
      --tokenizer-model /path/to/tokenizer.model \
      --workers 8 \
      --append-eod
      • output-prefix:输出文件的前缀
      • append-eod:是否添加 EOD Token?
      • 注意:还可以根据需要设置 split_sentences 参数,对文档进行拆分成 sentence 再做 tokenize
  • process_data.py 的核心处理逻辑如下:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    # split sentences in partition files
    if args.split_sentences and not split_sentences_present:
    processes = []
    for name in in_ss_out_names:
    p = multiprocessing.Process(target=partition.split_sentences, # TODO(ZJH): 构造完成数据 sentence 分隔的进程
    args=((name['partition'], name['sentence_split']),)) # TODO(ZJH): 参数是输入文件名和分隔结果的输出文件名
    p.start() # TODO(ZJH): 启动进程
    processes.append(p)

    for p in processes:
    p.join()

    if args.partitions == 1:
    return


    # encode partition files in parallel
    processes = []
    input_key = 'sentence_split' if args.split_sentences else 'partition' # TODO(ZJH): 根据是否走 sentence_split 来选择输入文件
    for name in in_ss_out_names:
    p = multiprocessing.Process(target=partition.process_json_file, # TODO(ZJH): 构造完成数据 encode 的进程
    args=((name[input_key], name['output_prefix']),)) # TODO(ZJH): 参数是输入文件名(上一步处理后的)和分隔结果的输出文件名
    p.start() # TODO(ZJH): 启动进程
    processes.append(p)

    for p in processes:
    p.join()

    if args.partitions == 1:
    return

Megatron-LM 训练过程梳理

  • 总入口(以 GPTModel 为例):
    1
    2
    3
    4
    5
    6
    7
    8
    9
    pretrain(
    train_valid_test_datasets_provider,
    partial(model_provider, gpt_builder), # TODO(ZJH): model_provider 调用 gpt_builder 构造模型
    ModelType.encoder_or_decoder,
    forward_step,
    args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
    extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,
    store=store,
    )

pretrain 函数是总入口,包含的核心参数如下

pretrain 参数一:train_valid_test_datasets_provider,负责管理数据
  • 输出返回迭代器,这个迭代器每个 Batch 将包含一个 micro-batch 数据
pretrain 参数二:partial(model_provider, gpt_builder),对应 model_provider 参数,负责构造并,返回模型对象
  • 返回对象的 __init__ 函数负责实现 模型结构定义

  • 返回的模型对象会实现一个 forward 函数

  • 该函数依次调用 _preprocess(),decoder() 和 _postprocess() 实现整体逻辑

    • _preprocess 负责处理输入层,包含位置编码等信息,返回 decoder 的输入
    • decoder 负责模型主要前向流程
    • _postprocess 负责处理输出层,包括 MTP 处理、 损失函数定义等,返回 损失函数值(lm_loss, 交叉熵损失)
      • 若需要执行 MTP 过程,执行 MTP 过程,同时若打开训练,则 MTP loss 在这里被计算(
        • 使用 mtp_num_layers 来表示 MTP 的深度,每深一层都会多预测一个 Token,每层对应交叉熵损失,然后乘以 loss_mask
        • 处理后的 MTP 损失使用 MTPLossAutoScaler(是 torch.autograd.Function 的子类,是 PyTorch 自定义算子的实现 ) 实现前向和反向传播
          1
          2
          3
          4
          5
          6
          7
          8
          9
          10
          11
          12
          13
          14
          15
          16
          17
          18
          19
          20
          21
          22
          23
          24
          25
          26
          27
          28
          29
          30
          31
          32
          33
          34
          35
          36
          37
          38
          39
          40
          41
          42
          43
          44
          45
          46
          47
          48
          49
          50
          51
          52
          53
          54
          55
          56
          57
          58
          59
          60
          61
          62
          63
          64
          65
          66
          67
          68
          # TODO(ZJH): 将 MTP 的每个 Token(t+2...t+k)的 loss 都添加到(通过特殊的自定义算子)主网络的计算依赖上,从而保证对主网络求梯度时,MTP 相关的梯度也能回传
          if self.mtp_process:
          mtp_labels = labels.clone()
          hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0) # TODO(ZJH): 将多个 hidden_state 拆开
          hidden_states = hidden_states_list[0] # TODO(ZJH): 主函数的输出
          if loss_mask is None:
          # if loss_mask is not provided, use all ones as loss_mask
          loss_mask = torch.ones_like(mtp_labels)
          for mtp_layer_number in range(self.config.mtp_num_layers): # TODO(ZJH): 每一层都计算 loss,每一层代表一个更深的未来 Token 预测目标
          # output
          mtp_logits, _ = self.output_layer( # TODO(ZJH): 每个 hidden_states 都要走 output_layer 得到 logits 后再计算损失
          hidden_states_list[mtp_layer_number + 1], # TODO(ZJH): hidden_states_list[0] 是主网络的 hidden_state
          weight=output_weight,
          runtime_gather_output=runtime_gather_output,
          )
          # Calc loss for the current Multi-Token Prediction (MTP) layers.
          mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group) # TODO(ZJH): MTP 目标逐步后移
          loss_mask, num_tokens = roll_tensor(
          loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group
          )
          mtp_loss = self.compute_language_model_loss(mtp_labels, mtp_logits)
          mtp_loss = loss_mask * mtp_loss
          if self.training:
          # TODO(shifangx): remove the use of parallel_state here
          # after moving loss logging to loss_func in pretrain_gpt.py
          MTPLossLoggingHelper.save_loss_to_tracker( # TODO(ZJH): for logging
          torch.sum(mtp_loss) / num_tokens,
          mtp_layer_number,
          self.config.mtp_num_layers,
          avg_group=parallel_state.get_data_parallel_group(
          with_context_parallel=True
          ),
          )
          mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers # TODO(ZJH): 根据参数和层数计算 scale,除以 mtp_num_layers 得到平均值,保证总的 MTP loss 量级(影响)不变
          if self.config.calculate_per_token_loss: # TODO(ZJH): 判断损失是否按照 Token 做平均
          # TODO(ZJH): MTPLossAutoScaler 是特殊的自定义算子,不改变第一个参数的值(输入即输出),求导时直接返回 第二个参数(loss)*scale 作为梯度
          hidden_states = MTPLossAutoScaler.apply( # TODO(ZJH): 经过这个自定义算子后,不会改变 hidden_states 的值(注意 hidden_states 始终是主网络的隐藏层),但对 hidden_states 计算梯度会直接返回 mtp_loss_scale * mtp_loss
          hidden_states, mtp_loss_scale * mtp_loss
          ) # TODO(ZJH): hidden_states 经过所有层后,最终得到的是所有 MTP 层 Token 的梯度(多个深度的 Token 一起)
          else:
          hidden_states = MTPLossAutoScaler.apply(
          hidden_states, mtp_loss_scale * mtp_loss / num_tokens
          )
          sequence_parallel_override = False
          if in_inference_mode and inference_context.materialize_only_last_token_logits:
          if inference_context.is_static_batching():
          hidden_states = hidden_states[-1:, :, :]
          else:
          if self.output_layer.sequence_parallel:
          # Perform the sequence parallel gather here instead of after the output layer
          # because we need to slice the last token logits from the full view of the
          # packed logits across all requests.
          # TODO(ksanthanam): Make the equivalent change in the `MambaModel` code after
          # merging in !3722.
          hidden_states = gather_from_sequence_parallel_region(
          hidden_states, group=self.model_comm_pgs.tp
          )
          self.output_layer.sequence_parallel = False
          sequence_parallel_override = True
          # Reshape [B, 1, H] to [1, B, H] -> extract each sample’s true last‐token hidden
          # state ([B, H]) -> unsqueeze back to [1, B, H]
          # (so that the output layer, which expects S×B×H, receives only the final token)
          hidden_states = inference_context.last_token_logits(
          hidden_states.squeeze(1).unsqueeze(0)
          ).unsqueeze(1)
          logits, _ = self.output_layer( # TODO(ZJH):主网络的 logits 计算
          hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output
          )
  • model_provider 函数(也是 pretrain 的入参)会返回一个模型 model

    • 这个返回的模型实现了 forward 函数,model.forward 函数整体返回值是一个 loss 值(主网络的 loss,但计算图上带着 MTP 所有深度上的 loss),该值是由 _postprocess 返回的值
pretrain 参数三:forward_step,输入参数包括模型,负责调用模型执行前向过程,并返回 loss 指针等
  • 返回的 loss 函数指针可以被调用,从而计算 loss

pretrain 的工作包括环境初始化,执行训练过程等

第一步:initialize_megatron(),初始化分布式环境,包括 TP,PP,DP 等的子进程组等
第二步:setup_model_and_optimizer(),定义模型架构,切割模型,完成 optimizer 初始化
第三步:build_train_valid_test_data_iterators(), 获取数据 iterator
第四步:train(),训练入口
  • 训练的入参包括上面得到的各种结果

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    def train(
    forward_step_func,
    model,
    optimizer,
    opt_param_scheduler,
    train_data_iterator,
    valid_data_iterator,
    process_non_loss_data_func,
    config,
    checkpointing_context,
    non_loss_data_func,
    ):
  • train_step():训练过程包含一个主要的 while 循环,每次走一个 train_step()

    1
    2
    def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func): # TODO(ZJH): 单步训练入口
    """Single training step."""
    • train_step 第一步:forward_backward_func(),完成一次前向和后向过程,是训练的核心函数,也最难

      • 实际上调用的函数 forward_backward_func 经过层层函数传递 train_step() <- train() <- megatron/core/pipeline_parallel/schedules.py,最终可追述到 schedules.py 文件的 get_forward_backward_func() 函数

        1
        2
        3
        4
        5
        6
        7
        8
        9
        10
        def get_forward_backward_func():
        pipeline_model_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
        if pipeline_model_parallel_size > 1: # TODO(ZJH): 若打开 PP
        if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None:
        forward_backward_func = forward_backward_pipelining_with_interleaving # TODO(ZJH): 若打开 interleaving pipeline 调度
        else:
        forward_backward_func = forward_backward_pipelining_without_interleaving # TODO(ZJH): 若关闭 interleaving pipeline 调度
        else:
        forward_backward_func = forward_backward_no_pipelining # TODO(ZJH): 没有 PP 的情况
        return forward_backward_func
      • 若打开 PP

        • 开启 interleaving pipeline:forward_backward_pipelining_with_interleaving
          • 负责实现对应的 1F1B 调度策略,函数内部像 forward_backward_no_pipelining 函数一样,会调用 forward_step 和 backward_step 两步完成前向后向过程和梯度的积累
        • 未开启 interleaving pipeline:forward_backward_pipelining_without_interleaving
          • 在 forward_backward_pipelining_with_interleaving 的基础上,增加了 interleaving 调度策略(实现则更为复杂),进一步优化气泡
        • 其他特殊的 Pipeline 并行调度策略,如 zero_bubble 的调度,实现都在这里新建函数就可以
      • 若为 没有打开 PP 的情况:调用同文件(schedules.py)下的 forward_backward_no_pipelining() 函数,下面是该函数的介绍:

        • 前向过程+后向过程函数为(config.overlap_moe_expert_parallel_comm 为 True):combined_1f1b_schedule_for_no_pipelining

          • 注:config.overlap_moe_expert_parallel_comm 为 True 表示 框架会尝试将专家并行所需的通信操作(如数据传输)与模型的计算操作(如其他层的前向 / 反向计算)重叠进行,而不是等通信完成后再执行计算

            1
            2
            3
            4
            5
            6
            7
            8
            9
            10
            11
            12
            13
            14
            15
            16
            17
            18
            19
            20
            if config.overlap_moe_expert_parallel_comm and not forward_only: # TODO(ZJH): 如果打开 overlap MoE 的专家并行通信(将专家并行(expert parallelism)中的通信操作与计算操作重叠执行),且包含 backward
            # TODO(ZJH): 当 config.overlap_moe_expert_parallel_comm 设为 True 时,框架会尝试将专家并行所需的通信操作(如数据传输)与模型的计算操作(如其他层的前向 / 反向计算)重叠进行,而不是等通信完成后再执行计算
            forward_data_store, total_num_tokens = combined_1f1b_schedule_for_no_pipelining( # TODO(ZJH): 1次前向+1次后向过程
            forward_step_func,
            data_iterator,
            model,
            num_microbatches,
            input_tensor,
            output_tensor_grad,
            forward_data_store,
            config,
            collect_non_loss_data,
            first_val_step,
            forward_only,
            no_sync_func,
            total_num_tokens,
            partial(check_first_val_step, first_val_step, forward_only),
            )
            else:
            # forward_step 和 backward_step 交替执行
          • 分开执行的函数分别为:forward_step 和 backward_step (这里会调用 for 循环完成多个 microbatches, forward_step 和 backward_step 在循环南北部)

            • 前置说明:microbatches - 1 个 microbatches 先调用,然后最后一个负责处理梯度同步等
              1
              2
              3
              4
              5
              6
              7
              8
              9
              10
              11
              12
              13
              14
              15
              16
              17
              18
              19
              20
              21
              22
              23
              24
              25
              26
              27
              28
              29
              30
              31
              32
              33
              34
              35
              36
              37
              38
              39
              40
              41
              42
              43
              44
              45
              46
              47
              48
              49
                  with no_sync_func(): # TODO(ZJH): 如果 no_sync_func 上下文管理器 是一个实际的同步禁用逻辑(比如禁用某些 IO 同步、锁机制等),则代码块会在 “不执行同步” 的环境中运行
              for i in range(num_microbatches - 1): # TODO(ZJH): 每个设备负责多个 microbatches,注意这里少一个
              output_tensor, num_tokens = forward_step( # TODO(ZJH): 前向过程,注意,这里是单个 microbatch 走一次
              forward_step_func,
              data_iterator,
              model,
              num_microbatches,
              input_tensor,
              forward_data_store,
              config,
              grad_finalize_pgs.cp.size(),
              collect_non_loss_data,
              is_first_microbatch=check_first_val_step(first_val_step, forward_only, i == 0),
              current_microbatch=i,
              )
              total_num_tokens += num_tokens # TODO(ZJH): 累加 Token 数
              if not forward_only:
              backward_step( # TODO(ZJH): 后向过程,梯度直接累加(forward 中已经对 loss/num_microbatches)
              input_tensor, output_tensor, output_tensor_grad, model_type, config
              )
              # Run computation for last microbatch out of context handler (want to
              # synchronize gradients).
              output_tensor, num_tokens = forward_step( # TODO(ZJH): 最后一个梯度单独处理,这个梯度的计算要在 context handler 之外,核心原因是确保最后一次梯度计算完成后能触发必要的同步操作,从而保证梯度的正确性和一致性
              forward_step_func,
              data_iterator,
              model,
              num_microbatches,
              input_tensor,
              forward_data_store,
              config,
              grad_finalize_pgs.cp.size(),
              collect_non_loss_data,
              is_first_microbatch=check_first_val_step(
              first_val_step, forward_only, num_microbatches == 1
              ),
              current_microbatch=num_microbatches - 1,
              )
              total_num_tokens += num_tokens # TODO(ZJH): 累加 Token 数
              if not forward_only:
              backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
              # TODO(ZJH): 梯度聚合
              if config.finalize_model_grads_func is not None and not forward_only:
              # Finalize model grads (perform full grad all-reduce / reduce-scatter for
              # data parallelism and layernorm all-reduce for sequence parallelism).
              config.finalize_model_grads_func( # TODO(ZJH): 梯度聚合操作 all_reduce + reduce-scatter ?
              [model],
              total_num_tokens if config.calculate_per_token_loss else None,
              grad_finalize_pgs=grad_finalize_pgs,
              )
          • forward_step() 执行过程:

            1
            2
            3
            4
            5
            6
            7
            8
            9
            10
            11
            12
            13
            14
            15
            16
            17
            18
            19
            20
            21
            forward_step 核心代码如下:

            with context_manager:
            if checkpoint_activations_microbatch is None:
            output_tensor, loss_func = forward_step_func(data_iterator, model)
            else:
            output_tensor, loss_func = forward_step_func(
            data_iterator, model, checkpoint_activations_microbatch
            )
            output_tensor, num_tokens = forward_step_calc_loss(
            model,
            output_tensor,
            loss_func,
            config,
            vp_stage,
            collect_non_loss_data,
            num_microbatches,
            forward_data_store,
            cp_group_size,
            is_last_stage,
            )
            • 其中,forward_step_calc_loss 核心代码如下:
              1
              2
              3
              4
              5
              6
              7
              8
              9
              10
              11
              12
              13
              14
              15
              16
              17
              18
              19
              if is_last_stage: # TODO(ZJH): 只有最后一个 stage 包含 loss,其他 stage 都不需要计算
              if not collect_non_loss_data:
              outputs = loss_func(output_tensor) # TODO(ZJH): 获取损失值,详情见 pretrain_gpt.py 的 loss_func 的返回值(中间使用 forward_step 作为 partial 封装)
              if len(outputs) == 3:
              output_tensor, num_tokens, loss_reduced = outputs
              # TODO(ZJH): 当 calculate_per_token_loss=True 时,损失计算会保留每个 token 的损失值(即按 token 粒度计算损失),通常用于需要获取单 token 损失的场景(如后续可能的梯度裁剪、损失分析等)
              if not config.calculate_per_token_loss: # TODO(ZJH): 当 calculate_per_token_loss=False 时,损失会被归一化(通常除以总 token 数和微批次数量),得到一个全局平均损失,这是训练中更常见的做法(避免损失值因批次大小不同而波动)
              output_tensor /= num_tokens # TODO(ZJH): 视情况看是否需要做 Token 粒度的归一化
              output_tensor /= num_microbatches # TODO(ZJH): 这里是单个 Batch,但除以 num_microbatches,是为了后续 backward 时梯度可以直接累加
              else:
              # preserve legacy loss averaging behavior (ie, over the number of microbatches)
              assert len(outputs) == 2
              output_tensor, loss_reduced = outputs
              output_tensor *= cp_group_size
              output_tensor /= num_microbatches
              forward_data_store.append(loss_reduced)
              else:
              data = loss_func(output_tensor, non_loss_data=True)
              forward_data_store.append(data)
          • backward_step() 执行过程:

            • 后向过程,可选择自定义的 backward 或 PyTorch 标准的官方实现,梯度直接累加(forward 中已经对 loss/num_microbatches)
              1
              2
              3
              4
              5
              if output_tensor[0].requires_grad:
              if config.deallocate_pipeline_outputs:
              custom_backward(output_tensor[0], output_tensor_grad[0]) # TODO(ZJH): 使用自定义的 backward
              else:
              torch.autograd.backward(output_tensor[0], grad_tensors=output_tensor_grad[0]) # TODO(ZJH): 直接使用 backward
    • train_step 第二步:optimizer.step()

      • 用梯度完成一次完整的参数更新
    • train_step 第三步:继续处理 loss 并上报

      • 注:调用完 optimizer 后,还要继续处理 loss 的原因是梯度更新不需要汇总 DP 的 loss,只有上报时需要聚合 所有 DP 的数据

附录:Megatron MTP 损失绑定函数的测试

  • MTP 损失绑定到 main_hidden_states 的方式是通过一个不修改值,但绑定梯度的自定义算子 MTPLossAutoScaler 实现:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    # megatron/core/transformer/multi_token_prediction.py
    class MTPLossAutoScaler(torch.autograd.Function): # TODO(ZJH): 相当于在实现自定义的 PyTorch 算子
    """An AutoScaler that triggers the backward pass and scales the grad for mtp loss."""

    main_loss_backward_scale: torch.Tensor = torch.tensor(1.0)

    @staticmethod
    def forward(ctx, output: torch.Tensor, mtp_loss: torch.Tensor): # TODO(ZJH): 前向过程,存储 loss,返回输入的原始值,不做任何计算
    """Preserve the mtp by storing it in the context to avoid garbage collection.

    Args:
    output (torch.Tensor): The output tensor.
    mtp_loss (torch.Tensor): The mtp loss tensor.

    Returns:
    torch.Tensor: The output tensor.
    """
    ctx.save_for_backward(mtp_loss)
    return output

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor): # TODO(ZJH): 后向过程,获取前向过程存储的 loss,乘以 main_loss_backward_scale 并返回
    """Compute and scale the gradient for mtp loss..

    Args:
    grad_output (torch.Tensor): The gradient of the output.

    Returns:
    Tuple[torch.Tensor, torch.Tensor]: The gradient of the output, scaled mtp loss
    gradient.
    """
    (mtp_loss,) = ctx.saved_tensors
    mtp_loss_backward_scale = MTPLossAutoScaler.main_loss_backward_scale
    scaled_mtp_loss_grad = torch.ones_like(mtp_loss) * mtp_loss_backward_scale
    return grad_output, scaled_mtp_loss_grad
  • MTPLossAutoScaler 算子测试代码:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    import torch

    # 待测试类(直接复用原代码)
    class MTPLossAutoScaler(torch.autograd.Function):
    main_loss_backward_scale: torch.Tensor = torch.tensor(1.0)

    @staticmethod
    def forward(ctx, output: torch.Tensor, mtp_loss: torch.Tensor):
    ctx.save_for_backward(mtp_loss)
    return output

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
    (mtp_loss,) = ctx.saved_tensors
    scaled_mtp_loss_grad = torch.ones_like(mtp_loss) * MTPLossAutoScaler.main_loss_backward_scale
    return grad_output, scaled_mtp_loss_grad

    def test_all_scenarios():
    # 初始化模型可训练参数(所有场景共用)
    x = torch.tensor([2.0], requires_grad=True)
    y = torch.tensor([3.0], requires_grad=True)

    # 定义测试场景配置(场景名、mtp_loss构造、scale系数)
    scenarios = [
    # 场景1:mtp_loss不依赖参数 + scale=1.0
    ("不依赖参数 + scale=1.0", lambda: torch.tensor([5.0], requires_grad=True), 1.0),
    # 场景2:mtp_loss不依赖参数 + scale=0.3
    ("不依赖参数 + scale=0.3", lambda: torch.tensor([5.0], requires_grad=True), 0.3),
    # 场景3:mtp_loss依赖参数(x+y) + scale=1.0
    ("依赖参数3*(x+y) + scale=1.0", lambda: 3*(x + y), 1.0),
    # 场景4:mtp_loss依赖参数(x+y) + scale=0.6
    ("依赖参数3*(x+y) + scale=0.6", lambda: 3*(x + y), 0.6),
    ]

    for scenario_name, mtp_loss_fn, scale in scenarios:
    # 重置梯度和scale系数
    x.grad = None
    y.grad = None
    MTPLossAutoScaler.main_loss_backward_scale = torch.tensor(scale)

    # 1. 计算模型输出和mtp_loss
    output = x * y # 模型输出(固定逻辑:x*y,梯度易验证)
    mtp_loss = mtp_loss_fn()

    # 关键修复:为mtp_loss保留梯度(无论是否为叶子张量)
    mtp_loss.retain_grad() # mtp_loss 依赖模型参数时(mtp_loss 非叶子张量),mtp_loss.grad 不存在,使用 retain_grad() 强制保留其梯度,方便后续查看

    # 2. 使用AutoScaler处理
    scaled_output = MTPLossAutoScaler.apply(output, mtp_loss)

    # 3. 构造总损失并反向传播
    total_loss = scaled_output.sum()
    total_loss.backward()

    # 4. 打印结果(保留1位小数,简洁清晰)
    print(f"=== {scenario_name} ===")
    # 确保grad存在(避免None报错)
    mtp_grad = mtp_loss.grad.item() # 注意:若不使用 mtp_loss.retain_grad(),则 mtp_loss 依赖模型参数时(mtp_loss 非叶子张量),mtp_loss.grad 不存在
    print(f"mtp_loss梯度: {mtp_grad:.1f}") # 验证scale是否生效
    print(f"参数x梯度: {x.grad.item():.1f}") # 验证是否受mtp_loss依赖关系影响
    print(f"参数y梯度: {y.grad.item():.1f}\n") # 验证是否受mtp_loss依赖关系影响

    if __name__ == "__main__":
    test_all_scenarios()

    # === 不依赖参数 + scale=1.0 ===
    # mtp_loss梯度: 1.0
    # 参数x梯度: 3.0
    # 参数y梯度: 2.0
    #
    # === 不依赖参数 + scale=0.3 ===
    # mtp_loss梯度: 0.3
    # 参数x梯度: 3.0
    # 参数y梯度: 2.0
    #
    # === 依赖参数3*(x+y) + scale=1.0 ===
    # mtp_loss梯度: 1.0
    # 参数x梯度: 6.0
    # 参数y梯度: 5.0
    #
    # === 依赖参数3*(x+y) + scale=0.6 ===
    # mtp_loss梯度: 0.6
    # 参数x梯度: 4.8
    # 参数y梯度: 3.8

NLP——技术报告解读-Qwen3

注:本文包含 AI 辅助创作

  • 参考链接:
    • 原始论文:Qwen3 Technical Report, Qwen, 20250514
    • Qwen3技术报告解读-包包算法笔记
    • Qwen3技术报告解读-刘聪NLP
    • Github 链接:github.com/QwenLM/Qwen3
      • github.com/QwenLM/Qwen3/tree/main/examples/llama-factory 文件夹下有 Qwen3 在 Llama-factory 框架下微调的方法

Paper Summary

  • 整体说明:
    • Qwen3 包含一系列 LLM,包括 Dense 和 MoE 模型,参数规模从 0.6B 到 235B 不等(这个模型 Size 的丰富度,可以称为地表最强开源模型了)
    • Qwen3 的一项关键创新是将思考模式(thinking mode,用于复杂的多步推理)和非思考模式(non-thinking mode,用于快速的上下文驱动响应)集成到一个统一框架中
    • Qwen3 引入了思考预算(thinking budget)机制,允许用户在推理过程中自适应地分配计算资源,从而根据任务复杂度平衡延迟和性能
    • 该模型基于包含 36 万亿 token 的大规模数据集进行预训练,能够理解和生成 119 种语言和方言的文本
    • Qwen3 在预训练和训练后模型的标准基准测试中均表现出色,涵盖代码生成、数学、推理和 Agent 相关任务
  • 利用旗舰模型的知识(蒸馏小模型),显著减少了构建小规模模型所需的计算资源,同时确保其具备高度竞争力
  • Qwen3 在多样化基准测试中取得了 SOTA 结果,包括代码生成、数学推理、智能体任务等,与更大的 MoE 模型和专有模型相比具有竞争力
  • 与上一代 Qwen2.5 相比,Qwen3 将多语言支持从 29 种扩展到 119 种语言和方言

Introduction and Discussion

  • 人类的目标:人工通用智能(Artificial General Intelligence, AGI)或人工超级智能(Artificial Super Intelligence, ASI)
  • 这项工作介绍了基础模型家族 Qwen 的最新系列 Qwen3
    • Qwen3 是一组开源 LLM ,在广泛的任务和领域中实现了 SOTA 性能
    • 发布了 Dense 和 MoE 模型,参数数量从 0.6B 到 235B 不等,以满足不同下游应用的需求
    • 旗舰模型 Qwen3-235B-A22B 是一个 MoE 模型,总参数为 235B,每个 token 激活的参数为 22B
      • 这种设计确保了高性能和高效的推理
  • Qwen3 引入了多项关键技术进步以增强其功能和可用性
    • 首先,它将两种不同的操作模式(思考模式和非思考模式)集成到单一模型中
      • 这使得用户无需在不同模型之间切换(例如从 Qwen2.5 切换到 QwQ (2024))即可切换模式
      • 这种灵活性确保开发者和用户可以高效地根据特定任务调整模型行为
    • 此外,Qwen3 引入了思考预算,使用户能够精细控制模型在执行任务时应用的推理努力水平
      • 这一能力对于优化计算资源和性能至关重要,能够根据实际应用中的复杂度调整模型的思考行为
    • 此外,Qwen3 在涵盖 119 种语言和方言的 36 万亿 token 上进行了预训练,有效增强了其多语言能力
      • 这种广泛的语言支持扩大了其在全球用例和国际应用中的部署潜力
    • 这些进步共同使 Qwen3 成为 SOTA 开源大语言模型家族,能够有效解决跨领域和跨语言的复杂任务
  • Qwen3 的预训练过程使用了约 36 万亿 token 的大规模数据集,经过精心筛选以确保语言和领域的多样性
    • 为了高效扩展训练数据,作者采用了多模态方法:Qwen2.5-VL (2025) 被微调以从大量 PDF 文档中提取文本
    • 作者还使用领域特定模型生成合成数据:Qwen2.5-Math (2024) 用于数学内容,Qwen2.5-Coder (2024) 用于代码相关数据
  • 预训练过程采用三阶段策略:
    • 1)通用阶段(S1) :所有 Qwen3 模型在 4,096 token 的序列长度上训练超过 30 万亿 token
      • 在此阶段,模型在语言能力和通用世界知识上完成了全面预训练,训练数据涵盖 119 种语言和方言
    • 2)推理阶段(S2) :为了进一步提升推理能力,作者优化了本阶段的预训练语料库,增加了 STEM、编码、推理和合成数据的比例
      • 模型在 4,096 token 的序列长度上进一步训练约 5 万亿高质量 token
      • 在此阶段,作者加速了学习率衰减
    • 3)长上下文阶段(S3) :在最后的预训练阶段,作者收集高质量的长上下文语料以扩展 Qwen3 模型的上下文长度
      • 所有模型在 32,768 token 的序列长度上训练了数千亿 token
      • 长上下文语料库包含 75% 长度在 16,384 到 32,768 token 之间的文本,以及 25% 长度在 4,096 到 16,384 token 之间的文本
  • 为了更好地将基础模型与人类偏好和下游应用对齐,作者采用了多阶段后训练方法,同时赋能思考(推理)和非思考模式
    • 在前两个阶段,通过长 CoT 冷启动微调和专注于数学与编码任务的强化学习 ,开发了强大的推理能力
    • 在最后两个阶段,将带有和不带有推理路径的数据合并为一个统一数据集进行进一步微调 ,使模型能够有效处理两种类型的输入,随后应用通用领域的强化学习以提升在广泛下游任务中的性能
    • 对于小模型,使用强到弱蒸馏(strong-to-weak distillation),利用从大模型的 Off-policy 和 On-policy 知识迁移来增强其能力
      • 从高级教师模型的蒸馏在性能和训练效率上显著优于强化学习

Architecture

  • Qwen3 系列包括 6 个 Dense 模型(Qwen3-0.6B、Qwen3-1.7B、Qwen3-4B、Qwen3-8B、Qwen3-14B 和 Qwen3-32B)和 2 个 MoE 模型(Qwen3-30B-A3B 和 Qwen3-235B-A22B)
    • 旗舰模型 Qwen3-235B-A22B 总参数为 235B,激活参数为 22B
  • Qwen3 Dense 模型的架构与 Qwen2.5 (2024) 类似,包括使用如下组件:
    • 分组查询注意力(Grouped Query Attention, GQA)(2023)
    • SwiGLU (2017)
    • 旋转位置嵌入(Rotary Positional Embeddings, RoPE)(2024)
    • RMSNorm (2023) 的预归一化(pre-normalization)
  • Qwen3 移除了 Qwen2 (2024) 中使用的 QKV 偏置(QKV-bias),并在注意力机制中引入了 QK 归一化(QK-Norm)(2023) 以确保 Qwen3 的训练稳定性
  • 模型架构的关键信息如表 1 所示
  • Qwen3 MoE 模型与 Qwen3 Dense 模型共享相同的基础架构,模型架构的关键信息如表 2 所示
    • 遵循 Qwen2.5-MoE (2024) 并实现了细粒度专家分割(fine-grained expert segmentation)(2024)
    • Qwen3 MoE 模型共有 128 个专家,每个 token 激活 8 个专家
    • 与 Qwen2.5-MoE 不同,Qwen3-MoE 设计排除了共享专家
    • Qwen3 采用全局批量负载均衡损失(global-batch load balancing loss)(2025) 以鼓励专家专业化
    • 这些架构和训练创新在下游任务中显著提升了模型性能
  • Qwen3 模型使用 Qwen 的分词器 (2023),它实现了 BBPE(Byte-level Byte-Pair Encoding)(2020; 2020; 2016),词表为 151,669

Pre-training

Pre-training Data

  • 与 Qwen2.5(2024)相比,Qwen3 显著扩大了预训练数据的规模和多样性
    • 作者收集了 36 万亿 token 的数据,覆盖 119 种语言和方言
      • 这些数据包含多个领域的高质量内容,例如代码、STEM(科学、技术、工程和数学)、推理任务、书籍、多语言文本以及合成数据
  • 作者对预训练数据做了如下扩展:
    • First,使用 Qwen2.5-VL 模型(2025)对大量 PDF 类文档进行文本识别,随后通过 Qwen2.5 模型(2024)对识别文本进行优化,从而获得额外的高质量文本 token,总量达到数万亿
    • Second,利用 Qwen2.5(2024)、Qwen2.5-Math(2024)和 Qwen2.5-Coder(2024)模型合成了数万亿不同格式的文本 token,涵盖教材、问答、指令和代码片段等,覆盖数十个领域
    • Third,通过引入更多语言数据进一步扩展了预训练语料库。与 Qwen2.5 相比,支持的语言数量从 29 种大幅增加到 119 种,显著提升了模型的跨语言能力
  • 作者开发了一套多语言数据标注系统,用于提升训练数据的质量和多样性
    • 该系统已应用于大规模预训练数据集,对超过 30 万亿 token 进行了多维度标注,包括教育价值、领域和安全性等
    • 这些细粒度的标注支持更高效的数据过滤和组合
    • 与以往研究(2023;2023;2024)在数据源或领域级别优化数据混合比例不同,Qwen3 通过在小规模代理模型上进行细粒度数据标签的消融实验,实现了实例级别的数据混合优化

Pre-training Stages

  • Qwen3 模型的预训练分为三个阶段:
    • 1)General Stage(S1) :在第一阶段,所有 Qwen3 模型在超过 30 万亿 token 上以 4,096 token 的序列长度进行训练
      • 此阶段模型已完全掌握语言能力和通用世界知识,训练数据覆盖 119 种语言和方言
    • 2)Reasoning Stage(S2) :为了进一步提升推理能力,作者优化了本阶段的预训练语料库,增加了 STEM、代码、推理和合成数据的比例
      • 模型以 4,096 token 的序列长度进一步预训练约 5 万亿高质量 token,并在此阶段加速学习率衰减
    • 3)Long Context Stage(S3) :在最后的预训练阶段,作者收集高质量的长上下文语料库,将 Qwen3 模型的上下文长度扩展到 32,768 token
      • 长上下文语料库中,75% 的文本长度介于 16,384 到 32,768 token 之间,25% 介于 4,096 到 16,384 token 之间
      • 遵循 Qwen2.5(2024)的方法,作者使用 ABF 技术(2023)将 RoPE 的基础频率从 10,000 增加到 1,000,000,同时引入 YARN(2023)和双块注意力(Dual Chunk Attention, DCA)(2024),在推理时实现序列长度容量的四倍提升
  • 与 Qwen2.5(2024)类似,Qwen3 基于上述三个阶段开发了超参数(如学习率 Scheduler 和 Batch Size)的扩展规律
    • 通过大量实验,作者系统研究了模型架构、训练数据、训练阶段与最优训练超参数之间的关系,最终为每个 Dense 或 MoE 模型预测了最优学习率和 Batch Size 策略

Pre-training Evaluation

  • Qwen3 系列的基础语言模型评估重点关注其在通用知识、推理、数学、科学知识、代码和多语言能力方面的表现。评估数据集包括 15 个基准测试:
    • 通用任务(General Tasks) :MMLU(2021a)(5-shot)、MMLU-Pro(2024)(5-shot,CoT)、MMLU-redux(2024)(5-shot)、BBH(2023)(3-shot,CoT)、SuperGPQA(2025)(5-shot,CoT)
    • 数学与 STEM 任务(Math & STEM Tasks) :GPQA(2023)(5-shot,CoT)、GSM8K(2021)(4-shot,CoT)、MATH(2021b)(4-shot,CoT)
    • 代码任务(Coding Tasks) :EvalPlus(2023a)(0-shot)(HumanEval(2021)、MBPP(2021)、Humaneval+、MBPP+ 的平均值)、MultiPL-E(2023)(0-shot)(Python、C++、Java、PHP、TypeScript、C#、Bash、JavaScript)、MBPP-3shot(2021)、CRUX-O of CRUXEval(2024)(1-shot)
    • 多语言任务(Multilingual Tasks) :MGSM(2023)(8-shot,CoT)、MMMLU(OpenAI,2024)(5-shot)、INCLUDE(2024)(5-shot)
  • 在基础模型的基线对比中,将 Qwen3 系列基础模型与 Qwen2.5 基础模型(2024)以及其他领先的开源基础模型进行了比较,包括 DeepSeek-V3 Base(2024a)、Gemma-3(2025)、Llama-3(2024)和 Llama-4(Meta-AI,2025)系列基础模型
    • 所有模型均使用相同的评估流程和广泛采用的评估设置,以确保公平比较
Summary of Evaluation Results
  • 基于整体评估结果,论文中总结了 Qwen3 基础模型的关键结论:
    • (1) 与之前开源的 SOTA Dense 和 MoE 基础模型(如 DeepSeek-V3 Base、Llama-4-Maverick Base 和 Qwen2.5-72B-Base)相比,Qwen3-235B-A22B-Base 在大多数任务中表现更优,且参数量显著更少
    • (2) 对于 Qwen3 MoE 基础模型,实验结果表明:
      • (a) 使用相同的预训练数据,Qwen3 MoE 基础模型仅需 1/5 的激活参数即可达到与 Dense 模型相当的性能;
      • (b) 由于架构改进、训练 token 规模扩大和更先进的训练策略,Qwen3 MoE 基础模型的性能优于 Qwen2.5 MoE 基础模型,且激活参数和总参数量更少;
      • (c) 即使仅使用 Qwen2.5 Dense 基础模型 1/10 的激活参数,Qwen3 MoE 基础模型也能达到相当的性能,显著降低了推理和训练成本
    • (3) Qwen3 Dense 基础模型的整体性能与更高参数规模的 Qwen2.5 基础模型相当
      • 例如,Qwen3-1.7B/4B/8B/14B/32B-Base 分别与 Qwen2.5-3B/7B/14B/32B/72B-Base 性能相当,尤其在 STEM、代码和推理基准测试中,Qwen3 Dense 基础模型的性能甚至超越了更高参数规模的 Qwen2.5 基础模型
  • 注:【此处省略一些评估细节】

Post-training

  • 图 1 展示了 Qwen3 系列模型的后训练流程,该流程围绕两个核心目标设计:
    • (1) 思维控制(Thinking Control) :通过整合“非思维模式(non-thinking mode)”和“思维模式(thinking mode)”,使用户能够灵活选择模型是否进行推理,并通过指定思维过程的 token 预算来控制思考深度
    • (2) 强到弱蒸馏(Strong-to-Weak Distillation) :旨在简化和优化轻量级模型的后训练过程
      • 通过利用大规模模型的知识,显著降低构建小规模模型所需的计算成本和开发工作量

Long-CoT Cold Start

  • 首先构建了一个涵盖数学、代码、逻辑推理和 STEM 问题的综合数据集,每个问题均配有已验证的参考答案或基于代码的测试用例
  • 该数据集用于长链思维(long Chain-of-Thought, long-CoT)训练的冷启动阶段
  • 数据集构建包含两阶段过滤流程:查询过滤(query filtering)和响应过滤(query filtering)
  • 查询过滤阶段:使用 Qwen2.5-72B-Instruct 识别并移除难以验证的查询(例如包含多个子问题或要求生成通用文本的查询),同时排除 Qwen2.5-72B-Instruct 无需 CoT 推理即可正确回答的查询,以确保仅包含需要深度推理的复杂问题
    • 使用 Qwen2.5-72B-Instruct 标注每个查询的领域,以保持数据集的领域平衡
  • 响应过滤阶段:保留验证集后,使用 QwQ-32B 为每个剩余查询生成 \(N\) 个候选响应
    • 若 QwQ-32B 无法生成正确解决方案,则由人工标注者评估响应准确性
    • 对于通过 Pass@\(N\) 的查询,进一步应用严格过滤标准,移除以下类型的响应:
      • (1) 最终答案错误;
      • (2) 包含大量重复内容;
      • (3) 明显猜测而无充分推理;
      • (4) 思维内容与总结不一致;
      • (5) 语言混合不当或风格突变;
      • (6) 疑似与验证集内容过度相似
  • 随后,使用精炼后的数据集子集进行推理模式的初始冷启动训练
    • 此阶段的目标是为模型奠定基础推理模式 ,而非过度强调即时推理性能,从而确保模型潜力不受限制,为后续RL阶段提供更大改进空间
    • 为实现这一目标,建议在此准备阶段尽量减少训练样本数量和训练步数

Reasoning RL

  • 推理强化学习阶段使用的 query-verifier 对需满足以下四个标准:
    • (1) 未在冷启动阶段使用;
    • (2) 冷启动模型可学习;
    • (3) 尽可能具有挑战性;
    • (4) 覆盖广泛的子领域
  • 最终收集了 3995 对 query-verifier ,并采用 GRPO(2024)更新模型参数
    • 实验表明,使用大的 Batch Size 和(每个查询)高 rollout 次数,结合 Off-policy 训练以提高样本效率,对训练过程有益
    • 此外,通过控制模型的熵(entropy)使其稳步增加或保持稳定 ,可以有效平衡探索与利用(exploration and exploitation),这对维持训练稳定性至关重要
      • 注:这里作者说模型的熵是稳步增加的,但大部分论文中给出的训练都是看到熵在下降的,作者应该是否做了一些特定的设计来提升模型在这方面的能力?
    • 结果表明,在单次 RL 训练过程中,无需手动干预超参数,即可实现训练奖励和验证性能的持续提升
      • 例如,Qwen3-235B-A22B 模型的 AIME’24 分数在 170 个 RL 训练步中从 70.1 提升至 85.1

Thinking Mode Fusion

  • 思维模式融合阶段的目标是将“非思维(non-thinking)”能力整合到已具备“思维(thinking)”能力的模型中
  • 这种方法使开发者能够管理和控制推理行为,同时降低部署独立模型以处理思维和非思维任务的成本和复杂性
  • 为实现这一目标,作者对推理强化学习模型进行持续 SFT ,并设计聊天模板以融合两种模式
  • 作者发现能够熟练处理两种模式的模型在不同思维预算下均表现稳定
Construction of SFT data
  • SFT 数据集结合了“thinking”和“non-thinking”数据
  • 为确保 Stage 2 模型的性能不受额外 SFT 影响 :
    • “thinking”数据通过使用 Stage 2 模型对 Stage 1 查询进行拒绝采样(rejection sampling)生成;
    • “non-thinking”数据则精心策划,涵盖编码、数学、指令遵循、多语言任务、创意写作、问答和角色扮演等多样化任务
    • 理解:因为 Stage 2 模型已经是训练好的了,为了增加思维链融合功能,额外的 SFT 是可能会导致模型出现问题的,这里使用 Stage 2 和 Stage 1 模型相关的数据来训练,从而保证 Stage 2 和 Stage 1 模型的原始能力?
  • 此外,作者采用自动生成的检查表评估“non-thinking”数据的响应质量,并特别增加低资源(low-resource)任务上的性能,作者特地增加了翻译任务的比例
Chat Template Design
  • 为更好整合两种模式并支持用户动态切换思维过程,作者设计了 Qwen3 的聊天模板(如表 9 所示)
    • 对于思维模式和非思维模式的样本,作者分别在用户查询或系统消息中引入 /think 和 /no_think 标志,使模型能够根据用户输入选择相应的思维模式
    • 对于非思维模式样本 ,作者在助手响应中保留空的思维块(thinking block)
      • 这一设计确保模型内部格式一致性,并允许开发者通过在聊天模板中拼接空的思维块来阻止模型进行思维行为
    • 默认情况下,模型以思维模式运行,因此作者添加了一些用户查询中不包含 /think 标志的思维模式训练样本
    • 对于更复杂的多轮对话,作者在用户查询中随机插入多个 /think 和 /no_think 标志,模型响应则遵循最后遇到的标志
Thinking Budget
  • 思维模式融合的额外优势在于,一旦模型学会以非思维和思维模式响应,它自然能够处理中间情况(即基于不完整思维生成响应)
  • 这一能力为实施模型思维过程的预算控制奠定了基础
    • 具体而言,当模型的思维长度达到用户定义的阈值时,手动停止思维过程,并插入停止思维指令:“Considering the limited time by the user, I have to give the solution based on the thinking directly now.\n</think>.\n\n”
    • 插入该指令后,模型基于已积累的推理生成最终响应
    • 值得注意的是,此能力并非通过显式训练获得,而是思维模式融合的自然结果

General RL

  • 通用强化学习阶段旨在广泛增强模型在多样化场景中的能力和稳定性
  • 为此,作者建立了一个覆盖 20 多项任务 的复杂奖励系统 ,每项任务均配备定制化评分标准,重点关注以下核心能力:
    • Instruction Following :确保模型准确解释并遵循用户指令,包括内容、格式、长度和结构化输出等要求,生成符合用户期望的响应
    • 格式遵循(Format Following) :除显式指令外,模型需遵循特定格式约定
      • 例如,通过 /think 和 /no_think 标志切换思维模式,并在最终输出中使用指定 token(如 <think> 和 </think>)分隔思维和响应部分
    • 偏好对齐(Preference Alignment) :针对开放式查询,偏好对齐侧重于提升模型的帮助性、参与度和风格,最终提供更自然和令人满意的用户体验
    • 智能体能力(Agent Ability) :训练模型通过指定接口正确调用工具
      • 在 RL rollout 过程中,模型可执行完整的多轮交互周期 ,并接收真实环境执行反馈 ,从而提升其在长视野决策任务中的性能和稳定性
    • 专项场景能力(Abilities for Specialized Scenarios) :在更专业的场景中,作者设计特定任务
      • 例如,在检索增强生成(Retrieval-Augmented Generation, RAG)任务中,引入奖励信号引导模型生成准确且上下文合适的响应 ,从而最小化幻觉风险
  • 为上述任务提供反馈时,作者采用三种不同类型的奖励:
    • (1)基于规则的奖励(Rule-based Reward) :广泛用于推理强化学习阶段,也适用于指令遵循(2024)和格式遵循等通用任务
      • 设计良好的基于规则的奖励可高精度评估模型输出的正确性,避免奖励破解(reward hacking)问题
    • (2)带参考答案的模型奖励(Model-based Reward with Reference Answer) :为每个查询提供参考答案,并提示 Qwen2.5-72B-Instruct 基于参考答案对模型响应评分
      • 该方法无需严格格式化即可灵活处理多样化任务,避免纯规则奖励可能导致的假阴性
    • (3)无参考答案的模型奖励(Model-based Reward without Reference Answer) :利用人类偏好数据训练奖励模型 ,为模型响应分配标量分数
      • 该方法不依赖参考答案,可处理更广泛查询,同时有效提升模型的参与度和帮助性

Strong-to-Weak Distillation

  • 强到弱蒸馏流程专为优化轻量级模型设计,涵盖 5 个 Dense 模型(Qwen3-0.6B、1.7B、4B、8B 和 14B)和 1 个 MoE 模型(Qwen3-30B-A3B)
    • 该方法在提升模型性能的同时,有效赋予其强大的模式切换能力
  • 蒸馏过程分为两个主要阶段:
    • (1) Off-policy 蒸馏(Off-policy Distillation) :在此初始阶段,作者结合教师模型在 /think 和 /no_think 模式下的输出进行响应蒸馏 ,帮助轻量级学生模型发展基础推理能力和思维模式切换能力,为后续 On-policy 训练阶段奠定基础
    • (2) On-policy 蒸馏(On-policy Distillation) :在此阶段,学生模型生成 On-policy 序列进行微调
      • 使用学生模型采样提示后,学生模型以 /think 或 /no_think 模式生成响应,并通过对齐其 logits 与教师模型(Qwen3-32B 或 Qwen3-235B-A22B)以最小化 KL 散度(Kullback-Leibler divergence)进行微调

Post-training Evaluation

  • For 全面评估指令调优模型的质量,采用自动化基准测试思维和非思维模式下的模型性能。这些基准分为以下几类:
    • 通用任务(General Tasks) :使用 MMLU-Redux(2024)、GPQA-Diamond(2023)、C-Eval(2023)和 LiveBench(2024-11-25)(2024)等基准
      • 对于 GPQA-Diamond,每个查询采样 10 次并报告平均准确率
    • 对齐任务(Alignment Tasks) :评估模型与人类偏好的对齐程度
      • 针对指令遵循性能,报告 IFEval(2023)的严格提示准确率;
      • 针对通用主题的人类偏好对齐,使用 Arena-Hard(2024)和 AlignBench v1.1(2023);
      • 针对写作任务,依赖 Creative Writing V3(2024)和 WritingBench(2025)评估模型的熟练度和创造力
    • 数学与文本推理(Math & Text Reasoning) :评估数学和逻辑推理能力,采用高阶数学基准 MATH-500(2023)、AIME’24 和 AIME’25(2025),以及文本推理任务 ZebraLogic(2025)和 AutoLogi(2025)
      • 对于 AIME 问题,每年试题包含 Part I 和 Part II,共 30 题,每题采样 64 次并以平均准确率作为最终分数
    • 智能体与编码(Agent & Coding) :测试模型在编码和智能体任务中的熟练度,使用 BFCL v3(2024)、LiveCodeBench(v5, 2024.10-2025.02)(2024)和 Codeforces Ratings from CodeElo(2025)
      • 对于 BFCL,所有 Qwen3 模型均采用 FC 格式评估,并使用 yarn 将模型部署至 64k 上下文长度以进行多轮评估
        • 部分基线来自 BFCL 排行榜,取 FC 和 Prompt 格式中的较高分;未在排行榜中报告的模型则评估 Prompt 格式
      • 对于 LiveCodeBench,非思维模式使用官方推荐提示,思维模式则调整提示模板以允许模型更自由地思考(移除限制 You will not return anything except for the program)
      • 为评估模型与竞技编程专家的性能差距,使用 CodeForces 计算 Elo 评分,每个问题最多生成八次独立推理尝试
    • 多语言任务(Multilingual Tasks) :评估四种任务:指令遵循、知识、数学和逻辑推理
      • 指令遵循使用 Multi-IF(2024)(覆盖 8 种关键语言);
      • 知识评估包含两类:区域知识通过 INCLUDE(2024)(覆盖 44 种语言)评估,通用知识通过 MMMLU(2024)(覆盖 14 种语言,排除未优化的约鲁巴语);
      • 数学任务采用 MT-AIME2024(2025)(涵盖 55 种语言)和 PolyMath(2025)(涵盖 18 种语言);
      • 逻辑推理使用 MlogiQA(覆盖 10 种语言,源自 2024)
  • 所有 Qwen3 模型采样超参数设置如下:
    • 在思维模式下的采样超参数设置为 temperature\(=0.6\)、top-p\(=0.95\)、top-k\(=20\);
      • 对于 Creative Writing v3 和 WritingBench,应用 presence penalty\(=1.5\) 以鼓励生成更多样化内容
      • 补充:presence penalty 会对已经生成过的词汇的概率值进行削弱
        • 值越大,削弱越厉害,生成的文本更加多样化;
        • 值越小,更容易生成重复的词,但文本会更加连贯
    • 非思维模式的采样超参数为 temperature\(=0.7\)、top-p \(=0.8\)、top-k \(=20\)、presence penalty\(=1.5\)
    • 两种模式的最大输出长度均设为 32,768 token,AIME’24 和 AIME’25 除外(扩展至 38,912 token 以提供充足思维空间)
      Summary of Evaluation Results
  • 从评估结果中,我们总结出最终确定的Qwen3模型的几个关键结论如下:
    • (1)我们的旗舰模型 Qwen3-235B-A22B 在思考模式和非思考模式下均展现出开源模型中的顶尖整体性能,超越了 DeepSeek-R1 和 DeepSeek-V3 等强基线模型
      • Qwen3-235B-A22B 与闭源领先模型(如 OpenAI-o1、Gemini2.5-Pro 和 GPT-4o)相比也具有高度竞争力,彰显了其深厚的推理能力和全面的通用能力
    • (2)我们的旗舰 Dense 模型 Qwen3-32B 在大多数基准测试中优于我们此前最强的推理模型 QwQ-32B,且性能与闭源的OpenAI-o3-mini 相当,表明其推理能力令人瞩目
      • Qwen3-32B 在非思考模式下的表现也极为出色,超越了我们此前的旗舰非推理 Dense 模型 Qwen2.5-72B-Instruct。
    • (3)我们的轻量级模型(包括 Qwen3-30B-A3B、Qwen3-14B 及其他较小的 Dense 模型)与参数规模相近或更大的开源模型相比,性能始终更优,证明了我们“Strong-to-Weak Distillation”方法的成功。
  • 注:【此处省略一些评估细节】

Discussion

The Effectiveness of Thinking Budget

  • 为验证 Qwen3 可通过增加思维预算提升智能水平,在数学、编码和 STEM 领域的四个基准上调整分配的思维预算
  • 如图 2 所示,Qwen3 展现出与分配思维预算相关的可扩展且平滑的性能提升
    • Thinking Budget 越大,效果越好
  • 若进一步将输出长度扩展至 32K 以上,模型性能有望在未来继续提升,作者将此探索留作未来工作

The Effectiveness and Efficiency of On-Policy Distillation

  • 通过比较蒸馏与直接强化学习在相同 Off-policy 蒸馏 8B 检查点后的性能和计算成本(以 GPU 小时计),评估 On-policy 蒸馏的有效性和效率
  • 注:为简化,以下仅关注数学和代码相关查询
  • 表 21 的结果表明,蒸馏在仅需约 1/10 GPU 小时的情况下,性能显著优于强化学习
    • 从教师 logits 蒸馏使学生模型能够扩展其探索空间并增强推理潜力,表现为蒸馏后 AIME’24 和 AIME’25 基准的 pass@64 分数较初始检查点有所提升;
    • 强化学习未带来 pass@64 分数的任何改进。
    • 这些观察凸显了利用更强教师模型指导学生模型学习的优势

The Effects of Thinking Mode Fusion and General RL

  • For 评估后训练中思维模式融合和通用强化学习的有效性,对 Qwen-32B 模型的各个阶段进行评估
  • 除前述数据集外,引入多个内部基准以监控其他能力:
    • CounterFactQA :包含反事实问题,模型需识别问题非事实并避免生成幻觉答案
    • LengthCtrl :包含带长度要求的创意写作任务,最终分数基于生成内容长度与目标长度的差异
    • ThinkFollow :包含随机插入 /think 和 /no_think 标志的多轮对话,测试模型能否根据用户查询正确切换思维模式
    • ToolUse :评估模型在单轮、多轮和多步工具调用过程中的稳定性,分数包括工具调用过程中的意图识别准确率、格式准确率和参数准确率
  • 结果如表 22 所示,可得出以下结论:
    • (1) Stage 3 将非思维模式整合至已通过前两阶段训练具备思维能力的模型中
      • ThinkFollow 基准分数 88.7 表明模型已具备初步模式切换能力,但仍偶有错误
      • Stage 3 还提升了模型在思维模式下的通用和指令遵循能力,CounterFactQA 提升 10.9 分,LengthCtrl 提升 8.0 分
    • (2) Stage 4 进一步强化模型在思维和非思维模式下的通用、指令遵循和智能体能力
      • ThinkFollow 分数提升至 98.9,确保准确模式切换
    • (3) 对于知识、STEM、数学和编码任务 ,思维模式融合和通用强化学习未带来显著提升
      • 相反,对于 AIME’24 和 LiveCodeBench 等挑战性任务,思维模式性能在这两个训练阶段后实际下降
      • 作者推测这种退化是由于模型在更广泛通用任务上训练,可能影响其处理复杂问题的专项能力
      • 在 Qwen3 开发中,作者选择接受这种性能权衡以增强模型的整体通用性

Future Work

  • 在不久的将来,作者的研究将聚焦于以下几个关键方向:
    • (1) 预训练扩展(Scale up pretraining) :作者将继续扩展预训练规模,使用质量更高、内容更多样的数据
    • (2) 架构优化(Improving model architecture) :改进模型架构和训练方法,以实现高效压缩、超长上下文支持等目标
    • (3) RL :增加计算资源投入,重点关注基于环境反馈的智能体强化学习系统,以构建能够处理需要推理时间扩展的复杂任务的智能体

Appendix

A.1 Additional Evaluation Results

A.1.1 Long-Context Ability
  • For 评估长上下文处理能力,在 RULER 基准测试 (2024) 中报告了结果(表 23)
  • 为实现长度外推(length extrapolation),使用 YARN (2023) 并设置缩放因子 \( \text{scaling_factor} = 4 \)
  • 在思考模式下,作者将思考预算(thinking budget)设为 8192 token,以减少对超长输入的冗余推理
  • 结果显示:
    • (1) 在非思考模式下,Qwen3 在长上下文处理任务中优于同规模的 Qwen2.5 模型
    • (2) 在思考模式下,模型性能略有下降
      • 作者推测思考内容对这些检索任务(无需依赖推理)帮助有限 ,甚至可能干扰检索过程
      • 未来版本将重点提升思考模式下的长上下文能力

A.1.2 Multilingual Ability

  • 表 24-35 展示了 Qwen3 系列模型在西班牙语、法语、葡萄牙语、意大利语、阿拉伯语、日语、韩语、印尼语、俄语、越南语、德语和泰语等多种语言中的详细基准得分
    • 这些结果表明,Qwen3 系列模型在所有评估基准中均表现优异,展现了强大的多语言能力
  • 为更广泛评估 Qwen3 的多语言表现,作者使用 Belebele (2023) 基准测试,覆盖 80 种优化语言(表 36 按语系列出)
  • 表 37 展示了 Qwen3 与其他基线模型在 Belebele 基准上的性能对比
  • 结果显示,Qwen3 在同等规模下与 Gemma 模型表现相当,同时显著优于 Qwen2.5

附录:Qwen3-MoE 源码

  • 参考博客:图解 Qwen3 MoE 模型源码
    • 包含一些流程图和源码,比较清晰

NLP——LLM对齐微调-VC-PPO

注:本文包含 AI 辅助创作

  • 参考链接:
    • (VC-PPO)What’s Behind PPO’s Collapse in Long-CoT? Value Optimization Holds the Secret, arXiv 20250303, ByteDance Seed

整体总结

  • 本文提出一种改进的 PPO 方法, VC-PPO(Value-Calibrated PPO)
  • 核心贡献:
    • Pretrained value:开始 RL 前先预训练价值网络
    • Decoupled-GAE:计算 Advantage (for Actor 损失)时和 计算 Target Reward(for Critic 损失)时使用不同的 \(\lambda\)

Motivation

  • LLM 在复杂推理任务(如数学、编程)中表现出色,尤其是通过生成长链思维(Long-CoT)来逐步推导答案
  • OpenAI的o1、DeepSeek-R1等模型都采用了这种“推理时扩展”策略
  • 一个发现:在Long-CoT任务中,PPO经常失效 ,表现为:
    • 模型输出长度迅速下降(注:通常是训练初期即开始大幅下降)
    • 验证集性能大幅退化
    • 无法有效利用长链思维
  • 注:这些现象与 PPO 在传统 RL 任务中的成功形成鲜明对比

问题诊断:PPO 失效的两大原因简单描述

  • 作者通过实验和分析,识别出 PPO 在 Long-CoT 任务中失效的两大根本原因
  • 注:导致模型输出长度迅速下降的直接原因是:前置 Advantage 被高估,详情见本文后面补充的附录

Value/Critic Initialization Bias

  • 在 RLHF 中,常用做法是用训练好的 奖励模型来初始化Value/Critic 模型
    • 注:这种做法则源于奖励模型和价值模型之间的表面相似性,因为两个模型都旨在预测关于响应的标量信息
  • 但奖励模型只在 <EOS> 处提供评分,对前面的 token 没有监督信号,导致它对前置 token 的打分偏低
  • 这种偏差在 GAE 中会被放大,导致前置 token 的 Advantage 被高估,进而促使模型倾向于生成短回答

奖励信号衰减(Reward Signal Decay)

  • 在 GAE 中,当 \(\lambda < 1\) 时,来自 <EOS> 的奖励信号会随传播距离指数衰减
  • 在 Long-CoT 任务中,序列长度可达数千 token,前置 token 几乎接收不到任何奖励信号
  • 这导致值模型无法有效学习,进而影响策略优化
  • 注:传统 RLHF 中常常使用 \(\lambda = 0.95\)
    • 这种做法源于传统的 RL 文献,其中 PPO 已在像 Mujoco 和 Atari 这样的环境中得到广泛测试
    • 在这些环境中,奖励会在轨迹上累积,导致高方差的回报,因此,使用 \(\lambda < 1\) 方差降低是必要的
    • 但这种方式会导致模型收敛缓慢
奖励信号衰减的数学推导
  • 定义:
    • \(V\):值函数
    • \(r_t\):即时奖励
    • \(\lambda\):GAE 平滑因子
  • 在标准 GAE 中,优势估计为:
    $$
    \hat{A}_t = \sum_{l=0}^{T-t-1} \lambda^l \delta_{t+l}, \quad \delta_t = r_t + V(s_{t+1}) - V(s_t)
    $$
  • Critic 目标为回报的估值:
    $$ R \approx V^{\text{target} }(s_t) = V_{old}(s_t) + A^{GAE}_t$$
  • 展开即可得到:
    $$
    V^{\text{target} }(s_t) =
    \begin{cases}
    \sum_{l=0}^{T-t-1} \lambda^l (r_{t+l} + V(s_{t+l+1}) - V(s_{t+l})) + V(s_t), & \lambda < 1 \\
    \sum_{l=0}^{T-t-1} r_{t+l}, & \lambda = 1
    \end{cases}
    $$
    • 上面式子的详细推导见:RL——强化学习中的方差与偏差 的 GAE 部分
    • 当 \(\lambda=1\) 时, Critic 目标即为累积奖励,无偏且稳定
    • 当 \(\lambda<1\) 时, Critic 目标引入了值函数自身的估计,可能不稳定
      • 当长度太大时,可以看到最终的奖励 \(r_T\) 几乎被淹没了,需要逐步透传到 Critic,所以 Critic 收敛性本身也会变慢

模型输出长度迅速下降 的原因详细分析

  • 现象描述:在标准的 PPO 训练中,模型本应生成长链思维(Long-CoT)来逐步推理答案
    • 但在实验中,训练刚开始不久,模型输出的平均长度急剧下降 ,随之而来的是验证集性能的崩溃(如图 1 所示)
  • 根据论文内容,可以得到一个清晰的因果链:
    • 1)初始阶段 :Critic 模型对前置 token 的估计偏低(来自奖励模型初始化)
    • 2)GAE 计算 :这种偏差被 GAE 放大,使前置 token 获得过高的优势值
    • 3)策略更新 :PPO 鼓励高优势值的动作,即鼓励模型多输出前置 token,少输出后置 token
    • 4)行为变化 :模型学会尽早结束生成,因为“早结束”意味着“多输出前置 token、少输出后置 token”
    • 5)恶性循环 :输出变短后,训练数据中的长序列减少,值模型更难学习长序列的价值,进一步强化短输出倾向

直接原因:前置 token 的优势被高估

  • 通过分析优势值(Advantage)与token位置的关联(如图 2 所示)发现:
  • 越靠前的 token,其优势值越高(正偏差越大);
    • 这种偏差导致模型倾向于更早结束生成 ,因为前置 token 被“鼓励”输出,而后置 token 被“惩罚”或忽略
    • 模型学会“尽早收尾”,从而输出长度急剧下降

根本原因之一:值初始化偏差

奖励模型的训练目标
  • 奖励模型只在 token 处给出评分(如正确=1,错误=-1);
  • 它对前面的 token 没有直接的监督信号;
  • 因此,奖励模型对前置 token 的评分偏低(因为信息不完整)
    • 理解:这里不够严谨
      • 其实需要看初始化值和大部分目标值的相对关系,本质应该是可能偏高也可能偏低才对
值模型从奖励模型初始化
  • 在 RLHF 中,常用做法是将训练好的奖励模型作为值模型的初始化;
  • 这导致值模型在初始阶段也对前置token的预期回报估计偏低
偏差在 GAE 中被放大
  • GAE 的优势估计公式为:
    $$
    \hat{A}_t = \sum_{l=0}^{T-t-1} \lambda^l \delta_{t+l}, \quad \delta_t = r_t + V(s_{t+1}) - V(s_t)
    $$
  • 由于前置token的 \(V(s_t)\) 被低估,而 \(V(s_{t+1})\) 相对较高(因为更接近 <EOS>),导致:
    • \(\delta_t\) 为正;
    • 这些正偏差在累加过程中被放大,最终使前置token的优势值显著偏高

根本原因之二:奖励信号衰减

GAE 中的奖励传播机制
  • 当 \(\lambda < 1\) 时(如默认的 0.95),来自 <EOS> 的奖励信号会随传播距离指数衰减:
    $$
    \text{传播到第 } t \text{ 个 token 的奖励信号} = \lambda^{T-t} \cdot r_{}
    $$
Long-CoT 任务的特点
  • 序列长度 \(T\) 可能达到数千token;
  • 前置 token 的 \(T-t\) 很大,\(\lambda^{T-t} \approx 0\);
  • 前置 token 几乎接收不到任何来自最终答案的奖励信号
对 Critic 模型的影响
  • 值模型难以学习到前置 token 的真实价值;
  • 值估计进一步失真,加剧了前置 token 的优势偏差

解决方案:VC-PPO

  • VC-PPO(Value-Calibrated PPO) 同时解决上述两个问题

创新1:Value-Pretraining

  • 目的:解决值初始化偏差,使值模型在训练开始前就与初始策略对齐
  • 方法流程
    • Step 1: 固定策略模型 ,使用初始策略(如SFT后的模型)生成大量 Response
    • Step 2: 使用 Monte-Carlo 回报(即 GAE \(\lambda=1.0\)) 作为值模型的目标,进行离线训练
    • Step 3: 训练至值损失和解释方差(explained variance)收敛
    • Step 4: 将该值模型作为后续 PPO 训练的初始值模型
  • 实验效果
    • 消除了前置 token 的优势偏差
    • 保留了 Long-CoT 的模式,避免输出长度崩溃

创新2:Decoupled-GAE(解耦 GAE)

  • 这是本文的核心创新点,目的是在策略优化和值函数优化中使用不同的 \(\lambda\) 值,以分别满足两者的不同需求
  • 注:使用不同的 \(\lambda\) 值更新时,需要证明其策略梯度还能准确,下文会证明这个事情
背景问题描述
  • 策略优化需要低方差 的梯度估计,因此希望使用较小的 \(\lambda\)(如 0.95);
  • 但值函数优化需要无偏的目标 ,因此希望使用 \(\lambda=1.0\),避免奖励信号衰减;
  • 传统 PPO 中,两者共用同一个 \(\lambda\),无法兼顾
方法流程
  • Step 1: 策略优化使用 \(\lambda_{\text{actor} } < 1.0\)(如 0.95),以降低梯度方差
  • Step 2: 值函数优化使用 \(\lambda_{\text{critic} } = 1.0\),以确保 Critic 目标无偏
  • Step 3: 两者的 GAE 计算独立进行 ,互不干扰
  • Algorithm1:
策略梯度的无偏性证明(待详细推导和理解)
  • 作者进一步证明,即使值函数是用不同 \(\lambda\) 训练的,将其代入策略梯度中仍然无偏:
  • 定义 n-step 回报为:
    $$
    G_{t:t+h} =
    \begin{cases}
    \sum_{l=0}^{h-1} r_{t+l} + \bar{V}(s_{t+h}), & t+h < T \\
    \sum_{l=0}^{T-t-1} r_{t+l}, & t+h = T
    \end{cases}
    $$
  • 则优势可写为:
    $$
    \hat{A}_t = (1-\lambda) \sum_{l=1}^{T-t-1} \lambda^{l-1} G_{t:t+l} + \lambda^{T-t-1} G_{t:T} - \bar{V}(s_t)
    $$
  • 即,具有任意 \(\lambda\) 的策略梯度可以重写如下:
    $$\begin{aligned} \mathbb{E}_t [\nabla_\theta \log \pi_\theta(a_t|s_t) A_t] &= \mathbb{E}_t \left[ \nabla_\theta \log \pi_\theta(a_t|s_t) \sum_{l=0}^{T-t-1} \lambda^l (r_{t+l} + \bar{V}(s_{t+l+1}) - \bar{V}(s_{t+l})) \right] \\ &= \mathbb{E}_t \left[ \nabla_\theta \log \pi_\theta(a_t|s_t) \left( (1-\lambda) \sum_{l=1}^{T-t-1} \lambda^{l-1} G_{t:t+l} + \lambda^{T-t-1} G_{t:T} - \bar{V}(s_t) \right) \right] \\ &= \mathbb{E}_t \left[ \nabla_\theta \log \pi_\theta(a_t|s_t) \left( (1-\lambda) \sum_{l=1}^{T-t-1} \lambda^{l-1} G_{t:t+l} + \lambda^{T-t-1} G_{t:T} \right) \right] \end{aligned} \tag{8}$$
    • 根据公式 8,代入任意价值函数不会给策略梯度引入额外的偏差
    • 鉴于大型语言模型所需的大量时间和计算资源,使用较小的 \(\lambda\) 来加快策略的收敛是可取的
    • 一个候选配置可以是 \(\lambda_{\text{policy}} = 0.95\) 和 \(\lambda_{\text{value}} = 1.0\)
  • 论文中提到:使用任意值函数不会引入额外偏差
    • 问题:这里并不是说明使用 Decoupled-GAE 前后的梯度一致,因为 \(G_{t:t+l}\) 中也包含了 \(\bar{V}(s_{t+l})\)
      • 此时 Actor 更新时最大化的目标(Advantage)本身已经发生了改变,目标变成了 最大化新的 \(\lambda\) 下的 Critic 值
    • 这里仅仅是证明了这 Actor 和 Critic 解耦的 \(\lambda\) 下,也有一个可以学习的目标值(形式上与原始的 策略梯度法/PPO 的结果是一样的)
      • 注:这也就是说形式上本身是不冲突的,但没有说明两者的目标是完全等价的

实验

Setting

  • 主要任务 :AIME、GPQA、Codeforces
  • Base 模型型 :Qwen2.5 32B
  • 冷启动 :使用少量 Long-CoT 格式样本进行 SFT
  • 奖励 :规则驱动的答案匹配(正确=1,错误=-1)
  • Baselines :PPO(\(\lambda=0.95\))和 GRPO(DeepSeek-R1 中使用)

对照 GRPO

  • VC-PPO vs GRPO
    模型 AIME 2024 pass@1 GPQA pass@1 CodeForces pass@1
    GRPO 38.9 49.4 12.6
    VC-PPO 48.8 48.8 12.8
  • VC-PPO 在 AIME 上显著优于 GRPO,达到 Qwen-32B 模型在该任务上的 SOTA

消融实验

  • 移除 Value-Pretraining :AIME pass@1 从 41.9 降至29.4
  • 移除 Decoupled-GAE :AIME pass@1 降至 29.4
  • \(\lambda_{\text{actor} }\) 调优 :\(\lambda=0.99\) 效果最佳,\(\lambda=1.0\) 效果最差

关键 Insight

  • 值模型与策略的对齐至关重要 :尤其是在 Long-CoT 任务中,值模型必须充分理解策略的生成模式,否则会破坏 CoT 结构
  • 值预训练不仅是值热身,更是知识注入 :它帮助模型理解哪些 token 更有价值,是提升性能的关键
  • 值优化对噪声更鲁棒 :值模型可以使用 \(\lambda=1.0\) 的无偏目标,而策略需要 \(\lambda<1.0\) 来降低方差,二者解耦是合理的

附录:证明 score function 的梯度性质

  • 核心思路:
    • 概率和的梯度为 0
  • 证明详情见:Math——KL中得分函数和路径梯度的理解

NLP——技术报告解读-DeepSeek-V3

注:本文包含 AI 辅助创作

  • 参考链接:
    • 原始论文:DeepSeek-V3 Technical Report, arXiv 202412 & 202502, DeepSeek-AI
    • GitHub:github.com/deepseek-ai/DeepSeek-V3

Paper Summary

  • 核心内容总结:
    • DeepSeek-V3 是一个大型 MoE 语言模型(671B-A37B),在 14.8T Token 上进行了训练
    • 除了 DeepSeek-V2 中的 MLA 和 DeepSeekMoE 架构外,还开创性地采用了无辅助损失(auxiliary-loss-free)的负载平衡策略,并设置了 Multi-token 预测(Multi-token prediction)训练目标以获得更强的性能
    • 性能:截止到发布,DeepSeek-V3 是最强的开源模型,并且实现了与领先的闭源模型(如 GPT-4o 和 Claude-3.5-Sonnet)相媲美的性能
    • 最大特点:
      • 性能强大且保持了经济的训练成本,其完整训练(包括预训练、上下文长度扩展和后训练)仅需 2.788M H800 GPU 小时
      • 训练过程非常稳定,在整个训练过程中,论文没有遇到任何不可恢复的损失尖峰或执行任何回滚
  • DeepSeek-V3 采用了:
    • 多头潜在注意力(Multi-head Latent Attention, MLA),from DeepSeek-V2
    • DeepSeekMoE 架构, from DeepSeek-V2
    • DeepSeek-V3 开创了一种无辅助损失的(auxiliary-loss-free)负载均衡策略,并设定了 Multi-token 预测训练目标以获得更强的性能
  • 论文在 14.8T 个多样化和高质量的 Token 上对 DeepSeek-V3 进行了预训练(随后进行了 SFT 和 RL)

Introduction and Discussion

  • LLMs 经历了快速的迭代和演进 (Anthropic, 2024; Google, 2024; OpenAI, 2024a),逐步缩小了通往通用人工智能(Artificial General Intelligence, AGI)的差距(以上都是闭源模型)
  • 开源模型,包括 DeepSeek 系列 (DeepSeek-AI, 2024a, 2024b, 2024c; 2024)、LLaMA 系列 (Al@Meta, 2024a,2023b)、Qwen 系列 (Qwen, 2023, 2024a, 2024b) 和 Mistral 系列 (2023; Mistral, 2024),也取得了显著进展,努力缩小与闭源模型的差距
  • 为了进一步推动开源模型能力的边界,论文扩大了模型规模并推出了 DeepSeek-V3,这是一个大型 MoE 模型,拥有 671B 参数,其中每个 Token 激活 37B 参数
  • 以前瞻性的视角,论文始终致力于实现强大的模型性能和经济高效的成本
    • 在架构方面,DeepSeek-V3 仍然采用 MLA (DeepSeek-AI, 2024c) 以实现高效推理,并采用 DeepSeekMoE (2024) 以实现经济高效的训练
    • 这两种架构已在 DeepSeek-V2 (DeepSeek-AI, 2024c) 中得到验证,证明了它们在保持强大模型性能的同时实现高效训练和推理的能力
    • 除了基本架构之外,论文还实施了两种额外的策略以进一步增强模型能力
      • 1)DeepSeek-V3 开创了一种无辅助损失的负载均衡策略 (2024a),旨在最小化因鼓励负载均衡而对模型性能产生的不利影响
      • 2)DeepSeek-V3 采用了 Multi-token 预测训练目标,论文观察到这可以提升在评估基准上的整体性能
  • 为了实现高效训练,论文支持 FP8 混合精度训练并对训练框架进行了全面的优化
    • 低精度训练已成为一种有前景的高效训练解决方案 (2022; 2019; 2017; 2023b),其发展与硬件能力的进步密切相关 (2024; 2022; 2023a)
    • 论文引入了一个 FP8 混合精度训练框架,并首次在超大规模模型上验证了其有效性
    • 通过对 FP8 计算和存储的支持,论文实现了加速训练和降低 GPU 内存使用
  • 至于训练框架,论文设计了 DualPipe 算法以实现高效的流水线并行(Pipeline Parallelism, PP),该算法具有更少的流水线气泡,并通过计算-通信重叠(Overlap)隐藏了训练过程中的大部分通信
    • 这种重叠确保了随着模型的进一步扩展,只要论文保持恒定的计算-通信比,论文仍然可以在节点间使用细粒度专家(fine-grained experts),同时实现接近零的全交换(all-to-all)通信开销
    • 论文还开发了高效的跨节点全交换通信内核,以充分利用 InfiniBand (IB) 和 NVLink 带宽
    • 论文精心优化了内存占用,使得在不使用昂贵的张量并行(Tensor Parallelism, TP)的情况下训练 DeepSeek-V3 成为可能
  • 结合这些努力,论文实现了高训练效率
  • 在预训练期间,论文在 14.8T 高质量和多样化的 Token 上训练 DeepSeek-V3
    • 预训练过程非常稳定,在整个训练过程中,论文没有遇到任何不可恢复的损失尖峰或不得不回滚的情况
  • 接下来,论文对 DeepSeek-V3 进行了两阶段的上下文长度扩展
    • 在第一阶段,最大上下文长度扩展到 32K
    • 在第二阶段,进一步扩展到 128K
  • 此后,论文进行了Post-training,包括对 DeepSeek-V3 基础模型进行 SFT 和 RL ,以使其与人类偏好对齐并进一步释放其潜力
    • 在后训练阶段,论文从 DeepSeek-R1 系列模型中蒸馏推理能力,同时仔细维护模型准确性和生成长度之间的平衡
  • 论文在全面的基准测试套件上评估了 DeepSeek-V3
    • 综合评估显示 DeepSeek-V3-Base 已成为当前最强的开源基础模型,尤其是在代码和数学领域
    • 其聊天版本也在一系列标准和开放式基准测试中优于其他开源模型,并实现了与领先闭源模型(包括 GPT-4o 和 Claude-3.5-Sonnet)相媲美的性能
  • 最后,论文再次强调 DeepSeek-V3 的经济训练成本,总结在表 1 中,这是通过论文对算法、框架和硬件的优化协同设计实现的
    • 在预训练阶段,在每万亿 Token 上训练 DeepSeek-V3 仅需 180K H800 GPU 小时,即在论文拥有 2048 个 H800 GPU 的集群上需要 3.7 天
    • 因此,论文的预训练阶段在不到两个月内完成,成本为 2664K GPU 小时
    • 加上上下文长度扩展的 119K GPU 小时和后训练的 5K GPU 小时,DeepSeek-V3 的完整训练仅花费 2.788M GPU 小时
    • 假设 H800 GPU 的租赁价格为每小时 2 美元,论文的总训练成本仅为 557.6 万美元
    • 请注意,上述成本仅包括 DeepSeek-V3 的官方训练,不包括先前在架构、算法或数据上的研究和消融实验相关的成本
  • 论文的主要贡献包括:
    • 架构:创新的负载均衡策略和训练目标 (Architecture: Innovative Load Balancing Strategy and Training Objective)
      • 在 DeepSeek-V2 的高效架构基础上,论文开创了一种无辅助损失的负载均衡策略 (auxiliary-loss-free strategy for load balancing),该策略最小化了因鼓励负载均衡而产生的性能下降
      • 论文研究了 Multi-token 预测 (Multi-token Prediction, MTP) 目标,并证明其有益于模型性能。它也可用于推测解码 (speculative decoding) 以加速推理
    • 预训练:追求极致的训练效率 (Pre-Training: Towards Ultimate Training Efficiency)
      • 论文设计了一个 FP8 混合精度训练框架 (FP8 mixed precision training framework),并首次在超大规模模型上验证了 FP8 训练的可行性和有效性
      • 通过算法、框架和硬件的协同设计 (co-design),论文克服了跨节点 MoE 训练中的通信瓶颈,实现了近乎完全的计算-通信重叠 (computation-communication overlap)
        • 这显著提高了论文的训练效率并降低了训练成本,使论文能够在没有额外开销的情况下进一步扩大模型规模
      • 仅以 2.664M H800 GPU 小时的经济成本,论文完成了 DeepSeek-V3 在 14.8T Token 上的预训练,产生了当前最强的开源基础模型
        • 预训练之后的后续训练阶段仅需 0.1M GPU 小时
    • 后训练:来自 DeepSeek-R1 的知识蒸馏 (Post-Training: Knowledge Distillation from DeepSeek-R1)
      • 论文引入了一种创新方法,将推理能力从长思维链 (long Chain-of-Thought, long CoT) 模型,特别是从 DeepSeek R1 系列模型之一,蒸馏到标准大语言模型中,尤其是 DeepSeek-V3
      • 论文的流程巧妙地将 R1 的验证和反思模式融入 DeepSeek-V3,并显著提高了其推理性能
      • 同时,论文也保持了对 DeepSeek-V3 输出风格和长度的控制
    • 核心评估结果总结 (Summary of Core Evaluation Results)
      • 知识 (Knowledge) :
        • (1) 在教育类基准测试如 MMLU、MMLU-Pro 和 GPQA 上,DeepSeek-V3 优于所有其他开源模型,在 MMLU 上达到 88.5,在 MMLU-Pro 上达到 75.9,在 GPQA 上达到 59.1
          • 其性能与领先的闭源模型如 GPT-4o 和 Claude-Sonnet-3.5 相当,缩小了开源模型与闭源模型在该领域的差距
        • (2) 在事实性基准测试方面,DeepSeek-V3 在 SimpleQA 和 Chinese SimpleQA 上均表现出优于其他开源模型的性能
          • 虽然它在英语事实知识 (SimpleQA) 上落后于 GPT-4o 和 Claude-Sonnet-3.5,但在中文事实知识 (Chinese SimpleQA) 上超越了这些模型,突显了其中文事实知识的优势
      • 代码、数学和推理 (Code, Math, and Reasoning) :
        • (1) DeepSeek-V3 在所有非长思维链的开源和闭源模型中,在数学相关基准测试上达到了 SOTA 性能。值得注意的是,它在特定基准测试上(如 MATH-500)甚至超过了 o1-preview,展示了其强大的数学推理能力
        • (2) 在代码相关任务上,DeepSeek-V3 成为代码竞赛基准测试(如 LiveCodeBench)中表现最佳的模型,巩固了其在该领域的领先地位
          • 对于工程相关任务,虽然 DeepSeek-V3 的表现略低于 Claude-Sonnet-3.5,但它仍然以显著优势领先所有其他模型,展示了其在各种技术基准测试上的竞争力
  • 剩余部分组织如下:
    • 详细阐述论文的 DeepSeek-V3 模型架构(第 2 节)
    • 介绍论文的基础设施,包括论文的计算集群、训练框架、对 FP8 训练的支持、推理部署策略以及论文对未来硬件设计的建议(第 3 节)
    • 描述论文的预训练过程,包括训练数据的构建、超参数设置、长上下文扩展技术、相关评估以及一些讨论(第 4 节)
    • 讨论论文在后训练方面的工作,包括 SFT 、 RL 、相应的评估和讨论(第 5 节)
    • 总结讨论 DeepSeek-V3 现有的局限性,并提出未来研究的潜在方向(第 6 节)

Architecture

  • 论文首先介绍 DeepSeek-V3 的基本架构,其特点是采用 MLA (DeepSeek-AI) 实现高效推理,以及采用 DeepSeekMoE (2024) 实现经济高效的训练
  • 然后,论文介绍一种 MTP 训练目标,论文观察到该目标能够提升模型在评估基准上的整体性能
  • 对于其他未明确提及的细节,DeepSeek-V3 遵循 DeepSeek-V2 (DeepSeek-AI) 的设置

Basic Architecture

  • DeepSeek-V3 的基本架构仍然在 Transformer (2017) 框架内
  • 为了高效推理和经济高效的训练,DeepSeek-V3 同样采用了 MLA 和 DeepSeekMoE,这两者已在 DeepSeek-V2 中得到充分验证
  • 与 DeepSeek-V2 相比,一个例外是:
    • 论文额外为 DeepSeekMoE 引入了一种 无辅助损失负载均衡 (auxiliary-loss-free load balancing) 策略 (Wang 等人),以减轻因努力确保负载均衡而导致的模型性能下降
    • 图 2 展示了 DeepSeek-V3 的基本架构,论文将在本节简要回顾 MLA 和 DeepSeekMoE 的细节
Multi-Head Latent Attention
  • 对于注意力机制,DeepSeek-V3 采用 MLA 架构
  • 令 \(d\) 表示嵌入维度,\(n_{h}\) 表示注意力头数,\(d_{h}\) 表示每个头的维度,\(\mathbf{h}_{t}\in\mathbb{R}^{d}\) 表示给定注意力层中第 \(t\) 个 Token 的注意力输入
  • MLA 的核心是对注意力 Key 和 Value 进行低秩联合压缩,以减少推理期间的 Key-Value 缓存 (Key-Value, KV cache):
    $$
    \begin{align}
    \boxed{\mathbf{c}_{t}^{KV} } &= \mathbf{W}^{DKY}\mathbf{h}_{t},\\
    \left[\mathbf{k}_{t,1}^{C};\mathbf{k}_{t,2}^{C};…;\mathbf{k}_{t,n_{h} }^{C}\right] = \mathbf{k}_{t}^{C} &=\mathbf{W}^{UK}\mathbf{c}_{t}^{KV},\\
    \boxed{\mathbf{k}_{t}^{R}} &= \text{RoPE}(\mathbf{W}^{KR}\mathbf{h}_{t}),\\
    \mathbf{k}_{t,i} &= \left[\mathbf{k}_{t,i}^{C};\mathbf{k}_{t}^{R}\right],\\
    [\mathbf{v}_{t,1}^{C};\mathbf{v}_{t,2}^{C};…;\mathbf{v}_{t,n_{h} }^{C}] = \mathbf{v}_{t}^{C} &=\mathbf{W}^{UV}\mathbf{c}_{t}^{KV}
    \end{align}
    $$
    • 其中 \(\mathbf{c}_{t}^{KV}\in\mathbb{R}^{d_{c} }\) 是 Key 和 Value 的压缩潜在向量;
    • \(d_{c}(\ll d_{h}n_{h})\) 表示 KV 压缩维度;
    • \(W^{DKY}\in\mathbb{R}^{d_{c}\times d}\) 表示下投影矩阵;
    • \(W^{UK},W^{UV}\in\mathbb{R}^{d_{h}n_{h}\times d_{c} }\) 分别是 Key 和 Value 的上投影矩阵;
    • \(W^{KR}\in\mathbb{R}^{d_{h}^{R}\times d}\) 是用于生成携带 旋转位置嵌入 (Rotary Positional Embedding, RoPE) (2024) 的解耦 Key 的矩阵;
    • RoPE(\(\cdot\)) 表示应用 RoPE 矩阵的操作;
    • [\(\cdot\);\(\cdot\)] 表示拼接
  • 对于 MLA,在生成过程中只需要缓存蓝色框内的向量(即 \(\mathbf{c}_{t}^{KV}\) 和 \(\mathbf{k}_{t}^{R}\)),这能在保持性能与标准 多头注意力 (Multi-Head Attention, MHA) (2017) 相当的同时,显著减少 KV 缓存
  • 对于注意力 Query ,论文也执行低秩压缩,这可以减少训练期间的激活内存:
    $$
    \begin{align}
    \mathbf{c}_{t}^{Q}&=W^{DQ}\mathbf{h}_{t},\\
    \left[\mathbf{q}_{t,1}^{C};\mathbf{q}_{t,2}^{C};…;\mathbf{q}_{t,n_{h} }^{C}\right]=\mathbf{q}_{t}^{C}&=W^{UQ}\mathbf{c}_{t}^{Q},\\
    \left[\mathbf{q}_{t,1}^{R};\mathbf{q}_{t,2}^{R};…;\mathbf{q}_{t,n_{h} }^{R}\right]=\mathbf{q}_{t}^{R}&=\text{RoPE}(W^{QR}\mathbf{c}_{t}^{Q}),\\
    \mathbf{q}_{t,i}&=\left[\mathbf{q}_{t,i}^{C};\mathbf{q}_{t,i}^{R}\right],
    \end{align}
    $$
    • 其中 \(\mathbf{c}_{t}^{Q}\in\mathbb{R}^{d_{c}^{\prime} }\) 是 Query 的压缩潜在向量;
    • \(d_{c}^{\prime}(\ll d_{h}n_{h})\) 表示 Query 压缩维度;
    • \(W^{DQ}\in\mathbb{R}^{d_{c}^{\prime}\times d}\),\(W^{UQ}\in\mathbb{R}^{d_{h}n_{h}\times d_{c}^{\prime} }\) 分别是 Query 的下投影和上投影矩阵;
    • \(W^{QR}\in\mathbb{R}^{d_{h}^{R}n_{h}\times d_{c}^{\prime} }\) 是用于生成携带 RoPE 的解耦 Query 的矩阵
  • 最终,注意力 Query (\(\mathbf{q}_{t,i}\))、 Key (\(\mathbf{k}_{j,i}\)) 和值 (\(\mathbf{v}_{j,i}^{C}\)) 被组合以产生最终的注意力输出 \(\mathbf{u}_{t}\):
    $$
    \begin{align}
    \mathbf{o}_{t,i}&=\sum_{j=1}^{t}\text{Softmax}_{j}(\frac{\mathbf{q}_{t,i}^{T}\mathbf{k}_{j,i} }{\sqrt{d_{h}+d_{h}^{R} } })\mathbf{v}_{j,i}^{C},\\
    \mathbf{u}_{t}&=W^{O}\left[\mathbf{o}_{t,1};\mathbf{o}_{t,2};…;\mathbf{o}_{t,n_{h} }\right],
    \end{align}
    $$
    • 其中 \(W^{O}\in\mathbb{R}^{d\times d_{h}n_{h} }\) 表示输出投影矩阵
DeepSeekMoE with Auxiliary-Loss-Free Load Balancing
Basic Architecture of DeepSeekMoE
  • 对于 前馈网络 (Feed-Forward Networks, FFNs),DeepSeek-V3 采用 DeepSeekMoE 架构 (2024)
  • 与传统的 MoE 架构(如 CShard (2021))相比,DeepSeekMoE 使用更细粒度的专家,并将一些专家隔离为共享专家
  • 令 \(\mathbf{u}_{t}\) 表示第 \(t\) 个 Token 的 FFN 输入,论文计算 FFN 输出 \(\mathbf{h}_{t}^{\prime}\) 如下:
    $$
    \begin{align}
    \mathbf{h}_{t}^{\prime} &=\mathbf{u}_{t}+\sum_{i=1}^{N_{s} }\text{FFN }_{i}^{(s)}\left(\mathbf{u}_{t}\right)+\sum_{i=1}^{N_{r} }g_{i,t}\text{ FFN}_{i}^{(r)}\left(\mathbf{u}_{t}\right),\\
    g_{i,t} &= \frac{g_{i,t}^{\prime} }{\sum_{j=1}^{N_{r} }g_{j,t}^{\prime} },\\
    g_{i,t}^{\prime} &= \begin{cases}s_{i,t},&s_{i,t}\in\text{Topk}(\{s_{j,t}|1\leqslant j\leqslant N_{r}\},K_{r}),\\ 0,&\text{otherwise},\end{cases} \\
    s_{i,t} &= \text{Sigmoid}\left(\mathbf{u}_{t}^{T}\mathbf{e}_{t}\right),
    \end{align}
    $$
    • FFN\({}_{i}^{(s)}(\cdot)\) 表示第 \(i\) 个共享专家;
      • \(N_{s}\) 表示共享专家数量
    • FFN\({}_{i}^{(r)}(\cdot)\) 表示第 \(i\) 个路由专家;
      • 和 \(N_{r}\) 表示路由专家数量
    • \(K_{r}\) 表示激活的路由专家数量;
    • \(g_{i,t}\) 是第 \(i\) 个专家的门控值;
    • \(s_{i,t}\) 是 Token 到专家的亲和度;
    • \(\mathbf{e}_{t}\) 是第 \(i\) 个路由专家的质心向量;
    • Topk\((\cdot,K)\) 表示包含为第 \(t\) 个 Token 和所有路由专家计算的亲和度分数中 \(K\) 个最高分的集合
  • 与 DeepSeek-V2 略有不同,DeepSeek-V3 使用 sigmoid 函数计算亲和度分数,并在所有选定的亲和度分数之间应用归一化以产生门控值
Auxiliary-Loss-Free Load Balancing
  • 对于 MoE 模型,不平衡的专家负载将导致 路由崩溃 (routing collapse) (2017),并在使用专家并行的情况下降低计算效率
    • 传统的解决方案通常依赖 辅助损失 (auxiliary loss) (2021;) 来避免负载不平衡,但过大的辅助损失会损害模型性能 (Wang 等人)
  • 为了在负载均衡和模型性能之间取得更好的权衡,论文开创了一种 无辅助损失负载均衡 (auxiliary-loss-free load balancing) 策略 (Wang 等人) 来确保负载均衡
    • 论文为每个专家引入一个偏置项 \(b_{i}\),并将其加到相应的亲和度分数 \(s_{i,t}\) 上以确定 top-K 路由:
      $$g^{\prime}_{i,t}=\begin{cases}s_{i,t},&s_{i,t}+b_{i}\in\text{Topk}((s_{j,t}+b_{j}|1\leqslant j\leqslant N_{r}),K_{r}),\\ 0,&\text{otherwise}.\end{cases}$$
    • 特别说明:偏置项仅用于路由
    • 将与 FFN 输出相乘的门控值仍然来自原始亲和度分数 \(s_{i,t}\)
  • 在训练期间,论文持续监控每个训练步骤的整个批次的专家负载
    • 在每个步骤结束时,如果某个专家对应的负载过重,论文将其偏置项减少 \(y\),如果其对应的专家负载不足,则将其偏置项增加 \(y\),其中 \(y\) 是一个称为 偏置更新速度 (bias update speed) 的超参数
    • 通过这种动态调整,DeepSeek-V3 在训练期间保持均衡的专家负载,并且比那些仅通过纯辅助损失来鼓励负载均衡的模型实现了更好的性能
Complementary Sequence-Wise Auxiliary Loss
  • 尽管 DeepSeek-V3 主要依赖无辅助损失策略进行负载均衡,但为了防止任何单个序列内出现极端不平衡,论文也采用了互补的序列级均衡损失:
    $$
    \begin{align}
    \mathcal{L}_{\text{Bal} }&=\alpha\sum_{i=1}^{N_{r} }f_{i}P_{i},\\
    f_{i}&=\frac{N_{r} }{K_{r}T}\sum_{t=1}^{T}\mathbb{I}\left(s_{i,t}\in\text{Topk}(\{s_{j,t}|1\leqslant j\leqslant N_{r}\},K_{r})\right),\\
    s^{\prime}_{i,t}&=\frac{s_{i,t} }{\sum_{j=1}^{N_{r} }s_{j,t} },\\
    P_{i}&=\frac{1}{T}\sum_{t=1}^{T}s^{\prime}_{i,t},
    \end{align}
    $$
    • 平衡因子 \(\alpha\) 是一个超参数,对于 DeepSeek-V3 将被赋予一个极小的值;
    • \(\mathbb{I}(\cdot)\) 表示指示函数;
    • \(T\) 表示序列中的 Token 数
    • 序列级均衡损失鼓励每个序列上的专家负载保持均衡
    • 理解:
      • 每次计算最小化 损失 \(\mathcal{L}_{\text{Bal} }\) 时,都可以看做是一个求解约束优化问题的过程
      • 约束优化问题为:
        $$
        \begin{align}
        \min_{P_i} \quad &\sum_{i=1}^{N_{r} }f_{i}P_{i} \\
        \text{s.t.} &\sum_i P_i = 1
        \end{align}
        $$
        • 求解约束优化问题时,系数 \(f_{i}\) 可以看做是固定值(是每个序列的统计值,不同序列该值不同)
        • 变量是 \(P_{i}\) ,需要满足一定约束
          $$ \sum_i P_i=1 $$
          • 该约束可以通过定义推导得到:
            $$ \sum_j P_i = \sum_i \frac{1}{T}\sum_{t=1}^{T}s^{\prime}_{i,t} = \sum_i \frac{1}{T}\sum_{t=1}^{T} \frac{s_{i,t} }{\sum_{j=1}^{N_{r} }s_{j,t} } = \frac{1}{T}\sum_{t=1}^{T} \sum_i \frac{s_{i,t} }{\sum_{j=1}^{N_{r} }s_{j,t} } = \frac{1}{T}\sum_{t=1}^{T} 1 = \frac{1}{T}T = 1$$
        • 直观上看,最小化 \(\sum_{j=1}^{D+1}f_iP_i\) 的解就是让概率 \(P_i\) 随着 \(f_i\) 变化, \(f_i\) 越小,则 \(P_i\) 应该越大
          • 这样才能才能满足最小化 \(\sum_{j=1}^{D+1}f_iP_i\)
        • 从梯度上看
          $$ \frac{\partial \mathcal{L}_{\text{ LB} }}{\partial P_i} = \alpha f_i $$
          • 对于 \(f_i\) 越大的组,其概率 \(P_i\) 下降的越多
          • 进一步理解:\(P_i\) 下降是通过调整模型 router 参数实现的,这会导致参数更新后下一轮中真实统计值 \(f_i\) 下降,最终会收敛到一个大家的真实分配统计值 \(f_i\) 都差不多相同的地方(这也就实现了所谓的均衡)
Node-Limited Routing
  • 与 DeepSeek-V2 使用的设备受限路由类似,DeepSeek-V3 也使用一种受限路由机制来限制训练期间的通信成本
    • 注:该方法在 DeepSeek-V2 中被称为 Device-Limited Routing
      • 对于 DeepSeek-V2,除了对被路由专家进行简单的 top-K 选择外,还额外确保每个 Token 的目标专家最多分布在 \(M\) 个设备上
      • 具体来说,对于每个 Token
        • 首先筛选出 \(M\) 个设备,这些设备中包含亲和度得分(affinity scores)最高的专家;
        • 随后,在这 \(M\) 个设备的专家中执行 top-K 选择
      • 在实际实验中作者发现,当 \(M \ge 3\) 时,这种设备受限路由能够实现良好的性能,其效果与无限制的 top-K 路由基本持平
  • 简而言之,论文确保每个 Token 最多被发送到 \(M\) 个节点 ,这些节点是根据分布在该节点上的专家的最高 \(\frac{K}{M}\) 个亲和度分数之和来选择的
    • 理解(为什么是选择依据是每个节点最高 \(\frac{K}{M}\) 个亲和度之和?):对于每个 Token,需要分发到共 \(K\) 个专家,在限定只能发送到 \(M\) 个节点的情况下,则每个节点平均需要负责这个 Token 的 \(\frac{K}{M}\) 个专家,所以根据每个节点的最高 \(\frac{K}{M}\) 个亲和度之和来选择最终的 \(M\) 个节点
  • 在此约束下,论文的 MoE 训练框架几乎可以实现完全的计算-通信重叠
    • 理解:因为限定了每个 Token 需要分发的节点数量,这种情况下可以大幅减少通信开销,避免一个 Token 被发送到太多的节点而造成通信开销太大的情况
No Token-Dropping
  • 由于有效的负载均衡策略,DeepSeek-V3 在整个训练过程中保持良好的负载均衡,DeepSeek-V3 在训练期间不丢弃任何 Token
  • 论文还实现了特定的部署策略以确保推理时的负载均衡,因此 DeepSeek-V3 在推理期间也不丢弃 Token

Multi-token Prediction

  • 受 Gloeckle 等人 (2024) 的启发,论文研究并为 DeepSeek-V3 设定了 MTP 目标,该目标将每个位置的预测范围扩展到多个未来 Token
    • 1)MTP 目标使训练信号更加密集,并可能提高数据效率
    • 2)MTP 可能使模型能够预规划其表示,以更好地预测未来 Token
  • 图 3 展示了论文的 MTP 实现
    • 与 Gloeckle 等人 (2024) 使用独立的输出头并行预测 \(D\) 个额外 Token 不同,论文顺序预测额外 Token ,并在每个预测深度保持完整的因果链
  • 注:上图是训练的流程,按照 Teacher Forcing 方式训练的,所以 MTP 头的输入是真实的 Token,在推理时,会使用主网络的输出作为第二个 Token 的输入,以此类推,详情见后文
MTP Modules
  • 具体来说,论文的 MTP 实现使用 \(D\) 个顺序模块来预测 \(D\) 个额外 Token
    • 注:从图3可以看出,并行头只在模型的最后一层(主模型的 Head 层前)上增加,而不是在 Transformer 网络前面的层增加
    • 注:加入的内容是单层的 Transformer 再加一个 Head
  • 第 \(k\) 个 MTP 模块包含一个共享嵌入层 \(\text{Emb}(\cdot)\)、一个共享输出头 \(\text{OutHead}(\cdot)\)、一个 Transformer 块 \(\text{TRM}_{k}(\cdot)\) 和一个投影矩阵 \(M_{k}\in\mathbb{R}^{d\times 2d}\)
  • 对于第 \(i\) 个输入 Token \(t_{i}\),在第 \(k\) 个预测深度(注:这里的预测深度是指往后预测的 Token 数,不是 Transformer 网络层深度)
    • 论文首先将第 \(i\) 个 Token 在第 \((k-1)\) 个深度的表示 \(\mathbf{h}_{i}^{k-1}\in\mathbb{R}^{d}\) 与第 \((i+k)\) 个 Token 的嵌入通过线性投影相结合:
      $$\mathbf{h}_{i}^{ \prime{k} }=M_{k}[\textrm{RMSNorm}(\mathbf{h}_{i}^{ {k}-1});\textrm{RMSNorm}(\textrm{Emb}(t_{i+k}))],$$
    • 其中 \([.;.]\) 表示拼接
    • 特别地,当 \(k=1\) 时,\(\mathbf{h}_{i}^{ {k}-1}\) 指的是主模型给出的表示
    • 特别注意:对于每个 MTP 模块,其嵌入层与主模型共享
    • 组合后的 \(\mathbf{h}_{i}^{ {k} }\) 作为第 \(k\) 个深度的 Transformer 块的输入,以产生当前深度的输出表示 \(\mathbf{h}_{i}^{ {k} }\):
      $$\mathbf{h}_{1:T-k}^{ {k} }=\textrm{TRM}_{k}(\mathbf{h}_{1:T-k}^{ \prime{k} }),$$
    • 其中 \(T\) 表示输入序列长度,\({}_{ {i:j} }\) 表示切片操作(包括左右边界)
    • 最后,以 \(\mathbf{h}_{i}^{ {k} }\) 作为输入,共享输出头将计算第 \(k\) 个额外预测 Token 的概率分布 \(p_{ {i+1+k} }^{ {k} }\in\mathbb{R}^{V}\),其中 \(V\) 是词汇表大小:
      $$p_{ {i+k+1} }^{ {k} }=\textrm{OutHead}(\mathbf{h}_{i}^{ {k} }).$$
  • 输出头 \(\textrm{OutHead}(\cdot)\) 将表示线性映射为 logits,随后应用 \(\textrm{Softmax}(\cdot)\) 函数来计算第 \(k\) 个额外 Token 的预测概率
    • 同样,对于每个 MTP 模块,其输出头与主模型共享
    • 论文保持预测因果链的原则与 EAGLE (Li 等人) 类似,但 EAGLE 主要目标是 推测解码 (speculative decoding) (2023; 2023),而论文利用 MTP 来改进训练
MTP Training Objective
  • 对于每个预测深度,论文计算一个交叉熵损失 \(\mathcal{L}_{\textrm{MTP} }^{ {k} }\):
    $$\mathcal{L}_{\textrm{MTP} }^{ {k} }=\textrm{CrossEntropy}(p_{2+ k:T+1}^{ {k} },t_{2+k:T+1})=\frac{1}{T}\sum_{i=2+k}^{T+1}\log p_ {i}^{ {k} }[t_{i}],$$
    • 其中 \(T\) 表示输入序列长度,\(t_{i}\) 表示第 \(i\) 个位置的真实 Token
    • \(p_{i}^{ {k} }[t_{i}]\) 表示由第 \(k\) 个 MTP 模块给出的 \(t_{i}\) 的相应预测概率
  • 最后,论文计算所有深度上 MTP 损失的平均值,并将其乘以权重因子 \(\lambda\) 以获得总体 MTP 损失 \(\mathcal{L}_{\textrm{MTP} }\),该损失作为 DeepSeek-V3 的额外训练目标:
    $$\mathcal{L}_{\textrm{MTP} }=\frac{\lambda}{D}\sum_{k=1}^{D}\mathcal{L}_{\textrm{ MTP} }^{ {k} }.$$
MTP in Inference
  • 论文的 MTP 策略主要旨在提高主模型的性能 ,因此在推理期间,我们可以直接丢弃 MTP 模块,主模型可以独立且正常地运行
  • 也可以将这些 MTP 模块重新用于推测解码,以进一步改善生成延迟
  • 理解:实际上,后续的其他工作来看,还是提高主模型本身的性能用的多
  • 额外补充:关于 MTP 推理的其他讲解
    • 参考链接:deepseek技术解读(2)-MTP(Multi-token Prediction)的前世今生 - 姜富春的文章 - 知乎
    • 基本思想:
      • 预测阶段(Predict) :通过 K 个头一次生成 K 个 Token 的预测
      • 验证阶段(Verify) :将 K 个 Token 组装成 K 个 <input,label> 对,并行地利用输入 Main Model 作为评估验证,如果输出 label 与 Main Model 一致,则接受该 token
      • 接受阶段(Accept) :最终接受满足 Main Model 的最大长度 tokens 作为输出
    • 但是,从上述 图3 来看,多头预估时,预测 \(t^{\prime}_3\) 的输入包括了 \(t_1\) 的 Transformer 最后一层表征和 \(t_2\) 的 Embedding(即 Teacher Forcing 的形式),所以推理时是无法直接使用的,流程同上述方法一样,但应该使用下图的形式:
  • 特别说明:虽然主头的预测和后面头的预测是串行的,但是实际上,后面的头进需要走一个 Transformer 层,速度是很快的,和动辄几十层的主网络比起来(比如 DeepSeek-V3 是 61 层),算是很快的了(几乎可以认为是并行了)

Infrastructures

Compute Clusters

  • DeepSeek-V3 在一个配备有 2048 个 NVIDIA H800 GPU 的集群上进行训练
  • H800 集群中的每个节点包含 8 个通过节点内的 NVLink 和 NVSwitch 连接的 GPU
  • 在不同节点之间,利用 InfiniBand (IB) 互连来促进通信

Training Framework

  • DeepSeek-V3 的训练由 HAI-LLM 框架支持,这是一个由论文的工程师从头开始精心打造的、高效且轻量级的训练框架
  • DeepSeek-V3 应用了 16 路流水线并行 (Pipeline Parallelism, PP) (2023b)、跨越 8 个节点的 64 路专家并行 (Expert Parallelism, EP) (2021) 和 ZeRO-1 数据并行 (Data Parallelism, DP) (2020)
  • 为了促进 DeepSeek-V3 的高效训练,论文实施了细致的工程优化
    • 首先,论文设计了 DualPipe 算法用于高效的流水线并行
      • 与现有的 PP 方法相比,DualPipe 具有更少的流水线气泡
      • 更重要的是,它在前向和后向过程中重叠了计算和通信阶段,从而解决了跨节点专家并行引入的沉重通信开销的挑战
    • 其次,论文开发了高效的跨节点 All-to-All 通信内核,以充分利用 IB 和 NVLink 的带宽,并节省专用于通信的流式多处理器 (Streaming Multiprocessors, SMs)
    • 最后,论文精心优化了训练期间的内存占用,从而使得论文能够在训练 DeepSeek-V3 时不使用昂贵的张量并行 (Tensor Parallelism, TP)
DualPipe and Computation-Communication Overlap
  • 对于 DeepSeek-V3,跨节点专家并行引入的通信开销导致了大约 1:1 的低效计算-通信比
  • 为了应对这一挑战,论文设计了一种名为 DualPipe 的创新流水线并行算法,该算法不仅通过有效重叠前向和后向计算-通信阶段来加速模型训练,而且还减少了流水线气泡
  • DualPipe 的关键思想是在一对独立的前向和后向块 (chunk) 内重叠计算和通信
    • 论文将每个块划分为四个组件:注意力 (attention)、 All-to-All 分发 (all-to-all dispatch)、MLP 和 All-to-All 合并 (all-to-all combine)
    • 特别地,对于一个后向块,注意力和 MLP 都进一步分为两部分:对输入的梯度 (backward for input) 和对权重的梯度 (backward for weights),类似于 ZeroBubble (2023a) 中的做法
    • 此外,论文还有一个 PP 通信组件
  • 如图 4 所示,对于一对前向和后向块,论文重新排列了这些组件,并手动调整了专用于通信与计算的 GPU SMs 的比例
    • 在这种重叠策略中,我们可以确保 All-to-All 通信和 PP 通信在执行过程中都能被完全隐藏
    • 颜色上看:
      • 绿色(B)表示 backward for input
      • 蓝色(W)表示 backward for weight
      • 橙色(F)表示 forward
    • 通信上看:
      • 绿色(B)表示 backward 的通信
      • 橙色(F)表示 forward 的通信
      • 紫色(PP)表示 PP 的通信
  • 鉴于高效的重叠策略,完整的 DualPipe 调度如图 5 所示
    • 它采用了双向流水线调度,同时从流水线的两端输入微批次 (micro-batches),并且大部分通信可以完全重叠
    • 这种重叠也确保了,随着模型的进一步扩展,只要论文保持恒定的计算-通信比,论文仍然可以在跨节点使用细粒度专家,同时实现接近零的 All-to-All 通信开销
  • 即使在通信负担较轻的更一般情况下,DualPipe 仍然展现出效率优势
    • 在表 2 中,论文总结了不同 PP 方法的流水线气泡和内存使用情况
    • 如表所示,与 ZB1P (2023b) 和 1F1B (2018) 相比,DualPipe 显著减少了流水线气泡,同时仅将峰值激活内存增加了 \(\frac{1}{PP}\) 倍
    • 尽管 DualPipe 需要保留两份模型参数 ,但这并不会显著增加内存消耗 ,因为论文在训练期间使用了较大的 EP 规模
    • 与 Chimera (2021) 相比,DualPipe 仅要求流水线阶段和微批次能被 2 整除,而不要求微批次能被流水线阶段整除
    • 此外,对于 DualPipe,气泡和激活内存都不会随着微批次数量的增加而增加
Efficient Implementation of Cross-Node All-to-All Communication
  • 为了确保 DualPipe 具有足够的计算性能,论文定制了高效的跨节点 All-to-All 通信内核(包括分发和合并),以节省专用于通信的 SMs 数量
    • 这些内核的实现与论文的 MoE 门控算法和集群的网络拓扑结构协同设计
  • 具体来说,在论文的集群中,跨节点 GPU 通过 IB 完全互连,节点内通信通过 NVLink 处理
    • NVLink 提供 160 GB/s 的带宽,大约是 IB (50 GB/s) 的 3.2 倍
    • 为了有效利用 IB 和 NVLink 的不同带宽,论文将每个 Token 限制为最多分发到 4 个节点,从而减少 IB 流量
    • 对于每个 Token ,当做出路由决策时,它将首先通过 IB 传输到其目标节点上具有相同节点内索引的 GPU
    • 一旦到达目标节点,论文将尽力确保它通过 NVLink 瞬时转发到承载其目标专家的特定 GPU,而不会被后续到达的 Token 阻塞
    • 通过这种方式,通过 IB 和 NVLink 的通信完全重叠,每个 Token 可以高效地平均在每个节点选择 3.2 个专家,而不会产生 NVLink 的额外开销
    • 这意味着,尽管 DeepSeek-V3 在实践中仅选择 8 个路由专家,但它可以将此数量最多扩展到 13 个专家(4 个节点 × 3.2 个专家/节点),同时保持相同的通信成本
    • 总体而言,在这种通信策略下,仅需 20 个 SMs 就足以充分利用 IB 和 NVLink 的带宽
  • 详细来说,论文采用了 warp 专业化技术 (2014),并将 20 个 SMs 划分为 10 个通信通道,在分发过程中下面的操作由各自的 warp 处理
    • (1) IB 发送
    • (2) IB 到 NVLink 转发
    • (3) NVLink 接收
  • 分配给每个通信任务的 warp 数量根据所有 SMs 上的实际工作负载动态调整,类似地,在合并过程中下面的步骤也有由动态调整的 warp 处理
    • (1) NVLink 发送
    • (2) NVLink 到 IB 转发和累加
    • (3) IB 接收和累加
  • 此外,分发和合并内核都与计算流重叠,因此论文也考虑了它们对其他 SM 计算内核的影响
    • 具体来说,论文采用了定制的 PTX (Parallel Thread Execution) 指令并自动调整通信块大小,这显著减少了 L2 缓存的使用以及对其他 SMs 的干扰
Extremely Memory Saving with Minimal Overhead
  • 为了减少训练期间的内存占用,论文采用了以下技术
  • RMSNorm 和 MLA 上投影的重计算 (Recomputation of RMSNorm and MLA Up-Projection)
    • 论文在反向传播期间重新计算所有 RMSNorm 操作和 MLA 上投影,从而无需持久存储它们的输出激活
    • 以微小的开销为代价,该策略显著减少了存储激活所需的内存
  • CPU 中的指数移动平均 (Exponential Moving Average in CPU)
    • 在训练期间,论文保留模型参数的指数移动平均 (EMA) 用于在学习率衰减后早期估计模型性能
      • EMA 参数存储在 CPU 内存中,并在每个训练步骤后异步更新
    • 这种方法使论文能够维护 EMA 参数,而不会产生额外的内存或时间开销
  • Multi-token 预测的共享嵌入和输出头 (Shared Embedding and Output Head for Multi-token Prediction)。通过 DualPipe 策略,论文将模型的最浅层(包括嵌入层)和最深层(包括输出头)部署在同一个 PP 排名 (rank) 上。这种安排使得 MTP 模块和主模型之间能够物理共享共享嵌入和输出头的参数和梯度。这种物理共享机制进一步提高了论文的内存效率

FP8 Training

  • 受低精度训练最新进展 (2022; 2023b) 的启发,论文提出了一个利用 FP8 数据格式的细粒度混合精度框架来训练 DeepSeek-V3
  • 虽然低精度训练前景广阔,但它通常受到激活、权重和梯度中异常值 (outliers) 存在的限制 (2024; 2024)
    • 尽管在推理量化方面取得了显著进展 (2022; 2023),但相对较少的研究证明了低精度技术在大规模语言模型预训练中的成功应用 (2024)
  • 为了应对这一挑战并有效扩展 FP8 格式的动态范围,论文引入了一种细粒度的量化策略:
    • 使用 \(1\times N_{c}\) 元素的 tile-wise 分组 (tile-wise grouping) 或 \(N_{c}\times N_{c}\) 元素的块级分组 (block-wise grouping)
      • 注:tile-wise 说明量化维度是按照块做的,不是整个张量统一量化(tile 粒度比张量粒度更细)
    • 相关的反量化开销在论文提高了精度的累加过程下得到了很大程度的缓解,这是实现精确 FP8 通用矩阵乘法 (General Matrix Multiplication, GEMM) 的关键方面
    • 为了进一步减少 MoE 训练中的内存和通信开销,论文以 FP8 格式缓存和分发激活,同时以 BF16 格式存储低精度的优化器状态
    • 论文在两个与 DeepSeek-V2-Lite 和 DeepSeek-V2 规模相似的模型上验证了所提出的 FP8 混合精度框架,训练了大约 1T 个 Token(更多细节见附录 B.1)
  • 与 BF16 基线相比,论文的 FP8 训练模型的相对损失误差始终低于 0.25%,这一水平完全在训练随机性的可接受范围内
Mixed Precision Framework
  • 基于低精度训练中广泛采用的技术 (2019; 2017),论文提出了一个用于 FP8 训练的混合精度框架
  • 在此框架中,大多数计算密集型操作以 FP8 精度进行,而一些关键操作则策略性地保持其原始数据格式,以平衡训练效率和数值稳定性
  • 整体框架如图 6 所示
    • 首先,为了加速模型训练,大多数核心计算内核,即 GEMM 操作,都以 FP8 精度实现
    • 这些 GEMM 操作接受 FP8 张量作为输入,并产生 BF16 或 FP32 的输出。如图 6 所示,与线性算子 (Linear operator) 相关的所有三个 GEMM,即 Fprop(前向传播)、Dgrad(激活反向传播)和 Wgrad(权重反向传播),都在 FP8 中执行
    • 该设计理论上比原始的 BF16 方法提高一倍的计算速度
    • 此外,FP8 的 Wgrad GEMM 允许激活以 FP8 格式存储,用于反向传播。这显著减少了内存消耗
  • 尽管 FP8 格式具有效率优势,但某些算子由于其对低精度计算的敏感性,仍然需要更高的精度
    • 此外,一些低成本算子也可以使用更高的精度,而对整体训练成本的开销可以忽略不计
    • 因此,经过仔细研究,论文为以下组件保留了原始精度(例如 BF16 或 FP32):嵌入模块、输出头、MoE 门控模块、归一化算子和注意力算子
    • 这些有针对性的高精度保留确保了 DeepSeek-V3 的训练动态稳定性
    • 为了进一步保证数值稳定性,论文以更高的精度存储主权重 (master weights)、权重梯度和优化器状态
    • 虽然这些高精度组件会产生一些内存开销,但它们的影响可以通过在论文的分布式训练系统中跨多个 DP 排名进行高效分片来最小化
Improved Precision from Quantization and Multiplication
  • 基于论文的混合精度 FP8 框架,论文引入了若干策略来增强低精度训练的准确性,重点关注量化方法和乘法过程
Fine-Grained Quantization
  • 在低精度训练框架中,由于 FP8 格式的有限动态范围(受限于其减少的指数位),溢出和下溢是常见的挑战
    • 作为标准实践,通过将输入张量的最大绝对值缩放到 FP8 格式的最大可表示值来将输入分布对齐到 FP8 格式的可表示范围 (2017)
  • 这种方法使得低精度训练对激活异常值高度敏感,这会严重降低量化精度
  • 为了解决这个问题,论文提出了一种细粒度量化方法,在更细粒度的级别上应用缩放
  • 如图 7 (a) 所示
    • (1) 对于激活,论文在 1x128 的 tile 基础上(即每个 Token 每 128 个通道)对元素进行分组和缩放;
    • (2) 对于权重,论文在 128x128 的块基础上(即每 128 个输入通道每 128 个输出通道)对元素进行分组和缩放
  • 这种方法通过根据更小的元素组调整缩放比例,确保量化过程能更好地适应异常值
  • 在附录 B.2 中,论文进一步讨论了当论文以与权重量化相同的方式对激活进行块级分组和缩放时出现的训练不稳定性
  • 论文方法中的一个关键修改是引入了沿 GEMM 操作内部维度 (inner dimension) 的每组缩放因子 (per-group scaling factors)
    • 此功能在标准的 FP8 GEMM 中并不直接支持
    • 但结合论文精确的 FP32 累加策略,它可以被高效实现
  • 值得注意的是,论文的细粒度量化策略与微缩放格式 (microscaling formats) 的思想高度一致 (2023a),而 NVIDIA 下一代 GPU(Blackwell 系列)的 Tensor Cores 已宣布支持具有更小量化粒度的微缩放格式 (2022b)
    • 作者希望论文的设计能为未来的工作提供参考,以跟上最新的 GPU 架构
Increasing Accumulation Precision
  • 低精度 GEMM 操作经常遭受下溢问题,其精度在很大程度上依赖于高精度累加,这通常以 FP32 精度执行 (2019; 2017)
    • 但论文观察到在 NVIDIA H800 GPU 上,FP8 GEMM 的累加精度仅限于保留大约 14 位,这显著低于 FP32 的累加精度
    • 当内部维度 K 很大时,这个问题会更加明显 (2023),这在大规模模型训练中增加批大小和模型宽度时是典型情况
    • 以两个随机矩阵的 GEMM 操作为例,其中 \(\textit{K}=4096\),在论文的初步测试中,Tensor Cores 中有限的累加精度导致最大相对误差接近 \(2%\)
  • 尽管存在这些问题,有限的累加精度仍然是一些 FP8 框架中的默认选项 (2022c),严重限制了训练精度
  • 为了解决这个问题,论文采用了提升到 CUDA Cores 以获得更高精度的策略 (2023)
    • 该过程如图 7 (b) 所示
    • 在 Tensor Cores 上执行 MMA(矩阵乘积累加)期间,中间结果使用有限的位宽进行累加
    • 一旦达到 \(N_c\) 个元素的间隔,这些部分结果将被复制到 CUDA Cores 上的 FP32 寄存器中,在那里执行全精度的 FP32 累加
    • 如前所述,论文的细粒度量化沿内部维度 K 应用每组缩放因子
    • 这些缩放因子可以作为反量化过程在 CUDA Cores 上高效地相乘,而只需最小的额外计算成本
  • 值得注意的是,这种修改降低了单个 warpgroup 的 WGMMA (Warpgroup-level Matrix Multiply-Accumulate) 指令发出率
    • 但在 H800 架构上,通常可以同时维持两个 WGMMA:当一个 warpgroup 执行提升操作时,另一个能够执行 MMA 操作
    • 这种设计使得两个操作能够重叠,保持了 Tensor Cores 的高利用率
    • 根据论文的实验,设置 \(N_c\) = 128 个元素,相当于 4 个 WGMMAs,是能够显著提高精度而不引入大量开销的最小累加间隔
Mantissa over Exponents(尾数优先于指数)
  • 与先前工作采用的混合 FP8 格式 (2022c; 2023b;) 相比,后者在 Fprop 中使用 E4M3(4 位指数和 3 位尾数),在 Dgrad 和 Wgrad 中使用 E5M2(5 位指数和 2 位尾数),论文在所有张量上采用 E4M3 格式以获得更高的精度
  • 论文将此方法的可行性归功于论文的细粒度量化策略,即 tile 和块级缩放
  • 通过在更小的元素组上操作,论文的方法有效地在这些分组元素之间共享指数位,从而减轻了有限动态范围的影响
Online Quantization
  • 张量级量化框架 (2022c; 2023b) 中采用了延迟量化 (delayed quantization),它维护先前迭代的最大绝对值历史记录以推断当前值
  • 为了确保准确的缩放因子并简化框架,论文为每个 1x128 激活 tile 或 128x128 权重块在线计算最大绝对值
  • 基于此,论文推导出缩放因子,然后将激活或权重在线量化为 FP8 格式
Low-Precision Storage and Communication
  • 结合论文的 FP8 训练框架,论文通过将缓存的激活和优化器状态压缩成更低精度的格式,进一步减少了内存消耗和通信开销
  • 低精度优化器状态 (Low-Precision Optimizer States)
    • 论文采用 BF16 数据格式而不是 FP32 来跟踪 AdamW (2017) 优化器中的一阶矩和二阶矩,而不会引起可观察到的性能下降
    • 然而,主权重(由优化器存储)和梯度(用于批大小累加)仍然保留在 FP32 中,以确保整个训练过程中的数值稳定性
  • 低精度激活 (Low-Precision Activation)
    • 如图 6 所示,Wgrad 操作以 FP8 精度执行
    • 为了减少内存消耗,一个自然的选择是以 FP8 格式缓存激活,用于线性算子的反向传播。然而,论文对几个算子进行了特殊考虑,以实现低成本的高精度训练:
      • (1) 注意力算子后的线性算子的输入
        • 这些激活也用于注意力算子的反向传播,这使得它对精度敏感
        • 论文为这些激活专门采用了一种定制的 E5M6 数据格式
        • 此外,这些激活在反向传播中将从 1x128 量化 tile 转换为 128x1 tile
        • 为了避免引入额外的量化误差,所有的缩放因子都是 2 的整数幂舍入缩放
      • (2) MoE 中 SwiGLU 算子的输入
        • 为了进一步降低内存成本,论文缓存 SwiGLU 算子的输入,并在反向传播中重新计算其输出
        • 这些激活也使用论文的细粒度量化方法以 FP8 格式存储,在内存效率和计算精度之间取得了平衡
  • 低精度通信 (Low-Precision Communication)
    • 通信带宽是 MoE 模型训练的关键瓶颈
    • 为了缓解这一挑战,论文在 MoE 上投影之前将激活量化为 FP8,然后应用分发组件,这与 MoE 上投影中的 FP8 Fprop 兼容
      • 与注意力算子后的线性算子的输入类似,此激活的缩放因子是 2 的整数幂
      • 类似的策略应用于 MoE 下投影之前的激活梯度
    • 对于前向和后向的合并组件,论文将它们保留在 BF16 中,以在训练管道的关键部分保持训练精度

Inference and Deployment

  • 论文将 DeepSeek-V3 部署在 H800 集群上,其中每个节点内的 GPU 使用 NVLink 互连,集群中的所有 GPU 通过 IB 完全互连
  • 为了同时确保在线服务的服务水平目标 (Service-Level Objective, SLO) 和高吞吐量,论文采用了以下将预填充 (prefilling) 和解码 (decoding) 阶段分离的部署策略
Prefilling
  • 预填充阶段的最小部署单元由 4 个节点(32 个 GPU)组成
  • 注意力部分采用 4 路张量并行 (Tensor Parallelism, TP4) 结合序列并行 (Sequence Parallelism, SP),以及 8 路数据并行 (Data Parallelism, DP8)
  • 其较小的 TP 规模(4)限制了 TP 通信的开销
  • 对于 MoE 部分,论文使用 32 路专家并行 (Expert Parallelism, EP32),这确保了每个专家处理足够大的批大小,从而提高了计算效率
  • 对于 MoE All-to-All 通信,论文使用与训练相同的方法:
    • 首先通过 IB 跨节点传输 Token ,然后通过 NVLink 在节点内 GPU 之间转发
    • 特别地,论文对浅层的稠密 MLP 使用 1 路张量并行以节省 TP 通信
  • 为了实现 MoE 部分中不同专家之间的负载平衡,论文需要确保每个 GPU 处理大致相同数量的 Token
  • 为此,论文引入了冗余专家 (redundant experts) 的部署策略,该策略复制高负载专家并冗余部署它们
    • 高负载专家是基于在线部署期间收集的统计信息检测出来的,并定期调整(例如,每 10 分钟)
    • 在确定了冗余专家集合后,论文根据观察到的负载,仔细地在节点内的 GPU 之间重新安排专家,力求在不增加跨节点 All-to-All 通信开销的情况下,尽可能平衡 GPU 间的负载
    • 对于 DeepSeek-V3 的部署,论文为预填充阶段设置了 32 个冗余专家
    • 对于每个 GPU,除了它原本承载的 8 个专家外,它还将承载一个额外的冗余专家
  • 在预填充阶段,为了提高吞吐量并隐藏 All-to-All 和 TP 通信的开销,论文同时处理两个计算工作量相似的微批次,将一个微批次的注意力和 MoE 与另一个微批次的分发和合并重叠起来
  • 最后,作者正在探索一种专家的动态冗余 (dynamic redundancy) 策略,其中每个 GPU 承载更多专家(例如,16 个专家),但在每次推理步骤中只激活 9 个
    • 在每层的 All-to-All 操作开始之前,论文实时计算全局最优的路由方案
    • 鉴于预填充阶段涉及大量计算,计算此路由方案的开销几乎可以忽略不计
Decoding
  • 在解码期间,论文将共享专家视为路由专家
    • 从这个角度来看,每个 Token 将在路由期间选择 9 个专家,其中共享专家被视为一个总是被选中的高负载专家
  • 解码阶段的最小部署单元由 40 个节点(320 个 GPU)组成
    • 注意力部分采用 TP4 结合 SP,以及 DP80,而 MoE 部分使用 EP320
    • 对于 MoE 部分,每个 GPU 仅承载一个专家,其中 64 个 GPU 负责承载冗余专家和共享专家
  • 分发和合并部分的 All-to-All 通信通过 IB 上的直接点对点传输进行,以实现低延迟
    • 论文还利用 IBSDA (2022a) 技术进一步最小化延迟并提高通信效率
  • 与预填充类似,论文基于在线服务的统计专家负载,以一定的间隔定期确定冗余专家集合
    • 但论文不需要重新安排专家,因为每个 GPU 只承载一个专家
  • 论文也在探索用于解码的动态冗余策略。然而,这需要更仔细地优化计算全局最优路由方案的算法以及与分发内核的融合以减少开销
  • 为了提高吞吐量并隐藏 All-to-All 通信的开销,论文也在探索在解码阶段同时处理两个计算工作量相似的微批次
    • 与预填充不同,注意力在解码阶段消耗的时间比例更大
    • 因此,论文将一个微批次的注意力与另一个微批次的(分发+MoE+合并)重叠起来
    • 在解码阶段,每个专家的批大小相对较小(通常在 256 个 Token 以内),瓶颈是内存访问而非计算
    • 由于 MoE 部分只需要加载一个专家的参数,内存访问开销很小,因此使用较少的 SMs 不会显著影响整体性能
    • 因此,为了避免影响注意力部分的计算速度,我们可以只分配一小部分 SMs 给(分发+MoE+合并)

Suggestions on Hardware Design

  • 基于论文 All-to-All 通信和 FP8 训练方案的实现,论文向 AI 硬件供应商提出以下芯片设计建议
Communication Hardware
  • 在 DeepSeek-V3 中,论文实现了计算和通信之间的重叠 ,以在计算期间隐藏通信延迟
    • 与串行计算和通信相比,这显著降低了对通信带宽的依赖
  • 但当前的通信实现依赖于昂贵的 SMs(例如,论文分配了 H800 GPU 中可用的 132 个 SMs 中的 20 个用于此目的),这将限制计算吞吐量
    • 而且使用 SMs 进行通信会导致显著的效率低下,因为 Tensor Cores 完全未被充分利用
  • 目前,SMs 主要为 All-to-All 通信执行以下任务:
    • 在 IB (InfiniBand) 和 NVLink 域之间转发数据 ,同时聚合来自单个 GPU、目的地为同一节点内多个 GPU 的 IB 流量
    • 在 RDMA 缓冲区(已注册的 GPU 内存区域)和输入/输出缓冲区之间传输数据
    • 执行归约操作以进行 All-to-All 合并
    • 在跨 IB 和 NVLink 域的分块数据传输到多个专家期间,管理细粒度的内存布局
  • 论文期望未来的供应商能够开发出将这些通信任务从宝贵的计算单元 SM 上卸载的硬件,作为 GPU 协处理器或类似 NVIDIA SHARP (2016) 的网络协处理器
    • 此外,为了降低应用程序编程的复杂性,作者希望这种硬件能从计算单元的角度统一 IB(横向扩展)和 NVLink(纵向扩展)网络
    • 通过这种统一的接口,计算单元可以基于简单的原语提交通信请求,轻松地在整个 IB-NVLink 统一域中完成诸如读、写、多播和归约等操作
Compute Hardware
  • Tensor Cores 中更高的 FP8 GEMM 累加精度 (Higher FP8 GEMM Accumulation Precision in Tensor Cores)
    • 在 NVIDIA Hopper 架构当前的 Tensor Core 实现中,FP8 GEMM 受到有限累加精度的困扰
    • 在基于最大指数通过右移对齐 32 个尾数乘积后,Tensor Core 仅使用每个尾数乘积的最高 14 位进行加法,并截断超出此范围的位(加法结果累加到寄存器中也采用 14 位精度)
    • 论文的实现通过将 128 次 FP8×FP8 乘法的加法结果累加到 CUDA core 中具有 FP32 精度的寄存器中,部分缓解了这一限制
    • 尽管这有助于实现成功的 FP8 训练,但这仅仅是由于 Hopper 架构在 FP8 GEMM 累加精度方面的硬件缺陷而做出的妥协
    • 未来的芯片需要采用更高的精度
  • 支持 tile 和块级量化 (Support for Tile- and Block-Wise Quantization)
    • 当前的 GPU 仅支持每张量 (per-tensor) 量化,缺乏对论文 tile 和块级量化等细粒度量化的原生支持
    • 在当前实现中,当达到 \(N_{C}\) 间隔时,部分结果将从 Tensor Cores 复制到 CUDA cores,乘以缩放因子,并添加到 CUDA cores 上的 FP32 寄存器中
    • 尽管结合论文精确的 FP32 累加策略,反量化开销得到了显著缓解,但 Tensor Cores 和 CUDA cores 之间频繁的数据移动仍然限制了计算效率
    • 论文建议未来的芯片通过使 Tensor Cores 能够接收缩放因子并实现具有组缩放的 MMA 来支持细粒度量化
    • 通过这种方式,整个部分和累加和反量化可以直接在 Tensor Cores 内部完成,直到产生最终结果,从而避免频繁的数据移动
  • 支持在线量化 (Support for Online Quantization)
    • 尽管论文的研究证明了在线量化的有效性,但当前的实现难以有效支持它
    • 在现有流程中,论文需要从 HBM(高带宽内存)中读取 128 个 BF16 激活值(先前计算的输出)进行量化,然后量化后的 FP8 值写回 HBM,只是为了再次读取用于 MMA
    • 为了解决这种低效率问题,论文建议未来的芯片将 FP8 转换和 TMA(Tensor Memory Accelerator)访问融合到单个融合操作中,这样量化可以在激活从全局内存传输到共享内存的过程中完成,避免频繁的内存读写
    • 论文还建议支持 warp 级的转换指令以加速,这进一步促进了层归一化和 FP8 转换的更好融合
      • 或者,可以采用近内存计算方法,将计算逻辑放置在 HBM 附近
    • 在这种情况下,BF16 元素在从 HBM 读入 GPU 时可以直接转换为 FP8,将片外内存访问减少大约 50%
  • 支持转置 GEMM 操作 (Support for Transposed GEMM Operations)
    • 当前的架构使得将矩阵转置与 GEMM 操作融合变得很麻烦
    • 在论文的工作流程中,前向传播期间的激活被量化为 1x128 的 FP8 tile 并存储
    • 在后向传播期间,矩阵需要被读出、反量化、转置、重新量化为 128x1 tile ,并存储在 HBM 中
    • 为了减少内存操作,论文建议未来的芯片能够在 MMA 操作之前直接从共享内存中对矩阵进行转置读取,以支持训练和推理中所需的那些精度
    • 结合 FP8 格式转换和 TMA 访问的融合,这一增强将显著简化量化工作流程

Pre-Training

Data Construction

  • 与 DeepSeek-V2 相比,论文通过提高数学和编程样本的比例来优化预训练语料库,同时将多语言覆盖范围扩展到英语和中文之外
  • 论文的数据处理流程经过改进,在保持语料库多样性的同时最大限度地减少了冗余
  • 受 (2024) 的启发,论文实施了文档打包方法(Document Packing Method)以保证数据完整性,但在训练期间并未引入跨样本注意力掩码
  • 最终,DeepSeek-V3 的训练语料库包含 14.8T 个高质量且多样化的 Token(使用论文的分词器)
  • 在 DeepSeekCoder-V2 (2024) 的训练过程中,论文观察到 Fill-in-Middle (FIM) 策略在使模型能够根据上下文线索准确预测中间文本的同时,并不会损害其下一个 Token 预测能力
    • 与 DeepSeekCoder-V2 保持一致,论文也在 DeepSeek-V3 的预训练中纳入了 FIM 策略
    • 具体来说,论文采用 Prefix-Suffix-Middle (PSM, Prefix, Suffix, Middle) 框架来结构化数据,如下所示:
      $$ \text{<|fim_begin|>} f_{\text{pre} } \text{<|fim_hole|>} f_{\text{suff} } \text{<|fim_end|>} f_{\text{middle} } \text{<|eos_token|>} $$
    • 该结构在文档级别应用,作为预打包过程的一部分
    • FIM 策略的应用率为 0.1,与 PSM 框架保持一致
    • 注:Fill-in-Middle(FIM,中间填充)是大语言模型的一种训练目标,核心是让模型根据上下文的 “前缀” 和 “后缀” 信息,预测并补全中间缺失的内容;这种设计旨在增强模型对文本全局逻辑的理解能力,尤其适用于需要双向参考上下文的场景
  • DeepSeek-V3 的分词器采用字节级 BPE (1999),并扩展了 128K Token 的词汇表
    • 论文对分词器的预分词器和训练数据进行了修改,以优化多语言压缩效率
    • 与 DeepSeek-V2 相比,新的预分词器引入了结合标点符号和换行符的 Token
      • 然而,当模型处理没有终止换行符的多行提示(特别是少样本评估提示)时,这种技巧可能会引入 Token 边界偏差 (2023)
      • 为了解决这个问题,论文在训练期间随机拆分一定比例的此类组合 Token ,使模型接触到更广泛的特例,从而减轻这种偏差

Hyper-Parameters

  • 模型超参数 (Model Hyper-Parameters)
    • 论文将 Transformer 层数设置为 61,隐藏维度设置为 7168
    • 所有可学习参数均使用标准差为 0.006 进行随机初始化
    • 在 MLA 中,论文将注意力头数 \(n_{h}\) 设置为 128,每个头的维度 \(d_{h}\) 设置为 128
    • KV 压缩维度 \(d_{c}\) 设置为 512, Query 压缩维度 \(d^{\prime}_{c}\) 设置为 1536
    • 对于解耦的 Query 和 Key ,论文将每个头的维度 \(d^{R}_{h}\) 设置为 64
    • 论文将除前三层之外的所有 FFN 替换为 MoE 层
    • 每个 MoE 层包含 1 个共享专家和 256 个路由专家,其中每个专家的中间隐藏维度为 2048
    • 在路由专家中,每个 Token 将激活 8 个专家,并且确保每个 Token 最多被发送到 4 个节点
    • 多 Token 预测深度 \(D\) 设置为 1 ,即除了精确的下一个 Token 外,每个 Token 还会预测一个额外的 Token
    • 与 DeepSeek-V2 一样,DeepSeek-V3 也在压缩潜在向量之后采用了额外的 RMSNorm 层,并在宽度瓶颈处乘以额外的缩放因子
    • 总结:在此配置下,DeepSeek-V3 总共包含 671B 参数,其中每个 Token 激活 37B 参数
  • 训练超参数 (Training Hyper-Parameters)
    • 采用 AdamW 优化器 (2017),超参数设置为 \(\beta_{1}=0.9\),\(\beta_{2}=0.95\),权重衰减 = 0.1
    • 在预训练期间,论文将最大序列长度设置为 4K,并在 14.8T Token 上对 DeepSeek-V3 进行预训练
    • 学习率调度方面
      • 首先,在前 2K 步期间将其从 \(0\) 线性增加到 \(2.2\times 10^{-4}\)
      • 然后,论文保持 \(2.2\times 10^{-4}\) 的恒定学习率,直到模型消耗完 10T 训练 Token
      • 随后,论文在 4.3T Token 内,按照余弦衰减曲线将学习率逐渐衰减到 \(2.2\times 10^{-5}\)
      • 在最后 500B Token 的训练期间,论文在前 333B Token 中保持 \(2.2\times 10^{-5}\) 的恒定学习率,并在剩余的 167B Token 中切换到另一个恒定学习率 \(7.3\times 10^{-6}\)
    • 梯度裁剪范数设置为 1.0
    • 论文采用批量大小调度策略(batch size scheduling strategy)
      • 在前 469B Token 的训练中,批量大小从 3072 逐渐增加到 15360,然后在剩余训练中保持 15360
    • 论文利用流水线并行将模型的不同层部署在不同的 GPU 上,对于每一层,路由专家将均匀部署在属于 8 个节点的 64 个 GPU 上
      • 问题:这里如何理解?
    • 至于节点限制路由 ,每个 Token 最多被发送到 4 个节点(即 \(M=4\))
    • 对于无辅助损失的负载平衡,论文在前 14.3T Token 中将偏置更新速度 \(\gamma\) 设置为 0.001,在剩余的 500B Token 中设置为 0.0
    • 对于平衡损失,论文将 \(\alpha\) 设置为 0.0001,仅用于避免任何单个序列内的极端不平衡
    • MTP 损失权重 \(\lambda\) 在前 10T Token 中设置为 0.3,在剩余的 4.8T Token 中设置为 0.1

Long Context Extension

  • 论文采用与 DeepSeek-V2 (2024) 类似的方法来使 DeepSeek-V3 具备长上下文能力
  • 在预训练阶段之后,论文应用 YaRN (2023) 进行上下文扩展,并执行两个额外的训练阶段,每个阶段包含 1000 步,以逐步将上下文窗口从 4K 扩展到 32K,然后再扩展到 128K
  • YaRN 配置与 DeepSeek-V2 中使用的配置一致,仅应用于解耦的共享 Key \(\mathbf{k}_{t}^{R}\)
  • 两个阶段的超参数保持相同
    • 尺度 \(s=40\),\(\alpha=1\),\(\beta=32\),缩放因子 \(\sqrt{t}=0.1\ln s+1\)
    • 在第一阶段,序列长度设置为 32K,批量大小为 1920
    • 在第二阶段,序列长度增加到 128K,批量大小减少到 480
  • 两个阶段的学习率均设置为 \(7.3\times 10^{-6}\),与预训练阶段的最终学习率相匹配
  • 通过这种两阶段的扩展训练,DeepSeek-V3 能够处理长达 128K 的输入,同时保持强大的性能
  • 图 8 显示,经过监督微调后,DeepSeek-V3 在 “Needle In A Haystack” (NIAH) 测试中取得了显著性能,在上下文窗口长度高达 128K 的范围内表现出一致的鲁棒性

Evaluations

Evaluation Benchmarks
  • DeepSeek-V3 的基础模型是在英语和中文占多数的多语言语料库上进行预训练的,因此论文在一系列主要以英语和中文为主的基准测试上评估其性能,同时也包括一个多语言基准测试
  • 论文的评估基于集成在论文 HAI-LLM 框架中的内部评估框架
  • 所考虑的基准测试分类并列出如下,其中带下划线的基准测试为中文,带双下划线的基准测试为多语言:
    • 多学科多项选择题 (Multi-subject multiple-choice) 数据集包括 MMLU (2020)、MMLU-Redux (2024)、MMLU-Pro (2024)、MMMLU (2024)、C-Eval (2023) 和 CMMLU (2023)
    • 语言理解和推理 (Language understanding and reasoning) 数据集包括 HellaSwag (2019)、PIQA (2020)、ARC (2018) 和 BigBench Hard (BBH) (2022)
    • 闭卷问答 (Closed-book question answering) 数据集包括 TriviaQA (2017) 和 NaturalQuestions (2019)
    • 阅读理解 (Reading comprehension) 数据集包括 RACE (2017)、DROP (2019)、C3 (2019) 和 CMRC (2019)
    • 指代消解 (Reference disambiguation) 数据集包括 CLUEWSC (2020) 和 WinoGrande (2019)
    • 语言建模 (Language modeling) 数据集包括 Pile (2020)
    • 中文理解与文化 (Chinese understanding and culture) 数据集包括 CCPM (2021)
    • 数学 (Math) 数据集包括 GSM8K (2021)、MATH (2021)、MGSM (2023) 和 CMath (2023)
    • 代码 (Code) 数据集包括 HumanEval (2021)、LiveCodeBench-Base (0801-1101) (2024)、MBPP (2021) 和 CRUXEval (2024)
    • 标准化考试 (Standardized exams) 包括 AGIEval (2023)。注意,AGIEval 包含英语和中文子集
  • 遵循论文之前的工作 (2024, 2024)
    • 对包括 HellaSwag、PIQA、WinoGrande、RACE-Middle、RACE-High、MMLU、MMLU-Redux、MMLU-Pro、MMMLU、ARC-Easy、ARC-Challenge、C-Eval、CMMLU、C3 和 CCPM 在内的数据集采用基于困惑度的评估;
    • 对 TriviaQA、NaturalQuestions、DROP、MATH、GSM8K、MGSM、HumanEval、MBPP、LiveCodeBench-Base、CRUXEval、BBH、AGIEval、CLUEWSC、CMRC 和 CMath 采用基于生成的评估
    • 此外,论文对 Pile-test 执行基于语言建模的评估,并使用 Bits-Per-Byte (BPB) 作为度量标准,以保证使用不同分词器的模型之间的公平比较
      • 理解:Pile-test 是大语言模型评估中常用的标准测试集,源于更庞大的通用文本数据集 The Pile
        • The Pile 由 EleutherAI 构建的大规模开源文本数据集,总规模约 800GB,涵盖 22 个不同来源的文本类型(如学术论文、网页文本、书籍、新闻、代码等),旨在为模型提供多样化、高质量的训练与评估数据,避免单一数据分布导致的 “过拟合评估”
        • Pile-test 是 The Pile 的测试子集,与训练集(Pile-train)严格划分,用于客观衡量模型在通用语言理解与生成任务上的泛化能力,由于其覆盖场景广,模型在 Pile-test 上的表现能更真实反映 “通用能力”,而非仅适配某类特定数据
      • 问题:这里 BPB 是什么含义?
      • 回答:Bits-Per-Byte(字节每比特,简称 BPB)是一种用于消除 “分词器差异” 影响、实现不同模型公平比较的度量标准,其核心是 “归一化模型的预测成本
        • 比如,计算模型的困惑度时,如果按照 Token 计算,可能是不公平的,不同模型的 Token 数量不一样,但是如果按照 BPB 计算,则能保证与 Tokenizer 无关
Evaluation Results
  • 在表 3 中,论文将 DeepSeek-V3 的基础模型与 SOTA 开源基础模型进行了比较,包括 DeepSeek-V2-Base (2024)(论文之前的发布)、Qwen2.5 72B Base (2024) 和 LLaMA-3.1 405B Base (2024)
  • 论文在内部评估框架下评估所有这些模型,并确保它们共享相同的评估设置
  • 请注意,由于过去几个月论文评估框架的变化,DeepSeek-V2-Base 的性能与论文之前报告的结果略有不同
  • 总体而言,DeepSeek-V3-Base 全面超越了 DeepSeek-V2-Base 和 Qwen2.5 72B Base,并在大多数基准测试中超过了 LLaMA-3.1 405B Base,基本上成为最强的开源模型
  • 从更详细的角度看,论文将 DeepSeek-V3-Base 与其他开源基础模型进行了单独比较
    • (1) 与 DeepSeek-V2-Base 相比,由于论文模型架构的改进、模型规模和训练 Token 的扩大以及数据质量的提高,DeepSeek-V3-Base 如预期般取得了显著更好的性能
    • (2) 与 SOTA 中文开源模型 Qwen2.5 72B Base 相比,DeepSeek-V3-Base 仅以一半的激活参数量,也在英语、多语言、代码和数学基准测试上展现出了显著优势,尤其是在中文基准测试上,除了中文多学科多项选择题任务 CMMLU 外,DeepSeek-V3-Base 也显示出比 Qwen2.5 72B 更好的性能
    • (3) 与最大的开源模型 LLaMA-3.1 405B Base(其激活参数量是 DeepSeek-V3-Base 的 11 倍)相比,DeepSeek-V3-Base 在多语言、代码和数学基准测试上也表现出更优的性能
    • 对于英语和中文语言基准测试,DeepSeek-V3-Base 显示出竞争性或更好的性能,尤其在 BBH、MMLU 系列、DROP、C-Eval、CMMLU 和 CCPM 上表现突出
  • 得益于论文高效的架构和全面的工程优化,DeepSeek-V3 实现了极高的训练效率
    • 在论文的训练框架和基础设施下,训练 DeepSeek-V3 每万亿 Token 仅需 180K H800 GPU 小时,这比训练 72B 或 405B 的稠密模型要便宜得多

Discussion

Ablation Studies for Multi-token Prediction
  • 在表 4 中,论文展示了 MTP 策略的消融结果
  • 具体来说,论文在两个不同规模的基线模型上验证了 MTP 策略
    • 在小规模上,论文在 1.33T Token 上训练了一个包含 15.7B 总参数的基线 MoE 模型
    • 在大规模上,论文在 540B Token 上训练了一个包含 228.7B 总参数的基线 MoE 模型
  • 在此基础上,保持训练数据和其他架构不变,论文为它们附加了一个深度为 1 的 MTP 模块,并训练了两个采用 MTP 策略的模型进行比较
  • 注意,在推理期间,论文直接丢弃 MTP 模块,因此比较模型的推理成本完全相同
  • 从表中我们可以观察到,MTP 策略在大多数评估基准上持续提升了模型性能
Ablation Studies for the Auxiliary-Loss-Free Balancing Strategy
  • 在表 5 中,论文展示了无辅助损失平衡策略的消融结果
  • 论文在两个不同规模的基线模型上验证了该策略
    • 在小规模上,论文在 1.33T Token 上训练了一个包含 15.7B 总参数的基线 MoE 模型
    • 在大规模上,论文在 578B Token 上训练了一个包含 228.7B 总参数的基线 MoE 模型
  • 这两个基线模型都纯粹使用辅助损失来鼓励负载平衡,并使用 sigmoid 门控函数和 top-K 亲和度归一化
    • 它们控制辅助损失强度的超参数分别与 DeepSeek-V2-Lite 和 DeepSeek-V2 相同
  • 在这两个基线模型的基础上,保持训练数据和其他架构不变,论文移除了所有辅助损失,并引入了无辅助损失平衡策略进行比较
  • 从表中我们可以观察到,与纯辅助损失方法相比,无辅助损失策略在大多数评估基准上持续实现了更好的模型性能
Batch-Wise Load Balance VS. Sequence-Wise Load Balance
  • 无辅助损失平衡与序列级辅助损失之间的关键区别在于它们的平衡范围:批次级与序列级
    • 与序列级辅助损失相比,批次级平衡施加了更灵活的约束,因为它不强制每个序列在域内平衡
    • 这种灵活性允许专家更好地专精于不同领域
  • 为了验证这一点,论文记录并分析了基于辅助损失的 16B 基线模型和无辅助损失的 16B 模型在 Pile 测试集上不同领域的专家负载
  • 如图 9 所示,论文观察到无辅助损失模型如预期那样表现出更大的专家专精模式(注:也就是说有辅助损失的模型得到的专家负载均衡结果更加平均)
  • 为了进一步研究这种灵活性与模型性能优势之间的相关性,论文额外设计并验证了一种批次级辅助损失,该损失鼓励每个训练批次而不是每个序列上的负载平衡
  • 实验结果表明,当达到相似水平的批次级负载平衡时,批次级辅助损失也可以实现与无辅助损失方法相似的模型性能
  • 具体来说,在论文使用 1B MoE 模型的实验中,验证损失分别为:2.258(使用序列级辅助损失)、2.253(使用无辅助损失方法)和 2.253(使用批次级辅助损失)
  • 论文在 3B MoE 模型上也观察到了类似的结果:使用序列级辅助损失的模型验证损失为 2.085,而使用无辅助损失方法或批次级辅助损失的模型验证损失均为 2.080
  • 此外,尽管批次级负载平衡方法显示出持续的性能优势,但它们在效率方面也面临两个潜在的挑战:
    • (1) 某些序列内或小批次内的负载不平衡
    • (2) 推理期间由领域偏移引起的负载不平衡
  • 第一个挑战通过论文使用大规模专家并行和数据并行的训练框架自然得到解决,这保证了每个微批次的规模足够大
  • 对于第二个挑战,论文也设计并实现了一个具有冗余专家部署的高效推理框架,如第 3.4 节所述,以克服它

Post-Training

Supervised Fine-Tuning

  • 论文精心策划了论文的指令微调数据集,包含了涵盖多个领域的 150 万条实例,每个领域都采用了针对其特定需求而定制的不同数据创建方法
  • 推理数据 (Reasoning Data)
    • 对于推理相关的数据集,包括那些专注于数学、代码竞赛问题和逻辑谜题的,论文通过利用内部的 DeepSeek-R1 模型来生成数据
      • 问题:DeepSeek-R1 的一些流程依赖着 DeepSeek-V3 添加一些 CoT 数据吧,两者之间目前从论文看起来是互相依赖的关系
    • 具体来说,虽然 R1 生成的数据表现出很高的准确性,但它存在一些问题,如过度思考、格式不佳和长度过长
    • 论文的目标是在 R1 生成的高精度推理数据和常规格式的清晰简洁的推理数据之间取得平衡
    • 为了建立论文的方法,论文首先为特定领域(例如代码、数学或通用推理)开发一个专家模型,该模型使用结合了 SFT 和 RL 的训练流程
      • 这个专家模型作为最终模型的数据生成器
    • 训练过程涉及为每个实例生成两种不同类型的 SFT 样本:
      • 第一种将问题与其原始回答配对,格式为 <问题, 原始回答>;
      • 第二种则结合了系统提示、问题以及 R1 的回答,格式为 <系统提示, 问题, R1 回答>
    • 系统提示经过精心设计,包含指导模型产生富含反思和验证机制的回答的指令
      • 在 RL 阶段,模型利用高温采样来生成融合了 R1 生成数据和原始数据模式的回答,即使在没有显式系统提示的情况下也是如此。经过数百个 RL 步骤后,中间的 RL 模型学会了融入 R1 的模式,从而战略性地提升了整体性能
    • 在完成 RL 训练阶段后,论文实施拒绝采样来为最终模型筛选高质量的 SFT 数据,其中专家模型被用作数据源
      • 这种方法确保了最终训练数据保留了 DeepSeek-R1 的优势,同时产生的回答简洁有效
  • 非推理数据 (Non-Reasoning Data)
    • 对于非推理数据,例如创意写作、角色扮演和简单问答,论文使用 DeepSeek-V2.5 来生成回答,并聘请人工标注员来验证数据的准确性和正确性
  • SFT 设置 (SFT Settings)
    • 使用 SFT 数据集对 DeepSeek-V3-Base 进行了两个 epoch 的微调
    • 使用了余弦衰减学习率调度
    • 学习率从 \(5 \times 10^{-6}\) 开始,逐渐降低到 \(1 \times 10^{-6}\)
    • 在训练期间,每个单独的序列由多个样本打包而成
      • 论文采用了样本掩码策略来确保这些样本保持隔离且相互不可见

Reinforcement Learning

Reward Model
  • 论文在 RL 过程中使用了基于规则的奖励模型 (Reward Model, RM) 和基于模型的 RM
  • 基于规则的 RM (Rule-Based RM)
    • 对于可以使用特定规则验证的问题,论文采用基于规则的奖励系统来确定反馈
      • 例如,某些数学问题具有确定性的结果,论文要求模型以指定的格式(例如,在一个框内)提供最终答案,从而允许论文应用规则来验证正确性
      • 类似地,对于 LeetCode 问题,我们可以利用编译器根据测试用例生成反馈
    • 通过在可能的情况下利用基于规则的验证,论文确保了更高水平的可靠性,因为这种方法能够抵抗操纵或利用(manipulation or exploitation)
  • 基于模型的 RM (Model-Based RM)
    • 对于具有自由形式标准答案的问题,论文依赖奖励模型来确定回答是否符合预期的标准答案
    • 相反,对于没有明确标准答案的问题,例如涉及创意写作的问题,奖励模型的任务是基于问题和相应的回答作为输入来提供反馈
    • 该奖励模型是从 DeepSeek-V3 SFT 检查点训练而来的
    • 为了增强其可靠性,论文构建了不仅提供最终奖励,还包含导致该奖励的思维链 (Chain-of-Thought) 的偏好数据
    • 这种方法有助于减轻特定任务中 Reward Hacking 的风险
GRPO (Group Relative Policy Optimization)
  • 与 DeepSeek-V2 (2024c) 类似,论文采用 GRPO(2024),它摒弃了通常与策略模型大小相同的评论家模型 (Critic Model),而是从组分数中估计基线
  • 具体来说,对于每个问题 \(q\),GRPO 从旧策略模型 \(\pi_{\theta_{old} }\) 中采样一组输出 \(\{o_{1}, o_{2}, \cdots, o_{G}\}\),然后通过最大化以下目标来优化策略模型 \(\pi_{\theta}\):
    $$
    \mathcal{J}_{GRPO}(\theta) = \mathbb{E} \left[ q \sim P(Q), \{o_{i}\}_{i=1}^{G} \sim \pi_{\theta_{old} }(O|q) \right] \frac{1}{G} \sum_{i=1}^{G} \left( \min \left( \frac{\pi_{\theta}(o_{i}|q)}{\pi_{\theta_{old} }(o_{i}|q)} A_{i}, \text{clip} \left( \frac{\pi_{\theta}(o_{i}|q)}{\pi_{\theta_{old} }(o_{i}|q)}, 1-\epsilon, 1+\epsilon \right) A_{i} \right) - \beta \mathbb{D}_{KL} \left( \pi_{\theta} || \pi_{ref} \right) \right), \\
    \mathbb{D}_{KL} \left( \pi_{\theta} || \pi_{ref} \right) = \frac{\pi_{ref}(o_{i}|q)}{\pi_{\theta}(o_{i}|q)} - \log \frac{\pi_{ref}(o_{i}|q)}{\pi_{\theta}(o_{i}|q)} - 1,
    $$
    • 其中 \(\epsilon\) 和 \(\beta\) 是超参数;\(\pi_{ref}\) 是参考模型;\(A_{i}\) 是优势 (Advantage),由每个组内输出对应的奖励 \(\{r_{1}, r_{2}, \ldots, r_{G}\}\) 推导得出:
      $$
      A_{i} = \frac{r_{i} - \text{mean}(\{r_{1}, r_{2}, \cdots, r_{G}\})}{\text{std}(\{r_{1}, r_{2}, \cdots, r_{G}\})}.
      $$
  • 论文在 RL 过程中融入了来自不同领域的提示,例如编码、数学、写作、角色扮演和问答
    • 这种方法不仅使模型更符合人类偏好,还提高了在基准测试上的性能,特别是在可用 SFT 数据有限的情况下

Evaluations

Evaluation Settings
  • 评估基准 (Evaluation Benchmarks)
    • 除了用于基础模型测试的基准之外,论文还在 IFEval (2023)、FRAMES (2024)、LongBench v2 (2024)、GPQA (2023)、SimpleQA (OpenAI, 2024c)、C-SimpleQA (2024)、SWE-Bench Verified (OpenAI, 2024d)、Aider、LiveCodeBench (2024)(2024年8月至11月的问题)、Codeforces、中国高中数学奥林匹克 (Chinese National High School Mathematics Olympiad, CNMO 2024) 和美国数学邀请赛 2024 (American Invitational Mathematics Examination 2024, AIME 2024) (MAA, 2024) 上进一步评估了指令模型
  • 对比基线 (Compared Baselines)
    • 论文对论文的聊天模型与几个强大的基线进行了全面评估,包括 DeepSeek-V2-0506、DeepSeek-V2.5-0905、Qwen2.5 72B Instruct、LLaMA-3.1 405B Instruct、Claude-Sonnet-3.5-1022 和 GPT-4o-0513。对于 DeepSeek-V2 模型系列,论文选择了最具代表性的变体进行比较
    • 对于闭源模型,通过它们各自的 API 进行评估
  • 详细评估配置 (Detailed Evaluation Configurations)
    • 对于包括 MMLU、DROP、GPQA 和 SimpleQA 在内的标准基准,论文采用了来自 simple-evals 框架的评估提示
      • 注:simple-evals 是一个轻量级的大语言模型评估框架,主要用于快速验证模型在特定任务上的基础能力,其Evaluation Prompts设计遵循简洁、直接的任务导向原则,通常针对具体任务类型构建标准化提示模板
    • 对于 MMLU-Redux,论文在零样本设置下使用了 Zero-Eval 提示格式 (Lin, 2024)
    • 对于其他数据集,论文遵循其原始的评估协议,使用数据集创建者提供的默认提示
    • 对于代码和数学基准,HumanEval-Mul 数据集总共包含了 8 种主流编程语言(Python、Java、C++、C#、JavaScript、TypeScript、PHP 和 Bash)
      • 使用 CoT 和非 CoT 方法来评估模型在 LiveCodeBench 上的性能,其数据收集自 2024 年 8 月至 11 月
      • Codeforces 数据集使用参赛者百分比来衡量
      • SWE-Bench verified 使用无代理框架 (Agentless Framework) (2024) 进行评估
      • 论文使用 “diff” 格式来评估与 Aider 相关的基准
    • 对于数学评估,AIME 和 CNMO 2024 在温度为 0.7 的情况下进行评估,结果取 16 次运行的平均值,而 MATH-500 则使用贪婪解码
    • 论文允许所有模型为每个基准输出最多 8192 个 Token
Standard Evaluation
  • 表 6 展示了评估结果,表明 DeepSeek-V3 是性能最好的开源模型
  • 此外,它与前沿的闭源模型(如 GPT-4o 和 Claude-3.5-Sonnet)相比也具有竞争力
  • 英文基准 (English Benchmarks)
    • MMLU 是一个广泛认可的基准,旨在评估大语言模型在不同知识领域和任务上的表现
      • DeepSeek-V3 展示了具有竞争力的性能,与顶级模型如 LLaMA-3.1-405B、GPT-4o 和 Claude-Sonnet 3.5 不相上下,同时显著优于 Qwen2.5 72B
      • DeepSeek-V3 在 MMLU-Pro(一个更具挑战性的教育知识基准)上表现出色,紧随 Claude-Sonnet 3.5 之后
      • 在 MMLU-Redux(一个带有修正标签的 MMLU 改进版本)上,DeepSeek-V3 超越了其同行
    • 在 GPQA-Diamond(一个博士级评估测试平台)上,DeepSeek-V3 取得了显著成果,仅次于 Claude 3.5 Sonnet,并以显著优势优于所有其他竞争对手
    • 在长上下文理解基准测试中,如 DROP、LongBench v2 和 FRAMES,DeepSeek-V3 继续展示其作为顶级模型的地位
      • DeepSeek-V3 在 DROP 的 3-shot 设置中取得了令人印象深刻的 91.6 F1 分数,超过了该类别中的所有其他模型
      • 在 FRAMES(一个需要在 100k Token 上下文中进行问答的基准测试)上,DeepSeek-V3 紧随 GPT-4o,同时以显著优势优于所有其他模型
        • 这证明了 DeepSeek-V3 在处理极长上下文任务方面的强大能力
      • DeepSeek-V3 的长上下文能力进一步通过其在 LongBench v2(一个在 DeepSeek V3 发布前几周才发布的数据集)上的最佳表现得到了验证
      • 在事实性知识基准 SimpleQA 上,DeepSeek-V3 落后于 GPT-4o 和 Claude-Sonnet,这主要归因于其设计重点和资源分配
      • DeepSeek-V3 分配了更多的训练 Token 来学习中文知识,导致在 C-SimpleQA 上表现卓越
      • 在指令遵循基准测试中,DeepSeek-V3 显著优于其前身 DeepSeek-V2 系列,突显了其理解和遵守用户定义格式约束能力的提升
  • 代码和数学基准 (Code and Math Benchmarks)
    • 编码对于大语言模型来说是一项具有挑战性且实用的任务,涵盖了以工程为重点的任务(如 SWE-Bench-Verified 和 Aider)以及算法任务(如 HumanEval 和 LiveCodeBench)
    • 在工程任务中,DeepSeek-V3 落后于 Claude-Sonnet-3.5-1022,但显著优于开源模型
      • 开源的 DeepSeek-V3 有望推动编码相关工程任务的进步
    • 通过提供对其强大功能的访问,DeepSeek-V3 可以推动软件工程和算法开发等领域的创新和改进,使开发人员和研究人员能够突破开源模型在编码任务中能力的界限
    • 在算法任务中,DeepSeek-V3 表现出卓越的性能,在 HumanEval-Mul 和 LiveCodeBench 等基准测试中优于所有基线
      • 这一成功归功于其先进的知识蒸馏技术,该技术有效地增强了其在算法任务中的代码生成和问题解决能力
    • 在数学基准测试中,DeepSeek-V3 展示了卓越的性能,显著超过了基线,并为非 o1 类模型设定了新的最先进水平
      • 具体来说,在 AIME、MATH-500 和 CNMO 2024 上,DeepSeek-V3 在绝对分数上比第二好的模型 Qwen2.5 72B 高出约 10%,这对于如此具有挑战性的基准测试来说是一个显著的差距
      • 这种卓越的能力突显了从 DeepSeek-R1 进行蒸馏技术的有效性,该技术已被证明对非 o1 类模型非常有益
  • 中文基准 (Chinese Benchmarks)
    • Qwen 和 DeepSeek 是两个对中文和英文都有强大支持的代表性模型系列
    • 在事实性基准 Chinese SimpleQA 上,DeepSeek-V3 超过了 Qwen2.5-72B 16.4 分,尽管 Qwen2.5 是在包含 18T Token (比 DeepSeek-V3 预训练的 14.8T Token 多 20%)的更大语料库上训练的
    • 在 C-Eval(一个代表性的中文教育知识评估基准)和 CLUEWSC(中文 Winograd 模式挑战赛)上,DeepSeek-V3 和 Qwen2.5-72B 表现出相似的水平,表明这两个模型都为具有挑战性的中文推理和教育任务进行了良好的优化
Open-Ended Evaluation
  • 除了标准基准测试,论文还使用大语言模型作为评判者对论文的模型在开放式生成任务上进行了评估,结果如表 7 所示
  • 论文遵循 AlpacaEval 2.0 (2024) 和 Arena-Hard (2024a) 的原始配置,它们利用 GPT-4-Turbo-1106 作为成对比较的评判者
  • 在 Arena-Hard 上,DeepSeek-V3 相对于基线 GPT-4-0314 取得了超过 86% 的惊人胜率,与 Claude-Sonnet-3.5-1022 等顶级模型表现相当
    • 这凸显了 DeepSeek-V3 的强大能力,尤其是在处理复杂提示(包括编码和调试任务)方面
  • DeepSeek-V3 实现了一个突破性的里程碑,成为第一个在 Arena-Hard 基准测试中超过 85% 的开源模型
    • 这一成就显著缩小了开源模型和闭源模型之间的性能差距,为开源模型在挑战性领域所能达到的水平设定了新标准
  • 类似地,DeepSeek-V3 在 AlpacaEval 2.0 上展示了卓越的性能,优于闭源和开源模型
    • 这证明了其在写作任务和处理直接问答场景方面的出色熟练度
  • 值得注意的是,它以 20% 的显著优势超过了 DeepSeek-V2.5-0905,突显了其在处理简单任务方面的实质性改进,并展示了其进步的有效性
DeepSeek-V3 as a Generative Reward Model
  • 论文将 DeepSeek-V3 的判断能力与 SOTA 模型(即 GPT-4o 和 Claude-3.5)进行了比较
  • 表 8 展示了这些模型在 RewardBench (2024) 中的性能(注:RewardBench 是用来评估 Reward Model 本身性能的)
  • DeepSeek-V3 达到了与最佳版本的 GPT-4o-0806 和 Claude-3.5-Sonnet-1022 相当的水平,同时超越了其他版本
  • 此外,DeepSeek-V3 的判断能力也可以通过投票技术得到增强
  • 因此,论文采用 DeepSeek-V3 并结合投票来为开放式问题提供自我反馈,从而提高对齐过程的有效性和鲁棒性

Discussion

Distillation from DeepSeek-R1
  • 论文基于 DeepSeek-V2.5 对从 DeepSeek-R1 进行蒸馏的贡献进行了消融研究
  • 基线是在短 CoT 数据上训练的,而其对比模型则使用上述专家检查点生成的数据
  • 表 9 展示了蒸馏数据的有效性,在 LiveCodeBench 和 MATH-500 基准测试上均显示出显著的改进
  • 论文的实验揭示了一个有趣的权衡:蒸馏带来了更好的性能,但也显著增加了平均回答长度
  • 为了在模型准确性和计算效率之间保持平衡,论文为 DeepSeek-V3 在蒸馏过程中仔细选择了最优设置
  • 论文的研究表明,从推理模型进行知识蒸馏为后训练优化提供了一个有前景的方向
    • 虽然论文目前的工作侧重于蒸馏数学和编码领域的数据,但这种方法显示出在需要复杂推理的各种任务领域中具有更广泛应用的潜力
    • 在这些特定领域展示的有效性表明,长 CoT 蒸馏对于增强其他认知任务中的模型性能可能很有价值
    • 在不同领域进一步探索这种方法仍然是未来研究的一个重要方向
Self-Rewarding
  • 奖励在强化学习中起着关键作用,引导着优化过程
  • 在通过外部工具易于验证的领域,例如某些编码或数学场景,强化学习表现出卓越的功效
    • 但在更一般的场景中,通过硬编码构建反馈机制是不切实际的
  • 在 DeepSeek-V3 的开发过程中,对于这些更广泛的上下文,论文采用了 Constitutional AI 方法 (2022),利用 DeepSeek-V3 自身的投票评估结果作为反馈源
    • 这种方法产生了显著的对齐效果,显著提升了 DeepSeek-V3 在主观评估中的性能
    • 通过整合额外的 Constitutional 输入,DeepSeek-V3 可以向 Constitutional 方向优化
  • 作者相信,这种将补充信息与大语言模型作为反馈源相结合的范式至关重要
    • 大语言模型作为一个通用的处理器,能够将来自不同场景的非结构化信息转化为奖励,最终促进大语言模型的自我改进
    • 除了自我奖励,论文还致力于发现其他通用和可扩展的奖励方法,以持续提升模型在一般场景中的能力
Multi-token Prediction Evaluation
  • DeepSeek-V3 不仅仅预测下一个单独的 Token ,而是通过 MTP 技术预测接下来的 2 个 Token
  • 结合推测解码 (Speculative Decoding) (2023; 2023) 的框架,它可以显著加速模型的解码速度
  • 一个自然的问题是,额外预测的 Token 的接受率如何
    • 根据论文的评估,在不同生成主题下,第二个 Token 预测的接受率在 85% 到 90% 之间,表现出一致的可靠性
    • 这种高接受率使得 DeepSeek-V3 能够实现显著提高的解码速度,提供 1.8 倍的 TPS(每秒 Token 数)

Limitations, and Future Directions

Limitations

  • 在承认其强大性能和成本效益的同时,论文也认识到 DeepSeek-V3 存在一些局限性,尤其是在部署方面
  • 首先,为了确保高效的推理,DeepSeek-V3 推荐的部署单元相对较大,这可能会给小型团队带来负担
  • 其次,尽管论文为 DeepSeek-V3 设计的部署策略已经实现了端到端生成速度超过 DeepSeek-V2 的两倍,但仍有进一步提升的潜力
  • 幸运的是,随着更先进硬件的发展,这些局限性有望得到自然解决

Future Directions

  • DeepSeek 始终坚持具有长期主义(longtermism)的开源模型路线,旨在稳步接近通用人工智能(Artificial General Intelligence, AGI)的终极目标
  • 未来,论文计划在以下几个方向进行战略性投入和研究:
    • DeepSeek 将持续研究和改进论文的模型架构,旨在进一步提高训练和推理效率,努力实现对无限上下文长度的高效支持
      • 此外,DeepSeek 将尝试突破 Transformer 的架构限制,从而推动其建模能力的边界
    • DeepSeek 将持续迭代训练数据的数量和质量,并探索纳入额外的训练信号源,旨在推动数据在更全面维度上的扩展
    • DeepSeek 将持续探索和迭代模型的深度思考能力,旨在通过扩展其推理长度和深度来增强其智能和问题解决能力
    • DeepSeek 将探索更全面、多维度的模型评估方法,以防止在研究过程中倾向于优化固定的基准测试集,这可能造成对模型能力的误导性印象并影响论文的基础评估

附录 B: Ablation Studies for Low-Precision Training

  • 图 10:BF16 与 FP8 训练的损失曲线比较(结果使用系数为 0.9 的指数移动平均(Exponential Moving Average, EMA)进行了平滑处理)

B.1: FP8 v.s. BF16 Training

  • 论文在两个不同规模的基线模型上验证了论文的 FP8 混合精度框架,并与 BF16 训练进行了比较
    • 在小规模上,论文训练了一个总参数量约为 16B 的 MoE 基线模型,使用了 1.33T token
    • 在大规模上,论文训练了一个总参数量约为 230B 的 MoE 基线模型,使用了约 0.9T token
  • 论文在图 10 中展示了训练曲线,并证明了通过论文的高精度累加和细粒度量化策略,相对误差保持在 0.25% 以下

B.2:Discussion About Block-Wise Quantization(分块量化)

  • 尽管论文的切片式(tile-wise)细粒度量化有效缓解了特征异常值(feature outliers)引入的误差,但它需要对激活量化进行不同的分组,即在正向传播中使用 1x128 的分组,在反向传播中则需要 128x1 的分组
    • 激活梯度也需要类似的处理过程
  • 一个直接的策略是像论文量化模型权重那样,对每 128x128 个元素应用分块(block-wise)量化
    • 这样,反向传播只需要进行转置操作
  • 论文进行了一项实验,将与 Dgrad 相关的所有张量都在分块基础上进行量化
    • 结果显示,计算激活梯度并以链式方式反向传播到浅层的 Dgrad 操作对精度高度敏感
  • 论文在总参数量约为 16B、训练了约 300B Token 的 MoE 模型上,激活梯度的分块量化会导致模型发散
    • 论文假设这种敏感性源于激活梯度在 Token 之间高度不平衡,导致了与 Token 相关的异常值 (2023)
    • 这些异常值无法通过分块量化方法有效管理

附录 C:16B 基于辅助损失和无辅助损失模型的专家专业化模式 (Expert Specialization Patterns of the 16B Aux-Loss-Based and Aux-Loss-Free Models)

  • 本节记录了 16B 基于辅助损失的基线模型和无辅助损失模型在 Pile 测试集上的专家负载
  • 如图 11 所示,无辅助损失模型在所有层中都倾向于表现出更强的专家专业化程度(Expert Specialization)
    • 问题:如何理解这里的 专家专业化程度(Expert Specialization)?
    • 理解:是指某些任务上,某类专家被激活的更多,而在另外的任务上,其他专家被激活的更多,专家体现出一定的专业化倾向(注:这不是我们想要的,因为会导致模型拟合能力下降,且单个序列上不利于 EP 负载均衡)
  • 图 11:无辅助损失模型和基于辅助损失的模型在 Pile 测试集的三个领域上的专家负载
    • 无辅助损失模型 比 基于辅助损失模型 表现出 更强的专家专业化模式
    • 相对专家负载(Relative Expert Load)表示实际专家负载与理论平衡专家负载之间的比率
1…789…65
Joe Zhou

Joe Zhou

Stay Hungry. Stay Foolish.

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