Venue: ICLR 2023 Conference

Reviews:

final_score: rating: null
review_content: 
I updated my original review, because I realized recently that I may have some concerns with theorem 1. 

review_title: Sorry for the mid-rebuttal update

final_score: rating: 8: accept, good paper
review_content: 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. 

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

* The multidiemsnional EMA is effectively integrated with Transformer stype attention through the prior framework of GAU and GRU based gating.

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

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

* Multiple experiments and ablations across various tasks in different modalities show the promise of the model. 

**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. 
strength_and_weaknesses: Strength:

1. Effective motivated mixture of multiple prior techniques. Introduction of better activation functions, expansion of EMA in a multidimensional form similar to S4. 

2. Strong empirical performance accross multiple modalities and tasks (speech classification, imagenet, LRA, language modeling, translation). 

3. Several relevant ablations are present (for example, removing EMA, or changing dimensions of EMA and so on). 


Weaknesses:


1. 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). 

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

Not necessarily strong weakness but additional ablations and experiments that could have been nice:

3. 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? 

4. Similar to 2. but what if S4 or other existing variants (instead of EMA) are used with attention. Can it offer even better performance?

5. 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)

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


clarity,_quality,_novelty_and_reproducibility: **Novelty/Originality**: Although aspects of these work have been presented before (for example SRU++ have combined recurrency and attention in a similar, but not same, manner. Chunking based attention isn't new. EMA is a old technique which is combined with strategies used in S4-like models -- leading to a simpler form of SSM. Adding GRU/LSTM style gating in between layers in Transformers have been done before too), the novelty is still significant (theoretical connection between gated single headed attention and multi headed attention is interesting, Multidimensional dampled EMA seems novel despite being a mixture of prior works, the setup is also more effective than SRU++ going by footnote 2, mixing effective techniques together (even if none are individually new) to show good empirical performance is still a significant contribution and so on)


**Clarity**: Good. 

**Quality**: High. 

**Reporducibility**: Decent. I am a bit unclear about the positional encoding being used in different tasks, but otherwise most hyperparameters are clarified and the model code is also released. 


**Questions**

1) Just to confirm if I am understanding correctly: non-chuncked MEGA still uses the nxn quadratic attention but nevertheless gets memory performance similar to BigBird. Is it because of having more EMA parts and single headed attentions?

2) What positional encoding are used in different tasks? Is it ROPE for all? 

3) Can you address this concern:

I had some concerns with Theorem 1. 

It seems that gamma is a function of single token representation. While the result to be recovered is a function of the whole sequence. So it's not clear if in theory gated single headed attention is indeed capable of approximating multi-headed attention unless I am misunderstanding something.

Particularly, say we have a sequence of hidden states:

$H = (h_1, h_2, \dots, h_n)$

Now,

$\gamma^i_t = U(h_t)$ where $U$ is some universal approximator (eqn 12), $\gamma_t$ is the gate at the $t^{th}$ position, $i$ indicates the head.  

And for single headed attention at position t can be formulated as:

$\tilde{h}_t = SHA(H, t)$ 

where $SHA(H, t) = \sum_{i=1}^n \alpha_{t,i} f_v(h_i)$

( $f_v$ is a function that makes value transformation; typically a linear layer, and $\alpha_{t,i}$ is the attention coefficient)

Similarly, we may have a multiheaded attention (MHA) based representation at position $t$, for head $i$.

$\hat{h}^i_t = MHA(H, t)$ 

The idea is to approximate MHA with gated SHA like this:

$\tilde{h}_t \odot \gamma^i_t = \hat{h}^i_t$

That is, we want: 

$SHA(H,t) \odot \gamma^i_t = MHA(H,t)$

$\implies  \gamma^i_t = MHA(H,t) \oslash SHA(H,t)$

Let's say $R(H,t) = MHA(H,t) \oslash SHA(H,t)$ (abbreviation)

The claim in the paper seems to be that since $\gamma^i_t$ is the result of an universal approximating function, it can theoretically recover $R(H,t)$. 

However, there seems to be a mistake. Since strictly speaking $\gamma^i_t$ only relies only on the information at position t ($h_t$) whereas $R()$ is a function of the whole sequence of hidden states even for the output representation at position i. 

More formally, since, we defined $\gamma^i_t = U(h_t)$, the claim in the paper boils down to saying:

$U(h_t) = R(H,t)$

That is, the claim boils down to saying that we can recover  $R(H,t)$ just from $U(h_t)$. But even if $U(.)$ is a universal approximator its true input differs from the input of the function it is trying to approximate. So I am not sure how we can get a guarantee for equivalence in theoretical capacity for Gated SHA and MHA. 


We can try to address this issue by removing the temporally-indexed formalism as:

$\gamma^i = \tilde{U}(H)$ where 

$\tilde{U}(H) = (U(h_1), U(h_2), ...., U(h_n)) = (\gamma^i_1, \gamma^i_2, ...., \gamma^i_n)$

But now even if $U(h_t)$ is a universal approximator of any continuous function $f(h_t)$ by virtue of $U()$ being a non-linear wide neural network over $h_t$, $\tilde{U}(H)$ can't be said to be a universal approximator of any function $f(H)$. This is because $\tilde{U}(H)$ is not simply a neural network over some vector $H$ but it is a position-wise repeated application of the neural network $U()$ over each token representations in the sequence $H$.  
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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.
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.

clarity,_quality,_novelty_and_reproducibility: The paper is well structured from formal definition, illustrations, experimental result, and analysis. All figures are helpful to understand the architecture and computational graph of MEGA.

Since large pre-trained models are most widely used, I wonder the effectiveness of MEGA can be also generalized to those models. In my understanding, EMA only calculates new representations based on the previous time positions. For example, BERT brought huge success from attending to tokens in both directions. In other words, I suspect EMA component may harms a gain from bidirectional interactions.

I couldn’t check Appendix B thoroughly but theoretical justification part (Secion 3.5) is very nice.

MEGA-chunk achieves good speed gain. However, it introduces additional hyperparameter (chuncking size) depending on different tasks. A method (or module) that can automatically decide this value might be needed.

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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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.
strength_and_weaknesses: + Improvement over S4 on a number of tasks, including language modeling.
+ Extensive experiments and hyperparameter tuning to show that this method has promise.

- 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.
- Some claims about the benefits of EMA over S4 are confusing.
- 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.

[1] Gu et al. On the the Parameterization and Initialization of Diagonal State Space Models. https://arxiv.org/abs/2206.11893
clarity,_quality,_novelty_and_reproducibility: **Clarity:** The paper is mostly well-written, but some discussion points seem obscured. For example, in section 3.2, the authors write "the computation of the convolutional kernel in S4 requires complex fast Fourier transformers... The EMA sub-layer... convolution kernel would be a Vandermonde product, which can be computed in an efficient and numerically stable way." But EMA also uses a (real) FFT to compute it, which has the same complexity (and nearly same line in PyTorch) as a complex FFT.

**Quality:** The results are positive, but I have a few concerns.

1. Baselines.
As the authors write, the EMA layer's convolution can be computed using a Vandermonde product. This is almost exactly the same setup as S4D [1], a diagonal version of S4 (see novelty section below for more detailed comparison).
This suggests a number of natural baselines, which the authors do not evaluate:
* Adding S4D to an attention layer
* Adding S4 to an attention layer
These can be added in exactly the same way as EMA is added to attention. I would like to see these baselines to evaluate where the lift is coming from (attention, S4/EMA, or both).

2. Hyperparameter sweeps.
Based on table 8, there appears to have been extensive hyperparameter tuning for LRA. The paper is missing a description of the sets of hyperparameters swept for MEGA. I would like to see the S4+attention baselines *with the same hyperparameter sweeps* (not the exact same hyperparameters) in LRA. I also wonder whether similar hyperparameter sweeps would result in better performance for the existing baselines.

3. MEGA-chunk appears to only have been evaluated on LRA. What are the results on the other tasks?

**Novelty:** The contribution of combining EMA with attention is a good idea, and I like it.

However, the paper presents EMA as a distinct contribution from previous sequence models like S4. The paper does not discuss differences with simpler models such as S4D (which is also computed with a vandermode kernel). This is concerning since parts of the EMA code appear to be lifted directly from the S4D code base without attribution.

The forward pass is almost exactly the same as S4D's implementation, with renamed variables
https://github.com/HazyResearch/state-spaces/blob/f72202696125ce10ee7e3745ade5b04692dd2b74/src/models/s4/s4d.py#L95

S4D forward pass:
```
""" Input u shape (B, H, L) """
L = u.size(-1)

# Compute SSM Kernel
k = self.kernel(L=L) # (H L)

# Convolution
k_f = torch.fft.rfft(k, n=2*L) # (H L)
u_f = torch.fft.rfft(u, n=2*L) # (B H L)
y = torch.fft.irfft(u_f*k_f, n=2*L)[..., :L] # (B H L)
```

EMA forward pass (rename u to x, dimension H to D):
```
seq_len, bsz, embed_dim = x.size()
...
# L x B x D -> B x D x L
x = x.permute(1, 2, 0)
...
k = self.kernel(seq_len)
fft_len = seq_len
...
k_f = torch.fft.rfft(k, n=2 * fft_len)
x_f = torch.fft.rfft(x, n=2 * fft_len)
# B x D x L
out = torch.fft.irfft(x_f * k_f, n=2 * fft_len)[..., s:s + seq_len]
```

The bidirectional implementation (lines 198-204) is also the same implementation as the bidirectional implementation of S4, with renamed variables and slightly different torch operations: https://github.com/HazyResearch/state-spaces/blob/f72202696125ce10ee7e3745ade5b04692dd2b74/src/models/s4/s4.py#L1499

S4:
```
if self.bidirectional:
    k0, k1 = rearrange(k, '(s c) h l -> s c h l', s=2)
    k = F.pad(k0, (0, L)) \
            + F.pad(k1.flip(-1), (L, 0)) \
```

EMA:
```
if self.bidirectional:
    k1, k2 = torch.split(k, [self.embed_dim, self.embed_dim], dim=0)
    # D x 2*L-1
    k = F.pad(k1, (kernel_size - 1, 0)) + F.pad(k2.flip(-1), (0, kernel_size - 1))
```

The two implementations also have the same error message:

S4:
```
if state is not None:
    assert not self.bidirectional, "Bidirectional not supported with state forwarding"
```

EMA (rename "state forwarding" to "incremental state"):
```
assert not self.bidirectional or incremental_state is None, 'Bidirectional EMA does not support incremental state'
```

**Reproducibility:** The code and final hyperparameters have been provided, but the hyperparameter sweeps are missing.

**Misc:** The code is not properly anonimized, it contains references to the "fairseq" library.

[1] Gu et al. On the the Parameterization and Initialization of Diagonal State Space Models. https://arxiv.org/abs/2206.11893
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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.

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

Overall, the paper is rich in its technical contributions, combining ideas to create an architecture for long-sequence modeling that achieves strong empirical results.


strength_and_weaknesses: Strength:
- 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.
- 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). 
- 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. 
- The Laplace attention function is a nice contribution in itself and could be re-used in other work independently of the core architecture. 

Weaknesses:
- 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.
- Minor ablation: How much does the method deteriorate if you do not use damping?
- 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.  

Minor spelling issues:
- TEext p8 -> Text
- Table 2 caption, remove parentheses around Long Range Arena?
clarity,_quality,_novelty_and_reproducibility: * Clarity: The paper does a great job. This is all the more impressive that there are several contributions.
* Quality: This is overall a high quality paper.
* Novelty: This scores quite high, with several contributions. I could see criticism that points out that many components already exist elsewhere but I believe finding the proper combination of attention + SSM/EMA like-bias is worthwhile and I see the Laplace function as the cherry on top.  
* Reproducibility: No problem here, the appendix is very helpful in this regard.
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.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: The paper proposes a new mechanism called the moving average equipped gated attention (MEGA) for addressing long-range modeling in transformers. MEGA is based on exponentially weighted moving averages (EMA) and shows improvements over various baselines on long range tasks. 

All reviewers agree that the technical contribution is well motivated and the empirical results are strong, showing significant improvement across a wide range of tasks/datasets. 

There has been some concerns in the reviews with respect to the novelty, clarity, theoretical results, and positioning of the the paper with respect to other closely relevant papers. However, most of the concerns have been resolved after extensive discussions with the authors.
(I would like to thank the reviewers and authors for the detailed discussions).
justification_for_why_not_higher_score: There are still some concerns with respect to clarity and positioning of the paper, which doesn't justify a higher score.
justification_for_why_not_lower_score: The reviewers who raise such concerns seem either satisfied with the author response and discussions with other reviewers, or do not strongly advocate for rejection.
Decision: Accept: poster