Flow Matching 原理与 MNIST 条件生成实践

1. Flow Matching 简介

Flow Matching 是一种连续归一化流 (Continuous Normalizing Flow, CNF) 生成模型。与扩散模型 (Diffusion Models) 不同，Flow Matching 直接在数据分布和噪声分布之间构建一条确定性的概率流路径，通过常微分方程 (ODE) 来描述这个变换过程。

1.1 核心思想

给定：

• \(x_0 \sim p_0\)：先验分布（如标准高斯噪声）
• \(x_1 \sim p_1\)：目标数据分布（如真实图片）
• \(t \in [0, 1]\)：连续时间变量

Flow Matching 学习一个向量场 \(v_\theta(x_t, t)\)，该向量场定义了从噪声到数据的 ODE：

\[\frac{dx_t}{dt} = v_\theta(x_t, t)\]

1.2 线性插值路径

本实现采用最简单的条件概率路径——线性插值：

\[x_t = (1 - t) \cdot x_0 + t \cdot x_1\]

其中 \(x_0\) 是噪声，\(x_1\) 是目标数据。

对应的真实向量场（速度）为：

\[v(x_t, t) = x_1 - x_0 = \frac{dx_t}{dt}\]

即向量场指向从噪声到目标数据的方向，且大小恒定。

1.3 与扩散模型的对比

特性	扩散模型 (DDPM)	Flow Matching
前向过程	逐步加噪（随机）	线性插值（确定性）
反向过程	去噪（随机 SDE）	ODE 积分（确定性）
损失函数	噪声预测 MSE	向量场预测 MSE
采样方式	逐步去噪	ODE 求解器（如 Euler 法）
理论框架	随机微分方程	常微分方程

2. 代码架构

2.1 整体结构

main.py
├── ConditionedDoubleConv  — 条件双卷积模块
├── Down                   — 下采样模块
├── Up                     — 上采样模块
├── ConditionalUNet        — 条件 U-Net 主模型
├── generate()             — 采样/推理函数
├── sample_images()        — 可视化保存函数
└── main (__main__)        — 训练循环

2.2 超参数配置

device      = "cuda" / "cpu"     # 自动检测
batch_size  = 256
lr          = 1e-4               # AdamW 学习率
epochs      = 100
num_classes = 10                 # MNIST 0-9

3. 核心组件详解

3.1 条件双卷积模块 (`ConditionedDoubleConv`)

class ConditionedDoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels, cond_dim):
        # 第一层：普通卷积 + GroupNorm + SiLU
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.norm1 = nn.GroupNorm(8, out_channels)
        # 第二层：拼接条件后卷积 + GroupNorm + SiLU
        self.conv2 = nn.Conv2d(out_channels + cond_dim, out_channels, 3, padding=1)
        self.norm2 = nn.GroupNorm(8, out_channels)

设计要点：

• 使用 GroupNorm 而非 BatchNorm：在小 batch 下更稳定
• 条件注入采用拼接 (concat) 方式：将时间嵌入和标签嵌入空间扩展后拼接到特征图上
• 激活函数使用 SiLU (Swish)：\(f(x) = x \cdot \sigma(x)\)，比 ReLU 更平滑

前向传播流程：

1. conv1 → GroupNorm → SiLU：提取基础特征
2. 将条件向量 cond 空间广播到与特征图相同尺寸
3. 沿通道维度拼接特征与条件：[x, cond]
4. conv2 → GroupNorm → SiLU：融合条件信息

3.2 下采样模块 (`Down`)

class Down(nn.Module):
    def forward(self, x, cond):
        return self.conv(self.maxpool(x), cond)  # MaxPool → DoubleConv

• 使用 MaxPool2d(2) 将空间尺寸减半
• 通过 ConditionedDoubleConv 融合条件信息并变换通道数
• 通道数变化：64 → 128 → 256

3.3 上采样模块 (`Up`)

class Up(nn.Module):
    def forward(self, x1, x2, cond):
        x1 = self.up(x1)                          # 双线性插值上采样
        # 处理尺寸不匹配（奇数特征图）
        x1 = F.pad(x1, [diffX//2, diffX-diffX//2, diffY//2, diffY-diffY//2])
        return self.conv(torch.cat([x2, x1], 1), cond)  # 拼接跳跃连接

• 使用 Upsample 双线性插值 2× 上采样
• 跳跃连接 (Skip Connection)：拼接编码器对应层的特征图 x2
• Padding 处理奇偶尺寸不一致问题（MNIST 28×28 必然出现）
• 通道数变化：384 → 128 → 64

3.4 条件 U-Net (`ConditionalUNet`)

class ConditionalUNet(nn.Module):
    def __init__(self):
        self.t_dim = 16      # 时间嵌入维度
        self.label_dim = 16  # 标签嵌入维度
        self.cond_dim = 32   # 总条件维度 = 16 + 16

        # 嵌入层
        self.time_embed = nn.Sequential(
            nn.Linear(1, 32), nn.SiLU(), nn.Linear(32, 16))
        self.label_embed = nn.Embedding(10, 16)

        # 编码器：1 → 64 → 128 → 256
        self.inc   = ConditionedDoubleConv(1, 64, 32)
        self.down1 = Down(64, 128, 32)
        self.down2 = Down(128, 256, 32)

        # 解码器：256 → 128 → 64 → 1
        self.up1 = Up(256+128, 128, 32)
        self.up2 = Up(128+64, 64, 32)
        self.outc = nn.Conv2d(64, 1, kernel_size=1)

网络架构图：

输入 x (B, 1, 28, 28)
       │
       ├── inc ──────────────► x1 (B, 64, 28, 28) ─────────────┐
       │                                                         │
       ├── down1 ────────────► x2 (B, 128, 14, 14) ──────────┐  │
       │                                                      │  │
       ├── down2 ────────────► x3 (B, 256, 7, 7)              │  │
       │                         │                             │  │
       │                    ┌────┴────┐                        │  │
       │                    │   up1   │◄── cat(x3, x2) ───────┘  │
       │                    └────┬────┘                           │
       │                         │                                │
       │                    ┌────┴────┐                           │
       │                    │   up2   │◄── cat(up1_out, x1) ─────┘
       │                    └────┬────┘
       │                         │
       └── outc ◄────────────────┘
              │
       输出 v_pred (B, 1, 28, 28)  ← 预测的向量场

条件嵌入

def forward(self, x, t, labels):
    t_emb = self.time_embed(t.view(-1, 1))     # (B,) → (B, 16)
    lbl_emb = self.label_embed(labels)          # (B,) → (B, 16)
    cond = torch.cat([t_emb, lbl_emb], dim=1)   # (B, 32)
    cond = cond.unsqueeze(-1).unsqueeze(-1)     # (B, 32, 1, 1)

• 时间嵌入：标量 \(t\) 通过两层 MLP (1→32→16) 映射到 16 维
• 标签嵌入：类别标签 (0-9) 通过 Embedding 层映射到 16 维
• 拼接后得到 32 维条件向量，用于指导生成过程

4. 训练过程

4.1 损失函数

Flow Matching 的训练目标是最小化预测向量场与真实向量场之间的 MSE：

\[\mathcal{L}(\theta) = \mathbb{E}_{t \sim U(0,1), x_1 \sim p_{data}, x_0 \sim \mathcal{N}(0,I)} \left[ \| v_\theta(x_t, t) - (x_1 - x_0) \|^2 \right]\]

代码实现：

# 1. 采样噪声和时间
noise = torch.randn_like(images)                  # x_0 ~ N(0,I)
t = torch.rand(images.size(0), device=device)     # t ~ U(0,1)

# 2. 线性插值构造 x_t
xt = (1 - t) * noise + t * images                 # x_t = (1-t)·x_0 + t·x_1

# 3. 预测向量场
vt_pred = model(xt, t, labels)                    # v_θ(x_t, t)

# 4. 计算损失
loss = F.mse_loss(vt_pred, images - noise)         # target = x_1 - x_0

4.2 训练循环要点

for epoch in range(epochs):
    for images, labels in train_loader:
        # 数据预处理：归一化到 [-1, 1]
        # transforms.Lambda(lambda x: 2*x - 1)

        # 梯度裁剪，防止训练不稳定
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

关键设计：

• 随机时间采样：\(t \sim U(0, 1)\)，每个 batch 的每张图采样不同的 \(t\)
• 梯度裁剪：max_norm=1.0，防止梯度爆炸
• 数据归一化：将 MNIST 从 [0, 1] 映射到 [-1, 1]，与噪声分布 \(\mathcal{N}(0, I)\) 对齐

5. 采样（推理）过程

5.1 Euler 法求解 ODE

训练完成后，通过数值积分求解 ODE 生成新样本：

\[\frac{dx}{dt} = v_\theta(x, t), \quad x(0) \sim \mathcal{N}(0, I)\]

使用 Euler 法（一阶 ODE 求解器）：

\[x_{t+\Delta t} = x_t + v_\theta(x_t, t) \cdot \Delta t\]

代码实现：

@torch.no_grad()
def generate(label, num_samples=16, num_steps=100):
    model.eval()
    x = torch.randn(num_samples, 1, 28, 28)      # x_0 ~ N(0,I)
    labels = torch.full((num_samples,), label, dtype=torch.long)
    dt = 1.0 / num_steps                          # 步长

    for i in range(num_steps):
        t = torch.full((num_samples,), i * dt)    # 当前时间
        x = x + model(x, t, labels) * dt          # Euler 步进

    return (x.clamp(-1, 1) + 1) / 2              # 映射回 [0, 1]

采样过程示意：

t=0.0:  x ~ N(0,I)  ← 纯噪声
  │     v_θ(x, 0.0) * dt
  │
t=0.5:  x ≈ 半噪半图
  │     v_θ(x, 0.5) * dt
  │
t=1.0:  x ≈ 干净图片  → clamp + 去归一化 → [0,1] 图像

• num_steps=100：100 步 Euler 积分，步长 dt = 0.01
• 最终 clamp(-1, 1) 防止数值溢出，再映射回 [0, 1] 用于可视化

6. 完整流程总结

                    【训练阶段】
    MNIST 数据 x_1          噪声 x_0 ~ N(0,I)
         │                       │
         └───────┬───────────────┘
                 │
         线性插值: x_t = (1-t)·x_0 + t·x_1
         真实向量场: v = x_1 - x_0
                 │
         ┌───────▼───────┐
         │ ConditionalUNet│ ← 条件：t, label
         │   预测 v_θ     │
         └───────┬───────┘
                 │
         Loss = MSE(v_θ, v)  ← 梯度下降


                    【推理阶段】
    噪声 x_0 ~ N(0,I) + 目标 label "7"
         │
         │  for t = 0 → 1 step dt:
         │    x_{t+dt} = x_t + v_θ(x_t, t, label) * dt
         ▼
    生成数字 "7" 的图片

7. 关键公式汇总

描述	公式
插值路径	\(x_t = (1-t)x_0 + t x_1\)
真实向量场	\(v = x_1 - x_0\)
条件嵌入	\(c = [\text{MLP}(t), \text{Embed}(y)]\)
训练损失	\(\mathcal{L} = \|v_\theta(x_t, t, y) - (x_1 - x_0)\|^2\)
Euler 采样	\(x_{t+\Delta t} = x_t + v_\theta(x_t, t, y) \cdot \Delta t\)

8. 实验结果

训练过程中每个 epoch 生成 10 个类别 × 8 个样本 = 80 张图片，保存为 flow_match_sample_{epoch}.png。

早期 epoch 生成的数字模糊，随着训练的进行：

• Epoch 1-5：数字轮廓逐渐出现，仍有较多噪声
• Epoch 10-20：数字清晰度明显提升
• Epoch 30+：生成质量趋于稳定，各数字形态清晰可辨

9. 扩展方向

1. 更高阶 ODE 求解器：使用 RK4 或 DPM-Solver 替代 Euler 法，可用更少步数达到相同质量
2. CFG (Classifier-Free Guidance)：在训练时随机丢弃标签条件，推理时通过引导强度控制生成
3. 更复杂的概率路径：使用 VP/VE 路径替代线性插值
4. 条件增强：加入文本、图像等其他模态的条件
5. 高分辨率生成：替换为更大 U-Net 或在潜空间 (latent space) 中进行 Flow Matching

附录：完整代码

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

class ConditionedDoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels, cond_dim):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.norm1 = nn.GroupNorm(8, out_channels)
        self.conv2 = nn.Conv2d(out_channels + cond_dim, out_channels, kernel_size=3, padding=1)
        self.norm2 = nn.GroupNorm(8, out_channels)

    def forward(self, x, cond):
        x = F.silu(self.norm1(self.conv1(x)))
        cond = cond.expand(-1, -1, x.size(2), x.size(3))
        x = torch.cat([x, cond], dim=1)
        return F.silu(self.norm2(self.conv2(x)))

class Down(nn.Module):
    def __init__(self, in_channels, out_channels, cond_dim):
        super().__init__()
        self.maxpool = nn.MaxPool2d(2)
        self.conv = ConditionedDoubleConv(in_channels, out_channels, cond_dim)

    def forward(self, x, cond):
        return self.conv(self.maxpool(x), cond)

class Up(nn.Module):
    def __init__(self, in_channels, out_channels, cond_dim):
        super().__init__()
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv = ConditionedDoubleConv(in_channels, out_channels, cond_dim)

    def forward(self, x1, x2, cond):
        x1 = self.up(x1)
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])
        return self.conv(torch.cat([x2, x1], dim=1), cond)

class ConditionalUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.t_dim = 16
        self.label_dim = 16
        self.cond_dim = self.t_dim + self.label_dim

        self.time_embed = nn.Sequential(
            nn.Linear(1, 32), nn.SiLU(), nn.Linear(32, self.t_dim)
        )
        self.label_embed = nn.Embedding(num_classes, self.label_dim)

        self.inc = ConditionedDoubleConv(1, 64, self.cond_dim)
        self.down1 = Down(64, 128, self.cond_dim)
        self.down2 = Down(128, 256, self.cond_dim)
        self.up1 = Up(256 + 128, 128, self.cond_dim)
        self.up2 = Up(128 + 64, 64, self.cond_dim)
        self.outc = nn.Conv2d(64, 1, kernel_size=1)

    def forward(self, x, t, labels):
        t_emb = self.time_embed(t.view(-1, 1))
        lbl_emb = self.label_embed(labels)
        cond = torch.cat([t_emb, lbl_emb], dim=1).unsqueeze(-1).unsqueeze(-1)

        x1 = self.inc(x, cond)
        x2 = self.down1(x1, cond)
        x3 = self.down2(x2, cond)
        x = self.up1(x3, x2, cond)
        x = self.up2(x, x1, cond)
        return self.outc(x)

@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 = ConditionalUNet().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}')
        sample_images(epoch + 1)

第31次生成结果：

参考：https://zhuanlan.zhihu.com/p/28731517852