ICML2026

Score Correction for Generative Models with Probabilistic Constraints

Shishang Wu, Bingjing Tang, Vinayak A Rao

Abstract

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.