深度学习 | Yadong's Blog

NormFormer：用额外归一化改进 Transformer 预训练

Thu, 09 Apr 2026 02:30:00 +0800

本文是关于 Meta AI 论文《NormFormer: Improved Transformer Pretraining with Extra Normalization》（）的深度精读笔记。这篇论文通过在 Transformer 架构中添加三处额外的归一化操作，有效缓解了 Pre-LN 架构中各层梯度分布不均的问题，在几乎不增加计算成本的前提下，显著提升了语言模型的预训练效率和下游任务表现。

1. 引言：Transformer 归一化问题的前世今生

1.1 从 Post-LN 到 Pre-LN 的演进

自 Vaswani 等人在 2017 年提出 Transformer 架构以来，层归一化（Layer Normalization） 就是其中不可或缺的核心组件。归一化层的放置位置虽然看似只是一个微小的工程决策，但实际上对模型的训练稳定性、收敛速度和最终性能有着深远的影响。

原始 Transformer 采用 Post-LN 架构，即将 LayerNorm 放在残差连接之后：

$$\text{PostLN}(x) = \text{LayerNorm}(x + \text{Sublayer}(x))$$

这一设计在 Transformer 的早期应用中被广泛使用，但随着模型规模的不断增大，研究者们逐渐发现了它的致命缺陷：后层参数的梯度范数远大于早层，导致梯度爆炸于深层，梯度消失于早层。这意味着在反向传播过程中，靠近输入端的层几乎无法获得有效的梯度信号，导致训练极度不稳定。

为了缓解这个问题，Pre-LN 架构 应运而生，即将 LayerNorm 移到子层的输入端：

$$\text{PreLN}(x) = x + \text{Sublayer}(\text{LayerNorm}(x))$$

Pre-LN 架构被 GPT-2、GPT-3 等里程碑模型所采用，成为大语言模型预训练的事实标准。它显著改善了训练稳定性，使得大规模模型的训练成为可能。问题解决了吗？答案是否定的。

1.2 Pre-LN 的隐患：反向失配与梯度失衡

Xiong 等人在 2020 年的研究中首次系统性地揭示了 Post-LN 的梯度问题。然而，NormFormer 的作者们进一步发现，Pre-LN 虽然解决了 Post-LN 的训练不稳定问题，但实际上引入了方向相反的梯度失配：

Post-LN：后层梯度 » 早层梯度（梯度消失）
Pre-LN：早层梯度 » 后层梯度（反向失配）

为什么梯度不均衡很糟糕？

用一个工厂流水线的类比来理解：

梯度 = 改进信号，梯度大的层在训练中改变快，梯度小的层改变慢
早层是上游，负责提取基础特征。如果上游改进太频繁（梯度过大），下游来不及适应，导致训练不稳定
后层是下游，负责学习任务特定的高级表示。如果下游改进太慢（梯度不足），即使获得好的基础特征，也无法有效地组织成好的上层表示

梯度分布不均 = 各层学习效率不同 = 整体训练收敛慢、最终性能不佳

在 Pre-LN 架构中，这种现象体现为：

早期层过度更新：由于梯度过大，早期层的参数在训练初期可能剧烈波动，导致学到的特征不够稳定。
后期层更新不足：深层网络中最靠近输出的层本应承担最重要的任务特定表示学习，但它们接收到的梯度信号却相对不足，导致学习速度缓慢。

关键观察：论文通过可视化展示了这一现象：

Pre-LN 的梯度分布呈现明显的递减趋势，与 Post-LN 的递增趋势恰好相反。但两种趋势都是问题——因为两者都导致梯度分布不均。

1.3 梯度分布可视化：问题的直观表现

为了更直观地理解三种架构的梯度分布差异，我们用下面的图示展示各层的梯度范数变化：

图的含义：

x 轴：网络的第几层（从 1 到 12）
y 轴：该层参数的梯度 L1 范数（梯度大小）
红线（Post-LN）：梯度从浅层的很小逐渐增大到深层，形成"阶梯上升"。这导致深层参数更新剧烈，早层参数更新缓慢
蓝线（Pre-LN）：梯度从浅层的很大逐渐减小到深层，形成"阶梯下降"。这导致早层参数更新剧烈，深层参数更新缓慢
绿线（NormFormer）：梯度在各层基本保持一致的水平，形成"平坦"的分布。这导致各层参数更新速度均衡

这张图直观地说明了为什么 NormFormer 能提升训练效率：通过在三个精心选择的位置添加归一化操作，它实现了梯度在各层的均衡分布，使得每一层都能以相近的速率学习，从而提高了整个网络的训练效率。

1.4 研究动机：能否让梯度在所有层间均衡分布？

面对这一发现，NormFormer 论文提出了一个自然而直接的研究问题：能否通过在 Transformer 的关键位置添加额外的归一化操作，使得各层的梯度范数趋于均衡？

这个问题的提出源于一个简单但深刻的直觉：归一化操作本质上是对激活值进行重新缩放（rescaling），它天然具备调节梯度流动幅度的能力。如果我们能在正确的位置插入归一化层，就有可能同时解决 Pre-LN 的早层梯度过大和后层梯度过小的问题。

在接下来的章节中，我们将看到 NormFormer 如何通过三处精心设计的改进，实现这一目标。

2. 核心方案：三处改进与架构演进

NormFormer 的核心思想可以用一句话概括：在 Pre-LN Transformer 的基础上，在三个关键位置添加额外的归一化操作，通过精心的梯度控制实现各层学习效率的均衡化。

这三个操作分别是：Post-Attention LayerNorm（注意力后归一化）、HeadScale（注意力头缩放）和 FFN Mid-LayerNorm（前馈网络中间归一化）。它们的协同作用可以用下式表示：

$$x_{l+1}^{\text{NormFormer}} = \text{NormFFN}(\text{NormScaledMHA}(x_l))$$

在深入讲解具体的改进点之前，我们先给出 NormFormer、Pre-LN 和 Post-LN 三种架构的全景对比，帮助读者快速理解各架构的特点。

2.0 架构总览

三种架构的特性对比

特性	Post-LN	Pre-LN	NormFormer
归一化位置	残差连接之后	子层之前	多点分布式
训练稳定性	差（需精细 warmup）	好	更好（支持更高学习率）
梯度分布	后层 » 早层	早层 » 后层	各层趋于均衡
头级控制	无	无	有（HeadScale）
FFN 内部归一化	无	无	有
额外参数量	-	基准	+0.4%
额外训练开销	-	基准	+2~6%

关键观察：NormFormer 的核心创新在于在三个精心选择的位置添加归一化操作，以极小的代价（0.4% 参数 + 2~6% 计算）实现了梯度分布的均衡化。这种"微创手术"式的改进策略在理论上可以解决 Pre-LN 的梯度失衡问题，同时保留其训练稳定性优势。

三处改进的位置总览

NormFormer 的三个改进分别位于 Transformer 块中的不同位置：

Post-Attention LayerNorm：在多头注意力输出后、残差连接前
HeadScale：在多头注意力的拼接（concat）操作前，对每个头乘以可学习标量
FFN Mid-LayerNorm：在前馈网络的激活函数后、第二个线性层前

我们将逐个讲解每个改进的工作原理、消融验证数据和实际贡献度。

2.1 改进点 A：Post-Attention LayerNorm

工作机制

核心设计： 在多头注意力的输出后、残差连接前，添加一个额外的 LayerNorm。

标准 Pre-LN 的注意力子层为：

$$\text{PreLN-MHA}(x) = x + \text{MHA}(\text{LN}(x))$$

NormFormer 将其修改为：

$$\text{NormScaledMHA}(x) = x + \text{LN}(\text{HeadScaleMHA}(\text{LN}(x)))$$

这里外层的 $\text{LN}(\cdot)$ 包裹注意力输出。这个额外的归一化层起到了下缩放（downscaling） 的作用：论文发现训练完成后，所有层的 Post-Attention LN 的缩放参数（gamma）都低于 1，这意味着它在系统性地降低注意力输出的幅度。

为什么这很重要？ 在标准 Pre-LN 中，注意力层的输出直接通过残差连接加到主干上。如果注意力输出的幅度过大，会导致残差分支主导信号传播，破坏信息在不同层之间的平衡传递。通过添加这一归一化层，NormFormer 能够自适应地控制每一层注意力输出的贡献幅度。

原理图和实现

简洁实现示例（PyTorch）：

attn_output = self.attention(self.layer_norm(x))
attn_output = self.post_attn_layer_norm(attn_output) # 新增
x = x + attn_output

消融实验数据

在 125M 模型上的消融结果：

配置	PPL	相比完整 NormFormer 的变化	贡献度
完整 NormFormer	15.88	基准	-
移除 Post-Attn LN	15.92	+0.04	最小

贡献度评估：Post-Attn LN 的单独贡献最小（+0.04 PPL 退化），但作为整体方案的补充，它有助于进一步稳定梯度流。

2.2 改进点 B：HeadScale

工作机制

核心设计： 在多头注意力的拼接（concat）操作前，对每个注意力头的输出乘以一个独立的可学习标量参数。

传统的多头注意力将所有头的输出直接拼接后通过输出投影矩阵：

$$\text{MHA}(Q, K, V) = \text{Concat}(h_1, h_2, ..., h_n) W^O$$

NormFormer 引入了 HeadScale 机制：

$$\text{HeadScaleMHA}(Q, K, V) = \text{Concat}(\gamma_1 \cdot h_1, \gamma_2 \cdot h_2, ..., \gamma_n \cdot h_n) W^O$$

其中 $\gamma_i$ 为可学习的标量参数，初始化为 1，确保训练初期与标准多头注意力完全一致。

原理图和关键发现

关键发现：

头级权重差异化：训练后的 $\gamma_i$ 值变化较大，不同头获得了不同的缩放权重，这表明模型学会了动态调整不同注意力头的重要性。
无单调性约束：$\gamma_i$ 与层深度之间没有明显的单调关系，说明 HeadScale 不是简单地对深层或浅层进行统一调节，而是在细粒度上优化每个头的贡献。
最大贡献度：在消融实验中，HeadScale 是三个操作中贡献最大的。

实现示例

head_scales = nn.Parameter(torch.ones(num_heads)) # 初始化为1
scaled_heads = [head_scales[i] * heads[i] for i in range(num_heads)]
attn_output = torch.cat(scaled_heads, dim=-1) @ W_o

消融实验数据

在 125M 模型上的消融结果：

配置	PPL	相比完整 NormFormer 的变化	相对贡献度
完整 NormFormer	15.88	基准	-
移除 HeadScale	16.22	+0.34	69% 的总改进

贡献度评估：HeadScale 是三个改进点中贡献最大的，单独贡献占总改进（0.49 PPL 相比基线）的约 69%。（注：三个改进的占比之和超过 100% 是因为存在协同效应，详见第 2.5 节）这表明对注意力头进行差异化加权是 NormFormer 最核心的创新。

与注意力头剪枝的联系：HeadScale 的思想与注意力头剪枝（Head Pruning）有一定的联系。Chen 等人在 2021 年的工作中使用类似的头级缩放进行模型压缩，而 NormFormer 将这一思想用于改进训练过程，目标不同但技术路线相似。

2.3 改进点 C：FFN Mid-LayerNorm

工作机制

核心设计： 在前馈网络（FFN）的第一个线性变换之后、激活函数之后，添加一个 LayerNorm。

标准 FFN 的计算流程为：

$$\text{FFN}(x) = \sigma(x W_1 + b_1) W_2 + b_2$$

NormFormer 将其修改为：

$$\text{NormFFN}(x) = x + \underbrace{\text{LN}_{\text{mid}}}_{\text{新增}}(\sigma(\underbrace{\text{LN}_{\text{pre}}}_{\text{原有Pre-LN}}(x) \cdot W_1 + b_1)) \cdot W_2 + b_2$$

其中 $\text{LN}_{\text{pre}}$ 是 Pre-LN 架构原有的归一化，而 $\text{LN}_{\text{mid}}$ 是 NormFormer 新增的 FFN 中间归一化——它位于激活函数 $\sigma(\cdot)$ 之后、第二个线性变换 $W_2$ 之前。

关键机制：自适应梯度抑制

这是解决梯度失配的核心机制。 论文的 Figure 4 & 5 展示了一个极为重要的发现：

早期层的 FFN LN gamma 参数系统性地小于后期层的。这意味着 FFN Mid-LayerNorm 自适应地减小了早期层全连接层输入的幅度，从而有效降低了早期层的梯度，缓解了 Pre-LN 固有的"早层梯度过大"问题。这种设计妙处在于：

无需显式约束 gamma 值
模型在训练过程中自动学习最优的梯度分配方案
早层自然获得较小的 gamma（抑制梯度），后层获得较大的 gamma（放大梯度）

数学直觉

归一化操作通过将激活值映射到零均值、单位方差的分布来工作。当早期层的 FFN 中间激活值幅度较大时，归一化层通过较小的 gamma 参数对其进行压缩，相当于在反向传播时减小了通过这些层的梯度流。这种自适应机制使得模型能够自动学习到最优的梯度分配方案。

实现示例

# 原始 FFN
h = activation(x @ W1 + b1)
output = h @ W2 + b2

# NormFormer
h = activation(x @ W1 + b1)
h = self.ffn_layer_norm(h) # 新增
output = h @ W2 + b2

消融实验数据

在 125M 模型上的消融结果：

配置	PPL	相比完整 NormFormer 的变化	相对贡献度
完整 NormFormer	15.88	基准	-
移除 FFN-LN	16.14	+0.26	53% 的总改进

贡献度评估：FFN-LN 是第二大贡献者，单独贡献占总改进的约 53%。这验证了前馈网络内部归一化对梯度均衡的重要性，是仅次于 HeadScale 的核心创新。

2.4 可选改进：ResScale

除了上述三个核心操作外，NormFormer 还提出了一个可选的 ResScale 操作：

$$\text{ResScale}(x) = \lambda_{\text{resid}} \odot x + \text{Sublayer}(\text{LayerNorm}(x))$$

其中 $\lambda_{\text{resid}}$ 是可学习的逐维度缩放参数，用于调节残差连接中主干信号和子层输出的相对权重。

重要警告： 论文实验表明，ResScale 仅在小模型（125M、355M 参数）上有效，在 1.3B 及以上规模的模型上反而会导致性能下降。因此，对于当前主流的大规模预训练场景，不建议使用 ResScale。这一发现也提醒我们，并非所有的归一化/缩放操作都是"越多越好"的——过度参数化在大模型上可能导致优化困难。

2.5 三大改进的协同效应

前四个小节分别介绍了三个改进点的单独贡献，现在我们看它们是如何协同工作的。

完整消融实验对比

论文在 125M 模型上进行了系统性消融。注意：此规模下的"完整 NormFormer"包含 ResScale（因为 ResScale 仅在小模型上有正收益），而对于 1.3B 及以上的模型，推荐配置不包含 ResScale：

配置	PPL	相比完整模型的变化	单独贡献度
完整 NormFormer（含 ResScale）	15.88	基准	-
移除 Post-Attn LN	15.92	+0.04	最小
移除 FFN-LN	16.14	+0.26	第二大
移除 ResScale	16.20	+0.32	中等
移除 HeadScale	16.22	+0.34	最大
增加 QKV 上的 3 个额外 LN	15.88	+0.00	无收益
基线 Pre-LN	16.37	+0.49	相比 Pre-LN 的总改进

关键结论

HeadScale 是绝对核心（+0.34 PPL，占总改进 69%）：注意力头的差异化加权是 NormFormer 最具影响力的创新。
FFN-LN 是第二支柱（+0.26 PPL，占总改进 53%）：前馈网络内部的自适应梯度抑制对整体改进也至关重要。
Post-Attn LN 是补充设计（+0.04 PPL，占总改进 8%）：单独看贡献最小，但与其他改进配合使用时能进一步稳定梯度流。
三个改进的互补性：总改进 0.49 PPL ≈ HeadScale(0.34) + FFN-LN(0.26) - 重叠。这说明三个改进虽然作用机制不同（头级权重、梯度抑制、层级下缩放），但它们的效果有部分叠加，共同作用于梯度分布均衡化。

反面教训：更多归一化不一定更好

论文还测试了一个反例：在 QKV 投影上额外添加 3 个 LayerNorm——这看似应该进一步改进梯度分布，但实际结果是：

困惑度无任何改进（+0.00 PPL）
训练速度反而降低 5%

启示：NormFormer 选择的三个位置（Post-Attn、HeadScale、FFN Mid）是经过精心设计的，不是简单的"到处加 LN"。随意添加更多归一化操作反而会浪费计算资源，甚至破坏优化过程。

小模型 vs 大模型的推荐配置

基于消融实验的结果，我们给出不同规模模型的推荐配置：

模型规模	推荐配置	预期改进	说明
125M-355M	Post-Attn LN + HeadScale + FFN-LN + ResScale	最大	ResScale 在小模型上有正贡献
1.3B+	Post-Attn LN + HeadScale + FFN-LN	稳定	不使用 ResScale，避免大模型优化困难

3. 实验结果深度分析

NormFormer 论文的实验设计非常全面，涵盖了因果语言模型（CLM）、掩码语言模型（MLM）、零样本评估、消融实验等多个维度。下面我们逐一分析关键实验结果。

3.1 学习率搜索：挑战 GPT-3 的默认设置

在正式实验之前，论文做了一项非常有价值的预实验：系统性的学习率搜索。结果出人意料地发现，在他们的数据集上，最优学习率比 GPT-3 论文建议的值高出 3-5 倍：

模型规模	GPT-3 建议学习率	实际最优学习率	倍数
125M	6e-4	3e-3	5x
355M	3e-4	1e-3	3.3x
1.3B	2e-4	6e-4	3x

为什么 NormFormer 能支持更高的学习率？

这个现象与我们在改进点 C（FFN Mid-LayerNorm）讨论的梯度缩放密切相关。通过在 FFN 第一个线性层后添加 LayerNorm，NormFormer 实现了分层的梯度缩放 — 前馈网络内部的特征被重新归一化，避免了极端的激活值。这种分层缩放机制使得早层的梯度幅度自动降低，即使在更高的学习率下，参数更新也不会过于剧烈，从而显著提升了训练稳定性。

工程启示

这一发现本身就具有独立的工程价值——针对自己的数据集进行学习率搜索可能带来显著的性能提升，不要盲目套用论文中的超参数。论文使用这些优化后的基线作为对比对象，确保了实验结果的公平性和说服力。

3.2 因果语言模型：稳定且一致的困惑度改进

在因果语言模型（Causal Language Model）预训练任务上，NormFormer 在所有模型规模上都取得了一致的困惑度（Perplexity）改进：

模型	参数量	基线 PPL	NormFormer PPL	改进幅度
125M	124.5M	21.09	20.11	-0.98
1.3B	1313.5M	12.21	11.94	-0.27
2.7B	2649.5M	10.92	10.55	-0.37

核心发现

1. 训练加速效果显著。 NormFormer-1.3B 达到基线相同困惑度的速度快了 24%。也就是说，使用 NormFormer，你只需要原来 76% 的训练时间就能获得相同质量的模型。对于动辄需要数千 GPU 小时的大规模预训练来说，24% 的训练时间节省意味着巨大的计算成本削减。

与改进点的关联：这个加速来自于梯度均衡带来的更高的单步训练效率 — 梯度分布均匀意味着每一层都在以最优速率学习，没有某些层学得太快而其他层跟不上的浪费。这正是改进点 B（HeadScale）和改进点 C（FFN-LN）共同实现的梯度重均衡的直接体现。

2. 大模型训练稳定性提升。 这可能是 NormFormer 最引人注目的工程价值：基线 2.7B 模型在 6e-4 学习率下训练发散（完全失败），而 NormFormer-2.7B 在相同学习率下可以稳定训练并取得最佳性能。这意味着 NormFormer 显著拓宽了大模型可用学习率的范围，降低了超参数调优的难度。

与改进点的关联：这与改进点 B（HeadScale）的头级权重调整密切相关 — 通过对注意力头进行细粒度控制，NormFormer 抑制了某些头过度主导信号的现象。在大模型中，这种头间差异往往更加突出，因此 HeadScale 的稳定化效果在大模型上更加显著。

3. 困惑度改进随模型规模变化。 125M 模型上的绝对改进最大（-0.98），而大模型上的绝对改进较小。但考虑到大模型本身的困惑度已经很低（基数效应），相对改进幅度仍然有意义。更重要的是，训练加速和稳定性提升在大模型上同样甚至更加显著。

总结

因果语言模型的实验充分验证了 NormFormer 的核心改进在实践中的效果：梯度均衡带来的训练加速，头级控制带来的稳定性提升，使得 NormFormer 成为一个具有重大工程价值的改进方案。

3.3 零样本任务评估：无需微调即见效果

NormFormer 在零样本（Zero-Shot）任务评估中展现了明显的优势，直接反映了预训练质量的提升：

任务	基线-1.3B	NormFormer-1.3B	基线-2.7B	NormFormer-2.7B
HellaSwag	58.5	60.5	-	-
WinoGrande	76.8	77.5	-	-
平均（多任务）	63.6	64.7	66.3	68.7

关键发现：

NormFormer-125M 达到 GPT-3 Large（1.3B 参数）零样本性能的速度快了 60%。
在 2.7B 规模上，平均零样本准确率从 66.3% 提升到 68.7%（+2.4 个百分点）。

3.4 掩码语言模型：GLUE 基准全面提升

论文在掩码语言模型（Masked Language Model）上的实验验证了 NormFormer 对编码器型模型（BERT 类）同样有效。在 GLUE 基准的所有 7 个任务上都取得了改进：

GLUE 任务	基线	NormFormer	提升
CoLA	74.3	82.6	+8.3
MNLI	85.9	86.3	+0.4
MRPC	84.6	86.0	+1.4
QNLI	91.6	91.9	+0.3
QQP	90.7	91.3	+0.6
RTE	66.4	67.9	+1.5
SST-2	92.9	93.8	+0.9
平均	83.77	85.69	+1.92

最令人印象深刻的是 CoLA 任务上的 +8.3 提升（从 74.3 到 82.6），说明 NormFormer 对语法特征学习的显著帮助。GLUE 平均分提升接近 2 个百分点，MLM 困惑度从 3.42 降低到 3.31，展现了强大的通用性。

3.5 消融实验：各组件贡献量化

在第 2 部分，我们已经讨论了 NormFormer 的三个核心改进点（HeadScale、Post-Attn LN、FFN-LN）。这里提供完整的消融实验细节，量化每个组件的具体贡献。论文在 125M 小模型（470 V100 GPU 小时）上进行了系统性消融。注意：此规模下的"完整 NormFormer"包含 ResScale（因为 ResScale 仅在小模型上有正收益），而对于 1.3B 及以上的模型，推荐配置不包含 ResScale：

配置	Perplexity	相比完整模型的退化
完整 NormFormer + ResScale	15.88	基准
移除 Post-Attn LN	15.92	+0.04
移除 FFN LN	16.14	+0.26
移除 ResScale	16.20	+0.32
移除 HeadScale	16.22	+0.34（影响最大）
增加 3 个额外 LN（QKV 上）	15.88	+0.00（无额外收益）
基线 Pre-LN	16.37	+0.49

关键结论：

HeadScale 贡献最大（移除后退化 +0.34），说明注意力头的差异化加权是 NormFormer 最核心的创新。
FFN LN 贡献第二（+0.26），验证了前馈网络内部归一化对梯度均衡的重要性。
Post-Attn LN 贡献最小（+0.04），但仍有正面效果。
更多归一化并不总是更好：在 QKV 投影上额外添加 3 个 LN 没有带来任何性能提升，反而使训练速度降低 5%。这证明了 NormFormer 选择的三个位置是经过精心设计的，不是简单的"到处加 LN"。

3.6 超参数鲁棒性验证

NormFormer 的另一个重要优势是其对超参数设置的鲁棒性。论文在 125M 模型上测试了多种超参数组合：

学习率	配置	基线 PPL	NormFormer PPL	差值
0.001	默认	16.80	16.33	-0.47
0.003	默认	16.37	15.88	-0.49
0.003	更长 warmup	16.50	16.06	-0.44
0.003	GPT-3 设置	16.29	15.88	-0.41

NormFormer 在所有超参数配置下都一致优于基线，改进幅度在 0.41-0.49 之间波动，方差极小。这意味着使用 NormFormer 不需要额外的超参数调优工作——只要基线能跑，NormFormer 就能带来稳定的改进。

3.7 Wikitext-103 验证

论文还在 Wikitext-103 数据集上进行了验证：

模型	最终 Perplexity	达到基线 PPL 所需步数
基线	18.70	100%
NormFormer	18.65	70%（节省 30% 训练时间）

NormFormer 仅需 70% 的训练步数就达到了基线的最终性能。虽然后 30% 的训练中 NormFormer 的改进趋于饱和，但论文指出这可能通过进一步的训练策略调优来改善。

3.8 计算开销分析

NormFormer 的工程吸引力在于其极低的额外开销：

指标	数值
额外参数量	+0.4%（不足 0.07% 实际额外参数）
额外内存开销	+2~6%
单步训练时间增加	+2~6%
推理开销	接近零

这些数字意味着，NormFormer 实质上是一个"免费的改进"——用不到 6% 的额外计算成本，换取 24% 的训练加速和可量化的性能提升。从性价比角度看，这非常划算。

4. 工程应用与落地分析

4.1 实现极度简单

NormFormer 的工程实现可以说是所有 Transformer 改进方案中最简单的之一。只需要在现有 Pre-LN Transformer 代码中做三处修改：

修改一：在 MultiHeadAttention 输出后添加 LayerNorm

# 原始 Pre-LN
attn_output = self.attention(self.layer_norm(x))
x = x + attn_output

# NormFormer
attn_output = self.attention(self.layer_norm(x))
attn_output = self.post_attn_layer_norm(attn_output) # 新增
x = x + attn_output

修改二：在 MHA concat 前对每个 head 乘以可学习标量

# 原始 MHA
attn_output = torch.cat(heads, dim=-1) @ W_o

# NormFormer
head_scales = nn.Parameter(torch.ones(num_heads)) # 初始化为1
scaled_heads = [head_scales[i] * heads[i] for i in range(num_heads)]
attn_output = torch.cat(scaled_heads, dim=-1) @ W_o

修改三：在 FFN 第一个线性层后添加 LayerNorm

# 原始 FFN
h = activation(x @ W1 + b1)
output = h @ W2 + b2

# NormFormer
h = activation(x @ W1 + b1)
h = self.ffn_layer_norm(h) # 新增
output = h @ W2 + b2

在 fairseq 框架中，这三个修改对应三个简单的命令行参数：

fairseq-train ... --scale-attn --scale-fc --scale-heads

4.2 与主流框架的兼容性

NormFormer 的设计具有极强的框架兼容性：

与 PyTorch 原生 Transformer 兼容： NormFormer 的三处修改都是在现有层之间插入标准的 LayerNorm 或可学习参数，不改变任何现有层的接口或行为。这意味着它可以无缝集成到任何基于 PyTorch 的 Transformer 实现中。

与 HuggingFace Transformers 兼容： 只需继承现有的注意力层和前馈网络层，在对应位置添加归一化操作即可。不需要修改分词器、数据加载器或训练循环。

与分布式训练框架兼容： NormFormer 添加的归一化层和可学习参数都是标准的 PyTorch 模块，完全兼容 DeepSpeed、Megatron-LM、FSDP 等主流分布式训练框架。归一化操作的计算和通信开销极小，不会成为分布式训练的瓶颈。

与不同归一化方式兼容： 虽然论文使用 LayerNorm 进行实验，但其设计思想对 RMSNorm（LLaMA 系列使用的归一化方式）同样适用。可以将 NormFormer 中的 LayerNorm 替换为 RMSNorm，在保持核心优势的同时获得 RMSNorm 的计算效率优势。

4.3 实际部署场景与建议

场景一：从头预训练大语言模型

这是 NormFormer 最适用的场景。如果你的团队正在从头训练一个数十亿参数的语言模型，添加 NormFormer 可以：

节省约 24% 的训练时间（以达到同等困惑度为标准）
支持使用更高的学习率而不发散，降低超参数调优成本
以不到 6% 的额外计算开销换取稳定的性能提升

场景二：中等规模模型的快速迭代

对于 125M-1B 参数规模的模型，NormFormer 的收益更加明显。在这个规模上，可以同时使用 ResScale 获得最大收益。特别适合需要快速迭代模型架构和训练策略的研究场景。

场景三：训练稳定性要求高的场景

如果你的训练任务容易出现发散（例如使用较大的学习率、较长的上下文、较大的 batch size），NormFormer 可以显著提升训练的鲁棒性。2.7B 模型在高学习率下的稳定训练就是一个很好的例证。

4.4 不适用场景

也需要诚实地指出 NormFormer 可能不太适用的场景：

已有预训练好的模型进行微调： NormFormer 的收益主要体现在预训练阶段。如果你只是微调一个现有模型，添加 NormFormer 需要重新预训练，成本远大于收益。
极大规模模型（>10B）： 论文最大的实验只到 2.7B，对于更大规模模型的效果尚未被验证。虽然理论上应该同样有效，但缺乏实证支持。
推理优化敏感的场景： 虽然 NormFormer 的额外推理开销极小，但在对推理延迟有极致要求的场景下（例如实时搜索排序），任何额外的计算都需要审慎评估。

4.5 成本收益分析

让我们做一个简单的成本收益计算。假设你正在训练一个 1.3B 参数的语言模型：

成本（额外开销）：

训练速度降低约 4%（1.3B 规模的典型值）
如果原始训练需要 10000 GPU 小时，NormFormer 版本需要约 10400 GPU 小时

收益：

达到相同困惑度仅需 7600 GPU 小时（节省 24%）
最终困惑度从 12.21 降低到 11.94
零样本平均准确率从 63.6% 提升到 64.7%
更强的训练稳定性，降低训练失败的风险

净收益： 即使考虑单步训练时间的增加，要达到基线相同性能仍然可以节省约 20% 的总训练时间。如果以固定的计算预算训练到收敛，则获得更好的最终性能。NormFormer 都是一个值得采纳的改进。

4.6 与后续工作的关系

NormFormer 发表于 2021 年底，此后 Transformer 归一化领域继续涌现了许多重要工作：

RMSNorm（Root Mean Square Layer Normalization）：去掉了 LayerNorm 中的均值中心化步骤，计算效率更高。被 LLaMA 系列广泛采用。NormFormer 的设计理念可以与 RMSNorm 无缝结合。
QK-Norm：对注意力中的 Query 和 Key 进行归一化，防止注意力得分过大。与 NormFormer 的 HeadScale 有互补作用。
HybridNorm（2025）：探索了在同一模型中混合使用 Pre-LN 和 Post-LN 的可能性，进一步细化了归一化位置的选择。
nGPT（2024）：提出了基于单位超球面上表示学习的归一化方案，代表了归一化研究的新方向。

这些后续工作并没有否定 NormFormer 的价值，反而证明了"在 Transformer 中优化归一化策略"这一研究方向的重要性。NormFormer 作为这一领域的先驱工作之一，为后续研究奠定了重要的理论和实验基础。

5. 总结与展望

5.1 核心贡献回顾

NormFormer 论文的核心贡献可以用三句话概括：

发现了问题： 系统性地揭示了 Pre-LN Transformer 中各层梯度分布不均的问题——早期层梯度过大、后期层梯度不足。
提出了方案： 通过在三个精心选择的位置（注意力输出后、注意力头拼接前、FFN 中间层）添加归一化操作，有效缓解了梯度失配。
验证了效果： 在多种任务（CLM、MLM、零样本）和多种规模（125M-2.7B）上，以不到 6% 的额外计算成本换取了 24% 的训练加速和一致的性能提升。

5.2 对工程实践的启示

NormFormer 给我们的最大启示不仅仅是"加几个 LayerNorm"这么简单，而是：

启示一：小改进，大回报。 在深度学习研究中，并非所有有价值的工作都需要颠覆性的架构创新。有时候，对现有架构的精细分析和微小调整就能带来显著的实际收益。NormFormer 的三处修改总共只增加了 0.4% 的参数量，却换来了 24% 的训练加速——这种高性价比的改进在工业界尤其受欢迎。

启示二：梯度分析是优化训练的利器。 NormFormer 的整个工作建立在对梯度分布的细致观察之上。通过可视化和分析各层的梯度范数，研究者找到了问题所在，并据此设计了针对性的解决方案。这提醒我们，在训练大模型时，不要只盯着损失曲线，还应该关注梯度的层间分布。

启示三：不是所有改进都能无限叠加。 消融实验表明，在三个位置之外继续添加归一化层不仅没有收益，反而降低了训练速度。ResScale 在大模型上甚至有害。这告诉我们，模型改进需要有度，过度设计反而可能适得其反。

5.3 未来展望

尽管 NormFormer 已经展示了令人信服的实验结果，但仍有一些开放的研究方向值得探索：

超大规模验证： 论文最大的实验只到 2.7B 参数，NormFormer 在 10B、100B 甚至更大规模模型上的表现如何？是否存在新的问题或需要调整的地方？
与新型归一化的结合： 将 NormFormer 的设计理念与 RMSNorm、QK-Norm 等新技术结合，是否能获得更大的收益？
多模态扩展： NormFormer 目前主要在语言模型上验证，在视觉 Transformer（ViT）、多模态模型（如 Flamingo、GPT-4V）中是否同样有效？
自适应归一化： 能否设计一种机制，让模型在训练过程中自动决定在哪些位置需要额外的归一化，而不是人工预设固定位置？

总而言之，NormFormer 是一项兼具理论深度和工程价值的优秀工作。它用最简洁的方式解决了一个被忽视但重要的问题，为大规模 Transformer 预训练提供了一个即插即用的改进方案。对于正在从事大模型预训练的团队来说，NormFormer 值得认真评估和尝试。

参考文献：

Shleifer, S., Weston, J., & Ott, M. (2021). NormFormer: Improved Transformer Pretraining with Extra Normalization.

Xiong, R., et al. (2020). On Layer Normalization in the Transformer Architecture. ICML 2020.

Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017.

Zhang, B., & Sennrich, R. (2019). Root Mean Square Layer Normalization. NeurIPS 2019.

UniMixer：统一推荐系统三大架构范式的缩放定律探索

Wed, 08 Apr 2026 10:15:00 +0800

本文是关于快手技术团队论文《UniMixer: A Unified Architecture for Scaling Laws in Recommendation Systems》（）的深度精读笔记。这篇论文从理论上揭示了推荐系统中注意力机制、TokenMixer 和因式分解机三大架构范式的内在统一性，提出了参数化的 UniMixing 模块及其轻量化变体 UniMixer-Lite，并在快手广告系统上验证了清晰的缩放定律。

1. 引言：推荐系统特征交互架构的三条技术路线

1.1 从人工特征到深度模型的演进

推荐系统的核心挑战之一是特征交互建模——如何从用户画像、物品属性、上下文信息等异构特征中提取有效的交叉模式，以预测用户行为。

回顾推荐系统的发展历程，特征交互的建模方式经历了几个关键阶段：

手工特征工程时代（2010 年前）：依赖领域专家设计交叉特征，如"用户年龄 × 物品类别"。这种方式的问题在于人力成本高、覆盖度有限、无法捕获高阶交互。
因式分解机时代（2010-2016）：FM（Factorization Machine）及其变体（FFM、DeepFM）通过学习特征的隐向量表示来自动建模二阶交互 $\hat{y} = w_0 + \sum_i w_i x_i + \sum_{i\lt j} \langle v_i, v_j \rangle x_i x_j$。这开创了自动化特征交互的先河，但受限于交互阶数和表达能力。
深度网络显式交叉时代（2017-2022）：DCN、xDeepInt、AutoInt 等工作尝试通过显式的交叉网络结构捕获高阶交互。但这些方法往往带来复杂的网络设计和有限的扩展性。
Token 化与大模型时代（2023-至今）：受 Transformer 和大语言模型启发，推荐系统开始将特征视为 Token 序列，引入注意力机制和 TokenMixer 等操作进行特征交互。这一范式转变使得推荐模型具备了类似 LLM 的缩放潜力。

1.2 三条路线的分化与困境

进入 Token 化时代后，推荐系统的特征交互架构逐渐形成了三条主要的技术路线：

路线一：基于注意力机制（Attention-based）

代表工作：AutoInt、HiFormer、FAT、HHFT
核心思想：通过 Self-Attention 机制让每个 Token 动态地关注其他 Token，使用 Token-specific 的 Q/K/V 投影
优势：理论表达能力强，能捕获任意 Token 对之间的交互
问题：$O(L^2)$ 的计算复杂度，在特征数量 $L$ 较大时计算成本过高；更关键的是，在异构特征场景下，注意力权重容易变得尖锐稀疏，导致梯度回传受阻、训练停滞

路线二：基于 TokenMixer

代表工作：RankMixer（TokenMixer）、TokenMixer-Large
核心思想：通过固定规则的矩阵（如 Split & Concat）对 Token 进行混合操作
优势：计算效率高，参数无关（parameter-free），支持深层堆叠
问题：基于规则的混合模式缺乏可学习性和场景适应性；强制要求 Token 数等于 Head 数（T=H），限制了交互模式的选择空间

路线三：基于因式分解机（FM-based）

代表工作：Wukong、FinalMLP、GDCN、FiBiNet
核心思想：通过特征向量的内积或双线性交互建模特征对交互
优势：参数效率高，可解释性较好
问题：显式低阶交互约束限制了 Scaling 性能提升，难以扩展到更高阶

这三条路线各有优劣，但在工业实践中往往是"选边站队"——一个团队通常只深耕其中一条路线。这种割裂带来了两个核心问题：

缺乏统一的理论视角：无法回答"这三种方法到底有什么本质区别和联系？"
无法系统性地比较扩展效率：各方法在各自的实验设置下报告结果，缺少公平的缩放定律对比

1.3 UniMixer 的破局思路

UniMixer 论文的核心贡献在于回答了一个根本性问题：这三种看似不同的特征交互方式，是否存在统一的数学框架？

答案是肯定的。论文揭示了一个优雅的统一结构：

$$\text{UniMixing}(X) = \text{reshape}\left(G(X, W_G) \cdot [\text{local patterns}],\ 1,\ L\right)$$

在这个框架下，注意力机制、TokenMixer 和因式分解机只是全局权重 $G$ 和局部模式的不同实例化。这一发现不仅具有理论美感，更带来了实际的工程价值——既然三者本质相同，我们就能设计出集三者优势于一体的新架构。

2. 来源元数据 (Metadata)

原文标题: UniMixer: A Unified Architecture for Scaling Laws in Recommendation Systems
原文链接:
来源: arXiv（快手技术团队）
作者: Mingming Ha, Guanchen Wang, Linxun Chen, Xuan Rao, Yuexin Shi, Tianbao Ma, Zhaojie Liu, Yunqian Fan, Zilong Lu, Yanan Niu, Han Li, Kun Gai
发表日期: 2026 年 4 月

3. 核心摘要 (Executive Summary)

UniMixer 提出了推荐系统特征交互的统一架构框架，其核心创新包括：

理论统一：将注意力机制、TokenMixer 和因式分解机归纳为同一数学框架的不同特例
参数化 TokenMixer：将基于规则的 Token 混合操作转化为等价的参数化矩阵乘法，使混合模式可在训练中端到端优化
UniMixer-Lite：通过基矩阵组合和低秩近似，在大幅压缩参数量和计算成本的同时提升模型性能
SiameseNorm：引入耦合双流归一化解决深层架构的训练稳定性问题
缩放定律验证：在快手广告系统上验证了清晰的幂律缩放关系，UniMixer-Lite 的缩放指数（0.142）显著优于 RankMixer（0.116）

4. 深度解读 (Deep Dive)

4.1 参数化 TokenMixer：从规则驱动到数据驱动

UniMixer 论文最精彩的理论贡献之一，是揭示了 TokenMixer 操作的矩阵本质。

4.1.1 TokenMixer 的本质是什么？

以 RankMixer 为例，其核心操作是 Split & Concat——将 Token 序列按某种规则拆分后重新拼接。这看起来只是一种数据搬运操作，但论文指出：

任何 TokenMixer 的 Split & Concat 操作，都等价于将一个排列矩阵（Permutation Matrix）乘以展平的输入嵌入。

具体来说，对于输入 $X \in \mathbb{R}^{T \times D}$（$T$ 个 Token，每个维度为 $D$），TokenMixer 操作可以表示为：

$$\text{TokenMixer}(X) = \text{reshape}\left(W^{\text{perm}} \cdot \text{flatten}(X),\ H,\ \frac{TD}{H}\right)$$

其中 $W^{\text{perm}}$ 是一个排列矩阵。这个发现至关重要，因为：

排列矩阵具有双随机性（doubly stochastic）——每行每列恰好一个 1
排列矩阵是稀疏的——大部分元素为 0
当 $T = H$ 时，排列矩阵是对称的

4.1.2 从排列矩阵到可学习权重

既然 TokenMixer 本质上是矩阵乘法，一个自然的问题是：为什么要使用固定的排列矩阵，而不让模型自己学习最优的混合权重？

这正是 UniMixer 的核心思路——将硬编码的排列矩阵替换为可学习的权重矩阵 $W$，同时通过正则化约束保持排列矩阵的良好性质。更重要的是，参数化使得 T=H 的约束被彻底解除——传统 TokenMixer 强制要求 Token 数等于 Head 数，而参数化权重矩阵可以自由选择任意维度。

具体的约束实施方式：

双随机性约束：通过 Sinkhorn-Knopp 迭代实现

$$S_k(W) = D_r^{-1} W D_c^{-1}$$

其中 $D_r$ 和 $D_c$ 分别是行和列的归一化矩阵，交替迭代直至收敛。

稀疏性约束：通过温度系数 $\tau$ 控制

$$W_{\tau} = \text{softmax}(W / \tau)$$

当 $\tau \to 0$ 时，权重趋向 one-hot 分布，恢复排列矩阵的稀疏性。

对称性约束：通过显式对称化

$$W_{\text{sym}} = \frac{W + W^T}{2}$$

4.1.3 温度退火训练策略

直接使用低温度训练可能导致梯度消失（softmax 输出接近 one-hot 时梯度极小）。论文提出了温度退火（Temperature Annealing）策略：

$$\tau_j = \max\left\{\tau_{\text{start}} - \frac{(\tau_{\text{start}} - \tau_{\text{end}}) \cdot j}{J},\ \tau_{\text{end}}\right\}$$

训练初期使用较高温度（$\tau_{\text{start}} = 1.0$），让模型在接近均匀分布的空间中自由探索；随着训练推进，逐步降低温度至 $\tau_{\text{end}} = 0.05$，使权重逐渐收敛到稀疏的最优混合模式。

这种"先探索、后收敛"的策略与模拟退火算法有异曲同工之妙，有效避免了过早陷入局部最优。

4.2 统一理论框架：三大范式的殊途同归

4.2.1 统一公式

论文的核心理论贡献是将三种架构范式统一到同一个框架下：

$$\text{UniMixing}(X) = \text{reshape}\left(G(X, W_G) \cdot [\text{local patterns}],\ 1,\ L\right)$$

不同方法的差异仅在于全局权重 $G$ 和局部模式的选择：

方法	全局权重 $G$	局部模式
Self-Attention	$\text{softmax}\left(\frac{(XW_Q)(XW_K)^T}{\sqrt{d}}\right)$	$XW_V$
Heterogeneous Attention	$\text{softmax}\left(\frac{(X\tilde{W}_Q)(X\tilde{W}_K)^T}{\sqrt{d}}\right)$	$X\tilde{W}_V$
TokenMixer	$G$（固定置换矩阵）	$X$（恒等映射）
FM	$XI(XI)^\top$	$Y$
UniMixer	$W_G$（可学习全局混合）	$\{W_B^i\}$（可学习块权重）

这个统一视角揭示了三个深刻的洞察：

注意力机制是数据依赖的动态混合：全局权重 $G$ 由输入 $X$ 通过 Query-Key 机制动态计算，因此不同输入有不同的混合模式。这赋予了注意力极强的表达能力，但也带来了 $O(L^2)$ 的计算成本。
TokenMixer 是数据无关的静态混合：全局权重 $G$ 是固定的可学习参数，与输入无关。这使得计算效率极高（权重可以预计算），但牺牲了对不同输入的适应性。
因式分解机是特征相似度驱动的混合：全局权重 $G = XI(XI)^\top$ 由特征向量的内积决定，本质上是基于特征相似度的混合。这在二阶交互上效率很高，但难以扩展到更高阶。

4.2.2 Kronecker 积分解与计算优化

完整的 UniMixing 操作涉及 $L \times L$ 的权重矩阵（$L$ 为 Token 数量），在特征数量较大时计算成本不可接受。论文利用排列矩阵的 Kronecker 积结构进行分解：

$$W^{\text{perm}} = G \otimes I$$

这意味着全局混合矩阵可以分解为全局模式 $W_G$ 和局部模式 $W_B$ 的组合，将计算复杂度从 $O(L^2)$ 降低到：

$$O\left(\frac{L^2}{B} + LB\right)$$

其中 $B$ 是块大小。这种"全局-局部"的分层结构是 UniMixer 兼顾表达能力和计算效率的关键。

4.3 UniMixer-Lite：效率与性能的帕累托前沿

4.3.1 设计动机

完整的 UniMixing 模块虽然理论上优雅，但在工业部署中仍面临参数效率的挑战——每个块都需要独立的全局和局部权重矩阵。UniMixer-Lite 通过两项关键技术将参数量大幅压缩：

4.3.2 基矩阵组合（Basis Composition）

对于局部权重 $W_B^{(i)}$（第 $i$ 个块的局部混合矩阵），UniMixer-Lite 不再为每个块独立学习权重，而是通过一组共享基矩阵的线性组合动态生成：

$$W_B^{(i)} = \sum_{\ell=1}^{b} \omega_\ell^{(i)} Z_\ell$$

其中 $\{Z_\ell\}_{\ell=1}^{b}$ 是 $b$ 个共享基矩阵，$\omega_\ell^{(i)}$ 是第 $i$ 个块对各基矩阵的组合系数。

这种设计的精妙之处在于：

基矩阵在所有块间共享，参数量从 $O(L/B \cdot B^2)$ 降到 $O(b \cdot B^2)$
组合系数 $\omega_\ell^{(i)}$ 允许每个块有自己独特的混合模式
基矩阵数量 $b$ 远小于块数量 $L/B$，实现了参数的高效复用

4.3.3 低秩近似（Low-Rank Approximation）

对于全局权重 $W_G$，UniMixer-Lite 采用低秩分解：

$$W_G \approx W_r = A_G B_G$$

其中 $A_G \in \mathbb{R}^{(L/B) \times r}$，$B_G \in \mathbb{R}^{r \times (L/B)}$，秩 $r \ll L/B$。

这将全局权重的参数量从 $O((L/B)^2)$ 降到 $O(r \cdot L/B)$，在实践中 $r$ 通常取 4-8 即可达到接近全秩的效果。

4.3.4 Sinkhorn-Knopp 的保障作用

值得注意的是，即使使用了低秩近似和基矩阵组合来压缩参数，UniMixer-Lite 仍然通过 Sinkhorn-Knopp 操作确保权重矩阵保持接近满秩的双随机性质。这种"先压缩、后修正"的设计避免了参数压缩带来的表达能力损失。

4.3.5 性能表现

实验数据来自快手广告投放场景，超过 7 亿用户样本、一年数据，包含数百个异构特征。任务为用户留存预测（次日回访）。结果令人印象深刻：

模型	参数量	FLOPs	AUC	ΔAUC	UAUC	ΔUAUC
Heterogeneous Attention	132.7M	1.68T	0.7446	baseline	0.7338	baseline
RankMixer	135.5M	1.68T	0.7493	+0.475%	0.7389	+0.511%
UniMixer-2B	101.5M	2.50T	0.7502	+0.566%	0.7400	+0.615%
UniMixer-Lite-2B	76.2M	2.60T	0.7514	+0.682%	0.7412	+0.739%
UniMixer-Lite-4B	84.5M	4.24T	0.7527	+0.814%	0.7425	+0.870%

几个关键发现：

UniMixer-Lite-4B 仅用 84.5M 参数，AUC 提升 +0.814%，显著优于 135.5M 参数的 RankMixer
在推荐系统领域，AUC 提升 0.1% 即被视为显著改进，0.8% 是非常大的提升
UniMixer 的 FLOPs 高于 RankMixer（2.50T vs 1.68T），这是参数化带来的计算开销——但考虑到参数量的大幅减少和性能的显著提升，这一权衡在工业场景中是值得的

4.4 Pertoken SwiGLU 与完整架构

在 UniMixing 完成特征交互后，模型通过 Pertoken SwiGLU 建模不同特征的异质性：

$$\text{pSwiGLU}(o_i) = W_{\text{down}}^i \left( (W_{\text{up}}^i o_i + b_{\text{up}}^i) \odot \text{Swish}(W_{\text{gate}}^i o_i + b_{\text{gate}}^i) \right) + b_{\text{down}}^i$$

每个 Token 拥有独立的 FFN 参数（$W_{\text{up}}^i$, $W_{\text{gate}}^i$, $W_{\text{down}}^i$），充分建模不同特征领域（用户画像、物品属性、行为序列等）的异质性。这与 NLP 中所有 Token 共享 FFN 参数形成了鲜明对比。

完整的 UniMixer 模型由以下部分组成：

Feature Tokenization → 异构特征按领域分组，投影为统一维度的 Token 表示
M 层 UniMixer Block（含 SiameseNorm） → 层叠的统一混合模块
Sparse-Pertoken MoE → 稀疏混合专家进一步增强表达能力
预测头 → 输出最终预测

4.5 SiameseNorm：解锁深层架构的训练稳定性

4.5.1 深层推荐模型的训练难题

随着推荐模型向更深的层数扩展，一个经典的矛盾浮出水面：

Pre-Norm（归一化在子层之前）：有利于梯度流动和训练稳定性，但可能导致深层表示退化——所有层的输出趋于相似
Post-Norm（归一化在子层之后）：理论上能产生更丰富的层间表示差异，但在深层网络中容易出现梯度消失或爆炸

这个问题在 NLP 领域已有大量研究，但推荐系统的特征异构性（数值型、类别型、序列型特征共存）使得问题更加复杂。

4.5.2 SiameseNorm 的双流设计

UniMixer 引入了 SiameseNorm，其核心思想是维护两条耦合的信息流（$\bar{X}_\ell$ 和 $\bar{Y}_\ell$）。具体更新规则为：

$$\tilde{Y}_\ell = \text{RMSNorm}(\bar{Y}_\ell), \quad O_\ell = \text{UniMixer}(\bar{X}_\ell + \tilde{Y}_\ell)$$$$\bar{X}_{\ell+1} = \text{RMSNorm}(\bar{X}_\ell + O_\ell), \quad \bar{Y}_{\ell+1} = \bar{Y}_\ell + O_\ell$$

其中 $\bar{X}_\ell$ 承担类似 Pre-Norm 路径的角色——每次更新都经过 RMSNorm，保证训练稳定性；$\bar{Y}_\ell$ 承担类似 Post-Norm 路径的角色——直接累加输出，保持层间表示的多样性。两条流在每一层通过 $\bar{X}_\ell + \tilde{Y}_\ell$ 的融合进行耦合交互，兼具两种归一化方案的优势。

这种设计使得 UniMixer 能够同时在模型深度和宽度两个维度上进行有效扩展，而不会遇到训练不稳定的瓶颈。

4.6 缩放定律：推荐系统的"Chinchilla 时刻"

4.6.1 为什么推荐系统需要缩放定律？

在自然语言处理领域，Kaplan et al. (2020) 和 Hoffmann et al. (2022, Chinchilla) 发现了模型性能与参数量之间的幂律关系，这一发现深刻地改变了 LLM 的训练策略——从"盲目堆大"转向"计算最优"。

但推荐系统的缩放定律研究相对匮乏。其原因在于：

推荐模型的架构多样性远高于 LLM（Transformer 一统天下），难以进行公平对比
推荐系统的特征异构性（稀疏 ID 特征 + 稠密数值特征）使得"参数量"的定义不如 LLM 清晰
工业界的保密性导致大规模缩放实验的结果难以公开

UniMixer 论文的重要贡献之一，就是在统一框架下，为推荐系统建立了可对比的缩放定律基准。

4.6.2 幂律关系

论文验证了 AUC 增益与参数量之间的幂律关系：

$$\Delta \text{AUC} = a \cdot \text{Params}^{\alpha}$$

其中 $\alpha$ 是缩放指数，反映了架构的参数效率。实验结果：

架构	Scaling Law 公式	缩放指数 $\alpha$
RankMixer	$\Delta\text{AUC} = 0.002718 \cdot \text{Params}^{0.116}$	0.116
UniMixer	$\Delta\text{AUC} = 0.003032 \cdot \text{Params}^{0.132}$	0.132
UniMixer-Lite	$\Delta\text{AUC} = 0.003767 \cdot \text{Params}^{0.142}$	0.142

UniMixer-Lite 的缩放指数比 RankMixer 高出 22.3%，这意味着在相同的参数预算增长下，UniMixer-Lite 能获得更大的性能提升。更直观地说：

如果将参数量翻倍，RankMixer 的 AUC 增益提升约 8.4%，而 UniMixer-Lite 的提升约 10.3%。

随着模型规模的持续增长，这种差距会被不断放大。

4.6.3 深层 Scaling 的关键差异

更引人注目的是深层 Scaling 的对比实验：

模型	AUC	趋势
RankMixer-2B	0.7478	—
RankMixer-4B	0.7467 (-0.107%)	性能退化 ↓
UniMixer-Lite-2B	0.7492	—
UniMixer-Lite-4B	0.7508 (+0.158%)	持续提升 ↑
UniMixer-Lite-8B	0.7509 (+0.165%)	持续提升 ↑

这是一个极其重要的发现：RankMixer 增加深度后性能反而下降（-0.107%），而 UniMixer-Lite 在 8B 规模仍然展现出清晰的提升趋势。这证明了参数化混合 + SiameseNorm 的组合确实解决了深层推荐模型的 Scaling 瓶颈。

4.7 消融实验：每个组件的贡献

论文提供了详尽的消融实验，量化了每个设计选择的贡献：

设置	AUC	ΔAUC	影响程度
完整 UniMixer	0.7485	—	—
去除温度系数	0.7468	-0.165%	显著
去除模型预热	0.7476	-0.086%	显著
去除对称性约束	0.7479	-0.057%	中等
去除分块特异权重	0.7480	-0.044%	轻微

关键结论：

温度系数是最关键的组件（-0.165%），这验证了"先探索后收敛"的退火策略对于找到最优混合模式至关重要
模型预热同样重要（-0.086%），说明从高温初始化开始训练对避免局部最优有显著帮助
对称性约束和分块特异权重的影响相对较小，但仍然为正向贡献

5. 工程实践：从论文到生产的关键挑战

5.1 异构特征的 Token 化处理

工业推荐系统的输入特征高度异构——用户 ID（稀疏类别型）、用户年龄（数值型）、行为序列（变长序列型）、上下文时间（连续型）需要被统一为 Token 表示。UniMixer 的处理方式是将输入特征按领域组织（用户画像、物品特征、行为序列、Query 特征等），每个领域通过 Embedding 层转换为向量，再均匀划分为块并投影为 Token 嵌入：

$$x_i = W_i^{\text{proj}} E_{di:di+d} + b_i^{\text{proj}} \in \mathbb{R}^D$$

这种分组 Token 化方案使得不同类型的特征被映射为统一维度的 Token 向量，为后续的 UniMixing 操作提供了一致的输入格式。

5.2 Sinkhorn-Knopp 迭代的计算开销

双随机约束的 Sinkhorn-Knopp 迭代在理论上需要无穷步才能严格收敛，但实践中论文发现 5-10 次迭代即可达到足够的精度。这一开销在训练时是可接受的，而在推理时权重已经固定，无需额外迭代。

5.3 温度退火的超参数选择与冷启动策略

温度退火策略引入了三个超参数：起始温度 $\tau_{\text{start}}$、终止温度 $\tau_{\text{end}}$ 和退火步数 $J$。论文推荐的默认值为 $\tau_{\text{start}} = 1.0$、$\tau_{\text{end}} = 0.05$。一个实用的经验法则是将退火步数设置为总训练步数的 60-80%，让模型在训练的最后阶段以稳定的低温度进行精调。

对于数据不足的场景，论文还提出了"冷启动"策略：先用高温度完成一轮完整训练，然后用高温训练得到的权重作为初始化，再进行低温度的重训。这种两阶段方法可以在数据有限的情况下依然获得良好的稀疏权重。

5.4 快手广告系统的部署实践

UniMixer 和 UniMixer-Lite 已在快手的多个广告投放场景中完成部署。论文报告了在线 A/B 测试的结果，以 30 天累计活跃天数（CAD, Cumulative Active Days） 为核心评估指标：

D1-D30 的 CAD 平均提升超过 15%

同时离线 AUC 提升 +0.814%（UniMixer-Lite-4B vs 基线）。在推荐系统领域，AUC 提升 0.1% 就已经是显著的改进，0.8% 的提升幅度在工业界是极为罕见的。这一结果有力地证明了统一架构的实际业务价值。

6. 与相关工作的对比分析

6.1 与 TokenMixer-Large（字节跳动）的对比

TokenMixer-Large 是字节跳动提出的推荐系统大模型架构，与 UniMixer 形成了有趣的技术路线对比：

维度	TokenMixer-Large	UniMixer
核心贡献	工程优化——通过纯净架构和 MoE 扩展至 15B	理论统一——揭示三大范式的等价性
扩展策略	深度扩展（Mixing & Reverting + 跨层残差）	参数效率扩展（基矩阵组合 + 低秩近似）
混合方式	基于规则的 Split & Concat	参数化的可学习混合矩阵
训练稳定性	Rezero 初始化 + 辅助损失	SiameseNorm 双流设计
稀疏化	Per-token MoE	温度退火的稀疏权重
验证规模	最大 15B 参数	聚焦于 100M 级别的缩放定律

两篇论文实际上代表了推荐系统大模型的两种互补思路：

TokenMixer-Large：以"工程至上"为原则，通过极致的架构优化和 MoE 稀疏化，在既有的 TokenMixer 框架内推到了 15B 的参数规模
UniMixer：以"理论先行"为原则，通过统一框架发现了更高效的参数化空间，用更少的参数实现了更好的缩放效率

6.2 与 Attention-based 方法的对比

UniMixer 的统一框架揭示了一个有趣的事实：注意力机制可以被视为"输入依赖的 UniMixing"，而 UniMixer 使用的是"输入无关的可学习权重"。这意味着：

注意力机制的表达能力理论上更强（动态权重 vs 静态权重）
但 UniMixer 的计算效率更高（权重可预计算，无需在线计算 QK）
在推荐系统的实际场景中，UniMixer 的缩放效率反而更好——这可能说明推荐系统中"最优的混合模式"并不需要随输入动态变化

6.3 与 FM-based 方法的对比

因式分解机方法在 UniMixer 的框架下被理解为"全局权重由特征相似度决定"的特例。UniMixer 的改进在于：

FM 的全局权重 $G = XI(XI)^\top$ 完全由输入决定，没有可学习参数
UniMixer 的全局权重 $W_G$ 是可学习的，能发现数据中不直接由相似度反映的交互模式
UniMixer-Lite 的基矩阵组合可以看作是"参数化的 FM"——每个块的局部交互模式由共享基矩阵动态组合

7. 流程图 (Flowchart)

graph LR A["异构特征"] --> B["Token 化"] B --> C["UniMixing"] C --> D["Pertoken SwiGLU"] D --> E["SiameseNorm"] E --> F["Sparse MoE"] F --> G["深层堆叠"] G --> H["预测输出"] style C fill:#4ecdc4,color:#fff style E fill:#45b7d1,color:#fff style F fill:#f7dc6f,color:#333

UniMixing 内部结构：

graph LR X["输入 X"] --> WG["全局权重 W_G
低秩近似"] X --> WB["局部权重 W_B
基矩阵组合"] WG --> SK["Sinkhorn-Knopp
双随机约束"] WB --> SK2["Sinkhorn-Knopp
双随机约束"] SK --> MIX["UniMixing 交互"] SK2 --> MIX MIX --> OUT["输出"] style SK fill:#4ecdc4,color:#fff style SK2 fill:#4ecdc4,color:#fff style MIX fill:#45b7d1,color:#fff

8. 优缺点分析

8.1 优势

理论贡献突出：首次将推荐系统三大特征交互范式统一到同一数学框架下，为后续研究提供了清晰的理论基础和统一的对比基准。
参数效率卓越：UniMixer-Lite 以不到一半的参数量超越全参数基线，缩放指数提升 22.3%。这意味着在相同的计算预算下，UniMixer 能获得更大的性能收益。
设计优雅：温度退火 + Sinkhorn-Knopp 约束的组合，使得可学习权重能在保持排列矩阵良好性质的同时，找到最优的混合模式。这种"在约束空间中优化"的思路具有广泛的启发意义。
工业验证充分：在快手广告系统的多个场景中完成部署，离线 AUC +0.814%、在线 30 天 CAD +15% 的效果证明了方案的实际业务价值。

8.2 局限与待探索方向

缩放实验规模有限：与 TokenMixer-Large 的 15B 参数实验相比，UniMixer 的缩放验证主要集中在 100M 参数级别。UniMixer-Lite-8B 虽然仍有提升，但增量已经放缓（+0.165% vs +0.158%），更大规模下是否仍能保持优势有待验证。
计算成本权衡：UniMixer 的 FLOPs 显著高于 RankMixer（2.50T vs 1.68T），这意味着参数化带来的性能提升需要以额外的计算开销为代价。在延迟敏感的在线服务场景中，这一权衡需要仔细评估。
动态权重的缺失：UniMixer 使用的是输入无关的静态权重，而注意力机制的核心优势在于动态权重。能否在保持计算效率的前提下引入有限的动态性（如条件化的权重调制），可能是一个有价值的研究方向。
公开数据集验证缺失：论文仅在快手内部数据集上进行了验证，缺乏公开数据集上的对比实验，这限制了社区对结果的复现和验证。
多任务场景的验证：论文主要聚焦在用户留存预测任务上，在其他推荐任务（CTR、CVR、停留时长等）和多任务学习场景下的表现尚待验证。
与序列建模的结合：当前的 UniMixing 主要处理特征间的交互，如何与用户行为序列建模（如 DIN、SIM）有效结合，也是一个值得探索的方向。

9. 未来展望：推荐系统架构的"大一统"趋势

9.1 从架构统一到训练范式统一

UniMixer 完成了特征交互层面的架构统一。下一步自然是训练范式的统一——能否将 CTR 预估、召回、排序等不同阶段的模型统一到同一个框架下？Spotify 的 NEO 论文已经在搜索与推荐的统一上迈出了一步，而 UniMixer 的统一框架可能为推荐系统内部的多阶段统一提供理论基础。

9.2 缩放定律指导资源分配

UniMixer 建立的缩放定律对比框架，为工业界的模型选型和资源分配提供了量化依据。在实际决策中，团队可以根据缩放指数估算"投入 X 倍计算资源，预期获得 Y% 的性能提升"，从而做出更理性的投资决策。

9.3 参数化混合的更广泛应用

“将规则驱动的操作参数化，同时通过约束保持原始性质"这一思路，可以推广到推荐系统的其他组件——例如特征选择、样本加权、多目标融合等。这些传统上依赖人工规则或简单启发式的模块，都可能从参数化中获益。

9.4 深度与效率的帕累托前沿

SiameseNorm 解决了深层训练的稳定性问题，而 UniMixer-Lite 的基矩阵组合和低秩近似则压缩了参数量。未来的研究可以进一步探索：在给定的延迟预算下，如何在模型深度、宽度和稀疏度之间找到最优的帕累托前沿？

10. 总结

UniMixer 论文的价值不仅在于提出了一个新的推荐系统架构，更在于建立了一个统一的理论视角，让我们重新审视过去十年推荐系统特征交互技术的发展脉络。

回顾全文：

三条技术路线的殊途同归：注意力机制、TokenMixer 和因式分解机看似是三种截然不同的特征交互方式，但在 UniMixing 框架下只是全局权重和局部模式的不同选择
参数化带来的红利：将规则驱动的 TokenMixer 转化为可学习的参数化形式，配合温度退火和 Sinkhorn-Knopp 约束，实现了更高效的缩放
从理论到实践的闭环：统一框架不仅具有理论美感，更通过 UniMixer-Lite 和 SiameseNorm 的工程设计转化为实际的性能收益

在推荐系统向大模型时代迈进的过程中，UniMixer 提供了一个重要的启示：真正的扩展效率提升，往往不是来自简单地堆叠更多参数，而是来自对架构本质的深刻理解。当我们理解了"为什么这样做有效”，才能找到"如何做得更高效"的答案。

TokenMixer-Large: 突破工业级推荐系统的大模型扩展瓶颈

Wed, 18 Mar 2026 10:15:00 +0800

来源元数据 (Metadata)

原文标题: TokenMixer-Large: Scaling Up Large Ranking Models in Industrial Recommenders
原文链接:
来源: Arxiv (ByteDance 团队)
作者: Yuchen Jiang, Jie Zhu, Xintian Han, Hui Lu, Kunmin Bai, Mingyu Yang, Shikang Wu 等

核心摘要 (Executive Summary)

针对工业级推荐系统面临的大模型扩展瓶颈，本文提出了 TokenMixer-Large 架构，通过引入“Mixing & Reverting”操作、层间残差、辅助损失以及稀疏 Per-token MoE 等一系列创新，解决了深层网络中的梯度消失、MoE 稀疏化不足以及硬件利用率低等问题，在字节跳动的核心业务（电商、广告、直播）中成功扩展至百亿参数规模，并取得了显著的在线业务增长。

深度解读 (Deep Dive)

核心痛点

随着推荐系统大模型（DLRM）尝试向大规模参数扩展，现有的主流架构（如 RankMixer、Wukong、DHEN）在实际应用中暴露出多个严重瓶颈：

次优的残差设计: RankMixer 等架构通过 Mixing 操作改变了 Token 的维度和数量，导致前后残差连接时 Token 的语义无法对齐，限制了模型的表现上限。
不纯粹的模型架构: 由于历史迭代，推荐模型中通常保留了许多琐碎、访存密集型的底层算子（如 LHUC、DCNv2），导致整体模型的计算利用率 (MFU) 极低。
深层网络梯度更新不足: 传统的 TokenMixer 往往只有浅层配置（如 2 层），随着网络加深，梯度消失问题严重，难以保持训练稳定性。
MoE 稀疏化不足: 原有的 ReLU-MoE 设计局限于“稠密训练、稀疏推理”范式，并未降低训练成本，且动态激活机制对推理极不友好。
扩展性受限: 受限于上述原因，工业界之前的探索仅止步于 10亿（1B）参数级别。

方法论 (Methodology)

1. 整体设计思路

TokenMixer-Large 的设计哲学可以用一句话概括：以”纯净架构”为基座，通过深度残差与稀疏化实现工业级大模型的高效扩展。具体来说，团队遵循了以下三条核心设计原则：

架构纯净化 (Architecture Purification)：移除所有历史遗留的碎片化算子，仅保留高计算密度的矩阵乘法操作，最大化 GPU 的 MFU（Model FLOPs Utilization）
残差对齐化 (Residual Alignment)：通过 Mixing-Reverting 的对称设计，确保跨层残差连接的语义一致性，为深层网络训练铺平道路
稀疏高效化 (Sparse Efficiency)：采用 Per-token MoE 实现真正的”稀疏训练+稀疏推理”，在保持模型容量的同时大幅降低计算开销

2. TokenMixer-Large 与初代 TokenMixer (RankMixer) 的核心区别

在理解 TokenMixer-Large 的创新之前，我们必须先看清它对初代架构做了哪些大刀阔斧的”革命”。以下是四个核心差异点：

2.1 从“维度错位”到“Mixing & Reverting 绝对对齐”

初代 TokenMixer 的痛点：在进行 Token 混合（Mixing）时，将 $T$ 个 token 强行变为 $H$ 个。输入输出维度不匹配，导致直接加和时产生语义错位，无法实现有效的跨层残差连接。
Large 版本的解法：设计了高度对称的“双层结构”：
- Mixing 层：负责跨 Token 混合信息 ($T \rightarrow H$)。
- Reverting 层：专门将混合后的 Token 维度完美恢复到原始状态 ($H \rightarrow T$)。这种设计确保了输入和输出维度的绝对一致性，构建出平滑且语义对齐的深度残差通道。

# 伪代码演示
# 输入 X: [T, D], T为Token数, D为维度

# 1. Mixing 阶段
H = Split_and_Concat(X) # 将 T 个 token 混合为 H 个, 维度变为 [H, T*D/H]
H_next = Norm(pSwiGLU(H) + H)

# 2. Reverting 阶段
X_revert = Split_and_Concat_Back(H_next) # 将 H 个 token 还原为 T 个, 维度恢复为 [T, D]
X_next = Norm(pSwiGLU(X_revert) + X) # 语义严格对齐的残差连接

2.2 从”碎片化算子堆砌”到”纯净架构 (Pure Architecture)”

初代 TokenMixer 的痛点：由于历史迭代，模型中通常堆砌了许多细碎、访存密集型的底层算子（如 LHUC、DCNv2），导致整体模型在 GPU 上的计算利用率（MFU）极低。
Large 版本的解法：剥离所有底层低效交互算子，将 Post-LayerNorm 替换为 Pre-RMSNorm，使用 pSwiGLU 替换 pFFN。完全依靠堆叠纯净的 TokenMixer-Large Block 进行特征交叉，使得核心广告模型的 MFU 飙升至 60%。

关于组件替换的详细说明：

Pre-RMSNorm vs Post-LayerNorm：Pre-RMSNorm 将归一化操作放在子层的输入端而非输出端，省去了均值计算步骤，减少了约 30% 的归一化开销。更重要的是，Pre-Norm 结构使得残差通道中的梯度流动更加顺畅，有利于深层网络的训练稳定性。
pSwiGLU vs pFFN：pSwiGLU（Per-token SwiGLU）将传统的 ReLU 激活替换为 SiLU（Swish）门控线性单元，引入了乘法门控机制，增强了特征的非线性表达能力。其公式为 $\text{SwiGLU}(x) = (xW_1) \otimes \text{SiLU}(xW_2)$，相比传统 FFN 增加了约 50% 的参数量，但带来的效果提升远超参数增长。

2.3 从”浅层堆叠”到”深层跨层残差 (Inter-Layer Residuals)”

初代 TokenMixer 的痛点：随着网络加深（如从浅层的 2 层扩展到深层），极易发生梯度消失现象。
Large 版本的解法：采取了”组合拳”：
- 跨层残差与辅助损失：每隔 2-3 层引入跨层残差连接，并将底层输出与高层输出结合计算辅助损失（Auxiliary Loss）。
- Rezero 初始化：将 SwiGLU 中最后一个投影矩阵的初始化方差缩小为 0.01，使模块在训练初期接近恒等映射，极大提升了模型收敛的稳定性。

2.4 稀疏 Per-token MoE (Sparse-Pertoken MoE) 的进化

初代 TokenMixer 的痛点：原有的 ReLU-MoE 设计局限于“稠密训练、稀疏推理”范式，并未真正降低训练成本，且动态激活对线上推理极不友好。
Large 版本的解法：采用**“先扩大，后稀疏” (First Enlarge, Then Sparse)** 的端到端策略：
1. 将 Per-token SwiGLU 拆分为多个细粒度的专家（Expert）并进行稀疏激活。
2. 引入 门控值缩放 (Gate Value Scaling) 来解决稀疏化带来的梯度更新不足问题。
3. 加入 共享专家 (Shared Expert) 以稳定训练过程。

这使得模型能够真正实现”稀疏训练与稀疏推理”。在实验中，模型在仅激活一半参数（2.3B out of 4.6B）的情况下，FLOPs 下降近半，但取得了与稠密模型完全相同的业务增益（AUC +1.14%）。

与其他方法的对比分析

为了更全面地理解 TokenMixer-Large 的定位和优势，我们从架构设计、训练效率和扩展能力三个维度进行横向对比：

维度一：架构设计对比

对比项	DLRM-MLP	Wukong	DHEN	RankMixer	TokenMixer-Large
特征交叉方式	MLP 隐式交叉	双塔交叉网络	层次化显式交叉	Token Mixing	Mixing & Reverting
残差连接	简单残差	跨塔残差	层内残差	维度错位残差	语义对齐残差
归一化方式	BatchNorm	LayerNorm	LayerNorm	Post-LayerNorm	Pre-RMSNorm
激活函数	ReLU	ReLU	ReLU	FFN	pSwiGLU
稀疏化支持	无	无	无	ReLU-MoE	Per-token MoE

维度二：训练效率对比

对比项	DLRM-MLP	Wukong	RankMixer	TokenMixer-Large
GPU MFU	< 10%	~15%	~25%	~60%
碎片化算子	大量	中等	中等	无
训练稳定性	浅层稳定	一般	深层退化	深层稳定
训练范式	稠密	稠密	稠密训练稀疏推理	稀疏训练+稀疏推理

维度三：扩展能力对比

对比项	DLRM-MLP	Wukong	RankMixer	TokenMixer-Large
已验证最大参数量	~500M	~1B	~1B	15B（离线）/ 7B（在线）
Scaling Law 表现	早期饱和	有限提升	中等	持续提升
多业务验证	单一场景	有限场景	有限场景	电商/广告/直播

实验结果详细分析

Scaling Law 验证

TokenMixer-Large 在离线实验中展现出了清晰的 Scaling Law 特性：

500M -> 1B：CTCVR AUC 持续提升，验证了架构设计的有效性
1B -> 4B：引入 Sparse-Pertoken MoE 后，模型在激活参数仅为 2.3B 的情况下达到了与 4.6B 稠密模型相当的性能
4B -> 15B：离线实验表明 AUC 仍在持续提升，未出现明显的饱和趋势

消融实验关键发现

论文中的消融实验揭示了几个重要结论：

Mixing & Reverting 的必要性：去除 Reverting 层后，AUC 下降约 0.15%，证明语义对齐的残差连接对深层网络至关重要
跨层残差的贡献：去除跨层残差后，深层模型（>6 层）出现明显的训练不稳定，AUC 波动加剧
Rezero 初始化的作用：将初始化方差从标准值改为 0.01 后，训练初期的 loss 曲线更加平滑，最终收敛效果提升约 0.08% AUC
门控值缩放的影响：在 MoE 稀疏化场景中，移除门控值缩放会导致约 0.12% 的 AUC 损失，验证了其对缓解稀疏梯度更新不足的有效性

流程图 (Flowchart)

graph LR A[Sparse Features] --> B[Embedding & Tokenizer] B --> C[Grouped Tokens X] D["Mixing: Split & Concat"] --> E["SwiGLU + Norm"] E --> F["Reverting"] --> G["SwiGLU + Norm"] C --> D C -.->|Residual| G G --> H[Deep Layers + MoE] H --> I[Pooling & Prediction]

优缺点分析

优势

架构简洁高效：通过彻底移除碎片化算子，TokenMixer-Large 将 GPU MFU 提升至 60%，这意味着同样的硬件资源可以训练更大的模型。这种”少即是多”的设计理念在工业界具有重要的参考价值。
真正的稀疏训练+推理：不同于 ReLU-MoE 的”稠密训练、稀疏推理”，Per-token MoE 实现了端到端的稀疏化，使得训练成本和推理成本同时降低。这对于大规模在线服务的部署预算控制至关重要。
经过大规模工业验证：该架构已在字节跳动电商、广告、直播三大核心业务线上线验证，覆盖了推荐系统的主要应用场景，证明了其普适性和鲁棒性。
清晰的 Scaling Law：实验证明了模型在 500M 到 15B 参数范围内持续受益于规模扩展，为后续进一步扩展提供了明确的方向。

不足

Embedding 层优化不足：论文主要聚焦于排序模型（Ranking Model）的上层架构，对 Embedding 层的优化讨论较少。而在实际工业系统中，Embedding 层往往占据了模型参数量的绝大部分（通常超过 90%），如何高效地扩展 Embedding 仍是一个开放问题。
训练基础设施要求高：扩展至 7B-15B 参数规模需要大量的 GPU 资源和分布式训练框架支持。论文对多机多卡的并行策略、通信优化等工程细节披露有限，其他团队复现的门槛较高。
冷启动与长尾问题未涉及：文章主要关注整体指标（AUC、GMV）的提升，未讨论大模型在推荐系统冷启动场景和长尾物品推荐上的表现，而这些恰恰是工业推荐系统的核心痛点。
跨域泛化能力有待验证：虽然在字节跳动内部三个业务线均有验证，但不同公司的推荐系统在数据分布、特征工程、业务目标上差异巨大，该架构的跨域迁移能力尚需更多外部验证。

工程实践启示

对于正在探索推荐系统大模型化的团队，TokenMixer-Large 提供了以下工程实践启示：

1. 先做架构”减法”，再做规模”加法”

在盲目扩大模型参数之前，优先审视现有架构中的碎片化算子
统计各算子的 FLOPs 占比和延迟占比，找出”高延迟、低计算”的瓶颈算子
逐步替换为高计算密度的标准化组件（如将各类特征交叉算子统一为矩阵乘法）

2. 渐进式扩展策略

不要一步跳到超大规模，建议按照 500M -> 1B -> 4B -> 10B 的节奏逐步扩展
每个阶段都需要充分的离线实验和在线 A/B 测试验证
关注 Scaling Law 曲线的拐点，当 AUC 提升开始饱和时及时调整策略

3. 稀疏化是大模型落地的关键

对于超过 1B 参数的在线推理模型，MoE 稀疏化几乎是必选项
建议同时评估训练和推理两端的稀疏化方案，优先选择能同时降低两端成本的方案
门控值缩放和共享专家等稳定化技巧在实践中非常重要，不可省略

4. MFU 是核心效率指标

将 MFU 纳入模型迭代的核心监控指标
目标至少达到 40% 以上（TokenMixer-Large 达到了 60%）
低 MFU 往往意味着存在大量的访存瓶颈或通信开销，需要针对性优化

结论 (Conclusion)

TokenMixer-Large 验证了在去除历史碎片化算子后，”纯净架构+大规模堆叠”在推荐领域的有效性。模型在离线实验中成功扩展至 150亿 (15B) 参数，在线部署达到了 70亿 (7B) 参数。在字节跳动核心业务取得巨大收益：

电商: 订单量提升 1.66%，人均 GMV 提升 2.98%
广告: ADSS 提升 2.0%
直播: 收入增长 1.4%

关键代码/数据

核心数据对比 (电商场景 500M 规模基线对比):

模型	参数量	训练 FLOPs/Batch	CTCVR AUC 提升
DLRM-MLP	499 M	125.1 T	基线
Wukong	513 M	4.6 T	+0.76%
RankMixer	567 M	4.6 T	+0.84%
TokenMixer-Large 500M	501 M	4.2 T	+0.94%
TokenMixer-Large 4B SP-MoE	2.3B 激活	15.1 T	+1.14%

注：Sparse-Pertoken MoE 在激活仅一半参数（2.3B in 4.6B）的情况下，不仅显著降低了 FLOPs，还达到了与稠密模型完全相同的业务增益，实现了极高的性价比 (ROI)。

总结与展望

TokenMixer-Large 的成功为工业级推荐系统的大模型化树立了一个重要的里程碑。它证明了推荐系统同样遵循 Scaling Law，只要架构设计得当，参数规模的扩展能够持续带来业务增益。

未来值得关注的方向包括：

多模态融合：将文本、图片等多模态特征纳入 TokenMixer 框架，进一步提升推荐质量
在线学习与实时更新：探索大模型在在线学习场景下的高效更新策略
模型压缩与蒸馏：研究如何将大模型的知识高效蒸馏到轻量级模型中，服务于延迟敏感的场景
跨业务迁移学习：探索不同业务线之间的模型迁移和知识共享机制