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 缩放定律的相关研究必须解决以下问题：
- 架构必须与硬件特性对齐，最大化现代 GPU 上的 MFU 和计算吞吐量；
- 模型设计必须适配推荐数据的特性，例如数百个特征域之间的异构特征空间、个性化的跨特征交互。
为解决上述挑战，我们提出了面向硬件感知的模型设计方法 RankMixer。RankMixer 的核心设计基于两个可扩展组件：
- 多头 Token 混合：仅通过无参数算子实现跨 Token 的特征交互，该策略在性能和计算效率上均优于自注意力机制；
- 逐 Token 前馈网络：通过为不同的特征子空间建模分配独立的参数，大幅扩展了模型容量，同时解决了跨特征空间的主导性问题。这类前馈网络也与推荐数据的模式高度契合，实现了更优的缩放特性。
为进一步提升大模型的 ROI，我们将 per-token FFN 扩展为 SparseMoE 结构。通过针对不同数据，为每个 Token 仅动态激活一部分专家，我们可以在计算成本增加极少的前提下，显著提升模型容量。
RankMixer 采用了与 Transformer 类似的高并行架构，但克服了基于自注意力的特征交互的几个关键缺陷：训练效率低下、跨空间 ID 相似度建模时的组合爆炸问题，以及注意力权重矩阵带来的严重内存受限问题。同时，在相同 FLOPs 下，RankMixer 比原生 Transformer 具备更强的模型容量和学习能力。
在抖音推荐系统的生产部署中，我们验证了该架构可以将模型参数量扩大 100 倍以上，同时推理延迟相比此前的基线模型更短。这一成果的核心，是 RankMixer 架构实现了参数量增长与 FLOPs 的解耦，以及通过高 MFU 和工程优化，实现了 FLOPs 增长与实际成本的解耦。
本文的主要贡献可总结如下：
- 我们提出了一种遵循硬件感知设计理念的全新架构 RankMixer。设计了多头 Token 混合和 per-token FFN 策略，以高效捕捉异构特征交互；同时采用动态路由策略，提升了 RankMixer 中 SparseMoE 的可扩展性。
- 借助高 MFU 和性能优化，我们在不增加推理成本的前提下，将模型参数量扩大了 70 倍（包括提升 MFU 和量化优化）。
- 我们开展了丰富的离线和线上实验，在万亿级工业推荐数据集上探究了模型的缩放定律。RankMixer 模型已成功在抖音信息流推荐排序中实现全流量部署，实现了用户活跃天数提升 0.3%、应用使用时长提升 1.08% 的效果。
