[논문 리뷰] Multi-Head Attention: Collaborate Instead of Concatenate

Jean-Baptiste Cordonnier, Andreas Loukas외 1명|arXiv (Cornell University)|2020. 06. 29.

Advanced Neural Network Applications참고 문헌 27인용 수 76

한 줄 요약

이 논문은 협업형 다중 헤드 어텐션(collab MHA)을 도입하여 헤드 간 핵심/쿼리 투영을 공유하고 매개변수를 줄이며 NLP 번역, NLP 이해, 비전 과제에서 성능을 유지하거나 향상시키는 방법을 제시한다. 또한 표준 MHA로부터 텐서 분해를 통한 사후 재매개화(post-hoc reparameterization)도 보인다.

ABSTRACT

Attention layers are widely used in natural language processing (NLP) and are beginning to influence computer vision architectures. Training very large transformer models allowed significant improvement in both fields, but once trained, these networks show symptoms of over-parameterization. For instance, it is known that many attention heads can be pruned without impacting accuracy. This work aims to enhance current understanding on how multiple heads interact. Motivated by the observation that attention heads learn redundant key/query projections, we propose a collaborative multi-head attention layer that enables heads to learn shared projections. Our scheme decreases the number of parameters in an attention layer and can be used as a drop-in replacement in any transformer architecture. Our experiments confirm that sharing key/query dimensions can be exploited in language understanding, machine translation and vision. We also show that it is possible to re-parametrize a pre-trained multi-head attention layer into our collaborative attention layer. Collaborative multi-head attention reduces the size of the key and query projections by 4 for same accuracy and speed. Our code is public.

연구 동기 및 목표

• MHA 헤드 간의 핵심/쿼리 투영의 중복성을 규명하다.
• 헤드 간 투영을 공유하는 협업형 헤드 어텐션을 제안하다.
• 사전에 학습된 트랜스포머를 협업형 어텐션을 사용하도록 재매개화하는 방법을 보이다.
• NLP 및 비전 과제에서 매개변수 및 계산 효율성을 입증하다.
• 트랜스포머 아키텍처에서 바로 적용 가능한 대체에 대한 실용적 지침을 제공하다.

제안 방법

• 헤드 간 W_Q W_K에 대한 PCA로 중복성을 분석하여 연결된 헤드의 저랭크(low-rank) 동작을 보이다.
• 모든 헤드가 W_Q 및 W_K를 공유 차원 ˜Dk로 투영하고 혼합 벡터 m_i를 사용하는 협업형 헤드 어텐션(collab MHA)을 정의한다.
• 협업형 헤드를 표준 MHA의 특수한 경우로 관련시킨다(˜Dk = Nh dk).
• 테커/ CP(정규형) 텐서 분해를 사용하여 기존 학습된 어텐션 레이어를 협업형 어텐션으로 재매개화한다( W̃_Q, W̃_K, 그리고 혼합 행렬 M 도출).
• 어텐션 점수에서 콘텐츠/컨텍스트 분해 및 기여하지 않는 바이어스 항목을 통한 바이어스 처리 논의.
• 복잡도/매개변수 분석을 제공하고 가지치기와 비교한다.
• NMT, NLU(GLUE), 비전(ImageNet) 과제에서 검증하며, from-scratch 학습과 사후 재매개화(post-hoc reparameterization)를 포함한다.

실험 결과

연구 질문

• RQ1 어텐션 헤드가 공유 투영을 가능하게 하는 핵심/쿼리 투영의 중복을 학습하는가?
• RQ2 협업형 MHA가 매개변수를 줄이면서 표준 MHA의 성능에 필적하거나 이를 능가할 수 있는가?
• RQ3 전체 재학습 없이 사전 학습된 트랜스포머를 협업형 MHA로 변환하는 것이 가능한가(사후 재매개화)?
• RQ4 NLP 및 비전 과제에서 매개변수 및 속도 측면의 실용적 효율성 향상은 어떤가?
• RQ5 GLUE, ImageNet 같은 벤치마크에서 협업형 MHA가 모델 압축 및 파인튜닝 필요성에 어떤 영향을 미치는가?

주요 결과

• 헤드 간에 투영을 부분 공간 전체로 공유하는 경향이 있어 연결되었을 때 W_Q W_K가 저랭크가 된다.
• 협업형 MHA는 핵심/쿼리 투영(˜Dk)과 혼합 행렬 M을 Nh 헤드에 걸쳐 공유함으로써 매개변수 수를 줄인다.
• 표준 MHA를 collab MHA로 교체하면 정확도 손실 없이 NMT에서 핵심/쿼리 차원을 4배 감소시키면서 성능을 유지하거나 개선한다.
• collab MHA로 from-scratch 학습은 BLEU 손실 없이 최대 35% 더 적은 어텐션 매개변수를 산출하고 때로는 WMT14 EN-DE에서 BLEU가 향상된다.
• 텐서 분해를 이용한 사후 재매개화를 통해 사전 학습된 트랜스포머를 협업 어텐션으로 재매개화하면 정확도가 보존되거나 약간 감소하되 상당한 압축이 가능하며(예: GLUE에서 ˜Dk의 2×–3× 감소).
• 비전에서 collab MHA는 더 낮은 ˜Dk 값에서도 연결 기반 성능에 필적하거나 이를 상회하며 ImageNet에서 최소한의 정확도 감소로 압축을 가능하게 한다.