[논문 리뷰] Sinkformers: Transformers with Doubly Stochastic Attention

Michael E. Sander, Pierre Ablin외 2명|arXiv (Cornell University)|2021. 10. 22.

Advanced Neural Network Applications참고 문헌 59인용 수 14

한 줄 요약

이 논문은 자기주의성에서 표준 소프트맥스 정규화를 대체하여 이중 확률적 주의 행렬을 생성하는 싱크호른 알고리즘을 사용하는 트랜스포머의 변종인 싱크포머를 소개한다. 이 방법은 비전 및 자연어 처리 작업 전반에서 모델 성능을 햖थ하며, 3차원 형태 분류에서 뚜렷한 정확도 향상을 보이며, 무한 표본 근사에서 주의 메커니즘을 워셔스타인 경사 하강 흐름과 열 확산 과정으로 이론적으로 해석함으로써 이론적 통찰을 제공한다.

ABSTRACT

Attention based models such as Transformers involve pairwise interactions between data points, modeled with a learnable attention matrix. Importantly, this attention matrix is normalized with the SoftMax operator, which makes it row-wise stochastic. In this paper, we propose instead to use Sinkhorn's algorithm to make attention matrices doubly stochastic. We call the resulting model a Sinkformer. We show that the row-wise stochastic attention matrices in classical Transformers get close to doubly stochastic matrices as the number of epochs increases, justifying the use of Sinkhorn normalization as an informative prior. On the theoretical side, we show that, unlike the SoftMax operation, this normalization makes it possible to understand the iterations of self-attention modules as a discretized gradient-flow for the Wasserstein metric. We also show in the infinite number of samples limit that, when rescaling both attention matrices and depth, Sinkformers operate a heat diffusion. On the experimental side, we show that Sinkformers enhance model accuracy in vision and natural language processing tasks. In particular, on 3D shapes classification, Sinkformers lead to a significant improvement.

연구 동기 및 목표

• 싱크호른 알고리즘을 통해 행 방향 확률적 소프트맥스를 이중 확률적 주의로 대체함으로써 트랜스포머 성능을 향상시키는 것.
• 싱크포머를 워셔스타인 경사 하강 흐름과 열 확산 과정으로 해석함으로써 자기주의성 메커니즘에 대한 이론적 기반을 제공하는 것.
• 훈련된 트랜스포머에서 주의 행렬이 이중 확률적 행렬로 수렴하는지 경험적으로 검증하여, 싱크호른 정규화를 의미 있는 인덕티브 바이어스로 사용하는 데 정당성을 부여하는 것.
• 싱크포머가 3차원 형태 분류, 이미지 분류, 시퀀스 모델링을 포함한 다양한 작업에서 표준 트랜스포머보다 뛰어난 정확도를 달성함을 보여주는 것.
• 기존 트랜스포머 아키텍처에 소프트맥스를 싱크호른 정규화로 쉽게 통합할 수 있도록 하되, 계산 오버헤드를 최소화하는 것.

제안 방법

• 자신의주의에서 소프트맥스 연산을 싱크호른 알고리즘으로 대체하여, 비정규화된 주의 행렬 $ K^0 = \exp(C) $ 를 이중 확률적 행렬로 정규화하는 것.
• 싱크호른 알고리즘을 통해 반복적인 행 및 열 정규화를 수행하여 행과 열의 합이 모두 1이 되는 행렬로 수렴하는 것.
• 싱크호른 반복 횟수를 조절할 수 있는 하이퍼파ram터를 도입하여, 표준 트랜스포머(1회 반복)에서 완전히 수렴한 싱크포머 사이의 보간을 가능하게 하는 것.
• 딥 러닝 프레임워크 내에서 미분 가능하게 구현함으로써, 이중성과 GPU 호환성을 유지하는 것.
• 이중 확률적 성질을 사전 지식으로 활용하여 모든 토큰 간에 더 균형 잡히고 민주적인 주의를 유도하는 것.
• 이론적 분석 결과, 싱크포머는 무한소 근사에서 이산 워셔스타인 경사 하강 흐름과 대응하며, 평균장 근사에서 열 방정식으로 수렴함을 보여줌.

실험 결과

연구 질문

• RQ1 훈련된 트랜스포머에서 최적화 과정 중 주의 행렬이 자연스럽게 이중 확률적 행렬로 수렴하는가?
• RQ2 소프트맥스를 싱크호른 정규화로 대체하면 비전 및 자연어 처리 벤치마크에서 모델 정확도가 향상되는가?
• RQ3 이중 확률적 주의 메커니즘이 연속 근사에서 최적 운반 이론과 경사 하강 흐름과 어떻게 관련되는가?
• RQ4 무한 깊이 및 무한 표본 근사에서 싱크포머의 이론적 행동은 어떠한가?
• RQ5 표준 트랜스포머에 비해 실용적으로 효율적으로 구현 가능하며 계산 오버헤드가 최소한인가?

주요 결과

• 3D 형태 분류 벤치마크인 ModelNet40에서 싱크포머는 표준 트랜스포머보다 뚜렷한 정확도 향상을 보이며 기하학적 학습 작업에서 뛰어난 성능 향상을 입증한다.
• IWSLT’14 독일어-영어 번역 작업에서, 싱크포머는 35 에포크 후 중앙값 BLEU 점수 34.73을 기록하여 동일한 훈련 프rotocol에서 표준 트랜스포머의 34.68을 略로 뛰어넘는다.
• IMDb 감성 분류 작업에서 싱크포머는 표준 트랜스포머보다 성능 향상을 보이며, 시퀀스 모델링에서 더 나은 일반화 능력을 보여준다.
• 비전 트랜스포머를 사용한 고양이와 강아지 이미지 분류 작업에서, 싱크포머는 중앙값 테스트 정확도 79.5%를 기록했고, 표준 트랜스포머의 79.0%보다 높았다.
• MNIST에 대해 1층, 1헤드 자기주의성 모델에서, 패치 크기가 감소할수록 싱크포머가 트랜스포머를 능가하며, 저용량 환경에서 더 뛰어난 표현 능력을 보여준다.
• 싱크포머의 에포크당 훈련 시간은 표준 트랜스포머와 거의 동일했으며(3분 25초 대비 3분 20초) 계산 오버헤드가 최소임을 확인한다.