Der LAMB-Optimierer (Layer-wise Adaptive Moments for Batch training) ist ein stochastischer Optimierungsalgorithmus zum Training von tiefen neuronalen Netzen, der den Adam-Optimierer um einen schichtweisen Normalisierungsschritt erweitert. Er wurde 2019 von Forschern bei Google eingeführt (insbesondere von Yang You, Jing Li, Jonathan Hseu und anderen) und wurde entwickelt, um die effiziente Nutzung sehr großer Mini-Batch-Größen (z. B. 32.768 oder mehr) zu ermöglichen, ohne die Modellgenauigkeit zu beeinträchtigen oder eine erschöpfende Hyperparameter-Abstimmung zu erfordern. Dies erreicht er, indem er die Aktualisierungsgröße für jede Schicht basierend auf dem Verhältnis der Gewichtsnorm der Schicht zu ihrer Aktualisierungsnorm skaliert und so die Lernrate effektiv von der Skala der Parameter jeder Schicht entkoppelt.
LAMB war besonders einflussreich beim Training von Transformer (architecture)-basierten Modellen, einschließlich früher großer Sprachmodelle (LLMs) und Vision-Architekturen wie ResNet. Seine Übernahme in verteilte Trainingsframeworks wie TensorFlow (über tf.keras.optimizers.LAMB) und PyTorch (über die LAMB-Implementierung in Bibliotheken wie NVIDIA, Hugging Face und FairScale) machte ihn zu einem Standardwerkzeug für die Skalierung von Trainingsläufen. Durch die Ermöglichung größerer Batches reduziert LAMB die Wanduhrzeit, die zum Training modernster Modelle erforderlich ist, was für Organisationen wie OpenAI, Anthropic und Google DeepMind entscheidend ist, die auf massive Rechencluster angewiesen sind.
Motivation und Hintergrund
Das Training tiefer neuronaler Netze mit großen Mini-Batches reduziert die Anzahl der Parameteraktualisierungen pro Epoche und die Rechenzeit, aber die naive Skalierung der Batch-Größe führt oft zu schlechter Generalisierung und instabiler Konvergenz. Dies ist als "Large-Batch-Trainingsproblem" bekannt. Standard-Optimierer wie stochastischer Gradientenabstieg (SGD) mit Momentum oder Adam erfordern eine sorgfältige Anpassung der Lernrate, wenn die Batch-Größe zunimmt, und selbst dann verschlechtert sich die Genauigkeit oft. Der LAMB-Optimierer wurde entwickelt, um dies zu beheben, indem er den Optimierer robuster gegenüber Änderungen der Batch-Größe macht.
Die zentrale Erkenntnis hinter LAMB ist, dass verschiedene Schichten in einem tiefen Netzwerk sehr unterschiedliche Skalen von Gradienten und Gewichtsnormen aufweisen. Beispielsweise haben frühe Faltungsschichten kleine Gewichte, während spätere vollständig verbundene Schichten große haben. Eine einzige globale Lernrate in Adam kann entweder zu zu großen Aktualisierungen für einige Schichten führen (was Divergenz verursacht) oder zu zu kleinen für andere (was die Konvergenz verlangsamt). LAMB führt eine schichtweise adaptive Rate ein, die die Aktualisierung basierend auf dem Verhältnis der Gewichts- und Gradientennormen der Schicht normalisiert und so sicherstellt, dass sich jede Schicht relativ zu ihrer Größe stabil bewegt.
Algorithmusbeschreibung
LAMB kann als eine Variante von Adam mit einem zusätzlichen Normalisierungsschritt betrachtet werden. Sei \(\theta_t\) der Parameter zum Zeitpunkt \(t\) und \(g_t\) der Gradient der Verlustfunktion bezüglich \(\theta_t\). LAMB verwaltet erste und zweite Momente (\(m_t\) und \(v_t\)) der Gradienten, ähnlich wie Adam, mit exponentiellen Abklingraten \(\beta_1\) und \(\beta_2\) (typischerweise 0,9 und 0,999). Nach der Bias-Korrektur berechnet es eine Adam-Aktualisierung \(\frac{m_t}{\sqrt{v_t} + \epsilon}\).
Der entscheidende Unterschied ist das Vertrauensverhältnis \(\phi\): Für jede Schicht \(i\) gilt \(\phi_i = \frac{||\theta_{t,i}||}{||r_{t,i}||}\), wobei \(r_{t,i} = \frac{m_{t,i}}{\sqrt{v_{t,i}} + \epsilon}\) die Adam-Aktualisierung (ohne Lernrate) für diese Schicht ist und \(||\cdot||\) die L2-Norm bezeichnet. Die endgültige Aktualisierung für Schicht \(i\) ist \(\theta_{t+1,i} = \theta_{t,i} - \eta \cdot \phi_i \cdot r_{t,i}\), wobei \(\eta\) die globale Lernrate ist. Dieses Vertrauensverhältnis skaliert die Aktualisierung proportional zur Gewichtsnorm der Schicht, sodass kleine Schichten kleine (aber nicht vernachlässigbare) Aktualisierungen erhalten und große Schichten größere, aber stabile Aktualisierungen.
In der Praxis wird eine kleine Konstante (z. B. 1e-6) zum Nenner hinzugefügt, um eine Division durch Null zu vermeiden. Der Algorithmus enthält auch optionalen Gewichtsabfall (L2-Regularisierung), der in die Aktualisierung integriert ist, gemäß dem entkoppelten Gewichtsabfall-Ansatz, der in AdamW verwendet wird. Wenn das Vertrauensverhältnis für alle Schichten auf 1 gesetzt wird, reduziert sich LAMB auf den Standard-Adam (mit Bias-Korrektur).
Hyperparameter und Abstimmung
LAMB erbt die meisten Hyperparameter von Adam: \(\beta_1\) (Momentum), \(\beta_2\) (Varianzabfall), \(\epsilon\) (numerische Stabilität) und die Gewichtsabfallrate. Der primäre neue Hyperparameter ist die globale Lernrate \(\eta\), die für das Training mit großen Batches oft im Bereich 0,01-0,1 liegt, deutlich höher als typisch für Adam (z. B. 1e-3). Die Autoren fanden heraus, dass für sehr große Batches (z. B. 32.768 für BERT) eine Lernrate von 0,01 mit linearem Warmup über die ersten 10 % der Schritte gut funktioniert, und sie empfahlen auch die Verwendung eines Kosinus-Abfalls Lernratenplans über die verbleibenden Schritte.
Zusätzlich kann die Wahl von \(\beta_2\) die Stabilität beeinflussen; für Modelle mit spärlichen Gradienten kann ein höheres \(\beta_2\) (z. B. 0,99) verwendet werden. Die Autoren schlugen auch vor, dass die Batch-Größe proportional zur Lernrate skaliert werden kann (lineare Skalierungsregel), eine Richtlinie, die gut mit LAMB funktioniert. Wenn beispielsweise die Batch-Größe verdoppelt wird, kann auch die Lernrate verdoppelt werden, ohne Genauigkeitsverlust.
Leistung und Benchmarks
In der ursprünglichen Arbeit wurde LAMB an zwei Hauptaufgaben evaluiert: Training von ResNet-50 auf ImageNet (Bildklassifikation) und BERT (ein transformerbasiertes Sprachmodell) für maskierte Sprachmodellierung. Mit LAMB erreichten die Autoren eine ImageNet-Top-1-Genauigkeit von 76,0 % in nur 2.048 Iterationen mit einer Batch-Größe von 32.768, was der mit kleineren Batches (z. B. 256) in weit weniger Epochen erreichten Spitzenleistung entspricht. Für BERT trainierten sie das Modell auf die gleiche Genauigkeit (z. B. F1-Score von 1,0 auf SQUAD) in etwa 3,5 Minuten mit 1.024 TPUs, eine 10-fache Beschleunigung gegenüber früheren Methoden.
Anschließend wurde LAMB zum Standard-Optimierer für das Training von BERT-basierten Modellen in Googles internen Arbeitsabläufen. Die Arbeit berichtete, dass LAMB sowohl Adam als auch SGD mit Momentum übertraf, wenn Batch-Größen von 1.024 auf 65.536 skaliert wurden. Die Autoren zeigten auch, dass LAMB gut mit Gradient Clipping funktioniert (zur Verhinderung explodierender Gradienten) und mit gemischtem Präzisionstraining kompatibel ist, wie es auf moderner Hardware wie NVIDIA-GPUs und Google-TPUs verwendet wird.
Anwendungen im großskaligen Training
Die Hauptanwendung von LAMB liegt im verteilten Training, bei dem die Batch-Größe zu groß ist, um in den Speicher eines einzelnen Geräts zu passen. In solchen Konfigurationen werden Gradienten über mehrere GPUs oder TPUs mittels Datenparallelismus gemittelt. Beispielsweise verwenden OpenAI und Google DeepMind Optimierer analog zu LAMB, wenn sie große Transformer (architecture)-Modelle mit Sequenzlängen im Tausenderbereich trainieren. Obwohl neuere Optimierer wie LAMB (und sein Nachfolger LAMB2) vorgeschlagen wurden, bleibt LAMB eine zuverlässige Wahl in vielen Open-Source-Bemühungen, einschließlich des Trainings von Vision-Transformatoren und LLMs durch Forschungsgruppen und Unternehmen wie AI21 Labs und SambaNova.
Im Kontext des Machine learning auf Amazon Web Services (mit AWS Trainium-Hardware) wird LAMB in benutzerdefinierten Kernen für Effizienz unterstützt. Ebenso haben Intel und AMD LAMB auf ihren Beschleunigern benchmarkt. Die Fähigkeit des Optimierers, extreme Batch-Größen zu bewältigen, macht ihn wertvoll für das Vortraining von Modellen auf massiven Datensätzen, bei denen die Kosten einer einzelnen Epoche hoch sind und die Reduzierung der Epochen von größter Bedeutung ist.
Beziehung zu anderen Optimierern
LAMB ist Teil einer Familie adaptiver Optimierer, die SGD-Varianten, Adam und dessen Nachfolger wie AdamW (entkoppelter Gewichtsabfall) und LARS (Layer-wise Adaptive Rate Scaling) umfasst. LARS, 2017 von You et al. für das Training mit großen Batches von CNNs eingeführt, verwendet ein ähnliches schichtweises Vertrauensverhältnis, verwaltet jedoch keine zweiten Momente; es stützt sich auf erste Momente (Momentum) und Gradientennormen. LAMB kombiniert die Vorteile von LARS (schichtweise Skalierung) mit Adams adaptiven Lernraten pro Parameter, was es robuster für Modelle mit spärlichen Gradienten (wie Transformatoren) macht.
Ein weiterer eng verwandter Optimierer ist NVLAMB (von Nvidia), der Varianzreduktion integriert. LAMB bleibt jedoch einfacher und weit verbreitet. Für Sequence-to-Sequence (Seq2Seq)-Aufgaben mit Beam Search beeinflusst LAMB die Inferenz nicht direkt, hilft aber bei der Trainingskonvergenz, was indirekt die Sequenzdekodierung verbessert.
Erweiterungen und Varianten
Seit seiner Einführung wurden mehrere Varianten vorgeschlagen. LAMB2 (ebenfalls von Google) fügt einen Normalisierungsfaktor basierend auf der Gradientenvarianz hinzu, was die Stabilität für bestimmte Probleme verbessert. LARS mit Bias-Korrektur und andere Modifikationen sind ebenfalls üblich. In der Praxis implementieren viele Frameworks LAMB mit optionaler Bias-Korrektur für die Momente, was in den ersten Schritten vorteilhaft ist. Einige Implementierungen, wie in PyTorchs torch.optim.Lamb (im torch_optimizer-Paket), ermöglichen die Anpassung des Vertrauensverhältnis-Parameters oder die Verwendung einer benutzerdefinierten Lernrate pro Schicht.
Trotz des Aufkommens neuer Optimierer wie AdamW mit unterschiedlichen Skalierungsstrategien (z. B. 1cycle-Pläne) bleibt LAMB eine starke Baseline für das Training mit großen Batches. Die Forschungsgemeinschaft hat die Kombination von LAMB mit Data Augmentation und Gradient Clipping untersucht, um die Generalisierung weiter zu verbessern.
Praktische Überlegungen und Einschränkungen
Während LAMB in Umgebungen mit großen Batches hervorragend abschneidet, ist es nicht immer die beste Wahl für kleine Batch-Größen (z. B. unter 1.024). In solchen Regimen können Standard-Adam oder SGD mit Momentum einfacher und gleichermaßen effektiv sein. LAMB fügt einen Rechenaufwand für die Berechnung von Normen pro Schicht hinzu, der auf moderner Hardware vernachlässigbar ist, aber für Modelle mit vielen kleinen Schichten (z. B. U-Net-Architekturen) nicht trivial sein kann.
Eine weitere Einschränkung ist, dass das Vertrauensverhältnis von LAMB gelegentlich zu instabilem Training führen kann, wenn einige Schichten sehr kleine Gewichtsnormen (nahe Null) haben. Dies wird typischerweise durch Hinzufügen eines Epsilon-Terms zum Nenner und durch Verwendung von Gewichtsabfall gemildert, der verhindert, dass Gewichte gegen Null driften. Darüber hinaus erfordert LAMB eine sorgfältige Abstimmung der Lernrate und der Warmup-Schritte; ein unangemessener Plan kann zu Divergenz führen.
Der Speicherverbrauch ist ähnlich wie bei Adam (zwei Momentvektoren pro Parameter), sodass er nicht speicherhungriger ist. Für sehr große Modelle können Model Pruning oder Gradient Clipping zusammen mit LAMB verwendet werden, aber dies sind orthogonale Techniken.
Fazit
LAMB ist zu einem Eckpfeiler im Werkzeugkasten von Optimierungsalgorithmen für großskaliges Deep Learning geworden. Durch die Ermöglichung effektiven Trainings mit massiven Mini-Batches hat es die Entwicklung vieler Benchmark-Modelle beschleunigt und die Kosten für Experimente reduziert. Sein Prinzip der schichtweisen Anpassung hat nachfolgende Optimiererdesigns beeinflusst und bleibt eine praktische, gut verstandene Lösung für Praktiker, die sich den Herausforderungen des verteilten Trainings stellen müssen. Während Artificial intelligence weiter wächst, werden sich Optimierer wie LAMB wahrscheinlich weiterentwickeln, aber seine Kernideen von schichtweisem Vertrauen und adaptiven Momenten bleiben bestehen.