Jiahong 的个人博客

凡事预则立,不预则废


  • Home

  • Tags

  • Archives

  • Navigation

  • Search

RS——MTGR

  • 参考链接:
    • MTGR:美团外卖生成式推荐Scaling Law落地实践 - 美团技术团队

整体总结性说明

  • 样本组织思路:基本思路是将样本按照用户组织,具体来说,将同一用户的天粒度样本组织为同一条样本,同时通过掩码机制保证不会出现训练时数据/特征穿越问题
  • 输入信息 Token 化(未详细说明 Token 化的方案)

    输入信息Token化:我们将样本信息拆分成一个个Token。其中对于User Profile每一个特征都表示为一个Token,对于用户行为,则将每一个具体行为对应的item ID以及多个side info的Embedding表征拼接在一起再通过非线形映射组装成一个Token。同理,对于每一个曝光,我们则将item ID、对应的side info、交叉特征、时空等Context信息拼接在一起组成一Token

  • HSTU架构使用

    采用HSTU架构统一建模多个序列:我们利用HSTU架构针对Token化的输入进行统一编码,对比原方法,我们对每一层的输入都做了额外的LayerNorm以保证深层次训练的稳定性,考虑到不同与传统的LLM,我们的数据可以分组为多种类型的Token,因此,我们将LayerNorm拓展为Group LayerNorm,针对不同类别的Token采用不同的LayerNorm参数,以实现不同语义空间下Token的对齐


美团外卖的历史Scaling路径

  • 历史Scaling路径
  • 粉色框部分是之前的 Scaling 主要模块

特征处理

  • 保留全部 DLRM 特征
    • User Profile:用户统计信息,购买次数等
    • Context:广告位,时间等
    • UBS(User Behaviour Sequence):用户历史点击、曝光序列,包含每个 item(商家) 及其 Side Info
    • Target item:待打分的商家,包含商家ID、Side Info 以及基于统计的交叉特征信息
    • 注:博客中给出了丢弃交叉特征后,性能损失 CTR_GAUC -1.25pp;CTCVR_GAUC -0.91pp
  • 按照用户压缩数据:将同一用户的天粒度样本组织为同一条样本,同时通过掩码机制(后续介绍)保证不会出现训练时数据/特征穿越问题
  • 问题:按照用户压缩样本数据后会导致模型数据分布发生变化,减少模型对活跃用户的重视程度吗?
    • 回答:不会,因为后文会提到,每个原始样本还是有一个自己的损失,总的样本损失数还是没有变化的

模型结构

  • 整体架构如下所示:
  • 输入特征 Token 化:
    • User Profile:每一个特征都表示为一个Token
      • 问题:是否都是离散特征,直接就是 Token化的?
    • UBS:对每一个行为 item,将其 item ID,以及 多个 Side Info 的 Embedding 表征拼接在一起,然后通过非线形映射映射成一个Token(理解:每个 item 输出一个离散 Token)
      • 问题:非线性映射如何将连续向量映射成 Token?使用类似 CodeBook 机制吗?
    • 每一个候选目标,将其 item ID、对应的side info、交叉特征、时空等 Context 信息拼接在一起组成 Token
    • 补充原文描述:

      我们将样本信息拆分成一个个Token。其中对于User Profile每一个特征都表示为一个Token,对于用户行为,则将每一个具体行为对应的item ID以及多个side info的Embedding表征拼接在一起再通过非线形映射组装成一个Token。同理,对于每一个曝光,我们则将item ID、对应的side info、交叉特征、时空等Context信息拼接在一起组成一Token

  • 使用HSTU架构,并做以下改进
    • 使用 GroupLayerNorm 层替换原始 HSTU 中的 Norm 层(Group的含义是针对不同类别的 Token 采用不同的 LayerNorm 参数)
    • 在每一层输入时增加 GroupLayerNorm(注:增加 GroupLayerNorm 可提升0.17pp~0.18pp 的GAUC)
  • 混合动态掩码机制(注:使用掩码机制可提升 0.12pp~0.16pp 的GAUC):
    • 历史静态特征:包括 User Profile + 用户历史行为序列(User History Sequence),完全互相可见
    • 当日实时行为序列(Real-Time Sequence):使用 Casual Mask 机制
    • 目标 item(Targets):按照用户压缩组织样本后,这里会有多个 item,一方面要求 item 看不到发生于自己之前的当日实时行为序列,同时看不到自己之前的 Target item 信息(Casual Mask)
    • 掩码机制的详细示意图如下(注:其中示意图的括号有点歪,按照颜色来区分更准确):

      实时行为/曝光样本按照时间近至远进行排列:最近发生的曝光样本可以看到全部的实时行为,但较早发生的曝光样本只能看到部分实时行为
      对于Target 1,因为它出现的时间较晚,因此它可以看到全部的User Profile & Sequence以及Real-Time Sequence;对于Target 2,有一部分Real-Time Sequence出现在其之前,因此它只能看到部分Real-Time Sequence

  • HSTU 最终输出的利用:
    • 任务建模:实际上使用的是 HSTU 编码后的 embedding 信息直接进行 多任务训练(CTR,CVR等)
    • 注意:多个 targets 可以同时建模(一个样本包含多个 Target item 头,每个头都对应之前的一个曝光样本)每个头又会被分为多个任务

附录:生成式体现在哪里?

  • 唯一体现生成式的地方是 Token 化特征后,Item 训练有 Casual Mask 的机制,变相在当日实时 Item 和 Target Item 上都实现了类似 Next Item Prediction 的感觉了
  • Casual Mask 部分的序列不长(同一天内同一个用户的曝光数不会特别多),更像是一个 Prefix LM 架构
  • 实际上生成式的体现很弱,甚至像是没有,但是确实拿到了 Scaling Law 的效果

附录:Scaling Law效果如何?

  • Scaling Law实验结果
  • 实际上效果有点微弱,虽然随着模型增大,效果一直在涨,但涨幅非常有限

RS——LUM

LUM: large user model

  • 参考链接:
    • (LUM)Unlocking Scaling Law in Industrial Recommendation Systems with a Three-step Paradigm based Large User Model, 202502, Alibaba

整体思路

  • LLM :自回归大型语言模型(LLMs)的成功主要归功于其可扩展性,通常被称为“扩展规律”(Scaling Law)
  • LLM for RS :目前大量工作的思路是将推荐系统(RecSys)任务重新表述为生成问题来适应 LLMs
    • 这种端到端生成推荐(End-to-End Generative Recommendation,E2E-GR)方法往往优先考虑理想化目标(idealized goals)
    • 牺牲了传统基于深度学习的推荐模型(DLRMs)在特征、架构和实践方面提供的实际优势
    • 这种理想化目标与实际需求之间的差距引入了一些挑战和限制,阻碍了工业推荐系统中扩展规律的应用
  • 论文方案 :论文提出了一种大型用户模型(large user model,LUM),通过一个三步范式(three-step paradigm)解决这些限制
    • 做到满足工业环境的严格要求
    • 合理复用历史 DLRMs 的特征工程
    • 释放 Scaling Law 在推荐系统上的潜力
  • 离线+在线实验
    • 离线:实验表名,LUM 在性能上优于 SOTA 的 DLRMs 和 E2E-GR 方法,此外,LUM 展现出卓越的可扩展性,当模型参数扩展到70亿时,性能持续提升

关于 LUM 提出的背景和一些讨论

  • LLMs 显示出 Scaling Law,推荐系统(RecSys)领域也在积极探索利用类似 Scaling Law 的潜力
  • 早期研究指出,工业中广泛使用的传统基于深度学习的推荐模型(DLRMs)并未展现出与LLMs相同水平的可扩展性(Liu 等, 2020)。这种差异可以归因于生成模型与判别模型之间的根本区别
    • 生成模型需要更大的容量来有效捕捉数据的复杂联合概率分布 \(p(x,y)\)
    • 判别模型(如传统的 DLRMs)则专注于建模更简单的条件概率 \(p(y|x)\)
    • 因此,增加计算资源对判别模型的性能提升效果较弱(Chen 等, 2018; Li 等, 2020; Chen 等, 2021)
  • 一些研究尝试 将推荐系统中召回和排序任务重新表述为生成任务 ,以模仿生成模型,具体来说,这是通过将用户行为序列(User Behavior Sequences,UBS)作为语料库,并以端到端的方式在“next-item prediction”任务上训练 Transformer 架构来实现的(Chen 等, 2018; Li 等, 2020)
  • 端到端生成推荐方法(E2E-GRs)往往过度强调理想化目标,从而忽视了 DLRMs 在特征、架构和实践中的固有优势。这种建模范式中理想化目标与实际实现之间的差距,随后引发了一系列挑战和限制
    • (1) 生成训练与判别应用之间的不一致性:尽管 E2E-GRs 擅长捕捉数据中的复杂模式和分布,但在应用于特定判别任务(如点击率预测,其中校准能力和排序能力至关重要)时,往往表现出局限性(Wang 等, 2020)
      • 这种差异源于生成模型对数据生成过程的固有关注,而非精确的预测结果
      • 尽管生成模型擅长捕捉数据的固有分布,但它们可能无法始终满足判别应用的严格要求,这些应用需要高水平的准确性和特异性(Li 等, 2020; Chen 等, 2021; Chen 等, 2021)
    • (2) 效率挑战:在工业环境中,连续流式训练对高效率的需求,以及在线实时推理对低延迟的严格要求,为 E2E-GRs 的直接实施带来了重大挑战
      • 即使部署可行,这些效率限制仍会阻碍这些端到端模型的可扩展性(后文第5.4节会继续讨论)
    • (3) 缺乏灵活性: E2E-GRs 在适应业务需求的动态变化方面表现出明显的局限性,具体而言,整合新型行为数据(如退款行为或与新场景相关的行为)是一 item重大挑战
      • E2E-GRs 的固有结构要求,对输入模式的任何修改(包括新增元素)都会触发对整个模型重新训练的需求。这种刚性不仅增加了新特征整合的复杂性,还对系统的适应性施加了重大限制
      • 在工业环境中,快速灵活地响应变化条件通常至关重要,而从头开始重新训练模型既耗时又耗费资源,从而降低了 E2E-GRs 的实用性
    • (4) 有限的兼容性: E2E-GRs 使用原始 UBS 和 Transformer 架构构建。这种方法固有地限制了它们与现有工业知识(如显式特征工程和 DLRMs 的参数继承)的兼容性
      • 这些 E2E-GRs 在开发阶段与工业环境中的实际部署之间往往存在显著的性能差距(第5.3节)。这种差距在在线模型经过多年甚至数十年优化的应用中尤为明显。这一差异凸显了需要更强大的整合策略,以弥合理论进展与实际应用之间的差距
  • 论文重新思考了一个关键问题:如何在工业环境中有效利用生成模型来解锁 Scaling Law?
    • 这一问题促使论文重新审视使用 LLMs 的常见范式(图1),具体而言,传统 LLM 的范式是:
      • LLMs 首先在大量数据上以生成方式进行训练,使其能够吸收广泛的知识
      • 终端用户(End-users)随后通过多样化的 prompts 与这些模型交互,提出问题(可以同时做多个领域的知识问答)
      • 终端用户使用 LLM 的返回做决策(注:原文中称该决策为informed decision,可以直接翻译为有见识/学问的决策)
    • LLMs的这种 multi-step, generative-to-discriminative 特性为在实际应用中利用生成模型提供了基础框架
    • 基于这一见解,论文提出了一种用于工业用途的大型用户模型(LUM)的三步训练范式。这些步骤如下(见图1):
      • Step 1:知识构建 引入 LUM,通过生成学习进行预训练(基于 Transformer 架构)。该模型捕捉用户兴趣和物品之间的协作关系,从而构建一个全面的知识库
      • Step 2:知识查询 在此阶段,LUM 通过预定义的与用户特定信息相关的问题进行查询,提取 relevant insights
        • 直观上,这一过程可以基本概念化为一种“prompt engineering”,专门设计用于引出广泛的知识
      • Step 3:知识利用 LUM 在 Step 2 中获得的输出作为补充特征 ,被整合到传统的 DLRMs 中,以增强其预测准确性和决策能力
  • 总体而言,我们可以从三步范式中获益:
    • (1) 在 Step 1 中,LUM 的生成学习能够探索扩展规律,这对提升模型性能至关重要
    • (2) 论文范式的解耦设计消除了连续流式训练或服务的限制(进使用 prompt engineering 提取知识了)
      • 这种分离便于为 LUM 实施缓存策略,从而缓解效率限制
    • (3) Step 3 确保 DLRMs 能够满足实时学习、灵活性和兼容性的要求
      • 这是通过将先前学习的 LUM 与 DLRMs 整合来实现的,从而增强了它们对动态环境的适应性,并确保与现有系统的无缝集成
  • 另一个挑战是如何有效地将 Step 1 中学习到的数据联合分布 \(p(x,y)\) 迁移到下游判别任务中
    • 理想情况下,LUM 中封装的用户知识应与目标判别任务具有显著相关性
    • 为了实现这一点,论文为 UBS 引入了一种新颖的 token 化策略,其中每个物品被扩展为两个不同的 token :** condition token 和物品 token **(图2和第4.1.1节)
    • 随后,论文将 UBS 的自回归学习过程从“next-item prediction”重新定义为“next-condition-item prediction”。通过这一重新表述,我们可以在 Step 2 过程中通过指定各种条件,将相关知识从 LUM 无缝 trigger 到判别任务中
  • 最终结果说明:
    • 离线+在线实验
    • LUM 展现出与 LLMs 相似的 Scaling Law(PS:能够成功扩展到70亿参数,同时保持一致的性能提升)

相关工作(直译)

基于深度学习的推荐模型

  • 传统的 DLRMs 通常采用深度神经网络,可以大致分为两大类:
  • (1) Retrieval-Oriented Models :这些模型主要专注于为用户召回最相关的物品。它们通常采用双塔架构,如EDB(Bernard 等, 2018),或序列模型,如SASRec(Krizhevsky 等, 2013)、BERT4Rec(Zhang 等, 2019)和GRU4Rec(B),以有效捕捉用户与物品之间的相关性
  • (2) CTR(Click-Through Rate)预估模型 :这类模型旨在预测点击率。许多模型采用嵌入+多层感知机(MLP)架构。一些工作(如DeepFM(Chen 等, 2019)和xDeepFM(Wang 等, 2019a))旨在构建更复杂的交互信号,而其他工作(包括DIN(Yang 等, 2019)、DIEN(Xu 等, 2019)、SIM(Song 等, 2020a)和 TWIN(Choi 等, 2019))则探索了 UBS 的重要性。尽管这些方法在推荐系统中取得了显著改进,但它们通常无法随着计算资源的增加而扩展,并且未能充分利用大型基础模型的进展

生成推荐模型

  • 为了研究扩展规律,一些研究尝试类 LLM 架构,并通过端到端方式在 next-item prediction 任务上构建自回归 Transformer(Cheng 等, 2019; Wang 等, 2020b)
  • 然而,如前文所述,这些方法往往做出强烈且理想化的假设(strong and idealized),忽视了传统 DLRMs 在特征、架构和实践中的优势。这导致这些模型在应用于工业环境时出现一系列问题和限制
  • 论文提出了一种基于 LUM 的三步范式来解决这些限制。值得注意的是,也有工作尝试通过端到端(E2E)或多步训练利用 LLMs 中的开放世界知识构建基于内容的推荐模型(Chen 等, 2019; Chen 等, 2021b; Zhang 等, 2019; Xue 等, 2019)。然而,这些研究超出了论文当前研究的范围。论文的重点是开发一种可扩展的模型,捕捉协同过滤信号而非内容信号
  • 吐槽一下:这篇论文的引用咋都是22年以前的?

推荐系统基础只是

传统深度推荐模型(DLRMs)

  • 在推荐系统(RecSys)中,主要任务分为两类:召回和排序。论文以搜索场景为例进行讨论,但所提出的方法同样适用于其他工业应用场景
  • 问题定义:给定用户 \(u \in \mathcal{U}\) 、物品 \(i \in \mathcal{I}\) 和搜索词 \(s \in \mathcal{S}\)
    • 召回任务 :该任务的目标是从候选池中筛选出与用户 \(u\) 在查询 \(s\) 下的偏好相匹配的物品子集
      • 传统召回任务通常采用双塔架构,包括 user-query 塔 \(UEnc\) 和物品塔 \(IEnc\)(这些塔可以由任何合适的神经网络结构(如多层感知机MLP)组成)
      • user-query 塔将用户 \(u\) 和搜索词 \(s\) 编码为统一嵌入向量 \(\boldsymbol{e}_{us}^{\prime} = UEnc(us)\)
      • 物品塔则为物品 \(i\) 生成嵌入向量 \(\boldsymbol{e}_{i}^{\prime} = IEnc(i)\)
      • 最后,通过对比学习技术优化这两个塔的表示
    • 排序任务 :排序任务的目标是预测用户 \(u\) 在特定查询 \(s\) 下点击物品 \(i\) 的概率
      • 其数学形式为:
        $$\hat{y} = f(u,i,s)$$
        • \(\hat{y}\) 表示预测的 CTR
        • \(f\) 为 Embedding+MLP 架构

基于“next-item prediction”的端到端生成推荐(E2E-GRs)

  • 在 E2E-GRs 中,给定用户 \(u\) 的用户行为序列(UBS) \(B_u = \{i_1, i_2, …, i_L\}\)(其中 \(L\) 为序列长度),“next-item prediction”框架假设物品 \(i_k\) 出现的概率依赖于其之前的物品 \(\{i_1, i_2, …, i_{k-1}\}\)。因此,整个 UBS 的似然可以表示为:
    $$ p(i_1, i_2, …, i_L) = \prod_{j=1}^{L} p(i_j | i_1, i_2, …, i_{j-1}) $$
    • E2E-GRs 的自回归学习目标是优化分布 \(p_{\theta}(i_1, i_2, …, i_L)\),这一方法通常被称为“next-item prediction”

LUM 方法整体介绍

Step 1:通过预训练 LUM 构建知识

Tokenization
  • 论文将 UBS 的自回归建模从 “next-item prediction” 转变为 “next-condition-item prediction” 范式
  • 序列 \(B_u\) 中的任意 item \(i\) 被分解为两个不同的 token:condition token 和 item token
  • UBS 可以表示为交替 token 的序列:\(\{c_1, i_1, c_2, i_2, …, c_L, i_L\}\),其中 \(c_k\) 表示与 item \(i_k\) 相关联的 condition token(如图2所示)
  • 举例 :用户的行为可能涉及多个场景(如推荐和搜索场景)中的 item 交互。此时,condition token \(c_i\) 可以定义为场景 token,从而捕捉 item \(i\) 所处的特定环境。这种方法有助于细致理解用户在不同场景中的偏好和行为
  • 论文方法与 HSTU 等方法引入了额外的动作 token 的形式不同
    • HSTU 中,任意 item被扩展为 <item, action>,其中动作 \(a_k\) 与前面的(preceding) item \(i_k\) 相关联,而非后面的(subsequent) item(如图2所示)
    • HSTU 的策略无法基于不同条件预测 next item,也难以捕捉用户在不同方面的偏好
LUM 整体架构
  • LUM 的整体框架如图3(a)所示,采用分层结构,包含一个 Token Encoder 和一个 User Encoder
  • Token Encoder :LUM 的输入 token 具有异质性,主要分为两类:condition token \(c\) 和 item token \(i\)
    • 单个 token 可能携带多种属性特征(例如 item token 的ID、统计和内容特征)
    • Token Encoder 旨在将这些异质输入统一为 token 嵌入
    • 具体过程是先将每个 token 的特征或嵌入拼接,再通过投影层将其整合到共同的表示空间中。数学上,这一转换可表示为:
      $$
      e^t = \text{proj}(\text{concat}(f_1^t; f_2^t; f_3^t; …)), \quad t \in \{i, c\}
      $$
      • \(f_k^t\) 表示 item token \(i\) 或 condition token \(c\) 的特征,其中\(t \in \{i, c\}\)
        • 注意:一个 token 只能是 item token \(i\) 或 condition token \(c\) 中的一种,不能同时属于两种 token
      • \(e^t\) 为生成的 token 嵌入
      • \(\text{proj}\) 表示线性投影层
  • User Encoder :User Encoder 用于捕捉用户偏好和 item间的协同信息
    • 输入 token 序列 \(\{c_1, i_1, c_2, i_2, …, c_L, i_L\}\) 通过 Token Encoder 表示为 \(\{e_1^c, e_1^i, e_2^c, e_2^i, …, e_L^c, e_L^i\}\)
    • 如图3(a)所示,User Encoder 使用传统的自回归 Transformer 架构处理这些嵌入,最终输出为 \(o_k^c\),封装了输入序列的整合信息
Next-condition-item Prediction
  • “next-condition-item prediction”关注在给定特定条件下预测 next item
    • 注:这种方法仅对 condition token 的输出应用自回归损失以推断 next item。因此,next-condition-item prediction 的自回归似然可表示为:
      $$
      p(c_1, i_1, c_2, i_2, …, c_L, i_L) = \prod_{l=1}^L p(i_l | c_1, i_1, c_2, i_2, …, i_{l-1}, c_l)
      $$
  • 为优化工业应用中的 \(p_\theta(c_1, i_1, c_2, i_2, …, c_L, i_L)\),本研究采用了 InfoNCE 损失函数并引入了 packing 策略
  • InfoNCE 损失 :在工业应用中,item的词表可能达到数十亿,直接计算所有 item的生成概率是不现实的。为解决这一问题,本研究使用InfoNCE损失预测next-condition item。其数学表达式为:
    $$
    \text{Loss} = -\sum_{l=1}^L \log \left( \frac{\exp(\text{sim}(o_{l-1}^c, e_l^i))}{\exp(\text{sim}(o_{l-1}^c, e_l^i)) + \sum_{k=1}^K \exp(\text{sim}(o_{l-1}^c, e_k^i))} \right)
    $$
    • \(\text{sim}\) 为相似度函数
    • \(o_{l-1}^c\) 是自回归序列中 token \(e_{t-1}^c\) 对应的 Transformer 输出
    • 对于任意 item \(i\),同一批次中的其他 item作为负样本,\(K\) 为负样本数量,\(e_k^i\) 为第 \(k\) 个负 item 的嵌入
    • 注意:应用 InfoNCE 损失以后,不需要面对过大的词汇表了,仅关注 \(K\) 个负样本就行
  • packing 策略 :在实际应用中,用户行为序列的长度差异显著,大多数序列长度远小于预设的最大长度。若单独处理每个序列,计算效率较低。受GPT系列打包策略的启发,本研究将多个 UBS 合并为单一序列,以充分利用序列长度

Step 2:基于给定条件查询知识

  • 在 Step 1 中,论文构建了联合概率分布 \(p(c_1, i_1, c_2, i_2, …, c_L, i_L)\),Step 2 是从这一分布中提取相关知识
  • 原文4.1.1节的分词方法支持在不同条件下查询知识,给定查询条件 \(c_q\),可以计算条件概率 \(p(i_q | c_1, i_1, c_2, i_2, …, c_L, i_L, c_q)\),以确定用户对 item \(i_q\) 的兴趣(如图3(b)所示)
  • 这种触发条件知识的方法在生成模型与判别任务之间架起了桥梁,从而提升了工业应用的效果。以下是不同 condition token 的应用示例:
    • 示例1 :若 condition token 为场景(scenario) token,模型可推断用户在不同场景中的兴趣
    • 示例2 :若 condition token 为搜索场景中的 query token,模型可基于不同搜索查询推断用户兴趣
    • 示例3 :若 condition token 为类别(category) token,模型可判断用户在不同类别中的兴趣
  • LUM 允许在 \(\{f_1^c; f_2^c; f_3^c; …\}\) 加入额外条件特征来同时考虑多种条件
    • 即加入 Token Encoder 的输入中,从而编码上指定条件信息
    • 实验表明,整合多样条件能显著提升性能(见原文5.3节)
    • 本质上,这一过程可视为一种“prompt engineering”,旨在激发广泛的知识
  • 分组查询以提升效率 : 由于单个用户可能涉及多个查询(均基于同一 UBS),单独处理这些查询会导致效率低下
    • 在实际场景中,用户数量可能达到数十亿, <user, query> 对的数量会急剧增加
    • 为解决这一问题,本研究引入了分组查询策略以提升计算效率
    • 如图4所示,所有查询被拼接为单一序列,表示为 \(p(i_{q_1}, i_{q_2}, … | c_1, i_1, c_2, i_2, …, c_L, i_L, c_{q_1}, c_{q_2}, …)\)
    • 为确保推理过程准确且高效,采用掩码机制防止不同查询条件 \(c_{q_j}\) 间的注意力交互。这种方法使得不同查询的公共前缀 \(\{c_1, i_1, c_2, i_2, …, c_L, i_L\}\) 仅需计算一次,同时支持在不同条件下查询 item \(i_{q_j}\)
    • 实验表明,分组查询策略可将推理速度提升78%(见5.3节)

Step 3:DLRM中的利用知识

  • 在 Step 2 完成后,论文获得了一组 \(N\) 个 next-condition item \(\{i_{q_1}, i_{q_2}, …, i_{q_N}\}\),其对应输出为 \(o_{q_n}^i\)
    • 注:每个 item \(i \in I\) 通过 Token Encoder(4.1.2节)生成嵌入 \(e_i^i\),在经过 User Encoder 最终生成 \(o_{i}^i\)
    • 问题: \(o_{i}^i\) 是没有损失函数的,为什么可以直接使用?是不是应该使用 \(o_{i}^c\) 才合理?
  • 为增强现有DLRM,本研究提出两种策略(如图3(c)所示):
    • 直接特征融合 :将输出 \(o_{q_n}^i\) 和目标 item \(i\) 的嵌入 \(e_i^i\) 作为固定附加特征整合到 DLRM 中
    • 通过相似度匹配兴趣 :计算目标 item \(i\) 与用户兴趣的相似度 \(\text{sim}(o_{q_n}^i, e_i^i)\)(如公式2所示),以量化目标 item 与上下文的匹配程度
  • 形式上,对于召回任务,双塔模型可重构为:
    $$
    e_{us}^r = \text{UEnc}(us, \{o_{q_1}^i, o_{q_2}^i, …, o_{q_N}^i\}), \quad e_i^r = \text{IEnc}(i, e_i^i)
    $$
  • 对于排序任务,排序模型可表示为:
    $$
    \hat{y} = f(u, i, s, \{o_{q_n}^i, \text{sim}(o_{q_n}^i, e_i^i) | n=1, …, N\}, e_i^i)
    $$
  • 这一框架有效整合了上下文信息和 item嵌入,提升了推荐系统中召回和排序的性能

讨论

  • 基于 LUM 的三步范式解决了原文第1节提到的四个限制:
    • 解决限制1(生成训练与判别应用之间的不一致性) :通过生成到判别的设计,Step 1 构建联合分布 \(p(c_1, i_1, c_2, i_2, …, c_L, i_L)\),Step 2 通过“next-condition-item prediction”触发相关知识,Step 3 进行判别学习以满足应用需求
    • 解决限制2(效率挑战) : Step 3 的模型训练和服务天然支持工业应用的高效需求,且 Step 1 和 Step 2 的计算成本较低,因其结果可预计算和存储
    • 解决限制3(缺乏灵活性) :业务需求的动态变化可通过 Step 3 灵活处理,同时这些需求可统一为 condition token,支持 LUM 的持续训练以适应新需求
    • 解决限制4(有限的兼容性) : Step 3 的DLRM主干可设置为在线模型,便于利用现有工业知识,并持续受益于DLRM技术的进步

Experiments

Experiment Setup

  • 数据集 :本研究使用了三个公共数据集和一个工业数据集来评估所提出的 LUM 方法的性能。公共数据集包括两个基准数据集:MovieLens(包含1M和20M两个子集)和 Amazon Books[23, 25]。工业数据集来自淘宝电商平台。这些数据集的统计信息如表1所示

  • 基线模型 :为了全面评估 LUM 的性能,论文将其与多种 SOTA 模型进行比较

    • 对于传统的召回模型,论文使用基于双塔架构的 EDB[9]
    • 对于传统的排序模型,论文考虑了 DIN[26]、DIEN[25]、SIM[16]和 TWIN[3]
    • 此外,论文还比较了 LUM 与端到端生成推荐方法(E2E-GRs),特别是 HSTU[23]
    • 论文还纳入了传统的序列推荐模型 SASRec[10]作为基线,该模型采用 Transformer 架构建模用户行为序列(UBS)
  • 训练细节 :为了确保公平比较,论文遵循以下训练配置

    • 默认情况下,论文为 Transformer 风格的模型(包括 LUM、HSTU 和 SASRec)保持相似的配置,以确保模型规模可比,其他 DLRMs 的配置根据其原始论文的建议设置
    • 所有模型均使用相同的特征集从头开始训练
    • 公共数据集的序列长度设置为 256,工业数据集的序列长度设置为 4096
    • 对于 LUM,在公共数据集的排序任务中,Step 3 的 DLRM 主干配置为 SIM;在召回任务中配置为 EDB(一种双塔结构)
    • 工业数据集的主干设置为当前生产环境中部署的在线模型(理解:使用当前线上部署模型去推理工业数据集)

推荐任务性能

  • 公共数据集上的性能 :论文首先在公共数据集上评估 LUM 的性能,性能指标为 AUC
    • 结果如 表2 所示。从结果可以看出,LUM 在所有数据集上均实现了显著提升,这表明基于三步范式的 LUM 能够有效捕捉广泛的用户兴趣并提升 DLRMs 的预测准确性
    • 注:Note 0.001 absolute AUC gain is regarded as significant for the ranking task
  • 工业环境中的性能 :本节论文比较 LUM 与 DLRMs 和 E2E-GRs 在工业环境中的性能
    • 基线使用论文应用中的在线模型,排序任务采用 Embedding+MLP 架构,召回任务采用双塔架构
    • 排序任务评估指标是 AUC,召回任务评估指标是 Recall@K(R@K)
    • 对于 LUM,Step 3 的DLRM主干也设置为在线模型
    • 论文还纳入了传统 SOTA DLRMs(SASRec[10] 和 DIN[26])和 E2E-GR(HSTU[23])作为基线进行比较
    • 所有模型使用相同的特征集,序列长度设置为4096,并从头开始训练
    • 结果如表3所示。LUM 相对于最佳基线实现了显著提升,AUC 提高了 +0.0176,R@10 提高了 +0.0133,R@50 提高了 +0.0134。这些提升主要归功于论文提出的生成-判别设计范式

有效性评估

  • 本节论文探讨基于三步范式的 LUM 的优势
  • 对多种 DLRMs 的影响 :由于论文的范式采用解耦设计,LUM 可以在 Step 3 中与任何 DLRM 集成。为了评估 LUM 的有效性,论文比较了原始 DLRMs(记为Base)和增强后的 DLRMs(记为Base+ LUM)在排序任务上的性能
    • 结果如 表4 所示。结果表明,在 LUM 的辅助下,所有方法的性能均显著提升,提升幅度从 +0.0053 到 +0.0176 不等。这些发现凸显了 LUM 在提升多种 DLRMs 预测准确性方面的通用性和有效性
  • 工业应用中 warming up 设置的影响 :在工业应用中,在线模型通常包含复杂的特征工程,并每天在数十亿数据上持续训练
    • 现有的 E2E-GRs 往往忽略这些在线模型的知识,这对其性能提出了挑战,尤其是在在线模型已开发数年甚至数十年的应用中
    • LUM 具有出色的兼容性,可以在 warming up 设置下训练。本节论文通过实现 LUM 的多个版本,并与 E2E-GR(HSTU)和DLRM(在线模型)进行比较,探讨这种兼容性的影响。性能指标为工业数据集上的 AUC
    • 结果如表5所示(“scratch”表示使用相同特征集从头训练的模型,“feature”表示使用复杂特征工程,“param”表示从持续训练的在线模型初始化参数)
    • 从表5可以得出以下结论:
      • (1) 与 LUM 相比,warming up 设置下的模型(包括 LUM(feature)、 LUM(param)和 LUM(feature+param))在AUC上实现了显著提升,幅度从+0.0106到+0.028。这表明利用在线模型的现有知识具有重要价值
      • (2) 尽管HSTU在scratch设置下优于DLRM,但与DLRM(param+feature)相比仍存在较大性能差距(0.7334 vs. 0.7777)。这凸显了与在线模型兼容的重要性。E2E-GRs 的有限兼容性可能阻碍其在真实应用中的实际部署
  • 提出的分词策略的影响 :论文评估了提出的分词策略的有效性
    • 论文开发了 LUM(w/o condition token),并在 表5 中给出AUC
    • 结果显示,与 LUM(w/o condition token)相比,LUM 通过给定条件更好地理解 UBS,从而实现了更好的性能
    • 论文还评估了使用多条件(如场景条件和搜索词条件)的效果,记为 LUM(multi-conditions)
    • 从 表5 可以看出,添加更多条件可以进一步提升性能
  • 知识利用的影响 :论文评估了 Step 3 中知识利用的不同策略
    • 结果如表5所示。“direct feature”和“interest matching”分别对应原文第4.3节中详细描述的知识利用策略
    • LUM(direct feature)和 LUM(direct feature + interest matching)均显著优于DLRM,证明了所提策略的有效性
  • 打包和分组查询的影响 :打包和分组查询旨在加速 Step 1 和 Step 2 的处理
    • 表6展示了这些策略在工业环境中的效率提升,Step 1 和 Step 2 的处理速度分别提高了82%和78%

效率评估

  • 训练效率(Training Efficiency) :图5(a)展示了不同模型在工业环境中一天数据的训练时间成本
    • 对于 E2E-GR,论文遵循HSTU[23]的优化方法,从展示级训练到用户级
    • 对于DLRM,论文使用应用中的在线模型作为基线
    • 对于 LUM,DLRM主干也设置为在线模型
    • 由于 LUM 的 Step 1 和 Step 2 可以预处理(4.4节),因此在 Step 3 训练下游模型时不计入其时间成本
    • 序列长度设置为4096,LUM 和 E2E-GR 的模型规模从0.5亿到140亿参数不等
    • 所有模型在128个GPU上训练
    • 最高时间成本设置为24小时,因为实际中的持续训练要求模型在此时间内处理每日数据
    • LUM 的训练时间成本与DLRM相近,并且由于其解耦的三步范式,对模型规模不敏感,这一特性在工业环境中解锁了训练时的扩展潜力
    • 相比之下,E2E-GR 比 LUM 慢12倍到98倍
    • 所有 E2E-GR 模型(不同规模)均未满足在24小时内完成训练的要求
    • 为了匹配 LUM 的吞吐量,E2E-GR 需要12倍到98倍的GPU;为了满足24小时要求,需要2倍到18倍的GPU
  • Serving 效率(Serving Efficiency) :图5(b)和(c)展示了不同模型在线服务时的延迟
    • 对于 E2E-GR,论文遵循 HSTU 的 M-FALCON 实现[23]
    • DLRM 设置为在线模型,同时也是 LUM 的主干排序模型
    • LUM 的 Step 3 需要评估 Serving 延迟(Step 1 和 Step 2 可以预计算(第4.4节)
    • 在图5(b)中,序列长度设置为4096,LUM 和 E2E-GR 的模型规模从 0.5亿 到 140亿 参数不等
    • 延迟约束设置为小于30毫秒(ms),在论文的案例中,排序候选数量约为100
    • 总体而言,LUM 的延迟与模型规模无关,使得我们可以扩展 LUM 而不违反延迟约束。相反,即使对于小型模型(0.5亿参数),E2E-GR 也无法提供及时响应(小于30 ms)
    • 为了进一步研究,论文尝试减少 E2E-GR 的序列长度以满足延迟约束(图5(c))。遗憾的是,在使用140亿参数模型时,满足延迟约束(<30 ms)的最大序列长度仅为64,远低于典型设置(在论文的案例中缩小了64倍)
    • 这些结果表明,尽管 E2E-GR 在离线性能上可能表现出扩展性,但由于延迟约束,大规模模型在工业环境中的实际应用受到严重限制。相比之下,LUM 的解耦架构确保了延迟性能的一致性,使其成为实时工业应用中更可行和可扩展的解决方案

LUM 的扩展性

  • 遵循 Improving language understanding by generative pretraining 和 Language models are unsupervised multitask learners 的协议,论文检验 LUM 是否遵循类似的 Scaling Law,论文重点评估模型参数和序列长度的扩展性
    • 对于模型参数,为了评估模型参数的影响,论文保持序列长度为4096,训练模型规模从1900万到70亿参数不等
    • 对于序列长度,为了评估序列长度的影响,论文在固定参数规模为3亿的情况下,训练序列长度从256到8192不等的模型
    • 结果如图6所示,论文观察到明显的幂律扩展趋势,与之前的研究一致[17, 18, 23]。幂律扩展规律可以表示为:
      $$
      R_{P} = 0.0068 \cdot \log(P) + 0.1741 \\
      R_{L} = 0.0147 \cdot \log(L) + 0.2326
      $$
      • \(R_{P}\) 和 \(R_{L}\) 分别表示不同模型规模和序列长度下的 R@10 指标,\(P\) 表示模型规模,\(L\) 表示序列长度
    • 这些结果证实了 LUM 的强大扩展性,表明增加模型规模和序列长度可以持续提升模型性能。这一发现凸显了 LUM 在扩展时实现更高性能的潜力,使其成为大规模工业应用的有前景的方法

在线结果

  • 为了评估 LUM 在工业环境中的有效性,论文将其部署在中国最大的电商平台淘宝的搜索广告系统中
  • 如原文第4.4节所述,评估过程包括以下关键步骤:
    • (1) 离线阶段,论文首先预训练 LUM,并在各种条件下 pre-triggered 必要的知识
    • (2) 存储 LUM 在 Step 2 生成的响应,以便在线服务时直接使用。这种预计算有助于满足实时工业应用中的严格延迟约束
    • (3) 最后,论文进行了在线A/B实验测试 LUM 在排序任务中的表现。关键性能指标 CTR(点击率)和 RPM(每千次展示收入)分别显著提升了 2.9% 和 1.2%,这些发现凸显了 LUM 的实际效益,证明了其在大规模电商平台中提升用户参与度和业务成果的能力

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——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——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——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——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 还没收敛

Gradient Norm

  • 这个指标需要从底层 Infra 去统计并上报
  • 经验:Gradient Norm 一开始会比较大(可能会涨一下),然后逐步下降
  • 若训练崩溃,则 Gradient Norm 会突然变的非常大,所以 Gradient Norm 变大要非常小心

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

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——技术报告解读-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 模型源码
    • 包含一些流程图和源码,比较清晰
1…8910…67
Joe Zhou

Joe Zhou

Stay Hungry. Stay Foolish.

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