Perceiver IO: あらゆるモダリティに対応するスケーラブルで完全アテンション型のモデル

Perceiver IOは、テキスト、画像、音声、ビデオ、ポイントクラウドを含むあらゆるモダリティにおいて、高次元データに通常伴う二次的な計算量およびメモリコストをかけることなく、Transformerの自己アテンション（self-attention）メカニズムを利用可能にするスケーラブルなニューラルネットワークアーキテクチャです。計算の大部分を、固定サイズの小さな潜在空間（latent space）で行うことにより、Perceiver IOは入力サイズと自己アテンション層の複雑さの間の依存関係を排除します。

Transformerのスケーリング制限の克服

標準的なTransformerアーキテクチャは、自己アテンションメカニズムがすべての層においてすべての入力に対してペアワイズのドット積を必要とするため、計算量とメモリにおいてスケーリングが困難です。高次元データを扱うために、既存のモデルは通常、データをトークンに離散化するためのモダリティ固有の前処理（例：ViTは画像パッチを使用し、Wav2Vec2は音声用の特徴量エンコーダを使用する）に依存しています。

Perceiver IOは、入力そのものではなく、潜在変数のセットに対して自己アテンションを行うことでこれを解決します。入力は、これらの潜在変数とのクロスアテンション（cross-attention）にのみ使用されます。これにより、Transformerエンコーダの複雑さは入力サイズに対して線形に依存し、潜在アテンションは入力サイズに依存しないことが保証されます。Perceiver IOはさらに、出力に対しても同様のクロスアテンションメカニズムを使用することで、モデルが任意の出力形状を扱うことを可能にします。

コアアーキテクチャとワークフロー

すべてのPerceiverのバリエーションは、PerceiverModelクラスに基づいて構築されており、異なるデータ型を扱うために3つのオプションのコンポーネントを利用します：

Preprocessor: 生の入力を（テキスト、画像など）ベクトルに埋め込みます。
Decoder: 潜在変数の最終的な隠れ状態を、分類ロジットやオプティカルフローのような有用な形式にデコードします。
Postprocessor: デコーダの出力を特定の機能に変換します。主に自己符号化（auto-encoding）タスクで使用されます。

処理パイプライン

Input Stage: 入力はオプションでpreprocessorによって処理されます。
Cross-Attention: 潜在変数はクエリ（Q）を生成し、前処理された入力はキー（K）とバリュー（V）を生成します。これにより、高次元の入力を潜在空間にマッピングします。
Latent Processing: 繰り返し可能な自己アテンション層のブロックが、潜在的な埋め込みを更新します。出力は、潜在変数の「最後の隠れ状態」のテンソルです。
Output Stage: オプションのdecoderが別のクロスアテンション操作を実行します。ここでは、学習可能な埋め込みがクエリ（Q）を生成し、潜在変数がキー（K）とバリュー（V）を生成します。その後、postprocessorがこれらの出力を洗練させることがあります。

モダリティ固有の実装

テキスト分類

Perceiver IOは、生のUTF-8バイトを直接処理できるため、明示的なトークン化（WordPieceやBPEなど）の必要性がありません。テキスト分類の場合、PerceiverTextPreprocessorはバイトIDを埋め込み、絶対位置埋め込みを追加します。その後、PerceiverClassificationDecoderがクロスアテンションを使用して、潜在変数を分類ロジットにマッピングします。マスク付き言語モデリング（PerceiverForMaskedLM）では、PerceiverBasicDecoderが欠落しているバイトを予測するために、潜在変数を入力シーケンスの長さにマッピングし直します。

画像分類

ビジョンタスクの場合、モデルはPerceiverImagePreprocessorを使用します。著者らは3つの前処理方法をテストしました：

1x1畳み込み層と学習可能な1D位置埋め込みを使用したピクセルのフラット化。
固定の2D Fourier位置埋め込みを使用したピクセルのフラット化。
2D畳み込み + maxpool層と固定の2D Fourier位置埋め込みを使用した方法。

JFTで事前学習した後、PerceiverForImageClassificationConvProcessingバリエーションは、ImageNetにおいてトップ1精度が84.5%に達しました。注目すべきは、1D学習可能な位置埋め込みのみを使用したバージョンが72.7%の精度を達成したことであり、これは、このモデルが明示的な2D構造的知識なしでも動作できる能力を示しています。

オプティカルフロー推定

Perceiver IOは、2つの画像フレーム間の2Dピクセル変位を推定できます。PerceiverForOpticalFlowモデルは、各ピクセルの周囲に3x3パッチを抽出し、フレームをチャネル次元で結合し、空間次元をフラット化します。その後、PerceiverOpticalFlowDecoderが前処理された入力をクエリとして使用して、潜在変数を予測されたフローマップにデコードします。このアプローチは、AutoFlowデータセットで学習した際、SintelおよびKITTIベンチマークにおいて最先端の結果を達成しました。

マルチモーダル自己符号化

PerceiverForMultimodalAutoencodingモデルを使用すると、アーキテクチャは複数のモダリティにわたる結合分布を学習できます。例えば、Kinetics-700データセットにおいて、ビデオフレーム、音声サンプル、およびクラスラベルを同時に処理します：

Images: