[论文解读] Eigen Analysis of Self-Attention and its Reconstruction from Partial Computation

Srinadh Bhojanapalli, Ayan Chakrabarti+4 more|arXiv (Cornell University)|Jun 16, 2021

Advanced Neural Network Applications参考文献 22被引用 5

一句话总结

该论文提出了一种方法，仅通过计算部分点积注意力分数，即可重建变换器中的完整自注意力矩阵，其核心思想是利用注意力分数位于低维特征空间的观察结果。通过贪心算法选择关键的查询-键对进行精确计算，并利用学习到的线性变换重建其余部分，该方法将注意力计算成本降低了25%–45%，同时在MLM和NLI任务上的模型准确率仅出现轻微下降。

ABSTRACT

State-of-the-art transformer models use pairwise dot-product based self-attention, which comes at a computational cost quadratic in the input sequence length. In this paper, we investigate the global structure of attention scores computed using this dot product mechanism on a typical distribution of inputs, and study the principal components of their variation. Through eigen analysis of full attention score matrices, as well as of their individual rows, we find that most of the variation among attention scores lie in a low-dimensional eigenspace. Moreover, we find significant overlap between these eigenspaces for different layers and even different transformer models. Based on this, we propose to compute scores only for a partial subset of token pairs, and use them to estimate scores for the remaining pairs. Beyond investigating the accuracy of reconstructing attention scores themselves, we investigate training transformer models that employ these approximations, and analyze the effect on overall accuracy. Our analysis and the proposed method provide insights into how to balance the benefits of exact pair-wise attention and its significant computational expense.

研究动机与目标

• 研究变换器模型在自然语言输入下自注意力分数的低维结构特征。
• 确定注意力分数的变化是否能通过少量主成分在不同层和模型间被捕捉。
• 开发一种仅使用部分精确点积计算来近似完整注意力矩阵的方法。
• 评估在端到端训练过程中，部分注意力计算对下游模型性能的影响。
• 探索跨模型和层的共享特征空间是否可支持通用的重建方法。

提出的方法

• 对完整的注意力分数矩阵及单个查询行矩阵进行特征值分析，以识别主导主成分。
• 使用贪心算法选择一组查询-键对进行精确注意力计算，目标是最大化信息增益。
• 通过最小二乘优化学习重建矩阵 R，将选定的注意力分数映射为完整的注意力向量。
• 将部分计算与重建过程整合进变换器的训练流程，支持通过选择与重建过程的反向传播。
• 在训练过程中与模型权重联合优化重建矩阵 R，可选择在所有层间共享或为每层分别学习，以适应特定任务的分布。
• 在 BERT-base 上评估该方法，包括掩码语言建模预训练及在 MNLI 基准上的微调。

实验结果

研究问题

• RQ1 在自然语言输入上计算的自注意力分数是否表现出其变化的低维结构？
• RQ2 注意力分数的主成分在不同层、头以及变换器模型之间有多大的共享程度？
• RQ3 仅使用少量精确计算的注意力分数，能否通过线性重建方法以低误差重建完整矩阵？
• RQ4 在部分注意力计算与重建的设置下训练变换器，是否能在下游NLP任务中保持有竞争力的性能？
• RQ5 重建矩阵的选择（固定 vs. 可学习，共享 vs. 每层独立）如何影响模型准确率与效率？

主要发现

• 不同层和模型中的注意力分数在主特征空间上表现出显著重叠，表明存在共享的低维结构。
• 仅使用每个查询的16–32个精确分数即可实现完整注意力矩阵的低重建误差，且均方误差由学习到的重建矩阵最小化。
• 在反向传播中联合训练重建矩阵 R 优于使用固定且预计算的 R，尤其在 k 较小时效果更明显。
• 对于较小的 k（如 k=16），跨层共享 R 可提升性能；而对于较大的 k（如 k=32），每层独立的 R 表现更优。
• 该方法将注意力 FLOPs 最多降低 45%，同时在 MNLI 下游任务上保持了 97% 的基线准确率。
• 使用部分注意力训练的模型在 k=32 时于 MNLI 上达到 79.7% 的准确率，相比完整模型的 81.6%，展现出优异的准确率-效率权衡。