用矩阵实例具象化 RankMixer 核心机制

结合文档中 100M 参数模型的典型配置（T=16 个特征 Token、D=768 隐藏维度、H=16 个头），通过具体矩阵维度和运算过程，拆解 RankMixer 核心模块的矩阵操作逻辑，让抽象架构落地为可感知的数值流程。

一、前置约定（贴合文档实验配置）

• 批量大小 B=512（工业训练常见批量）；
• 特征组数量 N=4（用户画像组、物品属性组、行为序列组、交叉特征组）；
• 特征 Token 数 T=16（文档 4.3 节 100M 参数模型配置）；
• 隐藏维度 D=768（每个 Token 的固定维度）；
• 多头数 H=T=16（文档 3.3.1 节设定，保证 Token 数量不变）；
• Per-token FFN 隐藏层缩放系数 k=4（文档 3.3.2 节，隐藏维度 = k×D=3072）。

二、模块拆解 + 矩阵实例

1. 特征 Tokenization：异构特征 → 统一 Token 矩阵（核心：拼接 + 分割 + 投影）

步骤 1：特征组嵌入（异构输入矩阵）

每个特征组先通过独立嵌入层转换为固定维度嵌入矩阵，矩阵维度为 [B, d_i]（d_i 为组内嵌入维度）：

用户画像组（年龄、性别等）：E1 ∈ R^(512×256)（d1=256）；

物品属性组（视频 ID、标签等）：E2 ∈ R^(512×192)（d2=192）；

行为序列组（点击历史等）：E3 ∈ R^(512×224)（d3=224）；

交叉特征组（用户 - 物品匹配度等）：E4 ∈ R^(512×96)（d4=96）。

步骤 2：矩阵拼接（统一异构特征）

按列拼接 4 个特征组矩阵，得到全局嵌入矩阵 E_input，维度满足 Total_Dim = d1+d2+d3+d4 = 256+192+224+96=768：Einput​=Concat(E1,E2,E3,E4)∈R(512×768)

矩阵含义：每行对应 1 个样本，每列对应 1 个特征嵌入维度，共 512 个样本、768 个总嵌入维度。

步骤 3：矩阵分割（切分为固定长度片段）

按列分割 E_input 为 T=16 个等宽片段，每个片段维度 d=Total_Dim/T=768/16=48：Eseg,i​=Einput​[:,(i−1)×48:i×48]∈R(512×48)(i=1 16)

实例：第 1 个片段 E_seg,1 ∈ R^(512×48)（取 E_input 前 48 列），第 2 个片段 E_seg,2 ∈ R^(512×48)（取 49-96 列），直至 16 个片段。

步骤 4：矩阵投影（统一为 D=768 维度 Token）

通过投影矩阵 W_proj ∈ R^(48×768) 和偏置 b_proj ∈ R^768，将每个片段映射为 D=768 维度的 Token，对应文档公式（2）：Xi​=Proj(Eseg,i​)=Eseg,i​×Wproj​+bp​roj∈R(512×768)

矩阵运算逻辑：512×48 的片段矩阵 × 48×768 的投影矩阵 → 512×768 的 Token 矩阵，偏置 b_proj 广播到每个样本；

最终输出：16 个 Token 矩阵堆叠为输入 Token 矩阵 X0 ∈ R^(512×16×768)（维度：[B, T, D]）。

2. 多头 Token 混合：跨 Token 无参数交互（核心：拆分 + 重组）

步骤 1：矩阵拆分（头分割）

将输入 Token 矩阵 X0 ∈ R^(512×16×768) 按最后一维（D=768）拆分，每个 Token 拆为 H=16 个头，每个头维度 D/H=768/16=48：X0h=X0[:,:,(h−1)×48:h×48]∈R(512×16×48)(h=1 16)

实例：第 1 个头矩阵 X0^1 ∈ R^(512×16×48)（取每个 Token 的前 48 维），第 2 个头矩阵 X0^2 ∈ R^(512×16×48)（取 49-96 维）。

步骤 2：跨 Token 重组（全局特征融合）

对每个头矩阵 X0^h，按 Token 维度（第 2 维）重组，将 16 个 Token 的同一头特征拼接，得到跨 Token 头矩阵 S^h ∈ R^(512×(16×48)) = R^(512×768)：Sh=Concat(X0h[:,0,:],X0h[:,1,:],...,X0h[:,15,:])∈R(512×768)

矩阵含义：每行对应 1 个样本，每列对应 "所有 Token 的第 h 头特征"，实现跨 Token 的子空间信息融合；

实例：S^1 是将 16 个 Token 的第 1 头（48 维）拼接为 768 维，即每个样本的 S^1 行 = [Token1_头 1, Token2_头 1, ..., Token16_头 1]。

步骤 3：矩阵堆叠 + 残差融合

将 16 个头矩阵 S^1~S16 按 Token 维度堆叠，得到 S ∈ R^(512×16×768)，与原始 X0 残差连接并层归一化：Sn−1​=LN(S+X0)∈R(512×16×768)

核心优势：无参数矩阵重组替代自注意力的内积运算，避免异构特征距离计算失效，同时保持 GPU 并行效率。

3. Per-token FFN：独立子空间建模（核心：独立权重矩阵）

每个 Token 配备专属 MLP 权重矩阵，避免特征主导问题，对应文档公式（6）（7）：

步骤 1：独立权重矩阵设计

为 16 个 Token 分别分配 2 层 MLP 权重：

第 1 层权重 W1_t ∈ R^(768×3072)（k=4，隐藏维度 = 768×4=3072），偏置 b1_t ∈ R^3072；

第 2 层权重 W2_t ∈ R^(3072×768)，偏置 b2_t ∈ R^768；

实例：Token1 用 W1_1/b1_1，Token2 用 W1_2/b1_2，...，Token16 用 W1_16/b1_16，共 16 套独立参数。

步骤 2：矩阵线性变换（单 Token 建模）

对 S_{n-1} 中每个 Token 的矩阵 S_t ∈ R^(512×768)（t=1~16），独立执行 MLP 运算：V1t​=St​×W1t​+b1t​∈R(512×3072)V2t​=Gelu(V1t​)∈R(512×3072)Vt​=V2t​×W2t​+b2t​∈R(512×768)

实例：Token1 的变换过程：S_1 ∈ R^(512×768) → 乘 W1_1 ∈ R^(768×3072) → 加 b1_1 → Gelu 激活 → 乘 W2_1 ∈ R^(3072×768) → 加 b2_1 → V_1 ∈ R^(512×768)。

步骤 3：矩阵堆叠输出

将 16 个 Token 的输出矩阵 V_1~V_16 堆叠，得到 V ∈ R^(512×16×768)，与 S_{n-1} 残差连接并层归一化，作为当前 RankMixer 块的输出：Xn​=LN(V+Sn−1​)∈R(512×16×768)

4. 稀疏 MoE 扩展（可选）：动态专家选择

将 Per-token FFN 替换为 N_e=8 个专家矩阵，每个专家 E_j ∈ R^(768×768)（j=1~8），核心是动态路由权重矩阵：

步骤 1：路由权重计算

通过路由矩阵 h ∈ R^(768×8) 计算 Token 与专家的匹配权重，用 ReLU 激活（非负化）：Gt​=ReLU(St​×h)∈R(512×8)

实例：Token1 的 S_1 ∈ R^(512×768) × h ∈ R^(768×8) → G_1 ∈ R^(512×8)，每行对应 1 个样本对 8 个专家的激活权重。

步骤 2：专家加权融合

每个 Token 的输出为专家矩阵的加权和（稀疏激活，仅权重非零的专家参与）：Vt​=∑j=18​Gt​[:,j]×(St​×Ej​)∈R(512×768)

矩阵含义：S_t × E_j ∈ R^(512×768) 是 Token 与第 j 个专家的交互结果，G_t[:, j] 是对应的权重（广播到每个样本），加权和实现稀疏专家融合。

5. 输出层：平均池化（Token 聚合）

对最后一层 RankMixer 块的输出 X_L ∈ R^(512×16×768)，按 Token 维度（第 2 维）平均池化，得到全局特征矩阵：Ooutput​=MeanPool(XL​,dim=1)∈R(512×768)

矩阵运算：对每个样本的 16 个 Token（768 维）取平均值，最终每个样本用 768 维向量表示，用于 CTR 预测等任务。