Technische Samenvatting: De Edge of Stability vormt de leerprocessen selectief over de datadistributie

1. Probleemstelling

Bestaande literatuur over de Edge of Stability (EoS) behandelt dit primair als een globale eigenschap van optimalisatie, waarbij de grootste Hessiaanse eigenwaarde zichzelf zelf-stabiliseert nabij de drempelwaarde $2/\eta$ (waarbij $\eta$ de leersnelheid is). Hoewel de dynamiek van dit regime—gekenmerkt door teken-alternerende oscillaties en een trage corrigerende drift—goed gedocumenteerd is, blijft de functionele consequentie van opereren op deze drempel onduidelijk. Specifiek is het onbekend of de stabiliteitsrestrictie een uniforme voordelen biedt of dat het de leerprocessen differentieel vormgeeft over verschillende subsets van de trainingsdata. De centrale vraag is: Induceert het EoS-regime een selectieve leerbias, en zo ja, welke geometrische en dynamische eigenschappen bepalen welke datasubsets profiteren?

2. Methodologie

2.1 Experimentele Opzet en Prototype Taxonomie

De auteurs instantiëren een gecontroleerd experimenteel kader met behulp van een binaire classificatietaak op CIFAR-10 (Automobiel versus Vrachtwagen) met een MLP met twee verborgen lagen, getraind met full-batch gradiëntafdaling. Om de effecten van de datageometrie te isoleren, definiëren zij een prototype taxonomie die de trainingsdistributie partitioneert in vier groepen op basis van hun positie ten opzichte van de klasse-centra ($\mu_c$) en labelconsistentie, onafhankelijk van het getrainde model:

1. Inliers: Punten met een hoge dichtheid nabij het klasse-centrum met correcte labels.
2. Boundary: In-distributie punten nabij de inter-klasse grens (hoge labelambiguïteit).
3. Input-outliers: Geometrisch atypische inputs (ver van $\mu_c$) met correcte labels.
4. Output-outliers: Inputs met een hoge dichtheid die onjuiste labels toegewezen gekregen hebben (synthetisch geconstrueerd).

2.2 Branching Interventie (Causale Identificatie)

Om het effect van het blijven in het EoS-regime versus het verlaten ervan causaal te identificeren, maken de auteurs gebruik van een branching interventie:

• Baseline Branch: Gaat door met trainen bij de oorspronkelijke leersnelheid $\eta$, waardoor men in het EoS-regime blijft.
• Exit Branch: Bij het begin van de EoS (gedetecteerd wanneer de grootste Hessiaanse eigenwaarde $\lambda_1$ de $2/\eta$ bereikt), wordt de leersnelheid gehalveerd ($\eta \to \eta/2$). Dit verhoogt de stabiliteitsdrempel naar $4/\eta$, waardoor de optimizer uit het EoS-regime kan treden.
• Vergelijking: Door de divergentie in per-groep verliezen ($\ell_k$) tussen de twee branches na de interventie te vergelijken, isoleren de auteurs de causale impact van de EoS-stabiliteitsrestrictie.

2.3 Mechanisme Isolatie

De auteurs testen twee hypothetische noodzakelijke condities voor een subset om te profiteren van EoS:

1. Directional Alignment (Directionele Uitlijning): De geaggregeerde gradiënt van de subset moet uitlijnen met de grootste Hessiaanse eigenvector ($v_1$). Dit wordt getest door "random-direction" input-outliers te construeren waarbij de afstand tot het centrum behouden blijft, maar de richting gerandomiseerd is, waardoor coherente uitlijning wordt vernietigd.
2. Gradient Persistence (Gradiënt Persistentie): De subset moet een niet-verdwijnende gradiëntmagnitude behouden. Dit wordt getest door de Mean Squared Error (MSE) loss (waarbij gradiënten persistent zijn) te vergelijken met de Cross-Entropy (CE) loss (waarbij gradiënten verdwijnen voor zelfverzekerd geclassificeerde punten).

3. Belangrijkste Resultaten

3.1 Selectieve Leer Trade-off

De branching interventie onthult dat EoS selectief is, niet globaal. De stabiliteitsrestrictie herverdeelt de optimalisatie-inspanning ongelijkmatig:

• Begunstigden: Subsets waarvan de gradiënten uitlijnen met $v_1$ en persistent zijn (bijv. input-outliers onder MSE, output-outliers onder CE) vertonen een versnelde reductie van het verlies in de EoS-branch vergeleken met de exit-branch.
• Suppressors (Onderdrukkers): Subsets met orthogonale gradiënten of verdwijnende magnitudes (bijv. inliers en boundary points onder specifieke condities) vertonen een onderdrukte progressie in de EoS-branch vergeleken met de exit-branch.
• Deze trade-off repliceert over architecturen (CNN, ResNet), optimizers (SGD, Momentum) en klasse-paren.

3.2 Mechanisme: Alignment $\times$ Persistence

Twee condities zijn gezamenlijk noodzakelijk voor een groep om het EoS-voordeel te verkrijgen:

1. Alignment: De geaggregeerde gradiënt van de groep moet uitlijnen met de grootste Hessiaanse eigenvector ($v_1$).
   • Bewijs: Wanneer input-outliers in willekeurige richtingen worden verplaatst (waarbij afstand behouden blijft maar uitlijning wordt vernietigd), stort hun krommingseffect $(\nabla \ell_k \cdot v_1)^2$ met een orde van grootte in, en verdwijnt het EoS-voordeel.
2. Persistence: De groep moet een niet-verdwijnende gradiëntmagnitude behouden.
   • Bewijs: Onder de Cross-Entropy loss verzadigen (verdwijnen) de gradiënten voor zelfverzekerd geclassificeerde input-outliers. Bijgevolg verschuift het EoS-voordeel naar output-outliers, wier gradiënten actief blijven door labelinconsistentie, ondanks potentieel lagere uitlijning.

3.3 Geometrie Bepaalt de Begunstigde

De subset die profiteert van EoS wordt bepaald door de geometrische compositie van de datadistributie:

• In de standaard CIFAR-10 setup domineren input-outliers (verที่สุด van de centra) $v_1$ en profiteren zij.
• In een moeilijker klasse-paar (Kat vs. Hond) of met andere parameters voor de constructie van outliers, kunnen boundary points de meest verre groep worden, $v_1$ domineren en de primaire begunstigde worden.
• Natuurlijke Data: Op de natuurlijke CIFAR-10 distributie ontstaat bij de onset van EoS een positieve correlatie tussen de afstand tot het centrum en de uitlijning met $v_1$, wat bevestigt dat geometrische outliers van nature de onstabiele modus aansturen.

3.4 Implicaties voor Generalisatie en Robuustheid

Het paper presenteert voorlopig bewijs dat de training-tijd selectiviteit overgaat naar gedrag tijdens de testfase:

• Adversariële Robuustheid: EoS verbetert de robuustheid op boundary points alleen wanneer boundary points $v_1$ domineren. Als input-outliers domineren, presteert de EoS-branch slechter op boundary robuustheid omdat het optimalisatiebudget wordt afgeleid.
• OOD Generalisatie (Out-of-Distribution): EoS verbetert de OOD-prestaties op geëxtrapoleerde inputs alleen wanneer die inputs tijdens de training $v_1$ domineren.

4. Theoretisch Kader

De auteurs breiden het zelf-stabilisatie kader (Damian et al., 2023) uit van globale loss naar per-subset losses. Zij leiden af dat de cyclus-gemiddelde drift in het EoS-regime een extra term bevat die proportioneel is aan $-\nabla S$ (waarbij $S$ de scherpte/sharpness is). Het differentiële effect op een subset $k$ wordt beheerst door de selector:
$$Q_k = \langle \nabla \ell_k, \nabla S \rangle$$
Onder een single-mode benadering is deze selector proportioneel aan het gekwadrateerde projectie van de subset gradiënt op de grootste eigenvector:
$$C_k = (\nabla \ell_k \cdot v_1)^2 = \|\nabla \ell_k\|^2 \cos^2 \theta_k$$
Deze decompositie bevestigt dat zowel directionele uitlijning ($\cos^2 \theta_k$) als gradiënt magnitude persistentie ($\|\nabla \ell_k\|^2$) vereist zijn voor een subset om de EoS-drift te ervaren.

5. Betekenis en Claims

Het paper claimt het begrip van EoS te verschuiven van een scalaire stabiliteitsgrens naar een mechanisme dat de allocatie van leren over de datadistributie beheert.

• Herformulering van Flatness: De auteurs argumenteren dat "flatness" geen scalaire eigenschap van de oplossing is, maar een directionele eigenschap die wordt bepaald door de datageometrie. Het functionele voordeel van lage kromming hangt volledig af van welke subset de grootste Hessiaanse eigendirection domineert.
• Inductieve Bias: EoS fungeert als een inductieve bias die specifieke subsets selecteert voor optimalisatie op basis van hun geometrische relatie tot de onstabiele modus, in plaats van enkel te selecteren tussen equivalente minimizers.
• Verenigende Perspectieven: Het werk verbindt impliciete regularisatie in de parameterruimte (kromming-restricties) met inductieve bias over de datadistributie (welke voorbeelden worden geleerd).
• Beperkte Omvang: De auteurs erkennen beperkingen, waarbij zij opmerken dat de experimenten beperkt zijn tot kleine modellen en full-batch training vanwege de computationele kosten van Hessian-vector producten. Zij stellen dat of deze selectiviteitseffecten standhouden voor grotere datasets, mini-batch optimizers of representatie-ruimte prototypes, een openstaande vraag is.

Concluderend demonstreert het paper dat de Edge of Stability een selectief regime is waarbij de stabiliteitsrestrictie het leren versterkt voor subsets die uitlijnen met en persistent zijn langs de grootste Hessiaanse eigendirection, terwijl het anderen onderdrukt, waardoor de functionele eigenschappen van het model fundamenteel worden gevormd.