UniMixer：面向推荐系统扩展定律的统一架构标题：UniMixer：面向推荐系统扩展定律的统一架构单位：快手英文

英文标题：UniMixer: A Unified Architecture for Scaling Laws in Recommendation Systems
时间：2026.04

摘要

近年来，推荐模型的扩展定律受到了越来越多的关注，它刻画了推荐模型性能与参数量／FLOPs 之间的关系。目前，实现推荐模型扩展主要有三类主流架构，即基于注意力的方法、基于 TokenMixer 的方法和基于因子分解机的方法。这些方法在设计理念和架构结构上均存在根本差异。本文提出了一种面向推荐系统的统一扩展架构 UniMixer，旨在提高扩展效率，并建立一个能够统一主流扩展模块的理论框架。通过将基于规则的 TokenMixer 转换为等价的参数化结构，我们构建了一个广义参数化特征混合模块，使词元混合模式能够在模型训练过程中得到优化和学习。同时，广义参数化词元混合消除了 TokenMixer 中头数必须等于词元数的限制。此外，我们还建立了推荐系统统一扩展模块的设计框架，打通了基于注意力、基于 TokenMixer 和基于因子分解机的方法之间的联系。为了进一步提高扩展投入产出比，我们设计了轻量级 UniMixing 模块 UniMixing-Lite，在显著提升模型性能的同时，进一步压缩模型参数量和计算成本。其扩展曲线如图 1 所示。大量离线和在线实验验证了 UniMixer 卓越的扩展能力。

图 1：UniMixer／UniMixer-Lite 与 RankMixer 架构中 AUC 和参数量／FLOPs 之间的扩展定律。横轴采用对数尺度。

1 引言

大语言模型揭示了一种令人瞩目的现象：随着模型规模、数据量和计算资源的增加，模型性能会持续提升，即呈现扩展定律。大语言模型展现出的显著性能扩展现象，促使推荐系统领域开始探索适合推荐任务的扩展框架。近年来，研究者尝试设计可扩展模块，并将其堆叠为多层结构，以提升排序模型的复杂度，从而在模型性能与模型规模或计算成本之间建立扩展定律，例如参数量和 FLOPs 与性能之间的关系。

推荐系统基于大量多特征域的用户和物品特征预测用户行为，向用户呈现最相关的内容，从而提升用户对推荐结果的正向参与度。这些多特征域特征通常包括类别特征和稠密特征，它们往往具有更加动态的嵌入表示，并从多个角度捕获信息。自然语言处理领域中的所有词元共享统一的嵌入空间，而推荐任务的特征空间天然具有异构性。因此，学习异构特征交互是推荐任务与自然语言处理领域之间的一项根本差异。鉴于 Transformer 在大语言模型中取得了巨大成功，一个自然的思路是对 Transformer 模块进行改造，使其适应推荐任务，因为通常无法直接将 Transformer 模块用作推荐系统扩展定律的基础模块。

为解决异构特征交互问题，当前推荐模型的主流扩展架构大体可以分为三类：基于注意力的方法、基于 TokenMixer 的方法和基于因子分解机的方法。基于注意力的方法，例如 HiFormer、FAT 和 HHFT，为每个输入词元构建词元专属的查询、键和值投影。不同于基于注意力的方法，基于 TokenMixer 的方法，例如 RankMixer 和 TokenMixer-Large，采用基于规则的词元混合操作实现异构特征交互，从而避免计算两个异构语义空间之间的内积相似度。另一方面，基于因子分解机的方法，例如 Wukong 和 Kunlun，通过引入因子分解机模块，在每一层中计算输入嵌入之间的交互。这些框架建立在完全不同的扩展模块之上，却都展现出了提升模型性能的扩展能力。这引出了一个根本问题：能否为推荐系统构建一个统一的扩展模块，将现有主流扩展组件的优势结合起来？

为了建立这些扩展模块之间的联系，我们首先找到了基于规则的 TokenMixer 操作的一种参数化表达。通过进一步优化计算流程，我们推导出计算成本更低的 UniMixing 模块。基于这一设计和相关结果，我们提出了一个统一理论框架，将推荐系统中的主流扩展模块统一起来。此外，我们还设计了一个轻量级 UniMixing 模块，它能够利用现有主流扩展模块的优势，并获得最佳的参数效率和计算效率。我们希望这一统一架构能够帮助推荐系统领域迎来属于自己的“注意力时刻”。本文的主要贡献概括如下：

我们通过对基于规则的 TokenMixer 进行等价参数化，揭示了 TokenMixer 的特征交互模式。
我们提出了一个名为 UniMixer 的统一扩展框架，它打通了基于注意力、基于 TokenMixer 和基于 FM 的方法之间的差异与联系。通过优化计算流程，UniMixer 在训练和推理阶段均显著降低了计算复杂度和 GPU 显存消耗。
为了进一步降低模型参数量和计算成本，我们设计了轻量级 UniMixing 模块 UniMixing-Lite。它能够同时利用基于注意力和基于 TokenMixer 的架构优势，从而获得更高的扩展效率。
大量离线和在线实验验证了 UniMixer 卓越的扩展能力。

2 相关工作

目前，为超大规模推荐系统建立扩展定律主要有三种建模范式：基于注意力、基于 TokenMixer 和基于 FM 的方法。

基于注意力的框架

近期的推荐系统研究已将 Transformer 引入点击率预测。这一范式的核心挑战，是弥合词元序列的异构特性与语言建模所假设的序列组合性之间的差距。为此，异构注意力层被提出以解决异构特征交互问题，HiFormer 则通过将异构词元展平为单一向量表示来显式建模高阶交互。此外，特征域感知 Transformer（FAT）通过因式分解的上下文对齐和跨特征域调制，将特征域感知的交互先验注入注意力机制，进一步建立了 CTR 预测中的经验扩展定律。HHFT 通过交替堆叠用于保留特征域专属语义的异构 Transformer 模块，以及用于学习高阶交互的 HiFormer 模块，进一步验证了这些扩展性质。在动态用户行为建模方面，HSTU-V1／V2、MARM、OneTrans、Climber、Hyformer 和 LLaTTE 等方法利用注意力机制捕获长程时间依赖。这些方法表明，统一特征交互与序列行为建模具有建立更稳健扩展定律的潜力。

基于 TokenMixer 的框架

注意力机制虽然能够提供富有表达力的特征交互，但由于注意力分数计算具有二次复杂度，其计算成本可能高得难以承受。受到计算机视觉领域 MLP-Mixer 成功的启发，工业推荐系统开始转向词元混合架构，并产生了 RankMixer、Lemur 和 TokenMixer-Large 等先进模型。例如，RankMixer 使用静态、非参数化的词元混合操作替代动态注意力，在保持严格可比 FLOPs 的同时取得了具有竞争力的 CTR 预测性能。在此基础上，TokenMixer-Large 通过引入辅助残差连接和定制损失函数，将该架构扩展至 130 亿参数配置，并在不同模型维度上展现出显著的扩展定律。然而，目前仍存在一个关键缺口：现有词元混合算子的设计高度依赖经验规则，缺少一座通向传统 FM 方法或注意力方法的严谨理论桥梁。

基于 FM 的框架

早期基于 FM 的方法使用低阶成对建模来处理推荐系统中的特征交互，之后，特征域感知因子分解机将其推广到特征域专属和上下文敏感的交互。虽然这些模型具有较高的可解释性和效率，但其低阶交互能力从根本上限制了模型容量。为了解决这一问题，DeepFM、AutoInt 和 DCN 系列等多种基于神经网络的扩展方法结合 MLP 或 Transformer 注意力来捕获高阶交互。近期，Wukong 通过堆叠 FM 风格的交互模块和线性压缩层，展现出了良好的扩展性质。然而，基于 FM 的方法依赖显式低阶交互，当模型沿参数量和 FLOPs 维度扩大时，这种依赖仍会限制性能提升，与大语言模型中观察到的可预测扩展定律形成对比。

3 预备知识

考虑评分预测、点击率（CTR）预测和点击后转化率（CVR）预测等判别式推荐任务，它们通常被表述为监督学习问题。数据集定义为：

\mathcal{D}=\left\{(\mathbf{X}_1,y_1),\ldots,(\mathbf{X}_i,y_i),\ldots,(\mathbf{X}_N,y_N)\right\},

其中， $\mathbf{X}_i=[\mathbf{x}_i^{(1)},\mathbf{x}_i^{(2)},\ldots,\mathbf{x}_i^{(F)}]$ 包含 $F$ 个特征域；对于二分类问题，第 $i$ 个样本的标签 $y_i\in\{0,1\}$ ，对于回归问题则有 $y_i\in\mathbb{R}$ ； $N$ 是数据点数量。通常，输入特征 $\mathbf{X}=\{\mathbf{X}^C,\mathbf{X}^D\}$ 被划分为类别特征 $\mathbf{X}^C$ 和稠密特征 $\mathbf{X}^D$ ，分别使用 $|C|$ 和 $|D|$ 表示类别特征和稠密特征的数量。对于 CTR 和 CVR 预测任务，核心目标是建立一个模型来预测点击或转化概率 $\Pr(y_i=1\mid\mathbf{X}_i)$ 。

在推荐系统中，学习到的嵌入表示更加动态。与语言模型的输入词元不同，推荐特征空间天然具有异构性。因此，将大语言模型所使用的 Transformer 架构直接迁移到推荐建模并不合适。到目前为止，推荐领域的扩展定律主要通过以下三类基础模块及其变体建立。

异构注意力层

基于异构注意力的架构通常使用特征域专属的查询、键和值投影来实现异构特征交互。给定输入隐状态 $\mathbf{X}=[\mathbf{x}_1;\ldots;\mathbf{x}_T]\in\mathbb{R}^{T\times D}$ ，异构注意力层表示为：

\mathbf{Q}_h= \begin{bmatrix} \mathbf{x}_1\mathbf{W}_Q^{1h}\\ \vdots\\ \mathbf{x}_T\mathbf{W}_Q^{Th} \end{bmatrix} \in\mathbb{R}^{T\times d},\quad \mathbf{K}_h= \begin{bmatrix} \mathbf{x}_1\mathbf{W}_K^{1h}\\ \vdots\\ \mathbf{x}_T\mathbf{W}_K^{Th} \end{bmatrix} \in\mathbb{R}^{T\times d},\quad \mathbf{V}_h= \begin{bmatrix} \mathbf{x}_1\mathbf{W}_V^{1h}\\ \vdots\\ \mathbf{x}_T\mathbf{W}_V^{Th} \end{bmatrix} \in\mathbb{R}^{T\times d}. \tag{1}

其中， $\mathbf{W}_Q^{ih},\mathbf{W}_K^{ih},\mathbf{W}_V^{ih}\in\mathbb{R}^{D\times d}$ 分别是查询、键和值投影的词元专属权重。多头异构注意力层的输出计算如下：

\mathbf{O}_h=\operatorname{softmax}\left(\frac{\mathbf{Q}_h\mathbf{K}_h^{\mathsf T}}{\sqrt d}\right)\mathbf{V}_h\in\mathbb{R}^{T\times d}. \tag{2}

随后，将多头异构注意力的输出拼接起来，并送入线性投影，使输出维度与输入 $\mathbf{X}$ 对齐。

TokenMixer

基于 TokenMixer 的框架使用无参数、基于规则的混合操作进行特征交互。对于给定输入 $\mathbf{X}=[\mathbf{x}_1;\ldots;\mathbf{x}_T]$ ，TokenMixer 首先将每个输入词元 $\mathbf{x}_t$ 均匀划分为 $H$ 个头：

\left[\mathbf{x}_t^{(1)}\mid\mathbf{x}_t^{(2)}\mid\cdots\mid\mathbf{x}_t^{(H)}\right]=\operatorname{SplitHead}(\mathbf{x}_t). \tag{3}

随后，第 $h$ 个词元 $\mathbf{s}^h$ 可表示为：

\mathbf{s}^h=\operatorname{concat}\left(\mathbf{x}_1^{(h)},\mathbf{x}_2^{(h)},\ldots,\mathbf{x}_T^{(h)}\right)\in\mathbb{R}^{\frac{TD}{H}}. \tag{4}

TokenMixer 的输出为：

\mathbf{S}= \begin{bmatrix} \mathbf{s}_1\\ \vdots\\ \mathbf{s}_H \end{bmatrix} \in\mathbb{R}^{H\times\frac{TD}{H}}. \tag{5}

其中，要求 $H$ 与 $T$ 相等。因此，输入 $\mathbf{X}$ 与输出 $\mathbf{S}$ 的维度完全相同。

Wukong

基于 Wukong 的模型将因子分解机模块（FMB）和线性投影层的输出进行拼接，以扩展交互组件：

\begin{aligned} \operatorname{FMB}(\mathbf{X})&=\operatorname{reshape}\!\left(\operatorname{MLP}\!\left(\operatorname{LN}\!\left(\operatorname{flatten}(\operatorname{FM}(\mathbf{X}))\right)\right)\right),\\ \operatorname{FM}(\mathbf{X})&=\mathbf{X}\mathbf{X}^{\mathsf T}\mathbf{Y},\\ \operatorname{LCB}(\mathbf{X})&=\mathbf{W}\mathbf{X}. \end{aligned} \tag{6}

其中， $\mathbf{W}\in\mathbb{R}^{n\times T}$ 和 $\mathbf{Y}\in\mathbb{R}^{T\times r}$ 是可学习的投影矩阵。 $\mathbf{Y}$ 用于降低存储交互矩阵 $\mathbf{X}\mathbf{X}^{\mathsf T}$ 所需的内存。

本文致力于为推荐系统建立一个统一的结构基础，整合当前扩展模块的优势，进一步提高扩展投入产出比。

4 UniMixer

4.1 概述

本文建立了一个用于扩展推荐系统的统一模块，即 UniMixer 模块。它在统一理论框架下，将推荐领域的主流扩展模块统一起来，包括基于注意力的模块、基于 TokenMixer 的模块以及基于 Wukong 的方法。如图 2 所示，整体架构由特征词元化、 $M$ 个带有 SiameseNorm 和稀疏逐词元 MoE 的 UniMixer 模块组成。通过对基于规则的 TokenMixer 进行参数化，我们建立了基于注意力、基于 TokenMixer 和基于 Wukong 的方法之间的联系，使所提出的 UniMixer 能够同时具备这些方法的优势。此外，我们开发了轻量级 UniMixing 模块，在显著提升模型性能的同时进一步压缩模型参数量和计算成本。

图 2：用于建立推荐系统扩展定律的 UniMixer 架构。

4.2 特征词元化

根据输入特征域的语义类别，首先将输入特征 $\mathbf{X}$ 划分为 $N$ 个互不相交的特征域：

\mathbf{X}=\left[ \underbrace{\mathbf{x}_U^{(1)},\ldots,\mathbf{x}_U^{(n_U)}}_{\text{用户画像}}, \underbrace{\mathbf{x}_I^{(1)},\ldots,\mathbf{x}_I^{(n_I)}}_{\text{物品特征}}, \underbrace{\mathbf{x}_B^{(1)},\ldots,\mathbf{x}_B^{(n_B)}}_{\text{行为序列}}, \underbrace{\mathbf{x}_Q^{(1)},\ldots,\mathbf{x}_Q^{(n_Q)}}_{\text{查询特征}}, \ldots \right]. \tag{7}

每个特征域通过嵌入层转换为不同维度的嵌入向量：

\mathbf{e}_n=\operatorname{Embedding}(\mathbf{X}_{\mathrm{domain}})\in\mathbb{R}^{d_{\mathrm{domain}}}. \tag{8}

其中， $\mathbf{X}_{\mathrm{domain}}$ 表示一个特征域内的全部特征， $d_{\mathrm{domain}}$ 是与该特征域对应的嵌入维度。所得特征域嵌入被拼接为一个嵌入向量 $\mathbf{E}=[\mathbf{e}_1,\mathbf{e}_2,\ldots,\mathbf{e}_N]$ 。随后，将嵌入向量 $\mathbf{E}$ 均匀划分为适当数量的块，再使用如下词元专属线性层将每个块投影为一个词元嵌入：

\mathbf{x}_i=\mathbf{W}_i^{\mathrm{proj}}\mathbf{E}_{di:di+d}+\mathbf{b}_i^{\mathrm{proj}}\in\mathbb{R}^{D}. \tag{9}

其中， $\mathbf{W}_i^{\mathrm{proj}}\in\mathbb{R}^{D\times d}$ ， $\mathbf{b}_i^{\mathrm{proj}}\in\mathbb{R}^{D}$ 。之后，将 $\mathbf{x}_i$ 按列堆叠，即可得到输入隐状态 $\mathbf{X}\in\mathbb{R}^{T\times D}$ 。

4.3 UniMixer 模块

异构特征交互

如第 3 节所述，异构注意力通过使用词元专属的查询、键和值权重，解决两个异构语义空间之间的特征交互问题。然而，通过计算内积相似度得到的注意力模式通常带有对角占优先验。在训练早期，查询和键权重 $\mathbf{W}_Q^h$ 、 $\mathbf{W}_K^h$ 随机初始化，注意力权重的大小很大程度上由输入词元值 $\mathbf{X}$ 主导。这很容易使注意力权重集中在少数词元上，如图 3（a）所示。

图 3：（a）不同方法的全局混合权重。（b）基于规则的 TokenMixer 的等价参数化。

从图 3（a）可以看到，异构注意力的注意力权重尖锐且稀疏，这会给梯度反向传播带来风险，使查询和键权重难以训练，并可能导致训练停滞，图 3（a）中异构注意力权重的第 10 行和第 15 行即体现了这一现象。同时，在大规模异构特征输入下，此类注意力模式可能导致特征交互趋于均匀，即注意力分数变得非常小且缺乏区分度，进而可能产生噪声信号，掩盖关键的特征交互模式。

另一方面，无参数且基于规则的 TokenMixer 操作缺少可学习性和场景适应性，可能造成异构特征交互不足或错误。此外， $T=H$ 的要求进一步限制了异构特征交互模式的选择。通过深入分析 TokenMixer 操作，我们得到了一些有趣的发现，使 TokenMixer 操作的参数化成为可能。如图 3（b）所示，我们发现 TokenMixer 操作可以视为置换矩阵 $\mathbf{W}^{\mathrm{perm}}$ 与展平后的输入嵌入 $\operatorname{flatten}(\mathbf{X})\in\mathbb{R}^{TD}$ 的乘积，可表示为：

\operatorname{TokenMixer}(\mathbf{X})=\operatorname{reshape}\left(\mathbf{W}^{\mathrm{perm}}\operatorname{flatten}(\mathbf{X})\right). \tag{10}

其中， $\mathbf{W}^{\mathrm{perm}}\in\mathbb{R}^{TD\times TD}$ 是一个大型置换矩阵。附录 A 给出了具体的数值示例。一个自然的想法是对置换矩阵 $\mathbf{W}^{\mathrm{perm}}$ 进行参数化，使基于规则的 TokenMixer 可学习、可优化。然而， $O(T^2D^2)$ 的计算复杂度和 $O(T^2D^2)$ 的参数量是不可接受的。通过观察，我们发现了 TokenMixer 置换矩阵 $\mathbf{W}^{\mathrm{perm}}$ 的一些有趣性质，并将其概括如下。

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

可压缩性。 置换矩阵 $\mathbf{W}^{\mathrm{perm}}$ 可以等价分解为两个较小矩阵的 Kronecker 积，即 $\mathbf{W}^{\mathrm{perm}}=\mathbf{G}\otimes\mathbf{I}$ ，其中 $\mathbf{I}\in\mathbb{R}^{\frac{D}{T}\times\frac{D}{T}}$ 是单位矩阵， $\mathbf{G}\in\mathbb{R}^{T^2\times T^2}$ ， $\otimes$ 表示 Kronecker 积运算。
双随机性。 在置换矩阵 $\mathbf{W}^{\mathrm{perm}}$ 中，每行与每列的元素之和均为 1，即 $\sum_{p=1}^{TD}w_{pq}^{\mathrm{perm}}=1$ ， $\sum_{q=1}^{TD}w_{pq}^{\mathrm{perm}}=1$ 。
稀疏性。 $\mathbf{W}^{\mathrm{perm}}$ 的每一行／列都恰好包含一个非零元素。
对称性。 若 $T=H$ ，置换矩阵 $\mathbf{W}^{\mathrm{perm}}$ 是对称矩阵，即 $\mathbf{W}^{\mathrm{perm}}=\mathbf{W}^{\mathrm{perm}\mathsf T}$ ；若 $T\ne H$ ，则 $\mathbf{W}^{\mathrm{perm}}$ 是非对称矩阵。

根据 TokenMixer 置换矩阵的这些性质，通过参数化矩阵 $\mathbf{G}$ 和 $\mathbf{I}$ ，词元混合的参数量可以显著减少至 $O(T^4+(D/T)^2)$ ，而通常 $T$ 远小于 $D$ 。此外，TokenMixer 的参数化仍面临三个挑战：第一，在模型训练和推理过程中，直接使用参数化的 $\mathbf{G}$ 和 $\mathbf{I}$ 重构 $\mathbf{W}^{\mathrm{perm}}$ ，仍会产生大小为 $[TD,TD]$ 的中间变量，对 GPU 显存提出极高要求；第二，如何保证学习到的参数满足双随机性、稀疏性和对称性；第三，如何设计一个统一的推荐扩展模块，整合现有扩展模块的优势，并为推荐系统建立更高的扩展效率。

统一词元混合模块

受图 3 启发，在统一词元混合模块中，我们不再使用 $T$ 和 $D$ ，而是在置换矩阵中定义块及其块大小。块大小记为 $B$ ，块的数量为 $(L/B)^2$ ，其中 $L$ 是输入嵌入维度，并且能够被块大小 $B$ 整除。将 $\mathbf{G}$ 的参数化权重记为 $\mathbf{W}_G$ 。考虑到置换矩阵的稀疏性，并为了实现充分的异构特征交互，我们为每一行分配不同的参数化权重 $\mathbf{W}_B^i$ 。经过这一操作，每个块都具有不同的特征交互模式。随后，通过学习参数矩阵 $\mathbf{W}_G$ 和 $\mathbf{W}_B^i$ ，可以获得交互模式更丰富的置换矩阵 $\mathbf{W}^{\mathrm{perm}}，其表达式为：

\operatorname{UniMixing}(\mathbf{X})= \operatorname{reshape}\left( \left(\mathbf{W}_G\otimes\left\{\mathbf{W}_B^i\right\}_{i=1}^{L/B}\right) \operatorname{flatten}(\mathbf{X}),1,L \right). \tag{11}

其中， $\otimes$ 表示广义 Kronecker 积。

接下来，我们对式（11）的计算流程进行优化，以显著降低计算成本和 GPU 显存需求。首先，将嵌入向量 $\operatorname{flatten}(\mathbf{X})$ 均匀划分为 $L/B$ 个向量，每个向量的大小为 $B$ ，表示为：

\left[\mathbf{x}_1\mid\mathbf{x}_2\mid\cdots\mid\mathbf{x}_{L/B}\right] =\operatorname{Split}\left(\operatorname{flatten}(\mathbf{X}),\frac{L}{B}\right). \tag{12}

随后，分别将块权重 $\mathbf{W}_B^i$ 与对应的分块向量 $\mathbf{x}^{(i)}$ 相乘，得到如下局部特征交互向量：

\begin{aligned} \operatorname{reshape}\left(\mathbf{H},\frac{L}{B},B\right) &=\operatorname{reshape}\left( \left[\mathbf{x}_1\mathbf{W}_B^1\mid\mathbf{x}_2\mathbf{W}_B^2\mid\cdots\mid \mathbf{x}_{L/B}\mathbf{W}_B^{L/B}\right],\frac{L}{B},B\right)\\ &=\begin{bmatrix} \mathbf{x}_1\mathbf{W}_B^1\\ \vdots\\ \mathbf{x}_{L/B}\mathbf{W}_B^{L/B} \end{bmatrix}. \end{aligned} \tag{13}

因此，UniMixing 模块的输出为：

\operatorname{UniMixing}(\mathbf{X})= \operatorname{reshape}\left( \mathbf{W}_G\operatorname{reshape}\left(\mathbf{H},\frac{L}{B},B\right),1,L \right). \tag{14}

与直接使用重构矩阵 $\mathbf{W}^{\mathrm{perm}}$ 相比，这种经过优化的计算流程将计算成本从 $O(L^2)$ 降低到 $O(L^2/B+LB)$ ，并避免在计算过程中产生大型中间变量。附录 B 给出了式（11）计算流程优化的证明。根据式（13）和式（14）， $\mathbf{W}_B^i$ 控制块内交互模式， $\mathbf{W}_G$ 控制块间交互模式。对于维度为 $L$ 的嵌入输入，不再要求 $T=H$ 。与 TokenMixer 操作相比，UniMixing 模块具有更多样的局部和全局特征混合模式以及交互尺度，同时保留了可学习、可优化的优点。

为了使学习到的置换矩阵满足双随机性，我们使用 Sinkhorn-Knopp 迭代：先通过指数算子使 $\mathbf{W}_G$ 和 $\mathbf{W}_B^i$ 的所有元素均为正数，然后进行迭代归一化，交替缩放矩阵的行和列，使其元素之和为 1。此外，我们引入温度系数来控制参数矩阵的稀疏性。最后，分别使用 $(\mathbf{W}_G+\mathbf{W}_G^{\mathsf T})/2$ 和 $(\mathbf{W}_B^i+\mathbf{W}_B^{i\mathsf T})/2$ 满足参数矩阵的对称性约束。最终的约束权重为：

\begin{aligned} \widetilde{\mathbf{W}}_G&=\frac{\mathbf{W}_G+\mathbf{W}_G^{\mathsf T}}{2},& \widetilde{\mathbf{W}}_B^i&=\frac{\mathbf{W}_B^i+\mathbf{W}_B^{i\mathsf T}}{2},\\ \overline{\mathbf{W}}_G&=\operatorname{Sinkhorn\text{-}Knopp}\left(\frac{\widetilde{\mathbf{W}}_G}{\tau}\right),& \overline{\mathbf{W}}_B^i&=\operatorname{Sinkhorn\text{-}Knopp}\left(\frac{\widetilde{\mathbf{W}}_B^i}{\tau}\right). \end{aligned} \tag{15}

其中， $\tau$ 是温度系数。

随后，使用残差连接和归一化模块处理 UniMixing 模块的输出：

\mathbf{O}=\operatorname{RMSNorm}\left(\mathbf{X}+\operatorname{UniMixing}(\mathbf{X})\right). \tag{16}

异构特征交互的统一视角

观察式（1）中的 $\mathbf{V}_h$ 和式（13）中的 $\operatorname{reshape}(\mathbf{H},L/B,B)$ ，我们发现，如果将块数 $L/B$ 设置为 $T$ ，并令 $\mathbf{W}_V^{ih}$ 与 $\mathbf{W}_B^i$ 具有相同维度，那么 $\operatorname{reshape}(\mathbf{H},L/B,B)=\mathbf{V}_h$ 。这意味着，当 $\mathbf{W}_V^i=\mathbf{W}_B^i$ 时，UniMixer 的局部交互投影等价于异构注意力层的值投影。另一方面， $\mathbf{W}_G$ 的维度和作用与注意力权重相同，区别在于 $\mathbf{W}_G$ 需要满足双随机性、稀疏性和对称性。

表 1：统一理论框架下，基于注意力、基于 TokenMixer 和基于 FM 的方法之间的差异。

Wukong 的特征交互建立在 FM 组件之上。根据式（6）， $\operatorname{FMB}(\mathbf{X})$ 可以改写为 $\operatorname{reshape}(\operatorname{MLP}(\operatorname{LN}(\operatorname{flatten}(\mathbf{X}\mathbf{I}(\mathbf{X}\mathbf{I})^{\mathsf T}\mathbf{Y}))))$ ，其中 $\mathbf{I}$ 是维度适当的单位矩阵。现在关注其核心特征交互模块 $\mathbf{X}\mathbf{I}(\mathbf{X}\mathbf{I})^{\mathsf T}\mathbf{Y}$ 。在注意力模块中，当 $\mathbf{W}_Q=\mathbf{I}$ 、 $\mathbf{W}_K=\mathbf{I}$ ，并且值矩阵不依赖于隐状态输入 $\mathbf{X}$ ，即 $\mathbf{V}_h=\mathbf{W}_V=\mathbf{Y}$ 时，注意力机制就退化为 FM 模块。因此，基于注意力、基于 TokenMixer 和基于 Wukong 的架构可以统一到如下单一理论框架中：

\operatorname{UniMixing}(\mathbf{X})= \operatorname{reshape}\left( \underbrace{\mathbf{G}(\mathbf{X},\mathbf{W}_G)}_{\text{全局混合模式}} \underbrace{ \begin{bmatrix} \mathbf{x}_1\mathbf{W}_B^1\\ \vdots\\ \mathbf{x}_{L/B}\mathbf{W}_B^{L/B} \end{bmatrix}}_{\text{局部混合模式}},1,L \right). \tag{17}

其中， $\mathbf{G}(\mathbf{X},\mathbf{W}_G)$ 是异构特征交互投影，用于衡量词元到词元／块到块的交互强度。为了便于分析不同方法之间的差异和联系，我们考虑单头注意力设置。在统一理论框架式（17）下，它们之间的差异总结于表 1。对于自注意力、异构注意力和 FM，全局混合模式 $\mathbf{G}(\mathbf{X},\mathbf{W}_G)$ 通过计算两个词元之间的内积相似度得到。TokenMixer 的全局混合模式则与输入词元嵌入无关。

UniMixing-Lite

如图 3 所示，随着分块粒度变细，局部交互参数矩阵 $\mathbf{W}_B^i$ 的数量增加，全局交互参数矩阵 $\mathbf{W}_G$ 也会变大。这会造成局部交互模式冗余。同时，更大的全局交互矩阵也无法高效减少参数量。因此，我们基于 UniMixing 模块设计了一个轻量级 UniMixing 模块 UniMixing-Lite，以进一步降低模块的参数量和计算成本，从而提高模型的扩展效率。

为解决局部交互模式的冗余问题，我们引入一个由基矩阵组合而成的模块，动态生成每个块专属的局部混合权重。将 $\mathbf{W}_B^i$ 的一组基矩阵定义为 $\{\mathbf{Z}_\ell\}_{\ell=1}^b$ ，将这些基矩阵对应的块专属权重向量定义为 $\{\boldsymbol{\omega}^i\}_{i=1}^{L/B}$ ，其中 $b$ 是局部混合基权重的数量， $\boldsymbol{\omega}^i=[\omega_1^i,\ldots,\omega_b^i]$ 。此外，对于全局交互参数 $\mathbf{W}_G$ ，我们使用低秩近似进一步提高效率。于是，UniMixing-Lite 模块可以表示为：

\begin{aligned} \operatorname{UniMixing\text{-}Lite}(\mathbf{X}) &=\operatorname{reshape}\left( \mathbf{W}_r\operatorname{reshape}\left( \left[\mathbf{x}_1\mathbf{W}_B^{*1}\mid\cdots\mid \mathbf{x}_{L/B}\mathbf{W}_B^{*L/B}\right],\frac{L}{B},B \right),1,L\right),\\ \mathbf{O}&=\operatorname{RMSNorm}\left(\mathbf{X}+\operatorname{UniMixing\text{-}Lite}(\mathbf{X})\right). \end{aligned} \tag{18}

其中， $\mathbf{W}_r=\operatorname{Sinkhorn\text{-}Knopp}(\mathbf{A}_G\mathbf{B}_G)$ ， $\mathbf{W}_B^{*i}=\operatorname{Sinkhorn\text{-}Knopp}(\sum_{\ell=1}^b\omega_\ell^i\mathbf{Z}_\ell)$ ， $\mathbf{A}_G\in\mathbb{R}^{(L//B)\times r}$ ， $\mathbf{B}_G\in\mathbb{R}^{r\times(L//B)}$ 。 $r$ 是 $\mathbf{W}_G$ 低秩近似的秩。在 UniMixing-Lite 模块中，我们既保留了 TokenMixer 低参数量的全局交互模式，也保留了注意力针对异构特征的局部交互能力。因此，它能够同时利用基于注意力和基于词元混合的方法的优势。

逐词元 SwiGLU

在 UniMixing 模块之后，我们引入逐词元 SwiGLU，以建模不同词元之间的特征异构性。对于每个输入词元 $\mathbf{x}_i$ ，SwiGLU 表示为：

\operatorname{pSwiGLU}(\mathbf{o}_i)= \mathbf{W}_{\mathrm{down}}^i\left( (\mathbf{W}_{\mathrm{up}}^i\mathbf{o}_i+\mathbf{b}_{\mathrm{up}}^i) \odot\operatorname{Swish}(\mathbf{W}_{\mathrm{gate}}^i\mathbf{o}_i+\mathbf{b}_{\mathrm{gate}}^i) \right)+\mathbf{b}_{\mathrm{down}}^i. \tag{19}

其中， $\mathbf{W}_{\mathrm{up}}^i,\mathbf{W}_{\mathrm{gate}}^i\in\mathbb{R}^{B\times nB}$ ， $\mathbf{W}_{\mathrm{down}}^i\in\mathbb{R}^{nB\times B}$ ， $\mathbf{b}_{\mathrm{up}}^i,\mathbf{b}_{\mathrm{gate}}^i\in\mathbb{R}^{nB}$ ， $\mathbf{b}_{\mathrm{down}}^i\in\mathbb{R}^{B}$ ； $\mathbf{o}_i$ 是第 $i$ 个词元的 UniMixing 输出， $n$ 是超参数。

4.4 SiameseNorm

当前的 RankMixer 架构缺少针对深层架构的专门设计，这通常表现为沿模型深度扩展的效果有限。TokenMixer-Large 虽然尝试通过在 TokenMixer-Large 模块内部引入间隔残差和辅助损失来解决这一问题，但并未触及问题的根源。为了在模型深度增加时获得训练稳定性和性能增益，我们将 SiameseNorm 引入 UniMixer 架构，如图 2 所示。SiameseNorm 通过为每一层引入两个耦合流，解决了 Pre-Norm 与 Post-Norm 之间的矛盾。

在本节中，这两个耦合流记为 $\bar{\mathbf{X}}_i$ 和 $\bar{\mathbf{Y}}_i$ ，并由输入嵌入初始化，即 $\bar{\mathbf{X}}_0=\bar{\mathbf{Y}}_0=\mathbf{X}$ 。对于第 $\ell$ 个模块，SiameseNorm 执行如下更新：

\begin{aligned} \widetilde{\mathbf{Y}}_\ell&=\operatorname{RMSNorm}(\bar{\mathbf{Y}}_\ell),& \mathbf{O}_\ell&=\operatorname{UniMixer}(\bar{\mathbf{X}}_\ell+\widetilde{\mathbf{Y}}_\ell),\\ \bar{\mathbf{X}}_{\ell+1}&=\operatorname{RMSNorm}(\bar{\mathbf{X}}_\ell+\mathbf{O}_\ell),& \bar{\mathbf{Y}}_{\ell+1}&=\bar{\mathbf{Y}}_\ell+\mathbf{O}_\ell. \end{aligned}

对于第 $M$ 个 UniMixer 模块，将 $\bar{\mathbf{X}}_M$ 和 $\bar{\mathbf{Y}}_M$ 融合以生成最终表示：

\mathbf{X}_{\mathrm{output}}=\bar{\mathbf{X}}_M+\operatorname{RMSNorm}(\bar{\mathbf{Y}}_M). \tag{20}

4.5 UniMixer 训练策略

为了使参数矩阵 $\mathbf{W}_G$ 和 $\mathbf{W}_B^i$ 具有稀疏性，我们引入温度系数来控制其稀疏程度。然而，较小的温度会产生更稀疏的权重，同时也会使梯度变得稀疏、微弱，甚至不稳定。这可能增加训练难度，并使优化陷入局部最优。另一方面，我们的实验表明，权重参数的稀疏性会显著改善模型性能，如表 3 所示。因此，这种稀疏性不可或缺。

一种常用方法是在训练过程中采用线性温度退火：从相对较高的初始温度开始，例如 $\tau=1.0$ ，再随着训练迭代次数增加，将其逐步线性退火至 $0.05$ ，表示为：

\tau_j=\max\left\{\tau_{\mathrm{start}}-\frac{(\tau_{\mathrm{start}}-\tau_{\mathrm{end}})j}{J},\tau_{\mathrm{end}}\right\}. \tag{21}

其中， $\tau_j$ 是第 $j$ 个温度系数， $\tau_{\mathrm{start}}$ 和 $\tau_{\mathrm{end}}$ 分别是初始温度和最终温度， $J$ 是温度退火的迭代区间。当数据量不足时，线性退火可能导致模型在训练早期因温度系数较高而探索不足，或者在后期因温度系数较低而未达到最优。为解决这一问题，可以先使用较高的温度系数，例如 $\tau=1.0$ ，对模型训练进行冷启动；待模型得到充分训练后，再降低温度系数，例如 $\tau=0.05$ ，并使用高温模型的权重作为初始化，重新训练低温模型。

5 实验

本节开展了大量实验，将 UniMixer 架构与现有最先进方法进行性能比较，并回答以下问题：

Q1： UniMixer 架构的扩展效率是否优于最先进架构？
Q2： 在不同的全局和局部混合模式设置下，所提方法的性能如何变化？
Q3： 轻量级模块 UniMixing-Lite 能否进一步提高扩展效率？
Q4： 部署到真实在线系统后，UniMixer／UniMixing-Lite 能否在 A/B 测试中改善业务指标？

5.1 实验设置

数据集与评估指标

我们使用快手广告投放场景中的真实训练日志数据来建模用户留存，并开展离线和在线评估。该数据集包含一年内收集的超过 7 亿条用户样本，涵盖数百个异构特征，例如数值特征、ID 特征、交叉特征和序列特征。二元标签“用户留存 = 1／0”表示用户是否会在首次激活后的第二天返回快手应用。对于推荐模型的扩展评估指标，我们采用推荐系统中常用的两个指标，即 ROC 曲线下面积（AUC）和用户级 AUC（UAUC）来评估模型性能，并使用稠密参数量、FLOPs 和 MFU 评估模型效率。

基线与实验细节

我们将包含 2 个／4 个模块的 UniMixer／UniMixing-Lite 架构与以下具有代表性的最先进框架进行比较，并按建模范式对其分类：

基于注意力的架构： 异构注意力、HiFormer 和 FAT。它们使用特征域专属的查询、键和值投影实现异构特征交互。
基于 TokenMixer 的框架： RankMixer 和 TokenMixer-Large。它们使用基于规则的词元混合操作进行特征交互。
基于 FM 的框架： Wukong。它通过拼接 FMB 与线性投影层的输出扩展交互组件。

所有实验均在一个由 40 块 GPU 组成的混合分布式训练框架中进行。所有模型采用一致的优化器超参数：稠密部分和稀疏部分均使用 Adam 优化器，学习率设为 $0.001$ 。

5.2 性能比较（Q1）

我们使用参数量约为 1 亿的最先进扩展架构与 UniMixer 和 UniMixer-Lite 进行比较，以探索它们的扩展定律，并将异构注意力架构用作基础模型。表 2 给出了本文模型和最先进模型的主要性能结果。可以看到，在参数预算和计算成本更低的情况下，UniMixer 和 UniMixer-Lite 架构在多项指标上均显著优于其他最先进模型。

表 2：广告投放场景中，约 1 亿参数的 UniMixer 与最先进模型的性能和效率。

随后，在该广告投放场景中，除 UniMixer／UniMixer-Lite 外，RankMixer 的性能优于所有其他最先进模型。因此，我们选择这一最强的最先进模型，与 UniMixer／UniMixer-Lite 一同进行扩展定律比较。所有模型均使用相同数据集和一致的超参数训练。图 4 给出了它们关于参数量和 FLOPs 的扩展曲线。可以看到，随着参数量／FLOPs 增加，三个模型的 AUC 均呈现明显的幂律趋势。UniMixer-Lite 具有最佳扩展效率，性能提升曲线的斜率也更陡。

根据图 4 所示参数量与 AUC 之间的关系，RankMixer、UniMixer 和 UniMixer-Lite 在 AUC 与参数量／FLOPs 之间呈现出良好的扩展定律，可表述为：

\begin{aligned} \Delta\mathrm{AUC}_{\mathrm{RankMixer}}&=0.002718\,\mathrm{Params}^{0.116043},& \Delta\mathrm{AUC}_{\mathrm{RankMixer}}&=0.002022\,\mathrm{FLOPs}^{0.116635},\\ \Delta\mathrm{AUC}_{\mathrm{UniMixer}}&=0.003032\,\mathrm{Params}^{0.131973},& \Delta\mathrm{AUC}_{\mathrm{UniMixer}}&=0.002058\,\mathrm{FLOPs}^{0.125702},\\ \Delta\mathrm{AUC}_{\mathrm{UniMixer\text{-}Lite}}&=0.003767\,\mathrm{Params}^{0.141903},& \Delta\mathrm{AUC}_{\mathrm{UniMixer\text{-}Lite}}&=0.002338\,\mathrm{FLOPs}^{0.135327}. \end{aligned}

扩展定律中的两个常数里，扩展指数常数对性能增长的影响最为显著，是决定扩展效率的主导因素。UniMixer-Lite 展现出最强的扩展效率：无论关于参数量还是 FLOPs，它都获得了最大的扩展指数和系数。这表明它能够从模型容量增加中获得最多收益。

图 4：包含 2 个模块的 UniMixer／UniMixer-Lite 与 RankMixer 架构中，AUC 和参数量／FLOPs 之间的扩展定律。横轴采用对数尺度。

5.3 消融实验（Q2）

为了探索全局与局部混合权重的性质，以及 UniMixer 各模块对 AUC 增益的贡献，我们对多个 UniMixer 变体进行了消融实验，并测量它们相对于完整 UniMixer 模型的 AUC 变化。所有变体都在相似设置下训练。结果如表 3 所示：移除任何模块或违反任何参数约束都会导致性能下降，其中低温度系数和模型预热对整体性能的影响最为显著。

表 3：对参数量为 657 万的 UniMixer 各组件进行的消融实验。

5.4 UniMixing-Lite 模块的性能（Q3）

根据图 4 中的扩展趋势，可以看到本文提出的 UniMixing-Lite 架构具有最佳的参数效率和计算效率。这里，我们通过实验研究基矩阵集合 $\{\mathbf{Z}_\ell\}_{\ell=1}^b$ 的不同基数量 $b$ 、 $\mathbf{A}_G$ 和 $\mathbf{B}_G$ 的不同秩 $r$ ，以及不同 UniMixer 模块数量所产生的影响。如表 4 所示，随着基数量 $b$ 以及 $\mathbf{A}_G$ 、 $\mathbf{B}_G$ 的秩 $r$ 增加，模型性能也相应提高。然而，从参数效率来看，增加基数量 $b$ 带来的 AUC 增益高于增加秩 $r$ 。

为了观察经过 Sinkhorn-Knopp 操作的低秩近似 $\mathbf{A}_G\mathbf{B}_G$ 和基矩阵 $\{\mathbf{Z}_\ell\}_{\ell=1}^b$ 对全局与局部混合矩阵重构的影响，我们在一个包含 2 个模块的 UniMixer-Lite 架构中，对第一个 UniMixer 模块重构出的全局矩阵 $\overline{\mathbf{W}}_G$ 及前六个局部混合矩阵 $\overline{\mathbf{W}}_B^i$ 进行了可视化，并分别采用温度系数 $\tau=1$ 和 $\tau=0.05$ ，如图 5 所示。输入嵌入维度为 768，块大小为 6，因此 $\overline{\mathbf{W}}_G\in\mathbb{R}^{128\times128}$ ， $\overline{\mathbf{W}}_B^i\in\mathbb{R}^{6\times6}$ ，其中 $\mathbf{A}_G\in\mathbb{R}^{128\times16}$ ， $\mathbf{B}_G\in\mathbb{R}^{16\times128}$ 。

根据图 5，尽管模块使用了低秩近似和基矩阵，Sinkhorn-Knopp 操作仍能确保矩阵保持接近满秩。此外，对比图 5（a）（b）与图 5（c）（d）可以看到，采用较低温度系数的全局和局部混合矩阵，其交互分布比采用较高温度系数时更加尖锐。结合消融实验的结果可以得出结论： $\overline{\mathbf{W}}_G$ 和 $\overline{\mathbf{W}}_B^i$ 的稀疏性能够显著提升模型性能。

图 5：不同温度系数下 UniMixer-Lite 的 $\overline{\mathbf{W}}_G$ 和 $\overline{\mathbf{W}}_B^i$ 。（a） $\tau=1$ 时的 $\overline{\mathbf{W}}_G$ ；（b） $\tau=1$ 时的 $\{\overline{\mathbf{W}}_B^i\}_{i=1}^6$ ；（c） $\tau=0.05$ 时的 $\overline{\mathbf{W}}_G$ ；（d） $\tau=0.05$ 时的 $\{\overline{\mathbf{W}}_B^i\}_{i=1}^6$ 。

另一方面，根据表 4 可以看到，随着 UniMixer 深度增加，所开发的模型持续呈现出明显的规模扩展趋势；而当 RankMixer 模块不断堆叠时，RankMixer 的性能反而下降。包含 2 个和 4 个模块的 UniMixing-Lite 的扩展曲线如图 6 所示，这表明沿深度扩展比沿宽度扩展更加高效。

表 4：UniMixing-Lite 中basis数量、rank和 UniMixer block数量的影响。

图 6：RankMixer 以及包含 2 个和 4 个模块的 UniMixer／UniMixer-Lite 中，AUC 和参数量／FLOPs 之间的扩展曲线。

5.5 在线 A/B 测试结果（Q4）

为了验证所提出的 UniMixer 架构的在线性能，我们已将 UniMixer 和 UniMixer-Lite 部署到快手的多个广告投放场景中。在在线 A/B 测试中，我们在 30 天观察窗口内使用累计活跃天数（CAD）衡量用户参与度，并排除安装当天，即第 0 天。在多个场景中，D1–D30 的 CAD 平均提升超过 15%。

6 结论

本文为推荐系统扩展定律建立了一个统一扩展框架，它打通了基于注意力、基于 TokenMixer 和基于 FM 的方法之间的联系，使整合各类方法的优势成为可能。从所得扩展定律来看，与最先进架构相比，UniMixer-Lite 获得了最佳的参数效率和计算效率。我们已将该架构部署到快手的多个场景中，并取得了显著的离线和在线收益。

本文不再孤立地看待推荐系统中现有的扩展模块，例如异构注意力、TokenMixer 和 Wukong，而是建立了一个统一理论框架，为推荐系统的扩展设计提供指导。我们相信，这一统一架构能够帮助推荐系统领域迎来属于自己的“注意力时刻”。统一模块 UniMixer 可以作为专为推荐领域定制的基础模块，其适用范围还可以进一步扩展至用户行为序列建模和生成式推荐任务。

附录 A TokenMixer 等价变换的数值示例

给定如下输入隐状态 $\mathbf{X}\in\mathbb{R}^{2\times6}$ ：

\mathbf{X}=\begin{bmatrix} x_1&x_2&x_3&x_4&x_5&x_6\\ x_7&x_8&x_9&x_{10}&x_{11}&x_{12} \end{bmatrix},

其中 $x_i$ 是标量。该输入隐状态 $\mathbf{X}$ 经过 TokenMixer 操作后被转换为：

\operatorname{TokenMixer}(\mathbf{X})= \begin{bmatrix} x_1&x_2&x_3&x_7&x_8&x_9\\ x_4&x_5&x_6&x_{10}&x_{11}&x_{12} \end{bmatrix}.

TokenMixer 的输出可以展平为向量：

\operatorname{flatten}(\operatorname{TokenMixer}(\mathbf{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}]^{\mathsf T}. \tag{22}

另一方面，向量 $\operatorname{flatten}(\mathbf{X})$ 可以通过乘以一个 $12\times12$ 矩阵转换为 $\operatorname{flatten}(\operatorname{TokenMixer}(\mathbf{X}))$ ，表示为：

\underbrace{ \begin{bmatrix} 1&0&0&0&0&0&0&0&0&0&0&0\\ 0&1&0&0&0&0&0&0&0&0&0&0\\ 0&0&1&0&0&0&0&0&0&0&0&0\\ 0&0&0&0&0&0&1&0&0&0&0&0\\ 0&0&0&0&0&0&0&1&0&0&0&0\\ 0&0&0&0&0&0&0&0&1&0&0&0\\ 0&0&0&1&0&0&0&0&0&0&0&0\\ 0&0&0&0&1&0&0&0&0&0&0&0\\ 0&0&0&0&0&1&0&0&0&0&0&0\\ 0&0&0&0&0&0&0&0&0&1&0&0\\ 0&0&0&0&0&0&0&0&0&0&1&0\\ 0&0&0&0&0&0&0&0&0&0&0&1 \end{bmatrix}}_{\mathbf{W}^{\mathrm{perm}}} \underbrace{ \begin{bmatrix} x_1\\x_2\\x_3\\x_4\\x_5\\x_6\\x_7\\x_8\\x_9\\x_{10}\\x_{11}\\x_{12} \end{bmatrix}}_{\operatorname{flatten}(\mathbf{X})} = \underbrace{ \begin{bmatrix} x_1\\x_2\\x_3\\x_7\\x_8\\x_9\\x_4\\x_5\\x_6\\x_{10}\\x_{11}\\x_{12} \end{bmatrix}}_{\operatorname{flatten}(\operatorname{TokenMixer}(\mathbf{X}))}. \tag{23}

根据式（22）和式（23），该数值示例中的 TokenMixer 操作可以等价转换为矩阵乘法形式。此外，置换矩阵 $\mathbf{W}^{\mathrm{perm}}\in\mathbb{R}^{12\times12}$ 可以等价分解为以下两个小矩阵的 Kronecker 积：

\mathbf{W}^{\mathrm{perm}}= \underbrace{ \begin{bmatrix} 1&0&0&0\\ 0&0&1&0\\ 0&1&0&0\\ 0&0&0&1 \end{bmatrix}}_{\text{全局混合矩阵}} \otimes \underbrace{ \begin{bmatrix} 1&0&0\\ 0&1&0\\ 0&0&1 \end{bmatrix}}_{\text{局部混合矩阵}}.

附录 B UniMixing 模块的计算流程优化

将 $\mathbf{W}_G\in\mathbb{R}^{(L//B)\times(L//B)}$ 和 $\mathbf{W}_B^i$ 定义如下：

\mathbf{W}_G= \begin{bmatrix} w_{(1,1)}^G&\cdots&w_{(1,L//B)}^G\\ \vdots&\ddots&\vdots\\ w_{(L//B,1)}^G&\cdots&w_{(L//B,L//B)}^G \end{bmatrix},\qquad \mathbf{W}_B^i= \begin{bmatrix} v_{(1,1)}^i&\cdots&v_{(1,B)}^i\\ \vdots&\ddots&\vdots\\ v_{(B,1)}^i&\cdots&v_{(B,B)}^i \end{bmatrix}. \tag{24}

其中， $w_{ij}$ 和 $v_{ij}$ 均为标量。根据式（12）， $\operatorname{flatten}(\mathbf{X})$ 被均匀划分为 $L//B$ 个向量，可以改写为：

\operatorname{flatten}(\mathbf{X})= \left[\mathbf{x}_1\mid\mathbf{x}_2\mid\cdots\mid\mathbf{x}_{L/B}\right]^{\mathsf T}, \tag{25}

其中， $\mathbf{x}_i$ 是维度为 $B$ 的行向量。

根据 UniMixing 的原始表达式（11），项 $\left(\mathbf{W}_G\otimes\{\mathbf{W}_B^i\}_{i=1}^{L//B}\right)\operatorname{flatten}(\mathbf{X})$ 可以改写为：

\begin{aligned} &\left(\mathbf{W}_G\otimes\{\mathbf{W}_B^i\}_{i=1}^{L//B}\right)\operatorname{flatten}(\mathbf{X})\\ &=\begin{bmatrix} w_{(1,1)}^G\mathbf{W}_B^1&\cdots&w_{(1,L/B)}^G\mathbf{W}_B^{L/B}\\ \vdots&\ddots&\vdots\\ w_{(L/B,1)}^G\mathbf{W}_B^1&\cdots&w_{(L/B,L/B)}^G\mathbf{W}_B^{L/B} \end{bmatrix} \begin{bmatrix} \mathbf{x}_1^{\mathsf T}\\ \vdots\\ \mathbf{x}_{L/B}^{\mathsf T} \end{bmatrix}\\ &=\begin{bmatrix} w_{(1,1)}^G\mathbf{W}_B^1\mathbf{x}_1^{\mathsf T}+\cdots+w_{(1,L/B)}^G\mathbf{W}_B^{L/B}\mathbf{x}_{L/B}^{\mathsf T}\\ \vdots\\ w_{(L/B,1)}^G\mathbf{W}_B^1\mathbf{x}_1^{\mathsf T}+\cdots+w_{(L/B,L/B)}^G\mathbf{W}_B^{L/B}\mathbf{x}_{L/B}^{\mathsf T} \end{bmatrix} \in\mathbb{R}^{L\times1}. \end{aligned} \tag{26}

另一方面，可以得到如下表达式：

\begin{aligned} &\mathbf{W}_G\operatorname{reshape}\left( \left[\mathbf{x}_1\mathbf{W}_B^{1\mathsf T}\mid\cdots\mid \mathbf{x}_{L/B}\mathbf{W}_B^{(L/B)\mathsf T}\right],\frac{L}{B},B\right)\\ &=\begin{bmatrix} w_{(1,1)}^G&\cdots&w_{(1,L/B)}^G\\ \vdots&\ddots&\vdots\\ w_{(L/B,1)}^G&\cdots&w_{(L/B,L/B)}^G \end{bmatrix} \begin{bmatrix} \mathbf{x}_1\mathbf{W}_B^{1\mathsf T}\\ \vdots\\ \mathbf{x}_{L/B}\mathbf{W}_B^{(L/B)\mathsf T} \end{bmatrix}\\ &=\begin{bmatrix} w_{(1,1)}^G\mathbf{x}_1\mathbf{W}_B^{1\mathsf T}+\cdots+w_{(1,L/B)}^G\mathbf{x}_{L/B}\mathbf{W}_B^{(L/B)\mathsf T}\\ \vdots\\ w_{(L/B,1)}^G\mathbf{x}_1\mathbf{W}_B^{1\mathsf T}+\cdots+w_{(L/B,L/B)}^G\mathbf{x}_{L/B}\mathbf{W}_B^{(L/B)\mathsf T} \end{bmatrix} \in\mathbb{R}^{\frac{L}{B}\times B}. \end{aligned} \tag{27}

式（26）中的元素与式（27）中的元素满足：

w_{(i,1)}^G\mathbf{W}_B^1\mathbf{x}_1^{\mathsf T}+\cdots+w_{(i,L/B)}^G\mathbf{W}_B^{L/B}\mathbf{x}_{L/B}^{\mathsf T} =\left( w_{(i,1)}^G\mathbf{x}_1\mathbf{W}_B^{1\mathsf T}+\cdots+w_{(i,L/B)}^G\mathbf{x}_{L/B}\mathbf{W}_B^{(L/B)\mathsf T} \right)^{\mathsf T}. \tag{28}

由此可得：

\begin{aligned} &\left(\mathbf{W}_G\otimes\{\mathbf{W}_B^i\}_{i=1}^{L/B}\right)\operatorname{flatten}(\mathbf{X})\\ &=\operatorname{reshape}\left( \mathbf{W}_G\operatorname{reshape}\left( \left[\mathbf{x}_1\mathbf{W}_B^{1\mathsf T}\mid\cdots\mid \mathbf{x}_{L/B}\mathbf{W}_B^{(L/B)\mathsf T}\right],\frac{L}{B},B \right),L,1\right). \end{aligned}

由于 $\mathbf{W}_B^i$ 和 $\mathbf{W}_B^{i\mathsf T}$ 都是可学习参数，参数转置不会影响模型。因此，经过计算流程优化后的 UniMixing 模块可以表示为：

\operatorname{UniMixing}(\mathbf{X})= \operatorname{reshape}\left( \mathbf{W}_G\operatorname{reshape}\left( \left[\mathbf{x}_1\mathbf{W}_B^1\mid\mathbf{x}_2\mathbf{W}_B^2\mid\cdots\mid \mathbf{x}_{L/B}\mathbf{W}_B^{L/B}\right],\frac{L}{B},B \right),1,L\right). \tag{29}