Mega: Moving Average Equipped Gated Attention
Authors: Xuezhe Ma
Chunting Zhou
Xiang Kong
Junxian He
Liangke Gui
Graham Neubig
Jonathan May
Luke Zettlemoyer
Venue: ICLR 2023 poster (2023)
Keywords: Neural Architecture
Attention
Exponential Moving Average

Abstract: 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 Mega, 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 Mega 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 Mega achieves significant improvements over other sequence models, including variants of Transformers and recent state space models.


Published as a conference paper at ICLR 2023
MEGA: MOVING AVERAGE EQUIPPED GATED ATTENTION
Xuezhe Ma* ISI, USC
Chunting Zhou+ Meta Al
Xiang Kong LTI, CMU
Junxian He SJTU
Liangke Gui LTI, CMU
Graham Neubig LTI. CMU
Jonathan May ISI, USC
Luke Zettlemoyer Meta Al
ABSTRACT
The design choices in the Transformer attention mechanism, including weak induc- tive bias and quadratic computational complexity, have limited its application for modeling long sequences. In this paper, we introduce MEGA, simple, theoreti- cally 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 MEGA 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 lan- guage modeling, and image and speech classification, show that MEGA achieves significant improvements over other sequence models, including variants of Trans- formers and recent state space models.
INTRODUCTION
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 devel- oped, including convolutional neural networks (CNNs) (Kim, 2014; Strubell et al., 2017), recurrent neural networks (RNNs) (Goller & Kuchler, 1996; Hochreiter & Schmidhuber, 1997), Transform- ers (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 al., 2018, 8), language understanding (Devlin et al., 2019; Liu et al., 2019), image recogni- tion (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 (Bah- danau 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 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 (MEGA) to solve the two weaknesses simultaneously. The key idea is 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-dimensionall damped form of EMA with learnable coefficients ($3.1). and subsequently develop the moving average equipped gated attention mechanism by integrating the EMA with variant of the single-head gated attention (Hua 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
"Equal Contribution. Contact: xuezhema @1 .duu

---

Published as a conference paper at ICLR 2023
Table 1: Experimental results of Transformer (XFM), S4 and MEGA on five sequence modeling benchmarks, including long range arena (LRA), machine translation (WMT16 en-de), language modeling (WikiText- 103), image classification (ImageNet-Iee 1k), raw speech classification (SC-Raw).
LRA (Ace. 1)
WMT16 (BLEU
WT103 (PPL. 4)
ImageNet (Acc. f)
SC (Acc.
XFM
59.24
27.97
18.66
81.80
x
S4
85.86
20.95
97.50
MEGA
88.21
29.18
18.07
82.31
96.92
from the incorporated moving average mechanism, we further propose a variant of MEGA with linear complexity, named MEGA-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, MEGA significantly outperforms a variety of strong baseline models, in terms of both effectiveness and efficiency (54) (see Table 1).
BACKGROUND
We use X= (x1, X2, x,) E Raxd to denote a sequence of input representations with length Let Y = {, ,2....,,........................ .n) E ied be the sequence output representations of each layer with the same length as the input X. In this paper, we assume the representations of the input and output sequences have the same dimension
2.1 SELF-ATTENTION MECHANISM
The traditional self attentttentoon mechanism a function:
(QK") (2K"* T(X)
Y Autn())= =
((X)
V
(1)
where Attn Raxd Raxd is the self-attention function.
XW, bys K XW
bg. and V = -WW, + by are the sequences of queries, keys and values, with learnable parameters Wg: We, W. E Rdxd and by. bR. by Rdd. f(-) an attention function, e.g. the softmax function fsoftmax (-) (Bahdanau et al., 2015), or the recently proposed squared ReLU function fremz (-) (So et al., 2021; Hua et al., 2022). r(X) is scaling term, which is commonly set to T(X) = vd for fsottmax (-), or r(X) = n for fretue (.). The commonly used multi-head variant of attention performs the attention function à times in parallel.
We can define matrix A f( f( E RRax following (I), which called the attention matrix. Since it specifies pairwise dependency weights, the matrix in principle delivers a flexible and powerful mechanism 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(h??) time and space, which becomes significant bottleneck.
2.2 EXPONENTIAL MOVING AVERAGE (EMA)
The moving average is a classic approach for sequential data modeling, which has been widely used in time series data smooth short-term fluctuations and highlight long-term trends cycles. The Exponential Moving Average (EMA) (Winters, 1960; Hunter, 1986), a special case of moving average, applies weighting factors that decrease exponentially:
Q C x, + (1 -1-61 -x) xx) ② yt-1,
(2)
where a 1) 1)dd is the EMA coefficient representing the degree of weighting decrease, and O is the element-wise product. A higher a discounts older observations faster (see Figure 1).
Using EMA places strong inductive bias the learning of pairwise dependencies: the depen- dency weight between two tokens decreases exponentially over time with an input-agnostic decay factor a. This property favors local dependencies, and limits long-distanee 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 (FFT's (see Appendix for details).
2

---

Published as a conference paper at ICLR 2023
y -1
y
a(1 o)
a(1 o)
0)
0(1
0)
a(1
o)
3
-1-1
St
3
-1-1
Figure 1: Illustration of the exponential moving average (EMA) approach, which averages the input
values X with weights decaying exponentially over timesteps.
2.3 WHY COMBINE ATTENTION 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).
MOVING AVERAGE EQUIPPED GATED ATTENTION (MEGA)
In this section, we describe in detail our proposed method, moving average equipped gated attention (MEGA). We first introduce multi-dimensions damped EMA ($3.1), which is a key component combined with the single-head gated attention in MEGA ($3.2), and discuss the relationship between MEGA 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 MEGA block, including feed-forward and normalization layers ($3.3), and present MEGA-chunK, variant of MEGA that splits input sequences into fixed chunks, reducing time and space complexity from quadratic to linear ($3.4), At last, provide theoretical justification for the design of single-head gated attention ($3.5).
3.1 MULTI-DIMENSIONAL DAMPED 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. a in (2)) produces robust dependency modeling. Inspired by this, MEGA allows the damping of the influence of the previous time step:
ye ① + a ④ 0)00) Y1-1,
where 8 E (0, 1)d is the damping factor.
Multi-dimennsonall Damped EMA. To further improve the expressiveness of EMA, we introduce multi-dimennional variant of EMA. Concretely, first expand each dimension j 21.,..... d) of the input sequence X individually into h dimensions via an expansion matrix B E |Rdxk.
( B,xt.j u
(4)
where G3 R' is the j-th row of B. u. R² is the expanded h-dimensional vector for the j-th dimension at timestep Correspondingly, we extend the shape of a and from a one-dimensional vector swo-dimensional matrix, i.e. a, d E Rd where az 8, E R² denote the j-th row of and 6, respectively. Then, for each dimension j. the damped EMA is applied to the h-dimensional hidden space:
h." = a 2 u
(1 (1-ay (1 aj 2 о 9
Yt. = n h(
(5) where h.()) E 2 is the EMA hidden state for the j-th dimension at timestep t. n E Rdxh is the projection matrix map the h--dimensional hidden state back to 1-dimensional output Yt.j IR. N; E R is the j-th row of 77. The output Y from (5) is denoted as Y - EMA(X). Because we do not need to explicitly compute ()) to get the output yt.j. and the time and space complexity is similar the standard EMA in (2) (see Appendix for the details).
3

---

Published as a conference paper at ICLR 2023
output
Add Norm
0
Feedfooward
QKT
bail
7(X)
Gate
Gate
XL
Single-head Attentian
QCK
Mega Layer: EMA & Gated Attention
QRRK
Dense SiLU
Dense SiLU
EMA output x'
Input Embedding
Layer input
x
(a) Mega architecture.
(b) Mega layer.
(c) Single-head attention unit.
Figure 2: MEGA model architecture. (a) shows the overall architecture. (b) illustrates the gated
attention sub-layer equipped with EMA, while (c) displays the details of single-head attention unit.
3.2 MOVING AVERAGE EQUIPPED GATED ATTENTION
The gated attention mechanism MEGA adopts the Gated Recurrent Unit (GRU: Cho et al. (2014)) and Gated Attention Unit (GAU; Hua et al. (2022)) the backbone architectures, with an EMA-based sub-layer embedded into the calculation the attention matrix. Formally, we first use the output from (5) compute the shared representation in GAU:
X' = (M)(X)
E Rnxd
(6)
= Aillu W + b,)
E R
where X' can be regarded as the updated or contextual input, because it encodes contextual infor- mation through EMA. Z is the shared representation with z dimensions, with projection matrix W, E Rdx and bias term by R Osilu the self gelf-gated activation function (SiLU) (Ramachandran et al., 2017; Elfwing al., 2018). Following GAU, the query and key sequences are computed by applying per-dimension scalars and offsets Z, and the value sequence from the original X:
Q Rig ④ + 1gg
(8)
K O Z + μ4
R
(9)
V dalu XW, bo)
R
(10)
where Kgg Mg. Rk. by o is the expanded intermediate dimension for the value sequence. The output of attention is:
R= are the learnable scalars and offsets of queries and keys, respectively.
0 0=1 ( + brel V 7(X)
QKT
E R R
(11)
The graphical specification is displayed in Figure (c). biel [R=x= is the relative positional bias. We choose brel from existing approaches, including T5 (Raffel et al., 2020), RoPE (Su et al., 2021), TUPE (Ke et al., 2020) and ALIBi (Press al., 2021).
Subsequently, MEGA introduces the reset gate y, the update gate io and computes the candidate activation output H:
7 = duata X'W, + by,)
E R nxx
(12)
4 Øsigmoid (X'W_ + bxp)
Raxd
(13)
A dsnu X WW + (r o 0)U, + bn.)
E gn
(14)
The final output Y computed with the update gate ;p:
Y = 40 A bp) o X
E Baad
(15)
The graphical architecture of MEGA sub-layer is visualized in Figure 2 (b).
4

---

Published as a conference paper at ICLR 2023
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 frelue()) (-) 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 frelu (-) is that neither its range its gradient is bounded, leading to unstable model training (see Appendix C.I for details). To address this issue, we propose the Laplace attention function:
(1aplace pr, o) 0.5 x erf( ov2
(16)
where e-f is the error function. and are two coefficients that we adjust to approximate fielua, yielding μ /1/2 and /1/4x (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 y, the update gate sp and the candidate activation output H in (12-14) is reminiscent of GRUs (Cho al., 2014). The main difference is that in GRU the two gates are applied between the hidden states of the current and previous timesteps, while in MEGA they are applied between the outputs from EMA and gated attention sub-layers. In addition, the output gating mechanism (15) is similar to the gated residual connection proposed in Parisotto et al. (2020); Xu et al. (2020).
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). MEGA 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 A in (14). Experimental gains over Flash demonstrate the effectiveness of this design chice ($4.1).
The multi-dimensional damped EMA can seen as simplified variant of a state space model. From this perspective. MEGA is also closely related to S4 (Gu et al., 2022a), 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 MEGA applies diagonalization on the state matrix, which is similar concurrent work S4D (Gu et al., 2022b). The differences between the MEGA 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.3
MEGA BLOCKS
The MEGA layer (moving average equipped gated attention) is used drop-in-reellacement for regular attention in Transformer. It is followed by position-wise feed-forward networks (FFNs) and normalization layers compose one MEGA block. As the gated residual connection has already been included in (15), we omit the original residual connection and directly apply normalization layer to Y. Concretely,
Y = NormmMeeaa (X))
(17) where Y' is the output of the MEGA block. The overall architecture of a MEGA block is shown in Figure (a). In Transformer, the hidden dimension of FFNs usually set to dFFN 4d. To retain a similar model size with each Transformer block, we reduce the hidden dimension of FFN to dFEN = 2d and set the expanded dimension 2d for the value sequence in (10) throughout this paper, unless specified otherwise.
Y' = Norm((FN(()) + Y)
3.4 MEGA-CHUNK: MEGA WITH LINEAR COMPLEXITY
So far we have only focused on introducing stronger inductive bias into attention, which still has quadratic computational complexity. In this section, we propose MFGA--hunk, a variant of MEGA 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 where = n//c is the number of chunks.' 1 The attention operation
c. e.g. Q = .Q,
Keys and values are split in the same way.
5

---

Published as a conference paper at ICLR 2023
chunk-wise Attention
EMA sublayer
12)
12)
22
21
23
Figure 3: Illustration the MEGA-chunk model with two chunks of length 3.
in (II) is individually applied to each chunk, yielding linear complexity O(ke²) O(ne) 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 MEGA mitigates this problem by capturing local contextual information near each token, whose outputs are used as the inputs the attention sub-layer. As result, the effective context being exploited by chunk- wise attention can go beyond the chunk boundary. Figure illustrates the largest dependency length captured by one MEGA-chunk block.
3.5 THEORETICAL JUSTIHICATION OF SINGLE EANLE--EED GATED ATTENTION
Single-head gated attention has been empirically shown performant vanilla multi-head atten- tion Liu et (2021); ; Hua et al. (2022), without any discussions on its theoretical insights. 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:
Q Q()). K K(X),,; V = V(X)
(18)
where Q, KC, V are three transformations, such as linear projections. Let q Q= (--................. q.. be single query vector (q E RR)), and = A(q. qq. K) denote the corresponding attention weights of q, where A is the attention transformation, i.e. f(-) in (11). For multi-head attention, common implementation is to split the query into h heads across the model dimension: q = =q ").").....q" (
where (8) E Rd/h and E (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 (() K())) Then, the outputs of single-head
and multi-head attention are, respectively:
(() (1)
aaiT V()
OSHA a VV=
O OMHA
(19)
a" (h)
(N) (h)
is straightforward see that O is more expressive than OSHA, because OMHA leverages sets of attention weights.
In the single-head gated attention, we introduce gate vector > ((X) for each q. and the output of single-head gated attention is OShGa OSha @ Yy. The following theorem reveals the equivalence of SHGA and O MHA w.r.t expressiveness (proof in Appendix B):
Theorem Suppose the transformation & is a universal approximator Then, VX, By = ((X) (X) s.f.
OSHGA = OжHA
(20)
Theorem indicates that by simply introducing the gate vector, OSHGA as expressive as O 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).
EXPERIMENTS
To evaluate MEGA. 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

---

Published as a conference paper at ICLR 2023
Table (Long Range Arena) Accuracy on the full suite of long range arena (LRA) tasks, together with training speed and peak memory consumption comparison on the Text task with input length of 4K. indicates results replicated by us.
Models
ps
Text
Retrieval
Image
Pathfinder
Path-X
Avg.
Speed
Mem.
XFM
36.37
64.27
57.46
42.44
71.40
X
54.39
XFM
37.11
65.21
79.14
42.94
71.83
x
59.24
1x
1x
Reformer
37.27
56.10
53.40
38.07
68.50
X 50.67
0.8x
0.24x
Linformer
35.70
53.94
52.27
38.56
76.34
X
51.36
5.5x
0.10x
BigBird
36.05
64.02
59.29
40.83
74.87
X 55.01
1.1x
0.30x
Performer
18.01
65.40
53.82
42.77
77.05
X 51.41
5.7x
0.11x
Luma-256
37.98
65.78
79.56
47.86
78.55
61.95 X
4.9%
0.16x
S4-v1
58.35
76.02
87.09
87.26
86.05
10
80.48
S4-v2
59.60
86.82
90.90
88.65
94.20
96.35
86.09
S4-v21
59.10
86.53
90.94
88.48
94.01
96.07
85.86
4.8x
0.14x
MEGA 63.14
90.43
91.25
90.44
96.01
97.98
88.21 85.66
2,9x
0.31 x
MEGA-chunk
58.76
90.19
90.97
85,80
94.41
93.81
5.5x
0.13x
4.1
LONG-CONTEXT SEQUENCE MODELING
We begin our experiments with an evaluation on the Long Range Arena (LRA) benchmark recently introduced by Tay et al. (2021), which designed for evaluating sequence models under the long- context scenario. This benchmark has six tasks: ListOps (Nangia & Bowman, 2018), byte-level text classification (Text; Maas al. (201 ))) byte-level document retrieval (Retrieval; Radev et al. (2013)), image classification
sequences of pixels (Image; Krizhevsky et al. (2009)), 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 variety of data types and modalities.
Table 2 compares MEGA against several baselines, including Transformer and its efficient variants, and the slate-of-the-al S4 models (both version (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 MEGA has similar number of parameters with $4-v1. For each experiment, we report the average over 5 runs with different random seeds.
On all the six tasks, MEGA substantially outperforms all the baselines. We also evaluate MEGA-chunk on each task, by setting the chunk size e = 128 for all the tasks, except Path-X where c 4096. We observe that MEGA-chunk consistently performs well, particularly on the three language tasks. We also examine the speed and memory efficiency of MEGA on the byte-level classification task with the input length of 4K. MEGA-chunk is highly efficient, which is about 5.5 times faster and consumes only 13% as much memory as the vanilla Transformer. In addition, MEGA with full attention field is also much more efficient than Transformer, benefiting from single-head gated attention.
Analysis of Hulti-dimmensinnl Damped EMA. To demonstrate the effectiveness of the multi- dimensional damped EMA component in MEGA, we performs ablation studies two LRA tasks byte-level text classification (Text) and image classification on sequences of pixels (Image). We train MEGA models with EMA dimension EMA component. From the left figure in Figure 4, we see that without the EMA component, model aceuracy on both the two tasks declines rapidly. Meanwhile, the significant improvements of MEGA with single dimensional EMA (h 1) demonstrate the importance of EMA inductive bias.
{0, 1, 2, 4, 8, 16, 32), where h indicates removing the
Analysis of Chunk Size. We further analyze the impact of chunk size e on the two tasks, by varying E {16, 32, 64, 128, 256, 512, jo). where xo indicates the original MEGA without chunking. The right figure in Figure 4 shows that image data more sensitive to chunk size than text data. On the Text task, MEGA-Chin Chuuk- with even a small chunk size c 16 is able to achieve around 90% accuracy. On the Image task, MEGA-chunk with e = 16 achieves around 75% accuracy, which is still much better than the vanilla Transformer model.
TTe $4-v2 used larger model sizes and hetter-tuned syper-parammters than 14-v1. We have also experimented with SRU++ (Lei, 2021) on Pathfinder but failed converge on this dataset after tuning hyperparameters.
7

---

Published as a conference paper at ICLR 2023
90
90
Accuracy (%)
Accuracy (%)
Text
Image
85-
85
softmax
90.43
89.87
80
image text
relu²
90.08
90.22
75-
laplace
90.22
90.43
2
8
16
32
16
32
64
128
256
512
EMA dimenssio
Chunk size
Figure 4: Ablations on EMA dimension and chunk size.
Table 3: Attention functions.
Analysis of Attention Functions. Finally, we evaluate performance with different attention func- tions. 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 Image and also competitive result Text, being consistently better than relu?. In the following experiments we use softmax for language tasks and laplace for vision and speech ones.
4.2 RAW SPEECH CLASSIFICATION
To evaluate the capability of MEGA on the long-range modeling of speech signals, we apply MEGA 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 MEGA--hunk variant with 1000, since the computation of MEGA and Transformer can not fit in GPU memory. As shown in Table 4, our MEGA--hunk (base) model with 300K parameters able to achieve an accuracy of 96.92 that is slightly worse than 97.50 from the state-of-the-at method S4, while by adding 0.18M parameters our MEGA-chunk (big) model performs comparably well with S4.
4.3 AUTO-REGRESSIVE LANGUAGE MODELING
We evaluate MEGA on two established language modeling benchmarks - WikiText- 103 (Merity al., 2017) and enwik8 (Hutter, 2006). Wiki'Text-109 103 is word-level language modeling dataset containing 103M training tokens from Wikipedia articles. Following previous work (Baevski & Auli, 2018), adopt adaptive softmax and input embeddings and use a vocabulary of 260K tokens. Enwik8 is character-level language modeling benchmark that has 100M tokens of unprocessed Wikipedia articles and a vocabulary size of about 200. test time, we split the test data into segments and process each segment sequentially. In Table 5, we compare with previous nop-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-XI.) and S4 (Gu et al., 2022a). On both WikiText-100 and enwik8, we obtain very competitive results, outperforming baselines by a large margin while enjoying much faster (9x) inference speed compared to the Transformer model. MEGA 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 MEGA attention due to the use rotary positional embeddings for training (Su et al., 2021). The details are provided in in Appendix D.4.
4.4 NEURAL MACHINE TRANSLATION
To evaluate MEGA on sequencc-to-seqqueee modeling, we conduct experiments on standard machine translation benchmark, WMT 2016 English- German news translation (WMT 16). consisting of 4.5M sentence pairs of training data, following Ott al. (2018). The MEGA models closely follow the architecture of Transformer-base 6 encoder and decoder layers with model dimension d = 512.
English-German
Table 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) 8) scores with 5 different random seeds. MEGA -bbe significantly outperforms Transformer- by over 1.1 BLEU, We also report results of MEGA with the Laplace attention function, which slightly but consistently underperforms Softmax.
*Our $4 number is obtained by running the official S4 code and is a bit worse than the original reported number (98.32), due different data splits the order of cs. 1 sttdii is deterministic across machines.
8

---

Published as a conference paper at ICLR 2023
Table 4: (SC-Raw) Accuracy on
Table 5: (Language Modeling) PPL (L) on WikiText-103 and bpe (4) on enwik8.
Speech Commands.
SC-Raw
WikiTeext1003
enwik8
Model
#Param.
Acc.
Model
#Param.
PPL
Speed
#Param.
bpc
XFM $4
786K
XFM-adaptive
247M
18.66
5.6k l/s
300K
97.50
XFM-XL
257M
18.30
41M
1.06
MEGA (base)
300K
96.92
S4
249M
20.95
MEGA (big)
476K
97.30
MEGA
252M
18.07
48k t/s
39M
1.02
Table (ImageNet) Top-1 accuracy.
Table 7: (WMT'16) Test BLEU scores.
Model
Img. size
#Param.
Acc.
EN-DE
DE-EN
Model
Token.
Sacre.
Token.
Sacre.
ResNet-152 152
224
60M
78.3
VIT-B
384²
86M
77.9
XPM-base
27.30
XPM-hase
27.97
27.33
31.92
31.33
DeîT-B
2242
86M
81.8
MRGA-softmax
29.18
28.47
32,90
32.35
MEGA
2242
90M
82.3
MEGA-laplace
28.95
28.27
32.81
32.22
4.5
IMAGE CLASSIHICATION
To evaluate MEGA on a large-scale image classification task. we conduct experiments on the Imagenet- 1A (Deng al., 2009) dataset, which consists 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. MEGA 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 facili- tate the training process, including Cutmix (Yun et al., 2019), Mixup (Zhang et al., 2017). stochastic depth (Huang et al., 2016), repeated augmentation (Hoffer 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 MRGA. Exploring the optimal data augmentation and regularization methods for MEGA is an interesting direction for future work.
RELATED WORK
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 al., 2020), and relative positional biases (Su et al., 2021; Press et al., 2021). Another line research combines the attention mechanism with other neural architectures with intrinsic strong inductive bias, such 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 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 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 S4 can be regarded convolution transform with kernel size equal to the sequence length, MEGA is also relevant with CNNs with continuous kernels (Romero et al., 2022).
CONCLUSION
We have introduced MEGA, 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, MEGA is capable of incorporating stronger inductive biases into the attention mechanism. Moreover, the EMA approach enables the design MEGA-chunk, an efficient variant of MEGA with linear complexity. On five sequence modeling tasks across various data types, MEGA achieves impressive improvements over variety of strong baselines. This leads to a potential direction of future work apply MEGA for multi-modality modeling.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENT
We thank the anonymous reviewers for their comments. This material is based on research spon- sored by Air Force Research Laboratory (AFRL) under agreement number FA8750-19-1-100 The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation therein.
REFERENCES
Rami AI-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. Character-level language modeling with deeper self-attention. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 3159-3166, 2019.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2018.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. In International Conference on Learning Representations (ICLR), 2015.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprini 05X1v:2004.05/50, 2020.
Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. Syntax, Semantics and Structure in Statistical Translation, pp. 103, 2014.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Ekin Cubuk, Barret Zoph. Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops, pp. 702-703, 2020.
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G Carbonell, Quoc Le. and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1978-2988, 2019.
Zihang Dai, Hanxiao Liu, Quoc Le, and Mingxing Tan. Coatnet: Marrying convolution and attention for all data sizes. Advances in Neural Information Processing Systems, 34:3965-3977, 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 computer vision and pattern recognition, pp. 248-255. Ieee, 2009.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-4188 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR-2020), 2020.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. higmoid-weighte linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. Convolutional sequence to sequence learning. In International conference on machine learning (ICML-2017), pp. 1243-1252. PMLR, 2017.
10

---

Published as a conference paper at ICLR 2023
Christoph Goller and Andreas Kuchler. Learning task-dependent distributed representations by backpropagation through structure. Neural Networks, 1996., IEEE International Conference on, volume 1, pp. 347-352. IEEE, 1996.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances Neural Information Processing Systems, 33: 1474-1487, 2020.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations (ICLR-2022), 2022a.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. arXiv preprint 22Xi::2206.11866 2022b.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural compulation, 9(8):
1735-1780, 1997.
Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, and Daniel Soudry. Augment your batch: Improving generalization through instance repetition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8129-8138, 2020.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, (55:359-3666, 1989.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International Conference on Machine Learning (ICML-2022), pp. 9099-9117. PMLR, 2022.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In European conference computer vision, pp. 646 661. Springer, 2016.
Zhiheng Huang. Davis Liang, Peng Xu, and Bing Xiang. Improve transformer models with better relative position embeddings. In Findings of the Association for Computational Linguistics (EMNLLP-020), pp. 3327-3335, 2020.
J Stuart Hunter. The exponentially weighted moving average. Journal of quality technology, 18(4): 203-210, 1986.
Marcus Hutter. The human knowledge compression contest. 2006. URL http://prize: hutterl.nee net/.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-588) 2021.
Guolin Ke, Di He, and Tie-Yan Liu. Rethinking positional encoding in language pre-training. In International Conference on Learning Representations (IC1,R-2020), 2020.
Yoon Kim. Convolutional neural networks for sentence classification. Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (204),-20044), pp. 1746- 1751, Doha, Qatar, October 2014. Association for Computational Linguistics.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.0443 2020.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. Technical Report. University of Toronto, 2009.
Tao Lei. When attention meets fast recurrence: Training language models with reduced compute. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7633-7648, 2021.
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learning long-ran long-range spatial dependencies with horizontal gated recurrent units. In S. Bengio, H. Wal- lach, Larochelle, K. Grauman, N. Cesa-Bianche and R. Gameett (eds.). Advances in Neural Information Processing Systemx, volume 31. Curran Associates, Inc., 2018
11

---

Published as a conference paper at ICLR 2023
Hanxiao Liu, Zihang Dai. David So, and Quoc V Le. Pay attention mlps. Advances in Neural Information Processing Systems, 34:9204 9215, 2021.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 1907.1197..110000 2019.
Xuezhe Ma. Apollo: An adaptive parameter-wise diagonal quasi-newton method for nonconvex stochastic optimization. arXiv preprini arXiv:2009. / 13586, 2020.
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.
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 amnual meeting of the association for computational linguisties: Human language technologies, pp. 142-150, 2011.
Ali Madani, Bryan McCann, Nikhil Naik, Nitish Shirish Keskar, Namrata Anand, Raphael R Eguchi, Possu Huang, and Richard Socher. Progen: Language modeling for protein generation. bioRxiv, 2020.
Eddie McKenzie and Everette S Gardner Jr. Damped trend exponential smoothing: modelling viewpoint. International Journal of Forecasting. 26(4):661 (6))6661-665, 2010.
Harsh Mchta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint arXiv:2206. 13447, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017. URL https: //
openreview n forum? (
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, pp. 92-99, 2018.
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation: Research Papers, pp. 1-9, 2018.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-IILT 2019: Demonstrations, 2019.
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. In International conference machine learning, pp. 7487-7498. PMLR, 2020.
Sejun Park, Chulhee Yun, Jaeho Lee, and Jinwoo Shin. Minimum width for universal approximation. arXiv preprint arXiv:2006.0883: 2020.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. International Conference on Machine Learning (ICML-2018), pp. 4055 4064. PMLR, 2018.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference Learning Representations, 2021. URL https: /ttp::///ppnnne wpwprevierreeeeeeee/eeeee. net forum? 1d=
Matt Post. A call for clarity in reporting BLEU scores. In Proceedings of the Third Conference Machine Translation: Research Papers, pp. 186-191, Belgium, Brussels, 2018. Association for Computational Linguistics.
12

---

Published as a conference paper at ICLR 2023
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations (ICLR-2021), 2021.
Dragomir Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 47(4)) 19-944, 2013.
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), 2020.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. Mach. Learn. Res., 21 21(44)):1-67, 2020.
Prajit Ramachandran, Barret Zoph, and Quoe Le. Swish: self-gated activation function. arXiv preprint arXiv: 171v://100.05411, 7(1):5, 2017.
David W Romero, David M Knigge. Albert Gu, Erik J Bekkers, Efstratios Gavves, Jakub M Tomczak, and Mark Hoogendoorn. Towards general purpose enn for long range dependencies in nd. arXiv preprint arXiv:2206.nm 03333922000033888 2022.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers), pp. 1715-1725, 2016.
David So, Wojeiech Manke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoe V Le. Searching for efficient transformers for language modeling. Advances in Neural Information Processing Systems, 34:6010-6022, 2021.
Emma Strubell, Patrick Verga, David Belanger, and Andrew McCallum. Fast and accurate entity recognition with iterated dilated convolutions. In Proceedings of the 2017 Conference on Empirical Methods Natural Language Processing (EMNLP-2017), pp. 1670-2680) Copenhagen, Denmark, September 2017. Association for Computational Linguistics.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864, 2021.
Ivan Svetunkov. Complex exponential smoothing. Lancaster University (United Kingdom), 2016.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint arXiv:2009.06332, 2020.
Yi Tay, Mostafa Dehghami, 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. In International Conference on Learning Representations, 2021. URL https: // /openreview..r net htt////ormmddoo/ee drrrmmm/dd-yyye-grC2
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference Machine Learning, pp. 10347-10357. PMLR, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.0477 2020.
Pete Warden. Speech commands: A dataset for limited-vocabullry speech recognition. arXiv preprint arXiv:1804.03320, 2018
Peter R Winters. Forecasting sales by exponentially weighted moving averages. Management science, 6(3):324-342, 1960.
13

---

Published as a conference paper at ICLR 2023
Hongfei Xu, Qiuhui Liu, Deyi Xiong, and Josef van Genabith. Transformer with depth-wiss wiie Istm. arXiv preprint urXiv:2007.0225, 2020.
Dmitry Yarotsky. Error bounds for approximations with deep relu networks. Neural Networks, 94: 103-114, 2017.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 6023-600 0023--032, 2019.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint 1rXxvv/710.094/// 2017.
Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmenta- tion. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 13001-13008,
2020.
14

---

Published as a conference paper at ICLR 2023
APPENDIX: MEGA: MOVING AVERAGE EQUIPPED GATED ATTENTION
A EFFICIENT COMPUTATION OF MULTI-DIMENSIONAL DAMPED 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 and omit the dimension index j in the following formulations. We denote the initial hidden state as hg- The multi multi-dimensional damped EMA defined in (5) can be vectorized into the following formulation:
h, = = u +1 о 8) h
(21)
y: hh,
(22)
where a, 8, and v E R.. u, = Bx Rº and h, E h the EMA hidden state at timestep &.
Let's denote Φ 1 - @ 8. Then, unrolling the above two equations explicitly yields:
h - o ho + a @ fx
h x²² o ho 2 Bx + a ② Ax2
y1 & @ ho n 8x1
Y2 === n ho + an Bxi + 1 Bx²
This can be written into vectorized formula:
Y1 17 Φ"
@ ho + 17
Coop Bx + ... a
(23)
y = K x + T 8² ho
(24)
where is the convolution transform with kernel K. E R.: K (n' ( B), 1)" (6 @
(25)
In the proposed multi minti-diieesiooal damped EMA, KC can be efficiently computed by the Vandermonde product. With K provided, the output in (24) can be computed efficiently with FFTs.
B PROOF OF THEOREM 1
Proof We split y into heads in the same way as Q, K, and
(1)
==
(())
Then we have
a" a'vv x1)
OSHGA a 07=
a (1i) o (())
To prove Theorem 1, we need to find 7 such that a (1) (1) (0) a ( (i)
aa v a a Vi (4, 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 ( ylt) 2 v() VX.
C LAPLACE ATTENTION FUNCTION
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 22:
fretuz (V2) fluplace (v2)
(26)
S'elug (V2) faplace (V2)
(27)
The Eq. (26) delivers μ 11/2 and Eq. 27 subsequently provides = v1/4T. Figure 5 visualizes the two functions.
15

---

Published as a conference paper at ICLR 2023
100
2
Valid I accuracy (%)
laplace
80
60
Taplace
2
2
50
100
150
200
x
Epochs
Figure 5: Laplace vs. ReLU
C.I STABILITY: LAPLACE VS. SQUARED RELU
Besides performance improvements, we also investigate the stability of the two attention functions. We conduct experiments the LRA Pathfinder task with MEGA models with the functions. Figure presents the accuracy on the validation set across training epochs. We observe that Laplace much more stable than ReLU2
16

---

Published as a conference paper at ICLR 2023
D EXPERIMENTAL DETAILS
D.1 DETAILS OF MEGA IMPLEMENTATION
Uni-directional Bi-directional EMA. In all the experiments, bi-directional EMA is applied when MEGA 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 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, simply very simple learnable relative biases, i.e. one learnable parameter for each relative positional distance (e.g. -3, +4).
D.2 LONG RANGE ARENA (LRA)
For all tasks, we closely follow Tay al. (2020) for details such as data preprocessing, data split, etc. The hyper-parametert of MEGA models on these tasks are listed in Table 8.
Table 8: Hyper-parametert of MEGA models on LRA and raw speech classification tasks. BSZ is
batch size, LR is learning rate and WD is weight decay. BN, LN and SN refer to Batch Normalization, Layer Normalization and Scale Normalization.
Task
Depth
derx
Attn-FN
Norm
Pre-norm
BSZ.
LRR
Dropout
WD
Epochs
Listpps
80
160
64
160
16
softmax
LN
False
0.001
0.1
0.01
60
Text
128
256
64
256
16
softmax
Talse
0.001
0.01
50
128
256
64
Image
160
320
96
320
16
laplace
BN
True
50
0.01
0.0
0.02
200
Pathfinder
128
256
64
256
laplace
128
0.01
200
Puth-X
64
128
32
128
16
laplace
BN
Thue
128
0.01
0.0
0.01
100
SC-Raw (hase)
60
120
30
120
16
laplace
True
0.01
0.0
0.01
200
SC-Raw (big)
72
144
36
144
laplace
True
0.008
0.0
0.01
200
D.3 Raw SPEECH CLASSIFICATION
Following Gu et al. (2022a), we perform speech classification on the SC10 subset of the Speech Commands dataset (Warden, 2018), which is 10-class classification task. The chunk size of are listed in Table
MEGA-chunk is 1000. Other
D.4 LANGUAGE MODELING
Training details We use the data of Wiki 103 and enwik8 and their splits provided by Dai et al. (2019). 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, chj. We use [cl,ch] [2, 6] for WikiText-103 and ccl.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 MEGA-chunk ($3.4) for training and set the attention chunk size to be 1024 and 2048 for WikiText-103 and enwik8 respectively. To use 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 In Table 5, we report test results that use longer chunk sizes (attention lengths) of 2048 and 4096 for WikiText-103 and enwik8 respectively. MEGA 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 illustrated in Figure 3. On the other hand, due to the use of rotary positional embeddings, attention can performed on longer chunk sizes at test time than those seen during training. We hope these two types of length extrapolation are clear readers. We provide the
17

---

Published as a conference paper at ICLR 2023
18.3
18.8
PPL
PPL
18.6
18.2
18.4
18.2
18.1
18.0
2.5K 3.3K 4.9K 9.8K 16K 25K
49K
512
1024
2048
3072
Max context length
Attn context length
Figure 6: Ablation studies of using different context lengths and attention lengths WikiText-100
ablation studies on these two types of length extrapolation below, i.e. extrapolation to longer context by increasing input sequence lengths and extrapolation to longer attention lengths by increasing the chunk size.
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 maximum context tokens [2.5K, 3.3K, 4.9K, 9.8K, 16K, 25K, 49K]. 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, MEGA can extrapolate to longer context lengths. The plot shows that PPL. decreases as the context length increased and the improvements saturate when the context length longer than 25K. This consistent with the observations in Press al. (2021).
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, MEGA consistently improves as we increase the attention length although it only uses 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 9: Syyper-parameten of models for language modeling.
WikiText-103
enwik8
Batch Size GPUs
6144 x 24
8192
Optimizer
Adam W
A
Learning Rate
0.005
0.005
Adam-8
(0.9, 0.98) linear
(0.9, 0.98)
Learning Rate Decay
linear
Weight Decay
0.1
0.1
Dropout
0.3
0.1
Attention Dropout
0.1
0.0
FEN Hidden Dropout
0.1
0.0
Gradient Clipping
1.0
1.0
Warmup steps
24K
24K
Total updates
400K
400K
Decoder Layers
Model size
1024
512
FEN Hidden size
1536
1024
Shared Repr. size (2)
256
128
Value Seq. size (v)
1536
1024
EMA dimension (h)
16
16
Chunk size
1024
2048
Total Parameters
252M
39M
18

---

Published as a conference paper at ICLR 2023
D.5 MACHINE TRANSLATION
The WMT 2016 English-German dataset contains 4.5M parallel sentence pairs for training. We following the standard setting (Ott al., 2018), using Newstes12013 the validation 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" (Post, 2018). The hyper-parameters of Transformer and MEGA models are listed in Table 10.
Table 10: Hyper-parameters of models for machine translation.
XFM-Base
MEGA-Base
Batch Size x GPUs
8192 8
8192 x
Optimizer
AdamW
AdamW W
Learning Rate
0.0005
0.001
Adam-8
(0.9, 0.98) inv. sqrt
(0.9, 0.98)
Learning Rate Decay
linear
Weight Decay
le
0.05
Dropout Attention Dropout
0.1
0.15
0.1
0.1
FFN Hidden Dropout
0.1
0.1
Gradient Clipping
1.0
1.0
Label Smoothing
0.1
0.1
Warmup steps
4K
4K
Total updates
500K
500K
Encoder Layers
6
Decoder Layers
6
Model dimension
512
512
FFN Hidden dimension
2048
1024
Shared Repr. dimension (2)
128
Value Seq. dimension (v)
1024
EMA dimension (h)
Total Parameters
65M
67M
'signature; nrefs:) ma noo tok: smooth: vers on: 11 .5.
19

---

Published as a conference paper at ICLR 2023
D.6 IMAGE CLASSIFICATION
Hlyper-parameter are listed in Table 11. We closely follow Touvron et al. (2021) by reusing most of the their typer-parameters.
Table 1I: Ingredients and Syyper-paaamene of DeiT and MEGA.
DeïT-B
MEGA
Batch size
1024
1024
Optimizer
AdamW
Adam W
learning rate
0.001
0.002
Learning rate decay
cosine
cosine
Weight decay
0.05
0.05
Epochs
300
300
Warmup epochs
20
Label smoothing
0.1
0.1
Dropout
Stoch. Depth
0.1
0.2
Repeated Aug
Gradient Clip.
1.0
Rand Augment
9/0.5
9/0.5
Mixup prob.
0.8
0.8
Cutmix prob.
1.0
1.0
Erasing prob.
0.25
0.25
Num. Layers
12
12
Model size
768
768
FFN Hidden size
3072
1536
Shared Repr. size (z)
256
Value Seq. size (v)
1536
Total Parameters
86M
90M
20


Reviews:
[{"review_id": "gZCe1ZheSfp", "reviewer": "Reviewer_ZPfy", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes MEGA, a method for combining exponential moving average (EMA) with attention. The paper evaluates MEGA on a number of benchmark datasets for sequence modeling and finds positive results.", "summary_of_the_review": "Overall, this paper introduces an interesting idea that deserves more exploration. However, I have some concerns about the novelty, the merits of the evaluation, and the framing.", "strength_and_weaknesses": "+ Improvement over S4 on a number of tasks, including language modeling.\n+ Extensive experiments and hyperparameter tuning to show that this method has promise.\n\n- EMA appears to be almost identical to S4D, a diagonal version of S4, but the authors do not compare against simple baselines such as replacing EMA with S4 or S4D. Thus it is unclear where the benefit comes from.\n- Some claims about the benefits of EMA over S4 are confusing.\n- The code implementation of EMA seems lifted from S4D [1], with no discussion of the differences or benefits, or acknowledgment of S4D. This should be corrected in the paper.\n\n[1] Gu et al. On the the Parameterization and Initialization of Diagonal State Space Models. https://arxiv.org/abs/2206.11893", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "N6w1GmQybZd", "reviewer": "Reviewer_nvWf", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper propose moving average equipped gated attention mechanism (MEGA) to address  transformer’s weakness in long-range modeling such as weak inductive bias and quadratic computational complexity based on the idea of EMA (more specifically, multi-dimensional damped EMA). MEGA (and its efficient variant, MEGA-chunk) can replace multi-head attention part of the original transformer. MEGA outperforms all other baselines including vanilla transformer and S4 on various benchmarks including Long Range Arena and different modalities. MEGA-chunk underperforms MEGA but it is much efficient.", "summary_of_the_review": "Overall, the paper is well written with principled design and good empirical results. Although the method can be regarded as a combination of existing methods, I believe MEGA can be highly useful considering the impact of commonly used architectures like transformer. Nevertheless, more rigorous evaluation on more diverse scenario might be necessary.", "strength_and_weaknesses": "MEGA shows superior performance on diverse datasets, meaning its practical usefulness. The motivation of MEGA is well supported and the solution is not too much complex. However, there are still many not straightforward design choices so more ablations on them could improve the completeness of the paper. It would be great if the authors include a comparison table about computational complexity in terms of FLOPs, latency, the number of parameters, memory footprint, etc. Also, having results on different scale might give a new insight.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LvEhtl_ZaS", "reviewer": "Reviewer_Dc4o", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "* The paper adds better inductive bias to Transformer attention by using a linear recurrence in the form of exponential moving average to contextualize the queries and keys. \n\n* The paper improves simple damped EMA into a multidimensional damped EMA turning it effectively in a simplified form of S4 where a diagonal bounded weight matrix is being used instead of HiPPOs.\n\n* The multidiemsnional EMA is effectively integrated with Transformer stype attention through the prior framework of GAU and GRU based gating.\n\n* The attention function is (computationally) simplified from multiheaded attention to gated single headed attention which is shown to be theoretically as expressive. A new activation function is introduced too for attention. \n\n*  The above component forms the MEGA block. The attention can be also chunked leading to MEGA-chunk. MEGA/MEGA-chunk can then be normed and put through FFN blocks. \n\n* Multiple experiments and ablations across various tasks in different modalities show the promise of the model. \n\n**Sorry for mid-rebuttal update**: I just had some thoughts that raise some concerns about some of the claims (about Theorem 1) made in the paper. I wanted to put it here as soon as possible so that the authors can engage with it. I have reduced my original score (8) to (5) due to the concerns. I am willing to improve the score again once that's resolved. ", "summary_of_the_review": "The paper combines multiple prior techiques and extends several of them with novel elements and shows strong performance among multiple domains and tasks. There are some additional experiments I would be curious about and there are some unclearity about positional encodings, and concerns about the theoretical proofs but otherwise it's a solid paper. ", "strength_and_weaknesses": "Strength:\n\n1. Effective motivated mixture of multiple prior techniques. Introduction of better activation functions, expansion of EMA in a multidimensional form similar to S4. \n\n2. Strong empirical performance accross multiple modalities and tasks (speech classification, imagenet, LRA, language modeling, translation). \n\n3. Several relevant ablations are present (for example, removing EMA, or changing dimensions of EMA and so on). \n\n\nWeaknesses:\n\n\n1. If I understand correctly you are using using a unidirectional EMA (going from left to right). If so the tokens wouldn't necessarily have context from the right side. That may counter the justification that EMA mitigates the issues of chunking (it may mitigate for the left boundary of the chunk, but the right boundary will still remain unable to access the further right side context). \n\n2. I also have some concerns with the theoretical proof in Theorem 1 for the equivalence of gated SHA and MHA. See the next section for details. \n\nNot necessarily strong weakness but additional ablations and experiments that could have been nice:\n\n3. I would be curious for some additional experiments/investigation. For example, what if *only* the EMA is used along with FFNs like in a S4 sort of setup (adjusted accordingly to keep parameter count similar)? Can it outperform or get close to S4 in LRA? \n\n4. Similar to 2. but what if S4 or other existing variants (instead of EMA) are used with attention. Can it offer even better performance?\n\n5. While thereotically, the gating function has the potentially to recover the \"missing\" portion from multiheaded attention, that may not always translate to similar empirical performance. I wonder if there could be some empirical performance loss from using single-head gating vs multi-headed attention. (**retracted** since there are prior works showing empirical strength of gated single-headed attention)\n\n6. I am unclear about which positional encoding is being used in all the tasks. I understand that ROPE was used in language modelling but is it used everywhere else too? I was also wondering for a possible ablation to see if a positional encoding is even as much necessary anymore given that Multidimensional EMA may model some positional aspects already.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "CKfBQgldVbm", "reviewer": "Reviewer_PXbY", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper introduces Mega, a gated, single-headed attention mechanism that incorporates a (damped, multi-dimensional) moving average. Instead of being computed directly on the hidden state, the query and keys are computed on a EMA-transformed version of it. The gating mechanism is akin to GRUs (with different inputs). The motivation of this approach is to introduce the EMA's inductive bias for recency/short-term relationships into the attention mechanism. Another feature of attention is its quadratic complexity in the context length, which can pose problems to model long-range dependencies. Here, the authors suggest chunking Q, K, V into k chunks and applying attention individually to each chunk to reduce it to linear complexity. This idea is not new. However because Q and K are computed on EMA-transformed inputs the attention operation can use information beyond the immediate chunk, which is in contrast with other chunking approaches. Finally, the authors also introduce the Laplace attention function as an alternative to the softmax in their attention layer.\n\nThe authors then test Mega and Mega(Chunk) on the Long Range Arena benchmark and MEGA on several long-context sequence modeling tasks (NMT, image classification, raw speech classification), finding strong performance in general. Ablations are also performed on the EMA layer, the chunk size for Mega-chunk and the attention function used.\n\nOverall, the paper is rich in its technical contributions, combining ideas to create an architecture for long-sequence modeling that achieves strong empirical results.\n\n", "summary_of_the_review": "This is a great paper, introducing several worthwhile technical contributions. The Mega architecture provided has interesting features and achieves strong performance across many tasks. The paper is high quality overall. There are minor criticisms, but I believe they are far outweighed by the strengths. ", "strength_and_weaknesses": "Strength:\n- The paper does a great job at explaining the newly introduced architecture, explaining how its components differ from the transformer and how they are related to prior work. This is quite an achievement given several components are introduced.\n- The key technical contributions (EMA layer for Mega, chunking for Mega-chunk) are well motivated by the issues with the classic attention block (lack of inductive bias, quadratic complexity). \n- The empirical results support the strength of Mega across a wide variety of tasks. I particularly liked that Speed and Memory columns were available in table 2. \n- The Laplace attention function is a nice contribution in itself and could be re-used in other work independently of the core architecture. \n\nWeaknesses:\n- More ablations on gating: While the paper does a good job on the EMA/attention function side of things, I believe some experiments could also be run on the gating mechanism. I understand that space limitations might have prevented it in this version.\n- Minor ablation: How much does the method deteriorate if you do not use damping?\n- S4 comparison: In section 4.1, I would point out explicitly that S4-v2 performs similarly to Mega-chunk which is the relevant baseline when accounting for speed and memory. Similarly, it would be more fair to include the Mega (base) number in Table 1 for SC-Raw, especially since the replicated S4 number is already lower than the original.  \n\nMinor spelling issues:\n- TEext p8 -> Text\n- Table 2 caption, remove parentheses around Long Range Arena?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster