こんにちは、技術開発の三浦です。
技術との付き合い方って難しいと感じることがあります。本当は全て自分たちで出来るようになりたいと思うのですが、限られた時間の中ではなかなか難しいです。自分たちで頑張った方がいいところと誰かの力を借りた方がいいところの判断に、いつも悩んでしまいます。
さて、最近「Zero-Shot Text-to-Image Generation」という論文を読んでおりまして、今回は読んだ内容について、自分の理解を深めるためにも書いてみようと思います。
論文について
「Zero-Shot Text-to-Image Generation」は、最近話題になっている文章から画像を生成するOpenAIのAIシステム「DALL·E2」の前バージョン「DALL·E」を実現する技術について述べている論文です。
- Zero-Shot Text-to-Image Generation
- Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, Ilya Sutskever
- Submitted on 24 Feb 2021 (v1), last revised 26 Feb 2021
- https://arxiv.org/abs/2102.12092
DALL·Eのソースコードは一部を除き公開されていません。現在公開されているのは画像をトークン化し、トークン化された情報から画像を再構築するdiscrete variational autoencoderに関する部分だけのようです。
「Zero-Shot」について
まず論文のタイトルになっている「Zero-Shot」について調べてみました。機械学習のタスクにおいて「Zero-Shot Learning」というものがあるそうで、論文のタイトルの「Zero-Shot」もこの意味を含んだものだと思います。しかし、Zero-Shot Learningよりもさらに広い意味を含んでいるように感じました。
Zero-Shot Learningは機械学習において、モデルの学習時には存在しなかったクラスのラベルに対してもモデルの推論時に正しく分類出来るようにしよう、というタスクです。
例えば動物の画像を動物の種類で分類する時に、学習時には「ウマ」のデータはあるけど「シマウマ」のデータがなかった場合、推論時に現れる「シマウマ」の画像を正しく「シマウマ」と分類することは出来るのでしょうか。
もし人間が分類する立場だったら、事前に「白と黒のシマ模様があるウマがシマウマだよ」という情報と「ウマ」の姿を知っていれば、その人が仮に今までの人生で一度も「シマウマ」の姿を見たことがなかったとしても、事前の知識から「シマウマ」を分類することが出来ると思います。これを機械学習で実現しようという問題が「Zero-Shot Learning」です。
この論文で触れている「Zero-Shot」は、学習時には見たことがないような画像であっても適切な文章の情報を与えてあげれば、それらしい画像を生成することが出来る、という意味合いで使われているのだと思いました。例えばMS-COCOデータセットを学習時には使用していないのに、MS-COCOのキャプションデータからMS-COCOに含まれていてもおかしくないような画像を生成することが出来ています。
上の画像の1行目(Validation)と比較しても、2行目のDALL·Eによって生成された画像がとても自然なものであることが分かります。
また、以下のような普通のデータセットには存在しないような不思議な文章であってもそれらしい画像を生成出来てしまう点も「Zero-Shot」と言えるのかもしれません。
それ以外にも2Dのイラストを3Dに変換したり、何かの設計図を描いてもらったり、文章の投げ方によってはさまざまなタスクに対応出来ることが考えられます。これまでの機械学習では、特定のタスクに特化した手法やデータセットを用意して個別にモデルを設計して学習させる必要がありました。DALL·Eは所見のタスクに対してもある程度の対応が可能であると言え、これも「Zero-Shot」な特徴と言えるのだと思います。
「汎用的なAI」という概念がありますが、もしかしたらその実現に近づいているのかもしれない。そう感じました。
DALL·Eの構造
ここからはDALL·Eの構造について理解したことをまとめてみます。DALL·Eは大きく2つのパートで構成されています。
- 文章から画像トークン列を生成するTransformer
- 画像トークン列から画像を再構築するdiscrete variational autoencoder (dVAE)
図にすると以下のようになります。
画像トークン列は長さが32x32=1,024で、語彙数は8,132です。この離散値で構成されるベクトルから、dVAEのDecoderを通じて画像が再構築されます。実際はdVAEから出力されるのは画像のピクセルを生成する分布のパラメータで、それによって決まる分布から画像が生成されます。ですのでDALL·Eで生成される画像にはランダム性があります。
DALL·Eの特徴
モデルの構造自体は割とシンプルですが、DALL·Eの大きな特徴はそのスケールの大きさと膨大な学習データ、そしてそれを可能にした様々なテクニックにあるのだと思います。
モデルのスケールという点では、Transformerの部分は64のAttentionレイヤを持ち、それぞれ62のmulti-headで構成されていて、そのため学習対象のパラメータの数は120億にも及ぶそうです。この学習のためにNVIDIAのGPU V100を64基使用したとのことで、ここからもそのスケールの大きさをうかがうことが出来ます。
ちなみに以前私がこのブログでもご紹介したTransformerを使ってImage Captioningをした時 のモデルのパラメータ数は430万程度でした。この時はNVIDIAのV100を1基使用して1日程度学習を行いました。
また、使用したデータセットは2億5,000万にも及ぶ文章と画像のペアとのことです。
Sparse Transformer
DALL·EのTransformerはSparse Transformerというもので、Attentionの向け方に工夫がされ、長さの長いデータを取り扱うことが可能になっています。Attentionは入力されたデータ(文章であればトークンで構成された列)を構成するトークンごとに、入力データのトークンのどこに注意して特徴を取り込むかを学習させるテクニックなのですが、入力するデータの長さが長くなると、調べるトークンもAttentionを向けるトークンも増えるため、計算コストが増加してしまう問題があります。
しかし研究によればAttentionは常にすべての場所に向けられるものではなく、向けられる場所にはある程度の傾向があるそうです。例えばあるトークンのAttentionの先は直近数個のトークンだけであったり、周期的に繰り返される位置にあるトークンだけであったりします。
このように、全体の中でAttentionを向けるべき場所は限られていて、それを制御することで計算コストを大幅に減らすことが出来ます。これを実現するテクニックがSparse Transformerです。
DALL·EのTransformerでは、Attentionを行う際にいくつかのマスクを使用し、特定のトークンにAttentionが向かないようにしています。以下の図がマスクを図示したものになっています。DALL·EのTransformerには文章のトークン列に画像のトークン列を結合したデータが流れるのですが、以下の図では実際のデータよりも小さいサイズ(文章のトークン列の長さ：6、画像のトークン列の長さ：16)で表現されています。
マトリクスの上の行から各トークンに対するAttention時にマスクする(Attentionを向けない)箇所が可視化されています。
Sparse TransformerについてもOpenAIから論文が発表されているので、こちらも今後調べてみたいと思います。
まとめ
ということで、今回はDALL·Eの論文「Zero-Shot Text-to-Image Generation」を読んで理解した大まかな内容について、紹介させて頂きました。
論文には私の理解が追いついておらず、紹介出来なかった内容がまだまだたくさんあります。例えばdVAEに関することや、GPUメモリの消費を抑えた学習の方法などです。引き続き論文を読み込んで理解を深めていこうと思います。
