OpenAIが120億個のパラメータを持つテキスト-画像生成装置をどのようにして学習させたのか：DALL-E (ダリー)

深層学習 2021年03月30日

3つの要点
✔️ 120億パラメータの画像からテキストへの生成モデルと、2億5000万の画像キャプションデータセット
✔️ 大規模なモデルをトレーニングするためのいくつかのテクニック
✔️ MS-COCOキャプションにおいて、ゼロショットで90％のリアリティと精度を実現

Zero-Shot Text-to-Image Generation
written by Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, Ilya Sutskever
(Submitted 24 Feb 2021 (this version), latest version 26 Feb 2021 (v2))
Comments: Accepted to arXiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV); Machine Learning (cs.LG)

code：

はじめに

DRAWモデルが初めて画像生成にテキストを導入して以来、この分野では多大な改善がなされてきました。それ以来、GANの統合、自己注意、補助損失などの手法により、これらのモデルは視覚的忠実度の高い画像を生成し、ゼロショットの汎化能力を持つことが可能になりました。それにもかかわらず、これらのモデルは、誤ったオブジェクトの配置、不自然なブレンド、オブジェクトの歪みなどの難読化を起こしやすかったです。

自己回帰変換器は、モデルのサイズ、データ、計算量を適切にスケーリングした場合に、素晴らしい結果を示しています。これまでのテキスト画像生成モデルは、MS-COCOやCUBなどの比較的小さなデータセットで学習・評価されてきました。そのため、モデルサイズやデータセットをスケールアップすることで、性能が向上するかどうかを確認することは価値のあることでした。

スケーリングがどのように機能するかを確認するために、このケースでは、インターネットから取得した2億5000万のテキストと画像のペアに対して、OpenAIは120億のパラメータを持つモデルを学習しました。上の図に示すように、このモデルは非常に堅牢です。MS-COCOデータセットでは、学習ラベルを使用せずに優れたゼロショット画像を生成し、以前のカスタム学習されたモデルと競合しています。

大規模データセット

120億個のパラメータモデルを学習するために，2億5000万個の画像とテキストのペアを使用しています。いくつかはインターネットから取得し、いくつかはConceptual captionsとYFCC100Mデータセットから取得しています。このデータセットにはMS-COCOデータセット全体は含まれておらず、キャプションの付いていない検証用の画像も含まれています。

2段階の学習プロセス

ここでは、テキストと画像が1つのデータストリームとしてモデルを通過する自己回帰変換器を学習します。しかし、256x256の解像度の画像でも、256x256x3の長さのシーケンスを形成するため、高い計算量とメモリが必要になります。また、CNNとは異なり、自己注視では画像の局所的な特徴を効率的に捉えることができないため、モデルが視覚的に認識可能な画像を形成することができません。そこで、可変オートエンコーダー（dVAE）を用いて、画像を8192個の可能値を持つ32x32のグリッドトークンに圧縮します。

学習プロセスは符号化されたRGB画像の画像x、キャプションy、トークンzに対する変分下限(ELBO)を最大化するものとみなすことができます。これらの変数に対する合同対数尤度は次のようにモデル化できます。

ここではトークン(z)が与えられたとき、キャプション(y)は画像(x)から条件付きで独立していると仮定しています。そのため、下界は数学的に次のように与えられます。

p_θは画像トークンとキャプションが与えられたときに、dVAEデコーダが生成するRGB画像に関する確率分布であり、q_φは、dVAEエンコーダが生成する画像トークンに関する分布です。次にELBOのモデルをψ：Priorとします。ここでp_ψは、テキストトークンと画像トークンの合同分布を表します。上記の境界はβ=1の場合にのみ成り立つが、より大きな値が有益であることがわかっています。(DKLはKLダイバージェンス)

1) Visual Codebookの学習

この段階ではRGB画像のみを用いて変分法自動符号化器(dVAE)を学習します。前述のように、dVAEエンコーダは、RGB画像をK=8192Codebookベクトルを持つ32x32グリッドトークンにエンコードします。ELBOを最大化するためにADAMオプティマイザを用います。q_φは離散的な分布であるため、最大化のためには再パラメータ化の手法は役に立たないので、微分可能にするためにgumbel-softmax緩和手法を利用します。q_φに対する期待値を1/q_φ^Tに置き換え、パラメータTは緩和量を制御します。

学習時にTを1/16にアニーリングすることで、緩和検証ELBOを実際のELBOに近づけることができます。また、1x1の畳み込みをエンコーダの最後とデコーダの最初(緩和操作に近い)に行うことで、真のELBOの近似性が向上することがわかりました。さらに、エンコーダとデコーダのリセットブロックの出力アクティベーションを小さな定数でスケーリングすることで、より安定した学習を行うことができます。

2）事前情報の学習

事前分布p_ψは、120億パラメータのスパース変換モデルを用いてモデル化されます。argmaxを用いて、dVAEから32×32 = 1024 個の画像トークンをサンプリングします。各テキストと画像のペアはBPEエンコードされます。BPE符号化では256個のトークンを使用し、16,384個の語彙サイズとなります。さらに、画像のトークンも8192の語彙サイズでエンコードされています。テキストトークンとイメージトークンは連結され、一緒に処理されます。

使用されている変換器は、それぞれ62個のアテンションヘッドを持つ64個の自己注目層と隠れ次元64を持つデコーダのみのモデルです。3種類の自己注意マスクを使用します。キャプションには標準的なテキスト間アテンションマスクを、画像には行、列、コンボリューションアテンションマスクを使用します。

キャプションの長さは最大で256トークンです。テキストトークンと画像トークンは、2つのパディングトークンで区切られます。256の位置にはそれぞれ特別なパディング・トークンがあり、これは学習中に基本的に学習されます。学習にはADAMオプティマイザを使用し、目的は主に画像のモデリングであるため、テキストと画像のクロスエントロピー損失の加重和を、それぞれ1/8と7/8の重みで取っています。

トレーニング方法

1) Mixed Precision Trainingの使用の難しさ

GPUのメモリを節約するためには、Adam momentやactivationなどのパラメータを16ビット精度で格納することが望ましいです。しかし、10億個以上のパラメータを混合精度で学習することは非常に困難です。モデルが深く広くなると、活性化勾配の指数が16ビットフォーマットの範囲外になってしまいます。指数部分は一般的なGPUでは約5ビットを占めており、モデルは不十分でした。そのため、活性化のノルムは後の層になるほど小さくなり、0に丸められてしまいます。これはアンダーフローと呼ばれ、不安定性の最大の原因の1つでした。

この問題は、モデル内の各resblockに異なるグラディエントスケールを使用することで解決されます。上の画像を使って説明することができます。実線はフォワードパス、破線はバックワードパスを表しています。入力された各グラディエントは、スケール／フィルタリングされ、16ビットに変換された後、resblockから離れる前にスケール／フィルタリングが解除されます。フィルタ処理では、NaNまたはInfの値が0に設定されます。これがないと、非有限のイベントが発生した場合、先行するすべてのブロックのグラディエントが低下し、アンダーフローが発生します。

2)分散最適化

モデルの学習にはパラメータシェアリングを使用し、24GBのメモリを必要とする混合精度の学習を行います。各パラメータ配列はノード内の8つのGPUにシャード(分割)されています。各GPUでは、現在のアクティベーションが計算されると同時にAll-gatherを用いて次のパラメータブロックのパラメータシャードをプリフェッチします。同様にバックワードパスでは、現在のブロックのアクティベーションとグラディエントが計算され、同時にオールギャザーを使って前のブロックのパラメータシャードをプリフェッチします。ノード内の8つのGPUすべてが、オールギャザーされたパラメータに対するグラディエントの計算を終えると、reduce-scatterは8つのGPUすべてのグラディエントを平均化し、各GPUにはそのパラメータブロックのグラディエントのみを残します。

同一ノード上のGPU間の通信は十分に高速であるため、これらの重い計算とオーバーラップさせることができます。大きな課題は、異なるノードで計算されたグラディエントの平均化に遅延が生じるノード間通信です。この問題は、勾配圧縮技術を用いて解決しています。各GPUは、隣接するGPUとは独立して、自分のパラメータの低ランク係数を計算します。その後、エラーバッファに、低ランク係数から計算されたグラディエントと、同じノードのGPUで計算されたグラディエントの平均値との差が格納されます。これにより圧縮されていない大きなパラメータからの通信オーバーヘッドを、低ランク因子のための小さな通信演算に減らすことができます。

実験

モデルは、AttnGAN、DM-GAN、DF-GANのようないくつかの他のSOTA-モデルと比較されます。上図は、これらのモデルが生成したサンプルの比較を示しています。さらに、モデルの画像がどれだけリアルであるかを検証するために、人間のレビュアーに我々のモデルDALL-Eと他のモデルが生成した画像を検証してもらいました。

モデルはMS-COCOデータセットのキャプションを対象とした場合、リアルさと正確さの両方の点でDF-GANモデルを大きく上回ります。精度とはキャプションが画像とどの程度一致しているかを示すもので，93.3%の確率で一致しています。

しかし、モデルはCUBデータセット(上図)ではあまり良い結果を得られていません。モデルとベストパフォーマンスのモデルの間には、大きなギャップ(40FID)があります。ゼロショット・アプローチは、CUBデータセットの特殊な分布に適切にフィットしていないようで、この問題に対しては、微調整を行うことがオプションとして考えられます。

まとめ

OpenAIがDALL-Eモデルを発表し、生成された画像を公開したとき、DALL-Eについて多くの宣伝が行われましたが、それは当然のことでした（アボカドの形をしたアームチェアを覚えていますか？)以前のモデルと比較して、テキストから画像への変換機能によって生成された結果は、ゼロショットのタスクであっても非常に印象的です。GPTモデルと同様に、モデル、学習インスタンス、計算機をスケールアップし、適切な学習を行うことで、深層学習モデルの性能を大幅に向上させることができることを示しています。場合によっては、このモデルはキャプションに関連した優れた画像を生成し、「創造性」と呼ばれるようなヒントを与えてくれます。このようなテキスト-画像モデルは、実際の映画でシーンを生成したり、広告ポスター、(アームチェアのような)オブジェクトのデザインなど、様々な用途に利用できるようになるのもそう遠くはないでしょう。