Stable Diffusion 是如何工作的？

作为当前最流行的开源文本到图像生成模型，Stable Diffusion通过独特的扩散过程（Diffusion Process）和潜在空间编码（Latent Space Encoding）技术，实现了高效且高质量的图像生成。其核心工作机制可分解为四个关键阶段：前向扩散（Forward Diffusion）、噪声预测（Noise Prediction）、反向去噪（Reverse Denoising）和潜在空间解码（Latent Space Decoding）。以下从技术原理、数学基础和工程实现三个维度展开详细分析。

一、前向扩散：从清晰图像到噪声的渐进破坏

前向扩散是Stable Diffusion训练阶段的基础过程，其本质是通过逐步添加高斯噪声（Gaussian Noise）将原始图像转化为纯噪声。这一过程遵循马尔可夫链（Markov Chain）规则，每一步仅依赖前一步的状态。

1.1 数学建模

设原始图像为 ( x0 )，经过 ( T ) 步扩散后得到纯噪声 ( x_T )，每一步的噪声添加公式为：
[
q(x_t | x{t-1}) = \mathcal{N}(xt; \sqrt{1-\beta_t}x{t-1}, \betat \mathbf{I})
]
其中 ( \beta_t ) 为时间步 ( t ) 的噪声调度参数，( \mathcal{N} ) 表示正态分布。通过重参数化技巧，可直接从 ( x_0 ) 采样任意时间步 ( t ) 的中间状态：
[
x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, \quad \epsilon \sim \mathcal{N}(0, \mathbf{I})
]
这里 ( \bar{\alpha}_t = \prod{i=1}^t (1-\beta_i) )，该公式将扩散过程简化为单步计算，显著提升了训练效率。

1.2 工程实现

在代码层面，前向扩散通常通过预定义的噪声调度表（Noise Schedule）实现。例如，使用线性调度时，( \beta_t ) 从 ( 0.0001 ) 线性增长到 ( 0.02 )，共 ( 1000 ) 步。以下是一个简化的PyTorch实现示例：

import torch
def forward_diffusion(x0, T=1000, beta_start=1e-4, beta_end=2e-2):
    betas = torch.linspace(beta_start, beta_end, T)
    alphas = 1. - betas
    alpha_bars = torch.cumprod(alphas, dim=0)
    t = torch.randint(0, T, (x0.shape[0],)).to(x0.device)
    sqrt_alpha_bar = torch.sqrt(alpha_bars[t])[:, None, None, None]
    sqrt_one_minus_alpha_bar = torch.sqrt(1 - alpha_bars[t])[:, None, None, None]
    epsilon = torch.randn_like(x0)
    xt = sqrt_alpha_bar * x0 + sqrt_one_minus_alpha_bar * epsilon
    return xt, epsilon, t

此代码展示了如何从原始图像 ( x0 ) 生成带噪声的中间状态 ( xt )，并返回对应的噪声 ( \epsilon ) 和时间步 ( t )。

二、噪声预测：U-Net的核心作用

噪声预测是Stable Diffusion反向去噪过程的关键，由U-Net架构的神经网络完成。该网络需根据带噪声的图像 ( xt ) 和时间步 ( t ) 预测添加的噪声 ( \epsilon )。

2.1 U-Net架构设计

Stable Diffusion的U-Net包含以下核心组件：

• 编码器-解码器结构：通过下采样（Downsampling）和上采样（Upsampling）逐步提取多尺度特征。
• 残差连接（Residual Connections）：缓解深层网络的梯度消失问题。
• 注意力机制（Attention Mechanism）：在特定层引入交叉注意力（Cross-Attention），将文本条件（如提示词）嵌入到图像特征中。

2.2 时间步嵌入

由于扩散过程是时间相关的，需将时间步 ( t ) 编码为向量并注入网络。通常采用正弦位置编码（Sinusoidal Position Encoding）的变体：

def timestep_embedding(timesteps, dim=64):
    half_dim = dim // 2
    timesteps = timesteps.float()
    exponent = -torch.arange(0, half_dim, device=timesteps.device).float() * (4.0 / dim)
    emb = torch.exp(exponent) * timesteps[..., None]
    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
    return emb

此函数将离散的时间步映射为高维向量，使其可与图像特征进行拼接或加权融合。

2.3 损失函数

噪声预测的优化目标是最小化预测噪声 ( \hat{\epsilon} ) 与真实噪声 ( \epsilon ) 的均方误差（MSE）：
[
\mathcal{L} = \mathbb{E}{x_0, \epsilon, t} \left[ |\epsilon - \hat{\epsilon}\theta(xt, t)|^2 \right]
]
其中 ( \hat{\epsilon}\theta ) 为U-Net模型，( \theta ) 为可学习参数。

三、反向去噪：从噪声到图像的渐进重建

反向去噪是生成图像的核心阶段，通过逐步去除噪声将纯噪声 ( x_T ) 转化为清晰图像 ( x_0 )。这一过程依赖于训练好的U-Net模型。

3.1 去噪步骤

给定纯噪声 ( xT )，反向去噪通过以下公式迭代生成图像：
[
x{t-1} = \frac{1}{\sqrt{\alphat}} \left( x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \hat{\epsilon}\theta(x_t, t) \right) + \sqrt{\beta_t} z
]
其中 ( z \sim \mathcal{N}(0, \mathbf{I}) ) 为可选的随机噪声，用于增加生成多样性。在实际应用中，通常设置 ( z=0 ) 以实现确定性生成。

3.2 采样器选择

反向去噪的效率与质量高度依赖采样器（Sampler）的选择。常见采样器包括：

• DDPM（Denoising Diffusion Probabilistic Models）：标准采样器，需 ( T ) 步完成去噪。
• DDIM（Denoising Diffusion Implicit Models）：非马尔可夫采样器，可通过少量步骤（如50步）生成高质量图像。
• PLMS（Pseudo Linear Multistep）：改进的DDIM变体，进一步加速收敛。

以下是一个简化的DDIM采样实现：

def ddim_sample(model, x_T, T=1000, eta=0.0):
    betas = torch.linspace(1e-4, 2e-2, T)
    alphas = 1. - betas
    alpha_bars = torch.cumprod(alphas, dim=0)
    x = x_T
    for t in reversed(range(T)):
        t_ = torch.full((x.shape[0],), t, device=x.device)
        alpha_t = alphas[t]
        alpha_bar_t = alpha_bars[t]
        alpha_bar_t_prev = alpha_bars[t-1] if t > 0 else 1.0
        sigma_t = eta * torch.sqrt((1 - alpha_bar_t_prev) / (1 - alpha_bar_t)) * torch.sqrt(1 - alpha_bar_t / alpha_bar_t_prev)
        pred_noise = model(x, t_)
        if t == 0:
            x = torch.sqrt(alpha_bar_t_prev) * (x - torch.sqrt(1 - alpha_bar_t_prev) * pred_noise / torch.sqrt(1 - alpha_bar_t))
        else:
            z = torch.randn_like(x) if sigma_t > 0 else 0.0
            x = torch.sqrt(alpha_bar_t_prev) * (x - torch.sqrt(1 - alpha_bar_t_prev) * pred_noise / torch.sqrt(1 - alpha_bar_t)) + torch.sqrt(1 - alpha_bar_t_prev - sigma_t**2) * pred_noise + sigma_t * z
    return x

此代码展示了如何通过DDIM采样器在少量步骤内完成去噪。

四、潜在空间编码：VAE的压缩与解压

Stable Diffusion通过变分自编码器（VAE）将图像压缩到潜在空间，显著降低了计算复杂度。

4.1 VAE架构

VAE由编码器（Encoder）和解码器（Decoder）组成：

• 编码器：将512×512的RGB图像压缩为64×64的潜在特征图，压缩比达64倍。
• 解码器：将潜在特征图还原为原始图像。

4.2 训练目标

VAE的优化目标包括重构损失（Reconstruction Loss）和KL散度（KL Divergence）：
[
\mathcal{L}{\text{VAE}} = \mathbb{E}{x \sim p{\text{data}}} \left[ -\log p\theta(x|z) \right] + \beta D{\text{KL}}(q\phi(z|x) | p(z))
]
其中 ( q\phi(z|x) ) 为编码器，( p\theta(x|z) ) 为解码器，( p(z) ) 为标准正态分布。

4.3 工程意义

潜在空间编码使Stable Diffusion在64×64的分辨率下进行扩散过程，而非原始的512×512。这一设计将内存占用从约24GB降至约3GB，使得在消费级GPU（如NVIDIA RTX 3060）上运行成为可能。

五、实际应用建议

对于开发者部署Stable Diffusion，以下建议可提升效率与质量：

1. 硬件选择：优先使用具有足够显存（≥12GB）的GPU，如NVIDIA A100或RTX 4090。
2. 采样器优化：生产环境推荐使用DDIM或PLMS采样器，步数设置为20-50步。
3. 文本编码：采用CLIP模型的文本编码器，确保提示词与图像特征的精准对齐。
4. 超参数调优：调整噪声调度参数（如 ( \beta ) 的线性/余弦调度）以平衡生成速度与质量。

六、总结

Stable Diffusion的工作机制可概括为：通过前向扩散将图像转化为噪声，利用U-Net预测噪声，再通过反向去噪重建图像，最后通过VAE解码潜在空间特征。其核心优势在于潜在空间编码带来的高效计算，以及扩散模型提供的稳定训练过程。对于开发者而言，理解这一流程不仅有助于优化部署，也为定制化模型开发（如特定领域图像生成）提供了理论基础。