EquiReg: Equivariance Regularized Diffusion for Inverse Problems

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Tolooshams, Bahareh, Chandrashekar, Aditi, Zirvi, Rayhan, Mammadov, Abbas, Yao, Jiachen, Wang, Chuwei, Anandkumar, Anima
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866910036162248704
author Tolooshams, Bahareh
Chandrashekar, Aditi
Zirvi, Rayhan
Mammadov, Abbas
Yao, Jiachen
Wang, Chuwei
Anandkumar, Anima
author_facet Tolooshams, Bahareh
Chandrashekar, Aditi
Zirvi, Rayhan
Mammadov, Abbas
Yao, Jiachen
Wang, Chuwei
Anandkumar, Anima
contents 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 augmentation or on data with symmetries. EquiReg is particularly effective under reduced sampling and measurement consistency steps, where many methods suffer severe quality degradation. By regularizing trajectories toward the manifold, EquiReg implicitly accelerates convergence and enables high-quality reconstructions. EquiReg consistently improves performance in linear and nonlinear image restoration tasks and solving partial differential equations.
format Preprint
id arxiv_https___arxiv_org_abs_2505_22973
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle EquiReg: Equivariance Regularized Diffusion for Inverse Problems
Tolooshams, Bahareh
Chandrashekar, Aditi
Zirvi, Rayhan
Mammadov, Abbas
Yao, Jiachen
Wang, Chuwei
Anandkumar, Anima
Machine Learning
Artificial Intelligence
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 augmentation or on data with symmetries. EquiReg is particularly effective under reduced sampling and measurement consistency steps, where many methods suffer severe quality degradation. By regularizing trajectories toward the manifold, EquiReg implicitly accelerates convergence and enables high-quality reconstructions. EquiReg consistently improves performance in linear and nonlinear image restoration tasks and solving partial differential equations.
title EquiReg: Equivariance Regularized Diffusion for Inverse Problems
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2505.22973