[论文解读] SparseBERT: Rethinking the Importance Analysis in Self-attention

Shi Han, Jiahui Gao+5 more|arXiv (Cornell University)|Feb 25, 2021

Topic Modeling参考文献 39被引用 13

一句话总结

本文提出 SparseBERT，一种通过重新思考自注意力重要性的方法，发现注意力矩阵中的对角线元素最重要性最低，可被移除而不损失性能。通过使用 Gumbel-sigmoid 的可微分注意力掩码（DAM），模型端到端学习稀疏、结构化的注意力模式，在开发集上实现 91.3% 的稀疏度与 80.9% 的 GLUE 分数。

ABSTRACT

Transformer-based models are popularly used in natural language processing (NLP). Its core component, self-attention, has aroused widespread interest. To understand the self-attention mechanism, a direct method is to visualize the attention map of a pre-trained model. Based on the patterns observed, a series of efficient Transformers with different sparse attention masks have been proposed. From a theoretical perspective, universal approximability of Transformer-based models is also recently proved. However, the above understanding and analysis of self-attention is based on a pre-trained model. To rethink the importance analysis in self-attention, we study the significance of different positions in attention matrix during pre-training. A surprising result is that diagonal elements in the attention map are the least important compared with other attention positions. We provide a proof showing that these diagonal elements can indeed be removed without deteriorating model performance. Furthermore, we propose a Differentiable Attention Mask (DAM) algorithm, which further guides the design of the SparseBERT. Extensive experiments verify our interesting findings and illustrate the effect of the proposed algorithm.

研究动机与目标

• 在 BERT 预训练过程中，重新评估自注意力机制中不同注意力位置的重要性。
• 探究注意力矩阵中的对角线元素是否对模型性能真正必要。
• 开发一种可学习、可微分的机制，以发现高效、稀疏的注意力模式。
• 设计一种稀疏 Transformer 架构（SparseBERT），在保持性能的同时降低计算成本。
• 验证移除对角线注意力是否能提高稀疏度而不损害下游任务的准确性。

提出的方法

• 提出一种可微分注意力掩码（DAM）算法，通过软性、可微分的掩码实现端到端学习注意力模式。
• 引入 Gumbel-sigmoid 函数，在训练期间对二值注意力掩码进行可微采样，支持基于梯度的优化。
• 将注意力掩码的优化作为预训练过程的一部分，联合学习掩码与模型参数。
• 使用可学习的温度参数控制掩码的软性程度，通过直通估计实现离散注意力模式。
• 施加结构化稀疏性约束，以确保生成的注意力模式具有可解释性且高效（例如，局部、全局或混合模式）。
• 采用加权损失函数，结合模型准确率与稀疏性正则化，以平衡性能与效率。

实验结果

研究问题

• RQ1 自注意力矩阵中的对角线元素是否对模型性能真正重要，还是可以安全地被移除？
• RQ2 一种可微分、端到端的方法能否学习到与手工设计的稀疏注意力机制性能相当或更优的稀疏注意力模式？
• RQ3 在预训练过程中，注意力位置的重要性（例如，对角线、邻近位置、特殊标记）在不同注意力头中如何变化？
• RQ4 在不降低下游任务性能的前提下，通过移除对角线注意力，稀疏度最多可提升多少？
• RQ5 与两阶段剪枝相比，单阶段可微分掩码学习在性能与稀疏度控制方面表现如何？

主要发现

• 自注意力矩阵中的对角线元素最重要性最低，其移除在多个任务中均未导致性能下降。
• SparseBERT 在 91.3% 的稀疏度下，于开发集上实现 80.9% 的平均 GLUE 分数，证明了其高效性且未造成准确率损失。
• 所提出的 DAM 算法在不同头中学习到多样化、结构化的注意力模式，邻近位置与特殊标记（如 [CLS]、[SEP]）获得的注意力高于对角线。
• 移除对角线注意力使分段掩码的稀疏度从 70.4% 提升至 71.2%，在 DAM 采样掩码中最高可达 93.5%，且性能无下降。
• 单阶段可微分掩码学习在下游准确率方面优于两阶段剪枝，尽管其通过超参数 λ 对稀疏度的直接控制能力较弱。
• 理论分析证实，对角线元素在自注意力机制的通用近似能力中并非必要。