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 的参数规模仅能扩展至约 1B。

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

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

2 相关工作

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

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

3 方法

3.1 整体架构

如图 1 所示，TokenMixer-Large 框架的整体架构由三部分组成：1）Tokenization：将原始高维稀疏 one-hot 特征转换为稠密嵌入向量，再转换为维度对齐的语义 tokens；2）Token Mixing 和 Channel Mixing：采用 “混合&还原” 范式解决 RankMixer（TokenMixer）中的维度不匹配问题；3）Sparse-Pertoken MoE：是 RankMixer（TokenMixer）中 Pertoken- FFN / ReLU-MoE 的升级版本。最后，我们采用均值池化方法聚合输出 tokens，得到的输出表征将用于各类任务的预测。图1：TokenMixer-Large的架构。Raw tokens 包括所有原始特征以及序列聚合和提取的特征（如din/longer）。整个 Tokenmixer-Large 模型由多个Tokenmixer-Large Block 组成，每个模块的主干包括（Norm、Mixing、S-P MoE、Reverting、Norm、S-P MoE）和残差连接。

3.2 Tokenization

3.2.1 语义 Group-wise Tokenizer

工业级推荐系统通常包含用户特征、物品特征、序列特征（如覆盖用户短期(DIN)、长期(SIM)、超长期兴趣的行为历史(LONGER)）和交叉特征。第一步通过嵌入层将所有这些特征从高维稀疏空间投影到低维稠密空间，特征 $F_{i}$ 将被嵌入为不同维度的嵌入向量 $e_{i}$ ：

e_{i}= Embedding \left(F_{i}, d_{i}\right) \in \mathbb{R}^{d_{i}} \tag{1}

由于 TokenMixer-Large 是高度并行的架构设计，需要将变长的嵌入表征转换为维度对齐的特征 tokens，以方便后续计算。为实现这一目标，我们根据语义将不同维度的嵌入向量分组，对每个分组 ${G_{1}, ..., G_{T-1}}$ 内的嵌入向量进行拼接，再执行压缩和对齐操作。考虑到不同语义分组代表不同的特征属性，我们在压缩和对齐过程中采用不同的 DNN 映射，以保留其异构性：

X_{i}=\text{MLP}_{i}(\text{concat}[e_{l}, ..., e_{m}]), e_{l},...,e_{m}\in G_{i} \tag{2}

3.2.2 Global Token

除语义分组 tokens 外，我们还引入了专门用于封装全局信息的 global token。受 BERT 中[CLS] token 机制的启发，该全局 token 用于聚合全局信息并将其传播至其他 tokens：

X_{G}=\text{MLP}_{g}\left(\text{concat}\left[G_{1}, ..., G_{T-1}\right]\right) \tag{3}

TokenMixer-Large 的最终输入为全局 token 与增强 tokens 的组合：

X=\text{concat}\left[X_{G}, X_{0}, ..., X_{T-1}\right] \in \mathbb{R}^{T × D} \tag{4}

3.3 TokenMixer-Large 模块

TokenMixer-Large 采用多个 TokenMixer-Large 模块的堆叠结构，每个模块可分解为混合模块、逐 Token SwiGLU 和归一化三个组件。

3.3.1 Mixing & Reverting

RankMixer（TokenMixer）模块可表示为如下形式：

\left[...,\left[x_{t}^{(0)}, ..., x_{t}^{(H)}\right], ...\right]=\text{split}(X) \in \mathbb{R}^{T × H \times(D / H)} \tag{5}

H_{h}=\text{concat}\left[x_{1}^{(h)}, ..., x_{T}^{(h)}\right] \in \mathbb{R}^{T × D / H} \tag{6}

= \text{concat} \left[H_{1}, ..., H_{h}\right] \in \mathbb{R}^{H \times(T-D / H)} \tag{7}

H^{next}=\text{Norm}(\text{pSwiGLU}(H)+H) \in \mathbb{R}^{H × (T * D / H)} \tag{8}