Preventing Expert Collapse in MoE-dVLMs via Modality-Wise Norm Alignment¶

会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/Weixin-AI/TowerAlign
领域: 多模态VLM
关键词: 离散扩散视觉语言模型, 混合专家模型, 专家塌陷, 模态范数对齐, 结构保持缩放

一句话总结¶

针对离散扩散多模态模型中 MoE 路由极易出现的视觉 token 专家选择塌陷问题，本文发现根源在于视觉与文本 token 之间高达 167 倍的范数差异引发 Pre-Norm 残差溺水与表征惯性，提出了在投影层后通过全局仿射压缩对齐范数统计量的轻量方法 TowerAlign，以极低开销恢复了专家分工与注意力判别力。

研究背景与动机¶

离散掩码扩散大语言模型（dLLM）通过双向上下文建模与并行迭代去噪机制，在多模态理解（dVLM）领域展现出与自回归模型并驾齐驱的潜力。为了在保持激活成本可控的前提下进一步扩展模型参数容量，引入稀疏混合专家（MoE）架构构成了极为自然的技术演进路线。然而，将 MoE 拓展至离散扩散视觉语言模型（MoE-dVLM）时，原本在密集扩散模型中收敛良好的两阶段训练范式却遭遇了灾难性崩溃——模型在 ChartQA 和 DocVQA 等基准上的精度发生断崖式下跌。

这种失败的核心表象是严重的视觉 token 专家选择塌陷：尽管路由器的 Softmax 输出概率分布熵依然正常，甚至在部分层高于文本 token，但离散的 Top-k 路由决策却极度倾斜，绝大部分视觉 token 被贪婪分配给少数几个专家。直觉上常被用来对抗专家不平衡的路由端正则化手段（如放大 10 倍的负载均衡辅助损失、注入强 Gumbel-Softmax 随机扰动）对此完全失效，甚至进一步恶化了利用率熵并破坏了主任务学习。深入的奇异值分解（SVD）分析揭示了真正的病灶：在 2048 维隐空间中，视觉 token 的有效秩崩塌到了仅 2.1，而文本 token 则高达 1673.3。视觉 token 极度局限在一个狭窄的低维锥体流形内，使得下游路由器在几何上根本无法对其进行有效区分。

导致视觉表征低秩同质化的物理机制，源自深层视觉编码器与语言塔在模态接口处的巨大范数悬殊。视觉编码器输出的视觉 token 平均 L2 范数约为 60.47，而文本嵌入仅约为 0.36，两者相差近 167 倍。在现代语言模型普遍采用的 Pre-Norm 残差架构中，子层更新量被 RMSNorm 限制在文本尺度的量级，将微弱的注意力更新加到量级高达百倍的视觉残差路径上，引发了不可忽视的“注意力溺水”与“表征惯性”——视觉 token 在层间轨迹的有效角速度被压缩至仅约 1 度，几乎以直线穿透深层网络而未获得足够的跨模态多头交互。本文的核心 idea 是跳过对下游路由器的强行干预，直接在视觉投影层后引入结构保持的全局仿射压缩（TowerAlign），将视觉 token 的径向范数整体等比缩放至文本统计量水平，从源头根除注意力溺水并使路由器重新具备几何判别力。

方法详解¶

整体框架¶

针对 MoE-dVLM 中视觉 token 在 Pre-Norm 残差流内的几何退化，TowerAlign 构建了一个在投影层输出端即时介入的模态级范数校准管线。系统输入包含图像输入与双向掩码扩散文本序列（由提示词与加噪响应序列拼接而成），视觉特征先经由 ViT 抽取并通过双层 MLP 投影映射至语言嵌入维度。在将视觉 token 馈入 Pre-Norm 语言主干之前，TowerAlign 计算整批视觉 token 的全局平均范数，执行一次保持内部拓扑关系的全局标量缩放；校准后的视觉序列与文本嵌入共同进入各层 Pre-Norm 块，确保自注意力和 MoE FFN 的层级增量与残差基底保持平衡，并在去噪目标下端到端训练。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像与文本指令"] --> B["视觉特征抽取与投影<br/>SigLIP2 编码 + 2层 MLP"]
    B --> C["全局仿射压缩<br/>计算标量因子等比缩放"]
    C --> D["双向掩码拼接与多轮建模<br/>视觉与文本嵌入序列联合打包"]
    D --> E["Pre-Norm 语言塔平衡更新<br/>消除注意力溺水与表征惯性"]
    E --> F["Top-k 专家差异化路由<br/>64 专家中激活 Top-8 并迭代去噪"]

关键设计¶

1. 全局仿射压缩：保持内部拓扑的径向尺度校正

直接在单 token 粒度执行强制归一化（如 Token-wise RMSNorm）虽然消除了绝对模态差异，但会摧毁视觉表征原生的显著性信号。深入的视觉 Transformer 结构表明，深层 ViT 中普遍存在 3 到 5 个 L2 范数极高（常常大于 100）的 ViT 汇聚节点（Sink tokens），它们汇集了场景级全局语义，而绝大多数局部细节 token 的范数仅在 60 左右。如果对每个 token 分别除以其自身范数，就会抹平全局中枢与局部特征之间的天然级差，严重损害下游密集图表与文档问答能力。为此，TowerAlign 采用了保持相对结构不变的全局仿射压缩：计算当前批次内所有样本及视觉位置的平均 L2 范数，通过单一标量系数进行整批广播缩放：

\[\tilde{\mathbf{h}}_{v,b,i}^{(0)} = \lambda \cdot \mathbf{h}_{v,b,i}^{(0)}, \quad \lambda = \frac{\mu_{\text{target}}}{\frac{1}{BL_v}\sum_{b=1}^{B}\sum_{i=1}^{L_v}\|\mathbf{h}_{v,b,i}^{(0)}\|_2}\]

其中目标标量 \(\mu_{\text{target}} \approx 0.36\) 与文本预训练均值严格对齐。由于每一个视觉 token 乘上的都是同一个标量 \(\lambda\)，任意两个视觉 token 之间的相对模态范数比例 \(\|\tilde{\mathbf{h}}_{v,i}^{(0)}\|_2 / \|\tilde{\mathbf{h}}_{v,j}^{(0)}\|_2 = \|\mathbf{h}_{v,i}^{(0)}\|_2 / \|\mathbf{h}_{v,j}^{(0)}\|_2\) 得到严格保持。汇聚节点在进入路由器时仍保留相对更高的点积主导权，使得路由器能够根据显著性层级将全局与局部信息分别派发给专门的专家子群。

2. 消除 Pre-Norm 语言塔中的注意力溺水与表征惯性

该设计的核心在于恢复 Pre-Norm 变压器块中自注意力与 FFN 更新的有效角速度。在 Pre-Norm 架构中，子层更新 \(\mathbf{u}^{(l)} = \mathbf{h}^{(l)} + \text{SelfAttn}(\text{RMSNorm}(\mathbf{h}^{(l)}))\) 的增量由于经过归一化，其模态无关的更新步长常数被束缚在较小范围 \(C^{(l)}\)。当未经缩放的视觉残差底座高达 \(\alpha_v \approx 60.47\) 时，更新方向与残差向量的夹角变化率被残差主干完全淹没：

\[\tan\theta_{\text{eff}, v}^{(l)} = \frac{C^{(l)}\sin(\theta_{v}^{(l)})}{\|\mathbf{h}_{v}^{(l)}\|_2 + C^{(l)}\cos(\theta_{v}^{(l)})} \approx \frac{C^{(l)}}{\alpha_v} \approx 1^\circ\]

这种微弱的转向角使得视觉向量在深达数十层的网络中陷入近乎直线的表征惯性，几何方差无法展开，诱发理论上证明的内积向 \(1.0\) 聚拢的高范数坍缩定理。经 TowerAlign 将视觉残差底座拉平至 \(\alpha_v \approx 0.36\) 后，更新增量占残差路径的有效权重提升了两个数量级，视觉 token 能够随网络深度充分扩散至各维子空间，有效秩恢复至健康水平，使 Top-k 路由决策能够敏感地捕捉视觉位置间的细微语义差异。

3. 联合多轮掩码双向建模训练范式

在离散扩散多模态模型的微调阶段，通常面临单轮独立切片与完整多轮对话拼合两种策略抉择。不同于自回归模型严格依赖因果下三角掩码，扩散语言模型采用全局双向注意力机制。在全序列打包时，双向交互在理论上会导致早期轮次在去噪阶段感知到未来轮次的语义提示。然而实验证明，联合多轮拼合训练反而取得了显著优于单轮独立微调（Per-turn SFT）的表现。多轮联合输入在掩码迭代重建中强迫模型在单张输入图像的引导下构建跨轮次语义一致性约束，同时摊薄了长视觉序列在每一轮单独前向传递时的重复计算开销，形成了兼顾训练吞吐与深层推理稳定性的训练闭环。

损失函数 / 训练策略¶

模型采用基于 LLaVA 的经典两阶段训练协议： 1. 阶段 1（特征对齐预训练）：在 LCS-558K 图文对上仅优化双层 MLP 投影层与 TowerAlign 缩放系数，训练 1 个 epoch，学习率设为 \(1 \times 10^{-3}\)，采用余弦退火与 0.03 的 Warmup 比例。 2. 阶段 2（视觉指令微调）：在 LLaVA-NeXT 的 780K 单图多轮对话数据集上全参数微调，包含视觉编码器、投影层及 MoE 扩散语言塔。语言塔学习率为 \(1 \times 10^{-5}\)，视觉编码器微调学习率为 \(2 \times 10^{-6}\)。训练优化目标为离散掩码扩散重建的负对数似然损失：

\[\mathcal{L}_{\text{SFT}}(\theta) = -\mathbb{E}_{\tau \sim \mathcal{U}[0, 1], (\mathbf{x}_v, \mathbf{p}, \mathbf{y})}\left[\sum_{i=1}^{L_y} \log p_\theta\left(y^i \mid \mathbf{x}_v, \mathbf{p}, \mathbf{y}_\tau\right)\right]\]

其中 \(\mathbf{y}_\tau\) 为对真实应答以噪声率 \(\tau\) 独立采样的掩码损坏序列。模型无需依赖任何激进的辅助路由惩罚项，即可实现自然均衡的稀疏激活。

实验关键数据¶

主实验¶

主实验在统一硬件和数据预算下对比了 TowerAlign（基于 LLaDA-MoE-7B-A1B，每 token 仅激活 1B 参数）与经典密集扩散基准、自回归 VLM 的评测得分。

评测基准	指标	TowerAlign (7B-A1B)	LLaDA-VLM (8B Dense)	LLaVA-1.6 (7B AR)	表现分析
ChartQA	Accuracy	59.0	57.4	54.8	相比密集扩散提升 +1.6，大幅超越自回归基线
DocVQA (val)	ANLS	67.7	61.2	74.4	相比密集扩散大幅提升 +6.5，细粒度文字大幅增强
InfoVQA (val)	ANLS	25.8	25.5	37.1	与密集扩散持平 (+0.3)
RealWorldQA	Accuracy	58.9	57.4	-	真实场景感知超越密集扩散 +1.5
MMStar	Accuracy	46.7	46.0	-	多模态综合能力提升 +0.7
AI2D	Accuracy	70.5	70.2	66.6	科学图表推理达到 70.5，领先 LLaVA-1.6 达 +3.9
MMBench (EN)	Score	73.3	71.9	54.6	通用评测相比自回归 7B 提升 +18.7，超越密集扩散 +1.4
MMMU (val)	Accuracy	40.3	40.1	35.1	学科多模态推理全面领先密集 8B 与自回归 7B
MMMU-Pro (standard)	Accuracy	25.2	25.4	-	高难度基准上与密集扩散相当 (-0.2)
MMMU-Pro (vision)	Accuracy	15.1	15.7	-	纯视觉推理小幅微调 (-0.6)

消融实验¶

消融实验系统性对比了模态对齐算子、路由端正则化及微调输入方式对 MoE-dVLM 的影响。

实验配置	ChartQA	DocVQA	RealWorldQA	AI2D	MMBench	MMMU	说明
TowerAlign (Full)	59.0	67.7	58.9	70.5	73.3	40.3	全局仿射压缩，保持相对结构不变
w/o TowerAlign	4.7	13.8	30.6	58.7	61.4	36.7	移除模态对齐，遭遇严重专家塌陷与性能暴跌
TokRMS	19.4	26.7	36.7	48.2	58.0	31.2	单 token 独立 RMSNorm，抹平汇聚节点全局信号
强负载均衡 (LB loss 10x)	9.1	9.3	28.9	34.9	43.9	31.5	路由端强约束无法治本，干扰主损失导致崩盘
Gumbel 噪声 (τ=2.0)	8.1	6.8	16.7	28.9	48.2	25.4	强行引入路由随机性破坏了语义专精
单轮微调 (Per-turn SFT)	43.8	55.3	55.6	67.7	72.1	37.3	无法利用双向多轮上下文与更高计算吞吐

关键发现¶

移除 TowerAlign 导致特定图表与文档阅读任务全面失效（ChartQA 从 59.0 暴跌至 4.7，DocVQA 从 67.7 暴跌至 13.8），直接证明了上游几何范数对齐是 MoE 多模态扩散模型正常运转的关键前置条件。
采用逐 token 独立的 RMSNorm（TokRMS）虽然缓解了彻底崩溃，但 DocVQA 和 ChartQA 依然分别下降了 41.0 与 39.6 点，强力印证了 ViT 汇聚节点（Sink tokens）高范数特征作为全局语义信标的重要价值。
传统的 MoE 均衡技术（加大负载均衡权重、加 Gumbel 噪声）在几何崩塌场景下不仅无法增加有效利用率熵，反而加剧了优化扰动，在所有维度上全面劣于甚至不加干预的配置。

亮点与洞察¶

诊断视角清晰精辟：突破了以往只在路由端寻找专家不平衡原因的思维定势，敏锐指出视觉 token 的低秩收敛和模态范数极度不均衡才是导致 MoE 路由失灵的物理根源。
算子设计极简而高效：采用零额外可学习参数的全局仿射压缩标量计算，在规避注意力溺水的同时完全保留了视觉内部局部与全局的信息拓扑结构。
颠覆了常规正则化直觉：给出了强有力的反例证据，证明当表征输入存在严重几何退化时，盲目依赖路由端均衡损失只会加速主任务学习崩溃。

局限与展望¶

目标缩放均值 \(\mu_{\text{target}}\) 仍固定依赖于文本嵌入的经验统计量（0.36），未探索依据不同语言主干分布动态自适应学习的最佳策略。
实验评估主要建立在 7B-A1B 规格与单图多轮场景下，在更高分辨率多图输入、交织视频帧以及更大规模参数下的范数迁移规律有待进一步验证。
离散扩散生成在长文本 CoT 推理上的步数延迟与生成速度权衡，仍是限制其实用化部署的重要考量。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次揭示并严密论证了 MoE-dVLM 中视觉专家塌陷的几何根源，机制解释极具说服力。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 10 个主流基准对比、系统性的表征秩与路由熵分析，消融实验设计严丝合缝。
写作质量: ⭐⭐⭐⭐⭐ 逻辑推导层次分明，数学分析与实验现象互为印证，论点清晰直接。
价值: ⭐⭐⭐⭐⭐ 为后续开发大规模稀疏多模态扩散基础模型提供了至关重要的工程准则与理论基石。