Calibrating Generative Models

Henry Smith, Brian L. Trippe

Keywords: diffusion, fine-tuning, reward fine-tuning, stochastic optimal control, generative model, protein structure generation

TL;DR: We provide a method for calibrating generative models whose statistics (e.g. moments or class probabilities) differ from target values; we focus on diffusion generative models

Abstract: Generative models frequently suffer miscalibration, wherein class probabilities and other statistics of the sampling distribution deviate from desired values. We frame calibration as a constrained optimization problem and seek the minimally perturbed model (in Kullback-Leibler divergence) satisfying calibration constraints. To address the intractability of the hard constraint, we introduce two surrogate objectives: (1) the relaxed loss, which replaces the constraint with a miscalibration penalty, and (2) the reward loss, which we derive as a divergence to a known, approximate solution. We show how to apply these methods to neural-SDE models and find they solve synthetic calibration problems to high precision. Lastly, we demonstrate the practicality of the approach by calibrating a 15M-parameter protein structure diffusion model to match the distribution of secondary structure composition of natural proteins.

Submission Number: 148