Technische Samenvatting: Gesloten-Vorm Optimalisatie van de Laatste Laag

Probleemstelling

Diepe neurale netwerken worden doorgaans geoptimaliseerd met Stochastic Gradient Descent (SGD) of variaties daarvan, waarbij alle lagen uniform worden behandeld. Echter, onder een kwadratische verliesfunctie hebben de optimale gewichten voor de lineaire laatste laag van een neurale netwerk een gesloten-vorm oplossing (specifiek, een ridge-regressie-oplossing) gegeven de kenmerken geproduceerd door de backbone.

Standaard SGD behandelt de laatste laag als een set parameters die worden bijgewerkt via gradiëntstappen, en negeert deze analytische optimaliteit. Dit creëert een mismatch: terwijl de backbone-parameters ($\theta$) iteratieve optimalisatie vereisen, kunnen de gewichten van de laatste laag ($W$) theoretisch direct worden berekend. Bovendien leidt in stochastische omgevingen met mini-batches het naïef berekenen van de gesloten-vorm oplossing voor elke batch tot ernstig overfitting op die specifieke batch, wat instabiliteit veroorzaakt. Bestaande benaderingen die proberen hier gebruik van te maken (bijvoorbeeld in causale inferentie) vereisen vaak backpropagatie door de gesloten-vorm inverse, wat rekenkundig duur is, of ze ontkoppelen de laatste laag van de backbone, wat leidt tot optimalisatie-instabiliteiten.

Methodologie

De auteurs stellen een trainingsprocedure voor die de laatste laag behandelt als een impliciete, deterministische functie van de backbone-parameters, waarbij uitsluitend wordt geoptimaliseerd voor de backbone, terwijl ervoor wordt gezorgd dat de laatste laag op elk moment optimaal blijft.

1. Theoretische Grondslag (Deterministisch Geval)

Voor een vaste backbone-parameter $\theta$ minimaliseert de optimale laatste laag $W^\star(\theta)$ de geregulariseerde kwadratische verliesfunctie:
$$W^\star(\theta) = \arg\min_W \sum \|y_i - W\phi_\theta(x_i)\|^2 + \beta\|W\|_F^2$$
De auteurs maken gebruik van de Envelope-stelling om aan te tonen dat de gradiënt van de samengestelde verliesfunctie $L^\star(\theta) = L(W^\star(\theta), \theta)$ met betrekking tot $\theta$ identiek is aan de gradiënt van de oorspronkelijke verliesfunctie waarbij $W$ vaststaat op $W^\star(\theta)$:
$$\nabla_\theta L^\star(\theta) = \nabla_\theta L(W^\star(\theta), \theta)$$
Deze cruciale uitkomst impliceert dat men niet hoeft te backpropageren door de matrixinversie die nodig is om $W^\star(\theta)$ te berekenen. In plaats daarvan wisselt het algoritme af tussen:

1. Het berekenen van $W_{t+1} = W^\star(\theta_t)$ (gesloten-vorm update).
2. Het bijwerken van $\theta_{t+1} = \theta_t - \alpha \nabla_\theta L(W_{t+1}, \theta_t)$ (standaard backpropagatie door de backbone).

2. Praktische Aanpassing (Stochastische Mini-Batch)

In de stochastische setting veroorzaakt het berekenen van $W^\star$ op een enkele mini-batch overfitting. Om dit aan te pakken, introduceren de auteurs een proximale regularisator die de afstand tussen de huidige gesloten-vorm oplossing en de vorige schatting $W_t$ bestraft. Het doel wordt:
$$L_{B_t, W_t}^{\text{prox}}(W, \theta) = \sum_{(x,y) \in B_t} \|y - W\phi_\theta(x)\|^2 + \lambda \|W - W_t\|_F^2$$
De gesloten-vorm oplossing voor deze proximale verliesfunctie is:
$$W^\star_{B_t, W_t}(\theta) = (Y_t \phi_\theta(X_t)^\top + \lambda W_t)(\phi_\theta(X_t)\phi_\theta(X_t)^\top + \lambda I)^{-1}$$
Deze formulatie implementeert effectief een benaderde Kalman-filter op de laatste laag, waarbij informatie van eerdere batches wordt geaccumuleerd terwijl er wordt aangepast aan de huidige batch. Het algoritme update de backbone via SGD en past vervolgens de laatste laag opnieuw aan met behulp van deze proximale gesloten-vorm oplossing.

3. Uitbreiding naar Classificatie

Hoewel de methode primair is ontworpen voor regressie, wordt deze uitgebreid naar classificatie door one-hot labels te behandelen als doelen voor een kwadratische verliesfunctie. De uitvoer wordt geïnterpreteerd via een $\arg\max$-operatie. De auteurs merken op dat hoewel deze heuristiek werkt voor kleine tot gemiddelde aantallen klassen, deze geen normalisatie van kansen garandeert.

Belangrijkste Bijdragen

1. Optimalisatiekader: Een praktisch algoritme dat gesloten-vorm updates van de laatste laag integreert in standaard SGD-pijplijnen zonder backpropagatie door de matrixinversie-operatie te vereisen, gebruikmakend van de envelope-stelling.
2. Stochastische Regularisatie: Een proximale formulering die overfitting in mini-batch-omgevingen voorkomt door de laatste laag te regulariseren naar zijn vorige schatting, geïnterpreteerd als benaderde Kalman-filtering.
3. Theoretische Analyse:
   • NTK-regime: Bewezen convergentie naar een globaal minimum in de limiet van oneindige breedte van de Neural Tangent Kernel (NTK), ondanks de niet-convexiteit van de verliesfunctie in functieruimte.
   • Eén-stapsanalyse: Een theoretische bovengrens die de superioriteit van de gesloten-vorm stap boven een gradiëntafdaalstap kwantificeert, waarbij wordt aangetoond dat het voordeel gerelateerd is aan het conditiegetal (variantie van singuliere waarden) van de kenmerkmatrix.
4. Empirische Validatie: Aangetoonde verbeteringen ten opzichte van standaard SGD en naïeve gesloten-vorm ridge-regressie over diverse taken.

Experimentele Resultaten

De methode ("$\ell_2$ c.f. proximale ($\lambda$)") werd geëvalueerd tegenover standaard kwadratische verlies SGD/Adam en een naïeve gesloten-vorm ridge-baseline ("$\ell_2$ c.f. ridge ($\beta$)").

• Kwantumchemie (QM9): De proximale methode presteerde consequent beter dan standaard SGD, vooral bij kleine batchgroottes waar de naïeve ridge-baseline faalde om te convergeren. Het toonde robustheid ten opzichte van de keuze van de backbone-optimizer (SGD versus Adam).
• Fourier Neural Operators (FNO): Bij operatorleer voor de viskeuze Burgers-vergelijking bleef de proximale methode effectief over batchgroottes, terwijl de naïeve ridge-baseline faalde bij kleine batches. De methode verwerkte effectief de hoogdimensionale ruimtelijke vrijheidsgraden.
• Deep Feature Instrumental Variable (DFIV): In causale inferentie elimineerde de methode de noodzaak voor kostbare her-schatting van laatste lagen op het volledige dataset na training, en bereikte prestaties die vergelijkbaar waren met her-schattingstrategieën, terwijl het rekenkundig efficiënter was.
• Classificatie (CIFAR-100, ImageNet): Op CIFAR-100 bereikte de methode een nauwkeurigheid die vergelijkbaar was met of beter was dan Cross Entropy voor kleine tot gemiddelde batchgroottes. Echter, op ImageNet (groot aantal klassen) presteerde het slechter dan Cross Entropy, wat een beperking van kwadratische verliesfuncties voor classificatie met hoge kardinaliteit benadrukt.

Betekenis en Beweringen

Het artikel beweert dat het benutten van de gesloten-vorm oplossing voor de laatste laag een levensvatbare en effectieve strategie is om optimalisatie te versnellen en stabiliteit te verbeteren, met name bij regressietaken en regimes met kleine batches.

• Theoretische Betekenis: Het werk biedt de eerste analyse van de kritieke punten van de gewijzigde verliesfunctie in functieruimte en bewijst convergentie naar een globaal minimum in het NTK-regime, waarmee de niet-convexiteit die inherent is aan het probleem wordt aangepakt.
• Praktische Betekenis: De proximale formulering lost het overfitting-probleem op dat inherent is aan gesloten-vorm updates in mini-batches, waardoor de methode robuust is en compatibel met standaard trainingspijplijnen.
• Beperkingen: De auteurs stellen expliciet dat de methode momenteel beperkt is tot kwadratische verliesfuncties. Hoewel het goed presteert bij classificatie met weinig klassen, presteert het slechter dan Cross Entropy bij grootschalige classificatie (ImageNet), wat suggereert dat het aanpassen van de aanpak aan cross-entropy-verlies een noodzakelijke richting is voor toekomstig werk. De methode brengt ook een rekenkundige overhead met zich mee (matrixinversie) die schaalt met de dimensie van de laatste laag, waardoor het trager is dan standaard SGD voor zeer grote laatste lagen in regimes met kleine batches, hoewel deze kosten afnemen naarmate de backbone-grootte toeneemt.