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.

翻译：扩散模型代表了解决图像复原等逆问题的最先进方法。基于扩散的逆求解器通过引入似然项来引导先验采样，从而生成与后验分布一致的数据。然而，由于似然项难以精确计算，现有方法大多依赖各向同性高斯近似，这可能导致估计值偏离数据流形，产生不一致且质量较差的复原结果。本文提出等变性正则化（EquiReg）扩散模型，这是一种通用的即插即用框架，通过惩罚偏离数据流形的采样轨迹来改进后验采样。EquiReg 形式化地定义了流形偏好等变函数，该函数对位于流形上的样本表现出较低的等变误差，而对偏离流形的样本则产生较高的误差，从而将采样引导至解空间中保持对称性的区域。我们指出，此类函数在通过数据增强训练非等变模型或处理具有对称性的数据时会自然涌现。EquiReg 在采样步骤减少和测量一致性约束较强的场景下尤其有效，而许多现有方法在此类条件下会出现严重的质量退化。通过将采样轨迹正则化至数据流形附近，EquiReg 隐式地加速了收敛过程，并实现了高质量的重建。在线性与非线性图像复原任务以及偏微分方程求解中，EquiReg 均展现出持续的性能提升。