UniMixer: A Unified Architecture for Scaling Laws in Recommendation Systems

基本信息

字段 内容
标题 UniMixer: A Unified Architecture for Scaling Laws in Recommendation Systems
作者 Mingming Ha, Guanchen Wang, Linxun Chen, Xuan Rao, Yuexin Shi, Tianbao Ma, Zhaojie Liu, Yunqian Fan, Zilong Lu, Yanan Niu, Han Li, Kun Gai
机构 Kuaishou
年份 2026 (arXiv: 2604.00590v1, 2026-04-01)
方向 Scaling Law, Unified Token Mixer, Feature Interaction Architecture, Sparse-aware Design
场景 推荐排序模型的统一建模架构 + scaling law
arXiv https://arxiv.org/abs/2604.00590

特征的异构性决定了 token 间的交互关系天然就不应该是稠密的。把参数花在没意义的交互上是噪音。

一、论文简介

UniMixer 是快手 2026 年提出的推荐系统统一缩放架构。论文的核心贡献是建立了一个统一的理论框架,将推荐系统中三大主流建模范式(Attention-based、TokenMixer-based、FM-based)桥接在一起,并在此基础上设计了具有良好 scaling law 的模型架构。

二、背景与问题

背景:推荐系统的 Scaling Law

大语言模型(LLM)已经展现出清晰的 scaling law——模型越大、数据越多、算力越多,性能就越好。推荐系统社区也在探索类似的规律,希望通过增加参数量和计算量来持续提升模型效果。

问题:三大范式各有局限

目前推荐系统中有三种主流的特征交互建模方式,但它们在设计哲学和架构上差异很大,各有短板:

  1. Attention-based 方法参数量大、计算开销高,但交互能力强
  2. TokenMixer-based 方法高效但交互模式固定,缺乏自适应能力
  3. FM-based 方法简洁但表达能力有限

这三种方法之间缺乏统一的理论联系,无法互相借鉴优势。更关键的是,现有方法在 scaling 时表现不佳——单纯增加参数并不能持续带来性能提升。

解决思路

  1. 将 TokenMixer 的规则操作参数化,发现它等价于一个置换矩阵乘法
  2. 进一步证明 Attention-based、TokenMixer-based、FM-based 三种方法都可以统一到同一个框架下(Table 1)
  3. 基于这个统一框架设计 UniMixer,同时具备三种方法的优势
  4. 设计轻量版 UniMixing-Lite,进一步压缩参数和计算量

三、三种主流建模范式详解

1. Attention-Based(注意力机制)

代表方法:Heterogeneous Attention、Self-Attention、FiT

核心思路:通过 Query、Key、Value 投影计算 token 之间的注意力权重,动态决定哪些特征之间应该交互、交互强度多大。

以 Heterogeneous Attention 为例:

给定输入隐状态 X=[x1,,xT]RT×DX = [x_1, \dots, x_T] \in \mathbb{R}^{T \times D},每个 token xix_i 有自己专属的投影矩阵:

Qi=xiHiQ,Ki=xiHiK,Vi=xiHiVQ_i = x_i H_i^Q, \quad K_i = x_i H_i^K, \quad V_i = x_i H_i^V

其中 HiQ,HiK,HiVRD×DH_i^Q, H_i^K, H_i^V \in \mathbb{R}^{D \times D} 是第 ii 个 token 专属的权重。多头注意力计算:

Oh=softmax(QhKhTd)VhRT×dO_h = \text{softmax}\left(\frac{Q_h K_h^T}{\sqrt{d}}\right) V_h \in \mathbb{R}^{T \times d}

多头输出拼接后过线性投影对齐维度。

优点:每个特征域有独立的投影,能捕捉异构特征之间差异化的交互模式
缺点:参数量是标准 attention 的 TT 倍(TT = token 数),计算开销大

2. TokenMixer-Based(规则混合)

代表方法:RankMixer、TokenMixer-Large

核心思路:不使用注意力机制,而是通过固定的规则(split + concat + circular shift)来混合 token 之间的信息。

给定输入 X=[x1,,xT]X = [x_1, \dots, x_T],首先将每个 token xix_i 拆分为 HH 个 head:

[xi1xi2xiH]=Split(Block(xi))[x_i^1 | x_i^2 | \cdots | x_i^H] = \text{Split}(\text{Block}(x_i))

然后对第 kk 个 head 的子向量做 circular shift(循环移位 kk 位):

xk=concat(x1k,x2k,,xTk)2kx^{*k} = \text{concat}(x_1^k, x_2^k, \dots, x_T^k) \gg 2^k

最终输出:

TokenMixer(X)=reshape([x1x2||xH])\text{TokenMixer}(X) = \text{reshape}\left(\left[\frac{x^{*1}}{x^{*2}} \middle| \cdots \middle| \frac{x^{*H}}{}\right]\right)

其中要求 H=TH = T,输出维度与输入 XX 保持一致。

具体例子T=4,D=4,H=4T=4, D=4, H=4):

head0 head1 head2 head3
输入 [a1,b1,c1,d1][a_1, b_1, c_1, d_1] [a2,b2,c2,d2][a_2, b_2, c_2, d_2] [a3,b3,c3,d3][a_3, b_3, c_3, d_3] [a4,b4,c4,d4][a_4, b_4, c_4, d_4]
shift kk 位后 [a1,b1,c1,d1][a_1, b_1, c_1, d_1] [b2,c2,d2,a2][b_2, c_2, d_2, a_2] [c3,d3,a3,b3][c_3, d_3, a_3, b_3] [d4,a4,b4,c4][d_4, a_4, b_4, c_4]

按 token 重新拼回:x1=[a1,b2,c3,d4]x_1^* = [a_1, b_2, c_3, d_4],每个输出 token 混合了来自不同原始 token 的信息。

优点:无 Q/K/V 投影,无 softmax,计算极快
缺点:混合模式完全固定,无法根据输入自适应调整;要求 H=TH=T

论文的关键发现:这个操作等价于乘以一个置换矩阵 WpermW^{\text{perm}},而且 WpermW^{\text{perm}} 可以分解为全局矩阵和局部矩阵的 Kronecker 积:

Wperm=GLRT2×T2W^{\text{perm}} = G \otimes L \in \mathbb{R}^{T^2 \times T^2}

3. FM-Based(因子分解机)

代表方法:Wukong

核心思路:通过因子分解机(Factorization Machine)的方式建模特征交互,输出 FMB 的结果后接线性投影。

FM(X)=XXTV\text{FM}(X) = X \cdot X^T \cdot V

FMB(X)=reshape(MLP(LN(flatten(FM(X)))))\text{FMB}(X) = \text{reshape}(\text{MLP}(\text{LN}(\text{flatten}(\text{FM}(X)))))

其中 VRT×TV \in \mathbb{R}^{T \times T} 是可学习的投影矩阵,用于减少内存占用。最后通过线性投影上采样交互结果:

LCS(X)=HX\text{LCS}(X) = H \cdot X

优点:结构简洁,参数量小
缺点:交互能力有限,缺乏显式的 token 级别混合机制

三种方法的统一视角(Table 1)

方法 异构特征混合 局部混合模式 全局混合模式
Self-Attention XHX \cdot H softmax(QKTd)V\text{softmax}(\frac{QK^T}{\sqrt{d}}) \cdot V
Heterogeneous Attention XHiX \cdot H_i softmax(QKTd)V\text{softmax}(\frac{QK^T}{\sqrt{d}}) \cdot V
TokenMixer II(单位矩阵) GG(固定置换)
FM VV XXTX \cdot X^T

UniMixer 的统一公式:

UniMixer(X)=reshape(G(WLHL)V,)\text{UniMixer}(X) = \text{reshape}\left(G \cdot \left(W_L \cdot H_L\right) \cdot V, \dots\right)

其中 GG 是全局混合模式,WLW_L 是局部混合模式,HLH_L 是特征投影。

四、UniMixer 的关键设计

4.1 Feature Tokenization(特征分词)

将异构输入特征按语义域分组,每组映射为一个 token。输入特征 X\mathbf{X} 首先按 NN 个语义域分组:

X=Xdomain1用户特征Xdomain2商品特征XdomainN上下文特征\mathbf{X} = \underbrace{\mathbf{X}_{\text{domain}_1}}_{\text{用户特征}} \oplus \underbrace{\mathbf{X}_{\text{domain}_2}}_{\text{商品特征}} \oplus \cdots \oplus \underbrace{\mathbf{X}_{\text{domain}_N}}_{\text{上下文特征}}

每个特征域独立 embedding 后拼接为向量 E=[e1,e2,,eN]E = [e_1, e_2, \dots, e_N],再通过 token-specific 线性投影统一到 block 维度 BB

hi=HiprojEdi×dieiRBh_i = H_i^{\text{proj}} E_{d_i \times d_i} \cdot e_i \in \mathbb{R}^B

输入隐状态 SRT×BS \in \mathbb{R}^{T \times B} 通过堆叠 eie_i 按列获得。

4.2 UniMixer Block

核心模块包含三个部分:

a) 异构特征交互(Heterogeneous Feature Interaction)

论文发现标准 attention 的一个问题:QQKK 用随机初始化的共享权重投影后,注意力权重被 input token 的数值主导,而非学到有意义的交互模式(Figure 3a)。

UniMixer 的解决方案:用可学习的参数矩阵 WGW_G(全局)和 WLW_L(局部)直接参数化混合模式,不依赖 QKTQK^T 计算。

UniMixer(X)=reshape(WGreshape([x1W1,x2W2,,xBWB],LB,B),1,L)\text{UniMixer}(X) = \text{reshape}\left(W_G \cdot \text{reshape}\left(\left[x_1 W_1, x_2 W_2, \dots, x_B W_B\right], \frac{L}{B}, B\right), 1, L\right)

其中:

  • WGR(L/B)×(L/B)W_G \in \mathbb{R}^{(L/B) \times (L/B)}:全局混合矩阵,控制 token 之间的交互
  • WbRB×BW_b \in \mathbb{R}^{B \times B}:局部混合矩阵,控制 block 内部的交互
  • 通过广义 Kronecker 积 \otimes 组合全局和局部模式

b) 温度系数与 Sinkhorn 归一化

为了让参数矩阵满足置换矩阵的约束(双随机性、稀疏性、对称性),论文引入了温度系数 τ\tau

步骤1:对称化

WGWG+WGT2,WLWL+WLT2W_G \leftarrow \frac{W_G + W_G^T}{2}, \quad W_L \leftarrow \frac{W_L + W_L^T}{2}

步骤2:温度缩放 + Sinkhorn-Knopp 归一化

WGSinkhorn-Knopp(WGτ),WLSinkhorn-Knopp(WLτ)W_G \leftarrow \text{Sinkhorn-Knopp}\left(\frac{W_G}{\tau}\right), \quad W_L \leftarrow \text{Sinkhorn-Knopp}\left(\frac{W_L}{\tau}\right)

τ\tau 的效果:

  • τ=1.0\tau = 1.0 → 权重分布均匀,所有特征交互强度接近
  • τ=0.05\tau = 0.05 → 权重分布尖锐,只有少数关键交互被保留(块对角结构)

温度退火策略(Model Warm-Up)

τj=τinit(τendτinit)min(jJ,1)\tau_j = \tau_{\text{init}} \cdot \left(\frac{\tau_{\text{end}}}{\tau_{\text{init}}}\right)^{\min\left(\frac{j}{J},\, 1\right)}

其中 τinit\tau_{\text{init}} 为初始温度(较高,如 1.0),τend\tau_{\text{end}} 为最终温度(较低,如 0.05),jj 为当前训练步数,JJ 为退火总步数。

  • 训练初期:τ\tau 较高,保持平滑探索,避免过早陷入局部最优
  • 训练后期:τ\tau 逐渐降低,锁定最优的稀疏交互模式

c) SiameseNorm

借鉴 SiameseNorm 的设计,在 UniMixer block 中引入 Pre-Norm 和 Post-Norm 两条归一化路径:

对于第 ll 层 UniMixer Block:

X^l=RMSNorm(Xl1)\hat{X}_l = \text{RMSNorm}(X_{l-1})

Ol=UniMixer(X^l+El)O_l = \text{UniMixer}(\hat{X}_l + E_l)

Xl+1,1=RMSNorm(OlVl+Yl+bl)X_{l+1,1} = \text{RMSNorm}(O_l \cdot V_l + Y_l + b_l)

Xout=E0+RMSNorm(Yl+1)X_{\text{out}} = E_0 + \text{RMSNorm}(Y_{l+1})

其中 ElE_l 是输入 embedding,VlV_lYlY_l 是残差连接项。

4.3 UniMixing-Lite(轻量版)

为进一步压缩参数和计算量,设计了轻量版:

核心改动:

  1. 局部混合矩阵 WbW_b 用一组基向量的线性组合生成:

Wb=kαb,kAkW_b = \sum_{k} \alpha_{b,k} \cdot A_k

共享基矩阵 {Ak}\{A_k\},每个 block 只学系数 αb,k\alpha_{b,k}

  1. 全局混合矩阵 WGW_G 做低秩近似,用 Sinkhorn-Knopp + 基矩阵近似。

  2. UniMixing-Lite 的完整公式:

UniMixing-Lite(X)=reshape(H1reshape([x1H21,,xBH2B],LB,B),1,L)\text{UniMixing-Lite}(X) = \text{reshape}\left(H_1 \cdot \text{reshape}\left(\left[x_1 H_2^1, \dots, x_B H_2^B\right], \frac{L}{B}, B\right), 1, L\right)

其中 H1=Sinkhorn-Knopp(kαkAk)H_1 = \text{Sinkhorn-Knopp}\left(\sum_{k} \alpha_k A_k\right)H2bRB×BH_2^b \in \mathbb{R}^{B \times B}rr 是低秩近似的秩。

效果:参数量大幅减少,但性能几乎不降。

五、实验结果与发现

5.1 Scaling Law 对比(Q1)

在约 100M 参数的广告推荐场景下,UniMixer 和 UniMixer-Lite 在所有指标上超越 SOTA:

模型 AUC Δ\DeltaAUC UAUC Δ\DeltaUAUC Params FLOPs/Batch
Heterogeneous Attention 0.744577 0.733829 132.7M 1.68T
HiFormer 0.741685 -0.2892% 0.731086 -0.2743% 107.5M 1.37T
Wukong 0.744477 -0.0100% 0.733849 0.0020% 107.1M 1.40T
FAT 0.744883 0.0306% 0.734280 0.0451% 138.4M 1.83T
RankMixer 0.749329 0.4752% 0.738938 0.5109% 135.5M 1.68T
TokenMixer-Large 0.748410 0.3833% 0.737940 0.4111% 103.3M 1.27T
UniMixer-2-Blocks 67.5M 0.749770 0.5193% 0.739331 0.5502% 67.5M 2.07T
UniMixer-2-Blocks 101.5M 0.750238 0.5661% 0.739983 0.6154% 101.5M 2.50T
UniMixer-Lite-2-Blocks 42.4M 0.751121 0.6544% 0.740739 0.6910% 42.4M 2.17T
UniMixer-Lite-2-Blocks 76.2M 0.751401 0.6824% 0.741215 0.7386% 76.2M 2.60T
UniMixer-Lite-4-Blocks 38.2M 0.752327 0.7750% 0.742091 0.8190% 38.2M 1.26T
UniMixer-Lite-4-Blocks 84.5M 0.752718 0.8141% 0.742530 0.8701% 84.5M 4.24T

关键发现:

  • UniMixer-2-Blocks 67.5M 用约一半的参数就超过了 RankMixer(135.5M)和 TokenMixer-Large(103.3M)
  • UniMixer-Lite-4-Blocks 38.2M 仅用 38.2M 参数和 1.26T FLOPs 就达到了最高水平的 AUC,参数效率和计算效率都是最优的
  • 所有 UniMixer 变体在 AUC 和 UAUC 上均超过全部 baseline
  • Scaling law 拟合公式显示 UniMixer-Lite 的 scaling exponent 最大,意味着它从增加参数中获益最多

5.2 消融实验(Q2)

Table 3 展示了各组件的贡献(基于 UniMixer 8.57M):

去掉的组件 AUC 下降
温度系数 (Temperature Coefficient) -0.1489%
对称性约束 (Symmetry Constraint) -0.0489%
Block-Specific 局部混合权重 -0.0649%
Block Waste (无效 block 填充) -0.0716%
SiameseNorm → Post-Norm -0.0557%

发现:

  • 温度系数是影响最大的单一组件
  • 去掉任何约束都会导致性能下降,验证了置换矩阵约束的必要性
  • Model warm-up 也很关键,直接用低温训练会导致性能退化

5.3 温度系数的可视化分析

Figure 5 展示了不同温度下混合矩阵的形态:

  • τ=1\tau = 1:矩阵接近均匀分布,所有特征交互强度相近
  • τ=0.05\tau = 0.05:矩阵呈现清晰的块对角结构,模型学会了哪些特征域之间应该交互

这说明低温不是简单地让权重"变稀疏",而是让模型学到了有结构的、有意义的交互模式。

5.4 UniMixing-Lite 的 Scaling 表现(Q3)

  • 增加 base number(基向量数量)比增加 rank 更有效
  • 增加 block 深度比增加宽度更有效(深度 scaling 优于宽度 scaling)
  • UniMixer-Lite 2-Blocks 和 4-Blocks 的 scaling curve 都优于 RankMixer

5.5 线上 A/B 测试(Q4)

在快手广告投放系统的多个场景中部署:

  • 使用 CAD(Consecutive Active Days,连续活跃天数)作为核心指标
  • 30 天观察窗口内,CAD of Eli-D30 平均提升超过 1%
  • 验证了 UniMixer/UniMixer-Lite 在真实线上环境的有效性

六、讨论:混合矩阵稀疏性的影响

6.1 温度系数如何控制稀疏性

UniMixer 中混合矩阵的稀疏性完全由温度系数 τ\tau 控制。当 τ\tau 较高时(如 1.0),Sinkhorn-Knopp 归一化后的矩阵接近均匀分布——每个特征域和其他所有特征域的交互强度差不多。当 τ\tau 降低到 0.05 时,矩阵变得尖锐,呈现出清晰的块对角结构,只有少数特征域之间保留了强交互。

这不是简单的"稀疏 = 好"。论文 Figure 5 的可视化表明,低温下的稀疏模式是有结构的——模型学到了哪些特征域之间应该交互(比如用户画像和商品类别),哪些不应该(比如设备信息和商品价格)。这种结构化稀疏比随机稀疏或均匀稠密都更有效。

6.2 稀疏性与 Scaling Law 的关系

消融实验(Table 3)显示,去掉温度系数导致 AUC 下降 0.1489%,是所有组件中影响最大的。这说明稀疏性对 scaling 至关重要:

  • 没有稀疏约束时,模型倾向于让所有特征均匀交互,参数被浪费在无意义的交互上,增加参数量带来的边际收益递减
  • 有稀疏约束时,模型把参数集中在关键交互上,每增加一个参数都能被有效利用,scaling exponent 更大

image-20260526013628333

同时,参考上图,稀疏性不等同于上图中 attention weight 集中在某个 token,而是动态地集中在某些关键的 token 上,特征的异构性本身就决定了 token 之前的交互关系不是稠密的,过于稠密会增加很多噪音。

七、总结

UniMixer 的核心贡献不是提出一个全新的模型,而是发现了三种看似不同的推荐建模范式之间的内在联系,并基于这个统一视角设计了一个兼具三者优势的架构。温度系数、Sinkhorn 归一化、SiameseNorm 等设计都服务于一个目标:让模型在 scaling 时能持续获益,而不是简单堆参数。