[논문 리뷰] Axial Attention in Multidimensional Transformers

Jonathan Ho, Nal Kalchbrenner외 2명|arXiv (Cornell University)|2019. 12. 20.

Generative Adversarial Networks and Image Synthesis참고 문헌 18인용 수 363

한 줄 요약

이 논문은 Axial Transformers를 도입한다. 이는 축 주의(axial attention)를 사용하여 단일 텐서 축을 따라 맥락을 계산하는 자기회귀(self-attention) 모델로, 커스텀 커널 없이 ImageNet-32/64 및 BAIR 로봇 푸시에서 최첨단 성능을 달성한다.

ABSTRACT

We propose Axial Transformers, a self-attention-based autoregressive model for images and other data organized as high dimensional tensors. Existing autoregressive models either suffer from excessively large computational resource requirements for high dimensional data, or make compromises in terms of distribution expressiveness or ease of implementation in order to decrease resource requirements. Our architecture, by contrast, maintains both full expressiveness over joint distributions over data and ease of implementation with standard deep learning frameworks, while requiring reasonable memory and computation and achieving state-of-the-art results on standard generative modeling benchmarks. Our models are based on axial attention, a simple generalization of self-attention that naturally aligns with the multiple dimensions of the tensors in both the encoding and the decoding settings. Notably the proposed structure of the layers allows for the vast majority of the context to be computed in parallel during decoding without introducing any independence assumptions. This semi-parallel structure goes a long way to making decoding from even a very large Axial Transformer broadly applicable. We demonstrate state-of-the-art results for the Axial Transformer on the ImageNet-32 and ImageNet-64 image benchmarks as well as on the BAIR Robotic Pushing video benchmark. We open source the implementation of Axial Transformers.

연구 동기 및 목표

• 계산비용이나 메모리 소모가 지나치게 커지지 않으면서 고차원 데이터 텐서에 대해 자체 주의(self-attention) 기반의 자기회귀 모델을 개발한다.
• 데이터를 평탄화하지 않고 텐서 축을 따라 어텐션으로 확장하기 위해 축 주의(axial attention)를 도입한다.
• 반병렬 샘플링 절차를 통해 전체 컨텍스트 모델링을 가능하게 한다.
• 이미지 및 비디오 벤치마크에서 최첨단 결과를 입증한다.
• 용이한 채택을 위한 오픈 소스 구현을 제공한다.]
• method:[

제안 방법

• 다차원 텐서의 단일 축을 따라 어텐션으로 정의하고 다른 축은 보존하여 계산을 O(N^2)에서 O(N^{(d-1)/d})로 감소시킨다.
• 마스크 처리된 축 주의 블록과 마스크 처리되지 않은 축 주의 블록을 쌓아 독립 가정 없이 전체 자기회귀 컨텍스트를 구축한다.
• 효율적인 샘플링을 위한 행 단위 내부 디코더와 이전 행 및 채널의 정보를 통합하기 위한 외부 디코더를 사용한다.
• 추가적인 비마스크 행/열 주의 층을 사용하여 이전 채널에 조건부로 다중 채널 데이터를 모델링한다.
• 전체 데이터 텐서의 로그가능도(log-likelihood)의 편향 없는 추정치를 얻기 위해 임의의 채널 슬라이스에서 학습한다.
• Axial Transformer의 오픈 소스 구현을 제공한다.]
• research_questions:[

실험 결과

연구 질문

• RQ1 전체 결합 표현력을 보존하면서 계산 비용을 줄이기 위해 어텐션 메커니즘을 다차원 텐서에 일반화하는 방법은?
• RQ2 커스텀 커널이나 대량의 데이터 복사 없이 이미지와 비디오에 대한 효율적인 자기회귀 모델링을 가능하게 할 수 있는가?
• RQ3 마스크된/마스크되지 않은 축 주의를 결합하는 것이 모델링 용량과 샘플링 속도에 어떤 영향을 미치는가?
• RQ4 Axial Transformer가 기존의 자기회귀 모델과 비교해 표준 이미지 및 비디오 벤치마크에서 어떤 성능을 보이는가?
• RQ5 이전 채널/프레임에 조건부로 처리하여 다중 채널 데이터와 비디오를 효과적으로 처리할 수 있는가?

주요 결과

모델	ImageNet-32 (bits/dim)	ImageNet-64 (bits/dim)
Multiscale PixelCNN	3.95	3.70
PixelCNN/RNN	3.86	3.63
Gated PixelCNN	3.83	3.57
PixelSNAIL	3.80	3.52
SPN	3.79	3.52
Image Transformer	3.77
Strided Sparse Transformer		3.44
Axial Transformer + LSTM inner decoder	3.77	3.46
Axial Transformer	3.76 (3.758)	3.44 (3.439)

• 축 주의는 d차원 입력 텐서에 대해 표준 자기 주의에서 계산 및 메모리를 O(N^{(d-1)/d})의 비율로 감소시킨다.
• Axial Transformer는 여러 베이스라인과 비교하여 ImageNet-32 및 ImageNet-64에서 차원당 비트수(bits-per-dimension)에서 최첨단을 달성한다.
• 이 모델은 이전 자기회귀 방법들에 비해 BAIR Robotic Pushing 비디오 모델링을 크게 개선한다.
• 반병렬 샘플링은 대부분의 컨텍스트를 병렬로 계산하여 대형 텐서에 대해 실용적인 디코딩을 가능하게 한다.
• 발실험(Ablation) 결과는 내부 디코더를 LSTM으로 교체하면 학습이 느려지지만 일부 성능은 맞출 수 있으며, 전체 4층 내부 디코더는 성능과 학습 속도를 모두 향상시킨다.
• 채널 조건부 확장은 무거운 아키텍처 변경 없이 다중 채널 이미지와 비디오를 효과적으로 모델링한다.