Research map / Architectures and efficiency / Attention and transformers
Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces Megalodon, an improved Mega architecture with complex EMA, timestep normalization, normalized attention, and two-hop residual pre-norm, pretrained at 7B on 2T tokens against Llama2-7B with 32K context and evaluated on long-context, academic, and small/medium benchmarks.
Problem
Transformers scale quadratically and extrapolate poorly to long sequences, while sub-quadratic models (linear attention, SSMs, Mega) underperform Transformers in pretraining efficiency and downstream accuracy.
Approach
Extends Mega's damped EMA to the complex domain (CEMA), introduces a causal timestep normalization via cumulative mean/variance, normalized (QK-norm style) attention, and a pre-norm with two-hop residual configuration, combined with chunked attention for linear complexity.
Key results
- Megalodon-7B reaches training loss 1.70 on 2T tokens, between Llama2-7B (1.75) and Llama2-13B (1.67), and about 32% faster training throughput than Llama2-7B at 32K context.
- Megalodon-7B scores higher than Llama2-7B on all listed academic benchmarks, e.g., MMLU 49.8 vs 45.3 and ARC-c 53.1 vs 45.9.
- Megalodon reaches 25.4 test word-level perplexity on PG-19 with 1.3B params, versus 28.9 for Perceiver AR, and 17.23 PPL on WikiText-103 versus 18.07 for Mega.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace Llama2-7B with Megalodon-7B | 7B language model · MMLU | accuracy: 45.3 -> 49.8 | Helped |
| replace Llama2-7B with Megalodon-7B | 7B language model · BoolQ | accuracy: 77.4 -> 80.5 | Helped |
| replace Llama2-7B with Megalodon-7B | 7B language model · HellaSwag | accuracy: 77.2 -> 77.5 | Helped |
| replace Llama2-7B with Megalodon-7B | 7B language model · PIQA | accuracy: 78.8 -> 80.1 | Helped |
| replace Llama2-7B with Megalodon-7B | 7B language model · SIQA | accuracy: 48.3 -> 49.6 | Helped |
| replace Llama2-7B with Megalodon-7B | 7B language model · WinoGrande | accuracy: 69.2 -> 71.4 | Helped |
| replace Llama2-7B with Megalodon-7B | 7B language model · ARC-e | accuracy: 75.2 -> 79.8 | Helped |
| replace Llama2-7B with Megalodon-7B | 7B language model · ARC-c | accuracy: 45.9 -> 53.1 | Helped |
| replace Llama2-7B with Megalodon-7B | 7B language model · NaturalQuestions | accuracy: 25.7 -> 25.7 | No clear effect |
| replace Llama2-7B with Megalodon-7B | 7B language model · TriviaQA | accuracy: 58.5 -> 60.5 | Helped |
| instruction-finetune Megalodon-7B without RLHF | Megalodon-7B · MT-Bench | score: 6.17 -> 6.27 | Helped |
| use Megalodon instead of Mega | 90M image classifier · ImageNet-1K val | top-1 accuracy: 82.3 -> 83.1 | Helped |
Surprising findings
- Megalodon-7B is about 6% slower than Llama2-7B at 4K context, though faster at 32K context.
- Megalodon's training NLL is slightly worse than Llama2-7B during the first quarter of pretraining, attributed to the increased RoPE base.
Limitations
- Megalodon-7B is compared with Mistral-7B and Gemma-8B, which were trained on much larger data, so results are not directly comparable.
- Instruction tuning uses proprietary data and MT-Bench results are reported without detailed description.
Experimental setup
Datasets: Llama 2 pretraining mix (2T tokens) PG-19 WikiText-103 ImageNet-1K Speech Commands SC10 Long Range Arena Scrolls
Benchmarks: MMLU BoolQ HellaSwag PIQA SIQA WinoGrande ARC-e ARC-c NaturalQuestions TriviaQA NarrativeQA Qasper
Metrics: training loss NLL accuracy F1 word-level perplexity WPS MT-Bench score
Baselines: Llama2-7B Llama2-13B Mistral-7B Gemma-8B MPT-7B RWKV Mamba Mega S4 DeiT-B Block-Recurrent Transformer MegaByte
Builds on: Mega Llama 2 S4 Group Normalization Flash-Attention 2
Closest papers by meaning
- InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory (NeurIPS 2024 · similarity 0.87)
- Mini-Sequence Transformers: Optimizing Intermediate Memory for Long Sequences Training (NeurIPS 2024 · similarity 0.87)
- DeciMamba: Exploring the Length Extrapolation Potential of Mamba (ICLR 2025 · similarity 0.86)
- MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention (NeurIPS 2024 · similarity 0.86)
- IceFormer: Accelerated Inference with Long-Sequence Transformers on CPUs (ICLR 2024 · similarity 0.86)
- A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.