Multi-Head Normalization for Wide Vision Transformers¶

会议: ECCV 2026
论文: ECCV 原文
领域: 模型压缩
关键词: 宽架构ViT, 特征归一化, MH-RMSNorm, 训练稳定性, 扩散模型

一句话总结¶

针对宽架构 Vision Transformer 在高维特征空间中因少数离群极值主导全局归一化而导致的表征压缩与梯度消失问题，本文提出 Multi-Head RMSNorm (MH-RMSNorm)，将高维向量均匀划分为多个子块独立进行 RMSNorm，成功消除训练坍塌并将 ViT 稳定扩展至 4096 维。

研究背景与动机¶

扩展 Vision Transformer (ViT) 的参数规模一直是增强视觉表征能力的核心途径，通常包括增加网络深度（堆叠更多层）与扩展模型宽度（增加隐层特征维度）两个正交维度。得益于 LayerScale、ReZero 以及 DeepNorm 等结构稳定技术的突破，ViT 能够可靠地扩展至 48 层甚至上千层；然而在宽度方向上，现有视觉大模型几乎均停留在 1024 左右的隐层维度。一旦尝试进一步拓宽特征通道，宽架构 ViT 经常会遭遇严重的训练不稳定、优化停滞甚至完全崩溃。

深入探究高维空间中的统计特性可以发现，宽架构模型的失效并非源自架构本身的表达瓶颈，而是由全局特征归一化引发的“极值激活主导”（high-activation domination）所致。在高维特征向量中，极少数异常维度往往呈现出数倍于平均水平的极高激活值。无论是 LayerNorm 还是 RMSNorm，全局范数的计算都会被这几个离群维度牢牢绑架；归一化除以庞大的全局模长后，其余绝大多数正常维度的表征被强行压缩至趋近于零。在低精度（如 bfloat16/fp16）训练环境下，这些微小数值极易发生下溢截断，导致反向传播时对应的非主导维度梯度完全消失。更为严重的是，这种特征不平衡性会在深层网络中逐层指数级放大，最终引发整体优化的不可逆坍塌。

本文的核心洞见在于：既然极值激活在高维空间中依据极值理论不可避免，就应当打破跨维度的全局统计耦合，将统计池化限制在局部正交子空间内。核心 idea：将高维特征向量均匀切分为多个 head-aligned 的紧凑子块并对每个子块独立执行 RMS 归一化（Multi-Head RMSNorm），将离群极值对均方根统计的污染严格局限在单个子空间内，从而保护绝大多数维度的梯度传播与有效表达。

方法详解¶

整体框架¶

MH-RMSNorm 的设计理念是作为即插即用的轻量级归一化原语，无缝替换 ViT 与 DiT 中自注意力模块前以及 MLP 模块前的标准 LayerNorm 或 RMSNorm。给定输入的 token 嵌入向量，系统不再计算全通道维度的统一范数，而是将其按预设维度划分为 $H$ 个互不重叠的子向量；每个子向量独立计算局部的均方根（RMS）标量并完成缩放，随后沿特征维度重新拼接并施加逐 head 的可学习仿射变换。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 Token 嵌入向量<br/>x ∈ R^D"] --> B["多头特征切分<br/>划分为 H 个子块 x^(h) ∈ R^d"]
    B --> C["子空间独立 RMS 统计<br/>计算局部 RMS(x^(h)) 标量"]
    C --> D["局部独立归一化与仿射缩放<br/>x^(h) / RMS(x^(h)) ⊙ γ^(h)"]
    D --> E["特征拼接与输出<br/>重组为完整高维表征 x_out ∈ R^D"]

关键设计¶

1. 多头子空间切分：阻断离群极值的全局污染 在高维特征 $x \in \mathbb{R}^D$ 中，全局归一化将所有维度置于同一个统计池中。当某一个维度 $j \in S$ 出现模长为 $M$ 的离群极值（$M^2 \gg \sigma^2$）时，全局均方根被放大至 $\mathcal{O}(M)$ 量级，使得非主导维度 $i \notin S$ 的梯度衰减因子变为： $$ \frac{\partial \hat{x}_i}{\partial x_i} \approx \frac{1}{\mathrm{RMS}(x)} = \mathcal{O}\left(\frac{1}{M}\right) $$ 导致有效维度的反向传播梯度直接消失。MH-RMSNorm 将特征划分为 $H$ 个维度为 $d = D / H$ 的子向量 $x = [x^{(1)}, \dots, x^{(H)}]$。若极端激活集中在某一特定 head 内，其破坏效应被物理隔离在局部子空间，其余 $H - 1$ 个子空间维持正常的方差水平，从而保留了 $\frac{H-1}{H}$ 维度以上的梯度流。

2. 局部均方根归一化：兼顾计算效率与免均值偏移 结合大语言模型与现代 ViT 对简化归一化的偏好，MH-RMSNorm 移除了传统 LayerNorm 中的均值中心化与可加偏置项，仅利用局部的二阶矩实现尺度缩放： $$ \mathrm{RMS}\big(x^{(h)}\big) = \sqrt{\frac{1}{d} \sum_{i=1}^{d} \big(x_i^{(h)}\big)^2 + \epsilon} $$ 每个子向量完成标准化后，使用独立的缩放参数 $\gamma^{(h)} \in \mathbb{R}^d$ 进行校准： $$ \mathrm{MH\text{-}RMSNorm}(x)^{(h)} = \frac{x^{(h)}}{\mathrm{RMS}\big(x^{(h)}\big)} \odot \gamma^{(h)} $$ 这种设计在计算上与快速算子（如 FlashNorm）高度契合，既避免了跨长维度的全局规约同步开销，又消除了通道间不必要的均值扰动。

3. 固定子空间粒度伸缩策略：保持归一化统计容量一致 在模型宽度由 256 扩展到 4096 的全生命周期中，MH-RMSNorm 采用“固定子块维度 $d=256$、动态增减 head 数量 $H = D / 256$”的伸缩策略。极值理论表明，最大激活与均方根的比值与子空间维度的平方根 $\sqrt{d}$ 成正比；固定 $d=256$ 保证了无论网络扩展得有多宽，每个局部归一化单元内部的极值期望上限始终恒定，从根本上杜绝了随着模型变宽而出现的方差爆炸与累积退化。

实验关键数据¶

主实验¶

在 DeiT-III 训练框架下评估 ImageNet-1K 分类，并在 UperNet 框架下评估 ADE20K 语义分割性能；在 DiT-XL/2 上评估 ImageNet-256 条件图像生成表现。

任务 / 模型	架构宽度	Normalizer	主要指标	对比 baseline	提升幅值
ImageNet-1K (ViT-B)	768	MH-RMSNorm	83.9% Top-1	83.8% (RMSNorm)	+0.1%
ImageNet-1K (ViT-L)	1024	MH-RMSNorm	84.7% Top-1	84.5% (RMSNorm)	+0.2%
ImageNet-1K (ViT-H)	1280	MH-RMSNorm	85.5% Top-1	85.2% (RMSNorm)	+0.3%
ADE20K (ViT-H UperNet)	1280	MH-RMSNorm	51.0% mIoU	50.6% (RMSNorm)	+0.4% mIoU
ImageNet-256 (DiT-XL/2 800k)	1152	MH-RMSNorm	13.0 FID	13.8 (RMSNorm)	-0.8 FID
ImageNet-256 (DiT-XL/2 7M)	1152	MH-RMSNorm	1.95 FID	2.20 (RMSNorm)	-0.25 FID

消融实验¶

在固定 12 层的 ViT 架构下，针对极限超宽模型（宽度 4096）进行归一化家族、头数配置及替换位置的消融分析。

配置	Top-1 Accuracy (%)	FID ↓	说明 (Table 4 & 5 依据)
RMSNorm (global, $H=1$)	61.1	82.3	全局归一化在 4096 维遭遇严重训练崩溃
LayerNorm (global, $H=1$)	60.5	88.4	全局 LayerNorm 同样发生严重性能断崖
MH-LayerNorm ($d=256, H=16$)	84.2	36.1	引入多头分块后性能立即大幅恢复
MH-RMSNorm (shared $\gamma$)	84.0	35.9	共享仿射参数无法适应不同子空间的异构尺度
MH-RMSNorm ($d=512, H=8$)	84.2	-	子块适度放大仍保持稳定
MH-RMSNorm ($d=256, H=16$, default)	84.4	35.5	默认设计，分类与生成综合表现最优
MH-RMSNorm ($d=64, H=64$)	83.5	-	过度碎片化破坏了特征维度间的局部相关性

关键发现¶

超宽架构下的两极分化：当宽度从 1024 扩展至 4096 时，全局 LayerNorm 和 RMSNorm 的准确率从 82% 骤降至 60.5%~61.1%，FID 急剧恶化到 80 以上；而 MH-RMSNorm 性能单调上升，在 4096 维达到 84.4% Top-1，成功释放了宽架构的模型容量红利。
深层网络收益更显著：仅在浅层（1-4层）替换 MH-RMSNorm 仅能取得 83.1% Top-1，而在深层（9-12层）替换可达 83.9%，全层替换达到最佳的 84.2%（在 2048 维设定下），证明极值激活的不良影响沿深度具有累积放大效应。
双分支同等重要：仅替换 Attention 分支或仅替换 MLP 分支的效果几乎完全对称（两者皆为 84.0% Top-1 / 37.2~37.3 FID），表明该失效机理普遍存在于所有高维特征变换之后，需全局一致采用局部归一化。

亮点与洞察¶

极值理论驱动的病因诊断：将宽 ViT 的崩溃根因严格定位为高维统计下的最大值与均方根比值按 $\mathcal{O}(\sqrt{D})$ 恶化，为长久以来的 ViT 宽度扩展难题提供了清晰的数学解释。
零额外算力代价的即插即用原语：MH-RMSNorm 不改变模型参数量与前向浮点计算量（FLOPs），仅通过重新组织内存求和规约逻辑，即可将宽模型优化转危为安。
对生成扩散模型的普适有效性：在 DiT-XL/2 上无需改变任何生成结构即可把 50k 步生成的 FID 压至 1.95，Recall 从 0.57 提升到 0.59，展现了更强的模式覆盖率与更低的样本坍缩率。

局限与展望¶

通道重排不敏感性尚未探索：当前分块默认采用自然物理顺序切分，未探究特征维度置换或跨 channel 混洗（shuffle）对局部均方根稳定性的影响。
缺乏在超大规模型（如百亿参数 LLM）上的验证：论文集中在 12 层宽架构 ViT/DiT，尚未探究在万维隐藏层、千亿级大语言模型中切分 head-wise 归一化对长文本注意力稳定性的交互影响。
未来方向：探索动态自适应 chunk 大小，以及将 MH-RMSNorm 与稀疏注意力算子（FlashAttention）内核的融合实现。

评分¶

新颖性: ⭐⭐⭐⭐ 深刻揭示了宽 ViT 归一化失效机理并提出极简优雅的 Multi-Head 统计隔离方案。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 ImageNet-1K 分类、ADE20K 分割与 DiT 生成，包含完备的超宽扩展与维度消融。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰，图表实验分析直观有力，逻辑严密。
价值: ⭐⭐⭐⭐⭐ 为未来突破 ViT/DiT 宽度瓶颈提供了通用可靠的基础组件。

配置	Top-1 Accuracy (%)	FID ↓	说明 (Table 4 & 5 依据)
RMSNorm (global, \(H=1\))	61.1	82.3	全局归一化在 4096 维遭遇严重训练崩溃
LayerNorm (global, \(H=1\))	60.5	88.4	全局 LayerNorm 同样发生严重性能断崖
MH-LayerNorm (\(d=256, H=16\))	84.2	36.1	引入多头分块后性能立即大幅恢复
MH-RMSNorm (shared \(\gamma\))	84.0	35.9	共享仿射参数无法适应不同子空间的异构尺度
MH-RMSNorm (\(d=512, H=8\))	84.2	-	子块适度放大仍保持稳定
MH-RMSNorm (\(d=256, H=16\), default)	84.4	35.5	默认设计，分类与生成综合表现最优
MH-RMSNorm (\(d=64, H=64\))	83.5	-	过度碎片化破坏了特征维度间的局部相关性