Restrained Diffusion Bridges: Targeted Denoiser Training and Refinement

Diffusion models are often inaccurate in critical, data-sparse regions of the target distribution. Sampling data from restrained distributions, biased distributions used in classical enhanced sampling to concentrate sampling near chosen configurations, allows control over the location of training data, but alters the target distribution. We establish a remarkable cancellation that occurs when the forward noising process from restrained data is conditioned on its value at an appropriately chosen intermediate time. The result allows a pretrained model fine-tuned on restrained data to achieve high accuracy in any chosen region of the target distribution, regardless of how probable that region is, using standard denoising or flow matching objectives and without importance weights. When no unrestrained target data are available, Gibbs sampling allows inference from a model trained exclusively on restrained data.