OminiControl: Minimal and Universal Control for Diffusion Transformer

拡散ベースの手法は、画像生成のための強力なフレームワークとして台頭してきた[9, 28]。これらの手法は、テキストから画像への合成[28, 2, 33]、画像から画像への変換[32]、画像編集[20, 1]など、多様なタスクにおいて成功を収めている。 最近の進歩により、品質と効率の両面で大幅な改善が見られ、特に潜在拡散モデル[28]の導入が注目される。 生成能力をさらに向上させるため、大規模なトランスフォーマーアーキテクチャがこれらのフレームワークに統合され、DiT[23, 2, 13, 3]のような先進的なモデルが生まれている。 これらのアーキテクチャの革新を基に、FLUX[13]はトランスフォーマーベースの設計とフローマッチング目的関数[18]を組み合わせ、最先端の生成性能を達成している。

制御可能な生成は、拡散モデルの文脈で広く研究されてきた。 テキストから画像への変換モデル[28, 25]は条件付き生成の基礎を確立し、一方で画像などの追加の制御信号を組み込むための様々なアプローチが開発されてきた。 注目すべき手法には、拡散モデルにおいて空間的に整列された制御を可能にするControlNet [41]や、軽量なアダプターで効率を改善するT2I-Adapter [22]がある。UniControl [26]は、Mixture-of-Experts (MoE)を使用して異なる空間的条件を統合し、モデルサイズをさらに削減している。しかしながら、これらの手法は、ノイズ除去ネットワークの隠れ状態に空間的に条件情報を追加することに依存しており、本質的に被写体駆動型生成のような非空間的タスクに対する有効性が制限されている。IP-Adapter [39]は、追加のエンコーダーを通じてクロスアテンションを導入することでこの問題に対処し、SSR-Encoder [42]は画像条件付きタスクにおけるアイデンティティ保持をさらに強化している。 これらの進展[15, 19, 5]にもかかわらず、空間的に整列されたタスクと非整列タスクの両方に対する統一的な解決策は依然として見出されていない。

FLUX.1 [13]、Stable Diffusion 3 [28]、PixArt [2]などのアーキテクチャで採用されているDiffusion Transformer (DiT)モデル [23]は、ノイズの多い画像トークンを反復的に精製するためにトランスフォーマーブロックのデノイジングネットワークを使用している。

各トランスフォーマーブロックは、2種類のトークンを処理する：ノイズの多い画像トークン$\mathbf{X}\in\mathbb{R}^{N\times d}$とテキスト条件トークン$\mathbf{C}_{\text{T}}\in\mathbb{R}^{M\times d}$である。ここで、$d$は埋め込み次元、$N$と$M$はそれぞれ画像トークンとテキストトークンの数である（図2）。これらのトークンは隠れ状態$X$と$C_{\text{T}}$に埋め込まれ、トランスフォーマーブロック全体を通じて一貫した形状を維持する。

各DiTブロックでは、$X$と$C_{\text{T}}$を正規化した後、コアのMM-Attentionモジュール [30]によって処理される。このモジュールは、Rotary Position Embedding (RoPE) [34]を使用してトークン間の位置依存性を組み込む。2次元グリッドの位置$(i,j)$にあるトークンに対して、RoPEはクエリとキーの射影に回転行列を適用する：

ここで、$R(i,j)$は位置$(i,j)$における回転行列である。同様に、テキスト条件トークン$C_{\text{T}}$のクエリとキーの射影も同じ方法で定義され、FLUX.1ではすべてのテキストトークンの位置が$(0,0)$に設定される。

RoPEを適用した後、両方のトークンタイプからのクエリ、キー、バリューが連結され、統合された行列$Q_{Z}$、$K_{Z}$、$V_{Z}$を形成し、結合されたトークンセット$Z=[X;C_{\text{T}}]$を表現する。その後、MM-Attention操作は以下のように計算される：

これにより、注意機構を通じて画像とテキスト条件トークン間の相互作用が可能になる。

我々のアプローチでは、まず条件画像をモデルのVAEを通してエンコードし、ノイズ画像トークンと同じ潜在空間に投影して$\mathbf{C}_{\text{I}}\in\mathbb{R}^{N\times d}$を形成する。

ControlNet [41]やT2I-Adapter [22]のような従来の手法では、条件画像を空間的に整列させ、その隠れ状態をノイズ画像トークンの隠れ状態に直接加算することで組み込んでいる：

ここで、$H_{X}$は更なる処理のための結合された隠れ状態を表し、$H_{C_{\text{I}}}$は条件画像からの隠れ状態である。 この手法は空間的に整列したタスクに対しては効果的であるが、2つの主要な制限がある：(1)空間的に整列していないシナリオを扱う際の柔軟性に欠け、(2)空間的に整列している場合でも、隠れ状態の直接的な加算はトークン間の相互作用を制限し、モデルの性能を潜在的に制限する可能性がある。

対照的に、我々の手法では、空間的に整列していない制御タスクを可能にし、より大きな条件付け柔軟性を提供するために、条件画像トークンをテキストとノイズ画像トークンと同様に処理し、統一されたシーケンスに統合する：

ここで、$Z$はノイズ画像トークン$X$、テキストトークン$C_{\text{T}}$、および条件画像トークン$C_{\text{I}}$の連結されたシーケンスを表す。この統一されたアプローチにより、特殊な処理経路なしにマルチモーダル注意機構[30]に直接参加することが可能になる（図2に示す）。 直接加算アプローチよりも優れた結果を達成する。

比較分析により、我々のアプローチが直接加算法よりも優れた結果を達成することが示されている。これは図3の注意の可視化によって実証されている。 さらに、図4(a)の訓練曲線は、マルチモーダル注意手法が一貫して直接加算アプローチよりも低い損失値を達成していることを示している。 この統一されたシーケンスアプローチの有効性は、空間的に整列したタスクと空間的に整列していないタスクの両方で実証されており（図5）、多様な条件付き生成シナリオを扱う上での汎用性を強調している。

条件画像の統合には、条件画像とターゲット画像の間の効果的な相互作用を確保するために、位置情報の慎重な考慮が必要である。 従来の手法では、条件画像と出力画像の間の厳密な空間的整列を仮定することが多く、非整列タスクへの適用が制限されていた。 さらに、条件トークンとターゲットトークンの相対的な位置関係は、モデルの学習効率と汎化能力に大きな影響を与える可能性がある。 我々の統一されたシーケンス設計では、条件トークンが他のトークンと連結されているため、異なる位置エンコーディング戦略を探索する柔軟性がある。

FLUX.1のTransformerでは、各トークンに空間情報をエンコードするための対応する位置インデックスが割り当てられる。 512$\times$512の入力画像に対して、VAE [11] エンコーダーはまずそれを潜在空間に投影し、その後、潜在表現を32$\times$32のトークングリッドに分割する。各トークンには、$(i,j)$（ここで$i,j\in[0,31]$）という一意の二次元位置インデックスが与えられる。 このインデックス付けスキームは、潜在空間において元の画像の空間構造を保持し、一方でテキストトークンは$(0,0)$という固定の位置インデックスを維持する。

空間的に整列したタスクに対して、我々の初期アプローチは、条件トークンに元の画像の対応するトークンと同じ位置埋め込みを割り当てることであった。 しかし、被写体駆動生成のような空間的に非整列のタスクに対しては、我々の実験により、条件トークンの位置インデックスをシフトすることで収束が速くなることが明らかになった（図4(b)）。 具体的には、条件画像トークンのインデックスを$(i,j)$（ここで$i\in[0,31]$かつ$j\in[32,64]$）にシフトし、元の画像トークン$X$との空間的な重複がないようにしている。

我々が採用した統一的注意機構は、柔軟なトークン間相互作用を可能にするだけでなく、条件画像の影響を正確に制御することも可能にする。 具体的には、推論時に条件画像の効果を手動で調整できる方法を設計した。 与えられた強度因子$\gamma$に対して、$\gamma=0$を設定すると条件画像の影響が除去され、元の入力のみに基づく出力が得られる。 $\gamma=1$では、出力は条件画像を完全に反映し、$\gamma$が1を超えて増加すると、条件の効果がさらに顕著になる。

この制御可能性を実現するために、我々は元のMM-Attention操作にバイアス項を導入する。具体的には、式3を以下のように修正する：

ここで、bias$\gamma$は強度因子$\gamma$に基づいて条件トークンと他のトークン間の注意の重みを調整するように設計されている。 バイアス項は$(M+2N)\times(M+2N)$行列として構築され、$M$はテキストトークンの数、$N$はノイズ画像トークンと条件画像トークンの数をそれぞれ表す。行列は以下の構造を持つ：

この設計により、強度因子$\gamma$はノイズ画像トークンと条件画像トークン間の注意の重みにのみ影響を与え、テキストトークンおよびモダリティ内の相互作用に対する元の注意パターンを維持する。

被写体一貫性のある生成のためのモデルの訓練には、通常、ポーズ、照明、その他の属性に変化がある一方で、アイデンティティの一貫性を維持するペアの画像が必要である。 IP-Adapterのような以前の手法[39]では、条件付けとターゲットのペアに同一の画像を使用しており、これらのアプローチには効果的である。 しかし、我々のフレームワークでは、このセットアップは過学習を引き起こし、モデルが入力とほぼ同一の出力を生成する原因となる。

これらの制限を克服するため、我々は被写体のアイデンティティを保持しつつ自然な変化を取り入れた画像を特徴とするデータセットを開発した。 既存のデータセット[31, 12, 15, 17]は同様のニーズに対応しているが、品質または規模のいずれかに制約がある場合が多い。 そのため、我々は慎重に作成されたプロンプトから視覚的に関連する画像のペアを生成するFLUXの固有の能力を活用した新しい合成パイプラインを提案する。

我々のパイプラインはChatGPT-4oを利用して2万以上の多様な画像説明を生成し、これらがFLUXを導いて20万枚以上の画像を生成する（図6）。 生成された画像はChatGPT-4oのビジョン機能を使用して品質評価を受け、最終的なデータセットの一貫性と多様性の両方を確保している。 キュレーションされたデータセットと完全な生成パイプラインは公開される³³3詳細は補足資料に記載されている。データセットとコードはhttps://github.com/Yuanshi9815/Subjects200Kで入手可能である。

タスクと基本モデル。 我々は、条件付き生成タスクの2つのカテゴリーで我々の手法を評価する：空間的に整列したタスク（Cannyから画像への変換、深度から画像への変換、マスクベースのインペインティング、および色付けを含む）と被写体駆動型生成である。 我々の手法は、画像生成のための潜在整流フロートランスフォーマーモデルであるFLUX.1[13]を基盤として構築している。 デフォルトでは、空間的に整列したタスクの画像生成にFLUX.1-devを使用する。被写体駆動型生成タスクでは、より良好な視覚品質を生成する傾向が観察されたため、FLUX.1-schnellに切り替える。

実装の詳細。 我々の手法は、デフォルトのランクが4のLoRA[4]を使用して基本モデルを微調整する。モデルの元の能力を保持し、柔軟性を達成するために、非条件トークンを処理する際のLoRAスケールはデフォルトで0に設定される。

訓練。 我々のモデルは、バッチサイズ1で訓練され、8ステップにわたる勾配累積（実効バッチサイズ8）を行う。 我々は、セーフガードウォームアップとバイアス補正を有効にしたProdigy最適化器[21]を使用し、重み減衰を0.01に設定する。 実験は2台のNVIDIA H100 GPU（各80GB）で実施される。空間的に整列したタスクでは、モデルは50,000イテレーションで訓練され、被写体駆動型生成モデルは15,000イテレーションで訓練される。

ベースライン。 空間的に整列したタスクについては、我々の手法を元のControlNet[41]およびStable Diffusion 1.5上のT2I-Adapter[22]、さらにControlNetのFLUX.1実装であるControlNetPro[14]と比較する。 被写体駆動型生成については、IP-Adapter[39]と比較し、そのFLUX.1実装[37]を評価する。

評価指標。 我々は、空間的に整列したタスクと被写体駆動型生成の両方でモデルを評価する。 空間的に整列したタスクについては、生成品質と制御可能性の2つの側面を評価する。 生成品質は、視覚的忠実性についてFID[8]、SSIM、MAN-IQA[38]、およびMUSIQ[10]を用いて測定し、意味的一貫性についてはCLIPスコア[27]を用いる。 制御可能性については、エッジ条件付き生成では抽出されたエッジマップと入力エッジマップ間のF1スコアを計算し、他のタスクでは抽出された条件マップと元の条件マップ間のMSEを計算する（深度にはDepth Anythingを使用し、色付けには色チャンネル分離を使用するなど）。 被写体駆動型生成については、被写体特性の保持（アイデンティティ保持、素材品質、色忠実度、自然な外観）と要求された修正の正確性の両方を評価する5基準のフレームワークを提案し、すべての評価はGPT-4oのビジョン機能を通じて実施され、体系的な評価を確保する。 詳細な評価方法は付録B.1に記載されている。

評価プロトコル。 我々は2つのデータセットで評価を実施した。 空間的に整列したタスクについては、COCO 2017検証セット（5,000画像）を512$\times$512にリサイズし、タスク固有の条件と関連するキャプションをプロンプトとして使用し、固定シード42で評価する。 被写体駆動型生成については、DreamBooth[31]データセットから750のテキスト-条件ペア（30被写体 × 25プロンプト）を5つの異なるシードでテストし、被写体ごとに1つの選択された画像を条件として使用する。

空間的に整列されたタスク 表1に示すように、我々は5つの空間的に整列されたタスクにおいて、既存のアプローチと比較して我々の手法を包括的に評価した。 我々の手法は、Canny-to-image生成において0.38という最高のF1スコアを達成し、SD1.5ベースの手法であるControlNet [41]とT2I-Adapter [22]、さらにFLUX.1ベースのControlNetPro [14]を大きく上回っている。 一般的な品質指標に関しては、我々のアプローチはほとんどのタスクで一貫して優れており、特にSSIM [36]、MAN-IQA [38]、MUSIQ [10]スコアにおいて顕著に優れたパフォーマンスを示している。 デブラーや色付けなどの困難なタスクでは、我々の手法は大幅な改善を達成している：MSEはControlNetProと比較してそれぞれ77%と93%減少し、FIDスコア [8]はデブラーにおいて30.38から11.49に改善している。 CLIPスコア指標 [27]は、我々の手法がすべてのタスクにおいて高いテキストと画像の一貫性を維持していることを示しており、より良い制御と視覚的品質を達成しながら、意味的整合性を効果的に保持していることを示唆している。 図7に示すように、我々の手法は色付けタスクにおいてより鮮明な詳細とより忠実な色再現を生成し、エッジガイド生成とデブラーのシナリオにおいてより良い構造的忠実性を維持している。

被写体駆動型生成 図8は、既存のベースラインとの包括的な比較を示している。 我々の手法は、特にアイデンティティ保持と修正精度において優れたパフォーマンスを示している。 ランダムシードの平均では、我々の手法はIP-Adapter (FLUX)の57.7%に対して75.8%の修正精度を達成し、同時にIP-Adapter (SD 1.5)の29.4%に対して50.6%のアイデンティティ保持を維持している。 この優位性は最良のシードのシナリオでさらに増幅され、90.7%の修正精度と82.3%のアイデンティティ保持を達成し、最強のベースラインをそれぞれ15.8ポイントと18.0ポイント上回り、効果的な被写体忠実性編集を実証している。 これらの定量的結果は、付録B.1に示されるユーザー調査によってさらに裏付けられている。

パラメータ効率の比較 表2に示すように、我々のアプローチは既存の手法と比較して顕著なパラメータ効率を達成している。 12Bパラメータを持つFLUX.1モデルに対して、我々の手法は学習可能なパラメータをわずか14.5M（約0.1%）必要とするだけであり、これはControlNet（27.5%）やIP-Adapter（7.6%）と比較して大幅に少ない。 FLUX.1のオリジナルのVAEエンコーダーを使用する場合でも、我々の手法はわずか0.4%の追加パラメータで高い効率を維持しており、我々のパラメータ効率の高い設計の有効性を示している。

訓練データの影響。

被写体駆動型生成において、我々のモデルは被写体（例えばぬいぐるみや物体）の参照画像とテキスト記述を入力として受け取り、テキストガイダンスに従いつつ主要な特徴を保持しながら、同じ被写体の新しい画像を生成することを目指している。

セクション3.5で説明した我々のSubjects200Kデータセットの有効性を検証するため、我々はこのタスクに対する2つの訓練戦略を比較した。 第一のアプローチは従来のデータ拡張に依存し、元の画像にランダムなクロッピング、回転、スケーリング、およびコントラスト、彩度、色の調整を適用する。 第二のアプローチは我々のSubjects200Kデータセットを利用する。 図9に示すように、データ拡張で訓練されたモデルは、最小限の変更で入力条件を複製することしか学習していない。 1行目では、タコのぬいぐるみを明るい部屋の設定に単に配置し、その外観と姿勢を正確に維持している。 同様に、2行目では、窓際に配置するという指示にもかかわらず、黄色の目覚まし時計がほぼ同一の詳細で再現されている。 対照的に、我々のSubjects200Kで訓練されたモデルは、テキストプロンプトに忠実に従いながら、被写体の多様かつ一貫した視点を生成する能力を示している。

条件強度分析。 我々は定性的実験を通じて条件強度制御を評価した。図10は、強度因子$\gamma\in{0.25,0.5,0.75,1.0}$を変化させて生成した結果を示している。 結果は、$\gamma$がCanny-to-image生成のような空間的に整列したタスクと、被写体駆動型生成のような空間的に整列していないタスクの両方で生成プロセスを効果的に制御し、条件の影響に対する柔軟な制御を可能にすることを示している。

LoRAランクの影響。 我々はCanny-to-imageタスクに対して、異なるLoRAランク（1、2、4、8、16）で広範な実験を行った。 表3に示すように、我々の実験はLoRAランクを増やすことで一般的にモデルのパフォーマンスが向上し、ランク16が画像品質（FIDとSSIMで測定）、条件制御能力（F1スコアで測定）など複数の側面で最良の結果を達成し、同時にテキスト-画像の一貫性（CLIPスコアで測定）においても競争力のある性能を維持することを示している。 特筆すべきは、より小さなランク（例えばランク1）でも、モデルは競争力のある性能を示し、特にテキスト-画像の整合性では0.765という最高のCLIPスコアを達成しており、限られたパラメータでも我々のアプローチの効率性を示している。

条件付け深度。 FLUX.1のトランスフォーマーアーキテクチャは、異なるモダリティトークン（テキストと画像）に対して別々の正規化モジュールを使用する初期ブロックと、すべてのトークンに対して統一された正規化を共有する後期ブロックという2つの異なるタイプのブロックを特徴としている。 表3に示すように、実験では条件信号の統合をこれらの初期ブロックのみに制限すると、生成プロセスに対する制御可能性が不十分であることが明らかになった。 これは、条件信号がトランスフォーマースタック全体に影響を与えることを許可することが、出力に対する望ましい制御レベルを達成するために重要であることを示唆している。 特に、この発見は、UNetベースのアーキテクチャで効果的であった、主に初期ブロックに条件信号を挿入するという以前のアプローチ[14, 37, 39, 41, 22]が、FLUX.1のようなDiTベースのモデルには完全には適用できない可能性があることを示している。