RankMixer: Scaling Up Ranking Models in Industrial Recommenders

1. 研究背景与动机¶

大语言模型的成功激发了推荐系统领域对模型规模化扩展的兴趣，但工业推荐排序模型面临两大实际障碍：

1. 延迟与 QPS 约束：训练和推理成本必须严格控制在工业级延迟和高 QPS 要求范围内。
2. CPU 时代遗留架构的效率低下：现有排序模型中大量人工设计的特征交叉模块继承自 CPU 时代，无法充分利用 GPU 的并行能力，导致 Model Flops Utilization (MFU) 极低（通常仅个位数百分比），扩展性很差。

传统 DLRM 架构基于异构的人工特征交叉模块（如 DCN、DeepFM 等），其核心算子多为 memory-bound，GPU 并行度低，MFU 低。将参数量扩大后，计算成本近似线性增长，scaling law 带来的收益难以实现。

核心问题：如何设计一种硬件感知（hardware-aware）的架构，使推荐排序模型在现代 GPU 上实现高 MFU，从而将参数扩展与推理成本解耦，实现高效规模化。

2. RankMixer 整体架构¶

2.1 总体结构¶

RankMixer 由 $T$ 个输入 token 经过 $L$ 层 RankMixer Block 处理，最终通过 mean pooling 获得输出表示。每个 RankMixer Block 包含两个核心组件：

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

前向过程为：

$$S_{n-1} = \text{LN}(\text{TokenMixing}(X_{n-1}) + X_{n-1}),$$ $$X_n = \text{LN}(\text{PFFN}(S_{n-1}) + S_{n-1}),$$

其中 $\text{LN}(\cdot)$ 为 Layer Normalization，$X_n \in \mathbb{R}^{T \times D}$ 为第 $n$ 个 RankMixer Block 的输出，$X_0 \in \mathbb{R}^{T \times D}$ 由特征 token $\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_T$ 堆叠而成。

2.2 输入层与特征 Tokenization¶

推荐系统输入包含用户特征、候选特征、序列特征和交叉特征，涉及数百个异构特征字段。

语义分组 Tokenization：将特征按语义领域知识分组为若干语义连贯的 cluster，每组特征顺序拼接为一个 embedding 向量，再按固定维度 $d$ 切分为 $T$ 个 token：

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

其中 $\mathbf{e}_{\text{input}}$ 为拼接后的 embedding 向量，$N$ 为特征组数，$T$ 为切分后的 token 数，$\text{Proj}$ 将每段映射到 $D$ 维。

2.3 Multi-Head Token Mixing¶

为实现跨 token 的信息交换（即跨特征交互），引入多头 Token Mixing 模块。每个 token $\mathbf{x}_t$ 被均匀切分为 $H$ 个 head：

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

对于第 $h$ 个 head，将所有 token 的第 $h$ 个 head 部分拼接，构造 shuffled token $\mathbf{s}^h$：

$$\mathbf{s}^h = \text{Concat}\left(\mathbf{x}_1^h, \mathbf{x}_2^h, \ldots, \mathbf{x}_T^h\right).$$

输出 $S \in \mathbb{R}^{H \times \frac{TD}{H}}$，设 $H = T$ 以保持 token 数不变。该操作本质上是无参数的跨特征信息混合——通过将不同 token 的子空间重新组合，实现全局特征交互。

为什么不用自注意力？ 推荐系统中特征空间天然异构（用户 ID、物品 ID 等来自完全不同的语义空间），计算内积相似度本身没有意义；且 ID 空间规模达数百万至数亿，自注意力的二次复杂度和 attention 权重矩阵的 memory-bound 特性使其在推荐场景效率低下。Multi-Head Token Mixing 通过无参数的 shuffle 操作即可实现跨特征交互，且完全 compute-bound。

2.4 Per-Token FFN (PFFN)¶

传统 DLRM 和 DHEN 等模型在单一交互模块中混合处理所有特征空间，容易导致高频特征淹没低频和长尾特征。RankMixer 引入 逐 Token 独立参数的 FFN：

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

其中：

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

$\mathbf{W}_{\text{pffn}}^{t,1} \in \mathbb{R}^{D \times kD}$，$\mathbf{W}_{\text{pffn}}^{t,2} \in \mathbb{R}^{kD \times D}$，$k$ 为隐藏维度倍数。每个 token 拥有独立的权重矩阵和偏置，实现特征子空间的独立建模。

与 MMoE/Transformer 的区别：Per-Token FFN 中每个"专家"处理的是不同的输入（不同 token），而 MMoE 中所有专家共享同一输入；Transformer 中不同输入共享同一 FFN。RankMixer 同时分离了输入和参数。

整体参数和计算量公式（Dense 版本）：

$$\text{#Param} \approx 2kLTD^2, \quad \text{FLOPs} \approx 4kLTD^2.$$

2.5 Sparse MoE in RankMixer¶

为进一步提升扩展 ROI，可将 Per-Token FFN 替换为 Sparse MoE 结构。但 vanilla Sparse MoE 在 RankMixer 中表现退化，原因：

1. 均匀 top-k 路由：对所有 token 平等分配 expert 预算，浪费了低信息量 token 上的计算，饿死了高信息量 token。
2. 专家训练不充分：Per-Token FFN 本身已按 token 数乘以参数量，再叠加非共享 expert 会使 expert 数量爆炸，导致路由极度不均衡。

解决方案：

ReLU Routing：将常见的 Top-k + softmax 替换为 ReLU gate + 自适应 $\ell_1$ 惩罚：

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

其中 $N_e$ 为每个 token 的 expert 数。ReLU Routing 允许高信息量 token 激活更多 expert，稀疏性通过正则项控制：

$$\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}.$$

Dense-Training / Sparse-Inference (DTSI-MoE)：训练时使用两个 router $h_{\text{train}}$ 和 $h_{\text{infer}}$，正则项仅施加于 $h_{\text{infer}}$；推理时仅使用 $h_{\text{infer}}$。这样训练时每个 expert 都能得到充分梯度更新，推理时保持稀疏以降低成本。

3. 扩展策略¶

RankMixer 可沿四个正交维度扩展：Token 数 $T$、模型宽度 $D$、层数 $L$、Expert 数 $E$。

实验发现与 LLM scaling law 一致的结论：模型质量主要与总参数量相关，不同扩展方向（深度 $L$、宽度 $D$、token 数 $T$）在相同参数量下效果相近。从计算效率角度，增大宽度 $D$ 能生成更大的矩阵乘形状，从而实现更高 MFU，优于堆叠更多层。

最终配置：

• RankMixer-100M：$D=768, T=16, L=2$
• RankMixer-1B：$D=1536, T=32, L=2$

4. 实验¶

4.1 实验设置¶

• 数据集：抖音推荐系统的万亿级工业数据集，包含 300+ 特征字段（数值特征、ID 特征、交叉特征、序列特征），涉及数十亿用户 ID 和数亿视频 ID。数据量为每天数万亿条记录，实验在两周数据上进行。
• 评估指标：
• Finish AUC/UAUC：用户是否完整观看视频
• Skip AUC/UAUC：用户是否快速跳过视频

• AUC 提升 0.0001 即被视为统计显著

• Dense-Param：不含 sparse embedding 的 dense 参数量

• MFU：模型实际 FLOPs 消耗 / 硬件理论 FLOPs 峰值

4.2 与 SOTA 方法对比（~100M 参数，Table 1）¶

结论：RankMixer-100M 在所有指标上显著超越同参数量级的所有 baseline，且 FLOPs 较 HiFormer 和 Wukong 更低。DCN、RDCN、AutoInt、DHEN 等经典交叉结构模型在扩展到 100M 参数时存在参数量与计算成本的不平衡。RankMixer-1B 进一步将性能大幅提升。

4.3 Scaling Law 对比（Figure 2）¶

从参数量和 FLOPs 两个维度观察 AUC gain 的 scaling curve。RankMixer 在两个维度上均展现出最陡峭的 scaling law，持续优于 DLRM、DHEN、HiFormer、Wukong 和 MoE。Wukong 虽然参数 scaling 曲线较陡，但其计算成本增长更快，在 AUC vs FLOPs 维度上差距更大。

4.4 消融实验¶

Table 2: RankMixer-100M 组件消融

Multi-Head Token Mixing 移除后性能下降最大（-0.50%），因为每个 FFN 只能建模局部信息，失去了全局交互能力。Per-Token FFN 替换为共享 FFN 也显著下降（-0.31%），证明独立参数建模不同特征子空间的重要性。

Table 3: Token Mixing 路由策略对比

All-Concat-MLP（先拼接所有 token 再用大 MLP 处理）和 All-Share（不做切分、所有 token 共享同一输入送入 Per-Token FFN，类似 MoE）性能均显著下降。Self-Attention 性能仅略低于 Multi-Head Token Mixing，但参数增加 16%、FLOPs 增加 71.8%，效率远不如。

4.5 Sparse MoE 可扩展性（Figure 3）¶

在不同 sparsity ratio（1, 1/2, 1/4, 1/8 激活参数比例）下比较 RankMixer 变体：

• Dense-Training + ReLU Routing 的 SMoE RankMixer：在 sparsity 为 1/8 时几乎无 AUC 损失，同时推理时间节省 50%+、内存减少 8 倍以上。
• Vanilla SMoE RankMixer：随 sparsity 增大单调下降。
• 添加 Balance Loss 的 Vanilla SMoE：缓解退化但仍不及 DTSI + ReLU 方案，因根本问题在于 expert 训练不充分而非仅路由不均。

4.6 在线 Serving 成本（Table 6）¶

关键洞察： 1. FLOPs/Param ratio 降低 3.6 倍：参数增加 70 倍，FLOPs 仅增加 20 倍，即每参数的计算代价大幅降低。 2. MFU 提升 10 倍：从 4.47% 到 44.57%，从 memory-bound 转为 compute-bound。 3. FP16 量化：RankMixer 的大矩阵乘形状天然适合半精度推理，硬件峰值 FLOPs 翻倍。

三项因素共同作用，使得参数增加两个数量级而推理延迟基本不变（14.5ms -> 14.3ms）。

4.7 线上 A/B 测试¶

Table 4: 抖音信息流推荐 A/B 测试结果

低活跃用户的提升最为显著（Active Day +1.74%），表明 RankMixer 作为统一骨干架构具有强泛化能力。

Table 5: 广告场景 A/B 测试

RankMixer 已在抖音信息流推荐和广告两个核心个性化排序场景全量部署，整体活跃天数增长 0.3%，App 使用时长增长 1.08%。

5. 核心贡献总结¶

1. 硬件感知架构设计：Multi-Head Token Mixing（无参数跨特征交互）+ Per-Token FFN（独立参数特征子空间建模），实现高 GPU 并行度和 MFU。
2. Sparse MoE 扩展：DTSI + ReLU Routing 策略解决 expert 训练不足和路由不均问题，在 8 倍稀疏度下几乎无精度损失。
3. 工业验证：在抖音万亿级数据上验证 scaling law，1B Dense 参数模型全量上线，推理延迟与 16M 基线持平，业务指标全面提升。
4. 参数与推理成本解耦：通过高 MFU + 低 FLOPs/Param ratio + FP16 量化，实现参数规模两个数量级增长而推理成本不变。