Megalodon: Эффективное предварительное обучение и вывод LLM с неограниченной длиной контекста

Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length

April 12, 2024

Авторы: Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, Chunting Zhou

cs.AI

Аннотация

Квадратичная сложность и недостаточная экстраполяция длины ограничивают способность Трансформеров масштабироваться до длинных последовательностей, и хотя существуют субквадратичные решения, такие как линейное внимание и модели пространства состояний, они эмпирически уступают Трансформерам в эффективности предварительного обучения и точности задач на выходе. Мы представляем Мегалодон, нейронную архитектуру для эффективного моделирования последовательностей с неограниченной длиной контекста. Мегалодон наследует архитектуру Меги (экспоненциальное скользящее среднее с воротным вниманием) и дополнительно вводит несколько технических компонентов для улучшения его способностей и стабильности, включая комплексное экспоненциальное скользящее среднее (CEMA), слой нормализации временных шагов, нормализованный механизм внимания и преднормализацию с конфигурацией двухшагового остатка. В контролируемом сравнении с Llama2 Мегалодон достигает лучшей эффективности, чем Трансформер с масштабом 7 миллиардов параметров и 2 триллиона обучающих токенов. Мегалодон достигает значения функции потерь на обучении 1.70, находясь посередине между Llama2-7B (1.75) и 13B (1.67). Код: 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