一、 TokenMixer-Large

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

1. 推荐模型的 scaling laws 已受到广泛关注，但 Wukong 、HiFormer、 DHEN 等现有架构常因设计欠佳、硬件利用率低，限制了实际的可扩展性。我们此前在 RankMixer 论文中提出的 TokenMixer 架构，通过轻量级的 token mixing operator 替代 self-attention，实现了效果与效率的平衡；然而，该架构在 deeper configurations 中遭遇了关键瓶颈，包括 residual paths 不合理、梯度消失、MoE sparsification 不彻底、以及 scalability 受限等问题。

   本文提出 TokenMixer-Large，这是一款为超大规模推荐任务设计的、经系统性改进的架构。通过引入 mixing-and-reverting 操作、inter-layer residuals 和辅助损失函数，我们确保了 model depth 增加时梯度传播的稳定性。此外，我们融入了 Sparse Per-token MoE 机制，实现了模型参数的高效扩展。TokenMixer-Large 成功扩展到 7 billion 参数规模从而用于线上流量、，离线实验中更是达到 15 billion 参数。该模型目前已在 ByteDance 多个业务场景落地部署，取得了显著的性能提升：

   • 电商场景中订单量提升 1.66%、人均预览支付商品交易总额（GMV）提升 2.98%。

   • 广告场景中广告主满意度评分（ADSS）提升 2.0%。

   • 直播场景中营收增长 1.4%。

2. 推荐系统在当今互联网生态中占据核心地位，为电商、在线广告、内容信息流、直播等核心应用提供技术支撑。因此，工业界和学术界都在持续迭代与升级推荐模型及架构。随着深度学习的广泛应用，深度学习推荐模型（Deep Learning Recommendation Models: DLRMs）展现出卓越的效果。这类模型通常先通过 embedding layer，将高维稀疏的 user, item, and behavioral features 转换为低维稠密的 embeddings 向量；随后设计复杂的 feature interaction layers ，从而捕获潜在的 user-item relationships 。

   受大语言模型（Large Language Models: LLMs）成功的启发，feature interaction 领域的研究开始探索 DLRMs 的 scaling laws，即研究是否能通过增加模型参数和浮点运算量（FLOPs）来实现性能的持续提升。早期研究主要聚焦于单纯扩展 model width 或 interaction layer parameters，缺乏精细化的架构设计，导致效果受限。后续的 Wukong、HiFormer、DHEN 等研究尝试设计更精细的模型结构，并从 size、width 和 depth 等维度进行扩展，但往往忽视了与硬件适配的 co-design 和 optimization，最终性能表现欠佳。

   目前，工业界领先的 SOTA 模型架构是 TokenMixer ——一种高度简化的 Transformer 变体。该架构通过轻量级的 token mixing 操作替代 attention，大幅降低了计算复杂度。RankMixer （《Rankmixer: Scaling up ranking models in industrial recommenders》）将其作为骨干网络，验证了该结构在 ranking 模型中的有效性。通过硬件感知的 co-design ，TokenMixer 在保持 scaling 效率的同时，显著提升了 Model FLOPs Utilization: MFU，大幅改善了硬件使用效率。尽管 TokenMixer 已在工业界取得成功（如 RankMixer ），但我们发现其存在若干关键局限性：

   • Residual Design 不合理：RankMixer 通过 mixing 机制将 original tokens 转换为 new tokens，要求 new tokens 数量 $T^{\prime }$$T^\prime$ 与 original tokens 数量 $T$$T$ 保持一致，以确保原始语义信息的完整性和传播性。此外，在其 add & norm 操作中，RankMixer 直接将 mixing 之前的 tokens 和之后的 tokens 相加，这往往会引发语义错位（semantic misalignment），最终导致性能欠佳。

   • Model Architecture 冗余：受历史迭代过程影响，许多已部署的模型即便采用 TokenMixer 作为骨干网络，仍保留了大量零散的算子，如 LHUC（《Learning hidden unit contributions for unsupervised acoustic model adaptation》）、 DCNv2。这类算子受内存限制，计算密度低但内存访问和通信开销大，降低了模型的整体 MFU。

   • Deep Models 的梯度更新不足：在工业场景中，TokenMixer 的配置通常较浅（如 RankMixer 中仅 2 layers）。随着模型深度增加，实现训练稳定性和性能提升的难度显著增大，而当前的 TokenMixer 缺乏针对 deep 架构的专属设计。

   • MoE Sparsification 不充分：尽管 RankMixer 采用 Sparse MoE 来优化模型成本，但其依赖 "Dense Train, Sparse Infer" 的范式，无法降低训练成本。此外，所使用的 ReLU-MoE （《ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing》）存在 activation dynamism 的问题，推理时难以预测每一个 batch 所激活的 experts 数量，因此需要对 activated experts 采取截断策略或 fallback 策略，对 inference 不友好。

     RankMixer 采用的 Sparse MoE 仅在 Inference 阶段启用专家稀疏激活，而 Training 阶段仍会让所有专家全部参与计算，即训练时保持模型的密集化运行，推理时才通过路由选择部分专家激活以降低推理成本，这一运行逻辑直接定义了其 "Dense Train, Sparse Infer" 的范式特征。这种设计下，训练过程无法利用 MoE 的稀疏性优势，所有专家的参数都会被同步更新，硬件计算和内存开销与同规模的密集模型无异，因此无法降低训练成本，仅能在推理阶段通过稀疏激活优化资源消耗。

     activation dynamism 核心指不同 batch数据在推理时，被路由机制激活的专家数量、具体专家索引会呈现无规律的动态变化。ReLU-MoE 该问题的根源来自其基于 ReLU 激活的路由设计：ReLU-MoE 以 ReLU 作为路由函数的核心组件，路由对输入样本的特征分布高度敏感。而工业推荐场景的输入（如用户、商品、行为特征）具有高度稀疏性和分布随机性（不同 batch 的样本特征分布差异大）。这导致推理时，针对不同 batch 的数据，路由模块输出的激活专家选择结果不稳定 —— 激活的专家数量时多时少、具体激活哪几个专家无固定规律，无法提前预判每一个 batch 所需激活的专家数量和组合。

     推理阶段的硬件资源（如 GPU 显存、计算核心）是提前固定分配的，而 ReLU-MoE 的 activation dynamism 导致单个 batch 实际激活的专家数量不可控：

     • 若某一 batch 被激活的专家数量超过硬件提前分配的资源上限，会直接引发显存溢出、计算资源不足，导致推理中断；

     • 若无统一的激活专家数量约束，推理引擎无法提前做资源调度和计算优化，会出现资源利用率剧烈波动。

     因此必须通过截断策略（将单 batch 激活的专家数量强行限制在预设阈值内，超出部分直接舍弃）或 fallback 策略（当激活专家数量超出阈值时，回退到预设的通用专家 / 密集计算模式），来保证推理过程的稳定性，避免资源异常。

   • Scaling 的探索不足：受框架和训练效率的约束，RankMixer 的参数规模仅能扩展至约 1-Billion。

   为解决上述问题，我们设计并深入研究了 TokenMixer-Large。该框架包含 Tokenization、TokenMixer-Large Block、 Sparse-Pertoken MoE 等模块。其中，TokenMixer-Large Block 由 Mixing & Reverting、Pertoken-SwiGLU、 Residual & Normalization、Interval Residual & Residual Loss 等组件构成。通过堆叠多个 TokenMixer-Large Blocks，我们大幅提升了工业场景下模型的离线性能和在线性能。本文的研究贡献总结如下：

   • 提出了 TokenMixer-Large 架构。通过重新审视并重新设计 TokenMixer 中存在缺陷的 residual 机制，引入 Mixing & Reverting 操作。针对核心组件的大量消融实验表明，TokenMixer-Large 是一种更优的模型架构。

   • 我们发现，随着模型规模扩大，low-level 零散算子带来的增益可通过堆叠多个 TokenMixer Blocks 来实现。移除这类算子能大幅提升 model MFU 和硬件利用率。我们将这一 "pure model" 设计理念融入 TokenMixer-Large，将广告骨干网络的 MFU 提升至 60%。

   • 为了解决 deep models 中梯度更新不足问题，我们在 TokenMixer-Large Block 中融入 interval residuals、auxiliary loss 机制、以及 down-matrix small initialization 策略，促进 deep models 更好地收敛。

   • 我们将 Rankmixer (TokenMixer) 中的 relu-MoE 升级 Sparse-Pertoken MoE，它采用统一的 "Sparse Train, Sparse Infer" 范式。结合高效的算子设计以及 FP8 量化、Token Parallel 等工程优化手段， Sparse-Pertoken MoE 大幅降低了训练成本和推理成本。

   • 我们开展了模型 scaling laws 的综合研究，验证了 TokenMixer-Large 的优越性。我们在 Douyin 的广告和电商场景的离线实验中，成功将模型分别扩展至 15B 和 7B 参数；在线上流量实验中，相应扩展至 7B 和 4B 参数。目前，TokenMixer-Large 已在 ByteDance 多个线上场景部署，服务数亿用户的同时，实现了 ADSS 提升 2% 、 GMV 提升 2.98% 的线上收益。

1.1 相关工作

1. DLRMs 的核心在于设计有效的 feature interaction 架构。

   • Wide and Deep作为经典研究之一，通过集成式方法同时捕获 low-order and high-order 的 feature interactions。这一建模范式在 DeepFM、XDeepFM、PNN、DCN、DCNv2、FCN 等后续设计中得到体现。但作为早期设计，这类模型仍依赖CPU 时代效率低下、延迟较高的算子，无法充分发挥 GPU 的计算能力。

   • 在不同子结构进行集成的基础上，DHEN 和 Wukong 进一步验证了推荐系统中的 scaling laws，证明 "unified design, stacked layers" 的迭代式模式在推荐任务中也能实现显著的性能提升。

2. 作为基础模型，Transformer 在 GPU 时代被广泛证实具有强大的通用性（CV、NLP、推荐系统RS）和良好的 scaling-law 特性。

   • 在推荐领域，AutoInt、HiFormer 等模型通过融入 attention 机制来增强 feature interaction 能力。

   • RankMixer （《Rankmixer: Scaling up ranking models in industrial recommenders》）和MLPMixer （《Mlp-mixer: An all-mlp architecture for vision》）进一步将注意力机制简化为更轻量的 token-mixing 操作，降低了标准注意力机制的平方计算成本。

   • LONGER、HSTU、MTGR等研究也在不同程度上证明了基于 Transformer 的模型（及其变体）在 sequential modeling 生成式推荐场景中的 scaling 潜力。

1.2 方法

1.2.1 整体架构

1. 如 Figure 1 所示，TokenMixer-Large 框架的整体架构由三部分组成：

   • 1) Tokenization：将原始高维稀疏的 one-hot features 转换为 dense embeddings ，再进一步转换为维度对齐的 semantic tokens 。

   • 2) Token Mixing & Channel Mixing：采用 "Mixing and Reverting" 范式，解决 RankMixer (TokenMixer) 中存在的维度不匹配问题；

   • 3) Sparse-Pertoken MoE：是 RankMixer (TokenMixer) 中提及的 Pertoken-FFN/relu-MoE 的升级版本。

   最后，我们采用均值池化方法聚合 output tokens，输出的 representation 将用于各类任务的预测。

   这是 RankMixer V2，它的基本架构与 RankMixer 差别不大。更多的是细节上的优化。

   

1.2.2 Tokenization

1. Semantic Group-wise Tokenizer：工业级推荐系统通常包含 user features、item features、sequence features 和 cross features。其中，sequence features 包括：覆盖了用户短期兴趣（DIN）、长期兴趣（SIM）、超长期兴趣（LONGER）的行为历史。第一步通过 embedding layer 将所有这些特征从高维稀疏空间投影到低维稠密空间，特征 $F_i$$F_i$ 将被嵌入为维度为 $d_i$$d_i$ 的 embedding 向量 ${\overrightarrow{\mathbf{e}}}_i$$\mathbf{\vec e}_i$：

   $${\overrightarrow{\mathbf{e}}}_i=\text{Embedding}(F_i,d_i)\in {\mathbb{R}}^{d_i}$$

   由于 TokenMixer-Large 是高度并行的架构设计，需要将变长的 embedding representations 转换为维度对齐的 feature tokens，以便后续计算。为实现这一目标，我们根据 semantic meanings 对不同特性的 embeddings 进行分组，得到分组 $\{G_1,\cdots ,G_{T-1}\}$$\{G_1,\cdots,G_{T-1}\}$；然后将每个分组 $G_i$$G_i$ 内的 embedding 拼接后，执行压缩和对齐操作。考虑到不同 semantic groups 代表不同的特征属性，我们在压缩和对齐过程中采用不同的 DNN mappings ，以保留其异构性：

   $${\mathbf{X}}_i={\text{MLP}}_i\left(\text{concat}[{\overrightarrow{\mathbf{e}}}_{i_1},\cdots ,{\overrightarrow{\mathbf{e}}}_{i_n}]\right)\in {\mathbb{R}}^{1\times D},{\overrightarrow{\mathbf{e}}}_{i_1},\cdots ,{\overrightarrow{\mathbf{e}}}_{i_n}\in G_i,i\in \{1,2,\cdots ,T-1\}$$

   其中：$D$$D$ 为对齐后的维度，一共 $T-1$$T-1$ 个 groups。

   如何分组？作者并未详细说明。根据 RankMixer 的论文：

   • 对于序列特征，每一种序列都构成一个分组：

     • 短期序列，用一个 Sequence Module 来抽取短期兴趣 representation ${\overrightarrow{\mathbf{e}}}_{s_1}$$\mathbf{\vec e}_{s_1}$ ，对应于 $G_1$$G_1$。

     • 长期序列，用另一个 Sequence Module 来抽取长期兴趣 representation ${\overrightarrow{\mathbf{e}}}_{s_2}$$\mathbf{\vec e}_{s_2}$ ，对应于 $G_2$$G_2$。

     • 超长期序列，用第三个 Sequence Module 来抽取超长期兴趣 representation ${\overrightarrow{\mathbf{e}}}_{s_3}$$\mathbf{\vec e}_{s_3}$ ，对应于 $G_3$$G_3$。

   • 对于非序列特征，User features 、candidate features、 Cross Features 分别构成了各自的分组 $G_4,G_5,G_6$$G_4, G_5, G_6$ 。

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

   $${\mathbf{X}}_G={\text{MLP}}_g(\text{concat}[G_1,\cdots ,G_{T-1}])\in {\mathbb{R}}^{1\times D}$$

   TokenMixer-Large 的 final input 为 global token 与 enhanced tokens 的组合：

   $$\mathbf{X}=\text{concat}[{\mathbf{X}}_G,{\mathbf{X}}_0,\cdots ,{\mathbf{X}}_{T-1}]\in {\mathbb{R}}^{T\times D}$$

   在实验章节， Global Token 的 auc gain 较小，仅仅 0.02% 。

1.2.3 TokenMixer-Large Block

1. TokenMixer-Large 采用 multiple TokenMixer-Large blocks 的堆叠结构，每个 block 可分解为 mixing module 、 pertoken SwiGLU 、 normalization 等三个组件。

2. Mixing & Reverting：一个 RankMixer (TokenMixer) block 可以表示为如下的形式：

   $$\begin{array}{c}[\cdots ,[{\overrightarrow{\mathbf{x}}}_t^{(1)},{\overrightarrow{\mathbf{x}}}_t^{(2)},\cdots ,{\overrightarrow{\mathbf{x}}}_t^{(H)}],\cdots ]=\text{split}(\mathbf{X})\in {\mathbb{R}}^{T\times H\times (D/H)}\\ {\mathbf{H}}_h=\text{concat}[{\overrightarrow{\mathbf{x}}}_1^{(h)},{\overrightarrow{\mathbf{x}}}_2^{(h)},\cdots ,{\overrightarrow{\mathbf{x}}}_T^{(h)}]\in {\mathbb{R}}^{1\times (T\ast D/H)}\\ \mathbf{H}=\text{concat}[{\mathbf{H}}_1,{\mathbf{H}}_2,\cdots ,{\mathbf{H}}_H]\in {\mathbb{R}}^{H\times (T\ast D/H)}\\ {\mathbf{H}}^{\text{next}}=\text{Norm}(\text{pSwiGLU}(\mathbf{H})+\mathbf{H})\in {\mathbb{R}}^{H\times (T\ast D/H)}\end{array}$$

   第一行：对每个 ${\mathbf{X}}_i$$\mathbf X_i$ 的元素进行分组，分为 $H$$H$ 组，每组包含 $D/H$$D/H$ 个元素。

   第二行：将每个 ${\mathbf{X}}_i$$\mathbf X_i$ 的第 $h$$h$ 组拼接在一起，构成一个 $T\times D/H$$T\times D/H$ 维的向量。

   第三行：将所有 $H$$H$ 个 ${\mathbf{H}}_i$$\mathbf H_i$ 堆叠在一起，拼接为一个矩阵。

   原始 input 的 tokens 数量为 $T$$T$ ，而第 $i$$i$ 层混合后的 tokens 数量为 $H$$H$ 。如上所示，除非每层的 $H$$H$ 保持不变且 $T=H$$T = H$ 相等，否则 residual connections 无法在层间平稳传播。

   为克服这一限制，我们设计的 TokenMixer-Large 采用特殊的双层 TokenMixer 结构：

   • 第一层负责原始 tokens 之间的信息混合。

     $$\begin{array}{c}[\cdots ,[{\overrightarrow{\mathbf{x}}}_t^{(1)},{\overrightarrow{\mathbf{x}}}_t^{(2)},\cdots ,{\overrightarrow{\mathbf{x}}}_t^{(H)}],\cdots ]=\text{split}(\mathbf{X})\in {\mathbb{R}}^{T\times H\times (D/H)}\\ {\mathbf{H}}_h=\text{concat}[{\overrightarrow{\mathbf{x}}}_1^{(h)},{\overrightarrow{\mathbf{x}}}_2^{(h)},\cdots ,{\overrightarrow{\mathbf{x}}}_T^{(h)}]\in {\mathbb{R}}^{1\times (T\ast D/H)}\\ \mathbf{H}=\text{concat}[{\mathbf{H}}_1,{\mathbf{H}}_2,\cdots ,{\mathbf{H}}_H]\in {\mathbb{R}}^{H\times (T\ast D/H)}\\ {\mathbf{H}}^{\text{next}}=\text{Norm}(\text{pSwiGLU}(\mathbf{H})+\mathbf{H})\in {\mathbb{R}}^{H\times (T\ast D/H)}\end{array}$$

   • 第二层专门还原 mixed tokens 的维度。

     $$\begin{array}{c}[\cdots ,[{\overrightarrow{\mathbf{x}}}_1^{{}^{\prime }(h)},{\overrightarrow{\mathbf{x}}}_2^{{}^{\prime }(h)},\cdots ,{\overrightarrow{\mathbf{x}}}_T^{{}^{\prime }(h)}],\cdots ]=\text{split}\left({\mathbf{H}}^{\text{next}}\right)\in {\mathbb{R}}^{H\times T\times (D/H)}\\ {\mathbf{X}}_t^{\text{revert}}=\text{concat}[{\overrightarrow{\mathbf{x}}}_t^{{}^{\prime }(1)},{\overrightarrow{\mathbf{x}}}_t^{{}^{\prime }(2)}\cdots ,{\overrightarrow{\mathbf{x}}}_t^{{}^{\prime }(H)}]\in {\mathbb{R}}^{1\times D}\\ {\mathbf{X}}^{\text{revert}}=\text{concat}[{\mathbf{X}}_1^{\text{revert}},{\mathbf{X}}_2^{\text{revert}},\cdots ,{\mathbf{X}}_T^{\text{revert}}]\in {\mathbb{R}}^{T\times D}\\ {\mathbf{X}}^{\text{next}}=\text{Norm}(\text{pSwiGLU}({\mathbf{X}}^{\text{revert}})+\mathbf{X})\in {\mathbb{R}}^{T\times D}\end{array}$$

     注意：这里的残差连接是 ${\mathbf{X}}^{\text{revert}}$$\mathbf X^\text{revert}$ 和 $\mathbf{X}$$\mathbf X$ 。

   这种对称的 "mixing–reverting" 设计确保了模块间 input and output 的维度一致性，从而建立了从 initial input 到 deep network 的连续的信号通路。它实现了每层稳定的 residual connections ，并有效防止了梯度信息的丢失。与 RankMixer (TokenMixer) 的详细对比见 1.3 节。

   mixing 阶段实现了信息的 cross-token 融合。在 mixing-reverting 阶段引入了操作 pSwiGLU ，促进了信息融合。

3. Pertoken SwiGLU：RankMixer (TokenMixer) 引入了一种参数隔离的 feed-forward network: FFN 架构，即 per-token FFN，旨在建模不同 tokens 间的特征异构性。在 TokenMixer-Large 中，我们将per-token FFN 升级为 pertoken SwiGLU：

   $$\begin{array}{c}\text{pSwiGLU}(\cdot )={\text{FC}}_{\text{down}}(\text{Swish}({\text{FC}}_{\text{gate}}(\cdot ))\odot {\text{FC}}_{\text{up}}(\cdot ))\\ {\text{FC}}_i\left({\overrightarrow{\mathbf{x}}}_t\right)={\mathbf{W}}_i^t{\overrightarrow{\mathbf{x}}}_t+{\overrightarrow{\mathbf{b}}}_i^t,i\in \{\text{up, gate, down}\}\end{array}$$

   其中：

   • ${\text{FC}}_i(\cdot )$$\text{FC}_i(\cdot)$ 为 MLP ，$\{{\mathbf{W}}_{\text{up}}^t,{\mathbf{W}}_{\text{gate}}^t\}\in {\mathbb{R}}^{nD\times D}$$\left\{\mathbf W^t_\text{up}, \mathbf W^t_\text{gate}\right\} \in \mathbb R^{nD\times D}$ ，${\mathbf{W}}^t\in {\mathbb{R}}^{D\times nD}$$\mathbf W^t\in \mathbb R^{D\times nD}$；$\{{\overrightarrow{\mathbf{b}}}_{\text{up}}^t,{\overrightarrow{\mathbf{b}}}_{\text{gate}}^t\}\in {\mathbb{R}}^{nD}$$\left\{\mathbf{\vec b}_\text{up}^t, \mathbf{\vec b}_\text{gate}^t\right\} \in \mathbb R^{nD}$，${\overrightarrow{\mathbf{b}}}_{\text{down}}^t\in {\mathbb{R}}^D$$\mathbf{\vec b}_\text{down}^t\in \mathbb R^D$ 。$n$$n$ 为超参数。

   • $\text{Swish}(\cdot )$$\text{Swish}(\cdot)$ 为 Swish 激活函数：

     $${\text{Swish}}_{\beta }(x)=x\sigma (\beta x),\sigma (x)=\text{sigmoid}(x)=\frac{1}{1+\exp (-x)}$$

     当 $\beta =1$$\beta = 1$ 时，Swish 激活函数简化为 $\text{Swish}(x)=\frac{x}{1+\exp (-x)}$$\text{Swish}(x) = \frac{x}{1 + \exp(-x)}$ 。

   • ${\overrightarrow{\mathbf{x}}}_t\in {\mathbb{R}}^D$$\mathbf{\vec x}_t\in \mathbb R^D$ 为第 $t$$t$ 个 token。

   实验表明，Pertoken SwiGLU 比 pertoken FFN 在 AUC 上提升了 0.1% 。

   • SwiGLU 的门控相当于为每个 token 的特征分配 “动态权重”，能更精准地捕捉推荐场景中 “用户兴趣 - 物品特征” 的复杂非线性关联，提升模型的 CTR/CVR 预测准确性。

   • 普通 FFN 的 ReLU 激活函数会导致梯度传播中断，尤其在深层模型中，训练后期梯度几乎无法传递到浅层；SwiGLU 采用 Swish 激活，是处处可导的平滑非线性函数，且梯度始终非零，能让梯度更顺畅地在深层模型中传播。

4. Residual & Normalization：标准残差连接如 Figure 1 所示。此外，我们用更轻量的 RMSNorm 替代了标准的 LayerNorm。通过移除 LayerNorm 中的均值中心化（mean-centering）步骤， RMSNorm 在保留核心归一化能力的同时，大幅降低了计算开销。然而，在 RankMixer (TokenMixer) 的 Post-Norm 设计下，我们经常遇到数值爆炸和梯度不稳定的问题。为解决这些稳定性问题，我们采用 Pre-Norm 设计替代原有的 Post-Norm 架构，详细对比见附录 A.4。

   给定一个向量 $\overrightarrow{\mathbf{x}}=(x_1,\cdots ,x_N{)}^{\mathrm{\top }}$$\mathbf {\vec x} = (x_1,\cdots,x_N)^\top$：

   • LayerNorm：

     $$\begin{array}{c}\mu =\frac{1}{N}\sum _{i=1}^N{x}_i,{\sigma }^2=\frac{1}{N}\sum _{i=1}^N(x_i-\mu {)}^2\\ \widetilde{\overrightarrow{\mathbf{x}}}=\frac{\overrightarrow{\mathbf{x}}-\mu }{\sqrt{{\sigma }^2+ϵ}}\odot \overrightarrow{\gamma }+\overrightarrow{\beta }\end{array}$$

     其中：$ϵ>0$$\epsilon > 0$ 为一个小的正数防止数值异常；$\overrightarrow{\gamma },\overrightarrow{\beta }$$\vec\gamma,\vec\beta$ 为待学习的参数。

   • RMSNorm：

     $$\begin{array}{c}\text{RMS}(\overrightarrow{\mathbf{x}})=\sqrt{\frac{1}{N}\sum _{i=1}^N{x}_i^2}\\ \widetilde{\overrightarrow{\mathbf{x}}}=\frac{\overrightarrow{\mathbf{x}}}{\text{RMS}(\overrightarrow{\mathbf{x}})}\odot \overrightarrow{\gamma }\end{array}$$

     .

5. Inter-Residual & Auxiliary Loss：如 Figure 2 所示，除标准的 residual connections 外，我们还加入了 inter-residual connections（通常每隔 2 层或 3 层添加），以加强层间的信息流动，其主要作用是解决梯度消失问题。通过增强 lower-layer features 向 higher layers 的传递，加速了 early layers 中参数的收敛，缓解了 deep networks 中梯度的逐渐衰减。不建议在最后一层添加 interval-residual connections，因为最后一层的主要作用是提取高度抽象的、high-level 的特征，为后续的分类任务提供精准的支持，引入过多原始的 lower-level information 可能会干扰这一抽象过程，从而损害最终的任务性能。

   同时，我们将 lower-layers 计算得到的 logits 与 higher layers 的 logits 相结合，计算 joint loss ，形成轻量级的辅助损失（auxiliary loss）。该机制本质上让 lower layers 学会 "estimate deviations in higher-layer features"，从而增强其 feature representation 能力，有助于防止随着网络深度增加，lower-layer 训练不充分的问题，实现更充分的 parameter learning，让 deep model 的性能更稳定。

   这个 joint loss 如何设计？论文并未说明。读者猜测：就是标准的 logloss：

   • 对于 MainTower：基于 main logit 来计算 logloss。

   • 对于 AuxLoss1：Block-6 和 Block-4 的 output 都馈入到 a shared AuxTower1 中，然后将它们的 logit 相加，然后计算 logloss。

   • 对于 AuxLoss2：Block-6 、 Block-4 和 Block-2 的 output 都馈入到 a shared AuxTower2 中，然后将它们的 logit 相加，然后计算 logloss。

   

1.2.4 Sparse-Pertoken MoE

1. 为进一步提升 TokenMixer-Large 的成本效益，我们设计了 Sparse-pertoken MoE 机制，在 pertoken SwiGLU 的基础上，将每个 expert 扩展为多个 sub-experts 并进行稀疏激活。与 RankMixer 的 "dense training and sparse serving" 相比，该机制实现了 "sparse training and sparse serving" ，大幅降低了超大规模推荐模型的训练成本和推理成本。

2. First Enlarge, Then Sparse：在迭代策略上，我们采用 "first enlarge, then sparse" 的思路：先扩大模型规模以提升效果，再对 pertoken SwiGLUs 进行拆分并稀疏激活，以获得效率收益。Sparse-Pertoken MoE (S-P MoE) 可表示为：

   $$\begin{array}{c}\text{S-P MoE}(\cdot )=\sum _{j=1}^k{g}_j(\cdot )\times {\text{Expert}}_j(\cdot ),\text{if j is chosen}\\ {\text{Expert}}_j(\cdot )={\text{FC}}_{\text{down,j}}(\text{Swish}({\text{FC}}_{\text{gate,j}}(\cdot ))\odot {\text{FC}}_{\text{up,j}}(\cdot ))\\ {\text{FC}}_{i,j}\left({\overrightarrow{\mathbf{x}}}_t\right)={\mathbf{W}}_{i,j}^t{\overrightarrow{\mathbf{x}}}_t+{\overrightarrow{\mathbf{b}}}_{i,j}^t,i\in \{\text{up, gate, down}\}\end{array}$$

   其中：

   • $\{{\mathbf{W}}_{\text{up}}^t,{\mathbf{W}}_{\text{gate}}^t\}\in {\mathbb{R}}^{(nD/E)\times D}$$\left\{\mathbf W_\text{up}^t, \mathbf W_\text{gate}^t\right\} \in \mathbb R^{ (nD/E)\times D}$，${\mathbf{W}}_{\text{down}}^t\in {\mathbb{R}}^{D\times (nD/E)}$$\mathbf W_\text{down}^t \in \mathbb R^{D\times (nD/E)}$，$\{{\overrightarrow{\mathbf{b}}}_{\text{up}}^t,{\overrightarrow{\mathbf{b}}}_{\text{gate}}^t\}\in {\mathbb{R}}^{nD/E}$$\left\{\mathbf{\vec b}_\text{up}^t, \mathbf{\vec b}_\text{gate}^t\right\}\in \mathbb R^{nD/E}$ ，${\overrightarrow{\mathbf{b}}}_{\text{down}}\in {\mathbb{R}}^D$$\mathbf{\vec b}_\text{down}\in \mathbb R^D$ 。$E$$E$ 为每个 token 的 experts 数量。

     $n$$n$ 为超参数，表示在 Expert 中，先将 hidden representation 扩张 $n$$n$ 倍，再缩小 $n$$n$ 倍。

     注意：每个 Expert 的宽度需要除以 $E$$E$ 。因此，这是用一组小的 Experts 来代替一个大的 Expert（即，SwiGLU）。

   • ${\overrightarrow{\mathbf{x}}}_t\in {\mathbb{R}}^D$$\mathbf{\vec x}_t\in \mathbb R^D$ 为第 $t$$t$ 个 token 。

   • $k$$k$ 为 router $g(\cdot )$$g(\cdot)$ 的 top-k 数量。$g(\cdot )$$g(\cdot)$ 通过 $\text{softmax}(\cdot )$$\text{softmax}(\cdot)$ 将 selected experts 的可能性进行归一化，代表 activated experts 的概率之和为 1 。详细的设计理念和实验结果见附录 A.1。

     给定 token ${\overrightarrow{\mathbf{x}}}_t$$\mathbf{\vec x}_t$，假设有 $E$$E$ 个 experts，则 $g(\cdot )$$g(\cdot)$ 的定义为：

     $$g\left({\overrightarrow{\mathbf{x}}}_t\right)=\text{softmax}\left(\frac{\mathbf{W}{\overrightarrow{\mathbf{x}}}_t+\overrightarrow{\mathbf{b}}}{\tau }\right)\in {\mathbb{R}}^E$$

     其中：$E\ge k$$E\ge k$ 。

     训练和推断过程中使用 top-k 。但是，top-k 操作在训练过程中不可导。此时采用 Straight-Through Estimator (STE) 方案：

     • 前向传播：执行真实的 hard top-k 操作（选概率最高的 k 个 experts，其余置0`），保证训练时的激活模式与推理一致。

     • 反向传播：忽略 top-k 的离散性，梯度仅回传到选中的 $k$$k$ 个 experts 所对应的概率维度。

3. Shared Expert：目前，多款基于 MoE 架构的 LLMs 提及使用 shared experts 来提升训练的稳定性和效果，我们在实际应用中也采用了这一设计。值得注意的是，此处的 shared expert 仍为 per token 的，而非对所有 tokens 全局可见的 shared expert：

   $$\text{S-P MoE}(\cdot )=\sum _{j=1}^{k-1}{g}_j(\cdot )\times {\text{Expert}}_j(\cdot )+\text{SharedExpert}(\cdot )$$

   这段话的含义是：SharedExpert 的结构是每个 token 共享的，但是它的 input-output 是 per-token 的。也就是：

   • 给定 token $t$$t$ ，SharedExpert 的输入仅仅依赖于 $t$$t$ ，而不依赖于其它 tokens 。

   • 给定 token $t$$t$ ，SharedExpert 的输出仅用于 $t$$t$ ，而不输出给其它 tokens 。

   SharedExpert 代表了一个共享的 extractor，用于抽取某种 per-token 的、共享的属性。

   “所有 tokens 全局可见的”，这句话的意思是 SharedExpert 的输出被馈入给所有 tokens 。

4. Gate Value Scaling：在实际应用中，我们发现由于 router $g(\cdot )$$g(\cdot)$ 通常使用 softmax 将 router logits 转换为和为 1 的概率分布，这一 summing = 1 的约束会导致 SwiGLUs 的梯度更新不足。因此，我们在 $g(\cdot )$$g(\cdot)$ 前加入常数 $\alpha$$\alpha$，让 expert networks 的更新更充分：

   $$\text{S-P MoE}(\cdot )=\alpha \times \sum _{j=1}^{k-1}{g}_j(\cdot )\times {\text{Expert}}_j(\cdot )+\text{SharedExpert}(\cdot )$$

   根据实验部分的总结，$\alpha$$\alpha$ 通常设置为稀疏率的倒数。假设有 $E$$E$ 个 experts（包含 SharedExpert ），则这里的稀疏率为 $k/E$$k/E$ 。则 $\alpha$$\alpha$ 设置为 $E/k$$E/k$ 。

5. Down-Matrix Small Init：受 Rezero （《Zero-shot text-to-image generation》）启发，我们降低了 SwiGLU 中 final layer 下投影矩阵 ${\mathbf{W}}_{\text{down}}$$\mathbf W_\text{down}$ 的初始化方差：采用 xavier_uniform ，并将默认标准差超参数从 1 降至 0.01 。这一调整让训练初期$F(x)+x$$F(x)+x$ 中的 $F(x)$$F(x)$ 项接近 0 ，使模块在初始阶段表现为近似恒等映射，从而提升训练稳定性；同时也有助于缓解 SwiGLU 中上投影矩阵 ${\mathbf{W}}_{\text{up}}$$\mathbf W_\text{up}$ 和门控矩阵 ${\mathbf{W}}_{\text{gate}}$$\mathbf W_\text{gate}$ 的乘法交互带来的 output value 爆炸的问题。

6. 结合上述改进，我们的 MoE model 版本仅激活一半参数，就能实现与 fully activated dense model 近乎持平的性能。我们还从 Figure 7 中观察到，引入这些改进后，tokens 的负载达到了良好的平衡，无需额外的 load-balancing loss ，详细分析见附录A.2 。

   

1.2.5 Training/Serving Optimization

1. 我们研发了一套高性能算子，提升 TokenMixer-Large 在训练和推理中的效率。

2. 高性能自定义算子：整体工作流程如 Figure 3 所示。

   • MoEPermute 将 input 从 batch-first 格式转换为 expert-first 格式，确保每个 expert FFN 的 input 是连续的。

   • 随后，MoEGroupedFFN 算子（MoE-GroupedSwiglu and MoEGroupedGemm）通过单个kernel 计算所有 expert FFNs，减少算子调度（operator scheduling）的开销，提升设备利用率。

   • 最后，MoEUnpermute 计算 multiple activated experts 的输出的加权和，作为 final output。

   

   各算子在训练和推理中的耗时如 Table 1 所示，其中 MoEGroupedFFN 是延迟的主要贡献者；MoEGroupedFFN 在推理耗时中的占比更高，因为即使推理时采用 smaller batch sizes 也不会降低其内存访问量。

   

3. FP8 Quantization：推理阶段，我们采用 FP8 E4M3 post-training quantization；训练全程保持 bfloat16 精度。线上实验表明， FP8 serving 在不损失模型准确性的前提下，实现了 1.7 倍的加速。如 Figure 3 所示：

   • input 在 MoEPermute 中完成量化，所有 expert weights 均已被提前量化。

   • 随后，MoEGroupedSwiglu 融合 the quantization of outputs ，直接生成 quantized activations。

   • 同理，MoEGroupedGemm 采用 FP8 进行计算，输出结果为 bfloat16 格式。

4. Token Parallel：在分布式 multi-device environments 中扩展 TokenMixer-Large 架构会带来显著的瓶颈。

   • 对于训练而言，主要挑战是 GPU 的显存容量和整体训练时间。随着模型规模扩大，模型的参数和 activations 可能超出单台设备的显存上限。

   • 对于推理而言，问题更为突出： online inference with small batch sizes 使计算成为 memory-bandwidth-bound ，导致硬件利用率低、吞吐量差。

   简单的 model-parallel solution 会将 pertoken operation weights 分片到各设备，每次计算之前和之后都需要执行 "all2all" operation 以交换 sharding layouts，这使得每个 block 的通信步骤多达四次。为解决这些问题，我们提出 Token Parallel，一种专为 TokenMixer-Large 架构设计的 model-parallelism 策略，将模型参数（尤其是 pertoken operation weights）和计算任务在多个设备间进行划分。

   实现方式：Token Parallel 的实现将 data sharding 与 TokenMixer-Large block 的计算流程（Mixing -> Reverting -> a pertoken operation）对齐。以下所有张量均可视为分布式张量，我们用 Shard(dim) 表示张量在指定维度上被分片。

   

   采用 Token Parallel 后， output ${\mathbf{X}}_l$$\mathbf X_l$ 仍按 token dimension 被分片（Shard(token) ），并直接被馈入下一个 block ；而下一个 block 的 initial "split" operation 可基于该 layout 执行，从而省去了两次额外的通信步骤，使 $L$$L$ layers 的通信开销从 $4L$$4L$ 降至 $2L+1$$2L+1$（仅在最后需要 a final all_to_all 来恢复 the batch sharding ）。

   Token Parallel 策略在训练和推理中均带来了显著的性能提升。在 production serving 中，采用 4-way token parallelism 的 TokenMixer-Large 模型（ global batch size = 320），其吞吐量相比 non-parallelized baseline 提升 29.2%；通过将通信与计算重叠（如采用细粒度的 micro-batch 调度、in-kernel 的 computation-communication 重叠等技术），这一增益进一步提升至 96.6%。

1.3 实验

1. 数据集：本研究的离线实验主要基于 Douyin 的 main feed 上电商场景的真实训练数据集。该数据集来源于平台的线上日志和 user feedback labels，包含商品点击、转化、商品交易总额（Gross Merchandise Volume: GMV）等指标。训练集包含 500 余种特征（包括数值特征、ID 类特征、交叉特征、序列特征），覆盖数亿独立用户。经采样后，数据集每日约有 400 million 条记录，数据收集周期为两年。

   离线实验还包含 Douyin Ads 和 Douyin Live Streaming 场景的真实训练数据。经采样后，它们每日数据量分别达到 300 million 和 17 billion 条记录。

2. 评估指标：我们以 AUC 和用 User-Level AUC: UAUC 作为主要的效果指标，以 dense parameter 数量（不含 sparse embeddings 参数）、FLOPs （模型处理 a single batch of 2048 samples 所需的浮点运算次数）和 MFU 作为效率指标。

3. baselines：我们主要对比以下公认的 SOTA基线模型：

   • DLRM-MLP：采用 multi-layer perceptrons: MLP 来实现 feature interaction。

   • DCNv2："Deep & Cross Network" 的改进版本，具备更深的、更强的 feature-crossing 能力。

   • AutoInt：通过 multi-head self-attention 机制，自动地、显式地学习高阶 feature combinations。

   • HiFormer：对标准的 self-attention 进行改进，采用heterogeneous self-attention layers 和低秩近似（low-rank approximation ）。

   • DHEN：采用集成的思想，在单个 block 中融合多种 feature interaction modules （DCN/SelfAttention/FactorizationMachineBlock），并堆叠 multiple layers。

   • Wukong：将 factorization machine module: FMB 和 linear compression module: LCB 集成在一起，进行多层堆叠，并研究 feature interactions 的 scaling laws。

   • Group Transformer：我们针对建模 heterogeneous features 而设计的 Transformer variant（公司内部研究的版本），通过 grouping features 并聚合为少量的 tokens，降低 self-attention 机制的计算复杂度，并采用 pertoken Q/K/V/O 来增强模型的 representing heterogeneous features 的能力。

   • FAT：Group Transformer 的升级版本，在 pertoken Q/K/V/O 的基础上，采用正交基组合（orthogonal basis combinations）来提升模型性能。

   • RankMixer：以 TokenMixer block 为核心（替代 self-attention 机制），大幅降低计算成本。

   所有实验均在混合分布式训练框架（hybrid distributed training framework）中进行；电商场景使用 64 GPUs，Feed Ads 和 Live Streaming 使用 256 GPUs。该框架对 sparse parameters 采用异步更新，对 dense parameters 采用同步更新。所有模型的 optimizer hyperparameters 保持一致：dense and sparse parts 均采用 Adagrad optimizer，学习率分别设为 0.01 和0.05 。

1.3.1 与 SOTA 模型的对比

1. Table 2 的结果表明，TokenMixer-Large 性能优于所有 SOTA 模型，相比 MLP-500M baseline ，其 CTCVR 的 AUC 提升了 1.14%。

   • 值得注意的是，Wukong, Group Transformer, and RankMixer 等模型采用了 pertoken network 设计，因此它们的每个 batch（ batch size = 2048）的 FLOPs 远低于其他 baselines。

   • 此外，我们发现 Sparse-Pertoken MoE 在仅激活一半参数的情况下，性能可与 dense 版本持平，大幅提升了模型训练/推理的投资回报率（ROI）。

   HiFormer 也是 per-token FFN 的，但是它的 FLOPs 要大得多。

   

1.3.2 与 Rankmixer (TokenMixer) 的详细对比

1. 我们对 TokenMixer-Large 和 Rankmixer (TokenMixer) 进行了详细对比，对三个关键设计维度定义如下：

   • Standard Residual (SR)：表示 blocks 之间是否存在标准的 residual connections。

   • Original Token Residual (OTR)：表示原始 tokens 的语义信息是否被保留并传播至后续层。例如，若 an input of 40 tokens 被混合成 16 new tokens，且后续始终保持 16 tokens，那么 40 tokens 的原始语义信息无法一直保留至 final output。

   • Token Semantic Alignment in Residual (TSA)：表示残差操作 $F\left(x^{\prime }\right)+x$$F\left(x^\prime \right) + x$ 前后，tokens 的语义是否保持一致。在RankMixer 中，$F\left(x^{\prime }\right)$$F\left(x^\prime \right)$ 中的 $x^{\prime }$$x^\prime$ 对应 tokens after mixing，而 residual input $x$$x$ 来自 tokens before mixing ，导致语义错位。

   由于 Group Transformer 的每一层都不会改变 input tokens 的数量，且在执行 attention and pertokenFFN 的计算后， token_dim 上各位置的语义保持对齐，因此满足所有上述三个特性。

   相比之下，RankMixer 在 residual connections 中直接将 tokens before mixing 和 tokens after mixing 相加，破坏了 token_dim 上各位置的 semantic alignment ；此外，除非刻意设计 new tokens 数量与 original tokens 数量一致，否则RankMixer 的 original token residual 无法在网络中有效传播。

   如 Table 3 所示，TokenMixer-Large 精准解决了 RankMixer 中存在的这些问题，满足所有三个设计特性的同时，实现了最优性能。

   

1.3.3 Scaling Laws

1. Scaling Laws of SOTA Models：在 Figure 5 中，我们选取了性能最优的两个模型—— RankMixer 和 TokenMixer-Large，分析其 scaling laws 。结果表明，两个模型的 AUC 均随参数数量/ FLOPs 的增加而提升；整体而言，TokenMixer-Large 性能最优，且性能提升的斜率更陡。

   

2. Scaling Laws of TokenMixer-Large on Different Scenarios ：我们通过增加宽度 $D$$D$、深度 $L$$L$ 和 scaling factor $n$$n$ （用于 SwiGLU hidden expansion）来扩大模型规模，得出两个关键结论：

   • 多维度均衡扩展的收益更高：与 1B 参数以下的 RankMixer 早期实验不同，我们发现单独扩展任意一个维度（宽度、深度、 scaling factor ）虽能带来收益，但会逐渐遭遇瓶颈；当参数规模超过 1B 时，需要对所有维度进行均衡扩展，才能实现更优的收益。

   • 更大的模型需要更多数据实现收敛：如 Table 4 所示，模型从 30M 参数扩展至 90M 参数时，仅需 14 天的训练样本即可收敛；而从 500M 参数扩展至 2B 参数时，需要 60 天的样本才能完全收敛。

     

   我们还在不同的业务场景中验证了 TokenMixer-Large 的 scaling laws ，发现每个场景均存在显著的 scaling law 效应。最终，在 Feed Ads 、电商、Live Streaming 的离线实验中，我们成功将模型分别扩展至 15B、7B、4B 参数；在线上流量实验中，相应扩展至 7B、4B、2B 参数，实现了显著的 scaling 收益。

   

1.3.4 消融实验

1. TokenMixer-Large Block 的消融实验：如 Table 5 所示，我们对 TokenMixer-Large 中的不同模块进行了消融实验，包括 Tokenization 、mixing & reverting、pertoken SwiGLU 、residual connections 等。结果表明，移除任何一个模块都会导致性能下降，其中 "mixing & reverting" 和 Pertoken-SwiGLU 对整体性能的影响最为显著。

   • Pure Model Design：随着 TokenMixer-Large 参数数量的持续增加，大量的 small, I/O-bound, fragmented operators （如 DCN、LHUC ）带来的收益，均可由 TokenMixer-Large 自身实现。此外，由于 TokenMixer-Large 仅包含无参数（parameterless）的 mixing and reverting 操作，以及大量的 GroupedGemm 操作，其效率极高，详细分析见附录 A.3。

   • Normalization：我们对归一化的应用位置进行了全面分析，包括 pre-norm、post-norm、sandwich norm，详细分析见附录 A.4。

   • Mixing Strategy：最后，我们对 mixer 的 splitting 策略进行了详细研究，详细分析见附录 A.5。

   

2. Sparse-Pertoken MoE 的消融实验：如 Table 6 所示，我们对 Sparse-Pertoken MoE 的不同模块进行了消融实验，包括 shared expert、Gate Value Scaling 、 Down-Matrix Small Init，发现每个模块均对最终的性能提升有积极贡献。上述操作不会引入额外的参数或计算开销（FLOPs），因此是无代价（cost-free）的改进。

   此外，在总参数数量和激活参数数量均对齐的条件下，我们将 sparse-pertoken MoE 与 standard MoE 进行对比，发现 standard MoE 的性能显著下降。这是因为 sparse-pertoken MoE 相当于为 standard MoE 赋予了一种路由先验（routing prior），即每个 token 可激活的 experts 并非共享的，从而避免了训练初期 routers 难以学习的问题。

   注意：“在总参数数量和激活参数数量均对齐的条件下”，这意味着 sparse-pertoken MoE 与 standard MoE 具有不同的参数配置。

   • Gate Value Scaling：同时，我们对 Gate Value Scaling 进行了全面实验，发现最优的 scaling factor 因模型稀疏率（activated parameters/total parameters）不同而变化：稀疏率（sparsity ratios）为 1:2 的模型在 scaling factor = 2 时性能最优；而稀疏率为 1:4 的模型在 scaling factor = 4 时达到最优效果。scaling factor 与稀疏率呈反比关系。详细分析见附录 A.6。

   • Down Matrix Small Initialization：将 SwiGLU 中 last down matrix ${\text{FC}}_{\text{down}}$$\text{FC}_\text{down}$ 的初始化值设为${\text{FC}}_{\text{up}},{\text{FC}}_{\text{gate}}$$\text{FC}_\text{up}, \text{FC}_\text{gate}$ 层的 0.01 倍，可显著提升模型性能和训练稳定性，详细分析见附录 A.7 。

   

1.3.5 线上性能

1. 为验证 TokenMixer-Large 的有效性，我们已将其在 Douyin 的多个场景落地部署，包括信息流广告（ feed advertising）、直播电商（live streaming e-commerce）、普通直播（live streaming），这些部署覆盖了 Douyin 的大部分核心业务，服务数亿用户。针对每个业务场景，我们采用不同的关键性能指标：

   • Douyin Live-Streaming E-Commerce ：以商品交易总额（Gross Merchandise Volume: GMV）和订单量作为核心线上业务指标。

   • Douyin Feed Advertising：以广告主满意度评分（Advertiser Satisfaction Score: ADSS）和广告主价值（Advertiser Value: ADVV）作为主要线上业务指标。

   • Douyin Live Streaming (Non-E-Commerce)：以总支付金额（Total Payment Amount）作为核心线上业务指标。

   广告、电商、直播场景的线上基线模型分别为 RankMixer-1B 、RankMixer-150M、RankMixer-500M。我们将 TokenMixer-Large 相应扩展至 7B 、4B、2B 参数。如 Table 7 所示，结果表明：

   • 电商场景中，TokenMixer-Large 实现订单量提升 1.66% 、人均预览支付 GMV （per-capita preview payment GMV）提升 2.98% 。

   • 广告场景中 ADSS 提升 2.0% 。

   • 直播场景中营收增长 1.4%。

   

1.4 结论

1. 本文提出了 TokenMixer 的升级版本—— TokenMixer-Large 。我们深入分析了原始 TokenMixer 设计中的局限性，并对 deep layer modeling 和 Sparse-Pertoken MoE 进行了详细的研究与探讨。在工程层面，我们实现了一套高效的 MoE operators，以及专为 Pertoken experts 设计的 Token-Parallel 策略。TokenMixer-Large 已在 ByteDance 多个场景实现了显著的离线和线上性能提升，服务数亿用户。

二、附录

A.1 First Enlarge Then Sparse

1. MoE 的传统应用方式是在保持 activated experts 数量不变的前提下，直接增加总参数数量。但我们认为，这种增加 experts 数量的方式相当于在单一维度上扩展 Pertoken SwiGLU 的 hidden expansion ，其上限极易达到。因此，我们选择先设计性能最优的 dense 模型，再通过稀疏化来实现效率收益，同时将性能损失降至最低。

   为此，我们采用 "First Enlarge, Then Sparse" 的策略，具体而言，将 TokenMixer-Large 中 SwiGLU 内的所有 FC networks 细化为细粒度版本，再进行 sparse activation。 Figure 6 展示了 dense model 、1:2 sparsity MoE 、1:4 sparsity MoE 的具体形式。目前：

   • 在 1:2 sparsity 下，我们能实现 offline/online 性能近乎零损失。

   • 在 1:4 sparsity 下，性能仅有轻微下降。

   • 但随着 sparsity 的增加， SwiGLU 中的 GEMM 组件占比会进一步降低，因此更高的 sparsity gains 依赖于更大的模型规模和更高的 GEMM 占比。

   因此，我们选择投资回报率（ROI ）最高的 1:2 版本进行线上部署。能否在保持性能无损失的前提下实现超过 1:8 sparsity，仍在探索中。

   

A.2 Load Balance

1. Figure 7 上半部分展示了电商场景下 4.6B Sparse-Pertoken MoE 的 load balancing 示例。我们使用 one shared expert 和three routing experts ，共激活 two experts 。在 1:2 sparsity 下，该模型的总参数数量与 MoE 一致，且 offline/online 性能与 dense model (4.6B dense) 持平。如图所示，我们展示了每个 token 的 activation 状态，其中索引0-2代表 routing experts，索引 3 代表 shared expert。可见整体负载达到了良好的平衡。由于索引 3 为始终激活的 shared expert，其激活程度高于索引 0-2。

   然而，随着 sparsity 增加（从 1:2 至 1:8），Figure 7 下半部分显示模型的负载平衡出现一定程度的恶化。我们也尝试了 Switch Transformer （《Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity》）和 Z-loss （《St-moe: Designing stable and transferable sparse expert models》）的负载平衡方法，能带来一定收益；但考虑到目前上线的 1:2 sparsity 版本的投资回报率最高，且不存在负载平衡问题，该部分内容将在未来进一步探索。

   

A.3 Pure Model Design

1. 在我们的 baseline 模型中，我们手动对序列的 pooling results 进行过滤，并馈入到一个 DCN 模块进行特征交叉。DCN 模块的输出作为 additional tokens ，补充至骨干网络 TokenMixer 的输入中。但该设计存在一个主要问题：随着骨干网络 TokenMixer 复杂度的持续提升，前端 DCN 带来的增益逐渐降低。实验数据如 Table 9 所示：

   • 在 150M 参数的模型中， DCN 能带来 0.09% 的增益。

   • 而当模型参数扩展至 700M 时，移除 DCN 结构不再造成任何效果损失。

   除 DCN 外，我们还尝试将 DHEN、 LHUC 等结构与 TokenMixer-Large 融合，采用并行和串行两种配置方式。实际结果表明，随着TokenMixer-Large 参数规模的扩大，这些 small, numerous, and IO-bound fragmented operators 带来的收益，可完全由TokenMixer-Large 自身实现。此外，TokenMixer-Large block 的 Model FLOPs Utilization: MFU 可达 60%，表现优异。

   

A.4 Normalization and Position

1. RMSNorm 可表示为：

   $$\begin{array}{c}\text{RMS}\left(\overrightarrow{\mathbf{x}}\right)=\sqrt{\frac{1}{d}\sum _{i=1}^d{x}_i^2+ϵ}\\ \text{RMSNorm}\left(\overrightarrow{\mathbf{x}}\right)=\frac{\overrightarrow{\mathbf{x}}}{\text{RMS}\left(\overrightarrow{\mathbf{x}}\right)}\odot \overrightarrow{\gamma }\end{array}$$

   其中：$\overrightarrow{\gamma }$$\vec\gamma$ 为 normalization 的可学习参数。

   考虑到 input 经过 normalization layer 后，kernel matrix 的 bias 在数学层面失去实际意义，同时受 Llama 系列模型的设计启发，我们移除了所有 bias kernel matrices ，并将 TokenMixer-Large 中所有 Norm 由 LayerNorm 替换为 RMSNorm。实测结果显示，该改动在保持模型性能完全不变的前提下，实现了端到端吞吐量 8.4% 的提升。

   我们针对 normalization 的应用位置开展了全面的对比分析，涵盖前归一化（ Pre-Norm ）、后归一化（Post-Norm）和三明治归一化（Sandwich-Norm）三种方式，结果如 Table 10 所示。从效果来看：

   • Post-Norm能带来更优的性能表现，但极易引发梯度爆炸问题。

   • Pre-Norm 的效果虽略逊于Post-Norm，却能保障模型训练的稳定性。

   

A.5 Mixing Strategy

1. 我们针对 mixer 的 splitting 策略开展了详细的研究，结果如 Table 11 所示。研究发现，只要每个 newly mixed token 都包含全部 original token information，不同的 split-concat 策略对模型性能无显著影响。作为对照实验，我们尝试仅对一半 original token information 进行混合，模型性能出现了明显的下降。

   给定 tokens 的 embedding 的序列 ${\mathbf{X}}_1,\cdots ,{\mathbf{X}}_T$$\mathbf X_1, \cdots, \mathbf X_T$ 。

   • 垂直划分：即论文中介绍的方式，第 $i$$i$ 个 mixed token 包含 $X_{j,i},j\le 1\le T$$X_{j, i}, j\le 1\le T$ 。

   • 对角线划分：第一个 mixed token 包含：第一个 token 的第一个元素、第二个 token 的第二个元素、第三个 token 的第三个元素，...。

   • 随机划分：随机指定一个映射方式，从而进行划分。

   

A.6 Gate Value Scaling Analysis

1. 在将 dense model 转换为 MoE 架构的过程中，我们发现在 router function $g(x)$$g(x)$ 之前引入一个常数缩放因子 $\alpha$$\alpha$ ——通过放大 router’s output 以增强 selected (routed) experts 的 gradient updates ——能显著提升模型性能。在引入 a shared expert 的情况下，MoE formulation 可表示为：

   $$y=\sum _{i=1}^{k-1}g(x{)}_i\times {\text{Expert}}_i(x)+\text{SharedExpert}(x)$$

   其中：$g(x)$$g(x)$ 为 routing weights ；${\text{Expert}}_i$$\text{Expert}_i$ 为第 $i$$i$ 个 expert network；$\text{SharedExpert}$$\text{SharedExpert}$ 为 shared expert network 。

   由此，Sparse-pertoken MoE 的数学表达式为：

   $$y=\alpha \times \sum _{i=1}^{k-1}g(x{)}_i\times {\text{Expert}}_i(x)+\text{SharedExpert}(x)$$

   实际实验中我们发现，缩放因子（scaling factor）$\alpha$$\alpha$ 的取值至关重要：取值过大或过小都会导致模型性能下降。进一步的实验表明，$\alpha$$\alpha$ 的最优值与模型的 sparsity 密切相关，具体而言，将 $\alpha$$\alpha$ 设置为当前 sparsity ratio 的倒数时，模型能取得最佳性能，结果如 Table 12 所示。

   

   我们认为 Gate Value Scaling 操作具备有效性的原因，以及将 $\alpha$$\alpha$ 设为 sparsity ratio 倒数的核心逻辑如下： SwiGLU 中每个全连接层（fully connected: FC）的梯度更新，需要与原始 dense model 保持一致。我们所采用的 "first enlarge, then sparse" 策略，将 SwiGLU 中的一个 large kernel 拆分为多个 smaller kernels ——这一思路与 DeepSeek-MoE （《Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models》）中的 fine-grained expert 设计相似，但核心区别在于我们并未增加 activated parameters 的数量（ topk 值）。这意味着，模型的拆分粒度越细，每个 expert 被更新的概率就越低。因此，我们需要在 an expert 被激活时放大其梯度幅度，而 Gate Value Scaling 操作恰好实现了这一目标。

   另一种可能性是，Gate Value Scaling 提升了 $\alpha \times \sum _{i=1}^{k-1}g(x{)}_i\times {\text{Expert}}_i(x)$$\alpha \times \sum_{i=1}^{k-1} g(x)_{i} \times \text{Expert}_{i}(x)$ 的整体 weighted sum-pooling 的结果，进而增大了模型每层输出的方差。为此，我们尝试通过直接增大 Expert kernel 的初始方差（initial variance），从另一角度模拟放大输出 $\alpha \times \sum _{i=1}^{k-1}g(x{)}_i\times {\text{Expert}}_i(x)$$\alpha \times \sum_{i=1}^{k-1} g(x)_{i} \times \text{Expert}_{i}(x)$ 的效果。结果如 Table 13 所示，直接增大 Expert kernel 的初始方差不仅未提升性能，反而导致结果出现显著恶化。

   

A.7 Down-Matrix Small Initialization

1. 为了解决 deeper models 的收敛难题，除引入standard residual connections 、interval residual connections 和 residual loss (resloss) 外，我们还借鉴了 Rezero 的设计理念。

   本研究中我们采用 Xavier Normal Distribution 初始化，其计算公式为：

   $$\mathbf{W}\sim \mathcal{N}(0,\frac{2\cdot \text{scale}}{n_{\text{in }}+n_{\text{out }}})$$

   其中：$\mathbf{W}$$\mathbf W$ 为全连接层的 kernal 矩阵；$n_{\text{in}},n_{\text{out}}$$n_\text{in}, n_\text{out}$ 为 input 单元数量和 output 单元数量；$\text{scale}$$\text{scale}$ 为 initialization scale，它的默认取值为 1 。

   我们将 SwiGLU block 中 final down-projection matrix ${\text{FC}}_{\text{down}}$$\text{FC}_\text{down}$ 的 initialization scale 直接降至 0.01，而保持${\text{FC}}_{\text{up}},{\text{FC}}_{\text{gate}}$$\text{FC}_\text{up},\text{FC}_\text{gate}$ 层的 initialization scale 为 1。如 Figure 8 所示，该调整让 $F\left(\overrightarrow{\mathbf{x}}\right)+\overrightarrow{\mathbf{x}}$$F\left(\mathbf{\vec x} \right) + \mathbf{\vec x}$ 在训练初期更接近恒等映射，从而促进 deep models 的收敛。此外，SwiGLU 结构中的 ${\text{FC}}_{\text{up}}\left(\overrightarrow{\mathbf{x}}\right)\odot \text{Swish}\left({\text{FC}}_{\text{gate}}\left(\overrightarrow{\mathbf{x}}\right)\right)$$\text{FC}_\text{up}\left(\mathbf{\vec x}\right)\odot \text{Swish}\left(\text{FC}_\text{gate}\left(\mathbf{\vec x}\right)\right)$ 操作容易放大中间隐藏层（intermediate hidden layer ）的 activations 和梯度幅度，进而引发局部不稳定性，降低 ${\text{FC}}_{\text{down}}$$\text{FC}_\text{down}$ 的 initialization scale 也有助于约束这些 intermediate representations 的幅值，提升训练的稳定性。

   

   如 Table 14 所示，我们严格评估了在不同 scale（从 1 到 0.01 ）和不同位置（从 all layers 到 reverse ordering ）下的 small initialization。结果表明，专门对 ${\text{FC}}_{\text{down}}$$\text{FC}_\text{down}$ 应用 small initialization 实现了最佳性能。