「猫がギターを弾いている」と入力するだけで、Stable Diffusionはまるで画家のように、ギターを抱えた猫のイラストを描き出します。しかし、冷静に考えてみると不思議なことだらけです。モデルはどうやって「猫」と「ギターを弾く」という2つの概念を理解し、それらを1枚の画像に統合しているのでしょうか？ さらに言えば、「猫」は画像のどこに配置すべきで、「ギター」はどの位置に描くべきかを、どうやって決めているのでしょうか？
この「テキストの意味を画像の空間に落とし込む」という魔法を担っているのが、U-Net内部のCross-Attention層です。Cross-Attentionは、画像の潜在表現の各ピクセル位置が「自分はテキストのどの単語に注目すべきか」を学習する仕組みであり、テキストプロンプトと画像生成を結びつける橋渡し役を果たしています。
Cross-Attentionの仕組みを理解すると、以下のようなことが見えてきます。
- プロンプトエンジニアリング: なぜ特定のプロンプトの書き方が画像品質に影響するのかを本質的に理解できる
- Attention Map解析: 各テキストトークンが画像のどの領域に影響を与えているかを可視化し、モデルの内部動作を解釈できる
- ControlNet・Prompt-to-Prompt: Cross-Attentionを操作することで画像編集や構図制御を行う発展的手法の基盤を理解できる
本記事の内容
- Stable Diffusionの全体アーキテクチャの中でCross-Attentionが果たす役割
- 拡散モデルの前向き・逆過程の基礎と損失関数
- U-Netの構造（ResBlock + Attention Block）
- Cross-Attentionの数学的定式化（Q=画像、K/V=テキスト）
- CLIPテキストエンコーダの出力形式と意味
- Attention Mapの解釈と可視化
- Classifier-Free Guidanceによる条件制御
- PyTorchによるCross-Attention層の実装とAttention Mapの可視化
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
Stable Diffusionの全体像
テキストから画像が生まれるまでの流れ
Stable Diffusionの画像生成プロセスは、料理に例えると分かりやすいかもしれません。テキストプロンプトは「レシピ」、ランダムノイズは「生の食材」、U-Netは「シェフ」、VAEデコーダは「盛り付け」に相当します。シェフ（U-Net）がレシピ（テキスト）を参照しながら食材（ノイズ）を繰り返し加工して料理（画像）を仕上げていく — このレシピ参照の仕組みこそがCross-Attentionです。
もう少し技術的に全体の流れを整理しましょう。Stable Diffusionは大きく3つのコンポーネントから構成されます。
[テキストプロンプト] → [CLIPテキストエンコーダ] → テキスト埋め込み c ∈ R^{77×768}
↓ (Cross-Attention)
[ランダムノイズ z_T] → [U-Net] → [U-Net] → ... → [U-Net] → 潜在表現 z_0
(t=T) (t=T-1) (t=1) ↓
[VAEデコーダ]
↓
生成画像 x
- CLIPテキストエンコーダ: テキストプロンプトを77トークン×768次元の埋め込みベクトル列に変換します。これがCross-Attentionの「参照先」です
- U-Net（ノイズ予測ネットワーク）: ノイズが加えられた潜在表現から、そのノイズを予測します。内部のCross-Attention層がテキスト埋め込みを参照し、テキストの意味に合ったノイズ予測を行います
- VAEデコーダ: 低次元の潜在表現をピクセル空間の画像に復元します
なぜ潜在空間で拡散するのか
Stable Diffusionが画期的だった点は、拡散過程をピクセル空間ではなく潜在空間（latent space）で行うことにあります。ピクセル空間では $512 \times 512 \times 3 = 786{,}432$ 次元ですが、VAEで圧縮された潜在空間では $64 \times 64 \times 4 = 16{,}384$ 次元に圧縮されます。これは約48倍の圧縮であり、U-Netの計算量とメモリ使用量を大幅に削減します。
ただし、圧縮したからといって情報が失われるわけではありません。VAEは画像の本質的な特徴（構図、色彩、質感）を保持しながら、ピクセルレベルの冗長な情報を削ぎ落とします。むしろ、意味的に重要な特徴空間で拡散を行うことが、高品質な画像生成につながっています。
ここまでで、Stable Diffusionの3つのコンポーネントと潜在空間での拡散の意味がわかりました。次に、拡散モデルの数学的な基礎を確認し、U-Netが「何を学習しているのか」を明確にしましょう。
拡散モデルの基礎
前向き過程: 画像をノイズに変える
拡散モデルの基本的なアイデアは、きれいな水にインクを一滴ずつ垂らしていく過程を思い浮かべるとわかりやすいです。最初ははっきりとしたインクの模様（画像）が見えていますが、インクを垂らし続けるうちに水全体が均一に濁り（ガウスノイズ）、元の模様は完全に消えてしまいます。拡散モデルの前向き過程は、この「インクが広がっていく過程」を数学的にモデル化したものです。
元の潜在表現 $\bm{z}_0$ に対して、$T$ ステップにわたってガウスノイズを加えていく前向き過程（forward process）を次のように定義します。
$$ q(\bm{z}_t | \bm{z}_{t-1}) = \mathcal{N}(\bm{z}_t; \sqrt{1 – \beta_t}\, \bm{z}_{t-1},\; \beta_t \bm{I}) $$
ここで $\beta_t$ はノイズスケジュールと呼ばれるパラメータで、各ステップでどれだけのノイズを加えるかを制御します。$\beta_t$ は通常 $0 < \beta_1 < \beta_2 < \cdots < \beta_T < 1$ と小さい値から始まって徐々に増加するように設定されます。
この式が意味しているのは、「前のステップの値を少し縮小し（$\sqrt{1-\beta_t}$ をかけ）、そこにノイズ（$\beta_t$ の分散）を加える」ということです。縮小とノイズ追加を繰り返すことで、最終的に $\bm{z}_T$ はほぼ純粋なガウスノイズ $\mathcal{N}(\bm{0}, \bm{I})$ になります。
重要な数学的性質: 任意ステップへの直接サンプリング
前向き過程を1ステップずつ計算するのは非効率です。$\alpha_t = 1 – \beta_t$ および累積積 $\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s$ を定義すると、$\bm{z}_0$ から任意のステップ $t$ の状態を一度に計算できます。
$$ \bm{z}_t = \sqrt{\bar{\alpha}_t}\, \bm{z}_0 + \sqrt{1 – \bar{\alpha}_t}\, \bm{\epsilon}, \quad \bm{\epsilon} \sim \mathcal{N}(\bm{0}, \bm{I}) $$
この性質がなぜ重要かというと、学習時にランダムなステップ $t$ を選んで直接 $\bm{z}_t$ を計算できるからです。もし逐次的に $\bm{z}_0 \to \bm{z}_1 \to \cdots \to \bm{z}_t$ と計算しなければならなかったら、$T=1000$ ステップ分の計算が毎回必要になり、学習が極めて遅くなります。
逆過程: ノイズから画像を復元する
画像生成は、前向き過程を逆にたどる逆過程（reverse process）で行います。ガウスノイズ $\bm{z}_T \sim \mathcal{N}(\bm{0}, \bm{I})$ から出発して、1ステップずつノイズを取り除き、最終的に元の画像に対応する潜在表現 $\bm{z}_0$ を復元します。
逆過程は次のように定式化されます。
$$ p_\theta(\bm{z}_{t-1} | \bm{z}_t) = \mathcal{N}(\bm{z}_{t-1};\; \bm{\mu}_\theta(\bm{z}_t, t),\; \sigma_t^2 \bm{I}) $$
ここでニューラルネットワーク $\bm{\epsilon}_\theta$ がノイズを予測し、それを使って平均 $\bm{\mu}_\theta$ を計算します。DDPM（Denoising Diffusion Probabilistic Models）では、ノイズ予測に基づいて次のように平均を構成します。
$$ \bm{\mu}_\theta(\bm{z}_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( \bm{z}_t – \frac{\beta_t}{\sqrt{1 – \bar{\alpha}_t}}\, \bm{\epsilon}_\theta(\bm{z}_t, t) \right) $$
損失関数
拡散モデルの学習目標は非常にシンプルです。ランダムなステップ $t$ でノイズ $\bm{\epsilon}$ を加えた潜在表現 $\bm{z}_t$ に対し、ネットワークが予測したノイズ $\bm{\epsilon}_\theta(\bm{z}_t, t)$ と実際に加えたノイズ $\bm{\epsilon}$ のMSE（平均二乗誤差）を最小化します。
$$ \mathcal{L}_{\text{simple}} = \mathbb{E}_{t,\, \bm{z}_0,\, \bm{\epsilon}} \left[ \left\| \bm{\epsilon} – \bm{\epsilon}_\theta(\bm{z}_t, t) \right\|^2 \right] $$
$t$ は $\{1, 2, \ldots, T\}$ から一様にサンプルし、$\bm{\epsilon}$ は標準正規分布からサンプルします。この損失関数は変分下界（ELBO）の再パラメータ化から導かれるものですが、直感的には「加えたノイズを当てるゲーム」を学習していると理解できます。
ただし、ここまでの定式化は無条件の拡散モデルです。テキストプロンプトに従った画像を生成するためには、ノイズ予測がテキスト情報を参照できるようにする必要があります。Stable Diffusionでは、$\bm{\epsilon}_\theta(\bm{z}_t, t)$ を $\bm{\epsilon}_\theta(\bm{z}_t, t, \bm{c})$ に拡張し、テキスト埋め込み $\bm{c}$ を条件として受け取ります。
$$ \mathcal{L}_{\text{cond}} = \mathbb{E}_{t,\, \bm{z}_0,\, \bm{\epsilon},\, \bm{c}} \left[ \left\| \bm{\epsilon} – \bm{\epsilon}_\theta(\bm{z}_t, t, \bm{c}) \right\|^2 \right] $$
では、テキスト条件 $\bm{c}$ はどのようにしてU-Netに注入されるのでしょうか？ その答えを理解するために、まずU-Netの内部構造を見ていきましょう。
U-Netの構造
U字型のエンコーダ・デコーダ構造
U-Netはその名のとおりU字型のネットワークです。画像セグメンテーションのために提案されたアーキテクチャですが、拡散モデルでは「ノイズを予測する」ネットワークとして使われています。
U-Netの構造を大まかに描くと、次のようになります。
入力 z_t (64×64×4)
↓
[ダウンサンプルブロック] 64×64→32×32→16×16→8×8
↓
[ボトルネック] 8×8
↓
[アップサンプルブロック] 8×8→16×16→32×32→64×64
↓ ↑
└── スキップ接続 ──────────┘
↓
出力 ε_θ (64×64×4)
ダウンサンプルパス（エンコーダ）では、解像度を半分にしながらチャネル数を増やし、画像の抽象的な特徴を抽出します。アップサンプルパス（デコーダ）では、解像度を2倍に上げながらチャネル数を減らし、詳細な空間情報を復元します。スキップ接続は、エンコーダの各解像度レベルの出力をデコーダの対応するレベルに直接送ることで、ダウンサンプルで失われがちな細部の情報を保持します。
ResBlock + Attention Block
U-Netの各解像度レベルは、ResBlockとAttention Blockの組み合わせで構成されます。
ResBlock（残差ブロック）は畳み込み層に残差接続を加えた構造で、タイムステップ埋め込みを受け取ります。タイムステップ $t$ はSinusoidal Position Embeddingで埋め込まれ、MLPで変換された後、ResBlockに加算されます。これにより、U-Netは「今どのノイズレベルのデノイジングを行っているか」を認識できます。
$$ \bm{t}_{\text{emb}} = \text{MLP}(\text{SinusoidalEmb}(t)) $$
Attention BlockはResBlockの後に配置され、以下の2つのAttention層で構成されます。
- Self-Attention: 画像の潜在表現の中で、各空間位置が他の全位置との関係を計算します。これにより、遠く離れた画素同士でも長距離の依存関係を捉えることができます
- Cross-Attention: 画像の潜在表現がテキスト埋め込みを参照し、テキスト条件に沿ったノイズ予測を行います
擬似コードで各層の処理を表すと次のようになります。
def unet_block(z, t_emb, c):
z = ResBlock(z, t_emb) # 特徴変換 + タイムステップ条件
z = SelfAttention(z) # 画像内の空間的依存関係
z = CrossAttention(z, c) # テキスト条件の注入
return z
ここで重要なのは、タイムステップ $t$ とテキスト条件 $\bm{c}$ の注入方法が異なるという点です。タイムステップはResBlock内で加算によって注入されるのに対し、テキスト条件はCross-Attentionという注意機構を通じて注入されます。加算は全画素に同じ情報を一律に与えますが、Cross-Attentionは各画素が必要な情報を選択的に参照できるため、「猫」という単語の情報は猫が描かれるべき領域に、「ギター」という単語の情報はギターが描かれるべき領域に、それぞれ届けることが可能になります。
この選択的な情報伝達がStable Diffusionの画像品質を支えている重要なメカニズムです。では、Cross-Attentionがどのようにしてこの選択的参照を実現しているのか、その数学的な仕組みを詳しく見ていきましょう。
Cross-Attentionの役割
テキストと画像を結ぶ橋
Cross-Attention層は、本記事の中心テーマです。一般的なCross-Attentionでは、ある系列（ターゲット）が別の系列（ソース）の情報を参照しますが、Stable Diffusionでは具体的に以下のように対応付けられています。
- Query（質問者）= 画像の潜在表現 $\bm{z}$: 「自分（この画素位置）はどんな情報を必要としているか」
- Key（鍵）= テキスト埋め込み $\bm{c}$: 「このテキストトークンはどんな情報を持っているか」
- Value（値）= テキスト埋め込み $\bm{c}$: 「このテキストトークンが実際に渡す情報の中身」
人間が絵を描くときのことを想像してみてください。キャンバスの各部分を描いているとき、画家は頭の中でプロンプト（指示文）の各単語を参照しています。「猫」を描くべき領域では「猫」という単語に強く注目し、背景を描いているときは「公園」や「木」に注目しているでしょう。Cross-Attentionは、まさにこの画家の「注目の向け方」を数学的にモデル化したものです。
Self-AttentionとCross-Attentionの違い
U-Net内で使われるSelf-AttentionとCross-Attentionは、Query・Key・Valueの生成元が異なるという一点で区別されます。
Self-Attentionでは、Query・Key・Valueが全て同じ入力（画像の潜在表現 $\bm{z}$）から生成されます。
$$ \bm{Q} = \bm{W}_Q \bm{z}, \quad \bm{K} = \bm{W}_K \bm{z}, \quad \bm{V} = \bm{W}_V \bm{z} $$
つまり、画像内の各画素が「他の画素とどう関係しているか」を計算します。たとえば、猫の顔の領域が猫の体の領域と関連していることを学習します。
Cross-Attentionでは、Queryは画像から、Key・Valueはテキストから生成されます。
$$ \bm{Q} = \bm{W}_Q \bm{z}, \quad \bm{K} = \bm{W}_K \bm{c}, \quad \bm{V} = \bm{W}_V \bm{c} $$
つまり、画像の各画素が「テキストのどのトークンに注目すべきか」を計算します。猫が描かれるべき領域のQueryは、テキスト中の「猫」トークンのKeyと高いスコアを持ち、その結果として「猫」のValueが多く取り込まれます。
この違いを表にまとめると次のようになります。
| Self-Attention | Cross-Attention | |
|---|---|---|
| Query の生成元 | 画像潜在表現 $\bm{z}$ | 画像潜在表現 $\bm{z}$ |
| Key の生成元 | 画像潜在表現 $\bm{z}$ | テキスト埋め込み $\bm{c}$ |
| Value の生成元 | 画像潜在表現 $\bm{z}$ | テキスト埋め込み $\bm{c}$ |
| 注意マップの形状 | $(hw) \times (hw)$ | $(hw) \times L$ |
| 役割 | 画像内の空間的依存関係 | テキスト→画像への情報注入 |
ここで $h, w$ は潜在表現の空間サイズ、$L$ はテキストトークン数（CLIPでは77）です。
Cross-Attentionの直感的な理解ができたところで、次はその数学的な定式化を詳しく見ていきましょう。
Cross-Attentionの数学的定式化
入力と射影
Cross-Attentionの入力は、画像の潜在表現とテキスト埋め込みの2つです。
画像の潜在表現は、U-Net内部で畳み込みや前段のAttentionを経た後のテンソルです。形状は $(B, C, H, W)$ ですが、Attention計算のために空間次元を平坦化（flatten）して $(B, HW, C)$ に変換します。ここで $B$ はバッチサイズ、$C$ はチャネル数、$H \times W$ は空間サイズです。
平坦化した画像特徴を $\bm{z} \in \mathbb{R}^{N \times d_z}$ と書きます。$N = H \times W$ は空間位置の数、$d_z$ はチャネル数（=特徴次元）です。テキスト埋め込みは $\bm{c} \in \mathbb{R}^{L \times d_c}$ で、$L$ はトークン数、$d_c$ はCLIPの埋め込み次元です。
まず、3つの学習可能な重み行列 $\bm{W}_Q \in \mathbb{R}^{d_z \times d}$、$\bm{W}_K \in \mathbb{R}^{d_c \times d}$、$\bm{W}_V \in \mathbb{R}^{d_c \times d}$ を使って、Query・Key・Valueを生成します。
$$ \begin{align} \bm{Q} &= \bm{z}\,\bm{W}_Q \quad \in \mathbb{R}^{N \times d} \\ \bm{K} &= \bm{c}\,\bm{W}_K \quad \in \mathbb{R}^{L \times d} \\ \bm{V} &= \bm{c}\,\bm{W}_V \quad \in \mathbb{R}^{L \times d} \end{align} $$
ここで $d$ は射影後の次元（Attention次元）です。重要なのは、$\bm{W}_Q$ は画像特徴の次元 $d_z$ から射影し、$\bm{W}_K$ と $\bm{W}_V$ はテキストの次元 $d_c$ から射影するという点です。つまり、画像とテキストは元々異なる次元空間に存在しますが、線形射影によって同じ $d$ 次元の空間に写像されます。この共通空間で内積を計算することで、画像の各位置とテキストの各トークンの「関連度」を測ることができるのです。
Scaled Dot-Product Attention
Query、Key、Valueが生成されたら、Scaled Dot-Product Attentionの計算を行います。
まず、QueryとKeyの内積でAttentionスコアを計算します。
$$ \bm{S} = \bm{Q}\bm{K}^\top \in \mathbb{R}^{N \times L} $$
この行列 $\bm{S}$ の $(i, j)$ 成分 $S_{ij}$ は、$i$ 番目の画素位置のQuery $\bm{q}_i$ と $j$ 番目のテキストトークンのKey $\bm{k}_j$ の内積です。内積が大きいほど、その画素位置はそのテキストトークンと「関連性が高い」ことを意味します。
次に、スケーリングファクター $\sqrt{d}$ で割ります。
$$ \bm{S}_{\text{scaled}} = \frac{\bm{Q}\bm{K}^\top}{\sqrt{d}} $$
$\sqrt{d}$ で割る理由を説明します。QueryとKeyの各要素が平均0、分散1の独立な確率変数だと仮定すると、内積 $\bm{q}_i^\top \bm{k}_j = \sum_{k=1}^{d} q_{ik} k_{jk}$ の分散は $d$ に比例します。$d$ が大きいと内積の値が非常に大きくなり、softmaxの出力が一つの要素にほぼ全ての重みが集中する「尖った」分布になってしまいます。$\sqrt{d}$ で割ることで内積の分散を1に正規化し、softmaxが適度に滑らかな分布を出力できるようにします。
softmaxを行ごとに適用して、注意重み（Attention Map）を計算します。
$$ \bm{A} = \text{softmax}\left(\frac{\bm{Q}\bm{K}^\top}{\sqrt{d}}\right) \in \mathbb{R}^{N \times L} $$
$\bm{A}$ の各行は確率分布であり、全要素が非負で行ごとの和が1になります。$A_{ij}$ は「$i$ 番目の画素位置が $j$ 番目のテキストトークンにどれだけ注目しているか」を表す重みです。
最後に、注意重みでValueを加重平均して出力を得ます。
$$ \text{CrossAttn}(\bm{Q}, \bm{K}, \bm{V}) = \bm{A}\bm{V} \in \mathbb{R}^{N \times d} $$
出力の $i$ 行目は、全テキストトークンのValueの加重和です。
$$ \text{output}_i = \sum_{j=1}^{L} A_{ij}\, \bm{v}_j $$
注意重み $A_{ij}$ が大きいトークンの情報が多く取り込まれるため、各画素位置は「自分に関係のあるテキスト情報」を選択的に受け取ることができます。
1つの画素の視点で見る
数式を具体的にイメージするために、$i$ 番目の画素位置に着目してみましょう。
この画素のQueryベクトル $\bm{q}_i$ は、77個のテキストトークンのKeyベクトル $\bm{k}_1, \bm{k}_2, \ldots, \bm{k}_{77}$ それぞれとの内積を計算します。もしプロンプトが「a cat playing guitar」なら、猫が描かれるべき領域の画素のQueryは、「cat」に対応するKeyと高いスコアを持つでしょう。
softmaxを通すと、たとえば以下のような注意重みが得られます。
| トークン | a | cat | playing | guitar | … |
|---|---|---|---|---|---|
| 重み $A_{ij}$ | 0.02 | 0.45 | 0.08 | 0.10 | … |
最終出力は、これらの重みでValueの加重和を取ったものです。「cat」のValueが最も大きな重みで取り込まれるため、この画素位置には「猫」に関連する特徴情報が強く注入されます。
Multi-Head Cross-Attention
実際のStable Diffusionでは、Cross-AttentionもMulti-Headで計算されます。Multi-Headにすることで、異なるヘッドが異なる「視点」からテキスト情報を参照できるようになります。
$h$ をヘッド数、$d_h = d / h$ を各ヘッドの次元とすると、各ヘッド $i$ は独立した重み行列 $\bm{W}_Q^{(i)}, \bm{W}_K^{(i)}, \bm{W}_V^{(i)}$ を持ちます。
$$ \text{head}_i = \text{Attention}(\bm{z}\bm{W}_Q^{(i)},\; \bm{c}\bm{W}_K^{(i)},\; \bm{c}\bm{W}_V^{(i)}) $$
全ヘッドの出力を結合して、出力射影行列 $\bm{W}_O$ で合成します。
$$ \text{MultiHead}(\bm{z}, \bm{c}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\, \bm{W}_O $$
たとえば、あるヘッドは「猫」と「ギター」の空間的配置関係に注目し、別のヘッドは「弾いている」という動作の描写に注目する、というように、複数の観点からテキスト情報を引き出すことが可能になります。Stable Diffusion v1.5では8ヘッド、SDXLでは20ヘッドが使われています。
Cross-Attentionの数学的な仕組みがわかりました。次に、Key・Valueの元となるCLIPテキストエンコーダの出力がどのような構造をしているかを見ていきましょう。
CLIPテキストエンコーダ
テキストを数値ベクトルに変換する
Cross-AttentionのKey・Valueは、CLIPテキストエンコーダの出力から生成されます。CLIPは画像とテキストの対照学習によって事前学習されたモデルであり、テキストの意味を豊かなベクトル表現に変換する能力を持っています。
テキストエンコーダの処理の流れは以下のとおりです。
- トークナイズ: テキストプロンプトをBPE（Byte Pair Encoding）トークナイザでトークン列に変換します。先頭に
[SOT]（Start of Text）、末尾に[EOT]（End of Text）の特殊トークンが追加されます - パディング: トークン列を最大長77にパディングします。短いテキストは
[PAD]トークンで埋められます - トークン埋め込み: 各トークンを埋め込みベクトルに変換し、位置エンコーディングを加算します
- Transformer処理: 12層（SD 1.xの場合）のTransformer Encoderでコンテキストを考慮した埋め込みを生成します
最終的な出力は次の形状になります。
$$ \bm{c} = \text{CLIPTextEncoder}(\text{text}) \in \mathbb{R}^{77 \times 768} $$
77はコンテキスト長（固定）、768はCLIP ViT-L/14の埋め込み次元です。Stable Diffusion 2.xではOpenCLIPのViT-H/14が使われ、次元は1024に拡大されています。
各トークンの埋め込みが持つ意味
この出力 $\bm{c}$ の各行は、対応するテキストトークンのコンテキストを考慮した埋め込みベクトルです。重要なのは、単なる単語の意味だけでなく、Transformerの自己注意機構によって文全体のコンテキストを反映している点です。
たとえば「a cat playing guitar in a park」というプロンプトでは、「cat」のトークン埋め込みは、単に「猫」という概念だけでなく、「ギターを弾いている猫」「公園にいる猫」というコンテキスト情報も含んでいます。
Cross-Attentionは、この77個のトークン埋め込みを全てKey・Valueとして使用します。つまり、[SOT]、[EOT]、[PAD] トークンの埋め込みもCross-Attentionに渡されます。実験的には、[SOT] と [EOT] のトークンは文全体の要約的な情報を保持していることが知られており、Cross-Attentionはこれらのトークンからグローバルな情報を取得し、各単語トークンから局所的な情報を取得するという使い分けをしていると考えられます。
SD 1.x と SD 2.x のテキストエンコーダの違い
| Stable Diffusion 1.x | Stable Diffusion 2.x | |
|---|---|---|
| モデル | CLIP ViT-L/14 | OpenCLIP ViT-H/14 |
| 埋め込み次元 $d_c$ | 768 | 1024 |
| Transformer層数 | 12 | 23（末尾の層は使わない） |
| コンテキスト長 $L$ | 77 | 77 |
SD 2.xでは、OpenCLIPのテキストエンコーダの最終層ではなく途中の層（penultimate layer）の出力を使用しています。これは最終層がCLIPの対照学習の目的に特化しすぎており、拡散モデルに必要な豊かな表現が中間層により多く含まれているためです。
テキスト埋め込みの構造がわかったところで、Cross-Attentionが生成するAttention Mapを解釈する方法を見ていきましょう。
Attention Mapの解釈
Attention Mapとは何を表すか
Cross-Attentionの計算で得られるAttention Map $\bm{A} \in \mathbb{R}^{N \times L}$ は、画像のどの領域がテキストのどのトークンに注目しているかを示す「対応関係マップ」です。
$\bm{A}$ を転置して $\bm{A}^\top \in \mathbb{R}^{L \times N}$ として見ると、各テキストトークンがどの画像領域に影響を与えているかを読み取れます。$j$ 番目のトークンに対応する行 $\bm{A}_{:,j} \in \mathbb{R}^{N}$ を $H \times W$ の2次元マップに変形すれば、そのトークンの「注意の空間分布」を可視化できます。
たとえば「a cat playing guitar」というプロンプトの場合、以下のようなAttention Mapが期待されます。
- 「cat」トークン: 猫が描かれる画像領域に高い注意重みを持つ
- 「guitar」トークン: ギターが描かれる領域に高い注意重みを持つ
- 「playing」トークン: 猫とギターが接する領域（猫の手がギターに触れている部分）に注意が分散する
- 「a」トークン: 特定の領域に集中せず、比較的均一に分散する
解像度による違い
U-Netには複数の解像度レベルにCross-Attention層が配置されています。Stable Diffusion 1.5の場合、$64 \times 64$、$32 \times 32$、$16 \times 16$、$8 \times 8$ の各解像度でCross-Attentionが計算されます。
- 低解像度（$8 \times 8$, $16 \times 16$）: 大まかな構図やオブジェクトの配置を決定します。各「画素」が広い受容野を持つため、グローバルな対応関係が学習されます
- 高解像度（$32 \times 32$, $64 \times 64$）: オブジェクトの形状や細部のテクスチャを決定します。各画素がより局所的な領域を担当するため、きめ細かい対応関係が現れます
研究（Prompt-to-Prompt: Hertz et al., 2022）により、中間解像度のCross-Attention層が画像の構図に最も大きな影響を与えることが示されています。この知見は、Cross-Attention Mapを操作して画像編集を行う手法の基盤となっています。
Attention Mapの限界
Attention Mapは直感的で解釈しやすいものですが、いくつかの注意点もあります。
まず、Multi-Head Attentionの各ヘッドは異なるAttention Mapを持つため、「どのヘッドのマップを見るか」で解釈が変わります。通常はヘッド間で平均したAttention Mapが使われますが、これは情報の一部を失っている可能性があります。
また、Attention Mapは「テキストトークンからの情報の流れ」を示していますが、実際の画像生成はSelf-Attention、ResBlock、スキップ接続なども含めた複合的なプロセスです。Cross-Attentionだけで画像の内容が決まるわけではないことに留意する必要があります。
Attention Mapの解釈を理解したところで、Stable Diffusionの画像品質を大きく左右するもう一つの重要な仕組み — Classifier-Free Guidanceを見ていきましょう。
Classifier-Free Guidance
条件付き生成の品質を高めるテクニック
Cross-Attentionによってテキスト条件をU-Netに注入できるようになりましたが、単にCross-Attentionを入れただけでは、テキストプロンプトへの忠実度が十分でないことがあります。「猫がギターを弾いている」と入力しても、猫は描かれるがギターが抜け落ちる、ということが起こりえます。
この問題を解決するのがClassifier-Free Guidance（CFG）です。CFGの着想は意外にシンプルです。「条件に合った方向」への推定を増幅するために、「条件なし」と「条件あり」のノイズ予測の差を利用します。
まず、学習時にはランダムにテキスト条件を空（null条件、$\bm{c} = \emptyset$）に置き換えて学習します。たとえば10%の確率でテキストをドロップします。これにより、同じU-Netが条件付きの予測 $\bm{\epsilon}_\theta(\bm{z}_t, t, \bm{c})$ と無条件の予測 $\bm{\epsilon}_\theta(\bm{z}_t, t, \emptyset)$ の両方をこなせるようになります。
CFGの数式
推論時、最終的なノイズ予測は次の式で計算されます。
$$ \tilde{\bm{\epsilon}}_\theta(\bm{z}_t, t, \bm{c}) = \bm{\epsilon}_\theta(\bm{z}_t, t, \emptyset) + w \cdot \left( \bm{\epsilon}_\theta(\bm{z}_t, t, \bm{c}) – \bm{\epsilon}_\theta(\bm{z}_t, t, \emptyset) \right) $$
$w$ はguidance scale（ガイダンススケール）と呼ばれるハイパーパラメータです。この式を整理すると次のようになります。
$$ \tilde{\bm{\epsilon}}_\theta = (1 – w) \cdot \bm{\epsilon}_\theta(\bm{z}_t, t, \emptyset) + w \cdot \bm{\epsilon}_\theta(\bm{z}_t, t, \bm{c}) $$
$w$ の値による振る舞いの違いを理解しましょう。
- $w = 1$: 条件付き予測そのもの。CFGなしと同じです
- $w > 1$: 条件付き予測と無条件予測の差分を増幅します。テキストに忠実な画像が生成されますが、大きすぎると過飽和や不自然なアーティファクトが発生します
- $w = 0$: 完全な無条件生成。テキストを無視します
直感的には、CFGは「テキスト条件がある場合とない場合のノイズ予測の差」を増幅しています。この差分はまさに「テキスト条件がノイズ予測にもたらす影響」であり、Cross-Attentionを通じてテキスト情報が注入された分の効果です。$w > 1$ とすることで、Cross-Attentionの効果をブーストしているとも解釈できます。
Guidance Scaleの実践的な選び方
Stable Diffusion 1.5では $w = 7.5$ が一般的なデフォルト値です。SDXLでは $w = 5.0 \sim 9.0$ の範囲が推奨されています。
| $w$ の範囲 | 生成傾向 |
|---|---|
| 1.0 ~ 3.0 | 多様だが曖昧。テキストとの一致度が低い |
| 5.0 ~ 8.0 | バランスの良い生成。プロンプトに忠実かつ自然 |
| 10.0 ~ 20.0 | プロンプトに非常に忠実だが色が過飽和に。詳細が崩れることもある |
| 20.0 以上 | アーティファクトが多発。実用的ではない |
CFGが計算コストに与える影響も重要です。各デノイジングステップで、条件付きと無条件の2回のU-Netフォワードパスが必要になるため、推論時間はCFGなしの場合の約2倍になります。
ここまでで、Cross-Attentionの理論からCFGまで一通り理解できました。次に、これらの知識をPyTorchのコードで実装して、理論と実装の対応を確認しましょう。
PyTorchによる実装
Cross-Attention層の実装
まず、Stable Diffusionで使われるCross-Attention層をPyTorchでスクラッチ実装します。ここでは、2次元の画像特徴テンソルを入力として受け取り、テキスト埋め込みを参照するMulti-Head Cross-Attention層を構築します。
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
class CrossAttention(nn.Module):
"""Stable Diffusion向けCross-Attention層"""
def __init__(self, query_dim, context_dim, num_heads=8):
super().__init__()
self.num_heads = num_heads
self.head_dim = query_dim // num_heads
self.scale = self.head_dim ** -0.5
# Query: 画像特徴から生成
self.to_q = nn.Linear(query_dim, query_dim, bias=False)
# Key, Value: テキスト埋め込みから生成
self.to_k = nn.Linear(context_dim, query_dim, bias=False)
self.to_v = nn.Linear(context_dim, query_dim, bias=False)
# 出力射影
self.to_out = nn.Linear(query_dim, query_dim)
# Attention Mapを保存（可視化用）
self.attention_map = None
def forward(self, x, context):
"""
x: 画像特徴 (batch, channels, height, width)
context: テキスト埋め込み (batch, seq_len, context_dim)
"""
b, c, h, w = x.shape
# 空間次元を平坦化: (b, c, h, w) → (b, h*w, c)
x_flat = x.reshape(b, c, h * w).permute(0, 2, 1)
# Query, Key, Valueの生成
q = self.to_q(x_flat) # (b, h*w, query_dim)
k = self.to_k(context) # (b, seq_len, query_dim)
v = self.to_v(context) # (b, seq_len, query_dim)
# Multi-Head分割: (b, n, num_heads, head_dim) → (b, num_heads, n, head_dim)
q = q.reshape(b, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
k = k.reshape(b, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
v = v.reshape(b, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
# Scaled Dot-Product Attention
scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
attn = F.softmax(scores, dim=-1)
# Attention Mapを保存（ヘッド平均）
self.attention_map = attn.mean(dim=1).detach()
# Valueの加重和
out = torch.matmul(attn, v)
# ヘッドを結合: (b, num_heads, h*w, head_dim) → (b, h*w, query_dim)
out = out.permute(0, 2, 1, 3).reshape(b, h * w, -1)
# 出力射影
out = self.to_out(out)
# 空間次元を復元: (b, h*w, c) → (b, c, h, w)
out = out.permute(0, 2, 1).reshape(b, c, h, w)
return out
このコードのポイントを確認しましょう。to_q は画像特徴から、to_k と to_v はテキスト埋め込みからそれぞれQuery・Key・Valueを生成しています。これが「Q=画像、K/V=テキスト」というCross-Attentionの核心部分です。また、可視化のために self.attention_map にヘッド平均のAttention Mapを保存しています。
Attention MapのダミーデータでAttention Block全体の構築
次に、ResBlock + Self-Attention + Cross-Attentionを組み合わせたAttention Blockを実装します。
class SelfAttention(nn.Module):
"""Self-Attention層（画像内の空間的依存関係）"""
def __init__(self, dim, num_heads=8):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.scale = self.head_dim ** -0.5
self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
self.to_out = nn.Linear(dim, dim)
def forward(self, x):
b, c, h, w = x.shape
x_flat = x.reshape(b, c, h * w).permute(0, 2, 1)
# Q, K, Vを一括生成
qkv = self.to_qkv(x_flat).chunk(3, dim=-1)
q, k, v = [t.reshape(b, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3) for t in qkv]
scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
attn = F.softmax(scores, dim=-1)
out = torch.matmul(attn, v)
out = out.permute(0, 2, 1, 3).reshape(b, h * w, -1)
out = self.to_out(out)
out = out.permute(0, 2, 1).reshape(b, c, h, w)
return out
class AttentionBlock(nn.Module):
"""Self-Attention + Cross-Attentionを組み合わせたブロック"""
def __init__(self, dim, context_dim, num_heads=8):
super().__init__()
self.norm1 = nn.GroupNorm(8, dim)
self.self_attn = SelfAttention(dim, num_heads)
self.norm2 = nn.GroupNorm(8, dim)
self.cross_attn = CrossAttention(dim, context_dim, num_heads)
def forward(self, x, context):
# Self-Attention（残差接続付き）
x = x + self.self_attn(self.norm1(x))
# Cross-Attention（残差接続付き）
x = x + self.cross_attn(self.norm2(x), context)
return x
Self-AttentionではQ・K・Vが全て同じ入力 x から生成されるのに対し、Cross-AttentionではQが x から、K・Vが context（テキスト埋め込み）から生成されている点を確認してください。両方とも残差接続が使われており、Attention層を通過してもオリジナルの情報が保持されます。
Attention Mapの可視化
Cross-Attention層が各テキストトークンに対してどのような空間的注意パターンを生成するかを可視化しましょう。
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
# 再現性のためシードを固定
torch.manual_seed(42)
# ダミーデータの作成
batch_size = 1
height, width = 16, 16 # 潜在空間の空間サイズ
channels = 320 # SD 1.5のCross-Attention次元
context_dim = 768 # CLIPの埋め込み次元
seq_len = 77 # CLIPのコンテキスト長
# ダミーの画像潜在特徴（4チャネルの場合、チャネルを320に射影済みと仮定）
z = torch.randn(batch_size, channels, height, width)
# ダミーのテキスト埋め込み（CLIPの出力を模擬）
# トークン: [SOT] a cat playing guitar [EOT] [PAD] [PAD] ...
context = torch.randn(batch_size, seq_len, context_dim)
# Cross-Attention層を生成
cross_attn = CrossAttention(query_dim=channels, context_dim=context_dim, num_heads=8)
# フォワードパス
with torch.no_grad():
output = cross_attn(z, context)
# Attention Mapを取得: (batch, h*w, seq_len)
attn_map = cross_attn.attention_map[0] # (h*w, seq_len)
print(f"出力形状: {output.shape}")
print(f"Attention Map形状: {attn_map.shape}")
# 各トークンのAttention Mapを空間マップとして可視化
tokens = ["[SOT]", "a", "cat", "playing", "guitar", "[EOT]"]
token_indices = [0, 1, 2, 3, 4, 5]
fig, axes = plt.subplots(1, len(tokens), figsize=(18, 3))
for idx, (token, tidx) in enumerate(zip(tokens, token_indices)):
# トークンjに対する全画素の注意重みを取得し、2Dに変形
spatial_map = attn_map[:, tidx].reshape(height, width).numpy()
im = axes[idx].imshow(spatial_map, cmap='hot', interpolation='nearest')
axes[idx].set_title(f'"{token}"', fontsize=12)
axes[idx].axis('off')
plt.suptitle('Cross-Attention Map for Each Token', fontsize=14, y=1.02)
plt.colorbar(im, ax=axes, shrink=0.8, label='Attention Weight')
plt.tight_layout()
plt.savefig('cross_attention_map.png', dpi=150, bbox_inches='tight')
plt.show()
このコードでは、ダミーデータを使ってCross-Attention層のフォワードパスを実行し、各テキストトークンに対するAttention Mapを空間的に可視化しています。今回はランダムな重みで初期化した未学習のモデルなので、Attention Mapにはまだ意味のあるパターンは現れません。学習済みのStable Diffusionモデルであれば、「cat」トークンは猫の領域に、「guitar」トークンはギターの領域に、それぞれ注意が集中するパターンが観察されます。
ダミーデータではありますが、重要なのは次の点です。出力形状 (1, 320, 16, 16) は入力と同じ形状であり、Cross-Attentionが画像特徴の空間構造を保持していることがわかります。Attention Map形状 (256, 77) は 16×16=256 の画素位置それぞれが77トークンへの注意重みを持っていることを意味し、これが「各画素がどのテキストトークンに注目するか」の情報を含んでいます。
Classifier-Free Guidanceの実装
CFGの実装は非常にシンプルです。各デノイジングステップで条件付き・無条件の2回のフォワードパスを実行し、その結果を線形結合します。
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
def classifier_free_guidance(
model, z_t, t, text_embedding, null_embedding, guidance_scale=7.5
):
"""
Classifier-Free Guidanceによるノイズ予測
model: U-Net（ノイズ予測ネットワーク）
z_t: ノイズ付き潜在表現 (batch, channels, h, w)
t: タイムステップ
text_embedding: テキスト埋め込み (batch, 77, 768)
null_embedding: 空テキストの埋め込み (batch, 77, 768)
guidance_scale: ガイダンススケール w
"""
# 条件付きノイズ予測
noise_cond = model(z_t, t, text_embedding)
# 無条件ノイズ予測
noise_uncond = model(z_t, t, null_embedding)
# CFGによる合成
noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond)
return noise_pred
実装は数行で済みますが、この線形結合の効果は絶大です。guidance_scale を上げることで、テキスト条件に沿ったノイズ予測が増幅され、プロンプトに忠実な画像が生成されるようになります。
Guidance Scaleの影響を可視化
guidance scaleが予測ノイズに与える影響を可視化してみましょう。
import torch
import numpy as np
import matplotlib.pyplot as plt
# ダミーのノイズ予測を生成
torch.manual_seed(0)
noise_uncond = torch.randn(1, 4, 8, 8) # 無条件予測
noise_cond = noise_uncond + 0.3 * torch.randn(1, 4, 8, 8) # 条件付き予測
# 異なるguidance scaleでの結果を計算
scales = [1.0, 3.0, 7.5, 15.0, 30.0]
results = []
for w in scales:
guided = noise_uncond + w * (noise_cond - noise_uncond)
results.append(guided)
# チャネル0の空間マップを可視化
fig, axes = plt.subplots(1, len(scales), figsize=(18, 3.5))
for idx, (w, result) in enumerate(zip(scales, results)):
spatial = result[0, 0].numpy() # チャネル0
vmax = max(abs(spatial.min()), abs(spatial.max()))
axes[idx].imshow(spatial, cmap='RdBu_r', vmin=-vmax, vmax=vmax)
axes[idx].set_title(f'w = {w}', fontsize=12)
axes[idx].axis('off')
plt.suptitle('Effect of Guidance Scale on Noise Prediction (Channel 0)',
fontsize=13, y=1.02)
plt.tight_layout()
plt.savefig('cfg_scale_effect.png', dpi=150, bbox_inches='tight')
plt.show()
# 各scaleでのノイズ予測のノルムを計算
print("Guidance Scale vs Noise Norm:")
for w, result in zip(scales, results):
norm = result.norm().item()
print(f" w = {w:5.1f}: ||noise|| = {norm:.2f}")
このグラフからは、guidance scaleが大きくなるにつれてノイズ予測の振幅が増大していく様子が観察されます。$w=1.0$（CFGなし）では穏やかなパターンですが、$w=7.5$（標準的な値）ではコントラストが強くなり、$w=30.0$ では極端に振幅が大きくなっています。ノルムの値からも、guidance scaleに比例してノイズ予測の大きさが増加していることが定量的に確認できます。実際のStable Diffusionでは、過大なguidance scaleは色の過飽和やアーティファクトを引き起こすため、$w=5.0 \sim 10.0$ の範囲が実用的です。
Cross-Attention Mapの統計的分析
最後に、Cross-Attention Mapのエントロピーを計算して、各画素の注意がどれだけ「集中」しているかを分析します。
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
torch.manual_seed(42)
# モデルとデータの準備
height, width = 16, 16
channels = 320
context_dim = 768
seq_len = 77
cross_attn = CrossAttention(query_dim=channels, context_dim=context_dim, num_heads=8)
z = torch.randn(1, channels, height, width)
context = torch.randn(1, seq_len, context_dim)
with torch.no_grad():
_ = cross_attn(z, context)
attn_map = cross_attn.attention_map[0] # (h*w, seq_len)
# 各画素のAttention分布のエントロピーを計算
# H(p) = -Σ p_j log(p_j)
attn_np = attn_map.numpy()
epsilon = 1e-10
entropy = -np.sum(attn_np * np.log(attn_np + epsilon), axis=1)
entropy_map = entropy.reshape(height, width)
# 最大エントロピー（一様分布の場合）
max_entropy = np.log(seq_len)
fig, axes = plt.subplots(1, 2, figsize=(12, 5))
# エントロピーマップ
im0 = axes[0].imshow(entropy_map, cmap='viridis')
axes[0].set_title('Attention Entropy Map', fontsize=12)
axes[0].set_xlabel('Width')
axes[0].set_ylabel('Height')
plt.colorbar(im0, ax=axes[0], label='Entropy (nats)')
# エントロピーのヒストグラム
axes[1].hist(entropy.flatten(), bins=30, color='steelblue', edgecolor='white', alpha=0.8)
axes[1].axvline(max_entropy, color='red', linestyle='--', label=f'Max entropy = {max_entropy:.2f}')
axes[1].set_xlabel('Entropy (nats)', fontsize=11)
axes[1].set_ylabel('Count', fontsize=11)
axes[1].set_title('Distribution of Attention Entropy', fontsize=12)
axes[1].legend()
plt.tight_layout()
plt.savefig('attention_entropy.png', dpi=150, bbox_inches='tight')
plt.show()
print(f"エントロピーの統計:")
print(f" 平均: {entropy.mean():.3f}")
print(f" 最小: {entropy.min():.3f}")
print(f" 最大: {entropy.max():.3f}")
print(f" 最大理論値（一様分布）: {max_entropy:.3f}")
このエントロピー分析から、Cross-Attentionの注意パターンの「集中度」がわかります。エントロピーが低い画素は少数のトークンに強く注目しており、エントロピーが高い画素は多くのトークンに均等に注目しています。未学習のモデルではエントロピーが最大理論値に近い値をとりますが、学習済みモデルでは、オブジェクトが描かれる領域の画素は対応するテキストトークンにエントロピーが低い（集中した）注意パターンを示し、背景領域は比較的エントロピーの高い（分散した）パターンを示す傾向があります。このことから、Cross-Attentionがテキストの各概念を画像の適切な領域にマッピングする「空間的なルーティング」として機能していることがわかります。
まとめ
本記事では、Stable DiffusionにおけるCross-Attentionの仕組みを、全体アーキテクチャの中での位置付けから数学的定式化、Attention Mapの解釈、Classifier-Free Guidance、そしてPyTorch実装まで解説しました。
- Cross-Attentionの核心: Query=画像潜在表現、Key/Value=テキスト埋め込みとして、画像の各画素がテキストのどのトークンに注目すべきかを学習する。これにより、テキストプロンプトの意味が画像の適切な空間位置に注入される
- 数学的な仕組み: Scaled Dot-Product Attentionにより、画像特徴とテキスト特徴の類似度を計算し、softmaxで正規化した注意重みでテキスト情報を加重平均する。$\sqrt{d}$ によるスケーリングがsoftmaxの安定性を保証する
- Attention Mapの解釈: 各テキストトークンが画像のどの領域に影響を与えているかを可視化できる。解像度レベルによって粗い構図決定から細部の描画まで異なる役割を担う
- Classifier-Free Guidance: 条件付きと無条件のノイズ予測の線形結合により、テキストへの忠実度を制御する。guidance scaleはCross-Attentionの効果を増幅する役割を果たす
- Multi-Head構造: 複数のヘッドが異なる観点からテキスト情報を参照することで、テキストの多面的な意味を画像生成に反映する
Cross-Attentionの理解は、近年急速に発展しているPrompt-to-Prompt画像編集やControlNetなどの技術の基盤となります。これらの手法はCross-Attentionのマップを直接操作することで、画像の構図制御や局所的な編集を実現しています。
次のステップとして、以下の記事も参考にしてください。