Stable Diffusion 的完整流程#

分主题阅读#

• 条件生成与 Guidance：条件信息怎样改变去噪方向，guidance scale 又控制什么？
• 文本编码与条件注入：文本向量怎样进入图像生成网络？
• ControlNet：如何给预训练生成模型加入空间条件？
• LoRA 原理：如何用低秩增量减少需要训练的参数？

组件阅读路线 · Classifier Guidance → CFG → Latent Diffusion → SD 全栈

1. Latent Diffusion Model：为什么在隐空间做扩散#

1.1 像素空间扩散的瓶颈#

一张 $512\times512\times3$ 的 RGB 图像有 $512 \times 512 \times 3 = 786{,}432$ 个像素值。U-Net 每一步都要处理这 78 万多个值——计算量巨大、显存占用高。

实际数据：在像素空间训练 256×256 的 DDPM，一张 A100（80G）只能跑 batch_size=8；而在隐空间做同样分辨率的扩散，batch_size 可以到 64+。

1.2 Latent Diffusion 的解决方案#

Rombach et al., 2022 (LDM 论文) 提出：先用预训练的 VAE 把图像压缩到低维隐空间，在隐空间做扩散，最后再解码回像素空间。

压缩比例#

$$512 \times 512 \times 3 \xrightarrow{\text{VAE Encoder}} 64 \times 64 \times 4$$

空间维度缩小 $8\times$（$512/64=8$），通道从 3 变 4

维度对比：

1.3 为什么可以压缩这么多？#

VAE 编码器学到了一种去除高频冗余的压缩方式。自然图像中大量像素与相邻像素高度相关（想象一片蓝天，连续几百个像素的值几乎一样）。VAE 保留了语义关键信息，丢掉了人眼不敏感的高频细节。

JPEG 类比

JPEG 也是去掉高频信息来压缩——先 DCT 变换，再量化高频系数。但 VAE 学到的压缩是端到端可微的，且保留了重建所需的信息，而非固定的频率截断。可以理解为"学来的、更聪明的 JPEG"。

1.4 LDM 的三阶段流程#

1. 编码（VAE Encoder）：$x \in \mathbb{R}^{H\times W\times3} \rightarrow z \in \mathbb{R}^{h\times w\times4}$，把图像压到隐空间
2. 扩散（U-Net in latent space）：在 $z$ 上做加噪 → 去噪，只操作 $h \times w \times 4$ 的张量
3. 解码（VAE Decoder）：$z_0 \rightarrow \hat{x}$，把去噪后的隐变量恢复为像素图像

Scaling Factor 0.18215#

SD 1.x/2.x 的代码中随处可见 z = z * 0.18215。这个"magic number"$= 1/\text{std}(z_\text{train})$，即训练集经 VAE 编码后隐变量标准差的倒数。乘上它让 $z$ 的方差接近 1，与扩散过程假设的 $\mathcal{N}(0,I)$ 初始分布对齐。SDXL 起改用可学习的 scaling factor，不再硬编码。

**面试要点：**在隐空间做扩散的好处不仅仅是"快"——隐空间天然是更语义化的表示（高频噪声已被过滤），扩散模型可以更高效地学习数据的语义结构，而非浪费容量在像素级细节上。

2. Stable Diffusion 完整推理流程#

2.1 端到端 Pipeline（带具体维度）#

以 SD 1.5 生成 512×512 图像为例：

Step 1：文本编码#

prompt = "a cat sitting on a beach"
tokens = tokenizer(prompt)           # (1, 77)
text_emb = clip_encoder(tokens)      # (1, 77, 768)

Step 2：初始化随机噪声#

z_T = torch.randn(1, 4, 64, 64)     # 隐空间噪声

Step 3：迭代去噪（以 50 步 DDIM 为例）#

for t in scheduler.timesteps:        # [999, 979, ..., 19, 0]

    # CFG 需要跑两次 U-Net
    z_t_input = torch.cat([z_t, z_t])       # (2, 4, 64, 64)  拼成 batch
    text_input = torch.cat([null_emb, text_emb])  # (2, 77, 768)

    noise_pred = unet(z_t_input, t, text_input)  # (2, 4, 64, 64)
    noise_uncond, noise_cond = noise_pred.chunk(2)

    # CFG 公式
    noise_pred = noise_uncond + w * (noise_cond - noise_uncond)

    # Scheduler step（DDIM / PNDM / Euler 等）
    z_t = scheduler.step(noise_pred, t, z_t)     # (1, 4, 64, 64)

Step 4：VAE 解码#

z_0 = z_t / 0.18215                 # 反 scaling
image = vae.decode(z_0)              # (1, 3, 512, 512)
image = (image + 1) / 2             # [-1,1] → [0,1]

CFG 让计算量翻倍

每一步都要跑两次 U-Net——一次有条件（$c = \text{text\_emb}$），一次无条件（$c = \varnothing$）。实际中通过 batch=2 拼接一起前向，减少 kernel launch 开销，但计算量确实是 2 倍。如果用了 negative prompt，也不会额外增加开销（只是把 null embedding 替换为 negative embedding）。

2.2 完整推理伪代码#

def stable_diffusion_inference(prompt, neg_prompt="", steps=50, w=7.5):
    # 1. 文本编码
    cond_emb   = clip_encode(prompt)       # (1, 77, 768)
    uncond_emb = clip_encode(neg_prompt)   # (1, 77, 768)

    # 2. 初始化噪声
    latent = randn(1, 4, 64, 64) * scheduler.init_noise_sigma

    # 3. 去噪循环
    for t in scheduler.timesteps:
        latent_input = cat([latent, latent], dim=0)
        text_input   = cat([uncond_emb, cond_emb], dim=0)

        noise_pred = unet(latent_input, t, encoder_hidden_states=text_input)
        noise_u, noise_c = noise_pred.chunk(2)
        noise_pred = noise_u + w * (noise_c - noise_u)

        latent = scheduler.step(noise_pred, t, latent).prev_sample

    # 4. 解码
    latent = latent / 0.18215
    image  = vae.decode(latent).sample
    image  = (image.clamp(-1, 1) + 1) / 2
    return image

2.3 各组件参数量#

3. 采样器演进与噪声预测策略#

3.1 采样器（Scheduler）的角色#

扩散模型的训练只学了一件事：给定 $(x_t, t)$，预测噪声 $\epsilon$。但从 $x_T$（纯噪声）到 $x_0$（干净图像）的去噪轨迹由采样器决定。不同采样器 = 不同的 ODE/SDE 求解策略，直接影响生成质量和所需步数。

3.2 主流采样器对比#

3.3 DDIM：为什么可以跳步？#

原始 DDPM 使用逐步随机逆链，但马尔科夫性本身不禁止减少时间步；重新定义采样网格和对应转移也可实现 respacing。DDIM 重新推导了一个非马尔科夫的后验分布，使得 $x_{t-\Delta}$ 可以直接从 $x_t$ 计算得到，跳过中间步骤。

$$x_{t-\Delta} = \sqrt{\bar\alpha_{t-\Delta}} \cdot \hat{x}_0(x_t) + \sqrt{1-\bar\alpha_{t-\Delta}-\sigma_t^2} \cdot \epsilon_\theta(x_t, t) + \sigma_t \cdot z$$

$\sigma_t = 0$ 时完全确定性（DDIM），$\sigma_t > 0$ 时引入随机性。$\hat{x}_0$ 是用当前噪声预测反推的干净图估计。

DDIM 的两个关键性质

• 确定性采样：$\sigma_t=0$ 时同一个初始噪声 → 同一张图。这使得 latent 空间的插值和编辑成为可能
• 步数灵活：可以从 1000 步的 schedule 中均匀取 20–50 个子步（subsequence），质量损失很小

3.4 DPM-Solver++：为什么更快？#

把去噪看作 ODE $dx = f(x, t)\,dt$：DDIM 是一阶 Euler 法（每步用当前点的切线外推），DPM-Solver++ 是二阶/三阶多步法（利用前几步的 $\epsilon$ 预测做更精确的外推）。类比数值分析：Euler 法 → Runge-Kutta 法，相同精度需要的步数更少。

3.5 $\epsilon$-prediction 与 v-prediction#

噪声预测网络的输出可以是不同的"参数化目标"：

v-prediction 的直觉

高噪声端 $x_t$ 已接近噪声，预测 $\epsilon$ 本身未必困难；问题之一是转换 $\hat x_0=(x_t-\sigma_t\hat\epsilon)/\alpha_t$ 时会放大误差。$v=\alpha_t\epsilon-\sigma_t x_0$ 提供另一种参数化，在 $\alpha_t^2+\sigma_t^2=1$ 时可用 $\hat x_0=\alpha_t x_t-\sigma_t\hat v$ 转换。scheduler 必须与 checkpoint 的 prediction type 匹配；SD 2.x 也要区分具体 checkpoint，不能只凭系列名判断。

**面试要点：**采样器不改变模型参数，只改变去噪轨迹。选择采样器是 quality-speed tradeoff：少步高阶求解器可能明显降低 NFE，但具体质量须在同一模型和设置下比较。$\epsilon$/v/velocity prediction 是训练目标的不同参数化，影响收敛性和兼容性。

4. Stable Diffusion 各版本演进#

SD 3 的架构变革

SD 3 用 DiT（Diffusion Transformer） 替代了 U-Net，用 Flow Matching 替代了 DDPM 的噪声调度。这是扩散模型从"卷积时代"迈向"Transformer 时代"的标志——与 LLM 的 scaling law 对齐，更容易扩展到更大模型。

MMDiT（Multi-Modal DiT）让文本和图像 token 在同一个 Transformer 中做双向 attention，取代了 U-Net 中"先 self-attn 再 cross-attn"的分离设计。

4.1 为什么 SDXL 的 CFG 更低？#

SDXL 使用更大的 U-Net（2.6B vs 860M）和双文本编码器，模型本身的条件理解能力更强。因此不需要很高的 $w$ 就能生成高质量的条件图像。过高的 $w$ 反而容易导致过饱和。

**面试要点：**SD 版本演进的主线是——更大的模型、更强的文本编码器、更高的分辨率。SD 3 的架构转向（U-Net → DiT，DDPM → Flow Matching）是新的趋势，很可能成为面试热点。

5. 理解检查#

5.1 Q1CFG 的公式是什么？为什么 $w > 1$ 有效？#

参考答案

CFG 公式：$\tilde{\epsilon}_\theta = \epsilon_\theta(x_t, t, \varnothing) + w \cdot [\epsilon_\theta(x_t, t, c) - \epsilon_\theta(x_t, t, \varnothing)]$

$w > 1$ 有效的原因：$\epsilon_\theta(x_t, t, c) - \epsilon_\theta(x_t, t, \varnothing)$ 可以理解为"条件信号的方向"。$w > 1$ 放大了这个方向，相当于在 score function 层面增大了条件对数似然的梯度——让采样轨迹更强烈地偏向条件分布的高密度区域。从 quality-diversity tradeoff 角度看，$w > 1$ 牺牲了分布的多样性来换取单个样本更高的条件匹配度。

5.2 Q2Classifier Guidance 和 CFG 的区别？#

参考答案

5.3 Q3为什么在隐空间做扩散而不是像素空间？#

参考答案

三个层面的原因：

1. 计算效率：512×512×3 = 786K 维压缩到 64×64×4 = 16K 维（约 48 倍压缩），每步 U-Net 计算量大幅减少
2. 语义更纯净：VAE 编码器过滤掉了人眼不敏感的高频信息，隐空间更"语义化"。扩散模型不用浪费建模能力学习像素级噪声模式
3. 训练稳定性：隐空间维度低、分布更规律（有 KL 正则约束接近高斯），扩散模型更容易收敛

5.4 Q4Cross-attention 在 SD 中怎么工作？#

参考答案

Cross-attention 实现"文本 → 图像"的信息注入。Q 来自图像特征（$Q = W_Q z$），K 和 V 来自文本 embedding（$K = W_K c$，$V = W_V c$）。注意力权重矩阵 $\text{softmax}(QK^T/\sqrt{d})$ 的大小是 [空间位置数 × 文本token数]，它学会了每个空间位置应该关注哪些文本 token。例如，图像中"猫"所在的位置会对 "cat" token 有更高的注意力权重，从而接收到猫的语义信息来引导生成。

5.5 Q5Negative prompt 的原理是什么？#

参考答案

标准 CFG 用空 embedding $\varnothing$ 作为无条件基线。Negative prompt 把这个基线替换为负面条件 $c_\text{neg}$ 的 embedding。公式变为：

$\tilde{\epsilon} = \epsilon_\theta(x_t, t, c_\text{neg}) + w \cdot [\epsilon_\theta(x_t, t, c_\text{pos}) - \epsilon_\theta(x_t, t, c_\text{neg})]$

效果是引导方向从"无条件 → 正面条件"变成了"负面条件 → 正面条件"——不仅要靠近正面描述，还要远离负面描述。

5.6 Q6SD 的 VAE 和训练时的 VAE 是什么关系？#

参考答案

SD 的 VAE 是预训练的、独立的模型。先在大规模图像数据上单独训练 VAE（重建损失 + KL 正则 + 感知损失 + 对抗损失），确保编解码质量。然后固定 VAE 参数，用 VAE 编码器把所有训练图像编码到隐空间，在隐空间上训练 U-Net。推理时先去噪得到 $z_0$，再用 VAE 解码器还原为像素图像。VAE 在扩散训练中参数是冻结的。

5.7 Q7CFG scale 太大会怎样？为什么？#

参考答案

CFG scale 太大会导致：颜色过饱和、对比度异常高、出现纹理伪影、图像看起来不自然。原因是 $w$ 过大时，$\tilde{\epsilon} = \epsilon_\text{uncond} + w \cdot (\epsilon_\text{cond} - \epsilon_\text{uncond})$ 的值域远超训练时 $\epsilon \sim \mathcal{N}(0,I)$ 的范围（大部分元素应在 $[-3, 3]$）。过大的噪声预测值让 scheduler 计算出的 $z_{t-1}$ 偏离隐空间的正常分布，VAE 解码器在这种异常输入下产生视觉伪影。

5.8 Q8SD 各版本的主要区别？#

2022.08

SD 1.4 / 1.5

• CLIP ViT-L/14
• 860M U-Net
• 512×512
• 4 通道 VAE latent
• 社区生态最丰富的基线

→

2022.11

SD 2.x

• OpenCLIP ViT-H/14
• 768×768 支持
• 过滤 NSFW 数据
• 写实人像退步，生态受限
• 社区大多跳过此版本

→

2023.07

SDXL

• CLIP ViT-L + OpenCLIP ViT-bigG
• 2.6B U-Net
• 原生 1024×1024
• Base + Refiner 两阶段
• 质量飞跃，生态回暖

→

2024.03

SD 3 / SD 3.5

• CLIP × 2 + T5-XXL
• U-Net → MMDiT
• DDPM → Flow Matching
• 16 通道 VAE
• 架构范式转变

主线：更大模型 · 更强文本编码器 · 更高分辨率 · 架构从 U-Net 走向 Transformer

参考答案

• SD 1.5：CLIP ViT-L/14，860M U-Net，512×512。社区最广泛的基线

• SD 2.x：使用 OpenCLIP ViT-H/14，并包含不同分辨率和 prediction type 的 checkpoint。训练数据筛选有变化，但不能仅据此解释社区采用率或特定能力差异。

• SDXL：双文本编码器（CLIP + OpenCLIP），2.6B U-Net，原生 1024×1024。引入 refiner 两阶段架构，质量飞跃。

  • Base 模型：从纯噪声去噪到中间状态（时间步 1000→200），负责生成全局构图和语义
  • Refiner 模型：接收 Base 的输出（t≈200 的带噪 latent），继续在同一条 DDIM 轨迹上从 200→0 精细去噪，专门负责细节、纹理、锐化
  • 类比：Base 是"打草稿"，Refiner 是"精修"——分工比一个模型包办全程效果更好
  • Refiner 可以跳过（直接用 Base 输出），但加上后细节质量明显提升，特别是人脸和纹理

  Base vs Refiner：核心差异

  维度	Base	Refiner
  训练时间步	均匀覆盖 0 ~ 1000	集中在 0 ~ 200（低噪声段）
  擅长	从纯噪声重建全局结构与语义	接近干净时精修纹理、边缘、人脸
  文本编码器	CLIP ViT-L + OpenCLIP ViT-bigG	仅 OpenCLIP ViT-bigG（精修阶段语义依赖减少）
  额外条件	—	Aesthetic Score 注入，推理时设为高分引导精修质量

  两者架构基本相同，本质差异是训练数据分布不同——Refiner 从未见过高噪声输入，因此对低噪声区域的感知更精准。

• SD 3：架构重构——U-Net → MMDiT，DDPM → Flow Matching，三重文本编码器（+T5-XXL），16通道 VAE。代表了下一代扩散模型的技术方向

面试加分建议

回答版本演进题时，重点讲为什么要这么改，而不是罗列参数。例如：

• 为什么用双编码器？→ 一个擅长短文本理解（CLIP），一个擅长长文本和复杂描述（OpenCLIP ViT-bigG / T5）
• 为什么 SD3 转 DiT？→ Transformer 的 scaling law 比 U-Net 更好；DiT 可以统一处理文本和图像 token
• 为什么 SD3 用 Flow Matching？→ 直线路径比 DDPM 的弯曲路径采样效率更高

5.9 Q9ControlNet 怎么做到不破坏预训练模型？#

参考答案

ControlNet 复制 U-Net encoder 作为独立分支，通过 zero-convolution（权重和偏置初始化为 0 的 1×1 卷积）连接到原 U-Net。训练初期 zero-conv 输出全 0，ControlNet 对原模型无影响；随着训练推进，逐渐学到有意义的空间控制信号。原 U-Net 全程冻结；在相同基座的部分微调 checkpoint 上可复用，但需验证结构、条件接口与生成效果。

5.10 Q10LoRA 的原理？为什么在 SD 社区这么流行？#

参考答案

LoRA 将权重更新分解为低秩矩阵：$\Delta W = BA$，$B \in \mathbb{R}^{d \times r}$，$A \in \mathbb{R}^{r \times d}$，$r \ll d$。只训练 A、B，原模型冻结。原因是微调的有效维度远低于参数空间维度（intrinsic dimensionality）。

在 SD 社区流行的原因：(1) 文件小（几十 MB vs 几 GB），方便分享；(2) 可叠加（多个 LoRA 线性组合实现多风格/人物混合）；(3) 可合并到原权重，不增加推理延迟；(4) 训练快，消费级 GPU 即可。

5.11 Q11DDIM 和 DDPM 有什么区别？为什么 DDIM 可以用更少步数？#

参考答案

原始 DDPM 按训练噪声网格逐步采样；马尔科夫性并不要求固定使用全部训练时间步。DDIM 重新推导了非马尔科夫的后验分布，使得可以从 $x_t$ 直接跳到任意 $x_{t-\Delta}$——本质上把随机过程变成了确定性 ODE，可以在 1000 步 schedule 中任意取子序列（如 50 步），用更大步长求解。同一个训练好的模型可以用 DDPM 或 DDIM 采样（只是采样策略不同，模型参数一样）。

6. 实现练习：CFG 推理 + VAE 重参数化 + EMA#

6.1 Classifier-Free Guidance 推理实现#

CFG 概念在第 3-4 节已详细讲过，但面试可能要求写代码。核心是每步做两次前向，再线性组合：

import torch

@torch.no_grad()
def sample_with_cfg(unet, vae, scheduler, text_emb, uncond_emb,
                     latent_shape, guidance_scale=7.5, num_steps=50):
    """
    Stable Diffusion 的 CFG 推理完整流程（diffusers 兼容写法）

    Args:
        unet:       去噪网络（U-Net 或 DiT），预测 ε
        vae:        VAE，用 .decode() 把 latent 解回像素
        scheduler:  噪声调度器（DDIM / Euler / DPM-Solver 等）
        text_emb:   (B, 77, 768) 文本条件嵌入（CLIP 编码）
        uncond_emb: (B, 77, 768) 无条件嵌入（空文本 "" 的编码）
        latent_shape:   (B, 4, H/8, W/8)
        guidance_scale: CFG 强度 w（w=1 无引导，w=7.5 标准）
        num_steps:  推理步数（SD 1.x 典型 20-50）
    """
    device = text_emb.device

    # 1. 初始化调度器：从训练时的 1000 步均匀下采样出 num_steps 个时间点
    #    例如 num_steps=50 → timesteps = [981, 961, 941, ..., 21, 1]
    #    ★ num_steps 就是在这里被真正消费的 ★
    scheduler.set_timesteps(num_steps, device=device)

    # 2. 初始化纯噪声 latent；部分调度器（Euler 等）要求乘以 init_noise_sigma
    #    DDIM 下 init_noise_sigma = 1.0，不影响；Euler 下 ≠ 1，必须乘
    latents = torch.randn(latent_shape, device=device)
    latents = latents * scheduler.init_noise_sigma           # [B, 4, 64, 64]

    # 3. 把条件和无条件 embedding 拼在一起，一次前向算两个
    #    工程优化：避免两次独立前向传播，用 batch 翻倍换速度
    text_emb_combined = torch.cat([uncond_emb, text_emb])    # [2B, 77, 768]

    # 4. 去噪循环：恰好跑 num_steps 次
    for t in scheduler.timesteps:
        # 输入 latent 复制一份（无条件 + 有条件）
        latent_input = torch.cat([latents, latents])         # [2B, 4, 64, 64]
        # 某些调度器（Euler/DPM-Solver）需要对输入做尺度归一化
        # DDIM 下是 identity；为了通用性统一调用
        latent_input = scheduler.scale_model_input(latent_input, t)

        # 一次前向，同时得到无条件和有条件的噪声预测
        noise_pred = unet(latent_input, t, text_emb_combined) # [2B, 4, 64, 64]
        noise_uncond, noise_cond = noise_pred.chunk(2)

        # ★ CFG 核心公式 ★
        # ε_guided = ε_uncond + w × (ε_cond - ε_uncond)
        noise_guided = noise_uncond + guidance_scale * (noise_cond - noise_uncond)

        # 调度器走一步：x_t → x_{t-1}（diffusers 返回 SchedulerOutput，取 .prev_sample）
        latents = scheduler.step(noise_guided, t, latents).prev_sample

    # 5. VAE 解码回像素空间（除以 0.18215 反归一化 latent）
    images = vae.decode(latents / 0.18215).sample            # [B, 3, 512, 512], 范围 [-1, 1]
    return images

num_steps 是怎么进入采样的？——scheduler 的两步初始化

这是面试里很容易被追问的点：DDPM 训练用 $T=1000$ 步，推理却只跑 20–50 步，这两者怎么对应起来？答案在 scheduler.set_timesteps(num_steps) 这一行：

1. 训练时：$t$ 从 $\{0, 1, ..., 999\}$ 均匀随机采样，每个 $t$ 都有对应的 $\bar\alpha_t$
2. 推理时：调用 set_timesteps(50) 后，调度器会从 1000 步里均匀下采样 50 个时间点存进 scheduler.timesteps，例如 [981, 961, 941, ..., 21, 1]
3. 循环直接 for t in scheduler.timesteps，走完恰好 50 步——num_steps 就是在这里被"花掉"的

这也是为什么 DDIM 能用少步数：它的 ODE 形式允许跨越多个训练时间步一次走一大步，随机 DDPM 采样也可以通过 respacing 减少步数，区别不只在“能否跳步”。init_noise_sigma 和 scale_model_input 是为了兼容 Euler / DPM-Solver 这些非 DDIM 调度器——它们用的是连续时间 σ 而非离散 t，初始噪声和输入都要按 σ 缩放。DDIM 常见配置下初始 sigma 为 1、输入缩放为恒等；保留这两处接口便于兼容，但 prediction type、额外参数与模型输入约定仍要单独核对。

CFG 代码要点

• 两次前向合并为一次：cat([uncond, cond]) 后 batch 翻倍，一次 forward 拿到两个输出，再 chunk(2) 拆开。实际推理时这是标准做法（显存换速度）。
• guidance_scale 的含义：$w=1$ 时公式退化为纯条件预测；$w>1$ 放大"条件 vs 无条件"的差异，越大越遵循文本但多样性越低。
• Negative Prompt：把 uncond_emb 替换为负面文本的编码即可。原理不变——CFG 公式中 $\epsilon_{uncond}$ 变成了"远离负面描述"的方向。

6.2 VAE 重参数化技巧（完整实现）#

面试可能问"VAE 怎么训练？重参数化技巧的代码怎么写？"：

import torch
import torch.nn as nn
import torch.nn.functional as F

class VAE(nn.Module):
    """简化的 VAE，展示重参数化技巧"""

    def __init__(self, in_ch=3, latent_ch=4):
        super().__init__()
        # Encoder: 图像 → (μ, log_σ²)
        self.encoder = nn.Sequential(
            nn.Conv2d(in_ch, 64, 4, 2, 1),   # 256→128
            nn.ReLU(),
            nn.Conv2d(64, 128, 4, 2, 1),     # 128→64
            nn.ReLU(),
            nn.Conv2d(128, 256, 4, 2, 1),    # 64→32
            nn.ReLU(),
        )
        # 输出 μ 和 log_σ² （各 latent_ch 通道）
        self.fc_mu      = nn.Conv2d(256, latent_ch, 3, 1, 1)
        self.fc_log_var = nn.Conv2d(256, latent_ch, 3, 1, 1)

        # Decoder: latent → 图像
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(latent_ch, 256, 4, 2, 1),
            nn.ReLU(),
            nn.ConvTranspose2d(256, 128, 4, 2, 1),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 4, 2, 1),
            nn.ReLU(),
            nn.Conv2d(64, in_ch, 3, 1, 1),
            nn.Tanh(),  # 输出 [-1, 1]
        )

    def encode(self, x):
        h = self.encoder(x)
        mu      = self.fc_mu(h)       # [B, 4, 32, 32]
        log_var = self.fc_log_var(h)   # [B, 4, 32, 32]
        return mu, log_var

    def reparameterize(self, mu, log_var):
        """
        ★ 重参数化技巧 ★

        问题：z ~ N(μ, σ²) 是随机采样，不可导，梯度传不回 encoder
        解决：z = μ + σ × ε，ε ~ N(0, I)
              把随机性"外包"给 ε，μ 和 σ 是确定性的，可以反向传播
        """
        std = torch.exp(0.5 * log_var)   # σ = exp(log_σ²/2)
        eps = torch.randn_like(std)      # ε ~ N(0, I)
        z = mu + std * eps               # ★ 核心：z = μ + σε
        return z

    def forward(self, x):
        mu, log_var = self.encode(x)
        z = self.reparameterize(mu, log_var)
        x_recon = self.decoder(z)
        return x_recon, mu, log_var

# ====== VAE 损失函数 ======
def vae_loss(x_recon, x, mu, log_var):
    """
    VAE 损失 = 重建损失 + KL 散度

    重建损失：让 decoder 输出尽量接近原图
    KL 散度：让 encoder 的输出分布接近标准正态 N(0,I)
    """
    # 重建损失（MSE 或 L1）
    recon_loss = F.mse_loss(x_recon, x, reduction='sum')

    # KL 散度（有解析解！不需要采样估计）
    # D_KL(N(μ,σ²) || N(0,1)) = -0.5 × Σ(1 + log(σ²) - μ² - σ²)
    kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())

    return recon_loss + kl_loss

面试追问准备：(1) 为什么用 log_var 而不是直接输出 $\sigma$？—— log 空间更数值稳定，$\sigma$ 必须 > 0 但 log_var 可以是任意实数；(2) KL 散度为什么有解析解？—— 两个高斯分布的 KL 有闭式公式，不需要蒙特卡洛估计；(3) SD 的 VAE 和这里的区别？—— SD 用的是 KL-VAE（更大的 encoder/decoder + skip connection），且加了感知损失和对抗损失来提升重建质量。

6.3 EMA（指数移动平均）：平滑模型权重#

EMA 是扩散训练中常用的参数平滑方法；收益需用同一训练运行的原始与 EMA 权重比较，不能把不同模型或训练长度的 FID 当作 EMA 消融：

import torch
import copy

class EMA:
    """
    指数移动平均 (Exponential Moving Average)

    核心思想：维护模型参数的"平滑版本"
      θ_ema = decay × θ_ema + (1 - decay) × θ_current

    为什么有效？
      1. 训练后期参数在最优解附近震荡，EMA 平均掉震荡 → 更稳定
      2. 相当于集成了训练过程中多个 checkpoint 的参数
      3. DiT 论文实验：不用 EMA FID=9.62，用 EMA FID=2.27（↓76%）

    典型 decay = 0.9999（即每步只更新 0.01% 的新参数）
    """

    def __init__(self, model, decay=0.9999):
        self.decay = decay
        # 深拷贝一份模型参数作为 EMA 参数
        self.ema_model = copy.deepcopy(model)
        self.ema_model.eval()
        self.ema_model.requires_grad_(False)  # EMA 不计算梯度

    @torch.no_grad()
    def update(self, model):
        """每个训练 step 后调用"""
        for ema_p, cur_p in zip(self.ema_model.parameters(),
                                model.parameters()):
            # ★ 核心更新公式 ★
            ema_p.data.mul_(self.decay).add_(cur_p.data, alpha=1 - self.decay)
            # 等价于: ema_p = 0.9999 * ema_p + 0.0001 * cur_p

    def forward(self, *args, **kwargs):
        """推理时用 EMA 模型"""
        return self.ema_model(*args, **kwargs)

# ====== 使用方式 ======
model = DiT_XL_2()
ema = EMA(model, decay=0.9999)

for images, labels in dataloader:
    loss = compute_loss(model, images, labels)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    ema.update(model)        # ← 每步更新 EMA

# 推理时用 ema.ema_model 而非 model
samples = sample(ema.ema_model, ...)

EMA 面试要点

• decay 的选择：0.9999 是常见示例值；应结合更新频率、训练长度与 warmup 选择。越接近 1 → 平滑窗口越大 → 响应越慢但越稳定。0.999 训练早期可以更快跟上变化。
• **为什么扩散模型特别需要 EMA？**扩散模型的损失函数是对所有时间步 $t$ 的期望，梯度方差本身就大。EMA 能有效平滑这种方差导致的参数震荡。
• 和学习率 warmup 的关系：二者互补。warmup 稳定训练早期，EMA 稳定训练后期。
• **推理时用哪个模型？**用 EMA 模型。训练模型只用来计算梯度更新 EMA。

6.4 Perceptual Loss（感知损失）#

SD 的 VAE 训练不只用 MSE，还用了感知损失。面试可能问"什么是 perceptual loss"：

import torchvision.models as models

class PerceptualLoss(nn.Module):
    """
    感知损失 / LPIPS

    思想：不在像素空间比较，而在预训练 VGG 的特征空间比较。
    因为人眼对高层语义（结构、纹理）敏感，对像素级差异不敏感。

    MSE: 两张图每个像素差的平方和 → 像素对齐但可能模糊
    Perceptual: VGG 特征差异 → 结构一致但允许像素级变化 → 更清晰
    """
    def __init__(self, layers=['relu1_2', 'relu2_2', 'relu3_3']):
        super().__init__()
        vgg = models.vgg16(pretrained=True).features
        vgg.eval()
        for p in vgg.parameters():
            p.requires_grad = False

        # 提取指定层的特征
        self.blocks = nn.ModuleList()
        layer_map = {'relu1_2': 4, 'relu2_2': 9, 'relu3_3': 16}
        prev = 0
        for name in layers:
            idx = layer_map[name]
            self.blocks.append(vgg[prev:idx])
            prev = idx

    def forward(self, x, target):
        loss = 0.0
        for block in self.blocks:
            x = block(x)
            target = block(target)
            loss += F.mse_loss(x, target)  # 特征空间的 MSE
        return loss

面试补充：SD 的 VAE 训练用了三种损失的组合——重建损失（L1/L2）+ 感知损失（VGG 特征匹配）+ 对抗损失（patch discriminator，类似 GAN）+ KL 正则化。这四项平衡让 VAE 既能高保真重建，又保持潜在空间的正则性。

6.5 Gradient Checkpointing 实现#

训练大模型（尤其 DiT-XL 或更大）时的显存优化技巧：

import torch
from torch.utils.checkpoint import checkpoint

class DiTWithCheckpointing(nn.Module):
    """
    Gradient Checkpointing：用时间换显存

    正常前向：保存每一层的中间激活值（用于反向传播）
      → 28 层 DiT-XL 的激活值占 ~30GB 显存

    Checkpointing：前向时不保存中间值，反向时重新计算
      → 显存从 O(N) 降到 O(√N)，但计算量增加 ~33%
    """
    def __init__(self, blocks, use_checkpoint=True):
        super().__init__()
        self.blocks = blocks
        self.use_checkpoint = use_checkpoint

    def forward(self, x, cond):
        for block in self.blocks:
            if self.use_checkpoint and self.training:
                # ★ 包一层 checkpoint，前向不存激活
                x = checkpoint(block, x, cond, use_reentrant=False)
            else:
                x = block(x, cond)
        return x

# 显存对比（DiT-XL/2，batch_size=256）：
# 无 checkpoint：~40 GB
# 有 checkpoint：~18 GB（省 55%，训练慢 30%）

面试答题：训练大模型的显存优化四件套

(1) 混合精度（AMP）：FP16/BF16 → 显存减半，速度加快

(2) Gradient Checkpointing：不存中间激活 → 显存 O(N)→O(√N)，计算+33%

(3) 梯度累积：小 batch 累积梯度模拟大 batch → 总显存不变但有效 batch 更大

(4) 模型并行：TP/PP/SP 分到多卡 → 单卡显存降低

原始资料：Latent Diffusion、ControlNet、Progressive Distillation、FLUX.1-dev。表中经验步数与 guidance 仅作示例，须以具体 checkpoint 的说明为准。