从RankMixer到TokenMixer-Large,细数字节推荐团队架构演进之路!
TokenMixer-Large: Scaling Up Large Ranking Models in Industrial Recommenders
论文链接:https://arxiv.org/pdf/2602.06563
近年来,大语言模型(LLMs)领域的 Scaling Law(扩展定律)可谓是大放异彩,无论是 GPT-4 还是最新的各大模型,都在向我们昭示一个朴素的真理:大力出奇迹,模型参数和计算量越大,效果越好。那么,在工业界搜广推(搜索、广告、推荐)领域,尤其是极其核心的排序模型(Ranking Models),是否也能直接套用并享受到 Scaling Law 的红利?
在字节跳动内部的最新研究《TokenMixer-Large》中,给出了极其硬核的回答:可以,但需要对模型架构进行极其严苛且符合底层硬件特性的重构设计!
今天这篇文章,我们要拆解的正是字节跳动在推荐系统大模型领域的一项重量级工作。这篇论文不仅仅是理论的堆砌,而是实打实地在一个高达 150亿(15B) 参数离线规模、70亿(7B) 参数在线规模的工业级推荐系统中,验证了排序模型的扩展潜力。在此之前,许多顶会论文(例如 Wukong、HiFormer 和 DHEN)都尝试过将推荐系统推向更深、更宽的网络结构,以此来探索推荐场景的 Scaling Law。但是他们往往受限于“不够优雅的残差设计”、“无法兼顾效果与算力效率”,也就是忽视了与硬件架构的适配。最终的结果就是,这些模型往往连亿级规模都难以突破就遭遇了天花板。
字节此前已经提出了 TokenMixer 的基础架构,并应用在 RankMixer 模型上,它极大地简化了 Transformer 的架构,将原本计算量庞大的 Self-Attention 替换成了轻量级的 Token 混合(Mixing)操作。然而,随着尝试将 RankMixer 扩展到更深的层数、更大的维度,作者团队发现了之前设计中潜藏的诸多限制:
- 残差设计的非最优性(Sub-optimal Residual Design):原版 RankMixer 在混合 Token 前后强行进行跨纬度的 Add 操作,不仅切断了语义传播,还导致语义错位,无法构建稳定极深的残差网络。
- 底层算子的冗余与拖累(Impure Model Architecture):为了堆叠性能,历史模型往往堆积了诸如 DCNv2、LHUC 等大量内存密集型的小算子,这些算子严重拉低了系统的 MFU(Model FLOPs Utilization,模型算力利用率)。
- 深层网络的梯度更新不足(Insufficient Gradient Updates):网络变深后,浅层参数往往难以获取有效梯度,产生训练停滞,这就和 CV 领域的 ResNet 诞生前夜的情况极为相似。
- MoE 稀疏化的痛点(Inadequate MoE Sparsification):以往的模型大多是采用“密集训练,稀疏推理(Dense Train, Sparse Infer)”的伪稀疏方案,或者使用了传统的 ReLU-MoE 导致路由专家极不平衡,根本无法显著降低训练阶段的高昂成本。
- 分布式系统扩展受限(Limited Scaling Exploration):当跨越多卡进行分布训练时,老版本的结构在通信上极为吃亏,难以扩展到 1B 参数以上。
面对以上五座大山,TokenMixer-Large 应运而生。它系统地解决了这些核心瓶颈,将模型的 MFU 推向了极致的 60%,并且在抖音的主信息流广告、电商以及直播业务中取得了卓越的在线 A/B 测试收益。本文非常长,因为包含的技术细节极其爆炸,我们将从三个核心切面,抽丝剥茧般带大家领略字节在推荐大模型上的这一最新杰作。
1. Overall Architecture:走向一致与对齐的征途
在推荐系统中,我们面临的输入特征往往是一个庞杂的大杂烩,既包括了用户的离散 ID (User ID)、候选物料 ID (Item ID),也包括了丰富的序列特征(如短期兴趣 DIN、长期兴趣 SIM、超长序列 LONGER 等),以及各式各样的交叉特征和连续特征。
构建大型推荐模型的首要难题,就是必须处理好这些来源不一、维度各异的高维稀疏输入,并将其转换为对底层 GPU 计算完全友好的稠密格式。在这里,最经典的路径是经过 Embedding 层。假设第 iii 个特征 FiF_iFi 被映射,其维度可能是随意的 did_idi,它的表达如下式所示:
ei=Embedding(Fi, di)∈Rdi \begin{equation} e_{i} = \text{Embedding}(F_i,~d_i) \in \mathbb{R}^{d_i}\end{equation} ei=Embedding(Fi, di)∈Rdi
但是这引发了一个致命的问题:如果不同特征对应的 Embedding 长度极其悬殊,我们要如何将它们组合成对并行计算友好的矩阵? 如果每一个特征单独占据一个 Token,那么面对数百个特征时,系统将会被切分得极其细碎。这会导致每个 Token 所分配的参数和计算量太少,不但重要特征被淹没,GPU 内核更是吃不饱(under-utilization)。而如果只用一个大 Token 把所有特征暴力拼接,那整个模型就退化成了最普通的 DNN,失去了对异构特征空间分别细致刻画的能力。
1.1 语义分组 Token 化(Semantic Group-wise Tokenizer)
针对这个问题,TokenMixer-Large 采用了一种极其充满工业智慧的做法:引入领域知识,基于语义进行对齐合并。
具体而言,模型并不是将每一个 eie_iei 当作独立的实体,而是根据它们的业务语义(例如:所有的上下文特征属于一组,所有的商品属性特征属于一组),将这些不同长度的 Embedding 串联拼接成一个个的特征语义簇 {G1, ..., GT−1}\{G_1,~ ...,~G_{T-1}\}{G1, ..., GT−1}。为了进行 GPU 的规整并行计算,每个组 GiG_iGi 内的向量拼接后,会被立刻送入一个专用的 MLP 网络,映射为固定统一的隐层维度 DDD。
这一步可以提炼为以下核心公式:
Xi=MLPi(concat[el, ..., em]), el,...,em∈Gi \begin{equation} \mathrm{X}_{i} = \text{MLP}_{i}(\text{concat}[e_{l},~...,~e_{m}]) , ~~~ e_{l},...,e_{m} \in G_{i}\end{equation} Xi=MLPi(concat[el, ..., em]), el,...,em∈Gi
这么做可谓一举两得。第一,计算被严密地规整了,所有通过 MLP 的结果 Xi\mathrm{X}_{i}Xi 在维度上呈现出极其完美的齐整性,大大便利了后续操作;第二,因为针对每个语义组使用了不同的 MLPi\text{MLP}_{i}MLPi 进行映射,特征的异构性(Heterogeneity)被原汁原味地保留了下来,而不是在一锅乱炖中丢失了空间结构。
1.2 全局聚合:引入 Global Token
但是,到这一步为止,还缺了点什么。如果所有的语义特征都被分散并隔离在一个个平行的 Token 通道中,那么当模型层数还较浅时,这些局部 Token 之间是缺乏“全局视野”的。
在自然语言处理(NLP)领域的大模型中,BERT 引入了 [CLS] 标签的精妙设计来汇聚整句话的语义。受此启发,TokenMixer-Large 的作者同样在排列的 Token 前端,引入了一个专门刻画全局信息的 Global Token(全局 Token)。
其计算方式简单且粗暴,但却非常有效——将所有语义组 GiG_iGi 重新聚合在一起,通过一个专门的全局映射网络 MLPg\text{MLP}_{g}MLPg,萃取出包含统领全局信息的向量 XG\mathrm{X_{G}}XG:
XG=MLPg(concat[G1, ..., GT−1]) \begin{equation} \mathrm{X_{G}} = \text{MLP}_{g}(\text{concat}\bigl[G_{1},~...,~G_{T-1}\bigr]) \end{equation} XG=MLPg(concat[G1, ..., GT−1])
至此,所有的准备工作均已就绪。TokenMixer-Large 送入堆叠网络架构的最原始输入 X\mathbf{X}X 就是这个融合了统御全局信息和各自局部深度语义信息的矩阵,如下式所示:
X=concat[XG,X0, ..., XT−1]∈RT×D \begin{equation} \mathbf{X} = \text{concat}\bigl[\mathrm{X_{G}},\mathrm{{X}}_0,~...,~\mathrm{{X}}_{T-1}\bigr] \in \mathbb{R}^{T \times D} \end{equation} X=concat[XG,X0, ..., XT−1]∈RT×D
其中 TTT 为最终整齐划一的 Token 数量,DDD 为模型隐层的统一维度。凭借这种矩阵级的整齐划一和包含全局概览的设计,为后续的超大规模堆叠(Stacking)创造了极其优良的生态。这仿佛就是在地基挖掘阶段,已经为修建数百层摩天大楼布置好了最稳固的钢筋混凝土承重墙。
可以说,从特征的稀疏海洋,到规整密集、内涵丰富的特征张量,上述的 Tokenizer 语义分组+Global聚合 设计,奠定了模型兼具体系结构效率与业务表征能力的坚实基础。
在理解和部署了前置的输入表征转换之后,大厦的基座已然完备,接下来便要进入本文真正的重头戏——让推荐模型深度扩展成为可能的 TokenMixer-Large Block 结构核心机制。在下一个 Part,我们将详细展开这部分是如何针对残差、梯度消失和高效网络进行拆解和跨越的!
2. TokenMixer-Large Block 与特征交互重构
在解决了输入的表征与维度对齐之后,特征矩阵将进入特征交叉交互的核心区域:TokenMixer-Large Block。
如果说 Tokenization 是地基夯实的过程,那么 TokenMixer 的每一次迭代,就是在进行钢筋的强韧交错绑扎。传统的 RankMixer 等模型在推进到深层时,暴露了极为致命的瓶颈——残差断裂导致的语义错位。这一切,都源自于残差网络设计的非最优性。

2.1 糟糕的残差语义错位
在传统的 RankMixer(或者第一代 TokenMixer)中,特征的信息交互(Mixing)通常非常简单暴力:
[…,[xt(0),…,xt(H)],… ]=split(X)∈RT×H×(D/H)Hh=concat[x1(h),…,xT(h)]∈RT×D/HH=concat[H1,…,Hh]∈RH×(T⋅D/H)Hnext=Norm(pSwiGLU(H)+H)∈RH×(T∗D/H) \begin{align} &[\dots,[\mathrm{x}_{t}^{(0)}, \dots, \mathrm{x}_{t}^{(H)}],\dots] = \text{split}(\mathbf{X}) \in \mathbb{R}^{T\times H\times(D/H)} \\ &\mathrm{H}_{h} = \text{concat}[\mathrm{x}_{1}^{(h)}, \dots, \mathrm{x}_{T}^{(h)}] \in \mathbb{R}^{T\times D/H} \\ &\mathbf{H} = \text{concat}[\mathrm{H}_{1}, \dots, \mathrm{H}_{h}] \in \mathbb{R}^{H\times (T \cdot D/H)} \\ &\mathbf{H}^{\text{next}} = \text{Norm}(\text{pSwiGLU}(\mathbf{H}) + \mathbf{H}) \in \mathbb{R}^{H\times (T * D/H)} \end{align} […,[xt(0),…,xt(H)],…]=split(X)∈RT×H×(D/H)Hh=concat[x1(h),…,xT(h)]∈RT×D/HH=concat[H1,…,Hh]∈RH×(T⋅D/H)Hnext=Norm(pSwiGLU(H)+H)∈RH×(T∗D/H)
这里隐藏了什么致命缺陷呢?仔细观察上式可以发现:网络最开始输入的 Token 数量是 TTT,而在经过了 split\text{split}split 和重组变换进行“跨 Token 特征 Mixing”操作之后,新生出的 Token 数量却变为了 HHH。
此时,如果 T≠HT \neq HT=H,或者每一层的输出并未精心设计以保持恒等数量,深层的残差直连通道就彻底断掉了!
更恐怖的是,纵然某些模型设计强制拉平了维度并勉强加上了加和残差(由 X′=F(X′)+XX' = F(X') + XX′=F(X′)+X 强行相加),但在这个相加操作发生时,F(X′)F(X')F(X′) 指代的是混合交互后(Mixed)的全新表征,而 XXX 则是未经过这层切分与混合重组的原始形态,两者的“语义锚点”可以说是完全脱节的。在这个残差加和里,TTT 个新特征加上 TTT 个旧特征,不仅无法做到深层的平滑传播,还会因为局部空间意义崩塌致使深层梯度彻底混乱!这便是导致老版本排序模型无法做深的罪魁祸首。
2.2 优雅解方:Mixing & Reverting 双向奔赴对称设计
针对这种跨层级语义坍塌的痼疾,TokenMixer-Large 给出的解决方案是构建一种严格前后对称的设计——Mixing & Reverting(混合与还原机制)。
不仅要“拆包”进行跨 Token 的交叉特征混合,混完之后,还要把拆开的行李严密地“缝合还原”回它初来乍到时的形状和意涵空间。这本质上是在每一个网络层的内部强行构造一个完全闭环的空间投射,从而使得任意两层外部的残差衔接变成了真正可以无缝叠加的恒等(Identity)映射。
核心数学操作极为精巧,展现了深厚的张量变换功底:
[...,[xt(1),xt(2),…,xt(H)],...]=split(X)∈RT×H×(D/H)Hh=concat[x1(h),x2(h),…,xT(h)]∈R(T∗D/H)H=concat[H1,H2,…,HH]∈RH×(T∗D/H)Hnext=Norm(pSwiGLU(H)+H)∈RH×(T∗D/H)[...,[x′1(h),x′2(h),…,x′T(h)],...]=split(Hnext)∈RT×H×(D/H)Xtrevert=concat[x′t(1),x′t(2),…,x′t(H)]∈RDXrevert=concat[X1revert,X2revert,…,XTrevert]∈RT×DXnext=Norm(pSwiGLU(Xrevert)+X)∈RT×D \begin{align} &[...,[\mathrm{x}_{t}^{(1)}, \mathrm{x}_{t}^{(2)}, \dots, \mathrm{x}_{t}^{(H)}],...] = \text{split}(\mathrm{\mathbf{X}}) \in \mathbb{R}^{T\times H\times(D/H)} \\ &\mathrm{H_{h}} = \text{concat}[\mathrm{x}_{1}^{(h)}, \mathrm{x}_{2}^{(h)}, \dots, \mathrm{x}_{T}^{(h)}] \in \mathbb{R}^{ (T*D/H)} \\ &\mathrm{\mathbf{H}} = \text{concat}[\mathrm{H}_{1}, \mathrm{H}_{2}, \dots, \mathrm{H_{H}}] \in \mathbb{R}^{H\times (T*D/H)} \\ &\mathrm{\mathbf{H}^{next}} = \text{Norm}(\mathrm{pSwiGLU}(\mathrm{\mathbf{H}}) + \mathrm{\mathbf{H}}) \in \mathbb{R}^{H\times (T*D/H)} \\ &[...,[\mathrm{x'}_{1}^{(h)}, \mathrm{x'}_{2}^{(h)}, \dots, \mathrm{x'}_{T}^{(h)}],...] = \text{split}(\mathrm{\mathbf{H}^{next}}) \in \mathbb{R}^{T\times H\times(D/H)}\\ &\mathrm{X_{t}^{revert}} = \text{concat}[\mathrm{x'}_{t}^{(1)}, \mathrm{x'}_{t}^{(2)}, \dots, \mathrm{x'}_{t}^{(H)}] \in \mathbb{R}^{D}\\ &\mathrm{\mathbf{X}^{revert}} = \text{concat}[\mathrm{X_{1}^{revert}}, \mathrm{X_{2}^{revert}}, \dots, \mathrm{X_{T}^{revert}}] \in \mathbb{R}^{T\times D} \\ &\mathrm{\mathbf{X}^{next}} = \text{Norm}(\mathrm{pSwiGLU}(\mathrm{\mathbf{X}^{revert}}) + \mathrm{\mathbf{X}}) \in \mathbb{R}^{T\times D} \end{align} [...,[xt(1),xt(2),…,xt(H)],...]=split(X)∈RT×H×(D/H)Hh=concat[x1(h),x2(h),…,xT(h)]∈R(T∗D/H)H=concat[H1,H2,…,HH]∈RH×(T∗D/H)Hnext=Norm(pSwiGLU(H)+H)∈RH×(T∗D/H)[...,[x′1(h),x′2(h),…,x′T(h)],...]=split(Hnext)∈RT×H×(D/H)Xtrevert=concat[x′t(1),x′t(2),…,x′t(H)]∈RDXrevert=concat[X1revert,X2revert,…,XTrevert]∈RT×DXnext=Norm(pSwiGLU(Xrevert)+X)∈RT×D
通过观察上面这组详尽的公式和推演,我们可以拆解决出其两个核心层面的推进战略:
第一层级(The First Layer - Mixing):承担 Token 表征间的高阶非线性信息混合。在此阶段,系统并不关注局部细节被如何解构,它在意的只是不同特征族裔(也就是切分后的 HhH_hHh)之间能否产生充分的碰撞。
第二层级(The Second Layer - Reverting):担任秩序收拢者。利用 concat 指令把分散开来的混编信息,按原有 T×DT \times DT×D 的坐标和形制还原。于是这时候形成的 Xrevert\mathbf{X}^{revert}Xrevert 和未进层池先的 X\mathbf{X}X 是完全平行同构的。这完美确保了跨模块的首尾一致性(Dimensional Consistency)。
\ Mixing 切分策略的迷思(Ablation on Mixing Strategy)
在这个 Mixing 过程中,将原始特征向量切分的策略是否会导致截然不同的表现呢?TokenMixer-Large 的研究团队同样对此进行了刨根问底地探究(如表格所示)。

令人意想不到的实验结论是——不论是垂直切分(Vertical division)、对角线切分(Diagonal division)还是随机切分(Random division),它们最终拿到的线上指标都惊人地一致!(增益误差 Δ\DeltaΔAUC均为 +0.00%+0.00\%+0.00% 区间内波动)。只有当模型强行削减一半的原始 Token 表征使其信息出现严重缺漏时,效果才发生了 −0.08%-0.08\%−0.08% 的陡然下降。
这给算法工程师带来了一个异常朴素却颠覆行业的哲学:不用迷信所谓的金牌切分方案!只要保证每次混合新生出的特征涵盖了充分的全部原始信息,那么混合本身就足够健壮。
2.3 Pertoken SwiGLU 专属专家系统:异质性建模
如果仔细看过上述还原结构,必然会发现其中嵌套了 pSwiGLU 这一环节。
原版 RankMixer 使用了全局共享的参数网络(Parameter-all-shared FFN),这意味着不管是点击特征、图像特征本身极具鸿沟的异构语义,所有的 Token 都被“一视同仁”的单层感知机“洗刷”过去,这抹杀了特征固有的异质化特征属性(Feature Heterogeneity)。
所以 TokenMixer-Large 取而代之使用了参数绝对隔离分配机制——Pertoken SwiGLU(每个 Token 私有化专属的非线性前馈模块):
pSwiGLU(⋅)=FCdown(Swish(FCgate(⋅))⊙FCup(⋅)), \begin{equation} \mathrm{pSwiGLU(\cdot)} = \mathrm{FC_{\text{down}} \left({\rm Swish} \left(FC_{\text{gate}} \left(\cdot \right) \right)\odot FC_{\text{up}} \left(\cdot \right) \right)}, \end{equation} pSwiGLU(⋅)=FCdown(Swish(FCgate(⋅))⊙FCup(⋅)),
拆解来看,其中所有的底层矩阵映射都被约束在一个私有化通道下:
FCi(x)=Witxt+bit,i∈{up,gate,down} \begin{equation} \mathrm{FC_i}(\mathbf{x}) = \mathrm{W}_i^t \mathrm{x}_{t} + \mathrm{b}_i^t, \quad i \in \{\text{up},\text{gate},\text{down}\} \end{equation} FCi(x)=Witxt+bit,i∈{up,gate,down}
这并非是简单的 MMoE(多任务专家)堆叠。MMoE 所有专家共用的是同一种拼盘混合输入,而 TokenMixer-Large 的每个 pSwiGLU,仅仅接收并对齐属于它特有维度的 Token 切片,相当于分配了上百位分门别类的极其垂类的审稿工匠,从源头上保留了各类先验意涵特征所展现出的分形复杂度。
2.4 Pre-Norm 阻断梯度雪崩(Normalization 位置的抉择)
大参数量带来最恐怖的梦魇便是:极其容易发生的梯度爆炸、激活值溢出(Numerical Explosion),最终往往用一个无情的 NaN 把模型训练断送。
模型堆叠时,必须对特征值进行规范化(Normalization)。在历史诸多结构采用 Post-Norm 的惨痛教训后,TokenMixer-Large 在 Normalization 的设计上进行了详尽的比对。

表格中触目惊心的实验结果昭示:
在极其深邃的网络中,试图依靠 Post-Norm 获得的那一丁点潜在理论提升(+0.01%+0.01\%+0.01%),付出的代价是毁灭性的数值坍塌(→NaN\rightarrow \textbf{NaN}→NaN)。而 Sandwich-Norm 的过渡反抗反而会导致性能的负增长。
因此,TokenMixer-Large 一律采取极为保守抗压的 Pre-Norm。不仅如此,还将繁重的 LayerNorm 统统被置换为更加灵活干练的 RMSNorm,因为抛弃了中心去均值的沉重包袱,在严密维持大盘收敛稳定的基础上,使得运行效率在端到端(End-to-end throughput)竟然大幅提升了 8.4% 之多。
2.5 深渊深处的引路灯:Inter-Residual & Auxiliary Loss
当我们把排序模型堆叠到夸张的 15B 并纵深推入百层时,浅层隐马尔可夫链往往陷入沉睡——顶层反馈的梯度经过一路的衰减,到达前中层时如同泥牛入海,前置网络的参数由于常年得不到有效更新而被废弃,网络形同虚设了。这就是深度学习噩梦的“梯度消失”。
为了让前中层网络同样保持强悍的活力,TokenMixer-Large 团队大刀阔斧构建了 跳层中间态残差连接(Inter-Residual Connections) 以及挂载在前置层的 辅助损失监督训练(Auxiliary Loss)。

如上图清晰所示的流程结构:
Inter-Residual(跳层间残差):设计每隔两三个交互 Block,便在更高维度进行一次跳层直连传送。使得底层的特征保留原始张量可以绕过中层的深锁拦截,直达更高层级,彻底打通梯度的回溯坦途。(不过论文作者精妙指出,绝不能将这种跳层一直延续到最终输出层,因为最高层的决策塔必须要承担高度抽象提纯的隔离工作,否则底层的嘈杂原始信号会干扰最终的预估目标逻辑。)
Auxiliary Loss(前置辅助损失):更具巧思的是,这是一种强迫基层神经元具备“提前猜想高频逻辑”能力的训练约束。模型将被要求由一部分较低层的输出特征,去和顶层推演完毕的高阶产物去联合预估目标。这就等同于在模型前半部分就强行注入了监督梯度,“拉拽打通”了那些因为处于深坑而无法吸收更新梯度的前序参数池,进而令得整个大模型表现出在各个层面参数皆“极强收敛”的可怕表现。
通过 Reverting+Mixing 双重闭环设计、垂类的独立空间解析、精准抗爆的 Pre-Norm 控制,辅以彻底唤醒全部神经元的跨网络 Aux Loss 输血通路,这些严苛得近乎洁癖的模型底层重构,终于扫清了阻碍深度残差网络堆叠一切技术障碍。这也是最终推荐大模型能够从几百 M 的小布包规模向 15B 这等天量怪兽跨越所给出的至关重要的“工业秘钥”!
3. 迈向15B稀疏超大模型:Sparse-Pertoken MoE 进化之路
当架构层面消退了导致层级崩溃的残差和梯度炸弹隐患,我们就能够非常顺畅地将模型参数叠加到数亿乃至几十亿。然而伴随而来更可怕的现实灾难是——如果模型所有参数都被稠密计算并且在线上直接提供实时的推断响应(Serving),随之而来的算力和显存显卡开销,足以令大厂都面临“破产的财务危机”。
为了解决工业推荐模型的大杀器“推断算力瓶颈”,TokenMixer-Large 给出的解法是:全盘进化,以 Sparse-Pertoken MoE(稀疏切分私有专家网络)取代稠密前馈系统,落实“稀疏训练+稀疏推断”。
3.1 First Enlarge, Then Sparse 的裂变主张
怎么将前文讲过的密集的专属 Token 网络 pSwiGLU 过渡到稀疏参数呢?这里使用的是被称为**“先放大,后割裂”(First Enlarge, Then Sparse)**的设计主张。首先依靠扩大稠密网络模型容量吃到 Scaling Law 的收益,随后我们不再给这无尽的膨胀买单了,而是把这个重型网络横向给拆分成大量的小个子子网络,并将它们称作“稀疏专家簇”。
S-P MoE(⋅)=∑j=1kgj(⋅)⋅Expertj(⋅), if j is chosen,Expertj(⋅)=FCdown, j(Swish(FCgate, j(⋅))⊙FCup, j(⋅)),FCi,j(xt)=Wi,jtxt+bi,jt, i∈{up,gate,down} \begin{equation} \begin{aligned} \mathrm{S\textbf{-}P~MoE(\cdot)} &=\sum_{j=1}^{k} g_{j}(\cdot)\cdot \text{Expert}_{j}(\cdot),~ \textbf{if j is chosen},\\ \text{Expert}_{j}(\cdot)&= \mathrm{FC_{\text{down,}~j} \left({\rm Swish} \left(FC_{\text{gate},~j} \left(\cdot \right) \right)\odot FC_{\text{up,}~j} \left(\cdot \right) \right)}, \\ \mathrm{FC_{i,j}}(\mathbf{x}_t) &= \mathrm{W_{i,j}^t} \mathbf{x}_t + \mathrm{b_{i,j}^t}, ~ i \in \{\text{up},\text{gate},\text{down}\} \end{aligned} \end{equation} S-P MoE(⋅)Expertj(⋅)FCi,j(xt)=j=1∑kgj(⋅)⋅Expertj(⋅), if j is chosen,=FCdown, j(Swish(FCgate, j(⋅))⊙FCup, j(⋅)),=Wi,jtxt+bi,jt, i∈{up,gate,down}
这等同于为每一个 Token 部署了庞大的专属顾问团,但我们通过 gj(⋅)g_{j}(\cdot)gj(⋅) 的路由决策,在海量的处理意见中,强制约束每次仅仅挑出 Top-K (通常是少量的几个)个子专家的计算结果去激活,其它专家统统陷入休眠。此时,这个拥有百亿参数容量的巨无霸模型,每次给出一个请求所需的浮点计算(FLOPs)居然被限制到跟一个十亿级小模型相若的极小水准里!
3.2 共享与缩放:榨干 MoE 每一点能力
然而,MoE 的路由并不总是随人愿的。在训练推荐网络早期,这种门控路由极容易发生“专家塌陷与死亡”:即大家都抢着吃那几口数据,而其他专家处于未被利用彻底死亡的状态。
为了兜底绝大部分常规公用知识并赋予其强劲稳定的更新梯队,TokenMixer-Large 引进了所有激活专家之外的兜底共用共享网络(Shared Expert)。
S-P MoE(⋅)=∑i=1k−1gi(⋅)⋅Experti(⋅)+SharedExperti(⋅) \begin{equation} \begin{aligned} \mathrm{S\textbf{-}P~MoE(\cdot)} =\sum_{i=1}^{k-1} g_{i}(\cdot)\cdot \text{Expert}_{i}(\cdot) + \text{SharedExpert}_{i}(\cdot) \end{aligned} \end{equation} S-P MoE(⋅)=i=1∑k−1gi(⋅)⋅Experti(⋅)+SharedExperti(⋅)
不仅如此,原本的 Router 因为被 softmaxsoftmaxsoftmax 函数所绑架裹挟,所有的路由分布概率之和强制为 111。这等同于把每次输入原本庞大的向后梯度强行均分和削减了。作者团队在此大胆进行改写,直接在分布输出外挂载一个放大常量 Gate Value Scaling α\mathbf{\alpha}α!通过人为倍卡扩充其乘数来保证梯度得以满载向底流传。
S-P MoE(⋅)=α⋅∑i=1k−1gi(⋅)⋅Experti(⋅)+SharedExperti(⋅) \begin{equation} \begin{aligned} \mathrm{S\textbf{-}P~MoE(\cdot)} = \mathbf{\alpha} \cdot \sum_{i=1}^{k-1} g_{i}(\cdot)\cdot \text{Expert}_{i}(\cdot) + \text{SharedExpert}_{i}(\cdot) \end{aligned} \end{equation} S-P MoE(⋅)=α⋅i=1∑k−1gi(⋅)⋅Experti(⋅)+SharedExperti(⋅)
这项优化被证实和模型当前的稀疏程度呈现出极其精妙的反比例法则!参见下面表格详尽的参数寻优推测:当稀疏度为 1:2 时,α\alphaα 取 2 是金标定;而稀疏度达到 1:4 时最佳 α\alphaα 就变为了 4。每一次 Scaling,都在精准捕获梯度对立的张量反馈。

3.3 Down-Matrix Small Init 拯救死亡初始态
在包含乘法逻辑和门控组合(上提和挤压组合的)的 SwiGLU 专家内部:
为了让模型初始迭代呈现更为完美的等价恒等连接以减弱对极度深奥参数的影响爆破,研究团队从 Rezero 思路中摄取了绝佳创意。对于 SwiGLU 的最后一层收拢矩阵 W_down,其 Xavier 初始化偏差值不采用传统的 1 比例尺度,而是强行将其收拢到仅仅 0.01!


这犹如给新生未开智巨兽戴上的一道温柔枷锁——前期计算中,所有的特征非线性映射反馈项通通近乎为 0。这就保证了在残差 F(X)+XF(X)+XF(X)+X 的迭代中,前几万步训练基本就是在做完全彻底的直连更新!然后随着训练步数递进,这只枷锁便慢慢解禁放宽并允许复杂非线性信息渗透。这就极大程度上完美根除了深层 MoE 结合时最要命的在随机权重未训练好前互相打架造成崩塌重绘的问题!
4. 从实验室到生产极限:工程与评测大赏
解决了高纬度的所有算法设计难题,真刀真枪落地的工业系统如果计算跟不上也是妄谈。
TokenMixer-Large 在端到端训练与在线推断中,研发了量身定做的高性能算子系统以及专供大模型使用的内存/分布式交互协议。
4.1 高性能专用算子 + 混合量化并行方案
首先在硬件维度中,将内存碎片化存取的麻烦利用特制的 MoEPermute 先一步拉直统合,然后由全局一统的 MoEGroupedFFN 一步下发(省去极其拖累性能的操作调度排队时间)。并在推理测运用了近乎极致的精度压缩方案 FP8 E4M3。

更为震撼的是,为了让庞大的分布卡集群在参数同步时不被网络 I/O 拉爆通信节点,架构内部摒弃了常规模型层层数据全互换那极其蠢笨的张量穿行,开创了 Token Parallel。通过只按照 Token 分发并不变异形状的穿刺,这极大地避免了 All-To-All 高昂的数据交互成本。
4.2 Scaling Law 与实验狂欢的凯歌
所有精细的打磨最后都在惊人的离线收益与实验横评中呈现出磅礴的统治力。这里是与其他所有 SOTA 顶会推荐模型的全维参数与算力降维对拉展现:

无论 DLRM-MLP 这种鼻祖,亦或近期大热的 HiFormer 和 DCNv2,以及上一任同源模型 RankMixer——TokenMixer-Large(不管是单纯在稠密结构下的 500M 版本,或者拉爆的几层极大版本 7B 参数、包含 4B SP-MoE版)都在 CTCVR 的指标提升下大幅领跑(对比 MLP 获得了最高 +1.20% 和 +1.14% 的增长)。
而且这并非仅仅停留下小尺度,它极其平滑且陡峭地维持住了性能爬坡。

更深层次探究它对最初模型 RankMixer 究竟作了何种程度的碾压性迭代,下述极其直观的设计模块查核表展现了这一点。那些导致语义和残差不匹配的根源:诸如原始 Token 记忆遗失(OTR)、Token 跨越时语义不同构(TSA)这些被 RankMixer 遗留的隐患全都变成了一路畅行的翠绿勾记 ✅。

最终,不管是在最核心也最为严酷的抖音核心信息流广告(Feed Ads 15B),亦或是对金钱转换最敏感的抖音电商场景(E-commerce 7B)、抑或是抖音直播互动场景(Live Streaming 4B)——这套宏伟架构全部重现了属于它的模型扩大化曲线规律(Scaling Laws),且在每个场景都无缝贴合着大力出奇迹的美感!



随着模型体积攀爬至这等量级,喂进数据的贪婪程度亦在增加。论文展示了抖音直播场景下各个尺度大模型由于容量陡增带来的收敛所耗费天数的恐怖增加!想要完全训练一只这样的怪物,至少需要给它高达 60 天以上数据投给,这足以窥见百亿级工业推荐系统内藏的浩瀚壁垒(见下表)。

4.3 终局:线上的盛大狂想
在线 A/B 测试永远会剥开花里胡哨的技术,直抵能够对业务线产生多大的商业价值这唯一金标。
TokenMixer-Large 被直接置换到拥有几亿日活的最疯狂前线中,部署替换后:
- 在 Feed 广告流:它将最终客户和广告商之间的连接提升了,让系统在维持极高转化下狂飙了 +2.0% ADSS(广告商满意分)。
- 在 直播电商:直接带动大涨了 +1.66%的订单量(Order)、以及对平台抽成至关重要提升的 +2.98%的GMV!
- 在 全场景直播:也贡献了 **+1.4% 的在线充值打赏支付金(Pay)**收益。

这等巨大的经济回馈和模型体量,都在这套将硬件架构与异质特征交叉完美统一的代码里,留下了工业应用落地的伟大回音。推荐大模型究竟存不存在 Scaling Law?在《TokenMixer-Large》被提出的那一刻,答案已被牢牢雕刻。
更多推荐




所有评论(0)