[2204.06125] Hierarchical Text-Conditional Image Generation with CLIP Latents

yoheikikuta commented 2 years ago

論文リンク

公開日（yyyy/mm/dd）

2022/04/13

概要

DALL·E 2 の論文。論文ではこのモデルは CLIP の画像特徴量ベクトルから画像を生成するものなので unCLIP と命名されているが、その後の展開を見るに DALL·E 2 と呼ばれていることがほとんどであると思う。 CLIP のテキスト特徴量から prior モデルで CLIP の画像特徴量ベクトルを作り、その画像特徴量ベクトルから decoder モデルで画像を生成するというモデルになっている。 diffusion モデルがふんだんに使われており、prior, decoder の画像特徴量ベクトルから 64x64 の画像を作るところ, decoder で 64x64 → 256x256 と upsample するところ、decoder で 256x256 → 1024x1024 と upsample するところ、と合計 4 つの diffusion model が使われている。その性能は高く、テキストから生成した画像はかなりの質の高さである。ただ論文の記述があまり詳しくなく、公式実装もないので具体的な理解が難しいところが多い。

公式実装はないけど生成した画像は公開されてる: https://github.com/openai/dalle-2-preview プロジェクトページ: https://openai.com/dall-e-2/ 非公式実装: https://github.com/lucidrains/DALLE2-pytorch

yoheikikuta commented 2 years ago

DALL·E が最初出たときにかなりびっくりしたが、その後 GLIDE https://github.com/yoheikikuta/paper-reading/issues/63 などが出てきてテキストからの画像生成がさらに発展し、その後も勢いは留まることを知らずにこの DALL·E 2 や Imagen などが出てきた。この一年くらいでテキストからの画像生成が壁を超えて一気に性能向上した感がある。

このトピックは興味深くて論文をフォローしていたので、DALL·E 2 についても読んでみる。

yoheikikuta commented 2 years ago

主たる道具立てである CLIP や diffusion model に関してはこの前身とも言える GLIDE と変わりはない。どちらも OpenAI から出てる論文で、著者もかなり被ってるので、現状の道具で quality の高い自然言語からの画像生成をするためにアーキテクチャを考案したり技術的な工夫（高解像度の画像を生成するために upsampler diffusion model を準備したり）を詰め込んだ、という感じ。

論文を眺めと細かいところはあまり書いてない感じがしてやや気になるが、とりあえず読んでいく。

yoheikikuta commented 2 years ago

道具立ては先行研究で出揃ってるので、以下のアーキテクチャ図を理解するのが目的となる。

詳細に入る前にまずはざっくりでの理解。

点線の上の部分は CLIP の学習を表現している。これは知ってるので詳しく追う必要はないが、改めて思い出しておくと、画像とテキストのペアから text, image それぞれの特徴量を同じ次元で学習する。このそれぞれの特徴量ベクトルはペアのもの（例えば text が a photo of dog で image が犬の画像）のときに類似度が最大になるように学習されるもので、つまり同じようなベクトルが得られるようになっている。図でいうと、青色ベースのベクトルは CLIP のテキストの特徴量で茶色ベースのベクトルは CLIP の画像の特徴量を表している。

点線の下の部分は text-to-image の生成過程を表現している。まずは生成した画像のテキスト（a corgi playing a flame throwing trumpet）から学習した CLIP でテキストの特徴量を抽出する。次にこのテキストの特徴量を prior モデルによって CLIP の画像の特徴量ベクトルへと変換する。prior としては auto-regressive model と diffusion model のどちらかを採用する。diffusion model であればここはテキスト情報を condition にした conditional diffusion model を使うことによって、与えられたテキストにマッチするような画像の特徴量ベクトルを作ることになる。これは CLIP 特徴量ペアがあれば学習可能であり、学習したものを inference 時に使う。最後に生成した画像特徴量から diffusion upsampler model （図には描かれていないが、64x64 → 256x256 と 256x256 → 1024x1024 の二つの upsampler models を使用）で高解像度の画像を生成している。この部分を decoder と呼んでいる。これも仕組みは同じで画像情報を condition にした conditional diffusion model を使っており、これは元画像と CLIP の画像特徴量があれば学習できるので、学習したものを inference 時に使う。

画像生成部分は CLIP の画像特徴量ベクトルをから生成するので、このアーキテクチャ自体は unCLIP と名付けられている。ただ DALL·E の後継ということで DALL·E 2 と呼ばれていることが多い気がする。

yoheikikuta commented 2 years ago

上記の prior と decoder がどういうものかを詳しくみる前に、モデルとして conditional diffusion model を知っておく必要がある。これは paper-reading に書いてなかったのでここに記載しておくことにする。ここで書く内容は https://arxiv.org/abs/2105.05233 の Appendix H そのものである。

まずは記号の定義として、$x{0} ... x{T}$ が元画像からノイズ画像を timestep 毎に表現したもので、$y$ をラベル（例えば犬とか）とする。ここで $y$ をラベルという離散的なものにしているが、これは別に離散的なものに限る必要はなく分散表現などでもよい。 diffusion model ではよく Markovian で noise を加えるモデルを $q$ と書くが、これの conditional バージョンを $\hat{q}$ と書く。

やりたいことはこの conditional バージョンの場合で遷移確率を（利用可能な形で）定式化することである。

最初に以下のように仮定する。一つ目はラベルという情報を無視して単に元画像 $x_0$ の分布を考えれば $\hat{q}$ と $q$ で違いがないことを意味しているので自然である。二つ目は元画像 given でラベルが得られる確率であるが、これはラベル情報を知っているという前提であれば単に教師ラベルの情報であることを意味しているので自然である。三つ目は timestep を一つ進めるつまりは gaussian noise を加えるときに、ラベルで条件づけられたものとそうでないものは同一であることを意味している。これはノイズの乗せ方は context によらずに同一であるということで、これも自然である。コンテキストによってノイズの乗せ方を変えるということも原理的にはできるが、徒に複雑になってしまうので、最初に考えるのはこれでよいだろう。四つ目は Markovian process と仮定しているのでそれを意味している。

$$ \begin{align} \hat{q}(x_0) &:= q(x_0) \ \hat{q}(y| x0) &:= \text{Known labels per sample} \ \hat{q}(x{t+1}| xt, y) &:= q(x{t+1}| xt) \ \hat{q}(x{1:T}| x0, y) &:= \prod{t = 1}^{T} \hat{q}(x{t}| x{t-1}, y) \ \end{align} $$

ここまでは自然であり、ここから出発して reverse process の遷移確率がどのように書けるかを明らかにしていく。

まずは $\hat{q}$ が y で条件づけられていないときに $q$ と同じ振る舞いをすることを示す。これは上記の仮定を使って計算すれば示せる。これは joint distribution でも $\hat{q}(x_{1...T} | x0) = q(x{1...T} | x_0)$ でも straightforward な計算で示せる。

これを使って timestep t での確率分布も同様であることが示せる。

これらを使うことで、 $\hat{q} (x{t} | x{t+1}) = \frac{ \hat{q} (x{t+1} | x{t}) \hat{q} (xt) }{\hat{q} (x{t+1})} = \frac{ q (x{t+1} | x{t}) q (xt) }{q (x{t+1})} = q (x{t} | x{t+1}) $ と unconditional な reverse process では $\hat{q}$ と $q$ は同一であることが示される。

いよいよ conditional な reverse process で、これは straightforward な計算により以下のようになる。

分母は $x_t$ に依らないので定数と考えることができる。これは unconditional な遷移確率にさらに label に依存する $\hat{q} (y | x_t)$ を掛ける形になっている。つまり新たにノイズあり画像 $x_t$ を入力にしてラベルを予測するモデルがあればこの遷移確率を表現できることになる。

これはノイズあり画像からラベルを予測する classifier を別に準備すればよい（ノイズあり画像をインプットとするため典型的な pre-trained モデルが使えないという問題はある）。ラベル情報を取り込む他の方法としては diffusion model 自体をラベル情報ありなしで学習してそれを用いる classifier free guidance などの手法がある。これは diffusion model を学習するときにラベルの情報をトークンとして含めて conditional なものを扱えるモデルにしつつ、ゼロ埋めしたものも unconditional の場合として学習データに入れることで conditonal/unconditional な場合を合わせて学習し、サンプリング時にどちらに重きを置くかを重みづけて利用する。

yoheikikuta commented 2 years ago

conditional diffusion model が理解できていれば話は早い。この論文読むときというか最近の画像生成の話ではこれを理解しておくのは恐らく前提条件だと思うが、自分はちゃんと理解してなかったのでこのタイミングでフォローした、という状態である。

何はともあれアーキテクチャの中身をもう少し詳しく見ていく。

記号の定義として、画像 $x$ に対して CLIP で抽出した画像特徴量ベクトルを $z_i$、ペアとなるテキスト $y$ に対して CLIP で抽出したテキスト特徴量ベクトルを $z_t$ とする。

$P(x|y) = P(x,z_i | y) = P(x | z_i,y) P(z_i | y)$ である（$x$ が一つ決まれば $z_i$ は deterministic に決まるので最初の等号は成り立つ）。この prior $P(z_i | y)$ がテキスト情報から CLIP の画像特徴量を作る部分である。これは $y$ → CLIP で $z_t$ → prior モデルで $z_i$ という流れになる。 decoder $P(x | z_i,y)$ は画像特徴量とテキスト情報から画像を生成する部分である。上述の通りこの $z_i$ にはすでにテキスト情報が encode されてるので $z_i$ のみでも画像は作ることができ、テキスト情報 $y$ は optional となっている。

yoheikikuta commented 2 years ago

prior には AutoRegreesive (AR) model と diffusion model の 2 パターンを検討している。 AR model であれ diffusion model であれ、classifier-free guidance でサンプリングできるようにしており、学習時には text 情報を 10% ほどランダムで落として使っているとのこと。

AR model prior

AR model の方は学習の効率性のために PCA して次元削減して、オリジナルの画像の特徴量ベクトル 1024 次元のうち 319 次元の情報を扱っている。この 319 次元の量を特異値の絶対値の大きい順に並べてからそれぞれ 1024 次元のバケットに離散化して、これを予測するという問題設定にしている。モデルは causal attention mask を使った decoder Transformer である。

つまりこれは離散量である sequence を予測するという問題になっており、その意味では文章生成と同様である。そのためここでのモデルは causal attention mask を使った Transformer となっている。

テキスト情報は、caption と CLIP のテキスト特徴量ベクトルを入力の prefix に encode して使っている。さらに $z_i \cdot z_t$ という token も追加している。

このテキスト情報は、例えば GPT モデルでいう場合の prompt に相当している。この prompt 部分は full で attention を張っていいはずなので、上で causal attention mask と書いているが、それはあくまで予測する離散化した画像特徴量ベクトルの部分であり、prompt 部分同士は causal になっていないはず。

これは $z_i$ を予測するモデルなのに $z_i$ を使って $z_i \cdot z_t$ という token を使っているというのがどういうことか理解できない。この積を上位 50% からサンプルすると言ってるが、やはり言ってることが分からない。気持ちとしてはこの積が大きい $z_i$ が作れればよりテキストのコンテキストを捉えた画像特徴量であるが、多様性を考えるとそこそこの値くらいのものがあってもよさそう。なので、画像を生成するときは様々な $z_i$ を作り（これでも決定的なモデルじゃないのかな。予測 token を常に確率最大のものにするならそうだけど、翻訳でいくつか予測するみたいにそこを変えてるんかな）、そこから $z_i \cdot z_t$ の値が上位 50% だけに絞り、そこからサンプルしたものを使うことでいい感じの画像生成ができる、みたいな話なのかとは思う。しかし論文の記述が言ってることはよく分からん...

diffusion prior

こちらも画像特徴量ベクトル $z_i$ を予測するために decoder only Transformer を使っている。 t+1 番目の timestep における decoder only Transformer の入力は次のものを順番通り並べたものになっている。encode したテキスト、CLIP で抽出したテキストの特徴量ベクトル、diffusion timestep の embedding （何番目のタイムステップかという情報を入れてるものだと思うがこれ効くのかな？）、ノイズありの画像特徴量ベクトル $z^{(t)}_i$、$z^{(t)_i}$ の予測をするために使われた Transformer の最終層の embedding。

AR model の方でやっていた $z_i \cdot z_t$ という token を使うことはせず、sampling 時に 2 つの $z_i$ を生成して $z_t$ との内積が大きい方を採用するということをしている。これも上で書いたようにもともとが理解できてないのでよく分からんが、学習後にテキストが与えられてそれを表現する $z_i$ を作るときのサンプリング、ということでいいんかな？ diffusion process は確率的なのでサンプリングできるのはいいとして、この prior での出力は最終的に一つにしている？色々な画像生成ができるけどそこの乱数依存性は decoder の部分で担保しているのかな？

denoising diffusion probabilistic model では gaussian の平均において $x_t$ からずれる分みたいなものとして $\epsilon$ をモデル化するとしていたけど、この論文曰く直接的に $z_i$ を予測する方が結果が良かったとして以下の損失関数を使っている。

この $f_{\theta}$ が decoder only Transformer を表している。ここでやってることは、学習時に diffusion process としては画像特徴量ベクトル $z_i$ は既知なのでそれにノイズを加えていって標準正規分布に従うノイズにして、その reverse process をモデル化するということ。ノイズを加えていくときにテキストの情報は conditional に入れて conditional diffusion process としていて、ここは decoder only transformer の入力部分で表現されている。 classifier free guidance を使っているので、この入力部分のテキストの token は確率的に空にされて学習に使われている。

ただ、この損失関数の書き方だと任意の timestep に対して denoise した結果をもともとの $ zi $ との差分を取ってるので意味が分からない。これ $ || f\theta (z^{(t)}_i, t, y) - z^{(t-1)}_i || $ がやりたいことと思うので typo かな。

この論文はさらっとした記述しかないし公式実装もないので理解しきれないところがあるので疑問系ばかりのメモになってしまった... AR model prior のところも decoder only Transformer と思って読んでいたけど、こっちのほうは a Transformer と書いてあるだけなので翻訳モデル的に $z_t$ を入力に $z_i$ を生成するようになってるの？とか思ったりもするがよく分からない。もう少しちゃんと定式化するか公式実装を出すかどちらかにしてほしい。

yoheikikuta commented 2 years ago

decoder は GLIDE で提案された方法を改良したものになっている。 GLIDE では入力に CLIP のテキスト特徴量ベクトル $z_t$ は入れてなかったがこれを 4 つの extra token に project してGLIDE の text encoding と concat したと言っている。

prior と同様に classifier-free guidance を使っており、10% でランダムに CLIP embedding $z_t$ を 0 にして 50% でランダムにキャプションを落としている。キャプションを落とす確率が結構高いが、画像を生成するための情報としては CLIP の画像特徴量ベクトル（を prior で構築したもの） $z_i$ があるので、テキストに寄りすぎないように結構落としているという感じだと思う。

これでできる画像は 64x64 であり、これを upsample する 2 つの diffusion model を準備し、64x64 → 256x256 → 1024x1024 と高解像度の画像を生成する。このように多段で upsample して高解像度の画像を作るというのは GAN とかでもよく使われているのでそれを踏襲しているということだろう。あとは robustness を高めるために画像をぼかしたり random crop したりとかの工夫もしている。

yoheikikuta commented 2 years ago

ということで、これでもかというほど diffusion model を盛り込んだモデルになっていることがわかる。 prior も diffusion prior を使うとすると、prior で 1 つ, decoder で 64x64 の画像を作るところで 1 つ、decoder で 1024x1024 の画像に upsample するところで 2 つ、合計 4 つである。

なぜ diffusion model が良い性能を発揮するのか、という点については自分はまだ理解が及んでないけど、とにかくよく使われている。こんだけ使われたら diffusion model 自体の理解を深めるための研究とかも出てきそう。

yoheikikuta commented 2 years ago

学習では、データ量として、encoder では 650M の画像（とテキストのペア）、decoder, upsampler, prior では 250M の画像（とテキストのペア）を使っている。

各種モデルの hyperparameter は以下。 sampling variance method の DDIM は Denoising Diffusion Implicit Models https://arxiv.org/abs/2010.02502 のことで、ここではノイズの gaussian の分散を 0 にして決定的に画像を生成（ここでは upsample model を考えているので upsample だが）する手法のことを指している。

yoheikikuta commented 2 years ago

ということで結果だが、これがまたすごい。すごい。

この論文では色々な観点で前身の GLIDE と比較しているが、画像生成の評価はなかなか難しいなと思うのでその辺は省略する。 GLIDE と比べてどの指標でも良いというところまではいってないが、概ね上回っていて安定して美的感覚に優れた画像を生成できている、という感じ。

yoheikikuta commented 2 years ago

あとは embedding space で画像特徴量ベクトルを一方からもう一方へ連続的に変化して生成される画像がどう変わるかとか、テキストでも同じことをやったりしている。画像の変化はこの手の論文ではよくあるので割愛して、テキストの方はもう少し込み入ってるのでこちらを見てみる。手順としては以下のようにする。

スタートとなる画像 x_0 とそのテキスト y0 と、それらの CLIP embedding z{i0}, z{t_0}, を準備する
新しく target とするテキスト y_t とその CLIP embedding z_t を準備する
まずは text の embedding の差分を取って normalize する z_d = norm(zt - z{t_0})
これを使って、球面線形補間 z\theta = slerp(z{i_0}, z_d, \theta) を使い、元の画像特徴量ベクトルからテキスト特徴量ベクトルの差の方向に沿って変化させる
\theta は [0, 0.25~0.5] くらいで変化させて、生成画像は決定的になるように DDIM を使う

結果が以下。例えば一番下では a photo of a landscape in winter という元画像から、テキストの差分としてここでは冬から秋に変わるという方向を取得しており、それにそって特徴量を変化させて画像を生成すると確かに冬から秋に変わっていく様子が生成できているという具合である。

yoheikikuta commented 2 years ago

その他にも CLIP の embedding space を調べるものとして、typographic attacks についても言及している。これは例えばりんごに iPad と書いた紙を貼り付けると、モデルのクラス予測が iPad になってしまうという攻撃方法である。

この論文ではそのようにラベルが誤認識される画像を embed してそれを元に画像生成をする（diffusion model のランダム性によって元の特徴量のコンテキストは維持しつつ色々な画像が生成できる）と、異なるラベルと予測されている画像でも元の画像を再現できることを示している。 CLIP では画像とテキストが同じ特徴量空間に埋め込まれるので、iPad の画像が生成されてもよさそうなものだが、一部 apple 製品っぽい画像を生成したりもしている。また、Pizza に関しては typographic attacks があまり効いてないことも見つけている。これは学習画像で iPad という文字は比較的よく出てくるけど Pizza はそうではないということかな。

yoheikikuta commented 2 years ago

その他には prior を入れた方が結果がよいこと（GLIDE がそうであったように、テキストから画像を生成するということを考えると prior はなくても可能なのでその比較をしている）を示したり、AR prior と diffusion prior を比べると人間の評価で Photorealism, Caption Similarity, Diversity いずれでも diffusion prior の方がよいことを示したり、MS COCO のデータで Zero-shot FID を測定して先行研究よりよいことを示したりしている。この辺は評価が難しいところもあるので、論文ではそういう結果になっているということだけ理解しておいてあまり深くは考えないことにしておく。

unCLIP が GLIDE よりも常によいかというとそういうわけでもなくて、苦手なケースが存在する。それがテキストにおいて 2 つのオブジェクトをくっつけたりするようなケースであり、以下の a red cube on top of a blue cube のような場合に GLIDE の方が人間的な理解ではよさそうに見える。

それ以外にも看板に deep learning と書いてある、みたいなケースも文字がうまく書かれない問題が知られている。これは画像生成でありがちな問題点で unCLIP でも解決できてない。ただしこれについては後続の Google が出した Imagen https://arxiv.org/abs/2205.11487 では劇的に改善されていそうなので楽しみ。結構難しい問題だと思うけど、どのように解決したのだろう。

yoheikikuta commented 2 years ago

ということで一通り読んでみた。 diffusion model がふんだんに使われてテキストからの画像生成は急激に性能が上がった。その理由は分かってないが、今後どのように理解が進んでいくかは楽しみ。

論文の記述があっさりしていて公式実装もないので具体的にどうなっているか分からないところがちょこちょこあるのは辛いところだった。

yoheikikuta / paper-reading