Research map / Image generation / Fast diffusion sampling

Calibrating Generative Models to Feature Distributions with MMD Finetuning

NeurIPS 2026 · Image generation › Fast diffusion sampling · approach: Distributional control of diffusion · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes kCGM, which finetunes pretrained generative models to match target feature distributions by minimizing an unbiased MMD estimate via score-function gradients with leave-one-out baselines and KL regularization.

Problem

Pretrained generative models often produce individually plausible samples whose feature distributions miscalibrate relative to a target set; direct finetuning overfits and cannot select which features to match, and prior CGM only matches feature means or binned CDFs.

Approach

Use the unbiased finite-sample squared MMD between model and target feature samples as the objective, optimized with score-function gradients (not requiring differentiable feature maps or samplers), with a novel leave-one-out baseline for the coupled kernel coefficients and KL regularization to the pretrained model.

Key results

• On the 174-antibiotic target, kCGM achieves a better Morgan MMD^2-KL tradeoff than CGM and direct finetuning, and avoids the high invalid-molecule fraction seen with direct finetuning.
• In the Genie 2 protein secondary-structure experiment, kCGM's distribution-matching error is roughly 2x smaller than CGM's at every tested KL budget.
• For conditional DNA diffusion on DeepMEL2, kCGM achieves a better KL-MMD tradeoff than CGM against held-out AlphaGenome target distributions.

Surprising findings

• The dot-product kernel version of kCGM is mathematically equivalent to CGM-relax, and the leave-one-out baseline improves it, suggesting the baseline is a drop-in improvement for CGM.
• Direct finetuning reduces Morgan MMD but produces a high fraction of chemically invalid molecules.

Limitations

• MMD may be ineffective for very high-dimensional features due to polynomial power decay.
• Requires tractable sampling-trajectory log-probabilities, so does not directly apply to flow matching models.
• Experiments construct target features from full target samples rather than true partial distributional measurements.

Experimental setup

Datasets: GuacaMol 174 antibiotics CATH DeepMEL2

Benchmarks: Morgan fingerprint MMD scaffold MMD descriptor MMD CATH secondary structure symmetrized KL AlphaGenome ATAC-seq activity MMD

Metrics: MMD^2 KL to pretrained symmetrized KL fraction invalid molecules uniqueness

Baselines: CGM-relax direct finetuning CGM with CDF binning

Builds on: Calibrating Generative Models (CGM) Genie 2 G2PT score-function gradient estimators

Closest papers by meaning

• Kernel-Gradient Drifting Models (NeurIPS 2026 · similarity 0.85)
• Kernel-Based Evaluation of Conditional Biological Sequence Models (ICML 2024 · similarity 0.84)
• Flow Density Control: Generative Optimization Beyond Entropy-Regularized Fine-Tuning (NeurIPS 2025 · similarity 0.83)
• PQMass: Probabilistic Assessment of the Quality of Generative Models using Probability Mass Estimation (ICLR 2025 · similarity 0.83)
• TFG-Flow: Training-free Guidance in Multimodal Generative Flow (ICLR 2025 · similarity 0.83)
• A Unifying Information-theoretic Perspective on Evaluating Generative Models (AAAI 2025 · similarity 0.83)