A Training-Time Diagnostic for Generalization via the Log-Alignment Ratio

Titre : Un diagnostic en temps d'entraînement pour la généralisation via le rapport d'alignement logarithmique

Énoncé du problème
À mesure que les modèles de réseaux de neurones et les ensembles de données s'agrandissent, le coût computationnel des entraînements augmente, rendant l'étude des dynamiques d'entraînement et la conception d'outils de diagnostic peu coûteux de plus en plus précieux. Les praticiens ont besoin de métriques capables de signaler des problèmes tels que le surapprentissage, l'instabilité ou la saturation en utilisant uniquement des quantités disponibles pendant l'entraînement, sans recourir à des données de validation retenues coûteuses. Bien que des travaux antérieurs aient exploré les propriétés spectrales des matrices de poids (par exemple, rang effectif, rang stable) ou les théories de paramétrisation (par exemple, paramétrisation de mise à jour maximale, $\mu$P) pour comprendre la généralisation, les métriques existantes se concentrent souvent exclusivement soit sur les poids, soit sur les activations, ou nécessitent des opérations computationnellement coûteuses comme la décomposition en valeurs singulières (SVD) complète ou l'analyse en composantes principales (PCA).

Méthodologie
Les auteurs étudient le rapport d'alignement logarithmique (LAR), une métrique introduite initialement dans la théorie de la paramétrisation par Everett et al. [2024] pour mesurer l'alignement entre la matrice de paramètres d'une couche et ses activations d'entrée. L'article reformule théoriquement le LAR pour révéler un second sens, jusque-là inexploré : il mesure le chevauchement entre un spectre de poids ($p$) et un spectre d'activation ($q$).

• Définitions :
  • Soit $W$ une matrice avec des valeurs singulières $s_i$. Le spectre de poids $p$ est défini comme les valeurs singulières au carré normalisées : $p_i = s_i^2 / \sum s_j^2$.
  • Soit $X$ un lot d'entrées. Le spectre d'activation $q$ est défini comme les projections au carré normalisées des entrées sur les directions singulières de $W$ : $q_i = \sum_{x \in B} (v_i^\top x)^2 / \sum_{x \in B} \|x\|^2$.
• Reformulation : Les auteurs démontrent que le LAR peut s'exprimer comme suit :
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( \sum_{i} p_i q_i \right)$$
  Cette formulation montre que le LAR est une mesure du chevauchement logarithmique entre les distributions $p$ et $q$. Il peut également s'écrire en termes de covariance et de corrélation de Pearson :
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \cdot \text{cov}(p, q) + \frac{1}{n} \right)$$
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \sigma_p \sigma_q \text{corr}(p, q) + \frac{1}{n} \right)$$
  où $\sigma_p$ et $\sigma_q$ représentent la « dispersion » ou la concentration des distributions.
• Efficacité computationnelle : Le LAR est calculable à partir de quantités disponibles lors du passage avant (normes de matrices) avec une surcharge négligeable ($O(|B|m + |B|n + mn)$), ne nécessitant ni données retenues ni SVD complète.

Contributions clés

1. Reformulation distributionnelle : L'article démontre que le LAR n'est pas simplement une mesure d'alignement, mais un quantificateur du chevauchement entre les spectres de poids et d'activation. Il fournit des formes statistiques équivalentes (covariance, corrélation) qui rendent le comportement de la métrique interprétable en termes de quantités élémentaires.
2. Prédiction de la dimension effective : Les auteurs montrent que lorsqu'une matrice de poids concentre l'énergie sur $k$ directions singulières et que les activations s'alignent sur celles-ci, le LAR est lié à la dimension effective $k$ par la relation $k \approx n^2(1 - \text{LAR})$. Cette formule est vérifiée empiriquement dans des expériences de « grokking ».
3. Diagnostic de surapprentissage à grande échelle : Dans l'entraînement préalable à grande échelle (modèles de 3 milliards de paramètres), l'article établit que le LAR de la couche de déembedding se stabilise pendant la généralisation mais décline brusquement à l'approche du surapprentissage. L'écart du LAR par rapport à une ligne de base non surapprenante suit l'écart de généralisation.
4. Proxy de la dispersion spectrale : L'étude révèle que l'optimisation pousse $p$ et $q$ à être fortement corrélés. Par conséquent, le LAR agit comme un proxy de la variance (dispersion) de ces distributions. Lors de l'entraînement préalable, la direction singulière supérieure domine, rendant la contribution du composant supérieur ($L_1$) un indicateur principal de la signature pré-surapprentissage.

Résultats expérimentaux
Les auteurs valident ces résultats dans deux contextes distincts :

1. Expériences de Grokking (tâches algorithmiques de petite taille) :

   • Entraînement de Transformers à 2 couches sur 12 tâches d'opérations binaires.
   • Observation : Pendant la phase de mémorisation, le LAR est plus faible (indiquant des $p$ et $q$ diffus). Lorsque le modèle passe au « grokking » (généralisation), le LAR augmente à mesure que $p$ et $q$ se concentrent sur un ensemble partagé de directions.
   • Dimensionnalité : La dimension prédite par le LAR, $k_{pred} = n^2(1 - \text{LAR})$, correspond étroitement à $k_{95}$ (le nombre de composantes principales expliquant 95 % de la variance) aux points de contrôle finaux, avec une faible erreur quadratique moyenne (RMSE).
   • Différenciation : Les modèles de grokking atteignent des valeurs finales de LAR plus élevées ($>0,65$) par rapport aux modèles non-grokking ($\le 0,65$).

2. Entraînement préalable à grande échelle (modèles de langage de 3 milliards de paramètres) :

   • Entraînement de modèles basés sur Gemma-2 (3 milliards de paramètres) utilisant les optimiseurs Adam et Muon sur des ensembles de données allant de 200 millions à 100 milliards de tokens.
   • Généralisation vs Surapprentissage : Dans les régimes non surapprenants (par exemple, 100 milliards de tokens), le LAR monte à un pic puis décline lentement, se stabilisant. Dans les régimes de surapprentissage (ensembles de données plus petits), le LAR décline brusquement à l'approche de l'étape de surapprentissage.
   • Écart de généralisation : La différence entre le LAR d'une exécution spécifique et celui de la ligne de base de 100 milliards de tokens (différentiel LAR) reflète étroitement l'écart de généralisation (perte de validation moins perte d'entraînement).
   • Mécanisme : Le déclin du LAR avant le surapprentissage est principalement piloté par le composant singulier supérieur ($L_1$). Avec Adam, cela est dû à la diffusion des activations ($q_1$) ; avec Muon, cela implique à la fois les poids ($p_1$) et les activations.

Importance et affirmations
L'article affirme que le LAR offre un diagnostic pratique, peu coûteux en calcul et sans données pour la généralisation.

• Efficacité : Contrairement au rang effectif ou à d'autres mesures spectrales nécessitant une SVD, le LAR peut être calculé avec une surcharge négligeable lors du passage avant.
• Puissance prédictive : Il suit avec succès la transition de la mémorisation à la généralisation dans le grokking et sert de signal d'alerte précoce pour le surapprentissage dans l'entraînement préalable à grande échelle, permettant potentiellement des méthodes d'entraînement sans évaluation.
• Interprétabilité : En reliant le LAR à la dispersion des spectres de poids et d'activation, la métrique offre une vue unifiée de la manière dont les dynamiques d'optimisation (concentration vs diffusion des spectres) se rapportent à la généralisation du modèle.

Les auteurs reconnaissent des limites, notant que leurs preuves sont corrélationnelles, qu'ils se concentrent principalement sur la couche de déembedding, et qu'ils n'ont pas encore testé si la régularisation directe de la dispersion spectrale empêcherait le surapprentissage. Ils suggèrent que des travaux futurs devraient explorer le LAR pour d'autres types de matrices et son comportement sous des tailles de lots variables.