一、UniMixer [2026]

《UniMixer: A Unified Architecture for Scaling Laws in Recommendation Systems》

1. 近年来，推荐模型的 scaling laws 受到越来越多的关注，它决定了推荐系统性能与参数、计算量（FLOPs）之间的关系。当前，实现推荐模型 scaling 的主流架构主要有三种，分别是 attention-based 方法、TokenMixer-based 方法以及 factorization-machine-based 方法，这些方法在设计理念与架构结构上存在本质差异。本文提出一种面向推荐系统的统一的 scaling 架构—— UniMixer，旨在提升 scaling 效率，并构建一个统一了主流 scaling blocks 的理论框架。通过将 rule-based TokenMixer 转化为等价的参数化结构，我们构建了一个通用的参数化的 feature mixing 模块，使 token mixing 模式可在模型训练过程中被优化和被学习。同时，通用的参数化的 token mixing 移除了 TokenMixer 中要求 heads 数量与 tokens 数量相等的约束。此外，我们为推荐系统建立了统一的 scaling module 设计框架，打通了 attention-based 方法、TokenMixer-based 方法、以及 factorization-machine-based 方法之间的联系。为进一步提升 scaling ROI，我们设计了轻量级 UniMixing 模块—— UniMixing-Lite，在大幅压缩模型参数与计算成本的同时显著提升模型性能。scaling curves 如下图所示。本文通过大量离线与在线实验验证了 UniMixer 优异的 scaling 能力。

   

2. 大语言模型（Large language models: LLMs）展现出一个令人瞩目的现象：随着模型规模、数据量与计算资源的增加，性能稳步提升，这一现象被称为 scaling laws。LLMs 中显著的性能 scaling 效果启发了推荐系统领域，研究者开始探索适配推荐任务的scaling 框架。近年来，研究人员尝试设计 scaling 模块并多层堆叠，以提升 ranking 模型复杂度，从而实现模型性能与模型规模、计算成本（如参数、FLOPs ）之间的 scaling laws。

   推荐系统基于大量的 multi-field 的 user and item features，预测用户行为，为用户展示最相关的内容，提升用户对推荐结果的 positive engagements。这些 multi-field features 通常包含 categorical features 与 dense features，具备更动态的 embedding representations，可从多视角捕获信息。与自然语言处理（natural language processing: NLP）领域不同（在 NLP 中，所有 tokens 共享一个统一的 embedding space），推荐任务的 feature space 天然具有异构性（heterogeneous）。因此，learning heterogeneous features interactions 是推荐领域与 NLP 领域的根本区别。

   得益于 Transformer 在 LLMs 中的巨大成功，一个自然的思路是修改 Transformer 模块以适配推荐任务，因为直接将Transformer 模块作为推荐系统 scaling laws 的 fundamental block 通常不可行。为解决 heterogeneous feature interaction problem，当前推荐模型的主流 scaling 架构可分为三类：attention-based 方法、TokenMixer-based 方法以及 factorization-machine-based 方法。

   • attention-based 方法为每个 input token 构建 token-specific query, key, and value projections。如 HiFormer （《Hiformer: Heterogeneous feature interactions learning with transformers for recommender systems》）、FAT（《From scaling to structured expressivity: Rethinking transformers for ctr prediction》）、HHFT（《Hhft: Hierarchical heterogeneous feature transformer for recommendation systems》） 等等。

   • 与 attention-based 方法不同，TokenMixer-based 方法采用 rule-based token mixing 操作实现 heterogeneous feature interactions，避免计算两个 heterogeneous semantic spaces 之间的内积相似度。如 RankMixer（《Rankmixer: Scaling up ranking models in industrial recommenders》）、TokenMixer-Large （《Tokenmixer-large: Scaling up large ranking models in industrial recommenders》）等等。

   • 而 factorization-machine-based方法则通过引入因子分解机（Factorization Machine: FM）模块，建模每层 input embeddings 之间的 feature interactions。如 Wukong（《Wukong: Towards a scaling law for large-scale recommendation》）、Kunlun （《Kunlun: Establishing scaling laws for massive-scale recommendation systems through unified architecture design》）等等。

   这些框架基于完全不同的 scaling blocks 来构建，却均具备模型性能 scaling up 的能力。这引发了一个根本性问题：我们能否为推荐系统构建一个统一的 scaling module，融合现有主流 scaling components 的优势？

   为打通这些 scaling modules 间的联系，我们首先为 rule-based TokenMixer 操作建立 parameterized formulation。通过进一步优化计算流程，我们推导出计算成本更低的 UniMixing 模块。基于该设计与实验结果，我们提出一个统一的理论框架，整合推荐系统主流的 scaling modules。此外，我们设计了轻量级的 UniMixer 模块，结合现有主流 scaling blocks 的优势，实现最优的参数效率与计算效率。我们希望该统一架构能助力推荐系统领域迎来属于自己的 "attention moment"。

   本文主要贡献总结如下：

   • 通过对 rule-based TokenMixer 进行等价的参数化，揭示其 feature interaction 模式。

   • 提出统一的 scaling 框架 UniMixer，打通 attention-based 方法、TokenMixer-based 方法、FM-based 方法之间的差异与联系。通过优化计算流程，UniMixer 显著降低训练与推理阶段的计算复杂度与 GPU 内存消耗。

   • 为进一步减少模型参数与计算成本，设计轻量级的 UniMixing 模块—— UniMixing-Lite，可同时利用 attention-based 架构与 TokenMixer-based 架构的优势，实现更优的 scaling 效率。

   • 开展大量离线与在线实验，证明 UniMixer 具备优异的 scaling 能力。

1.1 相关工作

1. 当前，面向大规模推荐系统建立 scaling laws 的建模范式主要有三种：attention-based 方法、TokenMixer-based 方法、以及 FM-based方法。

2. Attention-Based Framework：近年来，推荐系统领域将 Transformer 适配用于 CTR prediction。该范式的核心挑战是弥合token sequence 的异构性与语言建模假设的 sequential compositionality 之间的差距。

   • 为此，《Hiformer: Heterogeneous feature interactions learning with transformers for recommender systems》提出 heterogeneous attention layer 解决 heterogeneous feature interaction 问题，并设计 HiFormer，将 heterogeneous tokens 展平为 a single vector representation，显式建模高阶交互。

   • 此外，Field-Aware Transformers: FAT 通过 factorized contextual alignment 与 cross-field modulation，将 field-aware interaction 的先验知识注入 attention 机制，进一步建立 CTR prediction 的经验 scaling law （《From scaling to structured expressivity: Rethinking transformers for ctr prediction》）。

   • HHFT 通过交替使用 heterogeneous Transformer blocks（保留 domain-specific semantics ）与 HiFormer blocks （学习高阶交互），验证了这些 scaling 特性（《Hhft: Hierarchical heterogeneous feature transformer for recommendation systems》）。

   • 此外，在 dynamic user behavior modeling 中，HSTUV1/V2 、MARM、OneTrans、Climber、Hyformer、LLaTTE 等方法利用 attention 机制捕获长程时序依赖。

   这些方法凸显了统一 feature interaction 与 sequential behavior modeling 以实现更鲁棒 scaling laws 的潜力。

3. TokenMixer-Based Framework：尽管 attention 机制具备强大的 feature interaction表达能力，但 attention score computation 的二次复杂度会带来高昂的计算成本。受计算机视觉领域 MLP-Mixer（《Mlp-mixer: An all-mlp architecture for vision》）成功的启发，工业级推荐系统出现了向 token-mixing 架构的范式转变，诞生了 RankMixer（《Rankmixer: Scaling up ranking models in industrial recommenders》）、Lemur（《Lemur: Large scale end-to-end multimodal recommendation》）、TokenMixer-Large（《Tokenmixer-large: Scaling up large ranking models in industrial recommenders》）等先进模型。

   • 例如，RankMixer 用静态的、无参数的 token-mixing 操作替代 dynamic attention，在保持相当的 FLOPs 的同时，实现了有竞争力的 CTR 预测性能。

   • 在此基础上，TokenMixer-Large 通过引入辅助的 residual connections 与定制的 loss functions，将该架构扩展至 13 Billion 参数规模，在 various model dimensions 上展现出良好的 scaling laws。

   尽管如此，当前 token-mixing 算子的设计仍高度依赖经验规则，缺乏与传统 FM-based 方法或 attention-based 方法的严谨理论桥梁。

4. FM-Based Framework：FM-based 的开创性方法采用低阶 pairwise modeling 来实现推荐系统的 feature interactions（《Factorization machines》），后续经过 Field-aware FMs 来泛化（《Field-aware factorization machines for ctr prediction》），可捕获 field-specific and context-sensitive interactions。这类模型具备高可解释性与高效性，但固有地受限于低阶交互能力。

   为解决该局限，DeepFM （《Deepfm: a factorization-machine based neural network for ctr prediction》）、AutoInt（《Autoint: Automatic feature interaction learning via self-attentive neural networks》）、DCN （《Deep & cross network for ad click predictions》、《Dcn v2: Improved deep & cross network and practical lessons for web-scale ctr prediction》）等多种神经网络扩展方法，融合 MLP 或 transformer attention 以捕获高阶交互。

   近期，Wukong （《Wukong: Towards a scaling law for large-scale recommendation》）通过堆叠 FM-style interaction blocks with linear compression，展现出良好的 scaling 特性。

   然而，FM-based 方法对显式低阶交互的依赖，仍限制了模型在参数量与 FLOPs 扩大时的性能提升，这与 LLMs 中观察到的 predictive scaling laws 形成对比。

1.2 预备知识

1. 考虑一类判别式推荐任务，如 rating 预测、点击率（click-through rate: CTR）预测、点击后转化率（post-click conversion rate: CVR）预测等，这类任务通常被建模为监督学习问题。数据集定义为 $\mathcal{D}=\{({\overrightarrow{\mathbf{x}}}_1,y_1),\cdots ,({\overrightarrow{\mathbf{x}}}_i,y_i),\cdots ,({\overrightarrow{\mathbf{x}}}_N,y_N)\}$$\mathcal D=\left\{\left(\mathbf {\vec x}_{1}, y_{1}\right), \cdots,\left(\mathbf {\vec x}_{i}, y_{i}\right), \cdots,\left(\mathbf {\vec x}_{N}, y_{N}\right)\right\}$，其中：

   • ${\overrightarrow{\mathbf{x}}}_i=[x_i^{(1)},x_i^{(2)},\cdots ,x_i^{(F)}]\in {\mathbb{R}}^F$$\mathbf{\vec x}_{i}=\left[x_{i}^{(1)}, x_{i}^{(2)}, \cdots, x_{i}^{(F)}\right] \in \mathbb R^F$ 包含 $F$$F$ 个 feature fields。通常，输入特征 $\overrightarrow{\mathbf{x}}=\{{\overrightarrow{\mathbf{x}}}^C,{\overrightarrow{\mathbf{x}}}^D\}$$\mathbf{\vec x} = \left\{\mathbf{\vec x}^{C}, \mathbf{\vec x}^{D}\right\}$ 分为 categorical features ${\overrightarrow{\mathbf{x}}}^C$$\mathbf{\vec x}^C$ 与 dense features ${\overrightarrow{\mathbf{x}}}^D$$\mathbf{\vec x}^{D}$，用 $|C|$$|C|$、$|D|$$|D|$ 分别表示 categorical features 数量与 dense features 数量。

   • $y_i$$y_i$ 为第 $i$$i$ 个样本的标签，$y_i\in 0,1$$y_{i} \in{0,1}$对应二分类问题，$y_i\in \mathbb{R}$$y_{i} \in \mathbb{R}$ 对应回归问题。

   • $N$$N$ 为样本数量。

2. 对于 CTR prediction 与 CVR prediction 任务，核心目标是建立模型从而预测 click 或 conversion 的概率 $\text{Pr}(y_i=1\mid {\overrightarrow{\mathbf{x}}}_i)$$\text{Pr}\left(y_{i}=1 \mid \mathbf{\vec x}_{i}\right)$。推荐系统中学到的 embedding representations 更具动态性。与语言模型 input tokens 不同，推荐系统中的 feature spaces 天然具有异构性。因此，直接将大语言模型所用的 Transformer 架构迁移至推荐建模并不合适。

   迄今为止，推荐领域的 scaling laws 主要通过三类 foundational blocks 及其变体实现。

3. Heterogeneous Attention Layer：Heterogeneous-attention-based 架构通常采用 field-specific query, key, and value projections 来实现 heterogeneous feature interaction。给定 input hidden states $\mathbf{X}=[{\overrightarrow{\mathbf{x}}}_1;\cdots ;{\overrightarrow{\mathbf{x}}}_T]\in {\mathbb{R}}^{T\times D}$$\mathbf X=\left[\mathbf{\vec x}_{1} ; \cdots ; \mathbf{\vec x}_{T}\right] \in \mathbb{R}^{T \times D}$，heterogeneous attention layer 的公式如下：

   $$\begin{array}{c}{\mathbf{Q}}_h=\left[\begin{array}{c}{\mathbf{W}}_Q^{1h}{\overrightarrow{\mathbf{x}}}_1\\ ⋮\\ {\mathbf{W}}_Q^{Th}{\overrightarrow{\mathbf{x}}}_T\end{array}\right]\in {\mathbb{R}}^{T\times d},{\mathbf{K}}_h=\left[\begin{array}{c}{\mathbf{W}}_K^{1h}{\overrightarrow{\mathbf{x}}}_1\\ ⋮\\ {\mathbf{W}}_K^{Th}{\overrightarrow{\mathbf{x}}}_T\end{array}\right]\in {\mathbb{R}}^{T\times d},{\mathbf{V}}_h=\left[\begin{array}{c}{\mathbf{W}}_V^{1h}{\overrightarrow{\mathbf{x}}}_1\\ ⋮\\ {\mathbf{W}}_V^{Th}{\overrightarrow{\mathbf{x}}}_T\end{array}\right]\in {\mathbb{R}}^{T\times d}\end{array}$$

   其中：${\mathbf{W}}_Q^{ih},{\mathbf{W}}_K^{ih},{\mathbf{W}}_V^{ih}\in {\mathbb{R}}^{D\times d}$$\mathbf W_{Q}^{i h},\mathbf W_{K}^{i h}, \mathbf W_{V}^{i h} \in \mathbb{R}^{D\times d}$分别为 query, key, and value projections 的 token-specific weights。

   这里 $T$$T$ 为 heterogeneous tokens 数量，$h$$h$ 为 head 编号。

   multi-head heterogeneous attention layer 的输出计算如下：

   $${\mathbf{O}}_h=\text{softmax}\left(\frac{{\mathbf{Q}}_h{\mathbf{K}}_h^{\mathrm{\top }}}{\sqrt{d}}\right){\mathbf{V}}_h\in {\mathbb{R}}^{T\times d}$$

   将 multi-head heterogeneous attention 的输出进行拼接后，通过线性投影使 output 维度与 input $\mathbf{X}$$\mathbf X$ 对齐。

4. TokenMixer：TokenMixer-based 框架采用无参数的、基于规则的 mixing 操作实现 feature interaction。给定 input hidden states $\mathbf{X}=[{\overrightarrow{\mathbf{x}}}_1;\cdots ;{\overrightarrow{\mathbf{x}}}_T]\in {\mathbb{R}}^{T\times D}$$\mathbf X=\left[\mathbf{\vec x}_{1} ; \cdots ; \mathbf{\vec x}_{T}\right] \in \mathbb{R}^{T \times D}$，TokenMixer 首先将每个 input token ${\overrightarrow{\mathbf{x}}}_t$$\mathbf{\vec x}_t$ 均匀切分为 $H$$H$ 个 heads：

   $$[{\overrightarrow{\mathbf{x}}}_t^{(1)}\mid {\overrightarrow{\mathbf{x}}}_t^{(2)}\mid \cdots \mid {\overrightarrow{\mathbf{x}}}_t^{(H)}]=\text{SplitHead}\left({\overrightarrow{\mathbf{x}}}_t\right)$$

   其中 ${\overrightarrow{\mathbf{x}}}_t^{(h)}\in {\mathbb{R}}^{D/H}$$\mathbf{\vec x}_t^{(h)}\in \mathbb R^{D/H}$ 为第 $h$$h$ 个 head 。

   然后，第 $h$$h$ 个 token ${\overrightarrow{\mathbf{s}}}_h$$\mathbf{\vec s}_h$ 为：

   $${\overrightarrow{\mathbf{s}}}_h=\text{concat}({\overrightarrow{\mathbf{x}}}_1^{(h)},{\overrightarrow{\mathbf{x}}}_2^{(h)},\cdots ,{\overrightarrow{\mathbf{x}}}_T^{(h)})\in {\mathbb{R}}^{(TD/H)}$$

   TokenMixer 的输出可以公式化为：

   $$\begin{array}{c}\mathbf{S}=\left[\begin{array}{c}{\overrightarrow{\mathbf{s}}}_1\\ ⋮\\ {\overrightarrow{\mathbf{s}}}_H\end{array}\right]\in {\mathbb{R}}^{H\times (TD/H)}\end{array}$$

   其中：其中要求 $H$$H$ 与 $T$$T$ 相等，使得输入 $\mathbf{X}$$\mathbf X$ 与输出 $\mathbf{S}$$\mathbf S$ 的维度一致。

5. Wukong：Wukong-based 的模型将一个因子分解机块（Factorization Machine Block: FMB）的输出与一个线性投影层的输出进行拼接，以提升 interaction component：

   $$\begin{array}{c}\text{FMB}(\mathbf{X})=\text{reshape}(\text{MLP}(\text{LN}(\text{flatten}(\text{FM}(\mathbf{X}))))),\text{FM}(\mathbf{X})=\mathbf{X}{\mathbf{X}}^{\mathrm{\top }}\mathbf{Y}\\ \text{LCB}(\mathbf{X})=\mathbf{W}\mathbf{X}\end{array}$$

   其中：

   • $\mathbf{W}\in {\mathbb{R}}^{n\times T},\mathbf{Y}\in {\mathbb{R}}^{T\times r}$$\mathbf W\in \mathbb R^{n\times T}, \mathbf Y\in \mathbb R^{T\times r}$ 为可学习的投影矩阵，$\mathbf{Y}$$\mathbf Y$ 用于降低存储 interaction matrix $\mathbf{X}{\mathbf{X}}^{\mathrm{\top }}$$\mathbf X \mathbf X^{\top}$ 的内存需求。

   • $\text{LN}(\cdot )$$\text{LN}(\cdot)$ 为 Layer Normalization。

6. 本文聚焦于为推荐系统建立一个统一的结构基础，融合当前 scaling blocks 的优势，进一步提升 scaling ROI。

1.3 UniMixer

1. 本文构建了用于推荐系统 scaling 的统一模块—— UniMixer block，在统一理论框架下整合了 attention-based 模块、TokenMixer-based 模块、Wukong-based 模块等推荐系统主流 scaling 模块。如 Figure 2 所示，整体架构包含 feature tokenization 、$M$$M$ 个带 Siamese norm and Sparse-Pertoken MoE 的 UniMixer blocks。通过对 rule-based TokenMixer 进行参数化，我们打通了 attention-based 方法、TokenMixer-based 方法、Wokong-based 方法之间的联系，使所提出的 UniMixer 同时具备这些方法的优势。此外，我们开发了一个轻量级 UniMixing 模块，进一步压缩模型参数与计算成本，同时显著提升模型性能。

   

1.3.1 Feature Tokenization

1. 根据 input feature fields 的 semantic categories，首先将 input features $\overrightarrow{\mathbf{x}}$$\mathbf{\vec x}$ 划分为 $N$$N$ 个不相交的 feature domains。

   $$\overrightarrow{\mathbf{x}}=[\underset{\text{User Profile}}{\underbrace{x_U^{(1)},\cdots ,x_U^{(n_U)}}},\underset{\text{Item Features}}{\underbrace{x_I^{(1)},\cdots ,x_I^{(n_I)}}},\underset{\text{Behavior Sequence}}{\underbrace{x_B^{(1)},\cdots ,x_B^{(n_B)}}},\underset{\text{Query Features}}{\underbrace{x_Q^{(1)},\cdots ,x_Q^{(n_Q)}}},\cdots ]$$

   每个 feature domain 通过 embedding layers 转化为不同维度的 embedding vectors：

   $${\overrightarrow{\mathbf{e}}}_n=\text{Embedding}\left({\overrightarrow{\mathbf{x}}}_n\right)\in {\mathbb{R}}^{d_{\text{domain}}}$$

   其中：${\overrightarrow{\mathbf{x}}}_n$$\mathbf{\vec x}_n$ 为 feature domain 内某个 feature 的 one-hot embedding，$d_{\text{domain}}$$d_\text{domain}$ 为该 feature domain 对应的 embedding 维度。

   将所有 obtained feature domain embeddings 拼接为一个 embedding 向量 $\overrightarrow{\mathbf{e}}=[{\overrightarrow{\mathbf{e}}}_1,{\overrightarrow{\mathbf{e}}}_2,\cdots ,{\overrightarrow{\mathbf{e}}}_N]$$\mathbf{\vec e}=\left[\mathbf{\vec e}_{1}, \mathbf{\vec e}_{2}, \cdots, \mathbf{\vec e}_{N}\right]$。与 RankMixer 类似，我们将 embedding 向量 $\overrightarrow{\mathbf{e}}$$\mathbf{\vec e}$ 均匀切分为合适数量的 blocks，再通过 token-specific linear layer 将每个 block 投影为 token embedding：

   $${\overrightarrow{\mathbf{x}}}_i={\mathbf{W}}_i^{\text{proj}}{\overrightarrow{\mathbf{e}}}_{(d\times i):(d\times i+d)}+{\overrightarrow{\mathbf{b}}}_i^{\text{proj}}\in {\mathbb{R}}^D$$

   其中：${\mathbf{W}}_i^{\text{proj}}\in {\mathbb{R}}^{D\times d},{\overrightarrow{\mathbf{b}}}_i^{\text{proj}}\in {\mathbb{R}}^D$$\mathbf W_i^\text{proj}\in \mathbb R^{D\times d}, \mathbf{\vec b}_i^\text{proj}\in \mathbb R^D$ 为第 $i$$i$ 个 block 的待学习的参数。$d$$d$ 为每个 block 的维度。

   按列堆叠 ${\overrightarrow{\mathbf{x}}}_i$$\mathbf{\vec x}_i$ 即可得到 input hidden states $\mathbf{X}\in {\mathbb{R}}^{T\times D}$$\mathbf X \in \mathbb{R}^{T \times D}$。

   • 如果不按照 feature domains 来组织 $\overrightarrow{\mathbf{e}}$$\mathbf{\vec e}$，而是随机组织，那么结果会怎样？可以做实验来验证。

   • “将 embedding 向量 $\overrightarrow{\mathbf{e}}$$\mathbf{\vec e}$ 均匀切分为合适数量的 blocks“，这一步其实就是 sparsify 操作，参考论文 SSRNet（《Beyond Dense Connectivity: Explicit Sparsity for Scalable Recommendation》）。

   • 注意：这里要求 $T\ll D$$T\ll D$ ，并且 $T$$T$ 能整除 $D$$D$ 。这也是为什么需要拆分为 blocks，而不是采用原始的 feature-level embedding 的原因。

1.3.2 UniMixer Block

1. Heterogeneous Feature Interactions：如 ”预备知识“ 章节所述，heterogeneous attention 通过采用 token-specific query, key, and value weights，解决两个 heterogeneous semantic spaces 的 feature interaction 问题。然而，通过内积相似度得到的 attention pattern 通常带有对角占优先验（diagonally dominant prior）。在训练初期，权重矩阵 ${\mathbf{W}}_Q^h,{\mathbf{W}}_K^h$$\mathbf W_{Q}^{h},\mathbf W_{K}^{h}$ 被随机初始化，attention weights （即 ${\mathbf{Q}}_h{\mathbf{K}}_h^{\mathrm{\top }}$$\mathbf Q_h\mathbf K_h^\top$）的大小很大程度上由 input token values $\mathbf{X}$$\mathbf X$ 主导，易导致 attention weights 集中在少数 tokens 上，如 Figure 3(a) 所示。

   

   由 Figure 3(a) 可见，heterogeneous attention 的 attention weights 尖锐且稀疏，给梯度反向传播带来风险，导致 query and key weights 训练困难，甚至可能停滞（如 Figure 3(a) 中 heterogeneous attention 的 attention weights 的第 10 行与第 15 行）。同时，在大规模 heterogeneous feature inputs 下，这类 attention 模式可能导致 feature interactions 趋同，即 attention scores 极小且缺乏区分度，可能产生噪声信号，掩盖关键 feature interaction 模式。

   另一方面，无参数的、基于规则的 TokenMixer 操作缺乏可学习性与场景适应性，可能导致 heterogeneous feature interactions 不足或错误。此外，要求 $T=H$$T=H$ 进一步限制了 heterogeneous feature interaction 模式的选择。通过深入分析 TokenMixer 操作，我们得到一些有趣的发现，使 TokenMixer 操作参数化成为可能。如 Figure 3(b) 所示，我们发现：TokenMixer 操作可视为置换矩阵（permutation matrix） ${\mathbf{W}}^{\text{perm}}$$\mathbf W^\text{perm}$与 flattened input embedding $\text{flatten}(\mathbf{X})\in {\mathbb{R}}^{(TD)}$$\text{flatten}(\mathbf X) \in \mathbb{R}^{(TD)}$的乘积，公式如下：

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

   其中：${\mathbf{W}}^{\text{perm}}\in {\mathbb{R}}^{(TD)\times (TD)}$$\mathbf W^\text{perm}\in \mathbb R^{(TD)\times (TD)}$ 为一个大型置换矩阵，附录 A 给出具体数值示例。

   一个自然的思路是通过参数化置换矩阵 ${\mathbf{W}}^{\text{perm}}$$\mathbf W^\text{perm}$，使 rule-based TokenMixer 具备可学习性与可优化性。但 $O\left(T^2{D}^2\right)$$O\left(T^{2} D^{2}\right)$ 的计算复杂度与 $O\left(T^2{D}^2\right)$$O\left(T^2 D^2\right)$ 的参数数量难以接受。通过观察，我们总结出 TokenMixer 置换矩阵 ${\mathbf{W}}^{\text{perm}}$$\mathbf W^\text{perm}$ 的关键特性。

   • 可压缩性：置换矩阵 ${\mathbf{W}}^{\text{perm}}$$\mathbf W^\text{perm}$ 可等价分解为两个更小矩阵的克罗内克积（Kronecker product），即 ${\mathbf{W}}^{\text{perm}}=\mathbf{G}\otimes \mathbf{I}$$\mathbf W^\text{perm}= \mathbf G\otimes\mathbf I$。其中 $\mathbf{I}\in {\mathbb{R}}^{\frac{D}{T}\times \frac{D}{T}}$$\mathbf I\in \mathbb R^{\frac{D}{T}\times \frac{D}{T}}$ 为 $\frac{D}{T}$$\frac{D}{T}$ 阶单位矩阵，$\mathbf{G}\in {\mathbb{R}}^{T^2\times T^2}$$\mathbf G\in \mathbb R^{T^2\times T^2}$ 为 $T^2$$T^2$ 阶方阵，符号 $\otimes$$\otimes$ 代表克罗内克积运算。

     克罗内克积 $\mathbf{A}\otimes \mathbf{B}$$\mathbf A\otimes \mathbf B$ ：把矩阵 $\mathbf{A}$$\mathbf A$ 的每个元素，整体乘以整个矩阵 $\mathbf{B}$$\mathbf B$，按位置拼成新大矩阵。例如：

     $$\begin{array}{c}\mathbf{A}=\left[\begin{array}{cc}{a}_{1,1}& a_{1,2}\\ a_{2,1}& a_{2,2}\end{array}\right],\mathbf{B}=\left[\begin{array}{cc}{b}_{1,1}& b_{1,2}\\ b_{2,1}& b_{2,2}\end{array}\right],\mathbf{A}\otimes \mathbf{B}=\left[\begin{array}{cc}{a}_{1,1}\mathbf{B}& a_{1,2}\mathbf{B}\\ a_{2,1}\mathbf{B}& a_{2,2}\mathbf{B}\end{array}\right]\end{array}$$

     .

   • 双随机性：置换矩阵 ${\mathbf{W}}^{\text{perm}}$$\mathbf W^\text{perm}$ 的任意行、列元素求和结果均为 1，满足行和、列和归一特性。

   • 稀疏性：该置换矩阵的每一行、每一列仅有唯一一个非零元素。

   • 对称性：若维度参数 $T$$T$ 与 $H$$H$ 相等，则置换矩阵 ${\mathbf{W}}^{\text{perm}}$$\mathbf W^\text{perm}$ 为对称矩阵，满足 ${\mathbf{W}}^{\text{perm}}={\left({\mathbf{W}}^{\text{perm}}\right)}^{\mathrm{\top }}$$\mathbf W^\text{perm}= \left(\mathbf W^\text{perm}\right)^\top$ ；若二者不等，则置换矩阵 ${\mathbf{W}}^{\text{perm}}$$\mathbf W^\text{perm}$ 为非对称矩阵。

   这里有一个前提：要求 $D\ge T$$D \ge T$ 并且 $D$$D$ 能被 $T$$T$ 整除。$T$$T$ 为 input hidden states 的 tokens 数量，$D$$D$ 为 input hidden states 的维度。

   

   根据 TokenMixer 置换矩阵的特性，通过参数化矩阵 $\mathbf{G}$$\mathbf G$ 与 $\mathbf{I}$$\mathbf I$，token mixing 的参数数量大幅降低，即 $O(T^4+{\left(\frac{D}{T}\right)}^2)$$O\left(T^{4}+\left(\frac{D}{T}\right)^{2}\right)$，其中 $T$$T$ 通常远小于 $D$$D$。此外，TokenMixer 参数化仍面临三大挑战：

   • 直接用参数化的 $\mathbf{G}$$\mathbf G$ 与 $\mathbf{I}$$\mathbf I$ 来重构 ${\mathbf{W}}^{\text{perm}}$$\mathbf W^\text{perm}$，在模型训练与推理过程中仍会产生 size 为 $[TD,TD]$$[TD, TD]$ 的中间变量，对 GPU 内存要求极高。

     因为需要生成 ${\mathbf{W}}^{\text{perm}}=\mathbf{G}\otimes \mathbf{I}$$\mathbf W^\text{perm} = \mathbf G\otimes \mathbf I$ 。

   • 如何保证学到的参数满足双随机性、稀疏性与对称性。

   • 如何设计融合了现有 scaling 模块的优势的 unified recommendation scaling module，为推荐系统建立更优的 scaling 效率。

2. Unified Token Mixing Module：受 Figure 3 启发，unified token mixing module 不再使用 $T$$T$ 与 $D$$D$，而是定义置换矩阵中的 block num 与 block size。 block size 记为 $B$$B$，block num 为 $(L//B{)}^2$$(L // B)^{2}$，其中 $L$$L$ 为 input embedding 维度（即，$\overrightarrow{\mathbf{e}}=[{\overrightarrow{\mathbf{e}}}_1,{\overrightarrow{\mathbf{e}}}_2,\cdots ,{\overrightarrow{\mathbf{e}}}_N]$$\mathbf{\vec e}=\left[\mathbf{\vec e}_{1}, \mathbf{\vec e}_{2}, \cdots, \mathbf{\vec e}_{N}\right]$ 的维度），可被 block size $B$$B$ 所整除。

   换成前面的符号，即：$L=T\times D$$L = T\times D$。

   将 $\mathbf{G}$$\mathbf G$ 的 parameterized weights 记为 ${\mathbf{W}}_G\in {\mathbb{R}}^{(L//B)\times (L//B)}$$\mathbf W_{G}\in \mathbb R^{(L//B)\times (L//B)}$。考虑置换矩阵的稀疏性并实现充分的 heterogeneous feature interactions，我们为 ${\mathbf{W}}_G$$\mathbf W_G$ 的每一行分配 distinct parameterized weight ${\mathbf{W}}_B^i\in {\mathbb{R}}^{B\times B}$$\mathbf W_{B}^{i}\in \mathbb R^{B\times B}$，一共 $L//B$$L//B$ 行。该操作使每个 block 具备不同的 feature interaction 模式，通过学习参数矩阵 ${\mathbf{W}}_G$$\mathbf W_{G}$与${\mathbf{W}}_B^i$$\mathbf W_{B}^{i}$，可得到交互模式更丰富的置换矩阵 ${\mathbf{W}}^{\text{perm }}$$\mathbf W^\text{perm }$，公式如下：

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

   其中：$\otimes$$\otimes$为广义克罗内克积（generalized Kronecker product）。

   ${\mathbf{W}}_G\otimes {\left\{{\mathbf{W}}_B^i\right\}}_{i=1}^{L//B}$$\mathbf W_G\otimes\left\{\mathbf W_B^i\right\}_{i=1}^{L//B}$ 和经典的克罗内克积不同，它的物理含义为：

   $$\begin{array}{c}{\mathbf{W}}_G\otimes {\left\{{\mathbf{W}}_B^i\right\}}_{i=1}^{L//B}=\left[\begin{array}{cccc}{W}_{G,1,1}{\mathbf{W}}_B^1& W_{G,1,2}{\mathbf{W}}_B^2& \cdots & W_{G,1,L//B}{\mathbf{W}}_B^{L//B}\\ W_{G,2,1}{\mathbf{W}}_B^1& W_{G,2,2}{\mathbf{W}}_B^2& \cdots & W_{G,2,L//B}{\mathbf{W}}_B^{L//B}\\ ⋮& ⋮& \ddots & ⋮\\ W_{G,L//B,1}{\mathbf{W}}_B^1& W_{G,L//B,2}{\mathbf{W}}_B^2& \cdots & W_{G,L//B,L//B}{\mathbf{W}}_B^{L//B}\end{array}\right]\end{array}$$

   即：每一列采用不同的 ${\mathbf{W}}_B^i$$\mathbf W_B^i$ 。

   也可以对 ${\mathbf{W}}_G$$\mathbf W_G$ 的每个元素采用不同的 ${\mathbf{W}}_B^{i,j}$$\mathbf W_B^{i,j}$ 。此时会引起参数爆炸，以及过拟合。

   接下来优化 $\text{UniMixing}(\mathbf{X})$$\text{UniMixing}(\mathbf X)$ 的计算流程，显著降低计算成本与 GPU 内存需求。

   • 首先将 embedding vector $\text{flatten}(\mathbf{X})$$\text{flatten}(\mathbf X)$ 均匀切分为 $L//B$$L // B$ 个向量，每个向量的 size 为 $B$$B$，表示如下：

     $$[{\overrightarrow{\mathbf{x}}}_1\mid {\overrightarrow{\mathbf{x}}}_2\mid \cdots \mid {\overrightarrow{\mathbf{x}}}_{\frac{L}{B}}]=\text{Split}(\text{flatten}(\mathbf{X}),\frac{L}{B})\in {\mathbb{R}}^{B\times (L//B)}$$

   • 然后将 block weights ${\mathbf{W}}_B^i\in {\mathbb{R}}^{B\times B}$$\mathbf W_{B}^{i}\in \mathbb R^{B\times B}$与对应 block-wise vectors ${\overrightarrow{\mathbf{x}}}^{(i)}\in {\mathbb{R}}^B$$\mathbf{\vec x}^{(i)}\in \mathbb R^{B}$ 分别相乘，得到 local feature interaction vector：

     $$\begin{array}{c}\mathbf{H}=[{\overrightarrow{\mathbf{x}}}_1{\mathbf{W}}_B^1\mid {\overrightarrow{\mathbf{x}}}_2{\mathbf{W}}_B^2\mid \cdots \mid {\overrightarrow{\mathbf{x}}}_{L//B}{\mathbf{W}}_B^{L//B}]\in {\mathbb{R}}^{B\times (L//B)}\\ \text{reshape}(\mathbf{H},L//B,B)=\left[\begin{array}{c}{\overrightarrow{\mathbf{x}}}_1{\mathbf{W}}_B^1\\ {\overrightarrow{\mathbf{x}}}_2{\mathbf{W}}_B^2\\ ⋮\\ {\overrightarrow{\mathbf{x}}}_{L//B}{\mathbf{W}}_B^{L//B}\end{array}\right]\in {\mathbb{R}}^{(L//B)\times B}\end{array}$$

   • 最后，UniMixing 模块的输出为：

     $$\text{UniMixing}(\mathbf{X})=\text{reshape}({\mathbf{W}}_G\text{reshape}(\mathbf{H},L//B,B),1,L)$$

     $\text{reshape}(\mathbf{H},L//B,B)\in {\mathbb{R}}^{(L//B)\times B}$$\text{reshape}(\mathbf H, L//B, B)\in \mathbb R^{(L//B)\times B}$ ，而 ${\mathbf{W}}_G\in {\mathbb{R}}^{(L//B)\times (L//B)}$$\mathbf W_G\in \mathbb R^{(L//B)\times (L//B)}$ ，因此二者可以执行矩阵乘法。

   与直接使用 reconstructed matrix ${\mathbf{W}}^{\text{perm }}$$\mathbf W^\text{perm }$相比，该操作将计算成本从 $O\left(L^2\right)$$O\left(L^{2}\right)$ 降至 $O(L^2/B+LB)$$O\left(L^{2} / B+L B\right)$，并避免计算过程中产生大型中间变量。这个计算流程优化的证明见附录 B。根据优化后的公式，${\mathbf{W}}_B^i$$\mathbf W_{B}^{i}$ 控制 block 内的交互模式，${\mathbf{W}}_G$$\mathbf W_{G}$ 控制 block 之间的交互模式。对于维度为 $L$$L$ 的 embedding inputs，不再要求 $T=H$$T=H$。与 TokenMixer 操作相比，UniMixing 模块具备更多样的 local and global feature mixing patterns and interaction scales，同时保留优势：它是可学习的、可优化的。

   为保证学到的置换矩阵满足双随机性，采用 Sinkhorn-Knopp 迭代，通过指数算子（exponent operator）使 ${\mathbf{W}}_G$$\mathbf W_{G}$与 ${\mathbf{W}}_B^i$$\mathbf W_{B}^{i}$ 的所有元素为正，再交替缩放 rows and columns 使其和为 1，进行迭代归一化。此外，引入温度系数控制参数矩阵的稀疏性。最后，用$({\mathbf{W}}_G+{\mathbf{W}}_G^{\mathrm{\top }})/2$$(\mathbf W_{G}+\mathbf W_{G}^{\top}) / 2$ 与 $({\mathbf{W}}_B^i+{\mathbf{W}}_B^{i\mathrm{\top }})/2$$(\mathbf W_{B}^{i}+\mathbf W_{B}^{i \top}) / 2$ 实现参数矩阵的对称约束。最终约束权重如下：

   $$\begin{array}{c}{\tilde{\mathbf{W}}}_G=\frac{{\mathbf{W}}_G+{\mathbf{W}}_G^{\mathrm{\top }}}{2},{\tilde{\mathbf{W}}}_B^i=\frac{{\mathbf{W}}_B^i+{\mathbf{W}}_B^{i\mathrm{\top }}}{2}\\ {\bar{\mathbf{W}}}_G=\text{Sinkhorn-Knopp}\left(\frac{{\tilde{\mathbf{W}}}_G}{\tau }\right),{\bar{\mathbf{W}}}_B^i=\text{Sinkhorn-Knopp}\left(\frac{{\tilde{\mathbf{W}}}_B^i}{\tau }\right)\end{array}$$

   当 $\tau$$\tau$ 较大（如 1.0）时：除法后元素差异被缩小，指数运算后分布比较均匀，最终的双随机矩阵元素值差异不大（接近均匀分布）。当 $\tau$$\tau$ 较小（如 0.05）时：除法放大了元素之间的原始差异，指数运算后大的元素更大，小的元素更小，最终的双随机矩阵变得尖锐：少数元素接近 1，其余接近 0。在极限情况下（$\tau \to 0$$\tau\rightarrow 0$），矩阵趋近于一个硬置换矩阵（每行每列只有一个 1，其余为 0）。

   实验章节表明：$\tau$$\tau$ 需要采用一个较小的值。

   其中 $\tau$$\tau$ 为温度系数。

   Sinkhorn‑Knopp 迭代（也称 Sinkhorn 缩放算法）是一种将任意正矩阵转化为双随机矩阵（doubly stochastic matrix）的经典数值方法。所谓双随机矩阵，是指一个方阵满足：所有元素非负（通常为正）、每行之和为 1、每列之和也为 1。

   算法基本步骤：给定一个初始的正矩阵 $\mathbf{A}\in {\mathbb{R}}^{n\times n},a_{i,j}>0$$\mathbf A\in \mathbb R^{n\times n}, a_{i,j}\gt 0$，Sinkhorn‑Knopp 迭代通过交替缩放行和列，使其逼近一个双随机矩阵：

   • 重复迭代直到收敛（或固定次数）：

     • 行归一化：将每一行除以其行和，使每行和为 1。

     • 列归一化：将每一列除以其列和，使每列和为 1。

   经过足够多的迭代后，矩阵 $\mathbf{A}$$\mathbf A$ 收敛到一个唯一的双随机矩阵（当 $\mathbf{A}$$\mathbf A$ 为正且连通时）。实际实现时，通常使用指数算子先对原始矩阵元素取指数（exp），以保证元素为正，然后再应用上述缩放迭代。

   随后用残差连接与归一化模块处理 UniMixing block 的输出：

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

   .

3. A Unified Perspective of Heterogeneous Feature Interaction：观察 ${\mathbf{V}}_h$$\mathbf V_{h}$ 的公式与 $\text{reshape}(\mathbf{H},L//B,B)$$\text{reshape}(\mathbf H, L//B, B)$ 的公式，我们发现若将 blocks 数量 $L//B$$L // B$ 设为 $T$$T$，且 ${\mathbf{W}}_V^{ih}$$\mathbf W_{V}^{i h}$与 ${\mathbf{W}}_B^i$$\mathbf W_{B}^{i}$ 的维度相同，则 $\text{reshape}(\mathbf{H},L//B,B)={\mathbf{V}}_h$$\text{reshape}(\mathbf H, L//B, B)=\mathbf V_{h}$。这意味着在 ${\mathbf{W}}_V^i={\mathbf{W}}_B^i$$\mathbf W_{V}^{i}= \mathbf W_{B}^{i}$ 的条件下，UniMixer 的 local interaction projection 等价于 value projection of the heterogeneous attention layer。另一方面，${\mathbf{W}}_G$$\mathbf W_{G}$ 的维度与作用和 attention weights 相同；此外， ${\mathbf{W}}_G$$\mathbf W_G$ 需要满足双随机性、稀疏性与对称性。

   Wukong 的 feature interaction 基于 FM 组件。$\text{FMB}(\mathbf{X})$$\text{FMB}(\mathbf X)$ 的表达式可改写为：

   $$\text{FMB}(\mathbf{X})=\text{reshape}\left(\text{MLP}\left(\text{LN}\left(\text{flatten}((\mathbf{X}\mathbf{I}){\left(\mathbf{X}\mathbf{I}\right)}^{\mathrm{\top }}\mathbf{Y})\right)\right)\right)$$

   其中：$\mathbf{I}$$\mathbf I$ 为合适维度的单位矩阵。现在我们聚焦核心的 feature interaction 模块$(\mathbf{X}\mathbf{I})(\mathbf{X}\mathbf{I}{)}^{\mathrm{\top }}\mathbf{Y}$$(\mathbf X \mathbf I)(\mathbf X \mathbf I)^{\top} \mathbf Y$。在 attention 模块中，当${\mathbf{W}}_Q=\mathbf{I}$$\mathbf W_{Q}=\mathbf I$、${\mathbf{W}}_K=\mathbf{I}$$\mathbf W_{K}=\mathbf I$，且 value 矩阵不依赖 hidden state input $\mathbf{X}$$\mathbf X$（即 ${\mathbf{V}}_h={\mathbf{W}}_V=\mathbf{Y}$$\mathbf V_{h}=\mathbf W_{V}=\mathbf Y$）时，attention 机制退化为 FM 模块。因此，attention-based 机构、TokenMixer-based 架构与 Wukong-based 架构可统一于单一理论框架：

   $$\begin{array}{c}\text{UniMixing}(\mathbf{X})=\text{reshape}(\underset{\text{Global Mixing Pattern}}{\underbrace{G(\mathbf{X},{\mathbf{W}}_G)}}\underset{\text{Local Mixing Pattern}}{\underbrace{\left[\begin{array}{c}{\overrightarrow{\mathbf{x}}}_1{\mathbf{W}}_B^1\\ ⋮\\ {\overrightarrow{\mathbf{x}}}_{L//B}{\mathbf{W}}_B^{L//B}\end{array}\right]}},1,L)\end{array}$$

   其中：$G(\mathbf{X},{\mathbf{W}}_G)$$G(\mathbf X, \mathbf W_G)$ 为 heterogeneous feature interaction projection，衡量 token-to-token/block-to-block 的 interaction 强度。

   为便于分析各类方法的差异与联系，我们考虑 single-head attention setting。在统一理论框架下（即，$\text{UniMixing}(\mathbf{X})$$\text{UniMixing}(\mathbf X)$），各类方法的差异总结于 Table 1。对于 self-attention、heterogeneous attention 与 FM，global mixing pattern $G(\mathbf{X},{\mathbf{W}}_G)$$G(\mathbf X, \mathbf W_{G})$ 通过计算两个 tokens 的内积相似度得到；而 TokenMixer 的 global mixing pattern 与 input token embedding 无关。

   

4. UniMixing-Lite：如 Figure 3 所示，block 粒度越细（即，$B$$B$ 越小），local interaction parameter matrices ${\mathbf{W}}_B^i$$\mathbf W_{B}^{i}$ 的数量越多（即，$L//B$$L//B$ 越大），global interaction parameter matrix ${\mathbf{W}}_G$$\mathbf W_{G}$ 越大，导致了冗余的 local interaction patterns。同时，更大的 global interaction matrix 在 reducing the number of parameters 上的效率较低。因此，基于 UniMixing block，我们设计了一个轻量级的 UniMixing 模块—— UniMixing-Lite，从而进一步减少 module parameters 与计算成本，提升模型的 scaling 效率。

   为解决 local interaction pattern 的冗余性问题，我们引入一个 basis-composed module，动态地生成 block-specific local mixing weight。定义 ${\mathbf{W}}_B^i$$\mathbf W_{B}^{i}$ 的 basis matrices 的集合为 $\{{\mathbf{Z}}_{\ell }{\}}_{\ell =1}^b$$\{\mathbf Z_{\ell}\}_{\ell=1}^{b}$，这些 basis matrices 上的 block-specific weight vectors 为 ${\left\{{\overrightarrow{\omega }}^i\right\}}_{i=1}^{L//B}$$\left\{\vec \omega^{i}\right\}_{i=1}^{L // B}$，其中 $b$$b$ 为 basis local mixing weight 的数量，${\overrightarrow{\omega }}^i=[{\omega }_1^i,\cdots ,{\omega }_b^i]$$\vec \omega^{i}=\left[\omega_{1}^{i}, \cdots, \omega_{b}^{i}\right]$。此外，我们对 global interaction parameter ${\mathbf{W}}_G$$\mathbf W_{G}$ 采用低秩近似，从而进一步提升效率。则 UniMixing-Lite 模块可表示为：

   $$\begin{array}{c}\text{UniMixing-Lite}(\mathbf{X})=\text{reshape}({\mathbf{W}}_r\text{reshape}([{\overrightarrow{\mathbf{x}}}_1{\mathbf{W}}_B^{\ast 1}\mid \cdots \mid {\overrightarrow{\mathbf{x}}}_{\frac{L}{B}}{\mathbf{W}}_B^{\ast \frac{L}{B}}],\frac{L}{B},B),1,L)\\ \mathbf{O}=\text{RMSNorm}(\mathbf{X}+\text{UniMixing-Lite}(\mathbf{X}))\end{array}$$

   其中：

   • ${\mathbf{W}}_r=\text{Sinkhorn-Knopp}({\mathbf{A}}_G{\mathbf{B}}_G)$$\mathbf W_r = \text{Sinkhorn-Knopp}(\mathbf A_G\mathbf B_G)$，${\mathbf{A}}_G\in {\mathbb{R}}^{(L//B)\times r},{\mathbf{B}}_G\in {\mathbb{R}}^{r\times (L//B)}$$\mathbf A_G\in \mathbb R^{(L//B)\times r}, \mathbf B_G\in \mathbb R^{r\times (L//B)}$ ，$r$$r$ 为 ${\mathbf{W}}_G$$\mathbf W_G$ 低秩近似（low-rank approximation ）的秩。

   • ${\mathbf{W}}_B^{\ast }=\text{Sinkhorn-Knopp}\left(\sum _{\ell =1}^b{\omega }_{\ell }^i{\mathbf{Z}}_{\ell }\right)$$\mathbf W_B^* = \text{Sinkhorn-Knopp}\left(\sum_{\ell=1}^b \omega_\ell^i \mathbf Z_\ell \right)$。

   UniMixing-Lite 模块同时保留了 TokenMixer 的低参数的 global interaction pattern 、以及 attention 的针对heterogeneous features 的 local interaction 能力，可同时利用 attention-based 方法与 token-mixer-based 方法的优势。

   同时对 ${\mathbf{W}}_G$$\mathbf W_G$ 和 ${\mathbf{W}}_B^i$$\mathbf W_B^i$ 进行简化。

5. Pertoken SwiGLU：在 UniMixing block 之后，与 TokenMixer-Large 类似，我们引入 pertoken SwiGLU 来建模不同 tokens 之间的 feature heterogeneity。对于每个 input token ${\overrightarrow{\mathbf{x}}}_i$$\mathbf{\vec x}_{i}$，SwiGLU 公式如下：

   $$\text{pSwiGLU}\left({\overrightarrow{\mathbf{o}}}_i\right)={\mathbf{W}}_{\text{down}}^i(({\mathbf{W}}_{\text{up}}^i{\overrightarrow{\mathbf{o}}}_i+{\overrightarrow{\mathbf{b}}}_{\text{up}}^i)\odot \text{Swish}({\mathbf{W}}_{\text{gate}}^i{\overrightarrow{\mathbf{o}}}_i+{\overrightarrow{\mathbf{b}}}_{\text{gate}}^i))+{\overrightarrow{\mathbf{b}}}_{\text{down}}^i$$

   其中：

   • ${\mathbf{W}}_{\text{up}}^i,{\mathbf{W}}_{\text{gate}}^i\in {\mathbb{R}}^{B\times (nB)}$$\mathbf W_\text{up}^i, \mathbf W_\text{gate}^i\in \mathbb R^{B\times (nB)}$，${\mathbf{W}}_{\text{down}}^i\in {\mathbb{R}}^{(nB)\times B}$$\mathbf W_\text{down}^i\in \mathbb R^{(nB)\times B}$，${\overrightarrow{\mathbf{b}}}_{\text{up}}^i,{\overrightarrow{\mathbf{b}}}_{\text{gate}}^i\in {\mathbb{R}}^{(nB)}$$\mathbf{\vec b}_\text{up}^i, \mathbf{\vec b}_\text{gate}^i\in \mathbb R^{(nB)}$，${\overrightarrow{\mathbf{b}}}_{\text{down}}^i\in {\mathbb{R}}^B$$\mathbf{\vec b}_\text{down}^i\in \mathbb R^B$ ，$n$$n$ 为一个超参数。

   • ${\overrightarrow{\mathbf{o}}}_i$$\mathbf{\vec o}_i$ 为第 $i$$i$ 个 token 的 UniMixing output。

1.3.3 SiameseNorm

1. 当前 RankMixer 架构缺乏针对 deep architectures 的专门设计，主要体现在 model depth 的 scaling 效果有限。尽管TokenMixer-Large 尝试通过在 block 内加入 interval residuals 与 auxiliary loss 来解决该问题，但未触及根本。为实现model depth 增加时的训练稳定性与性能提升，我们将孪生归一化（ SiameseNorm ）引入 UniMixer 架构，如 Figure 2 所示。如相关工作所述（《Siamesenorm: Breaking the barrier to reconciling pre/post-norm》），SiameseNorm 通过在每层引入两个耦合流（coupled streams），解决预归一化（Pre-Norm）与后归一化（Post-Norm）的矛盾。本节将这两个耦合流记为 ${\bar{\mathbf{X}}}_i$$\overline {\mathbf X}_{i}$ 与 ${\bar{\mathbf{Y}}}_i$$\overline {\mathbf Y}_{i}$，由 input embeddings 初始化 ${\bar{\mathbf{X}}}_0={\bar{\mathbf{Y}}}_0=\mathbf{X}$$\overline {\mathbf X}_{0}=\overline {\mathbf Y}_{0}=\mathbf X$。对于第 $\ell$$\ell$ 个 block，SiameseNorm 执行如下更新：

   $$\begin{array}{c}{\tilde{\mathbf{Y}}}_{\ell }=\text{RMSNorm}\left({\bar{\mathbf{Y}}}_{\ell }\right),{\mathbf{O}}_{\ell }=\text{UniMixer}({\bar{\mathbf{X}}}_{\ell }+{\tilde{\mathbf{Y}}}_{\ell })\\ {\bar{\mathbf{X}}}_{\ell +1}=\text{RMSNorm}({\bar{\mathbf{X}}}_{\ell }+{\mathbf{O}}_{\ell }),{\bar{\mathbf{Y}}}_{\ell +1}={\bar{\mathbf{Y}}}_{\ell }+{\mathbf{O}}_{\ell }\end{array}$$

   ${\bar{\mathbf{Y}}}_i$$\overline{\mathbf Y}_i$ 执行的是 pre-norm，${\bar{\mathbf{X}}}_i$$\overline{\mathbf X}_i$ 执行的是 post-norm。

   对于第 $M$$M$ 个 UniMixer block，我们融合 ${\bar{\mathbf{X}}}_{\ell }$$\overline{\mathbf X}_{\ell}$ 与 ${\bar{\mathbf{Y}}}_{\ell }$$\overline{\mathbf Y}_{\ell}$ 从而生成 final representation，公式如下：

   $${\mathbf{X}}_{\text{output}}={\bar{\mathbf{X}}}_M+\text{RMSNorm}\left({\bar{\mathbf{Y}}}_M\right)$$

   根据离线实验表面，PostNorm -> SiameseNorm 能带来 0.027% 的 AUC 提升。因此，它并不是核心设计。

1.3.4 UniMixer 训练策略

1. 为使parameter matrices ${\mathbf{W}}_G$$\mathbf W_{G}$ 与 ${\mathbf{W}}_B^i$$\mathbf W_{B}^{i}$ 具备稀疏性，我们引入温度系数来控制其稀疏程度。但温度越小，权重越稀疏，同时会导致梯度稀疏、微弱甚至不稳定，使训练困难，optimization 陷入局部最优。另一方面，我们的实验表明 weight parameters 的稀疏性对模型性能有显著正向影响，如 Table 3 所示。因此，该稀疏性不可或缺。

   常用方法是在训练过程中采用线性温度退火（linear temperature annealing）：从较高初始温度（如 $\tau =1.0$$\tau=1.0$）开始，随 training iterations 次数增加，线性退火至 $\tau =0.05$$\tau = 0.05$，公式如下：

   $${\tau }_j=max\{{\tau }_{\text{start}}-\frac{j}{J}\times ({\tau }_{\text{start}}-{\tau }_{\text{end}}),{\tau }_{\text{end}}\}$$

   其中：

   • ${\tau }_j$$\tau_j$ 为第 $j$$j$ 次迭代的温度系数，${\tau }_{\text{start}}$$\tau_\text{start}$、${\tau }_{\text{end}}$$\tau_\text{end}$ 分别为初始温度与最终温度。

   • $J$$J$ 为温度退火的迭代范围。

   数据量不足时，线性退火可能导致高温阶段探索不充分，或低温阶段优化效果不佳。为解决该问题，可先用较高温度系数（如$\tau =1.0$$\tau=1.0$）冷启动训练模型；模型训练充分后，降低温度系数（如 $\tau =0.05$$\tau=0.05$），以高温模型权重为 initialization，重新训练低温模型。

   根据实验部分的描述，训练策略对 UniMixer 模型性能的影响最大。众所周知，dnn 模型的训练策略、模型架构都对最终模型性能产生重大影响。那么，UniMixer 的优秀性能是来自于它的训练策略，还是来自于它的模型架构？

1.4 实验

1. 本节开展大量实验，对比所提出的 UniMixer 架构与 SOTA 方法的性能，并回答以下问题：

   • Q1：UniMixer 架构的 scaling 效率是否优于 SOTA 架构？

   • Q2：所提出的方法在 different settings of global and local mixing pattern 下性能如何变化？

   • Q3：轻量级模块 UniMixing-Lite 是否进一步提升 scaling 效率？

   • Q4：部署至真实在线系统后，UniMixer/UniMixing-Lite 是否在 A/B 测试中提升业务指标？

2. 数据集与评估指标：我们采用 Kuaishou 的广告投放场景的真实训练数据集日志，建模用户留存（user retention），开展离线与在线评估。数据集包含一年收集的超 0.7 billion 的用户样本，涵盖数值特征、ID 特征、交叉特征、序列特征等数百个 heterogeneous features。二元标签（User Retention = 1/0 ）表示 users’ first activation 的次日是否返回 Kuaishou application。推荐模型的 scaling evaluation 的指标采用行业常用的 AUC 、UAUC（User-Level AUC）以评估模型性能，用dense parameter 数量、FLOPs、MFU 来评估模型效率。

3. 基线与实验细节：将本文的 2-blocks/4-blocks UniMixer/UniMixing-Lite 架构与以下代表性 SOTA 框架对比，按建模范式分类：

   • Attention-Based 架构：Heterogeneous Attention（《Hiformer: Heterogeneous feature interactions learning with transformers for recommender systems》）、HiFormer（《Hiformer: Heterogeneous feature interactions learning with transformers for recommender systems》）、FAT（《From scaling to structured expressivity: Rethinking transformers for ctr prediction》），采用 field-specific query, key, and value projections 实现heterogeneous feature interaction。

   • TokenMixer-Based 框架：RankMixer（《Rankmixer: Scaling up ranking models in industrial recommenders》）、TokenMixer-Large（《Tokenmixer-large: Scaling up large ranking models in industrial recommenders》），采用 rule-based token mixing operation 实现 feature interaction。

   • FM-Based 框架：Wukong（《Wukong: Towards a scaling law for large-scale recommendation》），将 outputs of a FMB and a linear projection layer 拼接起来，从而提升 interaction component。

   所有实验在由 40 GPUs 组成的 a hybrid distributed training framework 下进行。所有模型使用一致的 optimizer hyperparameters：dense and sparse parts 均用 Adam 优化，学习率设为 0.001。

1.4.1 Performance Comparison (for Q1)

1. 采用参数约 100 million 的 SOTA scaling architectures 与 UniMixer、UniMixing-Lite 对比，探索其 scaling laws。heterogeneous attention 架构作为 base model。Table 2 展示本文模型与 SOTA 模型的主要性能结果。可以看到：

   • 在更小的参数预算与计算成本下，UniMixer 与 UniMixing-Lite 架构在多个指标上均显著优于其他 SOTA 模型。

     TokenMixer-Large 的效果还不如 RankMixer？有点奇怪。

   

   此外，该广告投放场景中，除 UniMixer/UniMixing-Lite 外，RankMixer 性能优于所有其他 SOTA 模型。因此，我们选择该最强的 SOTA 模型与 UniMixer/UniMixing-Lite 进行 scaling laws 的对比。所有模型在相同数据集、一致超参数下进行训练，其参数量与 FLOPs 的 scaling curves 如 Figure 4 所示。可以看到：

   • 随 number of parameters/FLOPs 增加，三个模型的 AUC 均呈现清晰的 power-law 趋势。

   • UniMixer-Lite 实现最优的 scaling 效率，提升斜率更陡。

   根据 Figure 4，RankMixer、UniMixer、UniMixing-Lite 的 AUC 与 Parameters/FLOPs 的良好的 scaling laws 公式如下：

   $$\begin{array}{c}\mathrm{\Delta }{\text{AUC}}_{\text{RankMixer}}=0.002718\times {\text{Params}}^{0.116043},\mathrm{\Delta }{\text{AUC}}_{\text{RankMixer}}=0.002022\times {\text{Params}}^{0.116635}\\ \mathrm{\Delta }{\text{AUC}}_{\text{UniMixer}}=0.003032\times {\text{Params}}^{0.131973},\mathrm{\Delta }{\text{AUC}}_{\text{UniMixer}}=0.002058\times {\text{Params}}^{0.125702}\\ \mathrm{\Delta }{\text{AUC}}_{\text{UniMixer-Lite}}=0.003767\times {\text{Params}}^{0.141903},\mathrm{\Delta }{\text{AUC}}_{\text{UniMixer-Lite}}=0.002338\times {\text{Params}}^{0.135327}\end{array}$$

   scaling laws 中的两个常数里，scaling exponent 常数对性能增长影响最大，是 scaling 效率的主导因素。UniMixer-Lite 展现出最强的 scaling 效率，在参数量与 FLOPs 上均取得最大的 scaling exponent 与 scaling coefficient，说明其从 increased model capacity 中获益最大。

   

1.4.2 Ablation Studies (for Q2)

1. 为探索 global and local mixing weights 的特性，以及 UniMixer 中各模块对 AUC gains 的贡献，我们对多种 UniMixer 变体开展消融实验，测量其相对 full UniMixer model 的 AUC 变化。所有变体在相似设置下训练。结果如 Table 3 所示，移除任意模块或违反参数约束（parameter constraints）均会导致性能下降，其中 low temperature coefficient 与 model warm-up 对整体性能影响最显著。

   model warm-up 就是正文章节提到的线性温度退火.

   

1.4.3 Performance of the UniMixing-Lite Module (for Q3)

1. 根据 Figure 4 的 scaling 趋势可见，本文 UniMixing-Lite 架构具备最优的参数效率与计算效率。本节开展实验，探究基矩阵集合 ${\left\{{\mathbf{Z}}_{\ell }\right\}}_{\ell =1}^b$$\left\{\mathbf Z_{\ell}\right\}_{\ell=1}^{b}$ 的不同数量 $b$$b$、不同的秩 $r$$r$（${\mathbf{A}}_G,{\mathbf{B}}_G$$\mathbf A_G, \mathbf B_G$ 的秩），以及不同 UniMixer block 数量的影响。如 Table 4 所示：

   • 随基矩阵数量 $b$$b$ 与 ${\mathbf{A}}_G,{\mathbf{B}}_G$$\mathbf A_{G}, \mathbf B_{G}$ 的秩 $r$$r$ 的增加，模型性能相应提升。

   • 但在参数效率方面，增加基矩阵数量 $b$$b$ 比增加秩 $r$$r$ 带来更高的 AUC gain。

   注意：RankMixer 的层数越多，效果反而下降了。

   

2. 为观察低秩近似 ${\mathbf{A}}_G{\mathbf{B}}_G$$\mathbf A_{G}\mathbf B_{G}$ 与带 Sinkhorn–Knopp 操作的基矩阵 ${\left\{{\mathbf{Z}}_{\ell }\right\}}_{\ell =1}^b$$\left\{\mathbf Z_{\ell}\right\}_{\ell=1}^{b}$ 对重构 global and local mixing matrices 的影响，在2-blocks-UniMixer-Lite 架构中，我们可视化温度系数 $\tau =1$$\tau=1$ 与 $\tau =0.05$$\tau=0.05$ 时，第一个 UniMixer block 的 reconstructed global matrix ${\bar{\mathbf{W}}}_G$$\overline{\mathbf W}_{G}$ 与前六个 local mixing matrices ${\bar{\mathbf{W}}}_B^i$$\overline {\mathbf W}_{B}^{i}$，如 Figure 5 所示。input embedding 维度为 768，block 大小为 B = 6，因此 ${\bar{\mathbf{W}}}_G\in {\mathbb{R}}^{128\times 128}$$\overline {\mathbf W}_{G} \in \mathbb{R}^{128 \times 128}$，${\bar{\mathbf{W}}}_B^i\in {\mathbb{R}}^{6\times 6}$$\overline {\mathbf W}_{B}^{i} \in \mathbb{R}^{6 \times 6}$，其中 ${\mathbf{A}}_G\in {\mathbb{R}}^{128\times 16}$$\mathbf A_{G} \in \mathbb{R}^{128 \times 16}$，${\mathbf{B}}_G\in {\mathbb{R}}^{16\times 128}$$\mathbf B_{G} \in \mathbb{R}^{16 \times 128}$。

   由 Figure 5 可见，尽管模块中使用了低秩近似与 basis matrices，Sinkhorn–Knopp 操作仍能保证矩阵接近满秩。此外，对比Figure 5(a)(b) and (c)(d)，更低的温度系数下的 global and local mixing matrices，interaction distributions 比更高的温度系数下更尖锐。由消融实验结果可知，${\bar{\mathbf{W}}}_G$$\overline {\mathbf W}_{G}$ 与 ${\bar{\mathbf{W}}}_B^i$$\overline {\mathbf W}_{B}^{i}$ 的稀疏性使模型性能大幅提升。

   

3. 另一方面，由 Table 4 可见，随 UniMixer 的深度的增加，所提出的模型持续呈现清晰的 scaling-up 趋势，而 RankMixer 随深度的增加出现性能下降。UniMixing-Lite with 2 blocks and 4 blocks 的 scaling curves 如 Figure 6 所示，说明沿 depth 缩放比沿 width 缩放更高效。

   

1.4.4 Online A/B Test Results (for Q4)

1. 为验证所提出的 UniMixer 架构的在线性能，我们将 UniMixer 与 UniMixing-Lite 部署至 Kuaishou 的多个广告投放场景。在线A/B 测试中，以 30 天观察窗口内的累计活跃天数（Cumulative Active Days: CAD）衡量 user engagement（排除 installation day，即 day 0）。在多个场景中，D1-D30 的 CAD 平均提升超 15%。

   没有详细的图表来说明？base model 是啥？有没有上线？这些都没讲。

1.5 结论

1. 本文为推荐系统的 scaling laws 建立了一个统一的 scaling 框架，打通了 attention-based 方法、TokenMixer-based 方法与FM-based 方法的联系，使融合各自优势成为可能。从得到的 scaling laws 可见，与 SOTA 架构相比，本文的 UniMixing-Lite 实现了最优的参数效率与计算效率。我们已将该架构部署至 Kuaishou 的多个场景，取得显著的离线与在线收益。

   本工作不再孤立看待推荐系统中现有的 scaling 模块（如 Heterogeneous Attention、TokenMixer、Wukong），而是建立统一理论框架，为推荐系统的 scaling design 提供指导。我们相信该统一架构能助力推荐系统领域迎来属于自己的 "attention moment"。这个统一模块，UniMixer，是专为推荐领域设计的 fundamental block，其适用性可进一步扩展至 user behavior sequence modeling 任务与 generative recommendation 任务。

二、附录

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

1. 给定如下 input hidden state $\mathbf{X}\in {\mathbb{R}}^{2\times 6}$$\mathbf X \in \mathbb{R}^{2 \times 6}$，其中 $x_i$$x_i$ 为标量：

   $$\begin{array}{c}\mathbf{X}=\left[\begin{array}{cccccc}{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{array}\right]\end{array}$$

   input hidden state $\mathbf{X}$$\mathbf X$ 经过 TokenMixer 操作后被变换为：

   $$\begin{array}{c}\text{TokenMixer}(\mathbf{X})=\left[\begin{array}{cccccc}{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{array}\right]\end{array}$$

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

   $$\text{flatten}(\text{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}]}^{\mathrm{\top }}$$

   另一方面，向量 $\text{flatten}(\mathbf{X})$$\text{flatten}(\mathbf X)$ 可以通过乘以一个 ${\mathbb{R}}^{12\times 12}$$\mathbb R^{12\times 12}$ 的矩阵被变换为 $\text{flatten}(\text{TokenMixer}(\mathbf{X}))$$\text{flatten}(\text{TokenMixer}(\mathbf X))$，可写作：

   $$\begin{array}{c}\underset{{\mathbf{W}}^{\text{perm}}}{\underbrace{\left[\begin{array}{ccccccccccccccc}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{array}\right]}}\underset{\text{flatten}(\mathbf{X})}{\underbrace{\left[\begin{array}{c}{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{array}\right]}}=\underset{\text{flatten}(\text{TokenMixer}(\mathbf{X}))}{\underbrace{\left[\begin{array}{c}{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{array}\right]}}\end{array}$$

   根据上式，本数值例子中的 TokenMixer 操作被等价变换为矩阵乘法形式。此外，置换矩阵 ${\mathbf{W}}^{\text{perm}}\in {\mathbb{R}}^{12\times 12}$$\mathbf W^{\text{perm}} \in \mathbb{R}^{12 \times 12}$ 可以被等价分解为下面两个小矩阵的克罗内克积（Kronecker product）：

   $$\begin{array}{c}{\mathbf{W}}^{\text{perm}}=\underset{\text{Global Mixing Matrix}}{\underbrace{\left[\begin{array}{cccc}1& 0& 0& 0\\ 0& 0& 1& 0\\ 0& 1& 0& 0\\ 0& 0& 0& 1\end{array}\right]}}\otimes \underset{\text{Local Mixing Matrix}}{\underbrace{\left[\begin{array}{ccc}1& 0& 0\\ 0& 1& 0\\ 0& 0& 1\end{array}\right]}}\end{array}$$

   .

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

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

   $$\begin{array}{c}{\mathbf{W}}_G=\left[\begin{array}{ccc}{w}_{(1,1)}^G& \cdots & w_{(1,L//B)}^G\\ ⋮& \ddots & ⋮\\ w_{(L//B,1)}^G& \cdots & w_{(L//B,L//B)}^G\end{array}\right],{\mathbf{W}}_B^i=\left[\begin{array}{ccc}{v}_{(1,1)}^i& \cdots & v_{(1,B)}^i\\ ⋮& \ddots & ⋮\\ v_{(B,1)}^i& \cdots & v_{(B,B)}^i\end{array}\right]\end{array}$$

   其中 $w_{(i,j)}$$w_{(i,j)}$ 和 $v_{(i,j)}$$v_{(i,j)}$ 均为标量。

   根据公式 $[{\overrightarrow{\mathbf{x}}}_1\mid {\overrightarrow{\mathbf{x}}}_2\mid \cdots \mid {\overrightarrow{\mathbf{x}}}_{L//B}]=\text{Split}(\text{flatten}(\mathbf{X}),\frac{L}{B})$$\left[\mathbf{\vec x}_1\mid \mathbf{\vec x}_2\mid \cdots \mid \mathbf{\vec x}_{L//B}\right] = \text{Split}\left(\text{flatten}(\mathbf X), \frac{L}{B}\right)$，$\text{flatten}(\mathbf{X})$$\text{flatten}(\mathbf X)$ 被均匀切分为 $L//B$$L//B$ 个向量，重写为：

   $$\text{flatten}(\mathbf{X})={[{\overrightarrow{\mathbf{x}}}_1\mid {\overrightarrow{\mathbf{x}}}_2\mid \cdots \mid {\overrightarrow{\mathbf{x}}}_{L//B}]}^{\mathrm{\top }}$$

   其中：${\overrightarrow{\mathbf{x}}}_i\in {\mathbb{R}}^{1\times B}$$\mathbf{\vec x}_i\in \mathbb R^{1\times B}$ 是维度为 $B$$B$ 的行向量。

   根据 UniMixing 的原始表达式 $\text{UniMixing}(\mathbf{X})=\text{reshape}(({\mathbf{W}}_G\otimes {\left\{{\mathbf{W}}_B^i\right\}}_{i=1}^{L//B})\text{flatten}(\mathbf{X}),1,L)$$\text{UniMixing}(\mathbf X) = \text{reshape}\left(\left(\mathbf W_G\otimes\left\{\mathbf W_B^i\right\}_{i=1}^{L//B}\right)\text{flatten}(\mathbf X),1,L\right)$，项 $({\mathbf{W}}_G\otimes {\left\{{\mathbf{W}}_B^i\right\}}_{i=1}^{L//B})\text{flatten}(\mathbf{X})$$\left(\mathbf W_G\otimes\left\{\mathbf W_B^i\right\}_{i=1}^{L//B}\right)\text{flatten}(\mathbf X)$ 可以被重写为：

   $$\begin{array}{c}({\mathbf{W}}_G\otimes {\left\{{\mathbf{W}}_B^i\right\}}_{i=1}^{L//B})\text{flatten}(\mathbf{X})=\left[\begin{array}{ccc}{w}_{(1,1)}^G{\mathbf{W}}_B^1& \cdots & w_{(1,L//B)}^G{\mathbf{W}}_B^{L//B}\\ ⋮& \ddots & ⋮\\ w_{(L//B,1)}^G{\mathbf{W}}_B^1& \cdots & w_{(L//B,L//B)}^G{\mathbf{W}}_B^{L//B}\end{array}\right]\left[\begin{array}{c}{\overrightarrow{\mathbf{x}}}_1^{\mathrm{\top }}\\ ⋮\\ {\overrightarrow{\mathbf{x}}}_{L/B}^{\mathrm{\top }}\end{array}\right]\\ =\left[\begin{array}{c}{w}_{(1,1)}^G{\mathbf{W}}_B^1{\overrightarrow{\mathbf{x}}}_1^{\mathrm{\top }}+\cdots +w_{(1,L//B)}^G{\mathbf{W}}_B^{L//B}{\overrightarrow{\mathbf{x}}}_{L/B}^{\mathrm{\top }}\\ ⋮\\ w_{(L//B,1)}^G{\mathbf{W}}_B^1{\overrightarrow{\mathbf{x}}}_1^{\mathrm{\top }}+\cdots +w_{(L//B,L//B)}^G{\mathbf{W}}_B^{L//B}{\overrightarrow{\mathbf{x}}}_{L//B}^{\mathrm{\top }}\end{array}\right]\in {\mathbb{R}}^{L\times 1}\end{array}$$

   另一方面，我们可以得到下面的表达式：

   $$\begin{array}{c}{\mathbf{W}}_G\text{reshape}([{\overrightarrow{\mathbf{x}}}_1{\left({\mathbf{W}}_B^1\right)}^{\mathrm{\top }}\mid \cdots \mid {\overrightarrow{\mathbf{x}}}_{L//B}{\left({\mathbf{W}}_B^{L//B}\right)}^{\mathrm{\top }}],\frac{L}{B},B)\\ =\left[\begin{array}{ccc}{w}_{(1,1)}^G& \cdots & w_{(1,L//B)}^G\\ ⋮& \ddots & ⋮\\ w_{(L//B,1)}^G& \cdots & w_{(L//B,L//B)}^G\end{array}\right]\left[\begin{array}{c}{\overrightarrow{\mathbf{x}}}_1{\left({\mathbf{W}}_B^1\right)}^{\mathrm{\top }}\\ ⋮\\ {\overrightarrow{\mathbf{x}}}_{L//B}{\left({\mathbf{W}}_B^{L//B}\right)}^{\mathrm{\top }}\end{array}\right]\\ =\left[\begin{array}{c}{w}_{(1,1)}^G{\overrightarrow{\mathbf{x}}}_1{\left({\mathbf{W}}_B^1\right)}^{\mathrm{\top }}+\cdots +w_{(1,L//B)}^G{\overrightarrow{\mathbf{x}}}_{L//B}{\left({\mathbf{W}}_B^{L//B}\right)}^{\mathrm{\top }}\\ ⋮\\ w_{(L//B,1)}^G{\overrightarrow{\mathbf{x}}}_1{\left({\mathbf{W}}_B^1\right)}^{\mathrm{\top }}+\cdots +w_{(L//B,L//B)}^G{\overrightarrow{\mathbf{x}}}_{L//B}{\left({\mathbf{W}}_B^{L//B}\right)}^{\mathrm{\top }}\end{array}\right]\in {\mathbb{R}}^{(L//B)\times B}\end{array}$$

   上式中的元素满足：

   $$w_{(i,1)}^G{\mathbf{W}}_B^1{\overrightarrow{\mathbf{x}}}_1^{\mathrm{\top }}+\cdots +w_{(i,L//B)}^G{\mathbf{W}}_B^{L//B}{\overrightarrow{\mathbf{x}}}_{L//B}^{\mathrm{\top }}={(w_{(i,1)}^G{\overrightarrow{\mathbf{x}}}_1{\left({\mathbf{W}}_B^1\right)}^{\mathrm{\top }}+\cdots +w_{(i,L//B)}^G{\overrightarrow{\mathbf{x}}}_{L//B}{\left({\mathbf{W}}_B^{L//B}\right)}^{\mathrm{\top }})}^{\mathrm{\top }}$$

   因此有：

   $$({\mathbf{W}}_G\otimes {\left\{{\mathbf{W}}_B^i\right\}}_{i=1}^{L//B})\text{flatten}(\mathbf{X})=\text{reshape}({\mathbf{W}}_G\text{reshape}([{\overrightarrow{\mathbf{x}}}_1{\left({\mathbf{W}}_B^1\right)}^{\mathrm{\top }}\mid \cdots \mid {\overrightarrow{\mathbf{x}}}_{L//B}{\left({\mathbf{W}}_B^{L//B}\right)}^{\mathrm{\top }}],\frac{L}{B},B),L,1)$$

   由于 ${\mathbf{W}}_B^i$$\mathbf W_B^i$ 和 ${\left({\mathbf{W}}_B^i\right)}^{\mathrm{\top }}$$\left(\mathbf W_B^{i}\right)^\top$ 都是可学习参数，参数的转置不会影响模型。因此，computation pipeline optimization 后的 UniMixing 模块可以写作：

   $$\text{UniMixing}(\mathbf{X})=\text{reshape}({\mathbf{W}}_G\text{reshape}([{\overrightarrow{\mathbf{x}}}_1{\left({\mathbf{W}}_B^1\right)}^{\mathrm{\top }}\mid \cdots \mid {\overrightarrow{\mathbf{x}}}_{L//B}{\left({\mathbf{W}}_B^{L//B}\right)}^{\mathrm{\top }}],\frac{L}{B},B),1,L)$$

   .