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 introduceMegalodon, an neural architecture for efficient sequence modeling with unlimited context length.Megalodoninherits the architecture ofMega(exponential moving average with gated attention) , and further introduces multiple technical components to improve its capability and stability, includingcomplex exponential moving average (CEMA) ,timestep normalizationlayer,normalized attentionmechanism andpre-norm with two-hop residualconfiguration.
In a controlled head-to-head comparison withLlama2,Megalodonachieves better efficiency than Transformer in the scale of 7 billion parameters and 2 trillion training tokens.Megalodonreaches a training loss of 1.70, landing mid-way betweenLlama2-7B (1.75)  and 13B (1.67) .
The improvements ofMegalodonover Transformers are robust throughout a range of benchmarks across different tasks and modalities.Code:https://github.com/XuezheMax/megalodon