[논문 리뷰] Eigen Analysis of Self-Attention and its Reconstruction from Partial Computation

Srinadh Bhojanapalli, Ayan Chakrabarti외 4명|arXiv (Cornell University)|2021. 06. 16.

Advanced Neural Network Applications참고 문헌 22인용 수 5

한 줄 요약

이 논문은 자연어 입력에서 자기주의 어텐션 점수의 저차원 고유공간에 기반해, 부분적으로 계산된 내적 어텐션 점수만을 사용하여 전체 자기주의 어텐션 행렬을 재구성하는 방법을 제안한다. 이는 어텐션 점수가 저차원 고유공간에 존재한다는 관찰에 기반한다. 탐욕 알고리즘을 통해 핵심 쿼리-키 쌍을 선택하고 나머지를 학습된 선형 변환을 통해 재구성함으로써, 모델 정확도가 MLM 및 NLI 작업에서 약간 감소하는 것과 맞물려 어텐션 계산 비용을 25–45% 감소시킨다.

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을 모델 가중치와 함께 동시에 훈련하여, 작업에 특화된 분포에 적응시키기 위해, 레이어 간 공유 또는 각 레이어별로 설정하기 위해.
• 마스크된 언어 모델링 사전 훈련과 MNLI 벤치마크에서의 파인튜닝을 통해 BERT-base에서 방법을 평가하기 위해.

실험 결과

연구 질문

• 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%였으며, 이는 강력한 정확도-효율성 트레이드오프를 보여준다.