← 返回 2026-05-11

均值模式尖鸣：用于千层扩散 Transformer 的均值-方差分离残差 Mean Mode Screaming: Mean--Variance Split Residuals for 1000-Layer Diffusion Transformers

Pengqi Lu 📅 2026-05-07 👍 238 2026-07-13 08:36
Transformer 训练稳定性 扩散模型 残差结构 深度学习理论
arXiv HuggingFace PDF

提出 MV-Split 残差，把均值与中心化信号分开增益，稳住千层 DiT 训练。

前置知识

扩散 Transformer (DiT)

用 Transformer block 替代 U-Net 的扩散模型骨干，输入为 VAE 潜空间 patch token，通过自注意力跨 patch 通信。

本文所有实验都建立在单流 DiT 上，理解其 Post-Norm 残差链 $X_{l+1}=\mathrm{RMSNorm}(X_l+f_l(X_l))$ 是看懂 MMS 触发链的前提。

Post-Norm 残差与零初始化残差写头

残差写头指把分支输出写回主干的投影层（如注意力输出 $W_O$、FFN 输出 $W_2$）。'零初始化' 让分支初始对主干无贡献，靠学习逐步打开。

MMS 的关键就是这些写头在零初始化下突然被一次性'打开'，触发均值相干梯度尖峰，是全文机制分析的核心对象。

行随机注意力 (Row-stochastic Attention) 与 Softmax Jacobian

行随机指每行和为 1 的注意力矩阵 $A\mathbf{1}=\mathbf{1}$。Softmax Jacobian 描述 logits 微扰如何映射到权重微扰，对数 logits 全等时 Jacobian 把常值分量映射到零空间。

论文两个核心命题都依赖这一性质：均值子空间被严格保留，中心化子空间被收缩；值向量均匀化时 Q/K 梯度被 Softmax Jacobian 零空间抑制。

研究动机

把扩散 Transformer (DiT) 推到数百乃至上千层时，训练常出现一种独特的失败模式：loss 在上万步内看似稳定，却在几十步内骤升至接近初始化的水平并永久无法恢复，全程不出现 NaN 或明显前向饱和。400 层未加稳定器的基线在默认学习率下甚至在第一个 10k 检查点之前就发散；降到 1/2 学习率也仅是把崩溃延后到约 26k–27k 步。表 1 中 400L Base (η) 与 400L Base (η/2) 全部标 '—'，说明这种失败是结构性的，不是简单调参能解决的。

本文的目标是作者希望(1) 给出这种深度 DiT 训练崩溃的精确机制刻画，使之区别于常规的爆炸/消失梯度叙事；(2) 诊断出进入崩溃的触发事件 MMS，从而能针对性干预；(3) 设计一种能真正稳定 400 层匹配对比、并把可训练深度推到 1000 层的残差结构方案。整体目标是把'千层 DiT 训不动'从经验观察升级为可解释、可监控、可结构修复的工程问题，并通过这一案例厘清深 Transformer 中均值与中心化信号需要被差异化处理这一普遍原则。

与已有工作不同的是，已有深度稳定器（ReZero、LayerScale、DeepNorm、Keel）都把残差分支当一个整体做各向同性缩放，未区分 token 空间的均值与中心化两个子空间。本文的核心切入是把 token 表征显式分解为 $\mu(X)=JX$（序列均值）与 $c(X)=PX$（中心化变差），并证明行随机注意力在两者上行为完全不同，从而把崩溃归因到一个具体的、秩 1 的均值相干梯度模式 $\Delta W_\mu$ 上，并据此设计分别增益两个子空间的残差。

核心方法

论文沿'先诊断、再设计'的路线推进。先在 token 空间做正交分解 $X=\mu(X)+c(X)$，利用 $J$（均值投影）和 $P=I-J$（中心化投影）证明行随机注意力对均值子空间是恒等的、对中心化子空间是受 $\mu_{\mathrm{eff}}(A)=\|PAP\|_2$ 控制的收缩算子；再证明残差写头梯度存在精确加法分解 $\nabla_W\mathcal{L}=\Delta W_\mu+\Delta W_c$，其中 $\Delta W_\mu$ 在 token 对齐时进入相干累积的 $O(T)$ 标度，而 $\Delta W_c$ 是扩散式求和。由此识别出 MMS：当 $\Delta W_\mu$ 因对齐尖峰放大并打开写头后，残差分支主导 mean，Q/K 又被值均匀化后的 Softmax Jacobian 零空间压制，网络锁死在均值预测的平凡解。基于这一机制，把 Post-Norm 合并 $X_{l+1}=\mathrm{RMSNorm}(X_l+F_l)$ 替换为子空间路由的 MV-Split 合并，使均值与中心化信号获得独立、可学习的增益。

与 LayerScale、ReZero 等'整枝缩放'不同，MV-Split 在残差合并点显式分离两个正交子空间：中心化路径用增益 $\beta$ 走标准残差 $PZ_l=PX_l+\beta\odot PF_l$，均值路径走带泄漏的整合器 $JZ_l=(1-\alpha)\odot JX_l+\alpha\odot JF_l$。这样小 $\alpha$ 能在不连带压缩中心化信号的前提下，专门压住导致崩溃的均值相干模式 $\Delta W_\mu$，并通过反向传播 (式 10) 给予两个梯度分量独立增益，是机制驱动的最小结构改动。

方法步骤详情

前向：$X_l$ 与 $F_l=f_l(X_l)$，先 $Z_l=X_l+\beta\odot(PF_l)+\alpha\odot J(F_l-X_l)$，再 $X_{l+1}=\mathrm{RMSNorm}(Z_l)$，$\alpha,\beta\in\mathbb{R}^D$ 逐块可学，图像/文本段用 $J_{\mathrm{seg}},P_{\mathrm{seg}}$。$J,P$ 解耦得式 9：中心化子空间走带 $\beta$ 的标准残差；均值子空间走 $\alpha$ 控制的 leaky 整合器（$0<\alpha_d\le 1$ 时按 $1-\alpha_d$ 收缩原均值再加新估计）。反向：$G_l=\partial\mathcal{L}/\partial Z_l$，$J,P$ 自伴正交，故 $\partial\mathcal{L}/\partial F_l=\beta\odot(PG_l)+\alpha\odot(JG_l)$，两子空间梯度独立增益，$\Delta W_\mu$ 被 $\alpha$ 压缩而 $\Delta W_c$ 不等比缩小。

技术新颖性

新颖性在三点：(1) 把'深度 DiT 崩溃'从现象描述升级为带命题级证明的机制刻画（Prop. 1、Prop. 2、Lemma 1），给出对齐放大律 (式 6) 把崩溃锁定到 $\Delta W_\mu$ 的 $O(T)$ 相干累积；(2) 区分'崩溃状态'与'进入崩溃的事件 (MMS)'，并通过 400 层基线的诊断曲线 (Figure 3 a-f) 给出从反向触发到正向锁定的完整因果链；(3) MV-Split 是在残差合并点对子空间做选择性增益的最小结构，与 LayerScale 类方法的本质差异在于'按子空间分别缩放'而非'按通道统一缩放'，且通过反向链 (式 10) 显式解开对 $\Delta W_\mu$ 的压制。

Empirical trajectory of a representative divergence event (400-layer).
Figure 3: Empirical trajectory of a representative divergence event (400-layer).

实验结果

400 层匹配对比 (ImageNet 256×256, FLUX.2 VAE + Qwen3-0.6B, 25 NFE, CFG 2.0)：未稳定器基线在 $\eta$ 与 $\eta/2$ 下均发散；LayerScale 50k 步 FID 2.90/IS 165.5；MV-Split 50k 步 FID 2.60/IS 185.5，是匹配最优稳定前沿，后裁剪梯度范数比 LayerScale 更高且避开未稳定器尖峰。1000 层 MV-Split 50k 步 FID 2.77/IS 217.3，验证可外推。Figure 6 显示 LayerScale 把 $G_{\mathrm{mean}}$ 与 $G_{\mathrm{ctr}}$ 一并压缩，MV-Split 仅压前者、维持后者在更高稳定带。Figure 4 在 spike 步 $t^\star=3400$ 验证式 6：斜率 0.13–0.17 跳到 0.96–0.99、$R^2$ 0.31–0.50 跳到 0.81–0.83，最大 $A^{-1}\approx 167$。

Stability and convergence across 400-/1000-layer DiT runs.
Table 1: Stability and convergence across 400-/1000-layer DiT runs.
Writer amplification at the gradient spike (400-layer Base run, t* = 3400).
Figure 4: Writer amplification at the gradient spike (400-layer Base run, t* = 3400).
Quality and optimizer stability over 80k steps (ImageNet 256×256).
Figure 5: Quality and optimizer stability over 80k steps (ImageNet 256×256).
Residual-writer gradient mode decomposition.
Figure 6: Residual-writer gradient mode decomposition.
查看结构化数据
任务指标本文基线提升
ImageNet 256×256 文本到图像生成 (400 层, 50k 步) FID-50K / Inception Score MV-Split 400L: FID 2.60 / IS 185.5 LayerScale 400L: FID 2.90 / IS 165.5 FID 相对降低 10.3%，IS 相对提升 12.1%；未稳定器基线在 50k 步前已发散无法参与对比
ImageNet 256×256 文本到图像生成 (1000 层, 50k 步) FID-50K / Inception Score MV-Split 1000L: FID 2.77 / IS 217.3 400L LayerScale 50k 步: FID 2.90 / IS 165.5 1000 层仍可比肩 400 层稳定前沿；FID 4.5% 提升、IS 31.3% 提升，且在 1000 层全 schedule 保持稳定
ImageNet 256×256 (400 层, 30k 步) FID-50K / Inception Score MV-Split 400L: FID 3.09 / IS 166.5 LayerScale 400L: FID 4.09 / IS 130.5 FID 24.4% 更低、IS 27.6% 更高；30k 步前差距已稳定建立
写头梯度对齐放大验证 (400 层 spike 步) A^{-1} vs (T-1) \hat\kappa 拟合斜率 / R^2 Spike 时 Attn WO 斜率 0.961 R^2 0.808；FFN W2 斜率 0.985 R^2 0.829；最大 A^{-1} ≈ 167 Pre-spike 斜率 0.169 (Attn) / 0.126 (FFN)，R^2 0.306 / 0.499 斜率提升约 5.7–7.8 倍，R^2 提升约 1.6–2.6 倍，证实 MMS 与对齐放大律 (式 6) 强相关

局限与改进

作者明确指出 1000 层是独立尺度的 scale validation 而非匹配前沿对比，因其用了不同的训练/后训练流程，所以不能直接与 400 层控制点比较；GenEval、DPG-Bench 仅作为校准放在 Appendix K.1，未声称 SOTA。FID/IS 用 25 NFE Euler + CFG 2.0 采样，结论对该协议敏感。结构上 MV-Split 仍依赖 RMSNorm 与 Post-Norm 残差链，是否能直接迁移到 Pre-LN 或带 AdaLN 调制的 DiT 变体未经验证；附录 B 提示标准初始化下崩溃以深度渐进方式出现，MV-Split 对该情形的防护尚未给出完整 ablation。本人补充：(1) 1000 层只用单一架构且没有其他 1000 层基线作为对照，无法排除'千层本身就能在某种初始化下稳定'的假设；(2) 对齐放大律 (式 6) 是经验观测 + 等量上界代理，未给出在 spike 处的严格先验界；(3) $\alpha,\beta$ 引入少量额外参数与可能的初始化敏感性，但论文未给详细消融。

独立分析的弱点

独立分析有三点明显不足：(1) 基线对比仍较薄，缺少 Qwen3-0.6B 之外的文本编码器与 FLUX.2 VAE 之外的潜空间，跨配置泛化未知；这可以通过补跑 SD-VAE-FT-EMA + CLIP-L 或换更大 Qwen3 验证。(2) 1000 层只有 MV-Split 一种方案，没有 1000 层 + LayerScale 或 1000 层 + DeepNorm/Keel 作为直接对照，'千层稳定'这一结论的归属模糊；建议补 1000 层 + LayerScale/DeepNorm 匹配实验以分离方法贡献与初始化贡献。(3) MMS 触发依赖零初始化残差写头 + 极深 Post-Norm 链，对 Pre-LN 或带 skip/AdaLN 的现代 DiT 是否同样存在未被验证；应给出一份与 Pre-LN/Post-LN、零/标准初始化交叉的消融，并把 MMS 的诊断准则 (Figure 3) 推广为运行时监测指标。

未来方向

作者提出的方向集中在 scale validation 与多任务校准 (Appendix K GenEval/DPG-Bench)。基于成果可延伸：(a) 把 MV-Split 与并行方案 DeepNorm/Keel 组合，进一步推到 ≥2000 层的语言/多模态 Transformer，检验均值-中心化分离思路是否跨模态通用；(b) 将 $\alpha,\beta$ 与 AdaLN 调制融合，使均值通道也获得 timestep 调节自由度，避免在非零初始化下出现'均值通道恒为常数'的副效应；(c) 用 MMS 的诊断曲线 (Figure 3) 做运行时稳定性监测与早停/降学习率的自适应调度器；(d) 探究式 6 的对齐放大律能否被显式正则化 (例如对抗式 token 抖动) 主动抑制，从根本上避免 $\Delta W_\mu$ 进入 $O(T)$ 相干区。

复现评估

复现友好度较好。代码已开源 (github.com/erwold/mv-split)，权重已发布 (huggingface.co/StableKirito/mvsplit-dit-1000l)，数据集为 ImageNet-2012 + 约 50k 精选图做后训练，文本编码器是 Qwen3-0.6B，VAE 是 FLUX.2 frozen。架构细节 (Post-Norm、2D RoPE、Q/K-Norm、SwiGLU、零初始化残差写头、$\alpha,\beta$ 逐块可学习向量、分段 $J_{\mathrm{seg}}/P_{\mathrm{seg}}$) 与训练配置在主文与 Appendix G 中给出，采样协议明确 (Euler, 25 NFE, CFG 2.0)。但 1000 层 + 80k 步的算力门槛非常高 (论文称 1000 层为 scale validation 跑)，对个人研究者是个障碍；400 层 80k 步对比也需大集群。$\alpha,\beta$ 初始化策略、是否对 $\alpha$ 做上限裁剪以保 leaky 性质等细节建议复现前与作者对齐。

论文图表

Text-to-image generation samples from our 1000-layer MV-Split DiT.
Figure 1: Text-to-image generation samples from our 1000-layer MV-Split DiT.

由 1000 层 MV-Split DiT 生成的文本到图像样例，展示 1000 层深度下架构的可训练性与生成质量。

作为全文'scale validation'的视觉证据，证明 MV-Split 在边界深度不仅能跑、还能生成有意义的图像，是动机和结论的直观锚点。

Baseline DiT and representative training diagnostics.
Figure 2: Baseline DiT and representative training diagnostics.

左图为单流 DiT 主干 (Post-Norm + Q/K-Norm + 零初始化残差写头 + 2D RoPE)；中图是 400 层未稳定基线与 400/1000 层 MV-Split 在前 10k 步的 loss 曲线，基线在 26k–27k 步附近 crash；右图是基线在 L0–L384 上 per-layer 能量比 $\rho_T=\|\mu(X)\|_F/\|c(X)\|_F$。

把'崩溃前最后稳态 → 突然 spike → 锁入均值主导'用一张图串起来，是进入 Section 3 失败动力学的必要导览。