Technische Samenvatting: Hoe Grokking te Temmen: Representatiegeometrie als een Controlesignaal

Probleemstelling

Grokking is een fenomeen in overgeparameteriseerde neurale netwerken waarbij modellen aanvankelijk trainingsdata memoriseren en pas na langdurige optimalisatie sterke generalisatie vertonen. Hoewel recente studies de link hebben gelegd tussen grokking en weight decay, feature compressie, circuitvorming en representatieleer-dynamiek, blijft de specifieke factor die de timing en het ontstaan van deze vertraagde generalisatie bepaalt, onvolledig begrepen. Een veelvoorkomend thema in de bestaande literatuur is dat geleerde representaties progressief comprimeren naar lagerdimensionale structuren. Echter, voorgaand werk behandelt de representatiegeometrie grotendeids als een emergente eigenschap in plaats van een direct controlesmechanisme. Dit artikel onderzoekt of het expliciet manipuleren van de geometrie van verborgen representaties kan dienen als een controlesignaal om het begin en de snelheid van grokking te beïnvloeden.

Methodologie

De auteurs introduceren Geometric Dimensionality Regularization (GeomDR), een regularisator op representatieniveau die de spectrale eigenschappen gebruikt om direct de effectieve dimensionaliteit van verborgen representaties tijdens de training te controleren.

Geometrische Karakterisering

Voor een verborgen laag representatiematrix $Z \in \mathbb{R}^{N \times d}$, centreert de methode eerst de data om $\tilde{Z}$ te verkrijgen. Vervolgens wordt de empirische covariantie-matrix $C = \frac{1}{N-1}\tilde{Z}^\top\tilde{Z}$ berekend. De effectieve dimensionaliteit ($D_{eff}$) wordt berekend met behulp van de participatieratio van de eigenwaarden ($\mu_i$) van $C$:
$$D_{eff} = \frac{(\sum_i \mu_i)^2}{\sum_i \mu_i^2}$$
Daarnaast wordt een lokale nabijheidsafstand ($\rho$) gedefinieerd als de gemiddelde Euclidische afstand tot de $k$ dichtstbijzijnde buren van genormaliseerde representaties.

De GeomDR-doelstelling

GeomDR legt een straf op aan de variantie die aanwezig is in richtingen buiten een doelsubruimte van dimensie $d^*$. Gegeven de eigenwaarden gesorteerd in aflopende volgorde ($\mu_1 \ge \mu_2 \ge \dots \ge \mu_d$), is de laag-specifieke loss:
$$L_{GeomDR}^{(\ell)} = \sum_{j=d^*+1}^{d} \mu_j^{(\ell)}$$
De totale trainingsdoelstelling combineert de taakverlies ($L_{task}$) en de geometrische regularisator:
$$L = L_{task} + \lambda(t)L_{GeomDR}$$
Cruciaal is dat de regularisator niet vanaf het begin actief is. Deze wordt geïntroduceerd na een initiële ongecontroleerde trainingsfase (activatiestap $t_s$) en de sterkte $\lambda(t)$ wordt geleidelijk verhoogd met behulp van een cosinus-ramp. Dit ontwerp bootst het natuurlijke grokking-traject na waarbij memorisatie voorafgaat aan generalisatie.

Experimenteel Protocol

Het raamwerk werd geëvalueerd op drie algoritmische taken die bekend staan om het vertonen van grokking:

1. Modulaire Additie: Het voorspellen van $(a + b) \mod 97$.
2. Modulaire Divisie: Het voorspellen van $(a \cdot b^{-1}) \mod 97$.
3. Permutatie Compositie: Het voorspellen van de compositie van twee permutaties in $S_5$.

Experimenten maakten gebruik van twee architecturen: een Multilayer Perceptron (MLP) en een Transformer. De studie varieerde systematisch de interventieschema's (sterkte $\lambda_{max}$ en ramp-duur $T_{ramp}$), de doel-dimensionaliteiten ($d^*$) en de interventie-starttijden ($t_s$).

Belangrijkste Bijdragen

1. Observatie van Voorafgaande Collaps: De auteurs demonstreren dat dimensionaliteitscollaps consequent voorafgaat aan de transitie van memorisatie naar generalisatie (grokking) in baseline-instellingen. Dit suggereert dat geometrische compressie niet louter een correlatie is, maar een mechanistische drijfveer van de transitie.
2. Introductie van GeomDR: Zij stellen GeomDR voor, een spectrale regularisator die direct de covariantie-spectra buiten een doelsubruimte onderdrukt, waardoor de effectieve dimensionaliteit van verborgen representaties wordt gecontroleerd.
3. Systematische Controle van Dynamiek: Het artikel biedt een uitgebreide empirische studie die aantoont dat geometrische interventies de vertraagde generalisatiedynamiek aanzienlijk kunnen veranderen. Afhankelijk van het interventieschema en de doel-dimensionaliteit kan GeomDR grokking versnellen, vertragen of volledig onderdrukken.

Resultaten

Versnelling van Grokking

In alle geëvalueerde settings veranderde GeomDR de grokking-dynamiek consistent:

• MLP-prestaties: Bij modulaire additie versnelde GeomDR grokking met factor wel 52 (reductie van 362K naar ~7K stappen). Vergelijkbare verbeteringen werden waargenomen voor modulaire divisie (39x) en permutatie compositie (~23x). In alle MLP-gevallen bereikte GeomDR een succespercentage van 100%, waar de baseline instabiel was of faalde.
• Transformer-prestaties: Hoewel de baseline grokking-dynamiek voor Transformers al sneller was dan die van MLPs, zorgde GeomDR nog steeds voor consistente versnelling (bijv. ~2x op modulaire additie), waarbij een succespercentage van 100% werd behouden.

Invloed van Hyperparameters

• Doel-dimensionaliteit ($d^*$): Matige tot grote doel-dimensionaliteiten ($16 \le d^* \le 64$) leverden de sterkste versnelling op. Agressieve compressie ($d^* \le 8$) vertraagde het leren vaak of maakte het instabiel.
• Interventie-timing ($t_s$): Er werd een "kritieke tijdsvenster" geïdentificeerd. Interventies die te vroeg werden toegepast (voordat nuttige representaties gevormd zijn) of te laat (nadat memorisatie rigide is geworden) waren minder effectief. Gemiddelde activatietijden (bijv. $t_s = 2000$ voor MLPs) produceerden de snelste en meest stabiele grokking.
• Schema: Gemiddelde regularisatiestrektes en ramp-duur presteerden over het algemeen het best; te sterke interventies konden de training destabiliseren.

Geometrische Correlaties

De studie bevestigde dat interventies die grokking versnelden, ook de onset van dimensionaliteitscollaps versnelden. In GeomDR-modellen werd de temporele kloof tussen dimensionaliteitscollaps en succesvolle generalisatie aanzienlijk kleiner vergeleken met baselines, wat wijst op een nauwere koppeling tussen geometrische reorganisatie en de opkomst van generalisatie.

Schaalbaarheid en Efficiëntie

De auteurs demonstreerden dat GeomDR effectief blijft, zelfs wanneer het covariantie-spectrum wordt geschat op basis van een willekeurige subset van representaties (bijv. 5%–25% van de data), wat een 2,4x–2,7x versnelling in runtime oplevert zonder de grokking-versnelling materieel te beïnvloeden.

Betekenis en Claims

Het artikel claimt dat representatiegeometrie dient als een effectief controlesignaal voor grokking. Door de covariantie-spectra van verborgen activaties direct te modificeren, kunnen onderzoekers systematisch beïnvloeden wanneer en hoe snel een netwerk overgaat van memorisatie naar generalisatie.

De auteurs stellen dat dit werk bewijs levert dat geometrische interventies een praktische aanpak bieden voor het bestuderen en beïnvloeden van vertraagde generalisatie. In plaats van representatieveranderingen passief te observeren, stelt de methode de directe manipulatie van geometrische eigenschappen mogelijk om hypothesen over leer-dynamiek te testen. De resultaten suggereren dat de relatie tussen representatiegeometrie en grokking robuust is over verschillende taken en architecturen, hoewel de omvang van de verbetering varieert (het is prominenter bij MLPs dan bij Transformers, waarschijnlijk door de sterkere inductieve biases van de laatste).

Het artikel concludeert dat, hoewel het niet vaststelt dat representatie-dimensionaliteit de enige causale mechanisme is, het aantoont dat geometrische interventies voldoende zijn om de timing van grokking te veranderen, wat een nieuwe experimentele methodologie biedt voor optimalisatie- en generalisatieonderzoek.