🤖 AI

Beyond Output Matching: Preserving Internal Geometry in NVFP4 LLM Distillatio

This paper identifies that output-matching-only quantization-aware distillation (QAD) fails to preserve the internal representational geometry of NVFP4 models, leading to performance degradation in reasoning and coding tasks, and proposes CKA-QAD, a method that aligns layer-wise Gram matrices to recover accuracy by maintaining internal structure.

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

Published 2026-06-05

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant master chef (the Teacher) who can cook a perfect, complex dish. Now, you want to teach a student chef (the Student) to cook the exact same dish, but with a major catch: the student is only allowed to use a tiny, low-quality set of measuring cups and spoons (this is NVFP4 quantization, a way to shrink the model's size for faster, cheaper use).

Usually, to teach the student, you'd just say, "Make sure the final taste of your dish matches mine exactly." In the world of AI, this is called Output Matching (or KL Divergence). You check the final result (the "logits" or probability of the next word) and tell the student, "If your dish tastes like mine, you're doing a good job."

The Problem: The "Fake" Match

The authors of this paper discovered a surprising flaw in this method.

They found that the student chef could cheat. The student could rearrange their ingredients and cooking steps in a completely weird, chaotic way, but somehow, the final taste still matched the master chef's dish.

The Reality: The student's internal kitchen (the internal geometry) was a mess. The ingredients were mixed in the wrong order, the heat was applied strangely, and the "flavor map" was broken.
The Illusion: Because the final taste was right, the teacher thought the student was learning. But when the student tried to cook a new, complex dish (like a math problem or code), they failed because their internal kitchen was disorganized.

The paper calls this "Output Matching masking Internal Degradation." The student looks good on the surface, but their brain is scrambled underneath. This was especially bad for models that had been trained to "think" deeply (Reinforcement Learning models), where the internal logic is crucial.

The Solution: CKA-QAD (The "Kitchen Blueprint" Check)

To fix this, the authors introduced a new teaching method called CKA-QAD.

Instead of just checking the final taste, they added a rule: "Your kitchen layout must look like mine."

They used a mathematical tool called CKA (Centered Kernel Alignment). Think of CKA as a blueprint scanner. It doesn't care if the student uses different pots or pans (it ignores minor scale changes); it cares if the relationship between the ingredients is the same.

If the master chef puts salt before pepper, the student must do the same.
If the master chef mixes flour and eggs in a specific pattern, the student must mimic that pattern.

By adding this "blueprint check" to the training, the student is forced to keep their internal kitchen organized, not just fake the final taste.

What Happened When They Tried It?

The researchers tested this on two different AI models (one huge hybrid model and one smaller reasoning model).

The Old Way (Just Taste): The student could mimic the final answer, but their internal "brain" drifted away from the teacher's. When asked to solve hard math or coding problems, they stumbled.
The New Way (Taste + Blueprint): The student learned to keep their internal kitchen organized while matching the taste.
- Result: The student didn't just get the right answers; they got them better on complex tasks like math (AIME25) and coding (LiveCodeBench).
- Efficiency: This new method only added a tiny bit of extra time (about 0.5%) and memory (about 7%) to the training process. It was a small price to pay for a much smarter student.

The Big Takeaway

The paper concludes that for AI models to work well when shrunk down (compressed), you can't just look at the final answer. You have to make sure the internal structure of the model stays true to the original.

In simple terms: If you want a tiny AI to be smart, don't just tell it "get the right answer." Tell it, "Think the right way, too." By checking how the model thinks (its internal geometry), not just what it says, we can make small, fast AI models that are actually capable of deep reasoning.

Technical Summary: Beyond Output Matching in NVFP4 LLM Distillation

1. Problem Statement

As Large Language Models (LLMs) are deployed in latency and memory-constrained production environments, there is a growing demand for low-precision inference formats, specifically NVFP4 (a 4-bit floating-point format with fine-grained block scaling). While Quantization-Aware Distillation (QAD) has emerged as a robust method to recover accuracy lost during post-training quantization (PTQ) by aligning a quantized student's output distribution with a full-precision teacher via KL divergence, this paper identifies a critical flaw in the prevailing paradigm.

The authors argue that output-level alignment alone is insufficient for faithful capability recovery. Their analysis reveals that standard QAD, which optimizes solely for KL divergence, can mask significant internal representational degradation. Specifically, while the student model may successfully match the teacher's output logits, its intermediate activation geometries undergo unintended reorganization (drift). This "shortcut learning" is particularly severe in models that have undergone Reinforcement Learning (RL) post-training, where the internal semantic hierarchies required for complex reasoning and coding tasks are compromised despite high output accuracy.

2. Methodology: CKA-QAD

To address the disconnect between output matching and internal geometry, the authors propose CKA-QAD, a method that augments standard QAD with a Centered Kernel Alignment (CKA) regularizer.

Core Mechanism

Instead of relying exclusively on logit-level distillation, CKA-QAD introduces a lightweight regularizer that aligns the layerwise Gram matrices of the student and teacher activations.

CKA Metric: The method utilizes linear CKA, which measures the similarity between centered activation matrices ( $\bar{H}_t$ and $\bar{H}_s$ ). CKA is chosen for its invariance to orthogonal transformations and isotropic scaling, making it robust to the non-stationary activation distributions and scale shifts inherent in NVFP4 quantization.
Objective Function: The total loss combines the standard Top-K KL divergence loss ( $L_{TopK-KL}$ ) with a dynamic CKA alignment term ( $L_{CKA}$ ).
$L_{total} = L_{TopK-KL} + \text{sg}\left(\frac{L_{TopK-KL}}{L_{CKA} + \epsilon}\right) L_{CKA}$
The use of a stop-gradient (sg) operator and dynamic scaling allows the CKA term to be balanced against the logit loss without manual hyperparameter tuning of a global weight $\lambda$ .

Implementation Details

Layer Selection: CKA is computed on the outputs of matched decoder blocks (post-attention/MLP/Mamba updates). The authors explicitly exclude pre-normalization activations, attention weights, and KV cache states to avoid artifacts from residual stream scales and autoregressive dependencies.
Engineering Optimizations: To manage the computational overhead of computing CKA across many layers:
1. Top-K Distillation: KL loss is computed only on the teacher's top 8192 logits.
2. Feature-Space CKA: CKA is calculated via feature covariance matrices ( $\bar{H}^\top \bar{H}$ ) rather than explicit token-token Gram matrices, reducing memory complexity.
3. Hook-Based Capture: Activations are captured via forward hooks, with teacher activations detached from the computation graph.

3. Key Contributions

Diagnosis of QAD Failure Modes: The paper provides a representation-level diagnosis showing that KL-only distillation can exacerbate representational drift, particularly in mid-depth Transformer FFN layers and Mamba state blocks. This drift correlates with performance bottlenecks on reasoning and coding tasks.
CKA-Guided NVFP4 QAD: The introduction of CKA-QAD, a method that explicitly preserves internal representational geometry during low-bit distillation. It restores alignment to near BF16 levels with minimal overhead.
Empirical Validation: Comprehensive evaluation on two distinct architectures:
- Nemotron 3 Nano: A hybrid Mamba–Transformer MoE model with RL post-training.
- Qwen3-4B-Thinking-2507: A compact Transformer reasoning model.

4. Experimental Results

The authors evaluated CKA-QAD against BF16 teachers, NVFP4 PTQ baselines, and standard KL-only QAD.

Representational Alignment

Nemotron 3 Nano: Standard QAD (KL-only) resulted in a significant drop in average CKA (0.958) compared to PTQ (0.993), with a severe collapse in the last layer (0.740). CKA-QAD restored average CKA to 0.994 and last-layer CKA to 0.985, nearly matching the BF16 teacher.
Qwen3-4B: Similar trends were observed, with KL-only QAD dropping average CKA to 0.98 and last-layer CKA to 0.86. CKA-QAD restored these to 0.99 and 0.97, respectively.

Downstream Accuracy

CKA-QAD demonstrated consistent improvements in complex reasoning and coding benchmarks:

AIME25 (Math Reasoning): Improved from 68.5% (KL-only) to 72.3% (CKA-QAD) on Qwen3-4B.
GPQA-D (Expert Knowledge): Improved from 59.5% to 61.1%.
LiveCodeBench-v5 (Code Generation): Improved from 57.9% to 59.8%.
Nemotron 3 Nano: Showed similar recovery, with AIME25 improving from 87.9% to 88.3% and LiveCodeBench from 68.9% to 70.3%.

Efficiency Overhead

The method introduces modest training costs:

Wall-clock time: +0.5% per step.
Peak VRAM: +7.0%.
Throughput: -0.5%.
Convergence: Steps to reach target KL were slightly reduced (-3.3%).

5. Significance and Claims

The paper positions CKA-guided representational alignment as a practical complement to output matching for quantized LLM recovery. The authors claim that:

Functional alignment is necessary but insufficient: Matching output distributions does not guarantee the preservation of the internal semantic topology required for deep reasoning.
Geometry preservation aids low-bit generalization: By preventing the student from learning "shortcut" mappings that satisfy the loss function but degrade internal structure, CKA regularization helps aggressive quantization retain fine-grained semantic hierarchies encoded during post-training (especially RL).
Robustness to Quantization Artifacts: Unlike pointwise feature losses (e.g., MSE), CKA's invariance to scaling and rotation makes it uniquely suited for NVFP4, where block-wise scaling induces non-stationary distributions.

The authors conclude that preserving internal geometry is critical for tasks requiring multi-step feature propagation, such as mathematical reasoning and code synthesis, and that CKA-QAD offers a viable path for high-fidelity low-bit LLM deployment with minimal engineering overhead.