字节在《Adaptive Domain Scaling for Personalized Sequential Modeling in Recommenders》提出了它们的关于序列建模的多场景建模实现ADS：

摘要

在像抖音这样的超级应用中，用户通常会在多个业务场景中展现出复杂的行为模式和多样化的意图，这对当前工业界的多领域推荐系统提出了巨大的挑战。为了缓解不同领域之间的差异，研究和工业实践通常强调采用复杂的网络结构以适应多样化的数据分布，而忽视了从多领域角度对用户行为序列的深入理解。

本文提出了自适应领域缩放（Adaptive Domain Scaling, ADS）模型，该模型全面增强了跨多个领域的目标感知序列建模中的个性化能力。具体而言，ADS包含两个主要模块：

1. 个性化序列表示生成（Personalized Sequence Representation Generation, PSRG）
2. 个性化候选表示生成（Personalized Candidate Representation Generation, PCRG）

这些模块通过动态学习不同领域下的用户行为序列item表示和候选目标item表示，促进了自适应用户意图理解，从而实现了量身定制的多领域学习。我们在一个公共数据集和两个十亿级别的工业数据集上进行了实验，广泛的结果验证了ADS的高效性和兼容性。

此外，我们在两个有影响力的业务场景（包括抖音广告平台和抖音电商服务平台）上进行了在线实验，两者均显示出显著的业务改进。目前，ADS已在字节跳动的多个推荐服务中全面部署，服务于数十亿用户。

1 引言

随着数字内容的指数级增长和互联网的广泛使用，推荐系统在提升用户体验和缓解信息过载方面发挥了至关重要的作用[29]。在实际应用中，为了提高用户留存率并促进商业利益，工业推荐的需求广泛分布在多个领域[12, 14]。例如，如图1所示，在抖音（全球最大的视频观看应用之一）中，主要领域包括短视频、直播预览和直播滑动，用户可以在这些场景中观看短视频、直播流，并享受电商和本地生活服务。此外，由于其亿级用户规模，不同的用户群体(如来自不同国家、不同性别、高活跃度或低活跃度的用户)也构成了不同的领域。由于不同领域的数据分布差异较大，这为推荐系统带来了显著的多领域建模问题[30]。

图1

为此，工业界的常见实践通常是构建一个共享底层-多头输出的模型结构，利用分离和统一混合建模的优势来处理多领域问题[4]。为了进一步改进这一点，最近的研究致力于构建精细的网络结构以增强多领域建模，例如领域级方法如星型拓扑自适应推荐器（STAR）[18]、渐进分层提取（PLE）[19]，以及实例级方法如自适应参数生成网络（APG）[25]、AdaSparse[27]等。然而，大多数现有方法都专注于复杂的特征交互网络结构，而多领域序列建模的方法却较少受到关注。

序列建模在工业推荐系统中扮演着至关重要的角色，其中最流行且有效的方法是基于target attention的方法，例如深度兴趣网络（DIN）[32]、特征协同网络（CAN）[2]和多头注意力（MHA）[20]。在抖音的广告系统中，从转化率预测（CVR）模型中移除target attention模块后，ROC曲线下面积（AUC）指标下降了超过0.2%，这是一个显著且不可接受的损失，证明了基于target attention的序列建模的有效性。尽管其重要性，现有target attention方法较少考虑多领域差异的影响，这为多领域建模留下了显著的空白。

一般来说，当前用户序列的target attention机制可以表述为一个典型的query-key-value建模范式：

\[g(Rep_{cand}, Rep_{seq}) \times Rep_{seq}\]

其中：

• $Rep_{cand}$表示需要预测点击/转化概率的候选target item的表示
• $Rep_{seq}$表示用户序列嵌入
• $g$计算任意序列item与target item对之间的注意力权重

如前所述，当前工业推荐系统通常遵循共享底层嵌入范式，这意味着：

1. 候选item和用户行为item的嵌入表完全共享，未考虑不同领域item和用户之间的差异；
2. 候选item作为不同key/value的共享查询，未考虑用户序列中多领域item的差异

因此，这从两个方面对当前多领域推荐系统提出了潜在挑战：

• 序列表示的个性化：在多领域中，相同item在不同用户序列中的表示是必要的。例如，新用户可能喜欢观看高点赞视频，而长期用户可能更关注他们关注的视频创作者。因此，相同的视频对不同用户领域表现出不同的吸引力，而其嵌入是不同用户序列中的共享表示，这阻碍了推荐系统捕捉用户的实际意图。

• 候选item的个性化：对于不同用户或同一用户序列中的不同item，候选item由于多领域影响具有不同的影响和功能。例如，用户在抖音商城中的购物行为可能主要受产品价格影响，而在抖音短视频场景中，视频创作者对内容偏好用户的影响更为显著。因此，相同的候选item应在不同领域中个性化，以适应用户序列中的不同历史项。

为了克服这些局限性并填补多领域target attention建模的空白，我们提出了自适应领域缩放（Adaptive Domain Scaling, ADS）模型，该模型充分挖掘了当前基于target attention的推荐系统的个性化建模能力，并在多领域任务中提供了更准确和自适应的意图理解能力。具体而言，ADS包含两个模块：

• 个性化序列表示生成（PSRG）
• 个性化候选表示生成（PCRG）

在PSRG中，我们设计了一种新颖的共享-私有结构，用于学习用户行为中的多领域项表示，旨在为序列items生成个性化表示，即相同item在不同用户序列中具有不同的表示。在PCRG中，候选item通过为不同序列items生成不同的目标候选表示，进一步增强了个性化建模能力。通过将领域相关信息作为生成结构的输入，多领域的影响被充分注入序列建模中，从而增强了用户意图理解能力。值得注意的是，ADS是一个高效的即插即用网络，可以轻松集成到现有推荐系统中。

本工作的贡献总结如下：

• 我们提出了自适应领域缩放（ADS）模型，这是一种有效的即插即用个性化网络结构，通过个性化target attention建模来实现多领域用户意图理解。我们在一个公共数据集和两个十亿规模的工业数据集上进行了广泛实验，结果验证了其优越性。
• 我们在框架中开发了个性化序列表示生成和个性化候选表示生成模块，从用户行为序列和候选target item的角度捕捉多领域特征，增强了当前target attention机制的多领域学习效果。
• 我们在字节跳动抖音的广告系统和电商系统中部署了ADS，分别带来了抖音广告系统总收入的1.09%提升和电商系统的0.79%提升。目前，ADS已在字节跳动的多个推荐系统中全面部署，服务于数十亿用户。

2 方法论

2.1 预备知识

2.1.1 问题定义

本文主要关注推荐系统中的排序建模任务，这是一个典型的二分类问题。以点击率（CTR）预测为例，概率 $\hat{y}$ 可以通过以下公式计算：

\[\widehat{y} = f(\mathbf{E}_U, \mathbf{E}_I, \mathbf{E}_O) \quad (1)\]

其中：

• $E(\cdot)$ 表示嵌入函数，将原始分类特征直接转换为嵌入，连续特征则先分桶再嵌入为稠密向量。
• $f$ 是基于多层感知机（MLP）的转换函数。
• $U$、$I$ 和 $O$ 分别表示用户侧、目标候选item侧和其他特征。用户侧特征通常包括：人口统计特征（如用户位置和语言）和行为特征（如用户的观看列表或购物列表）。item侧特征包括项的描述性特征（如类别、创作者等）。此外，其他特征 $O$ 通常包含context特征和user-item交叉特征。

2.2 提出的ADS模型

图2

ADS模型的结构如图2所示，整体由两个主要部分组成：

• 个性化序列表示生成（PSRG）：通过共享-私有学习结构生成动态行为item嵌入，使得相同item在不同领域中具有不同的表示
• 个性化候选表示生成（PCRG）：捕捉序列items的不同方面，并为每个序列items生成多个自适应查询（即候选item），以反映不同查询对多样化序列items的影响

通过PCRG和PSRG生成的自适应Q、K和V，目标感知建模机制（如MHA、DIN和CAN）可以轻松集成到该框架中，促进多领域场景中的兴趣捕捉。

2.2.1 个性化序列表示生成（PSRG）

当前大规模工业推荐系统通常采用共享嵌入层将原始ID和其他特征嵌入为dense向量。这种方式下，嵌入表中的特定item具有统一的嵌入表示，在不同用户序列中共享，忽略了多领域差异的影响。

PSRG的基本思想是：为用户行为序列中的每个item嵌入动态生成一个个性化层，使得原始共享表示可以在多领域中多样化。具体而言，我们使用领域相关特征嵌入 $E_D \in R^{d_D}$ 作为PSRG生成部分的输入，这些特征包括：

• 1. 显式领域指示特征(explicit-domain-indicator features)：区分样本所属的领域。例如，指示器范围为[0, 2]，表示抖音中的三个不同业务场景。
• 1. 隐式领域指示特征(implicit-domain-indicator features)：在推荐系统中，某些领域难以明确定义。例如，用户是否高活跃度。因此，需要结合工程师构建的统计特征以进一步捕捉和区分不同领域。

如图2所示，序列items的weight和bias生成过程旨在动态调整原始item嵌入。

序列权重生成网络（Sequence-Weight Gen-Net）

用户序列嵌入表示为：$E_S \in R^{T \times d_S}$

其中：

• $T$ 和 $d_S$ 分别表示用户序列长度和每个序列items的嵌入维度。

基于领域特征 $E_D$，权重生成过程包括：私有权重部分和共享权重部分，以捕捉多领域的共性和个性。私有部分通过两层MLP生成私有权重：

\[\mathbf{W}_{private} = Sigmoid(ReLU(\mathbf{E}_D \mathbf{W}_1^T + \mathbf{b}_1) \mathbf{W}_2^T + \mathbf{b}_2) \quad (2)\]

其中：

• $\mathbf{W}_1^T \in \mathbb{R}^{d_D \times d_h}$，$\mathbf{W}_2^T \in \mathbb{R}^{d_h \times (d_S \times d_S)}$：表示weight
• $\mathbf{b}_1 \in \mathbb{R}^{d_h}$，$\mathbf{b}_2 \in \mathbb{R}^{(d_S \times d_S)}$：表示bias
• $d_h$ 表示隐藏层维度

通过引入两层函数而非单层，不仅可以提高模型的表达能力，还能显著减少模型参数和计算成本。

基于 $\mathbf{W}{private}$，进一步定义一个全局权重 $\mathbf{W}{shared} \in \mathbb{R}^{(d_S \times d_S)}$，作为可学习矩阵在所有用户间共享。生成的权重定义为：

\[\mathbf{W}_{generated} = \eta * (\mathbf{W}_{shared} \odot \mathbf{W}_{private}) \quad (3)\]

其中：

• $\odot$ 表示逐元素乘积。由于引入了 $Sigmoid$，$\mathbf{W}{private}$ 的值范围为[0, 1]，因此引入缩放超参数 $\eta$ 以扩大 $\mathbf{W}{private}$ 的表达范围。

序列偏置生成网络（Sequence-Bias Gen-Net）

类似于权重生成过程，偏置生成可以通过以下公式实现：

\[\mathbf{b}_{generated} = ReLU(\mathbf{E}_D \mathbf{W'}_1^T + \mathbf{b'}_1) \mathbf{W'}_2^T + \mathbf{b'}_2 \quad (4)\]

其中：

• $\mathbf{W’}_1^T \in \mathbb{R}^{d_D \times d_h}$，$\mathbf{W’}_2^T \in \mathbb{R}^{d_h \times d_S}$：为weight
• $\mathbf{b’}_1 \in \mathbb{R}^{d_h}$，$\mathbf{b’}_2 \in \mathbb{R}^{d_S}$：为bias

通过生成的权重和偏置，PSRG可以通过以下公式实现：

\[\mathbf{E}_{S-personalized} = \mathbf{E}_S Reshape(\mathbf{W}_{generated})^T + \mathbf{b}_{generated} \quad (5)\]

其中：

• $Reshape$ 操作将一维向量形式的 $\mathbf{W}_{generated}$ 重塑为二维矩阵形式，形状为 $d_S \times d_S$。

2.2.2 个性化候选表示生成（PCRG）

除了序列的个性化建模，另一个重要部分是：对target item的多领域建模，target item通常在target attention中扮演query的角色。个性化候选item包括两个方面：

• 1.类似于序列表示，target item本身的表示也通过共享嵌入层嵌入，未在不同领域中个性化。
• 2.候选item在不同领域的序列items中扮演不同角色。例如，用户在抖音商城频道中的观看列表反映其购物兴趣，而在短视频频道中反映其内容偏好。

多查询生成网络（Multi-Query Gen-Net）

图3

为此，我们提出了多查询生成网络（如图3中间子图所示），在领域相关特征 $E_D$ 和原始target item嵌入 $E_Q \in R^{d_Q}$ 的指导下，生成多个查询：

\[\mathbf{E}_{Q-private} = ReLU(\left(\mathbf{E}_D \oplus \mathbf{E}_Q\right) \mathbf{W}_{q1}^T + \mathbf{b}_{q1}) \mathbf{W}_{q2}^T + \mathbf{b}_{q2} \quad (6)\]

其中：

• $\oplus$ 表示拼接操作
• $\mathbf{W}{q1}^T \in \mathbb{R}^{(d_D + d_Q) \times d_h}$，$\mathbf{b}{q1} \in \mathbb{R}^{d_h}$：表示第一层的w, b
• $\mathbf{W}{q2}^T \in \mathbb{R}^{d_h \times (T \times d_Q)}$，$\mathbf{b}{q2} \in \mathbb{R}^{T \times d_Q}$：表示第二层的w, b
• 隐藏层维度 $d_h « (T \times d_Q)$，使得计算成本可控。

分块查询生成（Chunked-Query Generation）

对于长序列（$T$ 为数百或更高）的情况，我们还设计了一种轻量级的分块查询生成方法以提高计算效率。如图3右侧子图所示，由于用户的相邻行为往往发生在同一领域，原始序列可以划分为 $G$ 个块，相邻item形成一个组。因此，生成的 $\mathbf{E}_{Q-private} \in \mathbb{R}^{(G \times d_Q)}$ 可以进一步重复为 $\mathbb{R}^{(T \times d_Q)}$，计算成本进一步降低。

最终生成的多查询

对应于多个私有查询 $\mathbf{E}{Q-private} \in \mathbb{R}^{(T \times d_Q)}$，我们使用原始查询 $\mathbf{E}_Q$ 作为共享基，即 $\mathbf{E}{Q-shared} = tile(\mathbf{E}_Q)$，其中 $tile$ 表示将 $\mathbf{E}_Q$ 重复 $T$ 次。最终生成的多查询通过残差方式获得：

\[\mathbf{E}_{Q-personalized} = Reshape(\mathbf{E}_{Q-private} + \mathbf{E}_{Q-shared}) \quad (7)\]

其中：

• $Reshape$ 操作将一维向量重塑为二维矩阵形式，形状为 $T \times d_Q$。

2.2.3 目标感知注意力与预测

通过上述个性化查询 $\mathbf{E}{Q-personalized} \in \mathbb{R}^{T \times d_Q}$ 和个性化序列items $\mathbf{E}{S-personalized} \in \mathbb{R}^{T \times d_S}$，target attention模块计算每个项的注意力权重，并在查询的指导下聚合序列。通常，个性化查询和项可以轻松集成到许多流行的注意力方法中，如多头target attention、DIN和CAN。

以多头target attention为例，对于每个头，候选item和序列items首先通过以下公式转换：

\[\mathbf{Q} = \mathbf{E}_{Q-personalized} \mathbf{W}_Q \quad (8)\] \[\mathbf{K} = \mathbf{E}_{S-personalized} \mathbf{W}_K \quad (9)\] \[\mathbf{V} = \mathbf{E}_{S-personalized} \mathbf{W}_V \quad (10)\]

其中：

• $\mathbf{W}_Q \in \mathbb{R}^{d_Q \times d_A}$，$\mathbf{W}_K$ 和 $\mathbf{W}_V \in \mathbb{R}^{d_S \times d_A}$，$d_A$ 表示target attention中的维度大小。第 $t$ 个查询-键对的注意力权重 $\mathbf{z}’[t]$ 可以通过以下公式计算： \(\mathbf{z}'[t] = \frac{\mathbf{Q}_t^T \mathbf{K}_t}{\sqrt{d_A}} \quad (11)\) 随后通过softmax操作对个性化权重进行归一化并聚合个性化序列：

\[\mathbf{z} = softmax(\mathbf{z}'), \quad \mathbf{s} = \sum_{t=1}^T (\mathbf{z}[t] \cdot \mathbf{V}_t) \quad (12)\]

通过序列建模输出 $\mathbf{s}$ 和其他特征嵌入（包括 $\mathbf{E}_U$、$\mathbf{E}_I$ 和 $\mathbf{E}_O$），拼接层和多个高层MLP将所有信息合并并输出预测结果，训练损失通过二元交叉熵函数计算：

\[\mathbf{E}_{all} = \mathbf{s} \oplus \mathbf{E}_U \oplus \mathbf{E}_I \oplus \mathbf{E}_O, \quad \hat{y} = MLP(\mathbf{E}_{all}) \quad (13)\]

3 实验

3.1 实验设置

数据集和实验设置
为了充分评估所提出的ADS模型，我们在一个公共数据集（即淘宝数据集）和两个来自抖音的十亿级工业数据集（即抖音广告平台和抖音电商平台）上进行了实验。表1报告了这三个数据集的统计信息。

• 淘宝数据集：淘宝数据集来自[33]，提供了淘宝用户行为数据，目前广泛用于序列建模方法[3]。前7天数据用于训练，其余用于测试。过滤了至少有200次交互和10次正向行为的用户，以及至少有10次交互的项。数据集中有9,439个类别，我们将每个类别视为一个领域。在该数据集中，页面浏览被视为负向交互，其他行为被视为正向标签（订单）。

• 抖音广告数据集：我们选择了抖音广告中的转化率（CVR）预测任务，并收集了从2022年12月14日到2023年3月10日的在线流量日志子集，共87天，17.3亿样本。前77天用于模型训练，其余10天用于评估。在抖音广告平台中，根据用户的外部行为，数据集可以分为三个主要领域：直播支付、直播订单和短视频购物，分别表示为领域1、2和3。

• 抖音电商数据集：我们选择了抖音直播电商服务中的两种用户购物行为（点击和订单）作为预测目标。收集了从2024年1月1日到3月1日的在线流量日志子集，共61天，25.2亿样本。前54天用于训练，最后一周用于验证。数据集涉及抖音直播中的两个典型场景：直播预览和直播滑动，分别表示为领域1和2。

对比方法和评估指标
为了全面比较所提出的ADS与现有方法，我们选择了几个代表性的SOTA模型，分为两类：

1. 基于DNN的方法：包括DNN、DeepFM[9]、DCNv2[23]、APG[25]、AdaSparse[27]、DFFM[10]、MaskNet[24]。
2. 基于target attention的骨干方法：包括DIN[32]、MHA[20]和CAN[2]，并在此基础上添加了最近的多领域嵌入学习方法进行比较，包括FRNet[21]和PEPNet[6]。

所有方法均基于Tensorflow实现，并通过交叉熵损失进行优化。采用Adam[13]优化器，初始学习率为0.00002。使用AUC指标评估排序模型的性能。此外，我们引入了相对改进（Imp.）[26]来衡量相对AUC增益，计算公式如下（随机策略的AUC值为0.5）： \(AUC\ Imp. = \left(\frac{AUC(MeasuredModel) - 0.5}{AUC(BaselineModel) - 0.5} - 1\right) \times 100\% \quad (14)\)

3.2 实验结果

整体性能
表2展示了不同方法在三个数据集上的对比结果。对于淘宝数据集，由于领域数量较多，仅提供了整体结果。为清晰起见，结果按四组进行报告：第一组列出基于DNN的方法，其余三组分别列出基于DIN、MHA和CAN的多领域target attention方法。从结果中可以得出以下几点观察：

1. 与基于DNN的方法相比，基于target attention的序列建模在排序模型中起着至关重要的作用。从表中可以看出，DIN、MHA和CAN在淘宝数据集上分别提升了0.54%、0.23%和0.17%，在抖音广告数据集上分别提升了0.22%、0.19%和0.04%，在抖音电商的点击预测任务中分别提升了0.49%、0.31%和0.11%。

2. 现有的多领域方法对排序模型普遍有积极影响。具体而言，在基于DNN的方法中，AdaSparse在淘宝和抖音广告数据集上优于基线。FRNet和PEPNet在不同组中也表现出改进的性能。

3. 所提出的ADS在基于DIN、MHA和CAN的组中始终表现最佳，显示出其高效性和兼容性。具体而言，在淘宝数据集中，ADS在DIN、MHA和CAN组中分别比第二好的方法提升了0.47%、0.20%和0.20%。在抖音广告数据集中，ADS比第二好的方法分别提升了0.20%、0.12%和0.14%。在抖音电商数据集中，ADS在点击预测任务中分别比第二好的方法提升了0.05%、0.07%和0.15%，在订单预测任务中分别提升了0.11%、0.15%和0.18%。此外，在两个工业数据集的每个领域中，ADS均显著优于对比方法，证明了其个性化target attention机制的优势。

3.3 消融实验与敏感性分析

消融实验
为了进一步评估ADS中两个模块（PCRG和PSRG）的性能，我们在抖音广告数据集上进行了消融实验。如表3所示，移除PCRG模块后，整体性能在基于DIN、MHA和CAN的方法中分别下降了0.06%、0.03%和0.23%。此外，移除PCRG和PSRG模块后，整体性能分别下降了0.23%、0.23%和0.29%。因此可以得出结论，个性化target item和个性化序列对ADS均有积极影响，验证了这些模块的有效性。

ADS中分块数量的敏感性分析
为了研究分块数量的影响，我们对ADS进行了敏感性分析。具体而言，我们通过改变每个分块中的项数（从[1, 2, 5, 10]）来研究性能模式，并从训练效率和模型性能两个方面进行分析。

• 训练效率模式：我们通过总结和比较不同分块下的模型参数和训练浮点操作（FLOPs）来评估模型训练效率，结果如图4所示。从图中可以明显看出，随着模型个性化程度的提高，模型参数和训练FLOPs持续增加。

• 性能模式：图5展示了通过改变每个分块中的项数得到的模型性能模式。具体而言，首先，与原始的DIN、MHA和CAN相比，ADS及其不同分块版本在点击和订单预测任务中均表现出明显的性能提升。此外，随着分块中项数的减少，模型性能持续提升，最个性化的模型（即无分块的ADS）表现最佳，表明考虑候选item的个性化特征具有重要意义。

总体而言，随着ADS个性化能力的提高，性能和训练成本均有所增加，但我们观察到，即使训练成本略有增加（如分块10的ADS与原始方法相比），模型性能仍能实现显著提升，因此实践者可以根据效果和效率的平衡灵活选择参数。

3.4 在线部署

ADS模型通过分片和数据并行策略在多GPU上分布式部署。为了减少存储，低频嵌入被消除。为了进一步提高GPU吞吐量，我们引入了密集计算异步策略，将计算图分为稀疏前向和密集计算部分，从而实现流水线效果，大大提高了训练和推理效率。得益于这些优化，离线训练资源与基线相同，即64个Nvidia A100。以抖音广告为例，训练时间成本从41.3小时略微增加到42.8小时（+3.6%）。在线延迟保持在30ms，无明显变化。

3.5 在线A/B实验

为了研究所提出的ADS在真实工业场景中的性能，我们在抖音的广告系统和电商系统中分别进行了仔细的在线A/B测试。

• 抖音广告：抖音广告的在线实验于2023年11月2日至11月8日进行，覆盖了74,079,729名抖音APP用户。选择了两个指标进行比较：每千次展示成本（CPM）和广告主价值（ADVV）。需要注意的是，部署场景是字节跳动广告的主要流量来源，基线非常强，ADVV或CPM的0.5%提升被认为是显著的。对比结果如表4所示。从表中可以看出，部署ADS后，整体CPM提升了0.52%，ADVV提升了1.00%，显示了ADS自适应序列建模的优势。此外，在抖音广告的两个主要领域（直播和短视频）中，ADS均优于基线并实现了一致的提升，表明其在领域感知序列建模中的有效性。

• 抖音电商：在线实验于2024年1月23日至1月29日在抖音电商平台进行，覆盖了508,926,918名抖音APP用户。实验结果如表5所示。选择了三个指标进行比较：每用户总交易额（GMV/U）、每用户订单数（Order/U）和每千次展示交易额（GPM），这些都是抖音电商中的重要商业指标。与抖音广告实验类似，该部署场景贡献了字节跳动最高的GMV，基线非常强，GMV的0.5%提升被认为是显著的。如表5所示，整体GMV/U、Order/U和GPM分别提升了0.79%、0.36%和0.89%。此外，在直播预览和直播滑动两个领域中均观察到一致的提升，所有提升均经过统计显著性检验（𝑝值<0.01），验证了其有效性。

#

https://arxiv.org/pdf/2502.05523

字节在《Real-time Indexing for Large-scale Recommendation by Streaming Vector Quantization Retriever》提出了流式向量量化检索器。我们来看一下它的实现：

摘要

检索器作为推荐系统中最重要的阶段之一，负责在严格的延迟限制下高效地为后续阶段选择可能的正样本。因此，大规模系统通常会使用一个简单的排序模型，依靠近似计算和索引来粗略地缩小候选规模。考虑到简单模型缺乏生成精确预测的能力，大多数现有方法主要集中于引入复杂的排序模型。然而，索引有效性的另一个基本问题仍未解决，这也成为了复杂化的瓶颈。在本文中，我们提出了一种新颖的索引结构：流式向量量化模型（Streaming Vector Quantization model），作为新一代的检索范式。流式VQ能够实时为item附加索引，赋予其即时性。此外，通过对可能变体的细致验证，它还实现了索引平衡和可修复性等额外优势，使其能够像现有方法一样支持复杂的排序模型。作为一种轻量级且易于实现的架构，流式VQ已在抖音和抖音极速版中部署，并取代了所有主要的检索器，带来了显著的用户参与度提升。

1 引言

在现代推荐系统中，我们不断面临爆炸性增长的语料库，因此由检索、预排序和排序阶段组成的级联框架变得非常普遍。在这些阶段中，检索器的任务是从整个语料库中区分候选样本，但给定的时间却最少。例如，在抖音中，检索器需要从数十亿条内容中筛选出数千个候选样本，而后续阶段只需将候选规模缩小10倍。

然而，扫描所有候选样本会带来极高的计算开销，因此检索阶段不得不依赖于索引结构和近似计算。具体来说，诸如乘积量化（Product Quantization, PQ [8]）和分层可导航小世界（Hierarchical Navigable Small World, HNSW [11]）等索引方法被提出。PQ通过创建“索引”或“聚类”来表示属于它们的全部内容。当一个聚类被选中时，其所有内容都会被检索出来。同时，用户侧和内容侧的信息被解耦为两个独立的表示，用户表示用于搜索相关聚类。这导致了一种“双塔”架构 [2, 7]，其中每个塔由一个多层感知机（MLP）实现。由于其显著降低计算开销的能力，这种方法在许多工业场景中得到了广泛应用。在下文中，我们将其称为“HNSW双塔”。

尽管HNSW双塔架构简单，但它存在两个缺点：

• （1）其索引结构需要定期重建，在此期间内容表示和内容索引分配是固定的。然而，在一个充满活力的平台上，新内容每秒都在提交，聚类语义也会因新兴趋势而变化，而这些变化在建模中被忽略了。此外，这种构建过程与推荐目标并不一致。
• （2）双塔模型很少提供用户-内容交互，因此生成的预测较弱。不幸的是，在大规模应用中，复杂的模型（如MLP）会带来难以承受的计算开销。

许多现有方法都聚焦于这些问题，并开发了新的索引结构。然而，这些方法主要设计用于支持复杂模型，而忽略了索引本身的关键问题。根据我们的实践经验，索引即时性和索引平衡性与模型复杂性同样重要。如果索引结构严重失衡，热门内容会集中在少数几个索引中，导致模型难以区分它们。例如，在深度检索（Deep Retrieval, DR [4]）中，我们从路径中收集了500𝐾条内容，而仅排名第一的路径就生成了超过100𝐾个候选样本，这严重降低了检索效果。

在本文中，我们提出了一种新颖的索引结构——流式向量量化（streaming VQ）模型，以提升检索器的能力。Streaming VQ具有独特的实时将内容分配到合适聚类的特性，使其能够捕捉新兴趋势。此外，我们还详尽地研究了每种变体，以确定实现索引平衡的最佳解决方案。Streaming VQ使得索引内的内容可区分，因此它能够在保持优异性能的同时生成更紧凑的候选集。尽管它主要关注索引步骤，但它也支持复杂模型和多任务学习。凭借这些创新机制，streaming VQ在抖音和抖音极速版中超越了所有现有的主流检索器。事实上，它已经取代了所有主要检索器，带来了显著的用户参与度提升。本文提出的模型的主要优势总结如下：

• 实时索引分配与自修复能力：内容在训练过程中被实时分配到索引中，并且索引能够自我更新和修复。整个过程无需中断步骤。
• 平衡的索引结构：Streaming VQ 提供了平衡良好的索引，这有助于高效地选择内容。通过一种合并排序（merge-sort）的改进，所有聚类都有机会参与到推荐过程中。
• 多任务学习的优秀兼容性：Streaming VQ 展现出与多任务学习的出色兼容性，并且能够支持与其他方法相同的复杂排序模型。
• 易于实现的特性：最后但同样重要的是，与近期的工作相比，Streaming VQ 以其易于实现的特性脱颖而出。它具有简单清晰的框架，主要基于 VQ-VAE [17] 的现成实现，这使得它能够轻松部署于大规模系统中。

2 相关工作

如前所述，由于扫描整个语料库的计算开销过高，各种索引结构被提出以在可接受的误差范围内近似选择候选样本。乘积量化（Product Quantization, PQ [8]）就是这样一个例子，它将内容聚集到聚类中。当某些聚类被选中时，属于这些聚类的所有内容都会被检索出来。可导航小世界（Navigable Small World, NSW [10]）通过逐步插入节点来构建图，形成节点之间的捷径以加速搜索过程。分层可导航小世界（Hierarchical Navigable Small World, HNSW [11]）提供了分层结构，能够快速缩小候选规模，因此被广泛采用，尤其是在大规模场景中。此外，还有一些基于树的方法 [6, 14] 和局部敏感哈希（Locality Sensitive Hashing, LSH）方法 [15, 16]，旨在近似选择候选样本。

在建模方面，迄今为止最流行且基础的架构是所谓的“双塔模型”，其主要源自 DSSM [7]。双塔模型将用户侧和内容侧的原始特征分别输入到两个独立的多层感知机（MLP）中，并获取相应的表示（嵌入）。用户对某个内容的兴趣通过这两个嵌入的点积来表示。由于它将内容和用户信息解耦，在服务阶段可以预先存储内容嵌入，并通过近似最近邻（Approximate Nearest Neighbor, ANN）方法搜索结果。

然而，解耦用户和内容信息会丢弃它们之间的交互，而这种交互只能通过复杂模型（如MLP）来实现。为了解决这个问题，基于树的深度模型（Tree-based Deep Models, TDM [25], JTM [24], BSAT [26]）提出了树状结构，以从粗到细的层次化方式搜索候选样本。在TDM中，内容被收集在叶子节点上，而一些虚拟的非叶子节点用于表示其子节点的整体属性。TDM采用复杂的排序模型，并通过注意力模块交叉用户和内容信息。考虑到HNSW本身已经提供了分层结构，NANN [1] 直接在HNSW上搜索候选样本，同样使用复杂模型。

另一种方法试图避免ANN算法所需的欧几里得空间假设。深度检索（Deep Retrieval, DR [4]）主要由等距层组成，将内容定义为“路径”，并使用束搜索（beam search）逐层缩小候选范围。与TDM和NANN相比，它更关注索引而非排序模型的复杂性。还有一些方法 [9, 12] 使用多索引哈希函数对内容进行编码。

尽管上述方法主要集中在模型复杂性上，BLISS [5] 强调了索引平衡的重要性。它通过迭代强制模型将内容映射到桶中，甚至手动将一些内容分配到尾部桶中以保证平衡。

将内容附加到索引本质上是将它们“量化”为可枚举的聚类。因此，可以考虑向量量化（Vector Quantization, VQ）方法。从引入可学习聚类的VQ-VAE [17] 开始，许多方法 [22, 23] 已经考虑在检索任务中使用它或其变体。在本文中，我们将VQ模型发展为一种以流式方式更新、保持平衡、提供灵活性且轻量级的索引方法，并将其命名为“流式VQ（streaming VQ）”。

3 流式向量量化模型（Streaming VQ）

通常，检索模型包括索引步骤和排序步骤：

• 检索索引步骤使用近似搜索从初始语料库中逐步缩小候选范围
• 检索排序步骤则为后续阶段提供有序的结果和更小的候选集

大多数现有方法都遵循这种两步范式。例如，最流行的双塔架构本质上利用HNSW来高效搜索候选样本。在特定操作轮次中，它首先通过排序模型对邻居节点进行排序（排序步骤），然后选择内容并丢弃其他内容（索引步骤）。同样，TDM和NANN模型也依赖于它们自己的索引结构（基于树/HNSW）。DR主要引入了一种可检索的结构，在实践中我们还需要训练一个排序模型来对结果进行排序，并为索引步骤提供用户侧输入嵌入。DR与其他方法的不同之处在于，DR的索引步骤和排序步骤是按时间顺序执行一次，而其他方法中这两个步骤是交替执行的。

图 1 提出的流媒体VQ模型的训练框架

本文提出的流式VQ模型也由两个按时间顺序执行的步骤组成。在图1中，我们展示了其完整的训练框架（注意，流式VQ可以扩展到多任务场景，但为了简单起见，我们暂时只考虑预测视频是否会被播放完成的“完成任务”）。在索引步骤中，我们采用双塔架构（原因将在第5.5节讨论），并通过独立的塔生成内容侧和用户侧的中间嵌入 v 和 u（图1中的深蓝色和黄色块）。首先，这两个中间嵌入通过一个辅助任务进行优化，该任务采用in-batch Softmax损失函数：

\[𝐿_{aux} = \sum_{o} -\log \frac{\exp(\mathbf{u}_o^T \mathbf{v}_o)}{\sum_{r} \exp(\mathbf{u}_o^T \mathbf{v}_r)},\]

其中：

• $o$ 和 $r$ 表示样本索引。

量化操作出现在内容侧：我们保持一组可学习的聚类（单任务版本为16𝐾，多任务版本为32𝐾），并分配𝐾个嵌入。当生成 v 时，它在聚类集中搜索最近的邻居：

\[k^*_o = \arg\min_k ||\mathbf{e}^k - \mathbf{v}_o||^2,\]

…(2)

\[\mathbf{e}_o = \mathbf{e}^{k^*_o} = Q(\mathbf{v}_o),\]

…(3)

其中：

• $Q(\cdot)$ 表示量化
• e ：所选聚类嵌入
• u ：用户侧嵌入

它们一起优化：

\[𝐿_{ind} = \sum_{o} -\log \frac{\exp(\mathbf{u}_o^T \mathbf{e}_o)}{\sum_{r} \exp(\mathbf{u}_o^T \mathbf{e}_r)}.\]

…(4)

搜索到的聚类作为输入内容的“索引”。这种内容-索引分配被写回参数服务器（Parameter Server, PS）。我们遵循标准的指数移动平均（Exponential Moving Average, EMA [17]）更新：聚类嵌入通过其所属内容的移动平均值进行更新，而内容而非聚类接收聚类的梯度。EMA过程在图1中用红色箭头表示。

检索排序步骤与检索索引步骤共享相同的特征嵌入，并生成另一组紧凑的用户侧和内容侧中间嵌入。由于在此步骤中，更复杂的模型优于双塔架构，因此可以使用交叉特征和3D用户行为序列特征。我们基于连接嵌入为每个任务预测一个独立的塔（头），并由相应的标签进行监督。详细的模型架构可以在第3.5节中找到。

在服务阶段，我们首先通过以下公式对聚类进行排序：

\[\mathbf{u}^T \cdot Q(\mathbf{v}).\]

然后，所选聚类的内容被输入到下一个排序步骤中，并生成最终结果。

以上介绍了所提出方法的基础框架，在本节的剩余部分，我们将详细阐述如何在几个特别关注的方面进行改进，包括索引即时性、可修复性、平衡性、服务技巧，以及如何与复杂模型和多任务学习集成。

3.1 索引即时性

现有检索模型的整体更新周期由候选扫描（检查哪些内容可以被推荐）、索引构建和模型转储组成。其中，主要成本来自索引构建。

对于所有现有检索模型，索引构建是中断的，这导致索引语义的即时更新被忽略。例如，在抖音中，由于我们有数十亿规模的语料库，构建HNSW大约需要1.5-2小时，执行DR中的M步需要1小时。在此期间，索引保持不变。然而，在一个快速发展的平台上，新兴趋势每天都在出现。这种情况不仅需要实时将新提交的视频分配到适当的索引中，还需要同时更新索引本身。否则，它们无法相互匹配，只会产生不准确的兴趣匹配和较差的性能。相反，我们的模型通过流式样本进行训练，内容-索引分配会立即决定并实时存储在PS中（键=内容ID，值=聚类ID），无需中断阶段，并且聚类嵌入通过优化目标强制适应内容。这赋予了它最重要的优势：索引即时性。

现在，在流式VQ中，索引构建变为实时步骤，因此我们已经克服了主要障碍。此外，我们将候选扫描设置为异步，因此整体模型更新周期等于模型转储周期，仅需5-10分钟。

即便如此，仍存在一个潜在问题：内容-索引分配完全由训练样本决定。由于热门内容频繁曝光，它们的分配得到了充分更新。然而，新提交的和不受欢迎的内容获得曝光或更新的机会较少，这进一步恶化了它们的表现。

为了解决这个问题，我们添加了一个额外的数据流——候选流（candidate stream）——来更新它们。与称为“曝光流”的训练流不同，候选流只是以等概率逐个输入所有候选内容。如图1（虚线黑色箭头）所示，对于这些样本，我们仅通过前向传播获取并存储内容-索引分配，以确保其与当前聚类集的语义匹配。由于这些样本没有真实标签，因此不计算损失函数或梯度。

3.2 索引可修复性

流式更新范式是一把双刃剑：由于我们放弃了索引重建，整个模型面临性能退化的风险。这种现象广泛存在于所有检索模型中，但通常通过重建操作来解决。现在对于流式VQ，我们需要在没有重建操作的情况下解决这个问题。

原始的VQ-VAE引入了两个损失函数：一个与 $𝐿_{ind}$ 相同，另一个强调内容-聚类相似性：

\[𝐿_{sim} = \sum_{o} ||\mathbf{v}_o - \mathbf{e}_o||^2.\]

…(6)

在计算机视觉领域 [3, 13]，模式很少变化，因此VQ类方法表现良好。然而，在大规模工业推荐场景中，内容自然会发生归属变化，但 $𝐿_{sim}$ 反而会锁定它们。

在我们早期的实现中，我们遵循了与原始VQ-VAE相同的配置，起初在线指标确实有所改善。然而，我们观察到模型退化：性能随着时间的推移逐渐恶化。随后我们意识到，在我们的平台上，由于全局分布漂移，作为内容的概括表示，聚类的语义每天都在变化。内容-索引关系并非静态，相反，内容可能在不同天内属于不同的聚类。不幸的是，$𝐿_{ind}$ 和 $𝐿_{sim}$ 都只描述了内容属于某个聚类的情况。如果它不再适合该聚类，我们不知道它应该属于哪个聚类。这就是性能退化的原因。

通过用 $𝐿_{aux}$ 替换 $𝐿_{sim}$，我们解决了这个问题。由于 $𝐿_{aux}$，内容嵌入可以及时独立地更新，然后 $𝐿_{ind}$ 根据内容表示调整聚类。经过这一修改后，我们成功观察到了持续的改进。我们将其总结为设计检索模型的原则：内容优先。内容决定索引，而不是相反。

3.3 索引平衡性

推荐模型应能够区分热门内容，并为后续阶段精确选择所需内容。具体来说，对于检索模型，我们希望它们将内容均匀分布在索引中，以便我们只需选择少数索引即可快速缩小候选集。这种特性称为“索引平衡性”。不幸的是，许多现有方法存在流行度偏差，未能提出有效的技术来防止热门内容集中在少数几个顶级索引中。为了缓解这种偏差，BLISS [5] 甚至强制将一些内容分配到尾部聚类。

注意到 $𝐿_{ind}$ 在平均情况下获得最小的量化误差。热门内容占据的曝光量远多于其他内容，因此最小化 $𝐿_{ind}$ 的最直接方法是将它们拆分并分配到尽可能多的聚类中，这自然会实现良好的平衡性。在我们的实现中，流式VQ确实采用了这一策略，并产生了令人惊讶的平衡索引分布（见第5.1节）。

为了进一步提高索引平衡性，我们修改了主要的正则化技术。设 $\mathbf{w}$ 为初步的聚类嵌入，我们在EMA中插入一个流行度项：

\[\mathbf{w}^{t+1}_k = \alpha \cdot \mathbf{w}^t_k + (1 - \alpha) \cdot (\delta^t)^\beta \cdot \mathbf{v}^t_j,\]

其中内容 $j$ 属于聚类 $k$，$t$ 表示时间戳，$\delta$ 表示内容出现间隔，如 [21] 中提出的。这里我们添加了一个超参数 $\beta$ 来调整聚类行为，较大的 $\beta$ 会促使聚类更关注不受欢迎的内容。然后，我们还更新记录聚类出现次数的计数器 $c$：

\[c^{t+1}_k = \alpha \cdot c^t_k + (1 - \alpha) \cdot (\delta^t)^\beta,\]

最终表示计算为：

\[\mathbf{e}^{t+1}_k = \frac{\mathbf{w}^{t+1}_k}{c^{t+1}_k}.\]

我们还在向量量化步骤中提出了“扰动”，即将公式(2)修改为：

\[k^*_o = \arg\min_k ||\mathbf{e}_k - \mathbf{v}_o||_2 \cdot r,\] \[r = \min\left(\frac{c_k}{\sum_{k'} c_{k'}/K} \cdot s, 1\right),\]

其中 $r$ 表示折扣系数，$s = 5$ 是一个阈值。这意味着如果整个聚类的曝光量少于平均值的 $1/s$ 倍，则在内容搜索其最近聚类时会被提升。这也有助于构建一个平衡良好的索引结构。

3.4 服务阶段的合并排序

内容的表示可能具有两种内在语义：个性化和流行度。我们希望根据内容的个性化而非流行度进行聚类。为此，我们显式地将内容表示解耦为个性化部分（嵌入）和流行度部分（偏差）。数学上，将公式(5)修改为：

\[\mathbf{u}^T \cdot Q(\mathbf{v}_{emb}) + v_{bias}.\]

通过这种方式，我们观察到同一聚类内的内容在语义上更加一致。所有训练损失函数也遵循相同的修改。

注意到在公式(11)中，即使同一聚类内的内容具有相同的 $Q(\mathbf{v}{emb})$，$v{bias}$ 也可以用于粗略排序。因此，我们提出了一种合并排序解决方案，以有效选择候选内容进入检索排序步骤。

图 2

如图2(a)所示，$\mathbf{u}^T \cdot Q(\mathbf{v}{emb})$ 提供了聚类排名，而 $v{bias}$ 提供了聚类内内容的排名。然后基于这两部分的总和进行合并排序。这确保所有聚类（即使是那些大小超过排序步骤输入的聚类）都有机会为最终结果提供候选内容，因此我们可以收集一个非常紧凑的集合（50𝐾，仅为DR排序步骤输入大小的10%）以超越其他检索器。

具体来说，我们在这里使用最大堆来实现k路合并排序（图2(b)）。聚类的内容首先独立排序并形成列表，这些列表被分成块（大小=8）。然后将这些列表构建成一个堆，由其头部元素初始化。在每次迭代中，我们从堆中弹出顶部元素，但取出其块中的所有元素。然后，来自同一列表的另一个块及其头部元素被添加到堆中。该策略在保持性能质量的同时有效减少了计算开销。更多细节请参见附录A。

3.5 模型复杂性

如前所述，在检索索引步骤和检索排序步骤中，我们分别评估16𝐾和50𝐾个聚类/内容。这一规模不再难以承受，因此我们可以使用复杂模型。在图3中，我们展示了索引和排序模型的两种架构：双塔架构和复杂架构。

双塔模型（图3左侧）遵循典型的DSSM [7] 架构。内容侧特征和用户侧特征分别输入到两个独立的塔（即MLP）中，并获得紧凑的嵌入。用户对该内容的兴趣通过这两个嵌入的点积计算。特别地，我们为每个内容添加了一个偏差项，并将其添加到最终得分中，以表示内容的流行度。在排序步骤中使用双塔模型的版本称为“VQ双塔”。

复杂版本（图3右侧）也将内容侧和用户侧特征输入以生成两个中间嵌入。然而，内容侧嵌入被输入到一个多头注意力模块 [19] 中作为查询，以提取非线性用户-内容交互线索，其中用户行为序列被视为键和值。然后，转换后的特征以及其他所有特征（包括交叉特征）被输入到一个深度MLP模型中以输出结果。使用复杂排序模型的版本称为“VQ复杂”。

理论上，这两种架构都可以部署在索引和排序步骤中。然而，在我们的实验中，复杂的索引模型并未带来改进。正如第5.5节所讨论的，复杂模型提供的非线性接口违反了欧几里得假设，并可能将聚类和内容划分到不同的子空间中，从而遗漏一些聚类。因此，我们将索引模型保持为双塔架构。

相反，对于排序步骤，复杂版本优于双塔版本。然而，它也带来了更多的计算开销。考虑到投资回报率（ROI），并非所有目标都部署为复杂版本。详细信息见第5.3节。

作为一个娱乐平台，抖音有许多热门话题和新兴趋势，这些内容集中在用户的近期行为序列中。然而，顶级话题已经被充分估计和分发，因此由热门话题主导的序列特征几乎无益于兴趣建模。为了解决这个问题，我们利用Trinity [20] 提供的统计直方图，过滤掉落在用户前5个次要聚类中的内容（填充更多内容以达到足够长度）。生成的序列倾向于长尾兴趣，并提供更多的语义线索。通过修改后的序列特征，某些目标得到了显著改善（见第5.3节）。

我们还在VQ复杂版本中添加了数十个特征以达到其最佳性能。仅添加特征或增加模型复杂性只能产生适度的结果。然而，通过结合这两种技术，我们获得了显著改进的结果。原因是，随着更多特征的加入，我们的模型能够实现高阶交叉并真正利用复杂性。

3.6 多任务流式VQ

尽管前面的讨论是基于单任务框架的，但流式VQ可以扩展到多任务场景。如图1所示，在索引步骤中，用户对每个任务都有独立的表示，但它们共享相同的聚类集。对于每个任务，我们同时计算 $𝐿_{aux}$ 和 $𝐿_{ind}$ 并传播梯度。

对于多任务版本，聚类表示需要针对不同任务进行专门化。具体来说，公式(7)和公式(8)被修改为：

\[\mathbf{w}^{t+1}_k = \alpha \cdot \mathbf{w}^t_k + (1 - \alpha) \cdot \prod_{p} (1 + h_{jp})^{\eta_p} \cdot (\delta^t)^\beta \cdot \mathbf{v}^t_j,\] \[c^{t+1}_k = \alpha \cdot c^t_k + (1 - \alpha) \cdot \prod_{p} (1 + h_{jp})^{\eta_p} \cdot (\delta^t)^\beta,\]

其中 $\eta$ 是另一个用于平衡任务的超参数，$h_{jp}$ 是内容 $j$ 在任务 $p$ 中的奖励。例如，如果视频未完成/完成，则 $h_{jp} = 0/1$。对于停留时间目标，它被设计为对数播放时间。注意，整个奖励始终大于1，因此聚类会倾向于产生更高奖励分数的内容。

检索排序步骤为所有任务共享特征嵌入，并训练各自的双塔或复杂模型。

表 1

---

4 检索模型的详细分析

图 3

在这里，我们将流式VQ与其他现有方法进行比较，并说明为什么它有利于大规模工业应用。

表1列出了我们关注的检索模型的7个方面，并逐一讨论：

• 索引是否面向推荐？ 在本文中，“面向推荐”衡量索引构建过程是否针对推荐目标进行了优化。在HNSW中，索引的构建不考虑其分配的任务。类似地，DR由于其M步骤也不是面向推荐的检索器。由于 $𝐿_{aux}$ 和 $𝐿_{ind}$ 都受到推荐目标的监督，流式VQ是面向推荐的。

• 索引步骤中的负采样方法：HNSW和NANN在索引步骤中不涉及负采样方法。TDM引入了一种随机负采样方法，选择同一层级的另一个节点作为负样本。特别地，DR具有隐式负采样：由于所有节点都通过Softmax归一化，当我们最大化其中一个节点时，其他节点等效地被最小化。然而，这种最小化未考虑样本分布，因此DR仍然严重受到流行度偏差的影响。在我们的实现中，流式VQ在索引步骤中保持双塔架构，因此我们可以直接采用 [21] 中引入的现成in-batch去偏解决方案。

• 流行度去偏：如上所述，DR无法避免热门内容集中在同一路径中。在我们的系统中，DR索引步骤后总共收集了500𝐾个候选内容，而排名第一的路径提供了100𝐾个。相反，由于第3.3节中提出的所有技术，流式VQ中的热门内容广泛分布在索引中。尽管大多数现有方法都关注复杂性，但我们认为流行度去偏是另一个被忽视但至关重要的问题。

• 构建索引的时间成本：在抖音中，我们需要1.5-2小时来设置HNSW，并需要1小时来执行DR的M步骤。在流式VQ中，索引在训练过程中实时构建和更新。

• 索引步骤的候选限制：这意味着我们可以处理多少候选内容作为输入。由于需要存储一些元信息（例如边），它受到单机内存的限制。作为最复杂的结构，HNSW只能存储170𝑀个候选内容。由于我们的语料库规模超过了这一限制，因此会定期随机丢弃一些内容。DR的结构（一个内容可以通过3条路径检索）大大简化，因此我们可以将阈值扩展到250𝑀。当前的流式VQ具有独占结构，因此理论上它可以存储比DR多3倍的候选内容（详细分析见附录B）。我们仅扩展到350𝑀，因为更多的候选内容可能会带来一些过时的信息。

• 排序步骤的节点接触：这里我们展示了系统中每种方法的实际设置，而不是上限。由于HNSW/TDM/NANN在分层结构中检索候选内容，对于它们来说，排序步骤的节点接触指的是它们计算的总次数，而对于DR/流式VQ，它表示排序列表的大小。为了公平比较，我们将NANN和流式VQ的节点接触次数设置为相同（见第5.4节）。注意，由于流式VQ具有平衡良好的索引结构，并且可以在聚类内精细选择内容，因此即使排序候选规模减少10%，它仍然优于DR。

• 适用的排序模型：使用复杂的排序模型总是会显著增加计算开销。众所周知，HNSW无法支持复杂架构。在抖音中，由于投资回报率（ROI）较低，DR在排序步骤中也使用双塔模型。其他检索模型使用复杂架构。

5 实验

在本节中，我们剖析了流式 VQ 的性能，包括聚类可视化和在线指标。然后，我们解释为什么我们更关注索引结构，而不是开发复杂的排序模型。我们还讨论了是否需要索引复杂化/多层 VQ。

5.1 平衡且不受流行度影响的索引

图 4

在图 4 中，我们通过统计直方图（上）和 t-SNE [18]（下）可视化索引分布。在直方图中，我们将聚类按其所属item数量进行聚合。从结果来看，大部分聚类包含的item数 ≤ 25K。考虑到我们有一个十亿规模的语料库和 16K 聚类，在理想的均匀分布下，每个聚类将分配到数万个item。我们得到的结果相当接近这种理想分布。

另一幅图描述了它们在二维空间中的聚合程度，颜色越深表示聚类越大。首先，所有点均匀覆盖整个区域，这意味着它们与其他聚类在语义上是不同的。然后，每个层级的点，尤其是大聚类的点，分散开来，甚至在局部也没有聚集。这表明索引结构能够抵抗流行度的影响。

因此，我们可以得出结论，流式 VQ 确实提供了平衡且不受流行度影响的索引。

5.2 工业实验环境

在本文中，所有实验均在我们的大规模工业应用中实施：抖音和抖音 Lite，用于视频推荐。作为一个娱乐平台，我们专注于提升用户参与度，即日活跃用户（DAUs）。由于用户被均匀分配到对照组和实验组，因此无法直接测量 DAUs。我们遵循 Trinity [20] 中的相同指标。我们计算实验期间用户的平均活跃天数作为平均活跃天数（AAD），平均活跃小时数作为平均活跃小时数（AAH），并将观看时间作为辅助指标。

由于检索器是作为单任务模型进行训练的，因此总是存在指标权衡。例如，优化完成目标的检索器可能通过简单地增强短视频的分发来实现，这将导致更多的展示次数（VV），但观看时间会下降。一般来说，一个可接受的上线应该在展示次数和观看时间上保持平衡（例如，增加 0.1% 的观看时间，但减少 0.1% 的展示次数）。一个更有效的检索器应该同时提高观看时间和展示次数。

在检索阶段，我们已经部署了数百个检索器。因此，我们更倾向于那些占据足够展示次数的检索器，这通过展示比例（IR）来衡量。IR 计算这个检索器贡献了多少展示次数，不进行去重。根据我们的经验，IR 是最敏感且最具预测性的指标。一般来说，如果它的 IR 提高了，我们就得到了一个更有效的检索器。

将检索模型升级为流式 VQ 涉及以下目标：停留时间（ST）、完成（FSH）、有效观看（EVR）、活跃停留时间（AST）、个人页面停留时间（PST）、旧候选停留时间（OST）、评论区停留时间（CST）和 Lite 停留时间（LST）。具体来说，停留时间目标衡量用户观看视频的时间，如果他/她观看了超过 2 秒，则记录为正样本。我们根据实际播放时间给正样本分配奖励。AST/PST/CST 描述了相同的信号，但出现在喜欢页面/个人页面/评论区，而不是信息流标签中。OST 和 LST 也建模了停留时间目标，OST 只是将 ST 应用于 1-3 个月前发布的候选item，而 LST 是专门为抖音 Lite 训练的。完成直接描述了视频是否被看完。有效观看是一个综合目标：它首先通过分位数回归预测观看时间等于 60%/70%/80%/90% 持续时间，然后通过加权和融合预测。

5.3 在线实验

在表 2 中，我们展示了在线性能，其中只列出了具有统计显著性的指标。首先，对于每个模型，两种变化（从 HNSW/DR 到 VQ 双塔，以及从 VQ 双塔到 VQ 复杂模型）在 IR 上都提供了显著的改进。正如前面所展示的，这表明了更好的内在有效性，通常指的是索引平衡、即时性等。

所有实验在观看时间、AAD 和 AAH 上都产生了显著的改进，或者至少具有竞争力的表现。我们可以得出结论，流式 VQ 是比 HNSW 和 DR 更好的索引结构（与 NANN 相比，见第 5.4 节），并且 VQ 复杂模型优于 VQ 双塔。然而，令人惊讶的是，仅索引升级就产生了令人信服的 AAD 增益。这表明，尽管大多数现有工作都集中在复杂性上，但索引的有效性同样重要。

对于完成目标，“*” 表示复杂模型的序列特征没有通过 Trinity 进行去偏。通过比较两行相邻的数据，去偏版本在所有指标上都优于另一个版本，这表明长尾行为为全面描述用户的兴趣提供了补充线索。 抖音和抖音 Lite 在 DAUs 方面已经有一个非常高的基线。此外，检索阶段对展示结果的影响已经被 IR 按比例减少。检索模型的变化多年来没有为 AAD 提供显著的好处。然而，通过流式 VQ 替代，我们在几次上线中见证了令人印象深刻的改进。这验证了流式 VQ 作为一种新型检索模型范式的潜力。

5.4 索引优先，还是排序优先？

为了更好地理解索引和排序步骤在大规模场景中所起的作用，我们还进行了在线实验，比较了基于 EVR 目标的 NANN [1]（最先进的检索模型）与所提出的方法。为了公平比较，我们确保 NANN 和 VQ 复杂模型具有完全相同的计算复杂性。请注意，NANN 和 VQ 复杂模型也使用了更多的特征。

表 3

在表 3 中，我们将“HNSW 双塔”作为基线，并列出其他模型的性能。VQ 双塔、NANN 和 VQ 复杂模型依次提供了越来越好的结果，通过观看时间/AAH 来衡量。从这些结果来看，NANN 似乎与两种 VQ 架构具有竞争力。然而，一方面，正如我们在第 5.2 节中所展示的，NANN 比其获得的观看时间失去了更多的 VV，这并不是非常有效。另一方面，在图 5 中，我们可视化了它们的展示分布（与 HNSW 双塔相比的相对差异），这也得出了不同的结论。

图 5

为了吸引用户，分发更多热门item（VV>1 亿）一直是一个捷径，因为它们很容易产生更多的观看时间和点赞次数。但一个更有效的系统能够精确匹配兴趣，因此小众话题可以获得更多展示机会。例如，添加更多特征也可以增强系统对不受欢迎item的理解，并改善它们的分发。从这个方面来看，两种 VQ 架构符合我们的期望：VQ 双塔将“1 万 - 10 万”的展示量提高了约 2%，同时将“1 亿 +”减少了 1%。此外，VQ 复杂模型甚至将近乎 5% 的“1 万 - 5 万”展示量提高了，同时将近乎 2% 的“1 亿 +”减少了。然而，NANN 保持了不变的分布，这表明它没有充分利用特征和复杂性。总之，VQ 复杂模型在观看时间和 AAH 上优于 NANN，同时减少了热门item的分发。因此，它是我们应用中的更好模型。

可以得出结论，仅仅复杂化排序模型是不足以充分利用模型结构和特征所提供的所有优势的。这是因为整个模型的性能受到索引步骤的限制。只有拥有先进的索引模型，复杂化才能实现其理想性能。因此，我们建议优先优化索引步骤，特别是在大规模场景中。

5.5 索引复杂化

正如第 3 节所展示的，我们也可以在索引步骤中使用复杂的模型。然而，它意外地提供了较差的结果。为了找出原因，我们进一步实施了以下变化：(1) 保持双塔头部，并根据公式（10）附加索引，确定item - 索引分配；(2) 将 e 和 v 输入复杂的模型，如图 3 所示，但不从它那里接收梯度；(3) 除了 e 和 v 之外，共享两个头部的所有其他特征嵌入和 DNN 参数。通过这种方式，我们将item中间嵌入和聚类嵌入强制到相同的语义空间，并尽可能相似。令人惊讶的是，它仍然给出了较差的结果。

图 6

为了理解这种现象，想象我们有两个正样本（item）及其聚类（图 6 中的蓝色圆圈，较深的一个表示聚类）。在双塔索引版本中（左），它遵循欧几里得假设，模型只产生近线性界面，因此聚类与其item保持在相同的子空间中。

#

https://arxiv.org/pdf/2501.08695

kuaishou在《HoME: Hierarchy of Multi-Gate Experts for Multi-Task Learning at Kuaishou》提出了一种新的HoME的多任务实现方式，我们来看一下实现：

摘要

本文介绍了快手短视频服务中遇到的实际问题及从中汲取的经验教训。在工业界，广泛使用的多任务框架是专家混合（Mixture-of-Experts, MoE）范式，该框架通常为每个任务引入一些共享和特定的专家，然后通过门控网络来衡量相关专家的贡献。尽管MoE取得了显著的改进，但在我们的迭代过程中，我们仍然观察到三个严重影响模型性能的异常现象：

• （1）专家崩溃：我们发现专家的输出分布存在显著差异，一些专家在使用ReLU激活函数时，超过90%的激活值为零，这使得门控网络难以分配公平的权重来平衡专家。
• （2）专家退化：理想情况下，共享专家旨在同时为所有任务提供预测信息。然而，我们发现一些共享专家仅被一个任务占据，这表明共享专家失去了其能力，退化为特定专家。
• （3）专家欠拟合：在我们的服务中，有数十个行为任务需要预测，但我们发现一些数据稀疏的预测任务往往忽略其特定专家，并将较大的权重分配给共享专家。原因可能是共享专家能够从密集任务中感知更多的梯度更新和知识，而特定专家由于行为稀疏，容易陷入欠拟合。

基于这些观察，我们提出了HoME，旨在实现一个简单、高效且平衡的多任务学习MoE系统。具体而言，我们进行了三项有见地的改进：

• （1）专家归一化和Swish机制，以对齐专家输出分布并避免专家崩溃。
• （2）层次掩码机制，以增强任务间的共享效率，减少占用问题并远离专家退化。
• （3）特征门控和自门控机制，以确保每个专家都能获得适当的梯度，从而最大化其有效性。

据我们所知，本文是首个专注于提高多任务MoE系统稳定性的工作，我们进行了广泛的离线和在线实验（离线平均提高0.52%的GAUC，在线平均每位用户的播放时间提高0.954%）和消融分析，以证明HoME的有效性。HoME已部署在快手的短视频服务中，每日为4亿用户提供服务。

1 引言

近年来，诸如抖音和快手等短视频应用迅速崛起；与其他平台不同，用户在这些平台上通常没有明确的意图，例如在谷歌上搜索关键词或在亚马逊上购买衣物/食品，而快手几乎扮演了一种娱乐角色，用户无需输入任何概念。如图1所示，当使用快手时，用户通常通过简单地上下滑动屏幕观看多个自动播放的短视频，有时会留下一些互动，例如长观看、评论等。隐式反馈的比例远大于其他场景。因此，快手能够成长为一个拥有全球4亿用户的大型应用的唯一原因是，我们的系统能够提供个性化和有趣的短视频，给用户带来满意的体验。为此，利用用户留下的稀少但多样的行为线索来准确捕捉他们的兴趣是根本任务。通常，常见的做法是将这种学习过程构建为多任务学习范式，以构建一个能够同时输出不同用户互动的多个估计概率的模型，并通过真实的用户行为日志来监督该模型。

图1 典型的快手多任务行为

作为一种典型的多任务解决方案，专家混合（MoE）的思想在工业界广泛用于实现参数软共享。最著名的方法是多重门控专家混合（MMoE），它由两个主要组件组成（如图2(a)所示）：

• 专家网络（expert network）——一组用于建模输入特征和隐式高级特征交叉的专家网络（例如，带有ReLU的MLP）
• 门控网络（gating network）——任务特定的门控网络（例如，带有Softmax的MLP），用于估计不同专家的重要性，以融合它们的输出用于相应任务

最近，一些工作通过引入任务特定专家（例如，CGC）或堆叠更多专家层（例如，PLE，AdaTT）扩展了专家网络，以增强MMoE系统的能力。在快手，我们之前的在线多任务模块配备了MMoE，与基线相比显著改善了我们的A/B测试指标。然而，在推出MMoE后，我们在过去几年中尝试了对多任务建模模块的几种不同更改。但都以失败告终，包括升级到两个或更多专家层，扩展更多共享专家，引入额外的特定专家等。因此，我们开始深入分析以确定可能阻碍我们迭代的潜在原因。不出所料，我们发现了三个严重影响多任务性能的异常现象。

图2 一个朴素的MMoE和专家崩溃问题在实践中出现的说明。如(b)所示，expert总是所有任务分配的门值最大，大多数情况下超过0.98。我们还注意到，expert6的输出要小得多。激活值比其他专家更稀疏，如(c)所示。这些现象表明，在真实数据流场景中， MMoE不稳定，容易崩溃，阻碍了专家之间的公平比较，影响了模型的性能

专家崩溃：我们首先检查了MMoE的门输出情况，并在图2(b)中显示了主要任务分配给6个共享专家的门权重。值得注意的是，所有门（gate）都向共享专家6分配了较大的权重，几乎忽略了其他共享专家。因此，我们接下来检查了共享专家的输出值分布，并观察到它们的显著差异。如图2(c)所示，专家1∼5的均值和方差处于相似水平，但专家6的均值小了100倍。这种不一致的输出分布导致门控网络难以分配公平的权重来平衡不同的专家，这进一步导致处于不同数值水平的专家相互排斥。此外，我们还发现专家输出有太多的0激活（即超过90%的输出），导致其平均导数较小，参数训练不足。

专家退化：在我们解决了上述严重的专家崩溃问题后，我们成功地将多任务模块升级为共享特定的MoE变体CGC。因此，我们好奇门控权重是否能达到预期结果，即所有任务门控网络都能为共享专家和其特定专家分配可感知的分数，以达到平衡状态。不幸的是，我们发现了另一个意外的专家退化现象（如图3所示）。在这里，我们展示了一些主要塔的门控机制的平均分数，我们观察到共享专家几乎不对所有任务做出贡献，而是退化为仅属于少数任务的特定专家。因此，这种观察表明，天真的共享和特定专家架构难以收敛到理想状态。

图3 在CGC中的专家退化问题，其中两个共享专家几乎被task2和task7所垄断，分别以特定的风格工作。

专家欠拟合：在我们进一步解决了专家退化问题并增强了共享专家对所有任务的效率后，我们发现一些特定专家被分配了较小的门值，使得相应任务仅依赖于共享知识，较少使用特定参数。实际上，我们的模型需要同时预测数十个不同的任务，它们的密度（即正样本率）也差异很大，而密集任务可能是稀疏任务的100倍，例如点击与收藏。与可以从多个密集任务接收多个梯度更新的共享专家相比，特定专家（specific-experts）容易陷入欠拟合，进一步导致稀疏任务更多地依赖共享专家，而忽略其特定专家，使得特定参数被浪费。如图4所示，任务6的门控网络向共享专家分配了较大的值，但忽略了其特定专家。

图4 专家欠拟合问题，其中task1和task6几乎只依赖共享专家，而忽略了自己的特定专家，较少利用特定专家网络。

为了解决这些异常现象并提高MoE范式模型的稳定性，我们提出了一种简单、高效且平衡的神经网络架构用于多任务学习：多重门控专家层次模型，简称HoME。具体而言，我们从三个角度提供了有见地和深入的解决方案：公平专家权重的值分布对齐，重新组装任务的层次元专家结构，以及增强稀疏任务专家和深层多层MMoE训练的门控网络：

专家归一化和Swish机制：为了平衡专家输出的方差并避免专家崩溃，我们首先为每个专家引入了归一化操作，将其输出投影到近似正态分布，即专家输出分布≈N(0, I)。然而，在这种设置下，我们发现直接执行归一化也会导致在ReLU函数后存在太多的0。原因可能是归一化专家输出的均值接近0，因此一半的输出将小于0，然后在ReLU下激活为0。为了缓解零导数梯度现象，我们使用Swish函数替换ReLU函数，以提高参数的利用率并加速训练过程。由于归一化和Swish设置，所有专家的输出可以对齐到相似的数值幅度，这有助于我们的门控网络分配可比较的权重。

层次掩码机制

为了减少专家占用问题并避免专家退化（也称为任务冲突跷跷板问题），本文提出了一种简单而有效的级联层次掩码机制（Hierarchy mask）来缓解这种冲突。具体而言，我们插入了一个预排序元专家网络来将不同任务分组，以扩展标准化的MoE系统。如图1所示，我们的短视频行为任务可以根据其先验相关性手动分为两个元类别：

• （1）被动的观看时间任务，例如长观看；
• （2）主动的互动任务，例如评论。

因此，我们可以预先建模粗粒度的元类别专家，然后支持每个任务的以下想法：每个任务不仅应具有完全共享的全局专家，还应具有部分共享的类别内专家。

特征门控和自门控机制

为了增强稀疏任务专家的训练，我们提出了两种门控机制，以确保它们能够获得适当的梯度以最大化其有效性：特征门控（Feature-gate）和自门控（Self-gate）机制。考虑到同一层的专家总是共享相同的输入特征，但不同的专家将接收不同的梯度。因此，相同的特征输入可能会导致多个专家参数优化的梯度冲突的潜在风险。为此，我们首先提出了特征门控机制，以私有化灵活的专家输入，以保护稀疏任务专家的训练。此外，最新的MMoE研究表明，更深的堆叠专家网络可以带来更强大的预测能力。然而，在我们的实验中，我们发现原始门控网络容易逐层稀释梯度，这对稀疏任务专家的训练不友好。为了确保顶层梯度能够有效地传递到底层并稳定更深层的MMoE系统训练，我们进一步设计了自门控机制，以残差方式连接相邻的相关专家。

主要贡献

我们的主要贡献如下：

• 我们深入分析了当前MoE系统的专家问题，并提出了我们的里程碑工作HoME。据我们所知，本文是首个专注于增强多任务MoE系统稳定性的工作，这将为其他研究人员探索更稳健的多任务MoE系统提供启示。
• 我们在快手短视频服务中进行了广泛的离线和在线实验。离线实验显示，所有预测任务都得到了显著改进，在线实验在快手和快手极速版应用中分别获得了0.636%和0.735%的播放时间改进。
• 我们的HoME已在快手的各种服务中广泛部署，每天支持4亿活跃用户。

2 相关工作

在本节中，我们简要回顾多任务学习的演进轨迹。多任务学习在使模型能够在各种研究领域（包括推荐系统[[2, 9, 37]]、自然语言处理[[6, 8, 30]]、计算机视觉[[12, 18, 26]]和普适计算[[13, 22]]）中感知多种信号方面发挥着越来越重要的作用。

早期工作

在早期，一些研究利用硬专家共享架构，通过多个任务特定塔（由相同的专家输出驱动）来实现最简单的多任务学习系统。这些架构包括：

• 共享底层架构（Shared-Bottom）[[3]]
• 混合专家模型（Mixture-of-Expert, MoE）[[17]]

后续发展

随后，提出了交叉拼接网络（Cross-Stitch Network）[[25]]和闸门网络（Sluice Network）[[29]]，以构建深度专家信息融合网络，生成任务特定的输入，从而实现软专家知识共享。

除了复杂的纵向深度专家交叉外，横向专家权重估计是另一种定制任务特定塔输入的方法。近年来提出的多门混合专家模型（Multi-Gate Mixture-of-Expert, MMoE）[[23]]通过多门机制为不同的专家分配不同的权重，以平衡不同的任务。

MMoE变体及其应用

随着基于神经网络的推荐系统的兴起，MMoE的变体方法在提升模型能力和准确性方面也发挥了重要作用。

2 相关工作

在本节中，我们简要回顾了多任务学习的演进轨迹。多任务学习在赋能模型感知多种信号方面发挥着越来越重要的作用，广泛应用于推荐系统 [2, 9, 37]、自然语言处理 [6, 8, 30]、计算机视觉 [12, 18, 26] 和普适计算 [13, 22] 等领域。

在早期，一些工作利用硬专家共享架构，通过多个任务特定的塔（tower）接收相同的专家输出，构建了最简单的多任务学习系统，包括共享底层（shared-bottom）[3] 和专家混合（Mixture-of-Experts, MoE）[17]。随后，提出了cross-stitch network[25] 和sluice network[29]，通过构建深度专家信息融合网络生成任务特定的输入，以实现软专家知识共享。除了复杂的垂直深度专家交叉外，水平专家权重估计是另一种定制任务特定塔输入的方式。近年来提出的多重门控专家混合（Multi-gate Mixture-of-Experts, MMoE）[23] 提供了一种多重门控机制，为不同专家分配不同权重，以平衡不同任务。

随着基于神经网络的推荐系统浪潮的兴起，MMoE 的变体方法在提升模型能力和准确性方面也发挥了重要作用。开创性工作来自 YouTube 排序系统 [40]，该系统通过不同的门控网络利用多个共享专家来建模真实的用户-物品交互。为了缓解任务冲突跷跷板（task-conflict seesaw）问题 [5, 31, 33]，MMoE 的变体 CGC [33] 和 PLE [33] 不仅利用了共享专家，还引入了额外的特定专家，以实现更灵活的专家共享。基于共享/特定的思想，许多 MMoE 变体被提出，包括：

• MSSM [11]：通过采用字段级和单元级特征选择机制扩展了 PLE 方法，以自动确定输入特征的重要性。
• AdaTT [19]：在 PLE 基础上利用自适应融合门控机制，建模特定专家和共享专家之间的复杂任务关系。
• STAR [31]：采用星型拓扑结构，包含一个共享专家网络和多个特定网络，以融合专家参数。
• MoLA [43]：借鉴了大型语言模型（LLM）中的低秩微调技术，设计了轻量级的低秩特定专家适配器，以替代复杂的特定专家。

3 方法论

在本节中，我们将介绍我们的模型HoME的各个组成部分。我们首先回顾了在工业级推荐系统（RecSys）中，混合专家系统（MoE）是如何工作的，包括特征工程、MoE神经网络的细节以及用于排序的预测分数。随后，我们提出了针对三个问题的解决方案：专家归一化和swish机制用于克服专家崩溃问题，层次掩码机制用于缓解专家退化问题，以及两种门机制用于解决专家欠拟合问题。

3.1 预备知识：工业推荐系统中的多任务学习

工业推荐系统采用两阶段设计：（1）数百个候选物品的生成[35, 36]；（2）候选物品的排序[7, 24, 40]，以选出几十个最顶部的物品推荐给用户。由于这两个阶段的目标不同，因此使用的技术也完全不同：生成过程侧重于用户侧特征建模和粗略的物品采样，而排序过程则侧重于用户和物品特征的融合以及细粒度的用户多交互拟合。因此，多任务学习模型通常被用于排序过程中，以估计特定用户-物品对的各种交互概率。为简洁起见，模型生成的概率通常有一个简短的名称（$xtr$），例如点击概率为$ctr$，有效观看概率为$evtr$，点赞概率为$ltr$，评论概率为$cmtr$等。

3.1.1 标签与特征

形式上，这种排序学习过程通常被组织为多二元分类的形式，每个用户-物品样本包含两类信息——监督标签和输入特征：

• 监督信号：该用户-物品观看体验的真实标签，例如点击$y_{ctr} \in {0, 1}$，有效观看$y_{evtr} \in {0, 1}$，点赞$y_{ltr} \in {0, 1}$，评论$y_{cmtr} \in {0, 1}$以及其他标签。

• 特征输入：MoE的输入旨在从多个角度描述用户和物品的状态，大致可以分为四类：

  • （1）ID和类别特征，我们使用简单的查找操作符来获取它们的嵌入，例如用户ID、物品ID、标签ID、是否为活跃用户、是否关注作者、场景ID等；
  • （2）统计特征，需要设计分桶策略将其离散化并分配ID，例如过去一个月观看的短视频数量、过去一个月的短视频观看时长等；
  • （3）反映用户短期和长期兴趣的序列特征，通常通过一阶段或两阶段的注意力机制建模，例如DIN [42]、DIEN [41]、SIM [27]、TWIN [4]；
  • （4）预训练的多模态嵌入，例如文本嵌入[10]、asr嵌入[39]、视频嵌入[21]等。

将所有这些结合起来，我们可以获得多任务训练样本（例如，标签为${y_{ctr}, y_{evtr}, \dots}$，输入为$\mathbf{v} = [v_1, v_2, \dots, v_n]$），其中$n$表示特征的总数。

3.1.2 基于混合专家（MoE）的XTR预测

给定训练用户-物品样本的标签 $y_{ctr}, y_{evtr}, \dots$ 和特征 $\mathbf{v}$，接下来我们利用多任务模块进行预测。具体来说，我们展示了一种广泛使用的共享/特定范式MoE变体CGC的细节如下：

\[\begin{aligned} \widehat{y}^{ctr} &= \text{Tower}^{\ ctr}\left(\text{Sum}\left(\text{Gate}^{ctr}(\mathbf{v}), \{\text{Experts}^{\ \{shared,ctr\}}(\mathbf{v})\}\right)\right), \\ \widehat{y}^{evtr} &= \text{Tower}^{\ evtr}\left(\text{Sum}\left(\text{Gate}^{evtr}(\mathbf{v}), \{\text{Experts}^{\ \{shared,evtr\}}(\mathbf{v})\}\right)\right), \\ \widehat{y}^{ltr} &= \text{Tower}^{\ ltr}\left(\text{Sum}\left(\text{Gate}^{ltr}(\mathbf{v}), \{\text{Experts}^{\ \{shared,ltr\}}(\mathbf{v})\}\right)\right), \end{aligned}\]

其中：

\[\begin{aligned} \text{Tower}(\cdot) &= \text{Sigmoid}\left(\text{MLP}_T(\cdot)\right), \\ \text{Experts}(\cdot) &= \text{ReLU}\left(\text{MLP}_E(\cdot)\right), \\ \text{Gate}(\cdot) &= \text{Softmax}\left(\text{MLP}_G(\cdot)\right), \end{aligned}\]

…(1)

其中：

• $\text{Expert}^{\ shared}: \mathbb{R}^{\mid\mathbf{v}\mid} \rightarrow \mathbb{R}^D$ 和 $\text{Expert}^{\ xtr}: \mathbb{R}^{\mid\mathbf{v}\mid} \rightarrow \mathbb{R}^D$ 分别是ReLU激活的共享和特定专家网络，
• $\text{Gate}^{\ xtr}: \mathbb{R}^{\mid\mathbf{v}\mid} \rightarrow \mathbb{R}^N$ 是对应任务的Softmax激活的门网络，
• $N$ 是相关共享和特定专家的数量，
• $\text{Sum}$ 用于根据门生成的权重聚合 $N$ 个专家的输出，
• $\text{Tower}^{\ xtr}: \mathbb{R}^D \rightarrow \mathbb{R}$ 是Sigmoid激活的任务特定网络，用于衡量相应的交互概率 $\widehat{y}$。

在获得所有估计分数 $\widehat{y}^{ctr}, \dots$ 和真实标签 $y^{ctr}, \dots$ 后，我们直接最小化交叉熵二元分类损失来训练多任务学习模型：

\[L = -\sum_{xtr \in \{ctr, \dots\}} \left( y^{xtr} \log(\widehat{y}^{xtr}) + (1 - y^{xtr}) \log(1 - \widehat{y}^{xtr}) \right),\]

…(2)

在在线服务中，常见的操作是设计一个可控的复杂方程，将XTRs组合为一个排序分数：

\[\text{ranking_score} = \alpha \cdot \widehat{y}^{ctr} + \beta \cdot \widehat{y}^{evtr} + \gamma \cdot \widehat{y}^{cmtr} + \dots,\]

…(3)

其中：

• $\alpha, \beta, \gamma$ 是超参数。实际上，公式(3)在工业推荐系统中非常复杂，涉及许多策略。我们仅展示一个简单案例。在接下来的部分中，我们将重点改进公式(1)中的多任务学习过程的稳定性。

---

3.2 专家归一化与Swish机制

尽管公式(1)中的原始MMoE系统取得了显著的改进，但仍然存在严重的专家崩溃问题。将专家的 $MLP_E$ 函数生成的表示记为 ${ z^{shared}, z^{ctr}, z^{evtr}, \dots }$，我们发现它们的均值和方差值存在显著差异。受Transformer的启发，归一化操作是成功训练非常深度神经网络的关键技术之一。我们为每个专家引入了批归一化（Batch Normalization）[16]，以支持HoME生成可比较的输出 $z_{norm} \in R^D$：

\[z_{norm} = \text{Batch_Normalization}(z) = \gamma \frac{z - \mu}{\sqrt{\delta^2 + \epsilon}} + \beta,\]

其中：

\[\begin{aligned} \mu &= \text{Batch_Mean}(z), \\ \delta^2 &= \text{Batch_Mean}\left((z - \mu)^2\right), \end{aligned}\]

(4)

其中：

• $z$ 是任意专家的 $\text{MLP}_E$ 输出，
• $\gamma \in R^D$ 和 $\beta \in R^D$ 是可训练的缩放和偏置参数，用于调整分布，
• $\epsilon \in R^D$ 是一个非常小的因子，用于避免除以零错误
• $\mu \in R^D$ 和 $\delta^2 \in R^D$ 是当前批次相同专家输出的均值和方差。

经过专家归一化后，$z_{norm}$ 的分布接近于标准正态分布 $N(0, I)$。因此，$z_{norm}$ 的一半值将小于0，并在ReLU激活下变为0，导致它们的导数和梯度为0，阻碍模型收敛。因此，我们使用Swish函数替换公式(1)中的ReLU，得到HoME的专家结构：

\[\text{HoME_Expert}(\cdot) = \text{Swish}\left(\text{Batch_Normalization}\left(\text{MLP}_E(\cdot)\right)\right),\]

…(5)

其中：

• $\text{HoME_Expert}(\cdot)$ 是我们HoME中使用的最终结构。

在归一化和Swish的设置下，所有专家的输出可以对齐到相似的数值范围，这有助于门网络（gating network）分配可比较的权重。为简洁起见，在接下来的部分中，我们仍使用 $\text{Expert}(\cdot)$ 来表示 $\text{HoME_Expert}(\cdot)$。

3.3 层次掩码机制

针对专家退化问题，已有许多工作提出了新颖的特定专家和共享专家架构来缓解任务冲突。然而，遵循特定和共享范式，我们发现共享专家退化问题仍然存在。我们认为，考虑任务之间的先验相关性可能是有益的，如图1所示；我们的预测任务可以分为两类，例如主动交互任务（如点赞、评论等）和被动观看时间任务（如有效观看、长观看等）。在本节中，我们提出了一种简单而有效的级联层次掩码机制，以建模任务之间的先验归纳偏差。具体来说，我们插入了一个预排序的元专家网络来分组不同的任务，这里包括三种元任务知识来支持我们的两类任务：

\[\begin{aligned} z^{inter}_{meta} &= \text{Sum}\left(\text{Gate}^{\ inter}_{\ meta}(\mathbf{v}), \{\text{Experts}^{\ shared,inter}_{\ meta}(\mathbf{v})\}\right), \\ z^{watch}_{meta} &= \text{Sum}\left(\text{Gate}^{\ watch}_{\ meta}(\mathbf{v}), \{\text{Experts}^{\ shared,watch}_{\ meta}(\mathbf{v})\}\right), \\ z^{shared}_{meta} &= \text{Sum}\left(\text{Gate}^{\ shared}_{\ meta}(\mathbf{v}), \{\text{Experts}^{\ shared,inter,watch}_{\ meta}(\mathbf{v})\}\right), \end{aligned}\]

(6)

其中：

• $z^{inter}{meta}, z^{watch}{meta}, z^{shared}_{meta}$ 是粗粒度的宏观元表示，用于提取：
  • （1）交互类任务的知识
  • （2）观看时间类任务的知识
  • （3）共享知识

在获得这些元表示后，我们接下来根据其对应的元知识和共享元知识进行多任务预测。具体来说，我们利用元知识构建了三种类型的专家：

• （1）根据 $z^{shared}_{meta}$ 构建的全局共享专家
• （2）根据 $z^{inter}{meta}$ 或 $z^{watch}{meta}$ 构建的局部共享专家
• （3）根据 $z^{inter}{meta}$ 或 $z^{watch}{meta}$ 构建的每个任务的特定专家

对于任务特定的门网络，我们直接使用共享元知识 $z^{shared}_{meta}$ 和相应类别的元知识的拼接来生成专家的权重。这里，我们以点击和有效观看交互为例：

\[\begin{aligned} \widehat{y}^{ctr} &= Tower^{ctr}\left(Sum\left(Gate^{ctr}(z^{inter}_{meta} \oplus z^{shared}_{meta}), \\ \{Experts^{shared}(z^{shared}_{meta}), \\ Experts^{inter,ctr}(z^{inter}_{meta})\}\right)\right), \\ \widehat{y}_{evtr} &= Tower^{evtr}\left(Sum\left(Gate^{evtr}(z^{watch}_{meta} \oplus z^{shared}_{meta}), \\ \{Experts^{shared}(z^{shared}_{meta}), \\ Experts^{watch,evtr}(z^{watch}_{meta})\}\right)\right), \end{aligned}\]

(7)

其中：

• $\oplus$ 表示拼接操作符，
• $\text{Experts}^{\ shared}$ 是所有任务的共享专家，
• $\text{Experts}^{\ inter}$ 和 $\text{Experts}^{\ watch}$ 是类别内任务的共享专家。

值得注意的是，HoME的第一层元抽象与PLE的主要架构差异，这是基于我们在快手真实多任务推荐场景中的观察（见图5）。基于HoME的先验语义划分的元专家网络，我们可以尽可能避免任务之间的冲突，并最大化任务间的共享效率。

图5

---

3.4 特征门与自门机制

针对专家欠拟合问题，我们发现一些数据稀疏任务的门生成权重往往会忽略其特定专家，而为共享专家分配较大的权重。原因可能是我们的模型需要同时预测20多个不同的任务，而这些密集任务的密度可能是稀疏任务的100倍以上。为了增强稀疏任务专家的训练，我们提出了两种门机制，以确保它们能够获得适当的梯度以最大化其有效性：特征门和自门机制。

对于特征门，其目的是为不同任务专家生成不同的输入特征表示，以缓解所有专家共享相同输入特征时可能出现的梯度冲突。形式上，特征门旨在提取每个输入特征元素的重要性，例如 $\text{Fea_Gate}: \mathbb{R}^{\mid\mathbf{v}\mid} \rightarrow \mathbb{R}^{\mid\mathbf{v}\mid}$，如果输入是 $\mathbf{v}$。然而，在工业推荐系统中，$\mathbf{v}$ 通常是一个高维向量，例如 $\mid\mathbf{v}\mid > 3000+$；因此，为元专家引入这些大矩阵是昂贵的。受LLM效率调优技术LoRA [15] 的启发，我们引入了两个小矩阵来近似生成元素重要性的大矩阵：

\[\text{Fea\_LoRA}(\mathbf{v}, d) = 2 \times \text{Sigmoid}\left(\mathbf{v}(BA)\right),\]

其中:

• $B \in \mathbb{R}^{\mid\mathbf{v}\mid \times d}$，$A \in \mathbb{R}^{d \times \mid\mathbf{v}\mid}$，$BA \in \mathbb{R}^{\mid\mathbf{v}\mid \times \mid\mathbf{v}\mid}$。

(8)

注意，我们在Sigmoid函数后应用了一个2×操作符，旨在实现灵活的放大或缩小操作。实际上，$\text{Fea_LoRA}$ 函数是生成私有化专家输入的有效方法。在我们的迭代中，我们发现它可以进一步通过多任务思想增强，即引入更多的 $\text{Fea_LoRA}$ 从多个方面生成特征重要性作为我们的 $\text{Fea_Gate}$：

\[\text{Fea\_Gate}(\mathbf{v}) = \text{Sum}\left(\text{Gate}_{fea}(\mathbf{v}), \{\text{Fea\_LoRA}_{\{1,2,\dots,L\}}(\mathbf{v}, \mid\mathbf{v}\mid/L)\}\right),\]

(9)

其中 $L$ 是一个超参数，用于控制 $\text{Fea_LoRA}$ 的数量，$\text{Gate}{fea}: \mathbb{R}^{\mid\mathbf{v}\mid} \rightarrow \mathbb{R}^L$ 用于生成权重以平衡不同 $\text{Fea_LoRA}$ 的重要性。注意，我们需要选择一个能被输入长度 $\mid\mathbf{v}\mid$ 整除的 $L$ 来生成 $\text{Fea_LoRA}$ 的维度。因此，我们的专家输入可以如下获得（这里我们展示了第一层元共享专家的输入 $\mathbf{v}^{shared}{meta}$）：

\[\mathbf{v}^{shared}_{meta} = \mathbf{v} \odot \text{Fea\_Gate}^{shared}_{meta}(\mathbf{v}),\]

(10)

其中 $\odot$ 表示逐元素乘积。通过这种方式，不同的专家拥有自己的特征空间，这可以减少梯度冲突的风险，从而保护稀疏任务。

此外，最新的MoE研究表明，更深层次的专家网络堆叠可以带来更强大的预测能力 [19, 33]。然而，在我们的实验中，我们发现原始的门网络容易逐层稀释梯度，尤其是对于稀疏任务专家的训练。除了专家输入级别的 $\text{Fea_Gate}$，我们还在专家输出级别添加了基于残差思想的自门机制，以确保顶层的梯度能够有效地传递到底层。具体来说，$\text{Self_Gate}$ 只关注其特定专家的输出。以观看时间元专家的输出为例：

\[z^{shared}_{meta,self} = \text{Sum}\left(\text{Self\_Gate}^{shared}_{meta}(\mathbf{v}), \{\text{Experts}^{shared}(\mathbf{v})\}\right),\]

其中：

\[\text{Self\_Gate}(\cdot) = \begin{cases} \text{Sigmoid}\left(\text{MLP}_G(\cdot)\right) & \text{如果只有1个专家}, \\ \text{Softmax}\left(\text{MLP}_G(\cdot)\right) & \text{其他情况}. \end{cases}\]

(11)

其中 $\text{Self_Gate}: \mathbb{R}^{\mid\mathbf{v}\mid} \rightarrow \mathbb{R}^K$，$K$ 是相关专家的数量，其激活函数为Sigmoid（如果只有1个专家），否则设置为Softmax。类似地，$z^{inter}{meta,self}$ 和 $z^{watch}{meta,self}$ 可以通过相同的方式获得，然后我们将相应的表示（例如 $z^{inter}{meta} + z^{inter}{meta,self}$）添加到下一层的支持中。详见第5节以获取HoME的细粒度细节。

4 实验

在本节中，我们首先在离线环境中将HoME与几种广泛使用的多任务学习方法进行比较。然后，我们通过修改模型进行了一些变体实验，以验证HoME的有效性。我们还测试了HoME超参数对专家数量和特征门LoRA数量的鲁棒性影响。此外，我们提供了模型专家网络门权重的图表，以展示HoME作为一个平衡系统的潜力。最后，我们将HoME推送到在线A/B测试中，以验证HoME能为快手带来多大的收益。

---

4.1 实验设置

我们在快手的短视频数据流场景中进行实验，这是快手最大的推荐场景，每天包括超过4亿用户和500亿条日志。为了公平比较，我们仅更改公式(1)中的多任务学习模块，并保持其他模块不变。具体来说，我们实现了MMoE [23]、CGC [33]、PLE [33]、AdaTT [19] 等模型变体作为基线。对于评估，我们使用广泛使用的排序指标AUC和GAUC [42] 来反映模型的预测能力。具体来说，在我们的短视频服务中，GAUC是最重要的离线指标。其主要思想是分别计算每个用户的AUC，然后加权聚合所有用户的AUC：

\[\text{GAUC} = \sum_{u} w_u \text{AUC}_u \quad \text{其中} \quad w_u = \frac{\#\text{logs}_u}{\sum_{i} \#\text{logs}_i},\]

(12)

其中 $w_u$ 表示用户日志的比例。

---

4.2 离线实验

主要实验结果如表1所示。注意，在离线评估中，AUC或GAUC的0.03%~0.05%提升足以为我们的业务带来显著的在线收益。我们首先展示了HoME_Expert在MMoE上的有效性，即MMoE。然后，我们将HoME与所有改进的基线进行比较，这些基线都配备了HoME_Expert，例如“CGC w/o shared”，这是CGC的变体，忽略了共享专家和所有门网络。此外，我们还为HoME实现了消融变体：“w/o fg2”和“w/o fg”分别忽略了第二层特征门和所有特征门；“w/o sg”忽略了所有自门；“w/o mask”保留了HoME架构，但所有专家都是共享的。我们有以下观察：

1. MMoE*显著优于原始MMoE方法，这表明我们的专家归一化和Swish机制可以克服专家崩溃问题，平衡专家输出，并鼓励专家网络承担应有的责任。
2. “CGC* w/o shared”可以看作是每个任务具有特定专家的Shared-bottom。MMoE比这种简单的“CGC w/o shared”解决方案弱（在我们的实验中参数增加了24%），这表明MMoE系统在真实的大规模数据流场景中容易退化。
3. *与“CGC w/o shared”相比，CGC没有显著改进**，这表明CGC的共享专家正在退化为某些特定专家。
4. 与MMoE相比，PLE和AdaTT*表现更好，这表明在解决专家崩溃问题后，堆叠多个专家网络层并增加模型参数是释放多任务模块潜力的一种有前途的方法。
5. HoME在所有任务中均显示出统计上的改进，同时引入更少的参数并取得了最佳结果，这表明我们的修改可以增强多任务MoE系统的稳定性并最大化专家效率。

对于HoME的消融实验：

1. “w/o fg-sg-mask”变体与MMoE*表现相当，而“w/o fg-sg”变体在所有任务中均显示出显著改进（例如，AUC提升+0.15%），这表明我们的层次掩码机制是一种强大且低资源的策略，可以在不引入大量额外参数的情况下缓解专家退化问题。
2. “w/o fg”变体比“w/o fg-sg”变体表现更好且更稳定，这表明在不同层专家之间添加残差连接有助于训练专家。
3. 与HoME和“w/o fg2”变体相比，我们发现第二层特征门可以增强模型能力，但第一层特征门显示出更稳健和更大的改进。原因可能是第一层是输入信息源，并在粗粒度元层中使用，其任务梯度冲突问题比第二细粒度层更严重。

---

4.3 超参数敏感性讨论

本节探讨了专家数量和特征门LoRA数量的超参数敏感性，以研究HoME的鲁棒性。对于专家数量，我们在“HoME w/o fg”变体下进行实验，因为第一层特征门是一个昂贵的参数消耗操作。从表2中，我们可以观察到HoME的扩展规律现象：仅通过引入更多专家，预测精度会随着参数数量的增加而稳步提高。这种现象也表明我们的HoME是一个平衡的MoE系统，可以释放所有专家的能力。

对于特征LoRA数量，我们在“HoME w/o fg2”变体下进行实验，该变体仅涉及第一层特征门，同时在表1中显示出显著改进。具体来说，在我们的实现中，更多的LoRA数量只会减少隐藏维度，而不会增加额外的参数，这可能会降低单个LoRA的能力。从表3中，我们可以观察到两个LoRA的变体显示出最佳结果，这表明在平衡LoRA数量和LoRA建模能力以提供更多增量信息方面存在瓶颈。

---

4.4 HoME状态讨论

图6展示了HoME的专家输出分布和门权重流。从中我们可以观察到，HoME实现了平衡的门权重均衡状态：

1. 根据特征门的热图（随机可视化的64个维度），我们可以得出结论，我们的特征门可以为每个专家实现灵活的元素级特征选择。
2. 所有共享和特定专家的输出在数值上对齐。此外，我们可以发现元共享专家的分布与特定专家的分布不同，这表明共享知识倾向于由元网络编码，而差异知识则被推送到特定专家编码。
3. 所有专家都发挥了预期的作用；共享和特定专家贡献了可感知的权重。

图6

---

4.5 在线A/B测试

在本节中，我们还将HoME推送到在线排序模型中，服务于三个短视频场景：快手单页/双页（如图1所示）和快手极速版单页。在我们的服务中，主要指标是观看时间指标，例如（平均）播放时间，它反映了用户在快手上花费的总时间，我们还展示了视频观看指标，它衡量了用户观看的短视频总量。年轻用户和总用户组的在线A/B测试结果如表4所示。实际上，播放时间约0.1%的提升在快手上是一个统计上显著的改进。我们提出的HoME在所有用户的三个场景中分别实现了+0.311%、+0.474%和+0.169%的显著改进，这是过去一年中最显著的改进。此外，我们可以观察到HoME在所有交互指标（如点击、点赞、评论等）上均实现了显著的业务收益，这表明HoME可以将多任务系统收敛到一个更平衡的均衡状态，而不会出现跷跷板现象。此外，我们发现稀疏行为任务的提升更大，这表明我们的HoME使所有共享或特定专家能够获得适当的梯度以最大化其有效性。

---

5 结论

在本文中，我们专注于解决多任务学习方法在快手短视频服务中的实际问题以及我们从中学到的经验教训，这是全球最大的推荐场景之一。我们首先发现，现有的广泛使用的多任务家族（即门控混合专家）容易受到几个严重问题的限制，这些问题限制了模型的预期能力。从专家输出中，我们发现了专家崩溃问题，即专家的输出分布显著不同。从共享专家的学习中，我们观察到专家退化问题，即一些共享专家仅服务于一个任务。从特定专家的学习中，我们注意到专家欠拟合问题，即一些稀疏任务的特定专家几乎不贡献任何信息。

为了克服这些问题，我们提出了三个有洞察力的改进：

1. 专家归一化和Swish机制，以对齐专家输出分布；
2. 层次掩码机制，以规范化任务之间的关系，最大化共享专家的效率；
3. 特征门和自门机制，以私有化更灵活的专家输入并连接相邻的相关专家，确保所有专家都能获得适当的梯度。

此外，通过在快手这一全球最大的短视频平台上进行的大量离线和在线实验，我们展示了HoME相比其他广泛使用的多任务方法带来的显著改进。我们的HoME已广泛部署在快手的各种在线模型中，每天为数亿活跃用户提供服务。

附录：

• https://www.arxiv.org/pdf/2408.05430

google在《Titans: Learning to Memorize at Test Time》提出了区别于Transformer的的一种新架构:Titans。我们来看一下它的实现，是否有前景：

摘要

在过去的十多年里，关于如何有效利用循环模型（recurrent models）和注意力机制（attentions）的研究已经非常广泛。循环模型的目标是将数据压缩到一个固定大小的内存中（称为hidden state），而注意力机制则允许模型关注整个上下文窗口，捕捉所有token之间的直接依赖关系。然而，这种更精确的依赖关系建模带来了二次方的计算成本(quadratic cost)，限制了模型只能处理固定长度的上下文。我们提出了一种新的神经长期记忆模块（neural long-term memory module），该模块能够学习记忆历史上下文，并帮助注意力机制在利用过去信息的同时关注当前上下文。我们展示了这种神经记忆模块具有快速并行化训练的优势，同时保持了快速的推理能力。从记忆的角度来看，我们认为注意力机制由于其有限的上下文但精确的依赖关系建模，起到了短期记忆的作用；而神经记忆（neural memory）由于其能够记忆数据的能力，起到了长期、更持久的记忆作用。基于这两个模块，我们引入了一系列新的架构，称为Titans，并提出了三种变体，以探讨如何有效地将记忆融入这一架构中。我们在语言建模、常识推理、基因组学和时间序列任务上的实验结果表明，Titans比Transformers和最近的现代线性循环模型更有效。此外，与基线模型相比，Titans能够有效地扩展到超过200万的上下文窗口大小，并在“大海捞针”任务中表现出更高的准确性。

1.介绍

Transformers，一种纯基于注意力机制的架构（Vaswani 等人，2017），已被牢牢确立为序列建模中的最先进模型，主要归功于其上下文学习能力和大规模学习能力（Kaplan 等人，2020）。Transformers 的核心构建模块——注意力模块——充当关联记忆模块（Bietti 等人，2024），它们学习存储key-value关联性（associations），并通过计算query（即搜索信号）和key（即上下文）之间的成对相似性来检索这些关联。因此，从设计上看，Transformer 的输出完全取决于当前上下文窗口中token的直接依赖关系。然而，这种精确的依赖关系建模带来了与上下文长度相关的二次方时间和内存复杂度。在复杂的现实任务中（例如语言建模（N. F. Liu 等人，2024）、视频理解（C.-Y. Wu 等人，2019）、长期时间序列预测（H. Zhou 等人，2021）），上下文窗口可能变得非常大，这使得 Transformers 在这些下游任务中的适用性面临挑战。

为了克服 Transformers 的可扩展性问题，最近的研究旨在设计不同变体的线性 Transformers（Kacham、Mirrokni 和 P. Zhong，2024；Katharopoulos 等人，2020；S. Yang、B. Wang、Shen 等人，2024），其中注意力机制中的 softmax 被核函数取代（详见 §2.1），从而显著降低了内存消耗。尽管线性 Transformers 具有高效性并能够扩展到更长的上下文，但与 Transformers 相比，它们的性能并不具有竞争力，因为核技巧使模型变成了线性循环网络，其中数据被压缩为矩阵值状态（Katharopoulos 等人，2020）。然而，这带来了关于线性循环（或线性 Transformers）模型的一个矛盾事实：一方面，我们使用这些线性模型来增强可扩展性和效率（线性与二次方复杂度），其优势在非常长的上下文中显现；另一方面，非常长的上下文无法被适当地压缩到一个小的向量值或矩阵值状态中（S. Wang，2024）。

此外，除了效率问题外，大多数现有架构——从 Hopfield 网络（Hopfield，1982）到 LSTM（Jürgen Schmidhuber 和 Hochreiter，1997）以及 Transformers（Vaswani 等人，2017）——在处理泛化、长度外推和/或推理（Anil 等人，2022；Qin、Y. Zhong 和 Deng，2024）时都面临挑战，而这些是许多复杂现实任务中不可分割的部分。尽管这些架构从人类大脑中汲取了灵感，但它们都缺少以下关键部分：

• （1）学习过程中的关键组件——例如短期记忆、长期记忆、元记忆、关注当前上下文等（Cowan，2008）；
• （2）这些组件如何作为可以独立运行的互联系统；以及/或
• （3）从数据中主动学习并记忆过去历史的抽象能力。

我们认为，在一个有效的学习范式中，类似于人类大脑，存在独立但相互关联的模块，每个模块都负责学习过程中至关重要的组件。

记忆视角

记忆（memory）是一种基本的心理过程，也是人类学习中不可分割的组成部分（Terry，2017）。如果没有一个正常运作的记忆系统，人类和动物将只能局限于基本的反射和刻板行为。因此，记忆一直是机器学习文献中许多开创性研究的灵感来源；例如，Hopfield 网络（Hopfield，1982）、LSTM（Jürgen Schmidhuber 和 Hochreiter，1997）以及 Transformers（Vaswani 等人，2017）。

从神经心理学文献中对记忆和学习的常见定义中汲取灵感（Okano、Hirano 和 Balaban，2000），大多数现有架构将记忆视为由输入引起的神经更新，并将学习定义为在给定目标的情况下获取有效且有用记忆的过程。从这个角度来看，循环神经网络（RNN）（Williams 和 Zipser，1989）可以被定义为具有向量值记忆模块 M（也称为hidden state）的模型，其主要步骤包括：

在时间 𝑡 给定新输入 $𝑥_𝑡$ 时，模型

• （1）使用函数 $𝑓(M_{𝑡−1}, 𝑥_𝑡) $ 更新记忆（带有压缩）；
• （2）使用函数 $𝑔(M_𝑡, 𝑥_𝑡)$ 检索输入的相应记忆（详见 §2.1）。

类似地，Transformers 可以被视为具有增长记忆和两个相似步骤的架构。即，key和value矩阵对充当模型的记忆，模型：

• （1）通过将key和value附加到记忆中来更新记忆（无压缩），
• （2）通过查找query向量与key向量的相似性来检索query向量的相应记忆，然后将其用于加权value向量以生成输出。

这种视角可以帮助我们更好地理解现有范式、它们的关键差异，并设计更有效的架构。例如，Transformers（Vaswani 等人，2017）和线性 Transformers（Katharopoulos 等人，2020）之间的主要区别在于记忆结构以及记忆更新步骤，其中：线性 Transformers 将历史数据压缩为固定大小的矩阵值记忆，而 Transformers 则保留所有历史数据（在上下文长度内）而不进行任何压缩。虽然线性 Transformers 和线性 RNN（包括状态空间模型）都在记忆更新步骤中压缩信息，但关键区别在于记忆的结构，其中线性 RNN（相对于线性 Transformers）使用向量值记忆（相对于矩阵值记忆）。因此，这种视角促使我们提出以下问题：

• （Q1）什么是良好的记忆结构？
• （Q2）什么是适当的记忆更新机制？
• （Q3）什么是良好的记忆检索过程？

重新审视我们对人类记忆的理解，它既不是一个单一的过程，也不服务于单一的功能（Cowan，2008）。事实上，记忆是一个系统的联合体——例如短期记忆、工作记忆（working memory）和长期记忆——每个系统服务于不同的功能，具有不同的神经结构，并且每个系统都能够独立运行（Willingham，1997）。这一事实促使我们提出：

• （Q4）如何设计一个包含不同互联记忆模块的高效架构。

最后，存储记忆是一个神经过程，需要对过去的抽象进行编码和存储。假设一个单一向量或矩阵（其参数以线性方式编码数据）足以存储长期历史可能过于简化。

• （Q5）是否需要深度记忆模块来有效存储/记住遥远的过去？

贡献与路线图

在本文中，我们旨在通过设计一个长期神经记忆模块来回答上述五个问题，该模块能够在测试时高效且有效地学习记忆。基于其设计，我们讨论了如何将其融入架构中。

神经记忆（§3）。我们提出了一种（深度）神经长期记忆模块，它（作为元上下文模型）学习如何在测试时将数据记忆/存储到其参数中。受人类长期记忆系统（Mandler，2014）的启发，

我们设计了这个记忆模块，使得违反预期的事件（即令人惊讶的事件： surprising）更容易被记住。为此，我们通过神经网络在关联记忆损失中对输入的梯度来衡量输入的“惊讶度（surprise）”（详见 §3.1）。为了更好地处理有限的内存，我们提出了一种衰减机制，该机制考虑了内存大小与数据惊讶度的比例，从而实现更好的内存管理。我们展示了这种衰减机制实际上是现代循环模型中遗忘机制的泛化（Dao 和 Gu，2024；Gu 和 Dao，2024；S. Yang、Kautz 和 Hatamizadeh，2024）。有趣的是，我们发现这种机制等同于使用小批量梯度下降、动量和权重衰减来优化元神经网络。基于张量化小批量梯度下降以使用更多矩阵乘法操作（Yu Sun 等人，2024），我们提出了一种快速且可并行化的算法来训练我们的深度神经长期记忆模块。

Titans 架构（§4）

在设计完长期神经记忆模块后，一个重要的问题是：如何高效且有效地将记忆融入深度学习架构中。我们提出了 Titans，这是一个由三个超头部（hyper-heads）组成的深度模型家族：

• （1）核心模块：该模块包含短期记忆，负责数据处理的主要流程（我们使用有限窗口大小的注意力机制）；
• （2）长期记忆模块：这一分支是我们的神经长期记忆模块，负责存储/记住遥远的过去；
• （3）持久记忆模块：这是一组可学习但与数据无关的参数，用于编码任务相关知识。

最后，作为概念验证，我们提出了 Titans 的三种变体，其中我们将记忆分别融入为：

• （i）一个上下文（context）
• （ii）层（layer）
• （iii）一个门控分支（gated branch）

实验结果（§5）

我们在语言建模、常识推理、记忆密集型任务、“大海捞针”任务、时间序列预测和 DNA 建模任务上进行了实验评估。我们观察到，Titans 架构在所有现代循环模型及其混合变体（结合滑动窗口注意力机制）的综合基准测试中均表现优异。此外，Titans 在相同上下文窗口下优于 Transformers，并在使用整个上下文的 Transformers 中表现出竞争力。这些结果是在 Titans 能够扩展到超过 200 万上下文窗口大小的情况下实现的，而 Transformers 则无法做到这一点。

2 预备知识

在本节中，我们将讨论本文中使用的符号和一些背景概念。我们令：

• $ x \in \mathbb{R}^{N \times d_{\text{in}}} $ 表示输入
• $ \mathbf{M} $ 表示神经网络（神经记忆模块：neural memory）
• $ \mathbf{Q} $、$ \mathbf{K} $、$ \mathbf{V} $ 分别表示注意力机制中的query、key和value
• $ M $ 表示注意力掩码（attention mask）
• $ S^{(i)} $ 表示：在对序列进行分段时，使用第 $ i $ 段。

在本文中，我们简化符号并使用下标来指代矩阵、向量或段中的特定元素。例如，我们令：

• $ S^{(i)}_j $ 表示第 $ i $ 段中的第 $ j $ 个 token。

唯一的例外是下标为 $ t $ 的情况，我们保留它来表示时间上的递归或神经网络在时间 $ t $ 的状态。

给定：神经网络 $ \mathbf{N} $ 和数据样本 $ x $，我们使用：

• $ \mathbf{N}(x) $（或 $ \mathbf{N}^*(x) $）表示带权重调整（或不带权重调整）的前向传播

此外，我们简化符号并使用：

• $ \mathbf{N}^{(k)} $ 表示神经网络的第 $ k $ 层

接下来，我们首先讨论注意力机制及其高效变体的背景，然后回顾现代线性 RNN，最后讨论这些架构的记忆视角，这促使我们设计了 Titans。

2.1 背景

注意力机制。Transformers（Vaswani 等人，2017）作为许多深度学习模型的实际骨干，基于注意力机制。给定：

• 输入 $ x \in \mathbb{R}^{N \times d_{\text{in}}} $

因果注意力机制基于输入依赖的key、value和query矩阵计算输出 $ y \in \mathbb{R}^{N \times d_{\text{in}}} $：

\[\mathbf{Q} = x \mathbf{W}_Q, \quad \mathbf{K} = x \mathbf{W}_K, \quad \mathbf{V} = x \mathbf{W}_V, \quad (1)\] \[y_i = \frac{\sum_{j=1}^i \exp\left(\frac{\mathbf{Q}_i^\top \mathbf{K}_j}{\sqrt{d_{\text{in}}}}\right) \mathbf{V}_j}{\sum_{\ell=1}^i \exp\left(\frac{\mathbf{Q}_i^\top \mathbf{K}_\ell}{\sqrt{d_{\text{in}}}}\right)}, \quad (2)\]

其中：

• $ W_Q, W_K, W_V \in R^{d_{in} \times d_{in}} $ 是可学习参数

尽管 Transformers 在召回能力上表现出色，但它们至少需要 $ N \times d $ 次操作来计算输出，导致内存消耗较大且对较长序列的吞吐量较低。

高效注意力机制。为了提高软注意力机制在长序列上的内存消耗和吞吐量，许多研究集中在注意力机制的 I/O 感知实现（Dao 2024；Dao, D. Fu 等人，2022），通过稀疏化注意力矩阵（B. Chen 等人，2021；Choromanski 等人，2021；Dai 等人，2019）、近似 softmax（Arora 等人，2024）或开发基于核的（线性）注意力机制（Aksenov 等人，2024；Kacham, Mirrokni 和 P. Zhong，2024；Schlag, Irie 和 Jürgen Schmidhuber，2021；S. Yang, B. Wang, Shen 等人，2024）来设计更高效的注意力机制。在本部分，我们重点关注后者，即线性注意力机制，其中标准注意力中的 softmax 被替换为替代核函数 $ \phi(\cdot, \cdot) $，使得 $ \phi(x, y) = \phi(x) \phi(y) $。因此，注意力可以写成：

\[y_i = \frac{\sum_{j=1}^i \phi(\mathbf{Q}_i^\top \mathbf{K}_j) \mathbf{V}_j}{\sum_{\ell=1}^i \phi(\mathbf{Q}_i^\top \mathbf{K}_\ell)} = \frac{\phi(\mathbf{Q}_i)^\top \sum_{j=1}^i \phi(\mathbf{K}_j) \mathbf{V}_j}{\phi(\mathbf{Q}_i)^\top \sum_{\ell=1}^i \phi(\mathbf{K}_\ell)}, \quad (3)\]

由于：

• 项 $ \sum_{j=1}^i \phi(K_j), \sum_{\ell=1}^i \phi(K_\ell) $ 在每一步中重复使用，因此吞吐量更高。

当选择核函数为单位矩阵时（Yutao Sun 等人，2023），上述公式可以写成递归形式：

\[\mathbf{M}_t = \mathbf{M}_{t-1} + \mathbf{K}_t^\top \mathbf{V}_t, \quad (4)\] \[y_t = \mathbf{Q}_t \mathbf{M}_t, \quad (5)\]

这使得线性注意力机制能够高效推理。

现代线性模型及其记忆视角。如前所述，可以将学习定义为获取有效且有用记忆的过程。基于此，可以将循环神经网络（RNN）的hidden state视为记忆单元，模型旨在将信息压缩到其中。因此，在一般形式的循环神经网络中，hidden state可以被视为记忆单元，递归过程可以分为记忆单元的读和写操作。即，令：

• $ x \in \mathbb{R}^{N \times d_{\text{in}}} $ 为输入
• $ \mathbf{M} \in \mathbb{R}^d $ 为记忆单元
• $ y \in \mathbb{R}^{d_{\text{in}}} $ 为输出

则循环神经网络的一般形式定义为：

\[\mathbf{M}_t = f(\mathbf{M}_{t-1}, x_t), \quad \text{写操作} \quad (6)\] \[y_t = g(\mathbf{M}_t, x_t), \quad \text{读操作} \quad (7)\]

其中：

• $ f(\cdot, \cdot) $ 是读操作，
• $ g(\cdot, \cdot) $ 是写操作。

注意，这里的 $ \mathbf{M}_t $ 下标表示记忆在时间 $ t $ 的状态。

从这一视角来看，线性 Transformers 的递归公式（见公式 4）等同于将键和值 $ (\mathbf{K}_t, \mathbf{V}_t) $ 加性地压缩并写入矩阵值记忆单元 $ \mathbf{M}_t $ 中。因此，在处理长上下文数据时，这种加性特性会导致内存溢出，显著损害模型性能。为了解决这一问题，研究集中在两个有前景的方向上：

• （1）添加遗忘机制：一些研究提出了线性模型的自适应（数据依赖）遗忘门机制，可以在需要时擦除记忆。例如，GLA（S. Yang, B. Wang, Shen 等人，2024）、LRU（Orvieto 等人，2023）、Griffin（De 等人，2024）、xLSTM（Beck 等人，2024）和 Mamba2（Dao 和 Gu，2024）等模型，后者还与离散化的传统状态空间模型（Gu 和 Dao，2024）相关联。
• （2）改进写操作：为了克服传统循环模型中记忆写操作的加性特性，Widrow 和 Hoff（1988）提出了 Delta 规则，在添加记忆（即键值对）之前，模型首先移除其过去的值。为了增强可并行化训练和扩展性，S. Yang, B. Wang, Yu Zhang 等人（2024）提出了一种快速并行化算法。最后，最近 S. Yang, Kautz 和 Hatamizadeh（2024）通过添加遗忘门改进了 DeltaNets。

记忆模块。记忆一直是神经网络设计的核心部分之一（Graves, Wayne 和 Danihelka，2014；JH Schmidhuber，1992；Jürgen Schmidhuber 和 Hochreiter，1997；J. Zhang 等人，2024）。将线性层视为键值（关联）记忆系统（key-value (associative) memory system）的思想可以追溯到快速权重程序（fast weight programs），其中动态快速程序被纳入循环神经网络中作为可写记忆（JH Schmidhuber，1992）。Hebbian（Hebb，2005）和 delta（Prados 和 Kak，1989）学习规则是快速权重程序中最流行的学习规则，已在各种研究中广泛探索（Irie, Schlag 等人，2021；Munkhdalai, Sordoni 等人，2019；Munkhdalai 和 H. Yu，2017；Schlag, Irie 和 Jürgen Schmidhuber，2021；JH Schmidhuber，1992；S. Yang, Kautz 和 Hatamizadeh，2024；S. Yang, B. Wang, Yu Zhang 等人，2024）。然而，所有这些模型都基于瞬时惊讶度，忽略了序列中的 token 流（见第 3.1 节），并且大多数模型缺乏遗忘门，导致内存管理不佳。

我们在附录 C 中进一步讨论了我们的架构与最近模型的联系。其他相关工作在附录 A 中讨论。

3 测试时的记忆学习

为了克服长期记忆的不足，并使模型能够学习、遗忘和检索信息，本节提出了一种神经长期记忆模块，这是一种在测试时学习记忆的元模型。

• 在 3.1 节中，我们首先讨论神经记忆的动机和设计。
• 在 3.2 节中，我们讨论如何通过快速且可并行化的训练使我们的架构设计受益。
• 在 3.3 节中，我们通过持久记忆模块增强我们的架构，其中使用可学习但与数据无关的参数来学习任务的元信息。

图1 关于如何并行训练神经记忆并使用矩阵乘法（matmuls）的图示说明。

3.1 长期记忆

为了设计一个神经长期记忆模块，我们需要一个能够将过去历史的抽象编码到其参数中的模型。一个例子是大语言模型（LLMs），它们被证明能够记忆训练数据（Leybzon 和 Kervadec，2024；Schwarzschild 等人，2024；Staab 等人，2024）。因此，一个简单的想法是：训练一个神经网络并期望它记忆其训练数据。然而，记忆化几乎总是被认为是神经网络中的不良现象，因为它限制了模型的泛化能力（Bayat 等人，2024），引发隐私问题（Staab 等人，2024），并导致测试时性能不佳。此外，训练数据的记忆化在测试时可能没有帮助，因为数据可能是分布外的。我们认为，我们需要一个在测试时学习如何记忆/遗忘数据的在线元模型。在这种设置中，模型学习的是一个能够记忆的函数，但它不会过度拟合训练数据，从而在测试时实现更好的泛化。

学习过程与惊讶度度量(Surprise Metric)。训练长期记忆的关键思想是：将其训练视为一个在线学习问题，我们的目标是将过去的信息 $ x_1, \ldots, x_{t-1} $ 压缩到长期神经记忆模块 $ \mathbf{M}_t $ 的参数中。如前所述，违反预期的事件（即令人惊讶的事件）对人类来说更容易被记住（Mandler，2014）。受此启发，模型惊讶度的一个简单定义可以是：其相对于输入的梯度。梯度越大，输入数据与过去数据的差异越大。因此，使用这种惊讶度评分，我们可以更新记忆如下：

\[\mathbf{M}_t = \mathbf{M}_{t-1} - \theta_t \nabla \ell(\mathbf{M}_{t-1}; x_t) \quad \text{（惊讶度）} \quad (8)\]

然而，这种惊讶度度量可能会导致错过在重大惊讶时刻之后的重要信息。也就是说，梯度在几次惊讶步骤后可能变得非常小，导致陷入平坦区域（即局部最小值），并错过序列某些部分的信息。从人类记忆的角度来看，一个事件可能不会在长时间内持续让我们感到惊讶，尽管它是值得记忆的。原因是初始时刻足够令人惊讶，足以在长时间内吸引我们的注意力，从而记住整个时间段。为了改进上述惊讶度度量（公式 8），我们将惊讶度度量分为：

• （1）过往惊讶度（past surprise），衡量最近过去的惊讶程度；
• （2）瞬时惊讶度（momentary surprise），衡量输入数据的惊讶程度：

\[\mathbf{M}_t = \mathbf{M}_{t-1} + S_t, \quad (9)\] \[S_t = \eta_t S_{t-1} \quad \text{（过去惊讶度）} - \theta_t \nabla \ell(\mathbf{M}_{t-1}; x_t) \quad \text{（瞬时惊讶度）} \quad (10)\]

有趣的是，这个公式类似于带有动量的梯度下降，其中：

• $ S_t $ 是动量项

因此，这里的动量充当了跨时间（序列长度）的惊讶度记忆。在这个公式中：

• 项 $ \eta_t $ 是一个数据依赖的惊讶度衰减（$ x_t $ 的函数），控制惊讶度随时间衰减的程度
• 项 $ \theta_t $ 则控制瞬时惊讶度应以数据依赖的方式纳入最终惊讶度度量的多少

这种数据依赖性在这个设计中尤为重要：虽然前一个 token 的惊讶度可能需要影响下一个 token 的惊讶度，但这只有在所有 token 都相关且处于同一上下文中时才有效。因此，数据依赖的 $ \eta $ 可以控制记忆是否需要：

• （1）通过设置 $ \eta_t \to 0 $ 忽略上一次的惊讶度（可能由于上下文的变化），
• （2）通过设置 $ \eta_t \to 1 $ 完全纳入上一次的惊讶度（可能因为 token 与其最近的过去 token 高度相关）。

目标。我们上述的惊讶度度量基于损失函数 $ \ell(\cdot; \cdot) $，这是我们的记忆模块在测试时学习的目标。也就是说，我们的记忆模块是一个元模型，它基于损失函数 $ \ell(\cdot; \cdot) $ 学习一个函数。

在本节中，我们重点讨论关联记忆，其目标是将过去的数据存储为k-V对。给定输入 $ x_t $，类似于 Transformers（Vaswani 等人，2017），我们使用两个线性层将 $ x_t $ 投影为key和value：

\[\mathbf{k}_t = x_t \mathbf{W}_K, \quad \mathbf{v}_t = x_t \mathbf{W}_V, \quad (11)\]

其中 $ W_K $ 和 $ W_V \in R^{d_{in} \times d_{in}} $。接下来，我们希望记忆模块能够学习键和值之间的关联。为此，我们定义损失函数如下：

\[\ell(\mathbf{M}_{t-1}; x_t) = \|\mathbf{M}_{t-1}(\mathbf{k}_t) - \mathbf{v}_t\|_2^2, \quad (12)\]

通过在元模型（记忆）的内循环中优化上述损失函数，模型学习如何在测试时记忆键和值之间的映射。需要注意的是，类似于元学习模型（Nichol，2018；Zintgraf 等人，2019），记忆的训练是在内循环中进行的，因此参数 $ \mathbf{W}_K $ 和 $ \mathbf{W}_V $ 是上述损失函数中的超参数。因此，在内循环中，我们优化记忆模块 $ \mathbf{M} $ 的权重，而在外循环中，我们优化整个架构的其他参数。

遗忘机制

当处理非常长的序列（例如数百万个 token）时，管理哪些过去信息应该被遗忘至关重要——即使使用深度或非常大的矩阵值记忆。为此，我们使用一种自适应遗忘机制，允许记忆遗忘不再需要的信息，从而更好地管理记忆的有限容量。具体来说，给定下一个 token $ x_t $，我们修改更新规则如下：

\[\mathbf{M}_t = (1 - \alpha_t) \mathbf{M}_{t-1} + S_t, \quad (13)\] \[S_t = \eta_t S_{t-1} - \theta_t \nabla \ell(\mathbf{M}_{t-1}; x_t), \quad (14)\]

其中：

• $ \alpha_t \in [0, 1] $ 是一个门控机制，灵活控制记忆；即决定应该遗忘多少信息。例如：

• 通过设置 $ \alpha_t \to 0 $，可以在不影响过去抽象的情况下更新记忆；
• 通过设置 $ \alpha_t \to 1 $，可以清除整个记忆

在本节后面，我们将展示这种权重衰减机制与现代 RNN 中的门控机制密切相关（Dao 和 Gu，2024；Orvieto 等人，2023）。

记忆架构

在本文中，我们专注于使用具有 $ L_M \geq 1 $ 层的简单多层感知机（MLP）作为长期记忆的架构。选择这种架构的主要原因是，我们希望集中精力更好地激励长期记忆的设计及其融入架构的方式。然而，我们的公式和架构设计为设计在数据记忆方面更有效和高效的神经架构开辟了新的研究方向。最近，有一些有前景的工作致力于设计此类架构（Berges 等人，2024；Cetin 等人，2024；J. Zhang 等人，2024），将这些架构融入我们的框架（即用此类架构替换简单的 MLP）可能是一个有趣的未来工作方向。

当使用向量值或矩阵值记忆（De 等人，2024；Orvieto 等人，2023；S. Yang, B. Wang, Shen 等人，2024）时，记忆模块会压缩过去的数据并将其拟合到一条线上。也就是说，从元学习或在线学习的角度来看（Yu Sun 等人，2024），使用矩阵值记忆 $ \mathbf{M} = \mathbf{W} \in \mathbb{R}^{d_{\text{in}} \times d_{\text{in}}} $ 等同于优化 $ \ell(\mathbf{W}{t-1}; x_t) = |\mathbf{W}{t-1} \mathbf{k}_t - \mathbf{v}_t|_2^2 $，这是一个在线线性回归目标，因此最优解假设历史数据的潜在依赖关系是线性的。另一方面，我们认为深度记忆模块（即 $ L_M \geq 2 $ 层）在实践中更有效。这与理论结果一致，即至少具有两层的 MLP 严格比线性模型更具表达能力（Hornik, Stinchcombe, and White, 1989）。在第 5.5 节中，我们展示了深度记忆模块在实际应用中的有效性。

---

记忆检索

在上面，我们讨论了如何设计和训练一个在测试时学习记忆的长期记忆模块。一个关键的问题是：如何从记忆中检索信息？我们简单地使用不更新权重的前向传播（即推理）来检索与查询对应的记忆。形式上，给定输入 $ x_t $，我们使用线性层 $ \mathbf{W}_Q $ 投影输入，即 $ \mathbf{q}_t = x_t \mathbf{W}_Q $，并通过以下方式从记忆中检索相应的（或有用的）信息 $ y_t $：

\[y_t = \mathbf{M}^*(\mathbf{q}_t). \quad (15)\]

图2：记忆作为上下文（MAC）架构。该架构包括三个分支：（1）核心分支，（2）上下文（长期）记忆分支，以及（3）持久记忆分支。核心分支将相应的长期记忆和持久记忆与输入序列连接起来。接下来，注意力机制在序列上执行，并决定哪些信息应存储在长期记忆中。在测试时，与上下文记忆对应的参数仍在学习，与核心分支对应的参数负责上下文学习，而持久记忆的参数负责存储任务知识，因此是固定的。

3.2 如何并行化长期记忆的训练

如上所述，我们的长期记忆模块的设计等同于通过优化关联记忆损失函数 $ \ell(\mathbf{M}{t-1}; x_t) = |\mathbf{M}{t-1}(\mathbf{k}_t) - \mathbf{v}_t|_2^2 $ 来训练一个元模型，使用带有动量和权重衰减的梯度下降。因此，理论上，长期记忆模块的训练需要 $ O(N) $ 的浮点运算（FLOPs），其中 $ N $ 是序列长度。然而，在实践中，我们需要并行化训练过程，并充分利用硬件加速器（例如 TPU、GPU），因此需要将过程张量化并使用更多的矩阵乘法（matmuls）。

接下来，我们展示如何通过小批量梯度下降、数据依赖的学习率和权重衰减来重新表述内循环中的权重计算，使其仅使用矩阵乘法和求和。我们基于 Yu Sun 等人（2024）的工作，该工作表明，使用小批量梯度下降（具有恒定学习率）优化的模型的前向传播可以通过矩阵乘法计算。我们可以将序列分割为大小为 $ b \geq 1 $ 的块，并将小批量梯度下降表示为：

\[\mathbf{M}_t = (1 - \alpha_t) \mathbf{M}_{t-1} - \theta_t \nabla \ell(\mathbf{M}_{t-1}; x_t) = \beta_t \mathbf{M}_0 - \sum_{i=1}^t \theta_i \frac{\beta_t}{\beta_i} \nabla \ell(\mathbf{M}_{t'}; x_i), \quad (16)\]

其中 $ t’ = t - \text{mod}(t, b) $，且 $ \beta_i = \prod_{j=1}^i (1 - \alpha_j) $。为了简化，我们专注于第一个块，即 $ t = b $，因此 $ t’ = 0 $。此外，我们解释当 $ \mathbf{M}_t = \mathbf{W}_t $ 是线性时的情况。对于具有 $ N_p \geq 2 $ 层的 MLP，过程类似。使用我们的损失函数，我们有：

\[\nabla \ell(\mathbf{W}_0; x_t) = (\mathbf{W}_0 x_t - x_t) x_t^\top \Rightarrow \sum_{i=1}^b \theta_i \frac{\beta_b}{\beta_i} \nabla \ell(\mathbf{W}_0; x_i) = \Theta_b \mathbf{B}_b (\mathbf{W}_0 \mathbf{X} - \mathbf{X}) \mathbf{X}^\top, \quad (17)\]

其中 $ \Theta_b = \text{diag}(\theta_1, \theta_2, \ldots, \theta_b) $，且 $ \mathbf{B}b $ 类似地定义在 $ \frac{\beta_b}{\beta_i} $ 上。需要注意的是，我们不需要存储所有 $ \Theta{kb} $ 和 $ \mathbf{B}_{kb} $（$ k = 1, \ldots, N/b $），而是为每个块存储这些矩阵，从而减少内存使用。接下来，我们扩展这种表示，以便还可以纳入动量项。在带有动量的小批量梯度下降中，如果我们看动量项，我们有：

\[S_t = \eta_t S_{t-1} - \theta_t u_t, \quad (18)\]

其中 $ u_t = \nabla \ell(\mathbf{M}_{t’}; x_t) $。需要注意的是，我们可以同时计算所有 $ u_t $，因此公式 (18) 是一个线性递归，其中 $ u_t $ 是输入，$ S_t $ 是隐藏状态，$ \eta_t $ 是输入依赖的转移值。因此，我们可以使用并行关联扫描（J. T. Smith, Warrington, and Linderman, 2023）来计算该块中的 $ S_t $。

参数作为块的函数

与其让参数 $ \alpha_t $、$ \theta_t $ 和 $ \eta_t $ 依赖于输入（即 token $ x_t $ 的函数），我们可以让它们成为块的函数。尽管这会降低表达能力，但这种表述可以帮助使训练更快。在这种情况下，我们在每个块中对 $ \alpha $、$ \theta $ 和 $ \eta $ 使用相同的值。因此，在公式 (17) 中，我们可以使用单个标量存储 $ \Theta $。类似地，我们可以使公式 (18) 更快。也就是说，当 $ \eta $ 和 $ \theta $ 在每个块内可学习但时间不变时，该方程变为线性时不变系统（LTI），可以通过全局卷积计算（Gu, Goel, and Re, 2022）。在我们的实验中，我们将这些参数作为 token 的函数。然而，这种简化（即作为块的函数）可能是未来工作的兴趣点，以便以更高效的方式训练更大的模型。

3.3 持久记忆

我们的长期记忆也可以被视为一种上下文记忆，这意味着输出完全依赖于上下文。因此，除了长期记忆外，我们还使用一组可学习但与输入无关的参数来充当任务相关的记忆。这种类型的记忆在文献中被称为持久记忆或元记忆（X. Dong 等人，2024；Sukhbaatar, Grave 等人，2019）。给定 $ N_p \geq 1 $，我们使用可学习参数 $ P = [p_1, p_2, \ldots, p_{N_p}] $ 并将其附加到序列的开头：即，给定上下文窗口大小为 $ N $，我们将输入修改为：

\[x_{\text{new}} = [p_1, p_2, \ldots, p_{N_p}] \parallel x, \quad (19)\]

其中 $ \parallel $ 表示连接操作。接下来，我们从三个角度讨论持久记忆的动机：

---

记忆视角

如前所述，我们的神经长期记忆是一种上下文记忆，其中所有参数都依赖于输入。然而，一个有效的记忆系统还需要与输入无关的参数来存储任务知识的抽象。也就是说，掌握一个任务需要记忆如何完成该任务的知识，而这些参数负责存储此类知识。

---

前馈网络视角

在 Transformer 架构中，注意力模块之后有全连接层，这些层被证明类似于注意力权重，但具有与数据无关的参数。即，Sukhbaatar, Grave 等人（2019）表明，将全连接层中的 ReLU 替换为 Softmax 可以产生类似注意力的权重，其中权重与数据无关：

\[FFN(x) = W_V \text{Softmax}(W_K x), \quad (20)\]

实际上，当 $ W_K $ 和 $ W_V $ 与输入无关时，它们的作用类似于注意力模块中的 $ K $ 和 $ V $ 矩阵。持久记忆权重预计具有相同的功能，这意味着在序列的开头部分使用它们会导致具有与输入无关的注意力权重（Sukhbaatar, Grave 等人，2019）。

---

技术视角

带有因果掩码的注意力机制对序列中的初始 token 具有隐式偏差，因此注意力权重几乎总是对初始 token 高度活跃，从而导致性能下降。从技术角度来看，序列开头的这些可学习参数可以通过更有效地重新分配注意力权重来缓解这种影响（Han 等人，2024；Xiao 等人，2024）。

---

总结：

• 持久记忆的作用：存储任务知识的抽象，与输入无关。
• 前馈网络的类比：持久记忆权重类似于注意力机制中的 $ K $ 和 $ V $ 矩阵，但具有与数据无关的特性。
• 技术优势：通过在序列开头引入可学习参数，持久记忆可以缓解注意力机制对初始 token 的偏差，从而提升模型性能。

持久记忆的引入为模型提供了任务知识的存储能力，并通过优化注意力权重的分配进一步提升了模型的性能。

图3

4 如何融入记忆？

一个重要但尚未解答的问题是：如何有效且高效地将设计的神经记忆融入深度学习架构中？如前所述，从记忆的角度来看，Transformers 中的键值对矩阵可以解释为关联记忆块。由于其依赖关系的精确建模以及有限的上下文窗口，我们将其解释为短期记忆模块，专注于当前上下文窗口大小。另一方面，我们的神经记忆能够从数据中持续学习并将其存储在其权重中，可以扮演长期记忆的角色。在本节中，我们旨在通过提出 Titans 的三种不同变体来回答上述问题。在后续实验中，我们将展示每种变体的优缺点，以及在超长上下文中的效率与有效性之间的权衡。

---

4.1 记忆作为上下文（Memory as a Context, MAC）

在第一种架构设计中（见图 2），我们将记忆视为当前信息的上下文。即，给定一个长序列 $ x \in \mathbb{R}^{N \times d_{\text{in}}} $，我们首先将序列分割为固定大小的段 $ S^{(i)} $（$ i = 1, \ldots, N/C $）。给定传入的段 $ S^{(t)} $，我们将其视为当前上下文，并将其过去的段视为历史信息。因此，设 $ \mathbf{M}{t-1} $ 为段 $ S^{(t)} $ 之前的长期记忆状态，我们使用输入上下文作为查询，从长期记忆 $ \mathbf{M}{t-1} $ 中检索相应的信息。即，我们检索与 $ S^{(t)} $ 对应的过去信息：

\[h_t = \mathbf{M}^*_{t-1}(\mathbf{q}_t), \quad (21)\]

其中 $ \mathbf{q}_t = S^{(t)} \mathbf{W}_Q $。接下来，我们将此历史信息与持久记忆参数一起作为注意力模块的输入序列：

\[\tilde{S}^{(t)} = [p_1, p_2, \ldots, p_{N_p}] \parallel h_t \parallel S^{(t)}, \quad (22)\] \[y_t = \text{Attn}(\tilde{S}^{(t)}). \quad (23)\]

整个序列的注意力图结构如图 3a 所示。然后，我们使用 $ y_t $ 更新长期记忆模块以用于下一个段和最终输出：

\[\mathbf{M}_t = \mathbf{M}_{t-1}(y_t), \quad (24)\] \[o_t = y_t \otimes \mathbf{M}^*_t(y_t). \quad (25)\]

需要注意的是，在上述过程中，我们通过前向传播更新 $ \mathbf{M}_{t-1} $ 的权重。

图4

该架构具有两个关键优势：

1. 注意力机制：通过同时拥有历史和当前上下文，能够决定在当前数据下是否需要长期记忆信息。
2. 记忆管理：注意力模块帮助长期记忆仅存储当前上下文中有用的信息，避免记忆溢出。
3. 测试时的行为：
   • 持久记忆参数固定，因为它们编码了任务知识，不应更改。
   • 注意力模块权重是上下文学习者。
   • 长期记忆模块在测试时仍在学习（记忆）信息，即神经记忆的权重在测试时也会更新，以编码过去的抽象。

---

4.2 门控记忆（Gated Memory）

在第二种变体中（见图 4），我们在一个分支中直接使用输入数据更新长期记忆，在另一个分支中使用滑动窗口注意力（SWA）：

\[\tilde{x} = [p_1, p_2, \ldots, p_{N_p}] \parallel x, \quad (26)\] \[y = \text{SW-Attn}^*(\tilde{x}), \quad (27)\] \[o = y \otimes \mathbf{M}(\tilde{x}), \quad (28)\]

其中 $ \text{SW-Attn}^* $ 是带有前缀的滑动窗口注意力（见图 3b）。需要注意的是，与之前的设计不同，我们不对输入数据进行分段。此外，我们简化符号并使用 $ \mathbf{M}(x) $ 表示序列中所有 token 递归后的记忆最终输出。在上述等式中，$ \otimes $ 可以是任何非线性门控。在我们的实验中，我们使用可学习的向量值权重对输出 $ y $ 和 $ \mathbf{M}(\tilde{x}) $ 进行归一化，然后应用非线性函数 $ \sigma(\cdot) $。

该设计的整体注意力掩码如图 3b 所示。在这种设计中，滑动窗口注意力充当精确的短期记忆，而神经记忆模块则充当模型的衰减记忆。这种架构设计也可以被视为一种多头架构，其中头的结构不同（X. Dong 等人，2024）。

4.3 记忆作为层（Memory as a Layer, MAL）

在第三种变体中，我们将神经记忆作为深度神经网络的一层（见图 5）。这种架构设计在文献中更为常见，其中混合模型将循环模型与完整或滑动窗口注意力机制堆叠在一起。给定输入 $ x $，我们有：

\[\tilde{x} = [p_1, p_2, \ldots, p_{N_p}] \parallel x, \quad (29)\] \[y = \mathbf{M}(\tilde{x}), \quad (30)\] \[o = \text{SW-Attn}(y), \quad (31)\]

其中 $ \text{SW-Attn} $ 是滑动窗口注意力。这种设计的主要缺点是模型的能力受限于每一层，因此无法充分利用注意力和神经记忆模块的互补数据处理能力。在我们的实验中，为了评估这种设计中的记忆，我们使用了类似于 H3（D. Y. Fu 等人，2023）的架构，其中我们将序列模型替换为我们的神经记忆模块（LMM）。

---

无注意力的记忆

尽管上述讨论中我们将 MAL 视为 LMM 和注意力机制的顺序组合，但 MAL 的一个简单变体是将 LMM 视为没有任何注意力机制的序列模型。从记忆的角度来看，如第 1 节所述，我们期望记忆系统的每个部分都能独立工作，即使其他组件受到干扰。因此，即使没有短期记忆（即注意力机制），长期记忆模块仍然应该是一个强大的模型。我们在实验中称这种变体为 LMM 或 Titans（LMM）。我们在附录 C 中提供了关于 Titans 与其他现代循环模型联系的更多讨论。

图5

---

4.4 架构细节

为了简洁和清晰，我们避免讨论实现细节，例如使用残差连接、线性层门控和归一化。在所有块中，我们使用残差连接。在我们的实现中，我们使用 SiLU(.) 激活函数（Elfwing, Uchibe, and Doya, 2018）作为计算查询、键和值的非线性激活，并使用 $ \ell_2 $-范数对查询和键进行归一化。

---

卷积

遵循最近的现代线性循环模型（Gu 和 Dao，2024；S. Yang, Kautz, and Hatamizadeh，2024），我们在每个查询、键和值投影之后加入一维深度可分离卷积层。虽然这些一维卷积对性能的影响不大，但它们已被证明可以提升性能，并且在计算上也很高效。

---

门控

我们还遵循最近的架构，在最终输出投影之前使用归一化和线性层门控（Mehta 等人，2023）。

---

定理 4.1

与 Transformers、对角线性循环模型和 DeltaNet 不同，这些模型都受限于 $ \text{TC}^0 $（Merrill, Petty, and Sabharwal, 2024），Titans 能够解决超出 $ \text{TC}^0 $ 的问题，这意味着 Titans 在状态跟踪任务中理论上比 Transformers 和大多数现代线性循环模型更具表达能力。

#

https://arxiv.org/pdf/2501.00663v1

meta Ins在《QuickUpdate: a Real-Time Personalization System for Large-Scale Recommendation Models》给出了它们的系统实现：

摘要

深度学习推荐模型在在线公司中扮演着重要角色，并且占据了用于训练和推理的AI基础设施的主要部分。这些模型的准确性高度依赖于它们在服务端的发布速度。提高模型更新延迟和更新频率的主要挑战之一是：模型大小（model size），这些模型size已经达到了TB级别，并且预计未来还会进一步增加。大的模型size导致了在分布式服务器中更新模型时的高延迟（和写入带宽）。我们提出了QuickUpdate，一个用于大规模推荐模型实时个性化的系统，它能够作为在线训练的一部分，可以高频率地进行模型发布，提供与完全新鲜模型相当的服务准确性。该系统采用了新技术来最小化所需的写入带宽，包括：优先参数更新、间歇性全模型更新、模型转换和宽松一致性（relaxed consistency）。我们使用真实世界的数据在Meta的一个最大生产模型上评估了QuickUpdate。结果表明，QuickUpdate提供了与完全新鲜模型相当的服务准确性，同时将平均发布的更新大小和所需带宽减少了超过13倍。它为实时服务生产模型提供了一个可扩展的解决方案，这在网络和存储带宽有限的情况下，否则是不可能大规模实现的。

1 引言

深度学习推荐模型（DLRM）在许多在线公司中被广泛使用。这些模型通过大规模数据进行训练，以学习用户和产品特征，从而在各种场景中提供个性化推荐。例如，Netflix [7] 和 YouTube [4] 为用户提供电影列表；Amazon [19] 和 Alibaba [20] 根据用户的搜索查询推荐相关产品；Google [3] 和 Meta [23] 则根据用户兴趣展示广告和内容。DLRM 在这些公司中占据了 AI 基础设施的主要部分。以 Meta 为例，DLRM 消耗了超过 80% 的机器学习推理周期 [8] 和超过 50% 的训练周期。

推荐模型有助于业务增长。例如，它们贡献了 Amazon 总购买量的 35% [8, 14]。由于这种广泛的业务影响，准确性成为大规模推荐模型的重要性能指标。特别是在 Meta 的业务中，设计检查点（checkpoint）和量化算法时要求准确性损失小于 0.01%[5]。这是一个非常狭窄的容差范围，表明了推荐模型及其准确性的重要性。

模型新鲜度 是个性化推荐模型准确性的关键因素 [4, 6, 9, 22, 25]。由于模型在高度动态的环境中运行推理，准确性可能会迅速下降。例如，每天都有新用户和物品注册到系统中，用户的兴趣可能会受到近期事件的影响。如果模型没有频繁更新，它将无法反映用户和产品的变化，从而导致准确性逐渐下降。为了进一步强调新鲜度的影响，图 3 展示了模型在数小时未刷新时的显著准确性损失。因此，为了将准确性保持在可接受的水平，推荐模型需要使用最新数据进行重新训练，并使用更新后的模型来服务实时推理。

保持推理模型新鲜的一种常见技术是在线训练。与每次从头开始重新训练模型不同，它使用实时流数据不断训练和优化模型。定期创建模型的快照并将其发布到位于不同地理区域的数百台服务器中。这些服务器随后利用该模型对在线查询进行实时预测。然而，更新服务模型会带来训练集群与分布式服务主机之间的延迟，导致模型刷新延迟，这主要是由于现代模型的规模庞大。

多年来，模型规模迅速增长，达到了 TB 级别，并包含数万亿个参数 [5, 10, 15]，以捕获数百万个sparse特征并提高模型准确性。有限的写入带宽在将如此大的模型传输到分布式服务器和存储时构成了挑战。因此，更新延迟可能会延长到数小时。如第 3 节详细讨论的那样，这种长时间的延迟可能会对准确性产生不利影响。

为了解决上述由大模型规模及其导致的更新延迟带来的挑战，我们提出了 QuickUpdate。QuickUpdate 采用以下设计元素来实现大规模 DLRM 的实时个性化：

1. 优先参数更新：在数百个地理分布的节点中完全更新所有服务模型的参数需要大量的网络和存储带宽，这构成了瓶颈。
   QuickUpdate 通过优先参数选择来最小化更新规模。它对服务模型中的特定参数进行排序和选择，同时从更新中修剪其余参数。这种方法显著减少了总体更新规模并缓解了带宽需求。
   参数排序算法在最小化更新规模时避免准确性下降至关重要。

2. 间歇性全模型更新：间歇性全模型更新是指在一系列连续的部分更新之后进行一次完整的模型更新。这些完整更新的主要目的是保持服务模型的长期准确性。每次部分更新后，服务模型会偏离训练模型，因为前者仍然使用过时的参数值。随着更多部分更新的进行，这种偏差会变得更大，从而导致潜在的准确性影响。为了提高准确性，间歇性地发布完整模型更新以限制服务模型与训练模型之间的差距。

3. 实时更新的模型转换：QuickUpdate 采用了几种模型转换技术来减少发布的模型规模，包括推理剪枝和量化。
   量化已在一些研究中成功实施 [11, 24, 27]，以在不牺牲准确性的情况下降低浮点精度。它有助于减少推理集群中的存储需求和通信成本。推理剪枝则应用于非常大的查找表。在 DLRM 中，实体（如用户或视频）及其对应的向量以查找嵌入表的形式存储。实际不活动的实体索引（或 ID）从服务平台中修剪，以显著减少服务模型的规模。

4. 简化的服务设计和宽松的一致性要求：在传统的服务设计中，模型在开始服务查询之前会完全加载到服务平台中，以保持强一致性。在这种设计中，每个推理请求都基于特定版本的模型权重执行，确保一致且可靠的结果。然而，这种方法由于使用额外的缓冲节点而带来了相当大的基础设施开销。
   在 QuickUpdate 中，我们通过放宽一致性要求引入了一种更高效的服务设计。权重直接在服务节点中更新，而不是使用缓冲节点。这消除了对额外基础设施的需求并减少了开销。然而，这种宽松的设计可能会导致嵌入表中的一些不一致性，因为它们可能包含新鲜和过时权重的混合。
   尽管嵌入表中可能存在不一致性，但我们的评估表明，服务模型的准确性并未受到影响，反而带来了准确性的提升。

我们使用真实世界的数据和 Meta 生产中部署的最大模型之一对 QuickUpdate 进行了评估。总体而言，我们的结果表明，QuickUpdate 能够提供与完全新鲜模型相当的服务准确性，同时将所需的写入带宽减少了超过 13 倍。它为实时服务生产 DLRM 提供了一个可扩展的解决方案，而这在网络和存储带宽有限的情况下是无法大规模实现的。QuickUpdate 通过利用新颖的技术实现了这一点，包括选择性发布每次更新的最重要部分，同时仍然结合低频的间歇性全模型更新以确保长期准确性。

2 背景

2.1 深度学习推荐模型（DLRM）

通常，深度学习推荐模型由dense和dense层组成，如图 1 所示 [5, 10, 26]。sparse层实际上是嵌入表，其中：每个嵌入表表示一个分类特征，表的每一行表示一个特定ID（例如用户 ID 或视频 ID）。嵌入表将每个 ID 转换为一个固定大小的浮点值向量，这些向量是可训练的。模型中其余可训练的部分称为dense层。

图 1

图 1 展示了数据在 DLRM 中的流动方式。sparse特征通过嵌入表进行转换；dense特征通过底部的dense层进行转换。转换后的特征随后被连接起来，并在顶部的dense层中进一步转换，以计算输入数据的可能性。

2.1.1 训练 DLRM

并行化是规模化训练推荐模型的主要方法 [5, 8]。sparse层和dense层可以采用不同的并行化逻辑。sparse层占整个模型大小的 99% 以上，可能达到数 TB 的规模。由于将所有sparse层存储在单个节点中是不可行的，因此采用模型并行化的方法将表分片到多个节点上。另一方面，dense层的规模足够小，可以容纳在每个节点中，因此它们被复制到各个节点上以利用数据并行化。

在 Meta，一个典型的训练集群包括 16 个节点，每个节点包含一个多插槽 CPU 和 8 个 GPU：

• sparse层在所有 GPU 上分片，在正向和反向计算期间进行“all-to-all”通信。
• dense层在所有 GPU 上复制，在反向传播期间使用“all-reduce”通信来聚合多个 GPU 中计算的梯度 [16]。

在训练期间，sparse层和dense层的新权重被同步计算和更新，以避免准确性下降。

2.1.2 服务 DLRM

为了以高吞吐量的方式高效地服务批量请求，通常使用 GPU 进行模型服务。在 Meta，服务节点位于专用的服务集群中。一个服务节点由主机 CPU 和附加的 GPU 组成。服务模型在服务节点之间复制，并使用数据并行化进行规模化模型服务。

• 广告嵌入表存储在单个 GPU 中，因为它们需要更高的读取吞吐量。
• 其他嵌入表存储在 CPU 中，CPU 通常具有更大的内存容量（例如 1.5 TB DRAM）。

为了存储嵌入表，使用紧凑的数据结构来最小化大小并以 GPU 友好的方式存储。特别是，嵌入表是连续存储的，每个嵌入表以行优先顺序存储在数组数据结构中。

为了刷新服务模型，使用额外的缓冲节点来避免暂停当前的服务节点（双buffer机制）。在新发布的模型加载到缓冲节点后，请求流量会切换到缓冲节点。

2.2 在线训练与离线训练

当服务模型需要使用实时数据流不断训练和更新时，会实施在线训练。在在线训练中，服务模型在提供预测的同时，会定期（以分钟到小时为单位）更新。训练在后台（通常在单独的集群中）继续运行以微调模型。训练模型发布到服务平台的速率对其生成的预测准确性有显著影响。

相比之下，离线训练不使用实时数据流进行训练，也没有严格的时间限制来训练模型并将其发布到服务平台。相反，它通常使用已经存储在数据存储中的大量数据。模型使用所有可用数据进行训练，当满足某些最优条件时训练停止，之后将其发布到服务平台。

选择在线训练还是离线训练取决于具体的使用场景。当模型需要及时更新时，会实施在线训练系统。典型的用例可能是广告、搜索和视频的 DLRM（例如 [13, 18, 21]），当环境高度动态且需要模型几乎实时更新（以分钟为单位）时。在线训练帮助这些模型整合最新数据并避免准确性下降。当业务需求不需要实时模型更新时，可以使用离线训练方法来更新模型。

2.3 优化器状态作为特征重要性度量

通常，大型 DLRM 可以使用数十万个特征进行训练；然而，其中一些特征及其对应的权重对准确性没有影响。这些特征可能属于不活跃的用户或内容 ID，或者某些其他特征可能无法提供训练信号。在模型发布时，维护所有这些参数会消耗额外的带宽，或者在服务平台中运行推理时会消耗额外的计算和 GPU 存储。

为了减轻这些不利影响，我们可以计算特征重要性，并相应地修剪那些实际上不影响准确性的特征和权重。

QuickUpdate 使用优化器状态（或梯度动量）来计算特征重要性。它属于基于梯度的特征重要性度量家族（例如 [2, 12, 18]）。优化器状态是梯度值的历史平均值，比梯度值更稳定，因为梯度值有时会在正值和负值之间振荡。优化器状态可以为我们提供以下指示：

1. 对准确性的影响：它实际上显示了在训练过程中某个特定参数被更新的频率和幅度。较高的梯度值表明对提高准确性有较大影响；如果这种影响在历史上持续存在，我们可以更有信心地推断该参数对准确性很重要；因此，优化器状态是特征对准确性影响的稳定指示。

2. 访问率：优化器状态为零（或接近零）表明该参数在训练过程中实际上没有被更新。这有两层含义。首先，这可能意味着该参数没有被访问过。这可能发生在一些不活跃的用户 ID 或过时的视频上。其次，如果实体是活跃的，接近零的值可能意味着从相关数据中可能没有重要的信号可以学习。例如，某个特定用户 ID 可能没有使用平台点击广告。

基于上述信息，QuickUpdate 使用优化器状态来完成两项不同的任务：

1. 推理剪枝：当发布完整模型时，QuickUpdate 使用优化器状态来减少完整模型更新的规模。在此阶段，QuickUpdate 关注优化器状态值的低尾部分，并修剪那些对准确性没有影响的参数。

2. 优先参数选择：当发布部分更新时，QuickUpdate 使用优化器状态来选择更重要的参数进行更新。在此阶段，QuickUpdate 关注优化器状态值的高尾部分，以更新更重要的参数。

2.4 推理剪枝

推理剪枝是为了在将完整模型快照发布到服务平台时减少模型的大小。剪枝特别针对查找嵌入表实施，并显著减少其大小（例如减少 50%）。由于查找表占 DLRM 大小的 99% 以上，剪枝可以在不影响准确性的情况下显著减少模型的大小。减少模型大小有助于消耗更少的带宽来发布模型更新；因此，更新可以以更短的延迟发布到数百个地理分布的集群中。此外，它还有助于更快地执行推理，因为计算中涉及的行数更少。

图 2

图 2 展示了一个查找表的示例，其中包含索引和对应的行。每个索引代表分配给用户、视频等的唯一 ID。每一行可以被视为一个由可训练的浮点值组成的向量，模型使用这些向量为用户生成个性化推荐。

直观上，推理剪枝算法识别出代表不活跃实体或无法提供训练信号以提高准确性的行。从数学上讲，这是通过使用优化器状态向量来实现的。具体来说，训练器可以为每一行提供一个优化器状态向量。优化器状态向量中的每个元素表示该行中对应元素的梯度动量。优化器状态向量中元素的平均值用于量化行的重要性值。如果行的重要性值接近零，则意味着该行的元素在训练过程中实际上没有被更新；因此，该行可以被剪枝。

图 2 还展示了剪枝前后查找表的变化。使用行重要性值，推理剪枝算法确定最不重要的索引，并在将完整更新发布到服务器之前将其剪枝。出于操作目的，原始索引值会被重新映射到新的索引。新索引会简单地按照它们在原始表中的出现顺序递增。

需要注意的是，在本文中，完整模型更新指的是执行推理剪枝后的模型快照。

3 动机

在本节中，我们通过真实世界的数据来展示开发 QuickUpdate 的动机。首先，我们展示了当模型一小时或更长时间未更新时，准确性如何显著下降。接着，我们讨论了模型规模扩大带来的影响。如果不修改模型发布方法，我们只能选择接受更长的更新延迟和逐渐下降的准确性，或者大量投资基础设施以保持更新延迟的一致性。最后，我们强调了无损模型更新的局限性，强调了优先更新的必要性。

准确性增益

更新完整的服务模型是一个耗时的过程，可能需要数小时。因此，用户的最新行为和兴趣（例如发布新动态或与特定内容互动）在几小时内不会反映在服务模型中，这可能会降低模型的准确性。

图 3 展示了在 Meta 的一个大规模模型中，随着模型更新延迟（1 到 7 小时），准确性如何下降。它将陈旧服务模型的准确性损失与完全新鲜模型进行了比较。结果显示，当模型更新延迟时，准确性损失显著增加，7 小时后损失超过 0.6%。

减少更新规模有助于加速模型更新并提高服务模型的准确性。

模型规模

近年来，DLRM 的规模显著增加。这些模型利用大量数据和参数来更好地理解用户兴趣和产品特征，从而提高准确性。这一进展催生了具有数万亿参数 [15] 和数 TB 模型大小 [5] 的复杂模型。此外，这一趋势预计在未来仍将持续。

随着模型规模的增加，由于传输所需的带宽增加，模型更新延迟也随之延长。如果不加以解决，预计这将导致未来模型新鲜度的下降。鉴于模型规模的持续扩大，单纯增加基础设施并不是一个可行的长期解决方案。因此，对大规模 DLRM 进行部分更新似乎是一个有前景的策略，旨在减少更新延迟，而无需增加更多基础设施。

无损模型更新

为了更好地理解模型随时间变化的比例，我们监控了更新的嵌入行，并据此计算了模型中被修改的平均比例。图 4 显示了模型随时间更新的百分比。很明显，模型的大部分在短时间内被更新。例如，在短短 10 分钟的时间间隔内，58% 的模型被更新。更新 58% 的模型是资源密集型的，需要比每小时更新完整模型更多的基础设施。这促使我们探索一种优先更新的方法，以显著减少更新规模。

图 4

4 系统概述

图 5 提供了 QuickUpdate 架构的概览。DLRM 系统由训练节点、服务节点和用于保存模型快照的远程存储组成。QuickUpdate 的发布逻辑主要在 UpdateSelector 和 UpdatePatcher 代理中实现，这两个代理分别部署在训练节点和服务节点中。

• UpdateSelector 负责决定模型的哪一部分应该更新，并在保存到远程存储之前对其进行量化。
• UpdatePatcher 根据执行的更新类型实现不同的修补策略。

以下部分提供了更多详细信息。

图 5 系统架构

4.1 更新什么

QuickUpdate 专注于对嵌入表执行部分更新，这些表通常占深度学习推荐模型的绝大部分（在我们的工作负载中超过 99%）。在这些模型中，每个表表示一个分类特征（例如用户、视频），表中的每一行对应于与该特征相关的特定 ID。

在我们的探索中，我们考虑了两种更新嵌入表的选项：

• 1. 更新选定表的所有行。
• 1. 更新所有表中的选定行（不同表的选定行索引可能不同）。

我们发现，以行级粒度进行更新可以在最小化整体更新规模的同时提高准确性。因此，QuickUpdate 决定服务端需要更新表中的哪些特定行。这种方法使 QuickUpdate 能够优先更新更有可能提高准确性的内容或用户 ID，从而确保更新策略的高效性和有效性。

对于模型中的dense层，QuickUpdate 执行完整更新。这是因为这些层的更新规模相对较小，针对这些层的任何优化对整体更新过程的影响不大。

4.2 UpdateSelector

QuickUpdate 的 UpdateSelector 组件在训练集群中实现。这是因为它需要从训练器中获取某些模型信息（例如参数值）以准备模型更新。

在在线训练期间，训练器以批次间隔运行。在每个训练间隔结束时，训练器将模型状态和优化器状态共享给 UpdateSelector。模型状态包括分片的嵌入表和dense参数值，而优化器状态包括梯度值及其动量。这些状态从 GPU 内存复制到主机 CPU 内存中。

UpdateSelector 使用优化器状态对 CPU 中的模型副本执行以下两项任务：

1. 优先参数选择：此任务的主要目标是仅更新模型参数的一小部分，同时最小化准确性的下降（与完整更新相比）。在此阶段，QuickUpdate 根据优化器状态值选择嵌入行，优先选择那些可能对准确性提升较大的行。
2. 推理剪枝：此任务在发布完整模型时执行。推理剪枝专注于sparse嵌入表，旨在减少完整模型更新的规模。在此阶段，QuickUpdate 识别低尾优化器状态值，并剪枝值接近零的嵌入行。这些行对模型准确性的影响可以忽略不计。

一旦更新（无论是完整还是部分）准备就绪，它们会经过量化以减少其大小。量化作为一种压缩方法，对模型准确性的影响可以忽略不计。量化后的更新随后存储在远程存储中，准备用于更新过程。

4.3 UpdatePatcher

UpdatePatcher 负责加载发布的快照并更新服务模型。它对部分和完整模型更新都采用了一种高效的非原子更新方法。在非原子更新过程中，多个线程可以访问模型参数，并逐步将参数修补到服务器中。这种方法允许多个线程并发修补参数，而无需锁定服务器或模型。因此，服务器可以在应用更新的同时继续对传入流量进行推理。这种方法确保了在更新过程中实时流量的高效且不间断的服务。

4.4 工作流程

图 6 展示了 QuickUpdate 的工作流程。为简化说明，我们仅展示了训练器、UpdateSelector 和一个服务节点中的时间尺度。模型的演化是一个可重复的模式，因此我们专注于一个周期，该周期进一步分为多个间隔。在周期 $ c $ 的每个间隔 $ i $ 开始时，UpdateSelector 可以访问完整模型 $ F_{c,i} $ 以确定模型的哪一部分应该更新。具体来说，首先会发布一个完整快照（即 $ F_{c,1} $）并加载到服务器中，然后连续的部分更新（$ P_{c,i} $ 其中 $ i > 1 $）会被发布并修补到完整快照中，以创建服务快照 $ S_{c,i} $。

图6 QuickUpdate工作流

将更多部分更新与服务模型合并可能会导致服务模型 $ S_{c,i} $ 与当前训练器状态 $ F_{c,i} $ 之间的偏差增大。这种偏差可能会导致准确性下降。因此，另一个完整的新鲜快照（即 $ F_{c+1,1} $）将被发布到服务集群，标志着当前周期的结束。服务端的模型演化可以表示如下：

\[S_{c,1} = F_{c,1} \\ S_{c,i} = M(S_{c,i-1}, P_{c,i}) \quad \text{对于} \quad 1 < i \leq I\]

其中：

• $ I $ 是一个周期中的间隔数
• $ M $ 是合并操作符。合并操作符简单地复制 $ P_{c,i} $ 的参数值并将其更新到 $ S_{c,i-1} $ 中。

5 设计

在本节中，我们讨论了设计选项及其对准确性指标的影响。我们首先定义了指导设计和评估的具体准确性指标。通过在整个设计过程中优先考虑准确性，我们的目标是创建一个有效的系统，在解决网络和存储带宽瓶颈的同时，提供高服务准确性。需要注意的是，QuickUpdate 是可配置的，并在生产环境中进行监控，以应对罕见的准确性下降情况。

5.1 准确性指标

二元交叉熵（Binary Cross Entropy）或熵（Entropy）[17] 是评估广告模型准确性的一个众所周知的综合指标。在本研究中，我们使用归一化熵（Normalized Entropy, NE），其定义为二元交叉熵除以一个常数。为了理解部分更新相对于完全新鲜快照和过时模型的表现，我们计算了 NE 的以下变体。为简化说明，我们从符号中省略了周期下标 $ c $。

1. NE 损失：它表示使用模型 $ S_i $ 而不是相应的完全新鲜模型 $ F_i $ 运行推理时的准确性下降。

   \[\text{NE}_{\text{loss}}(S_i) = \frac{\text{NE}_{S_i} - \text{NE}_{F_i}}{\text{NE}_{F_i}} \times 100 \quad (1)\]

   其中：

   • $ NE_{S_i} $ 和 $ NE_{F_i} $ 分别表示模型 $ S_i $ 和 $ F_i $ 的归一化熵。

2. NE 增益：它表示如果使用 $ S_i $ 进行推理而不是过时模型，可以预期的准确性提升：

   \[\text{NE}_{\text{gain}}(S_i) = \frac{\text{NE}_{S_i} - \text{NE}_{\text{stale}}}{\text{NE}_{\text{stale}}} \times 100 \quad (2)\]

   过时模型被认为是最近发布的完整模型 $ F_1 $。

3. NE 恢复：它表示模型 $ S_i $ 已达到的最大 NE 增益的百分比。我们假设如果可以使用完全训练的模型 $ F_i $ 进行推理，则可以实现最大 NE 增益。因此，NE 恢复定义为：

   \[\text{NE}_{\text{recovery}}(S_i) = \frac{\text{NE}_{\text{gain}}(S_i)}{\text{NE}_{\text{gain}}(F_i)} \times 100 \quad (3)\]

5.2 选择标准

为了优先更新能够带来更大准确性增益的行，我们需要一个在训练过程中保持稳定的可靠指标。虽然梯度向量可以作为标准，但其在正值和负值之间的振荡引入了数值不稳定性。相反，我们可以使用优化器状态向量（也称为动量），它提供了更稳定的度量。优化器状态向量表示特定行的历史梯度的平均平方和。通过将：

• $ OS^r_{c,i} $： 表示为模型在间隔 $ i $ 时行 $ r $ 的优化器状态向量
• $ \overline{OS^r_{c,i}}$： 表示为其元素的平均值，我们可以利用该度量作为行重要性的指示。

直观上，具有较大 $ \overline{OS}^r_{c,i} $ 值的行更有可能提高准确性。例如，这些行可能代表频繁使用平台点击广告的特定用户，或者代表具有高访问率的特定视频。除了给定间隔的 $ \overline{OS}^r_{c,i} $ 的大小外，跟踪其随时间的变化也可能很重要。这对于我们更倾向于优先选择相对于旧版本发生变化的参数的情况可能具有潜在的信息价值。基于这些直觉，我们评估以下选择标准：

1. 绝对优化器状态：

   \[\text{abs}(\overline{OS_{c,i}^r}) \quad \text{对于} \quad i > 1 \quad (4)\]

2. 增量优化器状态：

   \[\text{abs}(\overline{OS_{c,i}^r} - \overline{OS_{c,i-1}^r)} \quad \text{对于} \quad i > 1 \quad (5)\]

选择使用绝对优化器状态还是增量优化器状态作为选择标准取决于它们各自的优势和权衡。虽然绝对优化器状态提供了行对准确性影响的稳定和综合度量，但增量优化器状态捕捉了与前一个间隔相比影响的变化。然而，使用增量优化器状态需要额外的内存来存储前一个间隔的优化器状态。为了评估这些标准的影响，我们进行了实验，间隔长度为 30 分钟，更新规模为 10%。在发布完整快照并再训练一小时后，我们根据这两个标准发布了更新规模为 10% 的部分快照。然后，我们评估了与完全新鲜模型相比的服务准确性。表 1 中的结果（在多次此类实验中一致）表明，增量优化器状态实现了 100% 的 NE 恢复，而绝对优化器状态实现了 70% 的 NE 恢复。这意味着基于增量优化器状态选择行可以减少服务模型与相应完整快照之间的差异。

表1 不同选择标准的NE recovery

5.3 增量选择的基线

增量优化器状态是基于基线计算的。在计算增量优化器状态时，我们考虑了两种选择基线的方法：

1. 上一次更新时的模型状态：此选项将上一次更新时的模型状态作为基线。这与上一节中增量优化器状态的定义相同。

2. 上一次完整更新时的模型状态：如第 5.6 节所述，QuickUpdate 还利用间歇性完整更新。在此基线选项中，将上一次间歇性完整更新时的模型状态用作增量。在这种情况下，增量优化器状态定义为：

   \[\text{abs}(\overline{\text{OS}^r_{c,i}} - \overline{\text{OS}^r_{c,1}}) \quad \text{对于} \quad i > 1 \quad (6)\]

对于第一种选项，需要在每个训练间隔结束时保存基线，而对于第二种选项，只需在周期的第一个间隔中保存一个基线。因此，第一种选项提供了更新鲜的基线，但需要额外的计算资源将其保存在内存中。

我们通过实验检查了不同基线对服务准确性的影响，每个间歇性完整更新后跟随四个部分更新。我们评估了与完全新鲜模型相比的服务准确性（NE 恢复）。表 2 中的结果显示，使用上一次间隔的基线可以实现 3.11%（95.94% 对比 99.05%）更高的 NE 恢复。这表明，使用上一次更新的模型状态作为基线可以更好地反映最近的用户兴趣，因为它在每个更新间隔中都会刷新。此外，使用完整更新作为基线可能会优先选择在前一个间隔中重要但不再对准确性有贡献的参数。随着时间的推移，这些参数的优化器状态可能达到平稳状态，但由于其较大的增量优化器状态值，完整更新基线可能仍然会考虑它们。更频繁地刷新基线有助于消除对此类参数的优先选择，转而优先选择最近变化的参数，这些参数更有可能对准确性提升有贡献。

5.4 实时推理剪枝

如第 2.4 节所述，推理剪枝有助于减少服务模型的大小和所需的 GPU 数量。它实际上会剪枝那些不再活跃或对准确性影响可以忽略的行（或 ID），以减少嵌入表的大小。然后，剪枝后的表以 GPU 访问友好的方式紧凑地存储，以进一步减少服务集群中的大小。

剪枝仅在完整模型发布到服务集群时实施。对于后续的间隔，我们希望部分更新能够与完整模型更新中的剪枝表兼容。理想情况下，部分更新中的行 ID 应存在于剪枝表中。这有助于我们简单地更新现有行的值，而无需重新构建 GPU 中的表。然而，情况并非总是如此。由于部分更新的训练数据与完整模型更新不同，可能会出现某些行 ID 在部分更新中变得重要，而这些 ID 在服务端的剪枝表中不存在的情况。当发生这种情况时，一种简单的实现方法是将缺失的行插入服务端的剪枝表中，但这可能是资源密集型的，并且可能需要重新调整所有 GPU 上的嵌入表以确保可访问性和效率（例如，避免内存碎片）。

为了避免嵌入表的密集重新调整，我们探索了两种与部分更新兼容的推理剪枝策略：

1. 固定索引剪枝（见图 7a）：在此策略中，QuickUpdate 执行优先参数选择以选择要更新的候选行索引。然而，仅更新嵌入表中已存在的行，而剪枝的行保持不变。

2. 固定剪枝比例（见图 7b）：在此策略中，每次完整更新时从嵌入表中剪枝固定比例的行。当 QuickUpdate 执行优先参数选择时，它最多选择 $ X $ 个索引进行更新，其中 $ X $ 是服务平台上给定表中的总行数。这确保了表中的行数保持一致。

图 7

第一种策略避免了重新调整，因为只会进行行更新操作，而不会向嵌入表中插入新行。第二种策略通过使用行更新和索引重映射操作来避免重新调整。由于嵌入表的大小在第二种策略中不会改变，因此也避免了在 GPU 之间重新分片嵌入表的需要。

为了评估这两种剪枝策略，我们考虑了三种训练场景：1-无剪枝，2-固定剪枝索引，3-每个表固定剪枝比例。

我们的实验表明，剪枝导致的 NE 损失实际上可以忽略不计（<0.001%），且两种剪枝策略之间没有准确性差异。考虑到实现需求，我们选择了固定剪枝索引策略，因为其实现更简单。与固定剪枝比例策略不同，它不需要在每次新更新时更新索引映射。

6 评估

我们在 Meta 部署的最大推荐模型之一上评估了 QuickUpdate，使用真实世界的数据，并在类似于 [15] 的生产训练集群上进行训练。该模型是 [16] 中提出的 DLRM 模型的扩展，但其规模显著更大，达到 TB 级别。我们在所有实验中使用相同的预记录数据流，使实验可重复且可比较，并消除了由于时间数据变化导致的潜在结果偏差。模型最初使用几周的真实世界数据进行训练作为预热期，以达到稳定状态。对于准确性评估，我们评估了在训练数据之后的时间段内数据流的服务预测（即推理期间评估的数据未在之前的训练中使用）。

6.1 准确性

在本节中，我们比较了不同更新粒度对准确性的影响，并推导出最小完整快照频率，以确保 NE 损失不超过 0.01%。在这些实验中，我们在开始时发布一个完整快照，并继续发布具有不同粒度的部分更新。这些部分更新应用于完整服务快照之上，并使用相同的记录数据集进行准确性评估。

6.1.1 与过时模型相比的 NE 增益

我们首先比较 QuickUpdate 与过时模型的准确性，以量化准确性增益并验证在完整快照之上应用部分更新不会对准确性产生负面影响。这里的过时模型指的是最初发布的完整快照。图 8 显示了不同更新粒度（且无间歇性完整模型更新）相对于过时模型的 NE 增益。所有更新规模的 NE 增益均高于过时模型，并且 NE 增益随时间增加。5% 和 10% 的更新提供了非常相似的 NE 增益，但 1% 的更新返回的 NE 增益较少，表明一些重要的行未包含在 1% 的更新中。总体而言，这些趋势表明，即使在应用部分更新超过 10 小时后，也没有负面影响，并且与过时模型相比，准确性提高了 0.7%。

图 8

6.1.2 与完全新鲜模型相比的 NE 损失

在本节中，我们研究了使用部分更新发布的 QuickUpdate 模型的 NE 损失，与理想的完全新鲜服务模式进行比较。图 9 中的结果显示，使用 10% 更新时，NE 损失在整个 10 小时内低于 0.005%。使用 5% 更新时，NE 损失始终高于 10% 更新，但在超过 6 小时内仍低于 0.01%。随着训练周期的增加，NE 损失增加，因为服务模型与相应训练模型之间的差异增加。结果还展示了采用不同更新粒度对完整模型发布延迟的影响，同时确保 NE 损失保持在可接受的 0.01% 阈值以下。通过采用 10% 的粒度，我们可以有效地将完整模型发布的需求延迟超过 10 小时。同样，当使用 5% 的粒度时，我们可以将完整模型发布延迟 6 小时，同时仍将 NE 损失保持在可接受范围内。这突显了 5% 粒度下部分更新在捕获重要更新并在相当长的时间内保持模型准确性和新鲜度方面的有效性。

图 9

6.1.3 短期内的 NE 损失

为了分析短期内的 NE 损失，我们进行了一项评估，涉及四个连续的 10 分钟更新。检查的更新粒度为 5%、3% 和 1%。每次更新后，使用未见过的数据测量与完全新鲜模型相比的 NE 损失。图 10 显示了不同 10 分钟间隔内的变化，强调了流数据的波动性。然而，当在多个短时间间隔内平均时，NE 损失趋于稳定。正如预期的那样，结果显示，随着粒度的增加，NE 损失减少。最后一列显示的平均 NE 损失证实，5% 的粒度在我们的工作负载中会返回可接受的 NE 损失（平均而言）。

图 10

6.1.4 结论

准确性结果证明了 QuickUpdate 在采用 5% 更新粒度长达 6 小时的有效性，同时保持与完全新鲜模型相当的准确性水平，并确保 NE 损失低于 0.01% 的阈值。

此外，使用 QuickUpdate 的 5% 更新粒度允许在需要发布完整模型之前延迟 6 小时。这种延迟之所以可能，是因为部分更新成功捕获并整合了重要变化，从而生成了准确且最新的模型。

基于这些发现，QuickUpdate 默认每 6 小时触发一次间歇性完整模型发布，从而优化了准确性与更新频率之间的平衡。

6.2 分析长期行收敛性

在之前的分析中，我们的重点是基于准确性指标最小化部分更新粒度，并确定间歇性完整更新的适当频率。结果表明，使用 5% 粒度的部分更新持续 6 小时可以达到令人满意的准确性。

在本实验中，我们的目标是探索部分更新更新了模型中多少百分比的重要行。

为了确定重要行的代理，我们训练模型 6 小时（即与满意准确性相同的持续时间）。我们将重要行定义为训练模型中排名前百分之几的行，使得在 6 小时结束时发布这些行（而不是整个模型）将返回令人满意的准确性（即与完全新鲜模型相比差异低于 0.01%）。

图 11 显示了在 6 小时训练后，不同大小的单次更新与完全新鲜模型相比的 NE 损失。可以看出，单次 5% 的部分更新不足以将 NE 损失降低到可接受的 0.01% 阈值以下。然而，10% 的部分更新证明足以将 NE 损失降低到可接受的水平。这表明排名前 10% 的嵌入行是此时间窗口内重要行的良好代理。

图 11

为了了解这些重要行中有多少百分比被多个较小的 5% 更新覆盖，我们运行了 QuickUpdate 6 小时，并使用多个 5% 粒度的部分更新。在将所有这些更新合并为一个联合集后，我们观察到该集合涵盖了上述重要行的 70%，并总体覆盖了模型中所有行的 7.3%。因此，大部分重要行被连续的较小部分更新所覆盖。

6.3 带宽使用

在 QuickUpdate 中，更新大小是带宽使用的代理。带宽使用量取决于粒度、更新间隔和间歇性完整模型更新的频率。通常，这些参数是可配置的，并可能根据 DLRM 的类型和所需的准确性而变化。在本节中，我们评估了基于发布模型百分比的不同策略的带宽使用情况。详细信息如下并在图 12 中展示：

图 12

1. 基线 1：每小时发布一次完整模型。
2. 基线 2：每 10 分钟发布一次完整模型（未在图中显示）。
3. 5% 更新（默认策略）：每 10 分钟发布一次部分更新，粒度为 5%，每 6 小时发布一次间歇性完整更新（如 6.1 节所述）。
4. 10% 更新：与之前的策略类似，每 10 分钟发布一次部分更新，但粒度为 10%。每 6 小时发布一次间歇性完整更新。

为了比较这些策略，我们平均了消耗的带宽。结果显示，默认策略（5% 更新粒度，6 小时间歇性完整更新间隔）平均每小时写入模型大小的 43.6%，而策略 3（10% 更新）为 68.2%，基线 1 为 100%。基线 2 提供了与策略 2 和 3 相当的准确性，但需要每小时发布模型大小的 600%。

总体而言，使用默认策略，QuickUpdate 能够将消耗的带宽比基线 1 减少 2.3 倍，同时提供与完全新鲜模型相当的更好准确性。与基线 2 相比（由于网络和存储带宽限制，无法大规模实施），QuickUpdate 能够将所需带宽减少超过 13 倍，同时仍提供相当的准确性。

6.4 宽松一致性

传统上，服务模型以原子方式更新以保持一致的推理。这涉及将所有模型权重加载到缓冲节点中，这些节点随后成为计算推理的服务节点。然而，这种方法由于使用缓冲节点而资源密集。为了解决这个问题，QuickUpdate 放宽了一致性要求，并在执行推理查询的同时直接在服务节点中更新参数。

我们评估了在 QuickUpdate 中间歇性完整模型更新期间的 NE 恢复（与完全新鲜模型相比），作为已更新权重百分比的函数。如图 13 所示，放宽一致性可以在加载期间提高生产中的准确性。随着更多参数的加载，NE 恢复增加。我们的数据显示，通过修补 30% 的参数，我们可以捕获约 54% 的 NE 恢复。在修补仅 70% 的参数后，NE 恢复达到约 94%。

图 13

宽松一致性允许早期服务新鲜行（而不是等待整个模型更新），从而整体提高准确性。尽管在加载期间表的视图不一致（意味着不同的行可能属于不同的状态），服务一部分新鲜行已经导致准确性增加。NE 恢复随着时间的推移继续增长，直到整个模型更新完毕。

7 相关工作

异步或部分更新策略已在少数实时 DLRM 中实施 [13,18,21]。在 Kraken [21] 中，dense参数每隔几秒批量更新一次，而sparse参数在训练器中值发生变化时更新。这是一种无损参数更新，对于具有 1000-10000 亿参数 [15] 和地理分布式服务器的大型模型，可能会产生大量流量。Monolith [13] 主要专注于开发具有无冲突嵌入表的sparse特征系统。sparse参数可以在训练时以分钟级粒度更新，其值自上次同步以来发生变化。与 Kraken 类似，这是一种无损更新，可能会产生巨大的流量。总体而言，无损模型更新可能非常资源密集，如第 3 节所述。为了克服这个问题，QuickUpdate 可以执行优先参数选择，从而减少约 78% − 92% 的带宽，并且准确性损失可以忽略不计（< 0.01%）。在另一项研究中，Ekko [18] 被设计为一个高效的系统，用于将更新从训练模型广播到所有服务推理节点。为了快速更新服务模型中的较大嵌入表，他们使用了嵌入表更新中的sparse性和时间局部性。Ekko 系统与 QuickUpdate 正交，两者可以一起实施。在 QuickUpdate 中，我们优化了设计元素，如发布间隔、更新粒度和参数选择标准，以实现所需的准确性并最小化完整模型的发布。优先参数选择是我们在本文中使用的技术之一。在以往的研究中（例如 [1,2,12]），基于梯度的参数选择已在分布式训练系统中探索。Ekko [18] 进一步扩展了这一标准，并额外考虑了每个参数的请求频率和参数新鲜度作为选择标准。在 QuickUpdate 中，我们决定选择梯度动量的增量，这是一个比梯度本身更稳定的度量，并且它发布的参数能够返回与基线快照相比的最高准确性。

8 结论

QuickUpdate 是一个支持在线训练执行低延迟部分更新的系统，同时提供与完全新鲜模型相当的服务准确性。它为实时服务生产规模的 DLRM 提供了一个可扩展的解决方案。这一点尤其有价值，因为由于网络和存储带宽的限制，大规模实时服务此类模型具有挑战性。

QuickUpdate 通过利用创新技术实现了其可扩展性和准确性目标。其中一项技术是选择性发布每次更新的最重要部分，从而在保持准确性的同时减少整体更新规模。此外，QuickUpdate 以低频率结合间歇性完整模型更新，以确保长期准确性。这种选择性部分更新和间歇性完整更新的结合使 QuickUpdate 能够在低延迟服务和长期保持准确性之间取得平衡。

我们使用大规模个性化广告模型的真实世界数据对 QuickUpdate 进行了评估，结果表明 QuickUpdate 能够提供与完全新鲜模型相当的服务准确性，同时将所需的写入带宽减少超过 13 倍。

附录

• https://www.usenix.org/system/files/nsdi24-matam.pdf
• https://www.usenix.org/conference/nsdi24/presentation/matam