🚀🚀🚀 Improve Diffusion Image Generation Quality using Levenberg-Marquardt-Langevin

We introduce LML, an accelerated sampler for diffusion models leveraging the second-order Hessian geometry. Our LML imlpementation is completely compatible with the diffusers.

This repository is the official implementation of the ICCV 2025 paper: "Unleashing High-Quality Image Generation in Diffusion Sampling Using Second-Order Levenberg-Marquardt-Langevin"

Fangyikang Wang^1,2, Hubery Yin², Lei Qian¹, Yinan Li¹, Shaobin Zhuang^3,2, Huminhao Zhu¹, Yilin Zhang¹, Yanlong Tang⁴, Chao Zhang¹, Hanbin Zhao¹, Hui Qian¹, Chen Li²

¹Zhejiang University ²WeChat Vision, Tencent Inc ³Shanghai Jiao Tong University ⁴Tencent Lightspeed Studio

The intuition of our LML diffusion sampler

Schematic comparison between our LML method and baselines. While previous works mainly focus on intriguing designs along the annealing path to improve diffusion sampling, they leave operations at specific noise levels to be performed using first-order Langevin. Our approach proposes to leverage the Levenberg-Marquardt approximated Hessian geometry to guide the Langevin update to be more accurate.

The relation between optimization algorithms and MCMC sampling algorithms. We initially wanted to develop a diffusion sampler utilizing Hessian geometry, following the path of Newton-Langevin dynamics. However, this approach proved to be highly computationally expensive within the DM context. Drawing inspiration from the Levenberg-Marquardt method used in optimization, our method incorporates low-rank approximation and damping techniques. This enables us to obtain the Hessian geometry in a computationally affordable manner. Subsequently, we use this approximated Hessian geometry to guide the Langevin updates.

👨🏻‍💻 Run the code

1) Get start

Python 3.8.12
CUDA 11.7
NVIDIA A100 40GB PCIe
Torch 2.0.0
Torchvision 0.14.0

Please follow diffusers to install diffusers.

2) Sampling

first, please switch to the root directory.

CIFAR-10 sampling

For baseline, you can do CIAFR-10 sampling as follows, choose sampler_type within [ddim, pndm, dpm, dpm++, unipc]:

python3 ./scripts/cifar10.py --test_num 1 --batch_size 1 --num_inference_steps 10  --save_dir YOUR/SAVE/DIR --model_id xx/xx/ddpm_ema_cifar10 --sampler_type ddim

For our LML sampler, there is an additional $\lambda$ hyperparameter:

python3 ./scripts/cifar10.py --test_num 1 --batch_size 1 --num_inference_steps 10  --save_dir YOUR/SAVE/DIR --model_id xx/xx/ddpm_ema_cifar10 --sampler_type dpm_lm --lamb 0.0008

For the optimal choice of LML, we have:

	5 NFEs	6 NFEs	7 NFEs	8 NFEs	9 NFEs	10 NFEs	12 NFEs	15 NFEs	20 NFEs	30 NFEs	50 NFEs	100 NFEs
optimal value of lamb	0.0008	0.0008	0.001	0.001	0.001	0.0008	0.001	0.001	0.0005	0.0003	0.0001	0.00005

CelebA-HQ sampling

For baseline:

python3 ./scripts/celeba.py --test_num 1 --batch_size 1 --num_inference_steps 10  --save_dir YOUR/SAVE/DIR --model_id xx/xx/ldm-celebahq-256 --sampler_type ddim

For our LML:

python3 ./scripts/celeba.py --test_num 1 --batch_size 1 --num_inference_steps 10  --save_dir YOUR/SAVE/DIR --model_id xx/xx/ldm-celebahq-256 --sampler_type ddim_lm --lamb 0.005

SD-15 and SD-2b on MS-COCO sampling

python3 ./scripts/StableDiffusion_COCO.py --test_num 30002 --num_inference_steps 10  --save_dir YOUR/SAVE/DIR --model_id xx/xx/stable-diffusion-v1-5 --sampler_type dpm_lm --lamb 0.001

For the optimal choice of LML on MS-COCO, for NFEs of {5, 6, 7, 8, 9, 10, 12, 15}, we always choose $\lambda = 0.001$:

SD-15, SD-2b, SD-XL, and PixArt-$\alpha$ on T2i-compbench sampling

Before running the scripts, make sure to clone T2I-CompBench repository. Generated images are stored in the directory "save_dir/model/dataset_category/sampler_type/samples".

For baseline, you can do T2i-compbench sampling as follows, choose sampler_type within [ddim, pndm, dpm, dpm++, unipc] and model within [sd15, sd2_base, sdxl, pixart]:

python3 ./scripts/StableDiffusion_PixArt_T2i_Sampling.py --dataset_category color --dataset_path PATH/TO/T2I-COMPBENCH --test_num 10 --num_inference_steps 10 --model_dir YOUR/MODEL/DIR --save_dir YOUR/SAVE/DIR --model sd15 --sampler_type ddim

For our LML sampler, there is an additional $\lambda$ hyperparameter:

python3 ./scripts/StableDiffusion_PixArt_T2i_Sampling.py --dataset_category color --dataset_path PATH/TO/T2I-COMPBENCH --test_num 10 --num_inference_steps 10 --model_dir YOUR/MODEL/DIR --save_dir YOUR/SAVE/DIR --model sd15 --sampler_type dpm_lm --lamb 0.006

Use our LML diffusion sampler with ControlNet

canny

python3 ./scripts/control_net_canny.py --num_inference_steps 10  --original_image_path /xxx/xxx/data/input_image_vermeer.png --controlnet_dir /xxx/xxx/sd-controlnet-canny --sd_dir /xxx/xxx/stable-diffusion-v1-5  --save_dir YOUR/SAVE/DIR  --sampler_type dpm_lm --lamb 0.001

depth

python3 ./scripts/control_net_depth.py --num_inference_steps 10  --controlnet_dir /xxx/xxx/control_v11f1p_sd15_depth --sd_dir /xxx/xxx/stable-diffusion-v1-5  --save_dir YOUR/SAVE/DIR  --sampler_type dpm_lm --lamb 0.001

pose

python3 ./scripts/control_net_canny.py --num_inference_steps 10 --controlnet_dir /xxx/xxx/sd-controlnet-openpose --sd_dir /xxx/xxx/stable-diffusion-v1-5  --save_dir YOUR/SAVE/DIR  --sampler_type dpm_lm --lamb 0.001

LML sampling on FLUX

For baseline:

python3 ./scripts/FLUX_T2i_Sampling.py --dataset_category color --dataset_path PATH/TO/T2I-COMPBENCH --test_num 10 --num_inference_steps 10 --model_id YOUR/MODEL/DIR --save_dir YOUR/SAVE/DIR --sampler_type fm_euler

For our LML:

python3 ./scripts/FLUX_T2i_Sampling.py --dataset_category color --dataset_path PATH/TO/T2I-COMPBENCH --test_num 10 --num_inference_steps 10 --model_id YOUR/MODEL/DIR --save_dir YOUR/SAVE/DIR --sampler_type lml_euler --lamb 0.01

3) Evaluation

FID evaluation on CIFAR-10

[Coming Soon] ⏳
FID evaluation on MS-COCO

[Coming Soon] ⏳

T2I-compbench evaluation

Please refer to the T2I-CompBench guide. Create a new environment and install the dependencies for T2I-CompBench evaluation. For testing combinations of multiple models and samplers, we also provide a convenient one-click script. Place the script file in the corresponding directory of T2I-CompBench to replace the origin script. For example:

# BLIP-VQA for Attribute Binding
cd T2I-CompBench
bash BLIPvqa_eval/test.sh
      ||
      ||
      \/
cp evaluations/T2I-CompBench/BLIPvqa_test.sh T2I-CompBench/BLIPvqa_eval
cd T2I-CompBench
bash BLIPvqa_eval/BLIPvqa_test.sh 'save_dir'

The directory structure of 'save_dir' should satisfy the following format:

{save_dir}/model/dataset_category/sampler_type/samples/
                                              ├── a green bench and a blue bowl_000000.png
                                              ├── a green bench and a blue bowl_000001.png
                                              └──...

🪪 License

This project is licensed under the MIT License - see the LICENSE file for details.

📝 Citation

If our work assists your research, feel free to give us a star ⭐ or cite us using:

@article{wang2025unleashing,
  title={Unleashing High-Quality Image Generation in Diffusion Sampling Using Second-Order Levenberg-Marquardt-Langevin},
  author={Wang, Fangyikang and Yin, Hubery and Qian, Lei and Li, Yinan and Zhuang, Shaobin and Zhu, Huminhao and Zhang, Yilin and Tang, Yanlong and Zhang, Chao and Zhao, Hanbin and others},
  journal={arXiv preprint arXiv:2505.24222},
  year={2025}
}

📩 Contact me

Our e-mail address:

[email protected], [email protected], [email protected]