💻 computer science

CLHA: Cross-Layer Hessian Aggregation for Quantizing Weight-Shared Mixture-of-Experts Transformers

本文介绍了 CLHA，一种用于权重共享混合专家（Mixture-of-Experts）Transformer 模型训练后量化的跨层 Hessian 聚合方法，该方法通过结合共享层之间的 Hessian 统计信息来最小化总重构误差，在显著优于现有逐层校准方法的同时，还解决了 Hessian 估计中关键的实现陷阱问题。

原作者： Kunal Dhanda

发布于 2026-06-25

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Kunal Dhanda

原始论文采用 CC BY 4.0 许可（https://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

大局观：缩小一座巨大的图书馆

想象一座巨大的图书馆（大型语言模型），它太大了，无法放进一个小书架。为了让它能放得下，我们需要对书籍进行“压缩”。在人工智能领域，这种压缩被称为量化（Quantization）。这就像是将一张高清照片缩小成一张缩略图。如果你做得不好，图片就会变得模糊且无法辨认。

这篇论文研究的是一种特定类型的图书馆，叫做混合专家模型（Mixture-of-Experts, MoE）。这些图书馆很特别，因为它们不会针对每个问题都阅读所有的书；它们拥有“专家”（专门的区域），只有在需要时才会开启。

最近，工程师们发明了一个聪明的技巧，叫做跨层权重共享（Cross-Layer Weight Sharing, CLWS）。想象一下图书馆里两个不同的楼层（层 A 和层 B）使用完全相同的一套参考书作为他们的专家。这节省了大量的空间，因为你不需要为同一本书购买两份副本。

问题所在：
当人们试图缩小（量化）这些图书馆时，他们犯了一个错误。他们把共享的书籍视为仅属于第一层的书籍。他们观察了层 A 提出的问题，决定如何缩小这些书，并将这种缩小方案应用到了层 B 的书籍上。

但问题在于：在层 A 提问的人与在层 B 提问的人是不同的。随着你走上楼梯，问题的“氛围”也在发生变化。由于只观察了层 A，这种缩小计划对层 B 来说是错误的，导致图书馆失去了正确回答问题的能力。

解决方案：CLHA（“双重检查”系统）

作者提出了一种名为 CLHA（跨层 Hessian 聚合）的新方法。

类比：裁缝与双胞胎西装
想象一位裁缝正在为一对完全相同的双胞胎制作西装。

旧方法 (PLIC)： 裁缝测量了双胞胎 A，剪裁了布料，并假设双胞胎 B 的尺寸完全一样。但双胞胎 B 的姿势略有不同。这套西装完美契合 A，但在 B 身上却显得太紧了。
CLHA 方法： 裁缝同时测量了两个双胞胎。他们获取了 A 和 B 的测量数据，并将它们进行平均（给予在房间里出现频率更高的人更多的权重），从而创造出一个能完美适配两个人的“超级西装”版型。

用技术术语来说，论文指出你不应该只看一个层的数据。你必须结合两个层的“敏感度”（数学上称为 Hessian）。这创建了一个组合图谱，能够精确地告诉你如何缩小共享权重，使其在两个楼层都能表现良好。

特殊升级：CLHA-MP

论文还引入了一个“混合精度”版本，称为 CLHA-MP。

类比：VIP 休息室
在这些图书馆中，有两种类型的专家：

路由专家 (Routed Experts)： 它们只为特定的、罕见的问题打开大门。
共享专家 (Shared Experts)： 它们是“始终开放”的，处理最常见、日常的问题。

因为“共享专家”被使用的频率更高，所以它们对误差更加敏感。如果过度缩小它们，整个图书馆就会崩溃。

解决方法： CLHA-MP 为这些“始终开放”的专家提供了额外的空间（多出一位精度的比特）。这就像是给 VIP 休息室开了一扇稍宽的门，而保持普通门窄小一样。这微小的额外空间对图书馆的运作效果产生了巨大影响。

隐藏陷阱：“幽灵”钩子

作者还在用于构建这些模型的软件工具（PyTorch）中发现了一个隐蔽的 Bug。
类比： 想象两个人试图同时在同一个笔记本上写字。如果两人都使用没有标签的同一支笔，他们的笔记就会混在一起，你无法分辨谁写了什么。
在代码中，当两个层共享同一个专家时，软件的“前向钩子”（用于追踪活动的工具）会将两个层的数据混合在一起，从而损坏测量结果。作者通过创建一个“层感知”系统解决了这个问题，该系统会为笔记贴上标签，以便裁缝知道准确地分辨出哪个是哪个。

他们证明了什么？

他们在名为 WikiText-2 的数据集上测试了一个小型模型。

在 2-bit 压缩下（极小尺寸）： 旧方法 (PLIC) 让模型变得非常困惑（高困惑度/Perplexity）。
CLHA 方法： 让模型变得明显更聪明（提升了 4.3% 到 5.4%）。
CLHA-MP 方法： 通过给“始终开放”的专家多出一位精度的空间，让模型变得更聪明（提升了 18.6%）。

他们还进行了一项测试，人为加大了两个楼层之间的差异。旧方法的效果随着差异增大而变得越来越差，但新方法保持了稳定，证明了他们的“双重检查”系统在两个楼层差异巨大时是至关重要的。

总结

问题： 在层之间共享权重节省了空间，但标准的压缩工具忽略了第二层，从而导致错误。
修复方案： 结合两个层的数据，以创建一个更好的压缩计划 (CLHA)。
加分项： 给使用频率最高的专家提供一点额外的精度 (CLHA-MP)。
结果： 得到了更聪明、更小的模型，在缩小规模时不会失去其“脑力”。

技术摘要：CLHA —— 用于量化权重共享混合专家（MoE）Transformer 的跨层 Hessian 聚合

1. 问题陈述

最近的高效混合专家（MoE）Transformer 架构（如 Megrez2 中的架构）采用了跨层权重共享（CLWS）。在这种设计中，相邻的 Transformer 层（例如 $\ell$ 和 $\ell+1$ ）重用相同的路由专家权重张量，从而在保持每个 token 激活参数数量不变的同时，将参数占用空间减半。

然而，现有的**训练后量化（PTQ）**方法（包括 GPTQ 式的方法）将这些共享权重视为属于单个层。它们使用来自单一“主”层的激活统计数据来进行校准。这造成了根本性的不匹配：

残差流表示在层与层之间不断演变，这意味着激活层 $\ell_a$ 处共享专家的 token 分布与激活层 $\ell_b$ 处的分布不同。
现有方法忽略了次级层的激活统计数据，导致重建误差最小化效果不佳，尤其是在低比特（如 2-bit）机制下。

2. 方法论

2.1 理论基础：跨层 Hessian 聚合

论文形式化了量化跨越层 $\ell_1, \dots, \ell_n$ 共享权重 $W$ 的目标。正确的目标是最小化所有层使用该权重的总重建误差：
$L(W_q) = \sum_{i} \alpha_i \text{tr}[(W_q - W)H_{\ell_i}(W_q - W)^\top]$
其中 $H_{\ell_i}$ 是第 $\ell_i$ 层的 Hessian（激活协方差）， $\alpha_i$ 代表路由频率（各层的 token 比例）。

命题 1 确立了最小化该组合误差在数学上等价于应用带有组合 Hessian 的 GPTQ：
$H_{comb} = \sum_{i} \alpha_i H_{\ell_i}$
这与**逐层独立校准（PLIC）**形成对比，后者仅使用单个层的 Hessian ( $H_{\ell_a}$ ) 来最小化误差。

2.2 CLHA 算法

作者提出了 CLHA（跨层 Hessian 聚合），一种三阶段校准算法：

收集（Collection）： 对于每个校准序列，为每个活跃专家累积逐层 Hessian 和 token 计数。
聚合（Aggregation）： 对于每对共享专家 $(\ell_a, \ell_b)$ $(ℓ_{a}, ℓ_{b})$ ，计算组合 Hessian。论文评估了两种权重策略：
- 路由频率加权 ( $\alpha = n_a / (n_a + n_b)$ )： 根据每层实际接收到的 token 数量对 Hessian 进行加权。
- 等权重（Equal Weighting, $\alpha = 0.5$ ）： 作为一种正则化项，在校准数据有限或次级层 Hessian 估计存在噪声时非常有用。
量化（Quantization）： 使用聚合后的 $H_{comb}$ 应用标准的 GPTQ 列顺序更新。

2.3 CLHA-MP（混合精度）

意识到始终激活的共享专家会接收所有 $N$ 个 token（不像路由专家仅接收比例为 $K/E$ 的 token），论文指出其量化敏感度比路由专家高 $E/K$ 倍。CLHA-MP 为这些共享专家分配 $b+1$ 比特，同时保持路由专家为 $b$ 比特，由此产生的存储开销极小（仅占专家参数的 10–20%）。

2.4 实现修正：Hook 歧义性

论文识别出一个关键的工程陷阱：标准的 PyTorch 前向传播 hook 如果附加到共享的 ExpertFFN 模块上，则会对调用该模块的任何层触发。这会导致来自不同层的激活被混合到一个单一的 Hessian 累加器中，从而破坏了逐层归因。

解决方案： 作者提出了一种通过 MoELayer.forward() 签名传递的层感知回调机制。它显式地向回调函数提供层索引，从而确保即使底层权重模块是共享的，也能为 $H_a$ 和 $H_b$ 分别累积独立的 Hessian。

3. 核心贡献

理论命题： 形式化了共享权重的最优 Hessian 是各层 Hessian 的路由频率加权和。
CLHA 算法： 在 GPTQ 框架内实现了组合 Hessian 校准。
CLHA-MP： 针对高敏感度共享专家的混合精度扩展。
实现修复： 识别并解决了在共享模块架构中会静默破坏校准的“hook 歧义性”问题。
实证验证： 证明了 CLHA 的收益与跨层路由差异性（cross-layer routing divergence）相关联。

4. 实验结果

实验是在一个基于 WikiText-2 从头训练的 34.5M 参数 CLWS-MoE 模型上进行的。

2-bit 性能表现： 这是误校准效应会复合的关键区间。
- CLHA-rt（路由加权）相比 PLIC 降低了 4.3% 的困惑度（PPL）。
- CLHA-eq（等权重）降低了 5.4% 的 PPL。
- CLHA-MP 相比 PLIC 实现了 18.6% 的提升。
- 在 4-bit 和 3-bit 下，差异在噪声范围内（<1%）。
校准与测试误差的关系： 有趣的是，PLIC 在训练集上实现了更低的校准 Hessian 误差（CHE），因为它直接针对该分布进行优化。然而，CLHA 在测试集上的泛化能力更好，这表明组合 Hessian 比单层视角能更好地捕捉共享权重的真实分布。
路由锐度实验： 通过人工锐化路由 Logits（增加跨层差异性），作者观察到：
- 随着差异性增加，PLIC 的性能下降了 5%。
- CLHA 保持稳定（下降 <1%）。
- CLHA 与 PLIC 之间的性能差距从 +4.3%（自然路由）扩大到了 +8.0%（高差异性），证实了 CLHA 收益随跨层差异性增加而扩大的理论预测。

5. 重要性与结论

论文声称提供了针对跨层权重共享 MoE 模型的第一项 PTQ 分析。其重要性在于：

修正目标函数： 证明了将共享权重视为单层实体在数学上是次优的，且 Hessian 聚合是必要的修正。
收益的可扩展性： 表明随着模型规模变大以及路由模式变得更加专业化（即跨层差异性增加），CLHA 的益处可能会增长，使其成为未来生产级 CLWS-MoE 模型（如 Megrez2, DeepSeek, Qwen）的必备技术。
实际必要性： 强调了将现有 PTQ 工具直接应用于 CLWS 架构会存在根本性缺陷，并为未来该领域的任何校准方法提供了必要的修复方案。

作者对其绝对性能表现保持谦逊，指出其模型较小（34.5M）且训练步数有限，但强调其相对于标准校准方法的改进是稳健且具有理论依据的。