Technische Zusammenfassung: Geschlossen-formale Optimierung der letzten Schicht

Problemstellung

Tiefe neuronale Netze werden typischerweise mittels Stochastic Gradient Descent (SGD) oder dessen Varianten optimiert, wobei alle Schichten einheitlich behandelt werden. Unter einer quadratischen Verlustfunktion admitieren die optimalen Gewichte der linearen letzten Schicht eines neuronalen Netzes jedoch eine geschlossen-formale Lösung (speziell eine Ridge-Regression-Lösung), gegeben die von der Backbone-Schicht erzeugten Merkmale.

Standard-SGD behandelt die letzte Schicht als Menge von Parametern, die über Gradientenschritte aktualisiert werden, und ignoriert diese analytische Optimalität. Dies erzeugt eine Diskrepanz: Während die Backbone-Parameter ($\theta$) eine iterative Optimierung erfordern, könnten die Gewichte der letzten Schicht ($W$) theoretisch direkt berechnet werden. Darüber hinaus führt in stochastischen Settings mit Mini-Batches die naive Berechnung der geschlossen-formalen Lösung für jeden Batch zu schwerwiegendem Overfitting auf diesen spezifischen Batch und verursacht Instabilität. Bestehende Ansätze, die versuchen, dies zu nutzen (z. B. in der kausalen Inferenz), erfordern oft eine Backpropagation durch die geschlossen-formale Inverse, was rechenintensiv ist, oder sie entkoppeln die letzte Schicht von der Backbone-Schicht, was zu Optimierungsinstabilitäten führt.

Methodik

Die Autoren schlagen ein Trainingsverfahren vor, das die letzte Schicht als implizite, deterministische Funktion der Backbone-Parameter behandelt, wobei ausschließlich die Backbone-Schicht optimiert wird, während sichergestellt wird, dass die letzte Schicht in jedem Schritt optimal bleibt.

1. Theoretische Grundlage (Deterministischer Fall)

Für einen festen Backbone-Parameter $\theta$ minimiert die optimale letzte Schicht $W^\star(\theta)$ den regularisierten quadratischen Verlust:
$$W^\star(\theta) = \arg\min_W \sum \|y_i - W\phi_\theta(x_i)\|^2 + \beta\|W\|_F^2$$
Die Autoren nutzen das Hüllkurven-Theorem (Envelope Theorem), um zu zeigen, dass der Gradient des zusammengesetzten Verlusts $L^\star(\theta) = L(W^\star(\theta), \theta)$ bezüglich $\theta$ identisch ist mit dem Gradienten des ursprünglichen Verlusts, wobei $W$ auf $W^\star(\theta)$ fixiert ist:
$$\nabla_\theta L^\star(\theta) = \nabla_\theta L(W^\star(\theta), \theta)$$
Dieses kritische Ergebnis impliziert, dass man nicht durch die Matrixinversion backpropagieren muss, die zur Berechnung von $W^\star(\theta)$ erforderlich ist. Stattdessen alterniert der Algorithmus zwischen:

1. Berechnung von $W_{t+1} = W^\star(\theta_t)$ (geschlossen-formale Aktualisierung).
2. Aktualisierung von $\theta_{t+1} = \theta_t - \alpha \nabla_\theta L(W_{t+1}, \theta_t)$ (Standard-Backpropagation durch die Backbone-Schicht).

2. Praktische Anpassung (Stochastischer Mini-Batch)

Im stochastischen Setting verursacht die Berechnung von $W^\star$ auf einem einzelnen Mini-Batch Overfitting. Um dies zu adressieren, führen die Autoren einen proximalen Regularisierer ein, der die Distanz zwischen der aktuellen geschlossen-formalen Lösung und der vorherigen Schätzung $W_t$ bestraft. Das Ziel wird zu:
$$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$$
Die geschlossen-formale Lösung für diesen proximalen Verlust lautet:
$$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}$$
Diese Formulierung implementiert effektiv einen approximativen Kalman-Filter auf der letzten Schicht, der Informationen aus vorherigen Batches akkumuliert und sich gleichzeitig an den aktuellen Batch anpasst. Der Algorithmus aktualisiert die Backbone-Schicht via SGD und passt dann die letzte Schicht unter Verwendung dieser proximalen geschlossen-formalen Lösung neu an.

3. Erweiterung auf Klassifikation

Obwohl die Methode primär für Regression entwickelt wurde, wird sie auf Klassifikation erweitert, indem One-Hot-Labels als Ziele für einen quadratischen Verlust behandelt werden. Die Ausgabe wird über eine $\arg\max$-Operation interpretiert. Die Autoren stellen fest, dass diese Heuristik für kleine bis moderate Klassenanzahlen funktioniert, jedoch keine Normalisierung der Wahrscheinlichkeiten garantiert.

Hauptbeiträge

1. Optimierungsrahmen: Ein praktischer Algorithmus, der geschlossen-formale Aktualisierungen der letzten Schicht in Standard-SGD-Pipelines integriert, ohne Backpropagation durch die Matrixinversionsoperation zu erfordern, unter Ausnutzung des Hüllkurven-Theorems.
2. Stochastische Regularisierung: Eine proximale Formulierung, die Overfitting in Mini-Batch-Settings verhindert, indem sie die letzte Schicht in Richtung ihrer vorherigen Schätzung regularisiert, interpretiert als approximatives Kalman-Filtering.
3. Theoretische Analyse:
   • NTK-Regime: Nachgewiesene Konvergenz zu einem globalen Minimum im unendlich-breiten Limit des Neural Tangent Kernel (NTK), trotz der Nicht-Konvexität des Verlusts im Funktionsraum.
   • Ein-Schritt-Analyse: Eine theoretische Schranke, die die Überlegenheit des geschlossen-formalen Schritts gegenüber einem Gradientenabstiegsschritt quantifiziert und zeigt, dass der Gewinn mit der Konditionszahl (Varianz der Singulärwerte) der Merkmalsmatrix zusammenhängt.
4. Empirische Validierung: Nachgewiesene Verbesserungen gegenüber Standard-SGD und naiver geschlossen-formaler Ridge-Regression über diverse Aufgaben hinweg.

Experimentelle Ergebnisse

Die Methode („$\ell_2$ c.f. proximal ($\lambda$)") wurde gegen Standard-SGD/Adam mit quadratischem Verlust und eine naive geschlossen-formale Ridge-Baseline („$\ell_2$ c.f. ridge ($\beta$)") evaluiert.

• Quantenchemie (QM9): Die proximale Methode übertraf Standard-SGD konsistent, insbesondere bei kleinen Batch-Größen, wo die naive Ridge-Baseline versagte, konvergieren zu können. Sie zeigte Robustheit gegenüber der Wahl des Backbone-Optimierers (SGD vs. Adam).
• Fourier-Neurale Operatoren (FNO): Beim Operator-Lernen für die viskose Burgers-Gleichung blieb die proximale Methode über Batch-Größen hinweg effektiv, wohingegen die naive Ridge-Baseline bei kleinen Batches versagte. Die Methode bewältigte effektiv die hochdimensionalen räumlichen Freiheitsgrade.
• Deep Feature Instrumental Variable (DFIV): In der kausalen Inferenz eliminierte die Methode die Notwendigkeit einer kostspieligen Neu-Schätzung der letzten Schichten auf dem gesamten Datensatz nach dem Training und erreichte eine Leistung, die mit Neu-Schätzungsstrategien vergleichbar war, während sie rechnerisch effizienter war.
• Klassifikation (CIFAR-100, ImageNet): Auf CIFAR-100 erreichte die Methode eine Genauigkeit, die der von Cross Entropy entsprach oder diese bei kleinen bis moderaten Batch-Größen übertraf. Auf ImageNet (große Anzahl von Klassen) unterlag sie jedoch Cross Entropy, was eine Einschränkung des quadratischen Verlusts für Klassifikationen mit hoher Kardinalität hervorhebt.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass die Nutzung der geschlossen-formalen Lösung für die letzte Schicht eine gangbare und effektive Strategie ist, um die Optimierung zu beschleunigen und die Stabilität zu verbessern, insbesondere bei Regressionsaufgaben und Regimen mit kleinen Batches.

• Theoretische Bedeutung: Die Arbeit liefert die erste Analyse der kritischen Punkte des modifizierten Verlusts im Funktionsraum und beweist die Konvergenz zu einem globalen Minimum im NTK-Regime, wodurch die im Problem inhärente Nicht-Konvexität adressiert wird.
• Praktische Bedeutung: Die proximale Formulierung löst das Overfitting-Problem, das in geschlossen-formalen Mini-Batch-Aktualisierungen inhärent ist, und macht die Methode robust und kompatibel mit Standard-Trainingspipelines.
• Einschränkungen: Die Autoren stellen explizit fest, dass die Methode derzeit auf quadratischen Verlust beschränkt ist. Obwohl sie bei Klassifikation mit wenigen Klassen gut funktioniert, unterliegt sie Cross Entropy bei großskaliger Klassifikation (ImageNet), was darauf hindeutet, dass die Anpassung des Ansatzes auf Cross-Entropy-Verlust eine notwendige Richtung für zukünftige Arbeiten ist. Die Methode verursacht zudem einen rechnerischen Overhead (Matrixinversion), der mit der Dimension der letzten Schicht skaliert, was sie für sehr große letzte Schichten in Regimen mit kleinen Batches langsamer als Standard-SGD macht, wobei diese Kosten jedoch abnehmen, wenn die Backbone-Größe zunimmt.