RankMixer: Scaling Up Ranking Models in Industrial Recommenders

116 阅读11分钟

标题:RankMixer:工业级推荐系统中排序模型的规模化扩展
单位:字节跳动

摘要

大语言模型(LLM)的最新进展推动了推荐系统规模化扩展的研究热潮,但目前仍存在两个现实阻碍:第一,工业级推荐系统的训练与服务成本必须严格满足延迟约束和高 QPS 的需求;第二,排序模型中绝大多数人工设计的特征交叉模块都继承自 CPU 时代,无法充分发挥现代 GPU 的算力优势,导致模型浮点运算利用率(MFU)极低,可扩展性极差。

为此,我们提出了RankMixer,这是一种面向硬件设计的模型架构,专为统一、可扩展的特征交互架构量身打造。RankMixer 保留了 Transformer 架构的高并行性,同时用多头 Token 混合模块替代了二次复杂度的自注意力机制,实现了更高的运行效率。此外,RankMixer 通过逐 Token 前馈网络(Per-token FFNs),同时实现了对独立特征子空间的建模,以及跨特征子空间的交互能力。

为进一步提升大模型的投入产出比(ROI),我们基于稀疏混合专家(Sparse-MoE)结构对其进行扩展,实现了十亿级参数量的模型版本。同时我们采用了一种动态路由策略,解决了专家训练不充分、训练不均衡的问题。实验结果表明,在万亿级生产数据集上,RankMixer 展现出了卓越的规模化扩展能力。

通过用 RankMixer 替换此前多样化的人工设计、低 MFU 的模块,我们将模型的 MFU 从 4.5% 提升至 45%,在推理延迟基本保持不变的前提下,将线上排序模型的参数量扩大了两个数量级。我们在推荐和广告两大核心应用场景中通过线上 A/B 测试,验证了 RankMixer 的通用性。最终,我们在不增加服务成本的前提下,上线了全流量服务的 10 亿级稠密参数量 RankMixer 模型,实现了用户活跃天数提升 0.3%、应用内总使用时长提升 1.08% 的效果。

关键词: 缩放定律、排序模型、推荐系统

1 引言

推荐系统(RS)是信息分发流程中的核心环节。作为重要的机器学习应用场景,推荐系统基于海量的多域特征数据,预测用户对物品的行为偏好,这些数据包括各类统计值等数值型特征、用户与物品 ID 等类别型特征、用户行为特征和内容特征等。当前主流的推荐方法都基于深度学习推荐模型(DLRM),这类模型以特征输入的嵌入向量为基础,通过神经网络构建稠密交互层,灵活地捕捉特征交互。DLRM 中的稠密交互层是决定推荐系统性能的核心,学界也已提出了多样化的模型结构。

大语言模型通过扩大参数量实现了性能飞跃,受此驱动,规模化扩展 DLRM 以充分利用海量数据成为了行业的迫切需求。此前已有大量关于 DLRM 规模化扩展的研究成果:早期研究仅对特征交互层进行加宽或堆叠,并未修改核心结构,这种方式带来的性能收益十分有限,有时甚至会出现负向效果;后续的研究如 DHEN 和 Wukong,则聚焦于设计创新的深度神经网络结构,以提升规模化扩展的性能。

但在推荐场景中,通过扩大模型规模提升性能,面临着独特的现实挑战。与自然语言处理(NLP)或计算机视觉(CV)任务不同,工业级推荐系统必须严格遵守严苛的延迟约束,同时支撑极高的 QPS。因此,核心挑战在于找到模型效果与计算效率之间的最佳平衡点。

从发展历程来看,推荐系统中排序模型的架构始终由 CPU 时代的设计原则主导。这类模型通常依赖组合多种异构的人工设计特征交叉模块来提取特征交互,但在现代 GPU 上,其核心算子大多是内存受限而非计算受限,导致 GPU 并行度极差,MFU 通常仅为个位数百分比。此外,CPU 时代的模型计算成本与参数量大致成正比,因此缩放定律所提出的激进规模化扩展带来的潜在 ROI,在实践中很难实现。

综上,DLRM 缩放定律的相关研究必须解决以下问题:

  1. 架构必须与硬件特性对齐,最大化现代 GPU 上的 MFU 和计算吞吐量;
  2. 模型设计必须适配推荐数据的特性,例如数百个特征域之间的异构特征空间、个性化的跨特征交互。

为解决上述挑战,我们提出了面向硬件感知的模型设计方法 RankMixer。RankMixer 的核心设计基于两个可扩展组件:

  1. 多头 Token 混合:仅通过无参数算子实现跨 Token 的特征交互,该策略在性能和计算效率上均优于自注意力机制
  2. 逐 Token 前馈网络:通过为不同的特征子空间建模分配独立的参数,大幅扩展了模型容量,同时解决了跨特征空间的主导性问题。这类前馈网络也与推荐数据的模式高度契合,实现了更优的缩放特性。

为进一步提升大模型的 ROI,我们将 per-token FFN 扩展为 SparseMoE 结构。通过针对不同数据,为每个 Token 仅动态激活一部分专家,我们可以在计算成本增加极少的前提下,显著提升模型容量。

RankMixer 采用了与 Transformer 类似的高并行架构,但克服了基于自注意力的特征交互的几个关键缺陷:训练效率低下、跨空间 ID 相似度建模时的组合爆炸问题,以及注意力权重矩阵带来的严重内存受限问题。同时,在相同 FLOPs 下,RankMixer 比原生 Transformer 具备更强的模型容量和学习能力。

在抖音推荐系统的生产部署中,我们验证了该架构可以将模型参数量扩大 100 倍以上,同时推理延迟相比此前的基线模型更短。这一成果的核心,是 RankMixer 架构实现了参数量增长与 FLOPs 的解耦,以及通过高 MFU 和工程优化,实现了 FLOPs 增长与实际成本的解耦。

本文的主要贡献可总结如下:

  1. 我们提出了一种遵循硬件感知设计理念的全新架构 RankMixer。设计了多头 Token 混合和 per-token FFN 策略,以高效捕捉异构特征交互;同时采用动态路由策略,提升了 RankMixer 中 SparseMoE 的可扩展性。
  2. 借助高 MFU 和性能优化,我们在不增加推理成本的前提下,将模型参数量扩大了 70 倍(包括提升 MFU 和量化优化)。
  3. 我们开展了丰富的离线和线上实验,在万亿级工业推荐数据集上探究了模型的缩放定律。RankMixer 模型已成功在抖音信息流推荐排序中实现全流量部署,实现了用户活跃天数提升 0.3%、应用使用时长提升 1.08% 的效果。

2 相关工作

现代推荐系统以深度学习推荐模型(DLRM)为基础,而如何高效建模特征交互,是 DLRM 的核心关键。Wide&Deep 是该领域最早的代表性工作之一,它结合了逻辑回归(宽部分)和 DNN(深部分),分别捕捉低阶和高阶特征交互。DeepFM 是另一项重要成果,它将因子分解机(FM)与 DNN 进行了融合。此外,DeepCross 是残差网络的扩展,旨在隐式地学习自动特征交互。但研究发现,仅依靠 DNNs 学习高阶特征交互,存在极大的挑战。

显式交叉方法通过设计不同的算子,显式地捕捉高阶特征交互,代表性工作包括 PNN、DCN 及其后续版本 DCNv2、xDeepFM、FGCNN 和 FiGNN。AutoInt 和 Hiformer 则采用带残差连接的注意力机制,学习复杂的特征交互。DHEN 提出将多种交互算子进行组合。尽管这些新架构提升了预测精度,但同时也增加了模型的延迟和内存开销,且模型规模普遍较小。

Scaling law 已成为深度学习领域的核心主题,也是过去十年众多技术突破的关键催化剂,尤其在 NLP、CV 和多模态建模领域。Scaling law 描述了模型性能与模型尺寸、数据量、计算量等缩放因子之间的幂律相关性。近年来,推荐系统中的 Scaling law 也受到了研究者的广泛关注。已有研究探索了用户行为序列预训练、通用用户表征和线上召回的规模化扩展策略。Wukong 通过堆叠因子分解机模块(FMs)和线性压缩模块(LCB)学习特征交互。与之并行的研究中,Zhang 等人将序列推荐模型规模化扩展至 8 亿参数量。HSTU 则提升了聚焦于序列部分的生成式推荐系统(GRs)的缩放效果。

3 方法论

3.1 整体架构

image.png 图 1:RankMixer 模块的架构。一个 RankMixer 模块包含两个核心组件:多头 Token 混合层,以及基于稀疏混合专家(SMoE)的 PFFN。Token 混合模块将每个 Token 的嵌入向量切分为 H 个更小的部分(即 “头”),随后跨 Token 重组这些部分,生成全新的混合 Token。这一机制能够实现不同特征之间的信息交互。

RankMixer 的整体架构为:T 个输入 Token 经过 L 个连续的 RankMixer 模块处理,之后经过输出池化算子。每个 RankMixer 模块包含两个核心组件:(1)多头 Token 混合层;(2)逐 Token 前馈网络(PFFN)层,如图 1 所示。

首先,输入向量 einput\text{e}_{\text{input}} 被 Token 化为 T 个特征 Token x1,x2,...,xTx_1,x_2,...,x_T,每个 Token 代表一组语义连贯的特征向量。RankMixer 模块通过 L 层迭代优化 Token 表征,计算过程如下:

Sn1=LN(TokenMixing(Xn1)+Xn1),Xn=LN(PFFN(Sn1)+Sn1),(1)\begin{aligned} S_{n-1} & = \operatorname{LN}\left( \operatorname{TokenMixing} \left(X_{n-1}\right)+X_{n-1}\right), \\ X_{n} & = \operatorname{LN}\left(\operatorname{PFFN}\left(S_{n-1}\right)+S_{n-1}\right), \tag{1} \end {aligned}

其中,LN(⋅)\operatorname{LN(⋅)} 为层归一化函数,TokenMixing(⋅)\operatorname{TokenMixing(⋅)}PFFN(⋅)\operatorname{PFFN(⋅)} 分别为多头 Token 混合模块和逐 Token 前馈网络模块;XnRT×DX_n \in \mathbb{R}^{T \times D} 为第 n 个 RankMixer 模块的输出,X0RT×DX_0 \in \mathbb{R}^{T \times D}x1,x2,...,xTx_1,x_2,...,x_T 堆叠得到,D 为模型的隐藏层维度。最终的输出表征 ooutput\text{o}_\text{{output}} 由最后一层输出表征 XLX_L 的均值池化得到,用于计算不同任务的预测结果。

3.2 输入层与特征 Token 化

构建大规模推荐模型的第一步,是准备包含丰富信息的输入,包括:用户特征(用户 ID 及其他用户信息)、候选物品特征(视频 ID、作者 ID 等)、经序列模块处理的序列特征(用于捕捉用户的时序兴趣,输出 es\text{e}_{\text{s}}),以及用户与候选物品之间的交叉特征。所有特征都会被转换为不同维度的嵌入向量。

为了在后续阶段实现高效的并行计算,不同维度的嵌入向量必须被转换为维度对齐的向量,即特征 Token。我们将这一嵌入对齐过程称为Token 化。最简单的策略是为每个特征分配一个嵌入向量,但在特征数量通常达数百个的场景下,这种策略会带来诸多问题:Token 数量达数百个时,分配给每个 Token 的参数和计算量会被切割为极小的片段,导致重要特征的建模不充分,同时 GPU 核心的利用率低下;反之,Token 数量过少(例如单个 Token)会让模型结构退化为简单的深度神经网络,无法对多样化的特征空间进行独立表征,还会出现主导特征掩盖其他特征的风险。

为解决这些问题,我们提出了基于领域知识的语义感知 Token 化方法,将特征划分为多个语义连贯的聚类。这些分组后的特征会被顺序拼接为一个完整的嵌入向量 einput=[e1;e2;;eN]e_{\text{input}} = [e_1; e_2; \ldots; e_N],随后被切分为固定维度、数量合适的 Token。每个特征 Token xiRDx_i \in \mathbb{R}^{D} 都对应一组语义相近的特征嵌入集合。

xi=Proj(einput[d(i1):di]),i=1,...,T(2)x_{i}=\operatorname{Proj}\left(e_\text{input}[d \cdot(i-1): d \cdot i]\right), \quad i=1, ..., T \tag{2}

其中 einpute_{\text{input}} 为拼接后的嵌入向量,dd 为每个 Token 的固定维度,NN 为特征分组的数量,TT 为最终得到的 Token 数量,Proj\text{Proj} 函数将切分后的嵌入映射到 DD 维空间。

3.3 RankMixer 模块

3.3.1 多头 Token 混合

为实现 Token 之间高效的信息交互(这对特征交叉和全局信息建模至关重要),我们提出了多头 Token 混合模块。每个 Token 会被均匀切分为 HH 个头,Token xtx_{t} 的第 hh 个头记为 xthx_{t}^{h}

[xt(1)xt(2)...xt(H)]=SplitHead(xt)(3)\left[x_{t}^{(1)}\| x_{t}^{(2)}\| ...\| x_{t}^{(H)}\right]=\operatorname{SplitHead}\left(x_{t}\right) \tag{3}

由于推荐任务需要从不同视角进行建模,这些头可以被视为 Token xtx_{t} 在低维特征子空间的投影。Token 混合则用于融合这些子空间向量,实现全局特征交互。形式上,多头 Token 混合后,第hh个头对应的第hh个 Token shs^{h} 的构建方式为:

sh=Concat(x1h,x2h,...,xTh)(4)s^{h}= \operatorname{Concat} \left(x_{1}^{h}, x_{2}^{h}, ..., x_{T}^{h}\right) \tag{4}

多头 Token 混合模块的输出为 SRH×TDHS \in \mathbb{R}^{H ×\frac{T D}{H}},由所有重排后的 Token s1,s2,...,sHs_{1}, s_{2}, ..., s_{H} 堆叠得到。在本文工作中,我们设置 H=TH=T,以保证 Token 混合前后的 Token 数量一致,支持残差连接。

经过残差连接和归一化模块后,最终生成:

s1,s2,...,sT=LN(TokenMixing(x1,x2,...,xT)+(x1,x2,...,xT))(5)s_{1},s_{2},...,s_{T}=\operatorname{LN}\left( \operatorname{TokenMixing}\left( x_{1},x_{2},...,x_{T}\right) +\left(x_{1},x_{2},...,x_{T}\right) \right) \tag{5}

尽管自注意力机制在大语言模型中被证明非常有效,但我们发现它在推荐系统中并非最优选择。在自注意力机制中,注意力权重通过 Token 的内积计算得到,这种方式在所有 Token 共享统一嵌入空间的 NLP 任务中表现优异,但推荐任务中的特征空间本质上是异构的。在两个异构语义空间之间计算内积相似度,本身就存在极大的难度 —— 尤其是在推荐系统中,用户侧和物品侧的特征 ID 空间可能包含数亿个元素。

因此,将自注意力机制应用于这类多样化的输入,效果无法超越无参数的多头 Token 混合方法,同时还会消耗更多的计算量、内存 IO 操作和 GPU 显存。

3.3.2 逐 Token 前馈网络

此前的 DLRM 和 DHEN 模型,通常会在单个交互模块中混合来自多个不同语义空间的特征,这可能导致高频域特征占据主导,淹没低频或长尾信号,最终损害整体推荐质量。为此,我们提出了参数隔离的前馈网络架构,即逐 Token 前馈网络

在传统设计中,前馈网络的参数在所有 Token 之间共享,而我们的方法为每个 Token 配备专属的变换算子,实现了每个 Token 的参数隔离。对于第 tt 个 Token sts_{t},PFFN 的表达式为:

vt=fpffnt,2(Gelu(fpffnt,1(st)))(6)v_{t}=f_{pffn}^{t, 2}\left(\operatorname{Gelu}\left(f_{pffn}^{t, 1}\left(s_{t}\right)\right)\right) \tag{6}

其中:

fpffnt,i(x)=xWpffnt,i+bpffnt,i(7)f_{pffn}^{t, i}(x)=x W_{pffn}^{t, i}+b_{pffn}^{t, i} \tag{7}

是逐 Token 前馈网络的第 ii 层多层感知机(MLP),Wpffnt,1RD×kDW_{pffn}^{t, 1} \in \mathbb{R}^{D ×k D}bpffnt,1RkDb_{pffn }^{t, 1} \in \mathbb{R}^{k D}Wpffnt,2RkD×DW_{pffn}^{t, 2} \in \mathbb{R}^{k D ×D}bpffnt,2RDb_{pffn}^{t, 2} \in \mathbb{R}^{D}kk 是调整逐 Token 前馈网络隐藏层维度的超参数,Gelu()\operatorname{Gelu}(·) 为 Gelu 激活函数,stRDs_{t} \in \mathbb{R}^{D}为第 tt 个 Token。

我们将逐 Token 前馈网络模块总结为:

v1,v2,...,vT=PFFN(s1,s2,...,sT)(8)v_{1}, v_{2}, ..., v_{T}=\operatorname{PFFN}\left(s_{1}, s_{2}, ..., s_{T}\right) \tag{8}

其中:

PFFN(s1,s2,...,sT)=fpffnt,2(Gelu(fpffnt,1(s1,s2,...,sT)))(9)\operatorname{PFFN}\left( s_{1}, s_{2}, ..., s_{T}\right)=f_{pffn}^{t, 2}\left( \operatorname{Gelu}\left( f_{pffn}^{t, 1}\left( s_{1}, s_{2}, ..., s_{T}\right) \right) \right) \tag{9}

与全参数共享的前馈网络相比,逐 Token 前馈网络在计算复杂度不变的前提下,通过引入更多参数提升了建模能力。

需要重点强调的是,逐 Token 前馈网络与多门控混合专家(MMoE)的专家存在本质区别:每个逐 Token 前馈网络接收的是不同的 Token 输入,而 MMoE 中的所有专家共享同一个输入。与 MMoE(多个专家处理同一输入)和 Transformer(不同输入共享同一个前馈网络)不同,RankMixer 同时对输入和参数进行了拆分,这非常有利于学习不同特征子空间的多样性。

3.4 RankMixer 中的稀疏混合专家(Sparse MoE)

为进一步提升规模化扩展的 ROI,我们可以将每个 Token 的稠密前馈网络替换为稀疏混合专家(MoE)模块,从而在计算成本基本保持不变的前提下,提升模型容量。

但原生的稀疏 MoE 在 RankMixer 中会出现性能退化,原因包括两点:

  1. 固定 kk 专家路由:Top-k 选择机制对所有特征 Token 一视同仁,会在低信息 Token 上浪费预算,同时导致高信息 Token 的预算不足,阻碍了模型捕捉 Token 之间的差异;
  2. 专家训练不充分:逐 Token 前馈网络已经让参数量随 Token 数量成倍增长,再加入非共享的专家会让专家数量进一步爆炸,导致路由高度不均衡,专家训练效果极差。

我们结合两种互补的训练策略,解决了上述问题。

ReLU 路由 为了让 Token 能够灵活选择激活的专家数量,同时保持可微性,我们用 ReLU 门控加自适应 1\ell_{1} 正则项,替换了常用的 Top-k+softmax 机制。对于 Token siRdhs_{i} \in \mathbb{R}^{d_{h}} 的第 jj 个专家 ei,j()e_{i, j}(\cdot) 和路由网络 h()h(-),计算方式为:

Gi,j=ReLU(h(si)),vi=j=1NeGi,jei,j(si)(10)G_{i, j}=\operatorname{ReLU}\left(h\left(s_{i}\right)\right), \quad v_{i}=\sum_{j=1}^{N_{e}} G_{i, j} e_{i, j}\left(s_{i}\right) \tag{10}

其中,NeN_{e} 为每个 Token 的专家数量,NtN_{t} 为 Token 数量。ReLU 路由会为高信息 Token 激活更多专家,提升参数效率。稀疏性通过带系数 λ\lambda 的正则项 Lreg\mathcal{L}_{reg } 控制,使平均激活专家比例保持在预算附近:

L=Ltask+λLreg,Lreg=i=1Ntj=1NeGi,j(11)\mathcal{L}=\mathcal{L}_{task }+\lambda \mathcal{L}_{reg }, \quad \mathcal{L}_{reg }=\sum_{i=1}^{N_{t}} \sum_{j=1}^{N_{e}} G_{i, j} \tag{11}

稠密训练 / 稀疏推理(DTSI-MoE)

受相关研究启发,我们采用了两个路由网络 htrainh_{train}hinferh_{infer},且仅对 hinferh_{infer} 施加正则项 Lreg\mathcal{L}_{reg}。训练过程中会同时更新 htrainh_{train}hinferh_{infer},而推理阶段仅使用 hinferh_{infer}。实验结果表明,该策略让专家避免了训练不充分的问题,同时降低了推理成本。

3.5 规模化扩展方向

RankMixer 本质上是一个高并行、可扩展的架构,其参数量和计算成本可以沿四个正交的维度进行扩展:Token 数量 TT、模型宽度 DD、层数 LL 和专家数量 EE。对于全稠密激活的版本,单个样本的参数量和前向传播 FLOPs 可通过以下公式计算:

#Param2kLTD2,FLOPs4kLTD2(12)\#Param \approx 2 k L T D^{2}, \quad \text{FLOPs} \approx 4 k L T D^{2} \tag{12}

其中 kk 是调整前馈网络隐藏层维度的缩放比例。在稀疏 MoE 版本中,每个 Token 的有效参数量和计算量会进一步通过稀疏率 s=#ActivatedParam#TotalParams = \frac{\#\text{ActivatedParam}}{\#\text{TotalParam}} 进行缩放。

4 实验

4.1 实验设置

4.1.1 数据集与实验环境

离线实验基于抖音推荐系统的训练数据开展,这些数据来自抖音的线上日志和用户反馈标签。训练数据集包含超 300 个特征,包括数值型特征、ID 特征、交叉特征和序列特征,覆盖数十亿用户 ID 和数亿视频 ID,所有特征均被转换为嵌入向量。数据覆盖每日万亿级记录,实验基于两周内采集的数据开展。

所有实验在数百张 GPU 上,通过混合分布式训练框架执行:稀疏部分采用异步更新,稠密部分采用同步更新。所有模型的优化器超参数保持一致:稠密部分采用 RMSProp 优化器,学习率 0.01;稀疏部分采用 Adagrad 优化器。

4.1.2 评估指标

我们采用曲线下面积(AUC\text{AUC})和用户级曲线下面积(UAUC\text{UAUC})作为核心性能指标,采用参数量、FLOPs 和 MFU 作为效率指标,具体定义如下:

  • 完播/划走 AUC\text{AUC}/UAUC\text{UAUC} :完播标签(finish=1/0finish=1/0)或划走标签(skip=1/0skip=1/0)表示用户是否完整观看了视频,或在短时间内滑动到下一个视频。我们针对完播标签评估 AUC\text{AUC}UAUC\text{UAUC},在工业场景中,AUC\text{AUC} 提升 0.0001 即可视为具有统计显著性的正向收益。
  • 稠密参数量:模型稠密部分的参数量,不含稀疏嵌入参数。
  • 单批次训练 FLOPs:单批次 512 个样本经过模型前向传播所需的浮点运算次数,代表训练的计算成本。
  • MFU:模型浮点运算利用率,衡量模型对硬件提供的浮点运算能力的利用效率,计算方式为模型实际 FLOPs 消耗除以硬件的理论 FLOPs 峰值。

4.1.3 基线模型

我们与以下公认的业界 SOTA 基线模型进行对比:

  • DLRM-MLP:原生 MLP 特征交叉结构,作为实验基线;
  • DCNv2、RDCN:特征交叉模型的 SOTA 方案;
  • MoE:通过多个并行专家实现规模化扩展的方案;
  • AutoInt、Hiformer:结合异构自注意力层和低秩近似矩阵计算的方案;
  • DHEN:组合多种特征交叉模块并堆叠多层(DCN / 自注意力 / FM/LR)的架构;
  • Wukong:基于 DHEN 架构,通过因子分解机模块(FMB)和线性压缩模块(LCB)探究特征交互缩放定律的方案。

4.2 与 SOTA 方法的对比

为探究模型的规模化扩展方式,我们在 1 亿左右相近参数量的条件下对比了不同模型,以确定相同计算成本下性能最优的模型结构。

原文表 1 总结了我们的方法与基线模型的核心性能与效率结果,可以看到,RankMixer 在多个目标和指标上,显著优于其他 SOTA 模型。

image.png 我们对每个模型进行了详细分析:首先,简单将 DLRM 扩展到 1 亿参数量,仅能获得有限的收益,这凸显了针对推荐数据特性设计模型结构,对于提升规模化扩展性能的必要性。随后我们将 RankMixer 与 DCN、RDCN、AutoInt、DHEN 等经典交叉结构设计进行对比,发现这些模型存在参数量与计算成本失衡的问题,即使在参数量相对较小的情况下,FLOPs 已经非常高,这一设计缺陷限制了其性能表现。同时,尽管 RankMixer 取得了最优的性能,在扩展到 1 亿参数量时,其 FLOPs 在所有模型中仍处于相对适中的水平,体现了其在模型容量和计算负载之间的均衡设计。

我们还将 RankMixer 与 Hiformer、Wukong 等常用的 SOTA 规模化扩展模型进行了对比,发现在相近的参数设置下,RankMixer 不仅性能更优,计算需求也更低。

4.3 不同模型的缩放定律

原文图 2 展示了从参数量和 FLOPs 两个维度观察到的缩放定律曲线,RankMixer 模型无论是在参数量还是 FLOPs 维度,都展现出了最陡峭的缩放定律曲线,且性能始终优于其他模型。

尽管 Wukong 模型的参数量缩放曲线相对陡峭,但其计算成本的增长速度更快,因此在AUC\text{AUC}-FLOPs 曲线上,其与 RankMixer、Hiformer 的差距更大。此外,Hiformer 的性能略逊于 RankMixer,这反映出其依赖特征级的 Token 分割和注意力机制,影响了运行效率。DHEN 的缩放效果不理想,反映出其交叉结构的可扩展性有限。而 MoE 通过增加专家实现规模化扩展的策略,面临专家均衡性维持的挑战,最终导致缩放性能不佳。

具体而言,我们可以通过增加模型宽度 DD、特征 Token 数量 TT 和层数 LL,对 RankMixer 模型进行规模化扩展。实验中我们观察到了与大语言模型缩放定律一致的结论:模型质量主要与总参数量相关,不同的扩展方向(深度 LL、宽度 DD、Token 数量 TT)能带来几乎一致的性能表现。从计算效率的角度来看,更大的隐藏层维度会生成更大的矩阵乘法算子,因此相比堆叠更多层数,能实现更高的 MFU。因此,我们最终将 1 亿和 10 亿参数量模型的配置分别设置为(D=768D=768T=16T=16L=2L=2)和(D=1536D=1536T=32T=32L=2L=2)。

4.4 消融实验

image.png 我们在 1 亿参数量的 RankMixer 模型上,针对残差连接、多头 Token 混合模块开展了消融实验,结果如原文表 2 所示。可以看到,移除这些组件会导致模型性能显著下降:移除多头 Token 混合模块会丢失全局信息,每个前馈网络仅能对部分特征建模,无法实现特征交互;移除残差连接和层归一化也会导致性能恶化,降低训练稳定性,更易出现梯度爆炸或梯度消失问题。

我们进一步在原文表 3 中分析了 Token 混合策略,即从特征 Token 到前馈网络的路由策略。与多头 Token 混合对比的路由策略包括:

  1. 全拼接 MLP:将所有 Token 拼接后通过一个大型 MLP 处理,再切分为相同数量的 Token。性能下降体现了大矩阵学习的难度,以及局部信息学习能力的弱化。
  2. 全共享:不进行 Token 拆分,整个输入向量被共享并输入到每个逐 Token 前馈网络中,与 MoE 类似。性能出现显著下降,证明了与全共享输入相比,特征子空间拆分和独立建模的重要性。
  3. 自注意力:在 Token 之间应用自注意力机制实现路由。其性能略逊于多头 Token 混合,同时计算成本极高,证明了在数百个不同特征子空间之间学习相似度的难度。

image.png

4.5 稀疏 MoE 的可扩展性与专家均衡性

可扩展性:原文图 3 绘制了不同稀疏度下,稀疏 MoE 版本 RankMixer 的离线AUC\text{AUC} 收益。在高稀疏度下,结合稠密训练 - 稀疏推理与 ReLU 路由,是保留模型精度的核心。该方案让 RankMixer 的参数容量(和显存占用)扩展了 8 倍以上,同时 AUC\text{AUC} 几乎没有损失,推理吞吐量提升了 50%。

原生稀疏 MoE 的性能会随着激活专家数量的减少而单调下降,印证了我们提出的专家不均衡和训练不充分问题。加入负载均衡损失能缓解原生稀疏 MoE 的性能退化,但效果仍不及 DTSI+ReLU 路由的版本,因为问题的核心在于专家训练而非路由网络。这一结果验证了稀疏 MoE 是将 RankMixer 从当前 10 亿参数量,向未来 100 亿级部署扩展,同时不突破成本预算的核心路径。

专家均衡性与多样性:原生稀疏 MoE 通常存在专家不均衡的问题,这会导致部分专家训练不充分,最终出现 “专家死亡”(专家几乎从未被激活),仅少数固定专家被持续激活。原文图 4 展示了结合 DTSI 与 ReLU 路由后,有效解决了这一问题:稠密训练保证了绝大多数专家都能获得充足的梯度更新,避免了专家饥饿问题;ReLU 路由让 Token 之间的激活比例实现动态调整,激活比例会根据 Token 的信息含量自适应变化,与推荐数据多样化、高动态的分布特性高度契合。

4.6 线上服务成本

如何在参数量扩大两个数量级的同时,避免推理延迟爆炸?在实际系统中,延迟与吞吐量成反比,与服务机器资源成本成正比。与此前全量部署的 1600 万参数量模型(DLRM 与 DCN 结合的架构)相比,我们的 RankMixer 模型将参数量扩大了约 70 倍,达到 10 亿级。尽管参数量大幅提升,最终的推理延迟仍保持稳定,这得益于我们面向硬件的模型设计和优化策略。

参数量大幅提升时,延迟可通过以下公式分解:

Latency=#Param×FLOPs / Param ratioMFU×(Theoretical Hardware FLOPs)(13)\text{Latency} =\frac{ \#Param \times \text{FLOPs / Param ratio} }{\text{MFU} \times(\text{Theoretical Hardware FLOPs})} \tag{13}

如原文表 6 所示,两个数量级的参数量增长,通过 3.6 倍的 FLOPs / 参数量比下降、10 倍的 MFU 提升,以及基于量化的 2 倍硬件 FLOPs 提升,被逐步抵消。具体包括:

  1. FLOPs / 参数量比:表 6 第三行展示了每个参数所需的浮点运算次数。得益于模型设计,RankMixer 在参数量提升 70 倍的同时,FLOPs 仅提升了约 20 倍,FLOPs / 参数量比仅为基线的三分之一,实现了 3.6 倍的效率提升。换句话说,在相同的 FLOPs 预算下,RankMixer 可以容纳 3 倍于基线的参数量。
  2. 模型浮点运算利用率(MFU) :如表 6 所示,MFU 衡量了机器算力的利用效率。通过采用大矩阵乘法算子、优秀的并行拓扑(将并行的逐 Token 前馈网络融合为一个算子),降低内存带宽成本和开销,RankMixer 将 MFU 提升了近 10 倍,让模型从内存绑定型转变为计算绑定型。
  3. 量化:半精度(fp16)推理可以将 GPU 的理论硬件 FLOPs 峰值提升 2 倍。如前所述,RankMixer 的核心计算由多个大型矩阵乘法组成,非常适配半精度(fp16)推理。 image.png

4.7 线上性能

为验证 RankMixer 作为规模化推荐模型框架的通用性,我们在个性化排序的两大核心场景 —— 信息流推荐和广告中开展了线上实验,覆盖了个性化排序的主要应用场景。各场景的核心监控指标如下:

  • 信息流推荐:实验周期内用户的平均活跃天数(作为日活增长的替代指标)、应用累计停留时长、完播 / 点赞 / 评论总量;
  • 广告:模型质量指标(ΔAUC\Delta\text{AUC} 提升幅度)和收入指标(广告主价值 ADVV)。

表 4:抖音及抖音极速版 App 信息流推荐场景的线上 A/B 测试结果,所有指标提升均具有统计显著性。基于长期反向 A/B 测试,以及对长期反向实验的持续观测结果显示,10 亿参数量 RankMixer 模型带来的效果增益尚未收敛,本文所呈现的结果仍在持续优化提升中。 image.png 上述场景的基线模型均为 1600 万参数量的 DLRM+DCN 结合模型。我们用 10 亿参数量的 RankMixer 模型替换了基线的稠密部分,实现了 0.7% 的 AUC\text{AUC} 提升。我们在信息流推荐和广告场景均开展了线上 A/B 测试,其中信息流推荐场景长达 8 个月的 A/B 测试长期观测结果如原文表 4 所示,广告场景的结果如原文表 5 所示。

RankMixer 已在信息流推荐(10 亿参数量版本)和广告(10 亿参数量版本)两大个性化排序场景完成部署和评估,在所有核心业务指标上均取得了统计显著的正向提升。从表 4 中还可以看到,与其他活跃度层级的用户群体相比,RankMixer 对低活跃用户的提升幅度最大,活跃天数提升超 1.7412%,证明了模型强大的泛化能力。这些结果表明,RankMixer 作为统一的骨干网络,能够可靠地泛化到不同的应用场景。

5 结论

在本文中,我们提出了最新的 RankMixer 模型,该模型已在抖音信息流排序场景完成全量部署。它结合了面向异构特征交互的模型设计,以及面向服务效率的高并行化架构。实验证明了其卓越的性能和陡峭的缩放定律曲线。在抖音 App 完成全量部署后,模型实现了用户活跃天数提升 0.3%、应用使用时长提升 1% 的业务效果。