Stochastische Mittelung der Gewichte

Stochastische Gewichtsmittelung (SWA) ist eine Modellmittelungstechnik im Deep Learning, die die Generalisierung verbessert, indem sie die Gewichte eines neuronalen Netzwerks über mehrere Punkte entlang seiner Trainingsbahn mittelt. Im Gegensatz zum traditionellen Training, das nur die endgültigen Gewichte behält, sammelt SWA Gewichte in regelmäßigen Abständen während der späteren Phase des Trainings und berechnet deren arithmetisches Mittel. Dieses einfache Verfahren führt oft zu flacheren Lösungen in der Verlustlandschaft, die mit besserer Testleistung und Robustheit verbunden sind. SWA wurde von Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov und Andrew Gordon Wilson in ihrer Arbeit von 2018 mit dem Titel „Averaging Weights Leads to Wider Optima and Better Generalization“ eingeführt. Die Methode erfordert keine Änderungen an der zugrunde liegenden Modellarchitektur oder der Verlustfunktion, was die Integration in bestehende Trainingspipelines erleichtert.

Die Kernidee hinter SWA ist, dass während des Trainings mit stochastischem Gradientenabstieg (SGD) oder seinen Varianten die Modellparameter um einen Bereich mit geringem Verlust oszillieren. Die Mittelung dieser Parameter kann die Oszillationen glätten und eine Lösung ergeben, die in einem breiten, flachen Minimum liegt. Es wird angenommen, dass solche flachen Minima besser generalisieren als scharfe Minima, da sie weniger empfindlich auf Störungen und Verschiebungen zwischen Trainings- und Testverteilungen reagieren. SWA hat sich als verbessernd für die Genauigkeit und Kalibrierung über verschiedene Architekturen und Aufgaben hinweg erwiesen, einschließlich Bildklassifikation, Objekterkennung und Sprachmodellierung.

Mechanismus und Algorithmus

SWA arbeitet in zwei Phasen. In der ersten Phase wird das Modell mit einem Standard-Lernratenplan (z. B. Kosinus-Abklingen oder Stufenabfall) trainiert, um eine Region nahe einem Minimum zu erreichen. In der zweiten Phase wird das Training mit einer konstanten oder zyklischen Lernrate fortgesetzt, und die Gewichte werden am Ende jeder Epoche oder nach einer festen Anzahl von Iterationen gesammelt. Die SWA-Gewichte werden als laufender Durchschnitt aktualisiert: n_swa = n_swa + 1; w_swa = (w_swa * (n_swa - 1) + w) / n_swa, wobei w die aktuellen Gewichte und w_swa die gemittelten Gewichte sind. Das endgültige Modell verwendet w_swa.

Ein wichtiges Detail ist, dass die Statistiken der Batch-Normalisierung nach der Mittelung neu berechnet werden müssen, da der laufende Mittelwert und die Varianz der Batch-Normalisierungsschichten nicht zusammen mit den Gewichten gemittelt werden. In der Praxis wird nach der Berechnung der SWA-Gewichte ein Vorwärtsdurchlauf über die Trainingsdaten durchgeführt, um die Batch-Normalisierungsstatistiken zu aktualisieren. Dieser Schritt ist entscheidend für Modelle, die Batch-Normalisierung verwenden, wie ResNets.

Theoretische Motivation

Der Erfolg von SWA wird oft durch die Linse der Geometrie der Verlustlandschaft erklärt. Forschung von Izmailov et al. und nachfolgende Arbeiten haben gezeigt, dass SGD dazu neigt, zu Punkten nahe der Grenze einer Region mit geringem Verlust zu konvergieren, während die Mittelung über mehrere Punkte die Lösung in Richtung des Zentrums dieser Region bewegt. Dieser zentrale Punkt liegt typischerweise in einem flacheren Becken, das mit besserer Generalisierung verbunden ist. Die Verbindung zwischen flachen Minima und Generalisierung wurde ausführlich untersucht, mit Verbindungen zu PAC-Bayes-Schranken und der Schärfe der Verlustfunktion.

Eine andere Perspektive bringt SWA mit Bayes'scher Inferenz in Verbindung. Unter bestimmten Annahmen kann die Trajektorie von SGD als Markov-Kette betrachtet werden, die aus einer Posterior-Verteilung sampelt. Die Mittelung entlang dieser Trajektorie approximiert den Posterior-Mittelwert, was die Bayes-optimale Vorhersage für den quadratischen Fehlerverlust ist. Diese Interpretation verbindet SWA mit stochastischen Gradienten-Markov-Ketten-Monte-Carlo-Methoden, obwohl SWA einfacher ist und keine sorgfältige Abstimmung der Rauschskalen erfordert.

Varianten und Erweiterungen

Es wurden mehrere Varianten von SWA vorgeschlagen, um die Leistung weiter zu verbessern. Eine bemerkenswerte Erweiterung ist Stochastic Weight Averaging Gaussian (SWAG), eingeführt von Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov und Andrew Gordon Wilson im Jahr 2019. SWAG approximiert die Posterior-Verteilung der Gewichte, indem es eine Gauß-Verteilung unter Verwendung der ersten und zweiten Momente der gesammelten Gewichte anpasst. Dies ermöglicht Unsicherheitsschätzung und verbesserte Kalibrierung und kann für Bayes'sches Deep Learning verwendet werden.

Eine weitere Variante ist Fast Geometric Ensembling (FGE), das einen zyklischen Lernratenplan verwendet, um Gewichte aus verschiedenen Modi der Verlustlandschaft zu sammeln. SWA kann mit FGE kombiniert werden, um eine noch bessere Leistung zu erzielen. Darüber hinaus wurde SWA mit Techniken wie Lernratenplänen, Datenaugmentierung und Pruning integriert, um Robustheit und Effizienz zu verbessern.

Anwendungen

SWA wurde erfolgreich in einer Vielzahl von Bereichen angewendet. Im Bereich Computer Vision verbessert es die Genauigkeit von Bildklassifikatoren auf Datensätzen wie CIFAR-10, CIFAR-100 und ImageNet. Beispielsweise hat sich gezeigt, dass SWA die Testgenauigkeit eines PreAct ResNet-164 auf CIFAR-100 von 81,15 % auf 82,90 % ohne architektonische Änderungen steigert. Bei der Objekterkennung hilft SWA, das Training zu stabilisieren und die mittlere durchschnittliche Präzision zu verbessern.

In der Verarbeitung natürlicher Sprache wurde SWA verwendet, um Transformer und große Sprachmodelle effektiver zu trainieren. Es hat sich gezeigt, dass es Überanpassung reduziert und die Perplexität bei Sprachmodellierungsaufgaben verbessert. SWA kommt auch generativen Modellen wie GANs und Diffusionsmodellen zugute, indem es glattere Gewichtsaktualisierungen bietet, was zu stabilerem Training und höherer Qualität der Stichproben führen kann.

Über das standardmäßige überwachte Lernen hinaus wurde SWA für halbüberwachtes Lernen, Transferlernen und Domänenanpassung angepasst. Seine Einfachheit und Allgemeinheit machen es zu einem wertvollen Werkzeug im Werkzeugkasten des Maschinenlern-Praktikers.

Vergleich mit anderen Techniken

SWA wird oft mit anderen Regularisierungs- und Ensembling-Methoden verglichen. Im Gegensatz zu traditionellen Ensembles, die Vorhersagen aus mehreren unabhängig trainierten Modellen mitteln, mittelt SWA Gewichte, was rechnerisch günstiger ist und nur einen einzigen Trainingslauf erfordert. Dies macht es besonders attraktiv, wenn Modellspeicher oder Inferenzkosten eine Rolle spielen.

Im Vergleich zu Techniken wie Dropout oder Batch-Normalisierung modifiziert SWA nicht die Trainingsdynamik, sondern verarbeitet die Gewichte nachträglich. Es kann in Verbindung mit diesen Methoden verwendet werden. SWA unterscheidet sich auch von Adam und anderen adaptiven Optimierern, die pro Parameter Lernraten verwalten; SWA wird typischerweise auf SGD- oder Adam-Trajektorien angewendet.

Eine Einschränkung von SWA besteht darin, dass ein Zeitplan erforderlich ist, um zu bestimmen, wann mit dem Sammeln von Gewichten begonnen werden soll. Wenn das Sammeln zu früh beginnt, kann der Durchschnitt von suboptimalen Gewichten dominiert werden; wenn zu spät, verringern sich die Vorteile. In der Praxis ist eine einfache Faustregel, mit dem Sammeln zu beginnen, nachdem die Lernrate auf einen niedrigen Wert abgeklungen ist.

Praktische Überlegungen

Die Implementierung von SWA ist unkompliziert. Die meisten Deep-Learning-Frameworks wie PyTorch und TensorFlow bieten Dienstprogramme oder Beispiele für SWA. Die Hauptschritte sind: (1) Trainieren Sie das Modell mit einem Standardplan, (2) setzen Sie das Training mit einer konstanten oder zyklischen Lernrate fort, (3) sammeln Sie Gewichte in Intervallen, (4) berechnen Sie den laufenden Durchschnitt und (5) berechnen Sie die Batch-Normalisierungsstatistiken neu.

SWA hat sich als robust gegenüber Hyperparameterwahlen erwiesen. Die Lernrate für die zweite Phase wird typischerweise auf den Mindestwert des anfänglichen Zeitplans oder eine kleine Konstante gesetzt. Die Häufigkeit der Gewichtserfassung kann jede Epoche oder alle paar hundert Iterationen betragen. In der Praxis reicht es oft aus, Gewichte für 5 bis 20 Epochen zu sammeln.

Auswirkungen und Rezeption

Die Einführung von SWA hatte erhebliche Auswirkungen auf die Deep-Learning-Gemeinschaft. Es bot eine einfache, aber effektive Methode zur Verbesserung der Generalisierung ohne zusätzliche Rechenkosten während der Inferenz. Die Arbeit wurde weithin zitiert und hat zahlreiche Folgearbeiten zu Gewichtsmittelung und flachen Minima inspiriert. SWA ist heute eine Standardtechnik in vielen Trainingspipelines, insbesondere in der akademischen Forschung und bei Wettbewerben.

SWA hat auch die Entwicklung anderer Mittelungsmethoden beeinflusst, wie den exponentiell gleitenden Mittelwert (EMA), der häufig beim Training von generativen Modellen und großen Sprachmodellen verwendet wird. EMA ist eine Variante, die neueren Gewichten mehr Gewicht verleiht, und wird in der Praxis oft wegen seiner Einfachheit und Effektivität verwendet. SWA und EMA ergänzen sich, und einige Frameworks ermöglichen das Umschalten zwischen ihnen.

Einschränkungen und zukünftige Richtungen

Trotz seiner Vorteile hat SWA Einschränkungen. Es nimmt an, dass die Verlustlandschaft relativ glatt ist und dass die Mittelung von Gewichten sinnvoll ist, was möglicherweise nicht für alle Architekturen oder Aufgaben gilt. Beispielsweise erfordern Modelle mit Batch-Normalisierung eine sorgfältige Handhabung, und Modelle mit komplexen Parameterräumen (z. B. rekurrente Netzwerke) profitieren möglicherweise nicht so stark. Darüber hinaus liefert SWA keine Unsicherheitsschätzungen, sofern es nicht erweitert wird (z. B. SWAG).

Zukünftige Forschungsrichtungen umfassen die Entwicklung adaptiver Mittelungszeitpläne, die Kombination von SWA mit Bayes'schen Methoden und die Anwendung von SWA auf neue Bereiche wie bestärkendes Lernen und föderiertes Lernen. Während sich das Deep Learning weiterentwickelt, bleiben Gewichtsmittelungstechniken wie SWA ein grundlegendes Werkzeug zur Verbesserung der Modellleistung und -zuverlässigkeit.