# Calibrating Generative Models to Feature Distributions with MMD Finetuning

arXiv: https://arxiv.org/abs/2606.19496 · HTML: https://mlsanity.com/paper/2606.19496/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 47 references that are in the map.
- First posted on arXiv 17 Jun 2026.

## Builds on (most cited first, 23 of 23)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [How transferable are features in deep neural networks?](https://mlsanity.com/paper/1411.1792/) (2014)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models](https://mlsanity.com/paper/2305.16381/) (2023)
- [GuacaMol: Benchmarking Models for De Novo Molecular Design](https://mlsanity.com/paper/1811.09621/) (2018)
- [MMD GAN: Towards Deeper Understanding of Moment Matching Network](https://mlsanity.com/paper/1705.08584/) (2017)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [DiffusionNFT: Online Diffusion Reinforcement with Forward Process](https://mlsanity.com/paper/2509.16117/) (2025)
- [Fr'echet ChemNet Distance: A metric for generative models for molecules in drug discovery](https://mlsanity.com/paper/1803.09518/) (2018)
- [Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control](https://mlsanity.com/paper/2402.15194/) (2024)
- [Out of Many, One: Designing and Scaffolding Proteins at the Scale of the Structural Universe with Genie 2](https://mlsanity.com/paper/2405.15489/) (2024)
- [Diversity-Inducing Policy Gradient: Using Maximum Mean Discrepancy to Find a Set of Diverse Policies](https://mlsanity.com/paper/1906.00088/) (2019)
- [Distributional Diffusion Models with Scoring Rules](https://mlsanity.com/paper/2502.02483/) (2025)
- [DISCO Nets: DISsimilarity COefficient Networks](https://mlsanity.com/paper/1606.02556/) (2016)
- [Graph Generative Pre-trained Transformer](https://mlsanity.com/paper/2501.01073/) (2025)
- [Flow Density Control: Generative Optimization Beyond Entropy-Regularized Fine-Tuning](https://mlsanity.com/paper/2511.22640/) (2025)
- [DART-Eval: A Comprehensive DNA Language Model Evaluation Benchmark on Regulatory DNA](https://mlsanity.com/paper/2412.05430/) (2024)
- [Calibrating Generative Models to Distributional Constraints](https://mlsanity.com/paper/2510.10020/) (2025)
- [Learning Latent Graph Structures and their Uncertainty](https://mlsanity.com/paper/2405.19933/) (2024)

## Abstract (arXiv)

Generative models can produce individually plausible samples while deviating substantially from a target set in the distribution of key features. For example, a model pretrained on broad drug-like chemical space may generate molecules whose molecular features differ from those of a therapeutic class of interest, such as known antibiotics. Correcting such distributional miscalibration is challenging: direct finetuning on the target set can overfit and does not control which features are matched. To fill this gap, we introduce kernel Calibrating Generative Models (kCGM). kCGM minimizes a maximum mean discrepancy (MMD) between generated and target feature distributions using an unbiased score-function estimator, with KL regularization to remain close to the pretrained model. On a target set of 174 antibiotics, direct finetuning sacrifices chemical validity for feature-distribution matching, whereas kCGM improves target feature matching while increasing validity. We further demonstrate kCGM in protein and DNA generation tasks, showing it can adapt autoregressive, continuous-space diffusion, and discrete diffusion models using only feature-level supervision. Code is available at https://github.com/smithhenryd/cgm.