[논문 리뷰] Mega: Moving Average Equipped Gated Attention

Xuezhe Ma, Chunting Zhou외 6명|arXiv (Cornell University)|2022. 09. 21.

Topic Modeling인용 수 36

한 줄 요약

Mega는 지수 이동 평균(EMA)을 단일 헤드 게이트드 어텐션에 통합한 이동 평균 장치를 갖춘 게이트드 어텐션 메커니즘을 도입하고, 선형 복잡도 Mega-Chunk 변형을 통해 여러 모달리티에서 장거리 시퀀스 작업에 대해 강력한 성능을 발휘합니다.

ABSTRACT

The design choices in the Transformer attention mechanism, including weak inductive bias and quadratic computational complexity, have limited its application for modeling long sequences. In this paper, we introduce Mega, a simple, theoretically grounded, single-head gated attention mechanism equipped with (exponential) moving average to incorporate inductive bias of position-aware local dependencies into the position-agnostic attention mechanism. We further propose a variant of Mega that offers linear time and space complexity yet yields only minimal quality loss, by efficiently splitting the whole sequence into multiple chunks with fixed length. Extensive experiments on a wide range of sequence modeling benchmarks, including the Long Range Arena, neural machine translation, auto-regressive language modeling, and image and speech classification, show that Mega achieves significant improvements over other sequence models, including variants of Transformers and recent state space models.

연구 동기 및 목표

• 강한 귀납 편향과 유연한 어텐션을 결합해 장거리 시퀀스 모델링을 개선하려는 동기 부여.
• EMA를 강화한 단일 헤드 게이트드 어텐션 메커니즘 개발(GAU 기반).
• 시퀀스 처리를 선형 시간으로 확장하기 위한 Mega-chunk 변형 도입.
• 다양한 모달리티와 벤치마크에서 Mega의 효과 입증.
• 단일 헤드 게이트드 어텐션의 표현력을 이론적으로 정당화.

제안 방법

• 지역적이고 시간적으로 감소하는 의존성을 포착하기 위한 다차원 댐핑 EMA 도입.
• EMA를 단일 헤드 게이트드 어텐션 프레임워크(GAU 기반)에 내재화하여 EMA 보강 표현으로부터 쿼리, 키, 값을 도출.
• 소프트맥스의 대안으로 다른 모달리티를 위한 Laplace 기반 어텐션 함수를 채택.
• 훈련 안정화를 위해 Mega 블록 내에 게이트드 잔차 연결과 정규화 단계 사용.
• 고정 길이 청크를 처리하고 크로스 청크 컨텍스트를 EMA로 제공하여 선형 시간/공간 복잡도를 달성하는 Mega-chunk 제안.
• 언어, 비전, 음성에서 어텐션 함수(softmax, relu^2, laplace)를 비교하여 모듈별로 적합한 옵션 선택.

실험 결과

연구 질문

• RQ1 EMA를 게이트드 어텐션에 통합하면 표준 어텐션에 비해 로컬 및 장거리 의존성 모델링이 개선되는가?
• RQ2 EMA를 갖춘 단일 헤드 게이트드 어텐션이 다중 헤드 어텐션과 비견될 정도의 표현력을 달성할 수 있는가?
• RQ3 Mega-chunk 변형이 계산 비용을 선형으로 줄이면서 성능을 유지하는가?
• RQ4 Mega가 다양한 모달리티(text, image, speech)와 작업(LRA, MT, LM, Vision, Speech)에서 어떤 성능을 보이는가?
• RQ5 도메인 간 Mega에 가장 적합한 어텐션 함수는 무엇인가?

주요 결과

모델	LRA (정확도 ↑)	WMT16 (BLEU ↑)	WT103 (PPL ↓)	ImageNet (정확도 ↑)	SC (정확도 ↑)	평균 (↑)	속도	메모리
XFM	59.24	–	18.66	81.80	✗	–	–	–
XFM ${}$	37.11	65.21	79.14	42.94	71.83	59.24	–	–
Reformer	50.67	–	20.95	–	–	–	–	–
Linformer	51.36	–	52.27	–	–	–	5.5×	0.10×
BigBird	55.01	–	59.29	–	–	–	1.1×	0.30×
Performer	51.41	–	53.82	42.77	77.05	–	5.7×	0.11×
Luna-256	61.95	–	79.56	47.86	78.55	–	4.9×	0.16×
S4-v1	80.48	–	87.09	87.26	86.05	80.48	–	–
S4-v2	86.09	–	90.90	88.65	94.20	86.09	–	–
S4-v2 ${}$	85.86	–	90.94	88.48	94.01	85.86	4.8×	0.14×
Mega	88.21	29.18	18.07	82.31	97.30	88.21	2.9×	0.31×
Mega-chunk	85.66	90.19	90.97	85.80	94.41	85.66	5.5×	0.13×

• Mega는 모든 Long Range Arena 작업 및 관련 벤치마크에서 Transformer 및 S4 변형을 능가한다.
• Mega는 LRA에서 평균 정확도 88.21을 달성하고 S4 변형의 86.09–86.35 및 XFM의 59.24를 상회하며 선형 시간 변형으로 강력한 장거리 모델링을 보여준다.
• Mega-(chunk)은 예를 들어 LRA 텍스트 작업에서 Vanilla Transformer에 비해 5.5배 빠르고 13%의 메모리 사용 이점을 제공한다.
• 언어 작업에서 Mega 내의 소프트맥스 어텐션은 강력한 성능을 보이고, Laplace 어텐션은 비전 및 음성 작업에서 뛰어나다.
• Mega-base는 Transformer-base에 비해 WMT’16 영어-독일 BLEU 점수를 1.1포인트 이상 향상시킨다.
• ImageNet-1k에서 Mega는 DeiT-B보다 약 0.5%p 상위-1 정확도를 달성한다.