Calibrating Generative Models to Feature Distributions with MMD Finetuning

Technical Summary: Calibrating Generative Models to Feature Distributions with MMD Finetuning (kCGM)

Problem Statement
Generative models often produce samples that are individually plausible but exhibit a distributional misalignment regarding key features of interest. For instance, a model pretrained on broad chemical space may generate molecules that differ significantly in feature distribution from a specific therapeutic class, such as antibiotics. Correcting this via direct finetuning on the target set is problematic: it risks overfitting to small datasets, requires full target samples (which may be scarce), and fails to specify which specific features should be matched. Furthermore, existing methods often struggle with discrete-valued or black-box features (e.g., Morgan fingerprints) that are non-differentiable. The paper identifies a gap in algorithms capable of calibrating generative models to match a target feature distribution using only feature-level supervision, without requiring differentiability of the feature map or the sampling procedure.

Methodology: kCGM
The authors introduce kernel Calibrating Generative Models (kCGM), a finetuning method that minimizes the Maximum Mean Discrepancy (MMD) between the generated feature distribution and a target feature distribution, regularized by the Kullback-Leibler (KL) divergence to the pretrained model.

• Objective Function: The method optimizes an objective $L_\theta$ composed of two terms:

  1. MMD Loss: An unbiased finite-sample estimator of the squared MMD between the model-induced feature distribution $P$ (where $y = h(x)$ and $x \sim p_\theta$) and the empirical target distribution $Q$ (from target features $y^{targ}_{1:N}$).
  2. KL Regularization: A term $\lambda D_{KL}(p_\theta || p_{\theta_{pre}})$ to ensure the finetuned model remains close to the pretrained distribution, preventing catastrophic forgetting or mode collapse.

• Kernel Selection: Unlike previous approaches that might rely on dot-product kernels (which only match means), kCGM utilizes kernels that define a metric on probability distributions (zero only for identical distributions).

  • For continuous Euclidean features, an energy-distance kernel is used.
  • For discrete bit-vector features (e.g., Morgan fingerprints), a Tanimoto/Jaccard kernel is employed.
  • These choices avoid introducing additional bandwidth hyperparameters.

• Gradient Estimation: Since the feature map $h$ and the sampling process may be non-differentiable, kCGM employs score-function gradients (REINFORCE-style).

  • Leave-One-Out (LOO) Baseline: A critical technical contribution is the derivation of a novel LOO baseline for the MMD loss. Standard LOO baselines average coefficients from other samples in a batch, but MMD coefficients are coupled across samples. The authors derive a specific "MMD-LOO" estimator that approximates the average coefficient had a specific sample not appeared in the mini-batch, reducing variance without introducing bias.
  • Algorithm: The process involves sampling from the current model (without gradient tracking), computing feature vectors, calculating MMD and KL coefficients with LOO baselines, and updating model weights via score-function gradients.

Key Contributions

1. Formulation: The paper formulates feature-distribution matching as KL-regularized MMD finetuning using only target feature samples, eliminating the need for full target samples in the original sample space.
2. Unbiased Estimator: The authors derive an unbiased score-function gradient estimator for the MMD objective, enabling the use of non-differentiable feature maps and non-reparameterizable samplers (e.g., autoregressive models, discrete diffusion).
3. Generalization: The method generalizes Calibrating Generative Models (CGM). While CGM matches target feature means, kCGM matches the full distribution via MMD. The paper notes that using a dot-product kernel in kCGM recovers the CGM-relax objective.

Experimental Results
The authors evaluate kCGM across three domains: small-molecule generation, protein structure generation, and regulatory DNA generation.

• Small-Molecule Antibiotics:

  • Task: Calibrate a G2PT autoregressive graph model to match the feature distribution of 174 known antibiotics.
  • Features: Morgan fingerprints (local substructure), Murcko scaffolds (topology), and physicochemical descriptors.
  • Findings: kCGM achieves a superior tradeoff between matching the target feature distribution (measured by MMD) and distance from the pretrained model compared to CGM and direct finetuning. Crucially, while direct finetuning on the small antibiotic set led to a high proportion of chemically invalid molecules, kCGM improved feature matching while increasing chemical validity. kCGM also demonstrated cross-feature transfer, improving scaffold and descriptor matching even when trained only on fingerprints.

• Protein Structure Diversity:

  • Task: Finetune Genie 2 (a protein diffusion model) to match the secondary structure distribution (alpha-helices vs. beta-strands) of the CATH dataset.
  • Findings: The pretrained model produced mostly alpha-helices. kCGM reduced the disagreement with the target distribution by approximately 2x compared to CGM. kCGM showed a smooth tradeoff curve with the KL weight $\lambda$, whereas CGM's performance was less sensitive to $\lambda$, making optimization difficult.

• Regulatory DNA Generation:

  • Task: Finetune a discrete diffusion model to generate DNA sequences with realistic cell-type-specific activity profiles (predicted by AlphaGenome).
  • Findings: kCGM successfully matched the full distribution of activity profiles for 47 different cis-regulatory topics, outperforming CGM (which only matched means). The finetuned model generated sequences with activity profiles similar to real enhancers, whereas the pretrained model generated under-active sequences.

Significance and Claims
The paper claims that kCGM provides a robust alternative to direct finetuning, particularly when target datasets are small or when only specific aspects of a distribution need preservation. Its primary significance lies in its ability to:

• Align generative models to user-specified feature distributions using feature-level supervision only.
• Handle black-box and non-differentiable features (e.g., fingerprints, predictor outputs) that are common in scientific domains but incompatible with standard gradient-based finetuning.
• Maintain or improve sample quality (validity) while correcting distributional miscalibration, a common failure mode of direct finetuning.
• Be applicable across diverse architectures, including autoregressive models, continuous-space diffusion, and discrete diffusion models.

The authors note limitations, such as potential ineffectiveness for very high-dimensional features due to the curse of dimensionality in MMD hypothesis testing, and the requirement for tractable sampling-trajectory log-probabilities (excluding flow matching models). They suggest future work could extend this framework to settings where only distributional constraints exist without full samples in the original generative space.