Résumé Technique : N-vium – Transformer à Mélange de Sorties pour une Génération Exacte Accélérée

1. Énoncé du Problème

Les modèles de langage larges (LLM) autorégressifs standards traitent les tokens séquentiellement à travers une pile de couches de décodeur. Ce passage avant sert deux fonctions couplées : (1) construire des représentations Clé-Valeur (KV) à chaque couche pour l'attention future, et (2) projeter l'état caché final pour prédire le token suivant. Par conséquent, chaque token nécessite un passage avant complet à travers toutes les couches, indépendamment du fait que la distribution de prédiction ait déjà convergé.

Les approches existantes visant à améliorer l'efficacité de l'inférence, telles que les méthodes de Sortie Anticipée (Early Exit), tentent de découpler ces fonctions en attachant des têtes de prédiction aux couches intermédiaires. Cependant, ces méthodes traitent généralement les prédictions intermédiaires comme des approximations de la distribution de la dernière couche. Cela entraîne plusieurs inconvénients :

• Erreurs d'Approximation : Les prédictions intermédiaires sont souvent plus faibles que celles en profondeur complète, créant un décalage entraînement-test.
• Routage Heuristique : Les décisions sur le moment de sortir reposent souvent sur des seuils de confiance conçus à la main nécessitant un étalonnage.
• Caches KV Incomplets : Pour économiser le calcul, les méthodes de sortie anticipée omettent souvent le calcul des caches KV pour les couches contournées, nécessitant une copie approximative ou un saut, ce qui dégrade la qualité.
• Dégradation de la Qualité : La vitesse est fréquemment obtenue au détriment de la perplexité et de la qualité de génération.

L'article soutient que, bien que les observations empiriques (par exemple, via la "logit lens") montrent que les distributions de prédiction se stabilisent souvent avant la dernière couche, les méthodes actuelles échouent à exploiter cela pour une génération exacte sans perte de qualité.

2. Méthodologie : N-vium

Les auteurs introduisent N-vium, un cadre à mélange de sorties qui généralise strictement l'architecture standard du transformer. Au lieu de traiter les sorties anticipées comme des approximations, N-vium redéfinit la sortie du modèle comme une distribution de mélange apprise sur plusieurs sorties.

Architecture

• Mélange de Sorties : Le modèle partitionne les $L$ couches du socle en $N_V$ blocs. À la fin de chaque bloc (jonction $k$), un routeur détermine la probabilité d'échantillonner le token suivant à partir de la distribution $\pi_k$ de cette couche.
• Routage Adaptatif aux Tokens : Un MLP de routeur traite l'état caché $h^{(\ell_k)}_t$ à la jonction $k$ pour produire une probabilité $w_k$. La probabilité $p_k$ de sortir à la jonction $k$ est calculée comme suit :
  $$p_k(x_t) = w_k(h^{(\ell_k)}_t) \prod_{j=1}^{k-1} (1 - w_j(h^{(\ell_j)}_t))$$
  où $w_{N_V} = 1$ par convention.
• Distribution de Mélange : La distribution finale du token suivant est une somme pondérée de toutes les distributions de sortie :
  $$\pi_{mix}(x_{t+1} | x_{1:t}) = \sum_{k=1}^{N_V} p_k(x_t) \pi_k(x_{t+1} | x_{1:t})$$
• Tête Partagée : Toutes les sorties partagent la tête de projection finale de modélisation du langage (LM) $W_{lm}$, avec des MLP adaptateurs légers transformant les états cachés intermédiaires pour correspondre à la géométrie attendue par la tête.

Entraînement

• Objectif : Le modèle est entraîné pour minimiser la perte d'entropie croisée de la distribution de mélange conjointe $\pi_{mix}$, plutôt que d'optimiser les têtes individuelles de manière isolée pour correspondre à la dernière couche.
• Pénalité de Calcul : Un terme de perte auxiliaire $\mathcal{L}_{compute}$ est ajouté pour encourager le routeur à attribuer une masse de probabilité aux sorties plus précoces, orientant le modèle vers l'efficacité :
  $$\mathcal{L}_{total} = \mathcal{L}_{mix} + \beta \mathcal{L}_{compute}$$
  où $\mathcal{L}_{compute}$ est la profondeur de couche normalisée attendue.
• Échauffement : Une phase d'échauffement du routeur est employée pour prévenir l'effondrement du routeur (où tous les tokens sont routés vers la dernière couche) en pénalisant les écarts par rapport à une distribution de sortie uniforme initialement.

Inférence : Surcharge (Piggybacking)

Une innovation critique est la technique de Surcharge (Piggybacking) pour assurer une récupération exacte du cache KV sans saut de calcul :

1. Calcul Différé : Lorsqu'un token sort à la jonction $k$, le calcul pour les couches $k+1$ à $L$ est différé.
2. Regroupement Parallèle : Au fur et à mesure que les futurs tokens sont traités, leurs passages avant atteignent les couches différées. À ces profondeurs, les états cachés des tokens différés sont concaténés avec l'état du token actuel et traités en parallèle.
3. Résultat : Cela comble les états KV manquants pour tous les tokens différés en un seul passage parallèle, garantissant que le cache KV est complet, exact et cohérent avec le token échantillonné. Aucun calcul n'est sauté ; il est simplement réorganisé d'une exécution séquentielle vers une exécution par lots parallèle.

3. Contributions Clés

1. Échantillonnage de Mélange Exact : N-vium est la première méthode permettant un routage adaptatif à la profondeur avec un échantillonnage exact de la distribution entraînée, évitant les erreurs d'approximation inhérentes aux méthodes standards de sortie anticipée.
2. Récupération Complète du Cache KV : En différant et en regroupant les calculs des couches supérieures, N-vium récupère des caches KV complets et exacts, une fonctionnalité souvent perdue dans les autres méthodes de calcul adaptatif.
3. Cadre d'Entraînement : Les auteurs fournissent un pipeline d'entraînement pratique qui optimise directement l'objectif de mélange, incluant des stratégies pour orienter les modèles vers l'efficacité via une pénalité de calcul apprenable.
4. Évolutivité : La méthode est validée à travers des échelles allant jusqu'à 1,5 milliard de paramètres, démontrant que l'approche n'est pas limitée aux petits modèles.

4. Résultats Expérimentaux

Les auteurs ont pré-entraîné des modèles N-vium (spécifiquement "Quadrivium" avec 4 sorties) sur le jeu de données C4 en utilisant l'architecture LLaMA 2, les comparant à des bases denses appariées en paramètres et en données.

• Accélération vs Perplexité :
  • Le plus grand modèle (1,5 milliard de paramètres, 48 couches, largeur cachée 1536) a atteint une accélération de 57,9 % en temps réel par rapport à la base dense sans dégradation de perplexité (différence de 0 %).
  • Les modèles plus petits ont montré des accélérations constantes allant de 23 % à 41 % avec de légères améliorations de perplexité (par exemple, un modèle de 36 couches a montré une accélération de 36,4 % avec une amélioration de perplexité de 3,7 %).
• Frontière de Pareto : Dans les graphiques qualité-vitesse, N-vium occupe le quadrant "haut-droit" (plus rapide et qualité égale ou meilleure), tandis que des méthodes comme CALM (réimplémenté sous CALM*) tombent dans le "bas-droit" (plus rapide mais qualité pire) et LayerSkip se situe sur une ligne de perplexité constante en dessous de la base.
• Performance en Aval : Après un ajustement fin supervisé (SFT) sur le mélange Tulu 3, les modèles N-vium ont égalé ou dépassé les bases denses sur les benchmarks en aval (MMLU, ARC, HellaSwag, etc.), confirmant que les modèles restent orientables et ne surajustent pas à l'objectif d'entropie croisée au détriment des performances de tâche.
• Robustesse Matérielle : Les mesures d'accélération sur les GPU A100, RTX A6000 et RTX 3090 ont donné des résultats cohérents (~57 %), tandis que les GPU plus anciens de l'architecture Turing ont montré des gains inférieurs en raison de l'absence de support natif du bfloat16.

5. Signification et Revendications

L'article revendique que N-vium représente un changement de paradigme dans l'efficacité de l'inférence. Contrairement à la plupart des méthodes qui améliorent les tokens par seconde en réduisant les FLOPs par token (souvent via approximation), N-vium améliore l'efficacité en augmentant les FLOPs par seconde. Il y parvient en réorganisant le calcul séquentiel en passes par lots parallèles, exploitant ainsi mieux le parallélisme matériel sans rejeter aucun calcul.

Les auteurs soulignent que N-vium est une généralisation stricte du transformer standard : lorsque les poids de routage pour les sorties intermédiaires sont nuls, le modèle se comporte exactement comme un transformer standard. La méthode satisfait simultanément quatre critères critiques, ce que les auteurs affirment qu'aucune autre méthode existante ne fait :

1. Inférence Exacte : Échantillonnage de la distribution entraînée sans approximation.
2. Récupération Complète du Cache KV : Aucune couche sautée ni remplissage de cache approximatif.
3. Routage Appris : L'allocation du calcul par token est effectuée par un routeur entraîné, et non par des heuristiques.
4. Accélération au Moment du Décode : Réduction de la latence en temps réel lors de la génération autorégressive.

L'article conclut que, bien que l'entraînement de N-vium entraîne un coût FLOP plus élevé par token (en raison de l'évaluation des têtes à plusieurs jonctions), cette surcharge est justifiée pour le déploiement en production où les coûts d'inférence dominent, et la surcharge relative diminue à mesure que la taille du modèle augmente.