Beyond Output Matching: Preserving Internal Geometry in NVFP4 LLM Distillation

📄 arXiv: 2606.05682v2 📥 PDF

作者: Fangbo Tu, Junhua Zhao, Chi Liu, Xin Chen, Haifeng Wu, Jian Wan, Srinivasan Manoharan

分类: cs.AI, cs.LG

发布日期: 2026-06-04 (更新: 2026-06-05)

备注: 13 pages,1 figures

💡 一句话要点

提出CKA-QAD以解决低比特量化模型内部几何保留问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 量化感知蒸馏 CKA对齐 低比特量化 内部几何结构 模型蒸馏 自然语言处理 推理任务 编码任务

📋 核心要点

  1. 现有的量化感知蒸馏方法在低比特量化下可能导致内部表示的几何结构退化，影响模型性能。
  2. 本文提出CKA-QAD，通过引入CKA指导的表示对齐，保留内部几何结构，改善低比特量化模型的准确性。
  3. 在Nemotron 3 Nano和Qwen3-4B-Thinking-2507的实验中，CKA-QAD显著提高了表示对齐度和下游任务的准确性。

📝 摘要（中文）

随着大型语言模型在延迟和成本受限的生产环境中的广泛应用，对低精度推理的需求日益增长。量化感知蒸馏（QAD）通过训练量化学生模型，使其输出分布与冻结的高精度教师模型匹配，从而帮助恢复在低比特量化下丢失的准确性。然而，单纯的输出匹配可能掩盖内部退化。本文提出CKA-QAD，一种基于CKA的表示对齐方法，通过对齐层级Gram矩阵来保留内部几何结构，从而显著提高下游推理和编码任务的准确性。

🔬 方法详解

问题定义：本文旨在解决低比特量化模型在蒸馏过程中内部几何结构的退化问题。现有的QAD方法主要关注输出匹配，忽视了内部表示的保留，导致模型在推理和编码任务上的性能下降。

核心思路：CKA-QAD的核心思想是通过CKA（Centered Kernel Alignment）来指导表示对齐，确保在蒸馏过程中保留内部几何结构，而不仅仅是匹配输出。这种设计旨在改善低比特量化模型的整体性能。

技术框架：CKA-QAD的整体架构包括两个主要模块：首先是量化学生模型的训练，其次是通过CKA对齐层级Gram矩阵，以保持内部表示的一致性。该方法在训练过程中引入了轻量级的正则化项。

关键创新：CKA-QAD的主要创新在于引入了CKA作为指导工具，强调内部几何结构的保留。这与传统的仅关注输出匹配的QAD方法有本质区别，提供了一种新的思路来提升低比特量化模型的性能。

关键设计：在CKA-QAD中，损失函数结合了KL散度损失和CKA对齐损失，确保在训练过程中既考虑输出匹配，又保留内部表示的几何结构。网络结构上，采用了适应低比特量化的设计，确保模型在训练和推理阶段的高效性。

🖼️ 关键图片

fig_0

📊 实验亮点

实验结果表明，CKA-QAD在Nemotron 3 Nano和Qwen3-4B-Thinking-2507模型上显著提高了表示对齐度，并在下游推理和编码任务中提升了准确性，具体提升幅度未知，显示出该方法的有效性和实用性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、代码生成和其他需要高效推理的任务。通过改善低比特量化模型的性能，CKA-QAD可以在资源受限的环境中实现更高效的模型部署，具有重要的实际价值和未来影响。

📄 摘要（原文）

Demand for low-precision inference, including NVFP4-based approaches, has grown as large language models are increasingly deployed in latency and cost constrained production environments. Quantization-aware distillation (QAD) helps recover accuracy lost under low bit quantization by training a quantized student to match the output distribution of a frozen higher precision teacher via a KL-divergence loss. In this work, we first provide a representation level diagnosis of QAD: output matching alone can mask internal degradation, because many intermediate activation geometries can yield similar teacher-aligned logits. Using CKA, we show that KL-only QAD can reduce layerwise representational similarity relative to the BF16 teacher, with especially severe drift in RL-post-trained models. This drift correlates with downstream bottlenecks on reasoning and coding tasks, suggesting that low bit recovery requires preserving internal geometry rather than matching outputs alone. Motivated by this finding, we propose \textbf{CKA-QAD}, a CKA-guided representational alignment method for NVFP4 QAD and low bit LLM accuracy recovery. The method adds a lightweight regularizer that preserves internal representational geometry during distillation by aligning layerwise Gram matrices through CKA. Across Nemotron 3 Nano and Qwen3-4B-Thinking-2507, CKA-QAD substantially improves representational alignment and improves downstream reasoning and coding accuracy with modest training overhead. Our findings position CKA-guided representational alignment as a practical complement to output matching for quantized LLM recovery.