Résumé technique : Leviathan – Découplage des représentations d'entrée et de sortie dans les modèles de langage

1. Énoncé du problème

Les modèles de langage modernes de grande taille (LLM) utilisent généralement une architecture à poids liés (tied-weight), où une unique matrice d'embedding $E \in \mathbb{R}^{V \times D}$ remplit deux rôles : mapper les tokens d'entrée vers des représentations continues et projeter les états cachés vers l'espace du vocabulaire pour la prédiction de sortie. Ce couplage force une seule matrice à satisfaire simultanément deux objectifs distincts :

1. Représentation d'entrée : Nécessite de la régularité (smoothness), où les tokens sémantiquement apparentés occupent des régions voisines dans l'espace de représentation.
2. Classification de sortie : Nécessite de la discrimination, séparant efficacement tous les tokens à travers le vocabulaire.

Les auteurs émettent l'hypothèse que ce couplage impose une limitation architecturale. Bien que les modèles entièrement découplés (apprenant des matrices d'entrée et de sortie séparées) théoriquement découplent ces objectifs, ils s'avèrent inefficaces en pratique. Les expériences montrent que les modèles entièrement découplés sous-performent les bases à poids liés malgré une augmentation d'environ 50 % du nombre de paramètres, suggérant que les paramètres supplémentaires ne sont pas utilisés efficacement en raison de difficultés d'optimisation et d'un signal d'entraînement insuffisant pour les embeddings discrets, en particulier pour les tokens rares.

2. Méthodologie : L'architecture Leviathan

Leviathan introduit une architecture Transformer qui découple la représentation d'entrée de la prédiction de sortie sans encourir le coût prohibitif en paramètres des embeddings entièrement découplés.

2.1 Vectorisation des embeddings appris (LEV)

L'innovation centrale est la couche LEV, qui remplace la matrice d'embedding d'entrée dense par une fonction continue apprise $G: \{0, \dots, V-1\} \to \mathbb{R}^D$. Cette fonction mappe les indices de tokens vers des embeddings via un mécanisme compact de partage de paramètres :

• Indexation compositionnelle : Le vocabulaire est factorisé en $k$ composantes en utilisant une décomposition en base $b$, convertissant un index de token $i$ en un tuple $(i_1, \dots, i_k)$.
• Cahiers de codes partagés : Ces indices mappent vers des cahiers de codes partagés $C_1, \dots, C_k$, produisant une représentation initiale $z(i)$. Cela permet aux tokens rares de bénéficier des signaux de gradient transmis via les entrées partagées utilisées par des tokens structurellement similaires.
• Projection continue : L'initialisation est projetée et normalisée vers une coordonnée latente bornée $\tilde{z}_\ell \in [0, 1]^{d_{seed}}$.
• Expansion par splines : Chaque dimension est modélisée par une expansion de base B-spline quadratique univariée avec $\kappa nœuds$.
• Produit tensoriel : Les dimensions sont combinées via un produit tensoriel de rang $r$.
• Somme des têtes : L'embedding final est la somme des sorties de $h$ têtes parallèles après une projection linéaire.

2.2 Architecture et découplage

• Tête de sortie découplée : La couche LEV permet à la tête de projection de sortie de rester entièrement découplée (une matrice dense indépendante de l'entrée).
• Efficacité paramétrique : Puisque la couche LEV est significativement plus petite qu'une table d'embedding dense, la surcharge de la tête de sortie découplée est négligeable. À l'échelle de 1,2 milliard de paramètres, l'augmentation totale des paramètres n'est que de 0,2 %.
• Biais inductif : L'architecture impose un biais inductif de régularité sur l'espace des tokens, contrastant avec la géométrie discrète, potentiellement anisotrope, des tables de recherche standards.

3. Configuration expérimentale

• Modèles : Les expériences ont été menées sur des Transformers uniquement décodeurs à trois échelles : 200 M, 400 M et 1,2 M de paramètres.
• Bases de comparaison :
  • LIÉ (TIED) : Base standard à poids liés.
  • DÉCOUPLÉ (UNTIED) : Modèle entièrement découplé (611 M de paramètres à l'échelle de 400 M).
  • ALBERT : Factorisation appariée en paramètres d'entrée (compression linéaire) avec une tête de sortie découplée.
  • Ablations : MLP plat et MLP factorisé appariés en paramètres pour isoler les contributions de la non-linéarité et de la structure compositionnelle.
• Entraînement : Tous les modèles ont été entraînés sur The Pile (sans droits d'auteur) avec des hyperparamètres, des graines et des budgets de tokens identiques (jusqu'à 24,1 milliards de tokens pour les modèles de 1,2 M).
• Évaluation : Performance de modélisation du langage sur un ensemble de validation retenu de The Pile, généralisation en zéro-shot sur Wikitext-103, et benchmarks en aval (HellaSwag, ARC, PIQA, WinoGrande, LAMBADA).

4. Résultats clés

4.1 Performance de modélisation du langage

Leviathan surpasse constamment la base à poids liés à toutes les échelles, l'écart de performance s'élargissant à mesure que l'entraînement progresse et que l'échelle du modèle augmente.

• Échelle 1,2 M : Leviathan a réduit la perplexité de validation de 9,2 % (0,087 nats) par rapport à la base liée.
• Efficacité de l'échantillon : À l'échelle de 1,2 M, Leviathan a nécessité 2,1 fois moins de tokens d'entraînement pour atteindre la perte finale de la base liée.
• Dynamique d'entraînement : À plus petite échelle (200 M), Leviathan a initialement sous-performé avant de dépasser la base ; à 1,2 M, il a surpassé la base tout au long de l'entraînement.

4.2 Comparaison avec les alternatives

• Vs. Entièrement découplé : Le modèle entièrement découplé (UNTIED) a sous-performé la base liée tout au long de l'entraînement, confirmant qu'un découplage naïf sans paramétrisation efficace est sous-optimal.
• Vs. Factorisation linéaire : La base factorisée de style ALBERT a également sous-performé la base liée, indiquant que la compression linéaire seule ne peut pas récupérer les avantages représentationnels de la géométrie liée.
• Vs. Leviathan : Leviathan a surpassé à la fois UNTIED et ALBERT, démontrant que la génération continue non linéaire est supérieure à la fois au découplage naïf et à la factorisation linéaire dans les budgets de calcul optimaux selon Chinchilla.

4.3 Transfert en aval et généralisation

• Tâches en aval : À l'échelle de 1,2 M, Leviathan a amélioré les performances sur les six benchmarks en aval évalués. Notamment, il a atteint une réduction de 30 % de la perplexité sur LAMBADA.
• Transfert en zéro-shot : Sur Wikitext-103 (une distribution différente), Leviathan a montré des améliorations de perplexité croissantes avec l'échelle (2,4 % à 200 M, 12,3 % à 1,2 M), suggérant que ses représentations continues se transfèrent plus efficacement à travers les distributions de données.

4.4 Analyse stratifiée par fréquence

Les gains de performance sont fortement concentrés dans la longue traîne de la distribution des tokens :

• Tokens rares : Pour les tokens apparaissant moins d'une fois sur 10 millions, Leviathan a réduit la perplexité de 81 % à l'échelle de 1,2 M.
• Tokens fréquents : Pour les tokens apparaissant plus d'une fois sur 10 000, l'amélioration a chuté à près de zéro (et légèrement négative pour les plus fréquents), car les embeddings discrets sont suffisants lorsque les signaux d'entraînement sont forts.
• Mécanisme : La paramétrisation continue partage les gradients entre des tokens structurellement similaires, fournissant un signal robuste là où les recherches discrètes échouent en raison de la rareté des données.

4.5 Robustesse

• Taille du vocabulaire : Les résultats sont restés cohérents à travers des tailles de vocabulaire de 50k, 100k et 200k, avec des améliorations relatives d'environ 2 % partout.
• Ablations : Les gains ont été décomposés en : représentation continue (+1,13 %), structure compositionnelle (+0,37 %) et régularité séparable par spline (+0,62 %).

5. Signification et affirmations

L'article affirme que la pratique standard du couplage des poids impose une limitation architecturale qui n'est pas résolue par un simple découplage des poids. Au contraire, la solution optimale implique de découpler les représentations d'entrée et de sortie tout en imposant une régularité via une paramétrisation continue.

• Insight architectural : La géométrie représentationnelle des embeddings d'entrée compte. La paramétrisation continue est un moyen efficace en données et en calcul d'améliorer cette géométrie, en particulier pour les tokens rares où les embeddings discrets peinent.
• Efficacité : Leviathan atteint ces gains avec une surcharge paramétrique négligeable (0,2 % à l'échelle de 1,2 M), ce qui en fait une mise à niveau pratique pour les modèles de langage modernes.
• Extensibilité : Les avantages de la couche LEV s'accumulent avec l'échelle et les données, suggérant que à mesure que les modèles grandissent, l'avantage des représentations d'entrée continues et régulières sur les recherches discrètes devient plus prononcé.

Les auteurs concluent que, bien que le motif d'amélioration dépendant de la fréquence soit corrélational, la cohérence à travers les échelles, les vocabulaires et les variantes architecturales fournit des preuves robustes que la paramétrisation continue atténue efficacement les difficultés d'optimisation associées à l'apprentissage d'embeddings discrets pour les tokens rares.