RankMixer 架构拆解（字节） | YelloooBlue Blog 博客

type

Post

status

Published

date

Mar 23, 2026

slug

summary

RankMixer

字节抖音、广告等渠道均推全，值得学习。

Tokenization

Token Mixing

Per-token FFN + SMoE

RankMixer Block

模型架构

第一步：Tokenization

这一部分在整体框架图的左下角，原文里是这么说的：

即按照语义，例如：用户特征、物品特征、序列特征、上下文特征进行分组。

在这个图里面，我们的分组数T=4，因为我们将数百个特征按照语义分为了4组。

💡

这里这么做主要有两个原因，所以选择了一个折中方案

如果全部特征转为一个Token（向量），模型会退化为DNN。

如果每个特征一个Token（向量），会导致碎片过多。

此外、因为同一组的特征的语义相对接近，在后续交叉的时候也更加符合直觉。

这里的 Tokenizer 很简单，其实就是一个组内的特征各自进行Embedding或者数值化操作：

特征映射 (Embedding)： 模型接收这 100 个特征（比如用户 ID、视频类别等），它们会分别经过类似 nn.Embedding 或线性层，变成 100 个基础的特征 Embedding 向量。

按语义分组并拼接 (Group & Concat)： 根据人为设定的领域知识（Domain Knowledge），把这 100 个特征分成 4 组（比如：25个用户画像、25个视频属性、上下文特征等）。在代码层面，就是对每一组内部的 25 个 Embedding 在特征维度上进行 torch.cat，拼成 4 个极其宽的长向量。

投影降维 (Projection/MLP)： 因为拼接后的 4 个长向量维度可能各不相同，且为了后续的高效并行计算，模型会为每一个组配备一个映射网络（nn.Sequential 构建的 MLP）。这个 MLP 的作用是将拼接后的长向量投射（Project）到一个统一的隐藏层维度（例如 D）。

得到 Token： 最终，这 4 个组经过各自的 MLP 之后，就变成了 4 个维度完全一致（均为 D）的稠密向量。这 4 个向量，就是下一步用于交叉计算的 4 个“特征 Token”。

最终得到 T 组维度为 D 的向量。

第二步：Token Mixing

在开始介绍之前，我们需要知道，这一步是用来替代「自注意力机制」的。

为什么要替代呢，原文是这么说的：

大概意思就是说，如果在NLP等场景，特征只有一种形式，就是token（词元）。但推荐系统不一样，特征可能是用户的年龄、物品的类别等等..，他们都来自不同的空间，所以说推荐系统的token是具备异质性（heterogeneous）的。

💡

在语言模型（NLP）中，所有词元（Token）都在同一个“语义空间”里： 语言模型里所有的词（如“我”、“爱”、“中国”）都是从同一个庞大的词表（Vocabulary）里映射出来的。它们共享同一个 Embedding 矩阵（Unified embedding space），因此计算“苹果”和“香蕉”的内积相似度是有意义的，距离近代表它们语义相近。

在推荐系统中，特征是“异构”的（Heterogeneous）： 推荐系统的输入是五花八门的，比如“用户年龄（数值）”、“用户所在的城市（类别ID）”、“昨天看过的视频ID（序列）”等。这就像是把“身高”、“温度”和“经纬度”放在一起。为这些完全不同物理含义的特征计算“内积相似度”（即传统自注意力的做法）是非常不可靠且没有实际意义的。它们各自的 Embedding 是通过不同的表独立映射出来的，这就是为什么说它们处于“完全不同的语义子空间”。

在整体架构图的右下角我们能看到 Token Mixing的具体设计

这里左上角这一块就是上一步得到的得到 T 组维度为 D 的向量（token）。

不同颜色代表他们来自不同类型（语义空间）的特征。例如紫色是用户特征、绿色是视频特征。

💡

在这里其实 H = T，因为后续要进行残差连接，具体可以往下看。

In this work, we set H=T to maintain the same number of tokens after Token Mixing for residual connection.

Token Mixing 用一种完全没有参数的数据重排操作实现了特征交叉。

我们可以用一个通俗的例子来理解这段论文中的操作：假设我们有两个 Token，各包含 4 个维度的信息：

Token A（代表用户特征）： [A1,A2,A3,A4]

Token B（代表物品特征）： [B1,B2,B3,B4]

拆分与重组的过程是这样的：

切分（Split）： 把每个 Token 切分成 H 个头（比如切成两半）。

Token A 变成：前半段 [A1,A2]，后半段 [A3,A4]

Token B 变成：前半段 [B1,B2]，后半段 [B3,B4]

重组（Merge / Concat）： 把不同 Token 相同位置的“头”拼在一起，形成新的 Token。

新 Token 1： 包含所有 Token 的前半段 →[A1,A2,B1,B2]

新 Token 2： 包含所有 Token 的后半段 →[A3,A4,B3,B4]

💡

为什么这样就实现了特征交互（Feature Cross）？

重组后的「新Token1」 里面，不仅有用户的特征 (A1,A2)，同时也包含了物品的特征 (B1,B2)。接下来，这个「新Token1」 会被送进后续的前馈神经网络（Per-token FFN）中。在神经网络内部的矩阵乘法中，用户特征和物品特征会在同一个神经元里相乘、相加，并经过激活函数。通过将不同空间的特征强制塞进同一个向量里让网络去计算，模型就自然而然地学到了用户与物品的交叉特征。

💡

这种做法既避开了传统自注意力机制（Self-Attention）在异构空间中计算内积的不合理性，又避免了巨大的算力消耗，是一种极其高效的、专为推荐系统特征设计的交互方式。

残差连接

💡

现在我们有「老Token」和「新Token」，这里进行Add & Norm操作，类似Transformer的残差链接

第三步：Per-token FFN（PFFN）

对于不同分组的特征，每组有一个PFFN，具体是这样考虑的：

传统 Transformer 的做法： 无论输入的是什么 Token（比如“我”、“爱”、“你”），在经过注意力机制后，都会通过同一个前馈神经网络（FFN）。也就是说，所有 Token 共享同一套 FFN 的权重参数。

推荐系统的痛点： 推荐系统的 Token 属于完全不同的空间（用户空间、物品空间等）。如果让它们共享同一套 FFN 参数，那些出现频率极高的强势特征就会占据主导地位，挤压其他长尾特征的学习空间。

RankMixer 的做法： 既然 Token 代表不同的语义空间，那就给每个 Token 定制它专属的 FFN（参数互相独立）。

这种设计在不增加推理计算复杂度（FLOPs）的情况下，极大地增加了模型的参数量，让模型能更好地刻画不同特征子空间的差异性。

对于每个PFFN，内部是这么设计的：

不要被这个图吓到了，在最开始，其实FNN没有这么复杂，就是Transformer当中的FFN，但是我们生成式推荐要探索Scaling Up嘛，所以我们可以在FFN上进一步做文章。

稀疏混合专家系统（Sparse MoE）的作用是：在几乎不增加算力成本（也就是不增加推理延迟）的前提下，将模型的参数规模Scaling Up到十亿（1B）级别。

替换 Dense FFN： 随着模型想要进一步扩大（Scaling up），RankMixer 把上述每个 Token 专属的 FFN 升级成了 MoE 结构。MoE 意味着把一个 FFN 变成了多个“专家 FFN”，并在前面加了一个路由（Router）来决定用哪几个专家。

为什么叫“稀疏”： 虽然模型里有非常多的专家（参数量极大），但对于任何一条输入数据，路由器只动态激活极少部分的专家。这样一来，计算量（FLOPs）和普通的 FFN 差不多，但模型却拥有了远超以往的巨大容量来记忆和泛化。

RankMixer 的特殊优化： 传统的 MoE 经常会遇到“有些专家被撑死，有些专家被饿死（不被激活）”的问题。RankMixer 特别设计了 ReLU Routing（自适应决定每个Token激活的专家数） 和 DTSI（训练时激活更多专家，推理时激活少量专家） 机制，完美解决了专家训练不平衡的问题，保障了线上极低的推理延迟。

每个PFFN都是输入 D 维向量（token），输出也是 D 维度向量，跟上一步一样，是为了可作残差。

残差连接

💡

现在我们有「新Token」和「新新Token」，这里进行Add & Norm操作，类似Transformer的残差链接

第四步：堆叠RankMixer Block

这里也是和 Transformer的Blocks堆叠类似。

由于我们从「token」到「新token」到「新新token」的维度全程没有变化，这意味着我们可以堆叠这些操作，也就是 RankMixer Block。这里可以堆叠L层（超参数）用于Scaling Up。

第五步：预测头

在经过 L 层的 RankMixer 模块处理后，模型会通过一个输出池化操作（Output pooling operator）来处理最终层的表示向量。

原文明确指出，这里使用的是平均池化（Mean pooling），将多个特征 Token 融合为一个单一的输出表示（output）。即 H x D → D

输出多任务预测值： 这个经过平均池化后得到的单一全局向量，随后会被送入任务网络中（一般有几个任务就是几个MLP），用于计算不同具体任务的预测结果。正如架构图所示，这个向量最终会连接到多个任务头部，输出具体的概率值，比如预测用户是否完播（finish）、是否跳过（skip）、是否点赞（like）等目标。

总结（AI）

RankMixer 的设计确实可以被视为对 Transformer 架构在工业推荐系统场景下的一次深度定制与“魔改”。它保留了 Transformer 高度并行的骨架，但为了解决推荐特征的异构性以及严苛的线上延迟（MFU利用率）要求，对其核心组件进行了彻底的替换。

总结来说，RankMixer 对 Transformer 的借鉴与改造体现在以下几个核心维度：

1. 借鉴 Transformer 的“骨架”与基建

输入表示（Tokenization）： 就像 Transformer 处理自然语言需要先分词一样，RankMixer 也将输入统一为 Token。但不同于语言模型中共享词表的 Token，RankMixer 采用基于语义的分词（Semantic-based Tokenization），利用领域知识将数百个稀疏和稠密特征分组、拼接并映射为固定维度的“特征 Token”，以此作为网络的标准输入。

高度并行的块堆叠（Block Stacking）： RankMixer 继承了 Transformer 的高度并行性结构，通过堆叠多个相同的 RankMixer Block（即多层网络）来逐步提取高阶特征表示。

残差连接与层归一化（Add & Norm）： 这是 Transformer 保证深层网络训练稳定性的灵魂组件。RankMixer 在其每一层的核心操作前后，都严格保留了残差连接（Residual Connection）和层归一化（Layer Normalization）。消融实验也证明，去掉它们会导致极其严重的性能下降和训练不稳定。

2. 改造 Transformer 的“核心”以适应推荐场景

虽然骨架相似，但 RankMixer 抛弃了 Transformer 内部最核心的两个算子（自注意力与共享FFN），原因在于它们在推荐系统的异构特征空间中“水土不服”：

用“多头 Token 混合”替换“自注意力（Self-Attention）”：

痛点： Transformer 的自注意力依赖计算 Token 之间的内积相似度。但这在推荐系统中并不合理，因为用户 ID、年龄、物品 ID 属于完全不同的异构语义空间，强行计算内积既不可靠，又会导致极高的显存和算力开销（二次方复杂度）。
改造： RankMixer 设计了**无参数的多头 Token 混合（Multi-head Token Mixing）**操作。它将每个 Token 切分成多个“头”，然后跨 Token 重新拼接相同位置的头。这种纯数据的重排操作不仅以极低的计算代价实现了全局特征交互，还在效果上击败了自注意力机制。

用“专属 FFN”替换“共享 FFN”：

痛点： 传统 Transformer 中，所有 Token 都共享同一个前馈神经网络（FFN）参数。如果直接套用于推荐系统，高频特征（如全局热门物品）会主导网络权重，导致低频或长尾特征的信号被淹没。
改造： RankMixer 为每个特征 Token 分配了参数完全独立的 FFN（Per-token FFN）。这种“既切分输入，又切分参数”的设计，极大地增加了模型容量，保证了各个异构特征子空间能够独立、充分地学习。

3. 迈向 10 亿参数的终极形态（Sparse-MoE）

为了追求 Transformer 时代的大模型“缩放定律（Scaling Law）”，RankMixer 进一步将每个 Token 专属的 FFN 升级为稀疏混合专家系统（Sparse-MoE）。通过引入动态的 ReLU Routing（自适应决定激活的专家数量）和 DTSI（稠密训练，稀疏推理）策略，RankMixer 在不增加线上推理延迟（FLOPs 成本基本不变）的前提下，成功将参数规模扩展了 70 倍至 10 亿（1B）级别。

结论： 你可以将 RankMixer 理解为一个**“硬件感知（Hardware-aware）”的推荐版 Transformer**。它借用了 Transformer 的流水线和稳定性基建，但通过创新的 Token 混合与参数隔离机制，完美避开了原生 Transformer 处理推荐特征时的低效与误差，实现了算力利用率（MFU）从 4.5% 到 45% 的十倍跃升。

实践

部分开源实现

Reference

arxiv.org

zhuanlan.zhihu.com