特征交叉 —— TokenMixer-Large 精读笔记

文章链接：https://arxiv.org/pdf/2602.06563

机构：字节

发布时间：2026.02

Copyright (c) Wang-Luning. All Rights Reserved.

本文为RankMixer团队推出的升级版RankMixer-v2，主要是针对原始的RankMixer存在的一些问题进行深入优化，从而将模型做的更大、更深，其分别在离线和在线场景下部署了15B和7B规模的模型。

TokenMixer尽管实现了效率和效果的平衡，成为了很多推荐系统的backbone，但其仍存在如下缺点：

• 次优的残差设计：为了保证经过Token Mixer后的输出可以和原始输入做残差连接，必须使mixing后的token数量$H$$H$和原序列长度$T$$T$相等。且直接把mixing前后的token进行相加很容易导致语义上的不对齐，从而损害模型能力。

• 不纯的模型架构：由于推荐系统的迭代更新不是一蹴而就的，因此即使采用了RankMixer作为了backbone，仍存在很多遗留的碎片化的memory bound算子，使得总体的MFU仍然难以提升。

• 深层模型的梯度更新不充分：RankMixer采用了稀疏MoE来优化推理成本，为了确保各个expert的充分训练采用了稠密训练+稀疏推理的设计，这就导致训练成本仍非常高。另外，使用ReLU产生门控分数使得推理时很难预测每个token的激活专家数量t

特征的Tokenization：

总体上和RankMixer类似，也是将特征进行分组后映射为维数相同的tokens。为了更好地捕捉全局信息，借鉴BERT中的[CLS] token，进一步设计了一个global token来聚合全局信息，它由所有特征的embedding连接后的大向量通过一个MLP映射后得到。这个global token会被连接到正常token序列的开头，从而形成模型输入$X_0\in {\mathbb{R}}^{T\times D}$$X_0\in\mathbb R^{T\times D}$

模型结构：

• Mixing & Revert模块：

  Token Mixing操作会改变序列长度。原始的Token Mixer层输入长度为$T$$T$的序列，输出长度为$H$$H$的序列，除非人工指定$H=D$$H=D$否则无法实现输入和输出的顺利短接。

  为了去除这一限制，一个自然的想法就是对mix后长度和维数都产生了变化的token序列加一个“恢复”（revert）操作，使得其在进入下一个block之前恢复到原始的张量形状。具体而言，在token mixer接收长度为$T$$T$并输出长度为$H$$H$的序列后，进一步通过PFFN（实际上是Per-token SwiGLU，详见下文）层得到长度为$H$$H$的输出，最后经过一个额外的Revert模块来负责恢复原始的$T\times D$$T\times D$形状。这样就保证了一个block的输入和输入形状保持一致（张量维数变化只存在于block内部），实现深层网络中信息传递的连贯性。当然在这种设计下就不能再在token mixer和PFFN上分别加短接了，只有一个横跨整个block的短接路径，如上图所示。

• Per-token SwiGLU：

  将RankMixer中的Per-token FFN优化为Per-token SwiGLU（pSwiGLU）。扩展到MoE后，相应地展为MoE SwiGLU，也即每个expert是一个SwiGLU网络。另外，给每个token的MoE网络中添加一个共享专家（注意只是token自己的MoE网络中的共享专家，不和其他token共享），从而获得了LLM中它对于训练稳定性和模型效果的优势。

  在路由分数生成上，由于ReLU导致expert激活数量不可控，因此这里又改回了softmax+top-k的路由机制。由于softmax的归一化会使得有些专家SwiGLU梯度更新不充分，因此额外添加一个常数$\alpha$$\alpha$来让expert更新更高效：

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

其他改进：

• 将LayerNorm换成RMSNorm来减轻计算压力，将post-norm换成pre-norm来防止梯度不稳定问题。

• 改进了SwiGLU中的W_down权重的初始化方差，使得$F(x)+x$$F(x)+x$中的$F(x)$$F(x)$部分在训练初始阶段接近0，使得模型层在开始时几乎相当于恒等映射，这样能够提升训练稳定性。这种设计还能够缓解SwiGLU中up输出和gate输出乘积造成的输出爆炸。

• 除了每个block的短接外，额外设置一些横跨多个block的短接（一般横跨2-3层，如下图左所示）来强化信息流动，用于解决低层的梯度消失问题，并由于能够让低层特征更好地传递到深层因此可以加快低层的收敛。需要注意的是，最后一层通常不加跨层的短接，因为最后一层的作用是提取高度抽象的特征来用于后续的分类任务，因此如果混入低层信息的话可能会扰乱这种抽象。

  另外，还加了一个辅助loss来使得低层学会如何估计高层特征的偏移，从而增强它们的特征表达能力，避免低层的不充分训练。其将低层输出的logits和高层输出的logits结合来计算一个联合loss。

  

在系统部署上，针对该模型的特性定制了MoE算子、token并行策略，以及使用了FP8 PTQ等。