Megalodon : Prétraitement et inférence efficaces pour les grands modèles de langage avec une longueur de contexte illimitée

Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length

April 12, 2024

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

cs.AI

Résumé

La complexité quadratique et la faible extrapolation de longueur des Transformers limitent leur capacité à s'adapter à des séquences longues. Bien que des solutions sous-quadratiques comme l'attention linéaire et les modèles à espace d'état existent, elles sous-performent empiriquement les Transformers en termes d'efficacité de pré-entraînement et de précision sur les tâches en aval. Nous présentons Megalodon, une architecture neuronale pour la modélisation efficace de séquences avec une longueur de contexte illimitée. Megalodon hérite de l'architecture de Mega (moyenne mobile exponentielle avec attention gated) et introduit en outre plusieurs composants techniques pour améliorer ses capacités et sa stabilité, notamment la moyenne mobile exponentielle complexe (CEMA), la couche de normalisation temporelle, le mécanisme d'attention normalisé et la pré-normalisation avec une configuration résiduelle à deux sauts. Dans une comparaison directe contrôlée avec Llama2, Megalodon atteint une meilleure efficacité que le Transformer à l'échelle de 7 milliards de paramètres et 2 000 milliards de tokens d'entraînement. Megalodon atteint une perte d'entraînement de 1,70, se situant à mi-chemin entre Llama2-7B (1,75) et 13B (1,67). 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