节约每一分钱 从零开始的微预算扩散培训

摘要：随着生成式人工智能的扩展规律推动性能的提高，它们同时将这些模型的开发集中在拥有大量计算资源的参与者身上。我们以文本到图像（T2I）生成模型为重点，旨在通过展示大规模 T2I 扩散变换器模型的极低成本训练来打破这一瓶颈。由于变换器的计算成本随每幅图像中的斑块数量而增加，我们建议在训练过程中随机屏蔽多达 75% 的图像斑块。我们提出了一种延迟掩蔽策略，即在掩蔽前使用斑块混合器对所有斑块进行预处理，从而大大降低了掩蔽带来的性能下降，在降低计算成本方面优于模型降维。我们还结合了变换器架构的最新改进，如使用专家混合层，以提高性能，并进一步确定在微预算训练中使用合成图像的关键优势。最后，我们仅使用 3700 万张公开的真实和合成图像，就训练出了 11.6 亿个参数的稀疏变换器，经济成本仅为 1890 美元，并在 COCO 数据集上实现了 12.7 FID 的零镜头生成。值得注意的是，我们的模型在自动评估和以人为中心的评估中都取得了具有竞争力的性能，并实现了高质量的生成，同时成本比稳定扩散模型低 118 倍，比目前最先进的方法（成本为 28,400 美元）低 14 倍。我们还进一步研究了合成图像对性能的影响，并证明仅在合成图像上进行微预算训练就足以生成高质量数据。

Prompt: ‘Image of an astronaut riding a horse in {} style’. Styles: Origami, Pixel art, Line art, Cyberpunk, Van Gogh Starry Night, Animation, Watercolor, Stained glass

我们提供了四个预训练模型的检查点。 下表介绍了每个模型及其量化性能。

Model Description	VAE (channels)	FID	GenEval Score	Model's filename
MicroDiT_XL_2 trained on 22M real images	SDXL-VAE (4 channel)	12.72	0.46	dit_4_channel_22M_real_only_data.pt
MicroDiT_XL_2 trained on 37M images (22M real, 15 synthetic)	SDXL-VAE (4 channel)	12.66	0.46	dit_4_channel_37M_real_and_synthetic_data.pt
MicroDiT_XL_2 trained on 37M images (22M real, 15 synthetic)	Ostris-VAE (16 channel)	13.04	0.40	dit_16_channel_37M_real_and_synthetic_data.pt
MicroDiT_XL_2 trained on 490M synthetic images	SDXL-VAE (4 channel)	13.26	0.52	dit_4_channel_0.5B_synthetic_data.pt

图像生成： 这些检查点可与用于生成图像的 micro_diffusion 官方代码库一起使用。 首先将 micro_diffusion 代码作为 python 软件包安装 `pip install git+https://github.com/SonyResearch/micro_diffusion.git

接下来，使用以下简单步骤从最终模型生成分辨率为 512×512 的图像。

from micro_diffusion.models.model import create_latent_diffusion
model = create_latent_diffusion(latent_res=64, in_channels=4, pos_interp_scale=2.0).to('cuda')
model.dit.load_state_dict(torch.load(ckpt_path_on_local_disk))
gen_images = model.generate(prompt=['An elegant squirrel pirate on a ship']*4, num_inference_steps=30, 
                                    guidance_scale=5.0, seed=2024)

训练管道：所有四个模型的训练配置和计算预算几乎完全相同。我们从低分辨率到高分辨率逐步训练每个模型。我们首先在 256×256 分辨率的图像上对模型进行 280K 步的训练，然后在 512×512 分辨率的图像上对模型进行 55K 步的微调。端到端模型在 8×H100 机器上的训练时间估计为 2.6 天。我们的 MicroDiT 模型默认在骨干变压器架构之前使用补丁混合器。使用补丁混合器可显著减少屏蔽带来的性能下降，同时大大缩短训练时间。我们在两种分辨率的补丁混合器后屏蔽了 75% 的补丁。在使用遮罩进行训练后，我们使用 0 的遮罩率进行后续微调，以略微提高性能。

这些模型根据 Apache 2.0 许可发布。

Checkpoint：https://huggingface.co/VSehwag24/MicroDiT
Code：https://github.com/SonyResearch/micro_diffusion