RankMixer详解

RankMixer从动机到完整结构一、为什么需要 RankMixer工业推荐和搜索精排的输入不是一种统一的数据而是大量异构特征例如用户 ID、年龄、地域、活跃度候选物品 ID、作者、类别、标签用户历史行为序列时间、设备、位置等上下文用户和物品之间的交叉统计特征。这些特征来自完全不同的空间但精排模型需要把它们充分交叉最终预测 CTR、CVR、完播率等目标。传统排序模型通常堆叠 FM、DCN、Attention、MLP 等多种人工设计模块。这些模块能够工作但存在两个 Scaling 障碍模型扩大以后效果不一定稳定提升大量小算子和不规则计算不适合 GPU容易受内存访问限制MFU 较低。而工业精排还受到高 QPS 和严格延迟约束不能只是把模型无脑做大。因此 RankMixer 的目标是设计一个统一的、适合 GPU 的特征交互主干使推荐排序模型能够从千万参数继续扩大到亿级、十亿级参数。二、RankMixer 的整体结构整个模型可以概括成原始推荐特征 ↓ Embedding 与特征分组 ↓ 生成 T 个 D 维 Feature Token ↓ RankMixer Block × L ↓ Pooling ↓ 任务预测头 ↓ CTR / CVR / 完播率等每个 RankMixer Block 只有两个核心模块Multi-head Token Mixing Per-token FFN两者分工非常明确Token Mixing负责不同 Token 之间的信息交换Per-token FFN负责学习信息交换后的非线性关系三、输入如何变成 Token假设原始系统有数百个特征字段。RankMixer 不采用一个字段 一个 Token因为这样会形成几百个 Token每个 Token 对应的计算规模太小最终产生大量碎片化的小矩阵运算。它也不把所有特征全部压成一个 Token因为这样就会退化成普通 DNN无法保留不同特征空间的独立性。RankMixer 的做法是数百个原始字段 ↓ 按照业务语义分组 ↓ Embedding、拼接和切分 ↓ 统一投影到 D 维 ↓ 得到 T 个 Feature Token最终输入为其中BBatch SizeTToken 数量D每个 Token 的隐藏维度。论文最终使用的配置为模型Token 数 TToken 维度 DBlock 数 LRankMixer-100M167682RankMixer-1B3215362也就是说1B 模型会将每条样本整理为论文没有公布每个字段具体属于哪个 Token也没有给出自动确定 Token 数量的公式。它采用的基本原则是Token 要足够多以区分不同特征子空间但不能多到让计算过度碎片化。四、一个 RankMixer Block一个 Block 可以写成对应结构为输入 TokenX ↓ Multi-head Token Mixing ↓ 残差连接 LayerNorm ↓ Per-token FFN ↓ 残差连接 LayerNorm ↓ 输出 TokenX_next残差和 LayerNorm 的作用可以简单理解为保留原表示同时提高训练稳定性。Block 真正的创新重点是前面的Token Mixing和后面的Per-token FFN。论文消融中移除 Token Mixing 或将 Per-token FFN 改成共享 FFN都会使效果下降。五、Multi-head Token Mixing假设输入有 4 个 Token分别代表不同的特征子空间。每个 Token 是 8 维并被切成 4 份原来是Token 1A1 A2 A3 A4 Token 2B1 B2 B3 B4 Token 3C1 C2 C3 C4 Token 4D1 D2 D3 D4Token Mixing 按照切分位置重新组合变成混合 Token 1A1 B1 C1 D1 混合 Token 2A2 B2 C2 D2 混合 Token 3A3 B3 C3 D3 混合 Token 4A4 B4 C4 D4这样每个新的 Token 都包含了所有原始 Token 的部分信息。假设原来的四个 Token 分别代表A用户特征 B物品特征 C行为序列 D上下文特征那么混合后的每个 Token 就同时包含部分用户信息 部分物品信息 部分序列信息 部分上下文信息因此Token Mixing 完成了跨特征空间的信息交换。Token Mixing 本质上是什么它本质上是一次张量重排交换 Token 和 Head 维度再重新拼接RankMixer 设置HT所以输出仍然是[B,T,D]这个过程不需要Q、K、VAttention ScoreSoftmax可学习的 Mixing 参数。Token Mixing 自己不判断哪些 Token 更相似也不直接学习特征关系。它只是保证不同特征子空间的信息能够被放进同一个新 Token 中。六、为什么不用 Self-AttentionSelf-Attention 的核心是它通过 Token 之间的内积相似度决定信息如何传递。这在语言模型中很自然因为文本 Token 基本处于统一的语言语义空间中。但是推荐模型中的 Token 可能分别表示用户 ID商品价格-历史行为设备类型统计特征。这些 Token 高度异构用户 ID Token 与价格 Token 的向量内积不一定具有清晰的“相似度”意义。RankMixer 的思路是不再先学习两个异构 Token 是否相似而是直接让所有 Token 的一部分信息进行固定混合再交给后面的 FFN 学习具体关系。论文实验中Self-Attention 替代方案相较于 Token Mixing参数量增加约 16%、FLOPs 增加约 71.8%但效果仍略低。七、Per-token FFNToken Mixing 只负责重新组织信息真正学习特征交叉关系的是 FFN。假设混合后的一个 Token 为它进入自己的 FFNFFN 通常是例如公式为由于输入中已经同时包含用户、物品、序列和上下文的部分信息FFN 可以学习它们之间的非线性组合。例如用户兴趣和候选类别是否匹配用户历史行为与当前场景是否一致特定用户、物品和上下文同时出现时点击概率是否更高。为什么是 Per-token FFN普通 Transformer 中Token 1 ─┐ Token 2 ─┼→ 同一个 FFN Token 3 ─┤ Token 4 ─┘RankMixer 中Token 1 → FFN 1 Token 2 → FFN 2 Token 3 → FFN 3 Token 4 → FFN 4即不同 Token 使用不同参数。这样做有两个主要作用。第一不同特征子空间可以独立建模。虽然 Token Mixing 已经交换了信息但每个混合 Token 仍然对应一个不同的切分子空间独立 FFN 可以学习不同的组合模式。第二增加参数容量但不同比增加计算量。假设有 TTT 个 FFN每个 Token 只进入自己的 FFN并不是一个 Token 依次经过全部 T 个 FFN。因此参数量可以明显增加而单个 Token 仍只执行一次 FFN 计算。论文将共享 FFN 替换回 Per-token FFN 后模型获得了更好的异构子空间建模能力。八、一个 Block 到底做了什么一个 Block 的工作可以压缩为两步。第一步Token Mixing不同特征子空间 ↓ 切开并重新组合 ↓ 每个 Token 获得跨空间信息第二步Per-token FFN已经混合的 Token ↓ 独立的非线性变换 ↓ 学习具体的特征交叉关系因此Token Mixing 负责让特征发生交互​ Per-token FFN 负责理解和加工交互结果九、为什么还要堆叠多个 Block第一个 Block 输入的是原始 Feature Token经过一次 Mixing 和 FFN 后得到此时已经包含一次跨特征交互结果。第二个 Block继续处理它是在上一层已经交互和加工过的表示上再次进行 Mixing 和非线性变换。因此第一个 Block 原始特征之间发生交互 第二个 Block 交互后的表示继续发生交互不能机械地说第一层只学二阶、第二层只学三阶但随着 Block 堆叠模型通常能够表达更复杂、更抽象的特征组合。论文最终的 100M 和 1B 模型都只使用了两个 Block。作者发现在总参数量相近时增加深度 LL、宽度 D 或 Token 数 TT的效果比较接近但增加宽度能够产生更大的矩阵乘法更有利于提高 GPU 利用率因此最终采用了较浅但较宽的结构。十、最后如何预测经过多个 Block 后模型对全部 Token 进行 Pooling例如 Mean Pooling再将整体表示输入任务预测头多任务场景可以配置多个预测头共同的 RankMixer 主干 ↓ ┌────┼────┐ ↓ ↓ ↓ CTR CVR 完播率所以 RankMixer 主要负责生成统一的高质量特征交互表示最后的任务 Head 负责输出具体目标。十一、Sparse MoE 如何继续扩大 RankMixer普通版本中每个 Token 对应一个 FFNToken i → FFN iSparse MoE 版本把一个 FFN 扩展成多个 Expert┌→ Expert 1 Token i → Router ────├→ Expert 2 ├→ Expert 3 └→ Expert 4每个 Expert 本质上仍然是一个 FFN。输出为这样模型可以存储大量专家参数但每个样本只调用其中一部分从而把总参数容量与单次计算量部分解耦ReLU Routing普通 Top-K MoE 会固定激活 K 个专家例如所有 Token 都激活两个。但不同 Token 和不同样本的复杂程度不同简单输入可能只需要 1 个专家 复杂输入可能需要 34 个专家RankMixer 使用 ReLU RoutingRouter 输出大于零的 Expert 被激活小于等于零的被关闭。因此每个 Token 可以动态激活不同数量的专家。为了防止模型把所有 Expert 都打开训练时还会加入自适应稀疏约束控制平均专家激活比例。Dense Training、Sparse Inference稀疏 MoE 容易出现某些专家长期不被选择得不到充分梯度更新的问题。RankMixer 因此采用即训练阶段让专家获得充分训练 推理阶段只计算被路由选中的专家Dense Training 缓解专家训练不足和“死亡专家”问题Sparse Inference 则控制线上计算成本。论文实验显示Dense Training 与 ReLU Routing 的组合在较低专家激活比例下仍能较好地保持模型效果。十二、RankMixer 如何 ScalingRankMixer 有四个主要扩展方向T增加 Token 数D增加隐藏维度L增加 Block 层数E增加 Expert 数量普通 Dense RankMixer 主要通过 T,D,L 扩大Sparse MoE 版本可以进一步通过增加 E 扩大总参数容量。其主要计算由 Per-token FFN 中的大矩阵乘法构成Token Mixing 本身只是简单的数据重排。因此它比大量异构的小型交叉算子更适合 GPU 并行、算子融合和半精度计算。论文的线上部署对比为指标原线上模型RankMixer-1B参数量15.8M1.1BFLOPs107G2106GMFU4.47%44.57%延迟14.5 ms14.3 ms也就是说参数规模约扩大 70 倍但通过更低的 FLOPs/参数比、更高的 MFU、FP16 推理以及系统工程优化论文报告的线上延迟基本保持不变。这里不能简单理解成“参数增加完全不产生计算”而是模型结构与硬件优化共同抵消了成本增长。十三、最后总结RankMixer 的完整逻辑是数百个异构推荐特征 ↓ 整理成少量固定维度 Token ↓ Token Mixing 跨 Token 进行无参数的信息交换 ↓ Per-token FFN 分别学习不同子空间中的非线性交叉 ↓ 堆叠多个 Block 继续加工已经交互过的特征 ↓ Pooling 任务 Head 输出 CTR、CVR 等预测大规模版本再把Per-token FFN替换成Per-token Sparse MoE最终可以把 RankMixer 的核心概括为固定、低成本的 Token Mixing独立、可扩展的 Per-token FFN​Token Mixing 负责连接不同特征空间Per-token FFN 或 Sparse MoE 负责提供真正的模型容量。这也是 RankMixer 的主要思想不再堆叠越来越复杂的人工特征交叉算子而是用简单统一的 Mixing 骨架加上适合 GPU 的大 FFN/MoE把推荐排序模型真正做大。