TokenMixer-Large: Scaling Up Large Ranking Models in Industrial Recommenders

标题：TokenMixer-Large：工业级推荐系统中大规模排序模型的扩展方法
单位：字节跳动

摘要

推荐模型的缩放定律已得到广泛关注，但现有的 Wukong、HiFormer 和 DHEN 等架构常因设计欠佳、硬件利用率低的问题，限制了实际的可扩展性。我们此前提出的 TokenMixer 架构（在 RankMixer 论文中首次介绍）通过轻量级 token mixing 算子替代自注意力机制，实现了效果与效率的平衡；然而该架构在深层配置中遭遇了关键瓶颈，包括次优的残差路径、梯度消失、不完整的 MoE 稀疏化以及可扩展性受限。本文提出 TokenMixer-Large，这是一款为超大规模推荐任务设计、经系统性改进的架构。通过引入混合-还原（mixing-and-reverting）操作、层间残差连接和辅助损失，我们保证了模型深度增加时梯度传播的稳定性。此外，我们融入了 Sparse Per-token MoE，实现了高效的参数扩展。TokenMixer-Large 成功将线上推理模型扩展至 7B 参数，离线实验模型扩展至 15B 亿参数。目前该模型已在字节跳动多个业务场景落地部署，取得了显著的性能提升：电商场景订单量提升 1.66%、人均预览支付交易额（GMV）提升 2.98%；广告场景广告主满意度（ADSS）提升 2.0%；直播场景营收提升 1.4%。

关键词

缩放定律；排序模型；推荐系统

1 引言

推荐系统是当今互联网生态的核心组成部分，为电商、在线广告、内容信息流、直播等核心业务提供技术支撑。因此，工业界和学术界一直在持续迭代和升级推荐模型与架构。随着深度学习的广泛应用，深度学习推荐模型（DLRM）展现出优异的效果。这类模型通常先通过嵌入层将高维稀疏的用户、物品和行为特征转换为低维稠密嵌入向量，再设计复杂的特征交互层捕捉潜在的用户-物品关联关系。

受大语言模型（LLM）成功的启发，特征交互领域的研究开始探索深度学习推荐模型的缩放定律：研究是否能通过增加模型参数和浮点运算量（FLOPs）实现性能的持续提升。早期研究主要仅关注模型宽度或交互层参数的简单扩展，缺乏精心的架构设计，因此效果受限。后续的 Wukong、HiFormer 和 DHEN 等工作尝试设计更精细的模型结构并从规模、宽度和深度维度进行扩展，但往往忽略了与硬件的协同设计和优化，导致性能表现欠佳。

目前工业界领先的最优模型结构是 TokenMixer——Transformer 的高度简化变体。该架构通过轻量级 token mixing 操作替代注意力机制，大幅降低了计算复杂度。RankMixer 将其作为骨干网络，验证了该结构在排序模型中的有效性。通过面向硬件的协同设计，该架构在保持缩放效率的同时显著提升了模型浮点运算利用率（MFU），大幅提高了硬件使用效率。尽管 TokenMixer 已在工业界取得成功（如 RankMixer），但我们发现其存在若干关键局限性：

残差设计不合理：RankMixer 通过混合机制将原始 Token 转换为新 Token，要求新 Token 数量 $T'$ 与原始 Token 数量 $T$ 保持一致，才能保证原始语义信息的完整性和传播性。此外，在其 add & norm 操作中，RankMixer 直接将混合前后的 Token 相加，这往往会引发语义错位，最终导致性能表现不佳。
模型架构不纯净：由于历史迭代过程，许多已部署的模型即便采用 TokenMixer 作为骨干网络，仍保留了大量碎片化算子（如 LHUC、DCNv2）。这些算子属于内存受限型，计算密度低但内存访问和通信开销大，降低了模型的整体 MFU。
深层模型梯度更新不足：在工业场景中，TokenMixer 通常配置为浅层结构（如 RankMixer 中仅 2 层）。随着模型深度增加，实现训练稳定性和性能提升的难度显著增大，而当前的 TokenMixer 缺乏针对深层架构的专属设计。
MoE 稀疏化不充分：尽管 RankMixer 采用 Sparse MoE 优化模型成本，但其依赖 “稠密训练、稀疏推理” 范式，无法降低训练成本。此外，所使用的 ReLU-MoE 存在激活动态性问题，推理时难以预测每个批次激活的专家数量，因此需要对激活专家采取截断或回退策略，对推理不友好。
扩展探索受限：受框架和训练效率的约束，RankMixer 的参数规模仅能扩展至约 10B。

为解决上述问题，我们设计并深入探索了 TokenMixer-Large。该框架包含 Tokenization、TokenMixer-Large Block、Sparse-Pertoken MoE 等组件。其中 TokenMixer-Large Block 由 Mixing & Reverting、Pertoken-SwiGLU、残差与归一化、层间残差与残差损失等部分构成。通过堆叠多个 TokenMixer-Large 模块，我们大幅提升了工业场景下模型的离线和线上性能。本文的贡献总结如下：

提出 TokenMixer-Large 架构。通过重新审视并重新设计 TokenMixer 中存在缺陷的残差机制，引入混合 - 还原操作。针对核心组件的大量消融实验表明，TokenMixer-Large 是一种更优的模型架构。
发现随着模型规模的扩大，底层碎片化算子带来的收益可通过堆叠多个 TokenMixer 模块实现。移除这些算子能大幅提升模型浮点运算利用率和硬件利用率。我们将这一 “纯模型” 设计理念融入 TokenMixer-Large，将广告骨干网络的浮点运算利用率提升至 60%。
为解决深层模型梯度更新不足的问题，在 TokenMixer-Large 模块中融入层间残差连接、辅助损失机制和下矩阵小初始化策略，助力深层模型实现更好的收敛效果。
将 RankMixer（TokenMixer）中的 ReLU-MoE 升级为采用统一 “稀疏训练、稀疏推理” 范式的稀疏逐 Token 混合专家版本。结合高效的算子设计和浮点 8 位（FP8）量化、Token 并行等工程优化手段，稀疏逐 Token 混合专家模块大幅降低了训练和推理成本。
对不同模型的缩放定律开展全面研究，验证了 TokenMixer-Large 的优越性。在抖音广告和电商场景的离线实验中，我们成功将模型分别扩展至 150 亿和 70 亿参数；在线上推理实验中，模型参数规模分别达到 70 亿和 40 亿。目前 TokenMixer-Large 已在字节跳动多个线上业务场景部署，服务数亿用户的同时，实现了广告主满意度提升 2%、交易额提升 2.98% 的线上业务收益。

2 相关工作

深度学习推荐模型（DLRM）的核心在于设计有效的特征交互架构。经典工作 Wide&Deep 通过融合式方法捕捉低阶和高阶特征交互，这一建模范式体现在 DeepFM、XDeepFM、PNN、DCN、DCNv2 和 FCN 等后续设计中。但作为早期设计，这些模型仍依赖 CPU 时代效率低下、延迟较高的算子，无法充分利用图形处理器（GPU）的计算能力。在不同子结构融合的基础上，DHEN 和 Wukong 进一步验证了推荐系统中的缩放定律，证明 “统一设计、逐层堆叠” 的迭代模式在推荐任务中也能实现显著的性能提升。

作为基础模型，Transformer 在 GPU 时代被广泛证明具有强通用性（计算机视觉、自然语言处理、推荐系统）和良好的缩放定律特性。在推荐领域，AutoInt、HiFormer 等模型通过融入注意力机制增强特征交互能力；RankMixer 和 MLPMixer 进一步将注意力机制简化为更轻量级的 Token 混合操作，降低了标准注意力机制带来的二次计算成本。LONGER、HSTU、MTGR 等工作也在不同程度上证明了基于 Transformer 的模型（及其变体）在序列建模和生成式推荐场景中的扩展潜力。