[논문 리뷰] TensorCoder: Dimension-Wise Attention via Tensor Representation for Natural Language Modeling

Shuai Zhang, Peng Zhang외 4명|arXiv (Cornell University)|2020. 07. 28.

Topic Modeling참고 문헌 26인용 수 4

한 줄 요약

TensorCoder는 텐서 연산을 통해 쿼리-키 상호작용을 재정의함으로써 표준 Transformer의 자기주의 복잡도를 O(N²d)에서 O(Nd²)로 감소시키는 차원별 주의 메커니즘을 제안한다. 이는 PTB 마스크된 언어 모델링에서 13% 향상된 성능과 표준 Transformer 대비 3.5배 적은 FLOPs를 달성하며, 계산량을 줄임에도 불구하고 기계 번역에서 경쟁적인 성능 유지를 한다.

ABSTRACT

Transformer has been widely-used in many Natural Language Processing (NLP) tasks and the scaled dot-product attention between tokens is a core module of Transformer. This attention is a token-wise design and its complexity is quadratic to the length of sequence, limiting its application potential for long sequence tasks. In this paper, we propose a dimension-wise attention mechanism based on which a novel language modeling approach (namely TensorCoder) can be developed. The dimension-wise attention can reduce the attention complexity from the original $O(N^2d)$ to $O(Nd^2)$, where $N$ is the length of the sequence and $d$ is the dimensionality of head. We verify TensorCoder on two tasks including masked language modeling and neural machine translation. Compared with the original Transformer, TensorCoder not only greatly reduces the calculation of the original model but also obtains improved performance on masked language modeling task (in PTB dataset) and comparable performance on machine translation tasks.

연구 동기 및 목표

• 장문의 시퀀스에 대한 확장성 제약을 초래하는 표준 Transformer 자기주의의 이차적 계산 복잡도를 해결한다.
• 표준 토큰 기반 주의의 출력 기능을 유지하면서도 복잡도를 낮추는 차원 기반 주의 메커니즘을 개발한다.
• 텐서 연산과 컨볼루션을 활용해 차원 기반 주의를 Transformer 아키텍처에 통합한다.
• 마스크된 차원 기반 주의 메커니즘을 통해 디코더에서 자동재귀적 성질을 유지한다.
• 감소된 FLOPs와 메모리 사용량을 통해 저자원 환경에서의 사전학습 모델 효율적 배포를 가능하게 한다.

제안 방법

• 표준 쿼리-키 내적 주의(QKᵀ)를 복잡도를 O(N²d)에서 O(Nd²)로 감소시키는 차원 기반 주의 메커니즘(QᵀK)으로 대체한다.
• Khatri-Rao(KR) 곱을 사용해 주의를 제3차 텐서로 표현함으로써 암시적 및 명시적 출력 표현을 가능하게 한다.
• 텐서 표현에 컨볼루션 연산을 적용해 다차원 상호작용 특징을 추출하고 통합한다.
• 이전 위치로의 정보 흐름을 제한하는 마스크된 차원 기반 주의 메커니즘을 구축하여 자동재귀적 생성을 유지한다.
• 표준 다중헤드 주의를 대체하기 위해 차원 기반 주의를 인코더 및 디코더 레이어 양쪽에 통합한다.
• 텐서 기반 계산과 표현을 활용해 인코더 및 디코더 아키텍처 간 주의의 설계 및 구현을 통일한다.

실험 결과

연구 질문

• RQ1 차원 기반 주의 메커니즘이 자기주의의 이차적 복잡도를 줄이면서도 표현 능력을 유지할 수 있는가?
• RQ2 기존 Transformer 아키텍처와 호환되는 출력 형태를 갖는 차원 기반 주의는 어떻게 설계할 수 있는가?
• RQ3 텐서 기반 표현과 연산이 시퀀스 모델링에서 다차원 간 상호작용을 효과적으로 포착할 수 있는가?
• RQ4 제안된 마스크된 차원 기반 주의가 자동재귀적 디코딩에서 자동재귀적 생성 성질을 유지하는가?
• RQ5 차원 기반 주의의 감소된 복잡도가 장문의 시퀀스 및 저자원 자연어 처리 작업에서 효율성과 성능 향상에 얼마나 기여하는가?

주요 결과

• PTB 마스크된 언어 모델링 작업에서 TensorCoder는 표준 Transformer 대비 13% 높은 성능을 기록했으며, 추론 시간은 1/3으로 감소했다.
• WMT16 En-De 번역에서 TensorCoder는 BLEU 점수 33.1(표준 Transformer 대비 33.5)로 유사한 성능을 달성했지만, 주의 레이어의 FLOPs가 3.5배 감소했다.
• IWSLT16에서 TensorCoder는 표준 Transformer 대비 FLOPs를 1/3으로 줄여 강력한 효율성 향상을 입증했다.
• 차원 기반 주의 메커니즘이 계산 복잡도를 O(N²d)에서 O(Nd²)로 감소시켜 시퀀스 길이에 대해 선형 스케일링이 가능하게 했다.
• 메모리 및 계산 요구량을 크게 줄였음에도 불구하고 다양한 작업에서 경쟁적인 성능을 유지해 엣지 디바이스에 적합한 모델이 되었다.
• Khatri-Rao 곱과 텐서 컨볼루션의 활용은 주의 메커니즘 내에서 효과적인 특징 추출 및 통합을 가능하게 하여 낮은 비용으로도 높은 성능을 달성하는 데 기여했다.