特征交叉 —— RankMixer 精读笔记

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

机构：字节

发布时间：2025.07

参考文献：

• https://zhuanlan.zhihu.com/p/1937436392402683166

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

推荐系统的scaling面临如下挑战：

• 训练和推理开销必须在严格的延迟约束下，并满足高QPS要求

• 精排模型中的大部分人类手工设计的特征交叉模块都是CPU时代的产物，很多核心算子都是memory-bound的，难以在现代GPU上实现高MFU和高并行度

RankMixer是一种硬件感知的、专为推荐系统设计的类transformer特征交互模型架构，希望在GPU上能够大幅提升MFU和计算吞吐，同时保持强大的特征交叉能力。具体而言，其保留了transformer的高并行度，同时将二次复杂度的attention模块换成了一种多头的、无参数的token-mixing模块，来提升特征交叉的效率和效果。另外，RankMixer通过Per-token FFNs来实现对于各个特征和交叉特征的子空间使用独立的参数进行建模，避免高强度特征信号淹没低强度特征，其中的每个FFN可以进一步扩展为MoE来提升模型参数规模。

RankMixer模型最大可达1B的大小，实现了推荐系统的scaling up，且相比于传统的特征交叉模块可以将MFU从4.5%提升到45%。其提供的token mixing模块和PFFN的思想在后续很多做特征交叉或统一建模的工作中被作为基础设施来使用。

特征的Tokenization

原始输入中包含的用户特征、候选物品特征、序列特征（用户历史序列特征）、和候选物品等交叉特征等会被转化为若干长度不一embedding向量。如果每个特征都分配一个专属的embedding的话，数百个特征就会产生数百个tokens，使得每个token分配的计算都是小块的，导致重要特征难以充分被建模，GPU利用率也会降低。而如果token数量太少又会导致难以表示多样化的特征空间，可能会使得少量主导的特征盖过了其他特征。

因此，RankMixer采用了一种折中的办法：按照语义将特征分成若干组（例如将用户画像类的若干特征），然后将每个组中的各个特征的embedding连接成一个向量$e_i$$e_i$，总共$N$$N$个组embedding连接起来成为一个超大向量：$e_{\text{input}}=[e_1;e_2;\cdots ,e_N]$$e_{\text{input}}=[e_1;e_2;\cdots,e_N]$。希望最终总共产生数量适中的$T$$T$个tokens（也可以看出$T$$T$个heads，用来代表多个特征子空间）作为模型输入，则将$e_{\text{input}}$$e_{\text{input}}$均匀切成$T$$T$段，每一段再通过一个线性层映射成token的维数${\mathbb{R}}^D$$\mathbb R^D$，即可得到一个含有$T$$T$个维数为$D$$D$的token的序列用于模型输入：$X_0\in {\mathbb{R}}^{T\times D}$$X_0\in\mathbb R^{T\times D}$。每个token $x_i\in {\mathbb{R}}^D$$x_i\in\mathbb R^D$可以涵盖一组类似的特征的语义。

模型结构：

模型总共由$L$$L$个RankMixer block堆叠而成，用于处理输入的长度为$T$$T$的序列$X_0\in {\mathbb{R}}^{T\times D}$$X_0\in\mathbb R^{T\times D}$。

输入序列在每个block中首先经过Token Mixer进行特征交叉融合，然后经过PFFN来独立处理各个token。和transformer一样，这两个模块也都含有短接。第$n-1$$n-1$个block的总体变换如下：

• Multi-Head Token Mixer：

  用于替代普通transformer中的attention层。

  每个token $x_t$$x_t$沿hidden dim方向被切成$H$$H$段，每段对应一个head：

  $$[x_t^{(1)}||x_t^{(2)}||\cdots ||x_t^{(H)}]=\text{SplitHead}(x_t)$$

  这些heads可以看做将token映射到多个特征自空间后的表征。

  然后，将所有token在同一个head的表征分别连接起来（和all-to-all操作的逻辑一样），得到各个head的表征$s^h$$s^h$：

  $$s^h=\text{Concat}(x_1^{(h)},\cdots ,x_T^{(h)})$$

  总共可以产生$H$$H$个这样的表征，每个head的表征融合了序列中所有token在该子空间中的信息，由于每个token都代表了不同的特征组，因此这样也就实现了特征交叉融合。将每个head的表征看做一个“token”，堆叠起来即可得到一个长度为$H$$H$的新序列：$S\in {\mathbb{R}}^{H\times \frac{T\cdot D}{H}}$$S\in\mathbb R^{H\times \frac{T\cdot D}{H}}$，该序列也就是Token Mixing的输出序列。（为了确保输入输出的序列长度相等，使得能够进行残差连接，这里设置$H=T$$H=T$。）

   

  作者指出，self attention在NLP中用于信息交叉之所以有效是因为所有token都处于同一个embedding空间中，但推荐系统里各个特征空间的差异可能非常大，因此在两个异质的token之间算内积相似度意义不大。因此这里选用了无参数的Token Mixing策略效果也不比attention差，且能够消除掉attention计算带来的开销。

• Per-Token FFN（PFFN）：

  为了避免过度的特征融合导致高频特征信号掩盖了那些低频、长尾的特征信号，在FFN部分将每个token隔离使用独立的参数进行处理（因为输入序列长度固定为$T$$T$所以在模型结构上也可以静态地设置$T$$T$个FFN），也即Per-Token FFN。类似于强制每个token分别流向一个特定的expert（PFFN）的MoE结构，每个expert（PFFN）也只能看到并处理它对应的这个token。

  为了进一步增大参数量，还可以进一步将每个dense的PFFN替换成一个稀疏MoE网络（如果套用到上一段的MoE视角的话，整个模块就相当于一个“MoE的MoE”，也即宏观MoE的每个expert分别是一个小MoE）。然而直接套用传统MoE效果不好，主要原因包括：

  • 固定的top-k个expert选取：每个token都分配$k$$k$个expert处理，相当于将所有token都同等看待，这就导致在那些低信息量的token上浪费计算资源，同时高信息量大token无法充分处理，使得模型难以捕捉不同token的区别

  • 由于总expert数量过多（每个token都有$N_e$$N_e$个expert，总共$T\times N_e$$T\times N_e$个），容易导致路由不均衡问题，使得有些experts训练不充分

  基于这些问题，采用如下两种改进策略：

  • 将top-k+softmax路由策略换成ReLU作为gate进行路由，也即对每个expert的路由分数由ReLU函数输出，且不经过softmax归一化和top-k强制截断：

    $$\begin{array}{c}{G}_{i,j}=\text{ReLU}(h(s_i))\\ v_i=\sum _{j=1}^{N_e}{G}_{i,j}{e}_{i,j}(s_i)\end{array}$$

    这样一来，高信息量大的token可以激活更多的experts，而没有top-k的限制。同时不做softmax归一化这一点可能也能更好地保持原始信号强度。

    由于在这种设置下expert激活数量没了限制，因此训练中为了确保稀疏度（总共激活的expert数量），额外设置一个损失${\mathcal{L}}_{reg}$$\mathcal L_{reg}$来惩罚门控分数的总幅值，使得平均激活量在计算预算之内：

    $$\begin{array}{c}{\mathcal{L}}_{reg}=\sum _{i=1}^T\sum _{j=1}^{N_e}{G}_{i,j}\\ \mathcal{L}={\mathcal{L}}_{task}+\lambda {\mathcal{L}}_{reg}\end{array}$$

  • 稠密训练+稀疏推理：在训练阶段采用两个router：$h_{train},h_{infer}$$h_{train},h_{infer}$，且${\mathcal{L}}_{reg}$$\mathcal L_{reg}$只施加在$h_{infer}$$h_{infer}$上。可见，这样在训练中$h_{train}$$h_{train}$的路由专家数量没有抑制，因此它会倾向于激活更多专家，使得训练过程中更多专家被激活从而得到充分训练（当然这也会导致训练成本相对较高）。而推理时则只保留抑制过激活数量的$h_{infer}$$h_{infer}$，确保推理阶段的成本降低。

最后一层的输出$X_0$$X_0$沿序列长度方向做平均后即可得到模型的输出，它可以进一步被用于预测多种任务目标。