M EGA : M OVING A VERAGE E QUIPPED G ATED A TTENTION
Xuezhe Ma ∗ ISI, USC Chunting Zhou ∗ Meta AI Xiang Kong
LTI, CMU
Junxian He
SJTU
Liangke Gui LTI, CMU Graham Neubig LTI, CMU
Jonathan May
ISI, USC
Luke Zettlemoyer
Meta AI
A BSTRACT
The design choices in the Transformer attention mechanism, including weak inductive bias and quadratic computational complexity, have limited its application for modeling long sequences.
In this paper, we introduce M EGA , a simple, theoretically grounded, single-head gated attention mechanism equipped with (exponential) moving average to incorporate inductive bias of position-aware local dependencies into the position-agnostic attention mechanism.
We further propose a variant of
M EGA that offers linear time and space complexity yet yields only minimal quality loss, by efficiently splitting the whole sequence into multiple chunks with fixed length.
Extensive experiments on a wide range of sequence modeling benchmarks, including the Long Range Arena, neural machine translation, auto-regressive language modeling and image and speech classification, show that M EGA achieves significant improvements over other sequence models, including variants of Transformers and recent state space models
I NTRODUCTION
Designing a single unified model to capture long range dependencies in sequential data is a central and challenging problem in sequence modeling.
A number of different archtectures have been developed, including convolutional neural networks (CNNs)
(Kim, 2014; Strubell et al., 2017), recurrent neural networks (RNNs) (Goller & Kuchler, 1996; Hochreiter & Schmidhuber, 1997), Transformers (Vaswani et al., 2017) and recent state space models (SSMs) (Gu et al., 2022a; Mehta et al., 2022).
Among these models, Transformer (Vaswani et al., 2017) has stood out for its impressive empirical success on a wide range of language and vision tasks, including machine translation (Vaswani et al., 2017; Ott et al., 2018), language understanding (Devlin et al., 2019; Liu et al., 2019), image recognition (Dosovitskiy et al., 2020; Touvron et al., 2021) and genetic sequence modeling (Madani et al., 2020; Jumper et al., 2021), mainly because of the conceptually attractive attention mechanism (Bahdanau et al., 2015) which directly models interactions between each pair of input tokens.
However, there are two common drawbacks in this design: i) weak inductive bias ; and ii) quadratic computational complexity .
First, the attention mechanism does not assume prior knowledge of the patterns of dependencies between tokens (e.g. positional inductive bias), instead learning to predict the pairwise attention weights directly from data.
Second, the cost to compute and store the attention weights is quadratic in the input length.
Recent studies have shown the limitations of applying Transformers to longer sequences, both with respect to accuracy and efficiency (Tay et al., 2020).
In this work, we propose a moving average equipped gated attention mechanism (M EGA ) to solve the two weaknesses simultaneously.
The key idea is to incorporate inductive biases into the attention mechanism across the timestep dimension, by leveraging the classic exponential moving average (EMA) approach (Hunter, 1986).
EMA captures local dependencies that exponentially decay over time (see Figure 1), and has been widely used in time series data modeling (§2).
We introduce a multi-dimensional damped form of EMA with learnable coefficients (§3.1), and subsequently develop the moving average equipped gated attention mechanism by integrating the EMA with a variant of the single-head gated attention (Hua et al., 2022) (§3.2).
Theoretically, we show that the single-head gated attention is as expressive as the most commonly used multi-head attention (§3.5).
Benefiting
[Table]
from the incorporated moving average mechanism, we further propose a variant of M EGA with linear complexity, named M EGA -chunk, with minimal loss of contextual information (§3.4). Experimentally, on five sequence modeling tasks across various data types, including long-context sequence modeling, neural machine translation, auto-regressive language modeling, and image and speech classification, M EGA significantly outperforms a variety of strong baseline models, in terms of both effectiveness and efficiency (§4) (see Table 1).
B ACKGROUND
We use X = { x 1 , x 2 , . . .
, x n } ∈ R n × d to denote a sequence of input representations with length n .
Let Y =
{ y 1 , y 2 , . . .
, y n } ∈ R n × d be the sequence of output representations of each layer with the same length n as the input X .
In this paper, we assume the representations of the input and output sequences have the same dimension d .
2. S ELF -A TTENTION M ECHANISM
The traditional self-attention mechanism is a function:
[Equation]
where Attn : R n × d → R n × d is the self-attention function.
Q = X
[Equation]
b q , K = X W k + b k , and V =
X W v + b v are the sequences of queries, keys and values, with learnable parameters
W q , W k , W v ∈ R d × d , and b q , b k , b v ∈ R d . f ( · ) is an attention function, e.g. the softmax function f softmax ( · ) (Bahdanau et al., 2015), or the recently proposed squared ReLU function f relu 2 ( · ) (So et al., 2021; Hua et al., 2022).
τ ( X ) is a scaling term, which is commonly set to τ ( X ) = √ d for f softmax ( · ) , or τ ( X ) = n for f relu 2 ( · ) .
The commonly used multi-head variant of attention performs the attention function h times in parallel.
We can define a matrix A
= f ( QK T τ ( X ) )
∈ R n × n following (1), which is called the attention matrix .
Since it specifies pairwise dependency weights, the matrix A in principle delivers a flexible and powerful mechanism to learn long-distance dependencies with minimal inductive biases.
However, it is in practice a challenging task to recognize all the dependency patterns directly from data, particularly when processing long sequences.
Moreover, calculating A with h attention heads takes
O ( hn 2 ) time and space, which becomes a significant bottleneck.
2. E XPONENTIAL M OVING A VERAGE (EMA)
The moving average is a classic approach for sequential data modeling, which has been widely used in time series data to smooth out short-term fluctuations and highlight long-term trends or cycles.
The Exponential Moving Average (EMA) (Winters, 1960; Hunter, 1986), a special case of moving average, applies weighting factors that decrease exponentially:
[Equation]
(2) where α ∈ (0 , 1)
d is the EMA coefficient representing the degree of weighting decrease, and ⊙ is the element-wise product.
A higher α discounts older observations faster (see Figure 1).
Using an EMA places a strong inductive bias on the learning of pairwise dependencies: the dependency weight between two tokens decreases exponentially over time with an input-agnostic decay factor α .
This property favors local dependencies, and limits long-distance dependencies.
Despite the recurrent formulation in (2), the computation of EMA can be represented as n individual convolutions, which can be computed efficiently using fast Fourier transforms (FFTs) (see Appendix A for details).
[Table]
2. W HY C OMBINE A TTENTION WITH EMA?
As discussed in Sections 2.1 and 2.2, EMA and attention mechanisms each have their own limitations, despite their wide applications and impressive successes in sequence modeling.
By leveraging their properties to complement each other, we propose to embed an EMA into the calculation of the attention matrix A .
The resulting model enjoys the benefit from strong inductive bias, while maintaining the capacity to learn complex dependency patterns.
Moreover, this integration enables the design of a computationally efficient chunk-wise attention mechanism with linear complexity w.r.t sequence length (§3.4).
M OVING A VERAGE E QUIPPED G ATED A TTENTION (M EGA )
In this section, we describe in detail our proposed method, moving average equipped gated attention (M EGA ).
We first introduce multi-dimensional damped EMA (§3.1), which is a key component combined with the single-head gated attention in M EGA (§3.2), and discuss the relationship between
M EGA and three closely related models: GRU (Cho et al., 2014), Flash (Hua et al., 2022) and
S4 (Gu et al., 2022a).
Then, we describe the detailed architecture of each M EGA block, including feed-forward and normalization layers (§3.3), and present M EGA -chunk, a variant of M EGA that splits input sequences into fixed chunks, reducing time and space complexity from quadratic to linear (§3.4).
At last, we provide theoretical justification for the design of single-head gated attention (§3.5).
3. M ULTI - DIMENSIONAL D AMPED EMA
Damped EMA.
Previous studies (McKenzie & Gardner Jr, 2010; Svetunkov, 2016) have shown that relaxing the coupled weights of the previous and current observations ( α vs. 1 − α in (2)) produces robust dependency modeling.
Inspired by this, M EGA allows the damping of the influence of the previous time step:
[Equation]
where δ ∈ (0 , 1) d is the damping factor.
Multi-dimensional Damped EMA.
To further improve the expressiveness of EMA, we introduce a multi-dimensional variant of EMA.
Concretely, we first expand each dimension j ∈ { 1 , 2 , . . .
, d } of the input sequence X individually into h dimensions via an expansion matrix β ∈ R d × h :
[Equation]
(4) where β j ∈ R h is the j -th row of β , u ( j ) t ∈
R h is the expanded h -dimensional vector for the j -th dimension at timestep t .
Correspondingly, we extend the shape of α and δ from a one-dimensional vector to a two-dimensional matrix, i.e. α , δ ∈ R d × h , where α j , δ j ∈ R h denote the j -th row of α and δ , respectively.
Then, for each dimension j , the damped EMA is applied to the h -dimensional hidden space:
[Equation]
where h ( j ) t ∈
R
h is the EMA hidden state for the j -th dimension at timestep t .
η ∈ R d × h is the projection matrix to map the h -dimensional hidden state back to 1 -dimensional output y t,j ∈ R .
η j ∈ R h is the j -th row of η .
The output Y from (5) is denoted as Y ∆ =
EMA( X ) .
Because we do not need to explicitly compute h ( j ) t to get the output y t,j , and the time and space complexity is similar to the standard EMA in (2) (see Appendix A for the details).
[Figure]
[Table]
[Figure]
3. M OVING A VERAGE E QUIPPED G ATED A TTENTION
The gated attention mechanism in M EGA adopts the Gated Recurrent Unit (GRU; Cho et al.
and Gated Attention Unit (GAU; Hua et al.
as the backbone architectures, with an EMA-based sub-layer embedded into the calculation of the attention matrix.
Formally, we first use the output from (5) to compute the shared representation in GAU:
[Equation]
where X ′ can be regarded as the updated or contextual input, because it encodes contextual information through EMA.
Z is the shared representation with z dimensions, with projection matrix
W z ∈ R
d × z and bias term b z ∈ R z .
ϕ silu is the self-gated activation function (SiLU) (Ramachandran et al., 2017; Elfwing et al., 2018).
Following GAU, 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 :
[Equation]
where κ q , µ q , κ k , µ k ∈ R z are the learnable scalars and offsets of queries and keys, respectively.
v is the expanded intermediate dimension for the value sequence.
The output of attention is:
[Equation]
[Equation]
The graphical specification is displayed in Figure 2 (c).
b rel ∈ R n × n is the relative positional bias.
We choose b rel from existing approaches, including T5 (Raffel et al., 2020),
RoPE
(Su et al., 2021),
TUPE (Ke et al., 2020) and ALiBi (Press et al., 2021).
Subsequently, M EGA introduces the reset gate γ , the update gate φ , and computes the candidate activation output ˆ H :
[Equation]
The final output Y is computed with the update gate φ :
[Equation] [Equation]
The graphical architecture of a M EGA sub-layer is visualized in Figure 2 (b).
Laplace Attention Function.
As mentioned in Section 2.1, the softmax function is the most common choice for the attention function f ( · ) .
So et al.
(2021) recently introduced the squared ReLU function f relu 2 ( · ) via architecture search techniques, which has shown faster convergence speed and competitive generalization performance on language tasks (Hua et al., 2022).
However, one issue of f relu 2 ( · ) is that neither its range nor its gradient is bounded, leading to unstable model training (see
Appendix C.1 for details).
To address this issue, we propose the Laplace attention function:
[Equation]
where erf( · ) is the error function.
µ and σ are two coefficients that we adjust to approximate f relu 2 , yielding µ = (cid:112) 1 / 2 and σ = (cid:112)
1 / 4 π
(The derivations and visualization are provided in Appendix C).
Relation to and Differences from GRU, Flash and S4.
The computation of the the reset gate γ , the update gate φ , and the candidate activation output ˆ H in (12-14) is reminiscent of GRUs (Cho et al., 2014).
The main difference is that in a GRU the two gates are applied between the hidden states of the current and previous timesteps, while in M EGA they are applied between the outputs from EMA and gated attention sub-layers.
In addition, the output gating mechanism in (15) is similar to the gated residual connection proposed in Parisotto et al.
Xu et al.
The computation of the shared representation Z , together with the sequences of queries, keys and values in (7-10) are inspired from GAU in Flash (Hua et al., 2022).
M EGA integrates EMA into GAU by computing Z in (7) from the EMA output X ′ rather than the original input X , and combining the GAU output with X ′ for the candidate activation ˆ H in (14).
Experimental gains over Flash demonstrate the effectiveness of this design chice (§4.1).
The multi-dimensional damped EMA can be seen as a simplified variant of a state space model.
From this perspective, M EGA is also closely related to S4 (Gu et al., 2022a), a state space model with structured state matrices.
S4 leverages the HiPPO framework (Gu et al., 2020) to initialize its low-rank structured state matrices, and the computation of the convolutional kernel in S4 requires complex fast Fourier transformers (FFTs).
The EMA sub-layer in M EGA applies diagonalization on the state matrix, which is similar to a concurrent work S4D (Gu et al., 2022b).
The differences between the M EGA EMA sub-layer and S4D are: i)
S4D parameterizes state matrix in the complex field while EMA uses real numbers; ii) EMA restricts the diagonal elements in the range of (0 , 1) ; iii)
EMA does not rely on the HiPPO framework for parameter initialization.
3. M EGA B LOCKS
The M EGA layer (moving average equipped gated attention) is used as a drop-in-replacement for regular attention in Transformer.
It is followed by position-wise feed-forward networks (FFNs) and normalization layers to compose one M EGA block.
As the gated residual connection has already been included in (15), we omit the original residual connection and directly apply a normalization layer to Y .
Concretely,
[Equation]
(17) where Y ′ is the output of the M EGA block.
The overall architecture of a M EGA block is shown in Figure 2 (a).
In Transformer, the hidden dimension of FFNs is usually set to d FFN = 4 d .
To retain a similar model size with each Transformer block, we reduce the hidden dimension of FFN to d FFN = 2 d and set the expanded dimension v = 2 d for the value sequence in (10) throughout this paper, unless specified otherwise.
3. M EGA - CHUNK : M EGA WITH L INEAR C OMPLEXITY
So far we have only focused on introducing stronger inductive bias into attention, which still has quadratic computational complexity.
In this section, we propose M EGA -chunk, a variant of M EGA with linear complexity, which simply applies attention to each local chunk of fixed length.
Specifically, we first split the sequences of queries, keys and values in (8-10) into chunks of length c .
e.g. Q = { Q 1 , . . .
, Q k } , where k = n/c is the number of chunks.
1 The attention operation
[Table]
+ymScGCrJfFGYcGQiNP0e9ZmixPCxJZgoZm9FZIgVJsZmlLcheIsvL5Nmpeydl6u31WLtKosjB8dwAiXw4AJqcAN1aAABAc/wCm+Ocl6cd+dj3rriZDNH8AfO5w/Eho++</latexit> x (2)
[Figure]
in (11) is individually applied to each chunk, yielding linear complexity O ( kc 2 ) =
O ( nc ) w.r.t n .
However, this method suffers from the critical limitation of losing contextual information from other chunks.
Fortunately, the EMA sub-layer in M EGA mitigates this problem by capturing local contextual information near each token, whose outputs are used as the inputs to the attention sub-layer.
As a result, the effective context being exploited by chunk-wise attention can go beyond the chunk boundary.
Figure 3 illustrates the largest dependency length captured by one M EGA -chunk block.
3. T HEORETICAL J USTIFICATION OF S INGLE - HEAD G ATED A TTENTION
Single-head gated attention has been empirically shown as performant as vanilla multi-head attention Liu et al.
(2021); Hua et al.
(2022), without any discussions on its theoretical insights.
In this section, we provide theoretical justifications of the expressiveness of single-head gated attention.
To facilitate subsequent analysis, we simplify the notations of the multi-head attention.
Specifically, we denote the sequences of queries, keys and values as the outputs of input transformations:
[Equation]
18) where Q , K , V are three transformations, such as linear projections.
Let q ∈ Q = { q 1 , . .
. , q n } be a single query vector ( q ∈ R d ), and
a =
A ( q , K ) denote the corresponding attention weights of q , where A is the attention transformation, i.e. f ( · ) in (11).
For multi-head attention, a common implementation is to split the query into h heads across the model dimension: q = [ q (1) , . . .
, q ( h ) ] , where q ( i ) ∈ R d/h , and i ∈ { 1 , . . .
, h } is the query of the i -th head.
K and V are split in the same way.
The attention weight of the i -th head is a ( i ) =
A ( q ( i ) , K ( i ) ) .
Then, the outputs of single-head and multi-head attention are, respectively:
[Equation]
It is straightforward to see that O MHA is more expressive than O SHA , because O MHA leverages h sets of attention weights.
In the single-head gated attention, we introduce a gate vector γ =
G ( X ) for each q , and the output of single-head gated attention is O SHGA =
O SHA ⊙ γ .
The following theorem reveals the equivalence of O SHGA and O MHA w.r.t expressiveness (proof in Appendix B):
Theorem 1 Suppose the transformation G is a universal approximator.
Then, ∀ X , ∃ γ =
G ( X ) s.t.
[Equation]
Theorem 1 indicates that by simply introducing the gate vector, O SHGA is as expressive as O MHA .
In practice, G is commonly modeled by a (shallow) neural network, whose universality of approximation has been extensively studied (Hornik et al., 1989; Yarotsky, 2017; Park et al., 2020).
E XPERIMENTS
To evaluate M EGA , we conduct experiments on five benchmark sequence modeling tasks across various data types.
All the numbers with ‡ indicate results from the baseline models replicated by us.
More detailed descriptions, results and analysis are provided in Appendix D.
[Table]
4. L ONG -C ONTEXT S EQUENCE M ODELING
We begin our experiments with an evaluation on the Long Range Arena (LRA) benchmark recently introduced by Tay et al.
(2021), which is designed for evaluating sequence models under the longcontext scenario.
This benchmark has six tasks: ListOps (Nangia & Bowman, 2018), byte-level text classification (Text; Maas et al.
(2011)), byte-level document retrieval (Retrieval; Radev et al.
, image classification on sequences of pixels (Image; Krizhevsky et al.
, Pathfinder (Linsley et al., 2018) and its extreme long version (Path-X; Tay et al. (2021)).
These tasks consist of input sequences ranging from 1K to 16K tokens and span across a variety of data types and modalities.
Table 2 compares M EGA against several baselines, including Transformer and its efficient variants, and the state-of-the-art S4 models (both version 1 (Gu et al., 2022a) and version 2 (Gu et al., 2022b)).
To ensure fair comparison, we adjust the number of layers and model dimensions on each task so that
M EGA has similar number of parameters with S4-v1.
For each experiment, we report the average over 5 runs with different random seeds.
On all the six tasks, M EGA substantially outperforms all the baselines.
We also evaluate M EGA -chunk on each task, by setting the chunk size c = 128 for all the tasks, except Path-X where c = 4096 .
We observe that M EGA -chunk consistently performs well, particularly on the three language tasks.
We also examine the speed and memory efficiency of M EGA on the byte-level classification task with the input length of 4K. M EGA -chunk is highly efficient, which is about 5 .
5 times faster and consumes only 13 % as much memory as the vanilla Transformer.
In addition, M EGA with full attention field is also much more efficient than Transformer, benefiting from single-head gated attention.
Analysis of Multi-dimensional Damped EMA.
To demonstrate the effectiveness of the multidimensional damped EMA component in M EGA , we performs ablation studies on two LRA tasks — byte-level text classification (Text) and image classification on sequences of pixels (Image).
We train
M EGA models with EMA dimension h ∈ { 0 , 1 , 2 , 4 , 8 , 16 , 32 } , where h = 0 indicates removing the
EMA component.
From the left figure in Figure 4, we see that without the EMA component, model accuracy on both the two tasks declines rapidly.
Meanwhile, the significant improvements of M EGA with a single dimensional EMA ( h = 1 ) demonstrate the importance of EMA inductive bias.
Analysis of Chunk Size.
We further analyze the impact of chunk size c on the same two tasks, by varying c ∈ { 16 , 32 , 64 , 128 , 256 , 512 , ∞} , where ∞ indicates the original M EGA without chunking.
The right figure in Figure 4 shows that image data is more sensitive to chunk size than text data.
On the Text task, M EGA -chunk with even a small chunk size c = 16 is able to achieve around 90% accuracy.
On the Image task, M EGA -chunk with c = 16 achieves around 75% accuracy, which is still much better than the vanilla Transformer model.
[Figure]
Text
Image
[Table]
[Table]
Analysis of Attention Functions.
Finally, we evaluate performance with different attention functions.
Table 3 shows the accuracy of the three attention functions.
On text data softmax obtains the best accuracy, while on image data it performs the worst.
The laplace function achieves the best accuracy on Image and also competitive result on Text, being consistently better than relu 2 .
In the following experiments we use softmax for language tasks and laplace for vision and speech ones.
4. R AW S PEECH C LASSIFICATION
To evaluate the capability of M EGA on the long-range modeling of speech signals, we apply M EGA to classify raw speech (with length 16000), rather than using traditional preprocessing (e.g. convert to
MFCC features).
Following Gu et al.
(2022a), we perform speech classification on the SC10 subset of the Speech Commands dataset (Warden, 2018).
We experiment with the M EGA -chunk variant with c = 1000 , since the computation of M EGA and Transformer can not fit in GPU memory.
As shown in Table 4, our M EGA -chunk (base) model with 300K parameters is able to achieve an accuracy of 96.92 that is slightly worse than 97.50 from the state-of-the-art method S4, 3 while by adding 0.18M parameters our M EGA -chunk (big) model performs comparably well with S
. 4. A UTO - REGRESSIVE L ANGUAGE M ODELING
We evaluate M EGA on two established language modeling benchmarks — WikiText-103 (Merity et al., 2017) and enwik8 (Hutter, 2006).
WikiText-103 is a word-level language modeling dataset containing 103M training tokens from Wikipedia articles.
Following previous work (Baevski &
Auli, 2018), we adopt adaptive softmax and input embeddings and use a vocabulary of 260K tokens.
Enwik8 is a character-level language modeling benchmark that has 100M tokens of unprocessed
Wikipedia articles and a vocabulary size of about 200.
At test time, we split the test data into segments and process each segment sequentially.
In Table 5, we compare with previous top-performing models that are designed to take advantage of longer context, including Transformers (Baevski & Auli, 2018;
Al-Rfou et al., 2019)
(XFM-adaptive),
Transformer-XL (Dai et al., 2019)
(XFM-XL) and S4 (Gu et al., 2022a).
On both WikiText-103 and enwik8, we obtain very competitive results, outperforming baselines by a large margin while enjoying much faster (9 × ) inference speed compared to the
Transformer model.
M EGA can also naturally achieve length extrapolation at inference time to any sequences that are longer than those seen during training due to the recurrent design of the EMA layer.
In addition, we can extrapolate to a longer chunk size for M EGA attention due to the use of rotary positional embeddings for training (Su et al., 2021).
The details are provided in in Appendix D.4.
4. N EURAL M ACHINE T RANSLATION
To evaluate M EGA on sequence-to-sequence modeling, we conduct experiments on a standard machine translation benchmark, WMT 2016 English-German news translation (WMT’16), consisting of 4.5M sentence pairs of training data, following Ott et al.
The M EGA models closely follow the architecture of Transformer-base: 6 encoder and decoder layers with model dimension d = 512 .
Table 7 presents the BLEU scores on the test sets of WMT’16 from two directions:
EN → DE and
DE → EN.
For each experiment, we report the average of both tokenized and SacreBLEU (Post, 2018) scores with 5 different random seeds.
M EGA -base significantly outperforms Transformer-base by over 1 . 1 BLEU.
We also report results of M EGA with the Laplace attention function, which slightly but consistently underperforms Softmax.
[Table] [Table]
[Table] [Table]
4. I MAGE C LASSIFICATION
To evaluate M EGA on a large-scale image classification task
[Table]
we conduct experiments on the Imagenet1 k (Deng et al., 2009) dataset, which consists of 1.28M training images and 50K validation images from 1000 classes.
Top-1 accuracy on the validation set is reported in Table 6 to assess various models.
M EGA obtains about 0 .
5 % accuracy improvement over DeiT-B (Touvron et al., 2021).
We mostly follow DeiT’s approach of applying several data augmentation and regularization methods that facilitate the training process, including Cutmix (Yun et al., 2019), Mixup (Zhang et al., 2017), stochastic depth (Huang et al., 2016), repeated augmentation (Hoffer et al., 2020), Rand-Augment (Cubuk et al., 2020), and random erasing (Zhong et al., 2020).
These methods were highly tuned towards optimizing the performance of DeiT, which might be sub-optimal for M EGA .
Exploring the optimal data augmentation and regularization methods for M EGA is an interesting direction for future work.
R ELATED W ORK
To incorporate stronger inductive bias into the attention mechanism, one research direction focuses on developing advanced positional encoding methods, including absolute and relative positional embeddings (Vaswani et al., 2017; Huang et al., 2020; Ke et al., 2020), and relative positional biases (Su et al., 2021; Press et al., 2021).
Another line of research combines the attention mechanism with other neural architectures with intrinsic strong inductive bias, such as convolutional (Gehring et al., 2017;
Dai et al., 2021) and recurrence (Dai et al., 2019; Rae et al., 2020; Lei, 2021).
Recently, many advanced variants of Transformer models ( ‘xformers’ )
(Tay et al., 2020; 2021) have emerged to improve the time and memory efficiency.
Popular techniques include sparse attention patterns (Parmar et al., 2018; Beltagy et al., 2020; Kitaev et al., 2020), low-rank approximations of the attention matrix (Wang et al., 2020; Ma et al., 2021), and approximations through kernelization (Choromanski et al., 2020; Peng et al., 2021).
Although these models demonstrate better asymptotic complexity for long sequences, their efficiency gains are less prominent for moderate length sequences and their performance lags behind Transformers with regular attention.
In addition, as EMA and more general state space models such as S4 can be regarded as a convolution transform with kernel size equal to the sequence length, M EGA is also relevant with CNNs with continuous kernels (Romero et al., 2022).
C ONCLUSION
We have introduced M EGA , a simple, efficient and effective neural architecture used as a drop-in replacement for regular multi-head attention.
By leveraging the classic exponential moving average (EMA) approach, M EGA is capable of incorporating stronger inductive biases into the attention mechanism.
Moreover, the EMA approach enables the design of M EGA -chunk, an efficient variant of
M EGA with linear complexity.
On five sequence modeling tasks across various data types, M EGA achieves impressive improvements over a variety of strong baselines.
This leads to a potential direction of future work to apply M EGA for multi-modality modeling. A CKNOWLEDGEMENT
We thank the anonymous reviewers for their comments.
This material is based on research sponsored by Air Force Research Laboratory (AFRL) under agreement number FA8750-19-1-1000.
The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation therein
A PPENDIX : M EGA : M OVING A VERAGE E QUIPPED G ATED A TTENTION A E FFICIENT C OMPUTATION OF M ULTI - DIMENSIONAL D AMPED EMA
Note that the computation of the multi-dimensional damped EMAs of different dimensions are entirely independent of each other.
Without loss of generality, we set d = 1 and omit the dimension index j in the following formulations.
We denote the initial hidden state as h 0 .
The multi-dimensional damped EMA defined in (5) can be vectorized into the following formulation:
[Equation]
[Equation]
where α , δ , and η ∈ R h .
u t = β x t ∈ R h and h t ∈ R
h is the EMA hidden state at timestep t .
Let’s denote ϕ = 1 − α ⊙ δ .
Then, unrolling the above two equations explicitly yields:
[Equation]
This can be written into a vectorized formula:
[Equation]
[Equation]
where ∗ is the convolution transform with kernel K ∈ R n :
[Equation]
In the proposed multi-dimensional damped EMA, K can be efficiently computed by the Vandermonde product.
With K provided, the output y in (24) can be computed efficiently with FFTs.
B P ROOF OF T HEOREM
Proof We split γ into h heads in the same way as Q , K , and V :
[Equation]

Then we have
[Equation]

To prove Theorem 1, we need to find γ such that
[Equation]
where ⊘ is the element-wise divide operation.
Since G ( X ) is a universal approximator and Q , K ,
V and a are all transformed from X , γ can theoretically recover a ( i ) T V ( i ) ⊘ a T V ( i ) , ∀ X .
C L APLACE A TTENTION F UNCTION
To approximate the squared ReLU function with the Laplace function in (16), we need to select proper coefficients µ and σ .
We derive the values of µ and σ by solving the following two equations at x = √ 2 :
[Equation] [Equation]
The Eq.
(26) delivers µ = (cid:112) 1 / 2 and Eq.
27 subsequently provides σ = (cid:112) 1 / 4 π .
Figure 5 visualizes the two functions.
[Equation]
f ( x ) f = relu 2 f = laplace
Epochs
V a li d a cc u r a c y ( % ) relulaplace
Besides performance improvements, we also investigate the stability of the two attention functions.
We conduct experiments on the LRA Pathfinder task with M EGA models with the two functions.
Figure 5 presents the accuracy on the validation set across training epochs.
We observe that Laplace is much more stable than ReLU 2 .
D E XPERIMENTAL D ETAILS D. D ETAILS OF M EGA I MPLEMENTATION
Uni-directional vs. Bi-directional EMA.
In all the experiments, bi-directional EMA is applied when M EGA is used in the encoder mode while uni-directional one is used in the decoder mode with auto-regressive property.
Relative Positional Bias.
We only used ROPE (Su et al., 2021) for the language modeling tasks because we conducted experiments of length extrapolation to longer attention chunks at inference time D.4.
For other tasks, we simply use very simple learnable relative biases, i.e. one learnable parameter for each relative positional distance (e.g. -3, +4).
D. L ONG R ANGE A RENA (LRA)
For all tasks, we closely follow Tay et al.
(2020) for details such as data preprocessing, data split, etc.
The hyper-parameters of M EGA models on these tasks are listed in Table 8.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Table]
D. R AW S PEECH C LASSIFICATION
Following Gu et al.
(2022a), we perform speech classification on the SC10 subset of the Speech
Commands dataset (Warden, 2018), which is a 10-class classification task.
The chunk size of
M EGA -chunk is 1000.
Other hyper-parameters are listed in Table 8.
D. L ANGUAGE M ODELING
Training details We use the data of WikiText-103 and enwik8 and their splits provided by Dai et al.
At training time, we split the training data into segments; each segment contains m consecutive chunks, where the chunk size is the effective attention length.
m is a random integer variable uniformly sampled from [ cl, ch ] .
We use [ cl, ch ] =
[2 , 6] for WikiText-103 and [ cl, ch ] =
[2 , 4] for enwik8.
Other training hyperparameters including optimizer, learning rate scheduler and architecture are presented in Table 9.
Length extrapolation at inference time We employ M EGA -chunk (§3.4) for training and set the attention chunk size to be 1024 and 2048 for WikiText-103 and enwik8 respectively.
To use a longer
Mega attention length at inference time than the one used at training time (i.e. 1024 or 2048), we apply rotary positional embedding (Su et al., 2021) to the attention sublayer.
At test time, we split the test data into K segments and sequentially process each segment by m chunks, i.e. the maximum context length of each segment is # test tokens K .
In Table 5, we report test results that use longer chunk sizes (attention lengths) of 2048 and 4096 for WikiText-103 and enwik8 respectively.
M EGA can naturally extrapolate at inference time to sequences longer than those seen during training due to the recurrent design of the EMA layer.
That design enables the inputs of each chunk to access the historic context through EMA as illustrated in Figure 3.
On the other hand, due to the use of rotary positional embeddings, attention can be performed on longer chunk sizes at test time than those seen during training.
We hope these two types of length extrapolation are clear to readers.
We provide the
[Figure]
Attn
[Figure]
length
[Figure]
Ablations on context lengths First, we fix the chunk size to be 2048 and vary K within [100 , 75 , 50 , 25 , 15 , 10 , 5] corresponding to maximum context tokens of [2 . 5 K , 3 .
3 K , 4 .
9 K , 9 .
8 K , 16 K , 25 K , 49 K ] .
We plot the test PPL as we increase the context length in the left of Figure 6.
Although at training time, the maximum context length the model has seen is 6144, M EGA can extrapolate to longer context lengths.
The plot shows that PPL decreases as the context length is increased and the improvements saturate when the context length is longer than 25K.
This is consistent with the observations in Press et al.
Ablations on attention chunk sizes Next, we fix the context length to be 25K and increase the chunk size from 512 to 3072.
As shown in the right side of Figure 6, M EGA consistently improves as we increase the attention length although it only uses an attention length of 1024 during training.
This contradicts with the findings in Alibi (Press et al., 2021), which finds that rotary embeddings don’t generalize to longer lengths and result in higher PPL.
[Table] [Table]
[Table]
linear linear
[Table]
[Table] [Table]
Dropout 0.
FFN
[Table] [Table]
Clipping 1.
Warmup steps 24K 24K
[Table]
400K 400K
Decoder
[Table]
[Table] [Table]
size
[Table] [Table]
EMA
[Table]
h )
[Table]
size Total Parameters 252M 39M
D. M ACHINE T RANSLATION
The WMT 2016 English-German dataset contains 4.5M parallel sentence pairs for training.
We following the standard setting (Ott et al., 2018), using Newstest2013 as the validation set and
Newstest2014 as the test set.
The dataset is pre-processed following (Ma, 2020), using the scripts from FairSeq package (Ott et al., 2019).
4 We share the source and target vocabularies within the language pair, with 32K byte pair encoding (BPE) types (Sennrich et al., 2016).
We report the average of both tokenized and SacreBLEU 5 (Post, 2018).
The hyper-parameters of Transformer and M EGA models are listed in Table
[Table] [Table]
signature: nrefs:1|case:mixed|eff:no|tok:13a|smooth:
D.
I MAGE C LASSIFICATION
Hyper-parameters are listed in Table 11.
We closely follow Touvron et al.
(2021) by reusing most of the their hyper-parameters
[Table] [Table]