Technische Samenvatting: Verbetert Weight Decay de Trainingsstabiliteit?

Probleemstelling

In modern deep learning wordt weight decay door praktici vaak toegeschreven aan het "stabiliseren" van trainingsdynamica, een rol die afwijkt van de klassieke interpretatie als een statische regularisatiestraf voor slecht gestelde inverse problemen. Hoewel recente literatuur de rol van weight decay bij feature selectie en het kiezen van oplossingen heeft vastgesteld, blijven de specifieke mechanismen waarmee het trainingsstabiliteit beïnvloedt onderbelicht.

Trainingsstabiliteit wordt over het algemeen begrepen via drie onderscheiden maar gerelateerde concepten:

1. Dynamica in de parameter-ruimte: Het beheersen van kromming van de loss en de Edge of Stability (EoS).
2. Dynamica in de functie-ruimte: Het begrenzen van het spectrum van de Neural Tangent Kernel (NTK).
3. Signaalpropagatie: Het beheersen van pre-activatie schalen (bijvoorbeeld via $\mu$P).

De centrale vraag die wordt behandeld, is of weight decay training werkelijk stabiliseert, en zo ja, via welk mechanisme. Specifiek onderzoekt het artikel hoe weight decay de EoS beïnvloedt, waarbij de maximale eigenwaarde van de loss Hessian ($\lambda_{\max}(\nabla^2 L)$) typisch rond de instabiliteitsdrempel $2/\eta$ blijft hangen (waarbij $\eta$ de leersnelheid is). Naïeve heuristieken suggereren dat het toevoegen van weight decay $\gamma$ deze drempel simpelweg zou moeten verschuiven naar $2/\eta - \gamma$. Empirische observaties wijzen echter op complexere, niet-lineaire gedragingen die dit eenvoudige additieve model uitdagen.

Methodologie

De auteurs maken gebruik van een combinatie van empirische analyse en wiskundige modellering om de effecten van weight decay op EoS-dynamica te onderzoeken.

Empirische Opzet

• Architecturen: Experimenten worden uitgevoerd op Multi-Layer Perceptrons (MLP's) met ReLU-activaties en Convolutional Neural Networks (CNN's).
• Datasets: Een subset van CIFAR-10 (5.000 samples).
• Optimalisatie: Volledige-batch Gradient Descent (GD) met Mean Squared Error (MSE) loss.
• Maten: De studie volgt de evolutie van loss-scherpte ($\lambda_{\max}$), de progressieve scherptecoëfficiënt ($\alpha = -\langle \nabla L, \nabla S \rangle$) en globale interactietermen ($c_y = \langle \nabla S, \theta \rangle$).
• Beperkingen: De analyse is beperkt tot volledige-batch GD en kleine schaalarchitecturen vanwege de rekenkosten van Hessian-vectorproducten die nodig zijn voor scherp-schatting.

Theoretisch Kader

Op basis van het zelf-stabiliseringskader van Damian et al. [23] modelleren de auteurs gradient descent met weight decay als een ondergedempt harmonische oscillator.

• Lokale Dynamica: Ze analyseren het systeem in een lokale basis opgespannen door de top-eigenvector van de Hessian ($u$) en de gradient van scherpte ($\nabla S$).
• Globale Interactie: Cruciaal omvat het model globale interactietermen ($c_x = \langle u, \theta^* \rangle$ en $c_y = \langle \nabla S, \theta^* \rangle$) die voortvloeien uit de weight decay-straf.
• Afleiding: De discrete update-regels voor de afwijking van scherpte ($y_t$) en de instabiele richting ($x_t$) worden afgeleid, waarbij wordt aangetoond hoe weight decay zowel lokale wrijving als globale offsets introduceert.

Belangrijkste Bijdragen en Resultaten

1. Weight Decay vertraagt Progressieve Scherpte

Zowel bij MLP's als bij CNN's vertraagt het verhogen van de weight decay-parameter $\gamma$ robuust de snelheid van progressieve scherpte. Dit vertraagt het begin van de Edge of Stability (EoS) en maakt de scherptetraject tijdens de vroege training vlakker.

2. Architectuur-afhankelijke Fase-overgangen

Zodra het netwerk de EoS-regime betreedt, induceert weight decay divergerend gedrag op basis van de architectuur:

• CNN's: Weight decay werkt voornamelijk als een demper. Het vermindert de amplitude van scherpte-oscillaties, waardoor het systeem sneller tot rust komt, maar de stabilisatiedrempel blijft dicht bij de naïeve verwachting van $2/\eta - \gamma$.
• MLP's: Er treedt een opvallende fase-overgang op. Wanneer $\gamma$ een kritieke drempel $\bar{\gamma}$ overschrijdt, daalt het stabilisatieniveau van scherpte significant onder de theoretische $2/\eta - \gamma$-grens. Bijvoorbeeld, met $\eta=0.02$ en $\gamma=0.02$, stabiliseert de scherpte rond 80, terwijl de naïeve drempel $\approx 99.98$ voorspelt. Ondanks deze lagere drempel behouden de trainingsdynamica EoS-kenmerken (niet-monotoon verlies, aanhoudende oscillaties).

3. Mechanisme: Globale Uitlijning en Ineenstorting van de Limietcyclus

Het artikel identificeert de globale uitlijning tussen de parametervector en de scherptegradient als de drijvende kracht achter de fase-overgang.

• Het systeem wordt gemodelleerd als een oscillator waarbij de amplitude van oscillatie $\Delta$ afhangt van de term $(\alpha + \gamma^2 - \gamma c_y)$.
• Een kritieke bifurcatie treedt op wanneer de globale interactieterm $c_y$ een drempel $c_y^{crit} = \frac{\alpha}{\gamma} + \gamma$ overschrijdt.
• In MLP's: Voldoende weight decay zorgt ervoor dat $c_y$ deze kritieke drempel overschrijdt tijdens progressieve scherpte, wat leidt tot een "instorting van de limietcyclus". Het systeem lost zich los van de standaard EoS-grens en stabiliseert op een veel lagere kromming.
• In CNN's: De globale interactieterm blijft onder de kritieke drempel, wat de fase-overgang verhindert en resulteert in slechts gedempte oscillaties.

4. Stabiliteit in Functie-ruimte (NTK)

Het artikel toont aan dat de verlaagde scherptedrempel in de parameter-ruimte vertaalt naar stabiliteit in de functie-ruimte.

• Door de EoS-drempel te verlagen, beperkt weight decay de spectrale straal van de empirische Neural Tangent Kernel (NTK).
• Met behulp van een rank-1 NTK-benadering tonen de auteurs aan dat een lagere scherptedrempel (die de schalingsfactor $c_t^2$ vermindert) de verhouding $\alpha_t = \|v_t\|^2 / c_t^2$ verhoogt.
• Deze toename in $\alpha_t` verschuift de uitlijning van de doelvector naar de leidende eigenvectoren van de NTK, waardoor feature learning en kernel target alignment (KTA) worden verbeterd.

5. Herwaardering van Kromming als Stabiliteitsdiagnose

De auteurs concluderen dat krommingdrempels afgeleid van convexe of kwadratische heuristieken (bijvoorbeeld het verwachten van stabiliteit alleen wanneer $\lambda_{\max} < 2/\eta - \gamma$) onbetrouwbare diagnostische middelen zijn onder regularisatie. Training kan karakteristiek EoS-gedrag vertonen, zelfs wanneer de scherpte aanzienlijk lager is dan deze naïeve drempels.

Betekenis en Beweringen

Het artikel beweert het eerste wiskundige kader te bieden dat de "folklore"-observatie nauwkeurig modelleert dat weight decay trainingsstabiliteit verbetert. De betekenis hiervan ligt in:

1. Unificerende Perspectieven: Het verbindt dynamica in de parameter-ruimte (EoS), dynamica in de functie-ruimte (NTK) en signaalpropagatie, en toont aan hoe weight decay alle drie moduleert.
2. Mechanistisch Inzicht: Het gaat voorbij het zien van weight decay als een statische straf, en identificeert het als een dynamisch controlesignaal dat lokale geometrie koppelt aan globale parametervector-uitlijning.
3. Verfijning van Stabiliteitsdiagnoses: Het daagt de gemeenschap uit om voorbij eenvoudige krommingdrempels te kijken om trainingsstabiliteit te diagnosticeren, en benadrukt dat "stabilisatie" kan manifesteren als een fase-overgang naar een regime met lagere kromming in plaats van een eenvoudige verschuiving van de instabiliteitsgrens.
4. Architectuur-gevoeligheid: Het benadrukt dat de effecten van weight decay niet uniform zijn, maar kritiek afhangen van de interactie van de architectuur met globale uitlijningstermen.

De auteurs stellen expliciet dat hun bevindingen de intuïtie van de praktijk ondersteunen dat weight decay training stabiliseert, maar verduidelijken dat dit gebeurt via een complexe wisselwerking van lokale wrijving en globale interacties die het optimalisatielandschap bij de Edge of Stability fundamenteel verandert.