Technische Samenvatting: Hyperparameter-overdracht voor Dichte Associatieve Geheugens

Probleemstelling

Moderne deep learning is sterk afhankelijk van hyperparameter-overdracht (HP-overdracht), waarbij optimale instellingen (bijv. schalen voor initialisatie, leersnelheden) worden geïdentificeerd op kleine modellen en datasets en vervolgens worden overgedragen naar grotere schalen om prohibitief dure optimalisatie te vermijden. Hoewel theoretische kaders voor HP-overdracht bestaan voor feed-forward netwerken (bijv. gebaseerd op de $\mu$P-heuristiek), zijn deze niet ontwikkeld voor Dichte Associatieve Geheugens (DenseAMs).

DenseAMs vormen unieke uitdagingen voor bestaande overdrachtsmethoden vanwege drie specifieke architecturale kenmerken:

1. Gewichtsdeling: De energie-functie vereist dat gewichten worden gedeeld over lagen en binnen operaties, in tegenstelling tot standaard feed-forward netwerken.
2. Residudynamiek: De dynamische vergelijkingen die energie-afstijging beschrijven, bevatten van nature residu-connecties.
3. Stepe Activaties: Om superieure schaalwetten voor informatiestorage te bereiken, maken DenseAMs gebruik van snel piekende activatiefuncties (bijv. high-power ReLUs, softmax) die zeldzaam zijn in standaard feed-forward architecturen.

Deze factoren creëren een complex landschap waar standaard dimensionale analyse faalt om te voorspellen hoe hyperparameters moeten schalen met modelgrootte ($N, K$), datasetgrootte ($P$) en batchgrootte ($B$).

Methodologie

De auteurs starten het onderzoek naar $\mu$P-type HP-overdracht voor DenseAMs door expliciete voorschriften af te leiden voor het schalen van hyperparameters in het proportionele schaalregime, gedefinieerd als:
$$N \to \infty, \quad K = \kappa N, \quad P = \rho N, \quad B = \beta P$$
waarbij $\kappa, \rho, \beta$ vaste constanten zijn.

Modeldefinitie

De studie richt zich op een ondiepe DenseAM-updateregel die een enkele stap van gradient descent uitvoert op een energie-landschap:
$$f(x) = s_2 W^\top \sigma_C [s_1 W g(x) + b] + c$$
waarbij $W \in \mathbb{R}^{K \times N}$ gekoppelde gewichten zijn, $g$ een vaste niet-lineariteit is, en $\sigma_C$ een mogelijk gecentreerde activatiefunctie is. Het trainingsdoel is een denoising-verlies:
$$L(W) = \frac{1}{2B} \sum_{\mu=1}^B \| f_W(x_\mu + \epsilon_\mu) - x_\mu \|^2$$

Belangrijkste Technische Aanpassingen

De auteurs identificeren dat standaard parametrisaties falen voor DenseAMs met niet-lineaire activaties vanwege "door het gemiddelde veroorzaakte pieken" in de Gram-matrix van activaties. Om dit aan te pakken, introduceren zij centrering:

• Centreringoperator: Een projectiematrix $C = I_K - \frac{1}{K}\mathbf{1}_K\mathbf{1}_K^\top$ wordt toegepast op pre- en post-activaties.
• Noodzaak: Voor niet-lineaire activaties (ReLU, Softmax) met een niet-nul gemiddelde leidt het nalaten van centrering tot een rang-een piek in de activatie-Gram-matrix ($S^\top S$) die anders schaalt dan de bulk-eigenwaarden. Deze mismatch veroorzaakt numerieke instabiliteit en verhindert HP-overdracht. Centrering verwijdert deze piek, waardoor ervoor wordt gezorgd dat alle interne statistieken consistent schalen.

Afleidingsbenadering

De auteurs passen de $\mu$P-heuristiek aan door een dimensionale analyse uit te voeren om drie desiderata te voldoen:

1. Stabiliteit: Initialisatiestatistieken (gemiddelde/variatie) van pre-activaties ($Z$) en outputs ($F$) moeten $O(1)$ blijven naarmate $N, K \to \infty$.
2. Maximaliteit: Parameterupdates ($\Delta Z, \Delta F$) na $T$ stappen moeten $O(1)$ blijven.
3. Balans: De verschillende termen die bijdragen aan updates (voortkomend uit de twee verschijningen van gekoppelde gewichten $W$) moeten in grootte gebalanceerd zijn.

Door de gradiënten van de MSE-loss en de resulterende updates voor zowel SGD- als Adam-optimizers te analyseren, leiden de auteurs specifieke schaalwetten af voor initialisatieschalen ($s_1, s_2$) en leersnelheden ($\eta_W, \eta_b, \eta_c$).

Belangrijkste Bijdragen en Resultaten

1. Expliciete Schaalvoorschriften

Het artikel biedt een tabel met schaalregels (Tabel 1) voor verschillende activatiefuncties en optimizers. Belangrijke bevindingen zijn:

• Lineaire Activaties: In het proportionele regime geldt $s_1 \sim N^{-1/2}$, $s_2 \sim K^{-1/2}$, en $\eta_W \sim K$ (voor SGD).
• ReLU-type Activaties: Centrering is kritiek. Zonder centrering faalt HP-overdracht door instabiliteit. Met centrering gelden de lineaire voorschriften voor SGD en Adam. Adam blijft stabiel zelfs voor high-power ReLUs ($p=10$), terwijl SGD instabiel wordt bij grote leersnelheden.
• Softmax-activaties:
  • SGD-instabiliteit: SGD lijdt aan numerieke instabiliteit bij grote schalen door "rijversterkingsfeedback", waarbij rijen met grotere softmax-gewichten onevenredig grote updates ontvangen, wat leidt tot lokalisatie en divergentie. Het artikel biedt geen stabiele SGD-voorschrift voor softmax.
  • Adam-succes: Adam is schaal-invariant en vertoont succesvolle HP-overdracht voor gecentreerde softmax-DenseAMs met $s_1 \sim N^{-1/2}$, $s_2 \sim K^{1/2}$, en $\eta_W \sim 1$.

2. Empirische Validatie

De auteurs valideren deze theoretische bevindingen in meerdere settings:

• Proportionele Schaling: Experimenten op isotrope Gaussische data en MNIST (via plaquette-downsampling) tonen aan dat trainingsdynamiek (verliescurves, interne feature-covarianties) over schalen instort wanneer de afgeleide voorschriften worden gebruikt.
• Anisotrope Data: De voorschriften werken verrassend goed zelfs voor anisotrope en realistische data (MNIST), ondanks dat de theoretische afleidingen isotrope inputs veronderstellen.
• Centrering vs. Niet-gecentreerd: Empirische resultaten (Figuur 7, 8, 10, 11, 13) tonen duidelijk aan dat niet-gecentreerde modellen divergeren of falen om hyperparameters over te dragen, terwijl gecentreerde modellen consistente dynamiek en succesvolle overdracht vertonen.

3. Dynamische Inzakking

Naast eenvoudige verlies-overdracht demonstreert het artikel dat de voorgestelde parametrisatie leidt tot een inzakking van trainingsdynamiek. Interne statistieken, zoals de RMS van gewichtsupdates en feature-covarianties, blijven constant over verschillende modelgroottes ($N, K$) wanneer de schaalwetten worden toegepast.

Betekenis en Claims

Het artikel claimt het eerste werk te zijn dat het onderzoek naar hyperparameter-overdracht voor Dichte Associatieve Geheugens start. De betekenis ligt in:

• Overbruggen van een Kloof: Het breidt de theorie van hyperparameter-overdracht uit van standaard feed-forward netwerken naar energie-gebaseerde modellen met gekoppelde gewichten en stepe activaties.
• Praktische Nut: Het biedt uitvoerbare, theoretisch onderbouwde regels voor het schalen van DenseAMs, waardoor practitioners in staat worden gesteld grootschalige modellen (bijv. Energy Transformers, NRGPT) te trainen zonder dure grootschalige hyperparameter-optimalisatie.
• Architecturale Inzicht: Het identificeert centrering als een fundamentele vereiste voor de stabiliteit van DenseAMs met niet-lineaire activaties, een nuance die eerder niet was aangepakt in de context van HP-overdracht.

De auteurs blijven bescheiden wat de reikwijdte betreft, en merken op dat hun analyse isotrope data veronderstelt voor afleidingen (hoewel empirisch robuust anderszins) en dat een volledige analytische beschrijving van trainingsdynamiek bij oneindige schaal voor niet-lineaire DenseAMs onderwerp blijft voor toekomstig werk. Zij claimen niet HP-overdracht op te lossen voor alle moderne architecturen, maar specifiek voor de DenseAM-familie.