学习线性逆问题的归一化能量模型

# 学习线性反问题的归一化能量模型  
来源：https://arxiv.org/html/2605.15487  

###### 摘要  
生成扩散模型可以为成像中的反问题提供强大的先验概率模型，但现有实现存在两个关键局限性：(i)(i)先验密度是隐式表示的，(ii)(ii)它们依赖引入采样偏差的似然近似。我们通过引入一种新的基于能量的模型来解决这些挑战，该模型针对去噪进行训练，并采用基于协方差的正则化项，强制不同测量条件下的一致性。训练后的模型可以计算各种线性反问题的归一化后验密度，无需额外的重新训练或微调。除了保留扩散模型的采样能力外，这还实现了以前不可用的功能：能量引导的自适应采样（可动态调整调度）、无偏的Metropolis-Hastings校正步骤，以及通过贝叶斯规则对退化算子的盲估计。我们在多个数据集（ImageNet、CelebA、AFHQ）和任务（图像修复、去模糊）上验证了该方法，展示了与已有基线相比具有竞争力或更优的性能。代码可在 https://github.com/nzilberstein/Anisotropic-energy-Model 获取。  

机器学习，ICML  

## 1 引言  
生成扩散模型 (Sohl-Dickstein et al., 2015 (https://arxiv.org/html/2605.15487#bib.bib36); Ho et al., 2020 (https://arxiv.org/html/2605.15487#bib.bib23); Song and Ermon, 2019 (https://arxiv.org/html/2605.15487#bib.bib151); Song et al., 2021b (https://arxiv.org/html/2605.15487#bib.bib37)) 在图像生成建模方面取得了显著成功。除了采样，这类模型的主要用途是为反问题提供先验概率 (Daras et al., 2024 (https://arxiv.org/html/2605.15487#bib.bib92))，例如去模糊或图像修复。为此，当前解决方案分为两类：*贝叶斯*方法 (Kadkhodaie and Simoncelli, 2021 (https://arxiv.org/html/2605.15487#bib.bib46); Kawar et al., 2022 (https://arxiv.org/html/2605.15487#bib.bib49); Chung et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib54)) 将预训练无条件扩散模型中嵌入的先验与通过贝叶斯规则得到的测量似然相结合。虽然这种分离提供了灵活性，但似然项是难以处理的，因此计算解依赖于可能导致采样偏差的近似 (Chung et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib54); Bruna and Han, 2024 (https://arxiv.org/html/2605.15487#bib.bib129))。*回归*方法 (Saharia et al., 2022 (https://arxiv.org/html/2605.15487#bib.bib79); Liu et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib122); Delbracio and Milanfar, 2023 (https://arxiv.org/html/2605.15487#bib.bib124); Negrel et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib104); Elata et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib120)) 学习直接连接测量和干净信号的条件扩散模型。这类方法避免了似然近似，但代价是拥有一个独立于所考虑退化的唯一先验。此外，两种方法都是基于分数的，使得后验密度评估计算成本高昂。我们通过开发一种基于能量的模型 (EBM) 来解决这些挑战，该模型学习显式且归一化的先验和后验密度，同时提供对后验均值和后验样本的有效访问。单个训练好的模型可用于以一致的方式恢复被各种退化损坏的图像。这种统一的公式利用了线性反问题与各向异性去噪之间的联系，也解锁了我们下面详述的新颖能力。  

已有一些出版物将未归一化的 EBM 与无条件扩散模型相结合 (Du et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib128); Thornton et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib125))，利用显式能量值进行组合生成。虽然最近的提议使用能实现适当归一化的正则化项 (Guth et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib114); Yu et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib115); Plainer et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib126))，但这些方法目前仅限于各向同性噪声，无法正确处理解决线性反问题时常见的相关噪声。在这里，我们引入了各向异性协方差分数匹配 (A-CSM)，这是对 (Guth et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib114)) 中“双重分数匹配”框架的扩展，可以处理各向异性（有色）噪声。我们开发了一种新颖的正则化器，通过约束能量相对于噪声协方差的梯度来强制一致性。我们的主要贡献是：  

- • 一种从 Fokker-Planck 方程推导出的新颖训练目标，强制跨噪声协方差矩阵的能量一致性；  
- • 一种互补架构，通过新颖的噪声嵌入支持各向异性噪声调度；  
- • 实验证明，学习到的归一化密度实现了新能力：能量引导的自适应采样、无偏 MCMC 校正步骤以及盲反问题解；  
- • 在多个数据集（ImageNet、CelebA、AFHQ）和任务（图像修复、去模糊）上验证模型的实验，表明我们单一的统一模型可以达到与已建立的扩散求解器竞争或更优的性能。  

## 2 背景与相关工作  

### 2.1 扩散模型  
扩散模型 (Sohl-Dickstein et al., 2015 (https://arxiv.org/html/2605.15487#bib.bib36); Ho et al., 2020 (https://arxiv.org/html/2605.15487#bib.bib23); Song et al., 2021b (https://arxiv.org/html/2605.15487#bib.bib37)) 来源于两个连续时间马尔可夫过程：(i)(i) 逐渐用噪声破坏干净数据的前向过程，以及 (ii)(ii) 通过迭代去噪生成样本的反向过程。在本工作中，我们假设前向过程采用方差爆炸 (VE) 公式 (Song et al., 2021b (https://arxiv.org/html/2605.15487#bib.bib37))：  
x_t = x + σ_t v, v ∼ N(0, I)。  (1)  
对应的离散反向过程为：  
x_{t-1} = x_t + δσ_t^2 ∇ log p(x_t) + δσ_t v_t，  (2)  
其中 δσ_t^2 = σ_t^2 - σ_{t-1}^2，未知的得分函数 ∇_{x_t} log p(x_t) 由神经网络 s_θ(x_t, t) 近似，该网络通过去噪分数匹配 (DSM) (Vincent, 2011 (https://arxiv.org/html/2605.15487#bib.bib21); Raphan and Simoncelli, 2011 (https://arxiv.org/html/2605.15487#bib.bib141)) 训练。DSM 损失提供了 KL(p(x) ∥ p_θ(x)) 的上界，因此对于特定的损失加权等同于最大似然训练 (Song et al., 2021a (https://arxiv.org/html/2605.15487#bib.bib68))。值得注意的是，该模型提供了 p(x) 的*隐式*表示，通过一组针对所有噪声水平学习的得分函数来表示 (Kadkhodaie and Simoncelli, 2021 (https://arxiv.org/html/2605.15487#bib.bib46))。  

##### 超越各向同性噪声。  
虽然大多数扩散模型假设各向同性前向过程，但最近的一些研究探索了更一般的过程，包括具有向量值参数化的各向异性过程 (Daras et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib105); Hoogeboom and Salimans, 2023 (https://arxiv.org/html/2605.15487#bib.bib106); Bansal et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib107); Gerdes et al., 2024 (https://arxiv.org/html/2605.15487#bib.bib109); Chen et al., 2024 (https://arxiv.org/html/2605.15487#bib.bib108); Asthana et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib143))，以及可学习 (Bartosh et al., 2024 (https://arxiv.org/html/2605.15487#bib.bib110); Sahoo et al., 2024 (https://arxiv.org/html/2605.15487#bib.bib111)) 和高阶过程 (Dockhorn et al., 2021 (https://arxiv.org/html/2605.15487#bib.bib25); Singhal et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib113))。在这里，我们利用具有协方差调度 (Σ_t)_{t=0}^T 的各向异性前向扩散过程：  
x_t = x + Σ_t^{1/2} v, v ∼ N(0, I)，  (3)  
对应的反向过程为：  
x_{t-1} = x_t + δΣ_t ∇ log p(x_t) + δΣ_t^{1/2} v_t。  (4)  

### 2.2 线性反问题  
线性反问题包括从通过退化算子 H（假设为线性）获得的噪声测量 y 中估计未知信号 x：  
y = H x + σ v, v ∼ N(0, I)。  (5)  
虽然大多数解旨在获得点估计（通常是后验均值），但也可以考虑从后验 p(x|y) ∝ p(y|x) p(x) 中采样，其中 p(y|x) 是测量模型 (5 (https://arxiv.org/html/2605.15487#S2.E5))，p(x) 是扩散模型中隐式嵌入的先验。文献中存在两类解：*贝叶斯方法*，它通过扩散模型近似先验 p(x) 并在推理时与测量模型结合；以及*回归方法*，它通过有监督回归在配对数据 (x,y) 上直接近似后验 p(x|y)。  

##### 贝叶斯方法。  
这些方法通过使用贝叶斯规则，将反向过程 (2 (https://arxiv.org/html/2605.15487#S2.E2)) 以观测 y 为条件来生成后验样本：  
∇_{x_t} log p(x_t|y) = ∇_{x_t} log p(x_t) + ∇_{x_t} log p(y|x_t)。  (6)  
虽然第一项通常使用预训练的*各向同性*扩散模型计算，但第二项（所谓的“引导”）由于需要高维积分而难以处理：p(y|x_t) = ∫ p(y|x) p(x|x_t) dx。之前的方法 (Kadkhodaie and Simoncelli, 2021 (https://arxiv.org/html/2605.15487#bib.bib46); Kawar et al., 2022 (https://arxiv.org/html/2605.15487#bib.bib49); Chung et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib54)) 采用了近似。但即使是简单的高斯近似 (Chung et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib54); Song et al., 2022 (https://arxiv.org/html/2605.15487#bib.bib55))，计算成本也很高，因为引导项涉及学习到的得分的*雅可比矩阵*。除了这些引导方法，最近的研究探索了基于优化的方法 (Mardani et al., 2024 (https://arxiv.org/html/2605.15487#bib.bib6); Zilberstein et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib135); Zhu et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib41); Feng et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib95))，将后验采样视为随机优化；以及基于序列蒙特卡洛的方法 (Wu et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib136))，利用粒子方法。  

##### 基于回归的方法。  
第二类方法聚焦于直接学习条件得分 (6 (https://arxiv.org/html/2605.15487#S2.E6))。这类工作学习噪声测量 y 与干净信号 x 之间的扩散模型，通过将退化吸收到训练数据中，有效地将反问题视为条件生成 (Liu et al., 2023 (https://arxiv.org/html/2605.15487#bib.bib122); Delbracio and Milanfar, 2023 (https://arxiv.org/html/2605.15487#bib.bib124); Negrel et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib104); Hu et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib119))。虽然这些方法并未明确将退化算子 H 输入得分网络，但最近的一些论文 (Elata et al., 2025 (https://arxiv.org/html/2605.15487#bib.bib120); Terris et al., 2026 (https://arxiv.org/html/2605.15487#bib.bib121)) 提出了条件去噪器 E[x|y, H] 的退化感知参数化。然而，它们依赖于一种条件机制，需要在架构的每个层次上在环境空间和特征空间之间进行重复映射，增加了计算复杂度。  

## 3 通过各向异性去噪学习 EBM  
我们的目标是学习一个*单一*模型 p_θ(x|y)，该模型显式近似来自各种不同退化算子的测量 y 的后验*密度*，这与之前的基于分数的方法不同。学习后验密度不仅通过微分提供了条件得分 (6 (https://arxiv.org/html/2605.15487#S2.E6)) 的访问，而且还开启了我们在第 4 节 (https://arxiv.org/html/2605.15487#S4) 中描述的若干新应用。我们首先在第 3.1 节 (https://arxiv.org/html/2605.15487#S3.SS1) 中描述线性反问题与各向异性去噪之间的联系，以及如何利用这一点来学习能量模型。我们在第 3.2 节 (https://arxiv.org/html/2605.15487#S3.SS2) 中介绍我们的广义双重分数匹配损失，在第 3.3 节 (https://arxiv.org/html/2605.15487#S3.SS3) 中介绍我们的各向异性能量架构。最后，我们在第 3.4 节 (https://arxiv.org/html/2605.15487#S3.SS4) 中展示我们的能量模型如何同时用作去噪器和线性反问题的后验采样器。  

### 3.1 基于能量的线性反问题求解器  
我们利用线性退化与有色噪声之间的等价性，展示一种以不同噪声协方差为条件的基于能量的模型如何实现对归一化后验密度、后验均值和后验样本的访问。  

参考图注  
图 1：图像空间（左）和协方差空间（右）中可能路径的示意图。各向同性模型仅限于对角线的重新参数化协方差调度（橙色），而各向异性模型可以探索更广泛的路径。