Calibrating Generative Models
Henry D. Smith, Nathaniel L. Diamant, Brian L. Trippe
TL;DR
Calibrating Generative Models (CGM) tackles miscalibration in samples by casting calibration as a constrained KL minimization: minimize $D_{\mathrm{KL}}(p_{\theta}\|p_{\theta_{\text{base}}})$ subject to $\mathbb{E}_{p_{\theta}}[\mathbf{h}(\mathbf{x})]=\mathbf{h}^*$. It introduces two surrogate objectives, CGM-relax and CGM-reward, enabling tractable optimization with unbiased gradient estimators; CGM-relax uses a miscalibration penalty plus a KL penalty, while CGM-reward targets a maximum-entropy tilt $p_{\boldsymbol{\alpha}}$ and minimizes $D_{\mathrm{KL}}(p_{\theta}\|p_{\boldsymbol{\hat{\alpha}}_N})$. Across protein design, conditional image generation, and language tasks, CGM reduces calibration error across hundreds of constraints on models up to $10^9$ parameters with minimal degradation to sample quality. The work highlights residual challenges in rare-event calibration and notes the current framework's reliance on tractable likelihoods, pointing to future work to extend calibration to implicit models such as VAEs, GANs, and other non-likelihood-based frameworks.
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 closest model in Kullback-Leibler divergence satisfying calibration constraints. 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 one billion parameters, spanning applications in protein design, image generation, and language modeling.