Calibrating Generative Models to Distributional Constraints

This paper addresses the miscalibration of generative models by framing calibration as a constrained optimization problem and introducing two tractable fine-tuning objectives—the relax loss and reward loss—that effectively align sampling distributions with desired statistical constraints across diverse applications and large-scale models.

The Gist

Imagine you have a master chef (a Generative Model) who is incredibly talented at cooking. This chef has learned to cook thousands of dishes by tasting millions of recipes from a specific cookbook (the Training Data).

However, there's a problem. Even though the chef is skilled, they have developed some bad habits or "miscalibrations":

• They might only cook pasta, even though the cookbook had plenty of pizza and sushi.
• They might make every pizza look exactly the same, missing the variety of the original recipes.
• In a storybook context, they might always make the "doctor" character a man and the "nurse" a woman, even though real life is more balanced.

This paper introduces a new method called CGM (Calibrating Generative Models) to fix these habits without firing the chef or making them start from scratch.

The Core Problem: The "Recipe" vs. The "Result"

The authors explain that while the chef knows how to cook individual dishes well, the overall menu they produce doesn't match the desired statistics. Maybe you want 50% pasta and 50% pizza, but the chef keeps making 90% pasta.

In technical terms, the paper calls this miscalibration. The statistics of the chef's output (the sampling distribution) don't match the desired values (the constraints).

The Solution: A Gentle Nudge, Not a Rewrite

The authors propose a way to "fine-tune" the chef. They don't want to replace the chef's entire brain (which would be expensive and risky). Instead, they want to find the closest possible version of the chef that does follow the new rules.

They frame this as a math problem: "Find the new chef who is as similar as possible to the old chef (to keep their unique style) but who strictly follows the new menu proportions."

To solve this, they invented two specific "training techniques" (algorithms):

1. The "Relax Loss" Method (The Penalty Box)

Think of this as a strict diet coach.

• How it works: The coach tells the chef, "You can keep cooking your way, but every time you serve a plate that doesn't match the 50/50 rule, you get a penalty point."
• The Goal: The chef tries to minimize their penalty points while trying to stay as close to their original cooking style as possible.
• Why it's good: It's very flexible. It can handle hundreds of rules at once (e.g., "50% pasta, 30% pizza, 20% sushi, and also make sure no dish is too salty"). The paper shows this works even for massive chefs (models with billions of parameters).

2. The "Reward Loss" Method (The Reward System)

Think of this as a gamified training session.

• How it works: Instead of just punishing mistakes, the coach calculates a "reward score" based on how well the chef is doing. If the chef serves a balanced plate, they get points. If they serve a lopsided one, they get fewer points.
• The Goal: The chef tries to maximize their total score.
• The Catch: This method is great for simple rules but can get confused if you give it too many rules at once (like trying to balance 100 different ingredients simultaneously). The paper found it struggles when the number of constraints gets too high.

Real-World Tests: Where Did They Try This?

The authors tested their methods on three very different types of "chefs":

1. Protein Design (The Molecular Chef):

   • The Problem: These models were designing proteins (the building blocks of life) that looked too similar to each other and lacked the natural variety found in nature.
   • The Fix: They used the "Relax" method to force the model to generate a more diverse mix of protein shapes (like alpha-helices and beta-strands) that matched real biological data.
   • Result: The proteins became much more diverse and realistic, without becoming "broken" or unusable.

2. Image Generation (The Visual Artist):

   • The Problem: An image model trained to draw animals was obsessed with drawing lions and leopards, rarely drawing foxes or wolves, even though the training data had all of them.
   • The Fix: They calibrated the model to draw all six types of wild animals in equal numbers.
   • Result: The model started drawing foxes and wolves much more often. The images were still realistic, though the paper notes a slight trade-off where some images looked a bit "blended" (like a lion-wolf mix) as a side effect of the adjustment.

3. Language Models (The Storyteller):

   • The Problem: A large language model was writing children's stories where the "lawyer" was almost always a man and the "nurse" was almost always a woman, reinforcing real-world biases.
   • The Fix: They calibrated the model to ensure that for every profession, the gender of the character was balanced (50/50).
   • Result: The model started writing stories with female lawyers and male nurses just as often as the traditional pairings, without the stories becoming nonsensical or losing their quality.

The Bottom Line

The paper claims that CGM is a powerful, general-purpose tool. It allows us to take a pre-trained AI model and gently steer its overall output to match specific statistical goals (like fairness, diversity, or biological accuracy) without having to retrain the whole model from scratch.

• CGM-Relax is the heavy lifter: It handles hundreds of complex rules at once and works on the biggest models.
• CGM-Reward is the specialist: It's great for simpler tasks and avoids some of the complexity of tuning parameters, but it struggles with too many rules.

The authors emphasize that while this solves many calibration problems, it's not a magic wand. It still requires careful tuning, and in some extreme cases (like trying to generate extremely rare events), it might need a lot of computing power to get perfect results. But for most practical applications, it successfully brings the AI's output back in line with what we actually want.

Technical Summary

Technical Summary: Calibrating Generative Models to Distributional Constraints

Problem Statement

Generative models frequently suffer from miscalibration, where the statistics of the sampling distribution (e.g., class frequencies, secondary structure proportions, or gender representation) deviate systematically from desired values. This issue manifests across domains: image models exhibit mode collapse; language models reinforce societal biases; and protein/DNA models generate structures with atypical substructure frequencies. These errors often stem from dataset imbalances, suboptimal training dynamics, or post-hoc adjustments like low-temperature sampling.

The authors frame calibration as a constrained optimization problem. Given a base generative model $p_{\theta_{base}}(x)$, a statistic $h(x)$, and a target expectation $h^*$, the goal is to find a new distribution $p_\theta$ that satisfies $E_{p_\theta}[h(x)] = h^*$ while minimizing the Kullback-Leibler (KL) divergence to the base model:
$$p_{\theta^*} = \arg \min_{p_\theta} D_{KL}(p_\theta \parallel p_{\theta_{base}}) \quad \text{s.t.} \quad E_{p_\theta}[h(x)] = h^*$$
This formulation seeks the model closest to the base distribution that satisfies the calibration constraint. However, solving this exactly is intractable because both the objective and the constraint involve expectations that are difficult to compute for high-dimensional generative models.

Methodology

To address the intractability, the paper introduces two surrogate objectives for fine-tuning, termed CGM-relax and CGM-reward. Both methods rely on stochastic optimization using unbiased gradient estimates derived from samples.

1. CGM-relax (Relax Loss)

This approach replaces the hard constraint with a penalty term. The objective function is:
$$L_{relax}(\theta) = \| E_{p_\theta}[h(x)] - h^* \|^2 + \lambda D_{KL}(p_\theta \parallel p_{\theta_{base}})$$

• Constraint Violation ($L_{viol}$): The squared $\ell_2$ norm of the difference between the empirical statistic and the target. The authors derive an unbiased estimator for this term using a bias-corrected sample variance formula.
• KL Divergence ($L_{KL}$): Measured via Monte Carlo sampling.
• Optimization: The method uses importance sampling to compute unbiased gradients. Specifically, it samples from the current model $p_\theta$ (with stop-gradient) and reweights samples to estimate gradients of the KL term and the constraint violation. A hyperparameter $\lambda$ trades off between satisfying the constraint and staying close to the base model.

2. CGM-reward (Reward Loss)

This approach connects the calibration problem to the Maximum Entropy principle. The solution to the maximum entropy problem (minimizing KL to the base model subject to moment constraints) is an exponential tilt of the base model: $p_{\alpha^*}(x) \propto p_{\theta_{base}}(x) \exp(\alpha^{*\top} h(x))$.

• Estimating $\alpha^*$: Since the true maximum entropy distribution is intractable, the method first estimates the optimal reward parameters $\alpha^*$ by solving a dual optimization problem using $N$ samples from the base model.
• Loss Function: The fine-tuning objective is the KL divergence between the current model $p_\theta$ and the estimated maximum entropy distribution $p_{\hat{\alpha}_N}$:
  $$L_{reward}(\theta) = D_{KL}(p_\theta \parallel p_{\hat{\alpha}_N}) = D_{KL}(p_\theta \parallel p_{\theta_{base}}) + E_{p_\theta}[-r_{\hat{\alpha}_N}(x)] + C$$
  This is equivalent to a reward fine-tuning problem where the "reward" is derived from the constraints.
• Optimization: Similar to CGM-relax, it uses score-function gradient estimators with leave-one-out baselines to reduce variance.

Gradient Estimation

Both methods employ score-function gradient estimators (REINFORCE-style) combined with importance sampling. To handle the dependency of the expectation on $\theta$, the authors use a "stop-gradient" trick where samples are drawn from a fixed version of the current model, and gradients are computed via the density ratio (score function). Variance reduction is achieved using leave-one-out baselines.

Key Contributions

1. Formulation: Framing generative model calibration as a constrained KL minimization problem.
2. Algorithms: Proposing two practical fine-tuning algorithms (CGM-relax and CGM-reward) that approximate the solution to the constrained problem using stochastic optimization.
3. Scalability: Demonstrating that these methods can handle hundreds of simultaneous constraints and models with up to 9 billion parameters.
4. Theoretical Connection: Establishing the link between calibration, maximum entropy, and reward fine-tuning, showing that CGM-reward targets the maximum entropy solution while CGM-relax targets a regularized version.

Results

The authors evaluate the methods across three diverse applications:

1. Protein Design (Genie2 and ESM3):

   • Goal: Calibrate secondary structure content (alpha-helices and beta-strands) to match natural protein statistics (CATH dataset).
   • Constraints: Up to 99 simultaneous constraints (bivariate CDFs).
   • Outcome: CGM-relax reduced the distance to natural protein distributions by >4x for Genie2 and >2x for ESM3. CGM-reward achieved more modest improvements, partly due to the difficulty of estimating $\alpha^*$ with high-dimensional constraints. Both methods maintained biophysical plausibility, though with a slight increase in "design failures."

2. Image Generation (TarFlow):

   • Goal: Balance class proportions within the "wildlife" category of the AFHQ dataset (e.g., equalizing lions, tigers, foxes, wolves).
   • Constraints: 5 simultaneous constraints.
   • Outcome: CGM-relax reduced the total variation distance of class proportions from 0.306 to 0.101. While sample realism (FID) decreased slightly, the calibrated model produced more diverse and balanced generations. CGM-reward failed to reduce constraint violation significantly.

3. Language Modeling (Gemma-2-9B-IT):

   • Goal: Reduce gender bias in generated children's stories for various professions (e.g., ensuring equal probability of male/female lawyers).
   • Constraints: 8 profession-conditional constraints.
   • Outcome: Both methods reduced gender imbalance. CGM-relax reduced imbalance by more than threefold relative to the baseline, while CGM-reward showed a statistically significant but smaller improvement. The fine-tuned models maintained story quality (no degradation in log-probability).

Synthetic Experiments:

• Rare Events: CGM successfully upweighted rare modes in Gaussian mixtures, even when the base probability was as low as $10^{-3}$.
• High Dimensions: CGM-relax effectively handled up to $10^3$ constraints. In contrast, CGM-reward became infeasible for constraints $>30$ due to the difficulty of solving the empirical maximum entropy problem (Equation 7).

Significance and Claims

The paper claims that CGM provides a practical framework for calibrating high-dimensional generative models to meet distributional constraints without requiring changes to the model architecture or retraining from scratch.

• Robustness: CGM-relax is shown to be more robust than CGM-reward, particularly in high-dimensional constraint settings where estimating the maximum entropy parameters becomes intractable.
• Trade-offs: The methods offer a navigable trade-off between calibration accuracy and fidelity to the base model (controlled by $\lambda$).
• Limitations: The authors acknowledge that residual errors remain, especially for very rare events. Furthermore, the framework currently relies on models with tractable likelihoods (diffusion models, normalizing flows, autoregressive models), limiting its direct application to implicit models like GANs.

The work positions calibration as a distinct and fruitful research direction, separate from standard reward fine-tuning, by using constraints themselves as the driving signal for distributional alignment.