技術要約：Untwisting RoPE: 共有アテンションにおける周波数制御（DiTsにおいて）

問題提起

Fluxなどのモデルで使用されているMMDiT（Multi-modal Diffusion Transformer）アーキテクチャのような拡散トランスフォーマー（DiT）において、位置情報は回転式位置埋め込み（RoPE）を介してエンコードされています。RoPEは標準的な画像生成における空間構造のモデリングには不可欠ですが、**共有アテンション（shared-attention）**の設定（ターゲット画像の生成がリファレンス画像のトークンに注目する場合）におけるその挙動については、十分に理解されていません。

著者らは、この設定における決定的な失敗モードを特定しました。それは**リファレンスのコピー（reference copying）**です。モデルは、リファレンス画像のスタイル属性をターゲットプロンプトに転送する（例：雄牛のスタイルのキリンを生成する）代わりに、リファレンスの内容をほぼそのまま複製するという崩壊を起こします。これは、RoPEによって導入される強力な空間バイアスが、ターゲット画像のクエリを、空間的に整列したリファレンスのトークンに対して不釣り合いに注目させることで、意味的な類似性を上書きしてしまうために発生します。

これに対処するための既存の試み（アテンションの共有を後半のデノイジングステップに制限したり、同時期の研究であるAlignedGenのようにシフトされた位置インデックスを使用したりすること）は、粗い制御しか提供できないか、あるいは「ゴースト・サイクリック・インペインティング（ghost-cyclic inpainting）」（モデルがリファレンスをターゲットの継続物として扱う現象）のような新たなアーティファクトを導入してしまいます。

手法

本論文は、モデルの再学習やアーキテクチャの変更を行うことなく、意味のある共有アテンションを復元するために、RoPEの**周波数認識型変調（frequency-aware modulation）**を提案しています。

1. RoPEの周波数構造の分析

著者らは、RoPEが明確に異なる位置感度を持つ周波数成分に分解されるという、原理的な分析を提示しています。

• 高周波成分： 埋め込みの低インデックスの次元に対応します。これらは位置の変位（$\Delta$）に対して角度が急速に変化し、強力な**局所性（locality）**を強制します。これにより、空間的な距離が離れるにつれてアテンションが急激に減衰します。
• 低周波成分： 高インデックスの次元に対応します。これらは変位に対して緩やかに変化し、弱い位置バイアスを提供しながら、**グローバルな空間的一貫性（global spatial coherence）**を維持します。

共有アテンションのシナリオでは、高周波成分がアテンション計算を支配し、モデルが空間的に整列したリファレンストークンに固執し、その内容をコピーする原因となります。

2. 周波数認識型変調

これに対抗するため、著者らはアテンション計算中に**リファレンス画像のキー（$K_{ref}$）**の周波数帯域を選択的に変調する方法を導入しました。

• メカニズム： RoPEを無効化する（退化した生成を引き起こす）あるいはアテンションを一様にスケーリングするのではなく、この手法はRoPEのスペクトル全体にわたってスケーリング因子の滑らかな多項式補間を適用します。
• スケーリング戦略：
  • 減衰（Attenuation）： 高周波成分は、コピーを駆動する強力な位置バイアスを軽減するために、スケールダウン（係数 $s_{hf} \in (0, 1)$）されます。
  • 増幅（Amplification）： 低周波成分は、グローバルで意味論的なガイダンスに基づいたアテンションを促進するために、スケールアップ（係数 $s_{lf} > 1$）されます。
• タイムステップ・スケジューリング： 変調スケール（$s_{hf}$ および $s_{lf}$）は、デノイジングのタイムステップに従って線形に増加します。初期のタイムステップでは、グローバルな構造と広範なスタイルガイダンスを確立するために低い変調を用い、後半のタイムステップでは、微細なスタイルの詳細やテクスチャの転送を可能にするためにスケーリングを増加させます。

最終的な共有アテンションのためのアテンションキーは、以下のように定式化されます：
K = [K'_{img}^{tar} \oplus K'_{txt}^{tar}] \oplus s_{FA}^t \cdot K_{img}^{ref}
ここで、$s_{FA}^t$ はタイムステップ $t$ における周波数認識型変調スケールであり、リファレンスキーに対してのみ適用されます。

主な貢献

1. 共有アテンションにおけるRoPEの原理的分析： 本論文は、DiTにおける「リファレンス・コピー」現象が、アテンション共有全般の失敗ではなく、RoPEの高周波成分が厳格な空間的局所性を強制することによる特定の帰結であることを証明しています。
2. 周波数制御メカニズム： 著者らは、RoMAの周波数帯域を制御するゼロショットかつ学習不要な手法を導入しました。これにより、位置の整列と意味的な類似性の間の制御可能なトレードオフが可能になります。
3. アーティファクトの解決： この手法は、座標を単にシフトさせるのではなく、根本的な周波数構造に直接対処することで、従来の（シフトされたRoPEのような）手法で見られた「ゴースト・サイクリック・インペインティング」やコンテンツ漏洩の問題を解決します。
4. スタイル整合型生成の復元： このアプローチにより、DiTはリファレンスの内容を複製することなく、多様な意味的内容（例：異なる動物）を生成しながら、一貫した視覚的スタイルを維持することが可能になります。これは、UNetベースのモデルや重いファインチューニングを必要としていた能力です。

結果

Flux 1.-dev モデルを用いた評価において、本手法は以下の成果を示しました。

• 効果的なスタイル転送： モデルは、リファレンス（芸術作品や写真を含む）からスタイルの属性（色、質感、筆致）を、リファレンスの内容を複製することなく、ターゲットプロンプトへ正常に転送できます。
• 制御可能性： 変調スケールを調整することで、ユーザーはスタイル転送の度合いとコンテンツ保持の度合いを制御できます。高い $s_{hf}$ はリファレンスの構造（ポーズや形状）を保持しつつスタイルを転送し、低い $s_{hf}$ はより多くの構造的変化を許容します。
• ベースラインに対する優位性：
  • StyleAligned（UNetベース）と比較して： 本手法は、DiTにおけるコンテンツ漏洩や構造的不規則性を回避します。
  • AlignedGen（シフトされたRoPE）と比較して： 本手法は、周期的なコピーのアーティファクトやゴースト・サイクリック・インペインティングを回避し、高解像度においてもスムーズなスタイル転送を提供します。
  • IP-Adapter および B-LoRA と比較して： 本手法は追加の学習や画像エンコーダを必要とせず、より汎用的なゼロショット・ソリューションを提供します。

意義

本論文は、位置エンコーディングは中立的なコンポーネントではなく、生成モデルにおける**局所性（locality）と意味的関連性（semantic association）**のバランスを能動的に制御するものであると主張しています。この知見は、なぜUNetベースのモデル（畳み込みによる帰納バイアスに依存）のために開発されたアテンション共有戦略が、RoPEに完全に依存するDiTに直接適用した場合に失敗するのかを説明しています。

高周波のRoPE成分がコンテンツのコピーを駆動することを明らかにすることで、本研究はアテンションの流れを形成するための直接的で解釈可能な手段を提供しています。これは、アーキテクチャの変更や再学習なしに、内部表現への標的を絞った介入を通じて生成モデルを制御するという、より広い方向性を示唆しています。