ULTRA-HSTU

0 阅读5分钟

标题:重塑大规模推荐系统的Scaling Law曲线
单位:meta
链接:Bending the Scaling Law Curve in Large-Scale Recommendation Systems(2026.02)

摘要

通过序列模型从用户交互历史中学习,已经成为大规模推荐系统的基石。大语言模型近期展现出的扩展规律前景,引发了推荐领域对长序列建模与更深架构的研究热潮。然而,许多近期方法高度依赖交叉注意力来缓解序列建模的二次计算瓶颈,这可能限制自注意力所能带来的表示能力。

本文提出 ULTRA-HSTU,这是一种通过端到端模型—系统协同设计得到的新型序列推荐模型。ULTRA-HSTU 在输入序列、稀疏注意力机制和模型拓扑设计方面进行创新,从而显著提升模型质量和效率。全面的基准测试表明,相较传统模型,ULTRA-HSTU 的训练扩展速度快 5 倍以上,推理扩展速度快 21 倍,同时提供更优的推荐质量。该方案已全面投入大规模生产,每天服务数十亿用户,并在真实生产环境中使内容消费和互动指标显著提升 4%—8%。

1 引言

在 GPU 计算规模不断扩大的时代,基于 Transformer 的序列建模已成为推进大规模推荐研究的新范式。传统深度学习推荐模型(Deep Learning Recommendation Model,DLRM)专注于利用精心设计的人工特征,对用户与物品的特征交互进行建模。虽然有效,但当计算量用于更多特征交互或额外网络层时,这类模型无法高效扩展。

相比之下,基于 Transformer 的序列建模强调从原始用户行为序列中进行端到端学习,能够同时捕获长期偏好与短期意图,并且相对于计算量展现出良好的扩展规律:使用更长序列、在注意力层中执行更稠密的计算,以及增加堆叠注意力层的深度,都可以提升模型表现。

这一方向的代表性研究是分层序列转导单元(Hierarchical Sequential Transduction Unit,HSTU)。它引入了一种针对推荐系统定制的 Transformer 风格架构,可以直接从原始序列数据中高效学习用户兴趣。HSTU 的突出意义在于,它首次证明专门面向推荐系统的类 Transformer 方法具有良好扩展特性。此后,抖音、美团、阿里巴巴、小红书、Meta 和 LinkedIn 等主要工业平台广泛采用并继续推进序列建模范式,分别提出自己的架构创新。这些大型平台的广泛采用,凸显了序列建模对大规模推荐系统的有效性和影响力。

不过,包括 HSTU 在内的 Transformer 推荐模型,都因自注意力机制而具有 O(L2)O(L^2) 复杂度,其中 LL 是用户历史序列长度。当尝试对包含 O(10k)O(10\mathrm{k})O(100k)O(100\mathrm{k}) 个事件的用户历史建模时,这种二次扩展很快变得不切实际,尤其是在每天需要以亚秒延迟提供数十亿次推荐的环境中。

为缓解二次计算瓶颈,工业界此前的方法主要采用交叉注意力,只使用排序候选或截断后的用户历史作为查询,而不使用考虑完整用户历史的自注意力。另一些方法则局限于浅层架构,只使用 2—4 层注意力。这些策略从根本上偏离了大语言模型的实践。虽然它们能显著降低计算复杂度,却可能放弃强大自注意力机制和深层模型架构带来的收益。

我们的实验表明,在工业环境中,自注意力仍优于交叉注意力,在堆叠网络层或扩大计算量时尤其如此。这是本文的一项关键研究发现,也体现了本文与既有方案的重要区别:我们并不消除自注意力,而是受到 DeepSeek-V2 的启发,通过模型与系统协同优化,高效发挥自注意力的优势。

为改变超长用户历史建模的扩展效率,我们提出下一代 HSTU 模型 ULTRA-HSTU,其中包含一整套细致的模型与系统优化。这里,扩展效率正式定义为模型表现与计算成本之间拟合线性回归的斜率。在输入序列配置固定时,相较原始 HSTU,本文优化使推理扩展效率提高 21 倍以上,训练扩展效率提高 5 倍。这项进展有效改变了推荐系统的扩展曲线,使增加计算资源时模型质量能够更快提升。

image.png 图 1:总体表现。模型表现相对于训练 FLOP(左)和推理 FLOP(右)的扩展情况。相较普通 HSTU,ULTRA-HSTU 的训练scaling效率提高 5.3 倍以上,推理scaling效率提高 21.4 倍。

为验证所提方案,我们在服务数十亿用户的大规模生产环境中部署了 ULTRA-HSTU。该模型包含 18 层自注意力,在 16K 用户行为序列上使用数百块 H100 GPU 训练,使内容消费和互动指标提升 4%—8%,核心顶层指标提升 0.217%。这既体现了推荐领域序列建模的扩展潜力,也证实了本文方案的有效性。据我们所知,ULTRA-HSTU 是工业界部署过的最大序列模型之一,并显著改善了扩展效率。

ULTRA-HSTU 的技术创新概括如下。

输入序列优化。 我们提出两项互补设计来优化原始 HSTU 的输入序列处理。第一,在序列设计中有效合并物品表示与行为表示,并使用异构行为编码增强这一简化设计。第二,为缓解同步分布式训练中不同 Rank 序列长度不平衡导致的低效率,我们提出负载均衡随机长度(Load-Balanced Stochastic Length,LBSL)。该方法在随机长度采样时施加逐 Rank 计算负载约束,减少拖慢整体进度的 Rank,并将训练吞吐量提高 15%。

面向极高效注意力的模型—系统协同设计。 我们提供端到端模型—系统协同设计,消除常见的二次复杂度和算子开销,使 HSTU 中的自注意力能够用于生产环境的超长用户交互历史建模。

在模型侧,我们提出针对用户行为序列结构定制的半局部注意力(Semi-Local Attention,SLA)。它以 O((K1+K2)L)O((K_1+K_2)\cdot L) 的复杂度实现高效线性稀疏注意力,同时不牺牲模型质量;其中 K1K_1K2K_2 分别是局部窗口和全局窗口大小。SLA 相较基线模型将推理扩展效率提高 5 倍以上。

在系统侧,我们为 SLA 配置精细调优的硬件感知优化,消除实际瓶颈,并提高训练和推理中的硬件利用率。我们协同设计了一套面向推荐系统、覆盖 16/8/4 位格式的混合精度框架:大多数操作保留 BF16 以确保稳定性,主要 GEMM 使用 FP8 加速,推理通信则通过 INT4 嵌入量化减少流量。

我们还扩展 FlashAttention V3 的思想,构建支持 HSTU 基于 SiLU 的注意力和非标准掩码的自定义 SLA 算子,并针对 NVIDIA H100 和 AMD MI300 等异构 GPU 架构进行调优,以保持较高 GPU 利用率。此外,我们还引入开销极低的内存优化,显著减少 HBM 占用,使超长序列训练成为可能。这些协同设计组件合在一起,相较未使用系统优化的同一模型,使训练吞吐量提高 70%,推理吞吐量提高 50%。为了最大化端到端表现,我们关注在模型效果相同时,完成固定数量样本的训练或推理有多快,而不只是优化 GPU 利用率。

动态拓扑模型设计。 推荐模型的可扩展性不只取决于序列长度。通过堆叠额外层实现纵向扩展也会带来收益,特别是利用残差连接增加容量。不过,朴素堆叠采用 SLA 的 HSTU 会产生 O(DL)O(DL) 成本,其中 DD 是模型深度。

基于不同用户信号具有不同预测价值这一观察,我们提出两种新拓扑,把计算集中于最重要的信号:第一,注意力截断,先让前 N1N_1 层处理完整序列,再选出一个较短的高价值片段,只对该片段应用额外 N2N_2 层;第二,转导器混合(Mixture of Transducers,MoT),把异构行为信号作为多条序列交给不同转导器处理,再融合其表示。这样可以针对高价值信号分配容量与计算,而无须让所有信号在同一时间线中竞争。实验表明,两种拓扑都显著改善了性能—效率权衡,并进一步增强模型的扩展能力。

2 相关工作

传统工业规模推荐模型通常遵循 DLRM 框架,重点建模用户和物品的特征交互。过去几年,工业界训练大规模推荐模型的方式发生了范式转变。近期许多进展不再依赖用户—物品交叉特征,而是来自对用户交互历史的学习。DIN 是经典的短序列学习方法之一,SASRec 是用于推荐的传统 Transformer 实现,HSTU 随后被提出,借助目标感知预测在推荐任务中优于传统 Transformer 模型。通过捕获从原始用户交互历史中学到的隐式与显式信息,HSTU 消除了对人工用户—物品特征的依赖,并展现出良好的扩展规律。

沿着这一方向,本文专注于进一步改善扩展行为,目标是以更低计算成本获得更好的模型。与本文密切相关的是提高序列模型训练与推理效率的研究。随着原生稀疏注意力取得突破,线性稀疏注意力成为部署可扩展大模型的重点。

除稀疏注意力外,堆叠式目标到历史交叉注意力(Stacked Target-to-History Cross Attention,STCA)只执行以排序目标为查询的计算,显著降低了模型复杂度。但由于简化后的注意力机制不含自注意力,它会造成性能下降。虽然 STCA 实现了线性复杂度,却为提升性能引入了成本更高的注意力前投影,带来很大计算开销,因此在较短序列上捕获信息时效果较差。

3 背景

如图 2(a) 所示,典型推荐系统接收输入特征,并在多任务分类问题上训练。形式化地说,它学习一个多任务模型 M\mathcal{M},针对候选 xjx_j 在不同预测任务 yky_k 上输出概率,例如点赞、完整观看视频和评论:

y^k=M(X,xj)[0,1].\hat{y}_k=\mathcal{M}(\mathbf{X},x_j)\in[0,1].

系统再根据预测分数对候选排序。这里,X\mathbf{X} 是用户输入特征。我们通过最小化预测 y^k\hat{y}_k 与日志中真实标签 yky_k 之间的交叉熵损失来优化模型。全文使用 LL 表示一般序列长度,X\mathbf{X} 表示输入特征。大多数生成式排序范式把输入 X\mathbf{X} 建模为嵌入序列,再利用注意力层从这些序列嵌入中学习概率。

image.png 图 2:模型设计概览。(a)通用推荐模型设计。(b)采用action-aware设计的输入序列优化。(c)具有线性复杂度的Semi-local attention mask。(d)用于动态拓扑设计的注意力截断。

输入。 推荐器使用特征预处理器,把不同输入特征转换为嵌入序列。用户交互历史(User Interaction History,UIH)序列记录特定用户交互过的物品、相应行为(例如点赞、评论、完整观看视频)和上下文(例如时间戳)。原始物品 ID 及其多模态表示和行为类型,都通过嵌入表查询表示为 dd 维可学习嵌入。

对于用户 ii,将其 UIH 表示为 Xi={Ii,Ai}\mathbf{X}_i=\{\mathbf{I}_i,\mathbf{A}_i\},其中物品嵌入为 Ii={Ii,j}j=1LiRLi×d\mathbf{I}_i=\{\mathbf{I}_{i,j}\}_{j=1}^{L_i}\in\mathbb{R}^{L_i\times d},行为嵌入为 Ai={ai,j}j=1LiRLi×d\mathbf{A}_i=\{\mathbf{a}_{i,j}\}_{j=1}^{L_i}\in\mathbb{R}^{L_i\times d}LiL_i 是用户 ii 的 UIH 总长度。

非序列特征包括国家、用户语言等用户侧特征,以及物品原始 ID 等稀疏物品侧特征和物品点击率等稠密物品侧特征。用户侧特征可以汇总为上下文嵌入,并放在 UIH 序列开头;物品侧特征可以汇总为物品嵌入,并作为目标侧嵌入插入序列。

模型。 给定嵌入序列,现代推荐器会采用 Transformer 风格模型。典型架构 HSTU 通过以下修改,在推荐系统中显著优于普通 Transformer:

归一化:X=Norm(Z).(1)\text{归一化:}\quad \mathbf{X}=\operatorname{Norm}(\mathbf{Z}). \tag{1}
注意力前计算:U,Q,K,V=ϕ1 ⁣(f1(X)).(2)\text{注意力前计算:}\quad \mathbf{U},\mathbf{Q},\mathbf{K},\mathbf{V}=\phi_1\!\left(f_1(\mathbf{X})\right). \tag{2}
注意力:A=(ϕ2 ⁣(QKT)M)V.(3)\text{注意力:}\quad \mathbf{A}=\left(\phi_2\!\left(\mathbf{Q}\mathbf{K}^{\mathsf T}\right)\odot\mathbf{M}\right)\mathbf{V}. \tag{3}
注意力后计算:Y=f2 ⁣(Norm(A)U).(4)\text{注意力后计算:}\quad \mathbf{Y}=f_2\!\left(\operatorname{Norm}(\mathbf{A})\odot\mathbf{U}\right). \tag{4}
残差连接:Z=Y+Z.(5)\text{残差连接:}\quad \mathbf{Z}=\mathbf{Y}+\mathbf{Z}. \tag{5}

其中,\odot 表示逐元素乘积,f1f_1f2f_2 分别是注意力前投影与注意力后投影使用的 MLP,ϕ1\phi_1ϕ2\phi_2 是 SiLU 激活函数。公式(3)使用因果掩码 M\mathbf{M},以维持序列物品间的时间关系。输入嵌入 Z\mathbf{Z} 在传入后续操作前先归一化,每一层通过标准残差连接把输出 Y\mathbf{Y} 与前一层连接。

虽然 HSTU 在推荐系统中展现出良好扩展规律,我们认为,借鉴 DeepSeek-V2 面向 LLM 的做法,通过模型—系统协同设计还可以进一步优化扩展曲线。因此,我们在原始 HSTU 之上提出 ULTRA-HSTU;下文讨论的思想通常也适用于其他序列推荐注意力架构。

4 ULTRA-HSTU:极高效高性能序列编码器

为了重塑普通 HSTU 的扩展曲线,我们在三个关键方面进行显著改进:第一,输入序列优化从源头缩短有效序列长度;第二,面向推荐定制的稀疏注意力实现线性计算复杂度;第三,动态拓扑设计使模型能够在不让每一层都承担完整序列成本的情况下实现良好的深度扩展。

除理论复杂度下降外,ULTRA-HSTU 还针对大规模分布式训练和推荐推理环境,与硬件协同设计以获得实际效率。最终,相较普通 HSTU,ULTRA-HSTU 的推理扩展效率提高 21 倍以上,训练扩展效率提高 5 倍。图 2 展示总体架构,下面分别介绍各组件。

4.1 输入序列优化

首先,我们提出一种高效行为编码方法,将输入序列有效缩短两倍,从而使注意力计算效率提高四倍。普通 HSTU 会交错排列物品和行为,把用户 ii 的输入序列表示为 {Ii,1,ai,1,Ii,2,ai,2,,Ii,Li,ai,Li}\{\mathbf{I}_{i,1},\mathbf{a}_{i,1},\mathbf{I}_{i,2},\mathbf{a}_{i,2},\ldots,\mathbf{I}_{i,L_i},\mathbf{a}_{i,L_i}\}。这虽然能同时支持检索与排序阶段,但会使排序序列长度变为实际 UIH 的两倍。

直接合并行为和物品可能泄漏待预测候选的行为信息,因此,对于推荐系统中待排序的候选位置 jj,我们把其行为嵌入掩蔽为 ai,j=0d\mathbf{a}_{i,j}=\mathbf{0}_d。我们探索了不同的行为—物品嵌入合并方式,最终选择简单相加,把用户 ii 的序列输入表示为 Xi={xi,j}j=1Li\mathbf{X}_i=\{\mathbf{x}_{i,j}\}_{j=1}^{L_i},其中 xi,j=Ii,j+ai,j\mathbf{x}_{i,j}=\mathbf{I}_{i,j}+\mathbf{a}_{i,j}。我们推测,这种方法能使梯度更容易通过行为编码传播。

此外,ULTRA-HSTU 还从隐式信号、显式信号和用户上下文侧信息中构造异构行为编码,以增强行为嵌入。重要的是,这种设计在不牺牲模型质量的情况下,把序列长度降至普通 HSTU UIH 设计的一半,使 ULTRA-HSTU 在保持可扩展性的同时获得显著性能提升。

我们进一步设计负载均衡随机长度算法,将训练吞吐量提高 15%。随机长度(SL)会在训练阶段随机选择用户,并将其历史序列采样至预定义阈值 Lα/2L^{\alpha/2},其中 α(1,2]\alpha\in(1,2] 是可调超参数。这会把训练计算复杂度从 O(L2)O(L^2) 降至 O(Lα)O(L^\alpha),并已被证明能够在推理时泛化到完整序列长度。

然而,在分布式训练环境中,每个 Rank 独立执行采样,导致各 Rank 的输入与输出负载(即用户序列长度之和)存在显著差异。在同步分布式训练框架中,这种负载不平衡会严重降低训练效率。

因此,我们提出 LBSL,它是 SL 的一个变体,通过显式控制每个 Rank 的计算量来减少拖慢整体进度的 Rank。将一个批次内某 Rank 的负载定义为 uranknuγ\sum_{u\in\mathrm{rank}}n_u^\gamma,其中 nun_u 是请求 uu 的序列长度,γ(1,2)\gamma\in(1,2) 表示 HSTU 的超线性成本。负载均衡程度定义为给定 World Size 内最大 Rank 负载与最小 Rank 负载之比。

LBSL 分三个阶段运行。首先,使用标准随机长度进行短暂预热,以估计全局目标负载 ˉ\bar{\ell}。其次,执行受约束采样,自适应选择一个不采样集合,使每个 Rank 的实际负载尽可能接近 ˉ\bar{\ell};同时通过权重 pup_u、无放回加权采样和贪心填充,保留 SL 更倾向于不采样短序列的偏置。最后,以可配置间隔周期性重新校准 ˉ\bar{\ell},以跟踪生产序列长度分布的缓慢变化。

如果每个批次都重新校准,LBSL 的平均负载与标准 SL 相同,但会在不同 Rank 之间重新分配采样:重负载 Rank 采样更多,轻负载 Rank 采样更少,从而在不损害质量的情况下减少拖慢整体训练的 Rank。附录 C 的算法 1 给出细节。

4.2 面向效率的模型—系统协同设计

4.2.1 半局部注意力设计

我们提出一种名为半局部注意力(SLA)的新型稀疏注意力机制。它使注意力计算达到线性复杂度,并将 ULTRA-HSTU 的推理扩展效率显著提高 5 倍。普通 HSTU 使用完整因果自注意力掩码:

A(X)=ϕ2 ⁣(Q(X)K(X)T)MV(X).(6)\mathbf{A}(\mathbf{X})=\phi_2\!\left(\mathbf{Q}(\mathbf{X})\mathbf{K}(\mathbf{X})^{\mathsf T}\right)\odot\mathbf{M}\,\mathbf{V}(\mathbf{X}). \tag{6}

其中,MRL×L\mathbf{M}\in\mathbb{R}^{L\times L} 是因果注意力掩码,只有满足 jLij\le L-i 时才有 Mi,j=1M_{i,j}=1ϕ2\phi_2 在 HSTU 中为 SiLU。当模型扩大序列长度时,这种计算会产生二次成本。

在大规模推荐系统中,UIH 长度很快会累积并超过 10K,导致模型无法部署到真实排序环境。受到 LLM 和推荐系统中注意力天然稀疏、动态这一特性的启发,我们开发半局部注意力机制,同时关注长期与局部模式。

定义两个超参数:局部窗口大小 K1K_1 和全局窗口大小 K2K_2。局部窗口控制纳入注意力掩码的局部模式窗口长度;全局窗口关注最新的 UIH 注意力模式,以捕获用户长期兴趣。半局部注意力的掩码定义如下:

Mi,j={1,LK1i+jL,1,jK2 且 jLi,0,其他情况.(7)M_{i,j}=\begin{cases} 1, & L-K_1\le i+j\le L,\\ 1, & j\le K_2\ \text{且}\ j\le L-i,\\ 0, & \text{其他情况}. \end{cases} \tag{7}

image.png 图 3:注意力掩码。左图为完整因果自注意力掩码,右图为半局部注意力掩码。

这种设计把注意力计算复杂度降为线性的 O((K1+K2)L)O((K_1+K_2)\cdot L)。当大规模推荐系统中的序列长度 LL 超过 10K 时,它能显著提高模型效率。与 DeepSeek 原生稀疏注意力只使用局部窗口不同,第 5 节将表明,SLA 的局部窗口和全局窗口都不可或缺;在长期用户行为十分关键的推荐系统中,这一现象尤其明显。

4.2.2 系统优化

混合精度训练与推理。 大规模推荐模型同时受稠密计算和数据移动限制。稠密计算主要是通用矩阵乘法(GEMM),数据移动尤其包括服务中的嵌入查询和主机到设备传输。为实现 ULTRA-HSTU 的端到端高效运行,我们协同设计一套覆盖 16/8/4 位格式的推荐专用混合精度框架:大多数操作使用 BF16 保证稳定性,主要 GEMM 使用 FP8 加速,推理通信则使用 INT4 嵌入量化降低流量。离线和在线实验表明,该混合精度栈在保持模型准确率的同时,使训练吞吐量提高 10%,服务吞吐量提高 40%。

我们为 HSTU 开发了定制 FP8 栈,同时处理两个实际瓶颈:提高 NVIDIA H100 上 Tensor Core 的利用率,从而提升稠密计算实际达到的 TFLOP/s;减少 FP8 量化和缩放的开销,避免其受内存带宽限制。

每个 HSTU 层包含两个 GEMM:一个是注意力前投影,在注意力之前把输入嵌入 X\mathbf{X} 映射为 U,V,Q,K\mathbf{U},\mathbf{V},\mathbf{Q},\mathbf{K} 张量;另一个是注意力后投影,把经过归一化和门控的注意力输出转换为该层输出。两个 GEMM 都使用 FP8,其余所有操作保留 BF16,以在不损害数值稳健性的情况下提高吞吐量。

仅把 GEMM 切换到 FP8 并不高效。朴素 FP8 流水线还需要额外执行缩放、量化和布局准备,可能抵消预期加速。因此,我们开发融合算子,把逐行缩放计算和量化,与处理不规则嵌入的前置层归一化算子(公式(1)和(4))融合,避免额外遍历内存并降低量化开销。

我们还为注意力后投影开发高性能 Triton FP8 GEMM 算子。该路径需要将投影输出与二维残差张量累加,因此我们把残差累加直接融合到 GEMM 尾声。PyTorch GEMM 算子通常假设偏置是一维向量,无法高效支持这一需求。我们的 Triton FP8 算子原生支持二维偏置,并利用持久化调度、TMA、Warp 专门化和尾声流水线来保持高吞吐量,同时避免过度的寄存器压力。除 FP8 GEMM 外,混合精度框架还在服务时对嵌入移动使用 4 位量化,详见附录 D。

image.png 图 4:混合精度计算框架。scaling/quantization步骤同前置算子融合。

面向异构硬件的高效 SLA 算子。 注意力操作是 HSTU 的瓶颈。原始算子使用 Triton 实现 FlashAttention V2 算法。我们采用 FlashAttention V3 的算法设计,积极重叠数据移动与计算,并针对 HSTU 的非标准注意力——逐位置 SiLU 激活和 SLA 掩码——定制算子。

该设计同时在 NVIDIA H100 和 AMD MI300X 上实现,使异构服务和交付在两个平台上都比 FlashAttention V2 基线快 2 倍。在 H100 上,我们使用 FlashAttention 3 风格的流水线,为完整 HSTU 注意力和半局部 HSTU 注意力实现 CUDA 算子族。在 AMD MI300X 上,我们使用 Composable Kernel 实现对应算子。由于 MI300X 不具备 FlashAttention 3 在 H100 上使用的特性,例如 TMA 和 Warp 专门化异步执行,我们引入 MI300X 原生优化:感知 XCD 的调度以利用 8 Chiplet 拓扑、减少共享内存 Bank 冲突的 LDS 布局,以及通过调度屏障显式交错 VMEM 与 MFMA。最终相较 Triton 算子基线获得 2 倍加速。

以极低开销节省内存。 标准注意力实现在前向传播中造成很高 GPU 内存压力,这会成为超长序列训练的主要瓶颈。我们设计了以下优化,在节省内存的同时保持训练效率。

首先,为 ULTRA-HSTU 引入专门的选择性激活重物化。具体而言,我们不保存六个大型前向张量,而是在反向传播中以极少重计算恢复它们,包括复用已保存的层归一化统计量得到归一化后的 X\mathbf{X},重新执行 GEMM 恢复 U,Q,K,V\mathbf{U},\mathbf{Q},\mathbf{K},\mathbf{V},以及在融合门控归一化算子内部计算中间量 Y\mathbf{Y}。这比通用检查点方法轻量得多,相较完全不重计算激活的基线只增加 5% 开销。详细算法见附录中的代码清单 1。

其次,我们通过取消对 dU,dQ,dK,dVd\mathbf{U},d\mathbf{Q},d\mathbf{K},d\mathbf{V} 的梯度拼接消除这部分开销,减少反向传播中的内存流量与算子开销。总体而言,ULTRA-HSTU 每层内存减少约 67%,效率没有下降。在嵌入维度 512、批量大小 256、序列长度 3K、数据类型 BF16 的设置下,该技术将每层 HBM 用量从 7 GB 降到 2.3 GB。第三,我们在端到端训练中使用完全不规则张量实现,无须填充为稠密张量,从而显著减少内存用量。附录 E 给出了详细效率基准。

4.3 动态拓扑设计

除扩展序列长度外,深度扩展同样对模型表现至关重要。朴素堆叠 ULTRA-HSTU 层,并让每个 SLA 层都处理完整序列,会产生 O(DL)O(DL) 计算成本,其中 DD 是模型深度。在真实应用中,当 LL 扩展到 10K 时,即使已经采用线性稀疏注意力,继续堆叠更多层仍会引入巨大的训练、内存和推理成本;与此同时,大规模推荐系统依然需要在毫秒内处理数百万请求。由此产生一个自然问题:堆叠越来越多层时,是否真的需要让每一层都关注完整序列?

注意力截断。 考虑到用户近期交互历史的重要性,在使用完整长度为 LL 的序列堆叠 N1N_1 层 HSTU 后,我们从完整序列中选择长度为 LL' 的片段,并只在该 UIH 片段上再堆叠 N2N_2 层 HSTU。

选择 UIH 片段的方法很多,包括:截取最近长度为 LL' 的 UIH;在第一次随机长度处理后再次应用 SL,选出长度为 LL' 的序列;在前 N1N_1 层后插入压缩模块,把完整序列压缩到长度 LL'。实践中,我们发现直接截取最近的 UIH 片段能获得最佳模型表现,如图 2(d) 所示。

转导器混合。 推荐模型天然会处理多条输入序列,因为不同来源、不同类型的用户互动信号通常分别记录。若把所有用户信号聚合为单一输入序列,并交给统一编码器,就会把异构用户交互压缩到同一条时间线中。稀疏高价值互动会被稠密隐式信号稀释,所有信号还会竞争有限的序列容量。

为解决这一问题,我们提出 MoT 范式。MoT 使用独立转导器处理多条不同输入序列,再融合学到的用户嵌入。它使模型能在不同时间跨度上捕获不同类型的用户行为,从而对多样且稀疏的互动模式形成更细致、更有效的表示。

关键在于,MoT 允许为不同输入序列灵活分配计算资源。例如,可以为高价值序列分配更深层数和更大容量,同时减少对已经充分理解或较不重要序列的资源投入。这种有针对性的计算预算分配,使模型把容量集中到最有意义的用户交互上,从而改善整体推荐质量和效率之间的权衡。

两种拓扑设计都比普通 HSTU 获得明显更好的模型质量—成本权衡,并且彼此兼容,可以组合到同一个模型中。实际应用时,应根据系统最关注的是效率还是模型质量来选择拓扑。第 5 节的实验采用注意力截断,因为它简单且能提供强大的质量—效率权衡;MoT 的研究见附录 A。

image.png 图 5:扩展规律消融研究。左图比较 SLA 的训练扩展,中图比较 SLA 的推理扩展,右图比较注意力截断(AT)的推理扩展。

5 实验结果

全文使用归一化熵(Normalized Entropy,NE)衡量模型质量。NE 定义为模型交叉熵除以仅根据正标签平均频率进行预测时的交叉熵:

NE=1Ni=1N(yilogpi+(1yi)log(1pi))plogp(1p)log(1p).(8)\operatorname{NE}= \frac{-\dfrac{1}{N}\sum_{i=1}^{N}\left(y_i\log p_i+(1-y_i)\log(1-p_i)\right)} {-plogp-(1-p)\log(1-p)}. \tag{8}

其中,NN 是训练样本数,yi{0,1}y_i\in\{0,1\} 是样本 ii 的标签,pip_i 是模型对该样本的预测,并且 pˉ=i=1Nyi/N\bar{p}=\sum_{i=1}^{N}y_i/N。NE 越低,模型越好。具体而言,我们测量消费任务(例如完整观看视频)的 NE 改善和互动任务(例如分享)的 NE 改善,分别记为 C-NE 和 E-NE。本文沿用原始 HSTU 论文及内部最佳实践报告 NE。根据我们的经验和实验,AUC 等其他指标会与 NE 同方向、近似同幅度地变化,因此因篇幅限制不再报告。

我们将模型与多种强基线比较,并按能否建模短程或长程用户行为分类。短序列方法包括 DIN 和 SASRec,长序列方法包括普通 HSTU 与 STCA。我们还比较了一个内部优化的 Transformer;它使用额外投影和归一化来稳定训练,避免经典 Transformer 用于推荐系统时发生意外指标下降。

5.1 工业数据集基准

5.1.1 数据集

我们首先使用来自内部大规模真实推荐系统的工业规模生产数据集报告模型表现。该数据集由线上用户交互历史的一个子集构成,总计超过 60 亿个样本;每个样本都包含长度从 3,072 到 16,384 个事件的超长用户交互序列。

为保证时间一致性并防止未来数据泄漏,我们按时间顺序划分数据:最初 85% 用于训练,剩余 15% 用于评估。工业数据集的所有实验都使用 LBSL。例如,推理时原始序列长度为 16,384,应用 LBSL 后的训练序列长度约为 4,400。相较使用完整序列训练的模型,LBSL 只造成极小 NE 差异,却能显著加快训练。这也解释了为什么每个样本的推理 FLOP 高于训练 FLOP。训练与推理序列长度的详细比较见附录表 6。

5.1.2 总体表现

表 1 给出所有方法在序列长度上限为 3,072 时的结果。我们调节模型深度和参数,使所有方法的 FLOP 大致相当。ULTRA-HSTU 显著优于其他所有方法。高度依赖交叉注意力来实现线性复杂度的 STCA,因为缺少自注意力能力而不如 ULTRA-HSTU。需要注意,ΔNE\Delta\operatorname{NE} 为正表示模型更差。根据经验,0.03%—0.05% 的改善即视为显著,并能带来可观的线上指标收益。

image.png 表 1:工业数据集上的模型表现。

5.1.3 扩展规律

为分析扩展行为,我们固定 ULTRA-HSTU 与普通 HSTU 的输入序列设计,并在改进模型架构和拓扑设计后,报告 C-NE 表现及 TFLOP 对比。模型层数从 6 变化到 18,序列长度取 L{3072,8192,16384}L\in\{3072,8192,16384\},模型维度固定为 d=512d=512。表 2 给出详细模型表现和 TFLOP。

随着序列长度和层数增加,ULTRA-HSTU 的效率与 C-NE 指标显著改善。图 1 把相对于基线模型的 C-NE 收益表示为计算成本的线性回归,并分别绘制 ULTRA-HSTU 与普通 HSTU。通过比较拟合线性函数的斜率可知,ULTRA-HSTU 的训练扩展效率提高 5.3 倍,推理扩展效率提高 21.4 倍。

image.png 表 2:ULTRA-HSTU 在工业数据集上的扩展结果。

5.2 开源数据集

ULTRA-HSTU 和 STCA 等方法面向用户历史达到数万次交互的工业规模推荐系统。为证明本文方法在极长序列以外同样具有普适性,我们还在公开 KuaiRand 基准上评估,其中序列长度只有 256。表 3 表明,即使在短序列场景中,ULTRA-HSTU 仍以最低训练与推理计算成本取得最佳 NE。由于注意力前计算开销较高,STCA 难以适应较短序列。

image.png 表 3:KuaiRand 基准上的比较。

5.3 扩展研究消融实验

我们先简要介绍输入序列优化的影响,然后固定普通 HSTU 与 ULTRA-HSTU 的输入序列设计,对 SLA 和注意力截断的扩展效率进行消融。附录 F 详细介绍扩展规律分析方法。

5.3.1 输入序列优化

移除输入序列中的物品—行为交错后,序列长度减半;在 UIH 长度为 3,072 时,训练 FLOP 显著降低 32.5%,推理 FLOP 降低 63.5%。与此同时,相较基线,异构行为嵌入构造使 C-NE 改善 0.45%。当 World Size 为 512 时,LBSL 加速 15%,证明它能有效加快大规模序列模型训练。

5.3.2 半局部注意力

图 8 绘制启用与未启用 SLA 时,C-NE 随总 FLOP 的变化。启用 SLA 的模型相较普通 HSTU 展现出明显更好的扩展性,训练扩展效率提高 2.7 倍,推理扩展效率提高 5.1 倍。

SLA 中的局部窗口大小 K1K_1 和全局窗口大小 K2K_2 都不可或缺,这与只启用局部滑动窗口的原生稀疏注意力明显不同。我们还发现,全局窗口 K2K_2 比局部窗口 K1K_1 更重要。例如,令 K1=0K_1=0、只启用 SLA 的全局窗口时,C-NE 仅下降 0.03%;令 K2=0K_2=0、只启用局部窗口时,C-NE 会下降 0.35%。

5.3.3 动态拓扑设计

在普通 HSTU 中堆叠更多层会显著提高性能,但训练和推理成本难以承受。图 8 右图比较使用注意力截断堆叠 HSTU 层,与使用完整序列单纯堆叠 HSTU 层时的扩展曲线。

实验在以下设置中进行:推理时前 n1n_1 层序列长度为 3,072,使用 SL 后训练长度约为 1,110;后 n2n_2 层序列长度为 512;n1{3,6,9,12}n_1\in\{3,6,9,12\}n2{0,3,6,9}n_2\in\{0,3,6,9\}。序列越长,注意力截断越有效。模型训练启用 LBSL 后,在约 1,110 的序列长度上,注意力截断带来的效率节省还不够显著;但在推理长度为 3,072 的更长序列上,结果明显更好,推理扩展效率提高 3.4 倍。

5.4 线上 A/B 测试

我们通过多项严格的 30 天线上 A/B 测试,验证 ULTRA-HSTU 在每天覆盖数十亿用户的大规模生产视频服务平台上的效果。报告三类线上指标:第一,观看时长、完整观看视频等线上消费指标(C-Metric);第二,点赞、评论、分享等线上互动指标(E-Metric);第三,访问次数、日活跃用户等线上顶层指标。

我们将现有生产模型从普通 HSTU 升级为 ULTRA-HSTU。结果显示,ULTRA-HSTU 使线上消费指标显著提升 4.11%,互动指标根据互动类型提升 2%—8%。更值得注意的是,平台整体健康状况的重要指示指标——关键“顶层”指标——也获得显著改善。在我们的系统中,互动和消费指标即使只提升个位数百分比,也会被视为重大突破;顶层指标 1 和 2 分别提高 0.05% 和 0.01%,就已被 Meta 视为高度显著。

综合来看,这些结果有力证明了 ULTRA-HSTU 的有效性与潜力。据我们所知,这是该推荐平台测试过的最大模型,并产生了过去几年中最显著的影响之一。

image.png 表 4:相较生产基线的一个月线上收益。

6 结论

本文提出 ULTRA-HSTU,这是一种新颖的端到端模型—系统协同设计方法,可以显著提高推荐领域序列建模的扩展效率。贡献可概括为三点。

第一,作为关键研究发现,我们证明自注意力仍优于交叉注意力;扩大注意力层计算量和序列长度,仍会持续改善模型表现。第二,作为关键技术创新,我们提出多项模型与系统协同优化,包括输入处理中的 LBSL、半局部注意力、采用混合精度训练与推理的异构硬件算子优化,以及动态模型拓扑设计,最终使训练扩展效率提高 5 倍,推理扩展效率提高 21 倍。

第三,作为面向推荐行业的关键经验分享,我们把 ULTRA-HSTU 部署到了大规模生产环境。该模型具有 18 层自注意力,处理 16K 用户序列,并在数百块 H100 GPU 上训练,产生了显著影响。这证明扩大推荐序列模型是一条前景良好的方向,也验证了本文创新的有效性。

附录 A 拓扑设计的选择

注意力截断和 MoT 都比普通 HSTU 获得显著更好的模型质量—成本权衡,并且两种设计彼此兼容。下面详细说明各自优势。在实际应用中,应根据系统最关注效率还是模型质量来选择拓扑设计。

转导器混合。 MoT 在互动任务上带来显著 NE 收益,同时在训练和推理 FLOP 方面保持具有竞争力的节省。通过把异构信号解耦到专用模块中,MoT 缓解了这样的信号竞争:当多样输入信号被限制在一个序列长度有限的模块内时,它们会彼此争夺表示容量。

具体而言,我们使用两个专门的 HSTU 模块:一个处理互动事件,一个处理消费事件,在表 8 中分别记为 E-seq 和 C-seq。每个模块处理的序列都比对应的单 HSTU 模型短,但通过精心构造序列获得更丰富的信号表示。例如,专门的互动模块虽然序列更短,却可以避免与稠密消费信号竞争,从而捕获明显更丰富的互动历史。

我们还为每个模块定制计算资源分配,以进一步优化计算效率。更短的序列使注意力操作显著变轻,因而在训练和推理中都能节省大量 FLOP。

image.png 表 8:MoT 的表现。MoT 显著改善 E-NE 指标,其中 C-seq 表示消费序列,E-seq 表示互动序列。

注意力截断。 表 9 给出补充数据,比较只增加普通 HSTU 层数与增加注意力截断层数的情况。模型层数增加后,消费和互动任务的 NE 都显著改善,但推理 FLOP 负担也会增大。使用注意力截断可以明显改善模型质量与计算成本之间的权衡。

例如,将 3 层普通 HSTU 再堆叠 6 层注意力截断,与只使用 6 层普通 HSTU 比较:注意力截断可以取得相当的 C-NE 和更好的 E-NE,同时节省 3% 的训练 TFLOP 和 38% 的推理 FLOP。

image.png 表 9:序列长度为 3,072 时,普通 HSTU 与注意力截断的扩展比较。ΔNE\Delta\operatorname{NE} 为负表示模型更好。

附录 B 交叉注意力深度扩展的收益递减

表 5 表明,在模型深度扩展方面,自注意力比交叉注意力更强。当序列长度约为 3,072 时,增加交叉注意力层数后,模型表现在 9 层时趋于饱和;自注意力则会随层数增加持续改善模型质量。

image.png 表 5:交叉注意力与自注意力的Depth scaling。

附录 C 负载均衡随机长度算法

算法 1 给出 LBSL 的详细步骤。在第 5 节实验中启用 LBSL 后,不同原始序列长度所对应的训练与推理序列长度见表 6。

image.png 表 6:使用 SL 训练与不使用 SL 推理时的平均 UIH 序列长度。

算法 1:负载均衡随机长度

输入要求: World Size RR;预热步数 TwarmT_{\mathrm{warm}};重新校准间隔 TrecalT_{\mathrm{recal}};负载指数 γ(1,2)\gamma\in(1,2);SL 参数 α\alpha;SL 采样长度 SL\ell_{\mathrm{SL}};批量大小 bb

  1. 初始化目标负载 ˉ0\bar{\ell}\leftarrow0,并对所有 Rank r{1,,R}r\in\{1,\ldots,R\}r0\ell_r\leftarrow0
  2. 对训练步 t=1,2,t=1,2,\ldots 执行循环。
  3. 对所有 Rank r{1,,R}r\in\{1,\ldots,R\} 并行执行以下步骤。
  4. 接收本地批次 BrB_r 及其原始长度集合 {nu}uBr\{n_u\}_{u\in B_r}
  5. 计算 rr+StandardSL_Loads(Br;α;γ)\ell_r\leftarrow\ell_r+\operatorname{StandardSL\_Loads}(B_r;\alpha;\gamma),将其作为截断前代理负载。
  6. 如果 tTwarmt\le T_{\mathrm{warm}},执行下一步。
  7. BrB_r 应用 StandardSL(Br;α)\operatorname{StandardSL}(B_r;\alpha) 来截断样本。
  8. 否则,执行以下受约束采样步骤。
  9. UrU_r\leftarrow\varnothings0s\leftarrow0,其中 UrU_r 是不截断集合。
  10. 对所有 uBru\in B_r,计算权重 puSLWeight(nu;α)p_u\leftarrow\operatorname{SLWeight}(n_u;\alpha)
  11. 使用 pup_uBrB_r 进行无放回加权随机排列,得到 π\pi
  12. 按排列 π\pi 的顺序遍历 uu
  13. 如果 s+(nuγSLγ)ˉbSLγs+(n_u^\gamma-\ell_{\mathrm{SL}}^\gamma)\le\bar{\ell}-b\cdot\ell_{\mathrm{SL}}^\gamma,执行下一步。
  14. UrUr{u}U_r\leftarrow U_r\cup\{u\},并令 ss+(nuγSLγ)s\leftarrow s+(n_u^\gamma-\ell_{\mathrm{SL}}^\gamma)
  15. 结束条件判断。
  16. 结束对 π\pi 的遍历。
  17. 保留所有 uUru\in U_r,不进行采样截断。
  18. 对每个 uBrUru\in B_r\setminus U_r,应用与 StandardSL(u;α)\operatorname{StandardSL}(u;\alpha) 相同的采样规则。
  19. 结束预热条件分支。
  20. 结束各 Rank 的并行步骤。
  21. 如果 t=Twarmt=T_{\mathrm{warm}},执行下一步。
  22. 使用 All-Reduce 得到平均负载 ˉ1RTwarmr=1Rr\bar{\ell}\leftarrow\dfrac{1}{RT_{\mathrm{warm}}}\sum_{r=1}^{R}\ell_r
  23. 对所有 Rank r{1,,R}r\in\{1,\ldots,R\},令 r0\ell_r\leftarrow0
  24. 否则,如果 tmodTrecal=0t\bmod T_{\mathrm{recal}}=0,执行下一步。
  25. 使用 All-Reduce 得到平均负载 ˉ1RTrecalr=1Rr\bar{\ell}\leftarrow\dfrac{1}{RT_{\mathrm{recal}}}\sum_{r=1}^{R}\ell_r
  26. 对所有 Rank r{1,,R}r\in\{1,\ldots,R\},令 r0\ell_r\leftarrow0
  27. 结束重新校准条件判断。
  28. 结束训练循环。

附录 D 其他系统优化

混合精度服务。 在模型服务中,长序列下的稀疏嵌入特征可能主导主机到设备的传输时间。因此,我们把嵌入张量量化为 INT4,并在嵌入查询和传输路径中始终保持量化形式,以减少传输量、缓解通信瓶颈。

此外,我们使用分组 INT4。它利用各组专属的缩放因子,相较每行只使用单个缩放因子可以显著减少质量损失,同时仍提供可观的吞吐量收益。

附录 E 效率基准

本节提供基准测试,以评估第 4.2.2 节提出的系统优化。

E.1 混合精度基准

FP8 精度效率。 我们评估 FP8 对注意力前模块和注意力后模块中 GEMM 的性能影响,并在表 10 中汇总加速效果。两个模块之间的关键区别是 GEMM 的偏置格式:注意力前 GEMM 使用一维偏置,注意力后 GEMM 使用二维偏置。

由于 Torch 已经支持带一维偏置的 FP8 GEMM,我们直接在注意力前 GEMM 中使用 Torch 算子。相反,Torch FP8 GEMM 并不原生支持注意力后模块所需的二维偏置,因此我们开发了原生融合二维偏置的 Triton FP8 GEMM 算子,并在表 7 中报告其算子级效率。

image.png 表 7:带二维偏置(2D bias)的 FP8 GEMM 算子效率基准。在 NVIDIA H100 上,以不同的 m,n,km,n,k 报告 TFLOP/s。Bias-Fused FP8 使用原生支持二维偏置的 Triton 算子;Bias-Split FP8 使用 Torch FP8 GEMM 后再单独加偏置,因为前者不支持二维偏置。

表 7 对以下运算进行基准测试:

D=AB+C,\mathbf{D}=\mathbf{A}\mathbf{B}+\mathbf{C},

其中,ARm×k\mathbf{A}\in\mathbb{R}^{m\times k}BRk×n\mathbf{B}\in\mathbb{R}^{k\times n}CRm×n\mathbf{C}\in\mathbb{R}^{m\times n},矩阵维度 (m,k,n)(m,k,n) 与本文模型负载一致。我们尤其关注由不规则变长序列造成的较大首维 mm,因为它在实践中主导计算成本。

结果表明,相较使用 Torch FP8 GEMM 后再单独添加偏置,将二维偏置直接融合到 FP8 GEMM 中最多可加速 1.75 倍。这解释了为什么我们为注意力后模块实现 Triton 算子。

表 10 报告从 BF16 切换到 FP8 后,完整注意力前和注意力后部分的加速。强劲的端到端收益来自两个方面:第一,性能更高的 FP8 GEMM 算子,包括用于注意力后模块的二维偏置 Triton 算子;第二,把量化融合到 GEMM 之前的算子中,较单独执行量化减少了额外内存流量和算子启动开销。

image.png 表 10:标准量化 FP8 GEMM 与把量化融合到前置算子的 FP8 GEMM 的性能比较;报告各方法相对于 BF16 的加速比。

INT4 量化效率。 表 11 汇总 INT4 稀疏嵌入量化对模型服务效率的影响。对稀疏嵌入应用 4 位量化后,主机到设备的数据传输延迟降低约 40%,峰值每秒查询数提高 20% 以上。应用 4 位量化后,线上模型准确率差异可以忽略。

image.png 表 11:稀疏嵌入数据类型对模型服务效率的影响。在 3.5K QPS 下测量嵌入查询latency;峰值 QPS 是端到端latency预算为 80 毫秒时的最大 QPS。所有结果均采集自单台 H100 主机。

E.2 注意力算子基准

图 6 给出注意力算子效率基准,在 NVIDIA H100 和 AMD MI300 GPU 上,把我们的优化实现与 FlashAttention V2 基线比较。评估包含半局部注意力和因果注意力两种设置。

在 H100 上,对于因果注意力,ULTRA 实现在 16K 序列长度时可保持 520 TFLOP/s 以上,相较基线加速 1.64 倍。对于 SLA,在不同批量大小和序列长度上,本文算子始终达到更高吞吐量,最多加速 2.5 倍。

在 MI300 上,图 6 报告前向传播算子性能。ULTRA 算子相较基于 FlashAttention V2 的实现最多加速 1.51 倍。在小批量、16K 序列长度下,MI300 上 ULTRA 的吞吐量最高可达到 H100 上 ULTRA 的 0.92 倍。这些结果体现了我们为大规模模型实现高效 AMD 推理所做的针对性工作。

image.png 图 6:ULTRA 与基线在 H100 和 MI300 上的性能比较。ULTRA 使用 FlashAttention V3 风格算法,基线使用基于 FlashAttention V2 风格算法的 Triton 实现。(a)(d)(g)(h):H100 上完整注意力前向与反向算子的 TFLOP/s。(b)(e)(i)(j):在 H100 上,ULTRA SLA 相对于两个基线的加速;两个基线分别是 ULTRA 实现的完整注意力和基线实现的 SLA。(c)(f):在 MI300 上,ULTRA SLA 相对于三个基线的加速。

代码清单 1:采用激活重物化的 ULTRA-HSTU 伪代码

class HSTULayerFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, norm_w, linear_w_1, gated_norm_w, linear_w_2):
        normed_x = norm_forward(x, norm_w)
        u, v, q, k = silu_forward(addmm(normed_x, linear_w_1))
        attn = attention_forward(q, k, v)
        y = gated_norm_forward(attn, u, gated_norm_w)
        out = addmm(y, linear_w_2, bias=x)
        ctx.save_for_backward(
            x, u, attn, norm_w, linear_w_1, gated_norm_w, linear_w_2
        )
        return out

    @staticmethod
    def backward(ctx, dout):
        x, u, attn, norm_w, linear_w_1, gated_norm_w, linear_w_2 = (
            ctx.saved_tensors
        )
        y = gated_norm_forward(attn, u, gated_norm_w)  # 重新物化 y
        normed_x = norm_forward(x, norm_w)
        u, v, q, k = silu_forward(  # 重新物化 u、v、q、k
            addmm(normed_x, linear_w_1)
        )
        dy = addmm(dout, linear_w_2.T)
        d_linear_w_2 = addmm(y.T, dout)
        dx = dout
        dattn, du, d_gated_norm_w = gated_norm_backward(
            dy, attn, u, gated_norm_w
        )
        dq, dk, dv = attention_backward(dattn, q, k, v)
        duqkv = concat(du, dq, dk, dv)
        uqkv = concat(u, q, k, v)
        d_addmm = silu_backward(duqkv, uqkv)
        d_normed_x = addmm(d_addmm, linear_w_1.T)
        d_linear_w_1 = addmm(normed_x.T, d_addmm)
        dx, d_norm_w = norm_backward(d_normed_x, x, norm_w)
        dx += dout
        return dx, d_norm_w, d_linear_w_1, d_gated_norm_w, d_linear_w_2

附录 F ULTRA-HSTU 的扩展规律

本节分析本文方法相对于计算量的扩展规律。我们假设 NE 关于计算量的函数遵循如下幂律:

L(C)=αCβ.(9)\mathcal{L}(C)=\alpha C^{-\beta}. \tag{9}

这里假设当计算预算 CC\to\infty 时,NE 指标 L(C)0\mathcal{L}(C)\to0。通常,这会使我们系统性低估真实扩展规律指数,低估因子为:

β^=β(1LL),(10)\hat{\beta}=\beta^*\left(1-\frac{\mathcal{L}_\infty}{\mathcal{L}}\right), \tag{10}

其中,β^\hat{\beta} 是对真实参数 β\beta^* 的估计,L\mathcal{L}_\infty 是数据上的不可约误差。采用这一假设,是因为它使 α\alphaβ\beta 在双对数空间中保持线性,无须估计不可约误差项。如下文所示,该假设还确保对扩展改善程度的估计较为保守。

考虑两个模型之间的估计扩展比:

β^1β^2=β1β21L/L11L/L2=β1β2R,(11)\frac{\hat{\beta}_1}{\hat{\beta}_2} =\frac{\beta_1^*}{\beta_2^*} \frac{1-\mathcal{L}_\infty/\mathcal{L}_1}{1-\mathcal{L}_\infty/\mathcal{L}_2} =\frac{\beta_1^*}{\beta_2^*}R, \tag{11}

其中,RR 是扩展比估计的修正因子。扩展比表示模型 1 相对于模型 2 的扩展曲线改善了多少。当 L1<L2\mathcal{L}_1<\mathcal{L}_2 时,R<1R<1。如果模型 1 还具有更好的扩展性——在本文考察的计算区间内,损失更低的模型始终满足这一点——那么我们对扩展比的估计就是保守的。

扩展规律指数改善的含义。 扩展规律指数的提升乍看可能不大,但随着计算预算增加,其影响会复合增长。考虑扩展规律分别为 L1(C)=αCβ1\mathcal{L}_1(C)=\alpha C^{-\beta_1}L2(C)=αCβ2\mathcal{L}_2(C)=\alpha C^{-\beta_2} 的两个模型,其中 β1=kβ2\beta_1=k\beta_2,改善因子 k>1k>1。要让模型 2 达到模型 1 使用计算量 CC 时的相同损失,需要:

C2=Ck.(12)C_2=C^k. \tag{12}

这意味着,计算优势会随扩展比呈多项式增长。例如,扩展指数提高 2 倍,就意味着基线模型需要以平方速度增加计算量,才能匹配改进模型的表现。

ULTRA-HSTU 的总体扩展表现。 图 7 分别以训练和推理 FLOP 为自变量,绘制 ULTRA-HSTU 与 HSTU 的拟合计算扩展规律。结果显示,相较 HSTU,ULTRA-HSTU 相对于训练 FLOP 的扩展指数提高 2.08 倍,相对于推理 FLOP 的扩展指数提高 4.59 倍。

image.png 图 7:总体计算Scaling Law。比较 ULTRA-HSTU 与普通 HSTU 相对于训练 FLOP(左)和推理 FLOP(右)的扩展表现。ULTRA-HSTU 的训练和推理计算Scaling Law指数分别提高 3.09 倍和 3.52 倍。

译者注:原文正文报告的训练/推理扩展指数提升为 2.08 倍/4.59 倍,而图 7 标题报告为 3.09 倍/3.52 倍,两处数值不一致;此处均按原文保留。

半局部注意力的扩展表现。 接下来单独考察 SLA 对总体扩展改善的贡献。图 8 给出 SLA 相对于训练和推理 FLOP 的扩展表现。本文方法相对于训练 FLOP 和推理 FLOP 的扩展规律指数分别提高 1.39 倍和 1.69 倍。

注意力截断的扩展表现。 最后分析注意力截断方法的扩展行为。相对于训练 FLOP,扩展规律指数只获得较小提升;但图 8 表明,相对于推理 FLOP 的指数提高了 1.8 倍。

image.png **图 8:扩展规律消融研究。半局部注意力(SLA)使训练扩展指数提高 1.39 倍(左),推理扩展指数提高 1.69 倍(中)。注意力截断扩展规律:结合自注意力与注意力截断机制,使推理计算扩展效率提高 1.8 倍(右)。