DiTの「無用」に見えるテンプレートトークンが意味を保持する

導入

文生画像モデルでは、意味を担うのはプロンプト中の内容語だと考えがちです。たとえば「猫」「赤い車」「未来都市」といった語が、生成される画像の対象や属性を決めるように見えます。一方で、テンプレートや構造を表すテキストトークンは、入力形式を整えるだけの補助的な存在に見えるかもしれません。

しかし、この論文は Diffusion Transformer（DiT）の内部を因果的に調べることで、その直感を覆しています。入力時に意味を表しているトークンと、生成中に意味を維持しているトークンは必ずしも同じではない、というのが中心的な主張です。

主要なポイント

• テンプレートトークンは最初は情報をほとんど持たない。 研究では、プロンプト内容トークンと構造的テンプレートトークンを分けて分析している。テキストエンコーダの出力段階では、テンプレートトークンにはプロンプト固有の情報がほとんど含まれていない。

• DiT 内部で注意の受け皿になる。 著者らは、注意分解と、トークン範囲・ヘッド・層に対する標的介入を組み合わせた因果的解釈フレームワークを用いた。その結果、テンプレートトークンが DiT 内部で画像からテキストへの注意を強く集める sink として現れることを示した。

• 物体同一性の維持に因果的に関与する。 これは単なる相関ではない。介入実験により、テンプレートトークンが生成中の対象アイデンティティを保つうえで機能していることが示されている。

• 意味は直接転送されない。 さらに興味深いのは、意味がプロンプトトークンからテンプレートトークンへ直接渡されるわけではない点である。まずプロンプトの意味が画像 latent に注入され、その後、画像 latent からテンプレートトークンへ読み戻されるという間接経路が示唆されている。

• 剪定にもつながる。 この知見をもとに、著者らは訓練不要の剪定規則を提案した。プロンプトトークンへ最も強く注意するヘッドは比較的削除可能であり、それらを剪定することで注意 FLOPs を 20% 削減し、GenEval の低下は 1.4 ポイントにとどまった。

意義と影響

この研究の価値は、テンプレートトークンが有用だと指摘したことだけではありません。DiT の生成計算において、意味の注入、保持、伝播、視覚合成が異なるヘッドや層に分かれて組織化されていることを、因果的な観点から描いている点が重要です。

モデル解釈の面では、注意マップを眺めるだけでなく、特定のトークンやヘッドを操作して生成結果への影響を見ることの重要性を示しています。実用面では、内部メカニズムの理解が計算削減に直結しうることも示されました。今後、DiT の設計やプロンプトテンプレートの扱いを考える際、こうした「見えにくい意味の保持場所」は無視できない要素になりそうです。

出典：Hugging Face Daily Papers