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.

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…

The ablation studies on small/mederatel-scale benchmarks are in the author rebuttal.

We 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.

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.

The 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.

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.

Previous studies 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.

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.

We 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.

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.

Previous study showed that, without any fine-tuning, the context extension of Mamba is not as good as Megalodon.

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?

There 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.

Q2: How much compute is added over Mega/Mega-Chunk to incorporate the necessary changes?

The 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.

Q3: How is the complex moving average implemented?

The 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.

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.

W1: There are no ablation studies in moderate-scale language modeling at all

The ablation studies on small/mederate-scale benchmarks are in the author rebuttal.

W2: Figure1: the training sequence length is not the same, thus the perplexity comparison might be unfair.

Though 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.

W3: Lacking discussions & comparison to recent hybrid (local) attention and RNN models: e.g. Griffin, Jamba

Griffin and Jamba are two concurrent works and they are trained in different scales of model size and data.

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

We 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.

For 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.

Further improving Megalodon for retrieval-oriented tasks is an interesting and important direction for future work.

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?

Gated 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.

Q2: What's your opinion on Sliding Window Attention vs. Chunk Attention?

Sliding 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.

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.

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.

We 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.

W2: Section 3.2 is difficult to understand due to the lack of background and intuitive explanations.

As 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.

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.

Thanks 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.

In 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.

[1] Sequence Parallelism: Making 4D Parallelism Possible

[2] Lightseq: Sequence level parallelism for distributed training of long context transformers

[3] USP: A Unified Sequence Parallelism Approach for Long Context Generative AI

[4] Ring Attention with Blockwise Transformers for Near-Infinite Context

W4: The code link provided in the abstract does not work.

The anonymous link is a placeholder, not a real link. We will release the code in the final version.

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.

In 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. For 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.

For 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.

When we compare Megalodon-chunk with Mega-chunk in Table 6, the improvements are significant (87.62 vs. 85.66). In 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.

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.

W1 & Q1: Have the authors conducted ablation studies for small-scale models before moving to 7B, similar to the results in Appendix C?

The ablation studies on small-scale benchmarks are in the author rebuttal.

W2: Recent theoretical work [1] has shown that efficient versions of attention can limit the expressiveness of the model…

We 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.

For 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.

Further improving Megalodon for retrieval-oriented tasks is an interesting and important direction for future work.