Optimierung großer Batch-Größen für Deep Learning: BERT in 76 Minuten trainieren

Das Training großer tiefer Neuronaler Netze anhand umfangreicher Datensätze stellt eine enorme rechnerische Herausforderung dar. In jüngerer Zeit hat sich das Interesse an der Verwendung von stochastischen Optimierungsmethoden mit großen Batch-Größen zur Bewältigung dieses Problems gesteigert. Der prominenteste Algorithmus in diesem Forschungsbereich ist LARS, der durch die Anwendung schichtweise angepasster Lernraten ResNet auf ImageNet in wenigen Minuten trainieren kann. Allerdings zeigt LARS bei Aufmerksamkeitsmodellen wie BERT schlechte Leistungen, was darauf hindeutet, dass seine Leistungsgewinne nicht über alle Aufgaben hinweg konsistent sind. In dieser Arbeit untersuchen wir zunächst eine prinzipielle schichtweise Anpassungsstrategie, um das Training tiefer Neuronaler Netze unter Verwendung großer Minibatches zu beschleunigen. Mit Hilfe dieser Strategie entwickeln wir eine neue schichtweise adaptive Optimierungstechnik für große Batch-Größen, die LAMB genannt wird; anschließend liefern wir eine Konvergenzanalyse sowohl für LAMB als auch für LARS und zeigen die Konvergenz zu einem stationären Punkt in allgemeinen nichtkonvexen Szenarien. Unsere empirischen Ergebnisse demonstrieren die überlegene Leistung von LAMB bei verschiedenen Aufgaben, wie dem Training von BERT und ResNet-50, mit sehr wenig Hyperparameter-Tuning. Insbesondere beim BERT-Training ermöglicht unser Optimierer die Verwendung sehr großer Batch-Größen von 32868 ohne jegliche Verschlechterung der Leistung. Durch das Erhöhen der Batch-Größe bis zum Speicherlimit eines TPUv3-Pods kann die Trainingszeit für BERT von drei Tagen auf nur 76 Minuten reduziert werden (Tabelle 1). Die Implementierung von LAMB ist unter https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py verfügbar.