UniMixer：推荐系统缩放定律的统一架构 UniMixer: A Unified Architecture for Scaling Laws in Recommendation Systems
提出统一推荐系统三大主流架构的缩放框架UniMixer，提升参数和计算效率。
前置知识
缩放定律
缩放定律描述模型性能与模型规模（参数数量、FLOPs）或数据规模之间的幂律关系，即在一定范围内性能会随着模型规模的增长而稳定提升，这种关系通常呈现为幂律形式，其中指数决定了性能增长的效率。
本文核心目标就是在推荐系统中建立类似LLMs的缩放定律，理解这一概念是理解本文动机和实验设计的基础。
TokenMixer
TokenMixer是一种基于规则的非参数化特征混合操作，它通过将输入token均匀分割成多个头，然后按照固定规则重新排列这些分割后的特征来实现特征交互。该操作避免了注意力机制中计算两个异构语义空间内积相似度的开销，通过静态的token混洗实现特征交互，但缺乏可学习性和场景适应性。
TokenMixer是本文Unified方法的核心起点，理解其工作机制和限制对理解本文的参数化创新至关重要。
Sinkhorn-Knopp迭代
Sinkhorn-Knopp迭代是一种将任意正矩阵转换为双随机矩阵（每行每列和都为1）的算法，通过交替对行和列进行归一化来实现。具体过程为：先通过指数操作使矩阵元素为正，然后交替进行行归一化和列归一化，直到收敛。这样可以保证学到的参数矩阵满足双随机性约束。
本文使用Sinkhorn-Knopp迭代来约束UniMixer中学习的参数矩阵，确保它们满足双随机性，这是保持特征交互模式合理性的关键约束。
Kronecker积
Kronecker积是一种矩阵运算，对于两个矩阵，它们的Kronecker积是一个更大的矩阵。这种操作具有可压缩性，可以用来表示具有特定结构的复杂矩阵。在本文中，TokenMixer的置换矩阵被分解为两个小矩阵的Kronecker积，从而显著减少参数数量。
Kronecker积分解是本文参数化TokenMixer操作的核心技术手段，利用这一性质可以将原来很大的参数量显著减少，是实现高效可学习混合的基础。
异构特征交互
异构特征交互是指推荐系统中来自不同语义空间的特征（如用户画像、物品特征、行为序列等）之间的交互模式。与NLP中所有token共享统一嵌入空间不同，推荐系统的特征空间本质上是异构的，不同特征域具有不同的语义表示和分布特性，因此需要专门的交互机制来处理这种异构性。
异构特征交互问题是推荐系统与NLP领域的根本区别，也是三种主流架构需要解决的核心挑战，理解这一点是理解本文统一框架的前提。
研究动机
现有推荐系统的缩放架构存在三大范式，分别是基于注意力的方法（如HiFormer、FAT、HHFT）、基于TokenMixer的方法（如RankMixer、TokenMixer-Large）和基于分解机的方法（如Wukong、Kunlun）。这些架构在设计理念和结构上存在根本性差异，各自基于完全不同的基础模块，虽然都展示了缩放能力，但缺乏统一的理论框架。更重要的是，基于注意力的方法面临二次方计算复杂度的问题，TokenMixer方法虽然避免了注意力计算但基于固定规则缺乏可学习性且要求T等于H的限制，而基于FM的方法受限于低阶交互能力，无法像LLMs那样展示强缩放定律。以Kuaishou的广告投放场景为例，超过7亿用户样本的规模下，需要平衡性能提升和计算成本，现有方法在缩放效率上仍有改进空间。
本文的目标是本文的目标是建立一个统一的推荐系统缩放理论框架，能够桥接三种主流缩放模块（注意力、TokenMixer、FM）的差异和联系，同时结合它们的优势。具体来说，通过将基于规则的TokenMixer参数化，构造一个广义的参数化特征混合模块，使token混合模式能够在模型训练过程中被优化和学习。此外，设计轻量级模块UniMixing-Lite，在显著改善模型性能的同时进一步压缩模型参数和计算成本，提升缩放ROI（投资回报率）。最终希望这一统一架构能够帮助推荐系统社区实现类似attention moment的突破。
与已有工作不同的是，本文的独特切入角度在于深入分析TokenMixer操作的数学本质，发现其可以等价为置换矩阵与扁平化嵌入向量的乘积，进而利用置换矩阵的可压缩性、双随机性、稀疏性和对称性等性质进行参数化。与现有工作不同，本文不是简单地选择某一种架构或进行组合，而是建立一个统一的理论框架，在数学上证明注意力机制、TokenMixer和FM可以被视为同一框架下的特殊形式。这种从数学本质出发的统一视角，使得设计具有多种架构优势的混合模块成为可能。
核心方法
UniMixer的整体思路是从数学角度统一三种主流推荐系统缩放架构，首先通过分析TokenMixer操作的本质，发现它等价于一个置换矩阵变换，然后利用置换矩阵的结构性质进行参数化，使其可学习。在此基础上，设计包含全局混合和局部混合的统一token混合模块UniMixing，进一步通过计算管道优化降低复杂度。为了兼顾效率和性能，还设计了轻量级版本UniMixing-Lite，利用基矩阵组合和低秩近似进一步压缩参数。整体架构包括特征分词化、M个UniMixer块（使用SiameseNorm）、Per-token SwiGLU和任务塔，通过这个统一的框架同时获得注意力机制的可学习性、TokenMixer的计算效率以及FM方法的显式交互建模能力。
核心创新点是对规则基础的TokenMixer操作进行等价参数化。具体发现是TokenMixer操作可以表示为置换矩阵与扁平化嵌入向量的乘积，其中置换矩阵是一个大型置换矩阵。关键洞察是利用置换矩阵的可压缩性：置换矩阵可以分解为两个小矩阵的Kronecker积，其中一个是单位矩阵，另一个是较小的方阵。通过参数化这两个矩阵，将参数量从原来的平方级别显著减少。在此基础上，进一步通过优化计算管道，将复杂度从平方级别降低到更优的形式，其中B是块大小。另一个关键创新是引入温度系数控制参数矩阵的稀疏性，并通过Sinkhorn-Knopp迭代确保双随机性约束。这使得UniMixer能够同时利用三种架构的优势：注意力机制的可学习局部交互、TokenMixer的低成本全局混合、以及FM的显式交互模式。
方法步骤详情
方法的具体步骤如下：第一步是特征分词化，将输入特征按语义类别分为多个不相交的特征域（用户画像、物品特征、行为序列、查询特征等），每个特征域通过嵌入层转换为不同维度的嵌入向量，然后拼接成完整的嵌入向量。第二步将嵌入向量均匀分割成适当数量的块，每个块通过token特定线性层投影为token嵌入，得到输入隐藏状态矩阵。第三步是UniMixer块的核心混合操作，将扁平化嵌入均匀分割为多个大小为B的向量，然后块权重与对应块向量相乘得到局部特征交互向量，再通过全局权重进行全局混合：具体操作是先将局部混合结果重塑为特定形状，然后通过全局权重矩阵进行混合，最后重塑为一维。第四步应用约束和正则化，使用对称化操作实现对称性约束，通过温度系数的Sinkhorn-Knopp迭代实现双随机性和稀疏性控制。第五步是残差连接和归一化，将混合结果与输入相加后应用RMSNorm归一化。对于UniMixing-Lite，使用基矩阵组合和低秩近似进一步压缩参数：局部混合权重通过基矩阵的线性组合生成，全局混合权重通过低秩矩阵乘积生成。第六步使用SiameseNorm处理深度架构，引入两个耦合流，在每个块中进行交替更新以解决深度扩展的稳定性问题。最后通过Per-token SwiGLU处理token级别的特征异构性：对每个token的输出分别应用不同的SwiGLU变换，建模token级别的特征差异。
技术新颖性
技术新颖性体现在多个层面。首先，在理论层面，建立了统一的推荐系统缩放框架，证明注意力机制、TokenMixer和FM方法可以统一表示为同一框架下的不同形式：全局混合模式与局部混合模式的组合。当全局混合模式采用不同的形式（注意力、固定置换、内积）时，分别对应三种架构。其次，在技术实现层面，通过Kronecker分解将TokenMixer从固定规则转为可学习参数，并通过计算管道优化避免了大型中间变量的创建。再次，在训练策略层面，引入温度系数的线性退火策略：从高温度开始便于探索，逐渐退火到低温度增加稀疏性。这种温度退火策略平衡了训练稳定性和稀疏性带来的性能提升。最后，在架构设计层面，SiameseNorm的引入解决了深度架构的缩放问题，使得模型可以沿着深度维度高效扩展，而不会像RankMixer那样在增加块数时出现性能下降。
实验结果
核心发现主要体现在四个方面。第一，在用户留存预测任务上，UniMixer和UniMixing-Lite在更小的参数和计算预算下显著超越SOTA模型。以约1亿参数的规模为例，UniMixer-2-Blocks 101.5M达到AUC=0.750238，比最强基线RankMixer提升0.5661%，UAUC=0.739983提升0.6154%。更令人印象深刻的是，轻量级的UniMixing-Lite-4-Blocks 38.2M仅用38.2M参数就达到AUC=0.752327，比RankMixer提升0.7750%，参数量减少了约72%。第二，缩放定律分析显示三种模型都呈现清晰的幂律趋势，但UniMixing-Lite具有最陡峭的改进斜率。具体的缩放公式表明UniMixing-Lite的缩放指数0.141903显著高于RankMixer的0.116043，意味着它从增加模型容量中获益最多。第三，深度扩展实验表明UniMixing-Lite沿着深度缩放比沿着宽度缩放更高效。从2块增加到4块，UniMixing-Lite的AUC从0.749228提升到0.750803（提升0.1575%），而RankMixer从2块增加到4块时AUC从0.747772下降到0.746706（下降0.1066%），这验证了SiameseNorm对深度架构的有效性。第四，在线A/B测试在Kuaishou的多个广告投放场景中部署后，30天观察窗口内的累积活跃天数（CAD）平均增长超过15%，证明了方法的实用价值。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 用户留存预测 | AUC | UniMixer-Lite-4-Blocks 38.2M: 0.752327 | RankMixer: 0.749329 | +0.7750% |
| 用户留存预测 | AUC | UniMixer-2-Blocks 67.5M: 0.749770 | RankMixer: 0.749329 | +0.5193% |
| 用户留存预测 | AUC | UniMixer-Lite-2-Blocks 42.4M: 0.751121 | TokenMixer-Large: 0.748410 | +0.6544% |
| 用户留存预测 | UAUC | UniMixer-Lite-4-Blocks 38.2M: 0.742091 | RankMixer: 0.738938 | +0.8190% |
| 用户留存预测 | FLOPs/Batch效率 | UniMixer-Lite-4-Blocks 38.2M: 0.126T | RankMixer: 0.168T | -25% 计算量，性能更高 |
局限与改进
作者在论文中承认了一些局限性。首先，虽然通过计算管道优化将复杂度降低，但在某些配置下FLOPs仍然相对较高。例如，UniMixer-Lite-2-Blocks 76.2M的FLOPs/Batch达到2.60T，相比RankMixer的1.68T增加了约55%，这可能限制了在极端资源受限场景下的应用。其次，温度系数的选择和退火策略需要调优，作者提到当数据量不足时，线性退火可能导致早期探索不充分或后期优化不足，需要采用两阶段训练策略。第三，Sinkhorn-Knopp迭代增加了计算开销，虽然保证了双随机性约束，但也可能影响训练效率。基于我自己的观察，还有一些局限性值得注意：第一，论文中的实验主要在Kuaishou的广告投放数据集上进行，虽然规模大（超过7亿样本），但属于单一场景，在其他推荐场景（如电商、内容分发）的泛化能力有待验证。第二，虽然UniMixing-Lite在参数效率上表现优异，但基矩阵数量和低秩的选择对性能有影响，增加这些超参数调优的复杂度。第三，论文没有深入讨论不同特征域（用户画像、物品特征、行为序列等）对统一框架的敏感性，以及如何针对不同场景调整块大小。第四，在线A/B测试结果虽然展示了15%以上的CAD提升，但没有提供与基线模型的详细对比，无法量化相对于现有部署模型的实际增益。
独立分析的弱点
独立分析的弱点包括：第一，温度退火策略的两阶段训练虽然解决了线性退火的不足，但需要训练两次，增加了训练成本和工程复杂度，可以考虑自适应温度调整策略，根据梯度稀疏性或损失变化动态调整温度系数。第二，Sinkhorn-Knopp迭代的计算开销随着模型规模增长而增加，可以考虑使用近似算法或松弛约束，如使用正则化项替代硬约束，或者使用GPU优化的并行Sinkhorn算法。第三，基矩阵组合和低秩近似虽然压缩了参数，但可能限制表达能力，特别是对于高度非线性的特征交互模式，可以考虑自适应选择基矩阵数量，或者使用分层基矩阵结构来平衡表达能力和效率。第四，SiameseNorm虽然解决了深度扩展问题，但引入了额外的计算开销和内存占用，可以考虑轻量级变体或只在深层使用。第五，在极端异构特征场景下（特征域数量和类型差异很大），统一的块大小和混合模式可能不是最优的，可以考虑针对不同特征域设计自适应的混合策略。第六，论文缺乏对长尾用户或冷启动场景的专门分析，这些场景下的特征分布可能与整体数据不同，需要特殊的处理策略。
未来方向
作者提出的未来研究方向包括：将统一架构扩展到用户行为序列建模和生成式推荐任务。论文提到，UniMixer作为为推荐领域定制的基础块，其适用性可以进一步扩展。基于论文成果，还有多个值得探索的方向：第一，自动化架构搜索，基于统一框架自动搜索最优的混合模式（全局vs局部）、块大小、深度配置等，以适应不同的应用场景和资源约束。第二，动态混合模式，根据输入样本的特征分布或实时反馈动态调整混合策略，实现更精细化的个性化。第三，多任务学习的扩展，在统一框架下同时建模多个推荐目标（点击、转化、留存、时长等），探索共享混合模式vs任务特定混合模式的权衡。第四，与现有推荐技术的深度融合，如将UniMixer与图神经网络结合用于知识图谱增强的推荐，或与对比学习结合用于更好的表征学习。第五，跨域泛化研究，探索统一框架在不同推荐领域（电商、社交、内容、搜索）的迁移能力和适配策略。第六，可解释性增强，利用双随机性和稀疏性约束分析学习到的特征交互模式，为推荐决策提供更好的解释。第七，绿色AI优化，进一步探索在保持性能的同时降低计算复杂度和能耗的方法，如知识蒸馏、量化、剪枝等。第八，理论分析深化，建立更严格的缩放定律理论框架，解释为什么统一的混合模式能够产生更好的缩放指数。
复现评估
复现评估方面：数据集方面，论文使用的是Kuaishou内部广告投放数据，包含超过7亿用户样本和数百个异构特征，这是私有数据集，外部研究者无法直接获取。虽然论文描述了数据特征（用户画像、物品特征、行为序列、查询特征等），但缺乏公开可用的替代数据集进行验证。算力方面，所有实验在40个GPU的混合分布式训练框架上进行，训练大规模模型（如UniMixer-Lite-8-Blocks 19.2M参数）需要显著的计算资源，这对复现提出了较高要求。代码方面，论文没有明确说明是否开源代码，考虑到是工业界研究成果（Kuaishou Technology），代码可能不会完全公开。实验细节方面，论文提供了较为详细的超参数设置（学习率0.001、Adam优化器、温度范围等）、评估指标（AUC、UAUC、参数量、FLOPs、MFU）和消融研究配置，这有助于复现部分实验。然而，一些关键实现细节（如Sinkhorn-Knopp的具体迭代次数、基矩阵的初始化方法、SiameseNorm的具体实现）可能仍需补充。总体而言，复现难度为中等到偏高，主要障碍是私有数据集和可能的代码闭源，但详细的实验描述使得方法实现和部分验证是可行的。建议作者在未来工作中开源核心代码或提供至少在公开数据集上的验证结果。
论文图表