論文「Circuit Mechanisms for Spatial Relation Generation in Diffusion Transformers」の技術的サマリー

本論文は、テキストから画像を生成する拡散トランスフォーマー（DiT）において、複数の物体間の「空間的関係（例：A が B の上にある）」を正しく生成するメカニズムを、**機械的解釈性（Mechanistic Interpretability）**の観点から解明した研究です。

1. 研究の背景と問題設定

現在のテキストから画像生成（T2I）モデルは、単一物体の属性（色や形状）を正確に生成する能力は大幅に向上しましたが、複数の物体を配置し、それらの空間的関係を指定された通りに描画することには依然として課題を抱えています。
既存の研究では、この失敗の原因を「クロスアテンションの制御不足」や「テキストエンコーダーの限界」のどちらかに帰結させる傾向がありましたが、モデル内部の具体的な回路（Circuit）がどのように動作し、なぜ異なるアーキテクチャで異なる振る舞いを示すのかは不明瞭でした。

本研究では、この問題を制御された環境で分析するため、2 つの物体と 8 種類の空間関係（上下左右など）を指定する合成データセットを構築し、ゼロから DiT モデルを学習させることで、内部の計算メカニズムを解明しました。

2. 手法と実験設定

データセットとタスク

• タスク: 灰色の背景に、色と形状が異なる 2 つの物体を、テキストプロンプトで指定された空間関係（例：「赤い四角形が青い円の左下に位置する」）で配置する画像を生成する。
• モデル: PixArt アーキテクチャに基づいた DiT（Diffusion Transformer）を使用。
• 変数: モデルサイズ（Nano, Micro, Mini, B）と、テキストエンコーダーの種類を比較対象としました。
  1. T5 エンコーダー: 事前学習済みで、文脈を考慮した意味的なトークン埋め込みを生成。
  2. ランダムトークン埋め込み（RTE）: 意味を持たないランダムなベクトルに、位置符号（Positional Encoding）を付与したもの。
  3. 位置符号なし RTE: 意味も位置情報も持たない純粋なランダム埋め込み。

解析手法

• Attention Synopsis（アテンション要約）: 数百万枚のサンプリングステップにおけるクロスアテンションマップを、トークンカテゴリ（物体、関係語など）間で集約・平均化し、安定した情報フローを可視化・定量化する手法を開発。
• 重み空間スクリーニング: 画像を生成せずに、モデルの重み（Q/K 行列）と位置符号の関係から、空間関係情報を処理するアテンションヘッドを特定する手法。
• 因果的介入（Ablation & Manipulation）: 特定のヘッドを無効化したり、埋め込みベクトルに操作を加えたりすることで、発見された回路が生成結果に因果的に寄与していることを検証。

3. 主要な発見と結果

3.1. テキストエンコーダーによる回路の劇的な違い

同じタスクで高い精度を達成するモデルであっても、使用されるテキストエンコーダーによって、空間関係の生成メカニズムは全く異なることが判明しました。

A. ランダム埋め込み（RTE）を用いた場合のメカニズム

RTE-DiT は、2 段階のモジュール化された回路を使用します。

1. 空間関係ヘッド（Spatial Relation Head）:
   • 特定のクロスアテンションヘッド（例：Layer 2, Head 8）が、関係語（"above", "left" など）のトークンを読み取り、画像トークンの位置符号と相互作用します。
   • この相互作用により、関係に応じた**空間勾配（Spatial Gradient）**が生成されます（例："above" なら上部に高いエネルギーを持つ勾配）。
   • この勾配は、物体が配置されるべき「位置タグ」として機能します。
2. 物体生成ヘッド（Object Generation Head）:
   • 別のヘッド（例：Layer 4, Head 3）が、上記の「位置タグ」を読み取り、対応する形状トークン（"square", "circle" など）をその位置に描画します。
   • 特徴: 関係情報と物体属性の処理が明確に分離（Disentangled）されており、ロバスト性が高いです。

B. T5 エンコーダーを用いた場合のメカニズム

T5-DiT は、単一のトークンからの情報融合という異なるメカニズムを使用します。

1. 文脈的埋め込みの活用:
   • T5 の自己アテンションにより、関係語の情報は物体トークン（特に 2 番目の物体の形状トークン）の文脈的埋め込み内に統合されます。
   • 分散分割（Variance Partitioning）の分析により、物体トークンの埋め込みベクトルの中に、物体の属性だけでなく空間関係の情報も線形に符号化されていることが確認されました。
2. 単一トークンからの読み取り:
   • DiT は関係語そのものではなく、物体トークン（Shape2）から直接空間関係と物体属性の両方をデコードします。
   • 特定のヘッド（例：Layer 3, Head 7）が、物体トークン内の関係因子ベクトルと位置符号を相互作用させ、空間勾配を生成します。
   • 特徴: 回路はコンパクトですが、関係情報と物体情報が混在しており、分離されていません。

3.2. ロバスト性の違い（重要な発見）

両モデルはトレーニングデータ内（In-domain）では同様の高い精度を示しますが、プロンプトの微小な摂動に対する耐性に大きな差があります。

• RTE-DiT: プロンプトに余分な単語（例："the"）を追加しても、空間関係の生成精度は安定しています。関係情報が明示的なトークンとして処理されるため、文脈の変化に影響されにくいからです。
• T5-DiT: 余分な単語の追加や語順の変更など、わずかな摂動で空間関係の生成精度が急激に低下します。
  • 原因: T5-DiT は関係情報を物体トークンの文脈埋め込みに依存しているため、不要な単語がその埋め込みベクトルを歪め、結果として空間勾配が誤った方向に偏ってしまうからです。

4. 貢献と意義

1. 非可換関係の生成メカニズムの解明:
   物体間の非可換な関係（A が B の上、と B が A の上は異なる）をニューラルネットワークがどのように符号化・利用するかという長年の疑問に対し、DiT 内での具体的な回路（空間勾配の生成と読み取り）を初めて提示しました。
2. 拡散モデルにおける回路発見の手法論的貢献:
   サンプリングの反復性により解析が困難だった拡散モデルに対し、「Attention Synopsis」と「重み空間スクリーニング」という体系的な手法を提案し、DiT における通信パターンの特定を可能にしました。
3. テキストエンコーダー設計の重要性の再認識:
   空間関係生成の失敗は単に DiT の問題ではなく、テキストエンコーダーの設計がボトルネックとなり得ることを示しました。T5 のような事前学習済みエンコーダーは言語一般化に優れますが、空間関係のような構造的なタスクにおいては、明示的な関係トークンを扱う RTE のような設計の方が、ロバストな制御を可能にします。
4. 実世界への示唆:
   現実世界のプロンプト（自然な文章）では、T5-DiT のような「文脈に埋め込まれた関係情報」に依存するモデルは、ノイズ（不要な単語など）に対して脆弱である可能性が高いことを示唆しました。

結論

本論文は、DiT が空間関係を生成する際、テキストエンコーダーの性質に応じて「モジュール化された 2 段階回路」または「文脈融合型回路」という全く異なるメカニズムを採用することを発見しました。特に、事前学習済みエンコーダー（T5）を使用するモデルは、文脈の微小な変化に対して空間関係の生成が不安定になるという重要な限界を明らかにしました。これは、将来の T2I モデル設計において、単なる性能向上だけでなく、埋め込みモデルの設計と DiT の回路設計のバランスを考慮する必要性を浮き彫りにしています。