Learning Normalized Energy Models for Linear Inverse Problems

Generative diffusion models can provide powerful prior probability models for inverse problems in imaging, but existing implementations suffer from two key limitations: $(i)$ the prior density is represented implicitly, and $(ii)$ they rely on likelihood approximations that introduce sampling biases. We address these challenges by introducing a new energy-based model trained for denoising with a covariance-based regularization term that enforces consistency across different measurement conditions. The trained model can compute normalized posterior densities for diverse linear inverse problems, without additional retraining or fine tuning. In addition to preserving the sampling capabilities of diffusion models, this enables previously unavailable capabilities: energy-guided adaptive sampling that adjusts schedules on-the-fly, unbiased Metropolis-Hastings correction steps, and blind estimation of the degradation operator via Bayes rule. We validate the method on multiple datasets (ImageNet, CelebA, AFHQ) and tasks (inpainting, deblurring), demonstrating competitive or superior performance to established baselines.

Many imaging problems (restoring a blurry image, filling in a missing region of an image, sharpening a low-resolution scan) share a common structure: we observe a corrupted version of a picture and want to recover what was really there. Modern AI methods tackle this by learning what natural images "look like" and combining that knowledge with the corruption process to make a best guess. But these methods face a trade-off. The flexible ones, which can in principle handle any type of corruption, rely on approximations that introduce errors. The accurate ones must be retrained from scratch for every new type of corruption. We trained a single AI model that learns to assign a probability score to any image given any corruption, rather than just guessing the answer directly. A single key idea — enforcing mathematical consistency across all corruption types during training — lets one model handle many problems at once. This unlocks capabilities previous methods cannot offer: the model can automatically figure out what kind of corruption affected an image, decide on its own how to best reconstruct it, and produce reliable uncertainty estimates, which is useful in scientific and medical imaging, where knowing how confident a reconstruction is matters as much as the reconstruction itself.