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

Titel: Eine Diagnose zur Generalisierung während des Trainings mittels des Log-Alignment-Verhältnisses

Problemstellung
Da sich neuronale Netzwerkmodelle und Datensätze skalieren, steigen die Rechenkosten von Trainingsläufen an, was die Untersuchung von Trainingsdynamiken und die Entwicklung kostengünstiger Diagnosewerkzeuge zunehmend wertvoll macht. Praktiker benötigen Metriken, die Probleme wie Overfitting, Instabilität oder Sättigung allein anhand von Größen während des Trainings erkennen können, ohne auf teure zurückgehaltene Validierungsdaten angewiesen zu sein. Während frühere Arbeiten spektrale Eigenschaften von Gewichtsmatrizen (z. B. effektiver Rang, stabiler Rang) oder Parametrisierungstheorien (z. B. maximal update Parametrisierung, $\mu$P) untersucht haben, um Generalisierung zu verstehen, konzentrieren sich bestehende Metriken oft ausschließlich entweder auf Gewichte oder Aktivierungen oder erfordern rechenintensive Operationen wie eine vollständige Singulärwertzerlegung (SVD) oder Hauptkomponentenanalyse (PCA).

Methodik
Die Autoren untersuchen das Log-Alignment-Verhältnis (LAR), eine Metrik, die ursprünglich in der Parametrisierungstheorie von Everett et al. [2024] eingeführt wurde, um die Ausrichtung zwischen der Parametrisierungsmatrix einer Schicht und ihren Eingabeaktivierungen zu messen. Das Papier formuliert LAR theoretisch um, um eine zweite, bisher unerforschte Bedeutung aufzudecken: Es misst die Überlappung zwischen einem Gewichtsspektrum ($p$) und einem Aktivierungsspektrum ($q$).

• Definitionen:
  • Sei $W$ eine Matrix mit Singulärwerten $s_i$. Das Gewichtsspektrum $p$ ist definiert als die normalisierten quadrierten Singulärwerte: $p_i = s_i^2 / \sum s_j^2$.
  • Sei $X$ ein Batch von Eingaben. Das Aktivierungsspektrum $q$ ist definiert als die normalisierten quadrierten Projektionen der Eingaben auf die Singulärrichtungen von $W$: $q_i = \sum_{x \in B} (v_i^\top x)^2 / \sum_{x \in B} \|x\|^2$.
• Umformulierung: Die Autoren leiten her, dass LAR ausgedrückt werden kann als:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( \sum_{i} p_i q_i \right)$$
  Diese Formulierung zeigt, dass LAR ein Maß für die Log-Überlappung zwischen den Verteilungen $p$ und $q$ ist. Es kann auch in Bezug auf Kovarianz und Pearson-Korrelation geschrieben werden:
  $$\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)$$
  wobei $\sigma_p$ und $\sigma_q$ die „Streuung" oder Konzentration der Verteilungen darstellen.
• Rechnerische Effizienz: LAR ist aus Größen berechenbar, die während des Vorwärtsdurchlaufs verfügbar sind (Matrixnormen), mit vernachlässigbarem Overhead ($O(|B|m + |B|n + mn)$), erfordert keine zurückgehaltenen Daten oder eine vollständige SVD.

Hauptbeiträge

1. Verteilungsmäßige Umformulierung: Das Papier zeigt, dass LAR nicht lediglich ein Maß für die Ausrichtung ist, sondern ein Quantifizierer der Überlappung zwischen Gewichts- und Aktivierungsspektren. Es liefert äquivalente statistische Formen (Kovarianz, Korrelation), die das Verhalten der Metrik in Bezug auf elementare Größen interpretierbar machen.
2. Vorhersage der effektiven Dimension: Die Autoren zeigen, dass, wenn eine Gewichtsmatrix Energie auf $k$ Singulärrichtungen konzentriert und Aktivierungen mit diesen ausgerichtet sind, LAR über $k \approx n^2(1 - \text{LAR})$ mit der effektiven Dimension $k$ zusammenhängt. Diese Formel wird in Grokking-Experimenten empirisch verifiziert.
3. Diagnose für Overfitting im großen Maßstab: Bei der großskaligen Vortrainierung (3B-Parameter-Modelle) stellt das Papier fest, dass das unembedding LAR während der Generalisierung stabilisiert, aber stark abfällt, wenn Overfitting einsetzt. Die Abweichung von LAR von einer Nicht-Overfitting-Baseline verfolgt die Generalisierungslücke.
4. Proxy für spektrale Streuung: Die Studie findet heraus, dass die Optimierung $p$ und $q$ dazu bringt, hochkorreliert zu sein. Folglich wirkt LAR als Proxy für die Varianz (Streuung) dieser Verteilungen. Beim Vortrainieren dominiert die oberste Singulärrichtung, wodurch der Beitrag der obersten Komponente ($L_1$) ein primärer Indikator für das Pre-Overfit-Signal wird.

Experimentelle Ergebnisse
Die Autoren validieren diese Erkenntnisse in zwei unterschiedlichen Settings:

1. Grokking-Experimente (kleine algorithmische Aufgaben):

   • Training von 2-Schicht-Transformern auf 12 binären Operationsaufgaben.
   • Beobachtung: Während der Memorierungsphase ist LAR niedriger (was eine diffuse $p$ und $q$ anzeigt). Wenn das Modell zu „Grokking" (Generalisierung) übergeht, steigt LAR an, da sich $p$ und $q$ auf einen gemeinsamen Satz von Richtungen konzentrieren.
   • Dimensionalität: Die von LAR vorhergesagte Dimension $k_{pred} = n^2(1 - \text{LAR})$ stimmt bei finalen Checkpoints eng mit $k_{95}$ (der Anzahl der Hauptkomponenten, die 95 % der Varianz erklären) überein, mit einem niedrigen Root Mean Squared Error (RMSE).
   • Differenzierung: Grokking-Modelle erreichen höhere finale LAR-Werte ($>0.65$) im Vergleich zu Nicht-Grokking-Modellen ($\le 0.65$).

2. Großskalige Vortrainierung (3B-Sprachmodelle):

   • Training von auf Gemma-2 basierenden Modellen (3B Parameter) unter Verwendung von Adam- und Muon-Optimierern auf Datensätzen im Bereich von 200M bis 100B Tokens.
   • Generalisierung vs. Overfitting: In Nicht-Overfitting-Regimen (z. B. 100B Tokens) steigt LAR auf einen Peak und fällt dann langsam ab, um sich zu stabilisieren. In Overfitting-Regimen (kleinere Datensätze) fällt LAR stark ab, wenn sich der Overfitting-Schritt nähert.
   • Generalisierungslücke: Der Unterschied zwischen dem LAR eines spezifischen Laufs und dem Baseline-100B-Lauf (LAR-Differenz) spiegelt die Generalisierungslücke (Validierungsverlust minus Trainingsverlust) eng wider.
   • Mechanismus: Der Rückgang von LAR vor dem Overfitting wird primär durch die oberste Singulärkomponente ($L_1$) getrieben. Bei Adam wird dies durch die Diffusion der Aktivierungen ($q_1$) getrieben; bei Muon sind sowohl Gewichte ($p_1$) als auch Aktivierungen beteiligt.

Bedeutung und Behauptungen
Das Papier behauptet, dass LAR eine praktische, rechnerisch günstige und datenfreie Diagnose für Generalisierung bietet.

• Effizienz: Im Gegensatz zum effektiven Rang oder anderen spektralen Maßen, die eine SVD erfordern, kann LAR während des Vorwärtsdurchlaufs mit vernachlässigbarem Overhead berechnet werden.
• Vorhersagekraft: Es verfolgt erfolgreich den Übergang von Memorierung zu Generalisierung beim Grokking und dient als Frühwarnsignal für Overfitting bei großskaliger Vortrainierung, was potenziell trainingsmethoden ohne Evaluation ermöglicht.
• Interpretierbarkeit: Durch die Verknüpfung von LAR mit der Streuung von Gewichts- und Aktivierungsspektren bietet die Metrik eine einheitliche Sichtweise darauf, wie Optimierungsdynamiken (Konzentration vs. Diffusion von Spektren) mit der Modellgeneralisierung zusammenhängen.

Die Autoren erkennen Einschränkungen an und stellen fest, dass ihre Beweise korrelativ sind, sie sich primär auf die Unembedding-Schicht konzentrieren und sie noch nicht untersucht haben, ob eine direkte Regularisierung der spektralen Streuung Overfitting verhindern würde. Sie schlagen vor, dass zukünftige Arbeiten LAR für andere Matrixtypen und sein Verhalten unter variierenden Batch-Größen untersuchen sollten.