[論文レビュー] Scalable Adaptive Computation for Iterative Generation

Allan Jabri, David J. Fleet+1 more|arXiv (Cornell University)|Dec 22, 2022

Stochastic Gradient Optimization Techniques被引用数 13

ひとこと要約

この論文は、潜在トークンを用いてグローバル自己自己注意を集中させることで、入力データの次元から計算を分離するドメインに依存しないアーキテクチャ、再帰的インターフェースネットワーク（RINs）を導入する。これにより、高次元データの適応的かつスケーラブルな生成が可能になる。RINsは、カスケードやガイドラインを用いずに1024×1024にスケーリング可能な画像・動画拡散モデルにおいて最先端の性能を達成しており、2D/3D U-Netと比較して最大10倍のFLOPs削減を実現している。これは、再帰を活用した潜在自己条件付き処理により、より深いコンテキスト豊富なルーティングが可能になることによるものである。

ABSTRACT

Natural data is redundant yet predominant architectures tile computation uniformly across their input and output space. We propose the Recurrent Interface Networks (RINs), an attention-based architecture that decouples its core computation from the dimensionality of the data, enabling adaptive computation for more scalable generation of high-dimensional data. RINs focus the bulk of computation (i.e. global self-attention) on a set of latent tokens, using cross-attention to read and write (i.e. route) information between latent and data tokens. Stacking RIN blocks allows bottom-up (data to latent) and top-down (latent to data) feedback, leading to deeper and more expressive routing. While this routing introduces challenges, this is less problematic in recurrent computation settings where the task (and routing problem) changes gradually, such as iterative generation with diffusion models. We show how to leverage recurrence by conditioning the latent tokens at each forward pass of the reverse diffusion process with those from prior computation, i.e. latent self-conditioning. RINs yield state-of-the-art pixel diffusion models for image and video generation, scaling to 1024X1024 images without cascades or guidance, while being domain-agnostic and up to 10X more efficient than 2D and 3D U-Nets.

研究の動機と目的

• 画像や動画などの高次元的で冗長な自然データにおけるディープラーニングモデルの均一な計算の非効率性に対処すること。
• 入力サイズに比例する二次的複雑性を回避しつつ、データの内容に応じて計算を動的に割り当てるスケーラブルなアーキテクチャの開発。
• 再帰を活用してイテレーティブな生成タスクにおける深さのある表現的ルーティングを実現し、時間軸に沿ったバックプロパゲーションなしでコンテキスト伝搬を緩和すること。
• ピクセルレベルの拡散モデルにおけるカスケードやガイドライン依存を排除し、より効率的でドメインに依存しないアーキテクチャの設計。

提案手法

• 隠れユニットを、入力サイズに比例して線形に増加するインターフェーストークンと、固定サイズのコンactな潜在トークンに分解し、計算と入力レイアウトを分離する。
• インターフェーストークンから潜在トークンへの情報ルーティングとその逆を、クロスアテンションを用いて実現し、双方向のボトムアップおよびトップダウンフィードバックを可能にする。
• グローバル自己自己注意を小さな潜在トークンの集合に限定することで、入力サイズに比例する二次的複雑性を回避する。
• 潜在自己条件付き処理の導入：直前の推論ステップでの潜在活性化を再利用し、その後続のステップでコンテキストとして利用することで、時間軸に沿ったバックプロパゲーションなしで再帰を模倣する。
• RINブロックをスタックすることで、イテレーション間でコンテキストを蓄積するより深いネットワークを構築し、ルーティングの表現力を向上させる。
• 固定された計算グラフを持つ標準的な微分可能なネットワークとしてモデルを訓練・デプロイすることで、現代のハードウェア上で効率的な推論を可能にする。

実験結果

リサーチクエスチョン

• RQ1 入力サイズに比例する二次的複雑性を伴わず、高次元データに対して計算を適応的に割り当てるニューラルアーキテクチャを設計することは可能か？
• RQ2 イテレーティブな生成タスクにおいて、時間軸に沿ったバックプロパゲーションなしで再帰を活用してルーティングの深さを向上させることは可能か？
• RQ3 RINsのようなドメインに依存しないアーキテクチャは、高解像度画像・動画生成において、U-Netベースの拡散モデルをサンプル品質とFLOP効率の両面で上回ることができるか？
• RQ4 潜在自己条件付き処理は、イテレーティブな拡散設定において、どの程度ルーティングの深さとモデルの表現力を向上させるか？
• RQ5 より効率的なアーキテクチャを用いることで、カスケードやガイドラインを用いずにピクセルレベルの拡散モデルを1024×1024解像度にスケーリングすることは可能か？

主な発見

• RINsは、ガイドラインやカスケードを一切使用せず、64×64から1024×1024のあらゆる解像度でImageNetにおいて最先端のFIDスコアを達成した。
• 1024×1024の画像生成において、RINsは2Dおよび3D U-Netベースラインと比較して、1ステップあたりのFLOPsを最大10倍削減した。
• Kinetics600動画生成ベンチマークにおいて、RINsは先端的なアプローチを上回りながら、1ステップあたりのFLOPsを10倍削減した。
• 潜在自己条件付き処理により、イテレーティブなノイズ除去ステップ間でルーティングコンテキストを保持できるようになり、モデルの表現力が向上した。
• RINsは複数のモodalに対し強力な性能を維持しており、ドメインに依存せず、アテンションとMLPのみに依存しているため、広範な応用が可能である。
• 最小限のアーキテクチャ的変更で優れた性能を達成しており、適応的計算が標準的な拡散学習パイプラインに効率的に統合可能であることを示した。