Technische Samenvatting: Voorbij Output-matching in NVFP4 LLM Distillatie

1. Probleemstelling

Naarmate Large Language Models (LLMs) worden ingezet in productieomgevingen met beperkte latentie en geheugen, groeit de vraag naar precisie-arme inferentievormaten, specifiek NVFP4 (een 4-bit floating-point formaat met fijngestelde blokschaling). Hoewel Quantization-Aware Distillation (QAD) een robuuste methode is geworden om het verlies aan nauwkeurigheid te herstellen dat optreedt tijdens post-training quantisatie (PTQ) via KL-divergentie tussen de outputdistributie van een gekwantiseerde student en een full-precision teacher, identificeert dit artikel een kritiek gebrek in het heersende paradigma.

De auteurs stellen dat alignement op outputniveau alleen onvoldoende is voor getrouwe capaciteitsherstel. Hun analyse laat zien dat standaard QAD, dat uitsluitend optimaliseert voor KL-divergentie, een significante interne representatieve degradatie kan maskeren. Specifiek: hoewel het studentmodel succesvol de logits van de teacher kan matchen, ondergaan de intermediaire activatiegeometrieën een onbedoelde reorganisatie (drift). Dit "shortcut learning" is bijzonder ernstig in modellen die een Reinforcement Learning (RL) post-training hebben ondergaan, waarbij de interne semantische hiërarchieën die nodig zijn voor complexe redenering en coderingstaken in gevaar worden gebracht, ondanks een hoge output-nauwkeurigheid.

2. Methodologie: CKA-QAD

Om de discrepantie tussen output-matching en interne geometrie aan te pakken, stellen de auteurs CKA-QAD voor, een methode die standaard QAD aanvult met een Centered Kernel Alignment (CKA) regularisator.

Kernmechanisme

In plaats van uitsluitend te vertrouwen op distillatie op logit-niveau, introduceert CKA-QAD een lichtgewicht regularisator die de layerwise Gram-matrices van de student- en teacher-activaties uitlijnt.

• CKA-metriek: De methode maakt gebruik van lineaire CKA, die de gelijkenis meet tussen gecentreerde activatiematrices ($\bar{H}_t$ en $\bar{H}_s$). CKA is gekozen vanwege de invariantie tegen orthogonale transformaties en isotrope schaling, wat het robuust maakt tegen de niet-stationaire activatieverdelingen en schaalverschuivingen die inherent zijn aan NVFP4-quantisatie.
• Doelfunctie: De totale loss combineert de standaard Top-K KL-divergentie loss ($L_{TopK-KL}$) met een dynamische CCA-alignment term ($L_{CKA}$).
  $$L_{total} = L_{TopK-KL} + \text{sg}\left(\frac{L_{TopK-KL}}{L_{CKA} + \epsilon}\right) L_{CKA}$$
  Het gebruik van een stop-gradient (sg) operator en dynamische schaling maakt het mogelijk om de CKA-term te balanceren tegen de logit-loss zonder handmatige hyperparameterinstelling van een globale gewicht $\lambda$.

Implementatiedetails

• Laalselectie: CKA wordt berekend op de outputs van gematchte decoder-blokken (post-attention/MLP/Mamba updates). De auteurs sluiten expliciet pre-normalisatie activaties, attention weights en KV-cache staten uit om artefacten van residual stream schalen en autoregressieve afhankelijkheden te vermijden.
• Engineering Optimalisaties: Om de computationele overhead van het berekenen van CKA over vele lagen te beheersen:
  1. Top-K Distillatie: KL-loss wordt alleen berekend op de top 8192 logits van de teacher.
  2. Feature-Space CKA: CKA wordt berekend via feature covariantie-matrices ($\bar{H}^\top \bar{H}$) in plaats van expliciete token-token Gram-matrices, wat de geheugencomplexiteit vermindert.
  3. Hook-Based Capture: Activaties worden gevangen via forward hooks, waarbij de teacher-activaties worden losgekoppeld van de computationele graaf.

3. Belangrijkste Bijdragen

1. Diagnose van QAD-foutmodi: Het artikel biedt een representatieve diagnose die aantoont dat KL-only distillatie de representatieve drift kan verergeren, met name in mid-depth Transformer FFN-lagen en Mamba state-blokken. Deze drift correleert met prestatie-bottlenecks bij redenerings- en coderingstaken.
2. CKA-gestuurde NVFP4 QAD: De introductie van CKA-QAD, een methode die expliciet de interne representatieve geometrie behoudt tijdens low-bit distillatie. Het herstelt de alignment naar bijna BF16-niveaus met minimale overhead.
3. Empirische Validatie: Uitgebreide evaluatie op twee verschillende architecturen:
   • Nemotron 3 Nano: Een hybride Mamba–Transformer MoE model met RL post-training.
   • Qwen3-4B-Thinking-2507: Een compacte Transformer reasoning model.

4. Experimentele Resultaten

De auteurs evalueerden CKA-QAD tegenover BF16 teachers, NVFP4 PTQ baselines en standaard KL-only QAD.

Representatieve Alignment

• Nemotron 3 Nano: Standaard QAD (KL-only) resulteerde in een significante daling van de gemiddelde CKA (0.958) vergeleken met PTQ (0.993), met een ernstige instorting in de laatste laag (0.740). CKA-QAD herstelde de gemiddelde CKA naar 0.994 en de laatste-laag CKA naar 0.985, wat bijna overeenkomt met de BF16 teacher.
• Qwen3-4B: Vergelijkbare trends werden waargenomen, waarbij KL-only QAD de gemiddelde CKA deed dalen naar 0.98 en de laatste-laag CKA naar 0.86. CKA-QAD herstelde deze naar respectievelijk 0.99 en 0.97.

Downstream Nauwkeurigheid

CKA-QAD toonde consistente verbeteringen in complexe redenerings- en codering-benchmarks:

• AIME25 (Wiskundig Redeneren): Verbeterde van 68.5% (KL-only) naar 72.3% (CKA-QAD) op Qwen3-4B.
• GPQA-D (Expert Kennis): Verbeterde van 59.5% naar 61.1%.
• LiveCodeBench-v5 (Codegeneratie): Verbeterde van 57.9% naar 59.8%.
• Nemotron 3 Nano: Toonde een vergelijkbaar herstel, waarbij AIME25 verbeterde van 87.9% naar 88.3% en LiveCodeBench van 68.9% naar 70.3%.

Efficiëntie Overhead

De methode introduceert bescheiden trainingskosten:

• Wall-clock tijd: +0.5% per stap.
• Peak VRAM: +7.0%.
• Doorvoersnelheid (Throughput): -0.5%.
• Convergentie: Stappen om de doel-KL te bereiken waren licht gereduceerd (-3.3%).

5. Betekenis en Claims

Het artikel positioneert CKA-gestuurde representatieve alignment als een praktische aanvulling op output-matching voor de recovery van gekwantiseerde LLM's. De auteurs beweren dat:

• Functionele alignment is noodzakelijk maar onvoldoende is: Het matchen van outputverdelingen garandeert niet het behoud van de interne semantische topologie die vereist is voor diep redeneren.
• Geometrie-behoud helpt bij low-bit generalisatie: Door de student te voorkomen dat hij "shortcut" mappingen leert die weliswaar aan de loss-functie voldoen maar de interne structuur degraderen, helpt CKA-regularisatie agressieve quantisatie om fijne semantische hiërarchieën te behouden (vooral tijdens RL).
• Robuustheid tegen Quantisatie-artefacten: In tegenstelling tot pointwise feature losses (zoals MSE), maakt de invariantie van CKA tegen schaling en rotatie het uniek geschikt voor NVFP4, waar block-wise schaling niet-stationaire verdelingen induceert.

De auteurs concluderen dat het behoud van interne geometrie cruciaal is voor taken die multi-step feature propagatie vereisen, zoals wiskundig redeneren en codering, en dat CKA-QAD een levensvatbaar pad biedt voor high-fidelity low-bit LLM-implementatie met minimale engineering overhead.