EquiReg: Equivariance Regularized Diffusion for Inverse Problems

TL;DR

EquiReg addresses the intractability of likelihood terms in diffusion-based inverse problems by introducing a distribution-dependent equivariance regularization that reweights reverse-diffusion trajectories to stay closer to the data manifold. It casts the reverse process as a Wasserstein-2 gradient flow and augments it with a regularizer derived from a manifold-focused equivariance error, implemented via two plug-in losses: Equi and EquiCon. The approach is general and compatible with multiple diffusion solvers, showing substantial performance gains in image restoration (linear and nonlinear tasks) and PDE reconstruction, including large improvements in perceptual metrics (e.g., FID, LPIPS) and reduced relative ℓ2 errors for Helmholtz and Navier–Stokes problems. This work advances diffusion-based inverse solving by integrating global symmetry properties into trajectory regularization, yielding more faithful reconstructions with improved robustness across noise levels.

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.

Figures (20)

• Figure 1: Ground truth, measurement, baseline, and reconstruction. R1 (left-to-right): super-resolution and motion deblur. R2: box and random inpainting. R3: PDE solving and phase retrieval.
• Figure 2: Equivariance Regularized (EquiReg) diffusion for solving inverse problems. a) EquiReg regularizes the posterior sampling trajectory, addressing the errors introduced due to the posterior mean expectation. b) The EquiReg loss function is lower for on- or near-manifold (clean) data and higher for off- or far-from-manifold (blurred and missing pixel) data.
• Figure 3: Illustration of off-manifold posterior expectation. a) Manifold interpretation. While samples from $p({\bm x}_0 | {\bm x}_t)$ may lie on the manifold, their posterior expectation, being a linear combination, can fall off the manifold. b) Distribution interpretation. In a bimodal distribution, the posterior expectation lies in a low probability region.
• Figure 4: MPE function examples.
• Figure 5: EquiReg consistently improves performance across a range of measurement noise levels. Gaussian deblur using FFHQ $256 \times 256$ for $\sigma_{{\bm y}} = [0, 0.8]$.

Theorems & Definitions (11)

• Definition 3.1: Equivariance
• Definition 3.2: Approximate Equivariance
• Definition 3.3: Equivariance Error
• Proposition 4.1
• Proposition 4.2
• Definition 4.1: Distribution-Dependent Equivariance Error
• Definition 4.2: Manifold Constrained Distribution-Dependent Equivariance Error
• Lemma G.1
• proof
• Remark 1