Calibrating Generative Models to Distributional Constraints

Generative models frequently suffer miscalibration, wherein statistics of the sampling distribution—such as the fraction of generations in a given class—deviate from desired values. We frame calibration as a constrained optimization problem and seek the closest model in Kullback-Leibler divergence satisfying a calibration constraint. To address the intractability of imposing these constraints exactly, we introduce two surrogate objectives for fine-tuning: (1) the relax loss, which replaces the constraint with a miscalibration penalty, and (2) the reward loss, which converts calibration into a reward fine-tuning problem. We demonstrate that these approaches substantially reduce calibration error across hundreds of simultaneous constraints and models with up to nine billion parameters, spanning applications in protein design, image generation, and language modeling. Code is available at https://github.com/smithhenryd/cgm.

Generative models are often judged by whether their individual outputs look good, but they can still make systematic mistakes in the distribution of outputs they produce. For instance, large language models (LLMs) may generate text that reinforces gender-occupation stereotypes, and image generative models may produce images that over- or underrepresent particular classes. This paper studies how to “calibrate” a generative model: fine-tune it so that its generated samples match desired distribution-level statistics (e.g., equal gender proportions), while changing the original model as little as possible. We introduce two practical fine-tuning methods, called CGM-relax and CGM-reward, that turn this calibration goal into objectives that can be optimized using samples from the model. In experiments on protein design, image generation, and language modeling, the proposed methods reduce calibration errors across many constraints, including in models with up to nine billion parameters.