Stable Diffusion 1.5のUNet構造

Stable Diffusion 1.5(SD1.5)の心臓部であるUNetは、「今の潜在表現にどれだけどのようなノイズがあるか」をテキスト条件をヒントに予測するネットワークです。(ノイズを予測し除去していくと元画像＝生成画像が現れる仕掛けです)
本記事では、その内部構造を、実装の仕組みから計算の直感まで順を追って解説します。ステップごとに通っている構造がこのUNETとなります。

1. 全体の流れ:UNetは何を入出力しているのか

まず誤解しやすいポイントとして、UNetが直接処理しているのは「テキスト」ではなく「潜在空間のテンソル」です(テンソルとは、単に多次元の数値配列のことです)。テキストはCLIPのテキストエンコーダーによって77トークン×768次元の配列に変換され、内部のクロスアテンション層に条件(conditioning)として注入されます。

段階	UNetへの入力	テキストの役割
学習時	実画像をVAEで圧縮+ノイズ付加した潜在	クロスアテンションで条件付け
txt2img推論	64×64×4のランダムノイズ	クロスアテンションで条件付け(ここが主役)
img2img推論	元画像をVAEで圧縮+ノイズ付加した潜在	クロスアテンションで条件付け

どのケースでも、UNet自体がやっているのは「潜在にどれだけノイズが乗っているかを予測する」ことであり、テキストを直接画像に変換しているわけではありません。生成時は、ランダムノイズからスタートし、UNetを何十回も繰り返し通してノイズを少しずつ除去し、最後にVAEのデコーダーで画像に戻します。

2. 全体構造:エンコーダー・ミドル・デコーダーのU字型

UNetは大きく3つのパートに分かれます。

エンコーダー(Downブロック):解像度を落としながら特徴を抽出
ミドルブロック:最も低い解像度で処理
デコーダー(Upブロック):解像度を戻しながら、対応するDownブロックの出力をスキップ接続で受け取る

Downブロックで解像度を落としていくのは、画像の細部ではなく「全体としてどんな構図・意味を持つか」という大局的な特徴を再現するためです(この理由は本記事の6章で詳しく説明します)。

ベースチャンネル数は320で、channel_mult = [1, 2, 4, 4]により各段で320→640→1280→1280と増加します。ダウンサンプリングは3回行われ、最も低解像度のDownブロック4とミドルブロックだけは同じ解像度(8×8)で処理されます。

ブロック	チャンネル数	構成
Down1	320	2×ResBlock+Attention → Downsample
Down2	640	2×ResBlock+Attention → Downsample
Down3	1280	2×ResBlock+Attention → Downsample
Down4	1280	2×ResBlockのみ(Attentionなし)
Middle	1280	ResBlock→Attention→ResBlock
Up1〜4	1280→640→320	対応するDownの出力をスキップ結合し、3×ResBlock+Attention → Upsample

ResBlockの中身

端的に言うとResBlockは「入力をそのまま出力に足し戻しながら、そこに加えるべき変化量(残差)だけを2層のConvで学習するブロック」です。GroupNorm→SiLU→3×3 Convの後、タイムステップ埋め込み(sinusoidal embeddingをMLPで1280次元に射影したもの)をチャンネルごとに加算し、さらにGroupNorm→SiLU→Dropout→3×3 Convを通します。入出力のチャンネル数が変わる場合は、スキップパスに1×1 Convを挟んで残差接続します。

SpatialTransformer(Attention)ブロックの中身

GroupNorm→1×1 Convでチャンネルを整えた後、自己注意→クロスアテンション→GEGLUによるFeedForwardを通し、最後に1×1 Convで戻して残差接続します。クロスアテンションの部分でCLIP ViT-L/14のテキストエンコーダー出力(77トークン×768次元)が条件として注入され、ここがプロンプトの内容を反映させる要となる箇所です。

スキップ接続

デコーダー側の各Upブロックは、対応する解像度のエンコーダー出力をチャンネル方向にconcat(結合)してから処理します。concatとは値を混ぜ合わせる(足し算する)のではなく、2つの配列をチャンネル方向に単純に並べて1つの大きな配列にする操作です。これにより低レベルの空間情報が失われず、細かいディテールの再構成に寄与します。

3. パラメータとは何か

「UNetのパラメータ数は約8.6億個」という時のパラメータとは、Conv・GroupNorm・Attentionなどの層が持つ学習可能な重み・バイアスの数値のことです。学習を通じて調整され、推論時は固定された値として使われます。

これは、ブロックを流れる一時的な計算結果(テンソルの要素数)とは全くの別物です。

	実体	学習で変わる?	推論のたびに変わる?
テンソルの要素数	一時的な計算結果(アクティベーション)	いいえ	はい(毎回計算し直す)
パラメータ	Conv/Attentionなどの重み・バイアス	はい(学習中に更新)	いいえ(固定)

パラメータは「計算のレシピ(掛け算する係数の表)」であり、テンソルはそのレシピを使って計算した「材料」にあたります。

4. ブロックと層の関係

「ブロック」は、同じ解像度・チャンネル数の範囲で完結する複数の「層」をひとまとめにした、実装上の整理単位です。

UNet全体
 └─ DownBlock1(ブロック)
     ├─ ResNetLayer 0(層)
     ├─ AttentionLayer 0(層)
     ├─ ResNetLayer 1(層)
     ├─ AttentionLayer 1(層)
     └─ DownsampleLayer(層)
 └─ DownBlock2(ブロック、チャンネル数640)
 └─ ...

DownBlock1内の各層(Conv、GroupNorm、Attentionのto_q/to_k/to_vなど)は、それぞれ独立した重み配列を持っています。実際のモデルファイル(safetensors)を開くと、次のような名前付きの配列がずらっと並んでいます。

down_blocks.0.resnets.0.norm1.weight         shape: (320,)
down_blocks.0.resnets.0.conv1.weight         shape: (320, 320, 3, 3)
down_blocks.0.resnets.0.time_emb_proj.weight shape: (320, 1280)
down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q.weight shape: (320, 320)
down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_k.weight shape: (320, 768)

この単位でまとめる理由は、①チャンネル数・解像度が揃った区切りだから、②スキップ接続で「1ブロック分の出力」をまとめて渡す必要があるから、の2点です。

5. 重みとテンソルはどう干渉するのか

層の種類によって、重みとテンソルの掛け合わせ方(干渉のしかた)は異なります。

層の種類	重みの形	干渉のしかた
Conv(3×3)	(出力ch, 入力ch, 3, 3)	空間的な窓をスライドさせながら内積(重み共有)
GroupNorm	(ch,) が2本	チャンネルごとに単純な掛け算・足し算
Attentionのto_q/to_k/to_v	(ch, ch)	各空間位置のベクトルに行列積(全結合)
FeedForward	(ch, 4ch)など	同じく全結合の行列積

Conv層は「同じ小さな重み窓を画像全体で使い回しながら、新しい配列を1から作り出す」処理です。入力を書き換えているのではなく、入力はそのまま残したまま別の出力配列を生成しています。

6. なぜ解像度を小さくしたり大きくしたりするのか

UNetがU字型に解像度を上げ下げする理由は主に3つです。

少ない計算量で広い範囲を見るため:縮小した状態で同じ3×3フィルタを使うと、元の解像度換算でより広い範囲を一度に見ていることになります。
大局的な情報と局所的な情報の両方が必要だから:エンコーダーで「そこに何があるか」という大局的な特徴を捉え、デコーダーでそれを元に細部を復元します。
計算コストとメモリの節約:Attentionの計算量はマス数に応じて増えてしまうため、低解像度にしてから使うことでコストを抑えています。

最終出力は元の潜在と同じ形である必要があるため、デコーダーで元の解像度まで拡大し、その際に失われがちな細かい位置情報をスキップ接続で補っています。

7. 重み(フィルタ)は位置によらず同じなのに、なぜ均質にならないのか

Conv層の重みは画像のどこでも完全に同じものを使い回します(重み共有)。これは「同じ模様はどこにあっても同じように認識してほしい」という平行移動不変性のための設計です。

それでも出力が均質にならないのは、入力側の中身が場所ごとに全く違うからです。同じフィルタでも、平坦な領域では反応が小さく、エッジがある領域では反応が大きくなります。さらに以下の要因も多様性を生みます。

1つの層に複数のフィルタ(=出力チャンネル)があり、それぞれ違う特徴に反応する
非線形関数(SiLUなど)を挟むことで単純な線形和にならない
層を重ねるごとに見える範囲(受容野)が広がる
Attention層は位置ごとに全く異なる注目のしかたをする

加えて、ダウンサンプリングによって、1つの潜在位置がより広い元画像領域に対応するようになり、少ない計算量で広い空間範囲を処理できるようになります。

8. 最初のノイズ画像は「均質」なのか

ランダムノイズは「全部同じ値」ではなく、各マスごとに独立にサイコロを振って決めた、バラバラな値の集まりです。テレビの砂嵐のように、統計的な性質(平均・分散)はどこでも同じでも、個々の値は場所ごとに全く違います。

そのため、拡散の最初の1ステップ目から、Conv層の出力は場所ごとに異なります。UNetを繰り返し通すたびに、テキスト条件をヒントにしながら「ランダムな局所的ばらつき」の中から徐々に意味のある構造(輪郭・色のまとまり・テクスチャ)が浮かび上がっていきます。

9. 畳み込みの計算方法にはいくつか種類がある

同じ「窓で内積を取る」という定義でも、GPU上での計算アルゴリズムにはいくつかの方式があります。

方式	特徴
Direct convolution	窓を素直にスライドしながら計算。単純だが非効率
im2col + GEMM	窓の中身を並べ替えて巨大な行列積に変換。フレームワークのデフォルトに近い
Winograd convolution	数学的トリックで掛け算の回数自体を削減。小さいフィルタで有効
FFT-based convolution	フーリエ変換経由で計算。大きいフィルタで有利
Implicit GEMM	並べ替えを仮想的に行いながら行列積。cuDNNなど現代のGPUライブラリで主流

実務ではcuDNNなどのライブラリが入力サイズやGPUに応じて自動選択するため、通常は意識する必要はありません。

10. 生成にかかる時間を左右するもの

「窓を動かす」こと自体はGPU上では行列積に変換され並列計算されるため、逐次的な遅さではありません。実際に時間を左右するのは以下の要素です。

掛け算・足し算の総回数(チャンネル数や解像度が大きいほど増える)
Attention層のマス数の2乗に比例する計算コスト
GPUメモリとの読み書き速度(メモリ帯域)
ノイズ除去のステップ数ぶん、UNet全体を何十回も繰り返す必要があること(これが体感速度に一番効く)

まとめ

SD1.5のUNetは、①潜在空間のテンソルを入出力とし、②テキストと時刻情報を条件として内部に注入しながら、③解像度を上げ下げするU字構造とスキップ接続によって、大局的な意味理解と局所的なディテール復元を両立させるネットワークです。約8.6億個のパラメータは、Conv・GroupNorm・Attentionなど無数の層それぞれが持つ固定の重み配列の総和であり、これらがランダムノイズという「バラバラな出発点」に、学習で身につけた構造を段階的に彫り出していくことで画像が生成されます。