1,890 ドルで、12 億のパラメータ拡散モデルをゼロからトレーニングできます。

わずか 1,890 米ドルと 3,700 万枚の画像があれば、かなり優れた拡散モデルをトレーニングできます。

現在、視覚生成モデルはリアルな視覚コンテンツの作成に優れていますが、これらのモデルをゼロからトレーニングするコストと労力は依然として高くなります。たとえば、Stable Diffusion 2.1 には、A100 GPU で 200,000 時間かかりました。研究者が最も先進的な方法を使用したとしても、8×H100 GPU でのトレーニングには 1 か月以上かかります。

さらに、大規模なモデルをトレーニングすると、データセットにも課題が生じます。これらのデータは基本的に数億単位であるため、モデルのトレーニングにも課題が生じます。

高いトレーニングコストとデータセット要件により、大規模な普及モデルの開発には乗り越えられない障害が生じます。

現在、Sony AI およびその他の機関の研究者は、11 億 6,000 万のパラメーターを持つ優れた拡散モデル、つまりスパース変換器をトレーニングするためにわずか 1,890 ドルを費やしました。

• 論文アドレス: https://arxiv.org/pdf/2407.15811

• 論文タイトル: 1 ドルを伸ばす: マイクロ予算でのゼロからの普及トレーニング

• プロジェクト (近日公開予定) : https://github.com/SonyResearch/micro_diffusion

具体的には、この研究では、著者はテキストから画像への拡散モデルのための低コストのエンドツーエンドのパイプラインを開発し、SOTAよりもトレーニングコストを低くしています。このモデルは、数十億のトレーニング画像や独自のデータセットにアクセスする必要がなく、一桁以上高速です。

著者は、テキストから画像生成へのビジュアルトランスフォーマーに基づく潜在的な拡散モデルを検討しました。これは主に、この方法が設計が簡単で広く使用されているためです。計算コストを削減するために、著者らは、変換器の計算コストが入力シーケンス サイズ (つまり、画像あたりのパッチの数) に大きく依存することを利用しています。

この記事の主な目的は、トレーニング プロセス中に各画像を処理するためにトランスフォーマーによって使用される有効なパッチの数を減らすことです。これは、トランスフォーマーの入力層でいくつかのトークンをランダムにマスクすることで簡単に実現できます。

ただし、既存のマスキング方法では、特に入力パッチの大部分が拡散トランスによってまったく観察されない高いマスキング レートの場合、パフォーマンスを大幅に低下させることなくマスキング レートを 50% を超えて拡張することはできません。

マスキングによる大幅なパフォーマンスの低下を軽減するために、著者は遅延マスキング戦略を提案します。この戦略では、すべてのパッチが軽量のパッチ ミキサー (パッチ ミキサー) によって前処理されてから、拡散トランスフォーマーに転送されます。パッチ ミキサーには、拡散トランスにあるパラメーターの数の一部が含まれています。

単純なマスキング手法と比較して、パッチ混合後のマスキングにより、マスクされていないパッチが画像全体に関するセマンティック情報を保持できるようになり、既存の最良の手法と競合しながら、非常に高いマスキング レートで拡散変換器の信頼性の高いトレーニングが可能になります。高度なマスキングへ。

著者らはまた、同じ計算予算の下で、遅延マスキング戦略がダウンサイジング (モデル サイズの縮小) よりも優れたパフォーマンスを達成することを実証しています。最後に、著者らは、レイヤーごとのスケーリング、MoE を使用したスパース Transformer など、Transformer アーキテクチャの最近の進歩を組み込んで、大規模なトレーニングのパフォーマンスを向上させています。

著者が提案した低コストのトレーニング パイプラインにより、実験のオーバーヘッドが削減されます。実際の画像を使用することに加えて、著者らはトレーニング データセット内で他の合成画像を組み合わせることも検討しました。結合されたデータセットには 3,700 万枚の画像のみが含まれており、既存の大規模モデルのほとんどが必要とするデータよりもはるかに少ないデータです。

この結合されたデータセットで、著者は 1890 ドルのコストで 11 億 6,000 万パラメータのスパース トランスフォーマーをトレーニングし、COCO データセットでのゼロショット生成で 12.7 FID を達成しました。

この記事でトレーニングされたモデルは、安定した拡散モデルのわずか 1/118、現在の最先端の方法 (コスト計算) の 1/1 のコストでありながら、競争力のある FID と高品質の生成を実現していることは注目に値します。 28,400ドル) 15.

方法の紹介

計算コストを大幅に削減するために、パッチマスキングではバックボーントランスフォーマーに入力する前に入力パッチの大部分を破棄する必要があるため、トランスフォーマーはマスクされたパッチの情報を取得できません。マスキング率が高いと (マスキング率 75% など)、トランスの全体的なパフォーマンスが大幅に低下する可能性があります。 MaskDiT を使用した場合でも、この方法では入力層自体の画像パッチのほとんどが破棄されるため、単純なマスキングに比べて弱い改善しか観察できません。

すべてのパッチのセマンティック情報を保持する遅延マスキング

マスキング率が高いと、画像内の貴重な学習信号のほとんどが削除されるため、作成者は、入力層でマスクする必要があるのか​​と尋ねずにはいられません。計算コストが一定である限り、これは単なる設計上の選択であり、根本的な制限ではありません。実際、著者らは、既存の MaskDiT 手法とほぼ同じコストで、大幅に優れたマスキング戦略を発見しました。パッチは拡散トランスフォーマー内の重複しない画像領域から取得されるため、各パッチの埋め込みによって画像内の他のパッチからの情報は埋め込まれません。したがって、著者らは、マスクされていないパッチが画像全体から情報を埋め込めるように、マスクする前にパッチの埋め込みを前処理することを目指しています。彼らは前処理モジュールをパッチミキサーと呼びます。

パッチミキサーを使用して拡散トランスフォーマーをトレーニングする

著者は、パッチミキサーとは、個々のパッチ埋め込みを融合できるニューラル アーキテクチャであると考えています。変圧器モデルでは、この目標はアテンション層とフィードフォワード層の組み合わせによって自然に達成できます。そこで、著者らはパッチミキサーとして数層のみからなる軽量トランスフォーマーを使用しています。入力シーケンス トークンがパッチ ミキサーによって処理された後、それらはマスクされます (図 2e)。

^{図 2: パッチ シーケンスを圧縮して計算コストを削減します。拡散変換器のトレーニング コストはシーケンス サイズ (つまり、パッチの数) に比例するため、パフォーマンスを低下させることなくシーケンス サイズを減らすことが最善です。これは、b) より大きなパッチを使用する、c) パッチの一部をランダムにマスクする、または d) 単純なマスキングと追加の自動エンコーディング目標を組み合わせた MaskDiT を使用することによって実現できます。著者らは、3 つの方法すべてが、特にマスキング レートが高い場合に、画像生成パフォーマンスの大幅な低下につながることを発見しました。この問題を軽減するために、彼らは、パッチミキサーによる処理後にパッチをマスクする単純な遅延マスキング戦略を提案しました。彼らのアプローチは、パッチミキサーを使用することを除いて、すべての点で単純なマスキングに似ています。 MaskDiT と比較すると、彼らの方法は代理目的の最適化を必要とせず、計算コストは​​ほぼ同じです。}

マスクがバイナリマスク m であると仮定すると、作成者は次の損失関数を使用してモデルをトレーニングします:

ここで、M_ϕ はパッチミキサー モデル、F_θ はバックボーン トランスフォーマーです。 MaskDiT と比較して、提案された方法は全体の設計も簡素化し、追加の損失関数やトレーニング中に 2 つの損失間の対応するハイパーパラメータ調整を必要としないことに注意してください。推論中、このメソッドはパッチをマスクしません。

マスクされていない微調整

マスキング率が非常に高いと、拡散モデルが画像の全体的な構造を学習し、シーケンスサイズにトレーニングテストの分布シフトを導入する能力が大幅に低下するため、著者は次のことを検討しています。マスクされていない微調整のマスクされた事前トレーニングの後に少量。微調整により、パッチ マスキングの使用によって発生する生成アーティファクトを軽減することもできます。したがって、以前の作業では、特にサンプリングで分類子を使用しないブートストラップを使用する場合、マスキングによって急激に低下したパフォーマンスを回復することが重要でした。しかし、著者らは、マスクされた事前トレーニングを使用した場合でも、彼らの方法はベースラインのマスクされていない事前トレーニングと同等のパフォーマンスを達成するため、これが完全に必要なわけではないと主張しています。著者らは、このアプローチを大規模なトレーニングでのみ使用して、高度なパッチ マスキングによる未知-未知の生成アーティファクトを軽減します。

MoE とレイヤーごとのスケーリングを使用してバックボーン変圧器アーキテクチャを改善します

著者はまた、変圧器アーキテクチャ設計の革新を活用して、計算上の制約の下でモデルのパフォーマンスを向上させています。

彼らはハイブリッド エキスパート レイヤーを使用します。これは、トレーニング コストを大幅に増加させることなく、モデルのパラメーターと表現力を向上させるためです。彼らは、エキスパート選択ルーティングに基づいた簡素化された MoE レイヤを使用します。各エキスパートは、エキスパート間の負荷のバランスを取るための追加の補助損失関数を必要としないため、どのトークンをルーティングするかを各エキスパートが決定します。彼らはまた、層ごとのスケーリングも考慮しました。これは、大規模な言語モデルで一般的なトランスフォーマーよりも優れたパフォーマンスを発揮することが最近示されました。この方法では、変換ブロックの幅、つまり注目層とフィードフォワード層の隠れ層の寸法が線形に増加します。したがって、ネットワークのより深い層には、以前の層よりも多くのパラメータが割り当てられます。著者らは、ビジュアル モデルのより深い層ほど複雑な機能を学習する傾向があるため、より深い層でより高いパラメーターを使用するとパフォーマンスが向上すると考えています。著者らは、提案する拡散トランスの全体的なアーキテクチャを図 3 に示します。

^{図 3: この記事で提案する拡散トランスの全体的なアーキテクチャ。著者らは、入力画像内のすべてのパッチをマスクされる前に処理する軽量のパッチ ミキサーをバックボーン トランスフォーマー モデルに追加しました。現在の研究に続いて、著者らはアテンション レイヤーを使用してキャプションの埋め込みを処理し、それを調整に使用します。正弦波埋め込みを使用して時間ステップを表します。彼らのモデルはマスクされていないパッチのみをノイズ除去するため、拡散損失 (論文の式 3) はこれらのパッチについてのみ計算されます。彼らはバックボーントランスを修正して、個々のレイヤーでレイヤーごとのスケーリングを使用し、交互のトランスブロックでエキスパートレイヤーを混合しました。}

実験

この実験では、拡散トランス (DiT) の 2 つのバリエーション、DiT-Tiny/2 と DiT-Xl/2 を使用します。

図 4 に示すように、遅延マスキング手法は複数のメトリクスで優れたパフォーマンスを実現します。さらに、マスキング率が増加すると、パフォーマンスの差が広がります。たとえば、75% のマスキング率では、ナイーブ マスキングにより FID スコアが 16.5 に低下します (低いほど良い) が、私たちの方法では 5.03 を達成でき、これはマスキングなしの FID スコア 3.79 に近くなります。

表 1 は、層ごとのスケーリング手法が拡散変換器のマスキング トレーニングにおいてより優れたフィッティング効果があることを示しています。

さまざまなマスキング戦略を比較します。著者らはまず、私たちの方法をより大きなパッチを使用した戦略と比較します。パッチ サイズを 2 から 4 に増やすと、75% のパッチ マスキングに相当します。遅延マスキングと比較すると、他の方法のパフォーマンスは低く、それぞれ 9.38、6.31、26.70 FID、Clip-FID、Clip-score しか達成できません。比較すると、レイテンシ マスキングでは、それぞれ 7.09、4.10、および 28.24 FID、クリップ FID、およびクリップ スコアが達成されます。

下の図は、遅延マスキングとトレーニング コストを削減するためのモデル縮小の比較を示しています。著者らは、マスキング率が 75% に達するまでは、遅延マスキングが 3 つの指標のうち少なくとも 2 つにおいてネットワーク削減よりも優れていることを発見しました。ただし、マスキング レートが非常に高い場合、マスキングが遅れるとパフォーマンスが低下する傾向があります。これは、これらの比率でマスクされる情報の損失が高すぎることが原因である可能性があります。

表 5 に、モデル トレーニングのハイパーパラメーターの詳細を示します。トレーニングのプロセスは 2 つの段階に分かれています。

コストを計算します。表 2 は、トレーニング FLOP と経済的コストを含む、各トレーニング段階の計算コストの内訳を示しています。フェーズ 1 とフェーズ 2 のトレーニングでは、それぞれ総計算コストの 56% と 44% が消費されました。 8×H100 GPU クラスター上のモデルの合計クロック トレーニング時間は 2.6 日で、これは 8×A100 GPU クラスターの 6.6 日と同等です。

その他の結果については、元の論文を参照してください。

以上が1,890 ドルで、12 億のパラメータ拡散モデルをゼロからトレーニングできます。の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。