REVIEW 2 major objections 1 minor 16 references
Beyond Output Matching: Preserving Internal Geometry in NVFP4 LLM Distillation
T0 review · 2 major / 1 minor · reviewed 2026-06-28 · grok-4.3
Pith's one-line read Aligning layerwise Gram matrices via CKA during NVFP4 distillation restores internal geometry that output matching alone allows to drift.
desk verdict CKA regularizer for NVFP4 QAD is a reasonable practical idea but the abstract leaves causality and effect sizes unverified. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
CKA regularizer that aligns layerwise Gram matrices between quantized student and higher-precision teacher during distillation
What would settle it
Running the full CKA-QAD training loop and finding no increase in layerwise CKA or no gain in reasoning and coding accuracy would falsify the claim that geometry preservation drives the improvement.
Extended reading notes
Core claim
KL-only QAD reduces layerwise CKA similarity to the teacher, especially in RL-post-trained models, and this reduction correlates with downstream bottlenecks; explicitly aligning Gram matrices through an added CKA regularizer during NVFP4 distillation restores similarity and improves reasoning and coding accuracy.
Load-bearing premise
The observed correlation between lower CKA similarity and accuracy loss is causal, so that explicitly raising CKA via the regularizer will produce the accuracy gains.
Editorial extensions
If this is right
- Layerwise representational similarity to the BF16 teacher increases across the network.
- Downstream accuracy on reasoning and coding tasks rises relative to standard KL-only QAD.
- The method applies to NVFP4 quantization with only modest extra training cost.
- The gains appear across distinct model families including Nemotron and Qwen variants.
Reading between the lines
- The same Gram-matrix alignment idea could be tested in other compression settings such as pruning or knowledge distillation without quantization.
- Certain layers may be more critical to preserve for particular task families, suggesting selective rather than uniform CKA weighting as a possible refinement.
- If the drift is driven by post-training dynamics, applying CKA-QAD during the RL stage itself rather than only in the distillation phase might yield further gains.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper claims that KL-divergence output matching in quantization-aware distillation (QAD) for NVFP4 LLMs masks internal representational degradation, as shown by reduced layerwise CKA similarity to the BF16 teacher (especially severe in RL-post-trained models). This drift correlates with downstream bottlenecks on reasoning and coding tasks. Motivated by the diagnosis, the authors propose CKA-QAD, which adds a lightweight CKA regularizer to align layerwise Gram matrices, claiming it substantially improves representational alignment and downstream accuracy on Nemotron 3 Nano and Qwen3-4B-Thinking-2507 with modest overhead.
Significance. If the empirical results hold with proper controls and quantitative validation, the work would be moderately significant for low-precision LLM deployment: it identifies a concrete limitation of output-only distillation and offers a practical, low-overhead complement that targets internal geometry preservation. The use of CKA as a diagnostic and regularizer is a straightforward extension of existing representation similarity tools to the quantization setting.
major comments (2)
- [Abstract] Abstract: The central claim that adding the CKA regularizer (rather than other training factors) recovers downstream accuracy is not supported by any described ablation studies, controls, or isolation experiments; the text only reports that KL-only QAD reduces CKA similarity and that this correlates with bottlenecks, without establishing causality for the proposed method.
- [Abstract] Abstract: The abstract states that CKA-QAD 'substantially improves' reasoning and coding accuracy but supplies no effect sizes, baseline comparisons (e.g., vs. standard QAD or other regularizers), number of runs, variance, or statistical tests, leaving the magnitude and reliability of the claimed gains impossible to evaluate.
minor comments (1)
- [Abstract] Abstract: 'NVFP4' is introduced without a definition or citation to prior work on this specific quantization format.
Simulated Author's Rebuttal
We thank the referee for the detailed and constructive comments. We address each major comment below and indicate the revisions we will make to strengthen the manuscript.
read point-by-point responses
-
Referee: [Abstract] Abstract: The central claim that adding the CKA regularizer (rather than other training factors) recovers downstream accuracy is not supported by any described ablation studies, controls, or isolation experiments; the text only reports that KL-only QAD reduces CKA similarity and that this correlates with bottlenecks, without establishing causality for the proposed method.
Authors: We agree that the abstract does not explicitly describe ablation studies isolating the contribution of the CKA regularizer from other training factors. The manuscript presents direct comparisons of KL-only QAD versus CKA-QAD on both representational similarity (CKA) and downstream tasks across two models, which provides correlational support for the benefit of the added regularizer. However, we acknowledge that stronger causal evidence would require additional controls such as ablating the CKA loss weight or comparing against alternative regularizers. We will add these ablation experiments and corresponding discussion in the revised manuscript to better establish causality. revision: yes
-
Referee: [Abstract] Abstract: The abstract states that CKA-QAD 'substantially improves' reasoning and coding accuracy but supplies no effect sizes, baseline comparisons (e.g., vs. standard QAD or other regularizers), number of runs, variance, or statistical tests, leaving the magnitude and reliability of the claimed gains impossible to evaluate.
Authors: We agree that the abstract would benefit from quantitative details. The full manuscript reports specific accuracy improvements on reasoning and coding benchmarks for Nemotron 3 Nano and Qwen3-4B-Thinking-2507 relative to the KL-only baseline, along with CKA similarity metrics. To address the concern, we will revise the abstract to include effect sizes (e.g., absolute and relative accuracy gains), baseline comparisons, and any available statistics on variance or multiple runs. Where multiple runs were not performed due to computational cost, we will note this limitation explicitly. revision: yes
Circularity Check
No significant circularity in derivation chain
full rationale
The paper's chain consists of an empirical observation (KL-only QAD reduces CKA similarity, which correlates with task bottlenecks) followed by the introduction of a new regularizer (CKA-QAD) and its experimental validation on accuracy metrics. No step reduces a claimed prediction or result to its own inputs by construction, no fitted parameter is relabeled as a prediction, and no load-bearing premise depends on a self-citation chain. The contribution rests on new measurements and a proposed loss term whose effects are tested externally rather than being tautological.
Assumptions & free parameters
free parameters (1)
- CKA regularizer weight
assumptions (1)
- domain assumption CKA between layer activations is a meaningful proxy for representational geometry that affects downstream task performance.
Cite this review
Pith. "Pith review of Beyond Output Matching: Preserving Internal Geometry in NVFP4 LLM Distillation." pith.science (2026). https://pith.science/paper/JIMS7PW4
@misc{pith2026260605682,
author = {Pith},
title = {Pith review of: Beyond Output Matching: Preserving Internal Geometry in NVFP4 LLM Distillation},
year = {2026},
howpublished = {\url{https://pith.science/paper/JIMS7PW4}},
note = {Machine review of arXiv:2606.05682}
}
read the original abstract
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.
Figures
Reference graph
Works this paper leans on
-
[1]
Quantization-aware distillation for NVFP4 inference accuracy recovery
Meng Xin, Sweta Priyadarshi, Jingyu Xin, et al. Quantization-aware distillation for NVFP4 inference accuracy recovery. arXiv:2601.20088, 2026
-
[2]
Introducing NVFP4 for efficient and accurate low-precision inference
Eduardo Alvarez, Omri Almog, Eric Chung, Simon Layton, Dusan Stosic, Ronny Krashinsky, and Kyle Aubrey. Introducing NVFP4 for efficient and accurate low-precision inference. NVIDIA Technical Blog, 2025
2025
-
[3]
Similarity of neural network representations revisited
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In ICML, 2019
2019
-
[4]
Toy models of superposition
Nelson Elhage, Tristan Hume, Catherine Olsson, et al. Toy models of superposition. Transformer Circuits Thread, 2022
2022
-
[5]
Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In ICML, 2024
2024
-
[6]
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang, Jianhong Tu, et al. Qwen3 technical report. arXiv:2505.09388, 2025
work page Pith review arXiv 2025
-
[7]
A White Paper on Neural Network Quantization
Markus Nagel, Marios Fournarakis, Rana Ali Amjad, Yelysei Bondarenko, Mart van Baalen, and Tijmen Blankevoort. A white paper on neural network quantization. arXiv:2106.08295, 2021
work page Pith review arXiv 2021
-
[8]
Distilling the Knowledge in a Neural Network
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv:1503.02531, 2015
work page Pith review arXiv 2015
Show all 16 references
-
[9]
FitNets: Hints for thin deep nets
Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. FitNets: Hints for thin deep nets. arXiv:1412.6550, 2014
2014 arXiv
-
[10]
Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer
Sergey Zagoruyko and Nikos Komodakis. Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer. In ICLR, 2017
2017
-
[11]
Similarity-preserving knowledge distillation
Frederick Tung and Greg Mori. Similarity-preserving knowledge distillation. In ICCV, 2019
2019
-
[12]
Feature structure distillation with Centered Kernel Alignment in BERT transferring
Hee-Jun Jung, Doyeon Kim, Seung-Hoon Na, and Kangil Kim. Feature structure distillation with Centered Kernel Alignment in BERT transferring. Expert Systems with Applications, 234:120980, 2023
2023
-
[13]
Rethinking Centered Kernel Alignment in Knowledge Distillation
Zikai Zhou, Yunhang Shen, Shitong Shao, Linrui Gong, and Shaohui Lin. Rethinking Centered Kernel Alignment in Knowledge Distillation. In IJCAI, 2024
2024
-
[14]
American Invitational Mathematics Examination 2025
Mathematical Association of America. American Invitational Mathematics Examination 2025. 2025
2025
-
[15]
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R. Bowman. GPQA: A graduate-level Google-proof Q&A benchmark. arXiv:2311.12022, 2023. 12
2023 arXiv
-
[16]
LiveCodeBench: Holistic and contamination free evaluation of large language models for code
Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang, Ar- mando Solar-Lezama, Koushik Sen, and Ion Stoica. LiveCodeBench: Holistic and contamination free evaluation of large language models for code. arXiv:2403.07974, 2024. 13
2024 arXiv
Reviewed June 28, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.