[ArXiv'25] RankMixer: Scaling Up Ranking Models in Industrial Recommenders

发表于 2025-12-17 分类于 recommender-system 阅读次数：本文字数： 8.3k 阅读时长 ≈ 30 分钟

#摘要

近年来, 大型语言模型 (LLM) 的进展激发了人们对 Scaling Up 推荐系统的兴趣, 但仍存在两个实际障碍.

首先, 工业级推荐系统的 training 和 serving 成本必须满足严格的延迟上限和高 QPS 要求.
其次, 排序模型中大多数人为设计的特征交叉模块都沿用自 CPU 时代, 无法充分利用现代 GPU, 导致模型浮点运算利用率 (MFU) 低, 可扩展性差.

我们提出了 RankMixer, 一种面向硬件的模型设计, 旨在构建统一且可扩展的特征交叉架构. RankMixer 保留了 Transformer 的高并行性, 同时用 Multi-Head Token Mixing 模块取代了二次 Self Attention 机制, 从而提高了效率.

此外, RankMixer 利用 Per-Token Feed-Forward Networks (PFFN) 同时实现了对不同特征子空间的建模和跨特征空间交叉. 我们进一步将其扩展到 1B 参数, 并提出了稀疏矩阵迭代 (Sparse-MoE) 变体, 以获得更高的 ROI.

我们还采用了一种动态路由策略来解决专家训练的不足和不平衡问题.

实验表明, RankMixer 在万亿级生产数据集上展现出卓越的扩展能力. 通过将之前分散的手工构建的低平均使用率 (MFU) 模块替换为 RankMixer, 我们将模型 MFU 从 4.5% 提升至 45%, 并在保持大致相同推理延迟的情况下, 将在线排序模型的参数规模扩大了两个数量级. 我们通过在两个核心应用场景 (推荐和广告) 中进行在线 A/B 测试, 验证了 RankMixer 的通用性. 最后, 我们推出了 1B 稠密参数 RankMixer, 可在不增加服务成本的情况下实现全流量服务. 这使得用户活跃天数提高了 0.3%, 应用内总使用时长提高了 1.08%.

#1. 引言

推荐系统 (Recommender System) 在信息分发过程中至关重要. 作为一项重要的机器学习应用, 推荐系统基于大量的多字段特征数据预测用户对物品的行为, 这些特征包括:

数值特征 (例如各种统计数据)
类别特征 (例如用户和物品 ID)
用户行为特征以及内容特征 (Zhang, 2021 ; Lin, 2023)

目前最先进的推荐方法基于深度学习推荐模型 (Deep Learning Recommendation Models, DLRM), 该模型灵活地利用神经网络在输入特征 Embedding 之上构建稠密交叉层, 从而捕捉特征间的交叉. DLRM 中的稠密交叉层对推荐系统的性能至关重要, 并且已经提出了多种模型结构 (Covington, 2016 ; Zhou, 2018 ; Guo, 2017 ; Wang, 2021 ; Zhang, [nd]).

随着大语言模型 (LLM) 的发展, 以及参数增加带来的性能提升 (Achiam, 2023 ; Kaplan, 2020 ; Hoffmann, 2022), 扩展 DLRM 以充分利用海量数据已成为迫切需求. 以往的研究已在 DLRM 扩展方面取得了诸多成果, 但早期研究 (Zhang, 2024 ; Ardalani, 2022 ; Chitlangia, 2023) 仅仅扩展或堆叠特征交叉层, 而没有改变模型结构. 这种方法带来的收益有限, 有时甚至会产生负面影响 (Wang, 2017 ; Lian, 2018). 随后, 诸如 DHEN (Zhang, 2022) 和 Wukong (Zhang, [未注明日期]) 等后续研究致力于设计创新的深度神经网络 (DNN) 结构以提升其扩展性能. 然而, 利用模型规模来提升推荐系统的性能面临着独特的实际挑战. 与 NLP/CV 任务不同, 工业级推荐系统必须严格遵守严格的延迟约束, 并支持极高的每秒查询数 (QPS). 因此, 核心挑战在于找到模型有效性和计算效率之间的最佳平衡点.

从历史上看, 推荐系统中排序模型的架构深受 CPU 时代设计原则的影响. 这些模型通常依赖于组合异构的、手工设计的跨特征模块来提取特征交叉作用, 但其许多核心算子在现代 GPU 上属于 Memory-bound 而非 Compute-bound, 导致 GPU 并行性差, 模型浮点运算利用率 (Model Flops Utilization, MFU) 极低, 通常只有个位数百分比. 此外, 由于 CPU 时代模型的计算成本与参数数量大致成正比, 因此, 根据 Scaling Laws 进行激进扩展所带来的潜在 ROI 在实践中难以落地.

总之, 对 DLRM Scaling Laws 的研究必须克服以下问题:

模型架构应对齐硬件特点, 最大限度地提高现代 GPU 上的 MFU 和计算吞吐量.
模型设计必须利用推荐系统数据的特点, 例如异构特征空间和数百个字段之间的个性化跨特征交叉.

为了应对这些挑战, 我们提出了一种硬件感知的模型设计方法 - RankMixer. RankMixer 的核心设计基于两个可扩展的组件:

Multi-head Token Mixing 仅通过无参数算子即可实现跨 Token 特征交叉. 该策略在性能和计算效率方面均优于 Self Attention 机制.
Per-token Feed-Forward Networks (PFFN) 显著扩展了模型容量, 并通过为不同的特征子空间建模分配独立参数来解决特征空间支配问题. 这些 FFN 也与推荐数据模式高度契合, 从而实现了更好的扩展性. 为了进一步提升大规模模型的 ROI, 我们将 Per-Token FFN 模块扩展为稀疏专家混合 (MoE) 结构. 通过针对不同数据动态激活每个 Token 的专家子集, 我们可以在计算成本增加极小的情况下显著提升模型容量.

RankMixer 采用类似于 Transformer 的高度并行架构, 但克服了基于 self-attention 的特征交叉的几个关键局限: 训练效率低、建模跨空间 ID 相似性时组合爆炸、以及注意力权重矩阵导致的严重 Memory-Bound. 同时, 与原版 Transformer 相比, RankMixer 在相同的 FLOPs 下提供了更大的模型容量和学习能力.

在抖音推荐系统的生产部署中, 我们验证了在推理延迟比基线模型更短的情况下, 模型参数规模可以扩大 100 倍以上. 这得益于 RankMixer 架构能够将参数增长与浮点运算 (FLOPs) 解耦, 并通过高 MFU 和工程优化将 FLOPs 增长与实际成本解耦.

主要贡献可概括如下:

提出了一种名为 RankMixer 的新型架构, 它遵循模型设计需要对齐硬件的理念. 我们设计了多头 Token 混合和 Per-Token FFN 策略, 以高效地捕获异构特征交叉, 并使用动态路由策略来提高 RankMixer 中 SparseMoE 的可扩展性.
利用高 MFU 和性能优化的优势, 我们在不增加推理成本的情况下, 将模型参数放大 70 倍, 包括改进 MFU 和量化.
我们进行了大量的线上和线下实验, 并在万亿级工业推荐数据集上研究了该模型的扩展规律. RankMixer 模型已成功部署在抖音信息流推荐排名中, 实现了全流量服务, 活跃天数和应用使用时长分别提升了 0.3% 和 1.08%.

#2. 相关工作

现代推荐系统基于深度学习推荐模型 (DLRM), 如何有效地对特征交叉进行建模是 DLRM 的关键因素 (Guo, 2017 ; Huang, 2022 ; Qu, 2018 ; Wang, 2024b, 2021 ; Zhang, 2022).

Wide & Deep (Cheng, 2016) 是最早的研究成果之一, 它结合了逻辑回归 (Wide 部分) 和深度神经网络 (DNN, Deep 部分), 分别用于捕捉低阶和高阶特征交叉.
DeepFM (Guo, 2017) 是另一项成果, 它集成了因子分解机 (FM) 和深度神经网络 (DNN).
此外, DeepCross (Shan, 2016) 是残差网络 (He, 2016) 的扩展, 旨在隐式地学习自动特征交叉. 然而, 仅仅依靠深度神经网络 (DNN) 来学习高阶特征交叉已被证明极具挑战性 (Qu, 2018; Rendle, 2020).

显式交叉方法设计了不同的算子来显式地捕获高阶特征交叉, 例如:

PNN (Qu, 2018)、DCN (Wang, 2017) 及其后续版本 DCNv2 (Wang, 2021)、xDeepFM (Lian, 2018)、FGCNN (Liu, 2019) 和 FiGNN (Li, 2010).
AutoInt (Song, 2019) 和 Hiformer (Gui, 2023) 采用带有残差连接的注意力机制来学习复杂的交叉.
DHEN (Zhang, 2022) 提出将多个交叉算子组合在一起.

尽管提高了准确率, 这些新架构增加了延迟和内存消耗, 而且模型规模也相对较小.

Scaling Law 已成为深度学习的一个基本主题, 也是过去十年中众多突破的关键催化剂, 尤其是在自然语言处理 (NLP) (Kaplan, 2020; Hoffmann, 2022)、计算机视觉 (CV) (Dosovitskiy, 2020; Zhai, 2022) 和多模态建模 (Radford, 2021; Ramesh, 2021) 等领域. 它描述了模型性能与缩放因子 (例如模型大小、数据量和计算能力) 之间的幂律相关性 (Kaplan, 2020; Achiam, 2023; Touvron, 2023; Team, 2024). 近年来, 推荐系统中的 Scaling Law 问题引起了研究人员的广泛关注 (Guo, 2024). 相关研究探索了用户活动序列预训练 (Chitlangia, 2023)、通用用户表示 (Shin, 2023 ; Zhang, 2024) 以及在线检索 (Fang, 2024 ; Wang, 2024c) 的 Scaling Law 策略. Wukong (Zhang, [nd]) 通过堆叠 FM 和 LCB 来学习特征交叉.

与之正交的方面, Zhang (2024) 将序列推荐模型扩展到了 0.8B 的参数. HSTU (Zhai, [nd]) 增强了生成式推荐器 (GR) 的扩展效果, 而 GR 更关注序列部分.

#3. 方法

图 1. RankMixer 模块的架构. 一个 RankMixer 块由两个模块组成:

Multi-head Token Mixing 模块: Token 混合模块将每个 Token 的 Embedding 分割成 $H$ 个更小的部分 (head), 然后将这些部分在不同 Token 之间重新组合, 从而创建新的混合 Token. 这使得来自不同特征的信息能够相互交叉.
基于 SMoE 的 Per-Token FFN 模块

#3.1 总体架构

RankMixer 的整体架构由 $T$ 个输入 Token 组成, 这些 Token 经由 $L$ 个连续的 RankMixer 模块处理, 最后由一个输出池化算子进行处理. 每个 RankMixer 模块包含两个主要组件:

Multi-Head Token Mixing 层
Per-Token Feed-Forward Network (PFFN) 层

首先, 输入向量 $\mathbf{e}_{\text{input}}$ 被 Tokenize 为 $T$ 个特征 Token $\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_T$ , 每个特征 Token 代表一个连贯的特征向量. RankMixer 模块通过以下方式迭代地细化第 $L$ 层的 Token 表示:

$\mathbf{S}_{n-1} = \text{LN}~(~\text{TokenMixing}~(~\mathbf{X}_{n-1}~) + \mathbf{X}_{n-1})$

$\mathbf{X}_n = \text{LN}~(~\text{PFFN}~(~\mathbf{S}_{n-1}~) + \mathbf{S}_{n-1})$

其中:

$\text{LN}(\cdot)$ 为层归一化函数.
$\text{TokenMixing}(\cdot)$ 和 $\text{PFFN}(\cdot)$ 分别为 Multi-Head Token Mixing 模块和 PFFN 模块.
$\mathbf{X}_n \in \mathbb{R}^{T \times D}$ 为第 $n$ 个 RankMixer 模块的输出, $\mathbf{X}_0 \in \mathbb{R}^{T \times D}$ 由 $\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_T$ 堆叠而成
$D$ 为模型的隐藏维度.
输出表示 $\mathbf{o}_{\text{output}}$ 由最终层表示 $\mathbf{X}_L$ 的均值池化得到, 用于计算不同任务的预测结果.

#3.2 输入层和特征分词

构建大规模推荐模型的第一步是准备包含丰富信息的输入数据, 例如

用户特征 (包括用户 ID 和其他的用户信息)
候选特征 (包括视频 ID, 作者 ID 等)
序列特征 (通过序列模块处理, Zhou, 2018; Chai, 2025 以捕捉时间兴趣, 生成 $\mathbf{e}_s$ )
交叉特征 (用户与候选对象之间的交叉特征)

所有特征最终将被转换为不同维度的 Embedding.

为了在后续阶段实现高效的并行计算, 不同维度的 Embedding 必须转换为维度对齐的向量, 称为特征 Token (feature-token). 我们将这种 Embedding 对齐过程称为 Tokenization. 最简单的策略是为每个特征分配一个 Embedding, 但考虑到通常有数百个特征, 这会带来诸多挑战. 数百个 Token 不可避免地会将每个 Token 的参数和计算量减少到很小的片段, 导致重要特征建模不足, GPU 核心利用率低下. 相反, Token 数量过少(例如只有一个 Token)会将模型结构简化为简单的深度神经网络 (DNN), 无法清晰地表示不同的特征空间, 从而导致主导特征掩盖其他特征.

为了克服这些问题, 我们提出了一种基于语义的 Tokenization 方法, 该方法利用领域知识将特征分组为若干语义上连续的簇. 这些簇依次拼接成一个大的 Embedding 向量 $\mathbf{e}_{\text{input}} = [\mathbf{e}_1; \mathbf{e}_2; \ldots; \mathbf{e}_N]$ , 随后再被分割成适当数量的、固定维度的特征 Token. 每个特征 Token $\mathbf{x}_i \in \mathbb{R}^D$ 捕获一组 Embedding，表示相似的语义方面.

$\mathbf{x}_i = \text{Proj}\big( \mathbf{e}_{\text{input}}[ d \cdot (i - 1) : d \cdot i ] \big), \quad i = 1, \ldots, T.$

其中 $\mathbf{e}_{\text{input}}$ 是拼接后的 Embedding, $d$ 是每个 Token 的固定维度, $N$ 是特征组的数量, $T$ 是结果 Token 计数, Proj 函数将分割的嵌入映射到 $D$ 维度.

#3.3 RankMixer 模块

#3.3.1 Multi-head Token Mixing

为了促进跨 Token 的有效信息交换(这对特征交叉和全局信息建模非常重要), 我们引入了 Multi-head Token Mixing 模块. 每个 Token 均分为 $H$ 个头, Token $\mathbf{x}_t$ 的第 $h$ 个头表示为 $\mathbf{x}_t^h$ :

$\big[ \mathbf{x}_t^{(1)} \,\|\, \mathbf{x}_t^{(2)} \,\|\, \ldots \,\|\, \mathbf{x}_t^{(H)} \big] = \text{SplitHead}\big( \mathbf{x}_t \big)$

这些头可以看作是 Token $\mathbf{x}_t$ 在低维特征子空间中的投影, 因为推荐任务需要考虑不同的视角.

Token Mixing 用于聚合这些子空间向量, 以实现全局特征交叉. 形式上, Multi-head Token Mixing 混合后, 对应于第 $h$ 个头部的第 $h$ 个 Token $\mathbf{s}_h$ 如下:

$\mathbf{s}_h = \text{Concat}\big( \mathbf{x}_1^{h}, \mathbf{x}_2^{h}, \ldots, \mathbf{x}_T^{h} \big)$

Multi-head Token Mixing 模块的输出为 $\mathbf{S} \in \mathbb{R}^{H \times T \times \frac{D}{H}}$ , 它由所有打乱顺序的 Token $\mathbf{s}_1, \mathbf{s}_2, \ldots, \mathbf{s}_H$ 堆叠而成. 本文中, 我们设置 $H = T$ , 以在 Multi-head Token Mixing 后保持剩余连接的 Token 数量不变.

经过残余连接和归一化模块后, 我们可以生成:

$\mathbf{s}_1, \mathbf{s}_2, \ldots, \mathbf{s}_T = \text{LN}\big( \text{TokenMixing}(\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_T) + (\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_T) \big).$

尽管 Self Attention 机制在大规模语言模型中已被证明非常有效, 但我们发现它对于推荐系统而言并非最优. 在 Self Attention 机制中, 注意力权重是通过 Token 的内积来计算的. 这种方法在自然语言处理(NLP)中效果良好, 因为所有 Token 共享一个统一的嵌入空间.

然而, 在推荐任务中, 特征空间本质上是异构的. 计算两个异构语义空间之间的内积相似度是出了名的困难, 尤其是在推荐系统中, 用户和物品端的特征 ID 空间可能包含数亿个元素.

因此, 将 Self Attention 机制应用于如此多样化的输入, 其性能并不优于无参数的 Multi-head Token Mixing 方法, 并且会消耗更多的计算资源, 内存 I/O 操作和 GPU 内存.

#3.3.2 Per-token FFN

以往的 DLRM 和 DHEN 模型倾向于将来自许多不同语义空间的特征混合在一个交叉模块中, 这可能导致高频字段占主导地位, 淹没低频或长尾信号, 最终损害整体推荐质量. 我们提出了一种参数隔离 (parameter-isolated) 的前馈网络架构, 称为 Per-token FFN.

在传统设计中, FFN 的参数在所有 Token 之间共享, 但我们的方法使用专用的转换来处理每个 Token, 从而隔离每个 Token 的参数. 对于第 $t$ 个 Token $\mathbf{s}_t$ , PFFN 可以表示为:

$\mathbf{v}_t = f_{\text{pffn}}^{t,2}( \text{Gelu}( f_{\text{pffn}}^{t,1}(\mathbf{s}_t) ) )$

其中,

$f_{\text{pffn}}^{t,i}(\mathbf{x}) = \mathbf{x} \mathbf{W}_{\text{pffn}}^{t,i} + \mathbf{b}_{\text{pffn}}^{t,i}$

是 PFFN 的第 $i$ 层 MLP. $\mathbf{W}_{\text{pffn}}^{t,1} \in \mathbb{R}^{D \times kD}$ , $\mathbf{b}_{\text{pffn}}^{t,1} \in \mathbb{R}^{kD}$ , $\mathbf{W}_{\text{pffn}}^{t,2} \in \mathbb{R}^{kD \times D}$ , $\mathbf{b}_{\text{pffn}}^{t,2} \in \mathbb{R}^{D}$ . $k$ 是一个超参数, 用于调整每个 Token 的 FFN 的隐藏维度. $\text{Gelu}(\cdot)$ 是 Gelu 激活函数. $\mathbf{s}_t \in \mathbb{R}^D$ 是第 $t$ 个 Token.

我们将 PFFN 模块总结如下:

$\begin{align} \mathbf{v}_1, \mathbf{v}_2, \ldots, \mathbf{v}_T &= \text{PFFN}\big( \mathbf{s}_1, \mathbf{s}_2, \ldots, \mathbf{s}_T \big). \\ &= f_{\text{pffn}}^{t,2}\!\left( \text{Gelu}\!\left( f_{\text{pffn}}^{t,1}\big( \mathbf{s}_1, \mathbf{s}_2, \ldots, \mathbf{s}_T \big) \right) \right). \end{align}$

与参数全部共享的 FFN 相比, PFFN 通过引入更多参数来增强建模能力, 同时保持计算复杂度不变.

值得强调的是, PFFN 与 MMoE 专家的区别在于, 每个 PFFN 处理的是不同的 Token 输入, 而 MMoE 中的所有专家组共享相同的输入. MMoE 模型中是多个专家处理相同的输入, Transformer 模型中不同输入共享一个 FFN, RankMixer 同时分割输入和参数, 这有利于学习不同特征子空间的多样性.

#3.4 RankMixer 中的稀疏 MoE

为了进一步提高 ROI, 我们可以把每个 Token 的稠密 FFN 结构替换为 Sparse-MoE (Sparse Mixture of Experts) 块, 从而在计算成本基本不变的情况下提升模型容量.

然而, 在 RankMixer 中使用原始的 Sparse-MoE 会导致性能下降, 原因如下:

均匀的 $k$ -expert 路由. Top-k 选择会平等对待所有特征 Token, 将预算浪费在低信息 Token 上, 而高信息 Token 则被忽略, 这阻碍了模型捕捉 Token 之间的差异.
专家训练不足. PFFN 已经将参数乘以 Token 数; 添加非共享专家会导致专家数量大大增加, 结果会导致路由严重不平衡和专家训练不足.

我们结合两种互补的训练策略来解决上述问题.

ReLU 路由:

为了赋予 Token 灵活的专家数量并保持可区分性, 我们替换了常见的 Top k + softmax 为一个 ReLU 门以及自适应 ℓ1 惩罚 (Wang, 2024a).

给定令牌 $\mathbf{s}_i \in \mathbb{R}^{d_h}$ 和路由器 $h(\cdot)$ 的第 $j$ 个专家 $e_{i,j}(\cdot)$ .

$G_{i,j} = \text{ReLU}\big( h(\mathbf{s}_i) \big), \quad \mathbf{v}_i = \sum_{j=1}^{N_e} G_{i,j} \, e_{i,j}(\mathbf{s}_i).$

其中 $N_e$ 表示每个令牌的专家数量, $N_t$ 表示令牌的数量. ReLU 路由将激活更多专家来处理高信息量的令牌, 并提高参数效率. 稀疏性受以下因素控制 $\mathcal{L}_{reg}$ , 系数为 $\lambda$ , 使平均活跃专家比率接近预算.

$\mathcal{L} = \mathcal{L}_{\text{task}} + \lambda \, \mathcal{L}_{\text{reg}}, \quad \mathcal{L}_{\text{reg}} = \sum_{i=1}^{N_t} \sum_{j=1}^{N_e} G_{i,j}.$

稠密训练 + 稀疏推理（DTSI-MoE）:

受 Pan(2024)的启发, 本文采用了两个路由器 $h_{\text{train}}$ 和 $h_{\text{infer}}$ , 而 $\mathcal{L}_{\text{reg}}$ 仅应用于 $h_{\text{infer}}$ . 训练过程中会更新 $h_{\text{train}}$ 和 $h_{\text{infer}}$ , 而推理过程中仅使用 $h_{\text{infer}}$ . 结果表明, DS-MoE 能够避免专家模型训练不足的问题, 同时降低推理成本.

#3.5 Scaling Up 的方向

RankMixer 本质上是一个高度并行且可扩展的架构. 其参数数量和计算成本可以沿四个正交轴展开: Token 计数 $T$ , 模型宽度 $D$ , 层数 $L$ 和专家数量 $E$ .

对于全稠密激活版本, 单个样本的参数规模和前向浮点运算次数可以计算如下.

$\#\text{Param} \approx 2 k L T D^2, \quad \text{FLOPs} \approx 4 k L T D^2.$

$k$ 是用于调整 FFN 隐藏层维度的缩放比例. 在 Sparse-MoE 版本中, 有效参数量和每个 token 的计算量会进一步按照稀疏度比例进行缩放, 其中 $s = \frac{\#\text{Activated\_Param}}{\#\text{Total\_Param}}$ .

#4. 实验

#4.1 实验设置

#4.1.1 数据集和环境

离线实验使用了抖音推荐系统的训练数据. 这些数据来源于抖音的在线日志和用户反馈标签. 训练数据集包含 300 多个特征, 例如数值特征, ID 特征, 交叉特征和序列特征, 涉及数十亿个用户 ID 和数亿个视频 ID, 所有这些特征都被转换为词嵌入. 数据每天包含数万亿条记录, 实验所用数据收集周期为两周.

#4.1.2 评估指标

为了评估模型性能, 我们使用 AUC(曲线下面积)和 UAUC(用户级 AUC)作为主要性能指标, 并使用参数数量, FLOPs 和 MFU 作为效率指标, 具体如下.

完成/跳过 AUC/UAUC: finish=1/0 或 skip=1/0 标签表示用户是否在短时间内看完视频或切换到下一个视频. 我们评估此完成标签的 AUC 和 UAUC. AUC 增加 0.0001 可视为显著改进.

稠密参数: 稠密部分中的参数数量, 不包括稀疏嵌入参数.

训练 FLOPs/Batch: 运行一个包含 512 个数据的批次所需的浮点运算(FLOPs)次数, 代表训练的计算成本.

MFU: MFU(模型 FLOPs 利用率)是一个衡量模型如何有效地利用硬件提供的浮点运算的指标, 计算方法是将模型的实际 FLOPs 消耗量除以硬件的理论 FLOPs 容量.

#4.1.3 基线

我们与以下公认的最先进基线模型进行比较:

DLRM-MLP: 作为实验基线, 它是用于特征交叉的原始多层感知器(MLP).

DCNv2 (Wang, 2021), RDCN (Borisyuk, 2024): 特征交叉模型的最先进模型. MoE 通过使用多个并行专家来扩展规模.

AutoInt (Song, 2019), Hiformer (Gui, 2023): 结合了异构自注意力层和低秩近似矩阵计算.

DHEN (Zhang, 2022): 结合了不同类型的特征交叉块, 并堆叠了多层(DCN/自注意力/FM/LR).

Wukong (Zhang, [未注明日期]): 研究了特征交叉的缩放规律, 遵循 DHEN 的架构, 并结合了因子分解机块 (Factorization Machine Block, FMB) 和线性压缩块 (Linear Compress Block, LCB).

所有实验均在数百个 GPU 上进行, 采用混合分布式训练框架, 其中稀疏部分异步更新, 而稠密部分同步更新. 所有模型的优化器超参数保持一致. 对于稠密部分, 我们使用学习率为 0.01 的 RMSProp 优化器; 对于稀疏部分, 我们使用 Adagrad 优化器.

#4.2 与最先进方法的比较

TODO

为了探索如何 Scaling Up 模型, 我们比较了参数规模在 1 亿左右的相似模型, 以确定在相同的计算成本下哪种模型结构性能最佳.

表 1 总结了我们方法和基线方法的主要性能结果. 我们可以看到, RankMixer 在多个目标和指标上都显著优于其他最先进的模型.

我们仔细分析了每个模型.

首先, 简单地将 DLRM 模型扩展到 1 亿个参数只能带来有限的性能提升, 这凸显了设计针对推荐数据特征量身定制的模型以获得更佳扩展性能的必要性.
然后, 我们将 RankMixer 模型与其他经典的跨结构设计(例如 DCN、RDCN、AutoInt 和 DHEN)进行比较, 发现它们都存在模型参数和计算成本不平衡的问题.

即使参数规模相对较小, 这些模型的 FLOPs 也已经很高, 表明其设计存在缺陷, 限制了表格中的结果. 而 RankMixer 模型虽然性能最佳, 但在扩展到 1 亿个参数时, 其 FLOPs 在所有模型中仍然保持相对适中, 这体现了其在模型容量和计算负载之间取得平衡的方法.

我们还将 RankMixer 与几个常用的最先进的 Scaling Up 模型(如 Hiformer 和 Wukong)进行了比较, 发现, 在类似的参数设置下, RankMixer 不仅性能更好, 而且计算需求也更低.

#4.3 不同模型的 Scaling Laws

scaling_law

图 2. 不同模型的最终 AUC-gain 与 Params/FLOPs 之间的 Scaling Law. x 轴采用对数刻度.

图 2 展示了从参数大小和浮点运算次数(FLOPs)两个方面观察到的 Scaling Law 曲线. RankMixer 模型在参数大小和浮点运算次数方面都展现出最陡峭的 Scaling Law. RankMixer 模型始终优于其他模型.

尽管 Wukong 模型的参数曲线相对陡峭, 但其计算成本增长速度更快; 因此, 在 AUC 与 FLOPs 曲线上, 它与 RankMixer 和 Hiformer 之间的差距更大. 此外, Hiformer 的性能略逊于 RankMixer, 这反映出其对特征级 Token 分割和注意力机制的依赖影响了其效率. DHEN 的扩展性并不理想, 反映了其交叉结构的扩展性有限. 此外, MoE 通过增加专家来 Scaling Up 模型的策略在维持专家平衡方面带来了挑战, 导致其扩展性能欠佳.

具体来说, 我们可以通过增加宽度( $D$ )、特征 Token 数( $T$ )和层数( $L$ )来扩展 RankMixer 模型. 在我们的实验中, 我们观察到与 LLM Scaling Law 一致的结论: 模型质量主要与参数总数相关, 不同的扩展方向(深度 $L$ , 宽度 $D$ , Token 数 $T$ )几乎产生相同的性能. 从计算效率的角度来看, 更大的隐藏层维度会产生更大的矩阵乘法形状, 因此比堆叠更多层能获得更高的 MFU. 因此, 1 亿和 10 亿的最终配置分别设置为 ( $D = 768$ , $T = 16$ , $L = 2$ ) 和 ( $D = 1536$ , $T = 32$ , $L = 2$ ).

#4.4 消融实验

在 RankMixer-100M 模型中, 我们对残差连接和多头标记混合(Multi-Head Token-Mixing)进行了消融实验. 从表 2 可以看出, 移除这些组件会显著降低模型的性能. 移除多头标记混合会丢失全局信息, 因为每个前馈神经网络(FFN)仅对部分特征进行建模, 彼此之间没有交互. 移除残差连接和层归一化(LayerNorm)也会降低性能, 减少训练稳定性, 并增加梯度爆炸或梯度消失问题的可能性.

我们进一步分析了词元混合策略, 即表 3 中从特征词元到前馈神经网络(FFN)的路由策略. 与多头词元混合(Multi-Head Token Mixing)相比, 这些路由策略包括以下几种.

All-Concat-MLP: 将所有词元连接起来, 并通过一个大型多层感知器(MLP)进行处理, 然后再将其分割成相同数量的词元. 性能下降表明学习大型矩阵和削弱局部信息学习的挑战.
All-Share: 不进行分割, 整个输入向量被共享并输入到每个词元对应的 FFN 中, 类似于 MoE. 性能显著下降, 这表明特征子空间分割和独立建模相对于全共享输入的重要性.
自注意力机制: 在词元之间应用自注意力进行路由. 其性能略逊于多头词元混合机制, 且计算成本较高, 这表明在数百个不同的特征子空间中学习相似性的难度很大.

#4.5 Sparse-MoE 的可扩展性和专家平衡

图 3. RankMixer 变体在激活率逐渐降低(1, 1/2, 1/4, 1/8 位专家)下的 AUC 性能. 密集训练 + ReLU 路由的 SMoE 几乎保留了 1B 密集模型的所有准确率.

可扩展性:

图 3 绘制了 SMoE 的离线 AUC 增益与稀疏度的关系. 将稠密训练-稀疏推理(DTSI)与 ReLU 路由相结合, 对于在高稀疏度下保持准确率至关重要.

这使得 RankMixer 能够在参数容量(和内存占用)扩展 8 倍以上的情况下, AUC 几乎没有损失, 并且推理时间显著缩短(吞吐量提升 50%).

原始 SMoE 的性能随着激活的专家数量减少而单调下降, 这说明了我们发现的专家不平衡和欠训练问题. 添加负载均衡损失可以减少相对于原始 SMoE 的性能下降, 但仍然不如 DTSI + ReLU 版本, 因为问题主要出在专家训练上, 而不是路由上.

这验证了 Sparse-MoE 是将 RankMixer 从当前的 10 亿参数规模扩展到未来 100 亿规模部署且不超出成本预算的有效途径.

图 4. RankMixer 中不同标记的激活专家比例。

表 4. 抖音和抖音 Lite app 中信息流推荐场景的在线 A/B 测试结果显示，所有改进均具有统计学意义。根据长期反向 A/B 测试和长期反向实验的持续观察，RankMixer-1B 带来的提升尚未收敛，此处展示的结果仍在不断改进中。

表 5. RankMixer 在广告领域的在线提升

表 6. 在线模型部署和成本指标

专家团队的平衡性和多样性:

传统的稀疏学习模型（MoE）经常面临专家不平衡的问题，这会导致部分专家训练不足，最终导致“死亡专家”（几乎从未被激活的专家），只有少数固定专家始终处于激活状态。图 4 显示，将 DTSI（密集训练，稀疏推理）与 ReLU 路由相结合可以有效解决这个问题：密集训练保证了大多数专家都能获得足够的梯度更新，从而避免专家饥饿；ReLU 路由使得激活率在不同词元之间动态变化——图中所示的激活比例会根据词元的信息量自适应地变化，这与推荐数据的多样性和高度动态分布非常吻合。

#4.6 在线服务成本

如何防止推理延迟随着参数数量增加两个数量级而急剧上升？在实际系统中，延迟与吞吐量成反比，与服务机器资源的成本成正比。与我们之前完全部署的 1600 万参数模型（结构集成了 DLRM 和 DCN）相比，我们的 RankMixer 模型将参数规模扩大了约 70 倍，达到 10 亿倍。尽管参数数量显著增加，但由于我们采用了与硬件相匹配的模型设计和优化策略，最终的推理延迟仍然保持稳定。

当大幅增加模型参数时，延迟可以分解为以下公式：

$\text{Latency} = \frac{\#\text{Param} \times (\text{FLOPs/Param ratio})} {\text{MFU} \times (\text{Theoretical Hardware FLOPs})}$

如表 6 所示，参数增加两个数量级逐渐被 FLOPs/Param 比率降低 3.6 倍、MFU 增加 10 倍以及基于量化的硬件 FlOPs 改进 2 倍所抵消。

FLOPs/参数比率:
表 6 的第三行报告了每个参数所需的浮点运算次数(FLOPs). 得益于模型设计, RankMixer 在 FLOPs 仅增加约 20 倍的情况下, 参数数量增加了 70 倍, 其 FLOPs/参数比率仅为基线的三分之一, 效率提升了 3.6 倍. 换句话说, 在相同的 FLOPs 预算下, RankMixer 可以容纳比基线多三倍的参数.
模型浮点运算利用率(MFU):
如表 6 所示, MFU 表示机器计算的利用率. 通过使用大型 GEMM 形状, 良好的并行拓扑结构(将并行的逐标记 FFN 融合到一个内核中), 降低内存带宽成本和开销, RankMixer 将 MFU 提高了近 10 倍, 使模型从内存密集型转变为计算密集型.
量化半精度(fp16)推理将使 GPU 的理论峰值硬件浮点运算能力提高 2 倍:
RankMixer 中的主要计算包含多个大型矩阵乘法, 如前所述, 这些乘法非常适合半精度(fp16)推理.

#4.7 在线表现

略

#5 结论

本文介绍了我们最新开发的 RankMixer 模型, 该模型已全面部署于抖音 Feed 排名. 它融合了异构特征交叉的模型设计和高度并行化的架构, 从而显著提升了服务效率. 实验结果表明, 该模型性能优异, 且具有极强的扩展性. 全面部署于抖音 App 后, 其活跃天数和 App 使用时长分别提升了 0.3% 和 1%.