Beyond Perplexity: A Geometric and Spectral Study of Low-Rank Pre-Training

Résumé technique : Au-delà de la perplexité : une étude géométrique et spectrale du pré-entraînement de faible rang

Énoncé du problème
Le pré-entraînement des grands modèles de langage (LLM) est contraint par les coûts mémoire associés au stockage des poids, des gradients et des états de l'optimiseur de rang complet. Le pré-entraînement de faible rang est apparu comme une solution pour réduire ces coûts, les méthodes évoluant vers trois catégories : basées sur l'optimiseur (par exemple, GaLore, Fira), basées sur l'architecture (par exemple, CoLA, SLTrain) et basées sur les adaptateurs (par exemple, ReLoRA). Cependant, l'évaluation de ces méthodes reste limitée. Les pratiques actuelles reposent presque exclusivement sur la perplexité de validation issue de runs à graine unique. Les auteurs soutiennent que la perplexité est un mauvais proxy pour la qualité de la solution ; des modèles ayant une perplexité identique peuvent converger vers des régions fondamentalement différentes du paysage de perte, dotées de représentations internes distinctes. De plus, les comparaisons existantes échouent souvent à distinguer la capacité d'une méthode à affiner un modèle pré-entraîné de la capacité d'un modèle pré-entraîné avec cette méthode à se généraliser en aval.

Méthodologie
Les auteurs introduisent un cadre diagnostique pour caractériser les solutions trouvées par cinq méthodes de pré-entraînement de faible rang (GaLore, Fira, CoLA, SLTrain, ReLoRA) par rapport à un entraînement de rang complet, à travers trois échelles de modèles (60M, 130M, 350M). L'étude évalue ces méthodes en utilisant 16 métriques réparties sur quatre dimensions :

1. Paysage de perte 1-D (directions aléatoires) : Les auteurs perturbent les paramètres du modèle le long de 100 directions aléatoires pour mesurer la moyenne et la variance de la perte de validation. Ils dérivent des métriques scalaires pour la Rigueur (Sharpness, le taux de croissance de la perte s'éloignant du point de contrôle) et la Variance directionnelle (l'anisotropie de la surface de perte).
2. Paysage de perte 1-D (directions PCA Top-K) : Des perturbations sont appliquées le long des directions singulières principales des matrices de poids pour analyser la géométrie le long des axes les plus significatifs.
3. Interpolation 1-D : Les auteurs mesurent la perte le long de trajectoires linéaires entre les points de contrôle. Ils introduisent deux métriques spécifiques :
   • Hauteur de barrière entre points de contrôle consécutifs (CCBH) : Mesure la barrière entre les points de contrôle séquentiels d'une même méthode.
   • Hauteur de barrière inter-méthode (IMBH) : Une métrique novatrice quantifiant la séparation géométrique entre différentes méthodes d'entraînement au même stade d'entraînement. Cela révèle si différentes méthodes convergent vers le même bassin ou vers des solutions géométriquement distinctes.
4. Métriques spectrales et de rang : L'étude suit la structure spectrale des poids et des mises à jour, incluant le Rang effectif, le Rang stable, le Gap spectral et le Rang de seuil.
5. Similarité des activations : Les auteurs comparent les états cachés des modèles de faible rang avec les références de rang complet en utilisant la distance L2 des activations, la similarité cosinus et le CKA linéaire (Alignement du noyau centré).

Enfin, les auteurs testent si ces caractéristiques géométriques et spectrales peuvent prédire les performances en aval en mode zéro-shot mieux que la perplexité seule, en utilisant une régression linéaire avec validation croisée « leave-one-size-out » (LOSO) et « leave-one-method-out » (LOMO).

Contributions clés

1. Cadre diagnostique : L'article propose une suite d'évaluation complète incluant la métrique novatrice IMBH, qui quantifie la séparation géométrique entre les algorithmes d'entraînement plutôt que simplement entre les graines ou les époques.
2. Résultats empiriques inter-familles : L'étude établit que les méthodes de faible rang ne sont ni équivalentes à l'entraînement de rang complet, ni entre elles. Les résultats clés incluent :
   • Séparation géométrique : Les méthodes convergent vers des bassins géométriquement distincts, l'IMBH croissant de manière monotone à mesure que l'entraînement progresse.
   • Signatures spectrales : GaLore correspond étroitement à la structure spectrale des mises à jour de rang complet, tandis que CoLA ne le fait pas.
   • Divergence des représentations : Les activations de faible rang divergent de plus en plus de celles de rang complet dans les couches profondes. CoLA présente des activations orthogonales, tandis que Fira ne retrouve l'alignement que dans la dernière couche.
3. Prédiction améliorée des performances en aval : La perplexité de validation seule est un prédicteur incomplet des performances en aval. L'ajout de caractéristiques géométriques et spectrales à la perplexité améliore considérablement la précision de la prédiction (la corrélation de Pearson LOSO passe de 0,873 à 0,913).
4. Base de code unifiée : Les auteurs publient un pipeline open-source intégrant le pré-entraînement, le calcul des métriques et l'évaluation en aval pour faciliter les recherches futures.

Résultats

• Géométrie du paysage de perte : L'entraînement de rang complet s'installe dans des bassins plus rigides le long des directions aléatoires, mais dans des bassins plus plats le long de la direction PCA top-1. À l'inverse, les méthodes de faible rang comme Fira et SLTrain convergent vers des bassins plus plats le long des directions aléatoires, tandis que CoLA et ReLoRA atteignent les bassins les plus rigides. GaLore présente un profil mixte : modérément plat le long des directions aléatoires mais rigide le long de sa direction PCA top-1.
• Barrières inter-méthodes : Toutes les méthodes occupent des bassins géométriquement séparés. À l'échelle 60M, la barrière entre le rang complet et toute méthode de faible rang est plus élevée que les barrières entre les méthodes de faible rang. Fira converge vers le bassin géométriquement le plus proche du rang complet, tandis que SLTrain reste le plus isolé.
• Analyse spectrale : Le rang complet domine toutes les métriques de rang. GaLore suit de près le rang complet en termes de rang effectif et stable, mais présente un gap spectral plus faible. Fira et GaLore reproduisent le plus fidèlement la distribution spectrale des mises à jour de poids de rang complet.
• Alignement des activations : ReLoRA suit les représentations de rang complet de manière la plus étroite sur toutes les échelles. CoLA montre la plus grande déviation, les activations devenant orthogonales à la référence dans les couches profondes. Fira montre une dérive moyenne par couche élevée, mais une déviation étonnamment faible dans la dernière couche.
• Performances en aval : La perplexité de validation ne prédit pas de manière cohérente le classement des performances en aval. Par exemple, Fira atteint la perplexité la plus faible à toutes les échelles, pourtant CoLA mène dans plusieurs catégories en aval à 130M et 350M, bien qu'il ne se classe pas dans les deux premiers pour la perplexité.

Signification
L'article conclut que les méthodes de pré-entraînement de faible rang ne sont pas des approximations interchangeables de l'entraînement de rang complet ; elles modifient fondamentalement les solutions atteintes. L'étude démontre que la perplexité de validation offre une image incomplète de la qualité du modèle. En intégrant des métriques géométriques et spectrales, les chercheurs peuvent mieux prédire les performances en aval et comprendre les biais inductifs distincts introduits par différentes stratégies de faible rang. Le travail souligne que si certaines méthodes (comme Fira) convergent vers des bassins géométriquement proches du rang complet, d'autres (comme CoLA) occupent des vallées distinctes, suggérant que le choix de la méthode de faible rang dicte la région spécifique du paysage de perte dans laquelle le modèle s'installe.