Meta HSTU：统一的生成式推荐

一、背景

在之前的文章中我们介绍了字节的 HLLM 方案（​​​字节 HLLM：大模型与推荐系统结合的新范式​​），其中也提到了 Meta 的HSTU，本文中我们简单介绍一些这篇论文。

对应的论文为：[2402.17152] Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations

对应的代码库为：GitHub - facebookresearch/generative-recommenders

二、摘要

大规模推荐系统具有以下特征：依赖于高基数、异质特征，并且需每天处理数百亿用户行为。尽管在海量数据（含数千种特征）上进行训练，但多数工业界的深度学习推荐模型（DLRM）无法在计算规模上取得突破。受到 Transformer 在语言与视觉领域的启发，作者重新审视了推荐系统的基本设计选择，将推荐问题重新表述为生成式建模下的生成式任务（生成式推荐，Generative Recommender，GR），并提出了一种新型架构 —— HSTU（Hierarchical Sequential Transduction Units），专为处理高基数、非恒定的流式推荐数据场景。

在合成和公开数据集上的实验表明，HSTU 相比基线模型在 NDCG 指标上最多提升了 65.8%，且在处理 8192 长的序列时，比基于 FlashAttention2 的 Transformer 快 5.3x - 15.2x。基于 HSTU 的生成式推荐系统（包含 1.5T 参数量）在在线 A/B 测试中指标提升 12.4%，并已部署于拥有数十亿用户的互联网平台的多项服务中。更重要的是，生成式推荐的模型呈现出跨三个数量级的幂律增长特性，直到与 GPT-3/LLaMa-2 的规模相当。

三、引言

推荐系统作为在线内容平台与电子商务领域的核心技术，每日为数十亿用户提供个性化体验，发挥着关键作用。当前最先进的推荐方法已基于 DLRM 发展了约十年。DLRM 的特点在于其利用了异质特征，包括数值特征（如计数器和比率）、嵌入向量以及分类特征（如创作者 ID、用户 ID等）。由于每分钟都有新内容和产品不断加入，特征空间呈现出极高的基数，通常达到数十亿级别。为充分利用数万种此类特征，DLRM 采用多种神经网络来组合特征、转换中间表示并生成最终输出。

尽管采用了广泛的人工设计特征并基于海量数据进行训练，大多数 DLRM 在实际工业应用中仍面临计算效率低下的问题。这一局限性值得关注，且至今尚未得到有效解决。

受 Transformer 在语言与视觉领域取得成功的启发，作者重新审视了现代推荐系统中的基础设计选择。作者观察到，在处理亿级用户规模的推荐系统时，需克服三大挑战。

• 首先，推荐系统中的特征缺乏显式结构。尽管在小规模场景下已探索了序列化表达，但在工业级 DLRM 中，异质特征，如高基数 ID、交叉特征、计数器、比率等，扮演着关键角色。
• 其次，推荐系统采用的词表量达数十亿，且持续变化。与语言模型中使用的 10 万级静态词表相比，亿级动态词表带来了训练难题，并因需以 Target-aware 方式评估数万候选 Item，导致推理成本高昂。
• 最后，计算成本成为实现大规模序列模型的主要瓶颈。GPT-3 在数千块 GPU 上历经 1-2 个月，处理了总计 300B 个 Token。这一规模看似惊人，但与用户行为规模相比则相形见绌。最大的互联网平台每日服务数十亿活跃用户，这些用户每日与数十亿帖子、图片和视频互动。用户序列长度可达 10 的五次方。因此，推荐系统每日需处理的 Token 量，远超语言模型在 1-2 个月内处理的数量级。

本文中，作者将用户行为视为生成建模中的新型模态。作者的核心洞见包括：

• 在工业级推荐系统中，核心的排序与检索任务在适当的新特征空间下可被转化为生成建模问题；
• 这一范式使作者能够系统性地利用特征、训练及推理中的冗余，从而提升效率。

四、推荐系统作为生成式任务：从 DLRM 到 GR

4.1 在 DLRM 中统一异质特征空间

现代 DLRM 通常采用大量类别型（“稀疏”）和数值型（“密集”）特征进行训练。在 GR 中，作者将这些特征整合并编码为单一的时间序列，如下图 Figure 2 所示。

类别型（“稀疏”）特征：其示例包括：用户喜欢的物品、用户关注的特定类别（如“户外”）中的创作者、用户语言、用户加入的社区、请求发起的城市等。作者对这些特征进行序列化处理，具体步骤如下：首先，选择最长的时序数据，通常通过合并表示用户互动 Item 的特征来构建主时间序列。其余特征通常为随时间变化缓慢的时间序列，如人口统计信息或用户关注的创作者。通过保留每个连续片段的最早 Item 来压缩这些时间序列，然后将结果合并到主时间序列中。鉴于这些时间序列变化极为缓慢，此方法并未显著增加整体序列长度。

数值型（“密集”）特征：其示例包括：加权衰减计数器、比率等。例如，某一特征可能表示用户过去在特定主题匹配项上的点击率（CTR）。相较于类别特征，这些特征变化更为频繁，几乎每发生一次（User，Item）交互都可能发生变化。因此，从计算和存储角度来看，完全序列化这些特征是不可行的。然而，一个关键发现是，进行这些聚合操作所基于的类别特征（如 Item 主题、位置）已在 GR 中被序列化并编码。因此，可以使用足够表达能力的生成式架构并结合目标感知（[1706.06978] Deep Interest Network for Click-Through Rate Prediction）公式的情况下，移除 GR 中的数值特征。通过增加序列长度和 GR 中的计算量，能够有效地捕捉数值特征。

4.2 将排序和检索重新表述为生成式任务

给定一个按时间顺序排列的 n 个 Token 序列 x₀, x₁, ..., xn-1（其中 xi ∈ X），以及这些 Token 被观察到的时间点 t₀, t₁, ..., tn-1，生成式任务将该输入序列映射为输出 Token y₀, y₁, ..., yn-1：

• 其中 yi ∈ X ∪ {∅}），yi = ∅ 表示 yi未定义。
• 用 Φi ∈ Xc（Xc ⊆ X）表示系统提供给用户的 Content（例如，图像或视频）。由于新内容不断生成，Xc 和 X 不是固定不变的。
• 用 ai 表示用户对 Φi 做出某种 Action（例如，点赞、跳过、视频完成并分享），其中 ai ∈ X。
• 用 nc 表示用户已交互的 Content 总数。

在因果自回归设置下，标准的排序和检索任务可以定义为生成式任务（见如下图 Table 1）。可以得出以下观察结果：

• 检索。在推荐系统的检索环节，针对 Φi+1 ∈ Xc 学习一个分布 p(Φi+1 | ui)，其中 ui 表示用户在第 i 个 Token 处的表征。典型的目标是通过选择 arg maxΦ∈Xc p(Φ | ui) 来最大化某种奖励。这与标准的自回归设置有两点不同。

a.首先，对于 xi 和 yi 的监督信号不一定就是 Φi+1，因为用户可能对 Φi+1作出负面反应。

b.其次，当 xi+1 代表与非交互相关的分类特征（如人口统计信息）时，yi 是未定义的。

• 排序。在 GR 中，排序面临独特的挑战，因为工业推荐系统通常需要“Target-aware”的公式化表示。在这种情境下，目标Φi+1与历史特征的“交互”需要尽早发生，这在标准自回归设置中是不可行的，因为“交互”通常发生较晚（例如，通过 Encoder 输出后的 Softmax 函数）。作者通过在如下图 Table 1 中交错排列 Item 和 Action 来解决这一问题，使得排序可以表述为 p(ai+1 |Φ0 ,a0 ,Φ1 ,a1 ,...,Φi+1)（在分类特征之前）。在实践中，作者应用一个小型神经网络将 Φi+1 处的输出转换为多任务预测。重要的是，这使得能够在一次处理中对所有nc 个交互应用目标感知的交叉注意力机制。​

如下图 Table 10 所示为以前的序列推荐工作与 GR 的对比：

如下图 Figure 8 所示：

• 左图为传统的序列推荐模型：输入可以没有 Item  Action 信息 ai，或者将 ai 融合到 Item Content 信息 Φi  中，而输出都是 Item Content 信息 Φi。
• 右图为本文的 GR 模型：Action ai 和 Content Φi 同时作为输入和输出。其中

a.下一个 Action ai Token 预测任务可以理解为排序（Ranking）任务。

b.下一个 Content Φi 预测可以理解为检索（Retrieval）任务。

4.3 生成式训练

工业推荐系统通常采用流式训练模式，其中每个样本按其可用顺序依次处理。在此模式下，基于 SelfAttention 机制的生成式架构的总计算需求与

成比例，其中 ni 为用户  i 的序列长度（Token 数量），d 为 Embedding 维度。括号中的第一项源自 SelfAttention 机制，第二项则来自 MLP 层，其隐藏层大小为 O(dff) = O(d)。设 N = maxini，整体时间复杂度可简化为 O(N3d + N2d2)，对于推荐系统场景来说成本过高。

为了以可扩展方式在长序列上训练，作者转向生成式训练，从而将计算复杂度降低 O(N) 倍，如下图顶部所示。通过这种方式，Encoder 成本可以在多个目标间分摊。更具体地，当以速率 su(ni) 采样第 i 个用户时，总训练成本可以按

  缩放，通过将 su(ni) 设为 1/ni，可降至 O(N2d + Nd2)。在工业规模系统中实现此采样的一种方式是在用户请求或会话结束时生成训练样本，从而使得

。

五、用于 GR 的高性能 SelfAttention Encoder

为了适应工业级推荐系统中词表大且不固定的特性，并实现 GR 的规模化应用，作者提出了一种新的 Encoder 设计：HSTU，HSTU 由一系列通过残差连接相互关联的相同层堆叠而成。每一层包含 3 个子层：

• 公式 1：Pointwise 投影。
• 公式 2：空间聚合。
• 公式 3：Pointwise 变换。​

其中，fi(X) 表示 MLP；为降低计算复杂度并优化 Query Q(X)、Key K(X)、Value V(X) 及门控权重 U(X) 的 Batch 处理效率，f1 和 f2 均采用单一线性层，即 fi(X) = Wi(X) + bi, 并且会使用 Kernel 融合。ϕ1 和 ϕ2 代表非线性函数，此处选用 SiLU；Norm 为 LayerNorm，rabp,t 表示相对注意力偏置，其整合了位置（p）与时序（t）信息。完整的符号信息如下图 Table 9 所示：

如下图 Figure 3 所示，DLRM 实际可理解为 3 个主要阶段：

• 特征提取：检索分类特征的池化 Embedding 表征，代表作包含 Pointwise Attention 和 Target-aware Attention 池化。
• 特征交互：DLRM 中最关键的部分。HSTU 通过Norm (A(X)V(X)) ⊙ U (X) 实现 Attention 池化特征与其他特征的“交互”来取代 DLRM 的特征交互。
• 表征转换：DLRM 中通常通过 MoE（Mixture of Expert）来实现，以处理多样化、异质的群体。其关键思想是：提供不同用户的特化子网络以实现条件计算。HSTU 中的 Pointwise 实际上可以执行 MoE 中的门控操作。​

5.1  Pointwise 聚合 Attention

HSTU 采用新的 Pointwise 聚合（归一化）Attention 机制（Softmax Attention 在整个序列上计算归一化因子）。这一设计基于两个因素：

• 与 Target 相关的先验数据点的数量作为强特征，能够指示用户偏好的强度，而Softmax 归一化后很难捕获这一特征。这一点至关重要，因为需要同时预测用户参与的强度（如，在某一个 Item 上花费的时间）以及 Item 的相对顺序（如，预测最大 AUC 的排序）。
• 尽管 Softmax 激活在结构上对噪声具有鲁棒性，但在流式场景中处理非恒定词表时，其适应性较差。

如上述公式（2）和（3）所示，在 Pointwise 池化后需要进行 LayerNorm 以稳定训练过程。如下图 Table 2 所示，作者合成了非恒定词表的流式数据，在此场景下，Softmax Attention 与 Pointwise Attention 的差距达到 44.7%。

5.2 增加稀疏性

在推荐系统中，用户历史序列的长度通常呈 Skewed 分布，导致输入序列非常稀疏，特别是在长序列场景尤为明显，这种稀疏性可被利用以显著提升 Encoder 的效率。为此，作者开发了一种高效的 GPU Attention Kernel，以融合连续的 GEMM 操作，实质上是将 Attention 计算转换为可变大小的 Grouped GEMM 计算。因此，HSTU 中的 SelfAttention 机制变为 Memory Bound，并且在 Memory 访存方面按照

的比例扩展，其中 ni 表示样本 i 的序列长度，dqk 为 Attention 维度，R 为寄存器大小。这种方法本身可以带来 2x-5x 的吞吐提升。

作者进一步通过随机长度（Stochastic Length, SL）增加用户历史序列的稀疏性。推荐系统中用户历史序列的一个关键特征是：用户行为在时间维具有重复性。这提供了一个在不损害模型质量的前提下人为增加稀疏性的机会，从而显著降低按 比例扩展的成本。可以将第 j 个用户的历史表示为序列，其中 nc,j 是用户交互的 Content 数量。设，表示从原始序列构建的长度为 L 的子序列。SL 选择输入序列的方式如下所示：

该方法将 Attention 相关复杂度降低至 ，需要说明的是，由于训练成本往往远大于推理成本，因此将稀疏 Attention 机制应用于训练成本收益更显著。

如下图 Table 3 展示了不同序列长度和 α 值下，针对用户 30 天历史的稀疏度。

• 2.0 表示未应用稀疏注意力机制的基础稀疏情况。
• 蓝色下划线表示损失可忽略不计的情况。
• 较低的 α 值适用于更长的序列。​

5.3 最大限度地减少激活内存使用

在推荐系统中，大 Batch 对于训练吞吐量和模型质量至关重要。因此，激活内存的占用成为主要瓶颈，而在 LLM 场景，通常 Batch 不是特别大，参数占的内存更大。

相较于 Transformer 结构，HSTU 采用了更加简化及完全融合的设计，可以显著降低激活内存的使用。

• 首先，HSTU 将 Attention 机制外的线性层数量从 6 个减少到 2 个。
• 其次，HSTU 采用了 Kernel 融合，将多个 OP 融合为一个，将每层的激活内存使用量减少到 BF16 格式的 2d+2d+4hdqk +4hdv +2hdv = 14d。

作为对比，Transformer 在 Attention 后还有 FFN + Dropout（中间激活 3hdv），随后还有 LayerNorm 等，总的激活为 2d + 4dff +2d+1d = 5d+4dff，通常 dff=4d，hdv >= d，考虑其他部分，总的激活内存大概为 33d。因此，HSTU 实际可以支持超过 2x 的深度网络。

此外，由于词表也会占据大量内存。以 10B 词表，512 的 Embedding 维度及 Adam 优化器为例，在 FP32 精度下存储 Embedding 向量及优化器状态需要 10B*512*3*4 Byte=60TB 内存。其中 3 表示原始 Embedding 及一阶矩和二阶矩。为此，作者采用 Rowwise AdamW 优化器，并将优化器状态放在 DRAM 中，从而将 HBM 显存开销从 12 byte 降低到 2 byte。

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

作者提到的最后一个挑战是推荐系统在 Serving 阶段需要处理大量候选 Item。

• 在检索方面，Encoder 成本完全可以分摊，并且存在量化、哈希或束搜索、层次检索等手段。
• 在排序方面，则要面临可能多达数万个候选 Item。

为此，作者提出了一种名为 M-FALCON（Micro Batched-Fast Attention Leveraging Cacheable OperationNs）的算法，用于处理输入序列长度为 n 的 m 个候选 Item 的推理任务。在单次 Forward 中，M-FALCON 通过调整 Attention Mask 和相对位置偏置，同时处理 bm 个候选 Item。使得针对这 bm 个候选 Item 执行的 Attention 操作完全相同。这一改进将 Cross Attention 的计算成本从 O(bmn2d) 降低到 O((n+bm)2d)=O(n2d) ，前提是 bm 是远小于 n 的常数。可以选择将全部 m 个候选 Item 划分为 [m/bm] 个大小为 bm 的 Micro Batch，以利用 Encoder 级别的 KV Cache。总体而言，M-FALCON 使得模型复杂度能够与 DLRM 模型排序阶段的候选 Item 数量线性扩展。

六、实验

6.1 验证 HSTU Encoder 的归纳假设

6.1.1 传统序列场景

如下图 Table 4 所示，作者在广泛采用的 2 个推荐系统数据集（MovieLens 和 Amazon Reviews）上评估了 HSTU 性能，并将 SASRec 作为基线（值得注意的是，此处采用的评估方法与工业级应用场景存在显著差异）。“SASRec (2023)” 表示论文中报告的最佳 SASRec 配置；“HSTU” 表示使用了与 SASRec 相同的配置（如层数、Head 数等）。“HSTU-large”则表示更大规模的 HSTU Encoder（层数增加 4 倍，Head 数增加 2 倍）。结果表明：

• 在相同配置下，HSTU 凭借专为推荐进行的优化设计，显著超越了基线；
• 当规模扩大时，HSTU 的性能进一步提升。​

6.1.2 工业级流式场景

如下图 Table 5 所示，作者进一步在工业级流式场景下对比了 HSTU 及其变种相比标准 Transformer 模型的优势。总体而言，在小规模实验环境下，HSTU 以 1.5x - 2x 的更快速度实现了更高的质量，并减少了 50% 的 HBM 显存占用。

6.2 Encoder 效率

随机长度：如下图 Figure 4 与 Figure 5a 展示了 SL 对模型指标的影响。当 α=1.6 时，长度为 4096 的序列大部分时间可以缩减至 776 的长度，即移除了超过 80% 的 Token。即便稀疏率提升至 64%-84%，主要任务获得的 NE 下降不超过 0.002（0.2%）。表明适当的 α 值，SL 不会对模型质量产生负面影响，并能实现高稀疏度以降低训练成本。

Encoder 效率：如下图 Figure 5 对比了 HSTU 与 Transformer Encoder 在训练和推理环境下的效率。对于 Transformer，作者采用了 FlashAttention2。序列长度范围从 1,024 到 8,192，并在训练过程中应用 SL。在评估中，对 HSTU 和 Transformer 采用相同配置（d=512，h=8，dqk=64），并消融相对注意力偏置。作者在 H100 GPU 上以 BF16 比较了 Encoder 的表现。总体而言，HSTU 在训练和推理中的效率分别比 Transformer 高出 15.2x 和 5.6x。

6.3 GR 与 DLRM 在工业级流式场景中的对比

最后，作者在工业级流式场景中，将 GR 的端到端性能与最先进的 DLRM 基线进行了对比。如下图 Table 6 和 Table 7 所示，GR 不仅在离线测试中显著优于 DLRM，在 A/B 测试中也带来了 12.4% 的提升。

如下图 Figure 6 所示，作者进一步对比了 GR 和生产环境 DLRM 的效率，尽管 GR 模型计算复杂度高出 285x，但由于 HSTU 和新型 M-FALCN 算法，在 1024/16384 个候选 Item 排序时，分别获得了 1.50x/2.99x 的 QPS。

6.4 推荐系统的缩放法则

众所周知，在大规模工业场景中，DLRM 在特定的计算和参数规模下，其性能会趋于饱和。为了深入理解这一现象，作者对比了 GR 与 DLRM 的可扩展性。

鉴于特征交叉层对 DLRM 性能至关重要，作者采用了多种方法来扩展 DLRM 基线模型在 Ranking 场景的表现。具体实验包括 Transformer 架构、DHEN、以及生产环境中应用的 DCN 变体。对于检索基线模型，由于已采用残差结构，作者通过扩大隐藏层大小、Embedding 维度及层数来提升性能。

对于 HSTU ，则通过调整 HSTU 的超参数来扩展模型，这些超参数包括残差层数量、序列长度、Embedding 维度以及 Attention Head 数目。此外，作者还对检索任务中的数目进行了调整。

最终，所有主要指标，包括用于检索的 Hit Rate@100 和 Hit Rate 500，以及用于排序的 NE，在给定适当超参数的情况下，均经验性地遵循所用计算量的幂律关系。如下图 Figure 1 所示，作者在 3 个数量级的范围内观察到这一现象，直到能够测试的最大规模（序列长度  8192，Embedding 维度 1024，层数 24），此时所使用的总计算量接近 GPT-3 和 LLaMA-2 所使用的总计算量，作者首次证明了 LLM 的缩放法则同样可能适用于大规模推荐系统。

七、参考链接

1. ​​​https://arxiv.org/abs/2402.17152​​​
2. ​​​https://github.com/facebookresearch/generative-recommenders​​​​

本文转载自​​AI闲谈​​，作者： AI闲谈 ​​​​