RCG: 一种突破性的无条件图像生成方法

rcg

RCG: 无条件图像生成的新突破

在计算机视觉和机器学习领域,图像生成一直是一个极具挑战性的任务。近年来,各种条件生成模型取得了巨大进展,但无条件图像生成的性能一直落后。最近,来自麻省理工学院和Facebook AI Research的研究人员提出了一种名为RCG(Representative Conditional Generation)的新方法,在无条件图像生成方面取得了突破性进展。

RCG方法简介

RCG是一种自监督的表示生成方法,其核心思想是将无条件生成问题转化为条件生成问题。具体来说,RCG首先训练一个表示扩散模型(Representation Diffusion Model, RDM)来学习图像的高级语义表示,然后基于这些表示来条件化生成像素级图像。这种方法巧妙地利用了条件生成模型的优势,同时保持了生成过程的无条件性。

如上图所示,RCG的工作流程可以分为以下几个步骤:

使用预训练的视觉编码器(如MoCo v3)提取图像的高级语义表示。
训练RDM来学习这些表示的分布。
在生成阶段,首先从RDM采样得到表示,然后使用像素生成器(如MAGE)基于这些表示生成最终的图像。

这种设计使得RCG能够生成具有丰富语义信息和高视觉质量的图像,同时不依赖于任何类别标签或其他外部条件。

性能突破

RCG在ImageNet 256x256数据集上的无条件生成任务中取得了令人瞩目的成果。如下图所示,RCG的性能大幅超越了之前的无条件生成方法,甚至接近或超过了一些条件生成模型:

在FID(Fréchet Inception Distance)和IS(Inception Score)这两个广泛使用的图像质量评估指标上,RCG都取得了最优或接近最优的结果。特别是在使用类别引导采样(Classifier-Free Guidance, CFG)技术时,RCG的表现更加出色,FID可达2.15,IS高达253.4。这些数据表明,RCG生成的图像在多样性和真实性方面都达到了非常高的水平。

实现细节

RCG的实现主要包括两个关键组件:表示扩散模型(RDM)和像素生成器。研究者提供了多种不同配置的预训练模型:

RDM:
- 基于MoCo v3 ViT-B的无条件RDM
- 基于MoCo v3 ViT-L的无条件RDM
- 基于MoCo v3 ViT-B的类条件RDM
- 基于MoCo v3 ViT-L的类条件RDM
像素生成器:
- MAGE-B和MAGE-L
- DiT-XL/2
- ADM
- LDM-8

这些模型在不同的硬件配置和训练时长下都能达到优秀的性能,为研究者提供了丰富的选择。

生成示例

RCG生成的图像质量非常高,无论是在无条件生成还是类条件生成任务中都表现出色。以下是一些生成示例:

上图展示了RCG在无条件生成任务中的表现。可以看到,生成的图像涵盖了各种物体、动物和场景,细节丰富,结构合理,显示出极高的多样性和真实性。

这张图则展示了RCG在类条件生成任务中的能力。每一行代表一个特定的类别,生成的图像不仅符合类别特征,还展现出丰富的变化和细节。

技术细节与实现

环境配置

要复现RCG的结果,需要按照以下步骤配置环境:

克隆代码仓库:

git clone https://github.com/LTH14/rcg.git
cd rcg

创建并激活conda环境:

conda env create -f environment.yaml
conda activate rcg

下载预训练模型:
- VQGAN tokenizer
- MoCo v3 ViT-B和ViT-L编码器

训练过程

RCG的训练分为两个主要阶段:RDM训练和像素生成器训练。

RDM训练

使用4个V100 GPU训练MoCo v3 ViT-B表示扩散模型:

python -m torch.distributed.launch --nproc_per_node=4 --nnodes=1 --node_rank=0 \
main_rdm.py \
--config config/rdm/mocov3vitb_simplemlp_l12_w1536.yaml \
--batch_size 128 --input_size 256 \
--epochs 200 \
--blr 1e-6 --weight_decay 0.01 \
--output_dir ${OUTPUT_DIR} \
--data_path ${IMAGENET_DIR} \
--dist_url tcp://${MASTER_SERVER_ADDRESS}:2214

像素生成器训练

以MAGE-B为例,使用64个V100 GPU训练200个epoch:

python -m torch.distributed.launch --nproc_per_node=8 --nnodes=8 --node_rank=0 \
main_mage.py \
--pretrained_enc_arch mocov3_vit_base \
--pretrained_enc_path pretrained_enc_ckpts/mocov3/vitb.pth.tar --rep_drop_prob 0.1 \
--use_rep --rep_dim 256 --pretrained_enc_withproj --pretrained_enc_proj_dim 256 \
--pretrained_rdm_cfg ${RDM_CFG_PATH} --pretrained_rdm_ckpt ${RDM_CKPT_PATH} \
--rdm_steps 250 --eta 1.0 --temp 6.0 --num_iter 20 --num_images 50000 --cfg 0.0 \
--batch_size 64 --input_size 256 \
--model mage_vit_base_patch16 \
--mask_ratio_min 0.5 --mask_ratio_max 1.0 --mask_ratio_mu 0.75 --mask_ratio_std 0.25 \
--epochs 200 \
--warmup_epochs 10 \
--blr 1.5e-4 --weight_decay 0.05 \
--output_dir ${OUTPUT_DIR} \
--data_path ${IMAGENET_DIR} \
--dist_url tcp://${MASTER_SERVER_ADDRESS}:2214

评估与可视化

训练完成后,可以使用以下命令评估模型性能:

python -m torch.distributed.launch --nproc_per_node=8 --nnodes=8 --node_rank=0 \
main_mage.py \
--resume ${CHECKPOINT_PATH} \
--evaluate \
...其他参数与训练时相同...