字节RankMixer
背景
由于推荐场景特征的稀疏性和异构性，简单地加宽或堆叠特征交互层而不改变结构，这样做的收益是有限甚至是负面的。同时由于工业推荐系统的高QPS特性，必须兼顾效果和效率。这偏论文是字节在rank阶段探索模型scaing law的论文，论文提出了RankMixer——一种专为工业级推荐排序设计的新型架构。它摒弃了计算昂贵且语义不匹配的自注意力机制，转而采用无参数的 Token Mixing +动态稀疏 MoE的组合，在保持极低计算开销的同时，实现了对多源异构特征的高效全局交互。
模型结构
模型有L层RankMixer Block堆叠而成，每一层RankMixer Block包含两个关键组件：
- Multi-Head Token Mixing Layer：负责在不同 token 之间交换信息（类似 Transformer 中的自注意力机制）
- PFFN：对每个 token 独立地进行非线性变换（类似于 Transformer 的 FFN 层）。
特征token化
输入的特征需要处理为token，最简单的方式的是给每个特征分配相同的维度，一个特征就是一个token，但是这样造成重要特征得不到充分表达，稀疏特征训练不充分。因此论文对特征先进行了分组（相似特征划分为一组），一个组特征对应一个token
Token Mixing
token mixing没有采用self-attention的结构，而是采用了一个无参数的结构，计算过程如下（假设输入有T个Token，每个Token的向量维度为D）：
1. 切分头
将每个Token的D维向量均匀地切分成H个“头”，每个头的维度为D/H。在 RankMixer 的默认配置中，设置H=T。这样每个头的维度就是D/T。
2. 按头重组
提取每个头的内容重新组成一个token
例子
假设三个输入Token的向量为（T=3，D=6，H=T=3，每个头的维度 D/H=6/3=2）：
- Token 1: x1=[a1,a2,a3,a4,a5,a6]
- Token 2: x2=[b1,b2,b3,b4,b5,b6]
- Token 3: x3=[c1,c2,c3,c4,c5,c6]
第一步：切分头
每个Token被均分为3个头（每个头2维）：
- Token1:
头1: [a1,a2] 头2: [a3,a4] 头3: [a5,a6] - Token2:
头1: [b1,b2] 头2: [b3,b4] 头3: [b5,b6] - Token3:
头1: [c1,c2] 头2: [c3,c4] 头3: [c5,c6]
第二步：按头重组拼接
- 对于头1：收集所有Token的头1 → [a1,a2,b1,b2,c1,c2]，得到混合后的Token s1
- 对于头2：收集所有Token的头2 → [a3,a4,b3,b4,c3,c4]，得到 s2
- 对于头3：收集所有Token的头3 → [a5,a6,b5,b6,c5,c6]，得到 s3
本质理解
这个操作本质上是对输入张量进行了一种无参数的维度重排，让同一个token混入了所有token的信息，近似达到self-attention的效果。这是一种参数无关、计算高效的全局信息交换方式。它避免了自注意力在异构特征空间上计算相似度的困难和巨大开销，并行度极高。
残差连接+归一化
这一步是标准的 Transformer 风格模块设计，目的是稳定训练过程、加速收敛
\(s_1,s_2,...,s_T = \text{LN}\big(\text{TokenMixing}(x_1,x_2,...,x_T)\big)+(x_1,x_2,...,x_T)\)
PFFN
在传统的Transformer或DLRM中，所有Token共享同一个FFN（参数共享）。这在NLP中是合理的，因为所有Token（词）共享同一个语义空间。但在推荐系统中，Token来自高度异构的特征空间：
- Token A：代表“用户ID + 用户画像” → 语义空间：用户侧
- Token B：代表“视频ID + 内容标签” → 语义空间：物品侧
- Token C：代表“上下文环境 + 统计特征” → 语义空间：上下文
如果让它们共享同一套FFN参数，会导致 “特征淹没”：
- 高频、强势特征（如热门视频）的梯度主导参数更新
- 低频、长尾特征（如冷门兴趣）的信号被稀释
- 模型难以对不同的特征子空间进行专门化建模
Per-token FFN的解决方案是：为每一个Token分配一套独立的FFN参数。这样，每个Token可以在自己的专属参数空间中，对混合后的信息进行深度交叉。用来提升模型对异构特征的建模能力和对长尾信号的敏感度。
输入多头混合后的token，过一个两层的NN：
\(\mathbf{v}_t = f_{\text{pffn}}^{t,2}\Big(\text{Gelu}\big(f_{\text{pffn}}^{t,1}(\mathbf{s}_t)\big)\Big)\)
\(f_{\text{pffn}}^{t,i}(\mathbf{x}) = \mathbf{x}\mathbf{W}_{\text{pffn}}^{t,i} + \mathbf{b}_{\text{pffn}}^{t,i}\)
- \(\mathbf{W}_{\text{pffn}}^{t,i},\mathbf{b}_{\text{pffn}}^{t,i}\)：第 t 个 token、第 i 层的权重与偏置
- \(\mathbf{W}_{\text{pffn}}^{t,1}\in\mathbb{R}^{D\times kD},\ \mathbf{b}_{\text{pffn}}^{t,1}\in\mathbb{R}^{kD}\)
- \(\mathbf{W}_{\text{pffn}}^{t,2}\in\mathbb{R}^{kD\times D},\ \mathbf{b}_{\text{pffn}}^{t,2}\in\mathbb{R}^{D}\)
- D：token 原始维度；k：超参，用来控制 FFN 隐藏层大小
- \(\mathbf{s}_t\in\mathbb{R}^D\)：输入的第 t 个 token 向量
Sparse MoE in RankMixer：
这是 RankMixer 从“百亿级”向“千亿级”参数规模迈进的核心技术，它完美体现了论文“解耦参数增长与计算增长”的设计哲学。
- 基础版 Per-token FFN：每个 Token t配一个专属的、稠密的、完整的 FFN。这个 FFN 对所有输入该 Token 的数据都进行全参数计算。
- SMoE 变体：将每个 Token t 的那一个 FFN，替换成一个由 N个“专家”（每个专家本身是一个小FFN）组成的网络。然后通过一个动态路由机制，对于当前输入的 Token t，稀疏地选择其中最有能力的几个专家来激活，并融合它们的输出。
- 注意：每个 Token 拥有自己独立的一组专家，而不是所有 Token 共享一个全局的专家池。
传统 MoE 的问题：通常使用 Top-K + Softmax 路由。这有两个缺陷：
- 专家使用不灵活：它为每个 Token 强制选择固定数量的专家。但推荐系统中的 Token 信息量差异巨大（如“用户ID” Token 和“季节” Token），固定数量会导致信息量大的 Token 资源不够，信息量小的 Token 浪费资源。
- 难以优化：Softmax 产生的概率分布是稠密的，需要专门设计负载均衡损失来防止专家坍塌。
改进方案：用一个简单的 ReLU 函数替代 Softmax 来选择专家
优势：
- 每个token实际激活的专家数目不孤独，对于复杂特征会有更多的激活专家，简单特征有更少的激活专家
- serving时只需要计算激活的专家，大大降低推理延迟
参考资料
《RankMixer: Scaling Up Ranking Models in Industrial Recommenders》
浙公网安备 33010602011771号