[논문 리뷰] Sharing Attention Weights for Fast Transformer

Tong Xiao, Yinqiao Li외 3명|arXiv (Cornell University)|2019. 06. 26.

Topic Modeling참고 문헌 20인용 수 7

한 줄 요약

이 논문은 자기주도적 인코더-디코더 어텐션 메커니즘에서 인접한 레이어 간에 어텐션 가중치를 공유함으로써 트랜스포머 추론을 가속화하는 공유 어텐션 네트워크(San)를 제안한다. 은닉 상태의 수직적 재사용과 함께 공유 정책을 공동으로 학습함으로써, 캐시 기반 기준 대비 1.3배의 속도 향상을 달성하고 캐시 없이도 최대 16배의 속도 향상을 이룩했으며 BLEU 점수 저하가 최소에 그친다.

ABSTRACT

Recently, the Transformer machine translation system has shown strong results by stacking attention layers on both the source and target-language sides. But the inference of this model is slow due to the heavy use of dot-product attention in auto-regressive decoding. In this paper we speed up Transformer via a fast and lightweight attention model. More specifically, we share attention weights in adjacent layers and enable the efficient re-use of hidden states in a vertical manner. Moreover, the sharing policy can be jointly learned with the MT model. We test our approach on ten WMT and NIST OpenMT tasks. Experimental results show that it yields an average of 1.3X speed-up (with almost no decrease in BLEU) on top of a state-of-the-art implementation that has already adopted a cache for fast inference. Also, our approach obtains a 1.8X speed-up when it works with the extsc{Aan} model. This is even 16 times faster than the baseline with no use of the attention cache.

연구 동기 및 목표

• 자기회귀적 디코딩과 반복적인 도트프로덕트 어텐션 계산으로 인해 발생하는 트랜스포머 모델의 높은 추론 지연을 해결하기 위해.
• 인접한 레이어 간 어텐션 가중치 분포의 구조적 유사성을 활용하여 스택형 어텐션 레이어의 부정산 계산을 줄이기 위해.
• 휴리스틱 공유 규칙 대신, 공유 정책과 MT 모델을 종합적으로 최적화하는 공동 학습 프레임워크를 개발하기 위해.
• 가중치 공유를 통한 레이어 간 은닉 상태 재사용으로 메모리 프로파일을 최소화하기 위해.
• 기존 가속 기술(예: 어텐션 캐싱)과 조합했을 때도 유의미한 정확도 손실 없이 빠른 추론을 달성하기 위해.

제안 방법

• 자기어텐션과 인코더-디코더 어텐션 서브레이어 양쪽에서 인접한 레이어 간에 수직적으로 어텐션 가중치를 공유하는 공유 어텐션 네트워크(San)를 제안한다.
• 동일한 어텐션 가중치 매트릭스를 여러 개의 스택형 레이어에 걸쳐 사용하도록 어텐션 메커니즘을 재구성함으로써 부정산 계산을 감소시킨다.
• 공유 정책을 MT 모델과 함께 종단간(end-to-end)으로 학습하는 전략을 도입하여 네트워크가 최적의 공유 패턴을 자동으로 탐색할 수 있도록 한다.
• 공유 메모리 위치에 은닉 상태를 저장함으로써 효율적인 은닉 상태 재사용을 가능하게 하여 메모리 프로파일을 감소시키고 추론 속도를 향상시킨다.
• 기존 가속 기술(예: 어텐션 캐싱, Aan 모델)과 통합하여 상호 보완적이고 독립적인 향상을 보여주며, 이는 서로 다른 가속 기술 간의 보완적 효과를 입증한다.
• 기본으로 스케일드 도트프로덕트 어텐션 메커니즘을 사용하고, 레이어 간에 공유 가중치를 적용함으로써 모델 구조를 유지하면서도 효율성을 향상시킨다.

실험 결과

연구 질문

• RQ1 인접한 레이어 간 어텐션 가중치 공유가 번역 품질 저하 없이 추론 시간을 단축시킬 수 있는가?
• RQ2 MT 모델과 함께 공유 정책를 공동으로 학습할 경우, 휴리스틱 공유 전략보다 성능이 뛰어나고 수렴 속도가 더 빠른가?
• RQ3 기존의 가속 기법(예: 어텐션 캐싱 또는 Aan 모델)과 결합했을 때 San의 효과는 어떠한가?
• RQ4 훈련된 트랜스포머에서 인접한 레이어들이 유사한 어텐션 가중치 분포를 보이는 정도는 어느 정도이며, 이는 가중치 공유의 타당성을 뒷받침하는가?
• RQ5 San은 추론 속도를 유지하거나 향상시키면서도 메모리 사용량을 줄일 수 있는가?

주요 결과

• San은 이미 어텐션 캐시를 사용하는 최신 트랜스포머 구현 대비 평균 1.3배의 속도 향상을 달성했으며 BLEU 점수 저하가 거의 없었다.
• Aan 모델과 조합했을 경우 San은 1.8배의 속도 향상을 기록했으며, 서로 다른 가속 기술 간의 보완적 이점을 입증했다.
• San과 Aan을 함께 사용한 시스템은 캐시 기반 기준 없이도 기준 대비 최대 16배 빠른 속도를 기록했다.
• 훈련 과정에서 인접 레이어 간 어텐션 가중치 분포 간의 JS 발산이 크게 감소하여, 레이어들이 자연스럽게 유사한 어텐션 패턴으로 수렴한다는 것을 시사한다.
• San의 훈련 가능도는 기준 대비 略로 높아, 공유 제도 하에서도 모델이 효과적이고 안정적임을 시사한다.
• 이 방법은 구현이 단순하며, 기존 기술과 독립적이며 다양한 트랜스포머 변종에 직접 적용 가능하다.