技術的概要：拡散モデルの解釈のための残差化時空スパースオートエンコーダ

問題定義

テキストから画像への拡散モデルは、反復的なノイズ除去プロセスを通じて画像を生成するため、内部のニューラル層は単一の静的表現ではなく、活性化の軌跡を生成する。スパースオートエンコーダ（SAE）は、拡散活性化を解釈可能な特徴方向に分解するために最近適用されてきたが、既存のアプローチは通常、孤立した時間ステップにおける活性化を分析するか、時間を条件として扱うか、Diffusion Transformer などのアーキテクチャに SAE を適用するものである。

これらの既存手法の決定的な限界は、拡散活性化に内在する強い時相相関を明示的に考慮していない点にある。隣接する時間ステップにおける活性化は非常に相関しており、あるステップの信号の多くは、直前のステップから線形的に予測可能である。生で時相的に連結された活性化に対して直接 SAE を訓練することは、特徴の背後にあるスパースな構造を捉える代わりに、これらの冗長で線形的に予測可能なダイナミクスを学習することにモデル容量を浪費するリスクがある。

手法

著者は、線形的に予測可能なダイナミクスを残差構造から明示的に分離することにより、活性化軌跡からスパースな特徴を学習するためのフレームワークである残差化時空 SAEを提案する。

1. データ収集と前処理

• ソース: 活性化は、クラスターフリーガイダンス中の Stable Diffusion 1.5 の U-Net における条件分岐から収集される。具体的には、ミッドブロックのクロスアテンション層の残差寄与が抽出される。
• 軌跡の構築: 生成された画像について、活性化は $\ell$ 回のノイズ除去ステップごとにサンプリングされ（例：$\ell=10$）、活性化テンソルの系列 $\{a(\tau_i)\}$ を形成する。

2. 線形残差化

時相的冗長性を除去するため、本手法は連続する正規化された活性化ブロック間で線形予測子を適合させる。

• 正規化: 活性化は、学習セットの統計量を用いて時間ステップごとに正規化される。
• リッジ回帰: 各隣接する遷移 $i-1 \to i$ について、リッジ回帰モデル $(W_i, b_i)$ が適合され、直前の活性化 $\bar{a}_{q,i-1}$ から現在の活性化 $\bar{a}_{q,i}$ を予測する。
• 残差の計算: 残差 $r_{q,i} = \bar{a}_{q,i} - (W_i \bar{a}_{q,i-1} + b_i)$ は、直前のステップから線形的に予測不可能な活性化の成分を捉える。

3. 残差化軌跡に対する SAE の訓練

SAE への入力は、初期活性化と後続の残差からなる連結ベクトル $z_q$ である：
$$z_q = [\bar{a}_{q,0}, r_{q,1}, \dots, r_{q,T-1}]$$
BatchTopK SAEがこの残差化された表現に対して訓練される。目的は、スパース性を強制（バッチごとに上位 $K$ 個のアクティブな潜在変数のみを保持）しながら、正規化された入力を再構成することである。これにより、スパースな潜在変数は、線形的に予測可能な時相的ドリフトを超えて存在する構造を捉えることが促される。

4. 活性化空間における解釈

本手法の重要な構成要素は、学習された特徴を元の活性化空間にマッピングし、軌跡として分析することである。

1. デコーダ方向: SAE デコーダベクトルは当初、正規化された残差空間にある。
2. 非正規化: スケールを適用して正規化された残差空間へ戻す。
3. 線形ダイナミクスの統合: 残差方向は、適合された線形予測子 ($W_i$) を通じてマッピングされ、正規化された活性化空間における完全な軌跡方向を再構成する。
4. 最終マッピング: 結果を非正規化し、元の活性化単位におけるデコーダ軌跡 $\{\phi_{k,i}\}$ を得る。これにより、各潜在変数をノイズ除去時間を通じて発展する特徴軌跡として分析することが可能になる。

主要な貢献

1. 残差化時空 SAE: 拡散活性化軌跡を線形的に予測可能な成分と残差に分離し、冗長な時相的ダイナミクスのモデリングを避けるために後者に対して特に SAE を訓練する手法の導入。
2. 活性化空間デコーダ解釈: 残差化されたデコーダ方向を元の活性化空間にマッピングする導出により、各潜在変数を時空特徴軌跡として分析可能にする。
3. 経験的評価: Stable Diffusion 1.5 における包括的な評価により、これらの表現が、意味的に意味のある生成効果を持つ、時相的に構造化され、空間的に局所化された方向を捉えていることを実証。

結果と評価

再構成性能

• 残差化の優位性: 固定された軌跡レベルのスパース性予算の下で、残差化された変種は、平均二乗誤差（MSE）と説明分散（EV）において、非残差化の対応物（連結型および時間ステップごとの両方）を一貫して上回る。
• 効率性: 残差化連結モデルは、軌跡全体に対して単一のスパースコードを使用しながら、競争力のある再構成性能を達成するのに対し、時間ステップごとのベースラインは分散されたコードを必要とする。
• ベースライン比較: 本手法は、非残差化で成分正規化された軌跡を使用するMatryoshka SAEベースラインよりも再構成品質において優れている。

時空構造分析

• 時相的専門化: 潜在変数は明確な時相プロファイルを示す。「初期」の特徴は広範な時相的支持と高い時間ステップ間自己相似性を示す傾向があり、永続的な構造概念を捉えていることを示唆する。「後期」の特徴はより時間ステップ固有であり、時相的に安定性が低い。
• 空間的局所化:
  • 初期の特徴は、より広範な正の空間的支持（高い空間エントロピー）を示し、しばしば拡散的な領域や境界のような領域をカバーする。
  • 中期および後期の特徴は、より空間的に集中しており（低いエントロピー）、特定の物体のような領域や顕著な領域に局在する。

操作実験

• 単一特徴操作: 特定の潜在変数への介入は、目に見える生成効果を生み出す。初期の特徴は全体的な構造に影響を与えるのに対し、中期の特徴は変化を局所化できる（例：胴体に装飾品を追加する）。後期の特徴は、いくつかの例において明確な意味効果を示さなかった。
• 特徴転送: RIEBench データセットを用いて、著者はソースからターゲットへの特徴転送を行った。残差化連結モデルは、編集効率（ターゲットプロンプトへの CLIP 類似度の向上）とソース画像構造の保持（LPIPS 距離の最小化）の間の最適なバランスを達成した。編集効率において、非残差化および時間ステップごとの変種を上回った。

意義と主張

本論文は、残差化時空 SAEが、時相的に構造化された拡散活性化を研究するための有用なフレームワークを提供すると主張する。線形的に予測可能なドリフトを明示的に除去することにより、本手法はスパースな潜在変数が、生成プロセスの非冗長で意味的に有意義な構造に焦点を当てることが可能になる。

著者は、自らのアプローチが、孤立した時間ステップや非残差化された時相的表現ではなく、残差化された活性化軌跡上でスパース辞書を学習する点で先行研究と異なることを強調している。これらの潜在変数を軌跡として活性化空間に逆マッピングする能力は、ノイズ除去プロセスにおいて特定の機能がいつ、どこで発現するかについての理解を深めることを可能にする。

著者が指摘する限界:

• 実験は、Stable Diffusion 1.5 における単一の U-Net 活性化ソース（ミッドブロックのクロスアテンション）に焦点を当てている。
• 操作実験は定性的であり、本論文は最適化された画像編集性能を主張するものではなく、学習された特徴の解釈可能性と因果的影響を検証するために操作を用いている。