Diffusion-Sharpening: Fine-tuning Diffusion Models with Denoising Trajectory Sharpening
This repository provides implementation for the paper "Diffusion-Sharpeninging: Fine-tuning Diffusion Models with Denoising Trajectory Sharpening"
🔥 Features
- Trajectory-level optimization: Guides the diffusion process toward high-reward regions.
- No additional inference cost: Amortizes trajectory optimization into training.
- Generalized reward alignment: Works with any reward model, including differentiable and non-differentiable objectives.
📖 Introduction
Diffusion-Sharpeninging is a trajectory-level fine-tuning framework for diffusion models that autonomously optimizes sampling paths through reward-guided iterative refinement. By integrating self-alignment with a path integral approach, it eliminates the reliance on predefined datasets and amortizes inference costs into training, achieving superior efficiency in both convergence and inference.
🚀 Installation
To set up the environment and install dependencies, run the following:
git clone https://github.com/Gen-Verse/Diffusion-Sharpening
cd Diffusion-Sharpening
conda create -n Diffusion-Sharpening python==3.11.2
conda activate Diffusion-Sharpening
pip install -r requirements.txt📚 Repository Structure
Diffusion-Sharpening/
│── requirements.txt # Basic pip requirements
│── reward.py # Reward model implementation (CLIPScore, ImageReward, Compositional Reward, MLLM Grader)
│── mllm_template/ # Custom templates for MLLM Grader
│── train_sft_diffusion_sharpening_sdxl.py # Main script for SFT-Diffusion-Sharpening
│── train_rlhf_diffusion_sharpening_sdxl.py # Main script for RLHF-Diffusion-Sharpening
🔍 Reward Model Configuration
We utilize four reward models for Diffusion-Sharpening training, with detailed configurations provided in reward.py:
- CLIPScore – Measures text-image alignment.
- ImageReward – Evaluates human preference in image generation.
- Compositional Reward – Provides composition-aware feedback for structured image generation.
- MLLM Grader – Uses an MLLM API for overall image quality assessment. We employ GPT-4o in our experiments.
🏗️ SFT-Diffusion-Sharpening
Supervised fine-tuning method that refines the diffusion trajectory using pre-existing image-text datasets, enabling efficient optimization with any reward model. Example accelerate launch command:
# from scripts/sft.sh export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0" export DATASET_NAME="your_dataset_name" accelerate launch --main_process_port 10950 train_sft_diffusion_sharpen_sdxl.py \ --pretrained_model_name_or_path=$MODEL_NAME \ --dataset_name=$DATASET_NAME \ --enable_xformers_memory_efficient_attention \ --resolution=512 --center_crop --random_flip \ --proportion_empty_prompts=0.2 \ --train_batch_size=1 \ --gradient_accumulation_steps=4 --gradient_checkpointing \ --max_train_steps=10000 \ --use_8bit_adam \ --learning_rate=1e-06 --lr_scheduler="constant" --lr_warmup_steps=0 \ --mixed_precision="fp16" \ --report_to="wandb" \ --validation_prompt="a cute Sundar Pichai creature" --validation_epochs 5 \ --checkpointing_steps=5000 \ --output_dir="output_dir" \ --reward_model="clip" (or "image_reward", "compositional", "mllm")
If you wish to use MLLM Grader, you can set your api key in the training command with --api_key="your_api_key" and --base_url="your_base_url".
🏗️ RLHF-Diffusion-Sharpening Training Command
Online reinforcement learning method that eliminates the need for curated datasets by generating positive and negative sample pairs dynamically during training. Example accelerate launch command:
# from scripts/rlhf.sh export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0" export DATASET_NAME="your_dataset_name" accelerate launch --main_process_port 10950 train_rlhf_diffusion_sharpen_sdxl.py \ --pretrained_model_name_or_path=$MODEL_NAME \ --dataset_name=$DATASET_NAME \ --enable_xformers_memory_efficient_attention \ --resolution=512 --center_crop --random_flip \ --proportion_empty_prompts=0.2 \ --train_batch_size=1 \ --gradient_accumulation_steps=4 --gradient_checkpointing \ --max_train_steps=10000 \ --use_8bit_adam \ --learning_rate=1e-06 --lr_scheduler="constant" --lr_warmup_steps=0 \ --mixed_precision="fp16" \ --report_to="wandb" \ --validation_prompt="a cute Sundar Pichai creature" --validation_epochs 5 \ --checkpointing_steps=5000 \ --output_dir="output_dir" \ --reward_model="clip" (or "image_reward", "compositional", "mllm")
📊 Results
📜 Citation
If you find this work useful, please consider citing:
@article{tian2025sharpening, title={Diffusion-Sharpeninging: Fine-tuning Diffusion Models with Denoising Trajectory Sharpeninging}, author={Tian, Ye and Yang, Ling and Zhang, Xinchen and Tong, Yunhai and Wang, Mengdi and Cui, Bin}, journal={arXiv preprint arXiv:2502.12146}, year={2025} }