Lab 3 解析：条件图像生成（DiT + VAE + 隐空间扩散）

0. 本 lab 导读

Lab 2 把 flow matching 的公式链跑通了，但跑的是二维玩具数据、用的是一个三层 MLP、生成的是「随便什么样本」。Lab 3 要在这三件事上同时升级：

但要强调一件让人安心的事：训练循环的那四行，和你在 Lab 2 Q3.1 写下的一模一样——采 $z$、采 $t$、采 $x_t$、回归 $u_t(x\mid z)$。Lab 3 的 15 个 TODO 里只有两个（Q2.2 和 Q5.1）在动这条循环，而且动的只是「$z$ 从哪来」和「$y$ 要不要丢掉」。剩下的 13 个全部是架构工程：怎么把一个函数 $u^\theta:\R^{1024}\times[0,1]\times\set{0,\dots,10}\to\R^{1024}$ 用张量算子搭出来。

1. 骨架的三处升级：从 Lab 2 到 Lab 3

Part 0 和 Part 1 是「回收站」，把 Lab 1/2 的类原样搬了过来。但有三处改动是刻意的，看懂它们能省掉后面一半的调试时间。

1.1 升级一：Sampleable 变成 LabeledSampleable

Lab 2 的分布只需要能采样：sample(n) -> x。Lab 3 要做条件生成，所以采样必须同时返回标签：

class LabeledSampleable(ABC):
    @abstractmethod
    def sample(self, num_samples: int) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Returns:
            - samples: b d   (或 b c h w)
            - labels:  b
        """

注意标签的形状是 (b,)，是一维整数张量，因为它要直接喂给 nn.Embedding。而 IsotropicGaussian（源分布 $p_{\text{init}}$）仍然只实现无标签的 Sampleable——高斯噪声谈不上「类别」，这个区分是有意义的。

1.2 升级二：时间 $t$ 的形状从 (b, 1) 变成 (b,)

这是 Lab 2 → Lab 3 最容易踩的一个坑，而且它贯穿全 lab。Lab 2 里数据是 (b, dim)，时间约定成 (b, 1)，广播天然对齐。Lab 3 的数据是 (b, c, h, w)，如果还用 (b, 1)，广播就会把时间维当成通道维，结果完全错乱。

官方的解法是：对外统一把 $t$ 定成 (b,)，进到路径类内部再按数据的秩自己撑开：

class GaussianConditionalProbabilityPath(ConditionalProbabilityPath):
    def __init__(self, p_data, p_simple_shape, alpha, beta):
        ...
        # p_simple_shape=[1,32,32] 时，这行等价于 Rearrange('b -> b 1 1 1')
        # p_simple_shape=[2]       时，等价于 Rearrange('b -> b 1')
        self.rearrange_scalar = Rearrange(f'b -> b{" 1" * len(p_simple_shape)}')

    def sample_conditional_path(self, z, t):
        alpha_t = self.rearrange_scalar(self.alpha(t))   # (b, 1, 1, 1)
        beta_t  = self.rearrange_scalar(self.beta(t))    # (b, 1, 1, 1)
        return alpha_t * z + beta_t * torch.randn_like(z)

这个设计非常漂亮：同一份路径代码，既能作用在 (b, 2) 的玩具数据上（Sanity Check 2.4），也能作用在 (b, 1, 32, 32) 的图像上（Part 3），还能作用在 (b, 128, 4, 4) 的隐变量上（Part 5）——因为「秩」这个信息在构造时就被烘焙进 rearrange_scalar 里了。

1.3 升级三：MNISTSampler 与一处性能改写

官方的 MNISTSampler 长这样（简化）：每次 sample 都用 torch.randperm 取一批下标，然后对每一张图调用 PIL 的 Resize((32,32)) 和 ToTensor()。这在功能上没问题，但在 20000 步 × batch 256 的训练里，等于要做 512 万次单张 PIL 变换，会成为纯 CPU 瓶颈——GPU 大部分时间在等数据。

本页的实测版本把它改成了一次性预处理并缓存到 GPU：

class MNISTSampler(nn.Module, LabeledSampleable):
    def __init__(self):
        super().__init__()
        dataset = datasets.MNIST(
            root='./data', train=True, download=True,
            transform=transforms.Compose([
                transforms.Resize((32, 32)),
                transforms.ToTensor(),
                transforms.Normalize((0.1305,), (0.2891,)),
            ])
        )
        loader = torch.utils.data.DataLoader(dataset, batch_size=2048,
                                             shuffle=False, num_workers=4)
        xs, ys = [], []
        for xb, yb in loader:
            xs.append(xb); ys.append(yb)
        self.register_buffer('images', torch.cat(xs))                 # (60000, 1, 32, 32)
        self.register_buffer('labels', torch.cat(ys).to(torch.int64)) # (60000,)
        self.dummy = nn.Buffer(torch.zeros(1))

    def sample(self, num_samples: int):
        idx = torch.randint(0, self.images.shape[0], (num_samples,),
                            device=self.images.device)               # (b,)
        # 高级索引本身会返回副本；再 clone 一次，确保调用方的原地改写
        # （CFG 的 label dropout 会写 y[mask] = 10）不会污染缓存的数据集
        return self.images[idx].clone(), self.labels[idx].clone()

整个数据集常驻显存的代价是 $60000\times1\times32\times32\times4\ \text{B}\approx245$ MB，对 16 GB 显存完全不是问题，换来的是训练吞吐提升约一个数量级。

这个改动改变数学吗？只改变了一点，而且是往好的方向：randperm 是无放回抽样，randint 是有放回抽样。理论上我们要的是从 $\data$ 独立同分布采一个 mini-batch，有放回才是那个理想设定；无放回反而引入了 batch 内的负相关。对 SGD 的收敛性没有实际影响，工程上更简单。

1.4 顺带一提：Trainer 加了什么

Lab 3 的 Trainer 基类比 Lab 2 多了三样东西，都值得知道：

• 线性 warmup：前 warmup_steps（默认 500）步把学习率从 0 线性升到 lr。transformer 对初期的大梯度非常敏感，没有 warmup 很容易在前几百步崩掉。
• checkpoint 回调：每 ckpt_every 步存一次权重 + 采一批样本存图。强烈建议用起来——20000 步要跑几十分钟，能在第 1000 步就看出「样本是不是完全的噪声」，比跑完再看划算得多。
• model_size_b：打印模型大小。DiT 约 40.3 MiB，VAE 约 6.2 MiB，可以用来快速核对你的架构参数有没有写错（比如把 dim 写成 512 会让大小翻四倍）。

2. Part 2 · Classifier-Free Guidance

先把理论理顺，代码只有五行。Lecture 3-B 已经完整讲过引导，这里只重做一遍与代码直接对应的那条推导链。

2.0 从条件生成到 CFG：三步推导

第一步：条件生成本身不需要新理论。固定一个标签 $y$，把 $\data(\cdot\mid y)$ 当成新的数据分布，Lab 2 的整套 flow matching 原封不动就能用：

再把 $y$ 也放进期望里（而不是固定住），就得到了「一个网络同时学会所有类别」的目标。到这一步理论上就已经完事了：$u^\theta_t(x\mid y)$ 的最优解就是条件边际向量场 $u_t(x\mid y)$，用它做 ODE 采样得到的就是 $\data(\cdot\mid y)$。

第二步：但人们发现「更像那一类」比「正确采样」更讨喜。回忆高斯路径下向量场与得分的线性关系（Lab 2 Q3.3 的那条）：对 $a_t=\frac{\dot\alpha_t}{\alpha_t}$、$b_t=-\frac{\dot\beta_t\beta_t\alpha_t-\dot\alpha_t\beta_t^2}{\alpha_t}$，有

对条件得分用一次贝叶斯：$p_t(x\mid y)=\frac{p_t(x)p_t(y\mid x)}{p_t(y)}$，两边取 $\log$ 再对 $x$ 求梯度，$\log p_t(y)$ 与 $x$ 无关直接消失：

代回去：

第二项 $\grad\log p_t(y\mid x)$ 就是一个「带噪分类器」：它指向「让当前样本更像类别 $y$」的方向。把这一项的权重放大 $w$ 倍，就得到引导向量场：

第三步：消掉分类器。上面那条恒等式反解出 $b_t\grad\log p_t(y\mid x)=u_t(x\mid y)-u_t(x)$，代进去：

其中 $u_t(x)$ 被重写成 $u_t(x\mid\varnothing)$：把「无条件」当成一个特殊的类别 $\varnothing$。这就是 classifier-free guidance（无分类器引导）——不需要训分类器，只需要让同一个网络偶尔看不见标签。

训练侧的改动只有一句话：以概率 $\eta$ 把标签 $y$ 换成 $\varnothing$。这样同一份参数既见过「带标签的样本」（学 $u_t(x\mid y)$），也见过「不带标签的样本」（学 $u_t(x\mid\varnothing)$）。本 lab 取 $\varnothing\triangleq10$，即 MNIST 十个数字之后的第十一个槽位。

2.1 Q2.2 · CFGTrainer.get_train_loss

题目在问什么。 实现下面这个期望的单批次蒙特卡洛估计：

数学依据。 与 Lab 2 的条件 flow matching 定理逐字对应，只多了第二行。回归目标是条件向量场 $u_t^{\text{ref}}(x\mid z)$（可解析计算），而 $L^2$ 最优解是它在给定 $(x,t,y)$ 下的条件期望，也就是我们真正想要的边际条件向量场 $u_t(x\mid y)$。标签丢弃只是把 $y$ 的边缘分布从 $p(y)$ 改成了 $(1-\eta)p(y)+\eta\delta_\varnothing$，不影响这个论证——在 $y=\varnothing$ 的那一支上，条件期望自动收敛到 $\E[u_t(x\mid z)\mid x,t]=u_t(x)$。

参考实现。

class CFGTrainer(Trainer):
    def __init__(self, path, eta: float, null_label: int, eps: float = 0.001, **kwargs):
        assert eta > 0 and eta < 1
        super().__init__(**kwargs)
        self.eta = eta
        self.eps = eps
        self.path = path
        self.null_label = null_label

    def get_train_loss(self, batch_size: int) -> torch.Tensor:
        # Step 1: 从 p_data 采 (z, y)
        z, y = self.path.p_data.sample(batch_size)   # z: (b, c, h, w),  y: (b,)

        # Step 2: 以概率 eta 把标签替换成空标签
        xi = torch.rand(y.shape[0], device=y.device)             # (b,)
        y[xi < self.eta] = self.null_label

        # Step 3: 采 t 和 x
        #   注意 t 的形状是 (b,)，不是 (b,1,1,1)
        #   乘 (1 - eps) 是为了避开 t=1
        t = torch.rand(batch_size, device=z.device).to(z) * (1 - self.eps)  # (b,)
        x = self.path.sample_conditional_path(z, t)                        # (b, c, h, w)

        # Step 4: 回归
        ut_theta = self.model(x, t, y)                       # (b, c, h, w)
        ut_ref   = self.path.conditional_vector_field(x, z, t)  # (b, c, h, w)
        return torch.square(ut_theta - ut_ref).mean()

为什么这样写。

• xi = torch.rand(y.shape[0]) 而不是 torch.rand(1)。 丢弃必须是逐样本独立的：一个 batch 里应该同时存在带标签样本和空标签样本。如果整批一起丢，梯度的方差会大很多，而且 batch norm 类统计量（这里没有，但 VAE 里有 GroupNorm）会在两种模式间来回跳。
• t 的形状是 (b,)。 上一节已经详述：路径类内部会用 Rearrange('b -> b 1 1 1') 自己撑开。这一点在 Sanity Check 2.4 里同样成立（那里 p_simple_shape=[2]，撑成 b 1），所以同一份 trainer 代码在二维玩具数据和图像上都能跑——这正是它被设计成 (b,) 的原因。
• * (1 - self.eps) 是必须的，不是保险丝。 展开条件向量场： $$ u_t(x\mid z)=\left(\dot\alpha_t-\frac{\dot\beta_t}{\beta_t}\alpha_t\right)z+\frac{\dot\beta_t}{\beta_t}x $$ 本 lab 用 LinearAlpha/LinearBeta，即 $\alpha_t=t$、$\beta_t=1-t$，于是 $\frac{\dot\beta_t}{\beta_t}=\frac{-1}{1-t}$，在 $t\to1$ 时发散。虽然代数上 $u_t(x\mid z)$ 化简后等于 $z-\epsilon$（有限），但代码是按原式逐项算的，$t=1$ 会出现 $-\infty\cdot 0$，直接得到 nan。eps=1e-3 把采样区间截到 $[0,0.999]$，代价可以忽略（$t$ 是连续均匀分布，$\{1\}$ 是零测集），收益是训练不会莫名其妙 NaN。采样时同样要截：torch.linspace(0, 0.999, num_timesteps)。
• .to(z) 而不是只写 device=。 .to(z) 会把 dtype 也对齐到 z。如果你后面上 AMP / bf16 训练，这一行能省掉一次类型不匹配。
• .mean() 而不是 .sum()。 对所有元素取平均，等价于「每维平均平方误差」。这让损失的数值不随图像分辨率变化，学习率也就不用跟着分辨率调。记住这个约定，第 6 节讨论损失下界时要用。

2.2 Q2.3 · MLPConditionalVectorField.forward

题目在问什么。 用最朴素的方式把 $(x,t,y)$ 三样输入喂进一个 MLP：拼接。这是自检用的小模型，跑在二维高斯混合上。

数学依据。 没有什么深刻的数学，但有一个重要的表示问题：$y$ 是离散的类别，不能当成实数直接拼进去（那样会隐含「类别 3 比类别 1 大 2」这种荒谬的序关系）。标准做法是查表——把每个类别映射到一个可学习的向量：

参考实现。

class MLPConditionalVectorField(ConditionalVectorField):
  def __init__(self, dim: int, hidden_dim: int, class_dim: int, num_classes: int):
    super().__init__()
    self.mlp = MLP([dim + class_dim + 1, hidden_dim, hidden_dim, dim])
    self.class_embedding = nn.Embedding(num_classes + 1, class_dim)
    #                                   ^^^^^^^^^^^^^^^ 多出来的那一格就是空标签

  def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor):
      xyt = torch.cat([
          x,                          # (b, dim)
          self.class_embedding(y),    # (b, class_dim)
          t.unsqueeze(-1),            # (b, 1)
      ], dim=-1)                      # (b, dim + class_dim + 1)
      return self.mlp(xyt)            # (b, dim)

为什么这样写。

• num_classes + 1 里的那个 +1，就是 $\varnothing$。 这是 CFG 落地的全部工程代价——一行代码里的一个加一。Sanity Check 里 num_classes=3（三个高斯模式），所以嵌入表有 4 行，null_label=3 指向最后一行。Part 3 的 DiT 里则写成 n_classes=11（10 个数字 + 1 个 $\varnothing$），是同一件事的另一种写法。如果你写成 nn.Embedding(num_classes, class_dim)，训练会在第一次抽到空标签时抛 index out of range——这算是好消息，因为它会立刻暴露。
• t.unsqueeze(-1)。 $t$ 进来是 (b,)，拼接要求所有张量除拼接维外形状一致，所以要撑成 (b, 1)。这里的 $t$ 是原始标量，没有做 Fourier 编码——对二维玩具问题够用，但对图像就不够了，这正是 Q3.1 存在的理由。
• 输入维度 dim + class_dim + 1。 三段拼接后的总宽度。写错这个数字会在第一次前向时报矩阵乘法维度不匹配，属于「好错误」。

2.3 Sanity Check 2.4 · 在二维高斯混合上验证 CFG

Q2.2 + Q2.3 写完之后，题面给了一个只要三分钟的自检：三个模式的对称高斯混合（半径 2、标准差 0.2），$\eta=0.25$，$\varnothing=3$，训练 3000 步。务必先跑通这个再去碰 Part 3——它能在三分钟内验证你的训练循环是对的，而 DiT 要几十分钟。

3. Part 3 · 从零搭一个 Diffusion Transformer

这是本 lab 篇幅最大的一部分：五个 TODO 把一个完整的 DiT 搭出来。先看总装图，后面每一题都是它的一个零件：

其中 64 = (32/4)^2 是 token 数（patch size 取 4），256 是隐藏维度。整个网络实现的函数是 $u_t^\theta:\R^{1\times32\times32}\times[0,1]\times\set{0,\dots,10}\to\R^{1\times32\times32}$——输入输出同形，因为向量场和它作用的点住在同一个空间。

3.1 Q3.1 · FourierEncoder

题目在问什么。 把标量 $t\in[0,1]$ 映射成 $d$ 维向量：

其中频率 $w_i\sim\N(0,1)$ 在初始化时随机抽取（本实现里做成 nn.Parameter，即可训练）。

数学依据：为什么标量 $t$ 不能直接喂进网络？

参考实现。

class FourierEncoder(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        assert dim % 2 == 0
        self.half_dim = dim // 2
        self.weights = nn.Parameter(torch.randn(1, self.half_dim))   # (1, d/2)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        """ t: (b,)  ->  embeddings: (b, dim) """
        t = t.view(-1, 1)                            # (b, 1)
        freqs = 2 * math.pi * t * self.weights       # (b, 1) * (1, d/2) -> (b, d/2)
        cos_emb = torch.cos(freqs)                   # (b, d/2)
        sin_emb = torch.sin(freqs)                   # (b, d/2)
        return torch.cat([cos_emb, sin_emb], dim=-1) # (b, dim)

为什么这样写。

• t.view(-1, 1) 加 weights 形状 (1, half_dim)，靠广播一次得到 (b, half_dim)。 这是本题唯一的形状技巧：外积用广播实现，不需要 einsum。
• 2 * math.pi 不能省。 省掉它并不会让模型不能训（相当于把 $w$ 的尺度缩小 $2\pi$ 倍），但会让频率分布偏低，时间分辨能力下降。既然公式写了就照写。
• nn.Parameter 而非 register_buffer。 题面给的公式里 $w_i$ 是固定的随机数，但做成可训练参数只有好处：网络可以自己把频率调到任务需要的量级。注意这也意味着它会被 model.parameters() 计入，别对模型大小感到意外。
• assert dim % 2 == 0。 因为要对半分给 $\cos$ 和 $\sin$。

fe = FourierEncoder(dim=64)
emb = fe(torch.rand(8))
assert torch.allclose((emb ** 2).sum(-1), torch.full((8,), 32.0))  # 32 = 64 / 2

3.2 Q3.2 · Patchifier

题目在问什么。 把图像 (b, c, 32, 32) 变成 token 序列 (b, n, d)，其中 $n=(32/p)^2$。两步：一层卷积把空间降到 $32/p$、通道升到 $d$；一次 Rearrange 把两个空间维压平成序列维。

数学依据：为什么 kernel = stride = patch 的卷积就是 patch embedding？

参考实现。

class Patchifier(nn.Module):
  def __init__(self, img_size: int, patch_size: int, c_in: int, dim: int):
    super().__init__()
    assert img_size % patch_size == 0, "Image size must be divisible by patch size"
    self.net = nn.Sequential(
        nn.Conv2d(c_in, dim, kernel_size=patch_size, stride=patch_size),
        #   (b, c, H, W) -> (b, d, H/p, W/p)
        Rearrange('b d h w -> b (h w) d'),
        #   -> (b, n, d),  n = (H/p) * (W/p)
    )

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.net(x)

为什么这样写。

• c_in 是参数而不是硬编码的 1。 Part 5 会把同一个类用在 (b, 128, 4, 4) 的隐变量上（c_in=128、patch_size=1）。不要硬编码通道数，否则 Part 5 要重写。
• Rearrange('b d h w -> b (h w) d') 的顺序至关重要。 括号里的 (h w) 表示「$h$ 是慢变维、$w$ 是快变维」，即 token 按行优先排列：第 0 个 token 是左上角，第 1 个是它右边一格。Depatchifier 必须用完全对应的写法折回去，否则图像会被转置或打乱。
• assert img_size % patch_size == 0。 否则卷积会悄悄丢掉边缘像素（PyTorch 的 Conv2d 默认 floor），你会得到一个尺寸对不上的 token 数。

P = Patchifier(img_size=32, patch_size=4, c_in=1, dim=48)
x1 = torch.zeros(1, 1, 32, 32)
x2 = x1.clone(); x2[0, 0, 0:4, 0:4] = 1.0
d = (P(x1) - P(x2)).abs().sum(-1)[0]        # (n,)
assert (d > 1e-6).sum().item() == 1         # 只有 1 个 token 变了

3.3 Q3.3 · MHA + DiffusionTransformerLayer + DiffusionTransformer

本 lab 最重的一题，一次要写三个类。题面建议自顶向下想、自底向上写，这里按后者组织。

3.3.1 MHA：从零手写多头自注意力

题目在问什么。 实现

其中 $X\in\R^{n\times d}$ 是 token 序列，$H$ 是头数，$d_h=d/H$ 是每个头的维度。

参考实现。

class MHA(nn.Module):
  """ Multi-headed self-attention """
  def __init__(self, dim: int, heads: int):
    super().__init__()
    assert dim % heads == 0
    self.heads = heads
    self.head_dim = dim // heads
    self.scale = self.head_dim ** -0.5        # 1 / sqrt(d_h)
    self.qkv = nn.Linear(dim, 3 * dim, bias=False)
    self.proj = nn.Linear(dim, dim)

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    """ x: (b, n, d) -> (b, n, d) """
    # 1. 一次线性层同时算出 q, k, v
    q, k, v = self.qkv(x).chunk(3, dim=-1)                       # 各 (b, n, d)

    # 2. 把「头」折进 batch 维：(b, n, h*e) -> (b*h, n, e)
    q = rearrange(q, 'b n (h e) -> (b h) n e', h=self.heads)     # (b*h, n, e)
    k = rearrange(k, 'b n (h e) -> (b h) n e', h=self.heads)
    v = rearrange(v, 'b n (h e) -> (b h) n e', h=self.heads)

    # 3. 缩放点积
    scores = torch.einsum('b i e, b j e -> b i j', q, k) * self.scale   # (b*h, n, n)
    attn = torch.softmax(scores, dim=-1)                               # 沿 key 维归一化

    # 4. 加权求和 value
    out = torch.einsum('b i j, b j e -> b i e', attn, v)               # (b*h, n, e)

    # 5. 头折回特征维
    out = rearrange(out, '(b h) n e -> b n (h e)', h=self.heads)       # (b, n, d)

    # 6. 输出投影
    return self.proj(out)

为什么这样写。

• nn.Linear(dim, 3 * dim) 一次算 qkv。 数学上等价于三个独立的 nn.Linear(dim, dim)，但只调一次 GEMM，对 GPU 友好得多。chunk(3, dim=-1) 沿最后一维切成三份。
• bias=False 是有理由的。 qkv 之后紧接着的是 $\softmax(qk^\top)$，而 $q$ 上的常数偏置 $q\to q+b$ 会给所有 logits 加上 $b\cdot k_j$——这一项不被 softmax 消掉，所以并非完全冗余，但实践中它的作用可以被后面的 LayerNorm 与 proj 的 bias 覆盖。省掉能减参数、也是主流实现的做法。
• rearrange('b n (h e) -> (b h) n e')：把头折进 batch。 这是本题最漂亮的一步。多头注意力本质上就是「$H$ 个互不相干的单头注意力并行跑」，所以把 $H$ 塞进 batch 维之后，后面的点积、softmax、加权求和全都是普通的三维张量运算，不需要写四维 einsum。折回去时用 '(b h) n e -> b n (h e)'，einops 会自动保证与折进去时的顺序一致——这正是用 einops 而非手写 view/permute 的价值。
• softmax(dim=-1)：沿 key 维归一化。 scores 的形状是 (b*h, i, j)，$i$ 是 query 下标、$j$ 是 key 下标。归一化必须沿 $j$，含义是「第 $i$ 个 query 对所有 key 的注意力权重之和为 1」。写成 dim=-2 是最隐蔽的一个 bug：形状完全一样，不会报错，但语义变成「每个 key 被所有 query 分配的权重之和为 1」，训练会明显变差却看不出原因。
• 最后的 self.proj 不能省。 没有它，各个头的输出只是被拼在一起，头与头之间从不交换信息；输出投影负责把 $H$ 个头的结果混合成新的特征。

mha = MHA(dim=64, heads=8).eval()
x = torch.randn(2, 10, 64)
with torch.no_grad():
    mine = mha(x)
    q, k, v = mha.qkv(x).chunk(3, dim=-1)
    q = q.view(2, 10, 8, 8).transpose(1, 2)     # (b, h, n, e)
    k = k.view(2, 10, 8, 8).transpose(1, 2)
    v = v.view(2, 10, 8, 8).transpose(1, 2)
    ref = torch.nn.functional.scaled_dot_product_attention(q, k, v)
    ref = mha.proj(ref.transpose(1, 2).reshape(2, 10, 64))
print((mine - ref).abs().max())     # 实测：8.94e-08

3.3.2 DiffusionTransformerLayer：adaLN-Zero 条件调制

题目在问什么。 实现 DiT 论文里的 adaLN-Zero block：

数学依据：AdaLN 在做什么。 这一部分 Lecture 4 §6.1 已经完整讲过，这里只复述关键的一句：LayerNorm 先把每个 token 的特征向量标准化（减均值除标准差），故意抹掉「尺度」和「偏置」信息；AdaLN 紧接着用由条件 $c$ 决定的 $(\gamma,\beta)$ 把尺度和偏置重新写回去。于是时间和类别不是作为额外 token「混」在序列里，而是作为一个全局旋钮，逐通道地控制网络内部每一处激活的强弱。

写成 $x\cdot(1+\gamma)+\beta$ 而不是 $x\cdot\gamma+\beta$，是为了让 $\gamma=0$ 对应「恒等缩放」。这一点和下面的零初始化配合起来才有意义。

参考实现。

class DiffusionTransformerLayer(nn.Module):
  def __init__(self, dim: int, heads: int):
    super().__init__()
    # 仿射参数关掉：缩放平移全部交给条件网络产生
    self.norm1 = nn.LayerNorm(dim, elementwise_affine=False)
    self.norm2 = nn.LayerNorm(dim, elementwise_affine=False)

    # 由条件 c 一次预测 6 组调制参数；final_init=True 把最后一层权重和 bias 清零
    self.cond = MLP([dim, dim, 6 * dim], final_init=True)

    self.attn = MHA(dim, heads)
    self.ff   = MLP([dim, 4 * dim, dim])

  def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
    """ x: (b, n, d),  c: (b, d)  ->  (b, n, d) """
    # (b, 6d) -> 6 个 (b, d) -> 6 个 (b, 1, d)
    gamma1, beta1, alpha1, gamma2, beta2, alpha2 = \
        [p.unsqueeze(1) for p in self.cond(c).chunk(6, dim=-1)]

    # 注意力支路
    h = self.norm1(x) * (1 + gamma1) + beta1     # (b, n, d)
    x = x + alpha1 * self.attn(h)

    # 前馈支路
    h = self.norm2(x) * (1 + gamma2) + beta2     # (b, n, d)
    x = x + alpha2 * self.ff(h)
    return x

其中 MLP 的 final_init 是题面已经给好的：

class MLP(nn.Module):
  def __init__(self, dims, activation=torch.nn.SiLU, final_init: bool = False):
    ...
    if final_init:
      nn.init.zeros_(self.net[-1].weight)
      nn.init.zeros_(self.net[-1].bias)

layer = DiffusionTransformerLayer(dim=32, heads=4).eval()
x = torch.randn(4, 9, 32); c = torch.randn(4, 32)
print((layer(x, c) - x).abs().max())    # 实测：0.00e+00

3.3.3 DiffusionTransformer：位置编码 + 堆叠

题目在问什么。 加位置编码，然后串 depth 层。

数学依据：为什么非要位置编码？ 自注意力对 token 顺序是置换等变的：设 $P$ 是任意置换矩阵，则

证明是一行：$Q,K,V$ 都是逐 token 的线性映射，所以 $(PX)W=P(XW)$；而 $\softmax\big((PQ)(PK)^\top/\sqrt{d_h}\big)(PV)=\softmax(PQK^\top P^\top/\sqrt{d_h})PV=P\,\softmax(QK^\top/\sqrt{d_h})V$。也就是说，如果把 64 个 image token 随机打乱，注意力算出来的结果也只是同样打乱一遍——网络根本不知道哪个 patch 在左上角。对图像来说这是灾难性的。

解决办法是在输入上加一个与位置有关的向量，把置换对称性打破。本 lab 用可学习位置编码：token 数是固定的（图像尺寸固定），所以直接学一张 $n\times d$ 的表就行，不需要正弦编码那样的外推能力。

参考实现。

class DiffusionTransformer(nn.Module):
  def __init__(self, depth: int, n_tokens: int, dim: int, **layer_kwargs):
    super().__init__()
    # 自注意力对 token 顺序置换等变，必须显式注入位置信息
    self.pos_embed = nn.Parameter(torch.randn(n_tokens, dim) * 0.02)   # (n, d)
    self.layers = nn.ModuleList([
        DiffusionTransformerLayer(dim=dim, **layer_kwargs) for _ in range(depth)
    ])

  def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
    """ x: (b, n, d),  c: (b, d) """
    x = x + self.pos_embed.unsqueeze(0)      # (1, n, d) 广播到 (b, n, d)
    for layer in self.layers:
        x = layer(x, c)
    return x

为什么这样写。

• * 0.02 的初始化尺度。 位置编码是加到 patch embedding 上的，如果初始幅度和特征本身同量级（标准正态，尺度 1），会在训练早期严重干扰内容信息。0.02 是 GPT/ViT 系列的通用取值，含义是「一开始几乎不影响，让网络自己决定要多少位置信息」。用 torch.randn(n, dim) 不缩放也能训，但前几千步会明显更慢。
• nn.Parameter 而不是 buffer。 它是要学的。
• unsqueeze(0)。 (n, d) -> (1, n, d)，广播到 batch。实际上 PyTorch 的广播规则会自动在左边补 1，所以直接写 x + self.pos_embed 也对——但显式写出来更不容易读错，题面也特意提醒了「be careful about broadcasting」。
• nn.ModuleList 而不是 Python list。 用普通 list 装子模块，参数不会被 model.parameters() 收集，也不会随 .to(device) 迁移。症状是：模型大小打印出来只有几百 KB，训练时报「输入在 cuda 但权重在 cpu」，或者干脆损失一动不动。这是 PyTorch 新手最经典的坑之一。
• 不能用 nn.Sequential。 每层要吃两个输入 (x, c)，Sequential 只支持单输入。必须显式写循环。

3.4 Q3.4 · Depatchifier

题目在问什么。 Patchifier 的逆操作：把 (b, n, d) 变回 (b, c_out, 32, 32)。题面给了四步配方：归一化 → MLP 映到 $f p^2$ 维 → Rearrange 折回空间 → 输出卷积。

数学依据。 没有新数学，但有一个必须理解的对偶关系。Patchifier 把每个 $p\times p\times c$ 的块压成一个 $d$ 维 token；Depatchifier 要把每个 $d$ 维 token 展回一个 $p\times p\times f$ 的块，再由最后一层卷积把 $f$ 个通道压到 $c_{\text{out}}$。所以中间那个 MLP 的输出维度必须是 $f\cdot p^2$——这个数字不是随便定的，它由「一个 token 要负责多少个像素」唯一确定。

参考实现。

class Depatchifier(nn.Module):
  def __init__(self, img_size: int, patch_size: int, dim: int,
               final_dim: int, c_out: int):
    super().__init__()
    assert img_size % patch_size == 0
    h = w = img_size // patch_size
    self.net = nn.Sequential(
        nn.RMSNorm(dim, elementwise_affine=False),                       # (b, n, d)
        MLP([dim, 4 * dim, final_dim * patch_size ** 2]),                # -> (b, n, f*p*p)
        Rearrange('b (h w) (f ph pw) -> b f (h ph) (w pw)',
                  h=h, w=w, f=final_dim, ph=patch_size, pw=patch_size),  # -> (b, f, H, W)
        nn.Conv2d(final_dim, c_out, kernel_size=3, padding=1),           # -> (b, c_out, H, W)
    )

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.net(x)

为什么这样写。

• 开头的归一化。 DiT 主干最后一层的输出尺度是不受控的（残差累加），先归一化再映射能稳住训练。RMSNorm 和 LayerNorm 都行，前者省掉减均值这一步，稍快。同样要 elementwise_affine=False——理由和 adaLN 里一样，尺度信息紧接着由后面的 MLP 自由决定，不需要再多一组冗余参数。
• Rearrange 的模式必须与 Patchifier 严格对偶。 逐段对照：
  • b (h w)——把序列维拆回两个空间维，$h$ 慢变、$w$ 快变，与 Patchifier 的 'b d h w -> b (h w) d' 一致。
  • (f ph pw)——把特征维拆成「通道、块内行、块内列」三段，顺序是 $f$ 最慢、$pw$ 最快。
  • b f (h ph) (w pw)——把「块坐标」和「块内坐标」交织成最终的像素坐标：第 $(i\cdot p+a)$ 行来自第 $i$ 个块的第 $a$ 行。
  这个交织顺序是整个函数的灵魂。如果写成 (ph h) 而不是 (h ph)，得到的图像会是一个「马赛克重排」——所有块的第 0 行被拼在一起、所有块的第 1 行被拼在一起，看起来像被撕碎重贴。
• h=h, w=w, f=..., ph=..., pw=... 这些关键字参数不能省。 einops 无法从形状唯一地推断出 (h w) 到底怎么拆（$64=8\times8$ 也可以是 $4\times16$），必须显式告诉它。省掉会抛异常，属于「好错误」。
• 最后一层 Conv2d(kernel_size=3, padding=1)。 为什么不直接让 MLP 输出 $c_{\text{out}}\cdot p^2$ 一步到位？因为 patch 之间是独立展开的，块与块的边界处会有不连续（每个 token 只管自己那 $4\times4$ 个像素）。一层 $3\times3$ 卷积跨过边界做局部平滑，能明显减轻方块状伪影。final_dim=10 就是给这层卷积留的「工作空间」通道数。

3.5 Q3.5 · DiffusionTransformerFlowModel

题目在问什么。 把前四题拼成完整的 $u_t^\theta(x\mid y)$。

数学依据：为什么时间嵌入和类别嵌入是「相加」而不是「拼接」？ 两者都被编码到同一个 $d=256$ 维空间，相加之后送进 adaLN 的条件 MLP。相加的合理性在于：条件 MLP 的第一层是线性的，所以

也就是说相加后过线性层，与分别过线性层再相加，完全等价——这正是拼接后过一个（分块的）线性层能做到的事。相加只是把这个分块结构固定下来，省掉一半参数。代价是两路信号共用同一组坐标，可能互相干扰；但实践证明对「时间 + 类别」这种低信息量的条件完全够用。Stable Diffusion 3 和 DiT 原论文用的都是相加。

参考实现。

class DiffusionTransformerFlowModel(ConditionalVectorField):
  def __init__(self, img_size=32, patch_size=8, num_layers=12, c=1,
               dim=256, heads=4, final_dim=10, n_classes=11):
      super().__init__()
      # 0. 时间与类别嵌入
      self.time_embedder = nn.Sequential(
          FourierEncoder(dim),      # (b,) -> (b, dim)
          MLP([dim, dim, dim]),     # -> (b, dim)
      )
      self.y_embedder = nn.Embedding(n_classes, dim)   # 11 = 10 个数字 + 1 个空标签

      # 1. Patchifier
      self.patchifier = Patchifier(img_size=img_size, patch_size=patch_size,
                                   c_in=c, dim=dim)
      # 2. DiT 主干
      n_tokens = (img_size // patch_size) ** 2
      self.dit = DiffusionTransformer(depth=num_layers, n_tokens=n_tokens,
                                      dim=dim, heads=heads)
      # 3. Depatchifier
      self.depatchifier = Depatchifier(img_size=img_size, patch_size=patch_size,
                                       dim=dim, final_dim=final_dim, c_out=c)

  def forward(self, x, t, y):
    """ x: (b, c, h, w),  t: (b,),  y: (b,)  ->  (b, c, h, w) """
    c = self.time_embedder(t) + self.y_embedder(y)   # (b, dim)   <- 相加
    x = self.patchifier(x)                           # (b, n, dim)
    x = self.dit(x, c)                               # (b, n, dim)
    x = self.depatchifier(x)                         # (b, c, h, w)
    return x

为什么这样写。

• FourierEncoder 后面还要跟一个 MLP。 Fourier 编码只是把标量升成了一组三角函数值，它本身是固定的、无学习能力的特征。后面的 MLP([dim, dim, dim]) 才让网络能把这些特征组合成任务需要的时间表示。少了它，adaLN 的条件 MLP 要直接从原始三角特征学起，收敛会慢很多。
• n_classes=11 而不是 10。 又一次是 CFG 的那个 $\varnothing$。null_label=10，正好是最后一个索引。
• 局部变量 c 覆盖了参数名 c（通道数）。 这是官方代码的写法，在 forward 里没问题（参数 c 只在 __init__ 用），但读代码时容易困惑：forward 里的 c 是条件向量 (b, dim)，__init__ 里的 c 是通道数。
• 实际训练用的超参和默认值不同。 默认签名是 patch_size=8, num_layers=12, heads=4，但训练时传的是 patch_size=4, num_layers=8, heads=8。$p=4$ 意味着 $n=64$ 个 token（而不是 16 个），空间分辨率高一倍，代价是注意力的计算量是 $O(n^2)$，涨了 16 倍。对 MNIST 这个尺寸完全吃得消，效果明显更好。

3.6 训练 DiT：跑出来应该长什么样

训练配置：$\eta=0.35$、$\varnothing=10$、20000 步、batch 256、学习率 $4\times10^{-4}$、warmup 500 步。模型 40.324 MiB。单张 RTX 5080 上约 30 分钟。

4. Part 4 · 训练一个变分自编码器

七个 TODO，但难度分布很不均匀：Q4.1–Q4.6 是形状流水账（照配方搭积木，唯一的挑战是把通道数和分辨率对上），真正需要动脑的是 Q4.7 的损失函数。

先明确目标。VAE 由两半组成：

本 lab 的具体形状是 $x\in\R^{1\times32\times32}$（$d=1024$）、$z\in\R^{128\times4\times4}$（$k=2048$）。注意隐空间的元素数比图像还多——这在 latent diffusion 里很常见，压缩发生在空间维（$32\times32\to4\times4$，降 64 倍），而通道维膨胀了 128 倍。真正的收益是让扩散模型面对 $4\times4$ 的「图像」，注意力的 token 数从 64 降到 16。

两个 $\sigma$ 都被参数化成可学习的共享标量 nn.Parameter(torch.zeros(()))，存的是 $\log\sigma^2$ 而不是 $\sigma^2$——这样可以无约束地取任意实数，再用 exp 保证正性，避免了显式的正性约束。

4.1 Q4.1 · ResidualBlock

题目在问什么。 实现 $x\mapsto x+\text{Conv}_{1\times1}\big(\text{SiLU}(\text{Conv}_{3\times3}(\text{Norm}(x)))\big)$。

参考实现。

class ResidualBlock(nn.Module):
  def __init__(self, channels: int, act: nn.Module = nn.SiLU):
    super().__init__()
    self.norm  = nn.GroupNorm(1, channels)                        # 等价于图像上的 LayerNorm
    self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
    self.act   = act()
    self.conv2 = nn.Conv2d(channels, channels, kernel_size=1)

    # 第二个卷积零初始化 => 初始时整个块是恒等映射
    nn.init.zeros_(self.conv2.weight)
    nn.init.zeros_(self.conv2.bias)

  def forward(self, x: torch.Tensor):
    x_skip = x                    # (b, c, h, w)
    x = self.norm(x)
    x = self.conv1(x)
    x = self.act(x)
    x = self.conv2(x)
    return x_skip + x             # (b, c, h, w)

为什么这样写。

• nn.GroupNorm(1, channels) 为什么等价于 LayerNorm？ GroupNorm 把 $C$ 个通道分成 $G$ 组，每组内跨「通道 + 空间」计算均值方差。取 $G=1$ 就是把整张特征图（所有通道、所有像素）当成一个整体归一化——这正是 LayerNorm 作用在 (C, H, W) 上的定义。相比 BatchNorm 的好处是不依赖 batch 统计量，训练和推理行为完全一致，batch size 小的时候也稳。nn.LayerNorm 需要显式给出 [C, H, W]，而 VAE 里同一个 block 会用在不同分辨率上，用 GroupNorm 更省事。
• padding=1 配 kernel_size=3。 保持空间尺寸不变（$3\times3$ 卷积 + padding 1 是「same」卷积）。残差连接要求输入输出同形，少了 padding 会掉两个像素而报错。
• 第二个卷积是 $1\times1$。 它的作用是「通道混合」而非「空间聚合」——空间感受野已经由 conv1 提供了。这也是「$3\times3$ 提特征、$1\times1$ 做投影」的经典瓶颈结构简化版。
• 零初始化 conv2。 与 adaLN-Zero 同一个思想：初始时残差支路输出恒为 0，整个 block 是恒等映射。VAE 有 $4\times2=8$ 个残差块加上下采样，深度不算浅，这一手能明显稳住前几百步。

rb = ResidualBlock(channels=8).eval()
x = torch.randn(2, 8, 16, 16)
print((rb(x) - x).abs().max())     # 实测：0.00e+00

4.2 Q4.2 · AttnBlock

题目在问什么。 把空间位置当作 token，在特征图上做一次自注意力 + 前馈，各带残差。

数学依据：卷积的局限与注意力的补位。 卷积的感受野是局部且随深度线性增长的；即使堆到编码器最后一层（$4\times4$），每个位置也只能看到有限邻域。而 VAE 需要把整张图的信息压进 $z$，全局交互是必要的。自注意力提供的正是「任意两个位置一步直连」的能力，代价是 $O((HW)^2)$ 的复杂度——这也是为什么 Stable Diffusion 的 VAE 只在最低分辨率层放注意力。

参考实现。

class AttnBlock(nn.Module):
  def __init__(self, channels: int):
    super().__init__()
    self.reshape1 = Rearrange('b c h w -> b (h w) c')
    self.norm1 = nn.LayerNorm(channels)
    self.attn  = MHA(channels, heads=1)     # heads=1：保证任意通道数都能整除
    self.norm2 = nn.LayerNorm(channels)
    self.ff    = MLP([channels, 4 * channels, channels])

  def forward(self, x: torch.Tensor):
    b, c, h, w = x.shape
    x = self.reshape1(x)                    # (b, h*w, c)

    x_skip = x                              # 注意力 + 残差
    x = self.attn(self.norm1(x))
    x = x_skip + x

    x_skip = x                              # 前馈 + 残差
    x = self.ff(self.norm2(x))
    x = x_skip + x

    return rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)

为什么这样写。

• 「像素即 token、通道即特征」。 'b c h w -> b (h w) c' 把 $H\times W$ 个空间位置变成序列，通道数当作 token 维度。于是 Q3.3 写的 MHA 一个字都不用改就能复用——这是本 lab 一处很好的架构设计，值得注意。
• heads=1 不是偷懒。 MHA 里有 assert dim % heads == 0。VAE 的通道数依次是 16 / 32 / 64 / 128，如果写死 heads=4，16 通道那层每个头只有 4 维，表达力堪忧；写 heads=8 则 16 通道刚好整除但更窄。heads=1 对所有这些通道数都成立，且在 $4\times4=16$ 个 token 的规模上多头并没有优势。
• 用 x.shape 解包出 h, w 再传给最后的 rearrange。 必须的：einops 没法从 (b, hw, c) 反推 $h$ 和 $w$。
• Pre-norm 而非 post-norm。 归一化放在残差支路内部（x + f(norm(x))），而不是在残差之后（norm(x + f(x))）。前者的主干路径是一条无阻碍的恒等通路，梯度传播更好，是现代 transformer 的标准做法。

4.3 Q4.3 · EncoderBlock 与 Q4.5 · DecoderBlock

这两题结构完全镜像，一起讲。题目在问什么：两个残差块 + 一个注意力块 + （可选的）下采样 / 上采样。

class EncoderBlock(nn.Module):
  def __init__(self, in_channels: int, downsample_channels: Optional[int] = None):
    super().__init__()
    self.res1 = ResidualBlock(in_channels)
    self.res2 = ResidualBlock(in_channels)
    self.attn = AttnBlock(in_channels)
    if downsample_channels is not None:
      # stride=2 的卷积同时完成「降采样」和「换通道数」
      self.downsample = nn.Conv2d(in_channels, downsample_channels,
                                  kernel_size=3, padding=1, stride=2)
    else:
      self.downsample = None

  def forward(self, x):
    x = self.res1(x); x = self.res2(x); x = self.attn(x)
    if self.downsample is not None:
      x = self.downsample(x)
    return x


class DecoderBlock(nn.Module):
  def __init__(self, in_channels: int, upsample_channels: Optional[int] = None):
    super().__init__()
    self.res1 = ResidualBlock(in_channels)
    self.res2 = ResidualBlock(in_channels)
    self.attn = AttnBlock(in_channels)
    if upsample_channels is not None:
      self.upsample = nn.Sequential(
        nn.Upsample(scale_factor=2, mode='nearest'),
        nn.Conv2d(in_channels, upsample_channels,
                  kernel_size=3, padding=1, stride=1),
      )
    else:
      self.upsample = None

  def forward(self, x):
    x = self.res1(x); x = self.res2(x); x = self.attn(x)
    if self.upsample is not None:
      x = self.upsample(x)
    return x

为什么这样写。

• 下采样用 stride=2 的卷积，而不是池化。 池化是固定的（不可学习）降采样；stride=2 卷积让网络自己学「怎么降」，同时顺手把通道数从 in_channels 换成 downsample_channels，一步两用。kernel_size=3, padding=1, stride=2 作用在偶数尺寸上恰好减半：$\lfloor(32+2-3)/2\rfloor+1=16$。
• Optional[int] = None 的语义是「最后一块不改变分辨率」。 见下一题的形状表。
• 注意力放在残差块之后、采样之前。 这样注意力总是在「还没降维」的分辨率上做全局混合，信息损失更小。

4.4 Q4.4 · Encoder

题目在问什么。 初始卷积 → 每个 hidden channel 一个 EncoderBlock（除最后一个外都下采样）→ 归一化 + $1\times1$ 卷积输出 z_mean，外加一个标量 logvar。

形状流水账（hidden_channels = [16, 32, 64, 128]）。这是本题唯一真正要想清楚的东西：

参考实现。

class Encoder(nn.Module):
  def __init__(self, in_channels: int, hidden_channels: list[int]):
    super().__init__()
    self.init_conv = nn.Conv2d(in_channels, hidden_channels[0],
                               kernel_size=3, padding=1, stride=1)

    # 关键的一行：把 [16,32,64,128] 配对成 (16,32),(32,64),(64,128),(128,None)
    ch_in  = hidden_channels                 # [16, 32, 64, 128]
    ch_out = hidden_channels[1:] + [None]    # [32, 64, 128, None]
    self.blocks = nn.ModuleList([
        EncoderBlock(i, o) for i, o in zip(ch_in, ch_out)
    ])

    z_dim = hidden_channels[-1]
    self.z_mean = nn.Sequential(
      nn.GroupNorm(1, z_dim),
      nn.Conv2d(z_dim, z_dim, kernel_size=1, stride=1, padding=0),
    )
    self.logvar = nn.Parameter(torch.zeros(()))    # 标量！形状是 ()

  def forward(self, x: torch.Tensor):
    x = self.init_conv(x)
    for block in self.blocks:
      x = block(x)
    return self.z_mean(x), self.logvar

为什么这样写。

• ch_out = hidden_channels[1:] + [None] 是整题的题眼。 它把「$n$ 个通道数」变成「$n$ 个 (输入, 输出) 对」，最后一个的输出是 None，正好触发 EncoderBlock 里「不下采样」的分支。四个 block、三次下采样，$32\to16\to8\to4$。
• 为什么最后一块不下采样？ 因为 $4\times4$ 就是目标分辨率了。多一个 block 是为了在最终分辨率上再做两轮残差 + 一次全局注意力，把信息充分整合后再输出 $\mu_\phi$。
• self.logvar = nn.Parameter(torch.zeros(())) —— 注意是 torch.zeros(()) 不是 torch.zeros(1)。 前者形状是 torch.Size([])（真正的 0 维标量），后者是 torch.Size([1])。两者在广播时行为几乎一样，但测试里会检查 z_logvar.shape == torch.Size([])。语义上也确实是「一个数」：所有样本、所有维度共享同一个后验方差。
• 共享标量方差意味着什么？ 标准 VAE 里 $\sigma_\phi(x)$ 是输入相关的（每个 $x$ 有自己的不确定度）。这里退化成一个全局常数，等于说「编码器对所有输入的置信度一样」。这是 latent diffusion 里的常见简化（Lecture 4 §7.7 讨论过），好处是训练稳定得多，代价是模型在概率意义上更弱——但反正我们只是要一个「被轻度正则化过的自编码器」。

4.5 Q4.6 · Decoder

题目在问什么。 Encoder 的镜像：每个 hidden channel 一个 DecoderBlock（除最后一个外都上采样）→ 归一化 + $1\times1$ 卷积输出 x_mean + 标量 logvar。注意没有 init_conv——输入已经是 128 通道了。

形状流水账。关键在于 VAE.__init__ 里传的是 list(reversed(hidden_channels))，即 [128, 64, 32, 16]：

class Decoder(nn.Module):
  def __init__(self, out_channels: int, hidden_channels: list[int]):
    super().__init__()
    ch_in  = hidden_channels                 # [128, 64, 32, 16]
    ch_out = hidden_channels[1:] + [None]    # [64, 32, 16, None]
    self.blocks = nn.ModuleList([
        DecoderBlock(i, o) for i, o in zip(ch_in, ch_out)
    ])

    x_dim = hidden_channels[-1]              # 16
    self.x_mean = nn.Sequential(
      nn.GroupNorm(1, x_dim),
      nn.Conv2d(x_dim, out_channels, kernel_size=1, stride=1, padding=0),
    )
    self.logvar = nn.Parameter(torch.zeros(()))

  def forward(self, x: torch.Tensor):
    for block in self.blocks:
      x = block(x)
    return self.x_mean(x), self.logvar

为什么这样写。 结构与 Encoder 逐行对称，唯一要小心的是 reversed 是在 VAE.__init__ 里做的，不是在 Decoder 里：

class VAE(nn.Module):
  def __init__(self, data_channels: int, hidden_channels: list[int], beta: float = 0.1):
    super().__init__()
    self.beta = beta
    self._encoder = Encoder(data_channels, hidden_channels)                 # [16,32,64,128]
    self._decoder = Decoder(data_channels, list(reversed(hidden_channels))) # [128,64,32,16]

所以 Decoder 内部的写法和 Encoder 完全一样，不要在 Decoder 里再 reverse 一次——那会让通道数变成 $16\to32\to64\to128$，但输入是 128 通道，第一层就崩。

4.6 Q4.7 · VAE.compute_loss —— 本 Part 唯一需要动脑的一题

题目在问什么。 实现讲义那条完整的 $\beta$-VAE 损失（Lecture 4 §7.4 末尾给出，notebook 里标注为式 (83)/(85)，不同版本编号略有出入）：

数学依据：两项各自从哪来。

参考实现。

def compute_loss(self, z_mean, z_logvar, x_mean, x_logvar, x_true):
    eps = 1e-6

    # KL 项：把 (mu, sigma^2) 推向 (0, 1)
    kl_loss = self.beta * (z_mean.pow(2) + torch.exp(z_logvar) - z_logvar - 1).mean()

    # 重构项：高斯负对数似然（去掉常数）
    mse_term        = (x_true - x_mean).pow(2) / (torch.exp(x_logvar) + eps)
    confidence_term = x_logvar
    recon_loss      = (mse_term + confidence_term).mean()

    return kl_loss + recon_loss

以及 forward 里的重参数化：

def forward(self, x: torch.Tensor):
    z_mean, z_logvar = self.encode(x)                                # (b,128,4,4), ()
    z = z_mean + torch.exp(0.5 * z_logvar) * torch.randn_like(z_mean)  # 重参数化
    x_mean, x_logvar = self.decode(z)                                # (b,1,32,32), ()
    return z_mean, z_logvar, x_mean, x_logvar

4.7 训练 VAE：跑出来应该长什么样

配置：hidden_channels=[16,32,64,128]、$\beta=10$、5000 步、batch 64、学习率 $10^{-3}$。模型 6.150 MiB，在 RTX 5080 上几分钟就跑完。

5. Part 5 · 隐空间扩散

所有零件都造好了，这一 Part 只剩一件事：把 Part 3 的 DiT 搬进 Part 4 的隐空间。只有一个 TODO。

5.1 Q5.1 · LatentCFGTrainer.get_train_loss

题目在问什么。 与 Q2.2 的 CFGTrainer 唯一的区别：训练数据不是像素 $x$，而是 VAE 编码后的隐变量 $z$。而且 VAE 是冻结的（先单独训好，这一步不再更新它的参数）。

数学依据。 Latent diffusion 的整个想法可以写成一个分布的复合：

其中 $p_\theta(x\mid z)$ 是冻结的解码器，$p_{\text{latent}}^{\text{flow}}$ 是我们要训练的流模型在 $t=1$ 时的分布。训练目标是让 $p_{\text{latent}}^{\text{flow}}$ 匹配编码器诱导的聚合后验（aggregate posterior）

而「从 $p_{\text{latent}}$ 采样」的操作，恰好就是「采一张真实图 → 编码 → 重参数化采样」。所以 flow matching 的公式一个字都不用改，只需要把 p_data.sample() 换成这个两步过程。

参考实现。

def get_train_loss(self, batch_size: int) -> torch.Tensor:
    # Step 1: 从 MNIST 采样并编码到隐空间
    #         VAE 是冻结的，整个编码过程放在 no_grad 里
    with torch.no_grad():
      x, y = self.mnist.sample(batch_size)                # x: (b,1,32,32), y: (b,)
      z_mean, z_logvar = self.vae.encode(x)               # (b,128,4,4), ()
      # 用完整的后验采样（而非只取均值），让扩散模型见到编码器真实的隐分布
      z = z_mean + torch.exp(0.5 * z_logvar) * torch.randn_like(z_mean)

    # Step 2: 以概率 eta 把标签置为空标签
    xi = torch.rand(y.shape[0], device=y.device)
    y[xi < self.eta] = self.null_label

    # Step 3: 采 t 和 z_t
    t = torch.rand(batch_size, device=z.device).to(z) * (1 - self.eps)   # (b,)
    zt = self.path.sample_conditional_path(z, t)                         # (b,128,4,4)

    # Step 4: 回归
    ut_theta = self.model(zt, t, y)                          # (b,128,4,4)
    ut_ref   = self.path.conditional_vector_field(zt, z, t)  # (b,128,4,4)
    return torch.square(ut_theta - ut_ref).mean()

为什么这样写。

• torch.no_grad() 是必须的，不只是省显存。 如果不加，编码器的计算图会被保留，loss.backward() 会把梯度一路传回 VAE 的参数。虽然优化器里只注册了 DiT 的参数（self.opt = AdamW(self.model.parameters())），VAE 不会真被更新，但你会白白多花显存和时间。更糟的情况是：如果你的 trainer 写法不小心把 VAE 也塞进了优化器，VAE 会边训边动，扩散模型追着一个移动的目标学，训练直接发散。加 no_grad 是最稳的做法。
• 为什么采样而不是只用 z_mean？ 因为扩散模型要匹配的是聚合后验 $p_{\text{latent}}$，那是带噪声的。只用均值会让扩散模型学到一个比真实隐分布「更瘦」的分布，采样时生成的 $z$ 落在解码器没见过的低方差区域，重构质量下降。这一步的噪声也起到轻微的数据增广作用。
• path 的构造变了。 注意训练脚本里 GaussianConditionalProbabilityPath(p_data=None, p_simple_shape=[128, 4, 4], ...)——p_data 直接传 None，因为 trainer 不再通过 path 取数据（那条路被上面的两步采样取代了）。path 现在只负责三件事：p_simple（隐空间形状的高斯）、sample_conditional_path、conditional_vector_field。
• DiT 的超参变了：img_size=4, patch_size=1, c=128。 patch_size=1 意味着每个空间位置就是一个 token，共 $4\times4=16$ 个 token（像素空间是 64 个）。Patchifier 退化成一个 $1\times1$ 卷积（$128\to256$ 的逐位置线性投影）。这正是 Q3.2 里坚持不硬编码 c_in 的回报。

5.2 训练结果：一个必须解释清楚的现象

配置：10000 步、batch 256、$\eta=0.35$、学习率 $4\times10^{-4}$。模型 39.844 MiB。

6. 为什么隐空间的损失几乎不下降？

这一节是本页的核心结论之一。答案是：flow matching 的损失有一个不可约下界，而这个下界在隐空间里比在像素空间里高一个数量级。

6.1 不可约下界从哪来

6.2 把下界算出来

本 lab 用 LinearAlpha / LinearBeta，即 $\alpha_t=t$、$\beta_t=1-t$。把 $x_t=tz+(1-t)\epsilon$ 代入条件向量场：

所以回归目标就是 $z-\epsilon$，下界是 $\E_{t,x_t}\big[\Var(z-\epsilon\mid x_t)\big]$（这里按 .mean() 的约定，取逐维平均）。

6.3 两个空间的对比

像素空间为什么能降到 0.19？因为 MNIST 高度结构化：图像的绝大部分是纯黑背景，笔画是连通的粗线条，相邻像素强相关。上面那个 $\frac{\pi\sigma}{2}$ 的公式只对各维独立的白噪声成立；当各维强相关时，观测到 $x_t$ 就等于同时观测到了 1024 个带噪线索，它们联合起来对 $z$ 的约束远比单维的情形强，条件方差因此小得多。0.19 远低于 1.571，量化地说明了「MNIST 离白噪声有多远」。

隐空间为什么停在 1.94？因为 $\beta=10$（讲义约定下的 5）的 KL 项正是在主动把隐变量推成白噪声：均值推向 0、方差推向 1、各维之间的冗余被压掉（KL 到 $\N(0,I)$ 的散度对任何偏离独立标准正态的结构都要收费）。于是隐变量非常接近上面那个理想化的假设，损失几乎就卡在理论下界上。

反过来用这个公式还能反推隐变量的尺度：由 $\mathcal{L}^\star=\frac{\pi\sigma}{2}$ 且实测 $\mathcal{L}=1.9431\ge\mathcal{L}^\star$，得

即隐变量每维标准差不超过约 1.24（等号在模型达到最优时取到）。这个数与 Lecture 4 §7.7 提到的现象一致：实际系统里隐分布的尺度并不严格是 1，所以工程实现通常要在编码结果上乘一个全局缩放常数，把经验标准差调到 1 附近再交给扩散模型——因为噪声调度 $\alpha_t,\beta_t$ 都是按「数据方差约为 1」设计的。本 lab 省略了这一步，代价就是那个略高于 $\pi/2$ 的损失平台。

7. 如何验证自己的实现：一套写完就跑的检查

这个 lab 的训练动辄几十分钟，「写完全部代码再开训、跑完才发现形状错」是最贵的失败模式。下面这套检查全部是纯 CPU、秒级完成的，建议每写完一个模块就跑对应的那几条。括号里是本页实现的实测值。

7.1 不变量检查（能抓出绝大多数错误）

7.2 形状检查

本页的实现在这套检查上 22 项全部通过。注意几处「刻意做成小尺寸」的设计：用 img_size=16、num_layers=2、hidden_channels=[8,16,32,64] 就能在 CPU 上瞬间跑完，同时完整覆盖所有形状逻辑——不要用 img_size=32, num_layers=8 做单元测试，慢且没有额外收益。

7.3 推荐的验证顺序

1. Q2.2 + Q2.3 → Sanity Check 2.4（3000 步，约 1–3 分钟）。这一步同时验证训练循环、CFG 组合公式、类别嵌入。不通过就绝不要往下走。
2. Q3.1 → Q3.2 → Q3.3 → Q3.4，每写完一个立刻跑上表对应的行。特别是 MHA 的对拍，它一次能验四件事。
3. Q3.5 写完先跑一次前向 + model_size_b，确认输出形状和 40.324 MiB 这个量级。
4. 启动 DiT 训练，开 ckpt_every=1000，看第 1000/2000 步的图。3000 步还是纯噪声就停下来查。
5. Q4.1–Q4.6，跑形状链路和两个恒等映射检查。
6. Q4.7，跑三条损失性质检查。
7. 启动 VAE 训练（几分钟），跑插值图。插值图不平滑就别继续——Part 5 的质量完全依赖 VAE。
8. Q5.1，启动隐空间训练。看到损失只在 1.9x 徘徊时，不要慌，去看 checkpoint 采样图。

8. 训练实测数据汇总

硬件：单张 RTX 5080（16 GB），PyTorch 2.11.0 + CUDA 12.8。四段训练合计约 1.5 小时（含 MNISTSampler 的缓存改写；不改写的话仅 DiT 一段就要成倍时间）。原题面说「~15 A100 minutes」是指 DiT 那一段。

本 lab 小结

做完这三个 lab，你手上就有了一条完整的链路：Lab 1 的模拟器 → Lab 2 的 flow matching / score matching → Lab 3 的 CFG + DiT + VAE + latent diffusion。把 MNIST 换成更大的数据集、把 DiT 加深、把 VAE 换成带感知损失和对抗损失的版本，你得到的就是 Stable Diffusion 3 的骨架——Lecture 4 §6.3 逐条对照过这件事。

延伸阅读

• Lecture 3-B · 引导与条件生成 — Part 2 的理论出处：引导公式的推导、$w$ 与「锐化分布」的关系、为什么 CFG 不是在采样真实的条件分布。
• Lecture 4 · 网络架构与隐空间 — Part 3–5 的理论出处：AdaLN 的动机（§6.1）、ELBO 的两条推导路径（§6.4）、KL 的闭式解与重参数化（§7）、VAE 在 latent diffusion 里的工程现实（§7.7）。
• Lab 2 · Flow Matching and Score Matching — 本 lab 训练循环的来源；那里的「损失不收敛到 0」正是第 6 节所讲现象的原始版本。
• Scalable Diffusion Models with Transformers (Peebles & Xie, 2023) — DiT 原论文，题面里那张架构图的出处。adaLN-Zero 与其它三种条件注入方式（in-context、cross-attention、adaLN）的消融实验在 §4，值得一读。
• High-Resolution Image Synthesis with Latent Diffusion Models (Rombach et al., 2022) — Stable Diffusion 的原论文，Part 4/5 的 VAE + latent diffusion 结构的直接来源，包括「共享标量 logvar」这个简化。
• Classifier-Free Diffusion Guidance (Ho & Salimans, 2022) — CFG 原论文。$\eta$（论文里记作 $p_{\text{uncond}}$）取值的消融在 §4，本 lab 用的 0.35 偏大，论文推荐 0.1–0.2。
• Auto-Encoding Variational Bayes (Kingma & Welling, 2013) — VAE 与重参数化技巧的原始论文，Q4.7 的全部数学都在这里。
• Deconvolution and Checkerboard Artifacts (Odena et al., 2016) — 棋盘格伪影的经典分析，解释了 Q4.5 为什么用「最近邻 + 卷积」而不是转置卷积。
• Fourier Features Let Networks Learn High Frequency Functions (Tancik et al., 2020) — Q3.1 背后的理论：随机 Fourier 特征如何克服 MLP 的谱偏置，含核函数视角的完整分析。