[2505.22973] EquiReg: Equivariance Regularized Diffusion for Inverse Problems

Computer Science > Machine Learning arXiv:2505.22973 (cs) [Submitted on 29 May 2025 (v1), last revised 1 Mar 2026 (this version, v2)] Title:EquiReg: Equivariance Regularized Diffusion for Inverse Problems Authors:Bahareh Tolooshams, Aditi Chandrashekar, Rayhan Zirvi, Abbas Mammadov, Jiachen Yao, Chuwei Wang, Anima Anandkumar View a PDF of the paper titled EquiReg: Equivariance Regularized Diffusion for Inverse Problems, by Bahareh Tolooshams and 6 other authors View PDF HTML (experimental) 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 highlight that such functions naturally emerge when training non-equivariant models with augment...