ECT: 简化一致性模型

ECT通过一种简单而有原则的方法，实现了最先进的少步生成能力。仅需极少的调优成本，ECT就展现出了令人期待的早期结果，同时受益于训练FLOPs的扩展，持续提升其少步生成能力。

尝试你自己的一致性模型！你只需要稍微微调一下。:D

⭐ 更新 ⭐

更多内容正在烘焙中。🙃

2024.04.27 - 环境升级至Pytorch 2.3.0。
2024.04.12 - ECM现在可以在CIFAR10上使用1个模型步骤超越最先进的GAN，使用2个模型步骤超越最先进的扩散模型。检查点可用。

环境

你可以运行以下命令通过conda设置Python环境。将安装Pytorch 2.3.0和Python 3.9.18。

conda env create -f env.yml

数据集

准备符合EDM格式的数据集。参考这里。

训练

运行以下命令，在1小时内使用1块A100 GPU调优你的最先进的2步ECM，并超越一致性蒸馏(CD)。

bash run_ecm_1hour.sh 1 <PORT> --desc bs128.1hour

运行以下命令，以批量大小128和200k迭代次数运行ECT。建议使用2/4个GPU。

bash run_ecm.sh <NGPUs> <PORT> --desc bs128.200k

将NGPUs和PORT替换为用于训练的GPU数量和DDP同步的端口号。

评估

运行以下命令计算预训练检查点的FID。

bash eval_ecm.sh <NGPUs> <PORT> --resume <CKPT_PATH>

生成性能

FID评估

将ECT训练的模型作为ECM，我们在CIFAR10数据集上比较ECM的无条件图像生成能力与最先进的生成模型，包括流行的扩散模型（配合高级采样器）、扩散蒸馏和一致性模型。

方法	FID	NFE	模型	参数	批量大小	计划
Score SDE	2.38	2000	NCSN++	56.4M	128	~1600k
Score SDE-deep	2.20	2000	NCSN++ (2 $\times$ 深度)	> 100M	128	~1600k
EDM	2.01	35	DDPM++	56.4M	512	400k
PD	8.34	1	DDPM++	56.4M	512	800k
Diff-Instruct	4.53	1	DDPM++	56.4M	512	800k
CD (LPIPS)	3.55	1	NCSN++	56.4M	512	800k
CD (LPIPS)	2.93	2	NCSN++	56.4M	512	800k
iCT-deep	2.51	1	NCSN++ (2 $\times$ 深度)	> 100M	1024	400k
iCT-deep	2.24	2	NCSN++ (2 $\times$ 深度)	> 100M	1024	400k
ECM (100k)	4.54	1	DDPM++	55.7M	128	100k
ECM (200k)	3.86	1	DDPM++	55.7M	128	200k
ECM (400k)	3.60	1	DDPM++	55.7M	128	400k
ECM (100k)	2.20	2	DDPM++	55.7M	128	100k
ECM (200k)	2.15	2	DDPM++	55.7M	128	200k
ECM (400k)	2.11	2	DDPM++	55.7M	128	400k

$\mathrm{FD}_{\text{DINOv2}}$ 评估

由于DINOv2能够产生与人类视觉更一致的评估，我们使用最先进的开源表示模型DINOv2的潜在空间中的Fréchet距离来评估图像保真度，记为$\mathrm{FD}_{\text{DINOv2}}$。

使用dgm-eval，我们得到了与最先进扩散模型和GAN相比的$\mathrm{FD}_{\text{DINOv2}}$。

方法	$\mathrm{FD}_{\text{DINOv2}}$	NFE
EDM	145.20	35
StyleGAN-XL	204.60	1
ECM	198.51	1
ECM	128.63	2

即使不结合GAN等其他生成机制或像分数蒸馏这样的扩散蒸馏，ECT也能够比最先进的扩散模型更快地生成高质量样本，且比最先进的扩散模型和GAN更好。

检查点

更多检查点将稍后提供。

CIFAR10 $\mathrm{FD}_{\text{DINOv2}}$ 检查点。

联系

如果你想一起训练模型，请给我发邮件至zhengyanggeng@gmail.com！！！

如果你有其他问题或对合作感兴趣，欢迎联系我。在Twitter或微信上找我。:D

引用

@misc{ect,
    title  = {Consistency Models Made Easy},
    author = {Geng, Zhengyang and Luo, William and Pokle, Ashwini and Kolter, Zico},
    year   = {2024},
    url    = {https://gsunshine.notion.site/Consistency-Models-Made-Easy-954205c0b4a24c009f78719f43b419cc?pvs=4}
}