Stable Diffusion Architecture

1. The VAE: Perceptual Compression

The Variational Autoencoder (VAE) is trained with a combination of reconstruction loss, KL-divergence regularization, and a patch-based adversarial loss. Given an image $x \in \mathbb{R}^{H \times W \times 3}$, the encoder $\mathcal{E}$ maps it to a latent distribution:

Explanation: The encoder spatially downsamples the image by factor $f$ (typically $f=8$) to produce a compact latent tensor $z$ with $c=4$ channels. Here, $H, W$ are pixel height/width; $h, w$ are latent height/width; $f$ is the downsampling factor; $c$ is the number of latent channels. The VAE decoder $\mathcal{D}$ inverts this: $\tilde{x} = \mathcal{D}(z)$. The VAE training objective is:

Explanation: The VAE is trained to reconstruct faithfully (first term), keep the latent distribution close to a standard Gaussian (KL term), and produce perceptually sharp outputs via an adversarial patch discriminator (third term). Where $\mathcal{L}_{\text{rec}}$ is a pixel-level reconstruction loss (L1 or perceptual LPIPS), $D_{\text{KL}}$ is the KL divergence, $\mathcal{L}_{\text{adv}}$ is the PatchGAN discriminator loss, and $\lambda_{\text{KL}}, \lambda_{\text{adv}}$ are loss weights (typically $10^{-6}$ and $0.5$). ---

2. The Forward Diffusion Process

The forward (noising) process incrementally adds Gaussian noise to the latent $z_0$ over $T$ timesteps according to a fixed Markov chain:

Explanation: At timestep $t$, the noisy latent $z_t$ is a mixture of the original latent $z_0$ (scaled by $\sqrt{\bar{\alpha}_t}$) and pure Gaussian noise, where $\bar{\alpha}_t \to 0$ as $t \to T$. Here $\bar{\alpha}_t = \prod_{s=1}^{t}(1 - \beta_s)$ is the cumulative noise schedule product; $\beta_s$ is the variance schedule at step $s$ (often cosine or linear); $T$ is the total number of timesteps (typically 1000 during training). This allows sampling $z_t$ directly from $z_0$ in one step — no sequential iteration needed during training. ---

3. The Denoising U-Net and Conditioning

The U-Net $\epsilon_\theta$ is trained to predict the noise $\epsilon$ added to $z_0$ to produce $z_t$, conditioned on timestep $t$ and text embedding $\tau_\theta(y)$:

Explanation: The U-Net learns to estimate which noise was added to the latent at timestep $t$, guided by the text prompt embedding, so it can be subtracted to recover a cleaner latent. Where $\epsilon_\theta$ is the denoising U-Net with parameters $\theta$; $z_t = \sqrt{\bar{\alpha}_t} z_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon$ is the noisy latent; $\tau_\theta(y)$ is the CLIP text encoder output for prompt $y$; and the expectation is over all training images, noise samples, timesteps, and prompts. ---

4. Classifier-Free Guidance

Stable Diffusion uses Classifier-Free Guidance (CFG) at inference to amplify prompt adherence without a separate classifier:

Explanation: The final noise prediction is the unconditional prediction shifted in the direction of the conditional prediction by guidance scale $s$, pushing the sample toward prompt-consistent regions of the latent distribution. Where $\epsilon_\theta(z_t, t, \varnothing)$ is the unconditional noise prediction (null-text condition); $\epsilon_\theta(z_t, t, y)$ is the text-conditioned prediction; $s > 1$ is the guidance scale (typically 7.5); and $\tilde{\epsilon}_\theta$ is the guided noise estimate used for the denoising update. Higher $s$ increases prompt fidelity at the cost of diversity and potential oversaturation. ---

PyTorch Sketch: The Core Inference Loop

```python import torch from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler from transformers import CLIPTextModel, CLIPTokenizer

# Load components vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="vae") unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="unet") text_encoder = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="text_encoder") tokenizer = CLIPTokenizer.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="tokenizer") scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="scheduler") device = "cuda" guidance_scale = 7.5 num_inference_steps = 50

# 1. Encode text prompt prompt = "a photorealistic mountain lake at sunset" tokens = tokenizer([prompt, ""], padding="max_length", return_tensors="pt").to(device) text_embeddings = text_encoder(tokens.input_ids)[0] # [2, 77, 1024]

# 2. Initialize random latent latents = torch.randn(1, 4, 64, 64, device=device) # z_T ~ N(0, I) scheduler.set_timesteps(num_inference_steps)

# 3. Iterative denoising in latent space for t in scheduler.timesteps: latent_input = torch.cat([latents] * 2) # duplicate for CFG with torch.no_grad(): noise_pred = unet(latent_input, t, encoder_hidden_states=text_embeddings).sample

# Classifier-free guidance (Eq. 5) noise_uncond, noise_cond = noise_pred.chunk(2) noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond) latents = scheduler.step(noise_pred, t, latents).prev_sample

# 4. Decode latent → pixel space latents = latents / vae.config.scaling_factor with torch.no_grad(): image = vae.decode(latents).sample # [1, 3, 512, 512] image = (image / 2 + 0.5).clamp(0, 1) ``` The `vae.config.scaling_factor` (0.18215 for SD 1.x) normalizes the latent distribution to unit variance before denoising — a subtle but critical detail that causes NaN artifacts if omitted. ---

Cross-Attention: How Text Controls the Image

Text conditioning enters the U-Net via cross-attention layers inserted between the convolutional residual blocks. The spatial features of the U-Net serve as queries, and the CLIP text embeddings serve as keys and values: ```python

# Simplified cross-attention inside U-Net ResBlock class CrossAttention(nn.Module): def __init__(self, query_dim, context_dim, heads=8, head_dim=64): super().__init__() inner_dim = heads * head_dim self.to_q = nn.Linear(query_dim, inner_dim, bias=False) self.to_k = nn.Linear(context_dim, inner_dim, bias=False) self.to_v = nn.Linear(context_dim, inner_dim, bias=False) self.to_out = nn.Linear(inner_dim, query_dim) self.heads = heads self.scale = head_dim ** -0.5 def forward(self, x, context):

# x: [B, HW, C_spatial] context: [B, 77, C_text] Q = self.to_q(x) # spatial features → queries K = self.to_k(context) # text tokens → keys V = self.to_v(context) # text tokens → values

# Scaled dot-product attention attn = torch.softmax(Q @ K.transpose(-1, -2) * self.scale, dim=-1) return self.to_out(attn @ V) ``` This is why prompt token order and length matter: each spatial position in the U-Net attends to all 77 CLIP tokens, and the attention weights directly determine which text concepts influence which image regions.