Calibrating Generative Models to Distributional Constraints

Henry D. Smith; Nathaniel L. Diamant; Brian L. Trippe

Source-linked AI summary

Calibrating Generative Models to Distributional Constraints

Henry D. Smith, Nathaniel L. Diamant, Brian L. Trippe

arXiv:2510.10020v4Oct 11, 2025stat.MLcs.LGq-bio.BM

TL;DR

Generative models often produce samples whose distributional statistics deviate from desired values, motivating calibration at the distribution level. The paper formulates calibration as KL-nearest constrained optimization and introduces relax and reward fine-tuning losses. Across protein, image, and language applications, these methods reduce calibration error under many simultaneous constraints, while residual rare-event error and likelihood requirements remain limitations.

Problem
Generative models can be miscalibrated because statistics of their sampling distributions, such as class frequencies, deviate from desired values.
Method
The paper formulates calibration as KL-nearest constrained optimization and approximates it with CGM-relax and CGM-reward fine-tuning objectives.
Results
CGM consistently reduces calibration error across protein design, conditional image generation, and language modeling with hundreds of constraints and models up to nine billion parameters.
Takeaways & Limitations
The two objectives provide practical approaches for calibrating high-dimensional generative models while preserving generation quality.
Takeaways & Limitations
Residual calibration error remains, especially for rare events, and the framework depends on models with tractable likelihoods.

Abstract

from arXiv · show

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.

1. Introduction

The paper treats generative-model miscalibration as distributional statistics deviating from desired values and formulates correction as constrained optimization. It seeks the closest calibrated model in KL divergence and relates this goal to existing calibration and fine-tuning approaches.

Motivation: Generative models can systematically misrepresent desired distributional statistics, including class frequencies, demographic attributes, protein structures, and DNA subsequences.Examples span image generation, language modeling, synthetic biology, and DNA generation.
Problem formulation: Calibration requires fine-tuning a miscalibrated base model so its expected statistic matches a desired value.For class membership, this means adjusting the probability that generated samples belong to a specified class.
Problem formulation: The proposed optimization selects, among calibrated models, the one closest to the base model in KL divergence.KL divergence is chosen because it is tractable for several classes of generative models.
Related work: Unlike supervised calibration methods that match predicted probabilities to future empirical frequencies, this work constrains the generation distribution itself.Existing generative fine-tuning methods typically optimize rewards for individual samples rather than impose distribution-level constraints.
Related work: Prior distribution-level approaches include language-model fine-tuning, optimal-transport balancing for diffusion models, and a concurrent diffusion-model method.These works establish related directions for enforcing distributional constraints in generative models.

2. Calibrating Generative Models with CGM-relax and CGM-reward

CGM replaces intractable distributional calibration constraints with two stochastic fine-tuning objectives: a violation-penalizing relax loss and a reward-based loss derived from maximum entropy. The relax method remains applicable to higher-dimensional constraints, while reward fine-tuning depends on estimating the maximum-entropy parameters.

2. Calibrating Generative Models with CGM-relax and CGM-reward: CGM uses unbiased gradient estimates to optimize two surrogate objectives despite intractable expectations, requiring samples, model probabilities, and score-function gradients.The constraint statistic h need not be differentiable.
2.1. The Relax Loss: The relax loss penalizes squared constraint violation alongside KL divergence to the base model, with λ trading off calibration accuracy against proximity to that model.As λ decreases, the violation term dominates and the solution approaches the constrained optimum; finite λ must balance the two objectives.
2.2. The Reward Loss: The reward loss approximates calibration through the maximum-entropy solution, then minimizes KL divergence to the resulting exponentially tilted base model.Its objective coincides with reward fine-tuning using a reward r_α(x)=α^T h(x), up to a normalizing constant.
2.2. The Reward Loss: Estimating α* from N base-model samples yields a concave optimization problem solvable by convex methods, and the estimate converges to α* at a parametric rate as N grows.The reward method becomes more accurate when enough samples are used to estimate α*.
2. Calibrating Generative Models with CGM-relax and CGM-reward: In a one-dimensional Gaussian-mixture experiment, both methods closely approximate the maximum-entropy solution when appropriately tuned.For relax, λ controls the constraint–KL trade-off; for reward, increasing N improves recovery of α* and approximation quality.
2.4. Relationship Between Relax and Reward Losses: CGM-relax can handle up to 10^3 constraints, whereas CGM-reward is not well-defined beyond 30 constraints because estimating α* can become infeasible.In high-dimensional constraints, relax outperforms reward even when the latter uses the oracle α*; the relax method avoids estimating reward parameters.

3. Simulations: Determining when CGM Thrives and Struggles

The simulations test when CGM calibrates rare events and many simultaneous constraints, while examining hyperparameter trade-offs and failure modes. Both methods can solve calibration accurately in suitable settings, but their behavior diverges as events become rarer or constraints higher-dimensional.

Hyperparameter selection: The regularization parameter λ trades off constraint satisfaction against deviation from the base model, while larger N improves CGM-reward’s recovery of α*.With appropriate hyperparameters, both algorithms solve the calibration problem to high accuracy.
Calibrating rare events: 80–90% calibration-error reduction at base proportion 10^-2 and approximately 50% at 10^-3, with constraint satisfaction degrading beyond this rarity threshold.At 10^-3, fewer than one rare-mode sample is drawn per batch on average; larger batches may enable calibration of rarer events.
Calibrating rare events: For rare events below 10^-1, CGM-relax deviates farther from the base model than CGM-reward and the maximum entropy solution.
Calibrating rare events: CGM satisfies a fixed 0.8 rare-event target with high accuracy even when the base proportion is as small as 10^-3.
Scalability to high-dimensional models and constraints: In high-dimensional constraints, CGM-relax eliminates most constraint violation up to k=10^3 but incurs excess KL divergence that grows linearly with dimension relative to the maximum entropy solution.
Scalability to high-dimensional models and constraints: CGM-reward performs well below 10 constraints, whereas the empirical maximum entropy problem is infeasible with high probability for k>30; CGM-relax still outperforms it using oracle α*.

4. Case-studies with Diverse Models, Data, and Constraints

Across protein design, image generation, and language modeling, CGM calibrates diverse generative models under many constraints while exposing trade-offs between calibration, fidelity, and sample quality.

Overview: CGM-relax reduces calibration error across protein, image, and language-model applications, including hundreds of constraints and models up to nine billion parameters.The case studies span protein secondary-structure statistics, animal class proportions, and gender-profession balance.
Protein design: Nearly fivefold for Genie2 and twofold for ESM3-open, CGM-relax improves secondary-structure diversity relative to CATH distributions.The improvement is measured by symmetrized KL distance and increases design failures for ESM3-open slightly.
Protein design: 74% versus 14% failure rates show that increasing Genie2 sampling noise improves diversity less safely than CGM calibration.For ESM3, increasing temperature to τ = 1 produces a 97% failure rate.
Image generation: 0.306 to 0.101, CGM-relax reduces TarFlow animal-proportion total variation, although FID worsens from 15.9 to 21.0 because calibration introduces artifacts.Reported artifacts include out-of-class animals and blended animals; CGM-reward remains near the base model but fails to reduce constraint violation.
Language modeling: More than threefold, CGM-relax reduces gender imbalance relative to baseline but moves farther from the base model than CGM-reward.Lower λ improves constraint satisfaction while increasing symmetric KL distance; the least-regularized model remains below 1.7 symmetric KL.
Language modeling: Four of six held-out professions improve after conditional calibration, while two become more imbalanced.This evaluates transfer beyond professions explicitly included during fine-tuning.

5. Conclusion

CGM provides practical fine-tuning procedures that reduce distribution-level calibration error across diverse generative models while preserving generation quality, but residual error and likelihood requirements remain.

Conclusion: CGM consistently reduces calibration error under hundreds of simultaneous constraints in models with up to nine billion parameters while preserving generation quality.The applications cover protein design, conditional image generation, and language modeling.
Limitations: Residual calibration error remains, especially for rare events relevant to protein structure modeling.The authors state that the calibration problem is not yet solved.
Limitations: CGM currently depends on models with tractable likelihoods, leaving extension to GANs and other implicit models unresolved.This is identified as a broader open challenge.

A. Extended Discussion of Related Work

Related work spans distributional calibration, inference-time reweighting, and reward fine-tuning, while CGM supplies model-flexible stochastic objectives with unbiased gradient estimates for distribution-level constraints.

The calibration problem: Prior calibration methods are restricted to specific model classes, show weaker empirical performance, or impose constraints during training rather than fine-tuning.The comparison motivates a general fine-tuning framework for generative models.
The calibration problem: CGM-reward reduces Gemma gender miscalibration by roughly 81%, while CGM-relax reaches up to 94% reduction in constraint violation.CGM-relax outperforms Khalifa et al. (2021) particularly in constraint satisfaction.
The calibration problem: Go et al. (2023) obtain < 50% reduction in constraint violation, while Shen et al. (2024) report difficulty meeting desired class proportions with an optimal-transport method.These methods target related distributional constraints but differ in divergence or model-class applicability.
Distributional constraints beyond fine-tuning: Inference-time reweighting changes sampling without fine-tuning, whereas fine-tuning pays compute upfront and preserves the base model’s sampling cost for downstream use.This trade-off is advantageous when a model will support conditional generation.
Reward fine-tuning and conditional generation: CGM connects to reward fine-tuning because it minimizes KL divergence to an exponentially tilted base model, and conditional generation corresponds to a near-unit indicator constraint.The conditional-generation limit drives the variational parameter toward infinity.
Algorithmic foundations: Unbiased stochastic estimates of the relax and reward losses permit optimization using samples, model probabilities, and score gradients without requiring differentiable constraints.The stated regularity assumptions justify exchanging integration and differentiation, while propositions establish unbiasedness.

B.2. Unbiased Gradient Estimates

The section develops unbiased stochastic gradient estimators for CGM-relax and CGM-reward despite parameter-dependent sampling distributions. The reward estimator uses score-function gradients, while the relax estimator combines importance weighting with independent stop-gradient samples.

Setting the proposal equal to the current model makes importance weights equal to 1 while their gradients become score functions.
Leave-one-out averages act as zero-expectation control variates that reduce estimator variance without changing unbiasedness.
The reward-gradient estimator is unbiased for the gradient of the reward loss.
The relax estimator draws independent samples from a stop-gradient distribution and differentiates importance-weighted violation estimates to handle the non-expectation penalty.
The relax-gradient estimator is unbiased for the gradient of the relax loss.

C.3. Connection Between the Relax and Reward Losses

The relax loss is a KL-regularized penalty problem whose solution approaches the maximum-entropy calibration solution as λ decreases. This connects relax optimization to reward-based calibration while clarifying finite-λ constraint behavior and model-class approximation.

As λ → 0, the relax problem approaches the maximum-entropy calibration solution under the stated assumptions.
The relax formulation minimizes squared constraint violation plus λ times KL divergence from the base distribution.
The relaxed problem has a unique dual solution αλ, even when the constraints are redundant, and a unique primal optimum under strict convexity.
For finite λ, the constraint violation is bounded and decreases at rate O(λ) when the relevant regularity conditions hold.
For expressive generative-model classes, the relax optimizer is expected to approximate the distributional penalty solution, while the reward loss is a divergence to a maximum-entropy tilt.

C.4. Consistency and Asymptotic Normality

The appendix establishes consistency and asymptotic normality for the empirical estimator of the reward-loss parameters. These guarantees require regularity assumptions and reveal larger variance when the base constraint probability approaches 0 or 1.

The proofs rely on uniform convergence of empirical dual objectives and domination conditions that permit differentiation and integration to be exchanged.
The empirical estimator bα_N is consistent for α* as the sample size N grows under Assumptions C.1, C.2, and C.4.
Under stronger conditions, bα_N is asymptotically normal when 2α* lies in the natural parameter space.
In the binary-threshold example, estimator variance is greatest when the base probability h̄ is near 0 or 1.
The asymptotic variance is governed by the inverse covariance of h under the maximum-entropy distribution.

D.1. Continuous-time Diffusion Models

For continuous-time diffusion models, Girsanov’s theorem supplies density ratios and KL computations needed by CGM. Under terminal constraints and an endpoint-independence assumption, the calibrated solution shifts the terminal distribution while preserving conditional paths.

Euler–Maruyama approximates diffusion sampling, with stronger fidelity expected as the number of time bins T increases.
Girsanov’s theorem provides conditions and formulas for densities between diffusion models, enabling the KL computations required by CGM.
When constraints depend only on the terminal state, the maximum-entropy solution has a closed-form diffusion characterization under the stated regularity conditions.
The endpoint-independence assumption is necessary for retaining the tractable initial distribution p_init; without it, the solution’s initial distribution cannot generally equal p_init.
The calibrated path distribution changes the terminal marginal while preserving the conditional path distribution given the endpoint.

E. Comparison to Augmented Lagrangian Method

The augmented Lagrangian (AL) comparison adapts AL to stochastic generative-model calibration by alternating parameter updates with dual-variable updates. On synthetic experiments, CGM-relax performs comparably while requiring fewer hyperparameters.

Method: CGM-AL alternates stochastic gradient updates to model parameters with Lagrange-multiplier updates because the AL subproblem has no closed-form minimizer.The update frequency ℓ is an additional algorithmic hyperparameter.
Experimental setup: The synthetic comparison covers rare-mode reweighting by a multiplicative factor, rare-mode reweighting to a fixed value, and increasing-dimensional independent Gaussian-mixture constraints.The number of constraints and KL distance to the base model increase with dimension.
Results: CGM-relax performs comparably to CGM-AL across rare-event and increasing-constraint-dimension experiments, while avoiding AL’s extra dual-update-frequency hyperparameter.The authors therefore prefer CGM-relax for practical reasons, while considering CGM-AL viable.
Results: In the first rare-event setting, CGM-AL reduces calibration error more for small base probabilities, whereas CGM-relax performs better in the second fixed-value setting.The two methods perform comparably as constraint dimension increases.

F. Additional Experimental Details

The additional experimental details describe shared optimization and hardware settings for the reported experiments.

Scope: The appendix provides further details on model classes, constraint functions and targets, hyperparameters, architectures, training procedures, and additional pre- and post-calibration samples..
Common training details: All experiments use Adam with default momentum parameters, cosine learning-rate decay, and a single H100 GPU.Additional common training configurations are reported in Table 1.

F.1. Synthetic Data Experiments

The synthetic experiments examine calibration across rare events and increasing numbers of constraints, with additional implementation details for diffusion and protein models. They show strong calibration improvements but expose sampling and rarity limits.

Rare event: For rare-mode probabilities as small as 10^-3, CGM reaches the target but requires approximately 5 nats of deviation from the base model.Below 10^-3, CGM-relax becomes highly seed-dependent and CGM-reward remains at the base model.
Implementation details: The appendix also documents diffusion-model parameterization, sampling schedules, protein-quality metrics, secondary-structure annotation, and experimental averaging procedures.These details include continuous-time conversion for Genie2 and self-consistency RMSD evaluation.

F.4. Calibrating TarFlow

TarFlow calibration targets a uniform distribution over six wildlife animal classes using image annotations and evaluates both class balance and image quality. CGM-relax substantially improves calibration but worsens FID and sample realism, whereas CGM-reward fails to calibrate the base model.

Constraint definition: The calibration target is a uniform distribution over six wildlife animals, but the constraint uses a five-dimensional one-hot encoding to avoid a poorly conditioned dual problem.Classifier confidence scores are not used for fine-tuning.
Evaluation: TarFlow calibration uses 5 × 10^3 annotated samples for class proportions and evaluates FID using 5 × 10^4 generated samples.Results are reported over three replicates with two standard errors.
Evaluation: FID is computed on wildlife-class samples, but it is imperfect because it favors animal-class compositions resembling the training distribution.The authors therefore also evaluate CGM-relax on a maximum-entropy reweighting of the training dataset with equal wildlife-class probabilities.
Results: TV distance to the uniform animal distribution falls from .306 to .101 with CGM-relax, while FID rises from 15.9 to 21.0.CGM-reward leaves miscalibration and FID roughly unchanged from the base model.
Sample quality: The calibrated model produces more dogs or cats outside the wildlife class and more blended-animal images, indicating worsened sample realism after calibration.These changes are observed in random generations after CGM-relax calibration.

F.6.2. EXAMPLE GENERATIONS

The section provides example generations before and after fine-tuning with CGM-relax using λ = 0.1, illustrating short stories produced for prompts involving recurring roles, characters, and objects.

The displayed examples are identified as samples from Gemma-2-9B-IT and include generations associated with CGM-relax fine-tuning at λ = 0.1.
Example generations include stories about pilots, doctors, lawyers, teachers, and their interactions with people, animals, or objects.
The samples repeatedly use simple narratives involving named characters, concrete settings, and everyday activities such as flying, driving, playing, and helping.