DiT (Diffusion Transformer) 骨干网络详解

1. 为什么从 U-Net 换到 Transformer？

图像生成模型（扩散模型、Flow Matching）长期使用 U-Net 作为骨干。U-Net 的核心是 CNN 的编解码结构，靠卷积的局部感受野和下采样来捕获多尺度信息——这很自然，因为图片本来就是二维网格。

但 2023 年 DiT 论文提出了一个问题：如果生成模型的核心不是"去噪"，而是"理解全局结构"呢？

Transformer 的优势恰好在于：

• 全局感受野：Self-Attention 天然看到所有位置，不需要像 CNN 那样堆很多层才能扩大感受野
• 灵活的条件注入：adaLN 让条件信号精细化控制每一层的归一化行为，而不是简单 concat 到特征图上
• 可扩展性强：Transformer 的 scaling law 已经被验证，加大模型就能稳定提效

DiT 的做法很 ViT：把图片打成 patch，当 token 序列处理，条件（时间步、类别）通过 adaLN-Zero 注入每个 Transformer 块。

---

2. DiT 整体架构

2.1 数据流：图片 → Patch → 序列 → Patch → 图片

  (B, 1, 28, 28)          输入灰度图
        │
        │  Conv2d(1→256, kernel=4, stride=4)    ← Patch Embedding
        ▼
  (B, 256, 7, 7)          49 个 grid 特征
        │  flatten + transpose
        ▼
  (B, 49, 256)            patch token 序列
        │  + pos_embed (可学习, 正态初始化)
        ▼
  (B, 49, 256)            带位置编码的 token
        │
        │  ┌──────────────────────────────────┐
        ├──┤   DiTBlock × 8                   │
        │  │   • Multi-Head Self-Attention    │
        │  │   • Pointwise MLP                │
        │  │   • adaLN-Zero 条件调制          │ ← c = time(t) + label(y)
        │  └──────────────────────────────────┘
        ▼
  (B, 49, 256)
        │  LayerNorm → Linear(256 → 16)
        ▼
  (B, 49, 16)             每个 token → patch_size² × C 个像素值
        │  unpatchify (reshape + permute)
        ▼
  (B, 1, 28, 28)          输出：预测的向量场 v_θ

关键设计选择：

决策	选择	原因
Patch 大小	4×4	28/4=7，49 个 token；太小 token 太多，太大丢失细节
Token 维度	256	足够表达力，attention 计算量可控
Transformer 层数	8	太深过拟合 MNIST，太浅欠拟合
注意力头数	4	256/4=64 维/头，合理的 head dim
MLP 扩展比	4×	Transformer 标配

---

3. 核心组件拆解

3.1 Patch Embedding

ViT 的做法是用 Conv2d(kernel=patch_size, stride=patch_size) 一步完成"切块 + 投影"。这比手动切成 (B, N, P*P*C) 再 Linear 更高效：

self.patch_embed = nn.Conv2d(
    in_channels=1,          # MNIST 灰度图
    out_channels=256,       # hidden_dim
    kernel_size=4,          # patch_size
    stride=4,               # 无重叠
)
# (B, 1, 28, 28) → (B, 256, 7, 7)
# flatten(2) → (B, 256, 49) → transpose → (B, 49, 256)

49 个 token，每个 256 维，序列长度相当于一个短句子的 token 数，attention 矩阵只有 49×49，非常快。

3.2 条件嵌入：时间 + 标签 → 统一条件向量

DiT 需要两个条件信号：时间步 t（决定当前去噪/流形阶段）和类别标签 y（控制生成哪个数字）。

# 时间嵌入
t_emb = SinusoidalEmbedding(t, dim=256)  # 正弦编码 (B, 256)
t_emb = MLP(256 → 1024 → 256)(t_emb)     # MLP 变换 (B, 256)

# 标签嵌入
y_emb = nn.Embedding(10, 256)(labels)     # (B, 256)

# 融合
c = t_emb + y_emb                          # (B, 256)

为什么是相加而不是 concat？ concat 会让维度翻倍，后续 adaLN 的参数量跟着翻倍。相加是 DiT 原论文的做法，信息通过梯度自动解耦——时间和标签的 embedding 各自学会不冲突的表示。

正弦时间嵌入并非简单 t 标量送 MLP，而是用多频率的正余弦函数编码：

def timestep_embedding(t, dim, max_period=10000):
    half = dim // 2
    freqs = exp(-log(max_period) * arange(half) / half)  # 频率从低到高
    args = t[:, None] * freqs[None, :]                    # (B, half)
    return cat([cos(args), sin(args)], dim=-1)            # (B, dim)

低频分量捕获"当前处于生成过程的哪个大阶段"，高频分量捕获"精细的时间差别"。这比直接 Linear(t) 的信息量丰富得多。

3.3 adaLN-Zero：DiT 的灵魂

这是 DiT 与普通 Transformer 最大的区别。普通的条件注入（如 U-Net 把条件向量 concat 到特征图）是空间均匀的——每个位置被同样对待。adaLN 不是这样。

adaLN（adaptive LayerNorm）让条件向量 \(c\) 控制 LayerNorm 的 scale（缩放）和 shift（平移）：

\[\text{adaLN}(x, c) = \text{LayerNorm}(x) \cdot (1 + \gamma(c)) + \beta(c)\]

注意这里的 \(1 + \gamma\)：当 \(\gamma = 0\) 时退化为标准 LayerNorm，网络可以学一个"接近标准 LN"的偏移。

Zero 后缀的意思是：输出 \(\gamma, \beta\) 的 Linear 层权重和偏置全部初始化为 0。这就有了一个非常优雅的性质——训练开始时：

\[\text{gate} = 0 \quad \Rightarrow \quad x_{out} = x_{in} + 0 \cdot F(x) = x_{in}\]

每个 DiTBlock 在初始化时是一个恒等映射。 网络整体退化为：patch_embed → identity → unpatchify。从恒等映射开始，梯度通过条件嵌入反向传播，逐步"教会"每一层：在这个时间步、对这个类别，应该有多少注意力、做多少 MLP 变换。

这相当于一种极强的结构性正则化——网络不会一开始就乱改特征，而是"谨慎地"逐渐偏离恒等映射。实践中表现为训练极其稳定，不需要 warmup。

每个 DiTBlock 输出 6 组调制参数：

adaLN_modulation(c) → 6 × hidden_dim

  ┌─ scale_attn  (D,)    Attention 前的 LN 缩放
  ├─ shift_attn  (D,)    Attention 前的 LN 平移
  ├─ gate_attn   (D,)    Attention 残差连接的门控
  ├─ scale_mlp   (D,)    MLP 前的 LN 缩放
  ├─ shift_mlp   (D,)    MLP 前的 LN 平移
  └─ gate_mlp    (D,)    MLP 残差连接的门控

3.4 DiTBlock 完整前向过程

def forward(self, x, c):
    # c: (B, 256) —— 时间和标签的融合条件向量
    # x: (B, 49, 256) —— patch token 序列

    mod = self.adaLN_modulation(c)                # (B, 6*256)

    # 拆成 attention 和 MLP 各 3 组
    attn_mod, mlp_mod = mod.chunk(2, dim=-1)
    scale_a, shift_a, gate_a = attn_mod.chunk(3, dim=-1)  # 每组 (B, 256)
    scale_m, shift_m, gate_m = mlp_mod.chunk(3, dim=-1)

    # --- Attention 子层 ---
    # 1. 条件调制 LayerNorm（逐 token broadcast）
    normed = LayerNorm(x) * (1 + scale_a[:, None, :]) + shift_a[:, None, :]
    # 2. Multi-Head Self-Attention
    attn_out = self.attn(normed, normed, normed)[0]
    # 3. 门控残差
    x = x + gate_a[:, None, :] * attn_out

    # --- MLP 子层 ---
    # 1. 条件调制 LayerNorm
    normed = LayerNorm(x) * (1 + scale_m[:, None, :]) + shift_m[:, None, :]
    # 2. Pointwise FFN：256 → 1024 (GELU) → 256
    mlp_out = self.mlp(normed)
    # 3. 门控残差
    x = x + gate_m[:, None, :] * mlp_out

    return x

几个值得注意的细节：

1. [:, None, :] 的作用：调制参数是 (B, 256)，需要 unsqueeze 第 1 维变成 (B, 1, 256) 才能广播到 (B, 49, 256) 的每个 token。

1. LayerNorm 的 elementwise_affine=False：普通 LN 自带可学习的 \(\gamma, \beta\)，但这里 \(\gamma\) 和 \(\beta\) 由条件调制提供，LN 自己不需要再学一份——关掉避免冗余参数。

1. 1 + scale：scale 初始为 0（因为 adaLN_modulation 输出层零初始化），所以初始时 1 + 0 = 1，LayerNorm 行为完全标准。网络通过训练逐渐学到"偏离标准 LN 多少"。

3.5 多头自注意力

用的是 PyTorch 原生的 nn.MultiheadAttention，设 batch_first=True：

self.attn = nn.MultiheadAttention(
    embed_dim=256,      # hidden_dim
    num_heads=4,        # 每个头 64 维
    batch_first=True,   # (B, N, D) 而非 (N, B, D)
)

4 个注意力头，每头处理 64 维。49 个 token 两两互相看，attention 矩阵 \(49 \times 49\)，计算量可以忽略不计。这是 DiT 对比 U-Net 的核心优势——从第一层起每个 token 就能看到整个图片，不需要逐层扩大感受野。

3.6 Unpatchify：把 Token 序列拼回图片

最后一层 Linear 把每个 256 维 token 映射为 \(p^2 \times C = 16\) 个像素值：

self.final_linear = nn.Linear(256, 16)  # 16 = 4×4×1

然后通过纯 reshape + permute 操作把 49 个 token 拼成 28×28 的图像：

def unpatchify(self, x):
    # (B, 49, 16) → 拆成空间维度
    x = x.reshape(-1, 7, 7, 4, 4, 1)     # (B, h, w, p, p, c)
    # 交错排列 patch 内像素
    x = x.permute(0, 5, 1, 3, 2, 4)      # (B, c, h, p, w, p)
    # 合并为完整图像
    x = x.reshape(-1, 1, 28, 28)         # (B, c, H, W)
    return x

整个过程完全用 reshape + permute，不需要 einops 等第三方库。这种操作被称为 "fold"，是 patch embedding 的逆操作：

Patch Embedding (forward):      28×28 图片 → 7×7 × 4×4 patch → 49 tokens
Unpatchify     (reverse):      49 tokens → 4×4 × 7×7 grid → 28×28 图片

---

4. 与传统 U-Net 的条件注入对比

维度	U-Net (上一篇)	DiT (本文)
骨干类型	CNN 编解码	纯 Transformer
感受野	逐层扩大（局部→全局）	第一层起全局
条件注入方式	Concat 到特征图通道	adaLN-Zero 调制
条件粒度	粗粒度（整个特征图共享）	细粒度（逐通道调制、逐 token 应用）
多尺度处理	MaxPool + Upsample	无（固定分辨率）
跳跃连接	U-Net 标配 skip connection	无（残差流在 block 内部）
初始化	标准 xavier	零初始化 gate（恒等映射起点）
参数量	~1.35M	~10M

两者最关键的区别在条件注入：

U-Net 方式：
  条件 c: (B, 32, 1, 1)  →  expand 到 (B, 32, H, W) → concat([feat, c]) → Conv2d

DiT 方式：
  条件 c: (B, 256)  →  MLP → [scale, shift, gate] × 每层  →  精细调制 LayerNorm

U-Net 告诉每一层"这是什么条件和时间步"，方式是把条件向量硬拼到特征后面让卷积自己去琢磨。DiT 告诉每一层"在这个条件下，这个时间步，你应该这样归一化、你应该放这么多信息通过"，方式是直接逐通道调控归一化行为——信息量更大，控制力更强。

---

5. 代码架构

main.py
├── timestep_embedding()          正弦时间编码
├── DiTBlock                      adaLN-Zero Transformer 块 ×8
│   ├── Multi-Head Self-Attn      (4 heads, batch_first)
│   ├── Pointwise MLP             (256 → 1024 → 256, GELU)
│   └── adaLN_modulation          (256 → 6×256, 零初始化末层)
├── DiT                           主模型
│   ├── patch_embed               Conv2d(1→256, kernel=4, stride=4)
│   ├── pos_embed                 (1, 49, 256) 可学习参数
│   ├── time_embed                Sinusoidal → MLP → 256
│   ├── label_embed               Embedding(10, 256)
│   ├── 8 × DiTBlock              Transformer 骨干
│   └── norm_final + final_linear + unpatchify
├── generate()                    Euler ODE 采样 (100 steps)
├── sample_images()               保存网格图
└── __main__                      训练循环

整个文件不依赖 einops、timm、transformers 等第三方库，只有 torch + torchvision，可以直接复制运行。

---

6. 训练结果

MNIST 上训练 40 epoch（batch_size=256, lr=1e-4, AdamW, gradient clip=1.0）：

Epoch	Loss	生成质量
1	0.5477	模糊噪声
5	0.3078	数字轮廓初现
10	0.2653	形态可辨
20	0.2115	清晰稳定
30	0.1875	细节丰富
40	0.1777	生成质量良好

得益于 adaLN-Zero 的恒等映射初始化，训练从一开始就非常稳定——没有 loss 震荡、没有梯度爆炸。Loss 曲线平滑下降，不需要 warmup、不需要特殊的学习率调度。

---

附录：完整代码

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torchvision.utils import save_image
from torch.utils.data import DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size = 256
lr = 1e-4
epochs = 100
num_classes = 10


# ============================================================
# 正弦时间嵌入
# ============================================================
def timestep_embedding(t, dim, max_period=10000):
    """Sinusoidal timestep embedding (same as DiT/DDPM)."""
    half = dim // 2
    freqs = torch.exp(
        -math.log(max_period)
        * torch.arange(half, dtype=torch.float32, device=t.device)
        / half
    )
    args = t[:, None].float() * freqs[None, :]
    return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)


# ============================================================
# DiTBlock: adaLN-Zero Transformer 块
# ============================================================
class DiTBlock(nn.Module):
    """Transformer block with adaLN-Zero conditioning."""

    def __init__(self, hidden_dim, num_heads, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_dim, elementwise_affine=False)
        self.norm2 = nn.LayerNorm(hidden_dim, elementwise_affine=False)

        self.attn = nn.MultiheadAttention(
            hidden_dim, num_heads, batch_first=True
        )

        self.mlp = nn.Sequential(
            nn.Linear(hidden_dim, int(hidden_dim * mlp_ratio)),
            nn.GELU(approximate='tanh'),
            nn.Linear(int(hidden_dim * mlp_ratio), hidden_dim),
        )

        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_dim, 6 * hidden_dim),
        )
        # 零初始化末层 —— adaLN-Zero 的关键
        nn.init.zeros_(self.adaLN_modulation[-1].weight)
        nn.init.zeros_(self.adaLN_modulation[-1].bias)

    def forward(self, x, c):
        mod = self.adaLN_modulation(c)                     # (B, 6D)
        attn_mod, mlp_mod = mod.chunk(2, dim=-1)           # 各 (B, 3D)

        scale_attn, shift_attn, gate_attn = attn_mod.chunk(3, dim=-1)
        scale_mlp, shift_mlp, gate_mlp = mlp_mod.chunk(3, dim=-1)

        # Attention 子层
        x_norm = self.norm1(x) * (1 + scale_attn.unsqueeze(1)) + shift_attn.unsqueeze(1)
        x = x + gate_attn.unsqueeze(1) * self.attn(x_norm, x_norm, x_norm)[0]

        # MLP 子层
        x_norm = self.norm2(x) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
        x = x + gate_mlp.unsqueeze(1) * self.mlp(x_norm)

        return x


# ============================================================
# DiT: Diffusion Transformer 主模型
# ============================================================
class DiT(nn.Module):
    """Diffusion Transformer for flow matching on MNIST (28x28, 1 channel)."""

    def __init__(
        self,
        image_size=28,
        in_channels=1,
        patch_size=4,
        hidden_dim=256,
        depth=8,
        num_heads=4,
        mlp_ratio=4.0,
        num_classes=10,
    ):
        super().__init__()
        self.image_size = image_size
        self.in_channels = in_channels
        self.patch_size = patch_size
        self.hidden_dim = hidden_dim

        assert image_size % patch_size == 0
        self.num_patches = (image_size // patch_size) ** 2

        # Patch embedding
        self.patch_embed = nn.Conv2d(
            in_channels, hidden_dim,
            kernel_size=patch_size, stride=patch_size, bias=True,
        )

        # 可学习位置编码
        self.pos_embed = nn.Parameter(
            torch.randn(1, self.num_patches, hidden_dim) * 0.02
        )

        # 时间嵌入：正弦编码 → MLP
        self.time_embed = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 4),
            nn.SiLU(),
            nn.Linear(hidden_dim * 4, hidden_dim),
        )

        # 标签嵌入
        self.label_embed = nn.Embedding(num_classes, hidden_dim)

        # Transformer 块
        self.blocks = nn.ModuleList([
            DiTBlock(hidden_dim, num_heads, mlp_ratio)
            for _ in range(depth)
        ])

        # 最终输出层
        self.norm_final = nn.LayerNorm(hidden_dim, elementwise_affine=False)
        self.final_linear = nn.Linear(
            hidden_dim, patch_size * patch_size * in_channels,
        )

        self._init_weights()

    def _init_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Conv2d):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, std=0.02)
        nn.init.normal_(self.pos_embed, std=0.02)

    def unpatchify(self, x):
        c = self.in_channels
        p = self.patch_size
        h = w = self.image_size // p

        x = x.reshape(-1, h, w, p, p, c)       # (B, h, w, p, p, c)
        x = x.permute(0, 5, 1, 3, 2, 4)        # (B, c, h, p, w, p)
        x = x.reshape(-1, c, h * p, w * p)      # (B, c, H, W)
        return x

    def forward(self, x, t, labels):
        # Patch embed
        x = self.patch_embed(x)                  # (B, D, 7, 7)
        x = x.flatten(2).transpose(1, 2)         # (B, 49, D)

        # 位置编码
        x = x + self.pos_embed

        # 条件向量
        t_emb = self.time_embed(timestep_embedding(t, self.hidden_dim))
        y_emb = self.label_embed(labels)
        c = t_emb + y_emb

        # Transformer blocks
        for block in self.blocks:
            x = block(x, c)

        # 输出投影 + unpatchify
        x = self.final_linear(self.norm_final(x))
        x = self.unpatchify(x)
        return x


# ============================================================
# 采样（Euler ODE 求解器）
# ============================================================
@torch.no_grad()
def generate(label, num_samples=16, num_steps=100):
    model.eval()
    x = torch.randn(num_samples, 1, 28, 28, device=device)
    labels = torch.full((num_samples,), label, device=device, dtype=torch.long)
    dt = 1.0 / num_steps
    for i in range(num_steps):
        t = torch.full((num_samples,), i * dt, device=device)
        x = x + model(x, t, labels) * dt
    return (x.clamp(-1, 1) + 1) / 2


def sample_images(epoch):
    samples = []
    for label in range(10):
        samples.append(generate(label, num_samples=8))
    samples = torch.cat(samples, dim=0)
    save_image(samples, f'flow_match_sample_{epoch}.png', nrow=8)


# ============================================================
# 训练主循环
# ============================================================
if __name__ == '__main__':
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: 2 * x - 1)
    ])
    train_dataset = torchvision.datasets.MNIST(
        root='./data', train=True, download=True, transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    model = DiT(
        image_size=28, in_channels=1, patch_size=4,
        hidden_dim=256, depth=8, num_heads=4,
        mlp_ratio=4.0, num_classes=10,
    ).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            noise = torch.randn_like(images)
            t = torch.rand(images.size(0), device=device)
            xt = (1 - t.view(-1, 1, 1, 1)) * noise + t.view(-1, 1, 1, 1) * images
            vt_pred = model(xt, t, labels)
            loss = F.mse_loss(vt_pred, images - noise)
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total_loss += loss.item()

        print(f'Epoch [{epoch + 1}/{epochs}], Loss: {total_loss / len(train_loader):.4f}',
              flush=True)
        sample_images(epoch + 1)