Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length

Summary

The paper introduces Megaladon, an improvement on the existing technique Mega. This technique uses a diagonalizable complex moving average to allow for integration of information across a longer context efficiently.

Strengths

- The technique is validated against modern models across a large variety of benchmarks - The diagonalizable complex moving average seems like an easy win for efficient computation and improvement - The scale of the models produced is quite large and provides clear signal, while also producing useful model artifacts

Weaknesses

The evaluation is lacking in some key aspects. If these are remedied, the paper is strong. - There needs to be more comparison against MEGA, this seems critical to establish technical novelty beyond operating at a larger scale than the Mega paper. At scale, for the core NLP benchmarks, how does Megaladon compare Mega and Mega chunk? What changes actually make an improvement? Ablations in particualar around these design choices seem key to establishing the technical novelty, since the changes are somewhat small and it's not clear they buy you a large gap in improvement, or which pieces do. I have looked at the Imagenet Mega comparison, and those in the appendix, but this is not at scale, which is a critical element of the paper. - The paper spends much time comparing against models of shorter context lengths, while it is clear the unlimited inference context length is an advantage, Mega had this too. A compute matched setting should still be compared against. In particular, How does Llama trained with let's say a ~28k context (or whatever context yields the same 1.3x speed improvement of Megaladon) perform in comparison at the scale of the main experiments on the main benchmarks. The monotonic perplexity in Figure 5 is nice to see, but the perplexity other methods at their native or extended context lengths should be included, the minimum perplexity by megaladon might be higher than those achieved by other methods at shorter context lengths. - A commonly targeted use case of models like this is direct byte level modeling, it would be interesting to see how this method performs on that task. Many of the methods in the long context section were designed with this purpose explicitly in mind. This seems like an application which is important to test on, particularly if context extension is viable in this domain using this method, that would be very interesting. - Mamba state space models seem to exhibit naturally good context extension, since this method can be seen as a state space model it might be good to compare in the long sequence domain at a similar scale to the 3b Mamba-[1,2] models. While it might be infeasible to include all of the above, the first 2-3 seem of particular importance to establish the usefulness and technical novelty of this technique.

Questions

Why do not start to modeling improvements until well over 1T training tokens, is this a byproduct of the learning schedule or of something structural within the model? How much compute is added over Mega/Mega-Chunk to incorporate the necessary changes? How is the complex moving average implemented?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations should be addressed further and more directly in the paper.

Summary

Megalodon i.e. Mega2 improves over Mega by using (1) complex-valued EMA; (2) improved normalization schemes (e.g. timestep normalization, attention normalization, pre-norm with 2-hop residuals).

Strengths

- Timestep normalization is simple and reasonable. A highly-optimized CUDA kernel provided in this work could be very influential - it is possible to be as popular as layernorm/groupnorm in the future. - Gated attention is verified in large-scale training setting for the first time.

Weaknesses

- There are no ablation studies in moderate-scale language modeling at all - Figure1: the training sequence length is not the same, thus the perplexity comparison might be unfair. - Lacking discussions & comparison to recent hybrid (local) attention and RNN models: e.g. Griffin, Jamba - Long context evaluation is not comprehensive: the paper claims unlimited context length, while it is unclear about the effective context length. It is better to provide the Needle in the Haystack results

Questions

- Complex-valued linear recurrence is replaced by gated linear recurrence for many recent works: e.g. LRU (complex) -> RG-LRU (real gated). HGRN1 (complex) -> HGRN2 (real gated). S4 (complex) -> Mamba1 & 2 (real gated). Did you try real-valued gated linear recurrence layers? - What's your opinion on Sliding Window Attention vs. Chunk Attention?

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Summary

This is an empirincal paper. The paper presents MEGALODON, a neural architecture designed to overcome the quadratic complexity and weak length extrapolation of Transformers. By extensive experiments, this paper demonstrates MEGALODON's ability to efficiently handle unlimited context lengths and its superior performance across different tasks and modalities. The claimed key contributions of this paper is as follows:: 1) improves upon the MEGA architecture by adding the complex exponential moving average (CEMA), timestep normalization, normalized attention, and a pre-norm with two-hop residuals. 2. achieves better pretraining efficiency and downstream task accuracy than Transformers, specifically in handling long sequences with 7 billion parameters and 2 trillion training tokens. 3. outperforms Transformers across various benchmarks, including long-document comprehension, multi-turn conversation, and video generation.

Strengths

1. The performance of the proposed architecture is excellent, as demonstrated in Table 1. 2. Extensive experiments are conducted to evaluate the proposed methods. 3. The topic of this paper is important and critical for the LLM domain.

Weaknesses

1. The training curves of MEGALODON and LLAMA2 7B in Figure 1 cross at 750 billion training tokens. It would be interesting to see if they cross again with more tokens, such as at 6 trillion. 2. Section 3.2 is difficult to understand due to the lack of background and intuitive explanations. 3. Section 3.5 makes overclaims about 4D parallelism. This topic is well-explored and relevant, as discussed in [1, 2, 3], but these references are ignored. If the paper aims to highlight the benefits of 4D parallelism, it should include comparative experiments. 4. The code link provided in the abstract does not work. 5. The long context tasks in Section 4.3, as shown in Tables 2 and 6, are not fairly set up. In Table 2, the only fair baseline is LLAMA2-L, which performs better. This raises doubts about the proposed method's long-context ability since other models (Yarn, MPT, Xgen) use different training datasets, potentially limiting their long-context capability. For Table 5, the baselines (Transformer, Reformer, ...) also use different training datasets, making the comparisons misleading. Compared to MEGA, the improvements are marginal. Thue the compression of the long-ctx ability is not convincing to me. [1] Sequence Parallelism: Making 4D Parallelism Possible\ [2] Lightseq: Sequence level parallelism for distributed training of long context transformers\ [3] USP: A Unified Sequence Parallelism Approach for Long Context Generative AI

Questions

see Weaknesses

Rating

5

Confidence

5

Soundness

2

Presentation

3

Contribution

3

Limitations

see Weaknesses

Summary

The paper proposes Megalodon, which introduces three advancements over Mega: complex EMA, timestep normalization, and normalized attention. These advancements address the limitations of chunk-wise attention and architecture divergence across different tasks and data types. The new model is evaluated alongside Llama-2, both trained on the same public dataset, and demonstrates competitive and superior performance across a wide range of benchmarks.

Strengths

1. Clear motivations: All three improvements directly target the limitations of Mega. 2. The complex EMA is a novel approach. 3. The authors provide efficient parallelism.

Weaknesses

Recent theoretical work [1] has shown that efficient versions of attention (like the chunk-based method used in this paper) can limit the expressiveness of the model, particularly for reasoning tasks that involve long-range information. The paper evaluates Megalodon with long-context open-book QA tasks. How will Megalodon perform on PhoneBook lookup [2] with ICL, especially for phonebook lengths longer than 4K tokens? How will it perform on complex reasoning tasks requiring long context, such as 8-shot or 16-shot math problems with GSM8K or coding tasks on HumanEval? Empirically, it is not clear how CEMA improves expressiveness. It would be most direct to compare using CEMA versus using EMA on these tasks. The reviewer understands that the rebuttal period is short and is therefore not requiring most of these experiments to be added. [1] Yang, Kai, et al. "Do Efficient Transformers Really Save Computation?" Forty-first International Conference on Machine Learning. [2] Jelassi, Samy, David Brandfonbrener, and Sham M. Kakade. "Repeat After Me: Transformers are Better than State Space Models at Copying." Forty-first International Conference on Machine Learning.

Questions

Have the authors conducted ablation studies for small-scale models before moving to 7B, similar to the results in Appendix C? Including these ablation studies, if already performed, would help readers understand how the three designs impact performance.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper does not discuss its limitations. The authors believe there is no negative societal impact. In the paper checklist, justifications are required for answers marked "Yes," but the authors have deleted the justification for several items. For limitations, the authors claim they are discussed, but there is no justification provided. The anonymous link is not working.

Nearly There

Hi, Thank you for the thoughtful response. The ablations help, and knowing that the changes were critical to allowing Mega to achieve larger scales is good. Knowing that there are byte level classification present. For the monotonic perplexity plot, can you please add in a y-axis scale and the the perplexities of the other techniques at their native context length. The moving average makes it clear that generalization is possible, but it's critical to know the range of improvement we can actually expect. The actual improvement may not be that much though. With this final change, I will raise my score. Thank you