Variance-Tilted Diffusion Models for Diverse Sampling

Abstract: Diffusion models are typically sampled independently, even when the downstream objective is to obtain a diverse set of candidates. We introduce a variance-weighted batch distribution that favours collections of samples with large empirical spread after a prescribed linear feature map. The target is specified explicitly, and the sampler is derived as the corresponding Doob $h$-transform of independent diffusion dynamics. The resulting correction has a compact form: an interaction term that repels posterior denoised means, together with a curvature term that moves particles to the region of higher feature variance. This yields an interacting-particle sampler with a transparent probabilistic target rather than a heuristic repulsive drift.

• The paper introduces a principled variance-tilted diffusion framework that explicitly models batch diversity via a variance-weighted target distribution and Doob h-transform corrections.
• It details an analytic derivation of score corrections that combine repulsion of denoised means with a curvature term to drive diversity in the feature space.
• Experimental results using Stable Diffusion v1.5 show that the enhanced sampler generates more diverse outputs, with trade-offs in image fidelity and increased computational cost.

Variance-Tilted Diffusion Models: Explicit Sampling for Enhanced Diversity

Motivation and Problem Statement

Diffusion models have yielded substantial progress in generative modeling across modalities, especially for image synthesis and scientific applications. However, standard practice typically involves i.i.d. sampling, which is suboptimal for scenarios where diversity within a batch is prized (e.g., drug design, motion planning), as it fails to capture a spread of candidates across the manifold. Previous strategies introduced heuristic repulsive drifts or particle guidance mechanisms to encourage diversity but resulted in implicit, often intractable target distributions. The reviewed paper proposes a principled target-first approach: define an explicit batch distribution favoring variance in a prescribed feature space and derive the corresponding sampler analytically.

Methodological Formulation

Variance-Weighted Batch Distribution

The work introduces a variance-weighted target batch distribution, defined as:

$$T_A(x) \propto \operatorname{Var}_A(x) \prod_{i=1}^n p_0(x^{(i)}),$$

where $A$ is a linear operator mapping samples to a task-relevant feature space, and $\operatorname{Var}_A(x)$ is the empirical variance of outputs after applying $A$, enforcing batch diversity in this space. This explicit formulation contrasts with prior heuristic repulsion and enables controlled sampling from the desired distribution.

Doob h-Transform and Score Correction

The batch target is realized by the Doob h-transform, which adjusts the reverse diffusion score via:

$$\operatorname{Score}_{\text{target}}(x) = s_t(x) + \nabla_x \log h_f(x),$$

where $h_f(x)$ is the expected terminal feature variance conditioned on the noisy latent. Tweedie's identities provide closed-form expressions for conditional means and covariances, enabling analytical computation of $h_f(x)$ and its gradient. The correction decomposes naturally into:

• Repulsion of posterior denoised means: Promotes empirical feature spread by repelling particles post-denoising, measured in the feature space.
• Curvature/Laplacian term: Moves particles toward regions with higher posterior feature variability, operationalized via divergence or trace estimators and providing deeper diversity encouragement beyond local mean spread.

Implementation: Efficient Sampler

Sampling proceeds via discretized reverse-time SDEs, with each particle receiving both the base score and the Doob-based correction at each step. The leading correction term is computed efficiently using vector-Jacobian products, and the curvature term is approximated via finite difference methods, Hutchinson estimators, or autograd-based higher-order derivatives. The algorithm is compatible with standard diffusion architectures and score models, including classifier-free guidance.

Experimental Results

Using Stable Diffusion v1.5 as a backbone, two major configurations are evaluated:

• VT (Variance-Tilted): Uses only denoised-mean repulsion.
• VT + Divergence: Includes the curvature term for enhanced diversity.

Qualitative assessments with diverse prompts demonstrate that the VT-based samplers generate batches spanning a broader range of poses, backgrounds, and styles compared to standard CFG sampling. The inclusion of the divergence term further enriches diversity, corroborated by visual inspection across multiple cases.

A notable caveat is a trade-off between diversity and fidelity: increased diversity may slightly degrade image quality, especially when the curvature term is prominent. The computational burden of higher-order derivative estimation is also highlighted as a limiting factor.

Practical and Theoretical Implications

The target-first variance-tilted paradigm offers several advantages:

• Explicit Control: The diversity target is mathematically transparent, facilitating theoretical analysis and reliable practical tuning.
• Probabilistic Consistency: The sampler aligns precisely with the specified batch distribution, unlike heuristic approaches.
• Generalizability: The framework is compatible with generic linear diffusion models and can be adapted to various feature spaces (e.g., frequencies, masks).

From a practical standpoint, the method enables systematic generation of diverse candidate sets, which is crucial in domains requiring exploration across a search space. In drug design, for instance, direct diversity encouragement in the molecular feature space could mitigate risks associated with mode collapse and homogeneity-induced failures.

The theoretical formulation highlights the power of analytic score correction through the Doob h-transform, providing a pathway for extending diversity-driven sampling to broader classes of generative models (e.g., parameterizations beyond VE SDEs, joint tasks with domain-specific features).

Future Directions

Potential developments include:

• Quantitative Diversity Metrics: Integrating precise measurements of diversity beyond qualitative visual assessments.
• Adaptive Annealing of Corrections: Dynamically modulating the Doob correction to optimize the trade-off between diversity and fidelity.
• Computational Optimization: Reducing the cost of Laplacian and divergence estimators via either efficient numerical schemes or novel stochastic trace methods.
• Feature-Space Extensions: Exploring non-linear, learned, or domain-specific feature maps for bespoke diversity objectives.

The approach may inspire subsequent research on diverse batch sampling for applications such as autonomous agents, ensemble uncertainty quantification, and scientific discovery pipelines.

Conclusion

This paper establishes a rigorous, probabilistic framework for diversity sampling in diffusion models by introducing explicit variance-tilted batch distributions and deriving exact Doob-score corrected samplers. The methodology offers tractable and theoretically motivated diversity induction, outperforming standard i.i.d. sampling and heuristic repulsion both in transparency and experimental outcomes. The explicit distribution along with efficient analytical and numerical correction methods lay a foundation for future explorations in controlled generative diversity, opening avenues for both practical deployment and theoretical study in generative modeling and beyond.