Label-NTK Alignments and A Tighter Convergence Bound in the NTK Regime

Technische Samenvatting: Label-NTK Alignments en een Strakkere Convergentiegrens in het NTK-regime

Probleemstelling
Het Neural Tangent Kernel (NTK)-raamwerk heeft met succes de optimalisatie van overgeparameteriseerde neurale netwerken verklaard door trainingsdynamiek te karakteriseren als bij benadering lineair. Hoewel bestaande theorieën exponentiële convergentie naar een globaal minimum garanderen, zijn deze garanties vaak te pessimistisch in vergelijking met empirische waarnemingen. Specifiek hangen klassieke convergentiesnelheden kritiek af van de kleinste eigenwaarde ($\lambda_{\min}$) van de NTK-matrix. In realistische situaties is $\lambda_{\min}$ doorgaans extreem klein, wat leidt tot theoretische convergentiegrenzen die orde van grootte trager zijn dan daadwerkelijke trainingscurven. Bovendien gaan deze worst-case analyses impliciet uit van een scenario waarin de data-labelvector uitgelijnd is met de eigenvector die correspondeert met $\lambda_{\min}$, een situatie die in de praktijk zelden voorkomt. Dit artikel adresseert de kloof tussen theoretische grenzen en praktische trainingsdynamiek door te onderzoeken of een strakkere convergentiegrens kan worden afgeleid die rekening houdt met de specifieke relatie tussen data-labels en het NTK-eigenspectrum.

Methodologie
De auteurs stellen een nieuwe analyse voor die gebaseerd is op de empirische en theoretische observatie van "alignments" (uitlijningen) tussen datavectoren (labels en residuen) en het NTK-eigensysteem.

1. Empirische Observatie: Door experimenten met Multi-Layer Perceptrons (MLP's) en Convolutional Neural Networks (CNN's) over datasets zoals CIFAR-10, observeren de auteurs twee distincte schalingswetten:

   • Label-NTK Alignment: De kwadratische projectie van de labelvector $y$ op de $i$-de eigenvector $v_i$ schaalt met het kwadraat van de corresponderende eigenwaarde: $(v_i^\top y)^2 \sim \lambda_i^2$.
   • Residual-NTK Alignment: De kwadratische projectie van de initiële residuvector $r_0 = f(w_0) - y$ op de eigenvectoren schaalt lineair met de eigenwaarden: $(v_i^\top r_0)^2 \sim \lambda_i$.
     Deze relaties gelden vooral in het regime van kleine eigenwaarden, wat aangeeft dat labels en residuen verwaarloosbare componenten hebben in de richtingen van kleine eigenwaarden.

2. Theoretische Rechtvaardiging: De auteurs leveren een theoretisch bewijs voor de Label-NTK alignment onder milde data-aannames, waaronder eenheidslengte van de invoer, niet-degeneratie en de Lipschitz-continuïteit van de ground-truth data-labelfunctie. Door een dataset met twee punten te analyseren en dit te generaliseren naar $n$ punten, tonen ze aan dat de Lipschitz-eigenschap van de ware functie de labelvector dwingt om sterker uitgelijnd te zijn met grotere eigenwaarden, wat de waargenomen schalingsexponent van 2 precies voorspelt.

3. Afleiding van Nieuwe Grenzen: Door gebruik te maken van de Residual-NTK alignment als voorwaarde (Voorwaarde 4.2), leiden de auteurs een verfijnde convergentiegrens voor gradient descent af. In plaats van te vertrouwen op $\lambda_{\min}$, hangt de nieuwe grens af van de trace van een functie van de volledige NTK-matrix: $L(w_t) \sim \text{tr}[(I - \eta K)^{2t} K]$.

Belangrijkste Bijdragen

• Identificatie van Alignment-verschijnselen: Het artikel identificeert en kwantificeert Label-NTK en Residual-NTK alignments, en toont aan dat data-labels en residuen van nature uitgelijnd zijn met de top-eigenvectoren van de NTK, waardoor ze effectief de worst-case scenario's vermijden die in eerdere literatuur worden verondersteld.
• Strakkere Convergentietheorie: De auteurs stellen een nieuw convergentiestelling op (Stelling 4.5) waarbij het verval van de loss wordt bepaald door het volledige spectrum van de NTK in plaats van slechts de kleinste eigenwaarde. Dit resulteert in een grens die nauw overeenkomt met empirische trainingscurven, zoals aangetoond in Figuur 1.
• Verbeterde Generalisatiegrenzen: Door gebruik te maken van de Label-NTK alignment-eigenschap, leidt het artikel verbeterde generalisatiegrenzen af voor overgeparameteriseerde netwerken, en toont het aan dat de testloss bovenaf begrensd wordt door $O(\text{tr}[K]/n)$, wat strakker is dan eerdere resultaten die afhankelijk zijn van $\lambda_{\min}$.
• Tijdcomplexiteitsanalyse: Het artikel biedt een tijdcomplexiteitsanalyse onder de aanname van een power-law verval van eigenwaarden, en toont aan dat het aantal iteraties nodig om een doelloss $\epsilon$ te bereiken significant lager is dan de standaardgrens $O(\frac{1}{\eta \lambda_{\min}} \log \frac{1}{\epsilon})$.

Resultaten

• Empirische Validatie: Experimenten met MLP's en CNN's over meerdere datasets bevestigen dat de voorgestelde theoretische grens de daadwerkelijke trainingslosscurven veel nauwkeuriger volgt dan klassieke grenzen. De theoretische curve met de op trace gebaseerde grens neemt snel af en komt overeen met de empirische data, terwijl de klassieke grens bijna vlak blijft vanwege de tiny $\lambda_{\min}$.
• Schalingswetten: De experimenten tonen consequent een lineair verband tussen $\log((v_i^\top y)^2)$ en $\log(\lambda_i)$ met een helling van 2, en tussen $\log((v_i^\top r)^2)$ en $\log(\lambda_i)$ met een helling van 1, wat de voorgestelde schalingswetten valideert.
• Generalisatie: De afgeleide generalisatiegrens suggereert dat de testloss schaalt met de trace van de NTK, wat een gunstiger schaling biedt dan eerdere worst-case analyses.

Betekenis
Het artikel stelt dat de primaire betekenis ligt in het oplossen van de discrepantie tussen NTK-gebaseerde theoretische voorspellingen en praktische trainingsprestaties. Door aan te tonen dat de worst-case aanname (labeluitlijning met $\lambda_{\min}$) in de praktijk niet opgaat vanwege inherente data-eigenschappen (Lipschitz-continuïteit), bieden de auteurs een wiskundig rigoureuze verklaring voor de snelle convergentie die in deep learning wordt waargenomen. Het werk biedt een "scherper" theoretisch perspectief dat volledige spectrale informatie integreert, waardoor de kloof wordt overbrugd tussen geïdealiseerde lineariseerde dynamiek en de waargenomen efficiëntie van gradient descent in overgeparameteriseerde regimes. De auteurs merken op dat hoewel hun theorie beperkt is tot het NTK-regime (voldoende brede netwerken, lage leersnelheden) en niet de niet-lineaire feature learning-dynamiek zoals de Edge of Stability vastlegt, het een noodzakelijke basis biedt voor het begrijpen van optimalisatie in de meest wiskundig hanteerbare setting die beschikbaar is.