让低精度大模型保持内部结构,提升推理和编码能力。
Beyond Output Matching: Preserving Internal Geometry in NVFP4 LLM Distillation

- 用CKA度量层间表示相似性,发现仅匹配输出会掩盖内部退化。
- 在Nemotron 3 Nano和Qwen3-4B上,推理与编码准确率显著提升。
- 轻量正则项引导内部几何对齐,适合资源受限场景的模型压缩。
随着大语言模型在延迟与成本敏感的生产环境中部署增多,基于NVFP4的低精度推理需求激增。量化感知微调(QAD)通过使量化学生模型匹配冻结高精度教师模型的输出分布来恢复精度,采用KL散度损失。本文首次从表征层面诊断QAD:仅匹配输出会掩盖内部退化,因多种中间激活几何可产生相似对齐的输出。通过使用中心核相关性(CKA)分析,我们发现仅使用KL的QAD会降低各层表示相似性,尤其在强化学习后训练模型中退化严重。该退化与推理与编码任务的下游瓶颈相关,表明低比特恢复需保留内部几何结构而非仅匹配输出。为此,我们提出CKA-QAD——一种基于CKA的表征对齐方法,通过在蒸馏中对齐层间格拉姆矩阵来保持内部表示几何。在Nemotron 3 Nano和Qwen3-4B-Thinking-2507上,该方法显著提升表征对齐度,并以较小训练开销改善下游推理与编码性能。研究结果表明,基于CKA的表征对齐是低比特模型恢复的有效补充。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。