Wukong: Towards a Scaling Law for Large-Scale Recommendation

Buyun Zhang; Liang Luo; Yuxin Chen; Jade Nie; Xi Liu; Daifeng Guo; Yanli Zhao; Shen Li; Yuchen Hao; Yantao Yao; Guna Lakshminarayanan; Ellie Dingqiao Wen; Jongsoo Park; Maxim Naumov; Wenlin Chen

Source-linked AI summary

Wukong: Towards a Scaling Law for Large-Scale Recommendation

Buyun Zhang, Liang Luo, Yuxin Chen, Jade Nie, Xi Liu, Daifeng Guo, Yanli Zhao, Shen Li, Yuchen Hao, Yantao Yao, Guna Lakshminarayanan, Ellie Dingqiao Wen, Jongsoo Park, Maxim Naumov, Wenlin Chen

arXiv:2403.02545v4Mar 4, 2024cs.LGcs.AI

TL;DR

Recommendation models lack the scaling laws seen in large language models because existing upscaling mechanisms are inefficient for increasingly complex datasets. Wukong uses stacked Factorization Machines and a synergistic scaling strategy to capture any-order interactions. It outperforms state-of-the-art models on public and internal datasets while scaling across two orders of magnitude beyond 100 GFLOP/example.

Problem
Recommendation models do not yet exhibit LLM-like scaling laws because existing upscaling mechanisms are inefficient for increasingly complex datasets.
Method
Wukong uses stacked Factorization Machines with MLP transformations and dense scaling to capture any-order feature interactions.
Results
Wukong outperforms state-of-the-art models across six public datasets and maintains quality superiority while scaling across two orders of magnitude beyond 100 GFLOP/example.
Takeaways & Limitations
Wukong establishes a recommendation scaling law and serves as a scalable architecture across model complexities, tasks, and datasets.
Takeaways & Limitations
AutoInt+ and DCNv2 encountered training instability when scaled up, while MaskNet encountered out-of-memory errors that blocked further scaling.

Abstract

from arXiv · show

Scaling laws play an instrumental role in the sustainable improvement in model quality. Unfortunately, recommendation models to date do not exhibit such laws similar to those observed in the domain of large language models, due to the inefficiencies of their upscaling mechanisms. This limitation poses significant challenges in adapting these models to increasingly more complex real-world datasets. In this paper, we propose an effective network architecture based purely on stacked factorization machines, and a synergistic upscaling strategy, collectively dubbed Wukong, to establish a scaling law in the domain of recommendation. Wukong's unique design makes it possible to capture diverse, any-order of interactions simply through taller and wider layers. We conducted extensive evaluations on six public datasets, and our results demonstrate that Wukong consistently outperforms state-of-the-art models quality-wise. Further, we assessed Wukong's scalability on an internal, large-scale dataset. The results show that Wukong retains its superiority in quality over state-of-the-art models, while holding the scaling law across two orders of magnitude in model complexity, extending beyond 100 GFLOP/example, where prior arts fall short.

1. Introduction

Recommendation models have struggled to scale interaction modeling and sustain quality improvements as datasets and computational demands grow. Wukong addresses this gap with stacked Factorization Machines and demonstrates strong quality and scaling results.

Scaling challenge: Existing recommendation scaling mainly expands embedding tables, increasing parameters without improving interaction capacity and diverging from compute-focused hardware progress.This sparse-scaling trend can reach trillions of parameters while leaving complex feature interactions insufficiently addressed.
Motivation: Wukong targets dense scaling to improve quality with dataset size, compute, and parameter budgets through a unified architecture and synergistic strategy.The goal is effective upscaling and downscaling across varying dataset sizes and computational constraints.
Proposed approach: Wukong uses stacked Factorization Machines with MLP transformations to capture any-order feature interactions across successive layers.Each Factorization Machine captures second-order interactions, while transformed outputs become embeddings for later layers.
Empirical results: Wukong consistently outperforms state-of-the-art models across six public datasets in AUC.These evaluations target the low-complexity regime and indicate generalization across varied recommendation tasks and datasets.
Empirical results: On a large-scale internal dataset, Wukong maintains quality advantages while improving continuously across two orders of magnitude in complexity beyond 100 GFLOP/example.Prior models fall short in this regime.

2. Related Work

Related work spans dense interaction architectures and scaling strategies for recommendation models. Prior studies explore feature interactions, sparse scaling, user representation, and LLM-based recommendation, but leave broader scalability questions open.

Dense Interaction Architectures: Dense interaction models use diverse mechanisms, including attention, Factorization Machines, higher-order optimization, feature crossing, bilinear fusion, masking, and outer products.These approaches target interaction capture through architectures such as AutoInt+, DLRM, DeepFM, HOFM, DCNv2, FinalMLP, MaskNet, and xDeepFM.
Scaling up DLRS: Prior scaling work primarily studies sparse scaling, user representations, sequence modeling, or extensions of existing scaling laws.Reported scales include less than 0.1 PF-days of compute for user representation models and less than 0.8B parameters for sequence modeling.
Scaling up DLRS: Research on LLM-powered recommendation questions whether larger models necessarily improve results and calls for evaluations on more diverse datasets.These studies frame the need for additional evidence before general conclusions about scaling in recommendation.
Scaling up DLRS: Existing work identifies limits in prior sparse-scaling approaches, complementing Wukong’s focus on dense interaction scaling.The related literature therefore motivates studying scaling mechanisms beyond embedding-table expansion.

3. Design of Wukong

Wukong combines standardized embeddings with a stacked interaction architecture that captures progressively higher-order feature interactions. Factorization Machines provide explicit interactions, while compression and MLP transformations control representation size and computational cost.

Wukong targets high-order interaction capture and graceful quality scaling across dataset size, GFLOP/example, and parameter budgets.
The Interaction Stack applies parallel Factorization Machine and Linear Compression Blocks, then maps accumulated interactions to predictions through a final MLP.Successive layers capture exponentially higher-order interactions, with residual connections and layer normalization used across layers.
Its embedding layer standardizes outputs to global dimension d, combining categorical and dense inputs into X0 ∈ R^n×d for interaction processing.Less important features can use smaller embeddings that are grouped, concatenated, and projected to dimension d.
Each interaction layer can represent feature interactions from order 1 through 2^i, with FMB providing higher-order combinations and LCB preserving lower orders.The lower bound comes from LCB outputs, while the upper bound comes from combining the highest-order inputs through the FM.
Optimized FM projection reduces interaction computation from O(n^2d) to O(nkd) when k << n, exploiting the low-rank structure of XX^T.The approach reduces the stored interaction matrix from n × n to n × k and assumes d <= n.
Wukong treats embeddings as whole units and uses MLPs mainly to transform interaction results into embeddings for subsequent interactions.This embedding-wise design reduces computation relative to element-wise interaction architectures.

4. Implementation

Training high-complexity Wukong on large-scale data requires distributed execution and coordinated memory, compute, and communication optimizations.

Distributed training makes high-complexity Wukong feasible by sharding embeddings columnwise and applying FSDP to dense parameters.The sharding factor is tuned to fit memory while limiting redundancy.
FP16 and BF16 quantization reduces compute, memory, and communication overheads across embedding lookup and dense-parameter gradient paths.Automatic operator fusion is also used to improve training performance.

5. Overview of Evaluations

Wukong is evaluated on six public datasets and a large-scale internal dataset to assess effectiveness and scalability. It surpasses prior methods on public benchmarks and maintains quality and runtime advantages across internal complexity scales.

Wukong surpasses previous state-of-the-art methods across all six public datasets in the low-complexity evaluation.
On the internal dataset, Wukong consistently outperforms all baselines in model quality and runtime speed across complexity scales.The internal dataset contains 30 times more samples and 20 times more features than Criteo.
The internal evaluation also reports a better scaling trend for Wukong than for baseline models and includes component-level ablations.

6. Evaluation on Public Datasets

Wukong achieves state-of-the-art recommendation results across six public datasets, evaluated against established academic and industry models using AUC and LogLoss.

6.1.1. Datasets: The six datasets cover app usage, multimodal micro-video, movie ratings, advertising clicks, and large-scale advertising feedback.The datasets include Frappe, MicroVideo, MovieLens Latest, KuaiVideo, TaobaoAds, and Criteo Terabyte.
6.2. Baselines and Metrics: The evaluation benchmarks Wukong against seven recognized models, including AFN+, AutoInt+, DLRM, DCNv2, FinalMLP, MaskNet, and xDeepFM.AUC is the topline reporting metric, while LogLoss is also reported, with higher AUC and lower LogLoss preferred.
6.1.2. Experimental Setup: For smaller datasets, the study uses BARS configurations where possible and additionally tests embedding dimension 128 alongside the framework default.Wukong-specific tuning includes the dropout rate and other configuration choices described in the evaluation setup.
6.1.2. Experimental Setup: For Criteo, one-pass training and extensive grid search support fair comparisons across Wukong and all baselines.The search process involved nearly 3000 individual runs under the larger-dataset training setup.
6.3. Results: Wukong achieves state-of-the-art AUC results across all six public datasets, demonstrating effectiveness across diverse recommendation tasks.Table 2 reports the best AUC and LogLoss results for each dataset; the accompanying analysis highlights Wukong’s AUC performance across all datasets.

7. Evaluation on an Internal Dataset

On a large-scale internal dataset, Wukong maintains quality advantages as model complexity and size increase, while ablations show that its interaction components and scaling strategy contribute unequally to performance.

Dataset and Setup: Wukong consistently outperforms baselines in quality and runtime speed across complexity scales on a dataset with 146B entries and 720 features.The dataset supports two tasks, with comparable results reported across both; Task1 is emphasized in the main text.
Quality vs. Compute Complexity: Over two orders of magnitude in complexity, Wukong outperforms all baselines by over 0.2% LogLoss and improves by approximately 0.1% for every quadrupling of complexity.Even DCNv2 requires a 40fold complexity increase to match Wukong’s quality.
Quality vs. Model Size: Wukong outperforms all baselines by roughly 0.2% across model-size scales while steadily improving beyond 637 billion parameters.The sparse embedding table remains fixed at 627B parameters across runs.
Model-Specific Scaling: Baseline scaling saturates or fails: DLRM beyond 31 GFLOP/example, FinalMLP beyond 36 GFLOP/example, AFN+ without improvement, and several models under instability or memory constraints.AutoInt+ and DCNv2 encountered training instability, while MaskNet encountered out-of-memory errors during further scaling.
Significance of Individual Components: Ablations show that removing FMB causes large degradation, while removing both LCB and the residual connection causes substantial degradation; either alone has a modest effect.The results imply that zero-padding FMB outputs and using a residual connection can simplify LCB.
Impact of Scaling Individual Components: Increasing Wukong layers, MLP size, k, and nF improves quality, while nL plateaus individually; jointly scaling k, nF, and nL is more effective than scaling them separately.The layer effect is attributed to capturing higher-order interactions.

8. Discussions

The discussion identifies unresolved limits in Wukong’s scalability, theoretical understanding, and applicability beyond recommendation, while noting serving challenges for highly complex models.

Serving highly scaled-up models creates real-time deployment challenges, motivating foundation-model training and distillation into smaller serving models.
Wukong’s ultimate scalability limit remains unknown because the required compute prevented experiments from reaching that regime.
The paper lacks a comprehensive theoretical account of Wukong’s principles, especially relative to transformer architectures with stacked dot-product structure.
Whether Wukong generalizes beyond recommendation to domains with heterogeneous input sources remains unexplored.

9. Conclusion

The paper presents Wukong as a recommendation architecture that establishes scaling across model complexity while retaining competitive quality across tasks and datasets.

Wukong scales up and down across two orders of magnitude in compute complexity while maintaining a competitive edge over state-of-the-art models.The authors position it as a backbone ranging from small vertical models to large foundational models across tasks and datasets.

Impact Statement

The supplied impact-statement passages primarily document experimental implementation choices, hyperparameter searches, model configurations, and complexity reporting rather than societal impacts.

The paper states that it has no societal consequences requiring specific emphasis.
Training uses separate dense and sparse optimizers, linear warmup, global batch size 131,072, ReLU activations, 128-dimensional embeddings, and FP32.
The experiments use extensive hyperparameter searches, including more than 3000 runs on the Criteo dataset for fair comparisons.
The evaluated models include AFN+, AutoInt+, DCNv2, DLRM, FinalMLP, MaskNet, xDeepInt, and Wukong configurations.
The supplied tables report public-dataset model complexity and size, detailed run-level compute and quality, and Wukong component comparisons.
Wukong transforms second-order factorization-machine interactions into representations for subsequent interaction layers, enabling more sufficient interaction orders than traditional factorization.

E. Scaling Law on Training Data Volume

Wukong’s quality improves with training data and compute, with larger models becoming more data-efficient and maintaining gains through 146B examples; the study also highlights architectural factors behind scaling.

All Wukong models improve quality through 146B training examples, with larger models showing steeper improvement trends.Larger models also require fewer samples to achieve the same quality improvement.
The dataset remains too small for the large model to converge, so the observed scaling behavior does not establish its ultimate limit.
Wukong uses bit-wise MLP projections and a pyramid shape, which the authors associate with heterogeneous-feature learning and fewer unnecessary computations.
Replacing Transformer-style FFNs with bit-wise MLPs improves LogLoss by 0.34% for V-projection and 0.65% after self-attention.
Combining both bit-wise MLP changes with a pyramid layer shape achieves 0.57% quality improvement, while Wukong exceeds scaled-up AutoInt+ by 0.08%.
The analysis reports that Wukong’s stacked architecture captures higher-order interactions more efficiently than approaches with linear complexity or costly outer products.