[論文レビュー] FIT: Far-reaching Interleaved Transformers

ting chen, Lala Li|arXiv (Cornell University)|May 22, 2023

Generative Adversarial Networks and Image Synthesis被引用数 5

ひとこと要約

FITは、グループ化されたデータトークンと適応的ラテントークンを用いて、局所的およびグローバル自己注意メカニズムをインタリーブする新しいトランスフォーマー・アーキテクチャを導入し、長時間系列における効率的で2次関数的な注意を実現する。局所的なウィンドウ付き注意と学習されたラテントークンの小さな集合に対するグローバル注意を組み合わせることで、FITは計算複雑性をO(L⁴/³)に低減し、16GBのメモリ内で160Kトークンの入力（例：6400×6400の画像）をエンドツーエンドで学習可能であり、モデル並列化や最適化を必要とせずに、高解像度画像タスクで最先端の性能を示している。

ABSTRACT

We present FIT: a transformer-based architecture with efficient self-attention and adaptive computation. Unlike original transformers, which operate on a single sequence of data tokens, we divide the data tokens into groups, with each group being a shorter sequence of tokens. We employ two types of transformer layers: local layers operate on data tokens within each group, while global layers operate on a smaller set of introduced latent tokens. These layers, comprising the same set of self-attention and feed-forward layers as standard transformers, are interleaved, and cross-attention is used to facilitate information exchange between data and latent tokens within the same group. The attention complexity is $O(n^2)$ locally within each group of size $n$, but can reach $O(L^{{4}/{3}})$ globally for sequence length of $L$. The efficiency can be further enhanced by relying more on global layers that perform adaptive computation using a smaller set of latent tokens. FIT is a versatile architecture and can function as an encoder, diffusion decoder, or autoregressive decoder. We provide initial evidence demonstrating its effectiveness in high-resolution image understanding and generation tasks. Notably, FIT exhibits potential in performing end-to-end training on gigabit-scale data, such as 6400$ imes$6400 images, or 160K tokens (after patch tokenization), within a memory capacity of 16GB, without requiring specific optimizations or model parallelism.

研究の動機と目的

• 標準のトランスフォーマーが長時間系列を処理する際の2次関数的計算複雑性に対処すること。
• モデル並列化を必要とせず、性能を損なわず、長時間系列における効率的かつスケーラブルな注意を実現すること。
• 多様なビジョンタスクに適応する統合アーキテクチャとして、エンコーダー、拡散デコーダー、自己回帰デコーダーとして機能するように設計すること。
• 適応的ラテントークンを用いたインタリーブされた局所的およびグローバルなトランスフォーマー層の有効性を検証し、長距離依存関係のモデリングを改善すること。

提案手法

• 入力データをトークンのグループに分割し、各グループを局所的自己注意（ウィンドウ付きコンテキスト内）で処理する。
• 各グループごとに、グローバルコンテキストを表す少数の学習可能なラテントークンを導入し、計算複雑性を低減したグローバル自己注意を可能にする。
• 局所的およびグローバルなトランスフォーマー層をインタリーブし、クロスアテンションによりデータトークンとラテントークン間の双方向情報伝達を実現する。
• 1回のフォワードパス内でデータトークンとラテントークンの反復的更新を実施し、局所的およびグローバル表現の深いつながりを確保する。
• 局所的およびグローバルな層の両方で標準的なフィードフォワードネットワークと注意メカニズムを維持し、入力範囲と注意パターンを除いては同一である。
• 各グループ内でクロスアテンションを活用し、データトークンからラテントークン、およびその逆方向への情報ルーティングを実現し、適応的計算を可能にする。

実験結果

リサーチクエスチョン

• RQ1 ラテントークンを用いたインタリーブされた局所的およびグローバルなトランスフォーマー層は、標準的またはスパースな注意メカニズムに比べ、より優れた長距離モデリングを達成できるか？
• RQ2 1グループあたりのラテントークン数は、高解像度画像タスクにおけるモデル性能と効率にどのように影響するか？
• RQ3 局所的およびグローバル層をインタリーブすることで、順次スタックする場合や局所的層のみを使用する場合に比べ、性能が向上するか？
• RQ4 FITは、モデル並列化や最適化を一切行わず、16GBのメモリ内でギガビットスケールの入力（例：6400×6400の画像）をエンドツーエンドで学習可能か？
• RQ5 FITは、ビジョン分野における自己回帰的、拡散的、エンコーダー的タスクの多様な文脈で、どれほど汎用的なアーキテクチャとして機能できるか？

主な発見

• FITは153Mパラメータで、ImageNet-64×64の自己回帰的モデリングにおいて、3.42ビット/次元という、より大きなモデルと同等の、ほぼ最先端の性能を達成した。
• ラテントークンの数を増やすことで性能が向上（nllおよびbpdが低下）するが、パラメータ数や学習速度にほとんど影響を与えないため、効率的なスケーリングが可能である。
• 局所的およびグローバル層のインタリーブは、順次スタックする場合や局所的層のみを使用する場合に比べ、優れた性能（例：ImageNetでは3.49 bpd）を示し、学習コストの増加は無視できるほど小さい。
• FITは、モデル並列化や特別な最適化を一切行わず、16GBのメモリ内で160Kトークンのシーケンス（例：6400×6400の画像）をエンドツーエンドで学習可能である。
• このアーキテクチャは、タスク間で優れた一般化性能を示し、高解像度ビジョンベンチマークにおいて、エンコーダー、拡散デコーダー、自己回帰デコーダーとして効果的に機能している。