Training-Free Distribution Adaptation for Diffusion Models via Maximum Mean Discrepancy Guidance

TL;DR

This paper tackles distribution mismatch in pre-trained diffusion priors when adapting to user-specific data without retraining. It introduces MMD Guidance, a training-free, inference-time mechanism that injects gradients of the squared Maximum Mean Discrepancy into the reverse diffusion (and in latent space for efficiency), with a prompt-aware extension via product kernels. The authors provide empirical evidence on synthetic GMMs and real image benchmarks (FFHQ, CelebA-HQ, SDXL/PixArt) showing improved distributional alignment (lower FD, KD, RRKE) while preserving fidelity, and they derive concentration bounds for the MMD gradient to support stability. The approach offers scalable, privacy-friendly domain adaptation for diffusion models with limited target data, and it opens avenues for combining divergences and extending to other modalities like video.

Abstract

Pre-trained diffusion models have emerged as powerful generative priors for both unconditional and conditional sample generation, yet their outputs often deviate from the characteristics of user-specific target data. Such mismatches are especially problematic in domain adaptation tasks, where only a few reference examples are available and retraining the diffusion model is infeasible. Existing inference-time guidance methods can adjust sampling trajectories, but they typically optimize surrogate objectives such as classifier likelihoods rather than directly aligning with the target distribution. We propose MMD Guidance, a training-free mechanism that augments the reverse diffusion process with gradients of the Maximum Mean Discrepancy (MMD) between generated samples and a reference dataset. MMD provides reliable distributional estimates from limited data, exhibits low variance in practice, and is efficiently differentiable, which makes it particularly well-suited for the guidance task. Our framework naturally extends to prompt-aware adaptation in conditional generation models via product kernels. Also, it can be applied with computational efficiency in latent diffusion models (LDMs), since guidance is applied in the latent space of the LDM. Experiments on synthetic and real-world benchmarks demonstrate that MMD Guidance can achieve distributional alignment while preserving sample fidelity.

Figures (21)

• Figure 1: Image generation via a text-conditioned latent diffusion model (LDM) with no guidance vs. our proposed MMD guidance. The LDM (Stable Diffusion-XL) following our proposed MMD guidance over 100 reference samples of "cat" and "dog" images could exhibit the visual format of the target distribution, but the unguided LDM's output samples differ in style from the target model.
• Figure 2: Comparison of MMD guidance with baselines on 100D Gaussian distributions, when guiding toward a user with 4 Gaussian components.
• Figure 3: Effect of mode proportions in MMD guidance.
• Figure 4: User's samples and generated data by unguided/guided LDMs on the FFHQ dataset.
• Figure 5: Qualitative comparison of reference set and MMD-guided image generation with SDXL.

Theorems & Definitions (10)

• Remark 1
• Theorem 1: Concentration of Cross Term in MMD Gradient
• Corollary 1: Gaussian RBF Kernel Concentration
• Theorem 2: Uniform concentration of the gradients
• Corollary 2
• proof
• Theorem 3: Concentration for Weighted Cross Term in Product Kernel
• proof
• Corollary 3: Product Kernel with Gaussian RBF
• proof