DiT 架构与条件调制#

分主题阅读#

• SD3 与 FLUX 架构案例：DiT 的条件交互如何在具体图像模型中展开？
• 视频生成的表示与注意力：时间维度给压缩、注意力和计算成本增加了哪些约束？

组件阅读路线 · U-Net → DiT · Patchify · FLUX · LoRA vs ControlNet · Temporal Attention · 3D VAE

1. 为什么要用 Transformer 替换 U-Net#

1.1 U-Net 的辉煌历史#

从 DDPM 到 Stable Diffusion 1.x / 2.x / XL，扩散模型的去噪网络几乎清一色使用 U-Net。这个 encoder-decoder + skip connection 的结构在像素级预测任务中表现极好——多尺度特征融合让模型同时兼顾全局语义和局部细节。

1.2 U-Net 的三大局限#

1. 结构不规整：encoder 各层通道数依次翻倍（如 320→640→1280→1280），decoder 镜像缩减，还有 skip connection 拼接。不同层的参数量、计算量差异很大，难以简单地"堆更多层"来扩大模型。

2. 扩展性差：想让 U-Net 变大，要么加深、要么加宽，但两种方式都受 skip connection 的牵制，改动牵一发动全身：

   • 加深：encoder 和 decoder 必须对称，加一层 encoder 就必须同时加对应的 decoder 层（实际增加 2 层），不能单独扩展某一侧。
   • 加宽：把某一层通道数翻倍，skip connection 要求对应的 decoder 层也必须跟着改，进而影响下一层的输入维度……整个对称结构都要联动修改，越改越麻烦。

   核心问题不是"做不到"，而是每次扩展都要手动重新设计整个对称结构，工程代价极高。SD-XL 的 U-Net 约 2.6B 参数已经是这种人工设计模式的极限附近。

3. 扩展规律需实测：DiT 论文展示了其模型族随训练计算增长的质量趋势；这不证明 U-Net 没有 scaling law，也不保证增大任意 Transformer 都会提升效果。

1.3 Transformer 的结构优势#

核心优势

结构简单统一：整个网络就是堆叠 $N$ 个完全相同的 block（Self-Attention + FFN），想变大就加 block 或加宽 hidden_dim。

成熟的 Scaling Law：从 ViT、GPT-3 到 LLaMA，Transformer 在视觉和语言领域都验证了"参数量 ↑ → 性能 ↑"。

硬件优化成熟：FlashAttention、Tensor Core、模型并行等针对 Transformer 的优化已经非常完善。

1.4 行业趋势#

**一句话总结：**U-Net 像是手工设计的精密零件，Transformer 像是标准化的乐高积木——后者更容易大规模拼装。

2. Vision Transformer (ViT) 回顾#

2.1 核心思路：图像 → 序列#

Dosovitskiy et al. (2020) 提出 ViT 的核心灵感很简单：把图像切成 patch，每个 patch 当成一个 token，然后直接用标准 Transformer 处理。

Patchify 过程#

以 ViT-B/16 为例：

• 输入图像：$224 \times 224 \times 3$（RGB）
• Patch size $P = 16$
• 切出的 patch 数量：$\frac{224}{16} \times \frac{224}{16} = 14 \times 14 = 196$ 个 token
• 每个 patch 展平：$16 \times 16 \times 3 = 768$ 维向量
• 线性投影到 hidden_dim $D$（ViT-B 中 $D=768$，刚好匹配）

$$\text{Patchify}: \mathbb{R}^{H \times W \times C} \rightarrow \mathbb{R}^{N \times (P^2 \cdot C)} \xrightarrow{\text{Linear}} \mathbb{R}^{N \times D}$$

其中 $N = \frac{H}{P} \times \frac{W}{P}$ 是 patch/token 数量

位置编码#

图像 patch 打乱后会丢失空间信息，所以需要给每个 token 加上位置编码。ViT 使用可学习的 1D 位置编码（长度为 $N$ 的可训练向量），也有后续工作使用 2D sinusoidal 编码。

标准 Transformer 处理#

位置编码之后就是标准 Transformer：每个 block 包含 Multi-Head Self-Attention (MHSA) 和 Feed-Forward Network (FFN)，中间用 LayerNorm 和残差连接。

# ViT 前向过程伪代码
x = image                           # [B, 3, 224, 224]
patches = patchify(x, P=16)         # [B, 196, 768]
tokens = linear_proj(patches)       # [B, 196, D]
tokens = tokens + pos_embed         # [B, 196, D]  加位置编码
tokens = prepend_cls(tokens)        # [B, 197, D]  加 [CLS] token

for block in transformer_blocks:
    tokens = block(tokens)           # MHSA + FFN

cls_out = tokens[:, 0]               # [B, D]  取 CLS 分类
logits = classifier_head(cls_out)    # [B, num_classes]

ViT 的 Scaling Law#

模型越大，准确率持续攀升，没有明显的收益天花板。这正是 DiT 后来敢用 Transformer 做扩散模型的信心来源。

ViT 的计算复杂度

Self-Attention 的复杂度为 $O(N^2 \cdot D)$，其中 $N$ 是 token 数量，$D$ 是 hidden dim。当 $N=196$（patch_size=16）时，$N^2 \approx 38000$，计算量可控。

但如果 patch_size=4（$N=3136$），$N^2 \approx 983$ 万，计算量急剧上升。这也是后来 DiT 选择在 latent space（分辨率更低）操作而非 pixel space 的重要原因——有效降低了 token 数量。

**面试提示：**面试官可能问"为什么不在 pixel space 直接用 ViT 做扩散？"答：pixel space 分辨率太高（256×256，patch_size=2 → 16384 个 token），self-attention 的 $O(N^2)$ 复杂度不可承受。在 latent space 操作（32×32）可以把 token 数降到 256，是计算效率和生成质量的平衡。

3. DiT (Diffusion Transformer) 架构#

论文：Peebles & Xie, "Scalable Diffusion Models with Transformers" (ICCV 2023, Meta)

核心思想：用 Transformer 替换 U-Net 作为扩散模型的去噪 backbone。在 latent space 中操作（使用预训练的 VAE encoder/decoder）。

3.1 输入处理流程#

DiT 工作在 latent space，输入是 VAE 编码后的隐变量 $z_t$：

1. VAE 编码：原始图像 $256\times256\times3$ → VAE encoder → latent $32\times32\times4$
2. Patchify：patch_size $p=2$，则 $\frac{32}{2}\times\frac{32}{2}=16\times16=256$ 个 token
3. 展平 + 投影：每个 patch $2\times2\times4=16$ 维 → 线性投影到 hidden_dim $D$
4. 加位置编码：可学习的位置 embedding
5. 通过 $N$ 个 DiT Block：包含 MHSA + FFN + 条件注入
6. 输出投影：线性层投回 $2\times2\times4\times2=32$ 维（$\times2$ 因为预测均值和方差）
7. Unpatchify：重组为 $32\times32\times4$ 的噪声预测

$$z_t \in \mathbb{R}^{32\times32\times4} \xrightarrow{\text{patchify}} \mathbb{R}^{256\times16} \xrightarrow{\text{proj}} \mathbb{R}^{256\times D} \xrightarrow{N \text{ blocks}} \mathbb{R}^{256\times D} \xrightarrow{\text{linear}} \mathbb{R}^{256\times32} \xrightarrow{\text{unpatchify}} \mathbb{R}^{32\times32\times8}$$

完整的 DiT 数据流（最后 8 = 4 通道均值 + 4 通道方差）

# DiT 前向过程伪代码
z_t = vae.encode(image) + noise      # [B, 4, 32, 32]
x = patchify(z_t, p=2)               # [B, 256, 16]
x = linear_proj(x)                   # [B, 256, D]
x = x + pos_embed                    # [B, 256, D]

t_emb = timestep_mlp(t)              # [B, D]  时间步嵌入
c_emb = label_embed(class_label)     # [B, D]  类别嵌入
cond = t_emb + c_emb                 # [B, D]  融合条件

for block in dit_blocks:
    x = block(x, cond)               # adaLN-Zero block

x = final_layer(x, cond)             # [B, 256, 2*p*p*C]
noise_pred = unpatchify(x)            # [B, 8, 32, 32]

3.2 DiT 模型规模#

对比：SD 1.5 的 U-Net 约 860M 参数 / 约 150 Gflops。DiT-XL/2 用更少参数和算力达到了更好的 FID。

DiT 的训练细节

训练数据：ImageNet 256×256，class-conditional 生成（1000 个类别）

扩散框架：标准 DDPM，$T=1000$ 步，线性噪声调度

VAE：使用 Stable Diffusion 预训练的 KL-VAE（$f=8$ 空间压缩）

优化器：AdamW，lr=$1\times10^{-4}$，无 weight decay

训练量：DiT-XL/2 训练了 7M 步（约 5.7 epoch），在 8×A100 上约 5 天

**关于 classifier-free guidance：**DiT 在推理时使用 CFG（guidance scale=1.5），将类别条件以 10% 概率随机替换为空标签来训练。这和 SD 系列的做法一致。CFG 使 FID 从 9.62 降到 2.27（DiT-XL/2）。

4. DiT Block 的内部结构（adaLN-Zero）#

4.1 四种条件注入方式对比#

DiT 论文系统比较了 4 种将条件（时间步 $t$、类别标签 $c$）注入 Transformer block 的方式：

为什么 adaLN-Zero 效果最好？

零初始化策略：缩放因子 $\alpha$ 初始化为 0，所以训练开始时每个残差分支的输出被乘以 0，整个 DiT block 退化为恒等映射 $f(x) = x$。

直觉：训练初期网络"什么都不做"，是一个安全的起点。然后随着训练推进，$\alpha$ 逐渐从 0 变大，网络慢慢学到有意义的变换。这类似 ResNet 中的 zero-init trick（何恺明等人在 ResNet 论文中的实践）。

**$\alpha = 0$ 是怎么实现的？**追溯到代码层面：adaLN_mlp 的最后一层 Linear 的权重和偏置全部初始化为 0。由于线性层计算为 $\text{output} = \text{input} \cdot W^T + b$，当 $W=0, b=0$ 时，无论输入 cond 是什么，MLP 输出恒为零向量。chunk 之后得到的 $\gamma_1, \beta_1, \alpha_1, \gamma_2, \beta_2, \alpha_2$ 全部为 0。代入公式：$h = (1+0)\odot\text{LN}(x)+0 = \text{LN}(x)$（特征正常保留），$x \leftarrow x + 0 \odot \text{SelfAttn}(h) = x$（残差分支被关闭）。整个 block 输出等于输入，即恒等映射。

4.2 adaLN-Zero 的具体计算流程#

每个 DiT block 内部包含一个 Self-Attention 子层和一个 FFN 子层，条件信息通过 MLP 生成 6 个调制向量：

$$\text{MLP}(\text{cond}) \rightarrow (\gamma_1, \beta_1, \alpha_1, \gamma_2, \beta_2, \alpha_2) \in \mathbb{R}^{6 \times D}$$

一个 MLP 同时输出 6 个 D 维向量，分别用于两个子层的调制

Self-Attention 子层#

$$h_1 = (1 + \gamma_1) \odot \text{LayerNorm}(x) + \beta_1$$

$$x \leftarrow x + \alpha_1 \odot \text{SelfAttn}(h_1)$$

FFN 子层#

$$h_2 = (1 + \gamma_2) \odot \text{LayerNorm}(x) + \beta_2$$

$$x \leftarrow x + \alpha_2 \odot \text{FFN}(h_2)$$

其中 $\odot$ 表示逐元素乘法。$\gamma$ 和 $\beta$ 调制 LayerNorm 的输出（类似 BatchNorm 中的 affine 参数，但由条件动态生成），$\alpha$ 缩放残差分支。

注意公式中的 $1+\gamma$：$\gamma$ 表示在"正常归一化"基础上的额外调制，初始值为 0 时代表"不做任何额外调制"（输出就是正常的 LayerNorm(x)）。若写成 $\gamma \odot \text{LayerNorm}(x)$，初始 $\gamma=0$ 会把特征直接清零，Attention 将没有有效输入。

# adaLN-Zero DiT Block 伪代码
class DiTBlock(nn.Module):
    def __init__(self, D, num_heads):
        self.norm1 = nn.LayerNorm(D, elementwise_affine=False)
        self.attn  = MultiHeadSelfAttention(D, num_heads)
        self.norm2 = nn.LayerNorm(D, elementwise_affine=False)
        self.ffn   = FeedForward(D, 4*D)
        self.adaLN_mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(D, 6 * D)       # 输出 6 个调制向量
        )
        # 关键：最后一层 linear 权重和偏置初始化为 0
        nn.init.zeros_(self.adaLN_mlp[-1].weight)
        nn.init.zeros_(self.adaLN_mlp[-1].bias)

    def forward(self, x, cond):
        # cond: [B, D]  (时间步+类别的嵌入)
        g1, b1, a1, g2, b2, a2 = self.adaLN_mlp(cond).chunk(6, dim=-1)

        # Self-Attention 子层
        h = (1 + g1.unsqueeze(1)) * self.norm1(x) + b1.unsqueeze(1)
        x = x + a1.unsqueeze(1) * self.attn(h)

        # FFN 子层
        h = (1 + g2.unsqueeze(1)) * self.norm2(x) + b2.unsqueeze(1)
        x = x + a2.unsqueeze(1) * self.ffn(h)
        return x

维度具体例子（DiT-XL/2）

$D=1152$，$N=28$ 层，每层输出 $6 \times 1152 = 6912$ 维调制参数。条件嵌入 $\text{cond} \in \mathbb{R}^{B \times 1152}$，序列长度 256 个 token。

adaLN 的 MLP 只有一层 Linear（$1152 \rightarrow 6912$），参数量约 $1152 \times 6912 \approx 8\text{M}$（每层），28 层共 $\sim$224M 调制参数。

5. DiT 为什么更容易 Scale#

5.1 Scaling Law 实验结果#

DiT 论文在 ImageNet 256×256 class-conditional 生成上做了系统的 scaling 实验：

*ADM 使用了 classifier guidance + 像素空间（Gflops 高出一个量级）

DiT-XL/2 用 119 Gflops 达到 FID 2.27，而 ADM 需要 1120 Gflops 才有可比性能。在相同算力预算下，DiT 始终优于 U-Net。

5.2 为什么 Transformer 比 U-Net 更好 Scale？#

原因 1：结构均匀

每个 DiT block 结构完全相同（MHSA + FFN + adaLN），想加深就复制一个 block。不存在 U-Net 中 encoder/decoder 不对称、skip connection 对齐的问题。

原因 2：计算密度高

Self-Attention 和 FFN 的核心操作都是密集矩阵乘法（GEMM），GPU 利用率极高。U-Net 中有大量 3×3 卷积（计算密度不如 GEMM）和 skip connection 的拼接操作（纯内存搬运，不贡献计算）。

原因 3：没有 skip connection 瓶颈

U-Net 的 skip 连接强制 encoder 第 $i$ 层和 decoder 第 $i$ 层通道数对齐，想单独扩大某部分很困难。DiT 没有这种硬性约束。

原因 4：和 LLM 的经验互通

Transformer 的训练超参数（学习率调度、warmup、weight decay、dropout、梯度裁剪）在 LLM 领域经过大量实验验证，可以直接复用到 DiT。U-Net 的训练经验相对独立，不能借鉴 NLP 社区的海量工程积累。

5.3 Patch Size 的影响#

DiT 的 patch_size 直接控制 token 数量（即序列长度），进而影响计算量和质量：

**面试关键点：**DiT-XL/2 的 "/2" 就是 patch_size=2 的意思。patch 越小 → token 越多 → self-attention 的计算量 $O(N^2)$ 增长很快 → 通常 $p=2$ 是质量和效率的最佳平衡点。

5.4 Gflops vs FID 的关系#

两个指标的含义

GFlops（Giga Floating Point Operations）：十亿次浮点运算，衡量模型跑一次前向传播需要多少计算量。模型越大、分辨率越高 → GFlops 越大。注意这是"工作量"而非"速度"，和硬件无关。

FID（Fréchet Inception Distance）：衡量生成图像质量的核心指标，越低越好。用 InceptionV3 分别提取真实图像和生成图像的特征分布（拟合为高斯），计算两个分布之间的距离。FID=0 表示生成分布与真实分布完全一致；典型参考值：DiT-XL/2 为 2.27，SD 1.4 约为 14，早期 GAN 通常在 30 以上。

DiT 论文中最有说服力的实验之一是：在相同 Gflops 预算下，DiT 总是比 U-Net 获得更低的 FID。即花同样的计算量，DiT 产出的图像质量更好。

具体来说，DiT-L/2（80 Gflops，FID=9.62）已经超过了 LDM-4 U-Net（104 Gflops，FID=10.56）。而 DiT-XL/2 仅用 119 Gflops 就达到了 FID 2.27。相比之下，像素空间的 ADM（U-Net）需要 1120 Gflops 才有可比性能——几乎 10× 的计算开销。

6. 理解检查（8 道）#

面试题 01 DiT 和 U-Net 的核心区别是什么？

结构层面：U-Net 是 encoder-decoder + skip connection 的不规则结构；DiT 是堆叠相同 Transformer block 的均匀结构。U-Net 输入输出保持分辨率，中间有下/上采样；DiT 通过 patchify 将空间展平为序列，全程在同一分辨率操作。

条件注入：U-Net 通常用 cross-attention 注入文本条件，时间步通过 adaptive group norm 注入。DiT 使用 adaLN-Zero 统一注入时间步和类别条件。

扩展性：DiT 更容易 scale——加层就是复制 block，有明确的 scaling law（FID 随参数量单调下降）。U-Net 扩展受限于 skip connection 的通道对齐约束。

性能：DiT-XL/2（675M，119 Gflops）在 ImageNet 256 上 FID=2.27，效率优于同量级 U-Net。

面试题 02 adaLN-Zero 是什么？为什么比 cross-attention 好？

机制：adaLN-Zero 是 Adaptive Layer Normalization 的变体。条件嵌入（时间步+类别）通过 MLP 生成 $(\gamma, \beta, \alpha)$ 三组向量：$\gamma, \beta$ 调制 LayerNorm 的 scale/shift，$\alpha$ 缩放残差分支的输出。

"Zero" 的含义：MLP 最后一层权重和偏置初始化为 0 → 训练开始时 $\alpha=0$ → 残差分支输出为 0 → 整个 block 是恒等映射。这提供了安全的初始化起点。

vs Cross-Attention：Cross-attention 需要额外的 K/V 投影矩阵（参数更多），且信息流是单向的（条件 → 主序列）。adaLN-Zero 参数更少但效果更好（FID 2.27 vs 5.75），因为它直接调制每层的统计特性，是一种更"全局"的条件注入方式。

面试题 03 DiT 为什么比 U-Net 更容易 scale？

四个原因：(1) 结构均匀——每个 block 完全相同，加深 = 复制 block，不受 skip connection 对齐约束。(2) 计算密度高——核心是矩阵乘法（GEMM），GPU 利用率高于 3×3 卷积。(3) 成熟的 scaling law——LLM 和 ViT 已经验证 Transformer 参数越大效果越好。(4) 工程经验可复用——FlashAttention、模型并行、学习率调度等 LLM 训练技巧可以直接迁移。

数据支持：DiT 论文中，从 DiT-S 到 DiT-XL，FID 从 68.4 单调下降到 2.27，没有出现 U-Net 常见的 scaling 收益递减现象。

面试题 04 视频生成和图像生成的核心差异？

维度：图像是 2D ($H \times W$)，视频是 3D ($T \times H \times W$)。额外的时间维度带来的不只是计算量线性增长——self-attention 的复杂度是 $O(N^2)$，token 数增加 $T$ 倍意味着计算量增加 $T^2$ 倍（full attention 下）。

核心挑战：(1) 时间一致性——帧间必须连贯，人眼对闪烁极度敏感。(2) 运动建模——需要理解物理世界的运动规律。(3) 数据稀缺——高质量视频-文本对远少于图像-文本对。

技术方案：3D VAE 做时空压缩，Temporal Attention 建模帧间关系，Factorized Attention 降低计算量。

面试题 05 视频生成中 Full 3D Attention 和 Factorized Attention 的区别？现在主流用哪个？

Factorized Attention：将 3D attention 分解为两步——Spatial Attention（每帧独立，帧内空间关系）和 Temporal Attention（每个空间位置独立，跨帧时间关系），交替进行。计算量从 $O((TH'W')^2)$ 降到 $O(T(H'W')^2 + H'W'T^2)$，节省倍数取决于 T 与每帧 token 数。代价是帧间信息需通过多层堆叠间接传播，时空一致性略差。早期 Open-Sora（STDiT）采用此方案。

Full 3D Attention：所有帧所有位置的 token 直接做 self-attention，时空建模最完整，视频一致性更好。早期因计算量大（无压缩下 token 数达数千甚至上万）被认为不可行，但现代 3D VAE（时间 4×、空间 8×，共 256× 压缩）已将 token 数压缩到可承受范围，配合 FlashAttention 完全可行。

当前主流是 Full 3D Attention：HunyuanVideo、CogVideoX、Wan2.1 均采用，并在论文中明确指出其时空一致性优于 Factorized 方案。Factorized 是 3D VAE 不够强时的历史妥协，头部模型已基本放弃。

面试题 06 3D VAE 和 2D VAE 有什么区别？

2D VAE（如 SD 的 KL-VAE）：使用 2D 卷积，只压缩空间（$H, W$ 各缩小 8×），不处理时间。对视频逐帧独立编码 → 帧间无法共享信息 → 时间一致性差。

3D VAE：使用 3D 卷积，同时压缩空间和时间。典型压缩率：空间 8×、时间 4× → 总计 256× 压缩。3D 卷积让相邻帧的 latent 互相关联，天然保持时间连续性。

Causal 3D VAE：时间方向只看过去帧不看未来帧（因果性），兼容自回归式长视频生成。实现方式是在时间维度使用单向 padding。

面试题 07 Sora 的 spacetime patches 是什么？

概念：直接对视频 latent 做 3D patchify，而非先空间 patch 再单独处理时间。用 3D patch size $(p_t, p_h, p_w)$ 将视频切成时空小块，每个块展平为一个 token。

优势：每个 token 天然包含空间和时间信息，不需要额外的 temporal attention 模块来建立时间关系——通过 self-attention 自动实现。同时支持可变分辨率和时长（token 数随输入变化，Transformer 天然支持可变长序列）。

示例：$16 \times 32 \times 32 \times 4$ 的 latent，用 $(2,2,2)$ 的 3D patch → $8 \times 16 \times 16 = 2048$ 个 token。

面试题 08 SD3 和 SD 1.x 的架构区别？

Backbone：SD 1.x 用 U-Net（~860M），SD3 用 MM-DiT (Multi-Modal DiT)（~2B/8B）。

条件注入：SD 1.x 通过 Cross-Attention 注入文本条件（单向信息流）；SD3 将图文 token 拼接后做 Joint Self-Attention（双向信息流）。

文本编码器：SD 1.x 只用 CLIP ViT-L；SD3 用 CLIP ViT-L + OpenCLIP ViT-bigG + T5-XXL 三个编码器，文本理解能力大幅提升。

采样方法：SD 1.x 用 DDPM/DDIM；SD3 用 Flow Matching（更直的采样轨迹，效率更高）。

效果：SD3 在文字渲染、构图理解、细节质量上全面超越 SD 1.x/2.x/XL。

7. 进阶讨论：DiT 可扩展性细节 + PyTorch 代码#

7.1 DiT 可扩展性的五个具体维度#

面试官问"DiT 和 UNet 相比，可扩展性到底体现在哪"时，需要能展开到具体维度，不能只说"加层就行"：

维度 1：深度扩展（加层）

DiT：每个 block 结构完全相同（MHSA + FFN + adaLN-Zero），加一个 block 只需要复制一份参数，输入输出维度不变（都是 $\mathbb{R}^{N \times D}$）。从 DiT-S（12层）到 DiT-XL（28层），只改一个超参数。

U-Net：encoder-decoder 必须对称，加一层 encoder 就必须加一层 decoder + 一条 skip connection。通道数要对齐（encoder 第 $i$ 层的输出通道 = decoder 第 $i$ 层的输入通道之一），改一层牵连整个结构。

# DiT 加深：改一个数字
dit_s  = DiT(depth=12, hidden_dim=384)   # DiT-S: 33M
dit_b  = DiT(depth=12, hidden_dim=768)   # DiT-B: 130M
dit_l  = DiT(depth=24, hidden_dim=1024)  # DiT-L: 458M
dit_xl = DiT(depth=28, hidden_dim=1152)  # DiT-XL: 675M

# U-Net 加深：要改一整套配置
unet = UNet(
    ch_mult=(1, 2, 4, 4),          # 每级通道倍数——改这里，decoder 也要跟着改
    num_res_blocks=2,               # 每级 ResBlock 数
    attention_resolutions=(32, 16, 8),  # 哪些分辨率加 attention
    channel_mult_attn=(1, 2, 4),    # attention 头数……
)  # 改任何一个都可能导致维度不匹配 crash

维度 2：宽度扩展（加 hidden_dim）

DiT：所有层共享同一个 $D$（hidden_dim），改大就是改一个数。attention head 数通常按 $D / 64$ 自动调整。

U-Net：各层通道数不同（如 320→640→1280→1280），改一层会通过 skip connection 传播到 decoder 对应层，进而影响后续所有层。手动调试工作量大。

维度 3：数据扩展（更大数据集）

Scaling Law 的存在：DiT 论文验证了在 ImageNet 上从 DiT-S 到 DiT-XL，FID 从 68.4 → 9.62 → 2.27，单调下降且无天花板。而 U-Net 在 ADM（554M）之后继续加大参数收益递减明显——SD-XL 的 2.6B U-Net 几乎是人工设计的极限。

和 LLM 的类比：GPT 系列证明 Transformer 的 loss 与参数量 $N$、数据量 $D$、计算量 $C$ 存在幂律关系（$L \propto N^{-0.076}$）。DiT 继承了同样的 scaling 性质，U-Net 没有这种可预测的 scaling law。

维度 4：硬件扩展（多卡并行）

DiT：Transformer 的并行策略在 LLM 领域已经非常成熟——张量并行（TP）、流水线并行（PP）、序列并行（SP）都可以直接复用。每个 block 结构相同，负载天然均衡。

U-Net：encoder 和 decoder 计算量不对称（encoder 分辨率逐渐降低，计算量递减；decoder 相反），skip connection 跨层传递数据导致流水线并行难以均匀切分。3×3 卷积的计算密度不如矩阵乘法，GPU 利用率更低。

维度 5：分辨率扩展（更大图像/视频）

DiT：分辨率变大 → token 数变多 → 序列更长，Transformer 天然支持可变长序列。配合 FlashAttention 和 RoPE/2D sinusoidal 位置编码，可以在不同分辨率间灵活迁移。Sora 就是利用这个特性处理不同宽高比和时长的视频。

U-Net：分辨率变化需要调整下采样/上采样层数，增加新的尺度级需要重新设计 skip connection 的连接方式。SD-XL 从 512→1024 分辨率就需要大量架构调整。

7.2 解释要点#

面试题 和 U-Net 相比，DiT 的可扩展性到底是些什么？

30 秒版本：DiT 可扩展性体现在五个维度——深度（加层=复制 block，不受 skip 约束）、宽度（改一个 hidden_dim）、数据（有明确 scaling law，FID 随参数单调下降）、硬件（LLM 的 TP/PP/SP 并行策略直接复用）、分辨率（可变长序列天然支持多分辨率）。

关键数据：DiT-XL/2（675M, 119 Gflops）FID=2.27，ADM U-Net（554M, 1120 Gflops）FID=2.07——DiT 用 1/10 算力达到同等质量。从 DiT-S 到 DiT-XL，FID 68.4→2.27，单调下降无天花板。

根本原因：U-Net 是手工设计的不规则结构，每次 scale 都是"重新设计"；DiT 是标准化的均匀结构，scale 只是"调参数"。这也是为什么 2024 年后所有头部模型（SD3、Flux、Sora、HunyuanVideo）都转向 DiT。

7.3 PyTorch 代码：完整的 DiT 图像训练流程#

面试可能要求手写训练循环。以下是一个完整的 DiT 训练流程（DDPM 方式），包含 VAE 编码、加噪、DiT 前向、损失计算：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# ====== 1. 超参数 ======
image_size   = 256
latent_size  = 32          # VAE 把 256→32
latent_ch    = 4           # VAE latent 通道数
patch_size   = 2           # DiT patch size
num_tokens   = (latent_size // patch_size) ** 2  # 16×16 = 256 tokens
hidden_dim   = 1152        # DiT-XL
depth        = 28
num_classes  = 1000        # ImageNet
T            = 1000        # 扩散步数
lr           = 1e-4
batch_size   = 256

# ====== 2. 噪声调度表（预计算） ======
betas = torch.linspace(1e-4, 0.02, T)
alphas = 1.0 - betas
alpha_bar = torch.cumprod(alphas, dim=0)       # ᾱ_t
sqrt_alpha_bar = torch.sqrt(alpha_bar)
sqrt_one_minus_alpha_bar = torch.sqrt(1 - alpha_bar)

# ====== 3. 数据加载 ======
transform = transforms.Compose([
    transforms.Resize(image_size),
    transforms.CenterCrop(image_size),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.5]*3, [0.5]*3),  # [-1, 1]
])
dataset = datasets.ImageFolder('imagenet/train', transform=transform)
loader  = DataLoader(dataset, batch_size=batch_size, shuffle=True,
                     num_workers=8, pin_memory=True, drop_last=True)

# ====== 4. 模型 ======
vae = load_pretrained_vae()           # 预训练 VAE（冻结）
vae.eval()
dit = DiT_XL_2(num_classes=num_classes)  # DiT-XL/2
optimizer = torch.optim.AdamW(dit.parameters(), lr=lr, weight_decay=0.0)
ema = EMA(dit, decay=0.9999)          # EMA 平滑

# ====== 5. 训练循环 ======
dit.train()
for epoch in range(num_epochs):
    for images, labels in loader:
        images = images.cuda()         # [B, 3, 256, 256]
        labels = labels.cuda()         # [B]  类别标签

        # Step 1: VAE 编码到 latent space（不计算梯度）
        with torch.no_grad():
            z = vae.encode(images)     # [B, 4, 32, 32]
            z = z * 0.18215            # SD 的 scaling factor

        # Step 2: 随机采样时间步 t 和噪声 ε
        t = torch.randint(0, T, (batch_size,), device='cuda')  # [B]
        noise = torch.randn_like(z)                             # [B, 4, 32, 32]

        # Step 3: 前向加噪 x_t = √ᾱ_t · z + √(1-ᾱ_t) · ε
        a = sqrt_alpha_bar[t].view(-1, 1, 1, 1).cuda()
        b = sqrt_one_minus_alpha_bar[t].view(-1, 1, 1, 1).cuda()
        z_t = a * z + b * noise        # [B, 4, 32, 32]

        # Step 4: DiT 前向（预测噪声和方差）
        noise_pred, var_pred = dit(z_t, t, labels)  # 各 [B, 4, 32, 32]

        # Step 5: 简单 MSE 损失（只对 ε 做回归）
        loss = F.mse_loss(noise_pred, noise)

        # Step 6: 反向传播
        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(dit.parameters(), max_norm=1.0)
        optimizer.step()
        ema.update()

**面试关键点：**训练流程和标准 DDPM 完全一致——区别只是把 U-Net 换成 DiT，在 latent space 操作。VAE 是冻结的预训练模型。EMA 对 DiT 非常重要（论文消融实验显示 EMA 可将 FID 从 9.62 降到 2.27）。

7.4 PyTorch 代码：DiT Block 的完整实现#

class DiTBlock(nn.Module):
    """DiT Block with adaLN-Zero conditioning"""
    def __init__(self, D, num_heads, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(D, elementwise_affine=False)
        self.attn  = nn.MultiheadAttention(D, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(D, elementwise_affine=False)
        self.mlp   = nn.Sequential(
            nn.Linear(D, int(D * mlp_ratio)),
            nn.GELU(approximate='tanh'),
            nn.Linear(int(D * mlp_ratio), D),
        )
        # adaLN-Zero: 条件 → 6 组调制参数
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(D, 6 * D),   # γ1, β1, α1, γ2, β2, α2
        )
        # Zero 初始化：训练开始时 block 是恒等映射
        nn.init.zeros_(self.adaLN_modulation[-1].weight)
        nn.init.zeros_(self.adaLN_modulation[-1].bias)

    def forward(self, x, cond):
        # cond: [B, D] — 时间步+类别嵌入
        gamma1, beta1, alpha1, gamma2, beta2, alpha2 = \
            self.adaLN_modulation(cond).chunk(6, dim=-1)  # 各 [B, D]

        # Attention 分支 + adaLN 调制
        h = self.norm1(x)
        h = h * (1 + gamma1.unsqueeze(1)) + beta1.unsqueeze(1)  # scale & shift
        h, _ = self.attn(h, h, h)
        x = x + alpha1.unsqueeze(1) * h  # α 缩放残差

        # FFN 分支 + adaLN 调制
        h = self.norm2(x)
        h = h * (1 + gamma2.unsqueeze(1)) + beta2.unsqueeze(1)
        h = self.mlp(h)
        x = x + alpha2.unsqueeze(1) * h

        return x

class DiT(nn.Module):
    """Simplified DiT for ImageNet class-conditional generation"""
    def __init__(self, latent_size=32, patch_size=2, in_ch=4,
                 hidden_dim=1152, depth=28, num_heads=16, num_classes=1000):
        super().__init__()
        self.patch_size = patch_size
        num_patches = (latent_size // patch_size) ** 2  # 256

        # Patchify: 2D conv 当线性投影
        self.proj = nn.Conv2d(in_ch, hidden_dim,
                              kernel_size=patch_size, stride=patch_size)
        self.pos_embed = nn.Parameter(
            torch.zeros(1, num_patches, hidden_dim))

        # 条件嵌入
        self.t_embed = nn.Sequential(          # 时间步
            nn.Linear(256, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )
        self.y_embed = nn.Embedding(num_classes, hidden_dim)  # 类别

        # N 个相同的 DiT Block
        self.blocks = nn.ModuleList([
            DiTBlock(hidden_dim, num_heads) for _ in range(depth)
        ])

        # 输出：预测 ε 和 Σ（各 patch_size² × in_ch）
        self.final_norm = nn.LayerNorm(hidden_dim, elementwise_affine=False)
        self.final_proj = nn.Linear(hidden_dim, patch_size**2 * in_ch * 2)

    def forward(self, z_t, t, y):
        # z_t: [B, 4, 32, 32], t: [B], y: [B]
        x = self.proj(z_t)              # [B, D, 16, 16]
        x = x.flatten(2).transpose(1,2) # [B, 256, D]
        x = x + self.pos_embed          # 加位置编码

        # 条件融合：时间步 + 类别 → 单个条件向量
        cond = self.t_embed(sinusoidal_encoding(t, 256)) + self.y_embed(y)

        # 通过所有 DiT blocks
        for block in self.blocks:
            x = block(x, cond)          # [B, 256, D]

        # 输出投影 + unpatchify
        x = self.final_norm(x)
        x = self.final_proj(x)          # [B, 256, p²·C·2]
        # Unpatchify: reshape back to [B, 8, 32, 32] → split → ε + Σ
        noise_pred, var_pred = x.chunk(2, dim=-1)
        return unpatchify(noise_pred), unpatchify(var_pred)

**面试提示：**DiT 的核心实现点：(1) Patchify 用 stride=p 的 Conv2d 实现 (2) adaLN-Zero 的 6 参数调制 (3) zero 初始化确保初始恒等映射 (4) 条件是 t_embed + y_embed 的加法融合。能手写 DiTBlock 就够了。

DiT 架构与视频生成

多模态算法面试准备系列

原始资料：DiT、SD3。文中的算力比较应区分 attention pair 数、FLOPs、NFE 与实际延迟；训练和 benchmark 结果仅对应各自实验配置。