EquiReg: Equivariance Regularized Diffusion for Inverse Problems
Bahareh Tolooshams, Aditi Chandrashekar, Rayhan Zirvi, Abbas Mammadov, Jiachen Yao, Chuwei Wang, Anima Anandkumar
TL;DR
EquiReg introduces a regularization framework for diffusion models that enhances inverse problem solutions by guiding sampling toward data manifolds, improving reconstruction quality especially with fewer steps.
Contribution
The paper proposes EquiReg, a novel plug-and-play regularization method that leverages equivariance to improve diffusion-based inverse problem solving.
Findings
Improves image restoration quality with fewer sampling steps.
Enhances convergence speed in diffusion-based inverse solutions.
Consistently outperforms existing methods in various inverse tasks.
Abstract
Diffusion models represent the state-of-the-art for solving inverse problems such as image restoration tasks. Diffusion-based inverse solvers incorporate a likelihood term to guide prior sampling, generating data consistent with the posterior distribution. However, due to the intractability of the likelihood, most methods rely on isotropic Gaussian approximations, which can push estimates off the data manifold and produce inconsistent, poor reconstructions. We propose Equivariance Regularized (EquiReg) diffusion, a general plug-and-play framework that improves posterior sampling by penalizing trajectories that deviate from the data manifold. EquiReg formalizes manifold-preferential equivariant functions that exhibit low equivariance error for on-manifold samples and high error for off-manifold ones, thereby guiding sampling toward symmetry-preserving regions of the solution space. We…
Peer Reviews
Decision·Submitted to ICLR 2026
1. Novelty: The paper proposed a novel reward gradient guidance that leads to on-manifold sampling. Instead of digging deep into the underlying data distribution, the reward discriminates on-manifold samples from off-manifold samples simply using symmetry arising from data itself or the training process. 2. Effectiveness: Numerical experiments demonstrates that EquiReg loss could be easily incorporated into gradient guidance-based diffusion inverse problem solvers, and achieve better performanc
1. Though it is mentioned in the main paper that various details are deferred to the appendix, the appendix is not included in the submission, which substantially affects readability. 2. Theoretical insights of the EquiReg loss is not sufficiently explored. It remains unknown how EquiReg loss could affect generation consistency. The authors mentioned some insights through Wasserstein gradient flow. It will be helpful if it could be further discussed. 3. Though it is natural to use symmetry when
- Introduces equivariance-based regularization as a proxy for manifold consistency, connecting geometric symmetry with probabilistic sampling in a fresh way. - The method is architecture-agnostic and simple to implement that can be directly added to existing diffusion frameworks (DPS, PSLD, SITCOM) without retraining.
- The link between low equivariance error and on-manifold behavior is intuitive but not rigorously proven. - The approach assumes an explicit group $G$ (e.g., rotation, reflection), which may not exist or be meaningful for all tasks. - The method relies on pre-trained MPE encoders, but how to systematically obtain or generalize them is not well discussed. - The paper frequently refers to an Appendix for details of tasks, proofs, and additional experimental results, but the Appendix is not prov
1. The story is not too complex, so the message is clear. 2. It is model-agnostic, so many users can adopt EquiReg for their own framework. 3. It generally improves fidelity metrics like PSNR and LPIPS.
This paper might be sacrificing the true strength of diffusion models for inverse problems. The reason diffusion models are used (like in DPS - Diffusion Posterior **Sampling**) is that they are **Samplers**, not mean-estimators. The advantage is that they can sample many different solutions that are all good. This paper has no consideration for this. However, this paper just says EquiReg helps getting good PSNR and SSIM. This is obvious. Equivariance is a somewhat classical regularization metho
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsNumerical methods in inverse problems · Advanced Neuroimaging Techniques and Applications · Markov Chains and Monte Carlo Methods
MethodsDiffusion