[論文レビュー] Object Representations as Fixed Points: Training Iterative Refinement Algorithms with Implicit Differentiation

Michael Chang, Thomas L. Griffiths+1 more|arXiv (Cornell University)|Jul 2, 2022

Domain Adaptation and Few-Shot Learning被引用数 5

ひとこと要約

本稿では、スロットアテンションなどの反復的精錬モデルを学習するための、バックプロパゲーションによるアンロールされた反復の代わりに、安定で微分可能な固定点定式化を用いた暗黙の微分を提案する。この手法により、メモリと時間計算量が一定のまま、最適化の安定性が向上し、検証損失が低減し、生成品質が向上する。コード変更はたった1行のみを要する。

ABSTRACT

Iterative refinement -- start with a random guess, then iteratively improve the guess -- is a useful paradigm for representation learning because it offers a way to break symmetries among equally plausible explanations for the data. This property enables the application of such methods to infer representations of sets of entities, such as objects in physical scenes, structurally resembling clustering algorithms in latent space. However, most prior works differentiate through the unrolled refinement process, which can make optimization challenging. We observe that such methods can be made differentiable by means of the implicit function theorem, and develop an implicit differentiation approach that improves the stability and tractability of training by decoupling the forward and backward passes. This connection enables us to apply advances in optimizing implicit layers to not only improve the optimization of the slot attention module in SLATE, a state-of-the-art method for learning entity representations, but do so with constant space and time complexity in backpropagation and only one additional line of code.

研究の動機と目的

• 反復的精錬モデルにおける潜在セット表現学習のためのアンロールされた反復を経由するバックプロパゲーションの不安定さと高い計算コストを解消すること。
• スロットアテンションのような反復的精錬アルゴリズムの固定点性質を活用し、前向きと後向きのパスを分離することで、暗黙の微分を可能にすること。
• アーキテクチャの変更なしに、SLATE やバニラスロットアテンションなどの最先端モデルの学習効率と安定性を向上させること。
• 暗黙の微分により、勾配クリッピングや初期学習率ウォームアップ、反復回数のハイパーパramータチューニングを必要とせずに、安定した学習が可能であることを示すこと。
• この手法がアーキテクチャに一般化可能であり、セグメンテーション品質を保持するとともに、下流のオブジェクト属性予測を向上させることを示すこと。

提案手法

• スロットアテンションモジュールを固定点反復として定式化し、最終的表現が z = f_x(z) の解として得られることを示し、固定点における暗黙の微分を可能にする。
• 陰関数定理を適用して、反復プロセスをアンロールせずに勾配を計算し、効率的な後向きパスのために一次のネウマン近似を用いる。
• 暗黙の勾配を損失計算に統合し、標準的なバックプロパゲーションと比較してたった1行のコード追加で、エンドツーエンド学習を可能にする。
• 前向きパス（アンロールされた反復）と後向きパス（暗黙の勾配計算）を分離することで、1回の後向きパスあたりのメモリと時間計算量を定数に保つ。
• 精錬関数 f のヤコビアンを用いて、ネウマン級数近似により暗黙の勾配を計算し、アンロールバックプロパゲーションによる勾配爆発を回避する。
• この手法を SLATE および元のスロットアテンションアーキテクチャの両方へ適用し、異なるデコーダーと学習設定において一般化性を検証する。

実験結果

リサーチクエスチョン

• RQ1 アンロールバックプロパゲーションによる勾配爆発を引き起こすスロットアテンションのような反復的精錬モデルの学習を、暗黙の微分が安定化できるか。
• RQ2 勾配クリッピングや初期学習率ウォームアップを必要とせずに、暗黙の微分がオブジェクト中心表現学習の最適化を改善できるか。
• RQ3 標準的バックプロパゲーションと比較して、暗黙の微分を用いた場合、精錬反復回数が学習の安定性と性能に与える影響はいかほどか。
• RQ4 暗黙の微分が、異なるネットワークアーキテクチャ間で生成マスクやセグメンテーション出力の品質を保持できるか。
• RQ5 標準的スロットアテンションと比較して、暗黙の微分が下流タスク（例：オブジェクト属性予測）においてどの程度性能を向上させるか。

主な発見

• 暗黙のスロットアテンションは、3つのベンチマークデータセットにおいて、バニラスロットアテンションや SLATE よりも顕著に低い検証損失を達成する。
• FID は最大30%低減され、画像再構成における平均二乗誤差（MSE）は最大50%低減される。
• 勾配クリッピングや初期学習率ウォームアップ、反復回数のチューニングを必要とせず、標準的学習とは異なり、この手法はそれらを不要にする。
• バックプロパゲーションの後向きパスにおいて、反復回数に関係なく、定数の空間的・時間的計算量を維持する。
• CLEVR データセットでは、7回の反復を用いた暗黙のスロットアテンションが、1回の反復を用いたバニラスロットアテンションよりも再構成MSEが2倍低い。
• 暗黙のスロットアテンションは、Locatello et al. [39] が提案した元の空間ブロードキャストデコーダを含む、さまざまなアーキテクチャにおいても高品質なセグメンテーションマスクを維持する。図8に示す。