Megalodon: Effizientes LLM-Vortraining und Inferenz mit unbegrenzter Kontextlänge

Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length

April 12, 2024

Autoren: Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, Chunting Zhou

cs.AI

Zusammenfassung

Die quadratische Komplexität und die schwache Längenextrapolation von Transformatoren begrenzen ihre Fähigkeit, auf lange Sequenzen zu skalieren. Obwohl subquadratische Lösungen wie lineare Aufmerksamkeit und Zustandsraummodelle existieren, unterperformen sie empirisch im Vergleich zu Transformatoren hinsichtlich der Effizienz beim Pretraining und der Genauigkeit bei nachgelagerten Aufgaben. Wir stellen Megalodon vor, eine neuronale Architektur für effizientes Sequenzmodellieren mit unbegrenzter Kontextlänge. Megalodon erbt die Architektur von Mega (Exponentieller gleitender Durchschnitt mit gesteuerter Aufmerksamkeit) und führt zusätzliche technische Komponenten ein, um ihre Leistungsfähigkeit und Stabilität zu verbessern, darunter komplexe exponentielle gleitende Durchschnittswerte (CEMA), Zeitschrittnormalisierungsschicht, normalisierte Aufmerksamkeitsmechanismen und Pre-Norm mit Zweischritt-Residualkonfiguration. In einem kontrollierten direkten Vergleich mit Llama2 erzielt Megalodon eine bessere Effizienz als der Transformer im Maßstab von 7 Milliarden Parametern und 2 Billionen Trainings-Token. Megalodon erreicht einen Trainingsverlust von 1.70, was sich zwischen Llama2-7B (1.75) und 13B (1.67) befindet. Code: https://github.com/XuezheMax/megalodon

English

The quadratic complexity and weak length extrapolation of Transformers limits their ability to scale to long sequences, and while sub-quadratic solutions like linear attention and state space models exist, they empirically underperform Transformers in pretraining efficiency and downstream task accuracy. We introduce Megalodon, a neural architecture for efficient sequence modeling with unlimited context length. Megalodon inherits the architecture of Mega (exponential moving average with gated attention), and further introduces multiple technical components to improve its capability and stability, including complex exponential moving average (CEMA), timestep normalization layer, normalized attention mechanism and pre-norm with two-hop residual configuration. In a controlled head-to-head comparison with Llama2, Megalodon achieves better efficiency than Transformer in the scale of 7 billion parameters and 2 trillion training tokens. Megalodon reaches a training loss of 1.70, landing mid-way between Llama2-7B (1.75) and 13B (1.67). Code: https://github.com/XuezheMax/megalodon