論文解説 : Stretching Each Dollar: Diffusion Training from Scratch ona Micro-Budget

arxiv : https://arxiv.org/abs/2407.15811
github : https://github.com/SonyResearch/micro_diffusion

ひとことまとめ

実写・合成あわせて3700万枚の画像を用い30万円(1,890ドル)でSD v1.5相当のモデルを作成した

概要

Text-to-Imageモデルは年に10億枚以上の画像を生成するのに用いられているが、学習にかかるコストは依然として高いままである。例えばStableDiffusion 2.1では200,000 A100 GPU hoursと非常に高いコストがかかっている。最新の手法であっても18,000 A100 GPU hours程度であり、H100 8台を1か月使用する必要がある。これらを高速化する手法としてMaskDiTがあげられるが、画像のマスク率が高いと急激に性能が悪化する。そこで、マスク前にトークンをミックスすることで75%のマスク率でも学習でき、SD v1.5相当の学習を短時間で高速に行うことができる。

DiT(Diffusion Transformer)の高速化

diffusion transformerの総計算量は 画像数×モデルのパラメータ数×入力シークエンス数(DiTの画像1枚のパッチ数) で表現することができる。モデルのパラメータ数は様々な生成を行うためには必要であり、この研究としてはこの数値は変更しないものとする。そのため、入力シークエンスを減らすことで高速化を目指す

(b) : 大きいパッチサイズ
大きいパッチサイズは1トークンにより広い領域のデータを入力するため、高すぎると急激に性能が悪化することが知られている。

(c) : パッチマスキング
MAE(Masked Auto Encoder)のように入力パッチの一部をtransformerに入力しないことで学習の高速化を行う。UNetでは画像の小領域を学習に使うことは良く行われている。しかし不連続な領域を学習できるのが大きく異なり、この効率性の良さから画像・言語の両方で広く用いられている

(d) : MaskDiT
マスク領域の表現学習を促進させるため、MaskDiTはマスク領域の再構成損失を計算(MAEと同じ損失を追加した)することで学習を高速化させた。この手法は軽量なdecoderを大きなtransformerの後に追加することで総合的な計算量を軽量化している。損失関数は以下のようになる。

Alleviating performance bottleneck in patch masking using deferred mask

計算量を劇的に減らすため、75%などの大量のパッチをマスクすると極端に性能が悪化する。MaskDiTでも悪化する傾向は変わらない。

高いマスク率は画像の価値のある信号を取り除いており、我々はパッチ化直後ではなく、レイヤを挟み情報を拡散した後に入れれば良いと考えた。しかしこのレイヤは計算量があまり増加しないように設計すれば基盤(GPUとか)の制約は無視できる。実際、提案手法はMaskDiTとほぼ変わらない計算量である。
パッチが重複せずに切り取られる場合は、それぞれのパッチはほかのパッチの情報を持っていない。そこで、マスク前に画像全体の情報を埋め込むpatch-mixerを追加する。この研究では、patch-mixerは数レイヤのtransformerを用いているが、任意のニューラルネットワークを利用可能である。

Unmasked finetuning
高いマスク率はdiffusion modelの全体構造が学習しづらく、学習時とテスト時で分布が大きく異なる。この事前学習後に小さいマスク比率でモデルを再学習させることでこれを軽減する。MaskDiTではこの事前学習後に生成すると形状が不安定で、微調整は必要と述べられている。しかしこの研究では必ずしも必要とは限らず、大規模学習での生成にのみ使用する

MoEとlayer-wise scalingを使用したバックボーンの改善
MoEは計算量の劇的な増加を防ぎながら、パラメータ数と表現力を向上させることができる。追加の損失関数を必要としないexpert-choice routingをベースとしたシンプルなMoEを導入した。
また、近年通常のtransformerより性能が高いといわれるlayer-wise scalingを導入した。これは、transformerのwidth(隠れ層のサイズ)を層が深くなるほど増やす手法である。一般的にtransformerは層が深いほど複雑な処理を行う傾向があるので、深い層に大きいパラメータ数を与えることで表現力が増え、性能の向上が期待できる。

実験

この論文では

1. 軽量なモデルでマスク率による影響などの調査

2. 大型モデルで既存の手法との性能比較

を行っている。

1.軽量なモデルでマスク率による影響などの調査

マスク率を0~87.5%に変化させた場合のnaive masking(Figure2.c,MAEと同じようにパッチをそのまま落とす)とdeferred masking(Figure2.e, パッチを落とす前に数層のtransformer層をはさむ)での性能さを検証した。

提案手法は全指標でnaive maskingの性能を上回った。特にマスク率が上がるほど性能差が大きく、75%のマスク率ではnaive maskingではFIDが16.5であるが、提案手法は5.03とマスクなしの場合の3.79と非常に近いことがわかる。

マスクの形状とpatch-mixerによる性能差

大きなpatch mixerはより高い性能を示すが、計算量が増え、相対的に費用がかかってしまう。この研究では計算量を減らすため小さいモデルを使用した。

また、マスクしない部分をより連続的にした場合の性能を検証した。

しかしながら、block sizeが増加するほど性能が減少することがわかった。そのため、この研究ではランダムにマスクする方法をとっている。

DiTの構造変更による性能改善

Layer-wise scaling

レイヤのwidth(隠れ層の大きさ)が一定の場合と層が深くなるほど大きくなる場合(同じ計算量)で比較した。

layer-wise scalingのほうがすべての指標でwidthが一定の場合より上回った。

Mixture-of-experts layers

DiT-Tiny/2にMoEを用いた場合、Clip-scoreは28.11から28.66と若干向上したが、FIDは6.92から6.98と若干悪化した。これは小さい学習ステップ(60k)で複数のexpertを学習させるのが難しかったためと予測され、大型モデルの学習では性能が大幅に改善することが確認できた。

ベースライン手法と提案手法の比較

Comparing different masking strategies

まず、大きいパッチ(Figure 2.b)の場合と比較した。
パッチサイズを2倍の4とし、提案手法のマスク率を75%とした場合、以下の表の結果となった。

$${\begin{array}{|c|c|c|c|} \hline \text{Method} & \text{FID} & \text{Clip-FID} & \text{Clip-score} \\ \hline \text{large patch size} & 9.38 & 6.31 & 26.70 \\ \hline \text{提案手法} & 7.09 & 4.10 & 28.24 \\ \hline \end{array}}$$

これにより、提案手法はパッチサイズを大きくする場合よりも性能が向上することがわかった。

また、Figure.2の他の手法とも比較したが、naive maskingはかなり性能が悪化し、MaskDiTはわずかな改善にとどまった。提案手法のpatch-mixerはマスクあり学習とマスクなし学習のギャップを大きく埋める手段であることが確認できた。

IsoFLOPs study - Why bother using masking instead of smaller transformers?

transformerを軽量化することで計算コストを削減することができるが、マスクあり学習と同じFLOPsでどちらのほうが性能が高いか調査した。

75%以下のマスク率の場合、多くの指標で提案手法が良い精度を表しているが、高いマスク率の場合性能が悪化している。これは高いマスク率の場合情報が少なすぎるためモデルの学習に影響があることが考えられる。

Deferred masking as pretraining + unmasked finetuning

提案手法は強い事前学習として考えることもでき、この後にマスクなしでファインチューニングすると、直接マスクなしで学習する場合より同じ計算量でより高い性能を実現できると分かった。
まず、75%のマスク率で学習し、その後ファインチューニングをした。この2つの学習の合計のisoflops(計算量)がマスクなし学習と同じになるように調整した。

グラフのように、マスクありで事前学習するとマスクなしで学習する場合よりも高い性能を示していることがわかる。

2.大型モデルで既存の手法との性能比較

Micro-budget Training of Large-scale Models

ここでは実際に使われているtext-to-imageモデルと比較した。画像は実際と合成の合計3700万枚を使用した。

学習プロセス

DiT-XL/2を用いて8エキスパートのMoEで学習した。学習は以下のような表を上から順に実行した。(512x512の学習は256x256の学習後の重みを利用した)

定量評価

提案手法は学習コストが非常に安く、学習使用画像数が少ないが、Stable-Diffusion-1.5に近い性能を示していることがわかる。

GenEvalでは、多くの設定でStable-Diffusion-1.5の性能を超えていることがわかる。

定性評価

提案手法はStable-Diffusion-1.5かそれ以上の生成品質を実現している。(論文には記載されていないが、合成データセットでMidjournyやStableDiffusionで生成した画像を使っているためアート系に強いのでは?)

まとめ

• DiTを高速に学習するマスキング手法を提案

• patch-mixerを導入することで75%などの高いマスク率での学習を安定化

• 提案手法は非常に少ない計算量でありながら、Stable-diffusion-1.5に近い性能を実現した