EquiReg: Equivariance Regularized Diffusion for Inverse Problems

2025-05-29

AI Summary

Diffusion models for inverse problems (e.g., image restoration) often resort to isotropic Gaussian approximations of the likelihood term due to its intractability, leading to posterior samples deviating from the underlying data manifold—resulting in inconsistent and unstable reconstructions. To address this, we propose an equivariance-regularized diffusion sampling framework. We first introduce a data-dependent equivariant error metric that geometrically distinguishes manifold-intrinsic from manifold-extrinsic samples. By reweighting diffusion trajectories, our method explicitly penalizes excursions outside the data manifold. Crucially, this manifold regularization is embedded directly into the Bayesian posterior sampling process—without modifying the diffusion model architecture. Extensive experiments across diverse linear and nonlinear image restoration tasks, as well as PDE-based inverse problems, demonstrate that our approach significantly outperforms state-of-the-art diffusion methods, achieving substantial improvements in reconstruction consistency, stability, and fidelity.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Learning with Manifolds
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Data quality aspects of human-annotated datasets
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

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.

Problem

• Improving diffusion models for inverse problems
• Reducing deviations from data manifold in reconstructions
• Enhancing stability and consistency in image restoration

Innovation

• Equivariance Regularized diffusion framework
• Reweights diffusion trajectories for consistency
• Penalizes deviations from data manifold