Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length
Authors: Xuezhe Ma
Xiaomeng Yang
Wenhan Xiong
Beidi Chen
LILI YU
Hao Zhang
Jonathan May
Luke Zettlemoyer
Omer Levy
Chunting Zhou
Venue: NeurIPS 2024 poster (2024)
Keywords: Mega
Efficient Architecture
Long Sequence Modeling
Unlimited Context Length
Primary area: deep_learning_architectures
TL;DR: Megalodon: Efficient Long-Context LLM Pretraining and Inference with Unlimited Context Length

Abstract: 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, an 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 LLAMA2-13B (1.67). This result is robust throughout a wide range of benchmarks, where MEGALODON consistently outperforms Transformers across different tasks, domains, and modalities.

MEGALODON: Efficient LLM Pretraining and Inference with Unlimited Context Length
Xuezhe Ma
Xiaomeng Yang"
Wenhan Xiong
Beidi Chen"
Lili Yu"
Hao Zhang
Jonathan May"
Luke Zettlemoyer*
Omer Levy'
Chunting Zhon
"Al at Meta
* University of Southern California
" "Carnegie Mellon University
University of California San Diego
Abstract
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 MEGA- LODON, 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 exponen- tial moving average (CEMA), timestep normalization layer, normalized attention mechanism and pre-norm with two-hop residual configuration. In controlled head-to-head comparison with LLAMA2, MEGALODON achieves better efficiency than Transformer in the scale of billion parameters and 2 trillion training tokens. MEGALODON reaches a training loss 1.70, landing mid-way between LLAMA2- 7B (1.75) and 13B (1.67). The improvements of MEGALODON over Transformers are robust throughout a range benchmarks across different tasks and modalities. Code: https: //githut cww ccm.//uuzhheaa///ggalooooo
Introduction
In many real-world applications, such multi-turn conversation, long-document comprehension. and video generation, large language models (LLMs) must efficiently process long sequential data, understand internal long-range dynamics, and generate coherent output. The Transformer archi- tecture (Vaswani et al., 2017), despite its remarkable capabilities. faces challenges with quadratic computational complexity and limited inductive bias for length generalization, making inefficient for long sequence modeling ( (Wang et al., 2024: Zhou et al., 2024). Even with recently proposed distributed attention solutions (Li al., 2023b; Liu et al., 2024), computing single training step of 7B parameter model over a 1M-token sequence is more than 100 times slower than performing the equivalent computation using 256 separate sequences 4K tokens each.
Techniques like efficient attention mechanisms (Tay et al., 2020; Ma et al., 2021) and structured state space models (Gu al., 2022a; Poli et al., 2023; Gu and Dao, 2023) have been introduced overcome these limitations, aiming to enhance scalability and performance. However, the practical application of these methods still falls short of Transformers (Tay et al., 2022: Gu and Dao. 2023). This work introduces an unlimited context model that outperforms the canonical Transformer architecture on real-world language modeling.
'Equal Contribution. Correspondence chuntinz mema.com
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

2.2
Llama2 7B
Llama2 13B
2.1
Megalodon 7B
2.0
Train Loss
1.9
1.8
1.7-
1.6
0
250
500
750
1000
1250
1500
1750
2000
Training Tokens (Billions)
Figure Negative log-likelihood for MEGALODON-7B, LLAMMA-7B and LLAMA2-132
Table Performance on standard academic benchmarks, compared to open-source hase models.
We reported model size, context length and total data tokens during model pretraining. indicates that the number was not reported in the original paper.
Model
Size
Context
MMLU
BoolQ
HellaSw
PIQA
SIQA
WinoG
Arc-e
NQ
Mamba
3B
0.6T
2K
26.2
71.0
71.0
78.1
65.9
68.2
41.7
RWKV
7B
1.1T
4K
70.8
77.3
68.4
74.9
46.1
MPT
7B
IT
4K
26.8
75.0
76.4
80.6
48.5
68.3
70.2
42.6
20.8
50.4
Mistral
7B
16K
60.1
83.2
81.3
82.2
47.0
74.2
80.0
54.9
23.2
62.5
Gemma
8B
6T
8K
64.3
83.2
81.2
81.2
51.8
72.3
81.5
53,2
23.0
63.4
LLAMA2
13B
21
4K
54.8
81.7
80.7
80.5
50.3
72.8
71.3
49.4
31.2
65.1
LLAMA2
2T
4K
45.3
77.4
25.7
MEGALODON
7B
2T
32K
49.8
80.5
77.5
80.1
49.6
71.4
79.8
53.1
25.7
60.5
We introduce MEGALODON, an improved MEGA architecture (Ma al., 2023), which harnesses the gated attention mechanism with the classical exponential moving average (EMA) (Hunter, 1986) approach ($2). further improve the capability and efficiency of MEGALODON on large-scale long- context pretraining. we propose multiple novel technical components. First, MEGAL.ODOS introduces complex exponential moving averuge (CEMA) component, which extends the multi-dimensional damped EMA in MEGA to the complex domain ($3.1). Then, MEGALODON proposes the timestep normalization layer, which generalizes the group normalization layer (Wu and He, 2018) to auto- regressive sequence modeling tasks allow normalization along the sequential dimension ($3.2). To improve large-scale pretraining stability, MEGALODON further proposes normalized attention ($3.3), together with pre-norm with two-hop residual configuration by modifying the iidely-adopted pre- and post-normalization methods ($3.4). By simply chunking input sequences into fixed blocks, as is done in MEGA-Chunk (Ma et al., 2023), MEGALODON achieves linear computational and memory complexity in both model training and inference.
Empirically, we demonstrate the potential of MEGALOODNN as general architecture for modeling long sequences, by evaluating its performance across multiple scales of language modeling, as well as downstream domain-specific tasks. Through direct comparison with LLAMA2, while controlling for data and compute, MEGALODON-77 significantly outperforms the state-of-the-art variant of Transformer used to train LLAMA2-7B (Touvron al., 2023) on both training perplexity (Figure 1) and across downstream benchmarks (Table 1). Evaluation on long-context modeling, including perplexity in various context lengths up to 2M and long-context context QA tasks in Scrolls (Parisotto et al., 2020) prove MEGALODON'S 's ability to model sequences of unlimited length. Additional experimental results on small/mediun-scale benchmarks, including IRAA (Tay al., 2021), ImageNet (Deng et al., 2009), Speech Commands (Warden, 2018), WikiText- (Merity et al., 2017) and PG19 (Rae et al., 2019), demonstrate the robust improvements of MEGALODON across scales and modalities.

---

2 Background: Moving Average Equipped Gated Attention (MEGA)
In this section, we setup notations, briefly review the key components in the MEGA architecture (Ma et al., 2023), and discuss the existing problems in MEGA.
Following the notations in MEGA, (2,, E Rnxd to denote the input and output sequences with length and assume the representations of the input and output sequences have the same dimension d.
X
Raxd and Y
2.1 Multi-dimennional Damped EMA
MEGA embeds an EMA component into the calculation of the attention matrix to incorporate inductive biases across the timestep dimension. Concretely, the multi-dimensional damped EMA first expands each dimension of the input sequence X individually into h dimensions via an expansion matrix 3 æd then applies damped EMA the h-dimensional hidden space. Formally, for each dimension j E 21,2..... d}:
h aj 3 u (1) + aj o o
o 0
Yt, = n h
where u E R is the expanded h-dimensional vector for the j-th dimension at timestep t. a E (0, 1)dx 8 e (0, 1)axh are the decaying and damping factors, respectively. In E RRA is the EMA hidden state for the j-th dimension at timestep 7 Iddx is the projection matrix to map the I--dimensional hidden state back to 1-dimensional output Yr, E R.
2.2 Moving Average Equipped Gated Attention
In the gated attention mechanism in MEGA, the output from EMA (1) is used to compute the shared representation (Hua et al., 2022), because it encodes contextual information through EMA. Subsequently, MEGA introduces the reset gate, the update gate, and computes the candidate activation with the update gate and the residual connection. The technical details are provided in Appendix A.
2.3 Existing Problems in MEGA
To reduce the quadratic complexity in the full attention mechanism, MEGA simply split the sequences of queries, keys and values in (14-16) into chunks of length c. The attention in (17) is individually applied to each chunk, yielding linear complexity O(ke²) = O(c).. Technically, the EMA sub-layer in MEGA helps capture local contextual information near each token, mitigating the problem of losing contextual information beyond chunk boundaries in the thunk-wise attention.
Despite the impressive successes of MEGA, it still suffers its own problems: i) the performance of MEGA with chunk-wise attention still fails behind the one with full attention, due to the limited expressiveness of the EMA sub-layer in MEGA. ii) for different tasks and/or data types, there architectural divergences in the final MEGA architectures. For example, different normalization layers, normalization patterns (pre-norm applied to different data types (see Ma et al. (2023) for details). iii) There are no empirical evidences showing that MEGA is scalable for large-scale pretraining.
post-norm) and attention functions
(f(-) in (17))
3 MEGALODON
To address the aforementioned problems of MEGA, in this section we describe the novel technical advancements of MEGALODON.
3.1 CEMA: Extending Multi-dimensional Damped EMA to Complex Domain
As discussed Ma et (2023), the EMA component can be regarded as a simplified state space model with diagonal state matrix. Directly inspired from Gu et al. (2022b), almost all matrices

---

feature
feature
feature
timestep
timestep
timestep
(a) Layer Norm
(b) Group Norm
(c) Timestep Norm
Figure 2: Normalization methods. The elements in blue or pink are the regions to compute means and variances. We omit the batch dimension for simplicity.
diagonalize over the complex plane, straight-forwars idea to improve EMA capability is to extend to work over the complex number system C. We propose the complex exponential moving average (CEMA), which re- re-writes Eq. (1):
h = a, (coss 0g + 1,) @ uf + (1 - aa Lyy& (coss 03 + iisn 0,) O ohh h
Yt.j Re(n, h()))
where a E Rdxh are the real number parameters same as in EMA. Different from EMA, 17 E d in CEMA are complex numbers. 0, E R* 21, d} the h arguments. To uniformly
space the h arguments over the period 2T,, we parameterize 03 as:
2 Wj,, Vk [1,2.
0g
where the learnable parameter E Rd depicts the d base angles. By decaying the absolute value of each hr CEMA preserves the decaying structure in kernel weights, which is key principle to success of convolutional models on long sequence modeling (Li et al., 2023c).
3.2 Timestep Normalization
Despite the impressive performance of Layer Normalization combined with Transformer, it obvious that layer normalization cannot directly reduce the internal covariate shift along the spatial dimension (a.k.a timestep sequential dimension) (Ioffe and Szegedy, 2015). Group Normalization (Wu and He, 2018) normalizes hidden representations both along the timestep dimension and subset of the feature dimension, which has obtained improvements over Layer Normalization on a range of computer vision tasks. However, it cannot be directly applied to Transformer on auto-regressive sequence modeling, due to the leakage of future information via the mean and variance across the timestep dimension.
In MEGALODON, we extend Group Normalization the auto-regressive case by computing the cumulative mean and variance. Formally, suppose an input sequence X {x1, {2,,.. Xn E Rxd
and & groups along the feature dimension with dg d/k elements per group. Then, the mean and variance of the first group at timestep t E 21,2.... n are:
1
#e x dg
Di.jo
of
ETi 2
t* dg
Figure illustrates Layer Normalization and Timestep Normalization. efficiently and precisely calculate the cumulative mean and variance in each timestep, provide hardware hardware-friendly friendly imple- mentation on modern hardware (GPU) (see Appendix B.
3.3 Normalized Attention in MEGALODON
Previous studies have investigated the saturation and instability issues in the original scaled dot- product attention (17). A number of novel techniques have emerged to modify the sealed dot-product attention, among which normalized attention mechanisms, such as (scaled-) cosine attention (Luo al., 2018 8; Liu et al., 2022) and QK-normalization (Henry et al., 2020), have stood out for the simplicity and effectiveness.

---

Output Y
FFN
FFN
Layer Norm
Layer Norm
Gate
Normalized Attention Unit
Attention
Attention
toRk V
CEMA output x'
Layer Norm
Timestep Norm
Layer input X
(a) Sketch Megalodon layer.
(b) Pre-Norm
(c) Two-hop Residual
Figure 3: Illustration of the MEGALODON architecture. Figure (a) shows a sketch of one MEGA- LODON layer. Figure (b) and (c) display the configurations of pre-norm and pre-norm with two-hop residual, respectively.
Directly inspired from these normalized attention mechanisms, we propose the normalized attention mechanism specifically defined for MEGA to improve its stability. Formally,
X CEMA(X)
Ruxd
(5)
Z |Z|
Z X'W, + b2s
Z'=
E R
(6)
Q = Kg @ Z + Ha
E R
(7)
K = K6 O Z' + Pt
E R
(8)
where Q and K are computed by using the normalized shared representation Z' instead of Z. Note that we remove the SiLU (Ramachandran et al., 2017) activation function dsiu in (13), because the normalization Z has incorporated non-linearity into Z'. Then the attention operation in (17) has been changed to:
= Ssoftmax (QK*) V
R
As we use learnable Kg. N2 in (7) and (8), we can remove the scaled term 7(X). In addition, found that with the normalized attention, the softmax function fsottmax obtains the best or at least comparable performance on different tasks and data modalities (see Appendix C). Hence, throughout this paper we use softmax as the default attention function.
3.4 Pre-Norm with Two-hop Residual
Normalization configurations normalization (Xiong et al.. 2020) has become the default normalization configuration because of its better convergence properties than post-normalization in the original Transformer architec- ture (Vaswani et al., 2017). However, extensive studies have investigated the instability issue of mre-oormalization when scaling up model size (Davis et al., 2021; Liu et al., 2022). Formally, Transformer- block in pre- noromalizzaio can be formulated as (shown in Figure 3 (b)):
crucial in stably training deep architectures, and pre-
Y Attention(Normm())) + X
Y FFN(Norm(Y)) + y
(10) where the output Y is the sum of the input X and the output of each component in one block. Hence, the range and/or variance of Y keeps increasing for deeper blocks, causing the instability issue. In
-FFN(Norm (Ý)) + Attention(Norm(XX
X
5

---

the original MEGA architecture, the update gate &p (19) is used for gated residual connection (21) to mitigate this problem (Parisotto et al., 2020; Xu et al., 2020). However, the update gate yp introduces more model parameters and the instability issue still exists when scaling up model size to 7 billion.
MEGALOODON introduces new configuration named pre-norm with two-hop residual, which simply re-arranges the residual connections in each block (shown in Figure 3 (c):
Y Attention(Norm( ()) + X
Y = FFN(Norm(Y)) + X
(11)
where the input X is reused as the residual connection of the FFN layer. Since Y is directly followed by a normalization layer, we remove the update gate &p and use standard residual connection. The graphical architecture of a MEGALODON sub-layer visualized in Figure 3 (a). Note that the Timestep Normalization is only applied before the attention layer. Before the FFN layer, we still use Layer Normalization. The reasons two-fold: Layer Normalization is faster than Timestep Normalization; ii) the output vector of each token from the attention layer is a mixture vectors from contextual tokens via attention weights. Hence, normalizing the attention output along the feature dimension is similar to indirectly normalize along the timestep dimension.
3.5 4-Dimensional Parallelism in Distributed LLM Pretraining
Efficient distributed training algorithm is essential to train a large-scale language model, and several parallelization mechanisms have been introduced. The three most commonly used parallelism strategies are data, tensor (Shoeybi et al., 2019) and pipeline parallelism (Huang et al., 2019). However, the 3-dimensional parallelism still insufficient to scale up the context length of L.LMs (Li et al., 2023b; Liu et al., 2024).
Benefiting from the chunk- attention in MEGALODDN,, new timestep/sqquence dimension, which is orthogonal to all the aforementioned three parallelism dimensions. In MEGALODON, the only communications between devices in one chunk-parallel group are the last hidden state of CEMA and the cumulative mean and variance of Timestep Normalization in each block. Using asynchronous communication, we can minimize the overhead of chunk parallelization by hiding the communication costs in the computation of other components inside the same block and/or other blocks.
efficiently parallelize it along the
4 Experiments
To evaluate the scalability and efficiency of MEGALODON on long-context sequence modeling, we scale up MEGALODON to 7-billion model size and apply it to large-scale language model pretraining on 2 trillion tokens. We also conduct experiments on small/mediun--ee scale sequence modeling benchmarks, including Long Range Arena (LRA) (Tay et al., 2021), raw speech classification on Speech Commands (Warden, 2018), image classification language-moded modellngg on WikiText-103 (Merity et al., 2017) and PG19 (Rae al., 2019). Empirically, MEGALOODON significantly outperforms all the state-of-the-arr baseline models on these tasks across various data modalities.
ImageNet-1K K (Deng et al., 2009), and
4.1 LLM Pretraining
Architectural Details In our MEGALODON-7B model, we adopt most of architectural hyperparam- eters from LLAMA2-7B to ensure fair comparison: MEGALODON-7B consists of 32 blocks, with feature dimension d = 4096. Following LLAMA2, we use the SwiGLU activation function (Shazeer, 2020) in the feed-forward layer, and rotary positional embedding (RoPE, Su et al. (2021)). We set the attention chunk size c = 4096, which is the same as the pretraining context length in LLAMA2. Benefiting from the attention gate (7 in (18)), we use much smaller number of attention heads h 4 in BEGGALODON-P7 comparing to / = 32 in LLAMA2-7B. In addition, we apply pre-norm with two-hop residual ($3.4), using Timestep Normalization ($3.2) and Layer Normalization (Ba et al., 2016), while LLAMA2 models apply pre- rre-nomaaiization with RMSNorm (Zhang and Sennrich, 2019).
Some results are provided in Appendix C. due to space limits.

---

Data and Pretraining Details We use the same mix of publicly available data from LLAMA2. ensuring that the model are trained on exactly the same 2-trillion tokens. We also use the same tokenizer LLAMA2, whose vocabulary size is 32K.
We trained MEGALODON-77B using the AdamW optimizer (Loshchilov and Hutter, 2019), with 31 = 0.9, 32 = 0.95, = le- 8. The learning rate is 3.5e -4 4 and cosine learning rate schedule is applied with warmup of 2500 steps. We use weight decay of 0.1 and gradient clipping of 1.0, and no dropout is applied during training. The context length in pretraining is 32K (4 attention chunks). The global batch size 4M tokens, and is distributed on 256 NVIDIA A100 GPUs (16K tokens per A100). We set data parallel size to 128, chunk parallel size to 2 and tensor parallel size to 1.
Data and Computation Efficiency We evaluate the efficiency MEGAL.ODON w.r.t both the data and computation perspectives. For data efficiency, display the negative log- log-likelihood (NLL) for MEGALODON-7B, LLAMA2-7B and LLAMA2-13B w.r.t processed tokens during training in Figure 1. MEGALODON-7B obtains significantly better (lower) NLL than LLAMA2-7B under the same amount of training tokens, demonstrating better data efficiency. Moreover, MEGALODON suffers less training spikes than the Transformer- mrnnsformer-basedd architecture LLAMA2. Note that at the first 1/4 the pretraining process (< 500B tokens), the NLL of MEGALODON- is slightly worse than LLAMA2-7B. BB. We found that the main reason is that increased the base of RoPE from 10, 000 in LLAMA2 to 100, 000 in MEGALODOM, which slows down model convergence at the beginning of the pretraining process. At the end, MEGAL.ODO reaches training loss of 1. 1.70, landing mid-way between LLAMAA-7B (1.75) and LLAMA2-133 13B (1.67).
For computation efficiency, conduct exper- iments of running LLAMA2-7B and MEGA- LODON-7B using the same amount of compu- tational resources and comparing their training speed under various context lengths. Specif- ically, we execute each experiment to train a model with global batch size 4M tokens dis- tributed on 256 NVIDIA A100 GPUs (16K to- kens per A100) and calculate the word/token per second (WPS) to measure the training speed. Figure illustrates the average WPS per device of LLAMA2-7B and MEGALODON-7B using 4K and 32K context lengths, respectively. For LLAMA2 models, we accelerate the computation of full attention with Flash-Attention V2 (Dao, 2024). Under 4K context length, MEGALODON-77B is slightly slower (about 6%) than LLAMA2-7B, due to the introduction of CEMA and Timestep Normalization. When we scale up context length to 32K, MEGAI.ODON-7B is significantly faster (about 32%) than LLAMA2-7B. demonstrating the com- putation efficiency of MEGALODON for long-context pretraining. addition, MEGALODON-7BB. which utilizes chunk parallelism ($3.5), achieves about 94% utilization of MEGALODON-77B44K.
x x1.40
Tokens Por Socord
X 1.0
Liema2-7B
Megalodon-7B 4K
Megalodon-78 32K
Llame2-7B 32K
Figure 4: Average WPS per device.
4.2 Short-Context Evaluation on Academic Benchmarks
We compare MEGALODON-TB
LLAMA2 models on standard academic benchmarks with short contexts (< 4K tokens), closely following the settings in LLAMM2 (Touvron et al., 2023). The benchmarks are grouped into the categories listed below:
Commonsense Reasoning (0-shot): HellaSwag (Zellers al., 2019), PIQA (Bisk al., 2020), SIQA (Sap et al., 2019), WinoGrande (Sakaguchi et al., 2021), ARC-e and -c (Clark et al., 2018).
World Knowledge (5-shot): NaturalQuestions (NQ, Kwiatkowski et al. (2019)) and TriviaQA (TQA, Joshi et al. (2017)).
Reading Comprehension (0-shot): BoolQ (Clark et al., 2019).
Popular aggregated results (5-shot): MMLU (Hendrycks et al., 2020).
Table summarizes the results of MEGAL
and LLAMA2 on these academic benchmarks, together with other open-source base models, including MPT (MosaicMI. 2023), RWKV (Peng al., 2023), Mamba (Gu and Dao, 2023), Mistral (Jiang et al., 2023) and Gemma (Mesnard et al., 2024). Pretrained on the same 2T tokens, MEGALODON-7B surpasses LLAMA2-7B across all the

---

Model
NaQA
Qasper
QMSum
Xgen
17.4
20.5
6.8
MPT
18.8
24.7
8.8
Valid PPL
Yarn
20.9
26.2
11.4
LLAMA2
18.8
19.8
10.1
LLAMA2-L'
23.5
28.3
14.5
MEGALODON
23.9
28.0
13.1
16K
128K 256K
512K
Table 2: Results on Scrolls. LLLAMA--L (Xiong al., 2023) continually trains LLAMA22 on 500B tokens for length extension.
Context Length
Figure PPL in various context lengths.
benchmarks, On some tasks, MEGALODON-7B achieves comparable or even better performance with ALAMA2-13B. Note that Mistral-7B 7B and Gemma-88 8B were pretrained much larger datasets than MEGALODON-7B, hence the results are not directly comparable.
4.3 Long-Context Evaluation
Perplexity over Long Sequences To demonstrate the capability of MEGALODON to make of very long contexts to improve next-token prediction, we start by conducting the evaluation of different context lengths. Concretely,
valid perplexity consists of 1,920 selected books. Each of these books contains sequences with at least 2M tokens. The validation dataset is constructed by first randomly shuffling all the files and then concatenating them. Figure 5 shows the perplexity (PPL) of the validation dataset in various context lengths ranging from 4K to 2M. We observe that the PPL decreases monotonically with context length, validating the effectivenss and robustness of MEGALODON on modeling extremely long sequences.
construct validation dataset which
Long-Context QA tasks in Scrolls Next, we evaluate MEGALODON on long-context open-book question answering (QA) tasks in the Scrolls dataset (Shaham et al., 2022), including Narra- tiveQA (Kočiský al., 2018), Qasper (Dasigi et al., 2021) and QMSum (Zhong et al., 2021). Following Xiong et al. (2023), we use simple prompt {CONTEXT] Q: {QUESTION] A: for all tasks, and evaluate 0-shot F1-score on NarrativeQA, 2-shot Fl-score on Qasper and 1-shot geometric-ROUGE QMSum. Table lists the results of MEGALODON-7BB. together with other open-source long-context models in the scale 7B, namely Xgen- 7B-8K (Nijkamp et al., 2023), MPT-7B-84 8K (MosaicML, 2023), YaRN-7B 128k (Peng al., 2024), LLAMA2-7B-4E (Touvron et al., 2023) and LLAMA2-7B-302 (LLAMA2-L. Xiong et al. (2023)). MEGALODON-7B obtains the best FI NarrativeQA, and competitive results with LLAMA2-7B Long. It should be noticed that LLAMA2-7B Long extends the context length of LLAMA2-7B from 4K to 32K by continually pretraining it additional 500B tokens from long-context data.
4.4 Instruction Finetuning
To evaluation the generalization capability of MEGA- LODON instruction following and alignment, We finetune the base model of MEGALODONN-P on a proprietary nnstruction-alignmer trolled setting. We did not apply any RLHF tech- niques further finetune it. Table 3 summarizes the performance of chat models in 7B scale on MT-Bench' MEGALODON exhibits superior per- formance on MT-Bench compared to Vicuna (Chi- ang al., 2023), and comparable performance LLAMA2-Chat, which utilizes RLHF for further alignment finetuning, We present some outputs from instruction finetuned MEGALODON Appendix D.
Table 3: MT Bench. Comparison of Chat models. LLAMA2-Chat utilizes RLHE
data under
Model
Size
MT-Bench
Vicuna
7B
6.17
LLAMA2-Chat
7B
6.27
Mistral-Instruct
7B
6.84
MEGALLODON
7B
6.27
'Geometric mean of ROUGE-1, 2 and http://ll..aa.glossssssss.aa http://ttteeeeeeeeee bench-eval

---

Table 4: (ImageNet-IK) Top-1 accuracy.
Table 5: (PG-19) Word-level perplexity.
Model
#Param.
Acc.
Model
# #Param.
Val
Test
ResNet-152
60M
78.3
Compressive Trans.
43,4
33.6
ViT-B
86M
77.9
Perceiver AR
975M
45.9
28.9
DeiT-B
86M
81.8
Block-Recurrent Trans.
1.3B
26.5
MEGA
90M
82.3
MEGAB YTE
1.3B
42.8
36.4
MEGALODON
90M
83.1
MEGALODON
1.3B
29.5
25.4
4.5
Evaluation on Medium-Scale Benchmarks
ImageNet Classification To evaluate MEGALODON on image classification task, we conduct experiments the Imagenet-1K (Deng al., 2009) dataset, which consists of 1.28M training images and 50K validation images from 1000 classes. We mostly follow DeiT's approach of applying several data augmentation and regularization methods that facilitate the training process, and adopt most the hyperparameters from Ma et al. (2023). For classification task, we replace the timestep normalization with the standard group normalization method. Top-I accuracy the validation is reported in Table 4 to assess various models. MEGALODON obtains about 1.3% accuracy improvement over DeiT-B (Touvron al., 2021), and 0.8%. improvement over MEGA (Ma et al., 2023).
Auto-regressive Language Modeling on PG-19 We also evaluate MEGALODON on auto-regressive language modeling on the medium-scale PG19 (Rae et al., 2019) datasets. We use the same vocabu- lary from Block-Recureed Transformer (Hutchins al., 2022) and adopt most of its hyper-parameters to train MEGALODON model with 1.3B parameters. Table 5 illustrate the word-level perplexity (PPL) MEGALODON on PG-19, together with previous state-of-the-art models, including Com- pressive Transformer (Rae et al., 2020), Perceiver AR (Hawthorne et al., 2022), Transformer (Hutchins et al., 2022) and MRGABYTE (Yu et al., 2024). MEGALLODON significantly outperforms all the baselines.
Conclusion
We have introduced MEGALODON, an improved MEGA architecture with multiple novel technical components, including complex exponential moving average (CEMA), the timestep normalization layer, normalized attention and pre-norm with two-hop residual configuration, improve its ca- pability, efficiency and scalability. Through direct comparison with LLAMA2, MEGALODON achieves impressive improvements on both training perplexity and across downstream benchmarks. Importantly, experimental results on long-contex, modeling demonstrate MEGALLODON' ability to model sequences of unlimited length. Additional experiments on small/medium--saa benchmarks across different data modalities illustrate the robust improvements of MEGALODON, which lead to potential direction of future work to apply MEGALODON for large-scale multi molti-modality pretraining.
Acknowledgments
We thank Sadhika Malladi, Zihao Ye, Dacheng Li and Rulin Shao for their helpful feedback and discussion during this work.
References
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.064500 2016.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2018.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense natural language. Proceedings of the AAAI conference on artificial intelligence, pages 7432-7439, 2020.

---

Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, lon Stoica, and Eric P. Xing. Vicuna: An open-source chathot impressing gpt-4 with 90%* chatgpt quality, March 2023. URL https:
//Imsys. http:////o//////2220-----------0
Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint
arXiv:1905. 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint
1ew::04435055 2018.
James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90))297-301 1965.
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G Carbonell, Quoc Le. and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978-2988, 2019.
Tri Dao. Flashattention-2: 2: Faster attention with better parallelism and work partitioning. In Interna- tional Conference on Learning Representations (ICIR-2024)) 2024.
Pradeep Dasigi, Kyle Lo. Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. A dataset of informattin-sseeking questions and answers anchored research papers. Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies (NAACL-2021), pages 4599-4610, Online, June 2021. Association for Computational Linguisties.
Jared Q Davis, Albert Gu, Krzysztof Choromanski, Tri Dao, Christopher Re, Chelsea Finn, and Percy Liang. Catformer: Designing stable transformers via sensitivity analysis. In International Conference on Machine Learning, pages 2489-2499. PMLR, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. leee, 2009.
Daniel Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Re. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations (ICLR-2023) 2023.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations (ICI.-20022, 2022a.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. arXiv preprint 22X1v:2206.1189, 2022b.
Stephen Hanson and Lorien Pratt. Comparing biases for minimal network construction with back- propagation. Advances in neural information processing systems, 1, 1988.
Curtis Hawthorne, Andrew Jaegle, Cătälina Cangea, Sebastian Borgeaud, Charlie Nash, Mateusz Malinowski, Sander Dieleman, Oriol Vinyals, Matthew Botvinick, Ian Simon, et al. General- purpose, long-context autoregressive modeling with perceiver ar. In International Conference on Machine Learning, pages 8535-8558. PMLR, 2022.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazcika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300, 2020.
Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen. Query-key normalization for transformers. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 4246-4253, 2020.
10

---

Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoe Le. Transformer quality in linear time. In International Conference on Machine Learning (ICML-2022), pages 9099-9117. PMLR, 2022.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le. Yonghui Wu, and zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
J Stuart Hunter. The exponentially weighted moving average. Journal of quality technology, 18(4): 203-210, 1986.
DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur, Block-recurrent transformers. Advances neural information processing systems, 35:33248-333 2022.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by
reducing internal covariate shift. In International conference on machine learning (ICML-2015), pages 448-456. pmlr, 2015.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310006825, 2023.
Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers). Association for Computational Linguistics, 2017.
William Kahan. Pracniques: further remarks on reducing truncation errors. Communications of the ACM, 8(1):00, 1965.
Tomáš Kočiský. Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The NarrativeQA reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317-328, 2018.
Alex Krizhevsky et Learning multiple layers of features from tiny images. Technical Report. University of Toronto, 2009.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
Bonan Li, Yinhan Hu, Xuecheng Nie, Congying Han, Xiangjian Jiang, Tiande Guo, and Luoqi Liu. Dropkey for vision transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), pages 22700-22709, June 2023a.
Dacheng Li, Rulin Shao, Anze Xie, Eric P Xing. Joseph E Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang. Lightseq:: Sequence level parallelism for distributed training of long context transformers. In Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ NeurlPS 2023), 2023b.
Yuhong Li, Tianle Cai, Yi Zhang. Deming Chen, and Debadeepta Dey. What makes convolutional models great on long sequence modeling? In International Conference on Learning Representa- fions (ICLR-20233, 2023c.
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learning long-range spatial dependencies with horizontal gated recurrent units. In S. Bengio, H. Wal- lach, H. Larochelle, K. Grauman,
Tesa-Bianchi, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near- infinite context. In International Conference Learning Representations (ICLR-2024) 2024.
11

---

Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the IEEF/CVE conference on computer vision and pattern recognition, pages 12009-12019, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Reprexentations, 2019.
Chunjie Luo, Jianfeng Zhan, Xiaohe Xue, Lei Wang, Rui Ren, and Qiang Yang. Cosine normalization: Using cosine similarity instead of dot product in neural networks. In 27th International Conference on Artificial Neural Networks (ICANN-2018), pages 382-391. Springer, 2018.
Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettle- moyer. Luna: Linear unified nested attention. Advances in Neural Information Processing Systems, 34:2441-2453, 2021.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In The Eleventh International Conference on Learning Representations. 2023.
Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pages 142-150, 2011.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations (ICLR-2017), 2017.
Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro- Ambrose Slone, Amélie Héliou, Andrea Tacchetti, Anna Bulanova, Antonia Paterson, Beth Tsai, Bobak Shahriari, Char- line Le Lan, Christopher A. Choquette-Choe Reid, Elena Buchatskaya, Eric Ni, Eric Noland, Geng Yan, George Tucker, George-Christian Mu- raru, Grigory Rozhdestvenskiy, Henryk Michalewski, Ian Tenney, Ivan Grishchenko, Jacob James Keeling, Jane Labanowski, Jean- -aptisst Lespiau, Jeff Stanway, Jenny Brennan, Jeremy Chen, Johan Ferret, Justin Chiu, Justin Mao-Jones, Katherine Lee, Kathy Yu, Katie Millican, Lars Lowe Sjoesund, Lisa Lec, Lucas Dixon, Machel Reid, Maciej Mikula, Mateo Wirth, Michael Sharman, Nikolai Chinaev, Nithum Thain, Olivier Bachem, Oscar Chang, Oscar Wahltinez, Paige Bailey, Paul Michel, Petko Yotov, Pier Giuseppe Sessa, Rahma Chaabouni, Ramona COmanescu, Reena Jana, Rohan Anil, Ross Mcllroy, Ruibo Liu, Ryan Mullins, Samuel L Smith, Sebastian Borgeaud, Sertan Girgin, Sholto Douglas. Shree Pandya, Siamak Shakeri, Soham De. Ted Kli- menko, Tom Hennigan, Vlad Feinberg, Wojeiech Stokowiec, Yu hui Chen, Zafarali Ahmed, Zhitao Gong, Tris Warkentin, Ludovic Peran, Minh Giang, Clément Farabet, Oriol Vinyals, Jeff Dean, Koray Kavukeuoglu, Demis Hassabis, Zoubin Ghahramani, Douglas Eck, Joelle Barral, Fernando Pereira, Eli Collins, Armand Joulin, Noah Fiedel, Evan Senter, Alek Andreev, and Kathleen Kenealy. Gemma: Open models based on gemini research and technology. 2024.
Clément Crepy, Daniel Cer, Daphne Ippolito. David
Austin,
MosaicML. Introducing mp:-7b: A new standard for open-source, commercially usable Ilms, 2023.
Nikita Nangia and Samuel Bowman. Listops: A diagnostic dataset for latent tree learning. In Proceed- ings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Student Research Workshop, pages 92 22-99, 2018
Erik Nijkamp, Tian Xie, Hiroaki Hayashi, Bo Pang, Congying Xia, Chen Xing, Jesse Vig, Semih Yavuz, Philippe Laban, Ben Krause, Senthil Purushwalkam, Tong Niu, Wojeiech Kryšciñski, Lidiya Murakhovs's ka, Prafulla Kumar Choubey, Alex Fabbri, Ye Liu. Rui Meng, Lifu Tu, Meghana Bhat, Chien-Sheng Wu, Silvio Savarese, Yingbo Zhou, Shafiq Joty, and Caiming Xiong. Xgen-7b technical report, 2023.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. International conference on machine learning, pages 7487-7498.
PMLR, 2020.
12

---

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style. nigh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. Rwkv: Reinventing rmns for the transformer era. arXiv preprint arXiv:2305. 13048, 2023.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN: Efficient context window extension large language models. In International Conference on Learning Representations (ICLR-2024). 2024.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International conference on machine learning (ICML-2023). PMLR, 2023.
Dragomir R Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 47(4):919-944, 2013.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL 05507.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modeling. International Conference on Learning Representations (ICLR-2020), 2020.
Prajit Ramachandran, Barret Zoph, and Quoc V Le. Swish: self-gated activation function. arXiv preprint arXiv:1710.0544 7(1):5,2011
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106,
2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions. arXiv preprint arXiv:1904.0m 09728, 2019.
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: Standardized CompaRison over long language sequences. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing (EMNL.-2022), pages 12007-12021, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05200 2020.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv: 1909:019000005: 2019.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864, 2021.
Yi Tay, Mostafa Dehghani. Dara Bahri, and Donald Metzler. 06732, 2020.
Efficient transformers: A survey. arXiv
preprint
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena A benchmark for efficient transformers. International Conference on Learning Representations, 2021. URL https: iereenrreieen.nn/ffoummmm did-///eoo grC2k.
Yi Tay, Mostafa Dehghani, Samira Abnar, Hyung Won Chung, William Fedus, Jinfeng Rao, Sharan Narang, Vinh Q Tran, Dani Yogatama, and Donald Metzler. Scaling laws vs model architectures: How does inductive bias influence scaling?, 2022.
13

---

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers distillation through attention. International Conference on Machine Learning, pages 10347-10357. PMLR, 2021.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXhv:2307.09288 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, K ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Xindi Wang. Mahsa Salmani, Parsa Omidi, Xiangyu Ren, Mehdi Rezagholizadeh, and Armaghan Eshaghi. Beyond the limits: A survey of techniques to extend the context length in large language models, 2024.
Pete Warden. Speech commands: dataset for limited-vocabulary speech recognition. arXiv preprint arXiv:1804.03209, 2018.
B. P. Welford. Note on method for calculating corrected sums of squares and products. Technomer- rics, 4(3):419-420, 1962.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV-2018), pages 3-19, 2018.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang. Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533, PMLR, 2020.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint ar arXiv:2309. X6039,,9..6639, 2023.
Hongfei Xu, Qiuhui Liu, Deyi Xiong, and Josef van Genabith. Transformer with depth-wise Istm. arXiv preprint ar arXiv:2007.06257, 2020.
Lili Yu, Dániel Simig, Colin Flaherty, Armen Aghajanyan, Luke Zettlemoyer, and Mike Lewis. Megabyte: Predicting million-byte sequences with multiscale transformers. Advances in Neural Information Processing Systems, 36, 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACI--2019). Association for Computational Linguistics, 2019.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan Awadallah, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, and Dragomir Radev. QMSum: A new benchmark for query-based multi-domain meeting summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACI-2021), pages 5905-5921, Online, June 2021. Association for Computational Linguistics.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but not robustly, 2024.
14

---

Appendix: MEGALODON: Efficient Long-Context LLM Pretraining and Inference with Unlimited Context Length
A Background: Moving Average Equipped Gated Attention
In the gated attention mechanism in MEGA, the output from EMA (1) used compute the shared representation (Hua et al., 2022) Z:
X' = (M)
E R2nd
(12)
Z Pallu (X'W, + ba)
E R
(13)
where X' can be regarded as the updated or contextual input, because it encodes contextual infor- mation through EMA. Then, the query and key sequences are computed by applying per-dimension scalars and offsets to Z, and the value sequence is from the original X:
Q Kig 3 + P
(14)
K KR O Z + μ4
R=x=
(15)
V Ositu W,, + bb)
R"xv
(16)
where Kig. M4. Kk,, f4 IR are the learnable scalars and offsets of queries and keys, respectively. is the expanded intermediate dimension for the value sequence. The output of attention is computed as follows:
QK" T(X)
0=
V
E R R
(17)
Subsequently, MEGA introduces the reset gate y, the update gate yp, and computes the candidate activation A and final output Y:
7 Iil (X' K'W, + by)
E R nX"
(18)
4p sigmoid (X'W, + bp)
E R nnd
(19)
H- dailu X W, + (7 @ 0)U, + bn)
E Rnd
(20)
Y o A + pp) X
E Raxd
(21)
with the update gate &p and the residual connection X.
B
Implementation Details
B.1
Efficient Fused CUDA Operators Implementation
Fused Attention We implemented fused attention operator to improve the efficiency, mainly by fusing the causal mask, softmax function and dropout operation (if necessary). The fused implementation reduces the IO costs from global memory for the attention weight. For attention dropout, we adopt the dropout-before-soff the dropout mask on the input attention matrix of the softmax function. Concretely, we fill the values of the attention matrix at dropout mask positions to -oo before feeding it into the softmax function. scheme comparing to the standard attention
scheme in DropKey (Li et al., 2023a), which applies
One important advantage of this cerouut-efore-sofffi dropout is that the computation of the gradients in back nuck-propagation is independent with the applied dropout mask.
Efficient FFT We also provide an efficient fused implementation of the FFTConv operator. Similar with the FlashConv in H3 (Fu et al., 2023), we fused the real number FFT (RFFT), its inverse (IRFFT) and implemented the Cooley-Tuker FFT algorithm (Cooley and Tukey, 1965) in the CUDA shared memory. Similar with the FlashConv in H3 (Fu et al., 2023), we fused the real number FFT (RFFT), its inverse (IRFFT) and the element wise multiplication, and implemented the Cooley-Tukey FFT algorithm (Cooley and Tukey, 1965) CUDA's shared memory. Our implementation is able to accommodate up 16K tokens in the limited shared memory of an A.00 GPU.
15

---

Table (Long Range Arena) Accuracy on the full suite of long range arena (LRA) tasks. Results of previous models are reported in Ma et al. (2023).
Models
ListOps
Text
Retrieval
Image
Pathfinder
Path-X
Avg.
Transformer
37.11
65.21
79.14
42.94
71.83
X
59.24 24
Reformer
37.27
56.10
53.40
38.07
68.50
X
50.67
Linformer
35.70
53.94
52.27
38.56
76.34
X
51.36
BigBird
36.05
64.02
59.29
40,83
74.87
X
55.01
Luna-256
37.98
65.78
79.56
47.86
78.55
X
61.95
S4
59.10
86.53
90.94
88,48
94.01
96.07
85.86
MEGA-chunk
58.76
90.19
90.97
85.80
94.41
93.81
85.66
MEGA
63.14
90.43
91.25
90.44
96.01
97.98
88.21
MEGALLDON-chunk
62.23
90.53
91.74
87.11
96.89
97.21
87.62
MEGALODON
63.79
90.48
91.76
89.42
98.13
98.17
88.63
Timestep Normalization For the TimestepNorm operator, we have efficient implementation to improve both its speed and numerical stability. To compute the cumulative mean and variance for each of the timesteps, our implementation distributed the threads in each CUDA block in both the timestep/sequence dimension and the feature dimension to balance the parallelism of the algorithm and the performance of the global memory access. To improve numerical stability, we used the Welford algorithm (Welford, 1962) to compute the cumulative mean and variance and the Kahan Summation (Kahan, 1965) to reduce the numerical error from summation.
B.2 Plus Reparameterization in Normalization Layers
In the normalization methods, two learnable parameters y and are introduced scale and shift the normalized value:
(22)
where μ and o² are the mean and variance of the input across the pre-defined dimensions. Initial- ization y and (9 is erucial for model performance and stability. The standard implementation of normalization layers, such as PyTorch (Paszke et al., 2019), initializes and / to vectors of ones and zeros, respectively, preserve the mean and variance of the normalized inputs at the beginning of training.
This standard implementation, however, suffers problem when weight decay regularization is applied to prevent overfitting (Hanson and Pratt, 1988). Technically, the weight decay regularization pushes the values of model parameters towards smaller magnitudes. In the context of normalization methods, weight decay pushes the values in y towards zero, which diverges from its initialization of one. This may prevent the model from learning the true scale of the data distribution, and may cause numerical stability issues as well.
To address this problem, used the plus reparameteriization" of the scale parameter
11
y = (7
(23)
where y is initialized to zero. Under weight decay, 7 remains centered around zero, resulting in desirable scale of around one.
C Experiments on Small-Scale Benchmarks
We conducted small-scale experiments on five benchmarks across various data modalities, including text, audio and image. To demonstrate the robustness of the MEGALODON architecture different tasks and data types, we used a single unified architecture with minimal architectural divergence in
SSiiilar idea in the blog: https://medium.cc com www.roooooooooooooooooooooe wwwer//nneeereerrriiiinnngoeeeee weight- -are-eeaaa--a-e-noooooaaiiaaai layer-normalization challenges-an and-a regarcamtteriztiioossotttiooooen
16

---

Table (SC-Raw) Accuracy.
Table 8: (WikiText-103) Word-level PPL.
Model
#Param.
Acc.
Model
#Param.
PPL
Transformer
786K
X
Transformer
247M
18.66
S4
300K
97.50
Transformer-
257M
18.30
MEGA
300K
96.92
S4
249M
20.95
MEGA (big)
476K
97.30
MEGA
252M
18.07
MEGALODON
300K
98.14
MEGALODON
252M
17.23
all the experiments; softmax attention function, rotary postional embedding, pre-norm with two-hop residual, and timestep Normalization (Group Normalization for classification). We adopt (almost) all the architectural and training hyperparameters from the corresponding experiments of the original MEGA (Ma al., 2023).
C.1 Long Range Arena (LRA)
Long Range Arena (LRA) benchmark (Tay et al., 2021) is designed for evaluating sequence models un- der the long-context scenario. They collect six tasks in this benchmark which ListOps (Nangia and Bowman, 2018), byte-level text classification (Text; Maas et al. (2011), byte-level document retrieval (Retrieval; Radev al. (2013)), image classification on sequences of pixels (Image: Krizhevsky et al. (2009)), Pathfinder (Linsley et al., 2018) and its extreme long version (Path-X; Tay et al. (2) These tasks consist input sequences ranging from IK to 16K tokens and span across variety of data types and modalities.
Table 6 compares MEGALODON against several baselines, including Transformer and its efficient variants, the state space model S4 (Gu et al., 2022a), and the original MEGA model. Following Ma al. (2023), also evaluate MEGALODON-chum for all the tasks, except Path-X where e = 4096. With chunk-wise attention, MEGALOODON-humm substantially outperforms MEGA-chunk all the six tasks. In addition, MEGALODON significantly narrows the gap between chunk-wise attention and full attention.
each task, by setting the chunk size 128
C.2 Raw Speech Classification
To evaluate the capability of MEGALODON the long-ran rongerange modeling of speech signals, apply MEGALODON to classify raw speech (with length 16000), rather than using traditional preprocessing (e.g. convert MFCC features). Following Ma al. (2023), perform speech classification on the SC10 subset of the Speech Commands dataset (Warden, 2018), and set attention chunk size c = 1000. As shown in Table 7, our MEGALODON model with 300K parameters achieves accuracy of 98. 14, which is significantly better than 97.50 from S4 and 96.92 from MEGA with similar model size.
C.3 Auto-regressive Language Modeling
We also evaluate MEGALODON auto-regressive language modeling the small-scale WikiText- 103 (Merity et al., 2017) datasets. Following Baevski and Auli (2018), we adopt adaptive softmax and input embeddings and use vocabulary of 260K tokens for Wiki'Text-10 -03x and train MEGALODON model with about 250M parameters with attention chunk size e 2048. Table 8 illustrate the word- level perplexity (PPL) of MEGALODON models, including Transformer (Baevski and Auli, 2018 8), Transformer-XL. (Dai et al., 2019), S4 (Gu al., 2022a) and MEGA (Ma et al., 2023). MEGALODON significantly outperforms all the baselines on both the two datasets.
WikiText-1001 103, together with previous state-of-the-ar
D Model Outputs from Instruction-finetuned MEGALODON
17

---

What the central dogma molecular biology? What processes are involved? Who named this?
The central dogma of molecular biology describes the flow of genetic information from DNA to RNA to proteins. The three processes involved in the central dogma are transcription, translation, and translocation.
The central dogma was first proposed by Francis Crick, James Watson, and Maurice Wilkins in the 1950s, and was later formalized by Crick in a 1958 paper titled "On the Central Dogma of Molecular Biology". The term "central dogmaa" coined by Crick describe the flow of genetic information from DNA to RNA to proteins.
Identify and fix one incorrect fact in your previous response.
The correct name of the person who proposed the central dogma of melecular biology is James Watson, not Maurice Wilkins, James Watson, along with Fruncis Crick, first progused the central dogma in the 1950s, and Maurice Wilkins was co recipient of the Nobel Prize in Physiology Medicine 1962 for his contributions the discovery the double helix structure of DNA.
will provide some mathematical equations or concepls, and -oy--o--nddrst terms. This could include providing step-by-step instructions for solving problem, demonstrating various techniques with examples everyday life or suggesting online resources for further study, My first request "I need help understanding how probability works."
Act a math teacher. I
your job to explain them
Probability is branch mathematics that deals with the likelihood event occurring. often expressed number between and I. where O represents an impossible event and I represents certain event.
example, if you were to roll fair six sided die, the probability of rolling would be 1/6, about 17%, This because there only one way for the die to land on I out of the six possible outcomes. To calculate probability, you can use formula:
Probability = Number of ways the event occur / Totta number of possible outcomes
For example, il you were to ask "What the probubility of lippiigg heads on a fair coin loss!", the answer would be 1/2, or 50%. This bocause there only one way the coin to land heads-up out of the possible outcomes (heads tails).
There are many different ways to interpret probubility, and it can be used in a variety of real real- world upplications such as gumbling. insurance, and statistics.
I hope this helps to give you basic understanding of probability. Let me know if you have any questions or need further clarification.
What are the differences between Riemannian geometry and euclidean geometry?
Riemannian geometry and Euclidean geometry are both branches of mathematics study properties of space, but they differ their assumptions and methods.
Faclidean geometry based
assumption that the space is flat flat Fuclidean space the simplest non-empty connected flat
metric space. It studies the properties of space using concepts such as distance, angle, and parallelism. It is named after the Groek mathematician Euclid, who developed the foundations of the field his book "Elements".
Riemarmian geumetry. the other hand, is based on the asumption that the space curved. It studies the properties of space using concepts such as distance, angle, and pamilelism, hut in a more general setting, and it is named after the German mathematician Remhard Riemann.
In Riemannian geometry, the concept of metric is used to measure the distance hetween two points in curved space. A metric is mathematical function that assigns a distance to each segment space. The most common metric Riemannian geometry is the distance function induced Riemannian metric. which is mathematical function that assigns distance to each segment in a space. In summary. Buclidean geometry is a special case of Riemanmian geometry, where the metric is the flat Buclidean metric. Riemanmian gevmetry
general field that studies the properties curved spaces, and used many other sciences.
mathematics, physies and
hope helps to give you basic understanding know if you have any questions or need further clarification.
differences between Riemannian geomelty and Euclidean geumetry. me
Evaluate the following movie reviews scale of I to 5, with I being very negative, 3 being neutral, and being very positive: 1. This movie released on Nov. 18, 2019, was phenomenal. The cinematography, the acting, the plot - everything was top-notch. 2. Never before have been disappointed with movic. The plot
predictable and the characters were oin-dimensi In my opinion, this movie is the worst one to have been released in 2022. 3. The movie was okay. There were some parts I enjoyed, but there were also parts that lackluster. This movie that was released 2018 and seems to quite ordinary. Return the answer as ISON array of integers.
[5, 1,
Update your previous reply by including the release date part of the JSON content.
('release_ date": "Nov 18, 2019*, "rating": 5), ("releaso_date": "2022", "rating": 1),
] ;ate"" "Feb 2018", "rating"; 3)
Figure 6: Model outputs from ciatructioo-iinetunee MEGALODON.
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: see the abstract
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: No theoretical results involved.
19

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: see appendix
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: appendix
Guidelines:
The answer NA means that paper does not include experiments requiring code.
20

---

Please see the NeurIPS code and data submission guidelines (https://nips.com httpa///guideee/eeuubbbiiiiaaae c for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips. cttpsscccppbbcccccccccc gwwide/Cooeeeeeee wwabbbmmssioooooooii.. for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if' applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: see appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
21

---

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: experiments and appendix Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /twprreeurppueeeee httppppllccccicccsiiiiieeecoeeea
Answer: [Yes]
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
22

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [Yes]
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi wwwswhhhh.oe...oo com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
23

---

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
24


Reviews:
[{"review_id": "7RjB2wmrvl", "reviewer": "Reviewer_U1i1", "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.", "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.", "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.\n\nThe anonymous link is not working.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Clear motivations: All three improvements directly target the limitations of Mega.\n2. The complex EMA is a novel approach.\n3. 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?\n\nEmpirically, it is not clear how CEMA improves expressiveness. It would be most direct to compare using CEMA versus using EMA on these tasks.\n\nThe reviewer understands that the rebuttal period is short and is therefore not requiring most of these experiments to be added.\n\n[1] Yang, Kai, et al. \"Do Efficient Transformers Really Save Computation?\" Forty-first International Conference on Machine Learning.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and constructive comments! We appreciate your positive feedback on the good motivation, novelty of Megalodon and its strong empirical results. We address your concerns and questions below and please let us know if you still have concerns after you read our response.\n> W1 & Q1: Have the authors conducted ablation studies for small-scale models before moving to 7B, similar to the results in Appendix C?\n\nThe ablation studies on small-scale benchmarks are in the [author rebuttal](https://openreview.net/forum?id=XlAbMZu4Bo&noteId=G72sD1BBxi).\n\n> W2:  Recent theoretical work [1] has shown that efficient versions of attention can limit the expressiveness of the model…\n\nWe conducted experiments to evaluate Megalodon on retrieval-oriented tasks, such as passkey retrieval. Similar to previous studies, due to chunk-wise attention, Megalodon under-performed on these retrieval-oriented tasks compared with full attention mechanism, but outperformed state-space models such as Mamba. \n\nFor example, without any fine-tuning for length extension, Megalodon completes the passkey retrieval task with up to 16K context length, while Mamba can only complete this task up to 4K context. Long-LlaMa2, which continually trains LlaMa2 on selected long-context data for length extension, is able to complete up to 32K context length. \n\nFurther improving Megalodon for retrieval-oriented tasks is an interesting and important direction for future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VtptzHPIr9", "reviewer": "Reviewer_73L7", "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::\n1) 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.\n2. achieves better pretraining efficiency and downstream task accuracy than Transformers, specifically in handling long sequences with 7 billion parameters and 2 trillion training tokens.\n3. outperforms Transformers across various benchmarks, including long-document comprehension, multi-turn conversation, and video generation.", "questions": "see Weaknesses", "limitations": "see Weaknesses", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The performance of the proposed architecture is excellent, as demonstrated in Table 1.\n2. Extensive experiments are conducted to evaluate the proposed methods.\n3. 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.\n2. Section 3.2 is difficult to understand due to the lack of background and intuitive explanations.\n3. 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.\n4. The code link provided in the abstract does not work.\n5. 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.\n\n\n[1] Sequence Parallelism: Making 4D Parallelism Possible\\\n[2] Lightseq: Sequence level parallelism for distributed training of long context transformers\\\n[3] USP: A Unified Sequence Parallelism Approach for Long Context Generative AI", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and constructive comments! We appreciate your positive feedback on the good motivation, novelty of Megalodon and its strong empirical results. We address your concerns and questions below and please let us know if you still have concerns after you read our response.\n\n> W1: 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.\n\nWe did not have sufficient computational resources to compare Megalodon and Transformer on a scale of 6T training tokens. But based on the trends of the learning curves in Figure 1, we believe that the gap between Megalodon and Transformer might not be narrowed with more training data.\n\n> W2: Section 3.2 is difficult to understand due to the lack of background and intuitive explanations.\n\nAs discussed in the beginning of Section 3.2, the motivation of Timestep Normalization is to perform normalization along the axis of time steps, which serves the motivation of reducing internal covariate shift along the spatial dimension, similar to Batch Normalization and Group Normalization. The effectiveness of Timestep Normalization has been shown in the [author rebuttal](https://openreview.net/forum?id=XlAbMZu4Bo&noteId=G72sD1BBxi).\n\n> W3: 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.\n\nThanks for pointing out the related work we missed. In fact, we have cited [2] in our submission. We will add the other two works in our final version. \n\nIn Section 3.5, we want to highlight the benefit of chunk-wise attention in parallel pre-training. As explored in [1-4], advanced algorithms have been developed for distributed computation of full attention. However, these algorithms involves significant communication costs. Benefiting from the chunk-wise attention in Megalodon, there are no communications of attention along the chunk/context parallel groups.\n\n[1] Sequence Parallelism: Making 4D Parallelism Possible\n\n[2] Lightseq: Sequence level parallelism for distributed training of long context transformers\n\n[3] USP: A Unified Sequence Parallelism Approach for Long Context Generative AI\n\n[4] Ring Attention with Blockwise Transformers for Near-Infinite Context\n\n> W4: The code link provided in the abstract does not work.\n\nThe anonymous link is a placeholder, not a real link. We will release the code in the final version.\n\n>W5: 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. 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.\n\nIn fact, the model that was trained on the same data with Megalodon is LlaMa2, not LlaMa2-L.  LlaMa2-L continually trained LlaMa2 on 500B selected long-context data for length extension. That is the partial reason why LlaMa2-L performs slightly better in Table 2.\nFor other long-context models, such as Yarn, Xgen and MPT, we agree that they are trained on different data and the comparison is not entirely fair. However, it is impractical to re-train all these models with the same data.\n\nFor the results in Table 6 on LRA, all the models (Transformer, Reformer, etc.) are trained on the same datasets for each task. Thus, the comparison in Table 6 is fair. \n\nWhen we compare Megalodon-chunk with Mega-chunk in Table 6, the improvements are significant (87.62 vs. 85.66). \nIn Table 6, the improvements of Megalodon over Mega with full attention are less significant than that of the chunk-wise attention models. It is reasonable because the benefits of Megalodon over Mega might be covered by the full attention mechanism.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DIePrdwi8k", "reviewer": "Reviewer_6wcE", "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).", "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?\n\n- What's your opinion on Sliding Window Attention vs. Chunk Attention?", "limitations": "NA", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\n- 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\n\n- Figure1: the training sequence length is not the same, thus the perplexity comparison might be unfair. \n\n- Lacking discussions & comparison to recent hybrid (local) attention and RNN models: e.g. Griffin, Jamba\n\n- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and constructive comments! We appreciate your positive feedback on the good motivation, novelty of Megalodon and its strong empirical results. We address your concerns and questions below and please let us know if you still have concerns after you read our response.\n\n> W1: There are no ablation studies in moderate-scale language modeling at all\n\nThe ablation studies on small/mederate-scale benchmarks are in the [author rebuttal](https://openreview.net/forum?id=XlAbMZu4Bo&noteId=G72sD1BBxi).\n\n> W2: Figure1: the training sequence length is not the same, thus the perplexity comparison might be unfair.\n\nThough the training sequence length is not the same, the Megalodon-7B model is using 4K attention chunk size, which is the same as the context length of full attention in LlaMa2-7B. Thus Megalodon-7B and LlaMa2-7B in Figure 1 were trained on the same data with similar flops/token, yielding relatively fair PPL comparison.\n\n> W3: Lacking discussions & comparison to recent hybrid (local) attention and RNN models: e.g. Griffin, Jamba\n\nGriffin and Jamba are two concurrent works and they are trained in different scales of model size and data.\n\n> W4: 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\n\nWe conducted experiments to evaluate Megalodon on retrieval-oriented tasks, such as passkey retrieval. Similar to previous studies, due to chunk-wise attention, Megalodon under-performed on these retrieval-oriented tasks compared with full attention mechanism, but outperformed state-space models such as Mamba. \n\nFor example, without any fine-tuning for length extension, Megalodon completes the passkey retrieval task with up to 16K context length, while Mamba can only complete this task up to 4K context. Long-LlaMa2, which continually trains LlaMa2 on selected long-context data for length extension, is able to complete up to 32K context length. \n\nFurther improving Megalodon for retrieval-oriented tasks is an interesting and important direction for future work.\n\n> Q1: 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?\n\nGated linear recurrence, such as Mamba, HGRN1 etc, incorporates input-dependent recurrence, which sacrifices efficiency compared to (complex) EMA and S4. We did not try gated linear recurrence mainly because Megalodon leverages chunk-wise attention mechanism, which might down-weight the necessity of input-dependent recurrence.\n\n> Q2: What's your opinion on Sliding Window Attention vs. Chunk Attention?\n\nSliding Window Attention is theoretically more powerful but less efficient than Chunk-wise Attention. Moreover, sliding window attention increases the difficulty of parallel training along the chunk/context parallel groups. It is an interesting direction for future work to efficiently incorporate sliding window attention into Megalodon.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kwu2F4bdi7", "reviewer": "Reviewer_xs5D", "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.", "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?\n\nHow much compute is added over Mega/Mega-Chunk to incorporate the necessary changes?\n\nHow is the complex moving average implemented?", "limitations": "Limitations should be addressed further and more directly in the paper.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The technique is validated against modern models across a large variety of benchmarks\n- The diagonalizable complex moving average seems like an easy win for efficient computation and improvement\n- 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.\n\n- 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.\nAt 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. \n\nI 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.\n\n- 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.\n\n- 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.\n\n- 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.\n\n\nWhile 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and constructive comments! We appreciate your positive feedback on the good motivation, novelty of Megalodon and its strong empirical results. We address your concerns and questions below and please let us know if you still have concerns after you read our response.\n\n> W1: 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…\n\nThe ablation studies on small/mederatel-scale benchmarks are in the [author rebuttal](https://openreview.net/forum?id=XlAbMZu4Bo&noteId=G72sD1BBxi).\n\nWe cannot successfully train Mega on a larger-scale setting, due to numerical instability. That is why we did not compare Megalodon with Mega with large scale experiments.\n\n> W2:  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.\n\nThe training flops of Megalodon-7B is similar with LlaMa2-7B with 4K context, because we also used 4K as the chunk size of chunk-wise attention. To keep the same 1.3x speed over Transformer with 32K context length, the context length of full attention should be at most 8K-10K.\n\n> W3: 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.\n\n[Previous studies](https://github.com/jzhang38/LongMamba) have investigated the PPL of Transformer and other architectures (such as Mamba) beyond their training context lengths. Without any fine-tuning for length extension, Transformer and Mamba did not achieve the monotonic PPL of Megalodon in Figure 5.\n\n> W4: 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. \n\nWe appreciate your suggestion on evaluating Megalodon on byte-level modeling tasks. In our experiments on PG-19 (Table 5), we compared Megalodon with MegaByte, which is one of the most advanced byte-level architectures. In addition, among the six tasks in LRA (Table 6), the three text-classification tasks are in byte level. These results indirectly demonstrate the effectiveness of Megalodon on byte-level modeling. We leave more advanced byte-level modeling with Megalodon to future work.\n\n> W5: 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.\n\n[Previous study](https://github.com/jzhang38/LongMamba) showed that, without any fine-tuning,  the context extension of Mamba is not as good as Megalodon.\n\n> Q1: 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?\n\nThere are different factors to impact the convergence speed at the beginning of training, including peak learning rate, learning rate scheduler and warmup steps. Thus we believe the final pre-trianing loss is more indicative.\n\n> Q2: How much compute is added over Mega/Mega-Chunk to incorporate the necessary changes?\n\nThe additional flops of Megalodon over Mega are from complex CEMA (over real-number EMA) and TImestep Normalization (over Layer/RMS Normalization). These additional flops are marginal (< 1%) compared to the total flops in Megalodon.\n\n> Q3: How is the complex moving average implemented?\n\nThe implementation of CEMA is similar to that of EMA. We use FFT to compute the outputs (see appendix A in Mega paper for details). To accelerate CEMA, we implemented separated CUDA kernels for the computation of the convolutional kernel and FFT.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Ablation studies on CEMA and Timestep Normalization

**Ablation on LRA**

We first conducted ablation studies on LRA to demonstrate the effectiveness of CEMA and Timestep Normalization components in Megalodon. The results are shown in the following table:

| Models                    |  ListOps (LN)     |  Text (SN)     |  Retrieval  (SN)    |  Image (BN)     |  Pathfinder (BN)  | Path-X (BN) |  Avg. |
| :--------- | :----------------: | :----------------: | :----------------: | :----------------: |  :----------------: | :----------------: | :----------------: |
| Mega (EMA)                     |   58.76   |   90.19   |   90.97   |   85.80   |   94.41   |   93.81   |   85.66   |
| Megalodon (CEMA)          |   61.13   |   90.58   |   91.51   |   87.32   |   96.11   |   96.98   |   87.37   |
| Megalodon (CEMA&TSN)|  62.25   |   90.50   |   91.76   |   87.16   |   96.85   |   97.21   |   87.63   |

In this study, we used the Mega architecture as the baseline, which uses different normalization layers for different tasks in LRA. The normalization layers for different tasks are in the brackets following the task names in the table headline. BN, LN and SN refer to Batch Normalization, Layer Normalization and Scale Normalization. 

The second row in the above table is the Megalodon architecture by replacing EMA with CEMA. For this architecture, we used the same normalization layers with the original Mega on different tasks. The third row is the Megalodon architecture with both CEMA and Timestep Normalization (TSN). Note that for this architecture, we used the same TSN for all the six tasks in LRA.

**Ablation on WikiText-103**

We then conducted ablation studies on auto-regressive language modeling on the moderate WikiText-103 dataset:

| Models                              |    #Param.   |     PPL    | 
| :-----------------------------| :------------: | :--------: |
| Mega (EMA&LN)               |       252M     |   18.07   | 
| Megalodon (CEMA&LN)   |       252M     |    17.63  |
| Megalodon (CEMA&TSN) |       252M    |    17.23  |

In this study, we also used the Mega architecture with EMA and Layer Normalization (LN) as baseline. The second row is Megalodon with CEMA and Layer Normalization, while the third row is the Megalodon with CEMA and Timestep Normalization (TSN).

From the two ablation studies, both CEMA and Timestep Normalization show improvements over the original Mega architecture.


# Effect of normalized attention, and pre-norm with two-hop residuals

The normalized attention and pre-norm with two-hop residuals are designed to improve stability of Megalodon in large-scale pre-training. Their effects on small-scale models are not significant.

Meta-review: This paper addresses long context length issue by introducing new architecture. 

Intially, three reviewers voted acceptance and one reviewer gave a borderline rejection with lack of experiments.

However, the authors successfuly addresss main concerns that reviewers raised, thus all reviewers gave positive scores.

AC also agree with the reviewers opinion, and moreover AC would like to encourage research on new architecture apporach for topic diversity in LLM regimes. 

So, AC recommens accepting this paper.
Decision: Accept (poster)