Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers

この論文は、テキストから画像を生成するマルチモーダル拡散トランスフォーマーにおいて深層でプロンプトの意味が失われる「プロンプト忘却」現象を特定し、初期層のプロンプト表現を後続層に再注入する学習不要な手法「プロンプト・リインジェクション」を提案することで、指示追従性や生成品質を向上させることを示しています。

要約

この論文は、最新の画像生成 AI（テキストから絵を描く AI）が抱えるある「忘れっぽさ」の問題を解決し、より指示通りに絵を描けるようにする新しい方法を提案したものです。

タイトル：「プロンプトの再注入（Prompt Reinjection）：マルチモーダル拡散トランスフォーマーにおける『指示忘れ』の解消」

これを、AI の仕組みを「料理」や「物語の伝達」に例えて、わかりやすく解説します。

---

1. 問題点：AI は「指示」を忘れがち（プロンプト・フォージティング）

最新の画像生成 AI（SD3 や FLUX など）は、従来の AI とは少し違います。

• 昔の AI：「料理のレシピ（テキスト）」を厨房の入り口で渡して、料理人（画像生成部分）がそれを見て料理を作る。レシピは最初だけ見ればよく、その後は料理人が勝手に進める。
• 最新の AI：「レシピ（テキスト）」と「料理（画像）」が同じ大きなキッチンで、一緒に調理の過程を歩みながら会話しながら進んでいくという形です。

しかし、ここに大きな問題がありました。
長い調理工程（AI の深い層）を進むにつれて、**「レシピの内容が薄れていってしまう」**のです。

• 例え話：
  あなたが「4 匹の犬が青い空の下で走っている」と指示を出したとします。
  最初のうちは AI は「4 匹」「青い」「犬」という情報をしっかり持っています。しかし、画像を完成させるための複雑な計算を何十回も繰り返す（AI の層を深くする）と、AI は「4 匹」や「青い」という細かい情報を忘れてしまい、結果として「3 匹の犬」や「赤い空」のような、指示と違う絵を描いてしまうことがありました。

これを論文では**「プロンプト・フォージティング（指示忘れ）」**と呼んでいます。特に「位置関係（右にある）」や「数（4 つ）」といった細かい指示が、一番忘れられやすいことがわかりました。

2. 原因：なぜ忘れるのか？

AI は「画像を綺麗に描くこと」に全力を注ぎます。その過程で、テキストの情報は「画像を作るためのヒント」として使われますが、「画像を作る」こと自体に直接の正解（正解の画像）がないため、テキストの情報は勝手に変化してしまい、元の意味が失われていくのです。

まるで、長い旅路の中で、出発時に持っていた「目的地の地図」が、道中の風景に気を取られて次第にボロボロになり、最後には「あ、どこに行くんだったっけ？」となってしまうようなものです。

3. 解決策：「再注入（Reinjection）」という魔法

そこで著者たちは、**「忘れないうちに、再び思い出させてあげよう」というアイデアを思いつきました。それが「プロンプト・リインジェクション（Prompt Reinjection）」**です。

仕組みのイメージ：

1. 浅い層（旅の初期）：AI がまだ「4 匹」「青い」という情報を鮮明に持っている、最初の段階のテキスト情報をコピーします。
2. 深い層（旅の後半）：AI が「4 匹」を忘れそうになっている、後半の工程で、**そのコピーした情報を「注入（リインジェクション）」**して、再び AI の頭に叩き込みます。

重要なポイント：
ただ単に情報を足すだけでは、AI が混乱してしまいます（「今、何の話をしてたっけ？」と）。
そこで、この方法は**「情報の形を合わせてから注入する」**という工夫をしています。

• 例え話：
  後半の AI の頭は「料理の味付け」に慣れている状態ですが、最初のレシピは「生野菜」のままです。これをそのまま入れると味が壊れます。
  そこで、**「後半の料理の味付けに合わせて、最初のレシピも少し味付け（統計的な調整）をしてから」**入れることで、AI が混乱せずに「あ、そうそう、4 匹の犬だったんだ！」と思い出せるようにします。

4. 効果：指示通りに描けるようになった！

この方法を実際に試したところ、驚くべき結果が出ました。

• 数え間違いの減少：「4 匹の犬」を「4 匹」で描けるようになりました。
• 位置関係の改善：「右にある赤い車」を、ちゃんと右側に描けるようになりました。
• 色や形の正確性：「青い空」「緑の芝生」など、色や素材の指定も守られるようになりました。

しかも、AI の学習（トレーニング）を一切行わず、絵を描く時（推論時）にこの「思い出させる作業」を挟むだけで実現できました。つまり、既存の AI モデルをそのまま使って、すぐに性能を上げられるという画期的な方法です。

5. まとめ

この論文が伝えたかったことはシンプルです。

「最新の AI は絵を描くのが上手だけど、長い工程の中で指示を忘れっぽくなっている。だから、**『忘れないうちに、最初の話（指示）を思い出させてあげる』**という簡単な作業を挟むだけで、AI はもっと指示通りに、正確で素晴らしい絵を描けるようになるよ」

これは、AI が「指示に従う力（インストラクション・フォロイング）」を劇的に向上させる、シンプルで効果的な「魔法の薬」のような発見と言えます。

技術概要

論文要約：Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers

この論文は、テキストから画像を生成するマルチモーダル拡散トランスフォーマー（MMDiT）において発生する**「プロンプト忘却（Prompt Forgetting）」現象を特定し、それを軽減するための新しい手法「Prompt Reinjection（プロンプト再注入）」**を提案したものです。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

---

1. 問題定義：プロンプト忘却（Prompt Forgetting）

近年の MMDiT（Stable Diffusion 3, FLUX, Qwen-Image など）は、テキストと画像の潜在表現を単一のトランスフォーマースタック内で統合的に処理し、双方向の相互作用を可能にすることで、複雑な指示への追従能力を向上させています。しかし、著者らは以下の問題点を発見しました。

• 現象: デノイジングプロセスが進行し、トランスフォーマーの層が深くなるにつれて、テキストブランチ（プロンプト）の表現が徐々に劣化し、プロンプトに含まれる微細な意味情報（属性、数、空間関係など）が失われる現象が発生します。
• 原因: 学習目的関数（損失関数）は画像の再構成（ノイズ予測）に対してのみ定義されており、テキストトークンに対する直接的な教師信号が存在しません。テキスト特徴は、画像生成への影響を通じて間接的にのみ更新されるため、深層になるほど意味的な保存性が低下し、プロンプト情報が「忘却」されてしまいます。
• 実証: 既存のモデル（SD3, SD3.5, FLUX）において、層ごとのテキスト特徴を分析した結果、以下の傾向が確認されました。
  • 局所的な意味構造の保存度（CKNNA）が層の深さとともに単調に低下する。
  • 深層におけるテキスト特徴の分布が収束し、識別可能性が失われる。
  • 層ごとのプロンプト属性復元能力（プローブ精度）が深くなるにつれて低下し、特に「空間関係」や「数」の情報が早期に失われる。

2. 提案手法：Prompt Reinjection（プロンプト再注入）

この忘却現象を軽減するために、著者は**学習不要（Training-free）かつ推論時（Inference-time）**に適用可能な手法「Prompt Reinjection」を提案しました。

• 基本的なアイデア: 浅い層（浅いブロック）で保持されている高忠実度のプロンプト特徴を、より深い層のトランスフォーマーブロックへ再注入（Reinjection）することで、深層における意味情報の欠落を補完します。
• 技術的詳細:
  1. 起源層（Origin Layer）の選択: 浅い層（通常はテキストエンコーダ出力の直後、または最初の急激な分布変化の直後の層）からテキスト特徴 $T^{(ori)}$ を抽出します。
  2. ターゲット層（Target Layer）への注入: 深層のテキスト特徴 $T^{(tgt)}$ に対して、起源層の特徴を残差接続として加算します。
  3. 分布アンカリングと幾何学的整列: 単なる加算では、層間の分布のズレ（スケール、オフセット）や座標系の回転が性能を低下させるため、以下の2段階の調整を行います。
     • Distribution Anchoring: Layer Normalization を用いて特徴を正規化し、注入後にターゲット層の統計量（平均・分散）に戻すことで、数値的な安定性を確保します。
     • Geometry Alignment: 正則 Procrustes 変換（直交回転行列 $R$）を用いて、起源層の特徴空間をターゲット層の多様体に整列させます。これにより、異なる深さでの特徴空間の回転ズレを補正します。
  • 数式: $T^{(added)} = \hat{T}^{(tgt)} + w \cdot \hat{T}^{(ori)}R$ （ここで $w$ は注入強度のハイパーパラメータ）。

3. 主要な貢献

1. 現象の定量的発見: MMDiT における「プロンプト忘却」を、CKNNA（局所構造保存）や層ごとのプローブ（属性復元精度）を用いて初めて体系的に定量化し、深層になるほど微細な情報が失われることを実証しました。
2. 学習不要な改善手法の提案: モデルの再学習や微調整（Fine-tuning）を一切行わず、推論時のみでプロンプト追従性を大幅に向上させる「Prompt Reinjection」を提案しました。
3. 広範なモデルへの適用性: SD3, SD3.5, FLUX, Qwen-Image など、異なるアーキテクチャを持つ主要な MMDiT モデルにおいて、一貫した性能向上を確認しました。

4. 実験結果

主要なベンチマーク（GenEval, DPG-Bench, T2I-CompBench++）および人間評価指標（HPSv2, ImageReward など）での評価結果は以下の通りです。

• 指示追従性の向上:
  • GenEval: SD3.5 で総合スコアが 6.48% 向上、FLUX で 5.64% 向上。
  • 特に「空間関係（Position）」や「数え上げ（Counting）」、属性結合（Color/Shape）といった、忘却の影響を受けやすいタスクで顕著な改善が見られました。
  • 例：「4 つの犬」や「右側の犬」などの指示が、ベースモデルでは無視されがちでしたが、本手法により正確に生成されるようになりました。
• 画像品質の維持:
  • 指示追従性の向上は、画像の美的品質や人間評価（HPSv2, ImageReward, PickScore）や CLIP スコアを低下させませんでした。むしろ、一部のモデルではわずかに向上しました。
• 計算コスト:
  • 追加の計算コストは非常に低く、1 つのトランスフォーマーブロックあたりの FLOPs 増加は約 8% 程度、レイテンシの増加も数ミリ秒レベルで、実用的なオーバーヘッドです。

5. 意義と結論

• 理論的意義: 従来の「テキストは固定された条件」として扱われていた U-Net 型モデルとは異なり、MMDiT においてテキスト特徴が動的に変化し、深層で意味を失うという根本的な課題を明らかにしました。
• 実用的意義: 大規模なモデルの再学習なしに、複雑なプロンプトへの追従性を劇的に向上させることができるため、既存の画像生成モデルの性能を即座に引き上げる実用的なソリューションとなります。
• 将来展望: 本手法は、テキストブランチに対する直接の教師信号（例：テキスト再構成損失）の導入や、層ごとの適応的な注入重みの学習など、さらなる研究の基盤を提供します。

総じて、この論文は MMDiT の内部動作に関する重要な洞察を提供し、プロンプト忘却という課題に対して、シンプルかつ効果的な「再注入」アプローチによって解決策を示した画期的な研究です。