[논문 리뷰] Talking-Heads Attention

Noam Shazeer, Zhenzhong Lan외 3명|arXiv (Cornell University)|2020. 03. 05.

Topic Modeling참고 문헌 9인용 수 49

한 줄 요약

토킹-헤드 어텐션(talking-heads attention)을 도입한 다중 헤드 어텐션의 변형으로, 소프트맥스 전후의 어텐션 헤드 차원에 학습된 선형 프로젝션을 적용하여 perplexity와 다운스트림 태스크 성능을 modest한 추가 파라미터로 개선한다.

ABSTRACT

We introduce "talking-heads attention" - a variation on multi-head attention which includes linearprojections across the attention-heads dimension, immediately before and after the softmax operation.While inserting only a small number of additional parameters and a moderate amount of additionalcomputation, talking-heads attention leads to better perplexities on masked language modeling tasks, aswell as better quality when transfer-learning to language comprehension and question answering tasks.

연구 동기 및 목표

• 전통적인 다중 헤드 어텐션에서 헤드별 프로젝션이 헤드가 너무 작을 때 표현력을 낮출 수 있는 한계를 지적하고 이를 보완할 방안을 모색한다.
• 학습된 프로젝션을 통해 어텐션 헤드 차원 간 정보 교환이 가능하도록 하는 메커니즘을 제안한다.
• 언어 모델링 및 이해 벤치마크에서의 사전 학습 및 미세 조정 태스크에 대한 실증적 이점을 입증한다.
• 일반적인 다중 헤드 어텐션과 비교한 새 메커니즘의 복잡도 및 분석을 제공한다.

제안 방법

• 소프트맥스 전후로 어텐션-헤드를 연결하는 두 개의 추가 학습 가능한 선형 프로젝션(P_l 및 P_w)을 도입하여 헤드 간 정보 교환을 가능하게 한다.
• TalkingHeadsAttention를 쿼리/키의 h_k, 로짓/가중치의 h, 값의 h_v의 세 가지 헤드 차원을 갖는 다중 헤드 어텐션의 일반화로 정의하고 두 가지 등가 표현(간결한 표현과 명시적 표현)을 제시한다.
• 계산 비용을 도출하고 Talking-헤드 프로젝션이 표준 비용에 항을 추가하지만 h_k, h, h_v가 적절히 선택될 때 이점이 있을 수 있음을 보인다.
• 표준 MHA와 Talking-Heads를 일반적인 쌍선형 다중 헤드 어텐션의 특수 사례로 연결하는 GBMA라는 통합 시각을 제공한다.
• 텍스트-투-텍스트 변환기(T5), ALBERT, 그리고 BERT 스타일 설정에서 다중 헤드와 토킹-헤드의 미리 학습 perplexity 및 다운스트럼 태스크를 비교하는 실험을 수행한다.

실험 결과

연구 질문

• RQ1 교차 헤드 프로젝션을 주입하면 표준 다중 헤드 어텐션에 비해 사전 학습 태스크의 모델 perplexity가 개선되는가?
• RQ2 토킹-헤드가 인코더/디코더 구성 요소 및 다양한 헤드 차원 구성에서 일관된 이점을 제공하는가?
• RQ3 토킹-헤드 어텐션을 전통적인 다중 헤드 어텐션과 비교할 때 복잡도 및 하드웨어 효율성의 트레이드오프는 무엇인가?
• RQ4 실무에서의 프로젝션 정밀도(logits/weights) 및 동적 대 정적 프로젝션과 함께 토킹-헤드가 어떻게 작용하는가?
• RQ5 BERT 및 ALBERT와 같은 널리 사용되는 아키텍처에 토킹-헤드의 이점이 전이되는가?

주요 결과

• 토킹-헤드 어텐션은 표준 다중 헤드 어텐션에 비해 마스킹된 언어 모델링 태스크에서 perplexity를 개선하며, 여러 헤드 크기 구성에서 이점을 보인다.
• softmax-헤드 차원 h를 늘리면 상당한 이점이 나타나는 반면, 매우 작게 설정된 키/값 차원(d_k, d_v)은 기존 MHA에서 성능 저하를 유발할 수 있지만 토킹-헤드가 이를 완화한다.
• 인코더 자체 어텐션에 토킹-헤드를 적용하면 디코더 어텐션에 적용했을 때보다 다운스트림 이점이 더 큰 경우가 많다.
• 세 가지 서로 다른 헤드 차원(h_k, h, h_v)을 구성하고 로짓- 및 가중치-프로젝션을 모두 사용할 때, 많은 실험에서 가장 큰 개선이 나타난다(예: T5 및 ALBERT에서).
• 동적 프로젝션은 사전 학습 perplexity를 줄일 수 있지만, 다운스트림 태스크에서는 일관되게 개선되지는 않았다.