テキストから画像への拡散モデルに条件付き制御を追加する (論文の閲覧)

Lvmin Zhang および Maneesh Agrawala、スタンフォード大学、arXiv、Cited:113、コード、論文

1 はじめに

追加の入力条件をサポートするために、事前学習済みの大規模拡散モデルを制御するための ControlNet と呼ばれるニューラルネットワークアーキテクチャを提案します。ControlNet はタスク固有の条件をエンドツーエンドで学習し、トレーニングデータセットが小さい (<50k) 場合でも学習は堅牢です。さらに、ControlNet のトレーニングは拡散モデルを微調整するのと同じくらい速く、モデルは個人のデバイスでトレーニングできます。あるいは、強力なコンピューティングクラスターが利用可能な場合は、モデルを大量 (数百万から数十億) のデータに拡張することができます。ControlNet を使用して安定拡散などの大規模拡散モデルを拡張し、エッジマップ、セグメンテーションマップ、キーポイントなどの条件付き入力を可能にした結果を報告します。これにより、大規模拡散モデルを制御する方法が強化され、関連アプリケーションがさらに容易になる可能性があります。

2. 問題提起

大規模なテキストから画像へのモデルの出現により、ユーザーが入力した短い説明的な手がかりだけから、視覚的に魅力的な画像を生成できるようになりました。テキストを入力して画像を取得した後、自然にいくつかの質問が生じるかもしれません。「このヒントベースのコントロールはニーズを満たしているでしょうか?」たとえば、画像処理では、明確に定義された問題定式化を伴う多くの確立されたタスクがあるとすると、これらの大規模なモデルを適用して、これらの特定のタスクを容易にすることができるでしょうか? 幅広い問題状況とユーザー制御を処理するには、どのようなフレームワークを構築する必要がありますか? 大規模なモデルは、特定のタスクで数十億の画像から得られる利点と機能を保持していますか?

これらの質問に答えるために、私たちはさまざまな画像処理アプリケーションを調査し、3 つの発見に到達しました。まず、タスク固有のドメインで利用できるデータの規模は、画像とテキストのドメインほど大きいとは限りません。多くの特定の問題 (オブジェクトの形状/法線、ポーズの理解など) の最大データセットサイズは、通常 100k を超えません。つまり、LAION5B の 5×104 分の 1 です。これには、大規模なモデルが特定の問題についてトレーニングされる場合に、過剰適合を回避し、一般化を維持するための堅牢なニューラルネットワークトレーニング方法が必要になります。

第 2 に、画像処理タスクにデータ駆動型ソリューションを使用する場合、大規模なコンピューティングクラスターが常に利用できるとは限りません。このため、許容可能な時間とメモリ空間内で特定のタスク (個人用デバイスなど) に合わせて大規模なモデルを最適化するには、高速トレーニング方法が重要になります。これには、事前にトレーニングされた重みを利用するだけでなく、戦略や転移学習を微調整することも必要です。

第三に、さまざまな画像処理問題には、問題定義、ユーザー制御、または画像注釈の形式が異なります。これらの問題に対処する際、画像拡散アルゴリズムは、ノイズ除去プロセスの制約、多頭注意アクティベーションの編集などの「手続き型」方法で調整できますが、これらの手作りのルール動作は主に人間の指示によって決まります。深度から画像への変換、ポーズから人間へのポーズなど、いくつかの特定のタスクを考慮すると、これらの問題は基本的に生の入力を解釈してオブジェクトレベルまたはシーンレベルで理解する必要があり、手作りの手続き型アプローチは実現可能性が低くなります。多くのタスクにわたる学習を可能にするソリューションは不可欠であり、エンドツーエンドの学習が必要です。

2. 全体論的思考

大規模モデルを支援する別のパスを使用した大規模モデル微調整のパラダイム。主なアイデアは、大規模モデルを修正し、大規模モデル内のデノイザーの構造の一部をコピーし、これらの構造をさまざまなタスク用にトレーニングすることです。条件一部のコンテンツの生成を制御できるエッジなどを使用します。

4. 方法

ControlNet は、事前トレーニングされた画像拡散モデルをタスク固有の条件で拡張するニューラルネットワークアーキテクチャです。セクション 4.1 で ControlNet の基本構造を紹介し、各部分の動機を説明します。セクション 4.2 では、安定拡散を例として、ControlNet を画像拡散モデルに適用する方法を詳しく説明します。セクション 4.3 では、学習目標と一般的なトレーニング方法について詳しく説明し、セクション 4.4 では、1 台のラップトップでのトレーニングや大規模なコンピューティングクラスターの使用など、極端な場合のトレーニングを改善するいくつかの方法について説明します。最後に、セクション 4.5 では、異なる入力条件を備えたいくつかの ControlNet 実装の詳細を説明します。

4.1 コントロールネット

ControlNet はさらに、ニューラルネットワークブロックの入力条件を操作することで、ニューラルネットワーク全体の動作を制御します。ここで、「ネットワークブロック」とは、「resnet」ブロック、「conv-bn-relu」ブロック、「マルチヘッドアテンション」など、ニューラルネットワークを構築するために頻繁に使用される単位としてグループ化されたニューラル層のグループを指します。ブロック、トランスブロック待ち。2D 特徴を例として、特徴マップ $R^{h×w×c} を指定すると、$ ，其中 ${h, w, c}$ 高さ、幅、チャネル数を表し、1 つは一連のパラメータ $Θ$ ニューラルネットワークブロック $F (\cdot; Θ)$ を入れてください $x$ は別の特徴マップ $y$ 形状を持ちます。 $y = F (x; Θ)$ 。図に示すように:
ここに画像の説明を挿入

零卷积是权重初始化为0

します $Θ$ のすべてのパラメータはロックされ、トレーニング可能なコピー $Θ_c$ 。コピー $Θ_c$ 外部条件ベクトル $c$ 一緒にトレーニングする。このペーパーでは、元のパラメータと新しいパラメータを「ロックされたコピー」および「トレーニング可能なコピー」と呼びます。このレプリケーションを採用する動機は、データセットが小さい場合の過剰適合を回避し、数十億の画像から学習された大規模なモデルで製品として使用可能な品質を維持することです。ニューラルネットワークブロックは、「ゼロ畳み込み」と呼ばれる独自のタイプの畳み込み層を介して接続されます。ゼロで初期化された $1 \times 1 つの$ 畳み込み層 (重みとバイアスを含む)。を使用します。 $Z (\cdot;\cdot)$ はゼロ畳み込み演算を表し、2 つのパラメーターインスタンス $Θ_{z1}、Θ_{z2} を$ ControlNet 構造を形成します。0 畳み込みを適用する場合、トレーニングの最初のステップでは、ロックされたレプリカからの情報が完全に使用されます。最初のトレーニングステップでは、ニューラルネットワークブロックのトレーニング可能でロックされたコピーのすべての入力と出力は、ControlNet がない場合の入力と出力と一致します。言い換えれば、最適化の前に ControlNet がいくつかのニューラルネットワークブロックに適用されるまでは、ディープニューラル機能には何の影響もありません。ニューラルネットワークブロックのパワー、機能、およびその結果として得られる品質は完全に保持され、さらなる最適化は (レイヤーを最初からトレーニングする場合と比較して) 微調整と同じくらい速くなります。

4.2 画像拡散モデルにおける ControlNet

例として安定拡散 [44] を使用して、ControlNet を使用してタスク固有の条件で大規模な拡散モデルを制御する方法を紹介します。Stable Diffusion は、数十億の画像でトレーニングされた大規模なテキストから画像への拡散モデルです。このモデルは本質的に、エンコーダー、中間ブロック、スキップ接続デコーダーで構成される U-Net です。エンコーダーとデコーダーの両方に 12 ブロックがあり、完全なモデルには 25 ブロック (中間ブロックを含む) が含まれています。これらのブロックのうち、8 ブロックはダウンサンプリングまたはアップサンプリングされた畳み込み層であり、17 ブロックはメインブロックであり、それぞれに 4 つの ResNet 層と 2 つのビジョントランスフォーマー (ViT) が含まれています。各 ViT には、いくつかの相互注意および/または自己注意メカニズムが含まれています。テキストは OpenAI CLIP を使用してエンコードされ、拡散時間ステップは位置エンコードを使用してエンコードされます。
ここに画像の説明を挿入

Stable Diffusion は、VQ-GAN のような前処理方法を使用して、安定化トレーニングのために 512 × 512 の画像データセット全体をより小さな 64 × 64 の「潜在画像」に変換します。これには、たたみ込みのサイズに一致させるために、ControlNet を使用して画像ベースの調整を 64×64 の特徴空間に変換する必要があります。私たちは小規模なネットワークを使用します $E (\cdot)$ カーネルサイズ $4 \times 4$ 、ステップサイズは $2 \times 2$ (ReLU によってアクティブ化され、チャネル数は 16、32、64、128、ガウス重みで初期化され、完全なモデルで共同トレーニングされます)、画像空間の条件 ci $c_i$ のサイズの特徴マップとしてエンコードされる $64 \times 64$ ．具体的には、ControlNet を使用して、Stable Diffusion の 12 個のエンコードされたブロックと 1 個の中間ブロックのトレーニング可能なレプリカを作成します。これらの 12 ブロックは 4 つの解像度 (64×64、32×32、16×16、8×8) に分散されており、各解像度には 3 つのブロックがあります。出力は、U-net の 12 個のスキップ接続と 1 個の中間ブロックに追加されます。SD は典型的な U-net 構造であるため、この ControlNet アーキテクチャは他の普及モデルでも同様に使用できる可能性があります。

4.3 トレーニング

与えられた画像 $z_0$ 、拡散アルゴリズムにより画像に徐々にノイズが追加され、ノイズの多い画像が生成されます $z_t$ 、ここで $tは$ ノイズを加える回数を表します。いつ $t$ が十分に大きい場合、画像は純粋なノイズに近づきます。タイムステップ $t$ 、テキストプロンプト $c_t$ およびタスク固有の条件 $c_f$ 、画像拡散アルゴリズムはネットワークノイズ画像への追加を予測するための $θ$ $z_t$ のノイズ:
$\mathcal{L} = \mathbb{E}_{z_0,t,c_t,c_f,\epsilon \sim N(0,1)}[||\epsilon-\epsilon_{\theta}(z_t,t,c_t,c_f)|| ^ 2_2]$
トレーニング中に、テキストキュー ct の 50% を空の文字列にランダムに置き換えます。これは、ControlNet が Canny エッジマップや人工的な落書きなどの入力条件付きマップからセマンティックコンテンツを識別するのに役立ちます。これは主に、SD モデルがヒントを認識しない場合に、エンコーダーがヒントの代わりに入力制御グラフからより多くのセマンティクスを学習する傾向があるためです。

4.4 トレーニングの改善

特に、コンピューティングデバイスが非常に制限されている (例: ラップトップ上)、または非常に強力である (例: 大規模な GPU を備えたコンピューティングクラスター上) という極端なケースにおいて、ControlNet トレーニングを改善するためのいくつかの戦略について説明します。)。私たちの実験でこれらの戦略が使用される場合は、実験設定で言及します。

小規模トレーニング: コンピューティングデバイスが制限されている場合、ControlNet と Stable Diffusion を部分的に切断すると、収束が速くなることがわかりました。図 3 に示すように、デフォルトでは、ControlNet を「SD Middle Block」および「SD Decoder Block 1、2、3、4」に接続します。デコーダー 1、2、3、4 を切断し、中間ブロックのみを接続すると、トレーニングが約 1.6 倍高速化できることがわかりました (RTX 3070TI ノートブック GPU でテスト)。モデルが結果と条件の間の妥当な関連性を示した場合、これらの切断された接続を再接続してさらなるトレーニングを行うことで、正確な制御が可能になります。

大規模トレーニング: ここで、大規模トレーニングとは、強力なコンピューティングクラスター (少なくとも 8 つの Nvidia A100 80G または同等の GPU) と大規模なデータセット (少なくとも 100 万のトレーニング画像ペア) の両方がある状況を指します。これは通常、Canny によって検出されたエッジマップなど、データがすぐに利用できるタスクに適しています。この場合、過剰適合のリスクが比較的低いため、ControlNet を十分な反復 (通常は 50,000 ステップ以上) トレーニングしてから、安定拡散のすべての重みのロックを解除してモデル全体を共同トレーニングできます。これにより、より問題に特化したモデルが作成されます。

4.5 実装

当社では、さまざまな画像ベースの条件を使用して大規模な拡散モデルを制御する、いくつかの ControlNet 実装を提供しています。

Canny Edges: Canny エッジ検出器 (ランダムなしきい値を使用) を使用して、インターネットから 300 万のエッジ画像キャプションペアを取得しました。モデルは、Nvidia A100 80G で 600 GPU 時間トレーニングされました。ベースモデルはStable Diffusion 1.5です。
Canny Edges (変更): 上記の Canny Edge データセットの画像解像度を並べ替え、1,000、10,000、50,000、500,000 サンプルを含むサブセットをサンプリングしました。同じ実験設定を使用して、データセットサイズの影響をテストします。明らかに、トレーニングデータが大きいほど、効果は高くなります。
ハフライン: 学習ベースのディープハフ変換を使用して Places2 から直線を検出し、BLIP を使用してタイトルを生成します。600,000 のエッジ画像とキャプションのペアを取得します。上記の Canny モデルを開始チェックポイントとして使用し、Nvidia A100 80G で 150 GPU 時間トレーニングしました。
HED 境界: HED 境界検出を使用して、インターネットから 300 万のエッジ画像キャプションのペアを取得しました。モデルは、Nvidia A100 80G で 300 GPU 時間トレーニングされました。ベースモデルはStable Diffusion 1.5です。
ユーザースケッチ: HED 境界検出と一連の強力なデータ拡張手法 (ランダムなしきい値処理、スケッチの割合のランダムマスキング、ランダムな形態学的変換、およびランダムな非最大値抑制) を使用して、画像から人間の落書きを合成します。私たちは、インターネットから 500,000 の落書きと画像とキャプションのペアを取得しました。上記の Canny モデルを開始チェックポイントとして使用し、Nvidia A100 80G で 150 GPU 時間トレーニングしました。より「人間らしい」合成方法も試しましたが、この方法は単純な HED 方法よりもはるかに遅く、大幅な改善は見られませんでした。
人間のポーズ: 学習ベースの姿勢推定方法を使用して、単純なルールを使用してインターネットから人間を「検索」します。人体を含む画像は、体全体のキーポイントの少なくとも 30% を検出する必要があります。ポーズ・画像・キャプションのペアを8万件取得しました。人間の骨格を含む可視化されたポーズ画像を学習条件として直接使用していることに注意してください。モデルは、Nvidia RTX 3090TI で 400 GPU 時間トレーニングされました。ベースモデルはStable Diffusion 2.1です。
人間のポーズ: 同じルールを使用して、上記の Openpifpaf 設定でインターネット上の人間を見つけるために、学習ベースの姿勢推定方法を使用します。200,000 のポーズ画像とキャプションのペアを取得します。人間の骨格を含む可視化されたポーズ画像を学習条件として直接使用していることに注意してください。モデルは、Nvidia A100 80G で 300 GPU 時間トレーニングされました。その他の設定は上記の Openpifpaf と同じです。