[논문 리뷰] Improving Transformers with Probabilistic Attention Keys

Tam Nguyen, Tan M. Nguyen외 6명|arXiv (Cornell University)|2021. 10. 16.

Topic Modeling인용 수 6

한 줄 요약

이 논문은 다중 가우시안 분포를 사용하여 표준 어텐션 키를 대체함으로써 어텐션 헤드 간의 부재를 줄이는 새로운 트랜스포머 아키텍처인 Transformer-MGK을 제안한다. 키를 확률적으로 모델링함으로써 학습과 추론 속도가 향상되고, 파라미터 수와 FLOPs가 감소하며, 더 적은 헤드 수로도 표준 트랜스포머와 비슷하거나 더 뛰어난 성능을 달성한다. 이는 Wikitext-103 및 Long Range Arena 벤치마크에서 입증되었다.

ABSTRACT

Multi-head attention is a driving force behind state-of-the-art transformers, which achieve remarkable performance across a variety of natural language processing (NLP) and computer vision tasks. It has been observed that for many applications, those attention heads learn redundant embedding, and most of them can be removed without degrading the performance of the model. Inspired by this observation, we propose Transformer with a Mixture of Gaussian Keys (Transformer-MGK), a novel transformer architecture that replaces redundant heads in transformers with a mixture of keys at each head. These mixtures of keys follow a Gaussian mixture model and allow each attention head to focus on different parts of the input sequence efficiently. Compared to its conventional transformer counterpart, Transformer-MGK accelerates training and inference, has fewer parameters, and requires fewer FLOPs to compute while achieving comparable or better accuracy across tasks. Transformer-MGK can also be easily extended to use with linear attention. We empirically demonstrate the advantage of Transformer-MGK in a range of practical applications, including language modeling and tasks that involve very long sequences. On the Wikitext-103 and Long Range Arena benchmark, Transformer-MGKs with 4 heads attain comparable or better performance to the baseline transformers with 8 heads.

연구 동기 및 목표

• 다중 헤드 어텐션 메커니즘에서 관찰되는 반복적인 어텐션 패턴 문제를 해결하기 위해.
• 모델 성능을 희생시키지 않고 트랜스포머의 계산 및 파라미터 오버헤드를 줄이기 위해.
• 각 헤드의 표현 능력을 향상시키기 위해 어텐션 키를 확률적 프레임워크로 모델링하기 위해.
• 다수의 부재 헤드를 하나의 더 표현력 있는 다중 가우시안 키 메커니즘으로 대체함으로써 효율적인 추론과 학습을 가능하게 하기 위해.
• 장시퀀스 모델링을 위한 선형 어텐션과의 통합을 위해.

제안 방법

• 어텐션 키를 다중 가우시안 모델(GMM)로 모델링함으로써 자기어텐션의 확률적 해석을 제안하고, 어텐션 점수를 GMM 하위 확률에 연결한다.
• 다중 헤드 어텐션 헤드의 표준 키 벡터를 다중 가우시안 성분의 혼합으로 대체하여, 한 개의 헤드가 동시에 여러 다릅니다 시퀀스 위치에 주목할 수 있도록 한다.
• 어텐션 계산을 가우시안 성분에 대한 가중 평균으로 유도하며, 어텐션 가중치는 GMM 하위 확률에 해당한다.
• 가우시안 성분의 평균, 분산 및 혼합 가중치를 최적화하는 미분 가능한 학습 절차를 도입한다.
• 기존 아키텍처 및 최적화 기법과의 호환성을 유지하면서 표준 트랜스포머 블록에 다중 가우시안 키를 통합한다.
• 장시퀀스 모델링을 위한 효율성 향상을 위해 확률적 키 표현을 기반으로 어텐션 메커니즘을 근사함으로써 선형 어텐션으로 확장한다.

실험 결과

연구 질문

• RQ1 어텐션 키의 확률적 모델링이 트랜스포머의 다중 헤드 어텐션 헤드 간의 부재를 줄일 수 있는가?
• RQ2 표준 키를 다중 가우시안으로 대체함으로써 각 어텐션 헤드의 표현 능력은 향상되고 모델 복잡도는 감소하는가?
• RQ3 표준 트랜스포머에 비해 더 적은 어텐션 헤드로 Transformer-MGK가 성능을 유지하거나 향상시킬 수 있는가?
• RQ4 장시퀀스 작업 및 자원이 제한된 환경에서 다중 가우시안 키 메커니즘이 어떻게 성능을 발휘하는가?
• RQ5 제안된 방법이 선형 어텐션 메커니즘으로 효율적으로 확장되어 스케일링 가능한 장문맥 모델링을 가능하게 하는가?

주요 결과

• Wikitext-103 언어 모델링 벤치마크에서 4개의 헤드를 사용한 Transformer-MGK가 표준 트랜스포머의 8개 헤드와 비슷하거나 더 뛰어난 성능을 달성한다.
• Long Range Arena 벤치마크에서 4개의 헤드를 사용한 Transformer-MGK가 표준 트랜스포머의 8개 헤드와 동등하거나 이를 초월하며, 특히 장문맥 작업에서 뛰어난 성능을 보인다.
• 헤드 간의 부재 감소로 인해 FLOPs와 파라미터 수가 감소하고, 학습 및 추론 속도가 향상된다.
• 다중 가우시안 키 메커니즘이 각 어텐션 헤드가 입력 시퀀스의 여러 다릅니다 부분에 동시에 주목할 수 있도록 하여 표현 다양성을 증가시킨다.
• 이 방법은 선형 어텐션과도 호환되어 매우 긴 시퀀스의 효율적 처리를 가능하게 한다.
• 실험 결과는 확률적 키 모델링이 다양한 NLP 작업에서 더 안정적이고 일반화 능력이 뛰어난 어텐션 패턴을 유도함을 보여준다.