用于可扩展多模态合成的自监督流匹配 (Self-Flow) [ICML'26]

Hila Chefer* · Patrick Esser*
Dominik Lorenz · Dustin Podell · Vikash Raja · Vinh Tong · Antonio Torralba · Robin Rombach
Black Forest Labs

此文件夹包含使用我们在 ImageNet 256×256 上训练的 [Self-Flow](https://bfl.ai/research/self-flow) 扩散模型生成图像的推理代码。 ## 概述 **Self-Flow**（用于可扩展多模态合成的自监督流匹配）是一个训练框架，它将流匹配（flow matching）目标与自监督特征重建目标结合在一起。 此推理代码允许您： 1. 加载 Self-Flow checkpoint（在 ImageNet 256x256 上预训练） 2. 生成 50,000 张图像以进行 FID 评估 可以使用 [ADM 评估套件](https://github.com/openai/guided-diffusion/tree/main/evaluations)对生成的样本进行评估。 ## 环境要求 ``` pip install -r requirements.txt ``` ## 快速开始 ### 下载 Checkpoint ``` python -c " from huggingface_hub import hf_hub_download hf_hub_download( repo_id='Hila/Self-Flow', filename='selfflow_imagenet256.pt', local_dir='./checkpoints' ) print('Downloaded!') " ``` ### 生成 5 万个样本（建议使用多 GPU） ``` torchrun --nnodes=1 --nproc_per_node=8 sample.py \ --ckpt checkpoints/selfflow_imagenet256.pt \ --output-dir ./samples \ --num-fid-samples 50000 ``` ### 单 GPU ``` python sample.py \ --ckpt checkpoints/selfflow_imagenet256.pt \ --output-dir ./samples \ --num-fid-samples 50000 \ --batch-size 64 ``` ## 命令行参数 | 参数 | 默认值 | 描述 | |----------|---------|-------------| | `--ckpt` | 必填 | 模型 checkpoint 的路径 | | `--output-dir` | `./samples` | 生成样本的输出目录 | | `--num-fid-samples` | `50000` | 要生成的样本数量 | | `--batch-size` | `64` | 每个 GPU 的 batch size | | `--num-steps` | `250` | 扩散采样的步数 | | `--mode` | `SDE` | 采样模式：`SDE` 或 `ODE` | | `--seed` | `31` | 用于可重复性的随机种子 | | `--cfg-scale` | `1.0` | 无分类器引导（Classifier-free guidance）缩放比例（1.0 = 无引导，如论文中所用） | ## 评估 生成的 `.npz` 文件可与 [ADM 评估套件](https://github.com/openai/guided-diffusion/tree/main/evaluations)一起使用，以计算 FID、IS、Precision 和 Recall。 ### 下载参考统计数据 ``` wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz ``` ### 运行评估 ``` python evaluator.py \ VIRTUAL_imagenet256_labeled.npz \ ./samples/samples_50000.npz ./samples ``` ## 模型架构 Self-Flow 模型基于 SiT-XL/2，具有以下规格 一项关键的架构修改是**逐 token 时间步条件化（per-token timestep conditioning）**，它允许每个 token 在训练期间具有不同的噪声水平。 ## 项目结构 ``` Self-Flow/ ├── sample.py # Main sampling script ├── checkpoints/ # Place model checkpoints here ├── requirements.txt # Python dependencies ├── README.md # This file └── src/ # Model and sampling implementations ├── model.py # SelfFlowPerTokenDiT model ├── sampling.py # Diffusion sampling utilities └── utils.py # Position encoding utilities ``` ## 训练详情 该模型使用以下配置进行训练： - **模型**：具有逐 token 时间步条件化的 SiT-XL/2 - **训练**：采用带有逐 token 掩码（25% 掩码率）的 Self-Flow - **优化器**：带有梯度裁剪（gradient clipping，max_norm=1）的 AdamW - **混合精度**：BFloat16 - **自蒸馏（Self-distillation））**：教师模型在第 20 层（EMA），学生模型在第 8 层 ## 致谢 此代码是基于以下项目构建的： - [REPA](https://github.com/sihyun-yu/REPA) - Representation Alignment for Generation - [SiT](https://github.com/willisma/SiT) - Scalable Interpolant Transformers ## BibTeX 如果您使用了这项工作，请引用： ``` @article{CheferEsser2026selfflow, title={Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis}, author={Hila Chefer and Patrick Esser and Dominik Lorenz and Dustin Podell and Vikash Raja and Vinh Tong and Antonio Torralba and Robin Rombach}, journal = {arXiv preprint arXiv:2603.06507}, year={2026}, } ```

标签：人工智能, 凭据扫描, 多模态合成, 扩散模型, 流匹配, 生成式AI, 用户模式Hook绕过, 索引, 自监督学习, 逆向工具