Robust and Fast Training via Per-Sample Clipping

Technische Zusammenfassung: Robustes und schnelles Training durch pro-Sample-Clipping

Problemstellung
Die Arbeit befasst sich mit dem stochastischen Optimierungsproblem $\min_{x \in \mathbb{R}^d} F(x)$, wobei $F(x) = \mathbb{E}_{\xi \sim \mathcal{D}}[f(x, \xi)]$. Während der stochastische Gradientenabstieg (SGD) der Standardlöser ist, beruhen seine Konvergenzgarantien typischerweise auf der Annahme, dass das Gradientenrauschen eine beschränkte Varianz aufweist (leichtschwänzige Verteilungen). Jüngste empirische und theoretische Arbeiten zeigen jedoch, dass das Gradientenrauschen in vielen Anwendungen des maschinellen Lernens schwerfällig sein kann (heavy-tailed), also endliche $p$-te Momente für $p \in (1, 2]$ besitzt, aber keine beschränkte Varianz aufweist. Unter solchen Bedingungen kann der Standard-SGD schlecht abschneiden. Bestehende robuste Alternativen, wie das Standard-Gradienten-Clipping (Clip-SGD) und Normalized SGD, versagen entweder dabei, die optimalen Komplexitätsschranken für Stichproben zu erreichen, oder erfordern eine spezifische, oft unpraktikable Parametertuning (z. B. Abhängigkeit von der Initialisierungslücke $\Delta_1$ oder spezifischen logarithmischen Faktoren der Ausfallwahrscheinlichkeit $\delta$).

Methodik
Die Autoren schlagen Per-Sample Clipped SGD (PS-Clip-SGD) vor, einen robusten Gradientenschätzer, der Clipping auf einzelne Stichproben innerhalb eines Mini-Batches anwendet, bevor die Aggregation erfolgt, anstatt den aggregierten Batch-Gradienten zu clippen.

Für einen Batch der Größe $n$ mit den Stichproben $\xi^{(1)}_t, \dots, \xi^{(n)}_t$ ist der Gradientenschätzer $G(x_t, \alpha, \beta)$ definiert als:
$$G(x_t, \alpha, \beta) = \frac{1}{n} \sum_{k=1}^n \gamma^{(k)}_t \nabla f(x_t, \xi^{(k)}_t)$$
wobei der Pro-Sample-Clipping-Faktor $\gamma^{(k)}_t := \min\left\{1, \frac{\alpha k^{1/\beta}}{|\nabla f(x_t, \xi^{(k)}_t)|}\right\}$ ist. Die Parameter $\alpha$ und $\beta$ steuern den Clipping-Schwellenwert, wobei $\alpha$ mit der Momentenbeschränkung $\sigma$ und $\beta$ mit dem Tail-Index $p$ zusammenhängt. Die Aktualisierungsregel folgt dem Standard-SGD: $x_{t+1} = x_t - \eta_t G(x_t, \alpha, \beta)$.

Wesentliche Beiträge

1. Optimaler Konvergenzrate im Erwartungswert: Die Autoren beweisen, dass PS-Clip-SGD eine optimale Stichprobenkomplexität für nicht-konvexe Optimierung unter schwerfällig verteiltem Rauschen erreicht. Speziell erreicht die Methode eine Konvergenzrate von $O(1/T)$ für die mittlere quadratische Gradientennorm, was den von Zhang et al. (2020) etablierten unteren Schranken entspricht. Dies verbessert die Ergebnisse von Clip-SGD und Normalized SGD, insbesondere da die Schranke für die Summe der quadratischen Gradientennormen gilt (ein stärkeres Kriterium als die Summe der Normen in einigen früheren Arbeiten), ohne dass die Kenntnis der Initialisierungslücke $\Delta_1$ für die Schrittweite erforderlich ist.
2. Konvergenz mit hoher Wahrscheinlichkeit: Das Paper etabliert Konvergenzgarantien mit hoher Wahrscheinlichkeit, die den Konvergenzraten im Erwartungswert bis auf polylogarithmische Faktoren in der Ausfallwahrscheinlichkeit $\delta$ und der Anzahl der Schritte $T$ entsprechen. Bemerkenswert ist, dass die Abhängigkeit von $\log(1/\delta)$ im Vergleich zum aktuellen Stand der Technik bei Clip-SGD verbessert wurde und die Methode keine Schrittweiten oder Clipping-Schwellenwerte erfordert, die explizit als Funktionen von $\delta$ abgestimmt werden müssen.
3. Einblick in die Gradientenakkumulation: Die Autoren hinterfragen die gängige Praxis, Gradienten-Clipping erst nach allen Akkumulationsschritten in der Gradientenakkumulation anzuwenden. Sie schlagen vor und validieren empirisch, dass die Anwendung von Clipping auf der Ebene des Mini-Batches (Mini-Batch Clipping) während der Akkumulation die Trainingsleistung verbessern kann, wobei der zusätzliche Rechenaufwand verschwindend gering ist – dies steht im Gegensatz zum allgemeinen Konsens, dass solches intermediäres Clipping Bias einführt.

Ergebnisse

• Theoretisch: Die theoretische Analyse zeigt, dass PS-Clip-SGD starke Konvergenzgarantien unter den Annahmen der $L$-Glattheit und endlicher $p$-ter Momente bietet. Die abgeleiteten Schranken sind hinsichtlich der Abhängigkeit von $\delta$ und der Stärke des Konvergenzkriteriums (quadratische Normen vs. Normen) enger gefasst als die für Normalized SGD und Clip-SGD.
• Quadratische Optimierung: In Experimenten zur Minimierung einer quadratischen Funktion mit Pareto-verteiltem Rauschen übertraf PS-Clip-SGD sowohl Clip-SGD als auch Normalized SGD in verschiedenen Rauschregimen ($p \in \{1.2, 1.5, 1.8\}$ signifikant, selbst mit ungetunten Hyperparametern. Experimente mit hoher Wahrscheinlichkeit bestätigten die theoretische sublineare Abhängigkeit von $\log(1/\delta)$.
• Bildklassifizierung (AlexNet auf CIFAR-100): Beim Training von AlexNet mit Momentum erreichte PS-Clip-SGD eine höhere Validierungsgenauigkeit (64,84 %) im Vergleich zu Vanilla SGD mit Momentum (58,90 %) und Clip-SGD mit Momentum (58,94 %). Diese Verbesserung blieb auch nach Berücksichtigung des etwa 1,3-fachen Rechenaufwands durch die Berechnung der pro-Sample-Gradientennormen bestehen. Die Autoren führen dies auf die effektive Kontrolle einer kleinen Anzahl von Ausreißer-Gradienten mit sehr großen Normen zurück, die das Standard-Batch-Level-Clipping nicht zu erkennen vermag.
• Sprachmodellierung (GPT-2): In einem Gradientenakkumulations-Szenario beim Training eines 124M-Parameter GPT-2 Modells übertraf Mini-Batch Clipping (Clipping pro Akkumulationsschritt) das Standard-Clip-SGD (Clipping erst nach vollständiger Akkumulation) über mehrere Lernraten hinweg, obwohl die vorherrschende Meinung ist, dass intermediäres Clipping schädlich ist.

Bedeutung und Behauptungen
Das Paper behaupten, die erste Arbeit zu sein, die Per-Sample-Clipping spezifisch als Methode zur Verbesserung der Trainingsleistung in Gegenwart von schwerfällig verteiltem Rauschen vorschlägt, im Gegensatz zu seiner etablierten Verwendung in der Differential Privacy. Die Bedeutung liegt darin, eine Methode bereitzustellen, die:

1. Optimale theoretische Konvergenzraten für nicht-konvexe Heavy-Tailed-Probleme erreicht.
2. Robust gegenüber Hyperparameter-Wahlen ist und selbst mit ungenau abgestimmten Einstellungen gut abschneidet.
3. Zeigt, dass Per-Sample-Clipping praktische Leistungssteigerungen bei Deep-Learning-Aufgaben (wie AlexNet) erzielen kann, die den Rechenaufwand überwiegen.
4. Einen Paradigmenwechsel in Gradientenakkumulationsstrategien nahelegt, indem sie zeigt, dass intermediäres Clipping vorteilhaft statt schädlich sein kann.

Die Autoren räumen Einschränkungen ein und merken an, dass ihre theoretische Analyse die Kenntnis der Momentenbeschränkung $\sigma$ und des Tail-Index $p$ erfordert, um optimale Parameter zu definieren, sowie dass die Annahme beschränkter $p$-ter Momente stärker ist als die Annahme beschränkter zentrierter $p$-ter Momente. Sie stellen zudem fest, dass während Per-Sample-Clipping für kleinere Modelle wie AlexNet effektiv ist, der Rechenaufwand es für sehr große Architekturen ohne weitere Optimierung unpraktisch machen könnte, wobei Mini-Batch-Clipping eine kosteneffiziente Alternative bietet.