HSTU

0 阅读6分钟

标题:行动胜于言语:面向生成式推荐的万亿参数序列转导器
单位:meta
链接:Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations(2024.02)

摘要

大规模推荐系统依赖高基数异构特征,并且每天需要处理数百亿次用户行为。尽管工业界的大多数深度学习推荐模型(Deep Learning Recommendation Model,DLRM)使用数千种特征和海量数据进行训练,它们的性能却无法随计算量有效扩展。

受到 Transformer 在语言和视觉领域成功经验的启发,我们重新审视推荐系统中的基础设计选择。我们在生成式建模框架下,把推荐问题重新表述为序列转导任务,并称之为“生成式推荐器”(Generative Recommender,GR);同时提出一种新架构 HSTU,用于处理高基数、非平稳的流式推荐数据。

在合成数据集和公开数据集上,HSTU 相较基线的 NDCG 最高提升 65.8%;在长度为 8,192 的序列上,其速度是基于 FlashAttention2 的 Transformer 的 5.3—15.2 倍。基于 HSTU、拥有 1.5 万亿参数的生成式推荐器在线上 A/B 测试中将指标提升了 12.4%,并已部署到某大型互联网平台的多个业务界面,为数十亿用户提供服务。

更重要的是,生成式推荐器的模型质量在三个数量级的范围内,经验上会按照训练计算量的幂律扩展,直至 GPT-3/LLaMA-2 的规模。这能够减少未来模型开发所需的碳足迹,并为推荐领域首批基础模型进一步铺平道路。

1 引言

推荐系统是在线内容平台与电子商务的典型基础设施,每天在为数十亿用户提供个性化体验方面发挥关键作用。大约十年来,最先进的推荐方法一直以 DLRM 为基础。DLRM 的特点是使用异构特征,包括计数器、比率等数值特征,嵌入,以及创作者 ID、用户 ID 等类别特征。由于每分钟都有新内容和新商品加入,特征空间具有极高基数,通常达到数十亿。为了利用数以万计的此类特征,DLRM 使用不同神经网络来组合特征、转换中间表示并合成最终输出。

尽管使用了大量人工设计的特征集,并在海量数据上训练,工业界的大多数 DLRM 仍无法随计算量良好扩展。这一局限值得关注,但一直没有得到解决。

image.png 图 1:历年训练深度学习模型所使用的总计算量。DLRM 的结果来自既有工作,GR 是本文已部署的模型。DLRM 和 GR 均在流式环境中持续训练,因此图中报告每年使用的计算量。

受到 Transformer 在语言和视觉领域成功经验的启发,我们重新审视现代推荐系统中的基础设计选择。我们发现,面向十亿用户的替代方案必须克服三项挑战。

第一,推荐系统中的特征缺乏显式结构。尽管小规模环境已经探索过序列化表述,但高基数 ID、交叉特征、计数器和比率等异构特征,在工业规模 DLRM 中仍发挥关键作用。第二,推荐系统使用持续变化的十亿规模词表。与语言模型约十万规模的静态词表不同,十亿规模动态词表会带来训练难题;同时,系统必须以目标感知方式考虑数以万计的候选,因此推理成本很高。

最后,计算成本是启用大规模序列模型的主要瓶颈。GPT-3 使用数千块 GPU,在一到两个月内总共对 3,000 亿个 Token 进行训练。这个规模看似惊人,但与用户行为规模相比则不然:最大的互联网平台拥有数十亿日活跃用户,他们每天与数十亿条帖子、图片和视频交互。用户序列长度可能达到 10510^5。因此,推荐系统每天需要处理的 Token 数量,比语言模型在一到两个月内处理的数量高出若干数量级。

本文把用户行为视为生成式建模中的一种新模态。我们的关键洞察是:第一,只要采用合适的新特征空间,工业规模推荐器中的核心排序和检索任务就能转化为生成式建模问题;第二,这种范式使我们能够系统地利用特征、训练和推理中的冗余来提高效率。借助这一新表述,我们部署的模型在计算复杂度上比此前最先进方法高出三个数量级,同时将核心指标提升了 12.4%,如图 1 所示。

本文首先在第 2 节提出生成式推荐器,这是一种取代 DLRM 的新范式。我们将 DLRM 中的异构特征空间序列化并统一起来;随着序列长度趋于无穷,新方法会逐渐逼近完整的 DLRM 特征空间。由此,主要推荐问题——排序和检索——可以在 GR 中重新表述为纯序列转导任务。更重要的是,模型可以按序列生成方式训练,从而在相同计算量下使用高出若干数量级的数据。

接着,我们处理训练与推理过程中的计算成本问题,提出一种新的序列转导架构:分层序列转导单元(Hierarchical Sequential Transduction Unit,HSTU)。HSTU 针对大型非平稳词表修改了注意力机制,并利用推荐数据集的特性,在长度为 8,192 的序列上,相较基于 FlashAttention2 的 Transformer 获得 5.3—15.2 倍加速。此外,通过一种利用微批次充分摊销计算成本的新算法 M-FALCON,我们可以在传统 DLRM 的相同推理预算下,服务复杂度高出 285 倍的 GR 模型,并获得 1.50—2.99 倍加速。

最后,我们在合成数据集、公开数据集以及某大型互联网平台的多个业务界面上验证所提技术,该平台拥有数十亿日活跃用户。据我们所知,本文首次表明,在生成式环境中,HSTU 这类纯序列转导架构能够在大规模工业场景中显著优于 DLRM。值得注意的是,我们不仅克服了传统 DLRM 已知的扩展瓶颈,还成功证明扩展规律同样适用于推荐领域,这可能代表推荐系统的“ChatGPT 时刻”。

2 将推荐表述为序列转导任务:从 DLRM 到 GR

2.1 统一 DLRM 中的异构特征空间

现代 DLRM 通常使用大量类别(“稀疏”)特征和数值(“稠密”)特征进行训练。在 GR 中,我们把这些特征整合并编码为单一统一时间序列,如图 2 所示。

类别(“稀疏”)特征。 这类特征包括用户点赞过的物品、用户在某一类别下关注的创作者(例如户外类)、用户语言、用户加入的社群、请求发起城市等。我们按以下方式将其序列化:首先选择最长的时间序列作为主时间序列,通常由表示用户交互物品的特征合并而成。其余特征一般是随时间缓慢变化的时间序列,例如人口统计信息或关注的创作者。我们只保留各连续片段中最早的一项来压缩这些序列,再将结果合并到主时间序列。由于这些时间序列变化十分缓慢,这种做法不会显著增加总体序列长度。

数值(“稠密”)特征。 这类特征包括经过加权和衰减的计数器、比率等。例如,某项特征可以表示用户过去对符合特定主题的物品的点击率。相较类别特征,这些特征的变化频繁得多,可能随每一次“用户—物品”交互而变化。因此,从计算和存储角度看,无法将它们完全序列化。

不过,一个重要观察是,用于计算这些聚合量的类别特征,例如物品主题和位置,已经在 GR 中完成序列化与编码。因此,只要序列转导架构具有足够表达能力,并结合目标感知表述,随着 GR 的总体序列长度和计算量增加,我们就可以移除数值特征,而由模型有效捕获它们。

image.png 图 2:DLRM 与 GR 的特征和训练流程对比。E,F,G,HE,F,G,H 表示类别特征,Φi\Phi_i 表示合并后主时间序列中的第 ii 个物品,Ψk(tj)\Psi_k(t_j) 表示在时刻 tjt_j 产生的第 kk 个训练样本。完整符号定义见附录 A。

2.2 将排序和检索重新表述为序列转导任务

给定按时间排列的 nn 个 Token x0,x1,,xn1x_0,x_1,\ldots,x_{n-1},其中 xiXx_i\in\mathcal{X},以及观察到这些 Token 的时刻 t0,t1,,tn1t_0,t_1,\ldots,t_{n-1},序列转导任务会将该输入序列映射为输出 Token y0,y1,,yn1y_0,y_1,\ldots,y_{n-1},其中 yiX{}y_i\in\mathcal{X}\cup\{\varnothing\}yi=y_i=\varnothing 表示 yiy_i 未定义。

我们用 ΦiXc\Phi_i\in\mathcal{X}_c 表示系统提供给用户的内容,例如图片或视频,其中 XcX\mathcal{X}_c\subseteq\mathcal{X}。由于新内容不断产生,Xc\mathcal{X}_cX\mathcal{X} 都是非平稳的。用户可以通过某个行为 aiXa_i\in\mathcal{X} 响应 Φi\Phi_i,例如点赞、跳过或“完整观看视频并分享”。用户交互过的内容总数记为 ncn_c

这样,在因果自回归环境中,标准排序与检索任务就可以定义为序列转导任务。

image.png 表 1:将排序和检索表述为序列转导任务。为简洁起见,省略其他类别特征。附录 B.2 将 GR 与传统序列推荐器进行比较。

检索。 在推荐系统的检索阶段,我们学习 Φi+1Xc\Phi_{i+1}\in\mathcal{X}_c 上的分布 p(Φi+1ui)p(\Phi_{i+1}\mid\mathbf{u}_i),其中 ui\mathbf{u}_i 是 Token ii 处的用户表示。典型目标是选择 argmaxΦXcp(Φui)\arg\max_{\Phi\in\mathcal{X}_c}p(\Phi\mid\mathbf{u}_i),以最大化某种奖励。

这与标准自回归设置有两点不同。第一,xix_i 的监督信号 yiy_i 不一定是 Φi+1\Phi_{i+1},因为用户可能会对 Φi+1\Phi_{i+1} 作出负面响应。第二,如果 xi+1x_{i+1} 表示人口统计信息等与互动无关的类别特征,则 yiy_i 未定义。

排序。 GR 中的排序任务具有独特挑战,因为工业推荐系统往往要求“目标感知”表述。在这种环境中,目标 Φi+1\Phi_{i+1} 与历史特征需要尽早发生“交互”;标准自回归设置则无法做到这一点,因为其中的交互发生得很晚,例如在编码器输出后的 Softmax 中。

我们通过交错排列物品与行为来解决这一问题,使排序任务可以表述为 p(ai+1Φ0,a0,Φ1,a1,,Φi+1)p(a_{i+1}\mid\Phi_0,a_0,\Phi_1,a_1,\ldots,\Phi_{i+1}),暂不考虑类别特征。实践中,我们使用一个小型神经网络,把 Φi+1\Phi_{i+1} 位置的输出转换为多任务预测。重要的是,这使我们能在一次前向传播中,对全部 ncn_c 次互动应用目标感知交叉注意力。

2.3 生成式训练

工业推荐器通常在流式环境中训练,每个样本在产生后依次得到处理。在这种设置下,Transformer 等基于自注意力的序列转导架构,其总计算需求按下式扩展:

ini(ni2d+nidffd),\sum_i n_i\left(n_i^2d+n_id_{\mathrm{ff}}d\right),

其中 nin_i 是用户 ii 的 Token 数,dd 是嵌入维度。括号内第一项来自自注意力。这里假设其扩展因子为 O(n2)O(n^2),因为大多数次二次算法都会牺牲质量,而且实际运行时间不如二次算法。第二项来自逐位置 MLP 层,其隐藏层大小为 O(dff)=O(d)O(d_{\mathrm{ff}})=O(d)。令 N=maxiniN=\max_i n_i,总体时间复杂度可化为 O(N3d+N2d2)O(N^3d+N^2d^2),对推荐场景而言成本过高。

为了以可扩展方式在长序列上训练序列转导模型,我们从传统的曝光级训练转向生成式训练,将计算复杂度降低 O(N)O(N) 倍,如图 2 顶部所示。这样,编码器成本可以由多个目标共同摊销。

具体而言,当以速率 su(ni)s_u(n_i) 对第 ii 位用户采样时,总训练成本按下式扩展:

isu(ni)ni(ni2d+nid2).\sum_i s_u(n_i)n_i\left(n_i^2d+n_id^2\right).

su(ni)=1/nis_u(n_i)=1/n_i,即可把成本降至 O(N2d+Nd2)O(N^2d+Nd^2)。在工业规模系统中,一种实现方式是在用户请求或会话结束时产生训练样本,从而得到 s^u(ni)1/ni\hat{s}_u(n_i)\propto1/n_i

3 面向生成式推荐的高性能自注意力编码器

为了让 GR 扩展到具有大型非平稳词表的工业推荐系统,我们提出一种新的编码器设计——分层序列转导单元 HSTU。HSTU 由多个相同层堆叠而成,各层通过残差连接相连。每一层包含三个子层:逐位置投影、空间聚合和逐位置变换:

U(X),V(X),Q(X),K(X)=Split ⁣(ϕ1 ⁣(f1(X))).(1)\mathbf{U}(\mathbf{X}),\mathbf{V}(\mathbf{X}),\mathbf{Q}(\mathbf{X}),\mathbf{K}(\mathbf{X}) =\operatorname{Split}\!\left(\phi_1\!\left(f_1(\mathbf{X})\right)\right). \tag{1}
A(X)V(X)=ϕ2 ⁣(Q(X)K(X)T+rabp,t)V(X).(2)\mathbf{A}(\mathbf{X})\mathbf{V}(\mathbf{X}) =\phi_2\!\left(\mathbf{Q}(\mathbf{X})\mathbf{K}(\mathbf{X})^{\mathsf T}+\operatorname{rab}_{p,t}\right)\mathbf{V}(\mathbf{X}). \tag{2}
Y(X)=f2 ⁣(Norm ⁣(A(X)V(X))U(X)).(3)\mathbf{Y}(\mathbf{X}) =f_2\!\left(\operatorname{Norm}\!\left(\mathbf{A}(\mathbf{X})\mathbf{V}(\mathbf{X})\right)\odot\mathbf{U}(\mathbf{X})\right). \tag{3}

其中,fi(X)f_i(\mathbf{X}) 表示 MLP。为降低计算复杂度,我们为 f1f_1f2f_2 各使用一个线性层,即 fi(X)=WiX+bif_i(\mathbf{X})=\mathbf{W}_i\mathbf{X}+\mathbf{b}_i;同时使用融合算子,将查询 Q(X)\mathbf{Q}(\mathbf{X})、键 K(X)\mathbf{K}(\mathbf{X})、值 V(X)\mathbf{V}(\mathbf{X}) 和门控权重 U(X)\mathbf{U}(\mathbf{X}) 的计算进一步批处理。ϕ1\phi_1ϕ2\phi_2 表示非线性函数,二者均采用 SiLU;Norm\operatorname{Norm} 表示层归一化;rabp,t\operatorname{rab}_{p,t} 表示融合位置 pp 与时间 tt 信息的相对注意力偏置。完整符号见表 9。

HSTU 编码器可以用单一模块化块替换 DLRM 中的异构模块。我们观察到 DLRM 实际包含三个主要阶段:特征提取、特征交互和表示变换。

特征提取会读取类别特征经过池化的嵌入表示。其最先进的实现可以统一概括为成对注意力和目标感知池化,而 HSTU 层能够覆盖这些操作。

特征交互是 DLRM 中最关键的部分。常用方法包括因子分解机及其神经网络变体、高阶特征交互等。HSTU 允许经过注意力池化的特征通过 Norm(A(X)V(X))U(X)\operatorname{Norm}(\mathbf{A}(\mathbf{X})\mathbf{V}(\mathbf{X}))\odot\mathbf{U}(\mathbf{X}) 直接与其他特征“交互”,从而替代专门的特征交互模块。这一设计源于用学习得到的 MLP 逼近点积的困难。由于 SiLU 作用于 U(X)\mathbf{U}(\mathbf{X}),上述表达式也可以解释为 SwiGLU 的一种变体。

表示变换通常使用混合专家模型(MoE)与路由机制,以处理多样的异构人群。其核心思想是让不同子网络专门服务不同用户,从而执行条件计算。HSTU 中的逐元素乘积在相差一个归一化因子的情况下,可以近似执行 MoE 所使用的门控操作。

image.png 图 3:DLRM 与 GR 的关键模型组件对比。左侧是完整的 DLRM 设置,右侧是简化后的 HSTU。

3.1 逐位置聚合注意力

HSTU 采用一种新的逐位置聚合归一化注意力机制。作为对比,Softmax 注意力会在整个序列上计算归一化因子。这一设计由两个因素驱动。

第一,与目标相关的先前数据点数量,是表示用户偏好强度的重要特征,而在 Softmax 归一化后很难捕获这一信息。这一点至关重要,因为我们既要预测互动强度,例如用户在某个物品上花费的时间,也要预测物品间的相对顺序,例如为最大化 AUC 而预测排序。第二,虽然 Softmax 激活从设计上对噪声具有稳健性,但它不太适合流式环境中的非平稳词表。

公式(2)给出了逐位置聚合注意力机制。需要强调的是,为稳定训练,必须在逐位置池化后使用层归一化。我们可以通过遵循狄利克雷过程生成的合成数据来理解这一设计,该过程在非平稳词表上生成流式数据,详见附录 C。在这种设置下,如表 2 所示,Softmax 与逐位置注意力的差距最高可达 44.7%。 image.png

3.2 利用稀疏性并通过算法提高稀疏性

在推荐系统中,用户历史序列长度通常服从偏斜分布,因而输入序列较为稀疏,超长序列场景尤其如此。可以利用这种稀疏性显著提高编码器效率。为此,我们开发了一个高效 GPU 注意力算子,以类似内存高效注意力和 FlashAttention 的方式融合前后相接的 GEMM,同时执行完全不规则化的注意力计算。该方法实质上将注意力计算转换为不同规模的分组 GEMM,详见附录 G。

因此,以内存访问次数衡量,HSTU 中自注意力的复杂度为 Θ(ini2dqkR1)\Theta(\sum_i n_i^2d_{qk}R^{-1}),其中 nin_i 是样本 ii 的序列长度,dqkd_{qk} 是注意力维度,RR 是寄存器大小。仅这一方法就能带来 2—5 倍吞吐量提升。

我们还通过随机长度(Stochastic Length,SL)算法进一步提高用户历史序列的稀疏性。推荐中的用户历史具有一个关键特征:用户行为在时间上具有重复性,因为用户行为会在其交互历史的多个尺度上反复出现。这使我们有机会在不损害模型质量的情况下人为提高稀疏性,从而显著降低按 Θ(ini2)\Theta(\sum_i n_i^2) 扩展的编码器成本。

将用户 jj 的历史表示为序列 (xi)i=0nc,j(x_i)_{i=0}^{n_{c,j}},其中 nc,jn_{c,j} 是该用户交互过的内容数。令 Nc=maxjnc,jN_c=\max_j n_{c,j},并令 (xik)k=0L(x_{i_k})_{k=0}^{L} 表示从原序列构造、长度为 LL 的子序列。SL 按以下方式选择输入序列:

{(xi)i=0nc,j,nc,jNcα/2,(xik)k=0Ncα/2,nc,j>Ncα/2,概率为 1Ncαnc,j2,(xi)i=0nc,j,nc,j>Ncα/2,概率为 Ncαnc,j2.(4)\begin{cases} (x_i)_{i=0}^{n_{c,j}}, & n_{c,j}\le N_c^{\alpha/2},\\ (x_{i_k})_{k=0}^{N_c^{\alpha/2}}, & n_{c,j}>N_c^{\alpha/2},\quad\text{概率为 }1-\dfrac{N_c^\alpha}{n_{c,j}^2},\\ (x_i)_{i=0}^{n_{c,j}}, & n_{c,j}>N_c^{\alpha/2},\quad\text{概率为 }\dfrac{N_c^\alpha}{n_{c,j}^2}. \end{cases} \tag{4}

这将注意力相关复杂度降至 O(Ncαd)=O(Nαd)O(N_c^\alpha d)=O(N^\alpha d),其中 α(1,2]\alpha\in(1,2]。关于子序列选择的更完整讨论见附录 F.1。值得强调的是,将 SL 用于训练可以形成高性价比的系统设计,因为训练通常比推理需要多得多的计算。

image.png 表 3:随机长度(SL)对序列稀疏性的影响。

表 3 给出了一个具有 30 天用户历史的代表性工业规模配置中,不同序列长度和 α\alpha 取值下的稀疏性。原表中带下划线和蓝色高亮的设置,对模型质量几乎没有负面影响。标记为 α=2.0\alpha=2.0 的行是不使用 SL 的基础稀疏情形。对于更长的序列,可以采用更小的 α\alpha,直至本文测试的最长序列 8,192。

3.3 最小化激活内存用量

在推荐系统中,大批量对于训练吞吐量和模型质量都至关重要。因此,与通常使用小批量训练、主要受参数内存限制的大语言模型不同,激活内存用量成为推荐模型扩展的主要瓶颈。

相较 Transformer,HSTU 使用更简洁、完全融合的设计,显著降低激活内存用量。首先,HSTU 将注意力之外的线性层从六层减少到两层,这与使用逐元素门控减少 MLP 计算的近期工作一致。其次,HSTU 将多项计算积极融合为单个算子,包括公式(1)中的 ϕ1(f1())\phi_1(f_1(\cdot)),以及公式(3)中的层归一化、可选 Dropout 和输出 MLP。

在 bfloat16 下,这种简化设计把每层激活内存用量降至 2d+2d+4hdqk+4hdv+2hdv=14d2d+2d+4hd_{qk}+4hd_v+2hd_v=14d。作为比较,Transformer 在注意力后使用前馈层与 Dropout,其中间状态大小为 3hdv3hd_v;随后是由层归一化、线性层、激活、线性层和 Dropout 构成的逐位置前馈块,其中间状态为 2d+4dff+2d+1d=4d+4dff2d+4d_{\mathrm{ff}}+2d+1d=4d+4d_{\mathrm{ff}}。在标准假设 hdvdhd_v\ge ddff=4dd_{\mathrm{ff}}=4d 下,再计入输入、输入层归一化和 QKV 投影,总激活状态为 33d33d。因此,HSTU 的设计可以支持深度超过两倍的网络。

此外,用于表示词表的大规模原子 ID 也需要大量内存。若词表规模为 100 亿、嵌入维度为 512,并使用 Adam 优化器,仅以 FP32 存储嵌入和优化器状态就需要 60 TB 内存。为缓解内存压力,我们采用逐行 AdamW 优化器并把优化器状态放在 DRAM 中,将每个浮点数的 HBM 用量从 12 字节降至 2 字节。

3.4 通过成本摊销扩展推理

最后一项挑战是推荐系统在服务阶段需要处理大量候选。本文重点研究排序;对于检索,编码器成本可以完全摊销,而且无论是利用量化、哈希或分区的最大内积搜索(MIPS),还是通过束搜索或分层检索处理的非 MIPS 情形,都已有高效算法。

在排序阶段,候选数量最多可达数万。我们提出 M-FALCON(Microbatched-Fast Attention Leveraging Cacheable OperatioNs)算法,用于在输入序列长度为 nn 时,对 mm 个候选执行推理。在一次前向传播中,M-FALCON 通过修改注意力掩码与 rabp,t\operatorname{rab}_{p,t} 偏置,并行处理 bmb_m 个候选,使这 bmb_m 个候选执行的注意力操作完全相同。

这会将交叉注意力成本从 O(bmn2d)O(b_mn^2d) 降至 O((n+bm)2d)=O(n2d)O((n+b_m)^2d)=O(n^2d);后一等式成立于相较 nn 而言 bmb_m 可视为较小常数时。我们还可以选择将全部 mm 个候选划分为 m/bm\lceil m/b_m\rceil 个大小为 bmb_m 的微批次,利用编码器级 KV 缓存在不同前向传播之间降低成本,或在不同请求之间降低尾延迟。详细讨论见附录 H。

总体而言,M-FALCON 使模型复杂度可以随传统 DLRM 排序阶段的候选数线性扩展。在典型排序配置下,我们成功以固定推理预算部署了复杂度高出 285 倍的目标感知交叉注意力模型,同时获得 1.5—3 倍吞吐量。

4 实验

4.1 验证 HSTU 编码器的归纳假设

4.1.1 传统序列环境

我们首先在两个常用推荐数据集 MovieLens 和 Amazon Reviews 上评估 HSTU。按照序列推荐文献中的设置,使用完全打乱和多轮训练。基线采用最先进的 Transformer 实现 SASRec。与近期工作一致,我们在整个语料库上报告 Hit Rate@K 和 NDCG@K。

表 4:在多遍、完全shuffle设置下,各种方法在公开数据集上的评估结果。 image.png

“SASRec (2023)”表示既有工作报告的最佳 SASRec 配方。标记为“HSTU”的行使用与 SASRec 相同的配置,包括相同的层数、注意力头数等;“HSTU-large”表示更大的 HSTU 编码器,其层数为四倍,注意力头数为两倍。结果表明:第一,在配置相同时,针对推荐任务优化设计的 HSTU 显著优于基线;第二,扩大 HSTU 后,性能还会进一步提高。

需要注意的是,此处使用的评估方法与工业规模设置存在显著差异,因为完全打乱和多轮训练通常不适用于工业界采用的流式环境。

4.1.2 工业规模流式环境

接下来,我们在流式环境中使用工业规模数据集,比较 HSTU、消融版 HSTU 和 Transformer。在本节余下内容中,排序任务统一报告归一化熵(Normalized Entropy,NE)。模型在 1,000 亿个等价 DLRM 样本上训练,每个任务使用 64—256 块 H100 GPU。由于排序采用多任务设置,我们报告主要互动事件“E-Task”和主要消费事件“C-Task”。在本文场景中,NE 降低 0.001 即视为显著,因为这通常能让数十亿用户的核心指标提升 0.5%。

检索设置与语言建模类似,因此报告对数困惑度。受资源限制,我们在较小规模下固定编码器参数,并网格搜索其他超参数:排序使用 l=3,n=2048,d=512l=3,n=2048,d=512,检索使用 l=6,n=512,d=256l=6,n=512,d=256

表 5:one-pass流式设置下,HSTU、消融版 HSTU 和 Transformer 在工业规模数据集上的评估。 image.png

结果首先表明,HSTU 显著优于 Transformer,在排序任务上尤其明显,这可能源于逐位置注意力和改进的相对注意力偏置。其次,消融版 HSTU 与完整 HSTU 之间的差距证实了各项设计的有效性。由于训练稳定性问题,基于 Softmax 的 HSTU 和 Transformer 的最优学习率比其他模型低约 10 倍。即使使用较低学习率和预归一化残差连接,标准 Transformer 在排序中仍频繁发生损失爆炸。

最后,HSTU 还优于 LLM 中常用的 Transformer 变体 Transformer++,后者使用 RoPE、SwiGLU 等技术。总体而言,在这一小规模设置下,HSTU 在质量更高的同时,实际运行速度快 1.5—2 倍,HBM 用量减少 50%。

4.2 编码器效率

随机长度。 图 4 和图 5(a) 展示随机长度对模型指标的影响。当 α=1.6\alpha=1.6 时,长度为 4,096 的序列在大多数时候会变成长度为 776 的序列,即移除超过 80% 的 Token。即使稀疏率提高到 64%—84%,主要任务的 NE 也没有下降超过 0.002(0.2%)。这说明,选择合适的 α\alpha 后,SL 不会损害模型质量,并能利用高稀疏性降低训练成本。附录 F.3 进一步证实,SL 显著优于现有的长度外推技术。

image.png 图 4:随机长度(SL)对指标的影响。左:n=4096n=4096;右:n=8192n=8192。完整结果见附录 F。

编码器效率。 图 5 比较 HSTU 与 Transformer 编码器在训练和推理环境中的效率。Transformer 使用最先进的 FlashAttention-2 实现。我们考察 1,024—8,192 的序列长度,并在训练期间应用 SL。评估中,HSTU 与 Transformer 使用相同配置,即 d=512,h=8,dqk=64d=512,h=8,d_{qk}=64;同时移除相对注意力偏置,因为第 4.1.2 节已经表明,即使不使用 rabp,t\operatorname{rab}_{p,t},HSTU 仍优于 Transformer。

我们在 NVIDIA H100 GPU 上使用 bfloat16 比较编码器级性能。总体而言,HSTU 在训练和推理中分别比 Transformer 最高快 15.2 倍和 5.6 倍。此外,第 3.3 节讨论的激活内存下降,使 HSTU 能够构建深度超过 Transformer 两倍的网络。

image.png 图 5:Encoder-level效率。比较 HSTU 与基于 FlashAttention2 的 Transformer 在训练(a、b)和推理(c)中的表现。

4.3 工业规模流式环境中的生成式推荐器与 DLRM

最后,我们在工业规模流式环境中,将 GR 的端到端性能与最先进的 DLRM 基线比较。GR 实现采用生产中的典型配置;DLRM 配置则凝聚了数百人在多年间的持续迭代。由于推荐系统的检索阶段会使用多个生成器,我们同时报告加入 GR 新检索源和用 GR 替换现有主要 DLRM 检索源的在线结果。

表 6:检索模型的离线与在线比较。 image.png

表 7:排序模型的离线与在线比较。 image.png

结果表明,GR 不仅在线下显著优于 DLRM,还在 A/B 测试中带来 12.4% 的提升。GR 建立在原始类别互动特征之上,而 DLRM 通常使用更多特征,其中大多数由原始信号手工构造。若只向 DLRM 提供 GR 所使用的同一组特征,即“DLRM(特征消融)”,DLRM 表现会显著下降。这说明 GR 可以借助自身架构和统一特征空间,有效捕获这些手工特征所表达的信息。

我们还通过与传统序列推荐器设置比较来验证第 2.2 节的 GR 表述。传统设置只考虑用户交互过的物品,即“GR(仅交互)”。该方案明显更差,其排序变体在主要消费任务上的 NE 比完整 GR 差 2.6%。考虑到基于内容的方法(包括语言模型)十分流行,我们还加入了只使用内容特征的 GR 基线“GR(仅内容)”。它与 DLRM、完整 GR 之间存在巨大性能差距,突显了高基数用户行为的重要性。

最后,图 6 比较 GR 与生产 DLRM 的效率。尽管 GR 模型的计算复杂度高出 285 倍,但在分别为 1,024 和 16,384 个候选打分时,借助 HSTU 和 M-FALCON,GR 的 QPS 仍分别提高到 1.50 倍和 2.99 倍。

image.png 图 6:最具挑战性的排序设置中的推理吞吐量对比。完整结果见附录 H.1。

4.3.1 推荐系统的扩展规律

众所周知,在大规模工业环境中,DLRM 的质量会在达到一定计算量和参数规模后饱和。为了更好理解这一现象,我们比较 GR 与 DLRM 的可扩展性。

由于特征交互层对 DLRM 性能至关重要,在排序设置中,我们尝试使用 Transformer、DHEN,以及生产中采用的带残差连接 DCN 变体来扩大 DLRM 基线。对于检索基线,由于其已经使用残差结构,我们通过扩大隐藏层、嵌入维度和层数来扩展模型。对于基于 HSTU 的 GR,我们调整 HSTU 的残差层数、序列长度、嵌入维度、注意力头数等超参数来扩大模型,并且在检索中调整负样本数。

图 7 给出了结果。在低计算量区间,由于使用手工特征,DLRM 可能优于 GR,这再次说明特征工程对传统 DLRM 的重要性。然而,GR 相对于 FLOPs 展现出明显更好的可扩展性,而 DLRM 的性能会进入平台期。无论按嵌入参数还是非嵌入参数衡量,GR 的扩展性也更好:GR 可以达到 1.5 万亿参数,而 DLRM 的性能在约 2,000 亿参数时已经饱和。

最后,在超参数选择适当时,所有主要指标——包括检索的 Hit Rate@100、Hit Rate@500 和排序的 NE——在经验上都会按照所用计算量的幂律扩展。我们在三个数量级范围内观察到这一现象,直到能够测试的最大模型:序列长度 8,192、嵌入维度 1,024、24 层 HSTU。此时,由于采用标准流式训练设置,按 365 天归一化的总计算量已接近 GPT-3 和 LLaMA-2 的总训练计算量。

在合理范围内,与投入的总训练计算量相比,具体模型超参数的重要性较低。与语言建模不同,序列长度在 GR 中发挥的作用明显更大,因此必须同步扩大序列长度和其他参数。这或许是本文方法最重要的优势,因为我们首次表明,源自 LLM 的扩展规律也可能适用于大规模推荐系统。

image.png 图 7:大规模工业环境中 DLRM 与 GR 的可扩展性对比,包括检索(上图、中图)和排序(下图)。HR 提升 0.005、NE 降低 0.001 均代表显著改善。

5 相关工作

以往的序列推荐器会把用户交互简化为单一的同质物品序列。序列方法在工业规模下的应用主要是成对注意力,或作为 DLRM 一部分的序列编码器。为提高效率,也有研究使用多阶段注意力代替自注意力。检索领域还探索了把 ID 表示为 Token 序列的生成式方法。附录 B.1 对既有工作作了更完整的讨论。

由于自注意力具有 O(n2)O(n^2) 扩展因子,高效注意力一直是重要研究方向,代表性方法包括分解注意力、低秩近似等。近期研究还探索了序列转导问题的替代表述。HSTU 的逐元素门控设计尤其受到 FLASH 的启发。硬件感知表述已被证明可以显著降低内存用量,并改善实际运行时间。长度外推使在短序列上训练的模型能够泛化到更长序列,不过大多数工作侧重微调或改进偏置机制。本文受深度维度随机化方法启发,转而在长度维度引入随机性。

对大语言模型的关注推动了另一类工作:在预训练 LLM 之上,把不同推荐任务视为上下文学习、指令微调或迁移学习。LLM 中蕴含的世界知识可以迁移到下游任务,并改善零样本或少样本推荐。用户行为序列的文本表示也在中等规模数据集上展现出良好的扩展特性。不过,大多数推荐 LLM 研究集中于低数据量环境;在大规模设置下,它们尚未在 MovieLens 上超越协同过滤。

6 结论

本文提出生成式推荐器这一新范式,将排序和检索表述为序列转导任务,使其能够以生成方式训练。新颖的 HSTU 编码器设计使这一范式成为可能:在长度为 8,192 的序列上,HSTU 比最先进的 Transformer 快 5.3—15.2 倍;同时,系统还使用 M-FALCON 等新的训练和推理算法。

借助 GR,我们在使用更少推理计算的情况下部署了复杂度高出 285 倍的模型。GR 和 HSTU 使生产指标提升 12.4%,并且相较传统 DLRM 展现出更强的扩展性能。结果证实,用户行为是生成式建模中一种尚未得到充分探索的模态。正如论文标题所说:“行动胜于言语”。

本文对特征的大幅简化,使推荐、搜索和广告领域的首批基础模型成为可能,因为不同领域可以使用统一的特征空间。GR 的完全序列化设置还使推荐能够转化为端到端生成式问题。这两点都有助于推荐系统以更全面的方式协助用户。

影响声明

我们认为,本文工作具有广泛的积极影响。降低推荐、搜索和广告系统对大量异构特征的依赖,既能改善用户体验,也能让这些系统更加保护隐私。通过完全序列化表述,使推荐系统能够把用户长期结果归因到短期决策,有望减少互联网上不符合用户长期目标的内容,包括标题党和虚假新闻,并使平台激励机制更好地与用户价值保持一致。

最后,基础模型和扩展规律的应用,有助于减少推荐、搜索及相关用例在模型研究与开发中产生的碳足迹。

附录 A 符号说明

本文使用的关键符号汇总于表 8 和表 9。

image.png 表 8:符号表(下页续表)。

image.png 表 9:符号表(续)。

附录 B 生成式推荐器:背景与问题表述

许多读者可能更熟悉经典 DLRM。自 YouTube DNN 时代以来,DLRM 已经十分流行,并广泛应用于每一个大型在线内容平台和电子商务平台。DLRM 运行在异构特征空间之上,使用不同的神经网络,包括特征交互模块、序列池化或目标感知成对注意力模块,以及先进的多专家多任务模块。因此,第 2 节和第 3 节通过与经典 DLRM 的直接对比来概述生成式推荐器。本附录则从经典序列推荐文献出发,提供另一种观察视角。

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

B.1.1 学术研究(传统序列推荐器设置)

GRU4Rec 最早将循环神经网络(RNN)用于推荐场景。该工作在 RecSys Challenge 2015 和专有 VIDEO 数据集上使用门控循环单元(GRU)。两个数据集中都只把正向事件保留在输入序列中,即被点击的电商物品,或用户观看时长达到一定阈值的视频。进一步观察可知,在由检索和排序组成的经典工业两阶段推荐系统中,GRU4Rec 解决的任务主要对应检索任务。

Transformer、序列转导架构及其变体。 此后,序列转导架构,尤其是 Transformer 的进步,也推动了推荐系统的发展。SASRec 首次在自回归环境中应用 Transformer。它将是否存在评论或评分视为正反馈,从而把 Amazon Reviews、MovieLens 等经典数据集转换为正向物品序列,这与 GRU4Rec 类似。该方法使用二元交叉熵损失,把下一个“正向”物品定义为正目标,而从物品语料库 X=Xc\mathcal{X}=\mathcal{X}_c 中随机采样负目标。

后续大多数研究都建立在 GRU4Rec 和 SASRec 的类似设置上,例如 BERT4Rec 采用 BERT 的双向编码器设置,S3Rec 引入显式预训练阶段等。

B.1.2 作为 DLRM 组成部分的工业应用

序列编码器和成对注意力模块等序列方法,能够增强 DLRM 中的用户表示,因此已广泛用于工业环境。DLRM 通常使用相对较短的序列,例如 BST 使用 20,DIN 使用 1,000,TransAct 使用 100。与本文使用的 8,192 相比,这些长度低一到三个数量级。

尽管序列较短,大多数 DLRM 仍能成功捕获长期用户偏好,这可以归因于两点。第一,现代 DLRM 通常使用预计算用户画像或嵌入,或使用外部向量存储,二者都能有效延长回溯窗口。第二,系统通常会使用大量上下文侧、用户侧和物品侧特征,并利用因子分解机、深度交叉网络、混合专家等异构网络转换表示、组合输出。

与附录 B.1.1 中的序列设置不同,所有主要工业方法都会在“用户或请求—候选物品”对上定义损失。排序通常采用多任务二元交叉熵损失;检索则仍以双塔设置为主。近期研究尝试把下一个待推荐物品表示为子 Token 序列上的概率分布,这种设置有时也称为“生成式检索”,并通常使用束搜索从子 Token 解码物品。随着 GPU、自定义 ASIC 和 TPU 等现代加速器普及,混合 Logit 等先进可学习相似度函数也已被提出并部署,作为双塔和束搜索的替代方案。

从问题表述角度看,我们把上述工业工作都归入 DLRM,因为其模型架构、特征和损失,与附录 B.1.1 的学术序列推荐研究存在显著差异。还需要指出的是,在本文之前,工业界尚未成功应用完全序列化的排序设置,更没有达到十亿日活跃用户规模。

B.2 在生成式推荐器中把排序和检索表述为序列转导任务

下面讨论传统序列推荐器设置和 DLRM 设置的三项局限,以及 GR 如何从问题表述角度解决这些问题。

忽略用户交互物品之外的特征。 过去的序列表述只考虑用户明确交互过的内容,而 GR 之前的工业规模推荐系统会在大量特征上训练,以增强用户和内容表示。GR 通过两种方式解决这一限制:第一,压缩其他类别特征,并将其合并到主时间序列;第二,按照第 2.1 节和图 2 的讨论,使用目标感知表述,通过交叉注意力交互捕获数值特征。

我们的实验表明,忽略这些特征的传统“仅交互”表述会显著损害模型质量。表 6 和表 7 中标记为“GR(仅交互)”的结果显示,仅使用交互历史会使检索 Hit Rate@100 降低 1.3%,并使排序的 NE 变差 2.6%;第 4.1.2 节和第 4.3.1 节已经说明,NE 变化 0.1% 即具有显著意义。

用户表示在与目标无关的环境中计算。 第二个问题是,GRU4Rec、SASRec、BERT4Rec、S3Rec 等大多数传统序列推荐器都采用与目标无关的表述。对于目标物品 Φi\Phi_i,它们使用 Φ0,Φ1,,Φi1\Phi_0,\Phi_1,\ldots,\Phi_{i-1} 作为编码器输入来计算用户表示,再据此给出预测。

相较之下,工业界的大多数主要 DLRM 方法都以目标感知方式构建序列模块,可以把“目标”,即排序候选的信息纳入用户表示,包括 DIN、BST、TWIN 和 TransAct。GR 通过交错排列内容序列和行为序列,将两类方法的优势结合起来,使因果自回归环境能够应用目标感知注意力。表 10 对本文与既有工作进行了分类比较。大多数大规模工业推荐器因日志数据量巨大,必须在流式、单遍设置中训练。

image.png 表 10:排序设置下既有序列推荐方法与 GR 的比较;为完整起见,表中也包括 DLRM。

判别式表述把既有序列推荐方法限制在逐点设置中。 最后,传统序列推荐器从设计上就是判别式的。GRU4Rec、SASRec 等经典方法对 p(ΦiΦ0,a0,,Φi1,ai1)p(\Phi_i\mid\Phi_0,a_0,\ldots,\Phi_{i-1},a_{i-1}) 建模,即根据用户当前状态,建模下一项推荐物品的条件分布。

另一方面,标准推荐系统中存在两个概率过程:推荐系统向用户建议内容 Φi\Phi_i,例如图片或视频;用户再通过行为 aia_i 响应所建议的内容,其中行为可以是点赞、完整观看、跳过等的组合。

生成式方法需要对建议内容与用户行为组成的序列联合分布建模:

p(Φ0,a0,Φ1,a1,,Φnc1,anc1).p(\Phi_0,a_0,\Phi_1,a_1,\ldots,\Phi_{n_c-1},a_{n_c-1}).

第 2.2 节对此已有讨论。生成式推荐器可以对该分布进行建模,如表 11 和图 8 所示。下一个行为 Token aia_i 的预测任务,正是表 1 所述的 GR 排序设置;下一个内容 Φi\Phi_i 的预测任务则类似于适配到交错设置后的检索任务,其目标经过调整,以学习输入数据分布。

image.png 表 11:对 p(Φ0,a0,,Φnc1,anc1)p(\Phi_0,a_0,\ldots,\Phi_{n_c-1},a_{n_c-1}) 进行生成式建模。图 8 给出了示意图。

这一表述不仅能够正确建模数据分布,还可以通过束搜索等方法直接采样待推荐的物品序列。我们推测,它将优于传统列表式设置,例如确定点过程和强化学习;这类系统的完整表述与评估留待未来研究,第 6 节对此进行了简要讨论。

image.png 图 8:传统序列推荐器(左)与生成式推荐器(右)的比较。为便于对比,图中使用因果自回归设置下的序列推荐器,以及不含上下文特征的 GR。左侧方法在进入自注意力块之前,要么忽略行为类型 aia_i,要么使用 MLP 将其与物品信息 Φi\Phi_i 组合。

附录 C 评估:合成数据

如第 3.1 节所述,标准 Softmax 注意力具有归一化因子,因此难以捕获用户偏好强度,而这一信息对用户表示学习非常重要。在推荐场景中,系统除了预测物品的相对顺序,还可能需要预测互动强度,例如未来对某一主题产生正向行为的次数。

为理解这一行为,我们按照狄利克雷过程构造合成数据,在动态词表集合上生成流式数据。狄利克雷过程能够刻画用户互动历史中的“富者愈富”现象。合成实验设置如下:

  • 将 20,000 个物品 ID 分别随机分配到 100 个类别中,每个物品只属于一个类别。
  • 生成 1,000,000 条记录,每条长度为 128;前 90% 用于训练,后 10% 用于测试。为模拟流式训练,初始只提供 40% 的物品 ID,其余 ID 按等间隔逐步开放。例如,在第 500,000 条记录处,可采样的最大 ID 为 (40%+60%×0.5)×20,000=14,000(40\%+60\%\times0.5)\times20{,}000=14{,}000
  • 对每条记录,从 100 个类别中随机选择最多 5 个类别,并在这 5 个类别上随机采样先验 HcH_c。随后按以下狄利克雷过程,依次为每个位置采样类别:
  • n>1n>1 时,以概率 α/(α+n1)\alpha/(\alpha+n-1)HcH_c 中抽取类别 cc;以概率 nc/(α+n1)n_c/(\alpha+n-1) 抽取类别 cc,其中 ncn_c 是此前类别为 cc 的物品数;然后在流式约束下,从分配到类别 cc 的物品中随机采样一个物品。
  • α\alpha(1.0,500.0)(1.0,500.0) 上的均匀分布中随机采样。

结果见表 2。由于该数据集没有时间戳,我们始终从 HSTU 中移除 rabp,t\operatorname{rab}_{p,t}。相较标准 Transformer,HSTU 的 Hit Rate@10 提高超过 100%。更重要的是,用 Softmax 替换 HSTU 的逐位置注意力机制后,命中率也会显著下降,从而验证了类似逐位置注意力的聚合机制的重要性。

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

第 4.1.1 节主要使用最新训练配方,将 HSTU 与最先进的 Transformer 基线 SASRec 比较。本节进一步考察另外两种方法。

循环神经网络。 我们采用经典序列推荐方法 GRU4Rec,以帮助读者理解:在完整纳入最新建模和训练改进后,Transformer、HSTU 等自注意力模型与传统 RNN 相比表现如何。

自监督序列方法。 我们采用最具代表性的 BERT4Rec,研究通过 Cloze 目标实现的双向自监督,与 SASRec、HSTU 等单向因果自回归设置相比表现如何。

image.png 表 12:传统序列推荐器设置(multi-pass, full-shuffle)下,各种方法在公开数据集上的评估。与表 4 相比,额外加入 GRU4Rec 和 BERT4Rec 两个基线,以保证完整性。

表 12 的 BERT4Rec、GRU4Rec 在 ML-1M 和 ML-20M 上的结果沿用既有报告。由于使用采样 Softmax 损失,为公平比较各方法,我们保持负样本数不变:ML-1M 和 ML-20M 使用 128 个,Amazon Books 使用 512 个。

结果证实,当使用采样 Softmax 损失时,SASRec 仍是传统序列推荐设置中最具竞争力的方法之一;HSTU 则显著优于所评估的 Transformer、RNN 和自监督双向 Transformer。

附录 E 评估:传统 DLRM 基线

第 4 节使用的 DLRM 基线配置,来自数百名研究人员和工程师历经多年的持续迭代,并且高度近似 HSTU 和 GR 部署之前,某拥有数十亿日活跃用户的大型互联网平台所采用的生产配置。下面对这些模型作高层说明。

排序设置。 基线排序模型使用约一千种稠密特征和五十种稀疏特征。我们纳入了多种建模技术,包括混合专家模型、深度交叉网络的不同变体、目标感知成对注意力等序列推荐模块,以及特殊交互层上的残差连接。在扩展规律研究的低 FLOPs 区间,某些计算成本较高的模块会被简化,或替换为 DCN 等其他最先进变体,以达到所需 FLOPs。

出于保密要求,我们无法披露确切设置;但据我们所知,在充分纳入近期研究成果后,该基线代表当前已知效果最佳的 DLRM 方法之一。为验证这一点并便于读者理解,我们在表 7 中报告了一个使用完全相同特征、但仅采用 DIN、DCN 和 MMoE 等主要公开成果的典型设置“DLRM(DIN+DCN)”。图 9 展示组合架构。该设置在主要 E-Task 和 C-Task 上的 NE,分别比生产 DLRM 差 0.71% 和 0.57%,而 NE 变化 0.1% 即具有显著意义。

image.png 图 9:基线 DLRM 排序模型的高层架构,即表 7 中的“DLRM(DIN+DCN)”。该模型采用 DIN、DCN 和 MMoE 等主要公开工作。

检索设置。 基线检索模型采用标准的双塔神经检索设置,并混合使用批内采样与批外采样。输入特征既包括物品 ID、用户 ID 等高基数稀疏特征,也包括语言、主题、兴趣实体等低基数稀疏特征。系统使用带残差连接的多层前馈网络,将输入特征压缩为用户嵌入和物品嵌入。

特征与序列长度。 两个 DLRM 基线所用特征,严格包含所有 GR 候选模型使用的特征;其中也包括由不同序列编码器或成对注意力模块使用的主要用户交互历史。本文所有研究都满足这一点,包括扩展规律研究。

附录 F 随机长度

F.1 子序列选择

在公式(4)中,我们从完整用户历史中选择长度为 LL 的子序列,以提高稀疏性。经验结果表明,仔细设计子序列选择方法能够改善模型质量。定义 fi=tntif_i=t_n-t_i,表示用户与物品 xix_i 交互后经过的时间。我们对以下子序列选择方法进行离线实验:

  • 贪心选择:SS 中选择 fif_i 最小的 LL 个物品。
  • 随机选择:SS 中随机选择 LL 个物品。
  • 特征加权选择: 按加权分布 1fn,i/(j=1Lfj,i)1-f_{n,i}/(\sum_{j=1}^{L}f_{j,i}),从 SS 中选择 LL 个物品。

离线实验中,特征加权子序列选择获得了最佳模型质量,如表 13 所示。

image.png 表 13:随机长度采用不同子序列选择方法时的模型质量比较,以归一化熵(NE)衡量。

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

表 3 展示了在具有 30 天用户互动历史的代表性工业规模配置中,随机长度对序列稀疏性的影响。序列稀疏性定义为:1 减去所有样本平均序列长度与最大序列长度之比。为了更准确描述稀疏注意力的计算成本,我们还定义 s2s_2,即 1 减去注意力矩阵的稀疏度。作为补充,表 14 和表 15 分别给出 60 天与 90 天用户互动历史的结果。

image.png 表 14:在 60 天用户互动历史上,随机长度(SL)对序列稀疏性的影响。

image.png 表 15:在 90 天用户互动历史上,随机长度(SL)对序列稀疏性的影响。

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

我们进行了额外研究,以验证随机长度与语言建模中现有序列长度外推技术相比是否具有竞争力。许多现有方法通过修改 RoPE 实现序列长度外推。为与这些方法比较,我们训练一个不使用相对注意力偏置、但采用旋转嵌入的 HSTU 变体 HSTU-RoPE。

在 HSTU-RoPE 上评估以下序列长度外推方法:

  • 零样本: 应用 NTK-Aware RoPE 后,不进行微调,直接评估模型。
  • 微调: 应用 NTK-by-parts 后,对模型微调 1,000 步。

在包含相对注意力偏置、不使用旋转嵌入的 HSTU 上评估以下方法:

  • 零样本: 按最大训练序列长度截断相对位置偏置,然后直接评估模型。
  • 微调: 按最大训练序列长度截断相对位置偏置,微调 1,000 步后再评估模型。

image.png 图 10:随机长度(SL)对排序模型指标的影响。从左到右,n=[1024,2048,4096,8192]n=[1024,2048,4096,8192];为在因果掩码环境中实现target-aware cross attention,此处的 nn 是经过第 2.2 节所述interleaving算法处理后的长度。

image.png 表 16:随机长度(SL)与现有长度外推方法的比较。

表 16 报告在训练期间人为引入数据稀疏性的模型,包括随机长度、零样本和微调方法,与完整数据训练模型之间的 NE 差异。对零样本和微调方法,我们把稀疏性定义为训练时平均序列长度除以评估时最大序列长度。所有零样本和微调模型都在长度为 1,024 的序列数据上训练,再用长度为 2,048 和 4,096 的数据评估。为了选取合适的随机长度基线,我们采用能够产生相同数据稀疏度指标的 SL 设置。

我们认为,零样本和微调式序列长度外推不适合处理高基数 ID 的推荐场景。经验结果显示,随机长度显著优于微调和零样本方法。这可能源于庞大的词表规模:零样本和微调方法无法为较早出现的 ID 学到良好表示,因而不能充分利用更长序列中包含的信息。

附录 G 稀疏分组 GEMM 与融合相对注意力偏置

本节补充介绍第 3.2 节提出的高效 HSTU 注意力算子。该方法建立在内存高效注意力和 FlashAttention 之上,是一种节省内存的自注意力机制。它将输入划分为块,并在反向传播中避免物化巨大的 h×N×Nh\times N\times N 中间注意力张量。

通过利用输入序列的稀疏性,我们可以把注意力计算重新表述为一组形状不同、前后相接的 GEMM,并实现高效 GPU 算子来加速这一计算。相对注意力偏置的构造也会因内存访问而成为瓶颈。为解决该问题,我们将相对偏置构造与分组 GEMM 融合进同一个 GPU 算子,并在反向传播时使用 GPU 的高速共享内存累积梯度。尽管该算法需要在反向传播中重新计算注意力和相对偏置,它仍显著快于 Transformer 使用的标准方法,并且内存用量更少。

附录 H M-FALCON:利用可缓存操作的微批次快速注意力

本节详细介绍第 3.4 节讨论的 M-FALCON 算法,并在算法 1 中给出伪代码。M-FALCON 包含三个关键思想。

批量推理可以应用于因果自回归环境。 第 2.2 节已经说明,GR 中的排序任务采用目标感知表述。通常认为,在目标感知设置中必须逐物品推理;对于 mm 个候选和长度为 nn 的序列,其成本为 O(mn2d)O(mn^2d)。这里我们说明,这并非最优解。即使使用普通 Transformer,也可以修改自注意力的掩码,把这些操作批量执行,将成本降至 O((n+m)2d)=O(n2d)O((n+m)^2d)=O(n^2d)

图 11 给出了示意图。其中,图 11(a) 和图 11(b) 都使用因果自回归注意力掩码矩阵。关键区别是:图 11(a) 在因果训练中使用大小为 2nc2n_c 的标准下三角矩阵;图 11(b) 则修改大小为 2nc+bm2n_c+b_m 的下三角矩阵,把满足 i,j2nci,j\ge2n_ciji\ne j 的位置 (i,j)(i,j) 设为 False 或 -\infty,以阻止目标位置 Φ0,,Φbm1\Phi'_0,\ldots,\Phi'_{b_m-1} 彼此关注。

这样,自注意力块对 Φi,ai\Phi'_i,a'_i 的输出只依赖 Φ0,a0,,Φnc1,anc1\Phi_0,a_0,\ldots,\Phi_{n_c-1},a_{n_c-1},而不依赖 Φj\Phi'_j,其中 iji\ne j。换言之,使用修改后的注意力掩码对 2nc+bm2n_c+b_m 个 Token 执行一次前向传播,可以为最后 bmb_m 个 Token 得到这样的结果:它与分别执行 bmb_m 次、每次处理 2nc+12n_c+1 个 Token 的前向传播完全相同。在第 ii 次独立前向传播中,Φi\Phi'_i 位于从 0 开始计数的第 2nc2n_c 个位置,并使用标准因果注意力掩码。

微批次将批量推理扩展到大型候选集合。 排序阶段可能需要处理多达数万条候选。我们可以把全部 mm 个候选划分为 m/bm\lceil m/b_m\rceil 个大小为 bmb_m 的微批次,使 O(bm)=O(n)O(b_m)=O(n)。这样,在大多数实际推荐设置中,即使候选达到数万,仍能保持前述 O((n+bm)2d)=O(n2d)O((n+b_m)^2d)=O(n^2d) 的运行时间。

编码器级缓存支持在请求内部和请求之间共享计算。 最后,KV 缓存既可以在同一请求内部使用,也可以跨请求使用。以本文 HSTU 模型为例,K(X)\mathbf{K}(\mathbf{X})V(X)\mathbf{V}(\mathbf{X}) 可以在不同微批次和不同请求间完整缓存。

对于使用缓存的前向传播,只需为最后 bmb_m 个 Token 计算 U(X)\mathbf{U}(\mathbf{X})Q(X)\mathbf{Q}(\mathbf{X})K(X)\mathbf{K}(\mathbf{X})V(X)\mathbf{V}(\mathbf{X}),并复用长度为 nn 的序列化用户历史所对应的缓存 K(X)\mathbf{K}(\mathbf{X})V(X)\mathbf{V}(\mathbf{X})。同样,f2(Norm(A(X)V(X))U(X))f_2(\operatorname{Norm}(\mathbf{A}(\mathbf{X})\mathbf{V}(\mathbf{X}))\odot\mathbf{U}(\mathbf{X})) 也只需针对 bmb_m 个候选重新计算。

这会把缓存前向传播的计算复杂度降至 O(bmd2+bmnd)O(b_md^2+b_mnd)。即使 bm=nb_m=n,相较 O((n+bm)d2+(n+bm)2d)O((n+b_m)d^2+(n+b_m)^2d) 也能改善 2—4 倍。

image.png 图 11:M-FALCON 算法示意图。上:采用 GR 目标感知表述的模型训练。下:对 mm 个候选 Φ0,,Φm1\Phi'_0,\ldots,\Phi'_{m-1} 执行模型推理,这些候选被划分为 m/bm\lceil m/b_m\rceil 个微批次。虚线上方展示第一个微批次 Φ0,,Φbm1\Phi'_0,\ldots,\Phi'_{b_m-1} 的推理;计入 Φ0,a0,,Φnc1,anc1\Phi_0,a_0,\ldots,\Phi_{n_c-1},a_{n_c-1} 后,共有 2nc+bm2n_c+b_m 个 Token。自注意力经过修改,使 iji\ne jΦi\Phi'_i 不能关注 Φj\Phi'_j,图中以“×”标出。

算法 1:M-FALCON 算法

image.png

  1. 输入: 合并后的 Token 序列 x0,x1,,xn1x_0,x_1,\ldots,x_{n-1},例如 (Φ0,a0,,Φnc1,anc1)(\Phi_0,a_0,\ldots,\Phi_{n_c-1},a_{n_c-1}),此时 n=2ncn=2n_cmm 个排序候选 Φ0,,Φm1\Phi'_0,\ldots,\Phi'_{m-1};一个在因果自回归环境中训练、具有 bb 层和 hh 个注意力头的自注意力模型(如 HSTU 或 Transformer):f(X,cacheStates,attnMask)(X,updatedCacheStates)f(\mathbf{X},\text{cacheStates},\text{attnMask})\rightarrow(\mathbf{X}',\text{updatedCacheStates})。其中 X,XRN×d\mathbf{X},\mathbf{X}'\in\mathbb{R}^{N\times d}attnMaskRN×N\text{attnMask}\in\mathbb{R}^{N\times N};由于在 bb 层间缓存 K(X)\mathbf{K}(\mathbf{X})V(X)\mathbf{V}(\mathbf{X}),缓存状态位于两个 Rb×h×N×dqk\mathbb{R}^{b\times h\times N\times d_{qk}} 空间中。微批次大小为 bmb_m;为简化表述,假设 mmbmb_m 的整数倍。
  2. 输出: 全部 mm 个排序候选的预测 (a0,,am1)(a'_0,\ldots,a'_{m-1})
  3. numMicrobatches=(m+bm1)//bm\text{numMicrobatches}=(m+b_m-1)\mathbin{//}b_m
  4. attnMask=Ln+bm\text{attnMask}=\mathbf{L}_{n+b_m},其中 Ln+bm\mathbf{L}_{n+b_m} 表示下三角矩阵,下三角位置为 0,其余位置为 -\infty
  5. 对所有满足 i,jni,j\ge niji\ne j 的位置,令 attnMask[i,j]=\text{attnMask}[i,j]=-\infty,以阻止最后 bmb_m 个元素彼此关注。
  6. 执行 (a0,a1,,abm1),kvCachef(embLayer(x0,x1,,xn1,Φ0,,Φbm1),,attnMask)(a'_0,a'_1,\ldots,a'_{b_m-1}),\text{kvCache}\leftarrow f(\operatorname{embLayer}(x_0,x_1,\ldots,x_{n-1},\Phi'_0,\ldots,\Phi'_{b_m-1}),\varnothing,\text{attnMask})
  7. predictions=(a0,a1,,abm1)\text{predictions}=(a'_0,a'_1,\ldots,a'_{b_m-1})
  8. i=1i=1
  9. i<numMicrobatchesi<\text{numMicrobatches} 时,重复以下步骤。
  10. 执行 (abmi,abmi+1,,abm(i+1)1)f(embLayer(x0,x1,,xn1,Φbmi,,Φbm(i+1)1),kvCache,attnMask)(a'_{b_mi},a'_{b_mi+1},\ldots,a'_{b_m(i+1)-1})\leftarrow f(\operatorname{embLayer}(x_0,x_1,\ldots,x_{n-1},\Phi'_{b_mi},\ldots,\Phi'_{b_m(i+1)-1}),\text{kvCache},\text{attnMask})
  11. predictionspredictions+(abmi,abmi+1,,abm(i+1)1)\text{predictions}\leftarrow\text{predictions}+(a'_{b_mi},a'_{b_mi+1},\ldots,a'_{b_m(i+1)-1})
  12. ii+1i\leftarrow i+1
  13. 结束循环。
  14. 返回 predictions\text{predictions}

图 11 对算法 1 作了说明。M-FALCON 不仅适用于 HSTU 和 GR,也可以广泛用作其他目标感知因果自回归设置的推理优化算法,只要其基于自注意力架构即可。

H.1 推理吞吐量评估:采用 M-FALCON 的生成式推荐器与 DLRM

如第 3.4 节所述,M-FALCON 会并行处理 bmb_m 个候选,在推理时将计算成本摊销到全部 mm 个候选上。为理解这一设计,我们在相同硬件设置下比较 GR 与 DLRM 的吞吐量,即每秒打分的候选数 QPS。

图 12 和图 13 表明,由于批量推理可以摊销成本,在候选数 mm 达到某一区间之前,GR 的吞吐量会随排序阶段候选数次线性扩展。在本文案例中,这一边界为 m=2048m=2048。这证实,批量推理在因果自回归环境中至关重要。

由于注意力复杂度按 O((n+bm)2)O((n+b_m)^2) 扩展,仅使用多个微批次就能提高吞吐量;在微批次之上加入缓存,还能进一步消除冗余的线性计算和注意力计算。如图 13 所示,两者结合后,相较于只使用一个微批次、满足 bm=m=1024b_m=m=1024 的基线,最多可额外加速 1.99 倍。

总体而言,凭借高效 HSTU 编码器设计和 M-FALCON,尽管基于 HSTU 的 GR 在 FLOPs 上复杂 285 倍,它在大规模生产环境中的吞吐量仍最多比 DLRM 高 2.99 倍。

image.png 图 12:端到端推理吞吐量。在大规模工业环境中比较 DLRM 与采用 M-FALCON 的 GR。本图与图 6 相同,为方便阅读在此再次给出。

image.png 图 13:M-FALCON 的端到端推理吞吐量扩展。在 285 倍 FLOPs 的 GR 模型上,考察大批量设置;排序候选总数 mm 从 1,024 变化到 16,384,且 bm=1024b_m=1024