Fine-Tuning Diffusion Models via Intermediate Distribution Shaping

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Anil, Gautham Govind, Haque, Shaan Ul, Kannen, Nithish, Nagaraj, Dheeraj, Shakkottai, Sanjay, Shanmugam, Karthikeyan
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866910038513156096
author Anil, Gautham Govind
Haque, Shaan Ul
Kannen, Nithish
Nagaraj, Dheeraj
Shakkottai, Sanjay
Shanmugam, Karthikeyan
author_facet Anil, Gautham Govind
Haque, Shaan Ul
Kannen, Nithish
Nagaraj, Dheeraj
Shakkottai, Sanjay
Shanmugam, Karthikeyan
contents Diffusion models are widely used for generative tasks across domains. Given a pre-trained diffusion model, it is often desirable to fine-tune it further either to correct for errors in learning or to align with downstream applications. Towards this, we examine the effect of shaping the distribution at intermediate noise levels induced by diffusion models. First, we show that existing variants of Rejection sAmpling based Fine-Tuning (RAFT), which we unify as GRAFT, can implicitly perform KL regularized reward maximization with reshaped rewards. Motivated by this observation, we introduce P-GRAFT to shape distributions at intermediate noise levels and demonstrate empirically that this can lead to more effective fine-tuning. We mathematically explain this via a bias-variance tradeoff. Next, we look at correcting learning errors in pre-trained flow models based on the developed mathematical framework. In particular, we propose inverse noise correction, a novel algorithm to improve the quality of pre-trained flow models without explicit rewards. We empirically evaluate our methods on text-to-image(T2I) generation, layout generation, molecule generation and unconditional image generation. Notably, our framework, applied to Stable Diffusion v2, improves over policy gradient methods on popular T2I benchmarks in terms of VQAScore and shows an $8.81\%$ relative improvement over the base model. For unconditional image generation, inverse noise correction improves FID of generated images at lower FLOPs/image.
format Preprint
id arxiv_https___arxiv_org_abs_2510_02692
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Fine-Tuning Diffusion Models via Intermediate Distribution Shaping
Anil, Gautham Govind
Haque, Shaan Ul
Kannen, Nithish
Nagaraj, Dheeraj
Shakkottai, Sanjay
Shanmugam, Karthikeyan
Machine Learning
Artificial Intelligence
Diffusion models are widely used for generative tasks across domains. Given a pre-trained diffusion model, it is often desirable to fine-tune it further either to correct for errors in learning or to align with downstream applications. Towards this, we examine the effect of shaping the distribution at intermediate noise levels induced by diffusion models. First, we show that existing variants of Rejection sAmpling based Fine-Tuning (RAFT), which we unify as GRAFT, can implicitly perform KL regularized reward maximization with reshaped rewards. Motivated by this observation, we introduce P-GRAFT to shape distributions at intermediate noise levels and demonstrate empirically that this can lead to more effective fine-tuning. We mathematically explain this via a bias-variance tradeoff. Next, we look at correcting learning errors in pre-trained flow models based on the developed mathematical framework. In particular, we propose inverse noise correction, a novel algorithm to improve the quality of pre-trained flow models without explicit rewards. We empirically evaluate our methods on text-to-image(T2I) generation, layout generation, molecule generation and unconditional image generation. Notably, our framework, applied to Stable Diffusion v2, improves over policy gradient methods on popular T2I benchmarks in terms of VQAScore and shows an $8.81\%$ relative improvement over the base model. For unconditional image generation, inverse noise correction improves FID of generated images at lower FLOPs/image.
title Fine-Tuning Diffusion Models via Intermediate Distribution Shaping
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2510.02692