Sintesi Tecnica: Oltre il Matching dell'Output nella Distillazione di LLM in NVFP4

1. Definizione del Problema

Con l'impiego di Large Language Models (LLM) in ambienti di produzione con vincoli di latenza e memoria, vi è una crescente domanda di formati di inferenza a bassa precisione, specificamente NVFP4 (un formato a virgola mobile a 4 bit con scaling a blocchi granulare). Sebbene la Quantization-Aware Distillation (QAD) sia emersa come un metodo robusto per recuperare l'accuratezza persa durante la quantizzazione post-training (PTQ) tramite l'allineamento della distribuzione dell'output di uno studente quantizzato con quella di un docente a precisione intera tramite divergenza KL, questo articolo identifica un difetto critico nel paradigma prevalente.

Gli autori sostengono che l'allineamento solo a livello di output è insufficiente per un recupero fedele delle capacità. La loro analisi rivela che la QAD standard, che ottimizza esclusivamente la divergenza KL, può mascherare una significativa degradazione della rappresentazione interna. Nello specifico, mentre lo studente può riuscire a corrispondere ai logit del docente, le sue geometrie di attivazione intermedia subiscono una riorganizzazione indesiderata (drift). Questo "apprendimento di scorciatoie" (shortcut learning) è particolarmente grave nei modelli che hanno subito il post-training tramite Reinforcement Learning (RL), dove le gerarchie semantiche interne necessarie per compiti complessi di ragionamento e coding vengono compromesse nonostante l'elevata accuratezza dell'output.

2. Metodologia: CKA-QAD

Per affrontare il disallineamento tra l'output e la geometria interna, gli autori propongono CKA-QAD, un metodo che aumenta la QAD standard con un regolarizzatore Centered Kernel Alignment (CKA).

Meccanismo Core

Invece di fare affidamento esclusivamente sulla distillazione a livello di logit, CKA-QAD introduce un regolarizzatore leggero che allinea le matrici di Gram per livello delle attivazioni dello studente e del docente.

• Metrica CKA: Il metodo utilizza la CKA lineare, che misura la similarità tra le matrici di attivazione centrate ($\bar{H}_t$ e $\bar{H}_s$). La CKA è stata scelta per la sua invarianza alle trasformazioni ortogonali e allo scaling isotropo, rendendola robusta alle distribuzioni di attivazione non stazionarie e agli shift di scala inerenti alla quantizzazione NVFP4.
• Funzione Obiettivo: Il loss totale combina lo standard Top-K KL divergence loss ($L_{TopK-KL}$) con un termine di allineamento CKA dinamico ($L_{CKA}$).
  $$L_{total} = L_{TopK-KL} + \text{sg}\left(\frac{L_{TopK-KL}}{L_{CKA} + \epsilon}\right) L_{CKA}$$
  L'uso di un operatore stop-gradient (sg) e dello scaling dinamico permette di bilanciare il termine CKA rispetto al loss dei logit senza dover regolare manualmente l'iperparametro globale $\lambda$.

Dettagli di Implementazione

• Selezione dei Livelli: La CKA è calcolata sulle uscite dei blocchi decoder accoppiati (post-attention/MLP/aggiornamenti Mamba). Gli autori escludono esplicitamente le attivazioni di pre-normalizzazione, i pesi dell'attenzione e gli stati della KV cache per evitare artefatti derivanti dalle scale del residual stream e dalle dipendenze autoregressive.
• Ottimizzazioni Ingegneristiche: Per gestire l'overhead computazionale del calcolo della CKA attraverso molti livelli:
  1. Top-K Distillation: Il loss KL è calcolato solo sui primi 8192 logit del docente.
  2. Feature-Space CKA: La CKA è calcolata tramite matrici di covarianza delle feature ($\bar{H}^\top \bar{H}$) piuttosto che tramite esplicite matrici di Gram token-token, riducendo la complessità di memoria.
  3. Capture tramite Hook: Le attivazioni sono catturate tramite forward hook, con le attivazioni del docente staccate dal grafo computazionale.

3. Contributi Chiave

1. Diagnosi dei Modi di Fallimento della QAD: L'articolo fornisce una diagnosi a livello di rappresentazione mostrando che la distillazione basata solo su KL può esacerbare il drift rappresentazionale, particolarmente nei livelli FFN dei Transformer di profondità media e nei blocchi di stato Mamba. Questo drift correla con colli di bottiglia nelle prestazioni su compiti di ragionamento e coding.
2. CKA-Guided NVFP4 QAD: L'introduzione di CKA-QAD, un metodo che preserva esplicitamente la geometria rappresentazionale interna durante la distillazione a basso bit. Esso ripristina l'allineamento a livelli vicini a BF16 con un overhead minimo.
3. Validazione Empirica: Valutazione completa su due architetture distinte:
   • Nemotron 3 Nano: Un modello ibrido Mamba–Transformer MoE con post-training RL.
   • Qwen3-4B-Thinking-2507: Un modello di ragionamento Transformer compatto.

4. Risultati Sperimentali

Gli autori hanno valutato CKA-QAD rispetto a docenti BF16, baseline PTQ NVFP4 e la standard QAD basata solo su KL.

Allineamento Rappresentazionale

• Nemotron 3 Nano: La QAD standard (solo KL) ha causato un calo significativo della CKA media (0.958) rispetto alla PTQ (0.993), con un collasso severo nell'ultimo livello (0.740). CKA-QAD ha ripristinato la CKA media a 0.994 e la CKA dell'ultimo livello a 0.985, eguagliando quasi il docente BF16.
• Qwen3-4B: Sono stati osservati trend simili, con la QAD basata solo su KL che ha abbassato la CKA media a 0.98 e l'ultimo livello a 0.86. CKA-QAD ha ripristinato questi valori rispettivamente a 0.99 e 0.97.

Accuratezza nei Task Successivi (Downstream)

CKA-QAD ha dimostrato miglioramenti costanti nei benchmark di ragionamento complesso e coding:

• AIME25 (Ragionamento Matematico): Migliorato dal 68.5% (solo KL) al 72.3% (CKA-QAD) su Qwen3-4B.
• GPQA-D (Conoscenza Esperta): Migliorato dal 59.5% al 61.1%.
• LiveCodeBench-v5 (Generazione di Codice): Migliorato dal 57.9% al 59.8%.
• Nemotron 3 Nano: Ha mostrato un recupero simile, con AIME25 che migliora dall'87.9% all'88.3% e LiveCodeBench dal 68.9% al 70.3%.

Overhead di Efficienza

Il metodo introduce costi di addestramento modesti:

• Tempo di esecuzione (Wall-clock time): +0.5% per step.
• VRAM di picco: +7.0%.
• Throughput: -0.5%.
• Convergenza: I passi per raggiungere il target KL sono stati leggermente ridotti (-3.3%).

5. Significato e Claim

Il paper posiziona l'allineamento della rappresentazione guidato da CKA come un complemento pratico al matching dell'output per il recupero di LLM quantizzati. Gli autori affermano che:

• L'allineamento funzionale è necessario ma non sufficiente: Far corrispondere le distribuzioni di output non garantisce la preservazione della topologia semantica interna necessaria per il ragionamento profondo.
• La preservazione della geometria aiuta la generalizzazione a basso bit: Impedendo allo studente di apprendere mappature di "scorciatoia" che soddisfano la funzione di loss ma degradano la struttura interna, la regolarizzazione CKA aiuta la quantizzazione aggressiva a mantenere le gerarchie semantiche fini codificate durante il post-training (specialmente l'RL).
• Robustezza agli Artefatti di Quantizzazione: A differenza delle perdite di feature pointwise (es. MSE), l'invarianza della CKA rispetto a scala e rotazione la rende unicamente adatta per NVFP4, dove lo scaling per blocco induce distribuzioni non stazionarie.

Gli autori concludono che preservare la geometria interna è critico per compiti che richiedono la propagazione di feature multi-step, come il ragionamento matematico e la sintesi di codice, e che CKA-QAD offre una via percorribile per il deployment di LLM a basso bit ad alta fedeltà con un overhead ingegneristico minimo.