もし生物情報科学専攻の学部生が "StableDiffusion" を理解しようとしたら 4 ~ViT & CLIP~

前回

cake-by-the-river.hatenablog.jp

今回は、前回紹介したTransformerの自然言語処理能力を画像処理の方面に応用した Vision Transformer (ViT) や、Contrastive Language-Image Pre-training (CLIP) について紹介していきます。

ViT

元論文：
arxiv.org

Vision Transformer (ViT) は Google が 2020 年に発表したモデルです。

自然言語処理の界隈は、BERT (Devlin et al., 2018) など Transformer ベースのモデルが用いられ、ラベルのないデータ（正解をヒトの手でキュレーションする必要がないので大量に得られる）を学習に利用する事前学習（Pre-training）と、比較的少量のラベル付きデータセットを用いて個別のタスクへモデルを適応させるファインチューニングを軸に、飛躍的に精度を向上させていました。一方で、画像処理の界隈では未だ CNN を基本としたアーキテクチャが主流であり、一部・全体を Attention に置き換える試みもあったようですが、上手く精度を上げられていませんでした。

そこで、著者の Dosovitskiy らは、Transformer の入力となる文字の埋め込み表現を画像（パッチ）の埋め込み表現に置き換えるシンプルな方法を利用しました。

この方法自体はこの論文が初めてというわけではなく、Cordonnier et al., 2020 では 2×2 の画像パッチを Self-Attention のモデルに入れるという非常に近いものが提案されていました。今回は、それよりもう少し大きな 16×16 のサイズ（中程度の解像度）のパッチを用い、大量のデータによる事前学習とファインチューニングによって高精度にしたということです。

ViT の工夫した点

ViT は「可能な限りオリジナルの Transformer に寄せた設計」をしたものだといい、確かに Transformer における Encoder 部分とほとんど同じですが、いくらか違う点があります。

画像のベクトル埋め込みでは、画像をパッチに分割したものを flatten し、線形変換を用いて Transformer の入力が受け付ける次元に変換します。埋め込みベクトルにTransformer で導入された Position Embedding（ただし2次元画像用に改良したもの）を付け加え、さらに "[class]" に対応するような分類用のベクトルも入力します。これは BERT など現在の自然言語処理の入力フォーマットで標準となっている方法で、入力する文章の最初に [class] という"単語"を挿入しておく（"[class] I want to ..." のようになる）ことで、その部分に当てはまる数値を文章の分類などに利用できるようにしているものです。今回は画像の分類を行いたいので、たくさんのパッチからなる"文章"の最初に [class] を入れておく、というわけです。

Transformer のレイヤー内では残差接続をした Multi-Head Attention や MultiLayer Perceptron (普通の FeedForward Network) を用いていますが、レイヤー正規化はそれらの主要な計算を行う前にしてあります（その方が良いらしい）。他にも、ReLUの代わりに 0 付近でなめらか（微分可能）となった GELU を用いるなど細かい違いもあるようです。

ViT は ImageNet の（今となっては）中程度のサイズのデータセットに対して、ResNet などに少し劣る精度を出し、若干がっかりな感じになりかけたそうですが、BiT ("Big Transfer", Kolesnikov et al., 2019) で導入された 3 億枚ほどの巨大データセットを事前学習に用いることで、およそ 90 %の精度を出し、SOTAに成功しています。結局は数がモノを言うということなのでしょうか…？

とにもかくにも、Transformer をそのまま画像分類タスクに応用した例として話題となりました。画像に対してAttentionを計算しているため、次に示すように画像のどの位置に注目しているのかを表示できるというのも解釈性が上がってCNNに比べて使いやすい点だと言えます。

CLIP

元論文：
arxiv.org

Contrastive Language-Image Pre-training (CLIP) は OpenAI が 2021 年に発表した手法で、初めて見るデータに対してもうまく適用できるようになることを目的として、”どの画像にどの説明文が合うか”という形の事前学習を提案しています。

なお、以下の説明では、次の記事も参考にしました。
techblog.exawizards.com

Zero-shot Learning

Zero-shot とは、学習のとき見たことがない概念の画像を入力したときにも、上手く分類タスクを実行できるか、という問題になります（定義は曖昧なんだとか）。ただし、ラベルの概念自体を全く知らなければ、「知らないもの」という以上の判定はできないので、文章などで上手くその概念のみを学習しておいている前提になります。

例えば、今まで見たことのない「マハールッカデヴァータ」の画像を提示されたときに、頭の中で「サンスクリット語に Mahā という単語があり、偉大な、という意味がある」だとか、「ヒンドゥー教では、インド神話で Rukkhadevatā という名の樹木の神霊が登場する」だとか、「原神というゲームのキャラクターのイラストには〇〇という特徴がある」などの色々な概念的知識を身に着けていれば、「この画像はマハールッカデヴァータかクラクサナリデビか？」といった質問にも答えられるようになる、というわけです。

このような問題に対し、自然言語処理を用いて取り組もうという研究自体は Mori et al., 1999 などずいぶん昔からあり、2013年の段階で DeViSE (Frome et al., 2013) という名の画像と文章の類似度を考えるモデルがすでに考案されていました。上の説明でも文章によりその概念を理解しているため、確かにこのアプローチは有効そうに見えます。

Contrastive Learning

Contrastive, すなわちデータ間にコントラストを付けるような学習手法は、Hadsell et al., 2006 を基に SimCLR (Chen et al., 2020) などで利用されてきました。より具体的には、トリミングや色の歪み、ブラーなどの軽い変換を施した（元は同じ）画像を 2 つ用意し、その画像ペア群 = positive pairs に対し、その潜在表現同士の類似度が近くなるように損失を求めて学習します。画像にラベルを付与する必要はないため、教師なし学習の一種である自己教師あり学習となります。

SimCLR の損失関数は実際、教師あり学習でクロスエントロピー関数のようなものを用います。通常のクロスエントロピー関数は、現在のデータ $(x_n, y_n)$ が各ラベル $c$ に従うと考えるときの評価値 $x_{n,c}$ の総和で正規化します（例えば PyTorch の nn.CrossEntropyLoss）。

$\displaystyle l_n = - w_{y_n} \log {\frac{\exp(x_{n,y_n})}{\sum_{c} \exp(x_{n,c})}}$

今回は正解となるラベルがないので、代わりにミニバッチ学習で入力された画像群がそれぞれ異なるラベルであると捉え、和の取り方に利用します。また、評価値は二つの画像の持つ潜在表現同士の類似度（コサイン類似度）を用います。個々のペアに対して損失関数が得られますが、実際に利用するのは positive pairs の値のみとすることで、「同じ画像からなる表現のみ近づける学習」が可能となるわけです。なお、expに（カノニカル分布で言う）温度の因子を加えて調節もできます。

$\displaystyle l(z_i, z_j) = - \log {\frac{\exp(sim(z_i, z_j) / T)}{\sum_{k\neq i} \exp(sim(z_i, z_k) / T)}}$

$\displaystyle sim(z_i, z_j) = \frac{z_i^\top z_j}{\|z_i\|\|z_j\|}$

$\displaystyle L = \sum_{(i,j)\ \mathrm{from\ same\ image}} l(z_i, z_j)$

CLIP の手法

CLIP はこれらを踏まえ、画像とその説明文の間の Contrastive Learning を行うことで、Zero-shot のような既存の枠組みにとらわれない画像にも適用できるようにしたものです。画像と説明文という異なる種類のデータに共通する（マルチモーダルと呼びます）潜在空間を学習するわけです。

著者らはもともと、画像をCNNに入力し、そのままTransformerにつなげることで説明文を予測するモデルを考えていましたが、上手く精度が上がらないでいました。その原因について、画像と文章を"正確に"対応付けようとしていたからだとしています。CLIP では画像と文章の組に対する Contrastive Learning にしており、正確に文章を予測しないようにすることで、実際には多種多様な形式の説明が可能である画像から本質的な表現のみを考えて類似性を捉えることが出来ると考えられます。

事前学習する際には、Text-Encoder (Transformer) と Image-Encoder (ResNet か ViT) がそれぞれ線形変換によって（マルチモーダルな）潜在空間のベクトルに変換されます。そして、SimCLR と同様に、それぞれのベクトルのコサイン類似度に対して、（温度付き）クロスエントロピーを取る形で損失が求まります。

CLIP では大量の (画像, 説明文) ペアが必要となりますが、既存のデータセットは数が少なかったり質が悪かったりで使えません。CLIP の論文では、Wikipedia に登場する 50 万ほどの単語を含む説明文と画像のWeb検索によって 4 億近くもデータを取得することに成功しています。もちろんこのデータの取得方法は活用方法によって色々と異なっても構いませんが、いずれにせよ大規模なデータセットの取得が重要なようです。

CLIP による Zero-shot 画像分類

こうして学習されたネットワークは、（Zero-shot なクラスへの）画像分類に利用できます。Text-Encoder に "A photo of 〇〇." という形の文を与え、出力される潜在表現をまとめておきます。そして Image-Encoder の出力する画像の潜在表現と類似度を計算することで、最も近かった〇〇に当てはまる単語がクラスとなるわけです。このような方法は、最初から特定のクラスを指定して学習したわけではないことから、Zero-shot の画像に対しても効果的に推測できるようです。

CLIPのZero-shot画像分類精度（https://openai.com/blog/clip/ より引用）

スケッチのバナナを例にすると、従来のモデルでは、学習セットでそれぞれ「スケッチの画像」や「バナナの画像」を学んだとしても、「スケッチのバナナの画像」は直接学んだことがないため、対応する"クラス"が存在せず、バナナであることを適切に判定できないことがあります。一方、CLIP では「スケッチの画像」と「バナナの画像」に対応する（マルチモーダル）潜在空間を記憶しているため、両者が絡んだ概念であっても空間的に近い概念を上手く抽出してくることが可能となった、という風に考えられます。

なお、細かい問題として、多義語の扱いがあります。画像の説明にある "right" が、権利を表すのか、右を表すのかは文脈に依りますし、"bat" の画像と言われたときに、野球のバットもコウモリも当てはまることになります。こうした場合に、"A photo of bat, a type of animal." のように入力文章（プロンプト）の改変を行うことで精度を上げられると言います。このように入力文を工夫することは、一般に Prompt Engineering という風に呼ばれています。

CLIP の活用

DALL·E 2 (Ramesh et al., 2022) や Stable Diffusion (Rombach et al., 2021) など、高精度に文章から対応する画像を生成するモデルにて、CLIP の Text-Encoder 部分がテキストから潜在表現を得るためのモジュールとして利用されています。Stable Diffusion では、以下の module.py の FrozenCLIPEmbedder などが対応しています。なお、BERT など他の Text-Encoder も、もちろん使うことが出来ます。なお、StableDiffusion を用いている NovelAI でも、この CLIP の Encoder を利用しているようでした。

github.com

The Model has been trained without crops, which means now you can generate arbitrary aspect ratio images. Lastly we’ve greatly improved the language understanding of the CLIP model.
— NovelAI (@novelaiofficial) 2022年10月3日

他の発展として、画像から適切な説明文を生成する BLIP (Li et al., 2022) が挙げられます。結構複雑なのでざっくりとしか説明しない（できない）ですが、ポイントは Multimodal mixture of Encoder-Decoder (MED) の仕組みと、入力データセットのフィルタリングのようです。

MED は 3 つの部分からなり、それぞれ (1) Unimodal Encoder, (2) Image-grounded Text Encoder, (3) Image-grounded Text Decoder と呼んでいます（図で左から順に並んでいる）。なお、図で同じ色で示されたレイヤーはすべて同じパラメータを持つようにしており、多くが共通した計算を行うことが分かります。

(1) は基本的に CLIP とほぼ同じで、ViT からなる Image-Encoder の潜在表現と、BERT による Text-Encoder の潜在表現のペアを Contrastive Learning します。(2) は画像からの潜在表現を query として Text-Encoder の Attention の方向性を揃えつつ、画像と文章のペアが positive かどうか、つまり文章が画像をしっかり説明できるものかどうかについて、クロスエントロピー損失を計算します。(3) も画像の潜在表現で方向性を決めつつ、文章生成のタスクを実行します（損失は穴埋め）。

BLIP は以下のサイトでお試し実行できます。また、BLIP を基に、イラスト用の単語なども追加するようにした CLIP-Interrogator も同様に実行できます（Localにも導入できます）。お手持ちの画像で遊んでみると面白いかもしれません。

huggingface.co

次回はいよいよ、拡散モデルの話題に入っていきます。

cake-by-the-river.hatenablog.jp