Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: This work aims at solving the alignment problem between the query and key distribution within the self-attention module. It is based on the hypothesis that the key and query following identical distribution can help to improve the feature learning for the attention module. 
Correspondingly, an alignment regularization is proposed to explicitly encourage self-attention to match the key and query within each head. Three kinds of alignment methods, including Adversarial training-based, Optimal transport-based, and, Bidirectional conditional transport-based, all of which can outperforms the baselines. Comprehensive experiments are performed on a variety of problem settings, which proves that the introduced alignment module can bring consistent improvements. Besides, the proposed method can work as a plug-and-play module to be added to many existing attention modules. 
main_review: I recommend accepting the paper. In general, this paper is well-written and easy to follow. It aims at studying the interactions between the keys and queries for attention modules. A variety of experiments on different problem settings with different proposed alignment methods are performed, and, the proposed method consistently outperforms the baseline, which proves that explicitly regularizing the key and query distributions can help to improve the feature learning for the attention-based methods.  I summarize the strengths and weaknesses as follows:

Strengths:
- This paper proposes a plug-and-play alignment module that can be applied to various attention-based models. 
- This paper proves that explicitly regularizing the distributions of the key and query feature can help the feature learning of the self-attention module. 
- Three alignment methods are introduced to regularizing the distributions of the keys and queries. 
- The experimental results are convincing. Experiments are performed in different language understanding settings. It shows that the proposed alignment can bring consistent improvements no matter which alignment method is adopted. 

Weaknesses:
- It lacks the experimental results showing that after being trained with the alignment loss, the distribution of the keys and queries are well-aligned. 
- In the baseline attention module, the keys and queries will be dot-producted, which means they should be of the same feature domain. Thus, it has implicitly introduced the regularization of their alignments. In comparison, the novel part of the proposed method is explicitly modeling the regularization. However, it lacks a detailed explanation on why explicitly align the keys and queries can bring improvements over the baselines. 
- Typos: Line-104 the tensor dimension should be [B, H, w, d]
limitations_and_societal_impact: N/A
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: The paper reviews the attention mechanism commonly used in deep neural networks. It proposes the alignment attention mechanism that match the query and key distributions in self-attention within each head. The paper regards the alignment attention as an unsupervised regularization in current attention. Three alignment methods have been introduced to explore the effectiveness of distribution matching. Extensive experiments are done to demonstrate the effectiveness and general applicability of alignment attention in domains of language understanding tasks, and even graph node classification and visual question answering. The alignment attention shows great potential of being a plug-and-play component to current attention models.
main_review: Main Review:

The paper rethinks the dominant self-attention mechanism in both deep neural networks and points out the alignments of query and key distributions in the current attention mechanism are largely unexplored. The motivation of align the projection of both query and key tokens is strong and meaningful as decreasing the risk of misaligned interactions between queries and keys. To conclude, the paper attempts to take distribution matching into consideration of attention mechanism.
Based on that motivation, the authors proposed three alignment methods attempting to align the attention and finally adopted adversarial training and conditional transport methods to align the distributions. The analyses of methods are theoretical and insightful.
The authors conducted extensive experiments in natural language understanding, GNN and VQA domains. The evaluations are not only in-domain dataset, but also out-domain ones to illustrate the generalization across domains. Adversarial attacks are taken into consideration to show the robustness of alignment attention. The experiments are very meticulous and rational to show the effectiveness of alignment attention as a plugin-and-play component.

Limits:

As is known, the multi-head attention mechanism is widely applied in vision tasks like classification, detection and segmentation. The authors have conducted detailed ablations to show effectiveness in NLP, GNN and VQA covering language, graph and multi-modal domains. I suggest supplements of some vision tasks to finalize the ablations would make the paper more completed 

Reason for weak rejection:

My concern about the paper is that do we really need alignment between key and query? As key and query are linear projections of input, the distribution of key and query should be well aligned. Thus I do not think the alignment story is reasonable and suggest weak rejection.

limitations_and_societal_impact: No
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper aims to improve the performances of current widely-used attention models by explicitly aligning the query and key distributions. To this end, the authors propose an alignment attention model, which can be easily incorporated into existing self-attention based models, including the pre-training models, e.g., BERT, to boost their performances while remaining efficient in memory and computation cost. To prove the effectiveness of the proposed approach, the authors conduct extensive experimental studies on natural language understanding tasks, graph attention network, and visual question answering task. The domain generalization ability and adversarial robustness of alignment attention are also studied.

> After Rebuttal:

Thanks the authors for the response. I'd like to keep my initial score.
main_review: Strengths:
1. The paper is well written. 
2. The paper gives extensive experiments to prove the effectiveness of the proposed approach in terms of various aspects, i.e., accuracy, uncertainty estimation, generalization across domains, and robustness to adversarial attacks.
3. The proposed approach achieves state-of-the-art results on several tasks while remaining efficient in memory and computation cost.

Weaknesses:
1. The contributions are limited.
- The idea of alignment learning has been explored in literature. Thus, the main contribution of this paper is the methodological contribution, i.e., the alignment attention model, which can align the query and key distributions. However, the motivation is unclear.

2. The motivation and the studied problem are unclear.
- Personally, it is unclear what problems of the conventional self-attention model can be solved by the proposed approach? In other words,  what are the problems that will cause if the distributions between query and key are not aligned?
- Specifically, in L44-45, the authors claim that "misalignment could lead to an undesired dependence between the attention weights of a token and its position in the input feature space", however, there is no evidence/experimental analysis that can support this claim. How to prove it? Although the extensive experiments have indeed demonstrated the effectiveness of the proposed approach, I would like to know why performing the alignment of query and key distributions can improve the performance. 

3. Some detailed analyses are missing.
- Can the proposed alignment attention be incorporated into the source-to-target attention component?
- In the discriminator-based modules, i.e., L137, why do you adopt the complex highway architecture to construct the discriminator? What are the performances of a single MLP (i.e., FC-Relu-FC) or a fully connected layer (i.e., FC)? Besides, to train the discriminator, what are the output dimension in the discriminator and the adopted loss, e.g., BCE/MSE?
- In Section 2.3, what are the performances of adopting the widely-used distance losses, e.g., KL Loss, to perform the alignment.
- In Equation (2), what are the performances of exchanging query and key distributions?
- If the approach brings errors, and what type of errors does it bring, and why? More importantly, are there errors caused by the approach when the model well aligns the distributions of query and key?
- The margins of the gains in all Tables are small and there appears to be no statistical significance testing. Statistical significant test and error ranges are highly appreciated.

4. Some parts of the paper can be made clearer for easier reading. 
- In all Tables, I encourage the authors to list several performances of recently published models to help the readers better understand the proposed approach.
- For the results of all base models in all experiments, it is better to clarify whether the results are reproduced by yourself or copied directly from the original paper.
- L48: 'emprical' -> 'empirical'; L180: 'is' -> 'are'; L269: 'with maximum' -> 'with a maximum'.

I highly encourage the authors to address these concerns and I am glad to increase my score if I see satisfying answers, thanks!

limitations_and_societal_impact: No.

- If the approach brings errors, and what type of errors does it bring, and why? More importantly, are there errors caused by the approach when the model well aligns the distributions of query and key?
- The margins of the gains in all Tables are small and there appears to be no statistical significance testing. Statistical significant test and error ranges are highly appreciated.
review_title: null

Meta-review: I recommend acceptance of the paper for the following reasons. As pointed out by reviewer E8gR, the premise of aligning key and value distributions might not be as important as it sounds at first for the purposes of improved learning. In addition, the analysis and ablations studies could be more thorough and convincing. Yet, experimentally the module works well. So despite potential limitations on the "intuition" and "analysis" of the module, there is value in knowing that such a module has empirical validation, especially for its potential impact on future analyses. There is also agreement from reviewers about this point and general support from two reviewers for the paper.

For these reasons, I recommend to accept the paper.
Decision: Accept (Poster)