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 与扩散模型的对比

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. 关键公式汇总

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