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

Technical Summary: MMD Guidance: Training-Free Distribution Adaptation for Diffusion Models via Maximum Mean Discrepancy Guidance

Problem Statement
Pre-trained diffusion models serve as powerful generative priors but often fail to align with user-specific target distributions, particularly in domain adaptation scenarios where only a few reference examples are available. Retraining these models is computationally infeasible for many users. Existing inference-time guidance methods, such as classifier guidance or classifier-free guidance, optimize surrogate objectives (e.g., class likelihoods) rather than directly minimizing the distributional discrepancy between generated samples and a target reference set. Consequently, there is a need for a training-free mechanism that can explicitly align diffusion model outputs with a target distribution defined by limited reference data.

Methodology
The authors propose MMD Guidance, a training-free adaptation method that augments the reverse diffusion process with gradients derived from the Maximum Mean Discrepancy (MMD).

1. Distribution Matching via MMD: The method minimizes the squared MMD between the empirical distribution of generated samples at timestep $t$ and the empirical distribution of a small reference dataset. MMD is selected over other divergences (like KL-divergence or Wasserstein distance) because it provides unbiased, low-variance estimates from limited samples without suffering from the curse of dimensionality, and its kernel-based formulation allows for efficient gradient computation.
2. Guidance Mechanism: During sampling, the standard update step is modified by subtracting the gradient of the squared MMD with respect to the current latent sample $z_t$. The update rule is:
   $$z_{t-1}^{(i)} = \text{sampler}(z_t^{(i)}, t, \epsilon_\theta) - \lambda_t \nabla_{z_t^{(i)}} \widehat{\text{MMD}}^2(\hat{P}_t, \hat{Q})$$
   where $\lambda_t$ controls the guidance strength. The gradient consists of an intra-batch term (promoting diversity among generated samples) and a cross-term (attracting samples toward the reference distribution).
3. Latent Space Implementation: To ensure computational efficiency, MMD Guidance is applied in the latent space of Latent Diffusion Models (LDMs). Reference images are encoded once using the VAE encoder, and guidance is computed on the compressed latent representations. This reduces dimensionality and improves the signal-to-noise ratio for MMD estimation.
4. Prompt-Aware Extension: For conditional generation, the method extends to a product kernel $k_{\text{multimodal}}([p, x], [p', x']) = k_{\text{text}}(p, p') \cdot k_{\text{image}}(x, x')$. This allows the guidance to jointly align with both textual semantics (via prompt similarity) and visual styles (via latent similarity), enabling prompt-aware distribution matching without retraining the denoiser.

Key Contributions

• Introduction of MMD Guidance, a training-free inference-time method for adapting diffusion models to target distributions specified by limited reference sets.
• Proposal of a prompt-aware adaptation extension using product kernels to jointly align text prompts and visual styles.
• Development of an efficient implementation operating in the latent space of LDMs.
• Theoretical guarantees showing that the MMD gradient concentrates around the true population gradient with high probability, even with limited reference samples.

Results
The authors evaluate MMD Guidance on synthetic Gaussian Mixture Models (GMMs) and real-world image benchmarks (FFHQ, CelebA-HQ, AFHQ) using models like Stable Diffusion v1.4, SDXL, and PixArt.

• Distribution Alignment: On synthetic GMMs, MMD Guidance successfully recovers target mixture proportions and modes from limited samples, outperforming baselines like No-Guidance, Classifier-Free Guidance (CFG), and Classifier Guidance (CG) in Fréchet Distance (FD) and Kernel Distance (KD).
• Image Generation: On FFHQ and CelebA-HQ, MMD Guidance achieves superior fidelity and distributional coverage (measured by FD, KD, Density, and Coverage) compared to training-free baselines and competitive performance against fine-tuning methods, despite requiring no training.
• Prompt-Aware Adaptation: In experiments with SDXL and PixArt, MMD Guidance successfully adapts the visual style of generated images (e.g., anime cats, Van Gogh dogs) to match reference styles while respecting text prompts, significantly reducing distribution discrepancy metrics.
• Efficiency: The method incurs minimal computational overhead. Runtime analysis shows that generating samples with MMD Guidance is only slightly slower than unguided sampling, making it scalable for large-scale LDMs.
• Robustness: The method performs well with as few as 10–50 reference samples and remains effective across varying domain gaps, though larger gaps require more reference samples or stronger guidance.

Significance and Limitations
The paper claims that MMD Guidance offers a practical, training-free solution for domain adaptation in diffusion models, reducing the computational and memory costs associated with fine-tuning. It is particularly useful in settings with limited reference data and rapid deployment requirements. The authors note that while the method performs well with limited data, extremely sparse scenarios may still suffer from gradient estimation variance. They also suggest that adaptive kernel selection and combining MMD with other divergences are potential directions for future work.