[2606.05682] 超越输出匹配：在NVFP4 LLM蒸馏中保持内部几何结构

# 在 NVFP4 LLM 蒸馏中保留内部几何结构

Source: https://arxiv.org/html/2606.05682

## 超越输出匹配：在 NVFP4 LLM 蒸馏中保留内部几何结构

Junhua ZhaoChi LiuXin ChenHaifeng WuJian WanSrinivasan Manoharan

###### 摘要

随着大语言模型越来越多地部署在延迟和成本受限的生产环境中，对低精度推理（包括基于 NVFP4 的方法）的需求不断增长。量化感知蒸馏（QAD）通过在训练量化学生模型匹配冻结的高精度教师模型的输出分布（使用 KL 散度损失）来帮助恢复低比特量化下损失的部分精度。在这项工作中，我们首先对 QAD 进行表示层面的诊断：仅匹配输出可能会掩盖内部退化，因为许多中间激活几何结构可以产生类似的教师对齐 logits。使用 CKA，我们表明仅使用 KL 的 QAD 会相对于 BF16 教师模型降低逐层表示相似性，在 RL 后训练模型中尤其严重。这种漂移与推理和编码任务的下游瓶颈相关，表明低比特恢复需要保留内部几何结构，而不仅仅是匹配输出。基于这一发现，我们提出 CKA-QAD，一种用于 NVFP4 QAD 和低比特 LLM 精度恢复的 CKA 引导表示对齐方法。该方法添加了一个轻量级正则化器，通过 CKA 对齐逐层 Gram 矩阵，从而在蒸馏过程中保留内部表示几何结构。在 Nemotron 3 Nano 和 Qwen3-4B-Thinking-2507 上，CKA-QAD 显著改善了表示对齐，并以适度的训练开销提高了下游推理和编码准确性。我们的研究结果将 CKA 引导的表示对齐定位为量化 LLM 恢复中输出匹配的实际补充。

## 1 引言

大语言模型（LLM）在延迟和内存受限环境中的快速部署加速了低精度推理格式的采用。在最近的进展中，NVFP4 [1 (https://arxiv.org/html/2606.05682#bib.bib1)] 是一种具有细粒度块缩放和扩展动态范围的 4 位浮点格式，相比 FP8 和 INT4 提供了显著的吞吐量和内存优势。然而，将 NVFP4 进行训练后量化（PTQ）通常会导致不可忽略的精度下降，特别是对于紧凑或推理专用架构。量化感知蒸馏（QAD）已成为一种高效解决方案，通过 KL 散度将量化学生模型的输出分布与全精度教师模型对齐。与量化感知训练（QAT）不同，QAD 避免了复制复杂的多阶段后训练流水线（例如监督微调和强化学习），并在不同模型家族中表现出显著的稳定性 [1 (https://arxiv.org/html/2606.05682#bib.bib1)]。

尽管 QAD 在恢复任务级精度方面取得了成功，但现有的 QAD 范式主要侧重于匹配输出级概率分布。这引出了一个关键问题：输出对齐能否保证忠实的内部表示恢复？在这项工作中，我们系统地分析了 NVFP4 量化模型在 QAD 期间的内部几何结构，并发现了显著的脱节。我们发现，仅 PTQ 本身就会导致逐层表示相似性明显下降（通过中心核对齐（CKA）[3 (https://arxiv.org/html/2606.05682#bib.bib3)] 量化）。令人惊讶的是，仅使用 KL 散度损失的标准 QAD 优化通常会加剧这种表示漂移，特别是在经历强化学习（RL）后训练的模型中。虽然基于 KL 的蒸馏将学生模型的输出分布与教师模型对齐，但学生模型的中间表示仍可能发生非预期的重组。这表明输出级对齐可能掩盖了 KL 目标单独无法捕获的内部表示退化。

我们假设，保留教师模型的内部表示几何结构对于稳健的低比特泛化非常重要，尤其是在依赖深层语义层次结构的复杂推理和编码任务上。为了解决这个问题，我们将我们的贡献框定为用于 NVFP4 QAD 和低比特 LLM 精度恢复的 CKA 引导表示对齐，并将其实例化为 CKA-QAD，这是 QAD 目标的一个轻量级扩展，使用 CKA 相似性显式正则化中间激活。通过联合优化 KL 散度（输出分布对齐）和 CKA（内部表示对齐），我们的方法阻止了以牺牲表示保真度为代价满足输出级损失的肤浅映射。关键是，CKA 对正交变换和各向同性缩放的不变性使其对激活空间中良性旋转和全局尺度变化的敏感度低于逐点特征损失，这对于在量化训练过程中分布可能发生变化的低比特激活是有利的。

我们在 Nemotron 3 Nano（一种混合 Mamba-Transformer MoE）和 Qwen3-4B-Thinking-2507（一种紧凑型 Transformer 推理模型）上评估了 CKA-QAD。我们的关键发现有三点：

- • QAD 中的表示漂移：我们量化并可视化了仅使用 KL 的蒸馏如何导致量化学生模型与高精度教师模型在 Transformer 和 Mamba 块上的 CKA 相似性降低，其中经过大量 RL 后训练的模型尽管输出分布与 BF16 教师模型接近，但表现出最明显的漂移。
- • CKA 引导的 NVFP4 QAD：我们为 NVFP4 QAD 引入了逐层 CKA 对齐损失，将内部对齐恢复到接近 BF16 水平，同时在匹配的 Qwen3-4B-Thinking-2507 训练配置中仅增加 0.5% 的步进时间开销和 7.0% 的峰值 VRAM 开销。
- • 下游精度：在 Qwen3-4B-Thinking-2507 上，CKA-QAD 在 AIME25（68.5% 到 72.3%）、GPQA-D（59.5% 到 61.1%）和 LiveCodeBench-v5（57.9% 到 59.8%）上的平均精度优于 QAD（仅 KL）。

这些结果表明，高效的 NVFP4 恢复受益于功能对齐和表示对齐的联合优化，促使 CKA 引导恢复成为低比特 LLM 蒸馏的一个实际方向。

本文其余部分组织如下。Section 2 (https://arxiv.org/html/2606.05682#S2) 回顾了 NVFP4 量化、QAD、表示相似性和蒸馏的背景与相关工作。Section 3 (https://arxiv.org/html/2606.05682#S3) 介绍了 CKA-QAD 目标和实现。Section 4 (https://arxiv.org/html/2606.05682#S4) 展示了实验设置、主要结果、效率分析和模型泛化研究。Section 5 (https://arxiv.org/html/2606.05682#S5) 和 6 (https://arxiv.org/html/2606.05682#S6) 讨论局限性并总结。

## 2 背景与相关工作

### 2.1 NVFP4 量化与量化感知蒸馏

低精度推理已成为在严格延迟和内存预算下部署大语言模型（LLM）的关键。NVFP4 格式 [1 (https://arxiv.org/html/2606.05682#bib.bib1),2 (https://arxiv.org/html/2606.05682#bib.bib2)] 扩展了 NVIDIA 的 microscaling FP4 设计，采用更小的量化块大小（16 个值，相比 MXFP4 的 32 个）、每块 FP8 E4M3 缩放因子实现细粒度自适应，以及第二级每张量 FP32 缩放以保持动态范围。这种设计相对于 FP8 将模型内存占用减少了约 1.8 倍，并且在实际 LLM 推理设置中可实现 2-3 倍的 token 生成吞吐量加速，使其对生产部署具有吸引力。

训练后量化（PTQ）在没有额外训练的情况下应用 NVFP4 校准，但通常会在紧凑或强化学习微调的模型上导致不可忽略的精度下降 [1 (https://arxiv.org/html/2606.05682#bib.bib1)]。量化感知蒸馏（QAD）通过训练量化学生模型匹配全精度 BF16 教师模型的输出分布来解决这个问题。设 zT\(x\),zS\(x\)∈R\|V\| 表示输入 x 的教师和学生 logits，其中 V 是词汇表。标准的温度缩放蒸馏目标是

LKD=Ex∼D\[τ2DKL\(σ\(zT\(x\)τ\)∥σ\(zS\(x\)τ\)\)\], (1) 其中 σ\(⋅\) 表示 softmax 函数，τ 是蒸馏温度。它可选地与针对真实标签的小交叉熵项相结合。与量化感知训练（QAT）[7 (https://arxiv.org/html/2606.05682#bib.bib7)] 相比，QAD 在现代 LLM 流水线中具有实际优势：它避免了重放复杂的多阶段后训练方案，如监督微调、RL 或模型合并；它对于部分或合成校准数据相对稳健；并且它直接使学生与 BF16 教师模型的输出分布对齐。尽管有这些好处，QAD 对模型内部表示几何结构的影响仍未得到充分研究。

### 2.2 表示相似性与中心核对齐（CKA）

理解内部表示在压缩或微调下如何演变需要能够捕捉几何相似性而非逐点激活相等性的度量。中心核对齐（CKA）[3 (https://arxiv.org/html/2606.05682#bib.bib3)] 是用于此目的的标准工具。给定来自教师模型的匹配激活 X∈RN×dT 和来自学生模型的 Y∈RN×dS，其中相同的 N 个 token，线性 CKA 为

CKA\(X,Y\)=HSIC\(X,Y\)HSIC\(X,X\)HSIC\(Y,Y\),HSIC\(X,Y\)=1\(N−1\)2tr\(KXHKYH\), (2) 其中 KX=XX⊤, KY=YY⊤, H=I−1N11⊤ 是中心化矩阵。等价地，对于中心化激活，CKA\(X,Y\)=‖Y⊤X‖F2/\(‖X⊤X‖F‖Y⊤Y‖F\)。CKA 不要求 dT=dS，并且对正交变换和各向同性缩放保持不变，因此对量化或微调网络中常见的特征旋转和幅度变化具有鲁棒性。尽管 CKA 已被广泛用于比较不同架构、宽度、深度和随机种子之间的表示，但它作为 NVFP4 QAD 和低比特 LLM 精度恢复的训练时指导的作用仍未得到充分探索。

### 2.3 特征蒸馏与 logit 蒸馏

知识蒸馏传统上在输出层面进行操作，通过 KL 散度或温度缩放交叉熵对齐教师和学生概率分布 [8 (https://arxiv.org/html/2606.05682#bib.bib8)]。为了迁移更丰富的结构知识，特征蒸馏方法对齐中间信号：FitNets [9 (https://arxiv.org/html/2606.05682#bib.bib9)] 最小化投影隐藏状态之间的特征回归损失，注意力迁移 [10 (https://arxiv.org/html/2606.05682#bib.bib10)] 匹配注意力图，相似性保持蒸馏 [11 (https://arxiv.org/html/2606.05682#bib.bib11)] 对齐实例级亲和结构。

CKA 风格的表示匹配也已被探索为蒸馏信号。先前的工作将 CKA 适配到 BERT 蒸馏中迁移特征结构 [12 (https://arxiv.org/html/2606.05682#bib.bib12)]，并重新审视 CKA 作为视觉知识蒸馏的表示对齐目标 [13 (https://arxiv.org/html/2606.05682#bib.bib13)]。因此，我们的贡献不是将 CKA 作为通用蒸馏损失引入。相反，我们研究 CKA 作为 NVFP4 QAD 的表示保持正则化器，其中教师和学生共享架构，但低比特量化和仅 KL 恢复仍可引起显著的内部几何漂移。

虽然在全精度场景中有效，但这些方法在低比特恢复中面临挑战。首先，NVFP4 的块级 E4M3 缩放和动态范围压缩导致非平稳激活分布，使直接激活匹配变得脆弱。其次，不考虑表示不变性的特征匹配目标可能迫使学生追求激活相等性而不是几何保持。因此，特征级蒸馏在 4 位推理恢复中的采用有限，logit 级 KL 仍是事实上的标准 [1 (https://arxiv.org/html/2606.05682#bib.bib1)]。

### 2.4 差距：输出对齐与内部几何结构

当前的 QAD 流水线是在一个隐含假设下运行的，即最小化输出级 KL 散度也保留了内部表示保真度。我们的实证分析揭示了这一假设是错误的：仅优化 KL 会加剧表示漂移，这通过 Transformer 和 Mamba 层中 CKA 分数的降低得到量化。这种漂移在 RL 后训练模型中尤其明显，学生模型可能学习满足蒸馏目标的捷径映射，同时以损害复杂推理和代码生成的方式重组内部特征空间。

叠加假说 [4 (https://arxiv.org/html/2606.05682#bib.bib4)] 为这个问题为何重要提供了一个视角：神经网络可以通过将特征打包到低干扰方向来表示比维度更多的特征，使得激活的几何结构在功能上承担着重要责任。NVFP4 QAD 中输出对齐与内部几何结构之间的特定脱节仍未得到充分探索，将 CKA 作为 4 位 LLM 恢复的实际对齐信号也是如此。我们的新颖之处在于将 CKA 直接集成到 QAD 目标中，作为 NVFP4 QAD 和低比特 LLM 精度恢复的 CKA 引导表示对齐，表明低比特蒸馏受益于通过 KL 进行功能对齐和通过 CKA 进行表示对齐的联合优化。

## 3 方法

### 3.1 总体框架

我们基于标准的量化感知蒸馏（QAD）流水线 [1 (https://arxiv.org/html/2606.05682#bib.bib1)]，该流水线通过 KL 散度将量化后的 NVFP4 学生模型与冻结的 BF16 教师模型对齐。虽然 QAD 有效地恢复了输出级功能精度，但我们的分析表明它常常引起内部表示漂移。为了缓解这个问题，我们在 QAD 目标中增加了一个中心核对齐（CKA）正则化项，显式保留中间激活的几何结构。

CKA 引导的对齐鼓励量化学生模型保留教师模型的内部表示几何结构，而不仅仅是匹配输出 logits。这在低比特场景中尤为关键，因为量化引起的尺度变化和裁剪使得逐点特征匹配（例如，基于 MSE 的对齐）变得脆弱。相比之下，CKA 在理论上对正交变换和特征缩放不变，从而对量化伪影具有鲁棒性。因此，CKA 是 KL 基 QAD 的自然补充：KL 对齐功能输出分布，而 CKA 保留中间表示几何结构，这对于在激进量化下进行稳健推理至关重要。

### 3.2 CKA 正则化蒸馏目标

让 zt\(x\)z