← 返回论文列表
🔬 Kuaishou Technology · 2026

UniMixer

UniMixer: A Unified Architecture for Scaling Laws in Recommendation Systems

作者
Mingming Ha, Guanchen Wang, Han Li, Kun Gai 等
机构
快手（Kuaishou Technology）
发表时间
2026年4月（arXiv:2604.00590v2）
核心贡献
统一推荐系统三大 Scaling 架构的理论框架 + 轻量版 UniMixer-Lite
⭐ 这篇论文做得好的地方——核心贡献速览

读完这篇论文，以下几个设计思路值得重点关注：

① 把 TokenMixer 看成矩阵乘法（等价参数化）
把 TokenMixer 的 reshape + transpose 操作等价成 $W_\text{perm} \cdot \text{flatten}(X)$，这是全文最关键的理论洞察。一旦有了这个视角，操作就变成了矩阵，后续所有的参数化、约束、压缩都有了统一的数学基础。这是一种"换角度看已有工作"的思路，而不是凭空设计新模块。
② 大矩阵拆成两部分（Kronecker 分解 + 可学习化）
$W_\text{perm} = G \otimes I_B$ 分解为块间（$W_G$）和块内（$W^i_B$）两层，分别替换成可学习浮点矩阵。收益双重：参数量从 $L^2$ 压缩到 $(L/B)^2 + (L/B) \cdot B^2$；计算复杂度从 $O(L^2)$ 降为 $O(L^2/B + LB)$，两步独立执行无需构造大矩阵。这是 Kronecker 积结构性质的直接利用，不是工程 trick。
③ 对矩阵施加三条结构约束（继承置换矩阵的性质）
不是让 $W_G$ 随便学，而是通过对称化 + 温度退火 + Sinkhorn 归一化，保持置换矩阵原本就满足的双随机性（信息不丢失不重复）、稀疏性（每块有明确去向，不退化成均值池化）、对称性（互换关系）。每条约束都有消融实验支撑：稀疏性（AUC ↓ 0.1645%）、热启动（↓ 0.0856%）、对称性（↓ 0.0573%）。
④ 统一三大 Scaling 范式（最大的理论贡献）
用同一个公式把 Attention、TokenMixer、FM 三条孤立路线统一了：$W_G$ 固定置换 = TokenMixer，$W_G$ 由内积动态计算 = Attention，局部投影固定 = FM。这让"三条路谁更好"的争论变成"同一框架下的不同特例"。UniMixer 是同时学习两层的一般解，三者都是它的特殊情况。
⑤ SiameseNorm——让深度 Scaling 真正可行
RankMixer 堆 4 层比 2 层 AUC 反而低 0.1066%，根本原因是 Pre-Norm 与 Post-Norm 的矛盾：Pre-Norm 训练稳定但弱化梯度，Post-Norm 表达力强但不稳定。SiameseNorm 用双流结构同时解决——一路每层归一化保稳定，另一路跳过归一化直接累积保梯度。有了它，UniMixer 堆 4 层比 2 层 AUC 再高 +0.1575%，8 层继续涨 +0.1647%。没有这个，前四条贡献的理论价值都无法在实验中体现。
📖 §1 背景与动机

论文链接：arXiv:2604.00590

一句话摘要：通过对 TokenMixer 的等价参数化，将推荐系统三大 Scaling 架构（Attention、TokenMixer、FM）统一到一个理论框架下，并提出更高效的 UniMixer 和 UniMixer-Lite。

三大 Scaling 范式

LLM 的成功揭示了 Scaling Laws 的巨大潜力——随着参数量和计算量的增长，性能持续提升。推荐系统社区也开始探索其自己的 Scaling 框架。目前有三条主流路线：

① Attention-Based（注意力机制）

HiFormer、FAT、HHFT 等。用 token 特有的 Q、K、V 投影实现异构特征交互。优点：表达能力强；缺点：attention score 计算是二次复杂度，且在推荐场景中异构特征会导致注意力矩阵呈对角线主导，训练初期梯度不稳定。

② TokenMixer-Based（Token 混合）

RankMixer、TokenMixer-Large 等。用静态无参数规则做 token 混合，避免了异构空间内积计算。优点：FLOPs 低、高效；缺点：没有可学习性，且严格要求 head 数 = token 数（$H=T$）。

③ FM-Based（分解机）

Wukong、Kunlun 等。用 FM Block 计算 embedding 之间的显式交互，配合线性压缩层堆叠。优点：可解释性好；缺点：低阶交互限制了 Scaling 性能上限。

核心问题

三大范式设计哲学截然不同，彼此孤立，缺乏统一的理论框架。能否构建一个兼具三者优点的统一 Scaling 模块？这是本文的出发点。

推荐系统 Scaling 的核心挑战

异构特征空间——推荐系统中用户特征、物品特征、行为序列等来自不同语义空间，不能像 NLP 中直接将 Transformer 移植过来（NLP 的所有 token 共享同一嵌入空间）。因此，每种方法都需要专门设计异构特征交互机制。
Figure 1: UniMixer/UniMixer-Lite 与 RankMixer 的 Scaling Law 曲线
Figure 1（论文原图）：UniMixer、UniMixer-Lite 与 RankMixer 的 AUC-参数量/AUC-FLOPs Scaling 曲线。UniMixer-Lite 的曲线斜率最大，意味着单位参数/算力提升 AUC 的幅度最高。
🏗️ §2 方法详解
整体架构流程：输入多域异构特征 → Embedding Layer → Token-Specific 线性投影 → $M$ 个 UniMixer Block（每个 Block 含 UniMixing/UniMixing-Lite + Pertoken SwiGLU + SiameseNorm） → Task Tower（多任务）→ 预测输出
Figure 2: UniMixer 整体架构图
Figure 2（论文原图）：UniMixer 完整架构图。左侧展示 UniMixing 模块结构（Local Mixing + Global Mixing），右侧展示 UniMixing-Lite 结构，上下以 SiameseNorm 双流连接。

2.1 特征 Tokenization

输入特征按语义类别分为多个域（User Profile、Item Features、Behavior Sequence、Query Features 等），每个域通过 Embedding 层得到该域的 embedding，拼接后等分为若干块，每块通过 token 特有的线性投影得到最终 token embedding：

$$x_i = W^i_{\text{proj}} E_{d_i:d_i+d} + b^i_{\text{proj}} \in \mathbb{R}^D$$
符号说明
  • $E$：所有域 embedding 拼接后的向量
  • $W^i_{\text{proj}} \in \mathbb{R}^{D \times d}$：第 $i$ 个 token 特有的投影矩阵
  • $D$：统一 token 维度（模型隐藏维）
  • $d_i$：第 $i$ 个分块在 $E$ 中的起始位置
  • $X \in \mathbb{R}^{T \times D}$：$T$ 个 token 的隐状态矩阵
💡 举例：特征 Tokenization 过程

假设推荐场景有 3 个域：用户档案（年龄、性别等，embedding 维 64）、商品特征（类别、价格等，embedding 维 128）、行为序列（过去点击的商品 ID，embedding 维 64）。

三个域 embedding 拼接：$E = [e_\text{user}, e_\text{item}, e_\text{behavior}] \in \mathbb{R}^{256}$

等分为 4 块（每块 64 维），每块分别通过各自的 $W^i_\text{proj} \in \mathbb{R}^{128 \times 64}$ 投影到统一的 128 维 token 空间，最终得到 $T=4$ 个 token，$X \in \mathbb{R}^{4 \times 128}$。

2.2 TokenMixer 的等价参数化（核心洞察）

这是全文最关键的理论发现。传统 TokenMixer 是一种静态无参的操作：将 $X \in \mathbb{R}^{T \times D}$ 中每个 token 按 head 拆分后重排，本质上等价于一个置换矩阵 $W_\text{perm} \in \mathbb{R}^{TD \times TD}$ 乘以 $\text{flatten}(X)$：

$$\text{TokenMixer}(X) = \text{reshape}(W_\text{perm} \cdot \text{flatten}(X))$$
符号说明
  • $W_\text{perm} \in \mathbb{R}^{TD \times TD}$：置换矩阵，每行每列恰好一个 1
  • $\text{flatten}(X)$：将 $X$ 展平为 $TD$ 维向量
  • $T$：token 数；$D$：每个 token 的维度
💡 置换矩阵是什么？

置换矩阵（Permutation Matrix） 是一种特殊的方阵：每行有且仅有一个 1，其余全为 0；每列同理。作用是把向量元素按规则换位置——本质就是一个"排列重组"操作。

最小例子：把 $[a, b, c]$ 变成 $[c, a, b]$（第3个挪到第1位）：

$$P = \begin{pmatrix} 0 & 0 & 1 \\ 1 & 0 & 0 \\ 0 & 1 & 0 \end{pmatrix}, \quad P \cdot [a,\ b,\ c]^T = [c,\ a,\ b]^T$$

第1行 $[0,0,1]$：新第1位 = 原第3位；第2行 $[1,0,0]$：新第2位 = 原第1位，以此类推。
TokenMixer 的操作（按 head 切分再重拼）恰好等价于这样一个置换矩阵左乘 $\text{flatten}(X)$。

置换矩阵的四条性质——以及 UniMixer 怎么用它们

① 稀疏性（Sparsity）
每行/列只有一个 1，每个输出只取一个输入的值，不叠加混合。

对论文的意义：UniMixer 把 0/1 换成可学习的浮点数，稀疏性不再严格，但通过温度系数 $\tau$ 做退火（从 1.0 降到 0.05），让权重逐渐集中到少数位置。消融实验：去掉温度系数（始终 $\tau=1.0$）AUC 下降 0.1645%，是所有因素中影响最大的一条。
② 双随机性（Doubly Stochastic）
"随机"是数学术语 = "归一化/和为1"，不是随机采样。普通 softmax 只有行和=1（行随机矩阵）；置换矩阵行和=1 且列和=1，所以叫"双随机"。
验证 $P$：行 $[0,0,1]$、$[1,0,0]$、$[0,1,0]$ 行和都=1；第1列 $[0,1,0]^T$、第2列 $[0,0,1]^T$、第3列 $[1,0,0]^T$ 列和也都=1 ✅

对论文的意义：双随机性保证信息既不重复放大也不丢失（能量守恒）。UniMixer 用 Sinkhorn-Knopp 迭代来保持这个约束——把可学习矩阵交替对行/列归一化直到收敛。
③ 可压缩性（Compressibility）

问题：置换矩阵太大了
假设有 $T=2$ 个 token，每个 token 维度 $D=3$（方便手算）。把 $X \in \mathbb{R}^{2 \times 3}$ 展平得到 $\text{flatten}(X) \in \mathbb{R}^{6}$，那么置换矩阵 $W_\text{perm}$ 就是 $6 \times 6$ 的。真实场景 $T=16, D=512$，矩阵大小是 $8192 \times 8192 \approx 6700$ 万个数——存都存不下。

关键发现：这个大矩阵可以拆开
置换矩阵做的事情其实有两层：① 哪个 token 换到哪个位置（token 间的搬运，由 $G \in \mathbb{R}^{T \times T}$ 描述）；② 一个 token 的 $D$ 个维度整体一起搬，内部顺序不变（由 $I_D$，即 $D \times D$ 单位矩阵描述）。
这两层可以用 Kronecker 积合并：$W_\text{perm} = G \otimes I_D$。

Kronecker 积是什么？——用数字说话
$A \otimes B$ 的含义：把矩阵 $A$ 的每一个元素，替换成"该元素 × 矩阵 $B$"这个子块，然后拼在一起。

设定：token 0 的特征 $[a_1, a_2]$，token 1 的特征 $[b_1, b_2]$（$D=2$）。
$G = \begin{pmatrix}0&1\\1&0\end{pmatrix}$（含义：第0行是 $[0,1]$ → 新 token 0 = 旧 token 1；第1行是 $[1,0]$ → 新 token 1 = 旧 token 0，即两者互换）

第一步：写出 Kronecker 积的结构
$G \otimes I_2$ 的含义 = 把 $G$ 的每个元素替换成「该元素 × $I_2$」的子块：

$G_{00}=0$：$0{\cdot}I_2 = \begin{pmatrix}0&0\\0&0\end{pmatrix}$
$G_{01}=1$：$1{\cdot}I_2 = \begin{pmatrix}1&0\\0&1\end{pmatrix}$
$G_{10}=1$：$1{\cdot}I_2 = \begin{pmatrix}1&0\\0&1\end{pmatrix}$
$G_{11}=0$：$0{\cdot}I_2 = \begin{pmatrix}0&0\\0&0\end{pmatrix}$

第二步：代入大矩阵（左上=第00块，右上=第01块，左下=第10块，右下=第11块）：

$G \otimes I_2 = \begin{pmatrix} 0 & 0 & 1 & 0 \\ 0 & 0 & 0 & 1 \\ 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \end{pmatrix}$

第三步：验证，令输入 $v=[a_1,a_2,b_1,b_2]^T$：

行1 $[0,0,1,0]$：输出 $= b_1$ ✓
行2 $[0,0,0,1]$：输出 $= b_2$ ✓
行3 $[1,0,0,0]$：输出 $= a_1$ ✓
行4 $[0,1,0,0]$：输出 $= a_2$ ✓

结果 $[b_1,b_2,a_1,a_2]^T$ ✅ — token 0 和 token 1 整体互换，内部维度顺序不变

参数量对比
直接存 $W_\text{perm} \in \mathbb{R}^{L \times L}$（$L = TD$）：参数量 $L^2 = (TD)^2$
用 $G \otimes I_B$，只需存 $G \in \mathbb{R}^{(L/B) \times (L/B)}$：参数量 $(L/B)^2$，$B^2$ 倍压缩（$B$ 为块大小，即每个 head 的维度 $d$）
以上面的例子为例：$L=12$，$B=d=3$，$G$ 只有 $4 \times 4 = 16$ 个参数，而 $W_\text{perm}$ 有 $12 \times 12 = 144$ 个，压缩 $B^2 = 9$ 倍。
在 UniMixing 中，块大小 $B$ 是超参数，$G$（即 $W_G$）的大小为 $\frac{TD}{B} \times \frac{TD}{B}$，参数量为 $\left(\frac{TD}{B}\right)^2$。

对论文的意义：UniMixer 把固定的 $G$ 换成可学习的 $W_G$，继续保持这个分解结构。UniMixer-Lite 进一步用低秩分解 $W_G \approx A_G B_G$（$r \ll T$）压缩 $W_G$ 本身，总参数再降约 78%。

④ 对称性（Symmetry）
置换矩阵在 head 数 = token 数（$H=T$）时是对称的，否则不对称。这是原始 TokenMixer 的一个硬限制，灵活性差。

对论文的意义：UniMixer 通过 $\tilde{W}_G = (W_G + W_G^T)/2$ 主动施加对称约束，让可学习矩阵也满足这一性质。消融实验：去掉对称化后 AUC 下降 0.0573%。
UniMixer 的核心创新：把固定的 0/1 置换矩阵换成可学习的浮点矩阵 $W_G$，同时用 Sinkhorn 保持双随机性、温度退火保持稀疏性、显式对称化保持对称性、Kronecker 分解降低参数量——四条性质全部继承并升级为可学习版本。
💡 举例：置换矩阵等价变换

给定输入 $X \in \mathbb{R}^{2 \times 6}$（2 个 token，每个维度 6），TokenMixer 操作（head=2，每 head 维度 $d=3$）将每个 token 按 head 切分后转置重排：

$X = [[x_1,x_2,x_3,x_4,x_5,x_6],[x_7,x_8,x_9,x_{10},x_{11},x_{12}]]$

$\text{TokenMixer}(X) = [[x_1,x_2,x_3,x_7,x_8,x_9],[x_4,x_5,x_6,x_{10},x_{11},x_{12}]]$

这等价于用一个 $12 \times 12$ 的置换矩阵 $W_\text{perm}$ 乘以 $\text{flatten}(X)$。而 $W_\text{perm} = G \otimes I_3$，其中：

  • $I_3$：$3 \times 3$ 单位矩阵（每个 head 内部的 3 个维度顺序不变）
  • $G$：$4 \times 4$ 置换矩阵，作用于 4 个"块"（block 0=[x1-3]、block 1=[x4-6]、block 2=[x7-9]、block 3=[x10-12]），描述块与块之间怎么换位：block 0↔block 0，block 1↔block 2，block 3↔block 3，即：

$G = \begin{pmatrix}1&0&0&0\\0&0&1&0\\0&1&0&0\\0&0&0&1\end{pmatrix}$（block 1 和 block 2 互换，对应 head0/head1 在两 token 间的转置）

$G$ 只有 $4 \times 4 = 16$ 个参数，而直接存 $W_\text{perm}$ 需要 $12 \times 12 = 144$ 个数。这正是 Kronecker 分解压缩参数的出发点。

2.3 UniMixing 统一混合模块

通过对 $G$ 和各块的局部投影矩阵 $W^i_B$ 进行参数化，UniMixing 模块将原来固定的置换变成了可学习的：

从 $W_\text{perm} = G \otimes I_B$ 到 UniMixing 的推导：
原始 TokenMixer 的置换矩阵有两层结构：① $G$（块间）：描述哪个块换到哪个位置，固定 0/1 置换；② $I_B$（块内）：每个块内部维度顺序不变，恒等映射。

UniMixing 做了两个"可学习化"替换：
  • 把固定的 $G \in \{0,1\}^{(L/B)\times(L/B)}$ → 换成可学习的浮点矩阵 $W_G \in \mathbb{R}^{(L/B)\times(L/B)}$（保持双随机约束），控制块与块之间的交互强度
  • 把固定的 $I_B$（恒等，无参）→ 换成可学习的 $W^i_B \in \mathbb{R}^{B\times B}$，每个块有各自独立的局部混合矩阵，控制块内特征如何交互
这样整个操作就从"固定置换"变成了"可学习的广义 Kronecker 积"：$W_\text{perm} = G \otimes I_B \;\longrightarrow\; W_G \otimes \{W^i_B\}$
$$\text{UniMixing}(X) = \text{reshape}\left(\left(W_G \otimes \{W^i_B\}_{i=1}^{L//B}\right)\text{flatten}(X),\ 1, L\right)$$
符号说明
  • $B$：块大小（Block Size），将 $\text{flatten}(X)$ 等分为 $L/B$ 个向量
  • $L$：输入 embedding 的总维度（$= T \times D$）
  • $W^i_B \in \mathbb{R}^{B \times B}$：第 $i$ 块的局部混合矩阵（控制块内交互模式）
  • $W_G \in \mathbb{R}^{(L/B) \times (L/B)}$：全局混合矩阵（控制块间交互强度）
  • $\otimes$：广义 Kronecker 积

Kronecker 积 $W_G \otimes \{W^i_B\}$ 按定义是分块结构——$W_G$ 的每个元素 $w_{ij}$ 对应一个块操作 $W^j_B$，因此天然可以分两步独立执行，而不需要构造完整的 $L \times L$ 大矩阵：

两步执行的推导：将 $\text{flatten}(X) \in \mathbb{R}^L$ 切成 $L/B$ 段，每段 $x_i \in \mathbb{R}^B$。利用 Kronecker 积的混合乘积性质 $(A \otimes B)(C \otimes D) = (AC)\otimes(BD)$，可以先对每段独立做局部变换，再对结果统一做全局变换——两步的顺序就是 Kronecker 积定义的直接展开：
  • 局部先算：每段 $x_i$ 乘以 $W^i_B$，共 $L/B$ 段，每段代价 $B^2$，总代价 $\frac{L}{B} \times B^2 = LB$
  • 全局后算：把 $L/B$ 段的结果拼成矩阵 $H \in \mathbb{R}^{(L/B)\times B}$，再用 $W_G \in \mathbb{R}^{(L/B)\times(L/B)}$ 作用于 $H$ 的行维，代价 $\left(\frac{L}{B}\right)^2 \times B = \frac{L^2}{B}$
总代价 $O(LB + L^2/B)$，远小于直接展开大矩阵的 $O(L^2)$（当 $B$ 取合适值时）。
1
局部混合（Local Mixing）：将 $\text{flatten}(X)$ 等分为 $L/B$ 段，每段 $x_i \in \mathbb{R}^B$ 与其对应的 $W^i_B$ 相乘： $H = [x_1 W^1_B \mid x_2 W^2_B \mid \cdots \mid x_{L/B} W^{L/B}_B]$，得到局部交互矩阵 $H \in \mathbb{R}^{(L/B) \times B}$。计算量 $O(LB)$。
2
全局混合（Global Mixing）：$W_G \in \mathbb{R}^{(L/B) \times (L/B)}$ 作用在局部交互矩阵的行维上，得到最终输出。计算量 $O(L^2/B)$。

总计算量从直接使用 $W_\text{perm}$ 的 $O(L^2)$ 降为 $O(L^2/B + LB)$，同时避免了中间大矩阵的 GPU 显存占用。

每次 forward 时对 $W_G$（及 $W^i_B$）执行三步串联处理
得到满足双随机性 + 稀疏性 + 对称性约束的最终矩阵 $\bar{W}_G$，用于 UniMixing 计算
$W_G$
可学习参数
→
① 对称化
$\tilde{W}_G$
→
② 除以 $\tau$
控制稀疏
→
③ Sinkhorn
双随机归一
→
$\bar{W}_G$
用于计算
$$\bar{W}_G = \text{Sinkhorn-Knopp}\!\left(\frac{\tilde{W}_G}{\tau}\right),\quad \bar{W}^i_B = \text{Sinkhorn-Knopp}\!\left(\frac{\tilde{W}^i_B}{\tau}\right)$$
1 对称化：$\tilde{W}_G = (W_G + W_G^T)/2$
为什么能对称？$\tilde{W}_{ij} = (w_{ij}+w_{ji})/2$，$\tilde{W}_{ji} = (w_{ji}+w_{ij})/2$，两者必然相等，所以 $\tilde{W}_{ij}=\tilde{W}_{ji}$ 对所有 $i,j$ 成立。
为什么要对称？置换矩阵里"block $i$ 换到 block $j$"与"block $j$ 换到 block $i$"是互换关系，权重天然相等。显式对称化让可学习的 $W_G$ 继承这一结构约束。（消融：去掉后 AUC ↓ 0.0573%）
2 除以温度系数 $\tau$（$1.0 \to 0.05$ 退火）
为什么要稀疏？稀疏的 $\bar{W}_G$ 每行只有一个大权重，趋近置换矩阵，每个 block 被明确分配到某位置；若不稀疏则退化为均值池化，失去特征选择能力。
为什么 $\tau$ 小 → Sinkhorn 后更稀疏？除以小 $\tau$ 放大了元素间差距（$(w_{ij}-w_{ik})/\tau$ 比 $w_{ij}-w_{ik}$ 大得多），Sinkhorn 归一化后大元素更压倒性，结果更接近 one-hot。类比：softmax 低温趋近 argmax。
3 Sinkhorn-Knopp 迭代 → $\bar{W}_G$（双随机归一化）
交替做：行归一化（每行 ÷ 行和）→ 列归一化（每列 ÷ 列和）→ 反复 3～10 次收敛，得到行和=1 且 列和=1 的双随机矩阵 $\bar{W}_G$。
注意：这是每次 forward 里临时执行的少量迭代，不是训练多轮。$W_G$ 本身只由反向传播更新一次；Sinkhorn 是可微操作，梯度正常从 $\bar{W}_G$ 反传回 $W_G$。
Figure 3: 不同方法的 global mixing weights 及 TokenMixer 等价参数化
Figure 3（论文原图）：(a) 不同方法的全局混合权重分布。Heterogeneous Attention 的权重矩阵呈对角线主导（某些行稀疏集中），容易导致梯度消失。(b) TokenMixer 的等价参数化：原始 TokenMixer 操作等价于置换矩阵，可通过 Kronecker 压缩降低参数量。

2.4 UniMixing-Lite（轻量版）

UniMixing 模块存在两个冗余问题：① 局部矩阵 $W^i_B$ 数量随块数线性增长，引入局部交互冗余；② 全局矩阵 $W_G$ 较大，参数效率不高。UniMixing-Lite 用以下方法解决：

基矩阵组合（Basis Composition）：定义 $b$ 个基矩阵 $\{Z_\ell\}_{\ell=1}^b$ 和每块的权重向量 $\omega_i$，局部混合矩阵变为：$W^{*i}_B = \sum_{\ell=1}^b \omega_{i\ell} Z_\ell$。所有块共享基矩阵，只学习各自的加权系数，大幅减少参数。

低秩近似（Low-Rank Approximation）：全局矩阵 $W_G$ 改用低秩分解：$W_r = \text{Sinkhorn}(A_G B_G)$，其中 $A_G \in \mathbb{R}^{(L/B) \times r}$，$B_G \in \mathbb{R}^{r \times (L/B)}$，$r \ll L/B$。
$$\text{UniMixing-Lite}(X) = \text{reshape}\left(W_r \cdot \text{reshape}\left([x_1 W^{*1}_B \mid \cdots \mid x_{L/B} W^{*L/B}_B],\ \frac{L}{B}, B\right),\ 1, L\right)$$
符号说明
  • $W_r = \text{Sinkhorn-Knopp}(A_G B_G)$：低秩近似的全局混合矩阵
  • $W^{*i}_B = \text{Sinkhorn-Knopp}(\sum_\ell \omega_{i\ell} Z_\ell)$：基矩阵加权合成的局部混合矩阵
  • $r$：低秩近似的秩，越小参数越少
  • $b$：基矩阵数量，越多表达能力越强
💡 举例：UniMixing-Lite 参数节省

设 $L=768$，$B=6$，则共有 $L/B = 128$ 个块。

UniMixing（原版）：局部参数 = $128 \times 6 \times 6 = 4608$，全局参数 = $128 \times 128 = 16384$，总计约 21K。

UniMixing-Lite（$b=4$，$r=16$）：基矩阵参数 = $4 \times 6 \times 6 = 144$，权重参数 = $128 \times 4 = 512$，低秩全局参数 = $128 \times 16 \times 2 = 4096$，总计约 4.7K，参数量减少约 78%。

更重要的是，UniMixing-Lite 同时保留了 TokenMixer 的轻量全局交互 和 Attention 的局部表达能力，兼得两家优点。

2.5 SiameseNorm（孪生归一化）

❌ Pre-Norm 的问题：梯度越深越弱
x_out = x + Module(RMSNorm(x))
RMSNorm 先把 $x$ 压缩到均值0、方差1，Module 的输出 $\delta$ 因此变得很小；残差连接 $x + \delta$ 里 $x$ 的量级远大于 $\delta$，Module 的贡献被稀释。

梯度反传时，残差路（$x$ 那条）畅通无阻，但 Module 参数拿到的梯度信号微弱——层数越多，深层 Module 几乎学不到东西。

✅ 训练稳定    ❌ 深层模块梯度衰减，Scaling 深度无收益
❌ Post-Norm 的问题：训练初期数值爆炸
x_out = RMSNorm(x + Module(x))
Module 直接作用在未归一化的原始 $x$ 上，训练初期 $x$ 量级不稳定，两者相加 $x + \text{Module}(x)$ 时没有任何"安全网"——在 Norm 保护到来之前，数值就已经可能爆炸。

梯度通过 $x + \text{Module}(x)$ 反传时不受控，深层网络梯度爆炸概率高，很难稳定收敛。

✅ 梯度信号强、表达力强    ❌ 初期不稳定，深层难收敛
这就是 RankMixer 堆 4 层反而比 2 层 AUC 低 0.1066% 的根本原因——它用 Pre-Norm，越深梯度越弱，新增的层几乎没学到东西。

SiameseNorm 的解法：让稳定性和梯度信号走两条独立的路。每层维护两条流，初始化 $\bar{X}_0 = \bar{Y}_0 = X$。下图展示单层内的完整数据流：

第 ℓ 层数据流（两条流共享同一个 UniMixer 输出 Oₗ）
Ȳₗ（副流）
① RMSNorm(Ȳₗ) → Ỹₗ
送入 UniMixer（与主流相加）
Oₗ = UniMixer(X̄ₗ + Ỹₗ)
② Ȳₗ + Oₗ → Ȳₗ₊₁
只加，不做任何归一化
Ȳₗ₊₁（副流输出）
✅ 梯度不被 Norm 截断，直通回浅层
⟷
X̄ₗ（主流）
直接作为 UniMixer 输入的一部分
Oₗ = UniMixer(X̄ₗ + Ỹₗ)
③ RMSNorm(X̄ₗ + Oₗ) → X̄ₗ₊₁
先残差连接，再归一化
X̄ₗ₊₁（主流输出）
✅ 每层归一化，训练不爆炸
M 层后融合输出：$X_\text{output} = \bar{X}_M + \text{RMSNorm}(\bar{Y}_M)$
主流（稳定的归一化特征）+ 副流归一化后的累积信号 → 最终表示
$$\tilde{Y}_\ell = \text{RMSNorm}(\bar{Y}_\ell),\quad O_\ell = \text{UniMixer}(\bar{X}_\ell + \tilde{Y}_\ell)$$ $$\underbrace{\bar{X}_{\ell+1} = \text{RMSNorm}(\bar{X}_\ell + O_\ell)}_{\text{主流：每层归一化，保稳定}},\quad \underbrace{\bar{Y}_{\ell+1} = \bar{Y}_\ell + O_\ell}_{\text{副流：只加不归一化，保梯度}}$$
效果：UniMixer 2→4→8 层，AUC 持续提升（+0.1575% → +0.1647%）；而 RankMixer 4 层比 2 层反降 0.1066%。SiameseNorm 是让深度 Scaling 真正可行的工程关键。

2.6 统一理论框架

UniMixer 的最大理论贡献是：将现有三大主流推荐 Scaling 架构统一为同一个公式：

$$\text{UniMixing}(X) = \text{reshape}\left(\underbrace{G(X, W_G)}_{\text{全局混合}} \cdot \underbrace{\begin{bmatrix}x_1 W^1_B \\ \vdots \\ x_{L/B} W^{L/B}_B\end{bmatrix}}_{\text{局部混合}},\ 1, L\right)$$
方法 局部混合（Local Mixing） 全局混合模式 $G(X, W_G)$ 核心特点
Self-Attention $X W^V$ $\text{softmax}((XW^Q)(XW^K)^T / \sqrt{d})$ 全局依赖于输入 $X$（动态权重）
Heterogeneous Attention $X\tilde{W}^V$（token-specific） $\text{softmax}((X\tilde{W}^Q)(X\tilde{W}^K)^T / \sqrt{d})$ token-specific 投影解决异构问题
TokenMixer $X$（恒等映射，无参） $G$（固定置换，与 $X$ 无关） 静态无参，高效但无法学习
FM（Wukong） $Y$（固定投影） $XI(XI)^T$（当 $W^Q=W^K=I$ 时退化为 FM） 显式低阶特征交互
UniMixing（本文） $W^i_B$（可学习 token-specific） $\bar{W}_G$（可学习，满足双随机性） 统一三大范式，兼具可学习性和高效性
统一视角的关键洞察：当 $W_G$ 固定为置换矩阵时 = TokenMixer；当 $W_G$ 由 $X$ 内积计算时 = Attention；当 $W_G = XI(XI)^T$，局部投影固定时 = FM。UniMixer 将这三者统一在一个参数化框架下，通过梯度下降同时学习全局和局部的最优混合模式。
Attention（异构）
token 内部改造：通过 $W_V^{ih}$ 对每个 token 的特征做线性变换，再加权聚合
token 间交互：$Q_h K_h^T$ 内积动态计算（每次 forward 依赖输入 $X$）
问题：注意力分数易被各域 embedding norm 主导，塌缩到固定列
TokenMixer
token 内部：恒等映射，不改变块内特征
token 间交互：按 head 固定拼接，完全不可学习，块间换位规则由 $H, T$ 的几何关系预先决定
问题：无参数，无法根据数据自适应调整交互模式
UniMixer（本文）
token 内部（$W^i_B$）：可学习，每块独立学习块内特征的最优重组方式
token 间（$W_G$）：可学习，学习块间最优交互强度；静态参数（训练后固定，不依赖 $X$），介于 TokenMixer（静态无参）和 Attention（动态有参）之间
✅ 既有可学习性，又避免了 Attention 的 norm 塌缩问题

2.7 训练策略：温度退火 + 热启动

稀疏性对性能至关重要，但低温度（高稀疏）会导致梯度稀疏不稳定。论文采用线性温度退火：

$$\tau_j = \max\left\{\tau_\text{start} - (\tau_\text{start} - \tau_\text{end}) \cdot \frac{j}{J},\ \tau_\text{end}\right\}$$
符号说明
  • $\tau_\text{start} = 1.0$：初始高温（矩阵均匀分布，易于探索）
  • $\tau_\text{end} = 0.05$：最终低温（矩阵稀疏，性能最优）
  • $j$：当前迭代步；$J$：退火的总迭代步数

当数据量不足时，还可采用热启动（Warm-Up）策略：先用高温（$\tau=1.0$）训练直到收敛，再以此作为初始化重新用低温（$\tau=0.05$）训练。

💡 举例：温度系数对矩阵稀疏性的影响

从 Figure 5 可以看出：当 $\tau=1.0$ 时，$\bar{W}_G$ 和 $\bar{W}^i_B$ 的元素分布较均匀（类似均匀置换，信息混合平滑）；当 $\tau=0.05$ 时，矩阵呈现明显的稀疏集中（类似接近置换矩阵，每行/列只有少数几个大权重），这种稀疏模式对应更精准的特征交互选择，显著提升 AUC。

消融实验显示，移除温度系数（始终使用 $\tau=1.0$）会导致 AUC 下降 0.1645%，是所有消融项中影响最大的。

📊 §3 实验结果

实验设置：使用快手广告投放场景的真实数据集（超过 7 亿用户样本，涵盖数百个异构特征，时间跨度一年）。任务是预测"用户次日留存"（二分类）。评估指标：AUC、UAUC（用户级 AUC）。

3.1 主结果对比

模型 AUC ↑ ΔAUC 参数量 FLOPs/Batch
Heterogeneous Attention（基线） 0.744577 – 132.7M 1.68T
HiFormer 0.741685 -0.2892% 107.5M 1.37T
Wukong（FM-based） 0.744477 -0.0100% 107.1M 1.40T
FAT 0.744883 +0.0306% 138.4M 1.83T
RankMixer（TokenMixer-based） 0.749329 +0.4752% 135.5M 1.68T
TokenMixer-Large 0.748410 +0.3833% 103.3M 1.27T
UniMixer-2-Blocks 67.5M（本文） 0.749770 +0.5193% 67.5M 2.07T
UniMixer-2-Blocks 101.5M（本文） 0.750238 +0.5661% 101.5M 2.50T
UniMixer-Lite-2-Blocks 42.4M（本文） 0.751121 +0.6544% 42.4M 2.17T
UniMixer-Lite-4-Blocks 84.5M（本文） 0.752718 +0.8141% 84.5M 4.24T
关键结论：UniMixer-Lite-4-Blocks（84.5M）以更少的参数量（相比 Heterogeneous Attention 的 132.7M）实现了 +0.8141% 的 AUC 提升，且 FLOPs 远低于注意力机制类方法。UniMixer-Lite 是参数效率和计算效率的双重赢家。

3.2 Scaling Law 曲线

三种模型（RankMixer、UniMixer、UniMixer-Lite）均呈现清晰的幂律 Scaling 关系：

$$\Delta\text{AUC}_{\text{UniMixer-Lite}} = 0.003767 \cdot \text{Params}^{0.141903}$$ $$\Delta\text{AUC}_{\text{UniMixer}} = 0.003032 \cdot \text{Params}^{0.131973}$$ $$\Delta\text{AUC}_{\text{RankMixer}} = 0.002718 \cdot \text{Params}^{0.116043}$$
符号说明
  • $\Delta\text{AUC}$：相对基线（Heterogeneous Attention）的 AUC 提升量
  • 系数（如 0.003767）：Scaling 起点的高低
  • 指数（如 0.141903）：Scaling 效率的核心指标，越大意味着每扩大一倍参数量，AUC 提升幅度越大

UniMixer-Lite 的 Scaling 指数（0.1419）比 RankMixer（0.1160）高出约 22%，这意味着模型越大，两者的差距越来越悬殊。

Figure 4: Scaling Laws 曲线
Figure 4（论文原图）：AUC 与参数量/FLOPs 的 Scaling Law 曲线（对数刻度）。三条曲线均符合幂律，UniMixer-Lite 斜率最大，随着参数量增加优势持续扩大。

3.3 消融实验

消融设置 ΔAUC（相对完整 UniMixer） 结论
w/o 温度系数（始终 τ=1.0） -0.1645% 影响最大，稀疏性至关重要
w/o 模型热启动（Warm-Up） -0.0856% 数据不足时热启动收益显著
w/o 对称约束 -0.0573% 对称性有效约束混合模式
w/o 块特有局部权重（共享 $W_B$） -0.0436% 块特异性提升局部表达力
SiameseNorm → Post-Norm -0.0273% SiameseNorm 改善深层训练稳定性
Figure 5: 不同温度系数下的混合权重矩阵可视化
Figure 5（论文原图）：UniMixer-Lite 中 $\bar{W}_G$ 和 $\bar{W}^i_B$ 在不同温度系数下的可视化。低温（0.05）时矩阵明显更稀疏集中，验证了温度退火策略的有效性。

3.4 线上 A/B 实验

线上部署结论：UniMixer 和 UniMixer-Lite 已在快手多个广告投放场景全量部署。30 天留存（CAD D1-D30）平均提升超过 15%。这是工业推荐系统中非常显著的在线收益。
Figure 6: 2 Blocks vs 4 Blocks 的 Scaling 曲线对比
Figure 6（论文原图）：UniMixer/UniMixer-Lite 的 2 Blocks vs 4 Blocks Scaling 曲线，与 RankMixer 对比。UniMixer-Lite 4 Blocks 持续优于 2 Blocks，证明深度 Scaling 比宽度 Scaling 更高效。

深度 Scaling 的优势

Block 数 UniMixer-Lite ΔAUC RankMixer ΔAUC
2 Blocks 基线 基线
4 Blocks +0.1575% -0.1066%（退化！）
8 Blocks +0.1647% –

RankMixer 增加深度反而性能下降，根本原因是缺乏针对深层架构的专门设计。而 UniMixer 通过 SiameseNorm 彻底解决了这一问题。

💡 §4 亮点、不足与启发

核心亮点

1. 理论统一框架

通过 TokenMixer 的等价参数化，第一次把 Attention、TokenMixer、FM 三大孤立范式统一在一个理论框架下，揭示了它们的本质区别仅在于全局/局部混合模式的选择方式。

2. 参数效率极高

UniMixer-Lite 42.4M 参数实现的 AUC 超过 RankMixer 135.5M 参数，参数量缩减约 70% 而性能反超。对工业部署成本意义重大。

3. 深度 Scaling 破解

SiameseNorm 首次让推荐系统实现随深度增加持续性能提升，解决了 RankMixer 等模型深度扩展退化的顽症。

4. 工业级验证

线上 A/B 实验 CAD +15%，来自快手真实广告系统，7 亿样本规模，结论可信度高。

潜在不足与局限

  • FLOPs 更高：UniMixer-Lite-4-Blocks 的 FLOPs（4.24T/Batch）显著高于 RankMixer（1.68T）和 TokenMixer-Large（1.27T），实际部署时延是重要考量。
  • 超参数较多：温度系数 $\tau$、退火策略、块大小 $B$、基矩阵数 $b$、低秩 $r$ 等超参数需要仔细调优，工程复杂度较高。
  • 单场景验证：实验场景仅为快手广告投放（用户留存预测），在其他推荐场景（如内容推荐、电商排序）的泛化性尚未验证。
  • Sinkhorn 迭代开销：Sinkhorn-Knopp 操作需要额外的迭代归一化，训练中会引入额外计算。

个人理解与启发

最值得借鉴的核心思想是「参数化等价替换」：找到固定规则操作（TokenMixer）背后等价的矩阵结构，然后把这个矩阵变成可学习的，同时利用结构属性（Kronecker 分解、稀疏性）保持高效性。这种思路本质上是在「规则先验」和「数据驱动学习」之间取平衡——先验结构保证效率，可学习性保证适应性。
对推荐系统 Scaling 的宏观启示：Attention 的全局动态权重（强表达力但训练难）、TokenMixer 的静态全局混合（高效但不可学习）、FM 的显式低阶交互（可解释但上限低）——这三条路各有 tradeoff，UniMixer 的统一框架说明它们不是互相排斥的，而是同一框架下的不同特例。未来的 Scaling 架构设计可以在这个统一框架内做更系统的探索。

与相关工作的关键区别

方法 全局混合 局部混合 深度 Scaling 是否统一理论
RankMixer 固定置换（无参） 无 ❌ 退化 ❌
HiFormer 动态 softmax（依赖 X） token-specific W_V 部分支持 ❌
Wukong FM 显式交互 固定 Y 支持 ❌
UniMixer（本文） 可学习双随机矩阵 $\bar{W}_G$ 可学习 token-specific $\bar{W}^i_B$ ✅ SiameseNorm ✅ 统一三大范式