wechat在《Introducing Context Information in Lifelong Sequential Modeling using Temporal Convolutional Networks》提出了TCN实现：

摘要

终身序列建模（Lifelong Sequential Modeling, LSM）在社交媒体推荐系统中的重要性日益增长。这一过程中的一个关键组件是注意力模块，它从序列中提取与候选item相关的兴趣表示。通常，注意力模块以pointwise方式工作，仅关注序列中单个item与候选item的相关性。然而，邻近item中的上下文信息对于更准确地评估每个item的重要性是有用的，这一点尚未被考虑。在本研究中，我们引入了一种新颖的网络，该网络采用时间卷积网络（Temporal Convolutional Network, TCN）为整个终身序列中的每个item生成上下文感知(context-aware)的表示。这些改进的表示随后被用于注意力模块中，以生成上下文感知的兴趣表示。基于这一TCN框架，我们提出了一个增强模块，该模块包括多个TCN层及其相应的注意力模块，以捕捉不同上下文范围内的兴趣表示。此外，我们还引入了一个轻量级的子网络，根据用户的基本profile特征创建卷积滤波器。这些个性化滤波器随后被应用于TCN层中，以替代原始的全局滤波器，从而生成更具用户特定性的表示。我们在一个公共数据集和一个专有数据集上进行了实验。结果表明，所提出的网络在预测准确性和在线性能指标方面优于现有方法。

1 引言

点击率（Click-Through Rate, CTR）预测是当今社交媒体平台上推荐系统的一项基本任务。其目标是预测：用户点击推荐内容的可能性。预测的准确性在很大程度上依赖于对用户兴趣与候选内容之间关系的理解。

近年来，深度神经网络（Deep Neural Networks, DNNs）在大多数场景中显著提高了CTR预测的准确性。成功的关键在于这些网络能够对用户的历史行为序列进行建模。这一过程的核心是注意力机制，它为序列中的每个item与候选item之间的相关性提供评分。这些评分被称为注意力分数，随后用于对序列进行加权求和，以生成候选item的最终兴趣表示。为了提高注意力机制的效率，已有大量研究工作展开[47, 48]。

随着用户行为的丰富化，用户历史行为序列的长度显著增加，甚至可能扩展到终身范围。序列长度的增加导致注意力机制的计算成本显著上升。在终身序列建模（Lifelong Sequential Modeling, LSM）中，一种有效减轻计算负担的方法是将注意力机制分为两个单元：通用搜索单元（General Search Unit, GSU）和精确搜索单元（Exact Search Unit, ESU）[32]。GSU的作用是从终身序列中筛选出与候选item最相关的item，随后ESU从GSU筛选出的item中提取用户兴趣表示。这种划分使模型能够处理更长的序列，这些序列包含更丰富的用户兴趣信息，从而进一步提高CTR预测的准确性。

然而，大多数先前的工作将注意力视为pointwise评分过程，通常分析候选item与序列中每个单独item之间的相关性。这种方法忽略了相邻item在理解用户意图方面提供的宝贵信息。这在当今的社交媒体平台（如TikTok、YouTube和微信视频号）上尤为重要，因为用户通常会连续消费一系列内容。例如，如图1所示，在连续消费流生成的用户行为序列中，用户对前三个item的观看时间非常长，但从第四个item开始，观看时间急剧下降。从pointwise角度来看，第三个item可能看起来是令人满意的，因为用户观看了很长时间。然而，从上下文感知的角度来看，第三个item可能并不那么受欢迎，因为它是用户愿意长时间观看的最后一个item。第三个item的某些属性可能导致用户失去兴趣，从而导致后续item的观看时间减少。因此，模型采用上下文感知的LSM至关重要，它考虑了序列中每个item的相邻item所包含的上下文信息。这种方法确保了对用户兴趣与候选item之间关系的更全面理解，并能提供更连续和相关的推荐结果。

图1

为了实现上下文感知的LSM，我们提出了上下文感知兴趣网络（Context-Aware Interest Network, CAIN）。CAIN首先通过对历史行为序列的时间轴进行卷积操作（称为时间卷积网络，Temporal Convolutional Network, TCN）[2]，提取序列中每个item的上下文感知表示。据我们所知，这是第一个将TCN引入LSM的网络。TCN具有两大优势：

• 与循环神经网络（Recurrent Neural Networks, RNN）[9, 20]和自注意力机制（Self-Attention）[41]等方法相比，TCN更轻量且计算效率更高；
• 此外，通过调整卷积的滤波器大小，可以轻松控制上下文长度。

TCN的输出表示随后被用于后续的注意力模块中，以提取与候选item相关的上下文感知兴趣表示，而不是使用原始的item表示。基于这一TCN框架，我们在CAIN中嵌入了多范围兴趣聚合器（Multi-Scope Interest Aggregator, MSIA）模块。MSIA模块包含多个堆叠的TCN层，这些层逐渐扩展输出表示的感受野。每一层的输出被发送到其对应的注意力模块中，以提取与候选item相关的不同上下文范围的兴趣表示。由于序列长度通过TCN层逐渐减少，后续卷积层的注意力计算成本也随之降低。最后，为了增强卷积操作的个性化能力，我们提出了个性化提取器生成（Personalized Extractor Generation, PEG）模块。该模块根据用户的基本资料特征为不同用户生成卷积滤波器。我们在所有TCN层中使用PEG模块生成的滤波器，而不是为所有用户使用统一的卷积滤波器。这使得TCN的输出表示更具用户特定性，从而进一步提高最终兴趣表示的代表性。

我们在一个公共数据集和从微信视频平台用户流量日志中收集的工业数据集上进行了广泛的实验。结果表明，与现有方法相比，所提出的CAIN在CTR预测准确性方面表现更优。此外，结果显示TCN框架和MSIA模块具有高度适应性，并在多种具有不同注意力设计的LSM基线模型上提供了性能提升。值得注意的是，CAIN在在线A/B测试中也取得了显著改进。这些发现证明了CAIN在复杂环境中增强CTR预测的有效性和鲁棒性。

2 相关工作

2.1 序列建模

基于深度学习的模型在工业应用中取得了显著进展，例如在线广告和推荐系统[7, 18, 26, 32, 34, 40, 42–44, 48]。用户历史行为序列的建模（称为序列建模，Sequential Modeling, SM）对于这些模型理解用户意图并实现个性化预测至关重要。在这一领域已有大量研究工作[4, 6, 13, 32, 47, 48]。

随着用户行为变得更加复杂，历史行为序列的长度显著增加。因此，近年来更多的工作集中在终身序列建模（LSM）上。Sim[32]和UBR4CTR[33]是两种引入两阶段框架来建模用户终身序列的方法，包括通用搜索单元（GSU）和精确搜索单元（ESU）。GSU从整个用户行为历史中检索与目标item最相关的前k个item，然后将这些item输入ESU进行后续的注意力计算。该框架严重依赖于预训练的嵌入，这可能会降低GSU和ESU阶段之间的一致性。为了解决这个问题，ETA[6]提出使用SimHash[5]检索相关item，并通过局部敏感哈希（Locality-Sensitive Hashing, LSH）在ESU中对item嵌入进行编码。SDIM[3]也被提出，为候选item和行为item生成哈希签名，然后收集具有匹配哈希签名的行为item来表示用户兴趣。这两种方法都允许两个阶段共享相同的嵌入以提高一致性。此外，TWINS[4]通过引入CP-GSU来增强一致性，CP-GSU不仅检索与目标相关的行为，还检索ESU认为重要的行为。此外，一些工作将两阶段框架升级为三级注意力金字塔[21]，以进一步增强阶段之间的一致性。

然而，大多数工作将LSM视为pointwise过程，仅关注序列中单个item与候选item之间的关系。它们忽略了序列中相邻item提供的上下文信息的重要性。在本文中，我们的目标是实现一种上下文感知的LSM，考虑序列中每个item的相邻item。

2.2 上下文感知建模

上下文感知建模方法在自然语言处理（NLP）、计算机视觉（CV）和语音识别（SR）领域得到了广泛应用。长短期记忆网络（LSTM）[20] 和门控循环单元（GRU）[9] 是经典的循环神经网络（RNN）模型，被广泛用于各种 NLP 和 SR 任务 [17, 27, 29, 45?, 46]。近年来，Transformer 模型和自注意力机制（Self-Attention）[41] 已成为 NLP 领域的基础组件，其特点是仅基于注意力机制的编码器和解码器。GPT [35] 和 BERT [11] 是构建在该模块上的两个著名模型。除了基于 RNN 和基于注意力的方法外，还存在其他上下文感知建模方法。时间卷积网络（TCN）[2] 就是其中之一。卷积神经网络（CNN）在 CV 任务中占据主导地位 [14, 15, 19, 23, 36–38]，因为卷积操作非常适合处理图像，通过自然地考虑目标像素的邻近像素来生成高级特征图。TCN 沿时间维度进行卷积操作，使生成的表示具有上下文感知能力 [8, 12, 24, 25, 30]。

研究人员已将 RNN 和自注意力机制引入序列建模（SM）中，用于点击率（CTR）预测。CA-RNN [28] 和 CRNNs [39] 是两种使用 RNN 预测用户历史记录中下一个item概率的方法。DEIN [47] 使用 GRU 提取每个用户的兴趣状态，并利用 AUGGRU 对目标item的兴趣演变进行建模。考虑到序列由会话组成，DSIN [13] 使用自注意力机制提取用户在每次会话中的兴趣，然后应用 Bi-LSTM 对用户兴趣在会话间的演变进行建模。

然而，尽管这些方法为 CTR 任务带来了一定的改进，但由于计算负担沉重，它们在扩展到终身序列建模（LSM）时面临挑战。因此，在 LSM 设置中，需要一种更轻量级且计算效率更高的方法来捕捉上下文信息。

3 预备知识

本文探讨上下文感知终身序列建模（Context-Aware Lifelong Sequential Modeling, LSM）的主题。与传统 LSM 不同，传统方法在注意力过程中通常将序列中的item视为孤立行为以评估其与候选item的相关性，而上下文感知 LSM 方法则考虑序列中每个item的上下文。本文定义的上下文由序列中每个item周围的相邻item组成，通常表示用户在该item前后直接交互过的item。

形式化定义

令：$\overset{\rightarrow}{𝐿𝐻} = \lbrace lh_1, lh_2, \cdots, lh_n \rbrace $表示用户的终身行为序列。以item $lh_t$（称为中心item）为中心的上下文可表示为：

\[\text{ctx}_t = \{𝑙ℎ_{t-c_l}, 𝑙ℎ_{t-c_l+1}, · · · , 𝑙ℎ_{t-1}, 𝑙ℎ_{t+1}, · · · , 𝑙ℎ_{t+c_l-1}, 𝑙ℎ_{t+c_l}\}, \tag{1}\]

其中 $c_l$ 表示上下文长度。

特征分类

本文模型中，每个用户包含三类特征：

• 基础画像特征：记为 {𝐵}；
• 短期行为序列：记为 $\overset{\rightarrow}{𝐻}= \lbrace h_1, h_2, \cdots, h_n \rbrace$
• 终身行为序列：记为 $\overset{\rightarrow}{𝐿𝐻} = \lbrace lh_1, lh_2, \cdots, lh_n \rbrace $

需注意：

• 终身和短期行为序列中的每个item（$lh_n$ 或 $h_n$）均包含其item ID 和附加信息（如用户观看时长）；
• 对于user-item对 $\langle u_i, v_i \rangle$，模型目标为预测用户 $u_i$ 对item $v_i$ 的点击率（CTR）：

\[p_i = P(y_i = 1 \,|\, u_i, v_i, B, \overset{\rightarrow}{𝐻}^t, \overset{\rightarrow}{LH}_t; \theta), \tag{2}\]

其中：

• $\theta$ 为模型参数。

损失函数

主网络通过交叉熵损失函数优化：

\[\mathcal{L}_{CTR} = -\frac{1}{BS} \sum_{i=1}^{BS} \left[ y_i \cdot \log(p_i) + (1 - y_i) \cdot \log(1 - p_i) \right], \tag{3}\]

其中：

• $y_i \in {0, 1}$ 表示用户真实反馈（点击=1，未点击=0）；
• $BS$ 表示训练批次中的样本对总数。

4 方法论

本文提出了一种新颖的上下文感知兴趣网络（CAIN），使模型在进行终身序列建模（LSM）时能够考虑用户历史序列中每个item的上下文信息。CAIN 通过时间卷积网络（TCN）计算序列中item的上下文感知表示，并将这些表示输入后续注意力模块（替代原始item表示），从而提取针对候选item的上下文感知兴趣表示。

基于 TCN 框架，我们进一步在 CAIN 中引入两个核心模块以提升性能：

1. 多范围兴趣聚合器（MSIA）：包含多个 TCN 层，用于提取不同上下文长度下的表示，并将其输入对应的注意力模块以提取不同上下文范围的兴趣表示；
2. 个性化提取器生成器（PEG）：包含轻量级网络，根据用户基础画像特征生成卷积滤波器。

图2 展示了 CAIN 的总体架构及其与传统 LSM 网络的对比。

图2

4.1 上下文信息提取

从终身序列中提取用户对候选item兴趣的常见方法是执行注意力过程。在此过程中，候选item（记为$v$）作为query项，与序列中的item $\langle v, lh_t \rangle$（其中$lh_t \in \mathbf{LH}$）形成配对。每个$\langle v, lh_t \rangle$对根据注意力模块中的网络分配注意力分数：

\[s_t = \text{Attn}(v^e, lh^e_t; \theta_a), \tag{4}\]

其中：

• $v^e$和$lh^e_t$分别表示$v$和$lh_t$的item表示
• $\theta_a$表示注意力模块参数
• $s_i$为注意力分数。这些分数将用于后续的检索或加权求和阶段。

值得注意的是，上述注意力过程主要关注$v$与$lh_t$之间的关系。然而，在现实场景中，用户行为往往具有连续性——用户会连续交互一系列item。这凸显了考虑上下文（即序列中每个中心item前后的相邻item）对全面理解行为及其与候选item相关性的重要性。

在提出的CAIN中，我们通过在注意力模块前引入时间卷积网络（TCN）[2]实现上下文感知的LSM。

4.1.1 表示提取

与循环神经网络（RNN）[9, 20]或自注意力[41]相比，TCN具有两大优势：

1. 轻量高效：TCN仅需单次矩阵乘法即可沿序列滑动卷积核，而RNN和自注意力计算量显著更大（尤其是RNN包含无法并行的操作）；
2. 显式控制上下文长度：RNN和自注意力的上下文长度通常隐式学习，可能影响泛化能力（实验部分将深入探讨）。

TCN操作可视为在序列上滑动的线性层。给定：上下文长度$cl$，卷积核$W_C$的尺寸为$2cl+1$（覆盖中心item$lh_t$两侧的上下文）。

令：$ctx^e_t$表示上下文$ctx_t$内的item表示，每个卷积窗口的计算为：

\[cr_t = \left( ctxe_t \cup lhe_t \right) \times W^C + b_c, \tag{5}\]

其中：

• $cr_t$为$lh_t$的上下文感知表示输出
• $b_c$为偏置项

对于序列首尾缺乏足够上下文项的元素，采用零填充（zero-padding）以满足指定上下文长度。

4.1.2 表示替换

经TCN处理后，原始item表示$lhe_t$被转换为：包含中心item及其上下文信息的$cr_t$。为获取用户对候选item的上下文感知兴趣表示，我们将注意力模块中的$lhe_t$替换为$cr_t$，公式(4)更新为：

\[s_t = \text{Attn}(ve, cr_t; \theta_a). \tag{6}\]

相似的，在其它过程中，我们将$lhe_t$替换为$cr_t$，该修改确保注意力模块考虑用户交互的全局上下文，从而生成更精准的上下文感知兴趣表示。

---

4.2 多范围兴趣聚合器

不同长度的上下文可提供多样化的序列洞察：

• 长上下文反映item对用户的深层影响（如兴趣迁移/衰减）
• 短上下文捕捉用户行为的即时变化

为提取多范围上下文下的兴趣表示，我们提出多范围兴趣聚合器（MSIA），其架构如图3所示。

4.2.1 堆叠层

通过堆叠TCN层可高效扩展上下文范围。随着网络加深，TCN层的感受野(receptive field)逐步增大。设：第$n$层的上下文长度为$cl_n$，当卷积步长=1时，第$n+1$层的上下文长度为：

\[cl_{n+1} = cl_n + fs_{n+1} - 1, \tag{7}\]

其中：$fs_{n+1}$为第$n+1$层的卷积核尺寸。

通过此方式，每层输出包含不同范围的上下文信息，且深层输出的表征性因TCN非线性操作而增强。

为降低深层计算成本，除首层外其他层的步长>1。这显著缩短输出序列长度，减少后续注意力模块的计算量。同时，步长>1加速上下文范围的扩展，以更少层数实现更大范围。首层步长固定为1以保证终身序列中每个中心item均有独立表示，而深层因关注整体上下文信息，允许合并输入表示。

4.2.2 独立注意力

为提取不同上下文范围对候选item的兴趣表示，我们为每层TCN输出配备独立参数的注意力模块（支持并行计算）。注意力模块分为两类：

1. 主注意力：仅应用于首层TCN输出，使用终身注意力金字塔（LAP）[21]提取兴趣表示$IR_1$（与基线方法兼容，也可替换为ETA [6]/SDIM [3]）；
2. 辅助注意力：应用于后续TCN层输出。因步长>1导致输出序列较短，且关注整体上下文信息而非单个item，采用带线性投影的目标注意力：

\[sa_t = \frac{(W^Q_n v^e) \cdot (W^K_n cr^n_t)^\top}{\sqrt{d}}, \tag{8}\] \[IR_n = \sum_{t=1}^{T_n} sa^n_t \cdot (W^V_n cr^n_t), \tag{9}\]

其中：

• $W^Q_n, W^K_n, W^V_n$为投影权重
• $d$为内部维度
• $cr^n_t$为第$n$层TCN输出

最终，MSIA模块将所有兴趣表示$IR_n$（含$IR_1$）拼接，形成多范围上下文的集成表示，平衡细粒度item信息与宏观上下文理解。

4.3 个性化提取器生成

传统卷积滤波器通常在不同输入间共享参数，其隐含假设是：输入数据服从相似分布。然而，当使用TCN处理用户历史序列时，该假设可能不成立。用户行为在不同类型用户间差异显著，且item对后续行为的影响程度也存在巨大差异。例如：

• 高活跃用户可能主要基于item内容决定观看行为，受先前展示item的影响极小（即使存在负面体验）；
• 低活跃用户的行为更易受外部因素影响，对给定item前的展示item更敏感。

为使上下文提取更具个性化，我们提出个性化提取器生成器（PEG）模块（受个性化冷启动模块POSO [10]启发）。PEG模块为每个用户生成专属卷积滤波器，适配其独特行为模式，其架构如图4所示。

4.3.1 滤波器生成

PEG模块包含轻量子网络，用于为每个用户生成个性化卷积滤波器$W^{PC}_n$。该子网络由两个全连接层构成，通过捕捉用户基础画像特征的交互关系，将输入映射到滤波器参数空间。生成过程形式化表示为：

\[O_n = \text{ReLU}(W^{P1}_n \times I + b^{P1}_n), \tag{10}\] \[W^{PC}_n = \text{Reshape}(W^{P2}_n \times O_n + b^{P2}_n), \tag{11}\]

其中：

• $W^{P1}_n, b^{P1}_n$：第一层全连接的权重与偏置
• $W^{P2}_n, b^{P2}_n$：第二层全连接的权重与偏置
• $I$：从集合$B$中提取的用户基础画像特征表示

PEG模块使用的基础画像特征包含：

• 人口统计信息：年龄、性别、地域、教育背景等
• 行为统计信息：展示item数、点击item数、用户最常交互作者（如最常观看作者）

实验部分将详细讨论这些特征对模型性能的影响。

4.3.2 个性化卷积

生成个性化滤波器后，我们将其替换各TCN层中的原始全局滤波器。尽管可采用”全局+个性化滤波器组合”方案，但实验表明完全替换方案效果更优。因此，公式(5)在CAIN中被改写为：

\[cr_t = \left( \text{concat}(ctx^e_t, lh^e_t) \right) \times W^{PC} + b_c. \tag{12}\]

通过个性化卷积滤波器，模型能更精准捕捉用户行为的独特模式，从而提升预测精度。

#

https://arxiv.org/pdf/2502.12634

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

Deepseek AI在《DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model》详细介绍了V2版的实现：

1.介绍

在过去的几年中，大型语言模型（LLMs）（Anthropic, 2023; Google, 2023; OpenAI, 2022, 2023）经历了快速发展，为通用人工智能（AGI）的曙光提供了初步的展望。通常，随着参数数量的增加，LLM的智能水平会显著提升，使其能够在各种任务中展现出涌现（emergent）能力（Wei et al., 2022）。然而，这种改进的代价是：更大的训练计算资源需求以及推理吞吐量的潜在下降。这些限制对LLM的广泛采用和利用构成了重大挑战。为了解决这一问题，我们提出了DeepSeek-V2，一个强大的开源混合专家（MoE）语言模型，其特点是通过创新的Transformer架构实现经济的训练和高效的推理。该模型总参数量为236B，每个token激活21B参数，并支持128K token的上下文长度。

我们在Transformer框架（Vaswani et al., 2017）中优化了注意力模块和前馈网络（FFNs），分别提出了多头隐注意力（MLA）和DeepSeekMoE。

1. 在注意力机制方面，多头注意力（MHA）（Vaswani et al., 2017）的键值（KV）缓存对LLM的推理效率构成了显著障碍。为了解决这一问题，研究者们探索了多种方法，包括分组查询注意力（GQA）（Ainslie et al., 2023）和多查询注意力（MQA）（Shazeer, 2019）。然而，这些方法在减少KV缓存的同时往往会牺牲性能。为了兼顾两者，我们引入了MLA，这是一种配备低秩键值联合压缩（low-rank key-value joint compression）的注意力机制。实验表明，MLA在性能上优于MHA，同时显著减少了推理过程中的KV缓存，从而提升了推理效率。

2. 在前馈网络（FFNs）方面，我们采用了DeepSeekMoE架构（Dai et al., 2024），该架构通过细粒度的专家分割和共享专家隔离，实现了更高的专家专业化潜力。与传统的MoE架构（如GShard（Lepikhin et al., 2021））相比，DeepSeekMoE展现出了显著优势，使我们能够以较低的成本训练强大的模型。在训练过程中，我们采用专家并行策略，并设计了补充机制来控制通信开销并确保负载均衡。

通过结合这两种技术，DeepSeek-V2在性能（图1(a)）、训练成本和推理吞吐量（图1(b)）方面均表现出色。我们构建了一个高质量、多来源的预训练语料库，包含8.1T token。与DeepSeek 67B（我们之前的版本）（DeepSeek-AI, 2024）使用的语料库相比，该语料库的数据量更大，尤其是中文数据，且数据质量更高。

• 首先我们在完整的预训练语料库上对DeepSeek-V2进行预训练。
• 然后，我们收集了1.5M个对话会话，涵盖数学、代码、写作、推理、安全等多个领域，用于对DeepSeek-V2 Chat（SFT）进行监督微调。
• 最后，我们遵循DeepSeekMath（Shao et al., 2024）的方法，采用组相对策略优化（GRPO）进一步对齐模型与人类偏好，生成DeepSeek-V2 Chat（RL）。

我们在广泛的中英文基准测试中评估了DeepSeek-V2，并将其与代表性的开源模型进行了比较。评估结果表明，即使仅激活21B参数，DeepSeek-V2仍然在开源模型中表现出顶级性能，成为最强的开源MoE语言模型。

• 图1(a)显示，在MMLU上，DeepSeek-V2仅以少量激活参数就达到了顶级性能。
• 如图1(b)所示，与DeepSeek 67B相比，DeepSeek-V2节省了42.5%的训练成本，减少了93.3%的KV缓存，并将最大生成吞吐量提升至5.76倍。

我们还在开放式基准测试中评估了DeepSeek-V2 Chat（SFT）和DeepSeek-V2 Chat（RL）。值得注意的是，DeepSeek-V2 Chat（RL）在AlpacaEval 2.0（Dubois et al., 2024）上达到了38.9的长度控制胜率，在MT-Bench（Zheng et al., 2023）上获得了8.97的综合评分，在AlignBench（Liu et al., 2023）上获得了7.91的综合评分。英文开放式对话评估表明，DeepSeek-V2 Chat（RL）在开源聊天模型中具有顶级性能。此外，AlignBench的评估表明，在中文方面，DeepSeek-V2 Chat（RL）超越了所有开源模型，甚至击败了大多数闭源模型。

图1 (a)不同开源模型中MMLU精度与激活参数的对比。(b) DeepSeek-67B（Dense）和DeepSeek-v2的训练成本和推理效率。

为了促进对MLA和DeepSeekMoE的进一步研究和开发，我们还向开源社区发布了DeepSeek-V2-Lite，这是一个配备MLA和DeepSeekMoE的小型模型。其总参数量为15.7B，每个token激活2.4B参数。关于DeepSeek-V2-Lite的详细描述见附录B。

在本文的其余部分，我们首先详细描述了DeepSeek-V2的模型架构（第2节）。随后，我们介绍了预训练工作，包括训练数据构建、超参数设置、基础设施、长上下文扩展以及模型性能和效率的评估（第3节）。接着，我们展示了对齐工作（alignment），包括监督微调（SFT）、强化学习（RL）、评估结果及其他讨论（第4节）。最后，我们总结了结论，探讨了DeepSeek-V2的当前局限性，并展望了未来的工作（第5节）。

2. 架构

总体而言，DeepSeek-V2仍然基于Transformer架构（Vaswani et al., 2017），其中每个Transformer块由一个注意力模块和一个前馈网络（FFN）组成。然而，对于注意力模块和FFN，我们设计并采用了创新的架构。对于注意力模块，我们设计了多头隐注意力（MLA），利用低秩键值联合压缩来消除推理时键值（KV）缓存的瓶颈，从而支持高效推理。对于FFN，我们采用了DeepSeekMoE架构（Dai et al., 2024），这是一种高性能的MoE架构，能够以较低的成本训练强大的模型。图2展示了DeepSeek-V2的架构示意图，本节将详细介绍MLA和DeepSeekMoE的细节。对于其他微小细节（例如层归一化和FFN中的激活函数），除非特别说明，DeepSeek-V2遵循DeepSeek 67B（DeepSeek-AI, 2024）的设置。

图2 DeepSeek-V2架构示意图。MLA通过显著减少生成的KV缓存来确保高效的推理，DeepSeekMoE通过稀疏架构以经济的成本训练强模型

2.1 多头隐注意力：提升推理效率

传统的Transformer模型通常采用多头注意力（MHA）（Vaswani et al., 2017），但在生成过程中，其庞大的键值（KV）缓存会成为限制推理效率的瓶颈。为了减少KV缓存，研究者提出了多查询注意力（MQA）（Shazeer, 2019）和分组查询注意力（GQA）（Ainslie et al., 2023）。这些方法需要更少的KV缓存，但其性能无法与MHA媲美（我们在附录D.1中提供了MHA、GQA和MQA的消融实验）。

对于DeepSeek-V2，我们设计了一种创新的注意力机制，称为多头隐注意力（MLA）。MLA配备了低秩键值联合压缩，不仅性能优于MHA，而且所需的KV缓存显著减少。以下我们将介绍其架构，并在附录D.2中提供MLA与MHA的对比。

2.1.1 预备知识：标准多头注意力

我们首先介绍标准MHA机制作为背景。设：

• $d$为嵌入维度
• $n_h$为注意力头（attention heads）的数量
• $d_h$为每个头的维度
• $h_t \in \mathbb{R}^d$为第$t$个token在注意力层的输入

标准MHA首先通过三个矩阵$W_Q$、$W_K$、$W_V \in \mathbb{R}^{d_h n_h \times d}$分别生成$q_t$、$k_t$、$v_t \in \mathbb{R}^{d_h n_h}$：

\[q_t = W^Q h_t, \quad (1) \\ k_t = W^K h_t, \quad (2) \\ v_t = W^V h_t. \quad (3)\]

然后，$q_t$、$k_t$、$v_t$ 将被切分为 $n_h$ 个头以进行多头注意力计算：

\[[q_{t,1}; q_{t,2}; \dots; q_{t,n_h}] = q_t, \quad (4) \\ [k_{t,1}; k_{t,2}; \dots; k_{t,n_h}] = k_t, \quad (5) \\ [v_{t,1}; v_{t,2}; \dots; v_{t,n_h}] = v_t, \quad (6) \\ o_{t,i} = \sum_{j=1}^t \text{Softmax}\ _j \left( \frac{q_{t,i}^T k_{j,i}}{\sqrt{d_h}} \right) v_{j,i}, \quad (7) \\ u_t = W_O [o_{t,1}; o_{t,2}; \dots; o_{t,n_h}], \quad (8)\]

其中:

• $q_{t,i}$、$k_{t,i}$、$v_{t,i} \in \mathbb{R}^{d_h}$ 分别表示第 $i$ 个注意力头的查询、键和值；
• $W_O \in \mathbb{R}^{d \times d_h n_h}$ 表示输出投影矩阵。

在推理过程中，所有键和值都需要被缓存以加速推理，因此MHA需要为每个token缓存 $2 n_h d_h l$ 个元素（$l$ 为层数）。在模型部署中，这种庞大的KV缓存是一个巨大的瓶颈，限制了最大batch-size和序列长度。

2.1.2 低秩键值联合压缩

MLA的核心是：通过低秩联合压缩键和值来减少KV缓存：

\[c^{KV}_t = W^{DKV} h_t, \quad (9) \\ k^C_t = W^{UK} c^{KV}_t, \quad (10) \\ v^C_t = W^{UV} c^{KV}_t, \quad (11)\]

其中:

• $c^{KV}_t \in \mathbb{R}^{d_c}$ 是键和值的压缩隐向量；
• $d_c (\ll d_h n_h)$ 表示KV压缩维度；
• $W^{DKV} \in \mathbb{R}^{d_c \times d}$ 是下投影矩阵(down-projection matrix)；
• $W^{UK}$ 和 $W^{UV} \in \mathbb{R}^{d_h n_h \times d_c}$ 分别是键和值的上投影矩阵(up-projection matrices)。

在推理过程中，MLA只需缓存 $c_t^{KV}$，因此其KV缓存仅为 $d_c l$ 个元素。此外，在推理过程中，由于 $W^{UK}$ 可以被吸收到 $W^Q$ 中，$W^{UV}$ 可以被吸收到 $W_O$中，我们甚至不需要显式计算键和值来进行注意力计算。图3直观地展示了MLA中的KV联合压缩如何减少KV缓存。

图3 多头注意（MHA）、分组查询注意（GQA）、多查询注意（MQA）和多头潜在注意（MLA）的简化说明。通过将键和值联合压缩成一个隐向量，MLA显著降低了推理过程中的KV缓存

此外，为了减少训练期间的激活内存，我们还对Query进行了低秩压缩，尽管这并不能减少KV缓存：

\[c^Q_t = W^{DQ} h_t, \quad (12) \\ q^C_t = W^{UQ} c^Q_t, \quad (13)\]

其中:

• $c^Q_t \in \mathbb{R}^{d’_c}$ 是查询的压缩隐向量；
• $d’_c (\ll d_h n_h)$ 表示查询压缩维度；
• $W^{DQ} \in R^{d’c \times d}$ 和 $W{UQ} \in R^{d_h n_h \times d’_c}$ 分别是查询的下投影和上投影矩阵。

2.1.3 解耦的旋转位置嵌入

我们计划为DeepSeek-V2使用旋转位置嵌入（RoPE）（Su et al., 2024），这与DeepSeek 67B（DeepSeek-AI, 2024）一致。然而，RoPE与低秩KV压缩不兼容。具体来说，RoPE对键和查询都是位置敏感的。如果我们对键 $k^C_t$ 应用RoPE，公式10中的 $W_{UK}$ 将与一个位置敏感的RoPE矩阵耦合。这样，$W_{UK}$ 在推理过程中无法再被吸收到 $W_Q$ 中，因为与当前生成token相关的RoPE矩阵会位于 $W_Q$ 和 $W_{UK}$ 之间，而矩阵乘法不满足交换律。因此，我们必须在推理过程中重新计算所有前缀token的键，这将显著降低推理效率。

作为解决方案，我们提出了解耦RoPE策略，该策略使用额外的多头查询 $q^R_{t,i} \in \mathbb{R}^{d^R_h}$ 和一个共享键 $k^R_t \in \mathbb{R}^{d^R_h}$ 来承载RoPE，其中 $d^R_h$ 表示解耦查询和键的每头维度。配备解耦RoPE策略后，MLA执行以下计算：

\[[q^R_{t,1}; q^R_{t,2}; \dots; q^R_{t,n_h}] = q^R_t = \text{RoPE}(W^{QR} c^Q_t), \quad (14) \\ k^R_t = \text{RoPE}(W^{KR} h_t), \quad (15) \\ q_{t,i} = [q^C_{t,i}; q^R_{t,i}], \quad (16) \\ k_{t,i} = [k^C_{t,i}; k^R_t], \quad (17) \\ o_{t,i} = \sum_{j=1}^t \text{Softmax}_j \left( \frac{q_{t,i}^T k_{j,i}}{\sqrt{d_h + d^R_h}} \right) v^C_{j,i}, \quad (18) \\ u_t = W^O [o_{t,1}; o_{t,2}; \dots; o_{t,n_h}], \quad (19)\]

其中：

• $W^{QR} \in \mathbb{R}^{d^R_h n_h \times d’_c}$ 和 $W^{KR} \in \mathbb{R}^{d^R_h \times d}$ 是生成解耦查询和键的矩阵；
• $\text{RoPE}(\cdot)$ 表示应用RoPE矩阵的操作；$[\cdot; \cdot]$ 表示拼接操作。

在推理过程中，解耦键也需要被缓存。因此，DeepSeek-V2需要的总KV缓存为 $(d_c + d^R_h) l$ 个元素。

为了展示MLA的完整计算过程，我们在附录C中整理并提供了其完整公式。

2.1.4 KV缓存对比

我们在表1中展示了不同注意力机制下每个token的KV缓存对比。MLA仅需要少量的KV缓存，相当于仅2.25组的GQA，但其性能优于MHA。

表1

2.2 DeepSeekMoE：以经济成本训练强大模型

2.2.1 基本架构

对于FFN，我们采用了DeepSeekMoE架构（Dai et al., 2024）。DeepSeekMoE有两个关键思想：将专家分割为更细粒度以实现更高的专家专业化和更准确的知识获取，以及隔离一些共享专家以减少路由专家之间的知识冗余。在激活专家参数和总专家参数数量相同的情况下，DeepSeekMoE能够大幅超越传统MoE架构（如GShard（Lepikhin et al., 2021））。

设 $u_t$ 为第 $t$ 个token的FFN输入，我们计算FFN输出 $h’_t$ 如下：

\[h'_t = u_t + \sum_{i=1}^{N_s} \text{FFN}^{(s)}_i (u_t) + \sum_{i=1}^{N_r} g_{i,t} \text{FFN}^{(r)}_i (u_t), \quad (20) \\ g_{i,t} = \begin{cases} s_{i,t}, & s_{i,t} \in \text{Topk}(\{s_{j,t} | 1 \leq j \leq N_r\}, K_r), \\ 0, & \text{否则}, \end{cases} \quad (21) \\ s_{i,t} = \text{Softmax}_i (u_t^T e_i), \quad (22)\]

其中：

• $N_s$ 和 $N_r$ 分别表示共享专家和路由专家的数量；
• $\text{FFN}^{(s)}_i (\cdot)$ 和 $\text{FFN}^{(r)}_i (\cdot)$ 分别表示第 $i$ 个共享专家和第 $i$ 个路由专家；
• $K_r$ 表示激活的路由专家数量；
• $g_{i,t}$ 是第 $i$ 个专家的门控值；
• $s_{i,t}$ 是token与专家的亲和度；
• $e_i$ 是第 $i$ 个路由专家在该层的中心点；
• $\text{Topk}(\cdot, K)$ 表示从第 $t$ 个token与所有路由专家的亲和度分数中选取最高的 $K$ 个分数。

2.2.2 设备限制路由

我们设计了一种设备限制路由机制，以限制MoE相关的通信成本。当采用专家并行时，路由专家将分布在多个设备上。对于每个token，其MoE相关的通信频率与其目标专家覆盖的设备数量成正比。由于DeepSeekMoE中的细粒度专家分割，激活的专家数量可能较大，因此如果采用专家并行，MoE相关的通信成本会更高。

对于DeepSeek-V2，除了简单的top-K选择路由专家外，我们还确保每个token的目标专家最多分布在 $M$ 个设备上。具体来说，对于每个token，我们首先选择 $M$ 个设备，这些设备中的专家具有最高的亲和度分数。然后，我们在这些设备上的专家中进行top-K选择。在实践中，我们发现当 $M \geq 3$ 时，设备限制路由能够实现与无限制top-K路由大致相当的良好性能。

2.2.3 负载均衡的辅助损失

在自动学习的路由策略中，我们考虑了负载均衡问题。首先，负载不均衡会增加路由崩溃的风险（Shazeer et al., 2017），导致某些专家无法被充分训练和利用。其次，当采用专家并行时，负载不均衡会降低计算效率。在DeepSeek-V2的训练过程中，我们设计了三种辅助损失，分别用于控制专家级负载均衡（$L_{\text{ExpBal}}$）、设备级负载均衡（$L_{\text{DevBal}}$）和通信均衡（$L_{\text{CommBal}}$）。

表2

表3

表4

表5

专家级均衡损失

我们使用专家级均衡损失（Fedus et al., 2021; Lepikhin et al., 2021）来减轻路由崩溃的风险：

\[L_{\text{ExpBal}} = \alpha_1 \sum_{i=1}^{N_r} f_i P_i, \quad (23) \\ f_i = \frac{N_r}{K_r T} \sum_{t=1}^T \mathbb{1}(\text{Token } t \text{ 选择专家 } i), \quad (24) \\ P_i = \frac{1}{T} \sum_{t=1}^T s_{i,t}, \quad (25)\]

其中，$\alpha_1$ 是一个超参数，称为专家级均衡因子；$\mathbb{1}(\cdot)$ 是指示函数；$T$ 表示序列中的token数量。

设备级均衡损失

除了专家级均衡损失外，我们还设计了设备级均衡损失，以确保不同设备之间的计算负载均衡。在DeepSeek-V2的训练过程中，我们将所有路由专家划分为 $D$ 个组 ${E_1, E_2, \dots, E_D}$，并将每个组部署在单个设备上。设备级均衡损失计算如下：

\[L_{\text{DevBal}} = \alpha_2 \sum_{i=1}^D f'_i P'_i, \quad (26) \\ f'_i = \frac{1}{|E_i|} \sum_{j \in E_i} f_j, \quad (27) \\ P'_i = \sum_{j \in E_i} P_j, \quad (28)\]

其中：

• $\alpha_2$ 是一个超参数，称为设备级均衡因子。

通信均衡损失

最后，我们引入了通信均衡损失，以确保每个设备的通信负载均衡。尽管设备限制路由机制保证了每个设备的发送通信是有界的，但如果某个设备接收的token比其他设备多，实际通信效率也会受到影响。为了缓解这一问题，我们设计了通信均衡损失如下：

\[L_{\text{CommBal}} = \alpha_3 \sum_{i=1}^D f''_i P''_i, \quad (29) \\ f''_i = \frac{D}{M T} \sum_{t=1}^T \mathbb{1}(\text{Token } t \text{ 被发送到设备 } i), \quad (30) \\ P''_i = \sum_{j \in E_i} P_j, \quad (31)\]

其中：

• $\alpha_3$ 是一个超参数，称为通信均衡因子。

设备限制路由机制的原则是确保每个设备最多向其他设备传输 $M T$ 个隐藏状态。同时，通信均衡损失用于鼓励每个设备从其他设备接收大约 $M T$ 个隐藏状态。通信均衡损失保证了设备之间的信息交换均衡，从而提高了通信效率。

2.2.4 Token丢弃策略

尽管均衡损失旨在鼓励负载均衡，但必须承认它们无法保证严格的负载均衡。为了进一步减轻因负载不均衡导致的计算浪费，我们在训练期间引入了设备级的token丢弃策略。该方法首先计算每个设备的平均计算预算，这意味着每个设备的容量因子为1.0。然后，受Riquelme et al. (2021)启发，我们在每个设备上丢弃具有最低亲和度分数的token，直到达到计算预算。此外，我们确保属于大约10%训练序列的token永远不会被丢弃。通过这种方式，我们可以根据效率需求灵活决定在推理期间是否丢弃token，并始终确保训练和推理之间的一致性。

3. 预训练

3.1 实验设置

3.1.1 数据构建

在保持与DeepSeek 67B（DeepSeek-AI, 2024）相同的数据处理阶段的基础上，我们扩展了数据量并提升了数据质量。为了扩大预训练语料库，我们探索了互联网数据的潜力并优化了清理流程，从而恢复了大量被错误删除的数据。此外，我们加入了更多的中文数据，旨在更好地利用中文互联网上的语料库。除了数据量，我们还关注数据质量。我们从各种来源丰富了预训练语料库的高质量数据，同时改进了基于质量的过滤算法。改进后的算法确保大量无益数据被移除，而有价值的数据则大部分被保留。此外，我们从预训练语料库中过滤掉了争议性内容，以减轻特定区域文化引入的数据偏差。关于该过滤策略影响的详细讨论见附录E。

我们采用了与DeepSeek 67B相同的分词器，该分词器基于字节级字节对编码（BBPE）算法构建，词汇量为100K。我们的分词预训练语料库包含8.1T token，其中中文token比英文多约12%。

3.1.2 超参数

模型超参数：我们将Transformer层数设置为60，隐藏维度设置为5120。所有可学习参数均以标准差0.006随机初始化。在MLA中，我们将注意力头数$n_h$设置为128，每头维度$d_h$设置为128。KV压缩维度$d_c$设置为512，查询压缩维度$d’_c$设置为1536。对于解耦查询和键，我们将每头维度$d^R_h$设置为64。根据Dai et al. (2024)，我们将除第一层外的所有FFN替换为MoE层。每个MoE层包含2个共享专家和160个路由专家，其中每个专家的中间隐藏维度为1536。在路由专家中，每个token激活6个专家。此外，低秩压缩和细粒度专家分割会影响层的输出规模。因此，在实践中，我们在压缩隐向量后使用额外的RMS Norm层，并在宽度瓶颈（即压缩隐向量和路由专家的中间隐藏状态）处乘以额外的缩放因子以确保训练稳定。在此配置下，DeepSeek-V2总参数量为236B，每个token激活21B参数。

训练超参数：我们使用AdamW优化器（Loshchilov and Hutter, 2017），超参数设置为$\beta_1 = 0.9$、$\beta_2 = 0.95$、$\text{weight_decay} = 0.1$。学习率采用预热和阶梯衰减策略（DeepSeek-AI, 2024）。初始时，学习率在前2K步从0线性增加到最大值。随后，在训练约60%的token后，学习率乘以0.316，在训练约90%的token后再次乘以0.316。最大学习率设置为$2.4 \times 10^{-4}$，梯度裁剪范数设置为1.0。我们还使用了批量大小调度策略，在前225B token的训练中，批量大小从2304逐步增加到9216，之后保持9216。我们将最大序列长度设置为4K，并在8.1T token上训练DeepSeek-V2。我们利用流水线并行将模型的不同层部署在不同设备上，每层的路由专家均匀分布在8个设备上（$D = 8$）。对于设备限制路由，每个token最多发送到3个设备（$M = 3$）。对于均衡损失，我们设置$\alpha_1 = 0.003$、$\alpha_2 = 0.05$、$\alpha_3 = 0.02$。我们在训练期间使用token丢弃策略以加速训练，但在评估时不丢弃任何token。

3.1.3 基础设施

DeepSeek-V2基于HAI-LLM框架（High-flyer, 2023）进行训练，这是我们工程师开发的高效轻量级训练框架。它采用了16路零气泡流水线并行（Qi et al., 2023）、8路专家并行（Lepikhin et al., 2021）和ZeRO-1数据并行（Rajbhandari et al., 2020）。由于DeepSeek-V2激活参数相对较少，并且部分算子被重新计算以节省激活内存，因此可以在不需要张量并行的情况下进行训练，从而减少通信开销。此外，为了进一步提高训练效率，我们将共享专家的计算与专家并行的all-to-all通信重叠。我们还为通信、路由算法和跨专家的融合线性计算定制了更快的CUDA内核。此外，MLA还基于改进版的FlashAttention-2（Dao, 2023）进行了优化。

我们在配备NVIDIA H800 GPU的集群上进行所有实验。H800集群中的每个节点包含8个GPU，节点内通过NVLink和NVSwitch连接。节点间使用InfiniBand互连以促进通信。

3.1.4 长上下文扩展

在DeepSeek-V2的初始预训练后，我们使用YaRN（Peng et al., 2023）将默认上下文窗口长度从4K扩展到128K。YaRN特别应用于解耦共享键$k^R_t$，因为它负责承载RoPE（Su et al., 2024）。对于YaRN，我们将缩放因子$s$设置为40，$\alpha$设置为1，$\beta$设置为32，目标最大上下文长度设置为160K。在这些设置下，我们可以预期模型在128K的上下文长度下表现良好。与原始YaRN略有不同，由于我们独特的注意力机制，我们调整了长度缩放因子以调节注意力熵。因子$\sqrt{t}$计算为$\sqrt{t} = 0.0707 \ln s + 1$，旨在最小化困惑度。

我们额外训练了1000步，序列长度为32K，批量大小为576个序列。尽管训练仅在32K的序列长度下进行，但模型在128K的上下文长度下仍表现出色。如图4所示，在“Needle In A Haystack”（NIAH）测试中，DeepSeek-V2在所有上下文窗口长度（最高128K）下均表现良好。

图4

3.2. 评估

3.2.1. 评估基准

DeepSeek-V2 是在双语语料库上进行预训练的，因此我们在英语和中文的一系列基准上对其进行了评估。我们的评估基于集成在 HAI-LLM 框架中的内部评估框架。包含的基准分类如下，其中带下划线的基准为中文：

• 多学科选择题数据集 包括 MMLU (Hendrycks et al., 2020)、C-Eval (Huang et al., 2023) 和 CMMLU (Li et al., 2023)。
• 语言理解和推理数据集 包括 HellaSwag (Zellers et al., 2019)、PIQA (Bisk et al., 2020)、ARC (Clark et al., 2018) 和 BigBench Hard (BBH) (Suzgun et al., 2022)。
• 闭卷问答数据集 包括 TriviaQA (Joshi et al., 2017) 和 NaturalQuestions (Kwiatkowski et al., 2019)。
• 阅读理解数据集 包括 RACE (Lai et al., 2017)、DROP (Dua et al., 2019)、C3 (Sun et al., 2019) 和 CMRC (Cui et al., 2019)。
• 指代消解数据集 包括 WinoGrande (Sakaguchi et al., 2019) 和 CLUEWSC (Xu et al., 2020)。
• 语言建模数据集 包括 Pile (Gao et al., 2020)。
• 中文理解与文化数据集 包括 CHID (Zheng et al., 2019) 和 CCPM (Li et al., 2021)。
• 数学数据集 包括 GSM8K (Cobbe et al., 2021)、MATH (Hendrycks et al., 2021) 和 CMath (Wei et al., 2023)。
• 代码数据集 包括 HumanEval (Chen et al., 2021)、MBPP (Austin et al., 2021) 和 CRUXEval (Gu et al., 2024)。
• 标准化考试 包括 AGIEval (Zhong et al., 2023)。注意，AGIEval 包含英语和中文子集。

根据我们之前的工作 (DeepSeek-AI, 2024)，我们对以下数据集采用基于困惑度（perplexity）的评估：HellaSwag、PIQA、WinoGrande、RACE-Middle、RACE-High、MMLU、ARC-Easy、ARC-Challenge、CHID、C-Eval、CMMLU、C3 和 CCPM；对以下数据集采用基于生成的评估：TriviaQA、NaturalQuestions、DROP、MATH、GSM8K、HumanEval、MBPP、CRUXEval、BBH、AGIEval、CLUEWSC、CMRC 和 CMath。此外，我们对 Pile-test 进行基于语言建模的评估，并使用 Bits-Per-Byte (BPB) 作为指标，以确保使用不同分词器的模型之间的公平比较。

为了直观地了解这些基准，我们在附录 G 中提供了每个基准的评估格式。

3.2.2. 评估结果

在表 2 中，我们将 DeepSeek-V2 与几个代表性的开源模型进行了比较，包括 DeepSeek 67B (DeepSeek-AI, 2024)（我们之前的版本）、Qwen1.5 72B (Bai et al., 2023)、LLaMA3 70B (AI@Meta, 2024) 和 Mixtral 8x22B (Mistral, 2024)。我们使用内部评估框架评估了所有这些模型，并确保它们共享相同的评估设置。总体而言，DeepSeek-V2 仅激活了 21B 参数，但在几乎所有基准上都显著优于 DeepSeek 67B，并在开源模型中达到了顶级性能。

进一步，我们详细比较了 DeepSeek-V2 与其他开源模型的表现：

1. 与 Qwen1.5 72B 的比较：Qwen1.5 72B 是另一个支持中文和英文的模型。DeepSeek-V2 在大多数英语、代码和数学基准上表现出压倒性优势。在中文基准上，Qwen1.5 72B 在多学科选择题任务上表现更好，而 DeepSeek-V2 在其他任务上表现相当或更好。需要注意的是，对于 CHID 基准，Qwen1.5 72B 的分词器在我们的评估框架中会遇到错误，因此我们未记录 Qwen1.5 72B 的 CHID 分数。

2. 与 Mixtral 8x22B 的比较：DeepSeek-V2 在英语基准上表现相当或更好，除了与英语常识知识密切相关的 TriviaQA、NaturalQuestions 和 HellaSwag。值得注意的是，DeepSeek-V2 在 MMLU 上优于 Mixtral 8x22B。在代码和数学基准上，DeepSeek-V2 与 Mixtral 8x22B 表现相当。由于 Mixtral 8x22B 并未专门针对中文数据进行训练，其中文能力远不及 DeepSeek-V2。

3. 与 LLaMA3 70B 的比较：DeepSeek-V2 的训练数据量不到 LLaMA3 70B 的四分之一。因此，我们承认 DeepSeek-V2 在基础英语能力上仍与 LLaMA3 70B 存在轻微差距。然而，即使训练数据和激活参数少得多，DeepSeek-V2 在代码和数学能力上仍与 LLaMA3 70B 相当。此外，作为双语模型，DeepSeek-V2 在中文基准上显著优于 LLaMA3 70B。

最后，值得一提的是，某些先前的研究 (Hu et al., 2024) 在预训练阶段引入了 SFT 数据，而 DeepSeek-V2 在预训练期间从未接触过 SFT 数据。

3.2.3. 训练和推理效率

训练成本：由于 DeepSeek-V2 为每个 token 激活的参数较少，且所需的 FLOPs 少于 DeepSeek 67B，理论上训练 DeepSeek-V2 比训练 DeepSeek 67B 更经济。尽管训练 MoE 模型会引入额外的通信开销，但通过我们的操作符和通信优化，DeepSeek-V2 的训练可以达到相对较高的模型 FLOPs 利用率（MFU）。在实际的 H800 集群训练中，每训练一万亿 token，DeepSeek 67B 需要 300.6K GPU 小时，而 DeepSeek-V2 仅需 172.8K GPU 小时，即稀疏的 DeepSeek-V2 比密集的 DeepSeek 67B 节省了 42.5% 的训练成本。

推理效率：为了高效部署 DeepSeek-V2 提供服务，我们首先将其参数转换为 FP8 精度。此外，我们还对 DeepSeek-V2 进行了 KV 缓存量化 (Hooper et al., 2024; Zhao et al., 2023)，进一步将其 KV 缓存中的每个元素平均压缩到 6 位。得益于 MLA 和这些优化，实际部署的 DeepSeek-V2 所需的 KV 缓存显著少于 DeepSeek 67B，因此可以支持更大的批量大小。我们基于实际部署的 DeepSeek 67B 服务的提示和生成长度分布，评估了 DeepSeek-V2 的生成吞吐量。在单个 8 H800 GPU 节点上，DeepSeek-V2 的生成吞吐量超过每秒 50K token，是 DeepSeek 67B 最大生成吞吐量的 5.76 倍。此外，DeepSeek-V2 的提示输入吞吐量超过每秒 100K token。

4. 对齐

4.1. 监督微调（SFT）

基于我们之前的研究（DeepSeek-AI, 2024），我们构建了包含 150 万条实例的指令微调数据集，其中 120 万条用于提升模型的有用性，30 万条用于提升安全性。与初始版本相比，我们提高了数据质量，以减少幻觉响应并增强写作能力。我们对 DeepSeek-V2 进行了 2 轮微调，学习率设置为 $5 \times 10^{-6}$。

对于 DeepSeek-V2 Chat（SFT）的评估，我们主要采用基于生成的基准测试，除了几个具有代表性的选择题任务（如 MMLU 和 ARC）。我们还对 DeepSeek-V2 Chat（SFT）进行了指令跟随评估（IFEval）（Zhou et al., 2023），使用提示级别的宽松准确率作为指标。此外，我们使用 2023 年 9 月 1 日至 2024 年 4 月 1 日的 LiveCodeBench（Jain et al., 2024）问题来评估聊天模型。除了标准基准测试外，我们还在开放式对话基准测试上进一步评估了模型，包括 MT-Bench（Zheng et al., 2023）、AlpacaEval 2.0（Dubois et al., 2024）和 AlignBench（Liu et al., 2023）。为了进行比较，我们还在我们的评估框架和设置中评估了 Qwen1.5 72B Chat、LLaMA-3-70B Instruct 和 Mistral-8x22B Instruct。对于 DeepSeek 67B Chat，我们直接参考了之前发布的评估结果。

4.2. 强化学习（RL）

为了进一步释放 DeepSeek-V2 的潜力并使其与人类偏好对齐，我们进行了强化学习（RL）以调整其偏好。

强化学习算法：为了节省 RL 的训练成本，我们采用了组相对策略优化（GRPO）（Shao et al., 2024），该方法省略了通常与策略模型大小相同的评论模型，而是从组分数中估计基线。具体来说，对于每个问题 $q$，GRPO 从旧策略 $\pi_{\theta_{\text{old}}}$ 中采样一组输出 ${o_1, o_2, \dots, o_G}$，然后通过最大化以下目标来优化策略模型 $\pi_{\theta}$：

\[J_{\text{GRPO}}(\theta) = \mathbb{E}\left[q \sim P(Q), \{o_i\}_{i=1}^G \sim \pi_{\theta_{\text{old}}}(O|q)\right] \frac{1}{G} \sum_{i=1}^G \min\left(\frac{\pi_{\theta}(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)} A_i, \text{clip}\left(\frac{\pi_{\theta}(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}, 1-\epsilon, 1+\epsilon\right) A_i\right) - \beta D_{\text{KL}}(\pi_{\theta} || \pi_{\text{ref}}),\]

其中：

• $\epsilon$ 和 $\beta$ 是超参数，
• $A_i$ 是优势值，通过每组输出对应的奖励 ${r_1, r_2, \dots, r_G}$ 计算：

\[A_i = \frac{r_i - \text{mean}(\{r_1, r_2, \dots, r_G\})}{\text{std}(\{r_1, r_2, \dots, r_G\})}.\]

训练策略：在初步实验中，我们发现针对推理数据（如代码和数学提示）的 RL 训练表现出与通用数据训练不同的独特特性。例如，模型的数学和编码能力可以在更长的训练步骤中持续提升。因此，我们采用了两阶段 RL 训练策略：首先进行推理对齐，然后进行人类偏好对齐。在第一阶段，我们训练了一个用于代码和数学推理任务的奖励模型 $RM_{\text{reasoning}}$，并使用其反馈优化策略模型：

\[r_i = RM_{\text{reasoning}}(o_i).\]

在第二阶段，我们采用多奖励框架，从有用性奖励模型 $RM_{\text{helpful}}$、安全性奖励模型 $RM_{\text{safety}}$ 和基于规则的奖励模型 $RM_{\text{rule}}$ 中获取奖励。最终响应 $o_i$ 的奖励为：

\[r_i = c_1 \cdot RM_{\text{helpful}}(o_i) + c_2 \cdot RM_{\text{safety}}(o_i) + c_3 \cdot RM_{\text{rule}}(o_i),\]

其中：

• $c_1, c_2, c_3$ 是对应的系数。

为了获得可靠的奖励模型，我们精心收集了偏好数据，并进行了严格的质量过滤和比例调整。我们基于编译器反馈获取代码偏好数据，基于真实标签获取数学偏好数据。对于奖励模型训练，我们使用 DeepSeek-V2 Chat（SFT）初始化奖励模型，并使用点对或对对的损失进行训练。在实验中，我们观察到 RL 训练能够充分挖掘和激活模型的潜力，使其能够从可能的响应中选择正确且令人满意的答案。

训练效率优化：在极大模型上进行 RL 训练对训练框架提出了高要求。我们实施了以下工程优化：

• （1）采用混合引擎，分别针对训练和推理采用不同的并行策略以提高 GPU 利用率；
• （2）利用 vLLM（Kwon et al., 2023）作为推理后端，加速推理速度；
• （3）精心设计模型卸载到 CPU 和加载回 GPU 的调度策略，以实现训练速度和内存消耗的近最优平衡。

4.3. 评估结果

标准基准测试评估：我们首先在标准基准测试上评估了 DeepSeek-V2 Chat（SFT）和 DeepSeek-V2 Chat（RL）。值得注意的是，DeepSeek-V2 Chat（SFT）在 GSM8K、MATH 和 HumanEval 评估中相比其基础版本有显著提升，这归因于我们的 SFT 数据中包含大量数学和代码相关内容。此外，DeepSeek-V2 Chat（RL）进一步提升了数学和代码基准测试的表现。

与其他模型的比较中，DeepSeek-V2 Chat（SFT）在几乎所有英语、数学和代码基准测试上均优于 Qwen1.5 72B Chat。在中文基准测试上，DeepSeek-V2 Chat（SFT）在多学科选择题任务上略低于 Qwen1.5 72B Chat，与其基础版本的表现一致。与最先进的开源 MoE 模型 Mixtral 8x22B Instruct 相比，DeepSeek-V2 Chat（SFT）在大多数基准测试上表现更好，除了 NaturalQuestions 和 IFEval。与最先进的开源模型 LLaMA3 70B Chat 相比，DeepSeek-V2 Chat（SFT）在代码和数学相关基准测试上表现相似，LLaMA3 70B Chat 在 MMLU 和 IFEval 上表现更好，而 DeepSeek-V2 Chat（SFT）在中文任务上表现更强。最终，DeepSeek-V2 Chat（RL）在数学和编码任务上相比 DeepSeek-V2 Chat（SFT）进一步提升了性能。

开放式生成评估：我们在开放式对话基准测试上进一步评估了模型。对于英语开放式对话生成，我们使用 MT-Bench 和 AlpacaEval 2.0 作为基准测试。评估结果显示，DeepSeek-V2 Chat（RL）相比 DeepSeek-V2 Chat（SFT）有显著优势，展示了 RL 训练在实现更好对齐方面的有效性。与其他开源模型相比，DeepSeek-V2 Chat（RL）在 MT-Bench 和 AlpacaEval 2.0 上均优于 Mistral 8x22B Instruct 和 Qwen1.5 72B Chat。与 LLaMA3 70B Instruct 相比，DeepSeek-V2 Chat（RL）在 MT-Bench 上表现相当，在 AlpacaEval 2.0 上显著优于后者。

对于中文开放式生成能力，我们基于 AlignBench 进行了评估。结果显示，DeepSeek-V2 Chat（RL）相比 DeepSeek-V2 Chat（SFT）有轻微优势。值得注意的是，DeepSeek-V2 Chat（SFT）显著优于所有开源中文模型，在中文推理和语言任务上大幅领先第二好的开源模型 Qwen1.5 72B Chat。此外，DeepSeek-V2 Chat（SFT）和 DeepSeek-V2 Chat（RL）均优于 GPT-4-0613 和 ERNIEBot 4.0，巩固了我们的模型在支持中文的顶级 LLM 中的地位。

4.4. 讨论

SFT 数据量：关于是否需要大量 SFT 数据的讨论一直存在争议。先前的工作（Young et al., 2024; Zhou et al., 2024）认为少于 10K 条 SFT 数据足以产生令人满意的结果。然而，在我们的实验中，如果使用少于 10K 条数据，我们在 IFEval 基准测试上观察到显著的性能下降。可能的解释是，语言模型需要一定量的数据来发展特定技能。尽管随着模型规模的增加，所需数据量可能会减少，但不能完全消除。我们的观察强调了为 LLM 提供足够数据以具备所需能力的关键性。此外，SFT 数据的质量也至关重要，尤其是在涉及写作或开放式问题的任务中。

强化学习的对齐税：在人类偏好对齐过程中，我们观察到开放式生成基准测试的显著性能提升，无论是 AI 还是人类评估者的评分。然而，我们也注意到“对齐税”现象（Ouyang et al., 2022），即对齐过程可能会对某些标准基准测试（如 BBH）的性能产生负面影响。为了缓解对齐税，我们在 RL 阶段在数据处理和训练策略改进上付出了巨大努力，最终在标准基准测试和开放式基准测试之间实现了可接受的权衡。

在线强化学习：在我们的偏好对齐实验中，我们发现在线方法显著优于离线方法。因此，我们投入了大量精力实现了一个在线 RL 框架来对齐 DeepSeek-V2。关于在线或离线偏好对齐的结论可能因不同情境而异，我们将在未来工作中进行更深入的比较和分析。

5. 结论、局限性与未来工作

本文介绍了 DeepSeek-V2，一个支持 128K 上下文长度的大型 MoE 语言模型。除了强大的性能外，它还以经济高效的训练和推理为特点，得益于其创新的架构（包括 MLA 和 DeepSeekMoE）。在实际应用中，相比 DeepSeek 67B，DeepSeek-V2 在显著提升性能的同时，节省了 42.5% 的训练成本，减少了 93.3% 的 KV 缓存，并将最大生成吞吐量提升至 5.76 倍。评估结果表明，仅激活 21B 参数的 DeepSeek-V2 在开源模型中达到了顶级性能，成为最强的开源 MoE 模型。

DeepSeek-V2 及其聊天版本具有其他 LLM 常见的局限性，包括预训练后缺乏持续知识更新、可能生成未经核实的信息（如未经验证的建议）以及可能产生幻觉。此外，由于我们的数据主要由中文和英文内容组成，模型在其他语言上的表现可能有限。在中文和英文以外的场景中，应谨慎使用。

DeepSeek 将持续投资于开源大模型，致力于逐步接近通用人工智能的目标。我们正在探索进一步扩展 MoE 模型的方法，同时保持经济高效的训练和推理成本。我们的下一步目标是在即将发布的版本中实现与 GPT-4 相当的性能。我们的对齐团队不断努力提升模型，旨在开发一个不仅有用而且诚实、安全的模型，最终目标是使模型的价值与人类价值观对齐，同时尽量减少人类监督的需求。通过优先考虑伦理和负责任的发展，我们致力于为社会创造积极和有益的影响。目前，DeepSeek-V2 仅支持文本模态。在未来计划中，我们打算使模型支持多模态，增强其在不同场景中的多功能性和实用性。

#

https://arxiv.org/pdf/2412.19437

追一科技在《ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING》提出了RoPE的方法：

摘要

位置编码最近在Transformer架构中显示出其有效性。它为序列中不同位置元素之间的依赖关系建模提供了有价值的监督。在本文中，我们：

• 首先研究了将位置信息整合到基于Transformer的语言模型学习过程中的各种方法。
• 接着，我们提出了一种名为旋转位置嵌入（Rotary Position Embedding, RoPE）的新方法，以有效利用位置信息。具体来说，所提出的RoPE通过旋转矩阵对绝对位置进行编码，同时将显式的相对位置依赖关系引入自注意力机制中。值得注意的是，RoPE具有一些有价值的特性，包括：序列长度的灵活性、随着相对距离增加而衰减的token间依赖性，以及为线性自注意力配备相对位置编码的能力。
• 最后，我们在多种长文本分类基准数据集上评估了增强后的Transformer（称为RoFormer）的性能。

实验结果表明，RoFormer始终优于其他替代方法。此外，我们还提供了理论分析来解释一些实验结果。RoFormer已经集成到Huggingface中：https://huggingface.co/docs/transformers/model_doc/roformer。

1 引言

词序对于自然语言理解具有重要价值。

• 基于循环神经网络（RNN）的模型通过沿时间维度递归计算隐藏状态来编码词的顺序。
• 基于卷积神经网络（CNN）的模型（如Gehring等[2017]）通常被认为是与位置无关的，但最近的研究（Islam等[2020]）表明，常用的填充操作可以隐式地学习位置信息。
• 近年来，基于Transformer（Vaswani等[2017]）构建的预训练语言模型（PLMs）在各种自然语言处理（NLP）任务中取得了最先进的性能，包括上下文表示学习（Devlin等[2019]）、机器翻译（Vaswani等[2017]）和语言建模（Radford等[2019]）等。

与基于RNN和CNN的模型不同，PLMs利用自注意力机制从语义上捕捉给定语料的上下文表示。因此，PLMs在并行化方面相比RNN取得了显著改进，并且与CNN相比，能够更好地建模更长的token内关系。

值得注意的是，当前PLMs的自注意力架构被证明是与位置无关的（Yun等[2020]）。基于这一观点，研究者们提出了多种方法将位置信息编码到学习过程中。

• 一方面，通*过预定义函数生成的绝对位置编码（Vaswani等[2017]）被添加到上下文表示中，而可训练的绝对位置编码（Gehring等[2017]、Devlin等[2019]、Lan等[2020]、Clark等[2020]、Radford等[2019]、Radford和Narasimhan[2018]）也被广泛使用。
• 另一方面，之前的工作（Parikh等[2016]、Shaw等[2018]、Huang等[2018]、Dai等[2019]、Yang等[2019]、Raffel等[2020]、Ke等[2020]、He等[2020]、Huang等[2020]）主要集中在相对位置编码上，通常将相对位置信息编码到注意力机制中。
• 此外，Liu等[2020]的作者提出从神经微分方程（Neural ODE，Chen等[2018a]的角度建模位置编码的依赖性，Wang等[2020]的作者提出在复数空间中建模位置信息。

尽管这些方法有效，但它们通常将位置信息添加到上下文表示中，因此不适合线性自注意力架构。

在本文中，我们提出了一种新方法，即旋转位置嵌入（Rotary Position Embedding, RoPE），将位置信息引入PLMs的学习过程中。具体来说，RoPE通过旋转矩阵对绝对位置进行编码，同时将显式的相对位置依赖关系引入自注意力机制中。值得注意的是，RoPE具有一些优于现有方法的特性，包括序列长度的灵活性、随着相对距离增加而衰减的token间依赖性，以及为线性自注意力配备相对位置编码的能力。在多种长文本分类基准数据集上的实验结果表明，增强后的Transformer（称为RoFormer）相比基线方法能够取得更好的性能，从而证明了RoPE的有效性。

简而言之，我们的贡献如下：

• 我们研究了现有的相对位置编码方法，发现它们大多基于将位置编码添加到上下文表示的分解思想。我们提出了一种新方法，即旋转位置嵌入（RoPE），将位置信息引入PLMs的学习过程中。其核心思想是通过将上下文表示与旋转矩阵相乘来编码相对位置，并具有清晰的理论解释。
• 我们研究了RoPE的特性，并表明其随着相对距离的增加而衰减，这符合自然语言编码的需求。我们还指出，之前的相对位置编码方法与线性自注意力不兼容。
• 我们在多种长文本基准数据集上评估了提出的RoFormer。实验结果表明，RoFormer始终优于其他替代方法。一些预训练语言模型的实验代码已在GitHub上开源：https://github.com/ZhuiyiTechnology/roformer。

本文的其余部分组织如下：在第2节中，我们建立了自注意力架构中位置编码问题的形式化描述，并回顾了之前的工作；在第3节中，我们描述了旋转位置编码（RoPE）并研究了其特性；在第4节中，我们报告了实验结果；最后，在第5节中对本文进行了总结。

2 背景与相关工作

2.1 初步知识

设 $S_N = \lbrace w_i \rbrace_{i=1}^N$ 为一个包含 N 个输入token的序列，其中：

• $ w_i $ 是第 $ i $ 个元素
• $ S_N $ 对应的词嵌入表示为 $ E_N = \lbrace x_i \rbrace_{i=1}^N $，其中 $ x_i \in R^d $ 是第i个token的d维词嵌入向量（不包含位置信息）。

自注意力机制首先将位置信息融入词嵌入中，并将其转换为查询（query）、键（key）和值（value）表示：

\[\begin{aligned} q_m &= f_q(x_m, m) \\ k_n &= f_k(x_n, n) \\ v_n &= f_v(x_n, n), \end{aligned}\]

…(1)

其中:

• $ q_m $、$ k_n $ 和 $ v_n $ 分别通过函数 $ f_q $、$ f_k $ 和 $ f_v $ 融入了第 $ m $ 和第 $ n $ 个位置的信息。

查询和键值用于计算注意力权重，而输出则是对值表示的加权和：

\[\begin{aligned} a_{m,n} &= \frac{\exp\left(\frac{q_m^\intercal k_n}{\sqrt{d}}\right)}{\sum_{j=1}^N \exp\left(\frac{q_m^\intercal k_j}{\sqrt{d}}\right)} \\ o_m &= \sum_{n=1}^N a_{m,n} v_n. \end{aligned}\]

…(2)

现有的基于Transformer的位置编码方法主要集中在选择合适的函数来构建公式（1）。

2.2 绝对位置嵌入

公式（1）的一个典型选择是：

\[f_{t: t \in \{q, k, v\}}(x_i, i) := W_{t: t \in \{q, k, v\}}(x_i + p_i),\]

…(3)

其中：

• $ p_i \in \mathbb{R}^d $ 是一个依赖于token $ x_i $ 位置的 $ d $ 维向量。

之前的工作（Devlin等[2019]、Lan等[2020]、Clark等[2020]、Radford等[2019]、Radford和Narasimhan[2018]）引入了一组可训练的向量 $ p_i \in {p_t}_{t=1}^L $，其中 $ L $ 是最大序列长度。Vaswani等[2017]的作者提出使用正弦函数生成 $ p_i $：

\[\begin{aligned} p_{i,2t} &= \sin\left(\frac{k}{10000^{2t/d}}\right) \\ p_{i,2t+1} &= \cos\left(\frac{k}{10000^{2t/d}}\right), \end{aligned}\]

…(4)

其中：

• $ p_{i,2t} $ 是 $ d $ 维向量 $ p_i $ 的第 $ 2t $ 个元素。

在下一节中，我们将展示我们提出的RoPE与这种正弦函数直觉相关。然而，RoPE不是直接将位置信息添加到上下文表示中，而是通过与正弦函数相乘来融入相对位置信息。

2.3 相对位置嵌入

Shaw等[2018]的作者对公式（1）应用了以下不同的设置：

\[\begin{aligned} f_q(x_m) &:= W_q x_m \\ f_k(x_n, n) &:= W_k (x_n + \widetilde{p}^k_r) \\ f_v(x_n, n) &:= W_v (x_n + \widetilde{p}^v_r), \end{aligned}\]

…(5)

其中：

• $ \widetilde{p}^k_r $、$ \widetilde{p}^v_r \in \mathbb{R}^d $ 是可训练的相对位置嵌入。注意，$ r = \text{clip}(m - n, r_{\text{min}}, r_{\text{max}}) $ 表示位置 $ m $ 和 $ n $ 之间的相对距离。

他们通过假设超出一定距离的相对位置信息无用，对相对距离进行了裁剪。Dai等[2019]的作者在保持公式（3）的形式下，提出将公式（2）中的 $ q_m^\intercal k_n $ 分解为：

\[q_m^\intercal k_n = x_m^\intercal W_q^\intercal W_k x_n + x_m^\intercal W_q^\intercal W_k p_n + p_m^\intercal W_q^\intercal W_k x_n + p_m^\intercal W_q^\intercal W_k p_n,\]

…(6)

其核心思想是：将绝对位置嵌入 $ p_n $ 替换为其正弦编码的相对对应项 $ \widetilde{p}_{m-n} $，同时将第三和第四项中的绝对位置 $ p_m $ 替换为两个与查询位置无关的可训练向量 $ u $ 和 $ v $。此外，$ W_k $ 被区分为基于内容的键向量 $ x_n $ 和基于位置的键向量 $ p_n $，分别表示为 $ W_k $ 和 $ \widetilde{W}_k $，从而得到：

\[q_m^\intercal k_n = x_m^\intercal W_q^\intercal W_k x_n + x_m^\intercal W_q^\intercal \widetilde{W}_k \widetilde{p}_{m-n} + u^\intercal W_q^\intercal W_k x_n + v^\intercal W_q^\intercal \widetilde{W}_k \widetilde{p}_{m-n}.\]

…(7)

值得注意的是，值项中的位置信息通过设置 $ f_v(x_j) := W_v x_j $ 被移除。后续工作（Raffel等[2020]、He等[2020]、Ke等[2020]、Huang等[2020]）遵循了这些设置，仅将相对位置信息编码到注意力权重中。然而，Raffel等[2020]的作者将公式（6）重新表述为：

\[q_m^\intercal k_n = x_m^\intercal W_q^\intercal W_k x_n + b_{i,j},\]

…(8)

其中:

• $ b_{i,j} $ 是一个可训练的偏置项。

Ke等[2020]的作者研究了公式（6）中的中间两项，发现绝对位置与词之间的相关性较弱。Raffel等[2020]的作者提出使用不同的投影矩阵对词或位置进行建模：

\[q_m^\intercal k_n = x_m^\intercal W_q^\intercal W_k x_n + p_m^\intercal U_q^\intercal U_k p_n + b_{i,j}.\]

…(9)

He等[2020]的作者认为，两个token的相对位置只能通过公式（6）中的中间两项完全建模。因此，绝对位置嵌入 $ p_m $ 和 $ p_n $ 被简单地替换为相对位置嵌入 $ \widetilde{p}_{m-n} $：

\[q_m^\intercal k_n = x_m^\intercal W_q^\intercal W_k x_n + x_m^\intercal W_q^\intercal W_k \widetilde{p}_{m-n} + \widetilde{p}_{m-n}^\intercal W_q^\intercal W_k x_n.\]

…(10)

对四种相对位置嵌入变体的比较（Radford和Narasimhan[2018]）表明，类似于公式（10）的变体在其他三种中效率最高。总的来说，所有这些方法都试图在自注意力设置下基于公式（3）的分解来修改公式（6），这是Vaswani等[2017]最初提出的。它们通常直接将位置信息添加到上下文表示中。与之不同，我们的方法旨在在某些约束下从公式（1）推导出相对位置编码。接下来，我们将展示通过旋转上下文表示融入相对位置信息，推导出的方法更具可解释性。

3 提出的方法

在本节中，我们讨论提出的旋转位置嵌入（Rotary Position Embedding, RoPE）。我们首先在第3.1节中形式化相对位置编码问题，然后在第3.2节中推导RoPE，并在第3.3节中研究其特性。

3.1 形式化

基于Transformer的语言建模通常通过自注意力机制利用各个token的位置信息。如公式（2）所示，$ q_m^\intercal k_n $ 通常用于在不同位置的token之间传递知识。为了融入相对位置信息，我们要求查询 $ q_m $ 和键 $ k_n $ 的内积由一个函数 $ g $ 表示，该函数仅以词嵌入 $ x_m $、$ x_n $ 及其相对位置 $ m - n $ 作为输入变量。换句话说，我们希望内积仅以相对形式编码位置信息：

\[\langle f_q(x_m, m), f_k(x_n, n) \rangle = g(x_m, x_n, m - n).\]

…(11)

最终目标是：找到一种等效的编码机制来解决函数 $ f_q(x_m, m) $ 和 $ f_k(x_n, n) $，以符合上述关系。

3.2 旋转位置嵌入

3.2.1 二维情况

我们从维度 $ d = 2 $ 的简单情况开始。在这些设置下，我们利用二维平面上向量的几何性质及其复数形式来证明（详见第3.4.1节），公式（11）的一个解为：

\[\begin{aligned} f_q(x_m, m) &= (W_q x_m) e^{i m \theta} \\ f_k(x_n, n) &= (W_k x_n) e^{i n \theta} \\ g(x_m, x_n, m - n) &= \text{Re}\left[(W_q x_m) (W_k x_n)^* e^{i (m - n) \theta}\right], \end{aligned}\]

…(12)

其中：

• $\text{Re}[\cdot]$ 表示复数的实部，$(W_k x_n)^*$ 表示 $(W_k x_n)$ 的共轭复数。
• $\theta \in \mathbb{R}$ 是一个预设的非零常数。

我们可以进一步将 $ f_{{q, k}} $ 写成矩阵乘法形式：

\[f_{\{q, k\}}(x_m, m) = \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} W_{\lbrace,q,k\rbrace}^{11} & W_{\lbrace,q,k\rbrace}^{12} \\ W_{\lbrace,q,k\rbrace}^{21} & W_{\lbrace,q,k\rbrace}^{22} \end{pmatrix} \begin{pmatrix} x_m^{(1)} \\ x_m^{(2)} \end{pmatrix},\]

…(13)

其中：

• $(x_m^{(1)}, x_m^{(2)})$ 是 $ x_m $ 在二维坐标中的表示。
• 类似地，$ g $ 可以视为一个矩阵，从而在二维情况下解决第3.1节中的形式化问题。

具体来说，融入相对位置嵌入非常简单：只需将仿射变换后的词嵌入向量旋转其位置索引的角度倍数，从而解释旋转位置嵌入的直觉。

3.2.2 一般形式

为了将二维结果推广到任意 $ x_i \in \mathbb{R}^d $（其中 $ d $ 为偶数），我们将 $ d $ 维空间划分为 $ d/2 $ 个子空间，并利用内积的线性性将它们组合起来，将 $ f_{{q, k}} $ 转化为：

\[f_{\{q, k\}}(x_m, m) = R^d_{\Theta, m} W_{\{q, k\}} x_m,\]

…(14)

其中:

\[R^d_{\Theta, m} = \begin{pmatrix} \cos m\theta_1 & -\sin m\theta_1 & 0 & 0 & \cdots & 0 & 0 \\ \sin m\theta_1 & \cos m\theta_1 & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos m\theta_2 & -\sin m\theta_2 & \cdots & 0 & 0 \\ 0 & 0 & \sin m\theta_2 & \cos m\theta_2 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \cdots & \cos m\theta_{d/2} & -\sin m\theta_{d/2} \\ 0 & 0 & 0 & 0 & \cdots & \sin m\theta_{d/2} & \cos m\theta_{d/2} \end{pmatrix}\]

…(15)

是旋转矩阵，其预定义参数为 $\Theta = {\theta_i = 10000^{-2(i-1)/d}, i \in [1, 2, …, d/2]}$。图（1）展示了RoPE的图形化说明。将RoPE应用于公式（2）中的自注意力机制，我们得到：

\[q_m^\intercal k_n = (R^d_{\Theta, m} W_q x_m)^\intercal (R^d_{\Theta, n} W_k x_n) = x_m^\intercal W_q R^d_{\Theta, n-m} W_k x_n,\]

…(16)

其中:

• $ R^d_{\Theta, n-m} = (R^d_{\Theta, m})^\intercal R^d_{\Theta, n} $。
• 注意，$ R^d_{\Theta} $ 是一个正交矩阵，这确保了在编码位置信息过程中的稳定性。

此外，由于 $ R^d_{\Theta} $ 的稀疏性，直接应用矩阵乘法如公式（16）所示在计算上并不高效；我们在理论解释中提供了另一种实现方式。

图1 Rotary Position Embedding(RoPE)的实现

与之前工作中采用的加法性质的位置嵌入方法（即公式（3）到（10））不同，我们的方法是乘法的。此外，RoPE通过旋转矩阵乘积自然地融入了相对位置信息，而不是在应用于自注意力时修改加法位置编码的扩展公式中的项。

3.3 RoPE的特性

长期衰减性：遵循Vaswani等[2017]，我们设置 $\theta_i = 10000^{-2i/d}$。可以证明，这种设置提供了长期衰减特性（详见第3.4.3节），这意味着内积会随着相对位置的增加而衰减。这一特性与直觉一致，即具有较长相对距离的token对应该具有较少的连接。

RoPE与线性注意力：自注意力可以改写为更一般的形式：

\[\text{Attention}(Q, K, V)_m = \frac{\sum_{n=1}^N \text{sim}(q_m, k_n) v_n}{\sum_{n=1}^N \text{sim}(q_m, k_n)},\]

…(17)

其中原始自注意力选择 $\text{sim}(q_m, k_n) = \exp(q_m^\intercal k_n / \sqrt{d})$。注意，原始自注意力需要计算每对token的查询和键的内积，其复杂度为 $ O(N^2) $。遵循Katharopoulos等[2020]，线性注意力将公式（17）重新表述为：

\[\text{Attention}(Q, K, V)_m = \frac{\sum_{n=1}^N \phi(q_m)^\intercal \varphi(k_n) v_n}{\sum_{n=1}^N \phi(q_m)^\intercal \varphi(k_n)},\]

…(18)

其中 $\phi(\cdot)$ 和 $\varphi(\cdot)$ 通常是非负函数。Katharopoulos等[2020]的作者提出 $\phi(x) = \varphi(x) = \text{elu}(x) + 1$，并首先利用矩阵乘法的结合性计算键和值的乘积。Shen等[2021]使用softmax函数在内积之前分别对查询和键进行归一化，这等价于 $\phi(q_i) = \text{softmax}(q_i)$ 和 $\phi(k_j) = \exp(k_j)$。有关线性注意力的更多细节，我们鼓励读者参考原始论文。在本节中，我们重点讨论将RoPE与公式（18）结合。由于RoPE通过旋转注入位置信息，这保持了隐藏表示的范数不变，因此我们可以通过将旋转矩阵与非负函数的输出相乘来将RoPE与线性注意力结合：

\[\text{Attention}(Q, K, V)_m = \frac{\sum_{n=1}^N \left(R^d_{\Theta, m} \phi(q_m)\right)^\intercal \left(R^d_{\Theta, n} \varphi(k_n)\right) v_n}{\sum_{n=1}^N \phi(q_m)^\intercal \varphi(k_n)}.\]

…(19)

值得注意的是，我们保持分母不变以避免除以零的风险，而分子中的求和可能包含负项。尽管公式（19）中每个值 $ v_i $ 的权重并未严格概率归一化，但我们认为计算仍可以建模值的重要性。

图 2

3.4 理论解释

3.4.1 二维情况下RoPE的推导

在 $ d = 2 $ 的情况下，我们考虑两个词嵌入向量 $ x_q $ 和 $ x_k $，分别对应于查询和键，以及它们的位置 $ m $ 和 $ n $。根据公式（1），它们的位置编码对应为：

\[\begin{aligned} q_m &= f_q(x_q, m), \\ k_n &= f_k(x_k, n), \end{aligned}\]

…(20)

其中:

• $ q_m $ 和 $ k_n $ 的下标表示编码的位置信息。

假设存在一个函数 $ g $，定义了由 $ f_{{q, k}} $ 生成的向量之间的内积：

\[q_m^\intercal k_n = \langle f_q(x_m, m), f_k(x_n, n) \rangle = g(x_m, x_n, n - m),\]

…(21)

我们进一步要求以下初始条件成立：

\[\begin{aligned} q &= f_q(x_q, 0), \\ k &= f_k(x_k, 0), \end{aligned}\]

…(22)

这可以理解为未编码位置信息的向量。在这些设定下，我们尝试找到 $ f_q $ 和 $ f_k $ 的解。首先，我们利用二维向量的几何意义及其复数形式，将公式（20）和（21）中的函数分解为：

\[\begin{aligned} f_q(x_q, m) &= R_q(x_q, m) e^{i \Theta_q(x_q, m)}, \\ f_k(x_k, n) &= R_k(x_k, n) e^{i \Theta_k(x_k, n)}, \\ g(x_q, x_k, n - m) &= R_g(x_q, x_k, n - m) e^{i \Theta_g(x_q, x_k, n - m)}, \end{aligned}\]

…(23)

其中：

• $ R_f $、$ R_g $ 和 $ \Theta_f $、$ \Theta_g $ 分别是 $ f_{{q, k}} $ 和 $ g $ 的径向和角度分量。将它们代入公式（21），我们得到以下关系：

\[\begin{aligned} R_q(x_q, m) R_k(x_k, n) &= R_g(x_q, x_k, n - m), \\ \Theta_k(x_k, n) - \Theta_q(x_q, m) &= \Theta_g(x_q, x_k, n - m), \end{aligned}\]

…(24)

对应的初始条件为：

\[\begin{aligned} q &= \|q\| e^{i \theta_q} = R_q(x_q, 0) e^{i \Theta_q(x_q, 0)}, \\ k &= \|k\| e^{i \theta_k} = R_k(x_k, 0) e^{i \Theta_k(x_k, 0)}, \end{aligned}\]

…(25)

其中：

• $ |q| $、$ |k| $ 和 $ \theta_q $、$ \theta_k $ 分别是 $ q $ 和 $ k $ 在二维平面上的径向和角度分量。

接下来，我们在公式（24）中设 $ m = n $，并考虑公式（25）中的初始条件：

\[\begin{aligned} R_q(x_q, m) R_k(x_k, m) &= R_g(x_q, x_k, 0) = R_q(x_q, 0) R_k(x_k, 0) = \|q\| \|k\|, \\ \Theta_k(x_k, m) - \Theta_q(x_q, m) &= \Theta_g(x_q, x_k, 0) = \Theta_k(x_k, 0) - \Theta_q(x_q, 0) = \theta_k - \theta_q. \end{aligned}\]

…(26)

一方面，从公式（26a）可以直接得到 $ R_f $ 的解：

\[\begin{aligned} R_q(x_q, m) &= R_q(x_q, 0) = \|q\|, \\ R_k(x_k, n) &= R_k(x_k, 0) = \|k\|, \\ R_g(x_q, x_k, n - m) &= R_g(x_q, x_k, 0) = \|q\| \|k\|, \end{aligned}\]

…(27)

这表明径向函数 $ R_q $、$ R_k $ 和 $ R_g $ 与位置信息无关。另一方面，从公式（26b）可以看出，$ \Theta_q(x_q, m) - \theta_q = \Theta_k(x_k, m) - \theta_k $ 表明角度函数不依赖于查询和键，我们设 $ \Theta_f := \Theta_q = \Theta_k $，并将 $ \Theta_f(x_{{q, k}}, m) - \theta_{{q, k}} $ 表示为位置 $ m $ 的函数，记为 $ \phi(m) $，从而得到：

\[\Theta_f(x_{\{q, k\}}, m) = \phi(m) + \theta_{\{q, k\}}.\]

…(28)

进一步，将 $ n = m + 1 $ 代入公式（24）并考虑上述方程，我们得到：

\[\phi(m + 1) - \phi(m) = \Theta_g(x_q, x_k, 1) + \theta_q - \theta_k,\]

…(29)

由于右边是一个与 $ m $ 无关的常数，$ \phi(m) $ 在连续整数输入下形成一个等差数列：

\[\phi(m) = m \theta + \gamma,\]

…(30)

其中 $ \theta, \gamma \in \mathbb{R} $ 是常数，且 $ \theta $ 非零。总结从公式（27）到（30）的解：

\[\begin{aligned} f_q(x_q, m) &= \|q\| e^{i (\theta_q + m \theta + \gamma)} = q e^{i (m \theta + \gamma)}, \\ f_k(x_k, n) &= \|k\| e^{i (\theta_k + n \theta + \gamma)} = k e^{i (n \theta + \gamma)}. \end{aligned}\]

…(31)

注意，我们没有对公式（22）中的 $ f_q $ 和 $ f_k $ 施加任何约束，因此 $ f_q(x_m, 0) $ 和 $ f_k(x_n, 0) $ 可以自由选择。为了使我们的结果与公式（3）可比，我们定义：

\[\begin{aligned} q &= f_q(x_m, 0) = W_q x_n, \\ k &= f_k(x_n, 0) = W_k x_n. \end{aligned}\]

…(32)

然后，我们在最终解中简单地设 $ \gamma = 0 $：

\[\begin{aligned} f_q(x_m, m) &= (W_q x_m) e^{i m \theta}, \\ f_k(x_n, n) &= (W_k x_n) e^{i n \theta}. \end{aligned}\]

…(33)

3.4.2 旋转矩阵乘法的计算高效实现

利用公式（15）中 $ R^d_{\Theta, m} $ 的稀疏性，$ R^d_{\Theta} $ 与 $ x \in \mathbb{R}^d $ 的乘法可以更高效地实现为：

\[R^d_{\Theta, m} x = \begin{pmatrix} x_1 \\ x_2 \\ x_3 \\ x_4 \\ \vdots \\ x_{d-1} \\ x_d \end{pmatrix} \otimes \begin{pmatrix} \cos m \theta_1 \\ \cos m \theta_1 \\ \cos m \theta_2 \\ \cos m \theta_2 \\ \vdots \\ \cos m \theta_{d/2} \\ \cos m \theta_{d/2} \end{pmatrix} + \begin{pmatrix} -x_2 \\ x_1 \\ -x_4 \\ x_3 \\ \vdots \\ -x_d \\ x_{d-1} \end{pmatrix} \otimes \begin{pmatrix} \sin m \theta_1 \\ \sin m \theta_1 \\ \sin m \theta_2 \\ \sin m \theta_2 \\ \vdots \\ \sin m \theta_{d/2} \\ \sin m \theta_{d/2} \end{pmatrix}.\]

…(34)

3.4.3 RoPE的长期衰减性

我们可以将向量 $ q = W_q x_m $ 和 $ k = W_k x_n $ 的条目成对分组，公式（16）中的RoPE内积可以写成复数乘法的形式：

\[(R^d_{\Theta, m} W_q x_m)^\intercal (R^d_{\Theta, n} W_k x_n) = \text{Re} \left[ \sum_{i=0}^{d/2-1} q_{[2i:2i+1]} k^*_{[2i:2i+1]} e^{i (m - n) \theta_i} \right],\]

…(35)

其中:

• $ q_{[2i:2i+1]} $ 表示 $ q $ 的第 $ 2i $ 到 $ 2i+1 $ 个条目。

记 $ h_i = q_{[2i:2i+1]} k^*{[2i:2i+1]}, S_j = \sum{i=0}^{j-1} e^{i (m - n) \theta_i} $，并设 $ h_{d/2} = 0 $ 和 $ S_0 = 0 $，我们可以使用Abel变换将求和重写为：

\[\sum_{i=0}^{d/2-1} q_{[2i:2i+1]} k^*_{[2i:2i+1]} e^{i (m - n) \theta_i} = \sum_{i=0}^{d/2-1} h_i (S_{i+1} - S_i) = -\sum_{i=0}^{d/2-1} S_{i+1} (h_{i+1} - h_i).\]

…(36)

因此，

\[\left| \sum_{i=0}^{d/2-1} q_{[2i:2i+1]} k^*_{[2i:2i+1]} e^{i (m - n) \theta_i} \right| = \left| \sum_{i=0}^{d/2-1} S_{i+1} (h_{i+1} - h_i) \right| \leq \sum_{i=0}^{d/2-1} |S_{i+1}| |h_{i+1} - h_i| \leq \left( \max_i |h_{i+1} - h_i| \right) \sum_{i=0}^{d/2-1} |S_{i+1}|.\]

…(37)

注意到，通过设置 $ \theta_i = 10000^{-2i/d} $，$ \frac{1}{d/2} \sum_{i=1}^{d/2} \mid S_i \mid $ 的值会随着相对距离 $ m - n $ 的增加而衰减，如图（2）所示。

4.实验

略

#

https://arxiv.org/pdf/2104.09864

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