🤖 AI Summary
Generative models often suffer from poor calibration, where predicted class probabilities diverge from empirical sampling statistics. This paper proposes a general constraint-based calibration framework: it minimizes the KL divergence between the calibrated and original model distributions, subject to multiple category- or statistic-specific constraints. Two scalable optimization objectives are introduced: (i) relaxed loss—incorporating calibration error as a regularization term—and (ii) reward loss—transforming constraints into differentiable reward signals for fine-tuning. Both support joint optimization over hundreds of constraints. The method preserves generation quality while substantially reducing calibration error—even for billion-parameter models. Extensive experiments across protein design, image generation, and language modeling demonstrate strong generalization and computational efficiency.
📝 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.