Australian National University

Qingqin Fang

#### An improvement architecture superior to the Transformer, proposed by Meta

<figure>
<img src="https://cdn-images-1.medium.com/max/852/0*pFztXZpCGW0GplEE" />
<figcaption>Image Generated by Canvas</figcaption>
</figure>

### Author

· Qingqin Fang (**ORCID:**
[0009--0003--5348--4264](https://orcid.org/0009-0003-5348-4264))

### Introduction

Recently, researchers from Meta and the University of Southern
California have introduced a model called Megalodon. They claim that
this model can expand the context window of language models to handle
millions of tokens without overwhelming your memory.

In previous studies on long texts, the focus was mainly on enhancing
performance through optimizations of the Transformer. However, lab data
from recent experiments show that Megalodon outperforms Transformer
models of similar size when processing large volumes of text, and many
researchers are beginning to view it as the potential successor to the
Transformer.

Transformer architecture is generally limited by its high complexity and
its weak ability to handle very long inputs. While there are simpler
solutions that try to address these issues, they usually don't perform
as well in terms of training speed or accuracy on practical tasks. The
emergence of Megalodon is specifically aimed at addressing the challenge
of handling infinite context, which makes it better suited for tasks
where understanding extensive context is crucial.

This post delves into Megalodon.

### Long Context Window

In long-text situations, "**context window**" indicates the number of
tokens a model can handle at one time. Expanding this window enables LLM
to engage in longer conversations, process a wider range of documents,
and learn more information.

In machine learning models, the **attention** mechanism is a special
structure that enhances the weights of certain parts of the input data
while reducing the weights of other parts. This mechanism automatically
learns and calculates the contribution of input data to the output data,
allowing the model to focus on different parts of the input with varying
levels of importance during task execution. This concept is somewhat
similar to human attention.

The diagram below shows the core parts and steps of the **QKV
attention** mechanism, an acronym for Query, Key, and Value, commonly
used in **self-attention** mechanisms, particularly in Transformer
models. Initially, input data is linearly transformed into three
segments: Query (yellow), Key (pink), and Value (green). Each Query is
then dot-multiplied with all Keys to calculate their similarity scores,
which are normalized using the Softmax function. This process assigns
attention weights to each Key relative to each Query, ultimately
producing the final output based on these weights.

<figure>
<img
src="https://cdn-images-1.medium.com/max/1024/0*0CGhyOD6zPwz2Nl3" />
<figcaption>Transformer-Attention Flow (scaled-dot-product). Link: <a
href="https://scholar.harvard.edu/sites/scholar.harvard.edu/files/binxuw/files/mlfs_tutorial_nlp_transformer_ssl_updated.pdf">From
Transformer to LLM: Architecture, Training and Usage</a></figcaption>
</figure>

However, self-attention mechanism compares each element in the input
sequence with all others using Scaled Dot-Product Attention. This means
that if the input size is doubled, the memory and computation time
needed to process the input will quadruple. Expanding the Transformer's
context window comes with significant costs. The challenge with
Transformers lies in their "**quadratic complexity**".

Meanwhile, Transformers also face issues with "**the lack of
position-related inductive bias**", which can impact their performance.
Unlike Recurrent Neural Networks (RNNs) or Convolutional Neural Networks
(CNNs), which naturally process data sequentially or have the ability to
recognize spatial structures, transformers rely on **position
embeddings** to incorporate information about the order of tokens. This
method involves adding positional information to the word embedding
vectors of each token, allowing the same word to have different
representations depending on its position in the sentence.

This artificially creates a position-related bias, but such a bias may
not be as effective as the natural positional sensitivity found in other
models. When dealing with tasks where positional context is crucial, for
example, the two *"I"* are located at different positions in the
sentence "*I think, therefore I am",* it should have the same output
representations, but the two I's don\'t. This can sometimes lead to
inefficiencies or require additional mechanisms to address the issue.

In 2022, the **Moving Average Equipped Gated Attention**(MEGA) model was
introduced as a significant enhancement to the traditional Transformer
architecture, particularly improving its performance in handling long
sequence data. [This
article](https://medium.com/@_init_/how-self-attention-with-relative-position-representations-works-28173b8c245a)
provides a more detailed explanation.

MEGA incorporates the **Exponential Moving Average** (EMA) into the
self-attention mechanism. Moving averages, which is a classic approach
for sequential data modelling, widely used in time series data to smooth
out short-term fluctuations and highlight long-term trends or cycles is
calculated for each element's Query and Key, smoothing data and updating
weights. These Queries and Keys reflect not only the current input
features but also the influence of historical data, enabling more
effective capture and utilization of long-distance contextual
information. EMA facilitates the flow of information across different
time steps, improving the model's understanding of long-range
dependencies within sequences.

Moreover, by dividing the computation of the attention mechanism into
smaller segments, MEGA substantially reduces computational complexity
from O(n²) to **approximately O(n)**. The use of EMA allows the MEGA
model to maintain a more coherent flow of information between local and
overall context, helping to prevent the loss of information during the
processing across blocks. This restructuring makes it possible to handle
longer data sequences more efficiently while preserving the integrity of
the information processed.

<figure>
<img
src="https://cdn-images-1.medium.com/max/1024/0*IhNa1lHvSi8YNuXh" />
<figcaption>Mega — model architecture, Source: Xuezhe Ma et al., 2022.
Link: <a
href="https://arxiv.org/pdf/2209.10655">https://arxiv.org/pdf/2209.10655</a></figcaption>
</figure>

### Megalodon

Building on the MEGA architecture, Megalodon has implemented a series of
enhancements and added several new technical components:

Firstly, Megalodon introduces the **Complex Exponential Moving Average
(CEMA)** component, a novel technique that extends the [Multidimensional
Damped Exponential Moving Average](https://arxiv.org/abs/2312.02549)
(EMA) from MEGA into the complex domain. The previous MEGA model was
based on real numbers, and Megalodon extends it to include imaginary
numbers, hence the term "complex domain" here. This enhancement boosts
the model's capability to process complex data, thereby improving
performance and stability.

Secondly, normalization is a crucial step in Transformer models,
involving adjusting data to a distribution with a mean of 0 and a
variance of 1. This reduces variations in data distribution during model
training, thus helping the model learn faster. The commonly used
normalization methods are:

- **Layer Normalization** adjusts the outputs by calculating the mean
  and standard deviation of all neuron outputs within a layer. This
  method effectively stabilizes the hidden states in recurrent networks
  but does not account for batch and sequential data properties,
  resulting in poor management of dependencies between time steps.
- **Group Normalization** normalizes neuron outputs by grouping them.
  This method is not affected by batch size and maintains good stability
  across a wide range of batch sizes. However, its performance heavily
  depends on how the channels are grouped, and it may inadvertently leak
  future information in autoregressive models due to normalization
  across different time steps.

To address these issues, researchers have proposed a new, more complex
method called **Timestep Normalization**. This technique extends group
normalization to autoregressive sequence modelling, allowing for
independent normalization at each time step. This enables the model to
effectively adjust step-by-step in sequence processing, ensuring that
dependencies between time steps do not compromise model performance.
This method helps the model handle sequential data more accurately,
improving processing effectiveness and efficiency.

<figure>
<img
src="https://cdn-images-1.medium.com/max/1024/0*V5PbIbNUwDAash4D" />
<figcaption>Normalization methods. The elements in blue or yellow are
the regions to compute means and variances. Source: Xuezhe Ma et al.,
2024. Link: <a
href="https://arxiv.org/pdf/2404.08801">https://arxiv.org/pdf/2404.08801</a></figcaption>
</figure>

Additionally, Megalodon has introduced a **normalized attention
mechanism**. This involves normalizing the queries (Query) and keys
(Key) before calculating the attention scores, which helps stabilize the
scores and addresses issues caused by inconsistent sizing, improving the
model's overall reliability and training stability.

These adjustments are achieved by enhancing common normalization
techniques. By breaking down the input sequence into parts, Megalodon
simplifies the computational process and memory usage during model
training and inference, making it more efficient. This allows Megalodon
to handle large amounts of data more effectively without using excessive
resources.

### Experiment and its Results

Researchers trained a 7-billion-parameter Megalodon model on the same 2
trillion pre-training data as used with the Llama2 model, allowing for a
completely fair comparison with Llama2--7B. The final results are shown
in the diagram below.

<figure>
<img
src="https://cdn-images-1.medium.com/max/1024/0*AJmNTzXPJVCQRkK8" />
<figcaption>Processed tokens during training for MEGALODON-7B, LLAMA2–7B
and LLAMA2–13B. Source: Xuezhe Ma et al., 2024. Link: <a
href="https://arxiv.org/pdf/2404.08801">https://arxiv.org/pdf/2404.08801</a></figcaption>
</figure>

It is evident from the direct comparison with Llama 2 that Megalodon not
only trains more efficiently but also surpasses the Transformer in
accuracy when handling tasks with 7 billion parameters and 2 trillion
training tokens. Specifically, Megalodon's training loss was 1.70, which
is between Llama2--7B (1.75) and 13B (1.67).

At the same time, researchers used the same computational resources to
run experiments with LLAMA 2--7B and MEGALODON-7B, comparing their
training speeds at different context lengths.

<figure>
<img src="https://cdn-images-1.medium.com/max/836/0*BGeSHoAz8e3dL6jV" />
<figcaption>Average word/token per second for MEGALODON-7B and
LLAMA2–7B. Source: Xuezhe Ma et al., 2024. Link: <a
href="https://arxiv.org/pdf/2404.08801">https://arxiv.org/pdf/2404.08801</a></figcaption>
</figure>

At a 4K context length, MEGALODON-7B was slightly slower (by about 6%)
than LLAMA 2--7B, due to the introduction of CEMA and timestep
normalization. However, when the context length was expanded to 32K,
MEGALODON-7B was significantly faster by 32% compared to LLAMA 2--7B,
demonstrating MEGALODON's computational efficiency in long-context
pre-training.

For further experimental details, please refer to the [full
paper](https://arxiv.org/pdf/2404.08801). Additionally, [the code for
this project has been made open
source](https://github.com/XuezheMax/megalodon), making it accessible to
a broader audience. The computational cost of running the 7-billion
parameter model is optimized to be feasible on personal computers,
enhancing its accessibility for individual researchers and developers.

### Future Thinking

So, can the Transformer be replaced by these new architectures?

As of now, the Transformer architecture still holds a leading position.
While Meta is exploring new architectures like Megalodon, it is also
actively improving its Transformer models and has recently released the
latest version of its open-source
LLM, [Llama-3](https://ollama.com/library/llama3).

One challenge that competitors to Transformers face is the need for
specialized hardware and software tools. There is already a large
library and ecosystem of tools available for training, fine-tuning, and
customizing Transformer models, suitable for various applications and
hardware configurations. Researchers have also developed optimized
software codes to enhance the performance of Transformer LLMs on
memory-constrained devices. Alternatives still need to keep up with
these advancements.

Meanwhile, other researchers are modifying the Transformer architecture
to reduce its memory and computational demands. For instance, Google
researchers recently published a paper on
[Infini-attention](https://arxiv.org/abs/2404.07143), aimed at providing
Transformers with an infinite context window without increasing memory
and computational complexity.

While new architectures like Megalodon bring exciting improvements,
particularly in handling longer sequences and reducing computational
demands, the traditional Transformer model still holds a dominant
position in the AI field due to its widespread application and
continuous development.

In the future, these new models could work alongside or gradually
integrate with the traditional Transformer framework to more effectively
address specific challenges. Whether these models will completely
replace Transformers, or rather complement them, depends on further
advancements and broader adoption of these innovative approaches in
practical applications.

Looking ahead, observing how these models evolve and are utilized across
various domains might reshape our understanding and implementation of
neural network architectures in large-scale artificial intelligence
systems. The journey from traditional Transformers to models like
Megalodon marks a dynamic and transformative phase in machine learning,
promising to enhance the capabilities and impact of AI technologies
globally.

### **References**

- Xuezhe Ma et al. (2024). Megalodon: Efficient LLM Pretraining and
  Inference with Unlimited Context Length. arXiv:2404.08801 \[cs.LG\].
  DOI: <https://arxiv.org/pdf/2404.08801>
- Xuezhe Ma et al. (2024). Mega: Moving average equipped gated
  attention. arXiv:2209.10655 \[cs.LG\]. DOI:
  <https://doi.org/10.48550/arXiv.2209.10655>

![](https://medium.com/_/stat?event=post.clientViewed&referrerSource=full_rss&postId=b4c7322602ba){width="1"
height="1"}


An improvement architecture superior to the Transformer, proposed by Meta  Author   · Qingqin Fang (
<strong>
 ORCID:
</strong>
0009–0003–5348–4264) Introduction   Recently, researchers from Meta and the University of Southern California have introduced a model called Megalodon. They claim that this model can expand the context window of language models to handle millions of tokens without overwhelming your memory.

Messaggi di Rogue Scholar

The longer the context, the better? Unlimited Context Length in Megalodon