Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper proposes a mixture-model perspective of multihead attention. Each key within each attention head is treated as a Gaussian component; with some assumptions, the softmax-based attention weights can be recovered as the posterior probability of the Gaussian mixture, conditioning on the query. The framework leads to a new parameterization of multihead attention, which is explored in transformers. It can be applied to several linear transformer variants, too. Experiments with text classification, language modeling show that the proposed method achieves same or better accuracy controlling the number of attention heads, and reduces computational overhead. Analysis suggests that it reduces the redundancy of multihead attention
main_review: 
- First, I like the Gaussian mixture perspective of multihead attention in general, but it could lead to cases that can be a bit awkward. For example, in causal attention, the number of components of the mixture changes over time steps. This could be avoided by, e.g., treating only the multihead part as a mixture, not the keys over timesteps. 
- Second, it can be confusing to weave into the narrative something that is actually not used at all, e.g., the norm constraints around Eq 5, and the M step around Eq 13. I suggest reworking these part and probably presenting them as connections to other models. 
- Third, some modeling details need to be clarified. Optional (A) for key’s parameterization, the notation W_{k_{jr}} seems to suggest it needs NM number of projection matrices. This seems  much more expensive than standard multihead attention. Similarly in option (B), why would it need N number of matrices at all? Further about the parameterizations of queries and values. The notations seem to suggest that they are never “multiheaded” as keys are. Could the authors confirm this? If this is true, it is worth clarifying.
- Adding onto the above, what does a MGK with 4 heads mean? Does it have 4 x seq_len number of Gaussian components for the keys?
- Additionally, it would be nice to include a time overhead comparison.
- Lastly about baselines and further experiments: the WikiText103 baseline seems pretty far behind, e.g., Baeski and Auli (2018). Could the authors explore MGK on top of a stronger baseline? Besides, [1] seems pretty related and worth comparing to. Given the weak LM baseline, I suggest the authors including other sequence modeling tasks such as machine translation.

Typos:
- Above 1.1, keyed -> key.
- Above Eq 4, unit matrix -> identity
- Above Eq 7, abusing -> overloading
- Below Eq 11, at learning -> at learning time.
- Start of Section 3, “full the number of heads” doesn’t parse for me.

References:

[1] https://arxiv.org/abs/1911.02150

[2] https://arxiv.org/abs/2103.02143
summary_of_the_review: Strengths:
- An interesting mixture model perspective of multihead attention
- It is nice to reduce compute cost with the same or better accuracy

Weaknesses:
- It can be confusing to introduce concepts that are never used in the model
- Several key details of the implementation need to be clarified.
- Weak LM baseline. Would be great to include more experiments.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper constructs a Gaussian mixture model with uniform prior and one Gaussian centered at each key and shows that the posterior distribution ie. the probability that a given query is generated by a given key matches the attention scores in the transformer architecture. In addition, it is proposed to describe each key as an additional Gaussian mixture and it is shown that the posterior is easy to compute and resembles the multi-head attention (given that the query remains the same for all heads). The paper proposes that this formulation can learn more diverse attention-heads and reduce the computation by replacing some heads with the Gaussian mixture heads which are easier to compute. Moreover, the authors show that it is straightforward to transfer the aforementioned formulation to linear transformers. Finally, the authors perform experiments on the LRA benchmark and language modelling experiments on Wikitext-103.
main_review: Strengths
-------------

- The paper provides a nice probabilistic view on the attention matrix that could be a new tool to study the transformer architecture.
- Language modeling experiments besides LRA are greatly appreciated
- The paper is well written and easy to follow

Weaknesses
------------------

The most significant weakness is that the provided experiments do not provide adequate evidence regarding any benefit of the proposed method.
1. In both LRA and Wikitext-103 all scores way too close to be able to judge whether any of the methods performs better.
2. Even if they do perform better (the only one could be sMGK on LRA-Retrieval) the model is sufficiently changed that more tests would be required. For instance, is the performance improvement coming from using the distance instead of the dot product? Is it coming from the key shifting? 
3. In all experiments, what is M? It is quite hard to compare the performance with different number of heads without having M for MGK and MLK.
4. Moreover, it would be interesting to see the performance of softmax and linear transformers with 4 heads on Wikitext-103. It is quite interesting in tables 1 and 2 that reducing the number of heads has minimal impact in performance.

Finally, it would be interesting and beneficial to the paper's main point to have a proposition regarding why parameterizing the attention matrix in this way would improve the diversity of the attention heads. The rank argument does not show diversity between attention heads given that the attention matrices can all be the same but have high rank. Moreover, when thinking about how diversity can be achieved with this formulation, it can also be measured and enforced with some type of loss or regularizer.

Suggestions
-----------------

- It would be much more informative to exchange FLOPS with wall-clock time and although number of parameters is important for large transformer models, so is the memory footprint in the GPU during training.
summary_of_the_review: As mentioned in the top of the weaknesses section, the main reason for my recommendation is that there is not sufficient evidence showing an improvement in performance with the proposed method.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The paper addresses the efficiency of Transformer models.
The paper provides a probabilistic view of self-attention in transformers and introduces a novel transformer model, Transformer-MGK, and its extensions.
In Transformer-MGK, each key of self-attention is modeled as a mixture of Gaussians whose mixture weights are estimated by the E-step of EM algorithm.
In the experiments, they show Transformer-MGK is comparable or better than Transformers on LRA benchmark and a language modeling benchmark.
In addition, the authors empirically show Transformer-MGK alleviates the head redundancy of conventional Transformer.
main_review: Strong points
 
1. The proposed method is empirically efficient in terms of the number of parameters and computational cost in terms of FLOPS.

2. Although there is previous work for the RBF kernel interpretation of self-attention [1], the proposed probabilistic interpretation of self-attention provides an interesting direction for future research.

3. This paper provides experimental results on LRA benchmark and a language modeling benchmark to show the effectiveness of the proposed method.

4. The paper also introduces the linear-attention variant of the proposed method.
 
[1] Yao-Hung Hubert Tsai et al. Transformer Dissection: An Unified Understanding for Transformer’s Attention via the Lens of Kernel.  EMNLP-IJCNLP 2019

##########################################################################

Weak points

1.  Although the authors argue that Transformer-MGK is introduced to improve the explanation power of each key, increase the representation of each attention head, and reduce the chance of learning redundant heads, this paper provides only empirical results and no theoretical explanation for these improvements.
(1) It is valuable to describe the computational/space complexity of the proposed method and compare them with those of conventional methods.
(2) Although I can somehow expect that the mixture of Gaussian increases the representation power of keys, I cannot imagine why the proposed method reduces redundant heads and promote the diversity in attention pattern.

2. Since pre-training and finetuning paradigm is out of scope of LRA benchmark, I suggest the authors to conduct experiments for standard NLP benchmarks, such as GLUE, SQuAD, SuperGLUE, or Winogrande.

 

summary_of_the_review: Overall, I vote for weak acceptance. I like the probabilistic interpretation of self-attention and the empirical results are strong. 
My major concern is about the clarity of the paper (please see weak point #1 below).
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes a transformer with a mixture of Gaussian keys (Transformer-MGK) in order to avoid training redundant attention heads in the transformer modules. The authors formulate the attention mechanism as a posterior distribution $p(t_j = 1 | q_i)$ (i.e., given the query $q_i$, how likely $q_i$ matches the key $k_j$), and extend this as the mixture of Gaussian keys where each key $k_j$ is modeled as a mixture of $M$ Gaussians with several assumptions. The paper introduces inference methods using the EM algorithm: Soft E-step and Hard E-step. The authors also propose a liner variant, Transformer-MLK, which can be more computationally efficient to long sequences. Empirical studies are performed to show the effectiveness of Transformer-MGK over the softmax Transformers. Results are shown on the long-range arena benchmark and language model on WikiText-103 corpus. Also, comparisons of the model complexity and computational cost and ablation study depending on different inference and learning techniques are also provided.
main_review:  Strong points

- The idea of introducing the mixture of Gaussian keys to increase the expression capacity of key embeddings and to allow the query to attend to more diverse positions is interesting.
- Overall, the paper is well written. Especially, in section 2, the formula explanations have a neat flow so that it is easy to follow.
- The proposed model shows comparable results to general softmax-based transformers with fewer parameters and FLOPs.
- In the experiment analysis, the visualization of the attention matrix and rank distribution supports that the paper's hypothesis that Transformer-GMK would learn diverse attention patterns.

Weak points

- The main concern about this paper is the experimental improvements are marginal. Moreover, although the method shows comparable to original transformers with fewer parameters and FLOPs, I suggest the authors to discuss whether there is any difficulty in the learning process using EM-based algorithms for practical usability; for example, how many learning epochs are required, whether it is sensitive to initial values, and an increase of computation required to calculate the E-step during the training process.
- Beyond the attention matrix analysis, in order to demonstrate the keys are well modeled as Gaussian mixtures, the learned $\pi_{jr}$,  $k_{jr}$, or $W^k_r$ which is the weight matrix projecting to $k_{jr}$ seems to be required to be analyzed.
- In the experiment section, the number of Gaussian mixtures is not noted. Also, comparing with baseline methods (Softmax and Linear heads) and proposing methods with half heads is, for me, unfair without considering the number of the mixtures.
-  I'm not convinced with the experimental results are reported at their convergence since the learning curves in Figures 1 and 4 end before converging.

Questions

- How does the proposed model result with the increasing number of mixtures?
- How many epochs are needed to converge via the EM algorithm?
- Also FLOPs of training time including the E-steps should be compared.
- Can this proposed approach be applied to deeper and large Transformers such as BERT and GPTs and show its effectiveness?

Comments

- I would like to recommend comparing the model complexity and computational cost with notations and formulas rather than empirical figures.
- Typo in Figure 3. The label of Softmax Transformer is wrong as Liner Transformer in the left figure.
- Typo in Table 4. The lower rows group: Transformer-MGK → Transformer-MLK
summary_of_the_review: Overall, I vote for "marginally below the acceptance threshold". I think the extension of attention keys to the mixture of Gaussian and exploring the learning methods with EM algorithms are somewhat technically sound. However, experimentally demonstrating its practical usefulness is required further. Hopefully, the authors can address my concern in the rebuttal period.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes Transformer with a Mixture of Gaussian Keys (Transformer-MGK), which replaces the multi-head self-attention module with a Gaussian mixture model. Specifically, the authors build the connection between self-attention and probabilistic framework. Based on their findings, they extend the self-attention to a Gaussian mixture attention. Compared with vanilla Transformer, Transformer-MGK has fewer parameters and accelerates training and inference with comparable performance. Various experiment results illustrate the advantage of Transformer-MGK.
main_review: Strengths:
The authors use the probabilistic framework to establish the connection between self-attention and gaussian distribution. Even though there are some assumptions: (i) the query and the key are normalized; (ii) all the standard derivations are identity; (iii) the prior distribution is uniform, I think these assumptions are easy to follow. Based on the connection, a natural extension is the mixture Gaussian model and they apply the EM algorithm to learn the posterior distribution. The experiment results are comprehensive.

Weaknesses:
1. In Table 1 and Table 2, the authors discover the influence of the number of heads. However, [1] claims that one head is indeed sufficient, which is not consistent with your results. I wonder whether the reason is the experiment settings.
2. In Table 1, Transformer-sMGK is better than Transformer-MGK in some cases but not in all cases. Can you provide some explanation about the reason? Also the situation in Table 2.
3. In Table 1 and Table 2, the authors start from 8 heads. Can you provide the results from the vanilla Transformer (12 heads) for better illustration?

[1] Michel, Paul, Omer Levy, and Graham Neubig. "Are Sixteen Heads Really Better than One?." Advances in Neural Information Processing Systems 32 (2019): 14014-14024.

summary_of_the_review: This paper builds the connection between self-attention and Gaussian distribution with some assumptions, which is novel for Transfomer model. The extension from original self-attention to mixture Gaussian model is natural and convincing to me. The experiment results are sufficient to illustrate the improvement of their models.
review_title: null

Meta-review: The paper proposes using a mixture of Gaussian models for transformer keys (MGK) so that the posterior distribution of key given query matches the attention scores in the transformer architecture under some assumption. Similarly but in reverse, the query given key under a MoG also matches transformer attention score. The paper proposes that this formulation can learn more diverse attention-heads and reduce the computation by replacing some heads with the Gaussian mixture heads which are easier to compute. Moreover, the authors show that it is straightforward to transfer their formulation to linear transformers. The authors perform experiments on the Long Range Arena benchmark and on Wikitext-103 where they show some improvements with less attention heads, while using up to 20% less FLOPs for softmax transformers and around mostly 20% but up to 80% less parameters for the worse linear transformer.

Reviewers generally find that the interpretation of transformers as mixture of Gaussian interesting, although it is based on normalization assumptions. The main objections from both positive and negative reviewers is the weakness of the empirical results. In summary, all the improvements in perplexity are less than 0.5, and accuracy improvements are less than 1. In the discussions, the authors claim that their comparisons are using fewer heads vs baselines with more heads, which is how the tables are presented, but upon closer look I find this unconvincing since the same head count comparisons can be reconstructed and still show very weak results. For example in table 1, acc(MGK4)=61.85 vs. acc(baseline4)=61.23, or in table 2 acc(LinearMG4)=55.7 vs acc(LK4-baseline)=55.61. To begin with, it would be better to compare the same number of heads for more head-to-head comparisons, so I find the whole argument to be misleading. The claims about FLOPS and memory has a similar flavor where the differences are small but the claims were big. A visual evaluation of figure 3 show that most reductions are around 10% in FLOPs or parameters for the better softmax model (note the axis does not start at 0). The retrieval task for the worse linear transformer is the only case where the FLOPs reduction seems significant. Given that EM is now required during training, I'd interpret the results as a negative for MGK given the marginal improvements achieved.

Several reviewers are unhappy with the strength of the empirical results while most reviewers gave favorable scores after the discussion where the authors insisted that some valid points are not valid. The authors requested the AC to look further into the one negative review due to the lack of a reviewer's response. After taking some time to look at the substance of the paper and the review, I recommend rejection due to the weakness of the results, the misleading presentation and discussions.
Decision: Reject