Source-linked AI summary
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
TL;DR
Recommendation scaling research lacks a unified treatment of attention-based, TokenMixer-based, and factorization-machine-based architectures. UniMixer parameterizes TokenMixer mixing and unifies these blocks, while UniMixing-Lite compresses parameters and computation; the authors report improved efficiency and gains in offline and online evaluations.
Problem
Existing recommendation scaling architectures use distinct attention-based, TokenMixer-based, and factorization-machine-based designs, motivating a unified scaling framework.
Method
UniMixer parameterizes rule-based TokenMixer operations to learn mixing patterns and establishes a framework connecting attention-based, TokenMixer-based, and factorization-machine-based methods.
Results
UniMixer-Lite achieved the best parameter and computational efficiency among the compared state-of-the-art architectures, while online CAD increased by more than 15% on average across scenarios.
Takeaways & Limitations
The framework provides a unified theoretical basis for scaling design and supports deployment across multiple recommendation scenarios.
Abstract
from arXiv · showhide
In recent years, the scaling laws of recommendation models have attracted increasing attention, which govern the relationship between performance and parameters/FLOPs of recommenders. Currently, there are three mainstream architectures for achieving scaling in recommendation models, namely attention-based, TokenMixer-based, and factorization-machine-based methods, which exhibit fundamental differences in both design philosophy and architectural structure. In this paper, we propose a unified scaling architecture for recommendation systems, namely \textbf{UniMixer}, to improve scaling efficiency and establish a unified theoretical framework that unifies the mainstream scaling blocks. By transforming the rule-based TokenMixer to an equivalent parameterized structure, we construct a generalized parameterized feature mixing module that allows the token mixing patterns to be optimized and learned during model training. Meanwhile, the generalized parameterized token mixing removes the constraint in TokenMixer that requires the number of heads to be equal to the number of tokens. Furthermore, we establish a unified scaling module design framework for recommender systems, which bridges the connections among attention-based, TokenMixer-based, and factorization-machine-based methods. To further boost scaling ROI, a lightweight UniMixing module is designed, \textbf{UniMixing-Lite}, which further compresses the model parameters and computational cost while significantly improve the model performance. The scaling curves are shown in the following figure. Extensive offline and online experiments are conducted to verify the superior scaling abilities of \textbf{UniMixer}.
1 Introduction
Recommendation systems are increasingly adapting scaling-law ideas from large language models by enlarging model size and computational cost. UniMixer unifies major recommendation scaling blocks, while UniMixing-Lite targets more efficient scaling and experiments evaluate their scaling abilities.
- Motivation: Scaling laws relate recommendation performance to model parameters and computational cost as ranking models become more complex.The motivation is inspired by scaling observed in large language models and by stacking scaling modules across multiple layers.
- Motivation: Recommendation models process heterogeneous categorical and dense features whose dynamic embeddings capture information from multiple perspectives.These multi-field user and item features support predictions intended to increase positive engagement with recommendations.
- Contributions: TokenMixer feature-interaction patterns are revealed by constructing an equivalent parameterization of its rule-based operation.This converts the rule-based mixing structure into a form suitable for analysis and further modeling.
- Contributions: UniMixer bridges attention-based, TokenMixer-based, and factorization-machine-based methods while reducing computational complexity and GPU memory consumption during training and inference.The framework is designed as a unified scaling architecture for recommendation systems.
- Contributions: UniMixing-Lite combines advantages of attention-based and TokenMixer-based architectures while reducing model parameters and computational cost.The lightweight module is designed to improve scaling efficiency.
- Evaluation: Extensive offline and online experiments evaluate the scaling abilities of UniMixer.The experiments are presented as evidence for the proposed architecture’s scaling performance.
2 Related Work
Recommendation scaling research has developed attention-based, TokenMixer-based, and factorization-machine-based paradigms. These approaches trade off heterogeneous feature modeling, computational efficiency, interpretability, and interaction order in different ways.
- Overview: Three main recommendation scaling paradigms are attention-based, TokenMixer-based, and factorization-machine-based methods.They are used to establish scaling laws for massive-scale recommendation systems.
- Attention-Based Framework: Attention-based methods adapt Transformers for click-through-rate prediction but must address heterogeneous feature interactions and the mismatch with language-modeling assumptions.Representative approaches include heterogeneous attention, HiFormer, and Field-Aware Transformers.
- TokenMixer-Based Framework: TokenMixer-based methods replace dynamic attention with parameter-free or static token-mixing operations to reduce the cost associated with quadratic attention computation.RankMixer, Lemur, and TokenMixer-Large exemplify this paradigm in industrial recommender systems.
- FM-Based Framework: Factorization-machine methods efficiently and interpretably model low-order pairwise interactions but are intrinsically limited in interaction order.DeepFM, AutoInt, and DCN variants extend this family with neural networks or transformer attention to capture higher-order interactions.
3 Preliminaries
Recommendation tasks are commonly formulated as supervised prediction over heterogeneous categorical and dense features. Their foundational interaction modules include heterogeneous attention, TokenMixer, and factorization-machine-based components, each with distinct computational and modeling properties.
- Task Formulation: Discriminative recommendation tasks such as rating, CTR, and CVR prediction are typically formulated as supervised learning problems.The dataset consists of feature-label pairs for prediction.
- Task Formulation: Recommendation inputs combine categorical and dense features, with labels representing binary classification or regression outcomes.For CTR and CVR, the objective is to predict the click or conversion probability conditioned on the input features.
- Foundational Scaling Blocks: Heterogeneous attention uses field-specific query, key, and value projections to model interactions among heterogeneous feature tokens.Its multi-head outputs are concatenated and linearly projected to align with the input dimension.
- Foundational Scaling Blocks: TokenMixer performs feature interaction through parameter-free, rule-based mixing after evenly splitting each token into heads.The operation requires the number of heads H to equal the number of tokens T, so input and output dimensions remain identical.
- Foundational Scaling Blocks: Wukong-based methods concatenate a factorization-machine block with a linear projection, using a low-rank projection to reduce interaction-matrix memory requirements.The factorization-machine interaction is represented through XX^T and then processed by the factorization-machine block.
- Motivation: The paper seeks a unified structural foundation that integrates existing scaling blocks to increase scaling return on investment.This motivation follows from the heterogeneous feature spaces and the three foundational block families used in recommendation scaling.
4 UniMixer
UniMixer unifies attention-, TokenMixer-, and FM-based recommendation scaling modules through parameterized token mixing and a shared theoretical framework. Its optimized mixing pipeline and UniMixing-Lite target lower computational cost, fewer parameters, and improved scaling efficiency.
- Overview: UniMixer unifies attention-based, TokenMixer-based, and Wukong-based scaling modules under a shared theoretical framework.The architecture uses feature tokenization, stacked UniMixer blocks, SiameseNorm, and Sparse-Pertoken MoE.
- Heterogeneous Feature Interactions: Heterogeneous attention can produce sharp, sparse weights that risk gradient backpropagation and make query and key training difficult.The attention pattern may also become less discriminative with large-scale heterogeneous feature inputs.
- Parameterized Token Mixing: Parameterized TokenMixer makes rule-based token mixing learnable, while removing its requirement that the number of heads equal the number of tokens.The original operation is represented using a large permutation matrix, whose structure supports parameterization.
- Parameterized Token Mixing: The permutation matrix can be compressed through a Kronecker-product decomposition and constrained by doubly stochasticity, sparsity, and symmetry.These properties reduce the parameterization burden while preserving the structural behavior of TokenMixer.
- Efficient Mixing: The optimized UniMixing pipeline reduces computational complexity from O(L^2) to O(L^2/B + LB) and avoids large intermediate variables.The block size B determines the block-based computation used by the optimized pipeline.
- UniMixing-Lite: UniMixing-Lite combines low-parameter global TokenMixer interaction with attention’s local heterogeneous-feature interaction capability.The module is designed to compress model parameters and computational cost while improving scaling efficiency.
5 Experiments
Experiments evaluate UniMixer and UniMixing-Lite against state-of-the-art attention-, TokenMixer-, and FM-based recommenders across offline scaling, ablations, matrix analysis, and online A/B testing. UniMixer-Lite shows the strongest scaling efficiency, while UniMixer benefits from increased depth and deployment improves CAD.
- Experimental setup: Experiments compare 2-block and 4-block UniMixer and UniMixing-Lite models with representative attention-, TokenMixer-, and FM-based SOTA frameworks.The evaluation uses a real-world Kuaishou advertising dataset with over 0.7 billion user samples and hundreds of heterogeneous features.
- Performance comparison: Under smaller parameter budgets and computational costs, UniMixer and UniMixing-Lite significantly outperform other SOTA models across multiple metrics.The comparison uses SOTA scaling architectures with approximately 100 million parameters.
- Scaling laws: UniMixing-Lite achieves the best scaling efficiency and a steeper AUC improvement slope as parameters and FLOPs increase.Its scaling laws are ΔAUC_UniMixer-Lite = 0.003767Params^0.141903 and ΔAUC_UniMixer-Lite = 0.002338FLOPs^0.135327, with the largest exponent and coefficient across both axes.
- Ablation studies: Removing any UniMixer module or violating parameter constraints degrades performance, with the low temperature coefficient and model warm-up having the largest effects.These findings come from ablations on the 6.57M-parameter UniMixer configuration.
- UniMixing-Lite analysis: Increasing basis number and low-rank rank improves performance, but increasing basis number yields a higher AUC gain per parameter than increasing rank.The study also analyzes reconstructed global and local mixing matrices under different temperature coefficients.
- Scaling dimensions: As UniMixer depth increases, performance continues scaling upward, whereas stacked RankMixer blocks degrade performance; scaling depth is more efficient than scaling width.The comparison includes UniMixing-Lite configurations with 2 and 4 blocks.
- Online A/B tests: Across multiple Kuaishou advertising scenarios, online A/B tests increased 30-day cumulative active days by more than 15% on average.CAD measures user engagement over a 30-day observation window excluding the installation day.
6 Conclusions
The paper establishes UniMixer as a unified scaling framework for recommendation systems and reports that UniMixer-Lite achieves the best parameter and computational efficiency among compared architectures.
- UniMixer bridges attention-based, TokenMixer-based, and FM-based methods within a unified scaling framework.
- UniMixer-Lite achieves the best parameter efficiency and computational efficiency compared with state-of-the-art architectures.
- Offline and online deployment across multiple Kuaishou scenarios yields significant gains.
- The framework is intended to guide scaling design and may extend to behavior-sequence modeling and generative recommendation tasks.
A A numerical example of equivalent transformation of TokenMixer
The numerical example shows that TokenMixer reorders token elements and that this operation can be represented as multiplication by a 12 × 12 permutation matrix.
- TokenMixer transforms the input sequence [x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12] into a reordered sequence.
- flatten(TokenMixer(X)) produces [x1, x2, x3, x7, x8, x9, x4, x5, x6, x10, x11, x12]T.
- The reordered flattened vector is obtained from flatten(X) by multiplying a 12 × 12 matrix.
- The permutation matrix W perm can be decomposed into the Kronecker product of smaller local and global mixing matrices.
B The computation pipeline optimization of the UniMixing module
The optimized UniMixing computation rewrites feature mixing using structured parameter matrices applied to evenly split input vectors, with learnable parameters replacing the original formulation's orientation.
- flatten(X) is evenly split into L//B vectors, each represented as a row vector of dimension B.
- The optimized UniMixing module is formulated with learnable parameters, and parameter transposition does not affect the model.