Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper looks at one specific aspect of transformer architectures - the normalization step that constrains the attention vectors in a probability simplex. The authors argue that this restriction has some limitations including limited information flow, makes the models biased towards local information at initialization and increases the sensitive to hyperparameters. To mitigate this, the authors propose an architecture that replaces the softmax in the attention mask with a simple normalization. Then, with several experiments on a synthetic dataset, the authors demonstrate that the proposed architectural change leads to models that are robust to hyperparameters, and less sensitive to biases in the data. 
main_review: The finding of the paper is interesting - the arguments made about the restrictive nature of the softmax operator makes sense. The proposal to replace the softmax with normalization is very simple and could have a lot of implications. The experiments conducted on the synthetic data are thorough and the discussions were made well. However, the authors could have shown gone beyond synthetic experiments and shown more results on real world datasets.
summary_of_the_review: While I enjoyed reading the paper, my biggest concern is lack of real world experiments. I appreciate the authors on conducting experiments on multiple runs, repeating the experiments with different model sizes, visualizing the results effectively and conducting carefully crafted simple experiments that illustrates the specific points being made. The carefully experiments were all designed well. However, for practitioners, a pressing question would be if these architectures can be used as a replacement of the traditional transformer models. Can this be used for training language models or benchmark NLP tasks?

To this end, it would have been really nice if authors had shown one experiment on real world datasets. I understand that some institutions might not have enough resources to train large-scale models, but a proof of concept at a decent scale on a realistic dataset would have been nice. Such an experiment would tell whether these findings would translate to large scale datasets and tasks. I would also like to point out that the authors have one experiment on protein-protein interaction graph prediction, but I am not aware of this benchmark and I am unable to gain much insights from this experiment aside from the fact that the model works.

While this issue is certainly concerning to me, I believe the findings in this paper are a reasonably good contribution that could help others in the community. It's just that the paper would have been really strong had the authors gone beyond synthetic dataset and done some benchmark NLP experiments.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The authors investigate the self-attention module architecture and provide a theoretical analysis for the limitation of using softmax function for normalization in the module. In this regard, an unconstrained normalization function is proposed as alternative in the transformer. The authors compare it with baseline transformer and several variants to validate the effect of the proposed normalization function in transformer.
main_review: Strengths:

The authors provide a theoretical analysis and proof on using softmax function in attention module for effect of gradient. I did not check the proof precisely. But analyzing neural network modules from a theoretical perspective is interesting and valuable to the community. 

The whole paper is well organized and written, and easy to follow. 

The authors give a comprehensive review and discussion to modern works related to transformer architectures and normalization and pooling methods.

The code is provided in the supplementary material, which would be helpful for reproducing the experiments in the paper.

Weaknesses:

The proposed NAP operation is quite related to layer normalization (LN). Instead learning element-wise scale and bias parameters in LN, scalar parameters are learned in NAP, which is supposed to function analogous to LN.  Can it be used as replacement of all the LN layers in the architecture?  If not considering variable length, is using LN operations only (i.e., replacing NAP as LN in the Figure 3) an effective variant compared to MTE and NAP architectures?

Most ablation studies and empirical analysis are conducted on synthetic tasks which are less convincing to demonstrate the effectiveness of the method. Although the authors present the experiments of exploring the proposed method into graph neural networks in PPI task and memory graph agent in reinforcement learning task, the performance improvements (in Fig. 9) are rather marginal compared to baseline methods.




summary_of_the_review: The authors theoretically analyze the limitation of using softmax function in attention module and propose an unconstrained surrogate in the paper. However, the empirical evidences shown in the paper are insufficient to validate the method.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper starts from the observation that current self-attention modules is sensitive to hyperparameter changing. The authors conjecture that this is due to the softmax operator in self-attention, and give some intuitive examples to support their hypothesis (eg., bias towards local information, gradient vanishing). To solve the problem, the authors propose to replace the softmax with a normalization operator. They show that normalized attention is more robust to hyperparameter in a synthetic dataset, as well as real-world settings.
main_review: Strengths:

+   The paper is clearly motivated by the problem that current self-attention is not robust to hyperparameter changing. The authors also give some intuitive examples on this, such as bias towards local information and gradient vanishing.
+   The paper designs synthetic tasks and does a through analysis of NAP and other baseline models in this setting. Then the authors also experiment in some real-world settings such as GNN and RL.
+   The visualization method of the results is clever. 
+   The experiment on finding the mode in the input data is inspiring, which implies that information is not effectively routing through all tokens in attention modules.
+   The paper is well-written and easy to follow.


Weaknesses:

-    A few related works are not cited or compared to. For example, [1] replace the softmax in self-attention with a token-wise Guassian kernel. [2] compares self-attention with fast weight programmers and proposes to remove the softmax. Several other works also adopt a similar design [3, 4]. Although these works are focusing on different problems instead of robustness to hyperparameter as in your paper, could you give some discussion on the relation between these prior arts and your model?
-    The authors did experiments in real-world settings (GNN and RL) in Sec 5.4, but the performance of NAP and BERT / MTE is somewhat similar. They all have the same validation results (99.1% in GNN, and 99.3%-99.4% in RL). The robustness to hyperparameters is also similar between MTE and NAP in RL task. Is it because the tasks are still too small in scale to show the difference between NAP and baseline models? 
-    The normalization on $l^{i, j}_m$ will force the attention weights $a^{i, j}_m$ to have a fixed variance. But for example, if the tokens are all very similar, the attention weights should be similar too (thus having a small variance). Is the normalized attention able to handle this situation?


Other comments:

-    Some other works have also focused on stabilize the training process of transformers [5]. Do they have any relation to your work? A short discussion on this line of work could be beneficial.




[1] Lu J, Yao J, Zhang J, et al. SOFT: Softmax-free Transformer with Linear Complexity[J]. arXiv preprint arXiv:2110.11945, 2021.

[2] Schlag I, Irie K, Schmidhuber J. Linear transformers are secretly fast weight programmers[C]//International Conference on Machine 
Learning. PMLR, 2021: 9355-9366.

[3] Shen Z, Zhang M, Zhao H, et al. Efficient attention: Attention with linear complexities[C]//Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2021: 3531-3539.

[4] Cao S. Choose a Transformer: Fourier or Galerkin[J]. arXiv preprint arXiv:2105.14995, 2021.

[5] Xiao T, Singh M, Mintun E, et al. Early convolutions help transformers see better[J]. arXiv preprint arXiv:2106.14881, 2021.
summary_of_the_review: The paper is well-motivated by the problem of sensitivity to hyperparameters and propose a sensible algorithm to solve the problem. The experiments on the synthetic task are sound. However, it seems lacking some convincing results on real-world dataset. And the relation to some prior works needs to be clarified.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The authors explore the implications of using softmax to implement attention mechanisms in NLP models, particularly transformers. They highlight the theoretical limitations of using softmax-attention and propose a normalization-based attention mechanism (NAP) that overcomes these constraints.. The authors conduct a set of synthetic experiments to compare different attention mechanisms’ performance and demonstrate that the proposed mechanism performs well compared to the traditional softmax attention while being  less sensitive to hyperparameters change.
main_review: Pros:
1. The paper is clearly written and easy to follow.
2. The implications of using softmax attention are thoroughly analyzed in theoretical and experimental settings.
3. The set of synthetic experiments conducted highlights the differences between attention mechanisms reasonably well and supports the authors’ claims.

Cons:
1. The authors do not apply their proposed mechanism to the established BERT benchmarks. Instead, they construct a set of synthetic tasks that aim to imitate real-world NLP tasks. Lack of real-world NLP benchmarks makes it harder to believe that the implications of the softmax use actually matter. As the authors mentioned in Appendix C, some softmax-attention constraints may actually be beneficial, and this can’t be confirmed or disproven without real world benchmarks. The synthetic tasks chosen are intuitively similar to the specific NLP problems, but in the end, it’s not clear how well the synthetic results will generalize.
2. The experimental results demonstrate the advantage of NAP over sum and max pooling, but the case of NAP vs NON is less obvious. For many tasks, NON performs adequately well in terms of final quality and hyperparams sensitivity (in figure 26 it seemingly outperforms NAP for sensitivity). The paper would benefit from providing clear evidence (in addition to Figure 7) that despite all the additional layer norms in MTE, the dynamic normalization of attention logits is still necessary.

Suggestions/questions:
1. It’d be great to add the learned aggregation weights mechanism into the comparison. The authors explicitly try to avoid depending on the sequence length but having such a comparison may help prove the hypothesis that convex limitations are to blame and moving away from the convex hull is beneficial. For example, Figure 7 may benefit from such comparison.
2. Due to the use of MTE architecture, it’s not clear how well NAP will perform as a drop-in replacement of the softmax attention. The paper would benefit from NAP experiments that are not MTE-based to confirm that NAP isn’t brittle and can be used in practice.
3. (Figure 7) The explanation here can be more convincing. Why is MTE worse than traditional BERT on this task? Is it because local information bias is actually beneficial? Why is NON underperforming here and how can it be confirmed (NON doesn’t underperform like this in other tasks)? Figure 23 from the Appendix is more convincing, so it may make sense to move it here
4. (Figure 9) It’d be great to move the NON, sum, and max results from the appendix here and add the explanation. 
5. (Style) Using RGB pixels for visualization purposes helps with the presentation but is not colorblind friendly. . It may be better to split the plots per-channel.  

summary_of_the_review: While the practical viability of the proposed method is still uncertain, the insights on using softmax attention from this paper are valuable. The proposed synthetic evaluation framework is also a good foundation for further discussions on softmax-attention alternatives.
review_title: null

Meta-review: This paper studies potential drawbacks in using softmax over attention in Transformers and evaluates other normalization approaches. Reviewers, while had been positive about the empirical analysis and the insights from the synthetic data experiments, agree that the paper lacks real world experiments/insights. I agree with that and believe the paper falls short in several areas.

1) Drawbacks of softmax: Paper states several generic drawbacks of softmax such as saturation issue leading to vanishing gradients. However the paper does not demonstrate if Transformer models used in practice suffer from this issue under standard training settings. Even the arguments that attention layer focuses on local information is quite vague and not well supported. Overall the analysis is quite weak without much concrete statements and demonstration in real settings. 

2) Experiments: The paper presents many synthetic experiments evaluating alternates to softmax varying from layer normalization to pooling. There are no experiments showing if the studied variations actually solve the issues discussed in earlier section. Finally due to the lack of any real world experiments (even small scale ones), it is not clear if the results apply in real world settings.

Overall I think the paper needs significant work in formalizing the drawbacks of using softmax in Transformers and demonstrating that the proposed solutions indeed solve this problem.
Decision: Reject