Yingqing Guo*,
Hui Yuan*,
Yukang Yang,
Minshuo Chen,
Mengdi Wang
Princeton University (*Equal Contribution)
NeurIPS 2024
Clone this repo:
git clone https://github.com/yukang123/GGDMOptim.git
cd GGDMOptim
Install required Python packages
(Optional) conda create -n guided_diffusion python=3.10
conda activate guided_diffusion
pip install -r requirements.txt
Computational Resource Requirements:
one NVIDIA 80G A100 GPU is recommended. But for the numerical simulations, the maximum GPU usage is less than 3GB; while for the image generation, the experiments could be accomodated within a smaller GPU memory (e.g., 24GB) by decreasing the batch size or image size.
Experiments on two different data distributions.
Prerequisites: Please download data (hyperparameters and dataset for pretraining) and checkpoints (pretrained checkpoints) from this huggingface repo and put them into the simulations folder (replace the current placeholder). Then
cd simulations.
Please check the main.py file to understand the meanings of different parameters. Please use tensorboard to visualize the results based on log files saved in logs folder.
(Optional) Pretraining Unconditional Diffusion Model:
python main.py --x_type linear_latent --pretrain --pre_lr 1e-4 --pre_num_episodes 20 --pretrain_data data/linear_latent/pretrain_65536.npy --pretrain_ckpt_folder checkpoints/pretrained_n
You can also generate new training data by setting --pretrain_data as None and trying different random seeds.
Comparisons between two different guidances
$G$ and$G_{loss}$ (Figure 4 of the paper)
Reward Function:
You may add --seed xxxx in the below commands (try seeds in
Algorithm 1: Gradient-Guided Diffusion for Generative Optimization
Please add --optimize to enable generative optimization. --interval means the expected increment of reward values per iteration, i.e.,
--cond_score_version v1)
python main.py --optimize --x_type linear_latent --cond_score_version v1 --func_type quadratic --pretrain_ckpt checkpoints/pretrained/linear_latent_pretrain_epoch_20.pth --generate_bs 32 --opt_rounds 1000 --interval 0.9
And you may save the generated samples by adding --save_samples.
--cond_score_version v2)
python main.py --optimize --x_type linear_latent --cond_score_version v2 --func_type quadratic --pretrain_ckpt checkpoints/pretrained/linear_latent_pretrain_epoch_20.pth --generate_bs 32 --opt_rounds 1000 --interval 0.2 --save_samples (optional)
Algorithm 2: Gradient-Guided Diffusion with Adaptive Fine-tuning
Please add --score_matching_finetune to finetune the score model with the score matching loss.
python main.py --optimize --x_type linear_latent --cond_score_version v1 --func_type quadratic --pretrain_ckpt checkpoints/pretrained/linear_latent_pretrain_epoch_20.pth --generate_bs 32 --opt_rounds 1200 --interval 0.9 --score_matching_finetune --sm_ft_lr 1e-6 --sm_ft_start_round 40 --save_samples (optional)
python main.py --optimize --x_type linear_latent --cond_score_version v2 --func_type quadratic --pretrain_ckpt checkpoints/pretrained/linear_latent_pretrain_epoch_20.pth --generate_bs 32 --opt_rounds 1200 --interval 0.2 --score_matching_finetune --sm_ft_lr 1e-6 --sm_ft_start_round 60 --save_samples (optional)
Algorithm 1 with
$G_{loss}$ for optimizing different functions (Figure 9 of the paper)
Reward Functions:
(a) --opt_rounds 1000)
Please set --func_type quadratic. When --interval 0.2 . When --use_theta_1 and set --interval 0.05.
(b) --opt_rounds 50)
Please set --func_type negnormplus.
When --interval 1.
When --use_b_1 and set --interval 1.
Please set --x_type unit_ball.
(Optional) Pretraining Unconditional Diffusion Model:
python main.py --x_type unit_ball --pretrain --pre_lr 1e-4 --pre_num_episodes 20 --pretrain_data data/unit_ball/pretrain_65536.npy --pretrain_ckpt_folder checkpoints/pretrained_n --seed 1
$G$ vs$G_{loss}$
The reward function is
python main.py --optimize --x_type unit_ball --cond_score_version v1 --pretrain_ckpt checkpoints/pretrained/unit_ball_pretrain_epoch_20.pth --generate_bs 256 --opt_rounds 25 --interval 0.2 --seed 1
python main.py --optimize --x_type unit_ball --cond_score_version v2 --pretrain_ckpt checkpoints/pretrained/unit_ball_pretrain_epoch_20.pth --generate_bs 256 --opt_rounds 25 --interval 0.2 --seed 1
Please try different --interval values in a large range (maybe 0.02-0.52) to control the guidance strength and get the
Codes are adapted from RCGDM.
Prerequisites: Please download reward_model.pth (reward model) from this huggingface repo and put them into the image_generation folder.
- cd image_generation
- python main.py --target 10 --guidance 100 --opt_steps 8 --repeat_epoch 4 --seed 5 --bs 5 --prompt 'fox'
You could tune the --target (eg., 2, 4) and --guidance to control the guidance strength, and also try out different --opt_steps to see different optimization processes. You may also specify different text prompts with detailed descriptions.
Additional Comments: The runtime of the above experiments may vary with different GPU conditions and batch sizes. But the relative increase in time cost between gradient-guided generation and unguided generation remains as shown in Table 1 of Appendix F.3 of the paper.
For help or issues about the github codes, please email Yukang Yang (yy1325@princeton.edu) and Yingqing Guo (yg6736@princeton.edu) or submit a GitHub issue.
If you find this code useful in your research, please consider citing:
@article{guo2024gradient,
title={Gradient Guidance for Diffusion Models: An Optimization Perspective},
author={Guo, Yingqing and Yuan, Hui and Yang, Yukang and Chen, Minshuo and Wang, Mengdi},
journal={arXiv preprint arXiv:2404.14743},
year={2024}
}