一、RankMixer [2025]

《RankMixer: Scaling Up Ranking Models in Industrial Recommenders》

  1. 大型语言模型(large language models: LLMs )的最新进展激发了人们对 scaling up推荐系统的兴趣,然而仍存在两个实际障碍:

    • 首先,工业级的推荐系统上的 trainingserving 成本必须遵守 strict latency 的限制和 high QPS 的需求。

    • 其次,ranking 模型中大多数的人工程设计的 feature-crossing 模块继承自 CPU 时代,无法充分利用现代 GPU,导致模型浮点运算利用率(Model Flops Utilization: MFU)低和可扩展性差。

    我们介绍了RankMixer,一种面向硬件感知(hardware-aware)的模型设计,旨在实现统一且 scalablefeature-interaction 架构。RankMixer 保留了 Transformer 的高并行性,同时用 multi-head token mixing 模块替换了平房复杂度的 self-attention 以提高效率。此外,RankMixer 通过 Per-token FFNs 保持了对 distinct feature subspaces 的建模、以及 cross-feature-space interactions 。我们进一步通过稀疏混合专家(Sparse-MoE)变体将其扩展到十亿级参数,以获得更高的投资回报率(ROI)。我们采用动态路由策略(dynamic routing strategy )来解决 experts training 不足和不平衡的问题。

    实验表明,在万亿级别的生产级数据集上,RankMixer 具有卓越的扩展能力。通过用 RankMixer 替换先前diverse handcrafted low-MFU modules,我们将模型 MFU4.5% 提升至 45%,并将 online ranking model 参数扩展了两个数量级,同时保持了大致相同的 inference latency。我们通过在两个核心应用场景(推荐 Recommendation 和广告 Advertisement)上的 online A/B tests 来验证了 RankMixer 的普适性。最终,我们为 full traffic serving 部署了 1B Dense-Parameters RankMixer,且未增加 serving 成本,这使得用户活跃天数(user active days)增加了 0.3%,应用内总使用时长(total in-app usage duration)增加了 1.08%

  2. 推荐系统(Recommender System: RS)在信息分发过程中至关重要。作为一个重要的机器学习场景,推荐系统基于大量的 multi-field feature data 来预测用户对 items 的行为,这些数据包括数值特征(例如各种各样的统计数据),categorical 特征(例如 user IDitem ID、用户行为特征、内容特征)。SOTA 的推荐方法是基于深度学习推荐模型(Deep Learning Recommendation Models: DLRMs),该类模型基于神经网络灵活地捕获特征交互(feature interactions )。这些 feature interactions 发生在 dense interaction layers 中;其中,dense interaction layers 位于 input embeddings 之上。DLRM 中的 dense interaction layer 对推荐系统的性能至关重要,并且人们提出了各种各样的模型结构。

    LLMsincreasing parameters 中获益的进展的驱动, scaling up DLRMs 以充分利用数据量是一项迫切需求。先前的研究在扩展 scaling up DLRMs 方面产生了许多成果。

    • 早期研究(《Understanding scaling laws for recommendation models》《Scaling generative pre-training for user ad activity sequences》《Scaling law of large sequential recommendation models》)仅仅是在不修改结构的情况下拓宽或堆叠 feature interaction layers 。这种方式获得的收益有限,有时甚至是负面的。

    • 随后的努力,例如 DHENWukong,专注于设计创新的 DNN 结构以提升 scaling 的性能。

    然而,在推荐系统中利用 model scale 来提高性能面临着独特的实际挑战。与 NLP 任务或视觉任务不同,工业级推荐系统必须遵守严格的 latency 约束并支持极高的 QPSqueries per second 每秒查询数)。因此,核心挑战在于找到模型有效性(model effectiveness )和计算效率(computational efficiency)的最佳平衡点。

    历史上,推荐系统中的 ranking model 架构是由 CPU 时代的设计原则塑造的。这些模型通常依赖于对 heterogeneous diverse handcrafted cross-feature modules 的组合从而来提取 feature interactions,但它们的许多核心算子(core operators)在现代 GPU 上是内存受限(memory-bound)而非计算受限(compute-bound)的,导致 GPU 并行性差和极低的 MFU(模型浮点运算利用率Model Flops Utilization),通常 MFU 仅为个位数百分比。此外,由于 CPU 时代模型的计算成本大致与参数数量成正比,scaling laws 所暗示的通过积极地 scaling 所带来的潜在 ROI 在实践中难以实现。

    总之,DLRMsscaling laws 上的研究必须克服以下问题:

    • 架构应与硬件对齐,在现代 GPU 上最大化 MFU 和计算吞吐量(compute throughput)。

    • 模型设计必须利用 recommendation data 的特性,例如异构的特征空间、以及数百个 fields 之间的个性化的 cross-feature interactions

    为了应对这些挑战,我们提出了一种硬件感知(hardware-aware)的模型设计方法,RankMixerRankMixer 的核心设计基于两个scalable 的组件:

    • Multi-head token mixing 仅通过无参数算子(parameter-free operator)就实现 cross-token feature interactions。该策略在性能和计算效率方面优于 self-attention 机制。

    • Per-token feed-forward networks (FFNs) 显著扩展了模型容量(model capacity),并通过为不同的 feature subspacemodeling 来分配独立的参数来解决 inter-feature-spacedomination 的问题。这些 FFNs 也与 recommendation data 的模式很好地契合,实现了更好的 scaling

    为了进一步提升 large-scale 模型的 ROI,我们将 per-token FFN modules 扩展为 Sparse Mixture-of-Experts (MoE) 结构。通过为不同的数据动态地在每个 token 仅激活 a subset of experts ,我们可以在计算成本 minimal increase 的情况下显著提高模型容量。RankMixer采用了类似于 Transformer 的高度并行架构,但克服了 self-attention based feature-interaction 的几个关键限制:训练效率低、modeling cross-space ID similarity 时的组合爆炸、以及 attention weights matrix 引起的严重的 memory-bound。同时,在相同 Flops 下,与 Vanilla Transformer 相比,RankMixer 提供了更大的模型容量和学习能力。

    modeling cross-space ID similarity 时的组合爆炸,这指的是:在推荐系统中,由于 ID 类特征数量巨大且来自不同语义空间,使用自注意力机制计算所有 feature pair 之间的相似性会导致计算和存储复杂度呈平方级增长,从而无法在实际系统中高效运行。RankMixer 通过去参数化的 token mixing 机制避免了这一问题,实现了高效且可扩展的 heterogeneous feature interaction建模。

    在抖音的推荐系统的生产部署中,我们证明了在与先前 baseline 相比延迟更短的情况下,将模型参数扩展 100 倍以上是可行的。这得益于 RankMixer 架构能够将 parameter growthFLOPs 解耦,并通过 high MFUengineering optimizationFLOPs growth 与实际成本解耦。

    主要贡献总结如下:

    • 我们提出了一种名为 RankMixer 的新架构,遵循硬件感知的模型设计理念。我们设计了 multi-head token mixingper-token FFN 策略来高效地捕获 heterogeneous feature interactions,并使用动态路由(dynamic routing)策略来提高 RankMixerSparseMoEscalability

    • 利用 high MFUperformance-optimization 的杠杆,我们在不增加推理成本的情况下将模型参数扩展了 70 倍,包括 improving MFUQuantization

    • 我们在万亿级别的工业级推荐数据集上进行了广泛的离线实验和在线实验,并研究了模型的 scaling lawRankMixer 模型已成功部署在抖音 Feed Recommendation Ranking 中用于 full-traffic serving ,分别实现了活跃天数(active days)和 App 使用时长(App duration)的 0.3%1.08% 的增长。

1.1 相关工作

  1. 现代推荐系统基于 Deep Learning Recommendation Models: DLRMs ,如何有效地建模 feature interactionsDLRMs 的关键因素。

    • Wide&Deep 是最早的努力之一,它结合了逻辑回归(wide part )和 DNNdeep part)来分别捕获低阶的和高阶的 feature interactions

    • DeepFM 是另一个成果,它集成了因子分解机(Factorization Machine: FM)和 DNN

    • 此外,DeepCross 是残差网络(residual network )(《Deep residual learning for image recognition》)的扩展,旨在隐式地学习自动的 feature interactions

    然而,仅依靠DNN 来学习高阶特征交互已被证明相当具有挑战性。显式的 cross methods 设计了不同的算子(operators )来显式地捕获高阶特征交互,例如 PNNDCN 及其后继者 DCNv2xDeepFMFGCNNFiGNNAutoIntHiformer 采用带有残差连接的 attention 机制来学习复杂的交互。DHEN 提出将多个 interaction operators 组合在一起。尽管提高了准确性,但这些新架构增加了模型 latency 和内存消耗,并且它们的模型规模相对较小。

  2. scaling law 已成为深度学习的一个基本主题,并且是过去十年中许多突破的关键催化剂,特别是在自然语言处理(NLP)、计算机视觉(CV)和多模态建模领域,这些定律描述了模型性能与 scaling factors(如模型大小、数据量和算力)之间的 power-law correlations。最近,推荐系统中的 scaling laws引起了研究人员的广泛关注。

    • 人们研究探索了用于 pre-training user activity sequences《Scaling generative pre-training for user ad activity sequences》)、general-purpose user representations《Scaling law for recommendation models: Towards general-purpose user representations》《Scaling law of large sequential recommendation models》)和 online retrieval《 Scaling laws for dense retrieval》《Scaling Laws for Online Advertisement Retrieval》)的 scaling 策略。

    • Wukong堆叠 FMLCB 来学习 feature interactions

    • 正交地,《Scaling law of large sequential recommendation models》 将一个序列推荐模型扩展到 0.8B 参数。

    • HSTU 提升了 Generative Recommenders: GRsscaling 效果。

1.2 方法论

1.2.1 整体架构

  1. RankMixer 的整体架构包括 T input tokens ,这些 tokens 经过 L 个连续的 RankMixer block 处理,最后是一个 output pooling operator。如 Figure 1 所示,每个 RankMixer block 有两个主要组成部分:

    • (1):一个 Multi-Head Token Mixing 层 。

    • (2):一个 Per-Token Feed-Forward Network: PFFN 层。

    首先,input vector einputtokenizeTfeature tokens x1,x2,,xT,每个 token 代表一个 feature vectorRankMixer block 通过以下方式对 token representations 进行反复地 refineL 层):

    Sn1=LN(TokenMixing(Xn1)+Xn1)Xn=LN(PFFN(Sn1)+Sn1)

    其中:

    • LN()layer normalization 函数,TokenMixing()PFFN() 分别是 multi-head Token Mixing 模块和 per-token FFN模块。

    • XnRT×D 是第 nRankMixer block 的输出;X0RT×Dx1,x2,,xT 堆叠而成,D 是模型的 hidden dimension

    output representation ooutput 源自 final layersrepresentations XLmean poolingooutput 将用于计算不同任务的 prediction

    注意,下图中的 H feature tokens ,在实际中选择 H=T ,使得在 mixing 前后保持 tokens 数量不变。

1.2.2 Input Layer and Feature Tokenization

  1. 构建 large-scale 推荐模型的首要任务是准备好具有丰富信息的 inputs,例如:

    • User features:包括 user ID 和其他的用户信息。

    • candidate features:包括 video IDauthor ID 等等。

    • Sequence Features:由 Sequence Module 进行处理以捕获时序的兴趣,产生 es

    • 以及 Cross Features :即 usercandidate 之间的交叉特征。

    所有特征都将被转换为具有 diverse dimensionsembeddings

    为了实现后续阶段的高效并行计算,必须将不同维度的 embeddings 转换为维度对齐的向量,称为 feature-tokens。我们将这种 embedding alignment 过程称为 tokenization。最简单的策略是为每个特征分配一个 embedding。考虑到通常有数百个特征,这会引入几个挑战:

    • 拥有数百个 tokens 不可避免地减少了分配给每个 tokenparameterscomputation ,使其成为 small fragments ,导致对重要特征的建模不足以及 GPU cores 的利用率不足。

    • 相反,tokens 太少(例如单个 token )会使模型结构退化为简单的深度神经网络(Deep Neural Network: DNN),无法可区分地表达 diverse feature spaces,这存在 dominant features 掩盖其他特征的风险。

    这里的挑战仅仅是针对 RankMixer 框架的。其它框架可能不存在这两个问题。因为 RankMixer 涉及到 Multi-Head Token MixingPer-Token Feed-Forward Network: PFFN

    • 如果 tokens 太多:Per-Token Feed-Forward Network: PFFN 针对每个 token 都有一个 FFN,因此总参数规模急剧膨胀。如果控制总参数规模不变,那么每个 FFN 的参数将急剧缩小。

    • 如果 tokens 只有一个:那么 Multi-Head Token Mixing 等价于 Identify 操作、Per-Token Feed-Forward Network: PFFN 退化为 FFN。整个模型就退化为简单的 DNN

    作者的解决方案为:先把所有 feature embeddings 拼接在一起,然后拆分为 Ttokens

  2. 为了克服这些问题,我们提出了一种基于语义的tokenization方法,利用 domain 知识将特征分组为几个语义连贯的 clusters。这些分组后的特征被串行地拼接成一个 embedding 向量 einput=[e1;e2;;eN] ,其中:en,1nN 表示每个 feature groupembeddingNfeature-groups 的数量;; 表示向量拼接。

    注意:因为要先拼接、再拆分、再投影。因此,feature embeddings {en}n=1N 的拼接顺序就非常重要。不同的拼接顺序会导致不同的拆分结果。这里按照语义连贯的顺序来拼接。作者并未给出 ”语义连贯“是如何定义的。读者猜测:将 user features 拼接在一起、将 item features 拼接在一起, 等等。

    随后 einput 被分割成适当数量的具有固定维度大小的 tokens 。每个 feature token xiRD 捕获一组 feature embeddings,代表一个 similar semantic aspect

    xi=Proj(einput[d×(i1):d×i]),i=1,,T

    其中: einput 是拼接后的 embedding 向量,d 是每个 token 的固定维度,T 是生成的 token 数量,Proj 函数将分割后的 embedding 映射到 D 维。

    假设 en 的维度为 dn ,则有:n=1Ndn=T×d

    注意,这里是按照固定的窗口 d 来拆分 token,而不是按照 feature groups 来拆分 token。是否可以参考 CNN 的做法,用一组 slide window (具有不同的步长)从而获得多组 inpu tokens ?可以试试。

1.2.3 RankMixer Block

  1. Multi-head Token Mixing:为了促进跨 tokens 的有效信息交换(这对于 feature crossglobal information modeling 非常重要),我们引入了 multi-head Token Mixing 模块。每个 token 被均匀地分成 H 个头,token xt 的第 hhead 表示为 xt(h)

    [xt(1)xt(2)xt(H)]=SplitHead(xt)

    这些 heads 可以看作是token xt 到低维 feature subspace 的投影,因为 recommendation 任务需要考虑不同的视角。

    Token Mixing 用于融合这些 sub-space vector 以实现 global feature interactions。形式上,经过 multi-head Token Mixing 后,对应于第 hheadsh (即,shuffled tokens)构造如下:

    s(h)=Concat(x1(h),x2(h),,xT(h))RTDH

    multi-head Token Mixing 模块的输出是 SRH×TDH,它由所有 shuffled tokens s(1),s(2),,s(H) 堆叠而成。在这项工作中,我们设置 H=T,以便在 Token Mixing 后保持相同数量的 tokens 用于 residual connection

    考虑到 H 必须能够整除 D ,因此 D 必须是 T 的整数倍。

    经过 residual connectionnormalization 模块后,我们可以生成:

    s1,s2,,sT=LN(TokenMixing(x1,x2,,xT)+(x1,x2,,xT))

    由于 H=T ,因此 SRT×D ,因此这里 TokenMixing(x1,x2,,xT)=(s(1),s(2),,s(H)) 。如果 HT ,那么可能需要先将 S 进行投影。

    尽管 self-attention 在大型语言模型中已被证明非常有效,但我们发现它对于推荐系统并非最优。在 self-attention 中,attention 权重是使用 tokens 的内积计算的。这种方法对于 NLP 很有效,因为所有 tokens 共享一个 unified embedding space。然而,在推荐任务中,特征空间本质上是异构的。计算两个异构语义空间(heterogeneous semantic spaces)之间的内积相似性非常困难——尤其是在推荐系统中,来自 user 侧和 item 侧的特征的 ID space 可能包含数亿个元素。因此,将 self-attention 应用于如此 diverse inputs,其性能并不优于无参数的 multi-head Token Mixing方法,并且消耗更多的计算、Memory IO 操作和 GPU memory usage

    TokenMixing 是可微的。因为它只是对输入进行固定的置换(permutation)操作。在深度学习中,置换操作(例如reshapetransposegather 等)都是可微的,因为它们的前向传播是确定的元素重排,反向传播时梯度可以按照相反的规则传播回原来的位置。

  2. Per-token FFN:先前的 DLRMDHEN 模型倾向于在单个 interaction 模块中混合来自许多不同的语义空间的特征,这可能导致high-frequency fields 主导并淹没低频或长尾的信号,最终损害整体的 recommendation 质量。我们引入了一种参数隔离的 feed-forward network 架构,称为 per-token FFN。在传统设计中,FFN 的参数在所有 tokens 之间共享,但我们的方法使用专用的变换(transformations)来处理每个 token,从而隔离了每个 token 的参数。对于第 ttoken stper-token FFN 可以表示为:

    vt=fpffnt,2(Gelu(fpffnt,1(st)))fpffnt,i(x)=Wpffnt,ix+bpffnt,i,i{1,2}

    其中:

    • fpffnt,i(x)per-tokenFFN 的第 iMLP

    • Wpffnt,1R(kD)×D,bpnnft,1RkD,Wpffnt,2RD×(kD),bpnnft,2RD 为待学习的参数。k 是一个超参数,用于调整 per-token FFNhidden dimension

      注意:这里的参数 Wpffnt,1,bpnnft,1,Wpffnt,2,bpnnft,2per-token 的,因此具有上标 t

    • Gelu()Gelu 激活函数。

    • stRD 为第 ttoken

    我们将 Per-token FFN 模块总结为:

    v1,v2,,vT=PFFN(s1,s2,,sT)

    其中:

    PFFN(s1,s2,,sT)=fpffnt,2(Gelu(fpffnt,1(s1,s2,,sT)))

    parameter-all-shared FFN 相比,per-token FFN 通过引入更多参数来增强建模能力,同时保持计算复杂度不变。

    这种思想在 Hiformer 中已经被引入:field-specific FFN

    值得强调的是,per-token FFNMMoE experts 不同,因为每个 Per-token FFN 看到的是 distinct token input,而 MMoE中的所有 experts 共享相同的输入。与 MMoE(许多 experts 处理相同的输入)和 Transformer(不同的输入共享一个 FFN )不同,RankMixer 同时分割 inputs 和参数,这有利于在不同 feature sub-spaces 中学习多样性。

 

1.2.4 Sparse MoE in RankMixer

  1. 为了进一步提高 scaling ROI,我们可以用 Sparse Mixture-of-Experts (MoE) blocks 来替换 the dense FFNs of each pertoken,这样模型的容量增长而计算成本大致保持不变。

    这段话可以这样理解:用 MoE 代替 FFN 可以带来模型容量的显著增加,因此 MOE 中的每个 expert 就相当于一个 FFN ;同时每次仅仅激活少数的几个 experts,从而使得计算量基本保持不变。

    实验部分缺乏这部分的消融分析。

    然而,原始的 Sparse-MoERankMixer 中性能下降,因为:

    • (i) uniform k-expert routingTop-k selection 平等地对待所有 feature tokens ,将预算浪费在 low-information tokens 上,而使 high-information tokens 缺乏资源,这阻碍了模型捕获 tokens 间的差异。

      这段话的意思是:top-kk 对每个 feature tokens 是相同的。理论上:如果对于 low-information tokens 采用较小的 k 、对于 high-information tokens 使用较大的 k ,则会效果更好。但是,一方面不知道哪些 tokenshigh/low information 的;另一方面也方便设置 per-tokenk 值。

    • (ii) expert under-trainingPer-token FFNs 已经将 expert 参数(每个 expert 就是一个 FFN)规模乘以了 #tokens ;添加 non-shared experts 进一步爆炸了 expert 数量,产生高度不平衡的 routing 和训练不佳的 experts

  2. 我们结合两种互补的 training 策略来解决上述问题:

    注意:这里有些符号与原始论文中的不一致。读者发现论文中符号容易混淆,因此做了修改。

    • ReLU Routing:为了赋予 tokens 灵活的 expert 数量并保持可微性,我们用 ReLU gate 上自适应的 l1 penalty《ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing》)替换了常见的 Topk+softmax。给定第 jexpert ej() 用于 token siRdhrouter h()

      {Gi,j}j=1Ne=ReLU(h(si))vi=j=1NeGi,jej(si)

      其中:Ne 是每个 tokenexperts 的数量。

      router 函数 h()、以及 expert 函数 {ej()}j=1Ne 都是 token-specific 的,而不是全局共享的。

      ReLU Routing 将为 high-information tokens 激活更多 experts ,并提高参数效率。稀疏性(Sparsity )由带有系数 λLreg 引导,该系数将平均 active-expert ratio 保持在预算附近:

      L=Ltask+λ×Lreg,Lreg=i=1Ntj=1NeGi,j

      Nt=Ttokens 数量。

    • Dense-training / Sparse-inference (DTSI-MoE):受 《Ads recommendation in a collapsed and entangled world》 的启发,我们采用了两个 routers htrainhinfer,并且 Lreg 仅应用于 hinferhtrainhinfer 在训练期间都会更新,而只有 hinfer 在推理时使用。事实证明,DS-MoE 使得 experts 不会遭受 under-training 的问题,同时降低了推理成本。

      双路由方式:

      {Gi,jtrain}j=1Ne=ReLU(htrain(si)){Gi,jinfer}j=1Ne=ReLU(hinfer(si))vi=j=1NeGi,jtrainej(si)L=Ltask+λ×Lreg,Lreg=i=1Nej=1NtGi,jinfer

      其中:htrain 用于确保任务的准确性,确保 experts 充分训练;hinfer 用于确保 routing 的稀疏性,控制 inference 成本。htrain,hinfer 采用相同的网络结构,但是采用独立的参数。Ltask 的梯度用于更新 htrainLreg 的梯度用于更新 hinfer

      虽然论文没有明确说明,但实践中通常会有隐式的对齐,确保 {Gi,jtrain}j=1Ne{Gi,jinfer}j=1Ne 的相似度尽可能高(如,cos 相似度)。例如:

      Lalign=MSE(htrain(si),hinfer(si))

      .

1.2.5 Scaling Up Directions

  1. RankMixer 本质上是一个高度并行和可扩展的架构。其参数数量和计算成本可以沿着四个正交轴来扩展:tokens 数量 T、模型宽度 D、层数 Lexperts 数量 E (就是前面描述的 Ne )。对于 full-dense-activated 版本,一个样本的参数和前向 FLOPs 可以计算为:

    Params2kLTD2,FLOPS4kLTD2

    k 是调整 FFNhidden dimension 的缩放比率。

    Sparse-MoE 版本中,每个 token 的有效的参数和计算量进一步按 sparsity ratio s=ActivatedParamsTotalParams 缩放。

1.3 实验

  1. 数据集:离线实验使用抖音推荐系统的训练数据进行。这些数据源自抖音的在线日志和 user feedback labels。训练数据集包括超过 300 个特征,例如数值特征、ID 特征、交叉特征和序列特征,涉及数十亿 user IDs 和数亿 video IDs。所有这些特征都被转换为 embeddings。数据覆盖每天数万亿条记录,实验是在两周内收集的数据上进行的。

  2. 评估指标:AUC (Area Under the Curve)UAUC (User-level AUC) 作为主要性能指标;Parameter countFLOPsMFU 作为效率指标。

    • Finish/Skip AUC/UAUCfinish=1/0skip=1/0 标签指示用户是否在短时间内看完了一个视频或滑到了下一个视频。我们评估此 finish labelAUCUAUCAUC 增加0.0001 可以被认为是一个显著且置信的改进。

    • Dense-Paramdense part 中的参数数量,不包括 sparse embedding 参数。

    • Training Flops/Batch:运行一个 512 batch-sizebatch 通过模型所需的 floating-point operations: FLOPs 次数,代表训练的计算成本。

    • MFUMFU(Model FLOPs Utilization) 是一个衡量模型如何有效利用硬件提供的浮点运算的指标,通过将模型的实际 FLOPs 消耗除以硬件的理论 FLOPs 容量来计算。

  3. baselines:我们与以下广泛认可的 SOTA 基线进行比较:

    • DLRM-MLP:原始 MLP 用于 feature crossing,这作为实验的 baseline

    • DCNv2, RDCNfeature cross 模型的 SOTA

    • MoE:通过使用多个并行 experts 进行 scale up

    • AutoInt, Hiformer :结合了 heterogeneous self-attention layerlow-rank approximation matrix computation

    • DHEN:结合了不同类型的 feature-cross blockDCN/self-attention/FM/LR )并堆叠了多个层。

    • Wukong:研究了 feature interactionscaling law,遵循带有 Factorization Machine Block (FMB)Linear Compress Block (LCB)DHEN 架构。

    所有实验都是在数百个 GPU 上进行的,采用 hybrid distributed training 框架,其中 sparse 部分异步更新,而 dense 部分同步更新。所有模型的 optimizer hyperparameters 保持一致。

    • 对于 dense 部分,我们使用 RMSProp 优化器,学习率为 0.01

    • sparse 部分使用 Adagrad 优化器。

    在论文中,sparse 部分和 dense 部分采用不同的优化器,这是工业级推荐系统训练中的常见做法。原因如下:

    1. 数据分布和更新频率不同

      • sparse 部分:参数数量巨大(通常有数十亿甚至数百亿的ID类特征),但每个训练样本只更新极少一部分 embedding 向量(即出现过的 ID 对应的 embedding)。这些 embeddings 的更新非常稀疏,而且每个 embedding 向量的更新频率与其在训练数据中出现的频率相关。

      • dense 部分:参数数量相对较少(几百万到几十亿),每个训练样本都会更新所有dense 参数。

    2. 优化器的适应性

      • 对于 sparse embedding,通常使用自适应优化器(如 AdagradAdamFTRL 等)来为每个特征(每个 ID )单独调整学习率。这是因为不同 ID 的出现频率差异很大,自适应优化器可以为低频 ID 设置较大的学习率,为高频 ID设置较小的学习率,从而更好地训练 sparse embedding

      • 对于 dense layer,通常使用 SGDRMSPropAdam 等优化器。在论文中,dense layer 使用了 RMSProp,这可能是因为 RMSProp 在推荐系统任务中表现稳定。

    3. 训练效率

      • sparse embedding layer 由于参数量巨大,通常存储在参数服务器(Parameter Server)或使用其他分布式存储方式,并采用异步更新。为了减少通信开销,sparse embedding 的更新可能采用延迟更新或更高效的优化算法。

      • dense layer 则可以使用同步更新,因为参数量相对较小,同步更新的开销不大。

    4. 收敛性考虑

      • sparse embeddingdense layer 通常具有不同的梯度分布和收敛速度。使用不同的优化器(甚至不同的学习率)可以分别调整两部分参数的学习过程,使训练更稳定、收敛更快。

1.3.1 与 SOTA 方法的比较

  1. 为了探索如何 scale up 模型,我们比较了参数大小相似(大约 100 million)的模型,以确定在相同计算成本下哪种模型结构表现最佳。我们的方法和基线性能的主要结果总结在 Table 1 中。我们可以观察到:RankMixer 在多个目标和指标上显著优于其他 SOTA 模型。

    我们仔细看一下每个模型:

    • 首先,仅仅将 DLRM 扩展到100 million 参数只能带来有限的收益,这强调了设计针对 recommendation data 特性的模型以获得更好 scaling 性能的必要性。

    • 然后,我们将 RankMixer 模型与其他经典的 cross-structure 设计(如 DCNRDCNAutoIntDHEN)进行比较,发现它们在模型参数和计算成本之间存在不平衡。即使参数规模相对较小,这些模型已经表现出较大的 FLOPs,表明设计缺陷限制了它们在 Table 中的结果。

    • 同时,尽管 RankMixer 实现了最佳性能;而且当扩展到100 million 参数时,其 FLOPs 在所有模型中仍然相对适中,反映了模型容量和计算负载的平衡方法。

    • 我们还将 RankMixer 与几种常用的 SOTAscaling-up 模型(如 HiformerWukong)进行了比较,发现在相似的参数设置下,RankMixer 不仅性能更好,而且计算需求更低。

    Hiformer 的性能排第二,而且参数规模、FLOPs/Batch 也适中。

1.3.2 不同模型的 Scaling Laws

  1. Figure 2 中,我们展示了从参数大小和 Flops 观察到的 scaling law 曲线。RankMixer 模型在参数和 FLOPs 方面都显示出最陡峭的 scaling lawRankMixer 始终优于其他模型。

    • 尽管 Wukong 模型表现出相对陡峭的参数曲线,但其计算成本增长更快;因此,在 AUCFLOPs 曲线上,它与 RankMixerHiformer 的差距更大。

    • 此外,Hiformer 的性能略逊于 RankMixer,反映了其在 feature level 上对 token segmentationAttention 的依赖影响了其效率。

    • DHENscaling 不理想,反映了其 cross structure 的有限的 scalability

    • 此外,MoE 通过添加 experts 进行扩展的策略在保持 expert balance 方面带来了挑战,导致 scaling 的性能不佳。

  2. 具体来说,我们可以通过增加宽度 (D)、feature tokenT)和层数 (L)来扩展 RankMixer 模型。在我们的实验中,我们观察到一个与 LLMscaling laws 共享的结论:模型质量主要与参数总数相关,并且不同的 scaling 方向(深度 L 、宽度 Dtokens T)产生几乎相同的性能。

    从计算效率的角度来看,更大的 hidden-dim 会产生更大的 matrix-multiply shapes,因此比堆叠更多层实现更高的 MFU。因此,100M1B 的最终配置分别设置为: (D=768,T=16,L=2) 和 (D=1536,T=32,L=2)。

1.3.3 消融研究

  1. RankMixer-100M 模型中,我们对 residual connectionsMulti-Head Token-Mixing 进行了消融研究。从 Table 2 中,我们可以看到:移除这些组件显著降低了模型的性能。

    • 移除 Multi-Head Token-Mixing 会丢失全局信息,因为每个 FFN 只建模部分特征而没有 interaction

    • 移除 residual connectionsLayerNorm 也会降低性能,降低训练稳定性,并使梯度爆炸或梯度消失问题更可能发生。

    这里缺乏 Sparse MoE 组件的消融研究:使用 Per-token Dense FFN 而不是 Per-token Sparse-MoE 的效果。

  2. 我们进一步分析了 token mixing 策略,即从 feature tokensFFNsrouting 策略,见 Table 3。与 Multi-Head Token Mixing 比较的 routing 策略包括:

    • All-Concat-MLP:将所有 tokens 拼接起来,通过一个大型 MLP 来处理,然后将它们分割成相同数量的 tokens 。性能下降显示了 learning large matrices 、以及削弱 local information learning 的挑战。

    • All-Share :不分割,整个 input 向量被共享并馈送到每个 per-token FFN,类似于 MoE。性能显著下降,这表明与 all- shared input 相比,feature subspace splitindependent modeling 的重要性。

    • Self-Attention :在 tokens 之间应用 self-attention 进行 routing 。其性能略逊于 Multi-Head Token-Mixing ,并且还遭受高计算成本的困扰,这表明了跨越数百个不同 feature subspaces 来学习相似性的困难。

1.3.4 Sparse-MoE Scalability 和 Expert Balance

  1. ScalabilityFigure 3 绘制了离线 AUC 增益与 Sparse-MoE 稀疏度的关系。

    • Dense-Training-Sparse-InferenceReLU routing 相结合,对于在 aggressive sparsity 下保持准确性至关重要,使 RankMixer 能够将参数容量(和内存占用)扩展 >8×,而 AUC 几乎没有损失,并且推理时间显著节省(吞吐量提高 +50%)。

    • 原始 Sparse-MoE 的性能随着 activated experts 数量的减少而单调下降,说明我们指出的 expert-imbalanceunder-training 问题。

    • 添加 load-balancing loss 相对于原始Sparse-MoE 减少了性能下降,但仍然不及 DTSI + ReLU 版本,因为问题主要在于 expert training 而不是 router。这验证了 Sparse-MoE 作为将 RankMixer 从当前的 1B 参数扩展到未来 10B 规模部署而不超出成本预算的路径。

  2. Expert balance and diversity:原始 Sparse-MoE 经常遭受 expert imbalance 的困扰,这反过来导致一些 experts under-trained ,最终导致 "dying experts"(几乎从未被激活的experts),并且只有少数固定的 expertsconstantly activatedFigure 4 显示,将 DTSIdense-training, sparse-inference )与 ReLU routing 相结合,可以有效地解决了这个问题:

    • dense-training 保证了大多数 experts 获得足够的梯度更新,防止 expert 饥饿。

    • ReLU routing 使得不同 tokensactivation ratio 是动态的——图中显示的 activation ratio 根据其信息内容自适应变化,这与 recommendation data 多样且高度动态的分布非常吻合。

1.3.5 Online Serving 成本

  1. 我们如何防止 inference latency 随着参数增加两个数量级而爆炸?在实际系统中,latency 与吞吐量成反比,与 serving machine resources 的成本成正比。与我们之前的 fully-deployed 16M-parameter model(采用集成了 DLRMDCN 的结构)相比,我们的 RankMixer 模型将参数扩展了约 70 倍至 1B。尽管参数显著增加,但由于我们 hardware-aligned 的模型设计和 optimization 策略,最终 inference latency 保持稳定。

    当大幅增加模型参数时,latency 可以分解为以下公式:

    Latency=Param Num×FLOPs/Param ratioMFU×(Theoretical Hardware FLOPs)

    Table 6 所示,两个数量级的参数增长被 3.6×decrease of FLOPs/Param ratio10×increase of MFUQuatization-based2× Hardware FlOPs improvements 所抵消。

    • FLOPs/Param ratioTable 6 的第三行报告了 per parameter 所需的 floating-point operations: FLOPs。由于模型设计,RankMixer 实现了 70 倍的参数增加,而 FLOPs 仅增加了约 20 倍,实现了仅为基线三分之一的 FLOPs/Param ratio3.6×efficiency gain。换句话说,在相同的 FLOPs 预算下,RankMixer 可以容纳比基线多三倍的参数。

    • Model FLOPs Utilization (MFU):同样 在Table 6 中显示,MFU 表示 machine computing 的利用率。通过使用 large GEMM shapes、良好的并行topology(将并行的 per-token FFNs 融合到一个 kernel 中)、减少内存带宽成本和开销,RankMixerMFU 提高了近 10× ,将模型从 Memory-bound 状态转变为 Compute-bound 状态。

    • Quantization:半精度(fp16inference 将使 GPU 的理论峰值硬件 FLOPS 提高 2×RankMixer 中的主要计算由几个大型矩阵乘法组成,如前所述,这些运算非常适合半精度(fp16)的 inference

1.3.6 Online 性能

  1. 为了验证 RankMixer 作为 scaling 的推荐模型框架的普适性,我们在个性化 ranking 的两个核心应用场景—— feed recommendationadvertising ——中进行了在线实验,涵盖了个性化 ranking 的主要用例。对于每个场景,我们监控以下关键性能指标:

    • Feed–RecommendationActive Days 是实验期间每个用户的平均活跃天数,是 DAU 增长的替代指标;Duration 衡量在 App 上的累计停留时间;Finish/Like/Comment 是用户的总完成播放、点赞、以及评论。

    • Advertising:我们报告模型质量指标(ΔAUC)和收入指标 ADVVAdvertiser Value)。

    在这些场景中,先前的基线是结合了 DLRMDCN16M 参数的模型。我们用 RankMixer-1B 模型替换了 Dense 部分,将 AUC 提高了 0.7%。我们在 Feed RecommendationAdvertising 上进行了 online A/B test 实验。Table 4 显示了 FeedRecommendation8 个月的 A/B test 的长期观察结果,Table 5 显示了 Advertising 的效果。

    RankMixer 在三个 personalised-ranking application 中部署和评估,包括 Feed RecommendationRankMixer-1B )和 AdvertisingRankMixer-1B )。RankMixer 在所有业务关键指标上都提供了统计显著的提升。我们还可以从 Table 4 中观察到,与其它 activeness level 的用户组相比,Low-active 用户的改进最大,Active Days 提高了超过 1.7412% ,这证明了模型的强大泛化能力。

    这些结果表明,RankMixer 作为一个 unified backbone ,能够可靠地泛化到不同的应用场景。

1.4 结论

  1. 在本文中,我们介绍了我们最新的 RankMixer 模型,该模型已完全部署在抖音 Feed ranking 中。它结合了用于 heterogeneous feature interactions 的模型设计和高并行化架构以实现 serving 效率。实验证明了其卓越的性能和陡峭的 scaling law 。在抖音 App 上全面部署后,它获得了 0.3%active days 增长、以及 1%App duartion 增长。