[논문 리뷰] Sparse Sinkhorn Attention

Yi Tay, Dara Bahri외 3명|arXiv (Cornell University)|2020. 02. 26.

Natural Language Processing Techniques참고 문헌 24인용 수 77

한 줄 요약

Sparse Sinkhorn Attention은 differentiable sorting 메커니즘을 도입하여 메모리 효율적인, 희소 어텐션을 생성하고 로컬 윈도우를 통한 준-글로벌 컨텍스트를 가능하게 하며 Vanilla Transformers와 경쟁력 있는 성능을 제공합니다.

ABSTRACT

We propose Sparse Sinkhorn Attention, a new efficient and sparse method for learning to attend. Our method is based on differentiable sorting of internal representations. Concretely, we introduce a meta sorting network that learns to generate latent permutations over sequences. Given sorted sequences, we are then able to compute quasi-global attention with only local windows, improving the memory efficiency of the attention module. To this end, we propose new algorithmic innovations such as Causal Sinkhorn Balancing and SortCut, a dynamic sequence truncation method for tailoring Sinkhorn Attention for encoding and/or decoding purposes. Via extensive experiments on algorithmic seq2seq sorting, language modeling, pixel-wise image generation, document classification and natural language inference, we demonstrate that our memory efficient Sinkhorn Attention method is competitive with vanilla attention and consistently outperforms recently proposed efficient Transformer models such as Sparse Transformers.

연구 동기 및 목표

• Dot-product 어텐션의 2차 메모리 비용을 줄이면서 희소 어텐션 출력을 학습한다는 동기를 부여한다.
• 입력 블록을 재배열하는 differentiable sorting 기반 어텐션 메커니즘을 제안하여 로컬 계산으로 준-글로벌 컨텍스트를 가능하게 한다.
• 자 autoregressive 디코딩을 위한 인과적 변형(Causal Sinkhorn Balancing)과 시퀀스를 잘라내는 동적 인코딩(SortCut)을 도입한다.
• Sinkhorn 어텐션과 표준 어텐션을 결합한 Mixture 모델을 탐구하여 성능을 향상시킨다.
• 언어 모델링, 이미지 생성, 문서 분류, 자연어 추론 등에서 효과를 입증한다.

제안 방법

• SortNet이라고 하는 매개변수화된 메타 정렬 네트워크를 도입하여 블록별 순열 행렬을 출력한다.
• 정렬 행렬을 이중 확률(polytope)인 더블리지스토캐스틱(Birkhoff) 다면체에 투사하기 위해 differentiable Sinkhorn 균형을 사용한다.
• 로컬 이웃을 보존하면서 준-글로벌 어텐션을 가능하게 하기 위해 개별 토큰이 아니라 블록을 정렬한다.
• 정렬된 블록 내에서 어텐션을 계산하고 키와 값에 공유 정렬 연산자를 적용한다.
• 이산 순열을 근사하기 위해 Gumbel 노이즈와 온도 파라미터를 도입한다(Gumbel-Sinkhorn).
• 필요시 Sinkhorn 어텐션과 표준 vanilla 어텐션을 혼합하는 Mixture 모델을 적용한다.
• (i) 누적합에 기초한 인과적 정렬 네트워크를 포함한 인과적 설정과 (ii) 향후 마스킹을 가진 인과적 Sinkhorn 균형을 확장한다.
• 정렬 후 예산(N_k)으로 시퀀스를 잘라 추가로 복잡성을 줄이는 SortCut을 제공한다.

실험 결과

연구 질문

• RQ1 Sparse Sinkhorn Attention이 메모리 복잡도를 줄이면서 Vanilla 어텐션에 비해 경쟁력 있는 성능을 달성하는가?
• RQ2 제안된 정렬 기반 희소 어텐션 변형들(Causal Sinkhorn Balancing, SortCut)이 인코딩/디코딩 효율성과 정확도에 어떤 영향을 미치는가?
• RQ3 Sinkhorn과 Vanilla 어텐션의 혼합이 다양한 작업에서 각 접근법을 단독으로 능가할 수 있는가?
• RQ4 실제로 Sparse Sinkhorn Attention의 메모리 복잡도는 Dense 어텐션 및 Sparse Transformers와 비교해 어떻게 되는가?
• RQ5 정렬 하이퍼파라미터(온도, Sinkhorn 반복 횟수)가 성능에 어떤 영향을 미치는가?

주요 결과

• Sparse Sinkhorn Attention은 자기 주의 메모리를 O(ell^2)에서 O(B^2 + N_B^2)로 감소시키고 SortCut을 사용하면 일부 설정에서 선형 시간 O(ell N_k)으로도 감소시킨다.
• 다양한 작업에서(Sort 네스팅, 언어 모델링, 픽셀 단위 생성, 문서 분류, NLI) Sinkhorn Transformers는 Vanilla Transformers와 동등하거나 이를 능가하고, 종종 Sparse Transformers 및 Local Attention을 능가한다.
• LM1B 하위 단어 모델링에서 Sinkhorn Transformers는 여러 블록 크기에서 Local Attention 및 Sparse Transformer를 능가하며, Sinkhorn Mixture가 최상의 성과를 달성한다.
• 이미지 생성 CIFAR-10에서 Sinkhorn Transformer는 테스트된 효율 모델들 가운데 가장 높은 바이트-당 차원(Bpd)을 달성한다.
• SortCut 인코더는 Vanilla Transformer와 비슷한 성능을 달성하면서도 메모리 사용량을 크게 줄일 수 있다.
• 적절한 분석에서 정렬 네트워크가 간단한 선형 형태에서 가장 잘 작동하며 Sinkhorn 정규화(N_k > 0)를 적용하는 것이 성능에 결정적임을 확인하였다.