Stable Diffusion: Latent Diffusion Models, Cross-Attention U-Net, and PyTorch Architecture

Latent Space Compression and Perceptual Redundancy

Pixel-space generative diffusion models (e.g., DDPM, Imagen) execute the forward and reverse Markov diffusion chains directly across high-resolution image grids $\mathbf{x}\in {\mathbb{R}}^{3\times H\times W}$. In pixel space, high-frequency imperceptible details (micro-textures, sensor noise) consume the overwhelming majority of model capacity and computational complexity ($\mathcal{O}(H\cdot W)$ in convolutions and $\mathcal{O}((H\cdot W{)}^2)$ in spatial self-attention).

Latent Diffusion Models (LDM) (Rombach et al., CVPR 2022; commercialized as Stable Diffusion) decouple the generative process into two orthogonal phases:

1. Perceptual Compression: An autoencoder $\mathcal{E}$ compresses high-dimensional images $\mathbf{x}\in {\mathbb{R}}^{3\times 512\times 512}$ into a lower-dimensional latent space $\mathbf{z}=\mathcal{E}(\mathbf{x})\in {\mathbb{R}}^{4\times 64\times 64}$ ($48\times$ spatial-channel reduction), discarding high-frequency imperceptible variance.
2. Semantic Generative Diffusion: A conditional U-Net ${\mathbf{ϵ}}_{\theta }$ learns the reverse diffusion process over the regularized continuous latent manifold ${\mathbf{z}}_t$, conditioned on multimodal prompts (e.g., text, bounding boxes, depth maps) via cross-attention.

Generative Architectures Comparison

Mathematical Foundations

Figure 1: Stable Diffusion training pipeline compressing pixels to latents via a frozen VAE and training a cross-attention U-Net denoiser.

1. Two-Stage Perceptual Autoencoding (VAE)

The autoencoder consists of an encoder $\mathcal{E}$ and decoder $\mathcal{D}$:

$\mathbf{z}=\mathcal{E}(\mathbf{x}),\hat{\mathbf{x}}=\mathcal{D}(\mathbf{z})=\mathcal{D}(\mathcal{E}(\mathbf{x}))$

To prevent arbitrarily high-variance latent representations, the autoencoder is trained under a Kullback-Leibler (KL) divergence penalty alongside perceptual (LPIPS) and patch-based adversarial objectives:

${\mathcal{L}}_{\text{VAE}}(\varphi ,\psi )=\Vert \mathbf{x}-\hat{\mathbf{x}}{\Vert }_1+{\lambda }_{\text{perceptual}}{\mathcal{L}}_{\text{LPIPS}}(\mathbf{x},\hat{\mathbf{x}})+{\lambda }_{\text{adv}}{\mathcal{L}}_{\text{PatchGAN}}(\hat{\mathbf{x}})+{\lambda }_{\text{KL}}{\mathcal{D}}_{\text{KL}}\left(q_{\varphi }(\mathbf{z}\mid \mathbf{x})\parallel \mathcal{N}(0,\mathbf{I})\right)$

Once trained, the VAE weights are completely frozen. The latent representations are normalized by an empirical scaling factor ${\sigma }_{\text{scale}}=0.18215$ (for SD 1.5) to ensure unit variance $\mathrm{Var}(\mathbf{z})\approx 1.0$:

${\mathbf{z}}_0={\sigma }_{\text{scale}}\cdot \mathcal{E}(\mathbf{x}),{\mathbf{x}}_{\text{recon}}=\mathcal{D}\left(\frac{{\mathbf{z}}_0}{{\sigma }_{\text{scale}}}\right)$

2. Forward Latent Diffusion Process

The forward diffusion process defines a Markov chain corrupting the clean latent ${\mathbf{z}}_0$ with Gaussian noise over discrete timesteps $t\in \{1,\dots ,T\}$ according to variance schedule ${\beta }_1,\dots ,{\beta }_T$:

$q({\mathbf{z}}_t\mid {\mathbf{z}}_{t-1})=\mathcal{N}\left({\mathbf{z}}_t;\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1},{\beta }_t\mathbf{I}\right)$

Using the reparameterization trick with ${\alpha }_t=1-{\beta }_t$ and ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$:

${\mathbf{z}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{z}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$

3. Cross-Attention Conditional U-Net Denoising Objective

The reverse denoising process parameterizes a noise-prediction network ${\mathbf{ϵ}}_{\theta }({\mathbf{z}}_t,t,{\tau }_{\theta }(\mathbf{y}))$ using a time-conditioned U-Net. The training objective is the simplified mean squared error (MSE) score-matching loss:

${\mathcal{L}}_{\text{LDM}}(\theta )={\mathbb{E}}_{{\mathbf{z}}_0,\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I}),t,\mathbf{y}}\left[{\Vert \mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{z}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ},t,{\tau }_{\theta }(\mathbf{y})\right)\Vert }_2^2\right]$

where ${\tau }_{\theta }(\mathbf{y})\in {\mathbb{R}}^{L\times d_{\text{text}}}$ represents the sequence of token embeddings generated by the frozen CLIP text encoder for text prompt $\mathbf{y}$.

Spatial Cross-Attention Layer Formulation:

Inside the intermediate layers of the U-Net, spatial feature maps ${\varphi }_i({\mathbf{z}}_t)\in {\mathbb{R}}^{C\times H_i\times W_i}$ are flattened into visual tokens ${\mathbf{F}}_i\in {\mathbb{R}}^{(H_i{W}_i)\times C}$ and projected against text features ${\tau }_{\theta }(\mathbf{y})$:

$\mathbf{Q}={\mathbf{F}}_i{\mathbf{W}}_Q^{(i)},\mathbf{K}={\tau }_{\theta }(\mathbf{y}){\mathbf{W}}_K^{(i)},\mathbf{V}={\tau }_{\theta }(\mathbf{y}){\mathbf{W}}_V^{(i)}$

$\mathrm{CrossAttn}({\mathbf{F}}_i,{\tau }_{\theta }(\mathbf{y}))=\mathrm{softmax}\left(\frac{\mathbf{Q}{\mathbf{K}}^{\top }}{\sqrt{d_k}}\right)\mathbf{V}$

Figure 2: Spatial Cross-Attention mechanism routing semantic text features to spatial U-Net pixel coordinates.

4. Classifier-Free Guidance (CFG) Mathematical Derivation

To enforce prompt alignment without training auxiliary classifier networks, Classifier-Free Guidance (Ho & Salimans, 2022) trains a single model unconditionally with probability $p_{\text{uncond}}\approx 0.1$ by substituting prompt $\mathbf{y}$ with empty null tokens $\varnothing$.

During inference, the score estimate is extrapolated along the guidance vector:

${\widehat{\mathbf{ϵ}}}_{\theta }({\mathbf{z}}_t,t,\mathbf{c})={\mathbf{ϵ}}_{\theta }({\mathbf{z}}_t,t,\varnothing )+s\cdot \left[{\mathbf{ϵ}}_{\theta }({\mathbf{z}}_t,t,\mathbf{c})-{\mathbf{ϵ}}_{\theta }({\mathbf{z}}_t,t,\varnothing )\right]$

where $s\ge 1.0$ is the guidance scale parameter.

Score-Matching Equivalence:

Using Tweedie’s formula, the noise estimate relates directly to the data log-likelihood score: ${\mathbf{ϵ}}_{\theta }({\mathbf{z}}_t,t,\mathbf{c})\propto -{\sigma }_t{\nabla }_{{\mathbf{z}}_t}\log p({\mathbf{z}}_t\mid \mathbf{c})$. Substituting into the CFG equation: