かつてFluxの技術記事におけるREPAを整理したことがある。
REPAは外部のCLIPモデルが必要という欠点があり、記事内ではその代替手段を示唆されていた。その代替手段は不明だったが、Self-Flowの技術が公開されたので読んでみる。
REPA
REPAは以下のような感じで外部のCLIP Encoderが必要であるが、2通りのlatentのcos類似度損失によってDiTの一部にToken間のCLIPの意味抽出を学習させる。
Self-Flow
反面、Self-Flowは以下のような感じである。REPAと似るが外部のCLIP Encoderが不要である。
REPAと違ってこれでDiTの一部(最初の8層)に意味抽出が学習できる根拠は何もなく感じる。
DiT深さの差、MLPの有無は以下の式から推測される。
DiT深さがD=28なら$28*0.3=8.4, 28*0.7=19.6$でl=8、k=20が求まる。
random MASKとtoken相関
Self Flowではt>sであればランダムに画面の5割以下により強いノイズを付与する。
一般的にLatent DiffusionはVAEによってVAEの中間層の分布はガウス分布であるlatentにガウスノイズを加える。ガウス分布にガウス分布を加えてもガウス分布のままだからこのようなDiTは一般にガウス超球面上の移動を学習する。
一方、VAEの代わりにVQVAEを使用した場合、latentは離散的なindexになる。この場合、拡散過程として離散的なindexにノイズを課す事はできないので、この状況において拡散過程はindexの値に徐々にランダムにMASKを掛けていくことになる。一般にこのようなMASKを一個ずつ取り除いくモデルの事を自己回帰モデルという。
画像生成モデルでの回帰モデルとしては例えばPartiが知られる。
一方、LLMにおいては学習データ(テキスト)にMASKを掛けこのMASK内の離散indexを予測する手法はBERT(2018)から普遍的に見られる。つまり、LLMにおいてはMASKを一個一個予測していくという事の方が普通である。
一般に入力の離散indexはVQVAEではcodebookによってガウス分布に変換される。このため一般的なLLMのTransformersにおいてもガウス分布(MASK)からガウス分布の移動を学習しているに過ぎない。
重要なのはこのrandom MASKの予測についてはtoken間の相関を強く学習している点にある。
つまり、Self Flowがtoken間の相関(意味抽出)をDiTの最初の8層に学習できるREPAと等価に見なせるのはランダムMASKのおかげなのか、というのが自分の考察である。
メリット
Self FlowのメリットはCLIP Encoderを用意しなくていい点ですがTextEncoder(VLM)にVisionEncoderが含まれることを考えると、取り立てて凄いメリットとは言えないように思いました。
これだけなら別にREPAでよくね、と思ってましたがCLIP Encoderを使用せずとも学習できるメリットは意味抽出EncoderのないaudioやvideoでもDiTの部分学習が可能な点があります。
逆にデメリットとしてはファインチューニングコードが難しく、単純な拡散モデルのlossによるファインチューニングでは破滅的忘却が起きやすくなるのでは？と思いました。