UniMixer: A Unified Architecture for Scaling Laws in Recommendation Systems¶

1. 核心问题与动机¶

推荐系统领域近年来借鉴LLM的Scaling Law思想，通过设计可堆叠的Scaling模块并增加模型深度来建立模型性能与参数量/FLOPs之间的Scaling Law。目前主要存在三种Scaling架构范式：

1. Attention-based：如HiFormer、FAT、HHFT，使用token-specific的Q/K/V投影实现异构特征交互
2. TokenMixer-based：如RankMixer、TokenMixer-Large，使用无参数的规则token mixing操作
3. FM-based：如Wukong、Kunlun，通过Factorization Machine Block建模特征交互

这三类方法在设计哲学和结构上存在本质差异，但都展现了Scaling能力。论文提出一个核心问题：能否构建一个统一的Scaling模块，融合三类方法的优势？

2. 方法详解¶

2.1 Feature Tokenization¶

输入特征 $\mathbf{X}$ 按语义类别划分为 $N$ 个不相交的特征域：

$$\mathbf{X} = \left[\underbrace{\mathbf{x}_U^{(1)}, \ldots, \mathbf{x}_U^{(n_U)}}_{\text{User Profile}}, \underbrace{\mathbf{x}_I^{(1)}, \ldots, \mathbf{x}_I^{(n_I)}}_{\text{Item Features}}, \underbrace{\mathbf{x}_B^{(1)}, \ldots, \mathbf{x}_B^{(n_B)}}_{\text{Behavior Sequence}}, \underbrace{\mathbf{x}_Q^{(1)}, \ldots, \mathbf{x}_Q^{(n_Q)}}_{\text{Query Features}}, \ldots\right]$$

每个特征域通过独立的embedding层映射为 $d_{\text{domain}}$ 维向量：

$$\mathbf{e}_n = \text{Embedding}(\mathbf{X}_{\text{domain}}) \in \mathbb{R}^{d_{\text{domain}}}$$

所有域的embedding拼接后均匀分块，通过token-specific线性层投影为token embedding：

$$\boldsymbol{x}_i = W_i^{\text{proj}} \mathbf{E}_{d_i:d_{i}+d} + \mathbf{b}_i^{\text{proj}} \in \mathbb{R}^D$$

2.2 TokenMixer的等价参数化¶

论文的关键发现是TokenMixer操作可以等价表示为一个置换矩阵 $W^{\text{perm}}$ 与展平输入的矩阵乘法：

$$\text{TokenMixer}(X) = \text{reshape}(W^{\text{perm}} \text{flatten}(X))$$

其中 $W^{\text{perm}} \in \mathbb{R}^{TD \times TD}$ 是一个大型置换矩阵。

置换矩阵 $W^{\text{perm}}$ 的性质：

• 可压缩性 (Compressibility)：$W^{\text{perm}} = G \otimes I$，其中 $I \in \mathbb{R}^{\frac{D}{T} \times \frac{D}{T}}$ 是单位矩阵，$G \in \mathbb{R}^{T^2 \times T^2}$，$\otimes$ 是Kronecker积
• 双随机性 (Doubly Stochastic)：每行每列之和为1
• 稀疏性 (Sparsity)：每行/每列恰好有一个非零元素
• 对称性 (Symmetry)：当 $T = H$ 时，$W^{\text{perm}}$ 是对称矩阵

2.3 UniMixing模块¶

基于上述发现，论文定义block size $B$，blocks数量为 $(L // B)^2$，其中 $L$ 是embedding维度。通过参数化矩阵 $W_G$（全局mixing）和 $\{W_B^i\}$（局部mixing），构建UniMixing模块：

$$\text{UniMixing}(X) = \text{reshape}\left(\left(W_G \otimes \{W_B^i\}_{i=1}^{L//B}\right) \text{flatten}(X), 1, L\right)$$

计算流水线优化：

先将 $\text{flatten}(X)$ 均匀分为 $L//B$ 个向量：

$$\left[\boldsymbol{x}_1 | \boldsymbol{x}_2 | \ldots | \boldsymbol{x}_{\frac{L}{B}}\right] = \text{Split}\left(\text{flatten}(X), \frac{L}{B}\right)$$

然后分别与block权重相乘得到局部特征交互：

$$\text{reshape}\left(H, \frac{L}{B}, B\right) = \begin{bmatrix} \boldsymbol{x}_1 W_B^1 \\ \vdots \\ \boldsymbol{x}_{\frac{L}{B}} W_B^{\frac{L}{B}} \end{bmatrix}$$

最终输出：

$$\text{UniMixing}(X) = \text{reshape}\left(W_G \text{reshape}\left(H, \frac{L}{B}\right), 1, L\right)$$

该优化将计算复杂度从 $O(L^2)$ 降至 $O(L^2/B + LB)$。

2.4 约束满足¶

为使学习到的参数矩阵满足双随机性、稀疏性和对称性：

• 对称性：$\tilde{W}_G = \frac{W_G + W_G^\top}{2}$，$\tilde{W}_B^i = \frac{W_B^i + W_B^{i\top}}{2}$
• 双随机性和稀疏性：通过Sinkhorn-Knopp迭代和温度系数控制

$$\hat{W}_G = \text{Sinkhorn-Knopp}\left(\frac{\tilde{W}_G}{\tau}\right), \quad \hat{W}_B^i = \text{Sinkhorn-Knopp}\left(\frac{\tilde{W}_B^i}{\tau}\right)$$

残差连接和归一化处理输出：

$$O = \text{RMSNorm}(X + \text{UniMixing}(X))$$

2.5 统一理论框架¶

论文建立了统一的特征交互理论框架（公式17）：

$$\text{UniMixing}(X) = \text{reshape}\left(\underbrace{G(X, W_G)}_{\text{Global Mixing Pattern}} \begin{bmatrix} \boldsymbol{x}_1 W_B^1 \\ \vdots \\ \boldsymbol{x}_{\frac{L}{B}} W_B^{\frac{L}{B}} \end{bmatrix}_{\text{Local Mixing Pattern}}, 1, L\right)$$

关键联系：

• 当 $L//B = T$ 且 $W_V^{1h} = W_B^i$ 时，UniMixer的局部交互投影等价于异构attention的value投影
• $W_G$ 的角色与attention权重相同，但需满足双随机性、稀疏性和对称性约束
• 当attention中 $W_Q = I, W_K = I, V_h = W_V = Y$ 时，attention退化为FM模块

2.6 UniMixing-Lite¶

为进一步降低参数量和计算成本，设计了轻量化的UniMixing-Lite模块：

局部交互的basis-composed策略：定义一组basis矩阵 $\{Z_\ell\}_{\ell=1}^b$ 和block-specific权重向量 $\{\boldsymbol{\omega}^i\}_{i=1}^{L//B}$：

$$W_B^{*i} = \text{Sinkhorn-Knopp}\left(\sum_{\ell=1}^{b} \omega_\ell^i Z_\ell\right)$$

全局交互的低秩近似：用 $A_G B_G$ 近似 $W_G$，其中 $A_G \in \mathbb{R}^{(L//B) \times r}$，$B_G \in \mathbb{R}^{r \times (L//B)}$：

$$W_r = \text{Sinkhorn-Knopp}(A_G B_G)$$

完整的UniMixing-Lite公式：

$$\text{UniMixing-Lite}(X) = \text{reshape}\left(W_r \text{reshape}\left(\left[\boldsymbol{x}_1 W_B^{*1} | \ldots | \boldsymbol{x}_{\frac{L}{B}} W_B^{*\frac{L}{B}}\right], \frac{L}{B}, B\right), 1, L\right)$$

2.7 Pertoken SwiGLU¶

每个token使用独立的SwiGLU前馈网络建模特征异构性：

$$\text{pSwiGLU}(\boldsymbol{o}_i) = W_{\text{down}}^i ((W_{\text{up}}^i \boldsymbol{o}_i + b_{\text{up}}^i) \odot \text{Swish}(W_{\text{gate}}^i \boldsymbol{o}_i + b_{\text{gate}}^i)) + b_{\text{down}}^i$$

2.8 SiameseNorm¶

引入SiameseNorm解决Pre-Norm和Post-Norm之间的矛盾，维护两个耦合流 $\bar{X}_\ell$ 和 $\bar{Y}_\ell$：

$$\bar{Y}_\ell = \text{RMSNorm}(\bar{Y}_\ell), \quad O_\ell = \text{UniMixer}(\bar{X}_\ell + \bar{Y}_\ell)$$ $$\bar{X}_{\ell+1} = \text{RMSNorm}(\bar{X}_\ell + O_\ell), \quad \bar{Y}_{\ell+1} = \bar{Y}_\ell + O_\ell$$

最终输出：

$$X_{\text{output}} = \bar{X}_M + \text{RMSNorm}(\bar{Y}_M)$$

2.9 温度退火训练策略¶

温度系数 $\tau$ 控制参数矩阵的稀疏度。采用线性温度退火：

$$\tau_j = \max\left\{\tau_{\text{start}} - \frac{(\tau_{\text{start}} - \tau_{\text{end}})j}{J}, \tau_{\text{end}}\right\}$$

从较高初始温度（如 $\tau = 1.0$）逐渐退火至 $\tau = 0.05$。也可采用两阶段策略：先用高温cold-start训练，再用低温重训。

3. 实验¶

3.1 实验设置¶

• 数据集：快手广告投放场景的真实logged数据，包含超过7亿用户样本，覆盖一年
• 标签：二元标签（User Retention = 1/0），表示用户激活后次日是否回访快手
• 特征：数百个异构特征（数值特征、ID特征、交叉特征、序列特征）
• 评估指标：AUC、UAUC（User-Level AUC）评估性能；参数量、FLOPs、MFU评估效率
• 训练设置：40 GPU混合分布式训练，Adam优化器，学习率0.001

Baselines：

• Attention-based：Heterogeneous Attention、HiFormer、FAT
• TokenMixer-based：RankMixer、TokenMixer-Large
• FM-based：Wukong

3.2 主实验结果（~100M参数）¶

Table 2: 约100M参数UniMixer与SOTA模型在广告场景的性能和效率对比

结论：

• 在更小的参数量和计算成本下，UniMixer和UniMixer-Lite均显著优于所有SOTA模型
• UniMixer-Lite-4-Blocks 38.2M仅用38.2M参数（约为RankMixer的28%）就超越了所有baseline
• 最优的UniMixer-Lite-4-Blocks 84.5M在AUC上领先RankMixer 0.3389%，UAUC领先0.3592%

3.3 Scaling Law对比¶

三个模型的Scaling Law公式：

$$\Delta \text{AUC}_{\text{RankMixer}} = 0.002718 \cdot \text{Params}^{0.116043}, \quad \Delta \text{AUC}_{\text{RankMixer}} = 0.002022 \cdot \text{FLOPs}^{0.116635}$$

$$\Delta \text{AUC}_{\text{UniMixer}} = 0.003032 \cdot \text{Params}^{0.131973}, \quad \Delta \text{AUC}_{\text{UniMixer}} = 0.002058 \cdot \text{FLOPs}^{0.125702}$$

$$\Delta \text{AUC}_{\text{UniMixer-Lite}} = 0.003767 \cdot \text{Params}^{0.141903}, \quad \Delta \text{AUC}_{\text{UniMixer-Lite}} = 0.002338 \cdot \text{FLOPs}^{0.135327}$$

UniMixer-Lite在参数和FLOPs上都具有最大的Scaling指数和系数，Scaling效率最优。

3.4 消融实验¶

Table 3: UniMixer 6.57M各组件消融

结论：

• 温度系数对性能影响最大（$\Delta$AUC -0.1645%），验证了参数矩阵稀疏性的重要性
• 模型warm-up排第二（-0.0856%），说明两阶段温度退火策略有效
• 所有组件去除或替换都导致性能下降，验证了各模块设计的必要性

3.5 UniMixing-Lite超参数分析¶

Table 4: Basis数量、Rank和Block数量的影响

Basis Number影响：

Rank影响：

Block Number影响：

结论：

• Basis数量 $b$ 增加对AUC提升的参数效率高于增加rank $r$
• RankMixer在增加depth时出现性能退化（4-Blocks比2-Blocks差），而UniMixer-Lite持续提升
• Depth scaling比width scaling更高效

3.6 在线A/B测试（Q4）¶

• 部署场景：快手多个广告投放场景
• 指标：Cumulative Active Days (CAD)，30天窗口（排除安装日）
• 结果：CAD D1-D30平均提升超过15%

4. 总结与评价¶

核心贡献： 1. 揭示了TokenMixer的置换矩阵性质，建立了等价参数化表示 2. 通过UniMixing模块统一了attention、TokenMixer、FM三大Scaling范式 3. UniMixing-Lite通过basis-composed和低秩近似策略实现了最优参数/计算效率 4. 在快手大规模线上系统验证，CAD提升超15%

优势：

• 理论贡献扎实，从TokenMixer置换矩阵的数学性质出发，自然地建立了三类方法的联系
• 实验充分，包括离线对比、Scaling Law拟合、消融实验和在线A/B测试
• 工业落地价值高，在更少参数下取得更好效果

局限：

• 实验仅在快手工业数据集上进行，缺少公开数据集验证
• FLOPs/Batch较高（UniMixer-Lite-4-Blocks 84.5M为4.24T），计算效率优势不如参数效率优势明显
• 未讨论与序列推荐/生成式推荐的结合