扩散 Transformer 与整流流

扩散 Transformer 与整流流

本节摘要:U-Net 不是扩散的秘诀。把它换成 Transformer、把噪声调度换成一条直线流,你就有了 SD3、FLUX,以及 2026 年几乎每一个文生图模型。本节追踪从 U-Net DDPM 到 DiT、MMDiT(SD3)、单+双流 DiT(FLUX)的演进;讲清整流流(Rectified Flow)为何用噪声与数据间的直线插值就能把采样从 1000 步压到 20 步;从零实现带 AdaLN 的 DiT 块和整流流训练循环(各不到 100 行);并区分 SD3、FLUX.1-dev/schnell、Z-Image、Qwen-Image 等模型的架构、参数量与授权。读完本节,你理解了 2026 生成图像栈的范式转移。

对应原课程:Phase 4 · Lesson 23 · diffusion-transformers-rectified-flow(原英文 phases/04-computer-vision/23-diffusion-transformers-rectified-flow/docs/en.md)。

学习目标

阅读完本节,你应当能够:

1. 追踪从 U-Net DDPM(第 10 节)到 DiT、MMDiT(SD3)、单+双流 DiT(FLUX) 的演进。
2. 解释整流流:为何噪声与数据间的直线轨迹让模型 20 步采样而非 1000 步。
3. 从零实现一个微型 DiT 块和整流流训练循环,各不到 100 行。
4. 按架构、参数量、授权区分各模型变体(SD3、FLUX.1-dev、FLUX.1-schnell、Z-Image、Qwen-Image)。

一、问题与直觉

第 10 节搭了一个用 U-Net 去噪器的 DDPM。这套配方在 2020~2023 主导:U-Net + beta 调度 + 预测噪声损失,产出了 Stable Diffusion 1.5、2.1 和 DALL-E 2。

2026 年每个 SOTA 文生图模型都已越过它。Stable Diffusion 3、FLUX、SD4、Z-Image、Qwen-Image、Hunyuan-Image——没有一个用 U-Net,都用扩散 Transformer(DiT)。SD3 和 FLUX 还把 DDPM 噪声调度换成整流流,把噪声到数据的路径拉直,让一致性或蒸馏变体能 1~4 步推理。

这次转移之所以要紧,是因为它是扩散图像生成变得可控、提示精准(SD3/SD4 解决了文字渲染)、生产级快的根本原因。理解 DiT + 整流流,就是理解 2026 的生成图像栈。

从 U-Net 到 Transformer

• DiT(Peebles & Xie, 2023)——用类 ViT 的 Transformer 在隐 patch 上替代 U-Net,经自适应层归一化(AdaLN)做条件。
• MMDiT(SD3, Esser 等, 2024)——文本和图像 token 两条流,各自独立权重,共享一个联合注意力。
• FLUX(Black Forest Labs, 2024)——前 N 块像 SD3 双流,后继块拼接共享权重(单流)以在更深时更高效。
• Z-Image(2025)——60 亿参数的高效单流 DiT,挑战「不计代价堆规模」。

整流流一段话讲清

DDPM 把前向过程定义为 x_t 逐步腐蚀的噪声 SDE,学到的反向是另一个 SDE,用 1000 小步解。

整流流定义干净数据与纯噪声间的直线插值:

x_t = (1 - t) * x_0 + t * epsilon,     t in [0, 1]
​

训一个网络预测速度 v_theta(x_t, t) = epsilon − x_0——沿干净数据到噪声的直线路径的前向方向(dx_t/dt)。采样时,把这个速度反向积分,从噪声走向数据。得到的 ODE 接近直线,故所需积分步数少得多。

SD3 把它叫整流流匹配(Rectified Flow Matching)。FLUX、Z-Image 和多数 2026 模型用同一目标。典型推理:2030 步 Euler(确定性),对比旧 DDPM 的 50+ 步 DDIM。蒸馏 / turbo / schnell / LCM 变体可压到 14 步。

AdaLN 条件

DiT 经自适应层归一化对时间步、类别/文本做条件:从条件向量预测 scale 和 shift,在 LayerNorm 后施加。比 U-Net 里的 FiLM 风格调制干净得多,是每个现代 DiT 的默认。

cond -> MLP -> (scale, shift, gate)
norm(x) * (1 + scale) + shift,然后残差加 * gate
​

SD3 与 FLUX 的文本编码器

• SD3 用三个文本编码器:两个 CLIP 模型 + T5-XXL,嵌入拼接后作文本条件喂给图像流。
• FLUX 用一个 CLIP-L + T5-XXL。
• Qwen-Image / Z-Image 变体用与自家基础 LLM 对齐的自研文本编码器。

文本编码器是 SD3/FLUX 比_SD1.5 更懂提示的一大原因——仅 T5-XXL 就有 47 亿参数。

无分类器引导依然成立

整流流换的是采样器,不是条件。无分类器引导(训练时以 10% 概率丢文本,推理时混合条件与无条件预测)在整流流上完全一样地工作。多数 2026 模型用引导尺度 3.5~5——比 SD1.5 的 7.5 低,因为整流流模型默认就更贴合提示。

一致性、Turbo、Schnell、LCM

四个名字,同一思想:把慢的多步模型蒸馏成快的少步模型。

• LCM(隐一致性模型)——训一个学生,从任意中间 x_t 一步预测最终 x_0。
• SDXL Turbo / FLUX schnell——对抗扩散蒸馏训出的 1~4 步模型。
• SD Turbo——OpenAI 式一致性模型适配到隐扩散。

任何新模型的生产服务都同时出货「全质量」检查点和「turbo / schnell」变体。Schnell(德语「快」,Black Forest Labs 惯例)跑 1~4 步,契合实时流水线。

2026 模型版图

FLUX.1-schnell 是 2026 开源默认,Z-Image 是效率领头羊,FLUX.2 和 SD4 是当前质量尖端的尖端。

为何这次范式转移要紧

DDPM + U-Net 能用。DiT + 整流流更好、更快、扩展更干净。这次转移平行于 NLP 里 RNN 到 Transformer:两者解的是同一问题,但 Transformer 能扩展,于是统治了一切。2026 年每个关于图像、视频、3D 生成的论文都用 DiT 形状的去噪器,且通常用整流流目标。U-Net DDPM 现在主要是教学性的(第 10 节)。

二、从零实现

步骤 1:带 AdaLN 的 DiT 块

import torch
import torch.nn as nn


class AdaLNZero(nn.Module):
    """
    带门的自适应 LayerNorm。从条件预测 (scale, shift, gate)。
    初始化使整个块起步即为恒等映射(零初始化)。
    """

    def __init__(self, dim, cond_dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
        self.mlp = nn.Linear(cond_dim, dim * 3)
        nn.init.zeros_(self.mlp.weight)
        nn.init.zeros_(self.mlp.bias)

    def forward(self, x, cond):
        scale, shift, gate = self.mlp(cond).chunk(3, dim=-1)
        h = self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        return h, gate.unsqueeze(1)


class DiTBlock(nn.Module):
    def __init__(self, dim=192, heads=3, mlp_ratio=4, cond_dim=192):
        super().__init__()
        self.adaln1 = AdaLNZero(dim, cond_dim)
        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.adaln2 = AdaLNZero(dim, cond_dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * mlp_ratio),
            nn.GELU(),
            nn.Linear(dim * mlp_ratio, dim),
        )

    def forward(self, x, cond):
        h, gate1 = self.adaln1(x, cond)
        a, _ = self.attn(h, h, h, need_weights=False)
        x = x + gate1 * a
        h, gate2 = self.adaln2(x, cond)
        x = x + gate2 * self.mlp(h)
        return x
​

AdaLNZero 因 MLP 权重初始化为零而起步为恒等映射,训练把块从恒等推开——这极大稳定了深层 Transformer 扩散模型。

步骤 2:微型 DiT

def timestep_embedding(t, dim):
    import math
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
    args = t[:, None].float() * freqs[None]
    return torch.cat([args.sin(), args.cos()], dim=-1)


class TinyDiT(nn.Module):
    def __init__(self, image_size=16, patch_size=2, in_channels=3, dim=96, depth=4, heads=3):
        super().__init__()
        self.patch_size = patch_size
        self.num_patches = (image_size // patch_size) ** 2
        self.patch = nn.Conv2d(in_channels, dim, kernel_size=patch_size, stride=patch_size)
        self.pos = nn.Parameter(torch.zeros(1, self.num_patches, dim))
        self.time_mlp = nn.Sequential(
            nn.Linear(dim, dim * 2),
            nn.SiLU(),
            nn.Linear(dim * 2, dim),
        )
        self.blocks = nn.ModuleList([DiTBlock(dim, heads, cond_dim=dim) for _ in range(depth)])
        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False)
        self.head = nn.Linear(dim, patch_size * patch_size * in_channels)

    def forward(self, x, t):
        n = x.size(0)
        x = self.patch(x)
        x = x.flatten(2).transpose(1, 2) + self.pos
        t_emb = self.time_mlp(timestep_embedding(t, self.pos.size(-1)))
        for blk in self.blocks:
            x = blk(x, t_emb)
        x = self.norm_out(x)
        x = self.head(x)
        return self._unpatchify(x, n)

    def _unpatchify(self, x, n):
        p = self.patch_size
        h = w = int(self.num_patches ** 0.5)
        x = x.view(n, h, w, p, p, -1).permute(0, 5, 1, 3, 2, 4).reshape(n, -1, h * p, w * p)
        return x
​

步骤 3:整流流训练

import torch.nn.functional as F

def rectified_flow_train_step(model, x0, optimizer, device):
    model.train()
    x0 = x0.to(device)
    n = x0.size(0)
    t = torch.rand(n, device=device)
    epsilon = torch.randn_like(x0)
    x_t = (1 - t[:, None, None, None]) * x0 + t[:, None, None, None] * epsilon

    target_velocity = epsilon - x0
    pred_velocity = model(x_t, t)

    loss = F.mse_loss(pred_velocity, target_velocity)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()
​

对比第 10 节 DDPM 的预测噪声损失:结构相同,目标不同。我们不预测噪声 epsilon,而预测速度 epsilon − x_0,它沿直线插值从数据指向噪声。

步骤 4:Euler 采样器

整流流是 ODE。Euler 法最简单,对训练良好的整流流模型,20+ 步时精度几乎与高阶求解器一样。

@torch.no_grad()
def rectified_flow_sample(model, shape, steps=20, device="cpu"):
    model.eval()
    x = torch.randn(shape, device=device)
    dt = 1.0 / steps
    t = torch.ones(shape[0], device=device)
    for _ in range(steps):
        v = model(x, t)
        x = x - dt * v
        t = t - dt
    return x
​

20 步。训练好的模型上产出可与 1000 步 DDPM 媲美的样本。

步骤 5:端到端冒烟测试

import numpy as np

def synthetic_blobs(num=200, size=16, seed=0):
    rng = np.random.default_rng(seed)
    out = np.zeros((num, 3, size, size), dtype=np.float32)
    yy, xx = np.meshgrid(np.arange(size), np.arange(size), indexing="ij")
    for i in range(num):
        cx, cy = rng.uniform(4, size - 4, size=2)
        r = rng.uniform(2, 4)
        mask = (xx - cx) ** 2 + (yy - cy) ** 2 < r ** 2
        colour = rng.uniform(-1, 1, size=3)
        for c in range(3):
            out[i, c][mask] = colour[c]
    return torch.from_numpy(out)
​

用整流流在其上训 TinyDiT。500 步后,采样输出应像淡淡的颜色斑块。

三、框架对比

用 FLUX / SD3 / Z-Image 做真实图像生成,diffusers 在统一 API 下提供全部:

from diffusers import FluxPipeline, StableDiffusion3Pipeline
import torch

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16,
).to("cuda")

out = pipe(
    prompt="a golden retriever surfing a tsunami, hyperrealistic, studio lighting",
    guidance_scale=0.0,           # schnell 训练时没用 CFG
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
out.save("surf.png")
​

三行。FLUX.1-schnell 四步。换模型 id 为 black-forest-labs/FLUX.1-dev 可在 20~30 步带 CFG 拿更高画质。

SD3:

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    torch_dtype=torch.bfloat16,
).to("cuda")
out = pipe(prompt, guidance_scale=3.5, num_inference_steps=28).images[0]
​

四、可复用产物

本节产出两个可复用文件(位于原课程 outputs/):

• prompt-dit-model-picker.md:按质量、延迟、授权约束,在 SD3、FLUX.1-dev、FLUX.1-schnell、Z-Image、SD4 Turbo 间挑。
• skill-rectified-flow-trainer.md:写出带 AdaLN DiT 和 Euler 采样的完整整流流训练循环。

五、练习

1. (简单) 在合成斑块数据集上训上述 TinyDiT 500 步,对比 10、20、50 步 Euler 产出的样本。
2. (中等) 加文本条件:把可学类别嵌入拼到时间嵌入上(按颜色分 10 个斑块「类别」)。用类别 0、5、9 采样,验证颜色匹配。
3. (困难) 计算等大小网络在等数据、等步数下,整流流版与 DDPM 版生成样本间的 Fréchet 距离(FID 代理),报告哪个收敛更快。

本节要点回顾

1. U-Net 不是扩散的秘诀——2026 每个 SOTA 文生图模型都用 DiT,SD3/FLUX 还用整流流换掉 DDPM 调度。
2. 架构演进:DiT(Transformer 替 U-Net,AdaLN 条件)→ MMDiT(文本/图像双流共享注意力,SD3)→ FLUX(前双后单流提效)→ Z-Image(高效单流)。
3. 整流流 = 直线插值——x_t=(1−t)x_0+t·ε,网络预测速度 v=ε−x_0,ODE 接近直线,20 步即可采样。
4. AdaLN 条件——从条件预测 scale/shift/gate,LayerNorm 后施加;零初始化使块起步为恒等,稳定深网。
5. 文本编码器要紧——SD3 三编码器(CLIP×2+T5-XXL),FLUX(CLIP-L+T5-XXL);T5-XXL 就 47 亿参数,这是 SD3/FLUX 更懂提示的主因。
6. CFG 依然成立——只换采样器不换条件;2026 引导尺度 3.5~5,比 SD1.5 的 7.5 低。
7. 少步蒸馏四名:LCM、Turbo、Schnell、一致性——1~4 步,生产实时。
8. 2026 版图:FLUX.1-schnell(开源默认)、Z-Image(效率)、FLUX.2/SD4(质量尖端),从 20 亿到 800 亿。
9. 范式转移平行于 RNN→Transformer——都解同一问题,但 Transformer 能扩展、能干净 scale,故统治;U-Net DDPM 现主要教学。

下一节进入 SAM3 开放词表分割——用提示驱动的分割模型,给任意物体画掩码,无论训练时见过没有。