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

要旨

Transformerの二次複雑度と弱い長さ外挿能力は、長いシーケンスへのスケーリングを制限しており、線形注意や状態空間モデルなどの準二次的な解決策が存在するものの、それらは事前学習効率と下流タスクの精度においてTransformerに劣ることが実証されています。本論文では、無制限のコンテキスト長を可能にする効率的なシーケンスモデリングのためのニューラルアーキテクチャ、Megalodonを紹介します。MegalodonはMega（指数移動平均とゲート付き注意）のアーキテクチャを継承し、複素指数移動平均（CEMA）、タイムステップ正規化層、正規化注意メカニズム、および2ホップ残差構成を持つ事前正規化など、その能力と安定性を向上させるための複数の技術的コンポーネントをさらに導入しています。Llama2との制御された直接比較において、Megalodonは70億パラメータと2兆トレーニングトークンのスケールでTransformerよりも優れた効率を達成します。Megalodonは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