Principle:NVIDIA NeMo Aligner DRaFT Plus Image Alignment

Knowledge Sources	NVIDIA_NeMo_Aligner
Domains	Multimodal, Image Generation, Diffusion Models, Alignment
Last Updated	2026-02-08 00:00 GMT

Overview

DRaFT+ (Differentiable Reward Fine-Tuning Plus) aligns diffusion models with human preferences by backpropagating through a reward model during the denoising process, using truncated backpropagation through the final diffusion steps and KL divergence to prevent mode collapse.

Description

DRaFT+ is a method for fine-tuning text-to-image diffusion models (such as Stable Diffusion and SDXL) to better align with human aesthetic preferences. It works by:

1. Differentiable reward optimization -- During training, the model generates images through its standard denoising process. The generated images are passed through a reward model (e.g., a CLIP-based model trained on human preference data), and the reward signal is backpropagated through the diffusion process to update the model parameters.

1. Truncated backpropagation -- Computing gradients through the entire denoising chain (e.g., 50 DDIM steps) would be prohibitively expensive. DRaFT+ uses truncated backpropagation through only the final truncation_steps of the denoising process. Earlier steps are run with torch.no_grad(), dramatically reducing memory and compute requirements while still providing meaningful gradient signal.

1. KL divergence regularization -- To prevent the fine-tuned model from collapsing to a narrow mode of high-reward images, DRaFT+ computes a Gaussian KL penalty between the noise predictions (epsilon) of the fine-tuned model and those of the original (frozen) base model. This penalty, weighted by kl_coeff, ensures the fine-tuned model stays close to the original distribution.

1. Annealed sampling -- At inference time, DRaFT+ supports annealed guidance, which interpolates between the base model's score function and the fine-tuned model's score function at each denoising step using a configurable weighing function (linear, power, step, etc.). This enables smooth control over the alignment-diversity tradeoff.

The framework supports both Stable Diffusion (SD) via MegatronSDDRaFTPModel and Stable Diffusion XL (SDXL) via MegatronSDXLDRaFTPModel, as well as full fine-tuning and PEFT (LoRA) approaches.

Usage

DRaFT+ is used when:

• You want to improve the aesthetic quality of a text-to-image diffusion model according to human preferences.
• You have a trained reward model (e.g., PickScore, HPSv2) that scores image-text alignment quality.
• You want to fine-tune while preventing mode collapse via KL regularization.
• You want to control the alignment-diversity tradeoff at inference time via annealed sampling.

Theoretical Basis

Training Objective: ${ℒ}=-{\mathbb{𝔼}}_{x_T\sim {𝒩}(0,I),c\sim D}[R(G_{\theta }(x_T,c),c)]+\lambda \cdot D_{\text{KL}}$

where $R$ is the reward model, $G_{\theta }$ is the diffusion generator, $c$ is the text conditioning, and $\lambda$ is the KL coefficient.

Gaussian KL Penalty (shared variance): $D_{\text{KL}}=\frac{1}{d}\sum _{t\in \text{trunc}}\Vert {ϵ}_{\theta }(x_t,t,c)-{ϵ}_{{\theta }_0}(x_t,t,c){\Vert }^2$

where $d$ is the dimensionality of the noise prediction, ${ϵ}_{\theta }$ is the fine-tuned model's noise prediction, and ${ϵ}_{{\theta }_0}$ is the frozen base model's noise prediction. The sum is over the truncated denoising steps.

Annealed Sampling:

At inference step $i$ out of $T$ total steps, the combined noise prediction is: ${ϵ}_{\text{combined}}=(1-w(i,T))\cdot {ϵ}_{{\theta }_0}+w(i,T)\cdot {ϵ}_{\theta }$

where $w(i,T)$ is a weighing function such as:

• Linear: $w(i,T)=i/T$
• Power: $w(i,T)=(i/T{)}^p$
• Step: $w(i,T)=\mathbb{𝟙}[i/T\ge \tau ]$

Related Pages

• Implementation:NVIDIA_NeMo_Aligner_MegatronSD_DRaFTP_Model
• Implementation:NVIDIA_NeMo_Aligner_MegatronSDXL_DRaFTP_Model