Flow Matching 教学项目

这个项目用中文 Jupyter notebook 讲解 Flow Matching 及其一步生成扩展算法，并采用统一的时间约定：

• x_0 是噪声分布中的样本。
• x_1 是数据分布中的样本。
• 时间 t 从 0 到 1，ODE 采样把噪声逐步流向数据。

当前包含两份教程：

• flow_matching_tutorial.ipynb：PyTorch 基础回顾、Flow Matching 数学推导、2D toy 数据实验，以及 MNIST 上的完整训练和采样流程。
• one-step_flow_models_tutorial.ipynb：Shortcut Models、MeanFlow、Pixel MeanFlow 的数学背景、算法细节、2D toy/MNIST 训练采样，以及 MNIST-FID 横向评测。

神经网络结构已抽到独立模块：

• mlp.py：2D toy 实验使用的条件 MLP。
• unet.py：MNIST 实验使用的轻量条件 UNet。

从零开始运行

1. 安装 uv

如果本机还没有安装 uv，可以使用官方安装脚本：

curl -LsSf https://astral.sh/uv/install.sh | sh

安装完成后，重新打开一个终端，检查 uv 是否可用：

uv --version

如果你使用 Homebrew，也可以这样安装：

brew install uv

2. 进入项目目录

git clone https://github.com/sdt201909/Flow_Matching_Tutorial_For_Beginner.git
cd Flow_Matching_Tutorial_For_Beginner

3. 安装项目依赖

项目使用 pyproject.toml 和 uv.lock 管理 Python 环境。运行：

uv sync

这会在当前目录创建 .venv/，并安装 PyTorch、TorchVision、JupyterLab、ipywidgets、Matplotlib 等依赖。

4. 启动 JupyterLab

uv run jupyter lab flow_matching_tutorial.ipynb

如果要运行一步生成算法章节：

uv run jupyter lab one-step_flow_models_tutorial.ipynb

打开 notebook 后，从上到下运行即可。第一次运行 MNIST 章节时，torchvision 会自动下载 MNIST 到 ./data。训练和可视化图片会保存到 ./outputs。

运行建议

Notebook 默认开启 FAST_MODE = True，用于在可接受时间内跑完整个流程。一步生成章节的默认训练步数已经比最初版本更长：2D toy 中 Shortcut 单独使用 SHORTCUT_TOY_STEPS=12000，MeanFlow / Pixel MeanFlow 使用 TOY_STEPS=3000；MNIST 默认 MNIST_STEPS=5000、MNIST_BATCH=128。完整运行会更久，但能更稳定地观察主要结构。

one-step_flow_models_tutorial.ipynb 的 2D toy 使用真正的 square checkerboard 采样和输入 Fourier features。当前版本不加入额外的 endpoint anchor，也不设计噪声到数据的确定性配对映射；每个 batch 都独立采样 x_0 与 x_1。Shortcut、MeanFlow、Pixel MeanFlow 都使用各自算法损失训练。Shortcut 采用官方式 dt_base=log2(1/d) 离散步长、bootstrap_every=8、EMA teacher、可选 bootstrap_dt_bias、合并 batch MSE；EMA 只作为 bootstrap teacher，默认评估/采样使用 raw model，与官方默认 use_ema=0 对齐。MNIST 条件 Shortcut 额外使用官方式 null label dropout、bootstrap_dt_bias=1 和 CFG 采样，并同时展示 1-step 与 8-step 结果：1-step 用来观察短训练下的一步分支，8-step 用来确认同一模型已学到可用的 Shortcut flow。MeanFlow 类算法的 r,t 采样遵循参考实现：采两个时间点排序成 r <= t，并按 flow_ratio 把一部分样本设成 r=t。pMF 使用官方式 pixel-space (u, v) 双 head：网络原始输出先转换成速度，训练使用 loss_u + loss_v，采样只用 u head。

MNIST 会 resize 到 32x32。生成模型训练时数据缩放到 [-1,1]，采样与可视化时再变回 [0,1]。一步生成章节中的 MNIST 模型使用 class-conditional SmallUNet，采样时默认用均衡类别标签生成 0 到 9。2D toy 为了数值稳定固定使用 AdamW；MNIST 也默认使用 AdamW，以减少 MeanFlow JVP 训练中的 loss 尖峰。当前 pMF 与 MeanFlow 是教学版横向对比，不是严格的一变量消融；小网络、短训练、无 CFG/perceptual auxiliary 的设定下，pMF 的局部指标不一定稳定优于 MeanFlow。如果想实验 hybrid Muon/AdamW，可以这样启动：

MNIST_USE_MUON=1 uv run jupyter lab one-step_flow_models_tutorial.ipynb

如果想看到更清楚的 2D checkerboard、更像数字的 MNIST 样本，或更稳定的 MNIST-FID 对比，可以在 notebook 的第一个代码单元里把：

FAST_MODE = True

改成：

FAST_MODE = False

然后重新运行训练相关单元。MacBook 上建议先用 fast mode，确认流程跑通后再增加训练步数；MNIST 的 FAST_MODE=False 会明显更久，但通常能得到更清晰的样本和更稳定的评测结果。

one-step_flow_models_tutorial.ipynb 会在 MNIST 上计算教学版 MNIST-FID。它使用一个本地训练的小型 MNIST 特征网络，而不是标准 ImageNet Inception-V3，因此适合比较本 notebook 里三个算法的相对表现，不应和论文中的 ImageNet FID 数值直接比较。

项目结构

.
├── README.md
├── LICENSE
├── pyproject.toml
├── uv.lock
├── flow_matching_tutorial.ipynb
├── one-step_flow_models_tutorial.ipynb
├── mlp.py
├── unet.py
├── data/       # 运行时自动生成，存放 MNIST
└── outputs/    # 运行时自动生成，存放可视化结果

data/、outputs/、checkpoints/、runs/ 以及常见模型权重文件不会提交到仓库；它们会在本地运行 notebook 时按需生成。

许可证

本项目采用 MIT License，详见 LICENSE。

常见问题

如果 Jupyter 里找不到当前环境，可以运行：

uv run python -m ipykernel install --user --name flow-matching-tutorial --display-name "Flow Matching Tutorial"

然后在 Jupyter 的 kernel 菜单里选择 Flow Matching Tutorial。

如果 MacBook 使用 MPS 时遇到个别算子报错，可以在第一个代码单元中把设备临时改成 CPU：

device = torch.device("cpu")

CPU 会慢一些，但本 notebook 的 fast mode 仍然可以跑通。