--- title: MMD Guidance for Distribution Alignment url: https://www.emergentmind.com/topics/mmd-guidance type: topic --- # MMD Guidance for Distribution Alignment Maximum Mean Discrepancy (MMD) Guidance is a training-free, distribution-matching methodology that steers generative models to align their outputs with a small reference dataset. It operates by injecting gradients of the empirical Maximum Mean Discrepancy—a kernel-based non-parametric statistical distance—into the inference procedure, most notably within reverse diffusion samplers, without updating model weights. MMD Guidance is suitable for unconditional, conditional, and prompt-aware sampling, achieves robust adaptation from limited reference data, and is computationally efficient in modern latent diffusion architectures [2601.08379]. ## 1. Definition of the MMD Objective and Its Gradient Let $\{x_i\}_{i=1}^B$ denote a batch of generated samples and $\{y_j\}_{j=1}^{N_r}$ a fixed reference set. Given a positive semi-definite kernel $k : \mathcal{X} \times \mathcal{X} \to \mathbb{R}$ (e.g. Gaussian RBF), the empirical squared Maximum Mean Discrepancy is \[ \widehat{\mathrm{MMD}^2}(\widehat{P},\widehat{Q}) = \frac{1}{B^2} \sum_{i,i'=1}^B k(x_i,x_{i'}) + \frac{1}{N_r^2} \sum_{j,j'=1}^{N_r} k(y_j,y_{j'}) - \frac{2}{B N_r} \sum_{i=1}^B \sum_{j=1}^{N_r} k(x_i,y_j) \] The gradient with respect to a generated sample $x_i$ is given by \[ \nabla_{x_i} \widehat{\mathrm{MMD}^2} = \frac{2}{B^2} \sum_{i'=1}^B \nabla_{x_i} k(x_i, x_{i'}) - \frac{2}{B N_r} \sum_{j=1}^{N_r} \nabla_{x_i} k(x_i, y_j) \] For a Gaussian RBF kernel $k(u,v) = \exp(-\|u-v\|^2 / (2 \sigma^2))$, the gradient specializes to \[ \nabla_{x_i} \widehat{\mathrm{MMD}^2} = - \frac{2}{\sigma^2 B^2} \sum_{i'=1}^B k(x_i, x_{i'})(x_i - x_{i'}) + \frac{2}{\sigma^2 B N_r} \sum_{j=1}^{N_r} k(x_i, y_j)(x_i - y_j) \] ## 2. Integration of MMD Guidance into Diffusion Sampling During the reverse diffusion process (e.g., DDPM/DDIM), each denoising step is modified by a small MMD gradient correction. At diffusion timestep $t$, the update for each sample $x_t^{(i)}$ is \[ x_{t-1}^{(i)} = \mathrm{Sampler}(x_t^{(i)}, t, \epsilon_\theta) - \lambda_t \nabla_{x_t^{(i)}} \widehat{\mathrm{MMD}^2}(\{x_t^{(i)}\}, \{y_j\}) \] where $\lambda_t$ is a guidance strength (typically constant or slowly decaying across timesteps). This can be performed directly in pixel space or—more efficiently—in the latent space of a pretrained Variational Autoencoder (VAE), as used in Latent Diffusion Models (LDMs): 1. Encode reference samples in the latent space: $z_j^{(r)} = \mathcal{E}(x_j^{(r)})$. 2. Sample $z_T^{(i)} \sim \mathcal{N}(0, I)$. 3. For $t = T$ to $1$, perform standard denoising and apply the MMD gradient correction. 4. Decode final latent to obtain the generated sample. ## 3. Kernel Choices and Prompt-Conditioned Extensions The reference-alignment objective can be tailored via kernel selection: - **Gaussian RBF**: $k(u,v) = \exp(-\|u - v\|^2 / (2\sigma^2))$, bandwidth $\sigma$ chosen by a grid (e.g. $\{1.25, 1.5, 2.0\}$) or proportional to latent scale. - **Polynomial kernel**: $(c + \langle u, v \rangle)^d$, $d \in \{2, 3, 4\}$. - **Product kernel** (for prompt-aware conditional generation): $k([p,z],[p',z']) = k_p(p,p') \, k_z(z,z')$, where $k_p$ measures prompt similarity (e.g. cosine or RBF in CLIP embedding space) and $k_z$ is the visual kernel. - **Guidance strength** $\lambda_t$: Default values are $\alpha \approx 10^{-4}$ in latent-space guidance, $\alpha \approx 10^{-2}$ in pixel-space. Prompt-aware adaptation weights the cross-term gradient by prompt similarity, ensuring samples are steered towards references matching the semantic intent. ## 4. Computational Efficiency and Latent-Space Implementation Operating in the latent space $\mathcal{Z}$ of a pretrained LDM confers multiple advantages: - Reduced dimensionality accelerates kernel computation and gradient evaluation. - Semantic compression yields better MMD estimation on structured features. - Memory overhead is minimal, since reference encodings are reused. - Runtime overhead for batch sizes up to 500 samples and 50 timesteps is only $10$–$15\%$ on consumer-class GPUs. Bandwidth selection and guidance strength can be optimized via a grid search on held-out reference data, targeting metrics such as Fréchet Distance (FD) and Kernel Distance (KD). ## 5. Experimental Evaluation Experimental benchmarks demonstrate that MMD Guidance achieves consistent and substantial improvements in distributional alignment and sample fidelity: - **Synthetic GMMs**: recovers desired mixture modes and proportions with as few as $50$–$200$ references, achieving the lowest FD and KD against classifier-guidance (CG) and classifier-free guidance (CFG) baselines. - **Mode-proportion correction**: can preferentially reproduce new mode proportions when the reference set is Dirichlet-reweighted. - **Real-world image adaptation** (FFHQ, CelebA-HQ): aligns generated samples to user-defined characteristics using $500$ references, reducing FD from $1221$ (no guidance) to $693$ (MMD guidance). - **Prompt-aware stylized image generation**: on Stable Diffusion XL and PixArt, guidance reduces FD, KD, Relative Reconstruction Kernel Error (RRKE), and increases coverage; adaptation occurs without network retraining. - **Reference set sensitivity**: performance improves quickly with $50$–$100$ references and plateaus, with robustness to kernel and guidance strength variations. ## 6. Advantages, Practical Considerations, and Significance MMD Guidance is fully training-free and provides direct distribution-aware alignment. It leverages the low-variance, consistent estimation properties of MMD, particularly under small reference sets, and its gradients are efficiently computable via modern hardware. MMD gradients can be injected into any generative process for which a differentiable kernel can be evaluated over its latent or output space, without requiring model finetuning or additional training steps. The framework generalizes to various adaptation scenarios: - **Prompt-aware adaptation** in conditional generative models, steering samples jointly with reference prompts. - **Distribution correction** or style transfer using limited reference data. - **Domain adaptation** in generative modeling pipelines. The principal advantage over classifier-based guidance methods is the direct minimization of RKHS distance to the true reference distribution, rather than a surrogate (e.g., classifier likelihood). This yields superior coverage, mode fidelity, and robustness to overfitting [2601.08379]. ## 7. Summary Table: MMD Guidance Properties | Property | Description | Reference Scenario | |------------------------------|----------------------------------------------------------------|----------------------------| | Training-free | No update to model weights; operates at inference | Domain adaptation, style | | Differentiable MMD | Empirical estimate and gradient in pixel or latent space | Any kernel metric | | Kernel flexibility | Gaussian RBF, polynomial, product kernels for conditional | Prompt-aware LDM | | Reference-efficient | Robust with $O(100)$ samples; minimal overfitting | Few-shot user adaptation | | Low computational overhead | Latent-space implementation adds $10$–$15\%$ runtime cost | Stable Diffusion XL | | Direct distribution matching | Aligns to reference data's empirical distribution in RKHS | Synthetic/real domains | | Sample fidelity | Preserves generative quality while achieving alignment | FFHQ, CelebA experiments | MMD Guidance formalizes distributional adaptation as direct kernel-mean matching, circumventing limitations of classifier- or discriminator-based guidance frameworks, and is widely applicable to various generation architectures. Source: https://www.emergentmind.com/topics/mmd-guidance