[論文レビュー] Axial Attention in Multidimensional Transformers

Jonathan Ho, Nal Kalchbrenner+2 more|arXiv (Cornell University)|Dec 20, 2019

Generative Adversarial Networks and Image Synthesis参考文献 18被引用数 363

ひとこと要約

本論文は Axial Transformers を紹介します。高次元データに対する自己回帰型自己注意モデルで、軸方向注意機構を用いて単一のテンソル軸に沿った文脈を計算します。カスタムカーネルを使わずに 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.

研究の動機と目的

• 高次元データテンソルに対して prohibitively compute or memory costs を避けた自己注意ベースの自己回帰モデルを開発する。
• データを平坦化することなく、テンソルの軸に沿って注意をスケールさせる軸方向注意を導入する。
• 半並列サンプリング手法を用いて全文脈モデリングを有効にする。
• 画像および動画ベンチマークで最先端の結果を示す。
• 容易に導入できるオープンソース実装を提供する。

提案手法

• 多次元テンソルの単一の軸に沿った注意として軸方向注意を定義し、他の軸を保持することで、計算を O(N^2) から O(N^{(d-1)/d}) に縮小する。
• 独立性の仮定を置かずに、マスク付き/マスクなしの軸方向注意ブロックを積み重ねて、完全な自己回帰の文脈を構築する。
• 効率的なサンプリングのために行単位の内部デコーダを使用し、前の行とチャネルから情報を組み込む外部デコーダを用いる。
• 追加の未マスク行/列の注意層を用いて、前のチャネルを条件付けることでマルチチャネルデータをモデル化する。
• 全データテンソルの対数尤度の不偏推定を得るために、ランダムなチャネルスライスで訓練する。
• Axial Transformers のオープンソース実装を提供する。

実験結果

リサーチクエスチョン

• RQ1 How can attention mechanisms be generalized to multidimensional tensors to reduce computational demands while preserving full joint expressiveness?
• RQ2 Can axial attention enable efficient autoregressive modeling for images and videos without custom kernels or heavy data copying?
• RQ3 What are the impacts of combining masked/unmasked axial attention on modeling capacity and sampling speed?
• RQ4 How does Axial Transformer perform on standard image and video benchmarks compared to prior autoregressive models?
• RQ5 Can the model effectively handle multi-channel data and video by conditioning on previous channels/frames?

主な発見

Model	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)

• Axial attention reduces computation and memory from standard self-attention by a factor of O(N^{(d-1)/d}) for a d-dimensional input tensor.
• The Axial Transformer achieves state-of-the-art bits-per-dimension on ImageNet-32 and ImageNet-64 compared to several baselines.
• The model significantly improves BAIR Robotic Pushing video modeling over prior autoregressive methods.
• Semi-parallel sampling computes most context in parallel, yielding practical decoding for large tensors.
• An ablation shows that replacing the inner decoder with an LSTM slows training but can match some performance, while the full four-layer inner decoder improves both performance and training speed.
• The channel-conditional extension effectively models multi-channel images and videos without heavy architectural changes.