隐空间与神经网络架构

0. 本讲导读

到目前为止，我们手里有一套完整的数学配方：

• Lecture 1 给出了流（flow）与扩散（diffusion）模型的定义：一个向量场 $\vf_t$ 通过 ODE $\dd{\Traj_t} = \vf_t(\Traj_t)\dd{t}$ 把简单分布 $\simple$ 推成数据分布 $\data$。
• Lecture 2 给出了训练算法：条件流匹配（conditional flow matching）损失 $\mathcal{L}_{\mathrm{CFM}}(\theta)=\E_{t,z,x}\norm{\vf_t^\theta(x)-\vf_t^{\text{target}}(x\mid z)}^2$， 用条件向量场当回归目标，却学到了边际向量场。
• Lecture 3 / 3b 给出了 score matching、SDE 扩展与 classifier-free guidance，让模型能听懂提示词。

这套配方里唯一还是「黑箱」的东西，就是 $\vf_t^\theta$ 这三个字母本身。它是一个神经网络，可它是什么样的神经网络？$t$ 是一个标量，$y$ 可能是一句话，$x$ 是一张 $1024\times 1024$ 的图；这三样东西怎么塞进同一个函数？输出又必须和 $x$ 一样大——这跟图像分类那种「越往后越窄」的网络完全不是一回事。

更要命的是尺度问题。一张 $600\times 1000$ 的彩色图，展平之后是

维。我们要在 $\R^{1.8\times 10^6}$ 上学一个向量场，还要沿着 ODE 反复调用它几十上百次。直接在像素空间干这件事，GPU 显存会爆炸。

本讲就是在补上这两块工程拼图：网络架构（§6.1，对应 U-Net 与 diffusion transformer）与隐空间（§6.2，对应 VAE），最后用 Stable Diffusion 3 和 Meta Movie Gen 两个真实系统（§6.3）把所有零件拼起来。附录 D 关于 VAE 的补充视角也一并讲透。

1. 从训练目标到具体网络

先把类型写清楚。带引导变量的向量场是一个函数

三个输入：数据点（或隐变量）$x\in\R^d$、时间 $t\in[0,1]$、条件 $y\in\mathcal{Y}$（类别标签、文本提示、或者 CFG 里的空标签 $\varnothing$）。$\theta$ 是参数。

为什么不能直接上 MLP

对二维玩具分布（前几讲实验里的那些同心圆、双月牙），一个 MLP 就够了：把 $x$、$y$ 的嵌入和 $t$ 拼接成一个向量，过几层全连接，输出 $\R^d$。讲义明确说了这一点：低维分布下 MLP 是充分的。

图像上为什么不行？三条理由，一条比一条致命：

1. 参数量爆炸。 一个从 $\R^d$ 到 $\R^d$ 的全连接层有 $d^2$ 个参数。$d=1.8\times 10^6$ 时 $d^2\approx 3.2\times 10^{12}$——一层就是 3 万亿参数。
2. 没有权重共享。 MLP 不知道「像素 $(i,j)$ 和 $(i,j+1)$ 是邻居」。图像的平移不变性、局部性这些先验全部要从数据里重新学，样本效率极低。
3. 输入输出必须同形。 这一点是流/扩散模型独有的困难。分类网络可以一路下采样，$224\times224\times3 \to 1000$，卷积栈可以「收窄」；而这里 $\vf_t^\theta(x)\in\R^d$ 必须和输入一样大。讲义原话：our flow-based modeling approach requires that our output $\vf_t^\theta(x)\in\R^d$ be just as large as its input。

2. 条件变量的嵌入：时间、类别与文本

网络的主干（U-Net 或 DiT）处理的是「图像形状」或「token 序列形状」的张量。而 $t$ 是一个标量，$y$ 可能是一串字符。所以第一步永远是嵌入（embedding）：把原始条件 $y_{\text{raw}}$ 和 $t$ 变成网络能吃的向量。

2.1 时间嵌入：为什么要用 Fourier 特征

幻灯片上的话很直白：时间只有一维，而其他变量都是高维的。 如果直接把标量 $t$ 拼到 $1.8\times10^6$ 维的输入上，它在第一层线性变换里的「话语权」大约是 $10^{-6}$，几乎必然被淹没。

更深一层的理由是谱偏置（spectral bias）：带 ReLU/SiLU 的 MLP 天然倾向于学低频函数，对输入的高频依赖学得极慢。而 $\vf_t^\theta$ 对 $t$ 的依赖恰恰是高频的——回忆 CondOT 路径 $\alpha_t=t,\ \beta_t=1-t$ 下条件向量场

当 $t\to 1$ 时系数 $\frac{1}{1-t}$ 会炸开。$t=0.99$ 和 $t=0.999$ 对应的目标场相差 10 倍。网络必须能分辨这两个非常接近的时间值。Fourier 特征（Fourier features）正是为解决这个问题被提出的 [Tancik et al. 2020]。

import math, torch, torch.nn as nn

class FourierTimeEmbedding(nn.Module):
    """TimeEmb(t) = sqrt(2/d) * [cos(2*pi*w*t) , sin(2*pi*w*t)]   ==>  ||TimeEmb(t)|| = 1"""
    def __init__(self, dim, w_min=1.0, w_max=1000.0):
        super().__init__()
        assert dim % 2 == 0
        half = dim // 2                                    # d/2
        # w_i = w_min * (w_max/w_min)^((i-1)/(d/2-1)),  i = 1..d/2   —— 对数均匀
        exps = torch.arange(half).float() / max(half - 1, 1)   # (half,)   从 0 到 1
        self.register_buffer("w", w_min * (w_max / w_min) ** exps)   # (half,)
        self.dim = dim

    def forward(self, t):                                  # t: (bs,)   取值于 [0, 1]
        ang = 2 * math.pi * t[:, None] * self.w[None, :]    # (bs, half)
        emb = torch.cat([ang.cos(), ang.sin()], dim=-1)     # (bs, dim)
        return emb * math.sqrt(2.0 / self.dim)              # (bs, dim), 每行范数 = 1


# 实践中还会再接一个小 MLP，把这组固定特征投影成可学习的条件向量
time_mlp = nn.Sequential(              # 输入 (bs, dim)
    nn.Linear(dim, 4 * dim), nn.SiLU(),
    nn.Linear(4 * dim, dim),           # 输出 (bs, dim) —— 这就是后文的 \tilde{t}
)

2.2 类别标签嵌入

当 $y_{\text{raw}}\in\mathcal{Y}\triangleq\{0,1,\dots,N\}$ 只是一个类别标签时，最省事的做法就是学一张嵌入表：为 $N+1$ 个可能取值各学一个向量 $\tilde y = E[y_{\text{raw}}]\in\R^d$，$E\in\R^{(N+1)\times d}$。这些嵌入参数被算作 $\vf_t^\theta$ 参数 $\theta$ 的一部分，跟着一起训练。

2.3 文本嵌入

文本要复杂得多，主流做法是依赖冻结的预训练模型，不从头学：

• CLIP（Contrastive Language-Image Pre-training）：在图文对上用对比损失训练，把图像和文本压进同一个嵌入空间——匹配的图文对靠得近，不匹配的推开。取 $y=\text{CLIP}(y_{\text{raw}})\in\R^{d_{\text{CLIP}}}$ 作为提示词的单向量表示。
• 但把整句话压成一个向量会丢掉结构信息（「戴帽子的狗」和「戴狗的帽子」可能撞在一起）。所以还会用预训练 Transformer 编码器（T5、UL2 等）得到一个序列表示 $$\text{PromptEmb}(y_{\text{raw}})\in\R^{S\times k},$$ $S$ 是 token 数，$k$ 是嵌入维度。序列表示让网络可以用 cross-attention「盯着」提示词的某个具体词。
• 把多种预训练嵌入拼起来同时用也很常见，取各家之长——SD3 和 Movie Gen 都是这么干的（§10）。

3. Diffusion Transformer：把图像变成 token 序列

回忆最基本的类型：一张图像就是一个张量 $x\in\R^{C_{\text{image}}\times H\times W}$，$C_{\text{image}}$ 是通道数（RGB 图 $C_{\text{input}}=3$），$H,W$ 是高和宽。Diffusion Transformer（DiT）的想法是：既然 Transformer 在语言上这么能扩展，就把图像也变成 token 序列，用标准注意力处理，最后再变回来。整条流水线是

下面用 $d$ 表示隐藏维度，$L$ 表示层数，$h$ 表示每层的注意力头数。

3.1 Patchify：切块 + 线性投影

patchify 本身只是一次张量重排，没有任何参数：把 $x\in\R^{C\times H\times W}$ 按 $P\times P$ 的方格切开，每个方格里的 $C P^2$ 个数拉直成一个向量。

再乘一个可学习矩阵得到最终的 patch 嵌入：

算个具体的。 Stable Diffusion 级别的隐张量 $x\in\R^{4\times 32\times 32}$，取 $P=2$：$C'=4\cdot 2^2=16$，$N=(32/2)^2=256$。于是 $\text{PatchEmb}(x)\in\R^{256\times d}$——256 个 token，跟一个短句子差不多长。如果换成 $64\times 64$ 的隐张量，$N=1024$；如果直接在 $512\times512$ 像素上做，$N=65536$，这个数字后面 §5 会用来算账。

def patchify(x, P):                                    # x: (bs, C, H, W)
    bs, C, H, W = x.shape
    x = x.reshape(bs, C, H // P, P, W // P, P)         # (bs, C, H/P, P, W/P, P)
    x = x.permute(0, 2, 4, 1, 3, 5)                    # (bs, H/P, W/P, C, P, P)
    return x.reshape(bs, (H // P) * (W // P), C * P * P)   # (bs, N, C') , C' = C*P*P

def depatchify(tok, C, H, W, P):                       # tok: (bs, N, C')
    bs = tok.shape[0]
    x = tok.reshape(bs, H // P, W // P, C, P, P)       # (bs, H/P, W/P, C, P, P)
    x = x.permute(0, 3, 1, 4, 2, 5)                    # (bs, C, H/P, P, W/P, P)
    return x.reshape(bs, C, H, W)                      # (bs, C, H, W)

3.2 三路输入汇合

把 §2 和 §3.1 拼起来，DiT 的三个输入分别是

注意三者的最后一维都被投影到了同一个 $d$——这正是「嵌入」这一步存在的意义。然后迭代地过 $L$ 层：

最后 depatchify 回图像形状：

这个张量就是模型的输出，即预测的速度场 $\vf_t^\theta(x\mid y)$。

3.3 注意力回顾：形状与那个 $\sqrt{d_h}$

缩放点积注意力。 给定查询 $Q\in\R^{N\times d_h}$、键 $K\in\R^{M\times d_h}$、值 $V\in\R^{M\times d_h}$，

softmax 按行作用。注意 $Q$ 的行数 $N$ 与 $K,V$ 的行数 $M$ 可以不同——这正是 cross-attention 的用武之地：$N$ 个图像 token 去查询 $M=S$ 个文本 token。

多头注意力。 设头数为 $h$，每头维度 $d_h=d/h$。对每个头学三个投影矩阵 $W_Q^{(j)},W_K^{(j)},W_V^{(j)}\in\R^{d\times d_h}$，定义

其中源序列 $z$ 的取法决定了注意力的种类：

拼接所有头再过输出投影 $W_O\in\R^{d\times d}$：

3.4 时间条件化：AdaLN 与 adaLN-Zero

现在到了 DiT 最关键、也最容易被一句话带过的地方。$\tilde t$ 怎么进网络？

最朴素的做法有两个：(a) 把 $\tilde t$ 当作一个额外的 token 拼进序列（in-context conditioning）；(b) 直接把 $\tilde t$ 加到每个 token 上。DiT 论文实测发现，一个更好的做法是调制（modulation）：用 $\tilde t$ 去生成归一化层的缩放和平移参数。这就是自适应层归一化（Adaptive Layer Normalization, AdaLN），思想来自 FiLM。

具体地，设 $g:\R^d\to\R^{2d}$ 是一个 MLP，令

给定 token 矩阵 $x\in\R^{N\times d}$ 和归一化算子 $\text{Norm}(\cdot)$（例如去掉了可学习仿射参数的 LayerNorm），定义调制归一化

其中 $\odot$ 是逐元素乘，并在 token 维度上广播（即所有 $N$ 个 token 共享同一组 $\gamma,\beta$）。

完整的 DiTBlock。 讲义给出的形式是：

其中 MLP 是逐 token 作用的前馈网络，$g_{\text{self}},g_{\text{cross}},g_{\text{MLP}}:\R^d\to\R^d$ 是可学习的门控（gating）参数（同样由 $\tilde t$ 生成）。三步都是残差形式：$x\leftarrow x+(\text{门控})\odot(\text{子层})$。最终输出 $x\in\R^{N\times d}$ 成为下一层的 $\tilde x_{i+1}$。

import torch, torch.nn as nn

class DiTBlock(nn.Module):
    """一个带 cross-attention 的 DiT block，adaLN-Zero 条件化。"""
    def __init__(self, d, n_heads, mlp_ratio=4.0):
        super().__init__()
        # 不带仿射参数的 LayerNorm：缩放/平移完全交给 AdaLN 生成
        mk_norm = lambda: nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.n1, self.n2, self.n3 = mk_norm(), mk_norm(), mk_norm()
        self.self_attn  = nn.MultiheadAttention(d, n_heads, batch_first=True)
        self.cross_attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
        hid = int(d * mlp_ratio)
        self.mlp = nn.Sequential(nn.Linear(d, hid), nn.GELU(), nn.Linear(hid, d))
        # 一个 MLP 一次产出 9 组参数: (gamma, beta, gate) x (self, cross, mlp)
        self.mod = nn.Sequential(nn.SiLU(), nn.Linear(d, 9 * d))
        nn.init.zeros_(self.mod[-1].weight)   # <<< adaLN-Zero
        nn.init.zeros_(self.mod[-1].bias)     # <<< 初始时 gamma = beta = gate = 0

    def forward(self, x, c, y):        # x: (bs, N, d)  c: (bs, d)  y: (bs, S, d)
        # c = time_mlp(TimeEmb(t)) + class_emb(y)   —— 时间与类别嵌入已相加
        g1, b1, a1, g2, b2, a2, g3, b3, a3 = self.mod(c).chunk(9, dim=-1)  # 各 (bs, d)

        def ada(norm, h, g, b):        # AdaNorm(h) = (1 + gamma) * Norm(h) + beta
            return norm(h) * (1 + g[:, None, :]) + b[:, None, :]           # (bs, N, d)

        h = ada(self.n1, x, g1, b1)                                        # (bs, N, d)
        x = x + a1[:, None, :] * self.self_attn(h, h, h, need_weights=False)[0]
        h = ada(self.n2, x, g2, b2)                                        # (bs, N, d)
        x = x + a2[:, None, :] * self.cross_attn(h, y, y, need_weights=False)[0]
        h = ada(self.n3, x, g3, b3)                                        # (bs, N, d)
        x = x + a3[:, None, :] * self.mlp(h)
        return x                                                           # (bs, N, d)

4. U-Net 与两种架构的取舍

U-Net 是 DiT 之外的另一条主线，一种特殊的卷积神经网络。它最早是为医学图像分割设计的，关键特征恰好就是我们在 §1 说的那条硬约束：输入和输出都具有图像的形状（通道数可以不同）。这使它天然适合参数化 $x\mapsto \vf_t^\theta(x\mid y)$——固定 $y,t$，输入是图像形状，输出也是。早期扩散模型文献（DDPM、ADM、SD 1.x/2.x）几乎全用 U-Net。

4.1 结构与张量形状

一个 U-Net 由一串编码器 $\mathcal{E}_i$、一串对应的解码器 $\mathcal{D}_i$、以及夹在中间的一个隐层处理块——讲义戏称为 midcoder $\mathcal{M}$ ——组成。以 $x_t\in\R^{3\times 256\times 256}$ 为例，讲义给出的路径是：

注意这里 U-Net 内部的「latent」和后面 §5 起要讲的 VAE 隐空间完全是两回事，只是名字撞了：前者是网络内部的中间特征图，后者是一个独立训练出来的、扩散模型真正工作的空间。

规律是：越往下走，通道数越多，空间分辨率越低。把典型的一条路径逐级写出来：

「通道翻倍、边长减半」这个搭配不是随意的：每下降一级，空间位置数变成 $1/4$，通道数变成 2 倍，特征图元素数变成 $1/2$，而卷积的计算量 $O(HWC^2k^2)$ 保持不变（$\frac14\cdot 2^2=1$）。也就是说金字塔的每一级花的算力一样多，这是一个刻意的设计。

讲义还强调了两个容易忽略的点：第一，输入 $x_t^{\text{input}}\in\R^{3\times 256\times256}$ 通常先经过一个预编码卷积把通道数抬起来，再进第一个编码器块；第二，编码器和解码器之间由残差连接（residual / skip connection）相连。

4.2 残差层与条件注入

U-Net 的基本积木是残差层，结构大致是

其中 $\text{MLP}(\tilde t,\tilde y)\in\R^{C}$ 沿空间维广播后逐通道相加。这是最简单的条件注入方式：加性偏置。它可以看成 FiLM 只用了平移分量的特例；很多实现（包括 SD 的 U-Net）会用完整的 FiLM，即同时给出缩放和平移 $h\leftarrow \gamma(\tilde t)\odot h+\beta(\tilde t)$——这与 DiT 的 AdaLN 是同一个想法，只不过归一化算子从 LayerNorm 换成了 GroupNorm/BatchNorm。

4.3 为什么跳连是必需的，而不是可选的

4.4 U-Net vs DiT

先把计算复杂度算清楚，这是两者最实质的区别。设图像边长为 $R$（即 $H=W=R$）。

• U-Net：一层 $k\times k$ 卷积在 $R\times R$、$C$ 通道上的乘加数是 $O(R^2C^2k^2)$。由于每下一级 $R^2$ 减 4 倍、$C^2$ 增 4 倍，金字塔各级开销相同，总量 $\Theta(R^2)$——与像素数成正比。
• DiT：token 数 $N=(R/P)^2$。每层的开销分两块：逐 token 的线性投影与 MLP 是 $O(Nd^2)$，注意力矩阵是 $O(N^2d)$。合起来 $$\Theta\!\left(\frac{R^2}{P^2}d^2+\frac{R^4}{P^4}d\right).$$ 第二项是分辨率的四次方。

两项的交叉点在 $N\approx d$。举个具体的：$d=1152$、$P=2$，隐张量 $64\times64$ 时 $N=1024\lt d$，投影项占主导，DiT 仍然「便宜」；换成 $128\times128$ 的隐张量，$N=4096\gt d$，注意力项开始主导，再往上就急剧变贵。这就是为什么即便有了 DiT，隐空间也一步都不能省。

5. 为什么必须去隐空间：一笔计算账

到目前为止我们一直在数据空间 $\R^d$ 上工作。但当分辨率往上走，直接建模的代价迅速变得不可承受。讲义给的例子：一张 $1024\times1024$ 的三通道图对应

幻灯片给的例子是 $600\times1000$：$d=1.8\times10^6$。视频还要再乘帧数 $T$。

5.1 为什么监督学习没事，扩散模型有事

幻灯片给出的答案分两条，讲义正文补充了第三条：

1. 输出维度和输入一样高。 图像分类的输出只有 1000 维，所以卷积栈可以一路收窄，越往后越省。而这里 $\vf_t^\theta(x)\in\R^d$——网络必须一路把 $10^6$ 维的信号带到最后。讲义原话：Unlike image classification, whose low-dimensional outputs allow for narrowing convolutional stacks, our flow-based modeling approach requires that our output be just as large as its input.
2. 要反复调用很多次。 分类是前向一次；采样是模拟 ODE/SDE，需要 50 到上千次网络前向。分辨率带来的开销要再乘以步数。
3. 要学的是分布，不是一个条件标签。 拟合 $p(y\mid x)$ 只需要抓住与标签相关的那点信息；拟合 $\data(x)$ 本身要求模型对 $\R^d$ 中每一个方向的统计结构都负责。而其中绝大多数方向（高频噪声、纹理相位）对人的感知毫无贡献。

于是核心问题变成讲义加粗的那一句：如何在合理的显存与算力预算内，对高维图像建模？

5.2 压缩能省多少：把数字算出来

讲义与幻灯片给出的真实系统数字与这个算例完全吻合：

6. 从自编码器到 VAE：ELBO 的完整推导

6.1 标准自编码器与它的致命缺陷

压缩的自然做法是自编码器（autoencoder, AE）：一个编码器 $\mu_\phi:\R^d\to\R^k$ 和一个解码器 $\mu_\theta:\R^k\to\R^d$，把原始数据 $x\in\R^d$ 映到隐变量 $z\in\R^k$ 再映回来，$k\ll d$。训练用重构损失

但这样得到的隐空间对生成建模不够用。讲义把这个问题称为 amenability to generative modeling：我们的最终目标是在隐空间里训一个生成模型，去拟合 $z=\mu_\phi(x),\ x\sim\data$ 诱导出的分布 $p_{\text{latent}}(z)$。可是重构损失对 $p_{\text{latent}}$ 只字未提，我们对它毫无控制。

6.2 变分自编码器：把确定性映射放松成分布

变分自编码器（variational autoencoder, VAE）的做法是放松编码器和解码器必须是确定性函数这一约束：改用条件分布。设编码器为 $q_\phi(z\mid x)$、解码器为 $p_\theta(x\mid z)$，最常见的取法是高斯：

各对象的类型是

都由神经网络参数化。编码和解码变成采样：

当 $\sigma_\phi(x)\equiv 0$ 且 $\sigma_\theta(z)\equiv 0$ 时，两个分布退化成 Dirac 测度，我们就回到了标准自编码器。所以 VAE 是 AE 的严格推广。

6.3 边际似然为什么算不出来

引入先验 $p_{\text{prior}}(z)=\N(0,I_k)$ 之后，解码器加先验就定义了一个数据空间上的生成模型

最大似然的标准做法是最大化 $\log p_\theta(x)$。问题在于这是一个 $k$ 维积分，$k$ 通常上万，没有闭式解。朴素的 Monte Carlo

方差高到完全不可用：从各向同性高斯里随机抽一个 $z$，解码出来恰好接近给定的 $x$ 的概率是天文数字级的小，绝大多数样本贡献 $p_\theta(x\mid z)\approx 0$。我们需要一个知道「$x$ 大概编码到哪儿」的提议分布——这正是编码器 $q_\phi(z\mid x)$ 的角色。

6.4 ELBO：两条推导路径

6.5 把 ELBO 拆成「重构 − 正则」

讲义把这两项分别记作

并用一个权重 $\beta\ge0$ 组合成 VAE 训练目标：

取 $\beta=1$ 时，$\mathcal{L}_{\text{VAE}}=-\E_{x\sim\data}[\text{ELBO}(x;\phi,\theta)]$，即最小化 VAE 损失就是最大化期望 ELBO。$\beta\neq1$ 的版本称为 $\beta$-VAE。

6.6 高斯解码器下重构项的具体形式

7. 高斯情形：KL 闭式解、重参数化与训练算法

7.1 KL 散度速查

对两个概率密度 $q,p$，Kullback–Leibler 散度定义为

它满足两条基本性质（讲义 Remark 30）：

7.2 两个对角高斯之间 KL 的闭式解

7.3 $\mathcal{K}(\alpha)$ 长什么样

公式的形状非常直观：$\KL$ 随均值的平方误差 $\norm{\mu_q-\mu_p}^2$ 单调增长；方差的贡献则通过 $\mathcal{K}(\alpha)=\alpha-\log\alpha-1$ 体现。

7.4 代入 $p_{\text{prior}}=\N(0,I_k)$：损失的最终形式

取 $\mu_p=0,\ \sigma_p^2=1$（即先验为标准正态），$q=q_\phi(\cdot\mid x)$：

展开成逐坐标的形式（$k$ 是隐维度）：

这就是所有 VAE 代码里那一行 0.5 * (mu**2 + var - logvar - 1).sum() 的来历。把它和 §6.6 的重构项合起来，得到讲义的完整 VAE 损失：

四项的作用讲义总结得很清楚：第一项是重构误差；第二项刻画解码器的不确定性——方差越小解码器越「自信」，但同时也让重构误差被更重地惩罚；第三、四项则强制隐分布的方差为 1、均值为 0，也就是逼近标准高斯。

7.5 重参数化技巧

剩下的问题是怎么优化。麻烦在于我们取期望所依赖的那个分布 $q_\phi(z\mid x)$ 本身依赖 $\phi$。

把重参数化代入 §7.4 的损失，得到最终可优化的形式：

（这里已把 $\sigma_\theta^2(z)$ 固定为常数 $\sigma^2$）。随机性只剩下 $\epsilon$，它的分布与 $\phi$ 无关，标准深度学习工具直接可用。

7.6 训练算法

import torch

def vae_loss(x, enc, dec, beta, sigma2):
    """x: (bs, d);  enc(x) -> (mu, logvar): (bs, k), (bs, k);  dec(z) -> (bs, d)"""
    mu, logvar = enc(x)                       # (bs, k), (bs, k)   —— 输出 log sigma^2
    eps   = torch.randn_like(mu)              # (bs, k)   随机源，与 phi 无关
    z     = mu + torch.exp(0.5 * logvar) * eps        # (bs, k)   重参数化
    x_hat = dec(z)                            # (bs, d)   解码器均值 mu_theta(z)

    # 重构项:  ||x - mu_theta(z)||^2 / (2 sigma^2)
    recon = ((x - x_hat) ** 2).flatten(1).sum(-1) / (2 * sigma2)       # (bs,)
    # KL 项:  0.5 * sum_j ( mu_j^2 + sigma_j^2 - log sigma_j^2 - 1 )
    kl    = 0.5 * (mu ** 2 + logvar.exp() - logvar - 1).sum(-1)        # (bs,)

    return (recon + beta * kl).mean()         # 标量

7.7 工程现实：VAE 在 latent diffusion 里到底是什么

讲义的四条实践备注，加上一点现实：

1. $\beta$ 的选取与 KL warm-up。 $\beta$ 太大会把隐变量压得太靠近先验，损害重构，甚至触发后验塌缩（posterior collapse）——编码器索性无视 $x$，输出 $q_\phi(z\mid x)\approx\N(0,I_k)$，KL 项直接归零而信息全部丢失。常见的稳定手段是 KL warm-up：从 $\beta=0$ 开始，在最初若干个 epoch 内逐步升到目标值。但讲义特别点明：在所有现代自编码器里，$\beta$ 的取值都非常小，$\beta\ll 1$。
2. 解码器方差。 学一个高斯解码器方差 $\sigma_\theta^2$ 在数值上很微妙，不加正则容易退化。为稳定起见，许多实现直接固定 $p_\theta(x\mid z)=\N(x;\mu_\theta(z),\sigma^2I_d)$，$\sigma^2$ 为常数，此时重构项就正比于 MSE（差一个常数）。
3. 超越像素 MSE 的重构损失。 逐像素高斯似然（即 MSE）给出的重构常常过度平滑。实践中会加感知损失（perceptual loss）：用一个预训练网络的特征空间上的距离，来提升锐度和语义保真度。
4. 对抗与混合目标。 为进一步提升视觉真实感，可以把 VAE 目标与对抗损失（VAE-GAN 风格，在解码结果上加判别器）结合。这通常能让输出更锐利，代价是额外的优化不稳定性和更多超参数。

8. VAE 的更多视角（讲义附录 D）

附录 D 换了三个角度重看同一个损失函数。这些视角不是花架子——第三个（amortization gap）直接回答了「既然 VAE 自己就是生成模型，为什么还要在它的隐空间里再训一个扩散模型」这个自然的疑问。

8.1 视角一：VAE 损失就是联合分布上的一个 KL

关键观察：编码器和解码器各自诱导出 $(x,z)$ 上的一个联合分布：

前者是「先抽一个真实数据，再编码」；后者是「先抽一个先验隐变量，再解码」。训练 VAE 可以理解为：让这两个联合分布尽量接近。

8.2 视角二：VAE 本身就是生成模型（链式法则与数据处理不等式）

既然解码器加先验定义了 $p_\theta(x)=\int p_\theta(x\mid z)p_{\text{prior}}(z)\ud z$，我们可以直接从 $z\sim p_{\text{prior}}$ 采样再解码。这样得到的样本有多接近真实数据？下面这条命题给出答案。

8.3 视角三：amortization gap ——为什么不到 VAE 为止

既然 VAE 自己就能生成，为什么还要费劲在它的隐空间里再训一个流/扩散模型？附录 D 给的答案是 amortization gap（摊销缺口）。

8.4 从联合 KL 里再把 ELBO 抠出来

9. 潜在扩散模型的完整配方

所有零件都在手上了。讲义 Remark 32 把它总结成一句话：照搬已有的训练配方，只不过直接在隐空间里做。

1. 数据。 取全部训练数据 $x_1,\dots,x_N$（例如互联网上的所有图片）。
2. 编码。 把所有图像转成隐向量。讲义正文的说法是训练时从 $q_\phi(z\mid x)$ 采样（这给隐数据加了一点噪声，起数据增广与正则的作用）；幻灯片给的简化配方是直接取均值预测 $z=\mu_\phi(x)$。两者都能用，后者更省事也更常见。
3. 隐数据集。 得到隐向量数据集 $z_1,\dots,z_N$，规模远小于原始高分辨率图像。
4. 隐扩散模型。 在这个数据集上训一个流/扩散模型——现在模型生成的是隐向量。
5. 解码。 采样得到 $z$ 之后，用解码器映回数据空间。讲义特别注明：这里取均值 $x=\mu_\theta(z)$ 而不是从 $p_\theta(\cdot\mid z)$ 随机采样，以避免噪声引入的伪影。

# ---------- 阶段一：训练 VAE，然后冻结 ----------
vae = train_vae(images)                 # 见 7.6
vae.requires_grad_(False).eval()        # <<< 冻结！后续不再更新

# ---------- 阶段二：把整个数据集编码成隐向量 ----------
@torch.no_grad()
def encode_all(vae, loader, scale):
    zs = [vae.encode(x)[0] * scale for x, _ in loader]   # 只取 mu: (bs, C, h, w)
    return torch.cat(zs)                                 # (N, C, h, w)

# ---------- 阶段三：在隐空间做条件流匹配（与 Lecture 2 完全一致）----------
def cfm_step(model, z1, y, p_uncond=0.1):    # z1: (bs, C, h, w) 来自数据的隐向量
    bs = z1.shape[0]
    t   = torch.rand(bs, device=z1.device)                       # (bs,)
    eps = torch.randn_like(z1)                                   # (bs, C, h, w)
    zt  = t.view(-1,1,1,1) * z1 + (1 - t.view(-1,1,1,1)) * eps   # CondOT 插值
    target = z1 - eps                                            # d/dt zt，直线路径
    y = drop_labels(y, p_uncond)                                 # CFG：随机置为空标签
    return ((model(zt, t, y) - target) ** 2).mean()

# ---------- 阶段四：采样 ----------
@torch.no_grad()
def sample(model, vae, y, shape, n_steps=50, w=4.0, scale=0.18):
    z = torch.randn(shape)                                       # (bs, C, h, w)
    for i in range(n_steps):
        t  = torch.full((shape[0],), i / n_steps)
        u  = (1 - w) * model(z, t, NULL) + w * model(z, t, y)     # CFG
        z  = z + u / n_steps                                      # Euler 步
    return vae.decode(z / scale)                                  # (bs, 3, H, W)

10. 案例研究：Stable Diffusion 3 与 Meta Movie Gen

最后把前面所有零件放进两个真实的、当前最先进的系统里。它们都用到了本课教的技术，外加一些为了扩大规模和处理复杂条件模态而做的架构增强。

10.1 Stable Diffusion 3

• 训练目标：与本课 Algorithm 4 完全相同的条件流匹配目标。SD3 的论文里广泛比较了各种流与扩散的替代方案，结论是 flow matching 表现最好。（他们对噪声的条件约定写法与本课不同，但只是记号差异，算法是同一个。）
• 调度：「直线」调度器，即 CondOT 路径。
• 引导：训练时按 CFG 的方式随机丢弃类别/文本标签；采样时引导权重取 2.0–5.0。
• 空间：在预训练自编码器的隐空间里做流匹配，正是 §9 的配方。
• 文本条件：同时用三种文本嵌入——CLIP 系列的嵌入（提供粗粒度、总揽式的整句表示）以及预训练的 Google T5-XXL 编码器输出的序列级嵌入。T5 的序列嵌入提供了更细粒度的上下文，使模型可以关注提示词中的具体成分。
• 架构：为了容纳这些序列上下文嵌入，作者把 DiT 扩展成 MM-DiT（multi-modal DiT）——不仅让图像 patch 之间互相注意，还让它们注意到文本嵌入，从而把 DiT 原本只支持的类别条件推广到序列上下文条件。文本和图像被贯穿整个网络地共同处理。
• 规模：最大模型 80 亿参数。
• 采样：50 步 Euler 模拟（即网络前向 50 次）。
• 数据集：LAION。

10.2 Meta Movie Gen Video

接下来是视频。数据不再是图像而是视频，即 $x\in\R^{T\times C\times H\times W}$，多出来的 $T$ 是时间维（帧数）。讲义的观察是：这个设定下的许多设计选择，都可以看作把图像领域已有的技术（自编码器、diffusion transformer 等）适配到多出来的时间维。

• 训练目标：同样是条件流匹配，同样用直线调度 $\alpha_t=t,\ \sigma_t=1-t$。
• 引导：classifier-free guidance。
• 空间：与 SD3 一样，在冻结的预训练自编码器隐空间里工作。对视频而言这一步比图像更关键——多一个时间维，显存压力是数量级的。
• 时序自编码器（temporal autoencoder, TAE）：把原始视频 $x_t'\in\R^{T'\times3\times H\times W}$ 映到隐张量 $x_t\in\R^{T\times C\times H'\times W'}$，压缩比为 $$\frac{T'}{T}=\frac{H}{H'}=\frac{W}{W'}=8,$$ 即时间和两个空间维各压 8 倍，时空位置数压缩 $8^3=512$ 倍（再乘上通道比 $3/C$）。
• 长视频的处理：提出时间分块（temporal tiling）——把视频沿时间切成若干段，每段单独编码，再把隐表示拼接起来。这样显存需求不随视频长度线性增长。
• 主干：DiT 风格的骨干网络，$x_t$ 沿时间和空间同时 patchify；得到的图像 patch 经过一个既做 patch 间自注意力、又与语言模型嵌入做交叉注意力的 transformer——与 SD3 的 MM-DiT 同源。
• 文本条件：三种嵌入各司其职——UL2 嵌入负责细粒度的、基于文本的推理；ByT5 嵌入负责字符级细节（例如提示词明确要求画面中出现某段文字时）；MetaCLIP 嵌入来自共享的图文嵌入空间。
• 规模：最大模型 300 亿参数，训练用了 6 144 张 H100。

本讲小结

延伸阅读

架构

• U-Net: Convolutional Networks for Biomedical Image Segmentation (2015) — 跳连式编解码器的原始论文，本来是做分割的，后来成了扩散模型十年的默认骨干。
• Attention Is All You Need (2017) — 缩放点积注意力与正弦位置编码的出处，§3.3 的公式全部来自这里。
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (2020) — ViT，patchify 这个操作的来源，DiT 直接继承。
• Scalable Diffusion Models with Transformers (2023) — DiT 原论文。想真正搞懂 adaLN-Zero 为什么优于其它三种条件化方案，读它的消融实验表。
• SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers (2024) — 在 DiT 骨干上系统比较各种插值/调度/目标的组合，是「架构固定、目标变化」的对照实验。
• FiLM: Visual Reasoning with a General Conditioning Layer (2018) — 「用条件变量生成逐通道缩放和平移」这一思想的源头，AdaLN 是它的直系后代。
• Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains (2020) — 讲义为时间嵌入引用的正是这篇；它解释了为什么低维输入必须先做频率展开。
• Diffusion Models Beat GANs on Image Synthesis (2021) — ADM，把 U-Net 的每一个设计选择（宽度、注意力分辨率、归一化）都做了消融，是 U-Net 调参的标准参考。

隐空间与自编码器

• Auto-Encoding Variational Bayes (2013) — VAE 原论文，ELBO 与重参数化技巧的出处，§6–§7 的全部内容都在这一篇里。
• High-Resolution Image Synthesis with Latent Diffusion Models (2022) — Stable Diffusion 的奠基论文，把「先训自编码器、再在隐空间做扩散」确立为标准范式；里面的自编码器就是 $\beta\ll1$ 加感知与对抗损失的那一版。
• Score-based Generative Modeling in Latent Space (2021) — LSGM，讲义并列引用的另一条隐空间扩散路线，它尝试端到端联合训练 VAE 与扩散模型，正好与 §9 的两阶段做法形成对照。

文本条件

• Learning Transferable Visual Models From Natural Language Supervision (2021) — CLIP，共享图文嵌入空间的来源。
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (2019) — T5，SD3 用的 T5-XXL 编码器出自这里。
• UL2: Unifying Language Learning Paradigms (2022) — Movie Gen 用于细粒度文本推理的嵌入。
• ByT5: Towards a Token-Free Future with Pre-trained Byte-to-Byte Models (2021) — 字符级嵌入，用来处理「画面里要出现这几个字」这类提示。
• Modeling Caption Diversity in Contrastive Vision-Language Pretraining (2024) — MetaCLIP，Movie Gen 的第三种文本嵌入。
• Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding (2022) — Imagen，最早系统论证「文本编码器越大越好、比扩散模型本身更重要」的工作。
• eDiff-I: Text-to-Image Diffusion Models with an Ensemble of Expert Denoisers (2023) — 组合多种预训练嵌入的代表工作之一，讲义在讨论多嵌入拼接时引用了它。

大规模系统

• Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (2024) — Stable Diffusion 3 技术报告，§10.1 的全部细节与 MM-DiT 结构图都出自这里；附录里有各种调度/目标的完整对照实验，是本讲最值得精读的一篇。
• Movie Gen: A Cast of Media Foundation Models (2024) — Meta Movie Gen 技术报告，TAE、时间分块、时空 patchify 的细节都在其中。讲义鼓励读者直接去读原报告获取更详尽的处理。