Résumé technique : Au-delà de l'appariement des sorties dans la distillation NVFP4 des LLM

1. Énoncé du problème

Alors que les modèles de langage de grande taille (LLM) sont déployés dans des environnements de production contraints en latence et en mémoire, la demande pour des formats d'inférence à faible précision, spécifiquement le NVFP4 (un format de virgule flottante 4 bits avec mise à l'échelle par blocs fine), est croissante. Bien que la Distillation avec Quantification (QAD - Quantization-Aware Distillation) soit apparue comme une méthode robuste pour récupérer l'exactitude perdue lors de la quantification post-entraînement (PTQ) via la divergence de KL entre la distribution de sortie d'un étudiant quantifié et celle d'un enseignant en pleine précision, cet article identifie une faille critique dans le paradigme prédominant.

Les auteurs soutiennent que l'alignement au niveau de la sortie est insuffisant à lui seul pour une récupération fidèle des capacités. Leur analyse révèle que la QAD standard, qui optimise uniquement la divergence de KL, peut masquer une dégradation significative de la représentation interne. Plus précisément, bien qu'un modèle étudiant puisse réussir à correspondre aux logits de sortie de l'enseignant, la géométrie de ses activations intermédiaires subit une réorganisation involontaire (dérive). Cet « apprentissage par raccourci » (shortcut learning) est particulièrement sévère dans les modèles ayant subi un post-entraînement par apprentissage par renforcement (RL), où les hiérarchies sémantiques internes nécessaires aux tâches complexes de raisonnement et de codage sont compromises malgré une haute précision de sortie.

2. Méthodologie : CKA-QAD

Pour remédier au décalage entre l'appariement de sortie et la géométrie interne, les auteurs proposent la CKA-QAD, une méthode qui augmente la QAD standard par un régularisateur d'Alignement par Noyau Centré (CKA - Centered Kernel Alignment).

Mécanisme central

Au lieu de se reposer exclusivement sur la distillation au niveau des logits, la CKA-QAD introduit un régularisateur léger qui aligne les matrices de Gram par couche des activations de l'étudiant et de l'enseignant.

• Métrique CKA : La méthode utilise la CKA linéaire, qui mesure la similitude entre les matrices d'activation centrées ($\bar{H}_t$ et $\bar{H}_s$). La CKA est choisie pour son invariance aux transformations orthogonales et à l'échelle isotrope, ce qui la rend robuste aux distributions d'activation non stationnaires et aux changements d'échelle inhérents à la quantification NVFP4.
• Fonction d'objectif : La perte totale combine la perte de divergence de KL Top-K standard ($L_{TopK-KL}$) avec un terme d'alignement CKA dynamique ($L_{CKA}$).
  $$L_{total} = L_{TopK-KL} + \text{sg}\left(\frac{L_{TopK-KL}}{L_{CKA} + \epsilon}\right) L_{CKA}$$
  L'utilisation d'un opérateur de stop-gradient (sg) et d'une mise à l'échelle dynamique permet d'équilibrer le terme CKA par rapport à la perte de logit sans réglage manuel de l'hyperparamètre global $\lambda$.

Détails de mise en œuvre

• Sélection des couches : La CKA est calculée sur les sorties des blocs décodeurs appariés (post-attention/MLP/mises à jour Mamba). Les auteurs excluent explicitement les activations de pré-normalisation, les poids d'attention et les états du cache KV afin d'éviter les artefacts provenant des échelles de flux résiduels et des dépendances autorégressives.
• Optimisations d'ingénierie : Pour gérer la surcharge de calcul de la CKA à travers de nombreuses couches :
  1. Distillation Top-K : La perte KL est calculée uniquement sur les 8192 meilleurs logits de l'enseignant.
  2. CKA dans l'espace des caractéristiques : La CKA est calculée via les matrices de covariance des caractéristiques ($\bar{H}^\top \bar{H}$) plutôt que par des matrices de Gram explicites jeton-à-jet, réduisant ainsi la complexité mémoire.
  3. Capture par Hooks : Les activations sont capturées via des hooks de propagation directe (forward hooks), les activations de l'enseignant étant détachées du graphe de calcul.

3. Contributions clés

1. Diagnostic des modes de défaillance de la QAD : L'article fournit un diagnostic au niveau de la représentation montrant que la distillation par simple KL peut exacerber la dérive de représentation, particulièrement dans les couches FFN (Feed-Forward Network) de profondeur intermédiaire et les blocs d'état Mamba. Cette dérive est corrélée aux goulots d'étranglement de performance sur les tâches de raisonnement et de codage.
2. QAD NVFP4 guidée par CKA : L'introduction de la CKA-QAD, une méthode qui préserve explicitement la géométrie de la représentation interne pendant la distillation à faible bit. Elle restaure l'alignement à des niveaux proches du BF16 avec un surcoût minimal.
3. Validation empirique : Évaluation complète sur deux architectures distinctes :
   • Nemotron 3 Nano : Un modèle hybride Mamba–Transformer MoE avec post-entraînement par RL.
   • Qwen3-4B-Thinking-2507 : Un modèle de raisonnement compact de type Transformer.

4. Résultats expérimentaux

Les auteurs ont évalué la CKA-QAD par rapport aux enseignants BF16, aux baselines PTQ NVFP4 et à la QAD standard (uniquement par KL).

Alignement de la représentation

• Nemotron 3 Nano : La QAD standard (KL uniquement) a entraîné une chute significative de la CKA moyenne (0,958) par rapport à la PTQ (0,993), avec un effondrement sévère de la dernière couche (0,740). La CKA-QAD a restauré la CKA moyenne à 0,994 et la CKA de la dernière couche à 0,985, rejoignant presque l'enseignant BF16.
• Qwen3-4B : Des tendances similaires ont été observées, la QAD par KL faisant chuter la CKA moyenne à 0,98 et la dernière couche à 0,86. La CKA-QAD a restauré ces valeurs à 0,99 et 0,97, respectivement.

Précision en aval (Downstream Accuracy)

La CKA-QAD a démontré des améliorations constantes sur les benchmarks de raisonnement complexe et de codage :

• AIME25 (Raisonnement mathématique) : Amélioration de 68,5 % (KL uniquement) à 72,3 % (CKA-QAD) sur Qwen3-4B.
• GPQA-D (Connaissances expertes) : Amélioration de 59,5 % à 61,1 %.
• LiveCodeBench-v5 (Génération de code) : Amélioration de 57,9 % à 59,8 %.
• Nemotron 3 Nano : A montré une récupération similaire, avec une amélioration de l'AIME25 de 87,9 % à 88,3 % et du LiveCodeBench de 68,9 % à 70,3 %.

Surcharge d'efficacité

La méthode introduit des coûts d'entraînement modérés :

• Temps de calcul (Wall-clock time) : +0,5 % par étape.
• VRAM de crête : +7,0 %.
• Débit (Throughput) : -0,5 %.
• Convergence : Les étapes pour atteindre la cible KL ont été légèrement réduites (-3,3 %).

5. Signification et affirmations

L'article positionne l'alignement de la représentation guidé par la CKA comme un complément pratique à l'appariement de sortie pour la récupération de la quantification des LLM. Les auteurs affirment que :

• L'alignement fonctionnel est nécessaire mais insuffisant : Faire correspondre les distributions de sortie ne garantit pas la préservation de la topologie sémantique interne requise pour le raisonnement profond.
• La préservation de la géométrie aide la généralisation à faible bit : En empêissant l'étudiant d'apprendre des correspondances de « raccourcis » qui satisfont la fonction de perte mais dégradent la structure interne, la régularisation CKA aide la quantification agressive à conserver les hiérarchies sémantiques fines encodées lors du post-entraînement (particulièrement le RL).
• Robustesse aux artefacts de quantification : Contrairement aux pertes de caractéristiques point par point (ex: MSE), l'invariance de la CKA à l'échelle et à la rotation la rend particulièrement adaptée au NVFP4, où la mise à l'échelle par bloc induit des distributions non stationnaires.

Les auteurs concluent que la préservation de la géométrie interne est cruciale pour les tâches nécessitant une propagation de caractéristiques multi-étapes, telles que le raisonnement mathématique et la synthèse de code, et que la CKA-QAD offre une voie viable pour le déploiement de LLM à faible bit de haute fidélité avec un surcoût d'ingénierie minimal.