[논문 리뷰] Fast Transformer Decoding: One Write-Head is All You Need

Noam Shazeer|arXiv (Cornell University)|2019. 11. 06.

Topic Modeling참고 문헌 5인용 수 26

한 줄 요약

이 논문은 다중 쿼리 어텐션을 제안하며, 이는 모든 어텐션 헤드 간에 키와 값이 공유되는 멀티헤드 어텐션의 변종으로, 점진적 디코딩 중 메모리 대역폭을 크게 감소시킨다. 이 방법은 품질 저하가 미미한 수준에서 최대 12배 빠른 추론을 달성하여 지연 민감한 응용 분야에 이상적이다.

ABSTRACT

Multi-head attention layers, as used in the Transformer neural sequence model, are a powerful alternative to RNNs for moving information across and between sequences. While training these layers is generally fast and simple, due to parallelizability across the length of the sequence, incremental inference (where such paralleization is impossible) is often slow, due to the memory-bandwidth cost of repeatedly loading the large "keys" and "values" tensors. We propose a variant called multi-query attention, where the keys and values are shared across all of the different attention "heads", greatly reducing the size of these tensors and hence the memory bandwidth requirements of incremental decoding. We verify experimentally that the resulting models can indeed be much faster to decode, and incur only minor quality degradation from the baseline.

연구 동기 및 목표

• 현대 하드웨어에서 인코딩의 점진적 디코딩에 따른 높은 메모리 대역폭 비용을 해결하기 위해.
• 모델 성능에 상당한 영향을 주지 않으면서 멀티헤드 어텐션의 키 및 값 텐서 크기를 줄이기 위해.
• 자기회귀적 생성 중 반복적인 메모리 액세스를 최소화함으로써 지연에 민감한 응용 분야에서 더 빠른 추론을 가능하게 하기 위해.
• 헤드 간에 키와 값을 공유하는 것이 표준 멀티헤드 어텐션과 경쟁 가능한 품질을 유지하는지 평가하기 위해.

제안 방법

• 모든 어텐션 헤드가 동일한 키 및 값 프로젝션을 공유하는 다중 쿼리 어텐션을 제안하며, 이는 파rameter 수와 메모리 대역폭을 감소시킨다.
• 각 헤드당 별도의 키 및 값 프로젝션을 사용하는 표준 멀티헤드 어텐션 메커니즘을 수정하여, 모든 헤드에 걸쳐 하나의 공통 프로젝션을 사용하도록 한다.
• 각 헤드당 동일한 쿼리 프로젝션을 사용하지만, 키 및 값 프로젝션은 모든 헤드 간에 공유되며, 이로 인해 K 및 V 텐서 크기가 O(hd)에서 O(d)로 감소한다.
• 공유된 K 및 V 행렬을 사용하는 표준 자기어텐션 계산을 활용하며, 동일한 어텐션 메커니즘을 유지하지만 메모리 프로파일이 감소한다.
• 표준 및 국소 어텐션 변종 모두에 이 방법을 적용하여 다양한 설정에서 성능을 평가한다.
• TPUv2 하드웨어에서 배치 추론과 점진적 디코딩을 적용하여 학습 및 추론 비용을 측정한다.

실험 결과

연구 질문

• RQ1 모든 어텐션 헤드 간에 키와 값을 공유함으로써 점진적 디코딩 중 메모리 대역폭을 줄일 수 있을까, 성능 저하가 크지 않은가?
• RQ2 제안된 다중 쿼리 어텐션은 표준 멀티헤드 어텐션 및 기타 감소된 어텐션 변종과 비교해 품질 면에서 어떻게 성과를 내는가?
• RQ3 다중 쿼리 어텐션은 자기회귀적 생성에서 추론 속도를 얼마나 향상시키는가?
• RQ4 이 방법은 기계 번역 및 언어 모델링 벤치마크에서 경쟁적인 성능을 유지하는가?
• RQ5 학습 시간이나 모델 크기를 늘리지 않고도 속도 향상을 달성할 수 있는가?

주요 결과

• TPUv2에서 다중 쿼리 모델은 토큰당 디코더 추론 시간을 47μs에서 3.8μs로 줄여 점진적 디코딩 단계에서 12.4배의 속도 향상을 달성했다.
• 토큰당 추론 비용은 기준 모델의 46μs에서 다중 쿼리 모델의 3.8μs로 감소했으며, 인코더 비용도 1.7μs에서 1.5μs로 감소했다.
• WMT14 EN-DE 번역 작업에서 다중 쿼리 모델은 빔 서치(4)를 사용해 BLEU 점수 28.5를 기록했으며, 기준 모델(28.4)보다 略적으로 높았다.
• 10억 단어 언어 모델링 벤치마크에서 다중 쿼리 모델은 퍼플렉서티 30.2를 기록했으며, 기준 모델(29.9)과 비교해 다소 열 劣했다.
• 구조적 변경에도 불구하고 학습 시간은 13.0μs/token(다중 쿼리)에서 13.2μs/token(기준)로 유사하게 유지되어 학습 오버헤드가 없음을 보여주었다.
• 모든 대안(즉, h, dk, 또는 dv 감소)보다 품질과 속도 면에서 뛰어나, 이는 메모리 대역폭 최적화 솔루션으로서의 효과성을 입증한다.