SP$^3$: Spherical Priors for Plug-and-Play Restoration¶

会议: NeurIPS 2026
arXiv: 2606.16396
OpenReview: https://openreview.net/forum?id=F9SLbhP7uz
代码: https://github.com/facebookresearch/sphere-encoder
领域: 图像恢复
关键词: 图像恢复, 即插即用先验, 球面自编码器, 半二次分裂, 无梯度优化

一句话总结¶

提出基于球面自编码器（Sphere Encoder）先验的免求导即插即用（PnP）图像恢复算法 SP$^3$，通过紧凑球面隐空间投影与半二次分裂闭式数据保真项交替迭代，在保持与扩散模型相当的真实感与高频细节的同时，将推理速度大幅提升 33 至 630 倍，并赋予单步即得可用结果的“随时”（anytime）恢复特性。

研究背景与动机¶

在计算摄影与逆问题成像中，从退化观测中重构高质量自然图像通常被建模为最大后验概率（MAP）估计问题。尽管数据保真项依据前向物理退化模型通常具备明确解析表达，但高维自然图像的先验分布极其复杂且难以直接刻画。早期的即插即用（Plug-and-Play, PnP）框架通过将优化算法中的先验近端算子替换为预训练高斯去噪器，实现了先验模型与逆问题求解的解耦；然而传统判别式去噪器倾向于输出过度平滑的最小均方误差（MMSE）均值解，难以恢复逼真的高频纹理。

随后涌现的生成模型尝试利用结构化隐空间作为图像先验，例如在生成对抗网络（GAN）、变分自编码器（VAE）或标准化流的隐空间中通过梯度回传求解 MAP。然而，VAE 存在严重的“后验空洞”（posterior hole）问题，导致反复编解码时快速偏离自然流形；更致命的是，直接在隐空间进行梯度下降需要对非线性生成器执行昂贵且数值不稳定的端到端反向传播。近年来基于扩散与流匹配（Diffusion/Flow）的零采样方法虽能生成极具真实感的纹理，却高度依赖推理阶段数十至上千步的分数函数梯度回传与轨迹采样，巨大的计算开销使得实时应用难以承受。

本文切入的角度在于重新审视自编码器的流形投影特性：近期提出的球面自编码器（Sphere Encoder, SE）通过将隐变量约束于均匀致密超球面上，消除了后验空洞，能够作为极为稳健的流形投影算子，且编解码天然融合了编码器对 OOD 退化输入的感知映射与解码器的流形生成能力。核心 idea：将球面自编码器的加噪球形投影作为自然图像流形先验，通过半二次分裂（HQS）将其与闭式数据保真近端算子解耦交替迭代，构建一种完全免求导、稳定收敛且具备“随时”（anytime）恢复能力的即插即用图像重构范式。

方法详解¶

整体框架¶

SP$^3$ 的核心思路是将退化观测值 $y$ 的恢复任务解耦为两个交替子步骤：在图像流形先验步中，利用球面自编码器（SE）对当前估计图像进行编解码流形投影；在数据保真步中，利用前向退化算子 $A$ 的线性结构通过闭式正则化逆解保证测量一致性。整个流程完全避免了在深度神经网络内部计算梯度或执行反向传播。

方法的整体 pipeline 如下所示：从退化测量 $y$ 初始化起始图像 $x_0$，在每轮迭代中，首先经过 SE 的编码器与球面加噪归一化算子将图像映射到紧凑球面隐空间，再由解码器投影回自然图像流形得到先验建议 $x_{\text{prior}}$；随后结合前向物理模型计算闭式近端解 $x_k$ 更新图像，迭代直到收敛或满足“随时”（anytime）计算预算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["退化测量 y & 初始估计 x0"] --> B["1. 编码与加噪球面归一化<br/>vk = S(E(xk-1), σ)"]
    B --> C["2. 解码流形重投影<br/>x_prior = D(vk)"]
    C --> D["3. 闭式数据保真近端更新<br/>xk = prox_λ(x_prior)"]
    D -->|未达到步数 K| B
    D -->|达到迭代预算或即时终止| E["高质量恢复图像 xK"]

关键设计¶

1. 编码与加噪球面归一化：构建密实隐空间流形投影

针对传统 VAE 隐空间分布不均、存在后验空洞导致多次迭代严重发散的痛点，SP$^3$ 引入球面自编码器的球面化隐空间设计。编码器 $E$ 首先将输入图像提取为隐向量，随后通过均方根归一化（RMS normalization）函数 $f(\cdot)$ 投影至超球面。为了保证隐空间被完全覆盖并增强投影稳定性，在隐向量上注入可控高斯白噪声 $e \sim \mathcal{N}(0, I)$ 后再次执行球面化： $$S(v, \sigma) = f(v + \sigma \cdot e)$$ 其中 $\sigma$ 控制隐空间噪声扰动尺度。这一加噪球面化操作不仅有效抵御了输入退化带来的分布偏移（OOD），还促使投影算子在球面流形高概率区域稳定停留，避免了确定性自编码器在连续循环重投影时出现的模式漂移。

2. 解码流形重投影：免梯度的紧凑流形映射

传统基于隐空间的逆问题求解方法大多只使用解码器 $D$，通过反向传播对隐变量 $v$ 计算梯度优化 $\min_v \|AD(v) - y\|_2^2$，这不仅计算极其耗时，而且完全丢弃了联合训练中编码器所学的丰富结构先验。SP$^3$ 将编码器与解码器组合为统一的流形投影算子 $P(\cdot) = D(S(E(\cdot), \sigma))$，直接用于近似半二次分裂中的先验子问题： $$x_{\text{prior}} = P(x_{k-1}) = D(v_k)$$ 实验表明，单次 $P$ 运算即可将模糊、噪声等严重退化样本大幅拉回自然图像流形（KID 指标在第 1 步即急剧下降），使得算法从第 1 轮迭代开始就具备“随时”（anytime）产出锐利、合理细节的能力，同时省去了昂贵的反向自动微分。

3. 闭式数据保真近端更新：轻量高效的物理约束对齐

纯粹无约束的流形投影虽然能不断提高生成图像的真实感知质量，但随着迭代步数增加会不可避免地导致图像内容偏离真实测量观测。SP$^3$ 在半二次分裂（HQS）框架下引入二次惩罚项解耦变量，对数据保真子问题推导出具有闭式解析解的近端算子： $$x_k = \text{prox}_{\lambda}(x_{\text{prior}}) = (A^\top A + \lambda I)^{-1} (A^\top y + \lambda x_{\text{prior}})$$ 由于许多常见退化算子 $A$（如高斯去噪、去模糊、超分辨率、掩膜补全）在线性代数或傅里叶频域下均为对角阵或块循环阵，该矩阵求逆可借助快速傅里叶变换（FFT）或逐像素加权在毫秒级闭式完成。通过将先验建议 $x_{\text{prior}}$ 与观测约束高效融合，在确保物理保真度的同时实现了极高的计算吞吐量。

损失函数 / 训练策略¶

SP$^3$ 本身为完全训练免调整（zero-shot / plug-and-play）的推理算法，直接使用预训练的 Sphere Encoder，其预训练包含三个自监督重构目标：针对轻微噪声隐向量的平滑 L1 与感知损失 $\mathcal{L}_{\text{L1+perc.}}(D(v_n), x)$、针对较大噪声隐向量的一致性保持损失 $\mathcal{L}_{\text{L1+perc.}}(D(v_N), \text{sg}(D(v_n)))$ 以及隐向量余弦相似度约束 $\mathcal{L}_{\text{cos. sim.}}(v, E(D(v_N)))$。

在推理执行时，算法仅需调整两个超参数：数据保真正则化权重 $\lambda$ 与隐空间注入噪声强度 $\sigma$。算法针对不同任务采用差异化初始化 $x_0$：去噪和去模糊采用伴随算子 $A^\top y$；修补任务采用基于卷积的高效中值滤波近似；超分辨率则采用双三次插值上采样。

实验关键数据¶

主实验¶

论文在 AFHQ-Cat（256px）与 CelebA（128px）数据集上，对六类含噪逆问题（高斯去噪、高斯去模糊、超分辨率、框状修复、随机像素修复、画笔修复）与多种 SOTA 方法进行评测，涵盖经典 PnP 去噪器（DPIR、RED）、扩散模型（DPS、DDRM、DiffPIR）与流匹配方法（PnP-Flow、OT-ODE、D-Flow）。所有方法均使用单张 NVIDIA L40S GPU 测试单图运行时间。

下表列出 AFHQ-Cat 数据集上代表性任务的主实验对比数据（KID 均乘 $10^3$）：

数据集 / 任务	指标	本文 (SP$^3$, 20步)	SOTA 扩散模型 (DPS)	流匹配基线 (OT-ODE)	PnP流模型 (PnP-Flow)
AFHQ 去模糊	KID $\downarrow$	2.56	2.69	6.74	64.01
AFHQ 去模糊	LPIPS $\downarrow$	0.118	0.178	0.155	0.446
AFHQ 去模糊	PSNR $\uparrow$	25.75	24.72	26.79	25.91
AFHQ 去噪	KID $\downarrow$	4.09	3.63	7.71	36.74
AFHQ 去噪	LPIPS $\downarrow$	0.174	0.159	0.146	0.259
AFHQ 去噪	PSNR $\uparrow$	22.84	26.38	27.32	29.49
AFHQ 框状修复	KID $\downarrow$	6.68	3.91	8.53	58.89
AFHQ 框状修复	LPIPS $\downarrow$	0.103	0.107	0.129	0.328
AFHQ 框状修复	PSNR $\uparrow$	25.12	25.85	22.95	25.57
AFHQ 超分辨率	KID $\downarrow$	3.47	3.50	6.40	98.44
AFHQ 超分辨率	LPIPS $\downarrow$	0.151	0.201	0.173	0.666
AFHQ 超分辨率	PSNR $\uparrow$	24.24	24.32	26.34	18.09
单图平均耗时	Time (秒) $\downarrow$	0.28s (1步仅0.016s)	17.77s ($\times 63$)	2.29s ($\times 8.2$)	3.20s ($\times 11.4$)

消融实验¶

针对隐空间优化策略与隐空间先验结构设计的有效性，论文在 AFHQ-Cat 去模糊任务上进行了深入消融分析（KID 均乘 $10^3$）：

配置 / 机制	PSNR (dB) $\uparrow$	LPIPS $\downarrow$	KID $\downarrow$	单图耗时 (s) $\downarrow$	说明
SP$^3$ 完整模型 (SE 编解码流形投影 + HQS)	25.7	0.12	2.64	0.28s	编解码协同投影，免求导且极速收敛
S-GD (仅解码器 + 软球面高斯约束梯度下降)	18.5	0.58	34.3	1.32s	缺乏编码器引导，难以投影到清晰流形
S-PGD (仅解码器 + 投影梯度下降约束)	23.1	0.26	5.49	1.17s	求解易陷入局部极值，测量一致性差
SP$^3$ (1 步随时恢复)	22.4	0.18	4.85	0.016s	单步直接前向推断即可得到可用锐利图像
SP$^3$ (w/o 隐空间加噪 $\sigma=0$)	24.2	0.21	8.92	0.28s	确定性映射在多步投影下逐渐出现漂移
替换先验为标准 VAE (LDM / FLUX VAE)	<15.0	>0.70	>120.0	-	VAE 后验空洞导致连续编解码急剧发散崩溃

关键发现¶

速度与质量的 Pareto 最优：SP$^3$ 在 20 步时的感知质量（KID/LPIPS）全面达到甚至超越强基线 DPS 和 OT-ODE，同时运行速度比 DPS 快 60 倍以上，比 D-Flow 快 630 倍以上，比 OT-ODE 快 8 倍。
独有的 Anytime 特性：得益于 SE 紧凑的流形表征，仅需 1 步前向传播（耗时仅 16 毫秒）即可消除严重伪影并生成逼真结构；随着步数增加（3步、5步、10步、20步），图像细节与物理测量一致性平滑提升。
编码器的不可或缺性：仅使用解码器的 S-GD 和 S-PGD 不仅耗时更长（需逐次计算网络反向梯度），而且重建图像失真大、真实感差，证实了全自编码器结构协同投影在逆问题中的核心价值。

亮点与洞察¶

流形投影代替去噪先验：突破了传统 PnP 对高斯去噪器或扩散反向过程的依赖，开创性地将均一密实的球面自编码器用作即插即用先验，用一步前向流形映射替代复杂去噪迭代。
彻底消除推理反向传播：将复杂的非线性生成先验步近似为端到端前向编解码，将数据保真步归结为闭式二次逆解，整个恢复流程完全免求导，带来数十至数百倍的运行速度提升。
自稳定与防漂移机制：揭示了超球面上注入随机高斯扰动再重新球面化的数学机制，巧妙利用随机平滑将迭代轨迹稳定在流形致密区域，彻底攻克了传统 VAE 循环编解码模式崩溃的难题。

局限与展望¶

失真与感知指标的内在权衡：由于流形投影强制将图像拉向真实高频样本分布，算法追求极高感知真实感的同时牺牲了一定的像素级保真度（如 PSNR 不及以过度平滑著称的 MMSE 类去噪方法）。
非后验采样器：当前算法通过交替优化求解 MAP 点估计，无法像扩散后验采样方法（如 DPS）那样通过多次随机采样自然遍历完整的后验概率分布并评估不确定性。
向复杂非线性逆问题的拓展：闭式保真算子目前高度依赖线性退化算子 $A$；未来探索如何将免求导流形投影推广至高阶非线性成像（如相位恢复、非线性失真校正）是一大关键方向。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将球面自编码器引入即插即用逆问题求解，构建了完全免求导的高效流形投影范式。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 6 种主流退化任务、2 个权威数据集、大量 SOTA 扩散/流匹配基准对比以及收敛性理论与实验支撑。
写作质量: ⭐⭐⭐⭐⭐ 动机清晰，数学推导严密自洽，图表与收敛轨迹表达直观易读。
价值: ⭐⭐⭐⭐⭐ 为轻量级、实时逆问题成像与移动端/边缘端高质量图像恢复提供了极具应用前景的新范式。

SP\(^3\): Spherical Priors for Plug-and-Play Restoration¶

一句话总结¶

研究背景与动机¶

方法详解¶

整体框架¶

关键设计¶

损失函数 / 训练策略¶

实验关键数据¶

主实验¶

消融实验¶

关键发现¶

亮点与洞察¶

局限与展望¶

相关工作与启发¶

评分¶

数据集 / 任务	指标	本文 (SP\(^3\), 20步)	SOTA 扩散模型 (DPS)	流匹配基线 (OT-ODE)	PnP流模型 (PnP-Flow)
AFHQ 去模糊	KID \(\downarrow\)	2.56	2.69	6.74	64.01
AFHQ 去模糊	LPIPS \(\downarrow\)	0.118	0.178	0.155	0.446
AFHQ 去模糊	PSNR \(\uparrow\)	25.75	24.72	26.79	25.91
AFHQ 去噪	KID \(\downarrow\)	4.09	3.63	7.71	36.74
AFHQ 去噪	LPIPS \(\downarrow\)	0.174	0.159	0.146	0.259
AFHQ 去噪	PSNR \(\uparrow\)	22.84	26.38	27.32	29.49
AFHQ 框状修复	KID \(\downarrow\)	6.68	3.91	8.53	58.89
AFHQ 框状修复	LPIPS \(\downarrow\)	0.103	0.107	0.129	0.328
AFHQ 框状修复	PSNR \(\uparrow\)	25.12	25.85	22.95	25.57
AFHQ 超分辨率	KID \(\downarrow\)	3.47	3.50	6.40	98.44
AFHQ 超分辨率	LPIPS \(\downarrow\)	0.151	0.201	0.173	0.666
AFHQ 超分辨率	PSNR \(\uparrow\)	24.24	24.32	26.34	18.09
单图平均耗时	Time (秒) \(\downarrow\)	0.28s (1步仅0.016s)	17.77s (\(\times 63\))	2.29s (\(\times 8.2\))	3.20s (\(\times 11.4\))

配置 / 机制	PSNR (dB) \(\uparrow\)	LPIPS \(\downarrow\)	KID \(\downarrow\)	单图耗时 (s) \(\downarrow\)	说明
SP\(^3\) 完整模型 (SE 编解码流形投影 + HQS)	25.7	0.12	2.64	0.28s	编解码协同投影，免求导且极速收敛
S-GD (仅解码器 + 软球面高斯约束梯度下降)	18.5	0.58	34.3	1.32s	缺乏编码器引导，难以投影到清晰流形
S-PGD (仅解码器 + 投影梯度下降约束)	23.1	0.26	5.49	1.17s	求解易陷入局部极值，测量一致性差
SP\(^3\) (1 步随时恢复)	22.4	0.18	4.85	0.016s	单步直接前向推断即可得到可用锐利图像
SP\(^3\) (w/o 隐空间加噪 \(\sigma=0\))	24.2	0.21	8.92	0.28s	确定性映射在多步投影下逐渐出现漂移
替换先验为标准 VAE (LDM / FLUX VAE)	<15.0	>0.70	>120.0	-	VAE 后验空洞导致连续编解码急剧发散崩溃

SP\(^3\): Spherical Priors for Plug-and-Play Restoration¶

一句话总结¶

研究背景与动机¶

方法详解¶

整体框架¶

关键设计¶

损失函数 / 训练策略¶

实验关键数据¶

主实验¶

消融实验¶

关键发现¶

亮点与洞察¶

局限与展望¶

相关工作与启发¶

评分¶

相关论文¶