One-sentence Summary

Researchers from Google, Google DeepMind, and Gatsby UCL propose learning classifier-free guidance schedules as a state-dependent function of diffusion time, conditioning, and the current noisy sample, framing the problem as density ratio estimation in which a discriminator estimates time-dependent log-density ratios and a lightweight generator predicts optimal guidance scales, and their approach outperforms both heuristic CFG schedules and prior methods for learning dynamic guidance on text-to-image generation benchmarks.

Key Contributions

• The paper presents a learned guidance schedule for text-to-image diffusion models that predicts an optimal guidance scale as a function of diffusion time, conditioning, and the current noisy sample, addressing the suboptimality of static and hand-designed classifier-free guidance schedules.
• The proposed method frames dynamic guidance as density ratio estimation and enforces a marginal consistency condition along the diffusion sampling trajectory; a discriminator estimates the time-dependent log-density ratio between true and guided marginal distributions while a lightweight MLP generator predicts the schedule, keeping inference overhead negligible.
• Experiments on MS-COCO 512x512 show that the learned scheduler outperforms constant CFG, heuristic dynamic schedules, and prior learned guidance approaches on human preference and text alignment metrics.

Introduction

Diffusion and flow matching models are leading generative approaches for high-fidelity image and video generation, with classifier-free guidance (CFG) playing a central role in practical text-to-image quality. However, constant CFG weights often cause artifacts and over-saturation, while heuristic dynamic schedules require dataset-specific tuning and remain agnostic to prompt complexity. Prior learned guidance schedules can adapt to conditioning and timestep, but a self-consistency-based approach improved sample quality without consistently surpassing constant CFG on text alignment. The authors instead enforce a weaker marginal consistency condition using an adversarial GAN-style framework: a discriminator estimates density ratios between true and guided marginals, and a generator learns timestep-, conditioning-, and sample-dependent guidance scales, yielding better text alignment and image quality than constant and prior learned guidance baselines.

Method

The authors propose a method to learn dynamic guidance weights for diffusion sampling by formulating the problem as a generative adversarial network based on density ratio estimation. The core goal is to enforce marginal consistency, ensuring that the distribution of guided particles matches the true data distribution at any given timestep.

Instead of using a static scalar for Classifier-Free Guidance, the authors parametrize the guidance weights $\omega$ω as a function of the source and target timesteps $(s, t)$(s,t), the conditioning signal $c$c, and the current noisy sample $x_s$xs​. This allows the model to finely adjust guidance strength along specific noise trajectories. To enforce marginal consistency $p_t^{s,\omega}(x_t) \approx p_t(x_t)$pts,ω​(xt​)≈pt​(xt​), they minimize the Kullback-Leibler divergence between the guided marginal distribution and the true marginal distribution. This divergence is rewritten in terms of the density ratio $r_t^\omega(x_t) = p_t(x_t) / p_t^{s,\omega}(x_t)$rtω​(xt​)=pt​(xt​)/pts,ω​(xt​).

To estimate this density ratio and optimize the guidance weights, the authors employ a discriminator and a generator. The discriminator $d_\phi(x_t, s, t, c)$dϕ​(xt​,s,t,c) is a convolutional network trained to distinguish between true marginal samples $x_t^{\text{real}} \sim p_t(x_t)$xtreal​∼pt​(xt​) and guided marginal samples $x_t^{\text{fake}}(\omega) \sim p_t^{s,\omega}(x_t)$xtfake​(ω)∼pts,ω​(xt​). It is optimized using a standard binary cross-entropy loss augmented with an $R_1$R1​ gradient penalty to ensure smooth ratio estimation. At convergence, the optimal discriminator logit directly recovers the log-density ratio.

The generator, implemented as a lightweight multi-layer perceptron, predicts non-negative guidance weights. To make the prediction computationally efficient and robust, the network does not take raw high-dimensional latents as input. Instead, it operates on a vector of summary statistics, including the log norm of the conditional velocity, the cosine similarity between the current latent and the guidance direction, the log distance to the data manifold, and Fourier embeddings of the timesteps. The generator is trained to minimize the discriminator output, supplemented by an $L_2$L2​ stabilizing penalty and a reward-based loss to push for better alignment.

The training strategy follows an adversarial paradigm using the Two-Time-Scale Update Rule, alternating updates between the discriminator and the generator. A crucial aspect of the training loop is the use of independent image-conditioning pairs drawn from the dataset. One pair is used to generate the real marginal target via the forward flow, while a completely independent pair is used to define the fake guided sample. This independent sampling prevents the discriminator from trivially memorizing the noise structure of a single trajectory, forcing it to evaluate whether the global distribution of guided particles matches the true data distribution at the specific timestep.

Experiment

The experiments evaluate a learned dynamic guidance schedule for text-to-image generation on MS-COCO at 512x512 resolution using frozen MMDiT models at 740M and 1.05B parameters, with baselines including unguided generation, constant CFG, heuristic schedules, and prior consistency-based learned guidance. The main results show that the proposed adversarial marginal-consistency approach improves text alignment, aesthetic quality, and human preference scores across both model sizes, while slightly increasing FID as a trade-off. Ablations validate that conditioning the guidance network on the noisy latent improves performance at a modest FID cost and that results are robust to the time-distribution parameter, with further studies on sampling steps, architecture, and noise. The authors conclude that the scheduler is effective but training-intensive and backbone-specific, limiting zero-shot transfer to new models.

On the small MMDiT-XS model, guidance improves CLIP, aesthetic, HPSv2, and PickScore over unguided generation. The adversarial marginal-consistency schedule is reported to achieve the highest CLIP and aesthetic scores, while the LIG heuristic reaches the lowest FID at the cost of lower alignment. These alignment gains transfer to the larger MMDiT-S model, with a modest FID increase attributed to adversarial and CLIP-reward optimization. Unguided generation trails all guided methods on alignment and human-preference metrics. The learned marginal-consistency schedule achieves the highest CLIP and aesthetic scores, while LIG has the lowest FID, showing a trade-off between realism and prompt alignment.

The marginal-consistency adversarial schedule improves text-to-image alignment and human preference metrics on the larger MMDiT-S model while incurring only a modest FID increase. Gains transfer from the smaller model, indicating that the learned dynamic guidance scales with model capacity. The trade-off reflects a shift toward semantic and aesthetic features over low-level Inception-v3 statistics. The learned schedule outperforms constant guidance and manually tuned heuristic schedules on CLIP, aesthetic, and HPSv2 metrics. The method introduces a modest FID increase, attributed to adversarial training and CLIP-reward optimization prioritizing semantic alignment and visual appeal over Inception-v3 statistics.

Across both MMDiT-XS and MMDiT-S, providing the guidance network with x_s improves CLIP, Aesthetic, HPSv2, and PickScore relative to conditioning only on s, t, and c. However, this added information consistently worsens FID for both model sizes. The results show a consistent trade-off between fidelity measured by FID and text alignment or human preference metrics. Adding x_s to the guidance network consistently improves CLIP, Aesthetic, HPSv2, and PickScore for both model sizes. The same conditioning change consistently increases FID, indicating a possible fidelity trade-off. The trend is consistent across MMDiT-XS and MMDiT-S.

MMDiT-XS performance remains broadly stable across δ values from 0.01 to 0.20. The largest δ tested gives the best FID, while the middle δ gives the best HPSv2 and near-best CLIP and aesthetic scores. PickScore shows only a slight decline as δ increases. Performance is largely insensitive to δ across FID, CLIP, aesthetic, HPSv2, and PickScore. δ=0.20 attains the lowest FID among the tested values, while δ=0.10 attains the highest HPSv2. CLIP and aesthetic scores remain almost flat across the tested δ values, and PickScore declines slightly as δ increases.

These experiments evaluate guided generation with MMDiT-XS and MMDiT-S models, comparing unguided outputs, constant guidance, learned marginal-consistency schedules, and heuristic schedules. Guidance consistently improves CLIP, aesthetic, HPSv2, and PickScore over unguided generation, and the learned adversarial marginal-consistency schedule transfers these gains to the larger model with only a modest FID increase. Supplying the guidance network with x_s further improves alignment metrics but worsens FID across both model sizes, showing a consistent fidelity versus alignment trade-off. Results are largely insensitive to the tested δ values, with the largest δ giving the best FID and an intermediate δ giving the best HPSv2.