EquiReg: Equivariance Regularized Diffusion for Inverse Problems

4411 words

22 minutes

EquiReg: Equivariance Regularized Diffusion for Inverse Problems

2025-08-25

Research

EquiReg

/

inverse

/

CV

/

PnP

/

Model

/

DDPM

/

PDE

/

DIFFUSION

/

Paper

一、前言#

本文章旨在深度剖析 Tolosbams 等人提出的《EquiReg: Equivariance Regularized Diffusion for Inverse Problems》一文。該研究引入了一種名為 EquiReg 的新型正則化框架，旨在解決擴散模型處理逆問題時的求解器謎團——圖像偏離的問題。在傳統採樣過程中，由於對偶然分佈的近似處理不當，會導致生成的解構品質異於數據所在的理想流形（data manifold）。

EquiReg 的核心創見在於，它利用了「分佈依賴的等變性誤差」（distribution-dependent equivariance error）作為一項關鍵指標。此誤差能夠有效地監控樣本點位於流形上（on-manifold）還是偏離流形（off-manifold）。基於此，EquiReg 透過懲罰那些違反對稱性而偏離有效軌跡的解，有效地生成了約束在數據空間內的解。

實驗結果強有力地證明了該框架的有效性。在圖像修復任務中，EquiReg 顯著提升了重構的感知品質。例如，在超解析度任務中，應用於 DPS 模型的 Equi-DPS 使得 FID 指標改善了 59%，在科學計算領域，該框架則提升了圖像斷層掃描（PDE）重建的準確性。

綜上所述，EquiReg 不僅是一個理論上堅實、且在實踐中表現卓越的「即插即用」（plug-and-play）工具，更為重要的是，它提供了一種通用的方法論，透過利用數據的幾何對稱性來增強先進擴散模型的穩定性與真實度，適用於廣泛的科學與影像應用場景。

二、擴散模型逆問題求解中的「偏離流形」挑戰#

問題簡述

逆問題旨在從不完整且帶有雜訊的觀測中恢復信號。其數學形式可表示為恢復信號 $x^*$ ，其中觀測值 $y = A(x^*) + \nu$ ，A 是已知的測量算子， $\nu$ 是未知雜訊。從貝氏推斷的角度來看，目標是從後驗分佈 $p(x|y)$ 中進行採樣，根據貝氏定理，該分佈正比於似然 $p(y|x)$ 與先驗 $p(x)$ 的乘積，即 $p(x|y) \propto p(y|x)p(x)$ 。近年來，擴散模型已成為一種極其強大的學習先驗 $p(x)$ 的工具，能夠捕捉複雜數據（如自然圖像）的豐富結構。

難以處理的偶然分數

擴散模型透過模擬一個反向去噪的隨機微分方程（SDE）來生成樣本。在解決逆問題時，這個過程需要被引導以符合觀測數據 $y$ 。這通常需要計算條件分數函數 $\nabla_{x_t} \log p_t(x_t|y)$ ，它可以被分解為先驗分數 $\nabla_{x_t} \log p_t(x_t)$ 和似然分數 $\nabla_{x_t} \log p_t(y|x_t)$ 。前者可由預訓練的擴散模型直接提供，但後者，即偶然分數項，在時間步 $t > 0$ 時是無法直接計算的 (computationally intractable)，因為它需要對所有可能的乾淨數據 $x_0$ 進行積分。

合同高斯近似及其後果

為了解決此問題，當前主流方法普遍採用一種權宜之計：用後驗分佈 $p_t(x_t|y_0)$ 近似為一個各向同性高斯分佈。這一假設使得可以利用 Tweedie 公式來估計後期期望值 $x_{0|t} = E_p[x_0|x_t]$ ，即在設定含噪狀態 $x_t$ 的情況下對應的「最可能」的乾淨數據。

「偏離流形」的後期期望」問題

然而，這種近似設定導致了嚴重不穩定和產生偽影的根源。此問題被稱為「偏離流形的後期期望」（Off-Manifold Posterior Expectation），其根本原因在於，用單一的 $x_{0|t}$ 這一線性還原元與真實數據流形的非線性、高維幾何結構之間存在著深刻的衝突。

根據流形假說（Manifold Hypothesis），真實世界的高維數據（如自然圖像）實際上分佈在一個嵌入於高維空間中的低維流形 $M$ 上。這個流形本質上是非線性的（Manifold Hypothesis），且其幾何形狀複雜。

當對處於非流形上的點進行線性組合（如平均）時，其結果很可能落在流形之外。論文中的圖3直觀地展示了這一點：對於一個非凸流形，其期望值恰好落在兩個組之間的一個低機率區域。

因此，當求解器使用這個偏離流形的點 $x_{0|t}$ 來計算似然梯度時，它實質上是在用一個「假想」或「分佈外」的點來引導生成過程。這個錯誤的引導信號會污染整個求解軌跡，使其逐漸偏離真實的數據流形，最終導致重建結果中出現不切實際的紋理和視覺偽影。這揭示了為何傳統方法在圖像修復、去模糊等關心紋理或細節的設計方面，而在流形內插等內在的幾何導航中，這也為 EquiReg 提出一個利用幾何屬性（對稱性）作為正則化的解決方案提供了強有力的動機。

三、EquiReg 框架：一對稱性驅動的流形正則化方法#

EquiReg 框架的設計旨在直接應對上述「偏離流形」的挑戰。它不是試圖更精確地近似後驗，而是引入一個植基於數據幾何對稱性的正則化項，以確保採樣的軌跡始終保持在流形附近。

3.1 理論基礎：Wasserstein 梯度流視角#

該研究首先將條件採樣的反向擴散過程重新詮釋為一個在樣本軌跡上施加小擾動 $\Phi_t(\cdot)$ 的 Wasserstein-2 梯度流。這一視角為採樣過程提供了基於優化的嚴謹理論框架。根據命題 4.1，這個過程的形式被明確定義。在實際應用中，由於對後驗的近似不準確，採樣器實質上是在施加一個被「污染」的梯度。

Distribution-dependent equivariance errors. While the literature has primarily studied the equivariance properties of functions for general inputs, we propose a new equivariance error for approximately equivariant functions under the support of an input data distribution (see Definition 4.1).

Definition 4.1 (Distribution-Dependent Equivariance Error). The equivariance error of the function $f: \mathcal{Z} \rightarrow \mathcal{X}$ under the group $G$ and under the distribution $p$ is $\mathcal{L}_{EE}(f, p) = \sup_{g} \mathbb{E}_{z \sim p}[\|S_g(f(z)) - f(T_g(z))\|]$ where $T_g: \mathcal{Z} \rightarrow \mathcal{Z}$ is the set of transformations under the symmetry group $G$ , $S_g$ is the transformation corresponding to $T_g$ operating the output space $S_g: \mathcal{X} \rightarrow \mathcal{X}$ , for $g \in G$ .

EquiReg 的核心思想是校正這個被污染的梯度。它引入一個正則化項 $\epsilon_t$ ，其作用是重新加權該梯度，降低來自不可信區域（偏離流形）的貢獻，同時放大來自可靠區域（在流形上）的貢獻，如論文中的公式4所示。

\tilde{\Phi}(\rho, t) = \beta_{F-t}[Z_t^{-1} \int \rho(x) \hat{\phi}_e(x, t) e^{\frac{\mathcal{R}(x)}{\hat{\phi}_e(x,t)}} dx + \frac{1}{2} \int \rho(x) \log \rho(x) dx], \quad(4)

3.2 核心概念：依賴分佈的等變性誤差#

為了設計出有效的正則化項 $\epsilon_t$ ，該研究引入了其核心的創新概念。首先，等變性 (Equivariance) 被定義為一個函數在輸入經過某種對稱群變換（如旋轉、反射）後，其輸出也隨著相應變換的性質。

論文的關鍵突破在於提出分佈依賴的等變性誤差 (Distribution-Dependent Equivariance Error)，如定義 4.1 所示。這與傳統的等變性概念不同，它不是衡量函數在所有可能輸入上的誤差，而是在一個特定數據分佈 $p$ 下的期望誤差。

基於此，論文定義了流形普適等變函數 (Manifold-Preferential Equivariant, MPE)：這類函數被特意設計成當測量對象在流形上的數據（即 $p_{data}$ ）具有很低的等變性誤差，而對於偏離流形的數據則具有很高的誤差，如圖 2b 和圖 4 所示。等變性誤差本身成為了衡量樣本與流形距離的代理指標 (proxy)。

3.3 MPE 函數的識別：訓練誘導與數據固有對稱性#

論文提出了兩種尋找 MPE 函數的策略：

訓練誘導等變性 (Training-Induced Equivariance)：MPE 屬性可以在標準神經網路架構（如自編碼器）中隱式地產生，特別是當它們在具有內在對稱性的數據集（如自然圖像）上訓練，並被引導到對稱性的數據增強時。該研究利用了潛在擴散模型 (LDM) 中預訓練的 VAE，並展示了其對於乾淨圖像的等變性誤差很低，而對於受損圖像的等變性誤差則高漲（圖 4a）。
數據固有等變性 (Data-Inherent Equivariance)：在物理學等領域，物理系統的解常被微分方程 (PDE) 本身就具有對稱性。在這些情境上訓練的神經網路，如傅立葉神經算子 (FNO)，可以學習到這些對稱性。該研究證明，訓練好的 FNO 對於 PDE 解數據同樣表現出 MPE 屬性（圖 4b）。

3.4 實際應用：EquiLoss 與 EquiCon 正則化器#

基於 MPE 的概念，論文設計了兩種具體的損失函數，可作為正則化項：

EquiLoss $\mathcal{R}_s(x) = \mathbb{E}_g[d(S_g(x_0), f(x_0))]$ ：其中 $f$ 是 MPE 函數（例如自編碼器 $E$ 或解碼器 $D$ ），該損失直接懲罰對等變性的違反。
EquiCon $\mathcal{R}_c(z_t) = \mathbb{E}_g[d(S_g(z_t), D_{0|t}(z_0))]$ ：這是一個更為精巧的損失，它強制執行一個重要的一致性，它確保一個經過變換的樣本在解碼、逆變換和重新編碼後，能夠恢復到其原始的潛在表示。這隱式地將樣本約束在由自編碼器定義的流形上。

這些正則化項的整合方式非常靈活，以 Equi-DPS 算法（算法1）為例，EquiReg 項被簡單地添加為一個額外的梯度下降步驟（第 8 行），緊隨標準的似然引導步驟（第 7 行）之後，使其成為一個極易於集成的即插即用模組。

EquiReg 對於那些可能產生非流形解的步驟特別有效。於那些方法在每一歩執行一個局部線性組合的估計。例如，在流形的鞍形點附近，這些操作可能會導致嚴重的錯誤。相比之下，EquiReg 透過對稱性施加了一個全局性的幾何約束，將整個求解過程引導回一個更合理的流形子空間。例如，在人臉重建中，如果圖像缺少某個特徵（如一隻帶有耳環的耳朵），基於對稱性的 EquiReg 能夠從數據庫中一個更正確的正常耳廓來恢復樣本。這種方式更為系統和整體，它調整了採樣過程使得它能夠得到一個保持對稱性的狀態空間的整體解，而不是僅僅被一個具有真實數據的最近鄰所吸引。這種方法不僅僅是為了保持對稱性的一致性，而是提供了一個基於數據對稱性的整體解決方案。

四、實驗驗證與性能分析#

該研究透過在圖像修復和 PDE 重建兩大類任務上的大量實驗，全面驗證了 EquiReg 框架的有效性。

下表總結了 EquiReg 在關鍵任務和指標上相對於基線模型的性能提升，直觀地展示了其帶來的顯著改進。

實驗模型	任務	數據集	指標	基線模型性能	EquiReg 性能	改善幅度
DPS	超解析度 (x4)	FFHQ	FID ↓	114.9	41.71	63.7%
ReSample	運動去模糊	FFHQ	LPIPS ↓	0.197	0.078	60.4%
ReSample	隨機修補	FFHQ	LPIPS ↓	0.115	0.047	59.1%
PSLD	方塊修補	FFHQ	LPIPS ↓	0.158	0.092	41.8%
FunDPS	亥姆霍茲 (逆問題)	PDE 數據	相對 $L_2$ 誤差 ↓	17.16%	15.91%	7.3%
FunDPS	納維-斯托克斯 (逆問題)	PDE 數據	相對 $L_2$ 誤差 ↓	8.48%	7.84%	7.5%

4.1 圖像修復任務：一項綜合評估#

根據論文中的表 1、表 2 和表 3，EquiReg 在多個層面展現了其優越性。

跨模型的普適性：EquiReg 在所有測試的求解器上都帶來了一致的性能提升，包括基於潛在空間的 ReSample 和 PSLD，以及基於像素空間的 DPS。這證明了該框架的通用性和適應性。
感知品質的顯著提升：一個重要的發現是，EquiReg 在感知指標（如 FID 和 LPIPS）上的改進尤為顯著，通常超過了在像素級指標（如 PSNR）上的提升。這有力地表明，EquiReg 生成的圖像不僅在像素上更接近真實值，更重要的是，它們在視覺上更真實、結構更連貫。
穩定採樣過程：在諸如圖像修補和超解析度等高度依賴結構先驗的任務中，EquiReg 的優勢尤為突出。在這些任務中，基線模型常常因產生模糊或不一致的內容而失敗。如圖 6 所示，PSNR 分佈直方圖顯示 EquiReg 減少了失敗案例的數量，這明顯提高了重建的穩定性。
對稱性的穩健性：圖 5 的分析顯示，即使在測量雜訊水平較高的情況下，EquiReg 的優勢依然保持甚至增強。這表明它同時也是一個強大的去噪正則化器。

這些在 FID 和 LPIPS 指標上持續且不成比例的巨大改進並非偶然。這些指標本身被設計用來捕捉自然圖像的統計特性和高層次結構，因此它們可以被視為衡量生成圖像是否「在流形上」的隱式標準。一個位於真實數據流形上的圖像，其深度特徵分佈應與真實圖像相似；而一個偏離流形的圖像則會產生異常的特徵。因此，當 EquiReg 大幅改善 FID 和 LPIPS 時，這直接證明了其正則化成功地將採樣軌跡約束在了數據流形附近，為論文的核心假設提供了強而有力的實驗證據。

4.2 科學應用：偏微分方程重建#

表 4 的結果顯示，EquiReg 的原理成功地從圖像領域遷移到了科學計算領域。分析表明，其性能提升在一些問題（從觀測結果恢復初始條件）中尤為顯著。這些問題是典型的非適定 (ill-posed) 問題，極度依賴強大的先驗知識。Equi-FunDPS 在亥姆霍茲和納維-斯托克斯逆問題上實現了約 7-8% 的相對 $L_2$ 誤差降低，這在該領域是一個實質性的進步。這一結果不僅驗證了「數據固有等變性」方法的有效性，也展示了 EquiReg 原理的廣泛應用潛力。

五、批評性評估與未來展望#

5.1 貢獻與意義總結#

該研究的主要貢獻可歸納為以下三點：

一個通用的正則化框架：EquiReg 提供了一個統一、純理論基礎堅實的框架來正則化擴散求解器，超越了以往那些臨時性或局部的修正方法。
一個新穎的流形代理指標：該研究引入了分佈依賴的等變性誤差和 MPE 函數的概念，這是一種巧妙且有效的方法，可以在不重建流形的情況下檢測並懲罰偏離流形的行為。
已驗證的頂尖性能：實驗證明，該框架在廣泛的任務、模型和數據領域均取得了顯著的性能提升。

其總體意義在於，為解決一個已知且根本性的技術難題，提供了一個實用、強大且易於實現的工具，從而提升了主流生成模型的性能。

5.2 已知局限性與未來研究方向#

作者在結論中坦誠地指出了該方法的局限性：

依賴先驗品質：EquiReg 只能正則化給定的擴散先驗，無法修復一個本身就有缺陷的先驗模型。
任務特定的設計：在應用時，需要為每個新任務手動選擇一個合適的對稱群（例如，為人臉選擇反射對稱，為 ImageNet 選擇旋轉對稱）和一個 MPE 函數。這在一定程度上限制了其完全自動化應用。

基於這些局限性，未來的研究方向也變得清晰。一個關鍵的探索方向是從數據中自動發現對稱性和 MPE 函數，這可能需要藉助深度學習或表示學習領域的技術。另一個方向是探索除等變性之外的其他全局數據屬性，看其是否也能作為有效的正則化器。

儘管 EquiReg 功能強大，但其對預定義對稱群和預訓練 MPE 函數的依賴也引入了新的潛在失效點：一個不恰當的選擇可能會產生反效果，將採樣器引導至一個錯誤的流形。例如，如果將反射對稱的正則化器應用於非對稱的數據集（如手寫體像），將會強加一個錯誤的先驗，從而透過強制不自然的對稱性來破壞重建結果。同樣，如果作為 MPE 函數的自編碼器訓練不佳或存在偏見，EquiCon 損失會將採樣過程引導至這個有缺陷的流形，而非真實的數據流形。

這種「流形誤導」的風險為該研究的結論增添了重要的批判性視角。它表明，應用 EquiReg 需要謹慎的領域知識，並對所選的對稱性和 MPE 函數進行驗證。這也為未來的研究提出了一個重要問題：如何設計對錯誤指定的對稱性或不完美的 MPE 模型更具穩健性的正則化器。