一、Wukong [2024]

《Wukong: Towards a Scaling Law for Large-Scale Recommendation》

1. scaling laws 在模型质量的持续提升中起着关键作用。遗憾的是，由于现有推荐模型的 upscaling 机制的效率低下，其尚未展现出类似大语言模型领域中观察到的 scaling laws。这一局限使其在适应日益复杂的真实世界数据集时面临显著挑战。本文提出一种完全基于堆叠因子分解机（stacked factorization machines）的高效网络架构，以及协同扩展策略（synergistic upscaling strategy），统称为 Wukong，旨在为推荐领域建立 scaling law 。 Wukong 的独特设计使其能够通过增加网络层数和宽度，捕获任意阶数的 interactions。我们在六个公开数据集上进行了广泛评估，结果表明 Wukong 在模型质量上始终优于 SOTA 的模型。此外，我们在内部大规模数据集上评估了 Wukong 的可扩展性，结果显示其不仅在质量上优于现有 SOTA 模型，还能在模型复杂度跨越两个数量级（超过 100 GFLOP/example）时保持 scaling law，而这是现有方法无法实现的。

   

2. 基于深度学习的推荐系统（Deep learning-based recommendation systems: DLRS）如今为广泛的在线服务提供支持。现代 DLRS 旨在处理混合的 continuous dense features （如日期）和 categorical sparse features（如用户点击历史）。每个 sparse feature 通过可训练的 embedding lookup table 转换为 dense embedding representation ，这些 dense embeddings 随后被馈入到 interaction component 中，以捕获特征之间的复杂交互。

   尽管现有模型在较小数据集上表现出良好的准确性，但其适应更大规模和更复杂数据集的能力，以及在模型 scale up 时持续提升质量的能力仍不明确。随着现代数据集呈指数级增长，这种可扩展性变得愈发重要。例如，当前的生产数据集可能包含数百亿个训练样本。此外，foundational models （《On the opportunities and risks of foundation models》）需要 scale ，以同时处理更复杂的多源输入（multiple complex input sources）。因此，构建能够有效扩展或收缩、适应不同数据集规模和计算约束的 DLRS 至关重要，这一特性即所谓的 "scaling law"。

   目前，DLRS up-scaling 的主要趋势是通过 sparse scaling，即扩大 embedding tables 的规模（更多行数和/或更高维度）以减少冲突（collision ）并提升表达能力。这使得 DLRS 的参数规模达到数万亿（《Learning to embed categorical features without embedding tables for recommendation》、《High-performance, distributed training of large-scale deep learning recommendation models》、《Persia: An open, hybrid system scaling deep learning-based recommenders up to 100 trillion parameters》），其中 embedding tables 占主导地位。然而，这种传统 up-scaling 存在实际缺陷：单纯扩大模型的 sparse component 并不能增强其捕获日益复杂的 interactions 的能力；此外，这一趋势与硬件发展方向明显背离，因为下一代 accelerators 的性能提升主要体现在计算能力上，而 embedding table lookups 无法利用这些能力。因此，单纯扩展 embedding table 会导致基础设施成本过高且 accelerator 利用率低下，尤其在分布式环境中。

   我们的工作旨在为推荐模型寻找一种替代的 scaling 机制，以建立类似大语言模型领域的 scaling law，即设计一个统一的架构，通过综合策略（synergistic strategy ）使其质量随数据集规模、计算量和参数预算（parameter budgets）的增加而持续提升。

   我们聚焦于扩展 interaction components （称为dense scaling），以缓解 sparse scaling 在质量和效率上的缺陷。然而，现有模型因各种原因无法从这一范式中受益。例如，DLRM 缺乏捕获高阶交互的能力；DCNv2 和 AutoInt+ 缺乏有效的 upscaling 策略，导致 scaling up 时收益迅速递减；此外，即使采用 residual connection、layernorm 、gradient clip 等现代技巧，扩展现有模型仍容易出现训练稳定性问题（《Improving training stability for multitask ranking models in recommender systems》）。

   为建立推荐模型的 scaling law，我们提出 Wukong，这是一种简单的交互架构，具有高效的 dense scaling 特性。受二进制指数（binary exponentiation）思想启发，我们的关键创新是使用一系列堆叠的因子分解机（Factorization Machines: FMs）来高效且可扩展地捕获任意阶数的特征交互。在设计中，每个 FM 负责捕获输入的二阶交互；FM 的输出通过 MLP 转换为新的 embeddings ，这种方式将对 interactions 的结果进行编码并作为下一层的输入。

   我们使用六个公开数据集和一个大规模内部数据集对Wukong进行了评估。结果表明，Wukong在所有公开数据集上的 AUC 均优于 SOTA模型，证明了其架构的有效性和跨广泛 recommendation tasks and datasets 的泛化能力。在内部数据集上，Wukong不仅在相当复杂度水平下质量显著优于现有模型，还能在模型复杂度跨越两个数量级（超过 100 GFLOP/example）时持续提升质量，而现有方法在此范围内已无法有效 scale up。

1.1 相关工作

1. Deep Learning Recommendation Systems (DLRS)：现有的 DLRS 具有相似的结构。典型的模型由 sparse component 和 dense component 组成。

   • sparse component 本质上是 embedding lookup tables ，将 sparse categorical features 转换为 dense embeddings 。

   • 而 dense component 则负责捕获这些 embeddings 之间的交互以生成 prediction 。

2. Dense Interaction Architectures：捕获特征之间的交互是 DLRS 有效性的关键，我们在此梳理一些先前的研究成果：

   • AFN+：将特征转换到对数空间（logarithmic space ）以捕获任意阶数的交互（《Adaptive factorization network: Learning adaptive-order feature interactions》）。

   • AutoInt+：使用多头自注意力机制（《Autoint: Automatic feature interaction learning via self-attentive neural networks》）。

   • DLRM 和 DeepFM：利用因子分解机（Factorization Machines: FM ）显式捕获二阶交互（《Deep learning recommendation model for personalization and recommendation systems》、《Deepfm: a factorization-machine based neural network for ctr prediction》）。

   • HOFM：优化 FM 以高效捕获高阶交互（《Higher-order factorization machines》）。

   • DCNv2：使用 CrossNet ，通过 stacked feature crossing 来捕获交互，可视为一种 elementwise input attention （《Dcn v2: Improved deep & cross network and practical lessons for web-scale learning to rank systems》）。

   • FinalMLP：采用双线性融合（bilinear fusion）来聚合两个 MLP streams 的结果，每个 stream 以 stream-specific gated features 为输入（《Finalmlp: An enhanced two-stream mlp model for ctr prediction》）。

   • MaskNet：采用一系列 MaskBlock 进行 interaction capture ，对 input 和 DNN 的 intermediate activations 应用 "input attention" （《Masknet: Introducing feature-wise multiplication to ctr ranking models by instance-guided mask》）。

   • xDeepFM：将 DNN 与 Compressed Interaction Network 结合，通过外积捕获交互并通过 element-wise summation 来压缩结果（《xdeepfm: Combining explicit and implicit feature interactions for recommender systems》）。

3. Scaling up DLRS：

   • 《Learning to embed categorical features without embedding tables for recommendation》、《High-performance, distributed training of large-scale deep learning recommendation models》、《Persia: An open, hybrid system scaling deep learning-based recommenders up to 100 trillion parameters》 提出了 sparse scaling 机制。

   • 《Scaling law for recommendation models: Towards general-purpose user representations》 专注于 scaling up user representation models，其报告的最大模型总计算量小于 0.1 PF-days 。

   • 《Scaling law of large sequential recommendation models》旨在改进用户侧的 sequence modeling ，最大模型参数少于 0.8B。

   • 此外，《Understanding scaling laws for recommendation models》 研究了 DLRM 的 scaling law ，该模型在我们的工作中作为基线，并在实验中进一步 scaled up。

   • 与之正交的是，《Breaking the curse of quality saturation with user-centric ranking》提出以用户为中心的 ranking 公式以提高可扩展性。

   • 《On the embedding collapse when scaling up recommendation models》 对 sparse scaling 的局限性进行了分析，与我们的工作形成互补。

   • 此外，VIP5（《Vip5: Towards multimodal foundation models for recommendation》）利用 LLM 中的现有 scaling laws 将多模态 LLM 应用于推荐，然而 《How can recommender systems benefit from large language models: A survey》指出需要进一步研究验证基于 LLM 的推荐是否“模型越大性能越好”，而 《Foundation models for recommender systems: A survey and new perspectives》 建议需在更多样化的数据集上进行评估以得出结论。

1.2 Wukong 的设计

1. 在设计Wukong的架构时，我们牢记两个目标：

   • (1)：有效捕获复杂的高阶 feature interactions 。

   • (2)：确保Wukong的质量随数据集规模、GFLOP/example 、以及 parameter budgets 的增加而优雅地扩展。

a. 概述

1. 在Wukong中，categorical features 和 dense features 首先经过 Embedding Layer ，将这些 inputs 转换为 Dense Embeddings。 如 Figure 2 所示，Wukong随后采用 Interaction Stack ，这是一组 unified neural network layers，用于捕获 embeddings 之间的交互。Interaction Stack 的设计灵感来自二进制指数（binary exponentiation ）的概念，允许每一层捕获指数级高阶的特征交互。Interaction Stack 中的每一层包含一个 Factorization Machine Block: FMB和一个 Linear Compression Block: LCB。FMB 和 LCB 独立地接收 input （来自于上一层），它们的 outputs 被集成（ensembled ）从而作为当前层的 output 。interaction stack 之后是一个 final 的 Multilayer Perceptron: MLP 层，将 interaction 结果映射为 prediction 。

   

b. Embedding Layer

1. 给定一个 multi-hot categorical input ，embedding table 将其映射到 dense embedding。该过程涉及一系列 lookups 操作，每个操作对应 input 中的一个 “hot” dimensions ，lookup 结果通过池化操作（通常为 summation ）聚合。

   在我们的设计中，所有由 Embedding Layer 生成的 embedding dimension 被标准化为 global embedding dimension $d$$d$。为适应不同特征的重要性差异，每个重要特征生成多个 embeddings。less important 特征分配更小的 underlying embedding dimensions。这些 smaller embeddings 被分组、拼接后通过 MLP 转换为 $d$$d$ 维 embeddings 。

   如何做？论文并未详细说明。这种 multi embeddings 可以提升模型的表达能力。为进行公平比较，所有其他 baseline 模型也应该这样做。

2. dense inputs 通过 MLP 转换为相同维度 $d$$d$ 的 latent embeddings，并与 categorical input 的 embedding outputs 拼接，得到大小为 ${\mathbf{X}}_0\in {\mathbb{R}}^{n\times d}$$\mathbf X_0\in \mathbb R^{n\times d}$ 的 output tensor，其中 $n$$n$ 是来自 dense part 和 sparse part 的 embeddings 的总数。${\mathbf{X}}_0$$\mathbf X_0$ 随后被馈入 Interaction Stack 进一步处理。

   如何做？论文也并未详细说明。读者猜测：是将所有 dense inputs 拼接在一起，构成一个向量。然后通过 MLP 进行映射。

   但是这里有个疑问：是否需要对 dense inputs 进行归一化、或者 batch normalize？

   值得注意的是，与 DCN（《Dcn v2: Improved deep & cross network and practical lessons for web-scale learning to rank systems》）等传统方法不同，我们将每个 embedding 向量视为一个整体单元（细节见后文），因此我们将 ${\mathbf{X}}_0$$\mathbf X_0$ 表示为 ${\mathbb{R}}^{n\times d}$$\mathbb R^{n\times d}$ 的矩阵，而非 ${\mathbb{R}}^{(nd)}$$\mathbb R^{(nd)}$ 的向量拼接。

c. Interaction Stack

1. interaction modules 堆叠 $l$$l$ 个相同的交互层（interaction layers），每层通过因子分解机（Factorization Machines: FMs）逐层捕获高阶的特征交互。

   每个交互层包含并行的两个模块：Factorization Machine Block: FMB 、Linear Compression Block: LCB。FMB 计算层内 input embeddings 的特征交互，LCB 则简单地对 input embeddings 进行线性压缩并 forwards 。FMB 和 LCB 的 outputs 随后被拼接起来。

2. 对于 stack 中的第 $i$$i$ 层，其结果可包含从 $1$$1$ 阶到 $2^i$$2^i$ 阶的任意阶特征交互。通过数学归纳法可简单证明：假设第 $i$$i$ 层的输入包含 $1$$1$ 阶到 $2^{i-1}$$2^{i−1}$ 阶的交互（第一层 $i=1$$i=1$ 时成立），由于 FMB 在给定 $o_1$$o_1$ 阶和 $o_2$$o_2$ 阶交互时生成 $(o_1+o_2)$$(o_1+o_2)$ 阶交互，因此第 $i$$i$ 层的输出必然包含 $1$$1$ 阶到 $2^i$$2^i$ 阶的交互。下界由 LCB 的输出保证，上界由 FM 对 input 中两个 $2^{i-1}$$2^{i−1}$ 阶 interactions 的交互而生成。

3. 为稳定训练，我们在层间采用 residual connections 和 layer normalization: LN。整体公式为：

   $${\mathbf{X}}_{i+1}=\text{LN}(\text{concat}({\text{FMB}}_i({\mathbf{X}}_i),{\text{LCB}}_i({\mathbf{X}}_i))+{\mathbf{X}}_i)$$

   根据 FMB 和 LCB 的具体配置，${\mathbf{X}}_i$$\mathbf X_i$ 与 ${\mathbf{X}}_{i+1}$$\mathbf X_{i+1}$ 的 embeddings 数量可能不同（常见于第一层）。此时，残差连接会通过线性压缩调整形状以匹配。

d. Factorization Machine Block (FMB)

1. FMB 包含一个 FM 和一个 MLP。

   • FM 用于捕获 input embeddings 的显式特征交互，输出为一个二维 interaction matrix ，其中每个元素表示 a pair of embeddings 之间的交互。

   • 该 interaction matrix 被展平，然后通过 MLP 被转换为一个尺寸为 ${\mathbb{R}}^{(d\times n_F)}$$\mathbb R^{(d\times n_F)}$ 的向量，最后被 reshape 为 $n_F$$n_F$ 个 embeddings 以供后续使用。

   具体操作如下：

   $$\text{FMB}({\mathbf{X}}_i)=\text{reshape}(\text{MLP}(\text{LN}(\text{flatten}(\text{FM}({\mathbf{X}}_i)))))\in {\mathbb{R}}^{n_F\times d}$$

   注意，这里有个 LN layer ，位于 flatten 之后。

2. Wukong的 FM 模块完全可定制：例如，在最基础版本中，我们采用 《Deep learning recommendation model for personalization and recommendation systems》 的 FM 设计，即对 embedding vectors 之间的所有 pairs 取内积，即：

   $$\text{FM}(\mathbf{X})=\mathbf{X}{\mathbf{X}}^{\mathrm{\top }}\in {\mathbb{R}}^{n\times n}$$

   更优化的 FM 设计见后续章节。

e. Linear Compress Block (LCB)

1. LCB 仅对 embeddings 进行线性重组而不增加交互阶数，这对确保各层交互阶数的不变性（invariance ）至关重要。具体而言，它保证第 $i$$i$ 层捕获的交互阶数范围为 $1$$1$ 到 $2^i$$2^i$。LCB 的操作可描述为：

   $$\text{LCB}({\mathbf{X}}_i)={\mathbf{W}}_L{\mathbf{X}}_i\in {\mathbb{R}}^{n_L\times d}$$

   其中：${\mathbf{W}}_L\in {\mathbb{R}}^{n_L\times n_i}$$\mathbf W_L\in \mathbb R^{n_L\times n_i}$ 是权重矩阵，$n_L$$n_L$ 是超参数（表示压缩后的 embeddings 数量），$n_i$$n_i$ 是第 $i$$i$ 层的 input embeddings 数量。

f. Optimized FM

1. 传统 FM 通过 pair-wise 内积计算交互矩阵，其计算和存储复杂度随 embeddings 数量呈二次方增长，这在具有数千特征的真实数据集上难以承受。 为在降低计算成本的同时有效捕获特征交互，我们采用类似《Feature fusion for the uninitiated》、 《Dot product matrix compression for machine learning》 的方案，利用 pair-wise 内积矩阵的低秩性质（这在许多真实数据集上已被观察到，《Dcn v2: Improved deep & cross network and practical lessons for web-scale learning to rank systems》）。

   当 $d\le n$$d\le n$ 时，内积交互矩阵 $\mathbf{X}{\mathbf{X}}^{\mathrm{\top }}$$\mathbf X\mathbf X^\top$ 是秩为 $d$$d$ 的矩阵（这在特征数量大于 embedding 维度的大规模数据集上常见）。因此，我们可通过将 $\mathbf{X}{\mathbf{X}}^{\mathrm{\top }}$$\mathbf X\mathbf X^\top$ 与可学习的投影矩阵 $\mathbf{Y}\in {\mathbb{R}}^{n\times k}$$\mathbf Y\in \mathbb R^{n\times k}$ 相乘（即计算 $\mathbf{X}{\mathbf{X}}^{\mathrm{\top }}\mathbf{Y}$$\mathbf X\mathbf X^\top\mathbf Y$），将输出矩阵大小从 $n\times n$$n\times n$ 压缩至 $n\times k$$n\times k$（$k$$k$ 为超参数），且理论上不损失信息。这减少了存储交互矩阵的内存需求。进一步利用结合律，先计算 ${\mathbf{X}}^{\mathrm{\top }}\mathbf{Y}$$\mathbf X^\top \mathbf Y$，可将计算复杂度从 $O(n^{2}d)$$O(n^2d)$ 降至 $O(nkd)$$O(nkd)$，其中 $k\ll n$$k\ll n$。

   读者猜测具体的做法是：先将 $\mathbf{X}$$\mathbf X$ 馈入一个 MLP 从而得到 $\left({\mathbf{Y}}^{\mathrm{\top }}\mathbf{X}\right)$$\left(\mathbf Y^\top \mathbf X\right)$ ，其中 $\mathbf{Y}$$\mathbf Y$ 就是 ${\mathbf{W}}^{\mathrm{\top }}$$\mathbf W^\top$ ，而 $\mathbf{W}$$\mathbf W$ 就是 MLP 的权重矩阵。然后计算 $\mathbf{X}{\left({\mathbf{Y}}^{\mathrm{\top }}\mathbf{X}\right)}^{\mathrm{\top }}$$\mathbf X \left(\mathbf Y^\top \mathbf X\right)^\top$ 即可得到结果。

   此外，为提升模型质量，可通过对 compressed input 进行 MLP 处理从而得到投影矩阵 $\mathbf{Y}$$\mathbf Y$ ，使其对 input 具有注意力机制。除非特别说明，我们在后续实验中默认使用 optimized FM。

   怎么做？作者也未详细说明。

g. 复杂度分析

1. 假设 Interaction Stack 的每层使用相同超参数，且 MLP 中最大的全连接层大小为 $h$$h$。

   • 对于第一层：

     • FMB 的时间复杂度为 FM 和 MLP 的总和：FM 部分为 $O(nkd)\simeq O(ndh)$$O(nkd)\simeq O(ndh)$，MLP 部分为 $O(nkh+h^2+n_{F}dh)\simeq O(ndh+h^2)$$O(nkh+h^2+n_Fdh)\simeq O(ndh+h^2)$ 。

     • LCB 的时间复杂度为 $O(n{n}_{L}d)\simeq O(ndh)$$O(nn_Ld)\simeq O(ndh)$。

   • 对于后续层，时间复杂度为 $O(n^{\prime }dh+h^2)$$O(n^\prime dh+h^2)$，其中 $n^{\prime }=n_L+n_F$$n^\prime=n_L+n_F$ 。

   因此，Wukong的总时间复杂度为 $O(ndh+l{n}^{\prime }dh+h^2)\simeq O(ndh\log (n)+h^2)$$O(ndh+ln^\prime dh+h^2)\simeq O(ndh\log(n)+h^2)$ 。

h. Scaling Wukong

1. 以下总结与 scale up 的主要超参数：

   • $l$$l$：Interaction Stack 的层数。

   • $n_F$$n_F$：FMB 所生成的 embeddings 的数量。

   • $n_L$$n_L$：LCB 所生成的 embeddings 的数量。

   • $k$$k$：optimized FM 中的 compressed embeddings 的数量。

   • MLP：FMB 中 MLP 的层数和全连接层大小。

   在 scaling up 过程中，我们首先增加 $l$$l$ 以捕获更高阶交互；随后调整其他超参数以提升模型捕获更广范围内交互的能力。

i. Wukong 的 Enhanced Effectiveness 的直觉

1. 与现有以 FM 为主要交互架构的工作相比，Wukong通过堆叠 FM 的创新方法显著增强了传统 FM 的能力，使其能够捕获任意阶交互，这对需要高阶推理的大规模复杂数据集尤为有效。尽管已有高阶 FM 的研究，Wukong 捕获高阶交互的指数级效率避免了 HOFM 的线性复杂度和 xDeepInt 中外积操作的高成本。

   MLP 在隐式地捕获交互方面存在局限性（《Latent cross: Making use of context in recurrent recommender systems》）。Wukong不同于依赖 MLP 进行 interaction capture的方法，而是主要利用 MLP 将交互结果转换为 embedding representations ，供后续 interactions 使用。这种对 MLP 的独特用法增强了模型处理和解释复杂异构特征的能力。 此外，Wukong将每个 embedding 视为一个独立单元，专注于 embedding-wise 交互，这与捕获 element-wise 的架构相比，显著降低了计算需求。

1.3 实现

1. 本节讨论在大规模数据集上有效训练高复杂度Wukong模型的实践方法。

2. 总体而言，分布式训练是使Wukong training 可行的必要条件。

   • 对于 embedding layer ，我们使用 Neo（《High-performance, distributed training of large-scale deep learning recommendation models》）和 NeuroShard（《Pretrain and search: Efficient embedding table sharding with pre-trained neural cost models》）提供的 column-wise sharded embedding bag implementation。

   • 在 dense 部分，我们通过采用 FSDP（《Pytorch fsdp: experiences on scaling fully sharded data parallel》）并调整 sharding factor ，平衡性能与内存容量，确保模型适合内存且不产生过多冗余。

3. 为提高训练效率，我们通过 automatic operator fusion 以提升训练性能。此外，我们积极应用 quantization 技术以同时降低计算、内存和通信开销。具体而言：

   • Wukong的 embedding tables 以 FP16 训练。

   • embedding lookup 结果在前向传播中以 FP16 通信，在反向传播中以 BF16 通信。

   • 反向传播中 dense 参数的梯度传输采用 BF16 quantization 。

1.4 评估概述

1. 我们使用六个公开数据集和一个内部数据集对 Wukong 进行评估，详情汇总于 Table 1 。评估结果分为两个部分呈现。

   • 首先，我们在六个公开数据集上评估 Wukong，重点展示其在低复杂度场景下的有效性。结果表明，Wukong在所有六个数据集上均超越先前的 SOTA方法，证明了其有效性。

   • 然后，我们在大规模内部数据集上评估 Wukong 的可扩展性。该数据集的样本量是最大公开数据集 Criteo 的 30 倍，特征数是其 20 倍。结果显示：

     • (1)：Wukong在模型质量和运行速度上始终优于所有基线模型，并在所有复杂度范围内保持这一优势。

     • (2)：与基线模型相比，Wukong展现出更好的 scaling 趋势。

     我们还进行了消融实验，以了解 Wukong 各组件的独立贡献和有效性。

   

1.5 公开数据集上的评估

1. 本节旨在展示Wukong在多种公开数据集上的有效性。除非另有说明，为与先前工作保持一致，我们使用 BARS benchmark 提供的预处理数据。

2. 数据集：

   • Frappe：是一个 app usage log 。该数据集预测用户在给定上下文下是否使用 app。

   • MicroVideo ：由 THACIL 提供的基于内容理解的数据集，包含用户与微视频的交互，日志包含多模态 embeddings 和传统特征。

   • MovieLens Latest：著名的用户对电影评分的数据集。

   • KuaiVideo：快手发布的竞赛数据集，用于预测用户对新微视频的点击概率，包含基于内容理解的 embeddings 、及其他 categorical and float features 。

   • TaobaoAds：包含淘宝广告 CTR prediction 的 8 天数据。

   • Criteo Terabyte：包含 24 天广告点击反馈数据，我们使用最后一天数据进行测试。

3. 基线模型：我们将Wukong与学术界和工业界广泛认可的七种 SOTA 模型进行基准测试，包括 AFN+、AutoInt+、DLRM、DCNv2、FinalMLP、MaskNet 和 xDeepFM。

4. 评估指标：

   • Area Under the Curve (AUC)：衡量模型在所有阈值下正确分类正负样本的能力，值越高越好。遵循推荐领域惯例，我们使用 AUC 作为超参数调优和报告的主要指标。

   • LogLoss：衡量 prediction 与 actual label 的距离，值越低越好。

5. 模型配置：

   • 对于除 Criteo 外的五个较小数据集，我们采用公开的 BARS 评估框架。在可能的情况下，直接使用 BARS 上搜索到的最佳模型配置，其余使用默认超参数。除框架提供的默认 embedding 维度外，我们进一步测试 128 维 embedding，并报告两种配置中表现更好的结果。对于Wukong，我们调优 dropout rate、optimizer settings 和 compression of LCB 以适应特征数量。

   • 我们利用较大的 Criteo 数据集评估模型在真实在线推荐系统中的性能，进行 one-pass training。鉴于新的 training setup ，我们使用 "Implementation" 章节所述的系统对所有基线模型和Wukong进行了广泛的网格搜索，以促进公平比较，该过程涉及近 3000 次独立运行。模型特定的搜索空间见附录 A。搜索得到的最佳模型超参数随后作为下一节的 base config 。

6. 结果汇总于 Table 2 。总体而言，Wukong在所有公开数据集中在 AUC 指标上达到 SOTA水平。这一结果证明了Wukong架构的有效性及其理解多样化数据集、跨广泛推荐任务泛化的能力。

   

1.6 内部数据集上的评估

1. 本节通过一个大规模数据集展示 Wukong的可扩展性，并深入理解其各组件的有效性。该数据集能够揭示小规模公开数据集未体现的模型特性。

2. 数据集：该数据集共包含 146B 条记录和 720 个不同特征，每个特征描述用户或 item 的属性。数据集关联两项任务：

   • Task 1：预测用户是否对 item 表现出兴趣（如点击）。

   • Task 2：预测是否发生转化（如点赞、关注）。

3. 评估指标：

   • GFLOP/example：Giga Floating Point Operations per example，量化模型训练的计算复杂度。

   • PF-days：总训练计算量，相当于单台机器运行 1 PetaFLOP/s for 1 day 的计算量。

   • #Params：模型参数数量（以参数个数衡量），sparse embedding table size 固定为 627B 参数。

   • Relative LogLoss：相对于固定基线的 LogLoss 改善率。我们选择 basic config 的 DLRM 作为基线，在该数据集上 0.02% 的 Relative LogLoss 改善被视为显著。我们报告 online training 最后 1B-window 的 Relative LogLoss。

4. 基线模型：我们沿用公开数据集上的基线设置，但 xDeepFM 因昂贵的外积操作与大规模数据集不兼容，即使在 minimal setups 下也持续导致内存不足问题，故未包含在结果中。

5. 训练：所有实验均使用 pilot study 中找到的最佳优化器配置：dense 部分使用 Adam 优化器（lr=0.04，beta1=0.9，beta2=1），sparse embedding tables 使用 Rowwise Adagrad 优化器（lr=0.04）。模型以 online training 方式进行训练和评估，所有 runs 的 embedding 维度固定为 160。

   我们以前一节 Criteo Terabyte评估中找到的最佳配置作为起点，逐步增加各模型的参数数量。所有实验的 global batch size 为 262,144，根据模型大小在 128 或 256 块 H100 GPU 上运行。

a. 实验结果

1. 两项任务的结果是差不多的，正文仅报告 Task 1 的结果，Task 2 的详细结果见附录 C。

2. Quality vs. Compute Complexity：如 Figure 1 所示，我们描绘了质量与计算复杂度的关系。经验公式：

   $$y=-100+99.56x^{0.00071}$$

   结果显示，Wukong 在各种复杂度水平下始终优于所有基线模型，LogLoss 改善超过 0.2%。值得注意的是：

   • Wukong 在模型复杂度跨越两个数量级（约每四倍复杂度带来 0.1% 改善）时保持 scaling law。

   • 而基线模型中，AFN+、DLRM 和 FinalMLP 在特定复杂度后趋于平台期；AutoInt+、DCNv2 和 MaskNet 无法进一步提升质量。即使是表现最佳的基线模型 DCNv2，也需要 40 倍复杂度提升才能达到 Wukong 的质量。

   

3. Quality vs. Model Size：如 Figure 3 所示，我们展示了模型质量与模型规模的相关性。与计算复杂度扩展趋势一致，Wukong在所有规模下均以约 0.2% 的优势持续优于所有基线模型，并在超过 637B 参数时保持稳定提升趋势。

   

4. Quality vs. Data Size：见附录 E。

5. Model-Specifc Scaling：在扩展过程中，我们为每个模型采用不同策略，每次运行的详细超参数设置见附录 C。各模型扩展过程总结如下：

   • Wukong：通过调优 "Scaling Wukong" 章节所述超参数进行 scale up。

   • AFN+：扩展 AFN 的 hidden layers 、ensemble DNN 、以及对数神经元的数量。结果显示 scaling up AFN 未提升模型质量。

   • AutoInt+：扩展 multi-head attention 和 ensemble DNN 。模型质量最初低于其他模型，但 scaling up 后显著改善。

   • DLRM：扩展 topMLP，结果显示质量在超过 31 GFLOP/example 后饱和。

   • DCNv2：同时扩展 Cross Network 和 Deep Network 。扩展 Cross Network 未带来质量提升。训练稳定性低于其他模型，需严格梯度裁剪。

   • FinalMLP：扩展 two MLP streams 和 Feature Selection 模块。质量在低复杂度区域提升，但超过 36 GFLOP/example 后饱和。

   • MaskNet：测试 Parallel MaskNet 和 Serial MaskNet，发现并行变体更优。降低初始压缩率（reduction ratio ）以确保模型可运行，逐步扩展 MaskBlock 数量、DNN 和压缩率。

b. 消融实验

1. 组件重要性：我们通过实验验证Wukong 的 Interaction Stack 中 FMB、LCB 和残差连接的重要性：将各组件结果置零（zero-padding ），观察质量变化。 如 Figure 4 所示：

   图中的描述错误：下图的 FCB 应该是 FMB。

   • 置零 FMB 导致质量大幅下降。

   • 单独置零 LCB 或残差连接仅引起质量小幅下降。

   • 而同时置零两者则导致显著退化。

   这表明，通过对 FMB 输出零填充并结合残差连接，LCB 可被简化。

   本质上是因为：LCB 、残差连接都提供了一阶信息，因此二者的功能重复。因此，可以去掉 LCB，然后将 FMB 的输出、全零（代表一个 fake LCB ）拼接起来，从而简化模型。

   

2. 单个组件扩展的影响：我们剖析扩展 Wukong各超参数对模型质量的贡献：从基础配置开始，逐步将每个超参数加倍。结果如 Figure 5 所示：

   • 增加 Wukong 层数 $l$$l$ 因捕获更高阶交互而显著提升质量。

   • 扩大 MLP 规模也带来显著性能提升。

   • 提升 $k$$k$ 和 $n_F$$n_F$ 有益，而 $n_L$$n_L$ 在基础配置下已趋于平台期。值得注意的是，同时扩展 $k$$k$、$n_F$$n_F$、$n_L$$n_L$ 比单独扩展带来更显著的质量改善。

   

1.7 讨论

1. 大型模型的实际 Serving：扩展到高复杂度模型对 real-time serving 提出显著挑战。潜在解决方案包括训练多任务 foundation model 以分摊成本，或将大模型的知识蒸馏到小型高效模型用于 serving 。

2. 局限性与未来工作：我们也注意到工作的局限性和待探索方向，可作为未来研究目标：

   • 理解Wukong可扩展性的具体极限是重要研究方向。由于计算需求巨大，我们尚未达到可观察极限的复杂度水平。

   • 尽管Wukong在各项评估中表现出优越质量，但其底层原理的全面理论分析（尤其是与具有 stacked dot product 结构的 Transformer 等架构的对比）仍需进一步探索。

   • Wukong在推荐领域之外的泛化能力（尤其是涉及类似推荐中 heterogeneous input data sources 的领域）有待进一步研究和理解。

1.8 结论

1. 我们提出了一种高效的网络架构 “Wukong”，并证明其在推荐领域建立了此前未被观察到的 scaling law：Wukong能够在计算复杂度跨越两个数量级的范围内高效扩展与收缩，同时保持对其他 SOTA 模型的竞争优势，使其成为跨广泛任务和数据集、从小型垂直模型到的大型 foundational models 的 backbone 的可扩展架构。

1.9 附录

a. Criteo 数据集上的模型特定网格搜索空间

1. 我们对 dense arch optimization 使用 Adam 优化器，对 sparse arch optimization 使用 Rowwise AdaGrad 优化器，并在最初 10% 的 training steps 中采用 linear warmup period 。我们使用 global batch size 为 $8\times 16384=131,072$$8\times 16384=131,072$。所有模型均使用 ReLU 作为激活函数。在 pilot experiments 中，我们发现 embedding 维度设为 128 时所有模型表现更佳，因此采用该值。所有 runs均使用 FP32 精度。由于数据集规模和模型大小的限制，我们使用 《High-performance, distributed training of large-scale deep learning recommendation models》 作为 sparse distributed training framework，dense synchronization 采用数据并行。

2. 为促进公平比较，我们在 Criteo 数据集上对通用超参数和模型特定配置进行了广泛的网格搜索（超过 3000 runs ）。

   • 对于所有模型，sparse learning rate 和 dense learning rate 分别在 $\{1e^{-3},1e^{-2},1e^{-1}\}$$\{1e^{−3},1e^{−2},1e^{−1}\}$ 范围内调优。

   • 对于所有模型中的 MLP ，hidden layers 数量在 {1,2,3,4} 范围内，layer sizes 在 {512,1024,2048} 范围内。

   • 为减少过大的搜索空间，我们在 optimizer 的超参数上执行 pilot experiments ，并确定 dense 部分学习率设为 $1e^{-3}$$1e^{−3}$、sparse 部分学习率设为 $1e^{-1}$$1e^{−1}$ 时对所有模型最佳，后续 runs 固定此学习率。

   以下为各模型的特定搜索空间：

   • AFN+：AFN hidden units 和 DNN hidden units 在所有 runs 中保持一致，遵循通用 MLP 搜索空间；对数神经元的数量在 {128,256,512,1024} 范围内。

   • AutoInt+：基于论文 《Autoint: Automatic feature interaction learning via self-attentive neural networks》 报告的最佳配置创建搜索空间，每个超参数额外考虑更大值。注意力层数在 {3,4} 范围内，注意力维度在 {256,512} 范围内，注意力头数在 {4,8} 范围内；DNN hidden units 遵循通用 MLP 搜索空间。

   • DCNv2：cross layers 数量在 1 ~ 4 之间；rank 搜索范围为 full-rank 或 512。

   • DLRM：bottom MLP layer 为两层，维度为 [512,256]。

   • FinalMLP：遵循公开基准设置（《BARS: towards open benchmarking for recommender systems》），将一个流的 FeatureSelection: FS 设为所有浮点特征，另一流的 FS 搜索 8 selected sparse features 之一。FS MLP 设为[800]。head 数量固定为 256。

   • MaskNet：测试 Parallel MaskNet 和 Serial MaskNet 。并行变体的 blocks 数量在 {1,8,16} 范围内，block 维度在 {64,128} 范围内。串行变体的层数在 {1,4,8} 范围内，layer size 在 {64,256,1024} 范围内。两种变体的压缩率均固定为 1。

   • xDeepInt：考虑 Compressed Interaction Network (CIN)，层数在 {3,4} 范围内，层维度在 {16,32,64} 范围内。

   • Wukong：bottom MLP layers 为两层，维度设为 [512,256] 。$l$$l$ 范围为 1 ~ 4。$n_F$$n_F$ 和 $n_L$$n_L$ 设为相同值，范围在 {8,16} 内。$k$$k$ 固定为 24。

b. 公开数据集上的模型复杂度 / 规模

1. 具体细节见 Table 3 。

   

c. Model-Specific Scaling-up Configurations

1. 具体细节见 Table 5。

   

d. Wukong 中的高阶交互分析

1. 传统 factorization machine 方法通过最小化如下公式（《Deep learning recommendation model for personalization and recommendation systems》）来解决二阶交互问题：

   $$min\sum _{(i,j)\in \mathcal{S}}{r}_{i,j}-{\mathbf{X}}^1{\left({\mathbf{X}}^1\right)}^{\mathrm{\top }}$$

   其中：

   • $r_{i,j}\in \mathbb{R}$$r_{i,j}\in \mathbb R$ 表示第 $j$$j$ 个用户对第 $i$$i$ 个商品的评分，$i=1,\cdots ,m;j=1,\cdots ,n$$i=1,\cdots,m; j=1,\cdots,n$ 。

   • $\mathbf{X}$$\mathbf X$ 表示 userand item representations (embeddings) ，上标 1 表示 embedding 包含一阶信息。

   这些 embedding 向量的内积为二阶交互的后续评分提供了有意义的 prediction 。在 Wukong中，这些有意义的交互随后通过 MLP 转换为二阶交互 representations ${\mathbf{X}}^2$$\mathbf X^2$ 。在第二层 FMB 中，通过残差和 LCB connection，$({\mathbf{X}}^1+{\mathbf{X}}^2){({\mathbf{X}}^1+{\mathbf{X}}^2)}^{\mathrm{\top }}$$\left(\mathbf X^1 + \mathbf X^2\right)\left(\mathbf X^1 + \mathbf X^2\right)^\top$ 的内积可生成从一阶到四阶的交互。以此类推，一个 $l$$l$ 层的 Wukong 通过最小化以下公式来求解问题：

   $$min\sum _{(i,j)\in \mathcal{S}}(r_{i,j}-\sum _{k\in 1,2,\cdots ,2^{l-1}}{\mathbf{X}}^k{\left({\mathbf{X}}^k\right)}^{\mathrm{\top }})$$

   因此，与传统因子分解方法相比，Wukong 能够通过更丰富的交互阶数解决 recommendation 问题。

e. Training Data 上的 Scaling Law

1. Figure 6 总结了 Wukong 在不同训练数据集规模（ one pass 训练）上的性能。

   • 与大语言模型中观察到的现象类似，我们发现大型模型具有更高的数据效率，即实现相同质量提升所需的样本更少。

   • 此外，我们发现所有 Wukong 模型在高达 146B 数据时均持续提升质量，且更大模型在质量提升趋势上更为陡峭。

   • 我们也注意到，本研究的数据集规模对于大型模型收敛而言仍不足，这将是未来研究的方向之一。

   

f. 与 Transformer-based 方法的比较

1. 我们强调 Wukong 与 AutoInt+ 等 Transformer-based 方法在 scaling 上的差异及背后的直觉：尽管 Wukong 的结构与 Transformer 类似，但存在以下架构差异：

   • 首先，Wukong 在 FMB 和每一层中使用 MLP（ bit-wise 操作）进行投影，而非 Transformer 中的 FFN（ embedding/position-wise 操作）。

   • 其次，Wukong 采用金字塔形（pyramid shape ）配置，而 Transformer 采用统一形状（uniform shape ）。

   我们假设投影方式的差异对质量起关键作用。这些 MLPs 在 flattened input embeddings 上操作，本质上为每个特征提供不同的投影矩阵。这帮助模型从 heterogeneous input features 中学习，这与 LLM 中 single embedding space 的设计形成对比。

   在效率方面，金字塔形配置使 Wukong 能够通过压缩每层使用的 embeddings 数量排除不必要的计算。

   虽然 Wukong 的每一层的 input embeddings 数量与 output embeddings 数量相同，但是：output embeddings 中仅有一半是来自于 FMB ，而另一半来自于 input embeddings 自身（如果用残差连接来简化 LCB ）。

   为验证这些假设，我们通过将 Wukong 的独特组件应用于 AutoInt + 进行实验，结论如下：

   • (1)：使用 bit-wise MLP 而非 FFN 进行 V-projection 可使 LogLoss 改善 0.34% 。

   • (2)：在自注意力后添加 bit-wise MLPs 可使 LogLoss 改善 0.65%。

   • (3)：结合两者并采用金字塔层形状（通过在第一层输出使用 LCB ）可实现 0.57% 的质量提升。与 scaled up Autoint+ 相比，Wukong在节省 90% FLOPs 的同时实现 0.08% 的质量提升。

   结果汇总于 Table 4。