Score Correction for Generative Models with Probabilistic Constraints

We introduce DualScore, a framework for correcting score functions so that score-based generative models satisfy probabilistic constraints. These constraints are specified through the marginal distribution of a stochastic transformation of the modeled variable. We formulate this as a constrained KL-minimization problem, and optimize its dual, parameterizing the infinite-dimensional dual variable with a neural network. This yields an additive correction to the original score function that can be used directly for sampling via stochastic gradient Langevin dynamics or reverse diffusion sampling, without modifying the base model. We evaluate DualScore on synthetic tasks and two real-world applications: regularized nonparametric maximum likelihood estimation and the incorporation of class-level constraints such as fairness into pretrained image diffusion models.

Generative AI systems learn to produce new outputs that reproduce patterns seen in their training data. But due to the complexity of these models, and because of subtle biases in training data, model outputs may not match what users know or require. For example, an image generator trained on historical data may produce images of scientists skewed towards certain demographics. More generally, one might want the distribution of some attribute of generated outputs — such as demographic representation — to match a desired target, even when that attribute is not directly observable or completely predictable. We introduce DualScore, a lightweight method that makes a targeted correction to a pretrained generative model so that the distribution of these attributes follows a desired target, while keeping the adjusted model as close as possible to the original. It learns a small correction to the paths the model follows to generate data, leaving the original model completely intact and avoiding the need for retraining. We show that DualScore works well when attributes are uncertain or noisy, a setting where simpler correction methods can give incorrect results. This gives practitioners a practical way to adapt powerful generative models so their outputs better match domain knowledge or policy requirements.