EquiReg: Equivariance Regularized Diffusion for Inverse Problems
Journal:
arXiv
Published Date:
May 29, 2025
Abstract
Diffusion models represent the state-of-the-art for solving inverse problems
such as image restoration tasks. In the Bayesian framework, diffusion-based
inverse solvers incorporate a likelihood term to guide the prior sampling
process, generating data consistent with the posterior distribution. However,
due to the intractability of the likelihood term, many current methods rely on
isotropic Gaussian approximations, which lead to deviations from the data
manifold and result in inconsistent, unstable reconstructions. We propose
Equivariance Regularized (EquiReg) diffusion, a general framework for
regularizing posterior sampling in diffusion-based inverse problem solvers.
EquiReg enhances reconstructions by reweighting diffusion trajectories and
penalizing those that deviate from the data manifold. We define a new
distribution-dependent equivariance error, empirically identify functions that
exhibit low error for on-manifold samples and higher error for off-manifold
samples, and leverage these functions to regularize the diffusion sampling
process. When applied to a variety of solvers, EquiReg outperforms
state-of-the-art diffusion models in both linear and nonlinear image
restoration tasks, as well as in reconstructing partial differential equations.
