tensorを形をみながらネットを潜ってみる話
このジャンルにおいて大きく大前提となってる考え方として、入力 -> 出力 と言ったものがあります。この入力と出力の間を行き来するものを大雑把にテンソル(Tensor)と呼びます。
なので、画像を入力として画像が出力される場合、この画像はテンソルです。実は画像は容易にテンソル化できます。テキストも然り、テキストも昨今耳にすることも増えてきたテキストトークンというテンソルになります。例えば単純な 0.5 という数字も、これもテンソル化が可能です。
用いるライブラリ等で呼び方はさまざまではありますが、配列、ベクトル、テンソル...といろんな言い方があります。要するに一つの、または複数の数値で何かを表したものです。
画像のテンソル化(h, w, c の形)
先に出ました画像をテンソル化することを考えると、通常カラー画像の場合、縦横をピクセルという単位で定義し、RGBチャンネルで色や明るさを表します。グレースケールと呼ばれる白黒画像ならRGBのように3つも必要とせずに1つのチャンネルで完結する場合もあります。
ここでは簡単にRGBの3チャンネルのカラー画像を想定してみていきましょう。例えばこの画像が何かファイルであったとします(jpg ファイルなどを思い浮かべてください)。この時、このファイルをいずれかのライブラリで読み込むと、大体は以下の形をしたテンソルとして表すことができます。
(h, w, c)
それぞれは、
- h: height(画像の高さ・縦幅・y軸)
- w: width(画像の幅・横幅・x軸)
- c: channel(画像のチャンネル数)
を表しています。座標を指し示すときは x, y の順で表すことが多いのですが、解像度にまつわることは h, w の順番で書かれてることが多いです。
例えばこのカラー画像 jpg ファイルが 640x480 の解像度だとすると、このテンソルの形は
(480, 640, 3)
ということになります。
NNにおけるテンソルの形と中間特徴
この形が重要なことが多く、例えばNN(ニューラルネットワーク)ではよく 特徴 だったり、中間特徴 というものを見聞きすることが多いのですが、これは大雑把に言って、いろんな処理を行なってそのテンソルの形をNNに都合のいいように変形させたものです。
また形自体は変形しないけど、順番が入れ替わってたり、全ての値にある処理をして同じ形に戻していたりと様々ですが、大体は元の形から空間次元を圧縮してチャンネル次元を拡張するということがよく用いられるテクニックです。
Stable Diffusionにおけるテンソルの圧縮と復元
何らかのGUIから操作をしている限りあまり意識することはないかもですが、あの有名なStable Diffusionの大元となっているLatent Diffusionの考え方としては(実際にはこの先に述べることを少しハックした形で使用されてます)、入力として (512, 512, 3) を受けとります。これは 512 x 512 x 3 で、合計 786,432 個の値が存在してます。
図1. Stable DiffusionのVAEにおけるテンソルの流れ
中間特徴空間ではこの入力がエンコーダー(sd-vae-encoder)というネットワークを潜って、(4, 64, 64) の形に変形します(PyTorchなどのライブラリではチャンネル、縦、横の順で操作することが多いです)。これは 4 x 64 x 64 で、16,384 の値に圧縮されてます。計算してみると約 2% に圧縮されたことになります。
この後に、よく拡散モデルの解説などで見るノイジング、デノイジング処理が行われますが、一旦それは横に置いておきまして、またこの (4, 64, 64) の形をしたテンソルをデコーダー(sd-vae-decoder)と呼ばれるネットワークを通すと、元の画像と同等の 意味 や 構造 を持つ画像が出力されるという仕組みです。
余談ですが、実際にStable Diffusionのsd-vaeと呼ばれるエンコーダーとデコーダーを直結して出力すると、ほぼほぼ入力と同じ画像が出力されます。また細かい話ですが、sd-vaeのデコーダーから出力されるのは (3, 512, 512) のテンソルで、これを画像ファイルとして保存するために配列次元の順番を入れ替えて、元あったように (512, 512, 3) へと戻します。この時、これをファイル保存するためにPyTorchからNumPyに渡すので、このNumPyオブジェクトになった (512, 512, 3) のテンソルはNumPy配列と呼ばれることが多いです。
なぜ圧縮された空間で処理を行うのか?
さて、ここで重要な疑問、**「なぜわざわざテンソルを2%まで圧縮するのか?」**に答えます。
それは、画像生成の核心である拡散処理(ノイジングとデノイジング)が、非常に**重い処理(heavy weight)**だからです。
元の 512 x 512 x 3 という広大な空間(約78万個の値)でこれらの複雑な処理を実行しようとすると、莫大な計算資源と時間が必要になります。
そこで、処理の舞台を 4 x 64 x 64 というたった約1万6千個の値で構成された圧縮空間、すなわち潜在空間(Latent Space)に移します。これがLatent Diffusionと呼ばれる考え方の核です。潜在空間で処理を行うことで、全体の計算負荷を劇的に下げ、軽量(light-weight) で高速な画像生成を実現しています。
この処理が可能なのは、SD-VAEが単なる エンコード-デコードではなく、入力画像を非常に高い精度で圧縮・解凍できる高度な圧縮解凍機だからです。様々な画像で学習されているため、そのマッピング(対応付け)は広大で、画像の「意味」や「構造」を完璧に潜在空間に閉じ込めることができます。
そして、圧縮空間に移行しても、64x64という空間次元を保持していることが重要です。ボクの経験から言っても、この空間次元はまだ空間情報を色濃く持っていると言えます。この豊かな情報こそが、デコーダー(sd-vae-decoder)が最終的に「猫が雲の上に座っている」といった意図した構図で画像を出力できる根拠となっています。
この「意図」を潜在テンソル((4, 64, 64))に正確に伝えるために、Cross Attention(クロスアテンション)などの技術を使ってテキストなどの 条件(Condition) を付与するわけです。
今回は、この 「テンソルの形が変わること」 と 「なぜその形が変わる空間で処理をするのか」 という部分に焦点を当てました。
