在上一章节《Stable Diffusion 3.5 开发指南（一）：模型获取与调用》中，我们详细介绍了 Stable Diffusion 3.5 的模型获取方法、API 调用流程以及手动调用步骤。本章将深入解析 Stable Diffusion 3.5 的核心生成机制——Flow Matching，帮助开发者从根本上理解模型的训练原理与推理过程，为后续的模型优化和自定义开发奠定基础。

1 Flow Matching 的理论概述与核心思想

1.1 基本概念

Flow Matching（流匹配）是一种基于连续归一化流（Continuous Normalizing Flows, CNFs）的生成建模方法，其核心在于学习一个时间依赖的向量场 $v_t(x)$，通过常微分方程（ODE）描述数据从简单分布（如噪声）到目标分布（如图像）的连续变换过程 。该方法避免了传统扩散模型中计算昂贵的模拟步骤，显著提升训练效率 。

1.2 核心思想

Flow Matching 的核心思想是将生成建模问题转化为“学习驱动数据沿连续概率路径流动的向量场”：

连续概率路径：在时间区间 $t\in [0,1]$ 上定义一条从简单分布（如高斯噪声 $x_0$）到目标数据分布 $x_1$ 的连续路径 $x_t$，刻画数据的演化过程。

时间依赖向量场：通过学习向量场 $v_t(x)$，描述每个时间点上数据点沿路径的变化方向和速度，即满足微分方程：

$$\frac{d{x}_t}{dt}=v_t(x_t,t)$$

简言之，Flow Matching 的目标是让模型学会“沿着连续路径流动”的正确向量场，以实现高效、一步生成数据的能力

---

2 数学推导与公式细节

2.1 流匹配（Flow Matching, FM）

条件概率路径：给定数据点 $x_1\sim p_1$ 和噪声 $x_0\sim p_0$，条件路径定义为：
$$x_t=t\cdot x_1+(1-t)\cdot x_0$$
条件向量场：对 $x_t$ 求导得目标向量场：
$$v=\frac{d{x}_t}{dt}=x_1-x_0$$

2.2 损失函数

FM 的损失函数通过最小化预测向量场与目标向量场的均方误差：
$${\mathcal{L}}_{\text{CFM}}={\mathbb{E}}_{t,x_1,x_0}{\Vert v_{\theta }(x_t,t)-(x_1-x_0)\Vert }^2$$

其中，$v_{\theta }$就是我们要训练的模型， $t\sim \mathcal{U}(0,1)$，$x_0\sim \mathcal{N}(0,I)$

2.3 推理过程

在流匹配（Flow Matching）模型中，推理（或称为采样）过程是通过在时间轴上逆向执行学习到的向量场来完成的。具体而言，我们将连续的时间区间 $t\in [0,1]$ 离散化为 $N$ 个等间隔的时间步，然后从先验分布（通常是标准高斯分布）中抽取初始样本，并按照学习到的向量场 $v_{\theta }$ 逐步迭代更新样本，最终生成目标数据分布的样本。

① 时间离散化

为了进行数值积分，我们将时间区间 ([0, 1]) 均匀分割为 ( N ) 份，得到一系列时间点：
$$t_0=0<t_1<t_2<\cdots <t_{N-1}<t_N=1$$
其中，每个时间步的间隔为：
$$\Delta t=\frac{1}{N}$$
因此，时间点可以表示为：
$$t_k=k\cdot \Delta t,k=0,1,\dots ,N$$

② 逆向采样过程

推理过程从 $t=0$开始，到$t=1$结束。在训练中，我们定义了从数据$x_1$ 到噪声$x_0$的前向过程$x_t=t{x}_1+(1-t)x_0$，而在推理时，我们需要执行反向过程：从噪声$x_0$出发，逐步逼近数据$x_1$。

初始化：从先验分布 $p_0=\mathcal{N}(0,I)$中采样初始噪声：
$$x_{t_0}\sim \mathcal{N}(0,I)$$

迭代更新：对于 $k=0,1,\dots ,N-1$，我们按照以下步骤更新样本：

1. 计算当前时间点$t_k$对应的向量场输出：
   $$v_{\theta }(x_{t_k},t_k)$$
2. 使用欧拉方法（或其他数值积分方法）更新样本，以逼近下一个时间点 $t_{k+1}=t_k+\Delta t$ 的样本：
   $$x_{t_{k+1}}=x_{t_k}+\Delta t\cdot v_{\theta }(x_{t_k},t_k)$$

终止：当迭代进行到$k=N-1$ 时，我们得到 $x_{t_N}=x_1$，即生成的样本。

综上所述，推理过程的算法可以描述如下：

算法：流匹配模型推理（采样）

1. 输入：

   • 训练好的向量场模型$v_{\theta }(x,t)$
   • 时间步数 $N$
   • 先验分布$p_0=\mathcal{N}(0,I)$

2. 初始化：

   • 设置时间步长$\Delta t=\frac{1}{N}$
   • 从$p_0$ 中采样初始噪声：$x_0\sim \mathcal{N}(0,I)$

3. 循环：对于$k=0$到 $N-1$：

   • 计算当前时间 $t_k=k\cdot \Delta t$
   • 计算向量场：$v_k=v_{\theta }(x_k,t_k)$
   • 更新样本：$x_{k+1}=x_k+\Delta t\cdot v_k$

4. 输出：生成的样本 $x_N$

• 数值积分方法：上述过程使用了最简单的欧拉方法进行数值积分。在实际应用中，可以根据需要选择更高阶的方法（如Heun方法、Runge-Kutta方法等）以提高精度和稳定性。
• 时间步数$N$：$N$的选择会影响生成样本的质量和计算效率。较大的 $N$通常能产生更准确的样本，但也会增加计算成本。
• 与训练的一致性：在训练中，我们通过最小化条件流匹配损失来学习向量场 $v_{\theta }$，使其在任意时间点$t$和目标$x_1$条件下能够预测正确的方向。推理过程则是该学习结果的实际应用，通过逆向积分将噪声转化为数据。

通过这种离散化的推理过程，流匹配模型能够从简单的噪声分布出发，生成符合复杂数据分布的样本。

2.4 直观类比：高中物理的速度问题

为了帮助理解流匹配的推理过程，我们可以用一个简单的高中物理例子来类比。

① 问题设定

假设一辆小车从起点 $x_0$ 出发，要到达终点 $x_1$。我们知道：

• 起点：$x_0=0$ 米（初始位置）
• 终点：$x_1=100$ 米（目标位置）
• 总时间：$T=10$ 秒

② 匀速运动的情况

如果是匀速直线运动：

• 速度：$v=\frac{x_1-x_0}{T}=\frac{100}{10}=10$ 米/秒
• 位置随时间变化：$x_t=x_0+v\cdot t=0+10t$

在任意时刻 $t$，我们都能准确知道小车的位置。

③ 变速运动的情况（类比流匹配）

但现实中，小车可能不是匀速运动。假设我们不知道具体的运动规律，但有一个智能速度计（类比我们的模型 $v_{\theta }$），它能够根据：

• 当前时间 $t$
• 当前位置 $x_t$

告诉我们当前瞬时速度 $v_{\theta }(x_t,t)$。

推理过程：

1. 时间离散化：把10秒分成10个时间段，每段 $\Delta t=1$ 秒

2. 初始状态：$t_0=0$ 秒，$x_0=0$ 米

3. 逐步更新：

   • 在 $t_0=0$ 秒时，速度计显示 $v_0=12$ 米/秒
     那么1秒后的位置：$x_1=x_0+v_0\cdot \Delta t=0+12\times 1=12$ 米

   • 在 $t_1=1$ 秒时，速度计显示 $v_1=9$ 米/秒
     那么2秒后的位置：$x_2=x_1+v_1\cdot \Delta t=12+9\times 1=21$ 米

   • 在 $t_2=2$ 秒时，速度计显示 $v_2=11$ 米/秒
     那么3秒后的位置：$x_3=x_2+v_2\cdot \Delta t=21+11\times 1=32$ 米

   • … 继续这个过程…

   • 在 $t_9=9$ 秒时，速度计显示 $v_9=8$ 米/秒
     那么10秒后的位置：$x_{10}=x_9+v_9\cdot \Delta t$

4. 最终结果：经过10次迭代，我们得到 $x_{10}\approx 100$ 米

④ 与流匹配的对应关系

物理问题	流匹配模型
起点 $x_0$	噪声样本 $x_0\sim \mathcal{N}(0,I)$
终点 $x_1$	目标数据样本 $x_1\sim p_1$
时间 $t\in [0,10]$	时间 $t\in [0,1]$
智能速度计 $v(t,x_t)$	神经网络 $v_{\theta }(x_t,t)$
位置更新 $x_{k+1}=x_k+v_k\cdot \Delta t$	样本更新 $x_{k+1}=x_k+v_{\theta }\cdot \Delta t$
总时间分段数 $N=10$	时间步数 $N$

⑤ 关键理解点

1. 模型的作用：我们的模型 $v_{\theta }$ 就像那个"智能速度计"，它不需要知道整个运动轨迹，只需要在任意时刻告诉我们"现在应该往哪个方向走，走多快"。

2. 训练的目标：在训练时，我们让模型学会在任意中间位置 $x_t$ 和时刻 $t$ 时，都能正确预测"应该朝向终点的速度"。

3. 推理的过程：就像我们一步步用速度计读数来推算小车位置一样，我们从噪声开始，一步步用模型的预测来更新样本，最终生成目标数据。

4. 为什么有效：因为模型在训练时已经学会了整个"运动规律"，所以在推理时，即使我们从完全不同的起点（噪声）开始，模型也能引导我们走到正确的终点（真实数据分布）。

这种类比让我们理解：流匹配的推理本质上是一个逐步导航的过程，模型在每个时间步提供方向指导，带领样本从简单的噪声分布"走"到复杂的数据分布

---

3 MNIST 数据集Flow Matching 代码实战

3.1 环境和依赖配置

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import os
from tqdm import tqdm

说明

• torch：PyTorch主库，用于深度学习计算。
• nn、nn.functional：神经网络模块与常用操作。
• DataLoader、datasets、transforms：数据加载与预处理。
• matplotlib.pyplot：可视化生成样本。
• tqdm：训练进度条显示。
• os：文件与目录操作。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

• 自动选择 GPU（可用时）或 CPU 进行训练和推理。

---

3.2 超参数配置

class Config:
    # 训练参数
    batch_size = 128
    learning_rate = 1e-3
    num_epochs = 100
    gradient_clip = 1.0

    # 模型参数
    image_size = 32
    in_channels = 1
    out_channels = 1

    # 保存路径
    save_dir = "flow_matching_mnist"
    model_save_path = "flow_matching_model.pth"

说明

• batch_size、learning_rate、num_epochs：训练参数。
• image_size、in_channels、out_channels：模型输入输出尺寸。
• save_dir、model_save_path：保存模型和生成样本路径。

os.makedirs(config.save_dir, exist_ok=True)

• 创建保存目录，如果已存在不会报错。

---

3.3 数据加载与预处理

def get_dataloaders():
    transform = transforms.Compose([
        transforms.Resize((config.image_size, config.image_size)),
        transforms.ToTensor(),
        transforms.Normalize((0.,), (1.,)),
    ])

    train_dataset = datasets.MNIST(
        root='./data',
        train=True,
        download=True,
        transform=transform
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=config.batch_size,
        shuffle=True,
        num_workers=2
    )

    return train_loader

说明

1. Resize：调整 MNIST 图片为 32x32。
2. ToTensor：转换为 PyTorch 张量。
3. Normalize：将像素归一化到 [0,1]（简单归一化）。
4. DataLoader：分批加载，支持多线程。

---

3.4 Flow Matching 模型定义

from diffusers import UNet2DModel

class FlowMatchingModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.unet = UNet2DModel(
            sample_size=config.image_size,
            in_channels=config.in_channels,
            out_channels=config.out_channels,
            layers_per_block=2,
            block_out_channels=(64, 128, 256, 512),
            down_block_types=("DownBlock2D", "DownBlock2D", "DownBlock2D", "DownBlock2D"),
            up_block_types=("UpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D"),
        )

    def forward(self, x, t):
        batch_size = x.shape[0]
        t = t * torch.ones(batch_size, device=x.device)
        return self.unet(x, t).sample

说明

• 使用 UNet 架构预测向量场 $v_{\theta }(x_t,t)$。
• forward 函数扩展时间步 $t$ 到 batch 大小，并返回预测向量场。

---

3.5 Flow Matching 损失函数

def correct_flow_matching_loss(model, x1):
    batch_size = x1.shape[0]

    t = torch.rand(batch_size, device=x1.device)
    t = t.view(-1, 1, 1, 1)  # 扩展到与x相同维度

    x0 = torch.randn_like(x1)

    xt = (1 - t) * x0 + t * x1

    target_vector_field = x1 - x0

    pred_vector_field = model(xt, t.squeeze())

    loss = F.mse_loss(pred_vector_field, target_vector_field, reduction='mean')

    return loss

说明

• 随机采样时间步 $t\sim \mathcal{U}(0,1)$。
• 插值：$x_t=(1-t)x_0+t{x}_1$。
• 目标向量场：$u_t=x_1-x_0$。
• 使用均方误差计算预测向量场与目标向量场差距。

---

3.6 训练函数

def train_model():
    model = FlowMatchingModel().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.num_epochs)

    train_loader = get_dataloaders()
    losses = []
    model.train()

    for epoch in range(config.num_epochs):
        epoch_loss = 0.0
        progress_bar = tqdm(train_loader, desc=f'Epoch {epoch + 1}/{config.num_epochs}')

        for batch_idx, (data, _) in enumerate(progress_bar):
            data = data.to(device)
            loss = correct_flow_matching_loss(model, data)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
            progress_bar.set_postfix({'Loss': f'{loss.item():.6f}'})

        scheduler.step()
        avg_loss = epoch_loss / len(train_loader)
        losses.append(avg_loss)

        print(f'Epoch {epoch + 1}/{config.num_epochs}, Average Loss: {avg_loss:.6f}, LR: {scheduler.get_last_lr()[0]:.8f}')

        if (epoch + 1) % 10 == 0:
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': avg_loss,
            }, os.path.join(config.save_dir, f'model_epoch_{epoch + 1}.pth'))
            generate_samples(model, epoch + 1, num_samples=16)

    torch.save(model.state_dict(), os.path.join(config.save_dir, config.model_save_path))

    plt.figure(figsize=(10, 5))
    plt.plot(losses)
    plt.title('Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.savefig(os.path.join(config.save_dir, 'training_loss.png'))
    plt.close()

    return model, losses

说明

• 使用 AdamW 优化器 + CosineAnnealingLR 学习率调度。
• 每个 epoch 循环训练数据，计算 Flow Matching 损失并反向传播。
• 每 10 个 epoch 保存模型并生成样本。
• 绘制训练损失曲线。

---

3.7 采样函数（欧拉法）

def generate_samples(model, epoch=None, num_samples=16, steps=100):
    """
    使用欧拉法采样生成图像
    参数:
        model: 训练好的 Flow Matching 模型
        epoch: 当前训练轮数（用于保存可视化图片）
        num_samples: 生成样本数量
        steps: 时间离散步数 N
    """
    model.eval()  # 推理模式，不计算梯度

    with torch.no_grad():
        # 1️⃣ 从标准正态分布初始化噪声 x0 ~ N(0,I)
        x = torch.randn(num_samples, config.in_channels, config.image_size, config.image_size, device=device)

        # 2️⃣ 离散化时间步 t ∈ [0,1]
        timesteps = torch.linspace(0.0, 1.0, steps + 1, device=device)

        # 3️⃣ 欧拉方法积分
        for i in range(steps):
            t = timesteps[i]  # 当前时间点
            dt = timesteps[i + 1] - timesteps[i]  # 时间步长 Δt

            # 4️⃣ 模型预测向量场 vθ(x_t, t)
            vector_field = model(x, t)

            # 5️⃣ 欧拉步更新: x_{t+dt} = x_t + vθ(x_t, t) * dt
            x = x + vector_field * dt

        # 6️⃣ 可视化生成样本
        fig, axes = plt.subplots(4, 4, figsize=(8, 8))
        for i, ax in enumerate(axes.flat):
            if i < num_samples:
                img = x[i].cpu().squeeze().numpy()
                ax.imshow(img, cmap='gray')
                ax.axis('off')

        title = f'Generated Samples (Epoch {epoch})' if epoch else 'Generated Samples'
        plt.suptitle(title)
        plt.tight_layout()
        filename = f'samples_epoch_{epoch}.png' if epoch else 'final_samples.png'
        plt.savefig(os.path.join(config.save_dir, filename))
        plt.close()

    return x

说明

• 从标准正态分布初始化 $x_0\sim N(0,I)$。
• 时间步从 0 到 1 离散化。
• 欧拉法积分：$x_{t+dt}=x_t+v_{\theta }(x_t,t)\cdot dt$。
• 生成 16 张图像进行可视化。

---

3.8 改进采样函数（Heun 方法）

def generate_samples_improved(model, epoch=None, num_samples=16, steps=100, method='euler'):
    """
    改进采样函数，支持欧拉法或 Heun 方法
    参数:
        method: 'euler' 或 'heun'
    """
    model.eval()

    with torch.no_grad():
        # 1️⃣ 初始化噪声
        x = torch.randn(num_samples, config.in_channels, config.image_size, config.image_size, device=device)

        # 2️⃣ 离散化时间
        timesteps = torch.linspace(0.0, 1.0, steps + 1, device=device)

        if method == 'euler':
            # 3️⃣ 欧拉法
            for i in range(steps):
                t = timesteps[i]
                dt = timesteps[i + 1] - timesteps[i]
                vector_field = model(x, t)
                x = x + vector_field * dt

        elif method == 'heun':
            # 3️⃣ Heun 方法（二阶 Runge-Kutta）
            for i in range(steps):
                t = timesteps[i]
                dt = timesteps[i + 1] - timesteps[i]

                # 3a️⃣ 预测步 (Euler预测)
                k1 = model(x, t)
                x_pred = x + k1 * dt

                # 3b️⃣ 校正步 (计算下一个向量场)
                k2 = model(x_pred, t + dt)

                # 3c️⃣ 更新状态，取平均作为更高精度的预测
                x = x + (k1 + k2) * dt / 2

        # 4️⃣ 可视化生成样本
        fig, axes = plt.subplots(4, 4, figsize=(8, 8))
        for i, ax in enumerate(axes.flat):
            if i < num_samples:
                img = x[i].cpu().squeeze().numpy()
                ax.imshow(img, cmap='gray')
                ax.axis('off')

        title = f'Generated Samples (Epoch {epoch}, {method})' if epoch else f'Generated Samples ({method})'
        plt.suptitle(title)
        plt.tight_layout()
        filename = f'samples_epoch_{epoch}_{method}.png' if epoch else f'final_samples_{method}.png'
        plt.savefig(os.path.join(config.save_dir, filename))
        plt.close()

    return x

说明

• 支持欧拉法和 Heun（二阶 Runge-Kutta）方法。
• Heun 方法可以提高数值积分精度，减少误差。

---

3.9 模型加载函数

def load_model(model_path):
    model = FlowMatchingModel().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    return model

• 加载保存的模型权重用于推理。

---

3.10 主函数

def main():
    model_path = os.path.join(config.save_dir, config.model_save_path)

    if os.path.exists(model_path):
        print("Loading pre-trained model...")
        model = load_model(model_path)
        print("Generating samples with pre-trained model...")
        generate_samples(model, num_samples=16)
        generate_samples_improved(model, num_samples=16, method='heun')
    else:
        print("Training new model...")
        model, losses = train_model()
        print("Generating final samples...")
        generate_samples(model, num_samples=16)
        generate_samples_improved(model, num_samples=16, method='heun')


if __name__ == "__main__":
    main()

说明

• 检查是否存在预训练模型：

  • 有：直接加载并生成样本。
  • 无：训练新模型并生成样本。

---

3.11 实验结果

---

4 总结

• 核心思想：Flow Matching 将生成建模为从噪声到数据的连续路径，通过学习时间依赖向量场 $v_{\theta }(x_t,t)$ 引导样本沿 ODE 运动。

• 损失函数：最小化预测向量场与目标向量场 MSE：

• 推理过程：

  1. 从噪声 $x_0\sim \mathcal{N}(0,I)$ 出发
  2. 时间离散化 $t_0,\dots ,t_N$
  3. 欧拉/Heun 更新：$x_{t+\Delta t}=x_t+v_{\theta }(x_t,t)\cdot \Delta t$
  4. 得到最终生成样本 $x_1$

• 关键理解：模型像“智能速度计”，在每个时间步给出方向和速度，引导噪声逐步逼近目标分布。

• 代码实现：

  • 数据：MNIST，归一化 32×32
  • 模型：UNet 预测向量场
  • 训练：随机采样 (t)、MSE 损失、AdamW + CosineAnnealingLR
  • 采样：支持欧拉法和 Heun 方法，生成样本可视化