跳转至

PRISM: Latent Composition Consistency for Single-Image Reflection Removal

会议: ECCV 2026
论文: ECCV 官方页面
PDF: ECCV Open Access
领域: 图像恢复
关键词: 单图反光去除 / 隐空间流匹配 / 组合一致性 / 层对比分离 / FLUX

一句话总结

PRISM 发现预训练 VAE 隐空间中透射层与反光层的互相关性较像素空间骤降十倍,据此将单图反光去除重构为隐空间线性分离任务,基于 FLUX 骨干实现单步流匹配同时恢复透射与反光分量,并结合跨样本反光交换的一致性约束与分块对比分离损失,在六大基准上全面刷新 SOTA。

研究背景与动机

透过玻璃窗等透明介质拍摄时,场景常被叠加的非目标虚影所污染。单图反光去除(Single-Image Reflection Removal, SIRR)旨在从混合图像 \(I = T + R\) 中分离出干净的透射层 \(T\) 并剥离反光层 \(R\)。然而该任务在数学上是严重病态的逆问题,给定一个观测图像存在无穷多组有效解。传统方法依赖手工先验、相对平滑假设或可学习残差项;近年基于深度学习的双流网络(如 DSRNet、DSIT、RDNet)虽然显式建模双分支交互,但均根植于 sRGB 像素空间。由于非线性相机响应曲线与色调映射的存在,像素级加性模型仅是粗糙近似,网络在分离语义特征的同时极易将反光误当残差吸收,导致实际反光分支退化为近乎全黑的无意义输出。

即便 RAW 传感器域具有更真实的物理线性叠加性,但严苛的硬件依赖使其无法应用于海量既有数字影像。更深层的矛盾在于:在 sRGB 像素空间中,透射层与反光层具有高度的统计纠缠。作者在 SIR2 数据集的 454 组真实图对上统计发现,透射与反光的像素余弦相似度均值高达 0.74;然而,将其映射到预训练生成模型(FLUX VAE)的隐空间后,两者余弦相似度骤降至 0.07。这种广义的解相关特性表明,预训练隐空间天然为图像层的正交分解提供了远优于像素空间的线性分离几何结构。

既然隐空间展现出卓越的互低相干性,反光去除便不再需要笨重的双分支迭代或复杂残差修补。本文的核心 idea 是:将单图反光去除重构为隐空间流匹配线性分离问题,通过单步速度场直接预测反光负偏置,同时利用隐空间低相关特性设计跨样本反光置换的循环组合一致性(LCC)与免真值监督的分块层对比分离(LCS)损失,实现透射与反光的高保真单次前向解耦。

方法详解

整体框架

PRISM 首先利用预训练 FLUX VAE 编码器 \(\mathcal{E}\) 将输入的退化混合图像 \(I\) 投影为隐空间表示 \(z_I = \mathcal{E}(I) \in \mathbb{R}^{16 \times \frac{H}{8} \times \frac{W}{8}}\)。基于隐空间近似加性假设 \(z_I \approx z_T + z_R\),模型以 FLUX.2 Klein (4B) 的多模态扩散 Transformer(MM-DiT)作为流匹配速度场网络 \(v_\theta\),直接从干净输入隐变量 \(z_I\) 出发执行单步前向传播。网络输出的速度向量恰好对应负反光分量,从而通过一次前向计算同时解出预测的透射隐变量 \(\hat{z}_T = z_I + v_\theta(z_I, 0)\) 和反光隐变量 \(\hat{z}_R = -v_\theta(z_I, 0)\),并通过 VAE 解码器 \(\mathcal{D}\) 还原为图像。

为防止网络退化为对场景特定信息的过拟合虚假分离,PRISM 在批次维度构建了辅助训练支路:将不同样本分离出的反光隐变量循环置换叠加到透射隐变量上构造合成混合隐变量,通过第二次前向推理施加循环一致性监督;同时利用分块池化提取局部特征,施加层对比分离损失,在无需显式反光像素真值的前提下稳固正交解耦。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入混合图像 I"] --> B["VAE 编码器 ℰ<br/>投影为隐变量 z_I"]
    B --> C["隐空间流匹配单步线性分离<br/>MM-DiT 预测速度场 v_θ(z_I, 0)"]
    C --> D["单步向量解析<br/>z_T_hat = z_I + v_θ, z_R_hat = -v_θ"]
    D --> E["隐空间交叉组合一致性<br/>跨样本轮换拼接合成 z_I_tilde 并二次解耦"]
    D --> F["分块层对比分离<br/>Patch-Pool 局部正负样本对比拉伸"]
    D --> G["VAE 解码器 𝒟<br/>重构透射图像 T_hat 与反光 R_hat"]

关键设计

1. 隐空间流匹配单步线性分离:摆脱双流冗余并实现双分量单步解耦

在像素空间中,为了同时估计透射和反光,现有先进方法通常需要设计对称或循环的双流网络(如双分支 LSTM 或交互式 Transformer),计算开销倍增且容易出现反光分支塌缩。针对这一痛点,PRISM 依据隐空间极低的特征相干性(余弦相似度仅 0.07),建立了隐变量近似线性分离假定 \(z_I \approx z_T + z_R\)。由此,从混合隐变量 \(z_I\) 输运至目标透射隐变量 \(z_T\) 的理想流匹配速度目标即为:

\[v^* = z_T - z_I \approx -z_R\]

这一数学特性带来了极其简洁优雅的推论:流匹配速度场的真实目标恰好与反光隐变量严格互为相反数。PRISM 采用 FLUX.2 Klein (4B) MM-DiT 骨干,并在条件输入端使用全零嵌入替代文本交叉注意力以实现无条件图像恢复。不同于常规扩散模型从高斯噪声开始多步迭代去噪,PRISM 严格遵循直线常速度场特性,在 \(t=0\) 处直接以未加噪的 \(z_I\) 为流起点,仅需单次前向传播便可同时解析出透射 \(\hat{z}_T = z_I + v_\theta(z_I, 0)\) 与反光 \(\hat{z}_R = -v_\theta(z_I, 0)\)。此设计不仅免去了独立的双流架构,而且完全避免了多步采样带来的高额推理延迟。

2. 隐空间交叉组合一致性:利用相干性洼地构造合成循环监督

仅靠透射真值的单向重建损失,容易引发退化解——网络可能将场景特有的结构细节混杂记录在 \(\hat{z}_T\) 与 \(\hat{z}_R\) 中,造成解耦不彻底。在像素空间中,由于强相关与光照重叠,将图 A 的反光硬叠加到图 B 的透射往往产生物理不自然的瑕疵图像;然而在低相干的 VAE 隐空间中,不同样本的透射与反光正交组合后解码,能生成高度逼真且语义合法的合成混合图。PRISM 敏锐利用该特性提出了 Latent Composition Consistency (LCC) 机制。

在一个批次的 \(N\) 个样本内,模型将样本 \(i\) 的第一阶段预测透射 \(\hat{z}_T^i\) 与相邻样本的反光 \(\hat{z}_R^{(i \bmod N) + 1}\) 组合,合成伪混合潜变量 \(\tilde{z}_I^i = \hat{z}_T^i + \hat{z}_R^{(i \bmod N) + 1}\)。随后将其送入 MM-DiT 进行第二次前向推理,得到解耦结果 \(\tilde{z}_T^i\) 与 \(\tilde{z}_R^i\)。为了防止平凡解和模式坍塌(即网络将所有样本的第一阶段预测收缩至相同常数),PRISM 在第一阶段目标上引入截断梯度算子 \(\texttt{sg}(\cdot)\),通过双向循环一致性损失驱动模型在任意隐空间组合下均保持稳定的分离能力:

\[\mathcal{L}_{\text{cycle}} = \frac{1}{N} \sum_{i=1}^{N} \frac{1}{2} \left( \|\tilde{z}_T^i - \texttt{sg}(\hat{z}_T^i)\|_2^2 + \|\tilde{z}_R^i - \texttt{sg}(\hat{z}_R^{(i \bmod N) + 1})\|_2^2 \right)\]

3. 分块层对比分离:规避非线性误差的免反光真值语义解耦

尽管隐空间加性假设具有优异的实用价值,但由于 VAE 编码器的内在非线性,严格的数学等式 \(\mathcal{E}(I - T) = \mathcal{E}(I) - \mathcal{E}(T)\) 并不成立。若直接构造显式反光回归目标强制拟合像素残差,必然会将非线性失真传播给速度场。为此,PRISM 提出了 Layer Contrastive Separation (LCS) 损失,通过对比学习引导透射分量与反光分量在特征流形上拉开语义距离。

针对反光分布在空间上的非均匀遮挡特性,PRISM 放弃全局平均池化,而是将隐空间特征图划分为无重叠的 \(p \times p\) 局部块(默认 \(p=4\)),对每个局部块独立平均池化并做 \(\ell_2\) 归一化。以二次循环解出的透射块 \(f_k(\tilde{z}_T)\) 作为锚点,正样本集合 \(\mathcal{P}\) 包含第一步透射预测块 \(\hat{z}_T\) 与真值透射块 \(z_T^{\text{gt}}\),负样本集合 \(\mathcal{N}\) 则汇集对应的反光预测块 \(\hat{z}_R\)、循环反光块 \(\tilde{z}_R\) 以及真值反光块 \(z_R^{\text{gt}}\)。通过温度系数 \(\tau=0.1\) 的 Patch-level InfoNCE 损失:

\[\mathcal{L}_{\text{LCS}} = - \mathbb{E} \left[ \frac{1}{K} \sum_{k=1}^{K} \log \frac{\sum_{z^p \in \mathcal{P}} \exp(\sigma_k(z^p))}{\sum_{z^p \in \mathcal{P}} \exp(\sigma_k(z^p)) + \sum_{z^n \in \mathcal{N}} \exp(\sigma_k(z^n))} \right]\]

该机制在不强求刚性像素值对齐的前提下,强力驱逐了透射表示中残留的反光语义。

损失函数 / 训练策略

PRISM 的端到端训练总损失由四部分加权构成:

\[\mathcal{L} = \lambda_{\text{lat}} \mathcal{L}_{\text{latent}} + \lambda_{\text{pix}} \mathcal{L}_{\text{pixel}} + \lambda_{\text{cycle}} \mathcal{L}_{\text{cycle}} + \lambda_{\text{LCS}} \mathcal{L}_{\text{LCS}}\]

其中,隐空间速度场损失为 \(\mathcal{L}_{\text{latent}} = \|v_\theta(z_I, 0) - v^*\|_2^2\);解码后的像素级保真度由 \(\mathcal{L}_{\text{pixel}} = \|\mathcal{D}(\hat{z}_T) - T\|_1 + \lambda_{\text{lpips}} \mathcal{L}_{\text{LPIPS}}(\mathcal{D}(\hat{z}_T), T)\) 提供监督。超参数设定为 \(\lambda_{\text{lat}} = 1\)、\(\lambda_{\text{pix}} = 1\)、\(\lambda_{\text{lpips}} = 2\)、\(\lambda_{\text{cycle}} = 1\)、\(\lambda_{\text{LCS}} = 0.1\)。

模型基于 PyTorch 在单张 NVIDIA A6000 GPU 上训练,批大小为 2,输入随机裁剪为 \(512 \times 512\),采用 bfloat16 混合精度与 AdamW 优化器,学习率经由余弦退火从 \(5 \times 10^{-5}\) 平滑衰减至 \(1 \times 10^{-6}\),总计训练 50,000 次迭代。训练混合了 PASCAL VOC 合成对(7,643 对)、Zhang 等人的真实对(90 对)与 Nature 真实对(200 对),按 0.6 : 0.2 : 0.2 比例平衡采样。

实验关键数据

主实验

PRISM 在六大权威真实反射基准数据集上与 ERRNet、IBCLN、LASIRR、YTMT、RobustSIRR、DSRNet、RRW、DSIT、RDNet 及基于扩散先验的最新方法 DAI (AAAI 2026) 进行了严密对比。

表 1 汇总了六大基准上的 PSNR 与 SSIM 评估结果。PRISM 在全部六个数据集上均刷新了 PSNR 的 SOTA 纪录,尤其在最具挑战的真实场景 Real 上超越 RDNet 达到 1.56 dB 的显著增益,在 Wild 上大幅提升 1.84 dB。

方法 Real (20) PSNR / SSIM Object (200) PSNR / SSIM Postcard (199) PSNR / SSIM Wild (55) PSNR / SSIM Nature (20) PSNR / SSIM SIR2 (454) PSNR / SSIM
ERRNet (CVPR'19) 22.89 / 0.803 24.87 / 0.896 22.04 / 0.876 24.25 / 0.853 20.58 / 0.756 23.41 / 0.874
IBCLN (CVPR'20) 21.86 / 0.762 24.87 / 0.893 23.39 / 0.875 24.71 / 0.886 23.57 / 0.786 24.08 / 0.875
LASIRR (ICCV'21) 23.34 / 0.812 24.36 / 0.898 23.72 / 0.903 25.73 / 0.902 23.45 / 0.808 24.18 / 0.893
YTMT (NeurIPS'21) 23.26 / 0.806 24.87 / 0.896 22.91 / 0.884 25.48 / 0.890 23.85 / 0.810 24.04 / 0.880
RobustSIRR (CVPR'23) 23.61 / 0.835 24.90 / 0.917 19.91 / 0.868 23.67 / 0.884 20.97 / 0.764 22.54 / 0.884
DSRNet (ICCV'23) 23.91 / 0.818 26.74 / 0.920 24.83 / 0.911 26.11 / 0.906 25.22 / 0.832 25.72 / 0.907
RRW (CVPR'24) 23.82 / 0.817 26.55 / 0.927 24.03 / 0.903 26.51 / 0.913 25.96 / 0.843 25.40 / 0.908
DSIT (NeurIPS'24) 25.22 / 0.836 27.27 / 0.932 25.58 / 0.922 27.40 / 0.918 26.77 / 0.847 26.50 / 0.919
RDNet (CVPR'25) 25.58 / 0.846 26.78 / 0.921 26.33 / 0.922 27.70 / 0.915 26.21 / 0.842 26.63 / 0.915
DAI (AAAI'26) 25.24 / 0.840 27.26 / 0.920 27.30 / 0.922 27.44 / 0.912 27.05 / 0.846 27.30 / 0.919
PRISM (本文) 27.14 / 0.853 27.61 / 0.925 27.85 / 0.906 29.54 / 0.932 27.35 / 0.853 27.95 / 0.918

为了验证野外真实未知分布下的零样本泛化能力,表 2 给出了在 OpenRR 1K 测试集(100 张未参与任何训练的真实反光图,分辨率 \(512 \times 512\))上的直接推理评测,测试设备为单张 RTX 4090 GPU。

方法 PSNR ↑ SSIM ↑ LPIPS ↓ DISTS ↓ NIQE ↓ 推理耗时 (ms) ↓ 参数量
YTMT 25.16 0.927 0.096 0.059 3.6241 75.68 58.47M
DSRNet 26.24 0.935 0.074 0.048 3.6573 288.01 144.65M
DSIT 26.01 0.903 0.083 0.051 3.5536 222.91 326.96M
RRW 25.55 0.929 0.086 0.055 3.6825 35.88 27.99M
RDNet 26.81 0.939 0.074 0.048 3.7713 102.26 314.03M
DAI 26.04 0.903 0.085 0.070 3.6197 138.70 1.30B
PRISM (本文) 27.23 0.893 0.064 0.041 3.5512 104.23 3.97B

消融实验

表 3 评估了循环一致性损失 \(\mathcal{L}_{\text{cycle}}\) 与层对比分离损失 \(\mathcal{L}_{\text{LCS}}\) 的正交增益。基线仅使用 \(\mathcal{L}_{\text{latent}} + \mathcal{L}_{\text{pixel}}\)。

\(\mathcal{L}_{\text{cycle}}\) \(\mathcal{L}_{\text{LCS}}\) Real (20) PSNR / SSIM Nature (20) PSNR / SSIM SIR2 (454) PSNR / SSIM 平均 PSNR / SSIM
\(\times\) \(\times\) 26.67 / 0.847 27.24 / 0.850 27.68 / 0.916 27.20 / 0.871
\(\checkmark\) \(\times\) 26.90 / 0.853 27.35 / 0.851 27.86 / 0.917 27.37 / 0.874
\(\times\) \(\checkmark\) 26.74 / 0.846 27.34 / 0.852 27.90 / 0.917 27.33 / 0.872
\(\checkmark\) \(\checkmark\) 27.14 / 0.853 27.35 / 0.853 27.95 / 0.918 27.48 / 0.875

表 4 探讨了层对比分离中不同空间特征表示粒度对解耦性能的影响。

空间表征策略 Real (20) PSNR / SSIM Nature (20) PSNR / SSIM SIR2 (454) PSNR / SSIM 平均 PSNR / SSIM 说明
Patch (\(p=2\)) 27.03 / 0.851 27.35 / 0.853 27.94 / 0.917 27.44 / 0.874 块过小带来局部对比噪声
Patch (\(p=4\), 默认) 27.14 / 0.853 27.35 / 0.853 27.95 / 0.918 27.48 / 0.875 兼顾局部细节与表征稳定性
Patch (\(p=8\)) 26.83 / 0.849 27.22 / 0.852 27.93 / 0.917 27.33 / 0.873 感受野过大损失空间异质线索
Global Avg. Pool 26.73 / 0.845 27.21 / 0.852 27.94 / 0.918 27.29 / 0.872 空间信息坍缩,表现最差

关键发现

  • 双向循环约束缺一不可:在循环一致性监督目标消融中,单独仅监督透射循环 \(\tilde{z}_T\)(平均 27.34 dB)或仅监督反光循环 \(\tilde{z}_R\)(平均 27.36 dB)性能均明显低于双向联合监督(平均 27.48 dB)。透射分支确保内容在噪声干扰下的鲁棒留存,反光分支则约束分离分量符合物理加性结构。
  • 感知-失真权衡(Perception-Distortion Tradeoff)的体现:PRISM 在感知类度量指标上全面领跑(OpenRR 1K 上 LPIPS 相对 RDNet 大幅降低 14.0%,DISTS 降低 14.6%,NIQE 降至 3.5512),在生成细节更自然纯净的同时,由于局部高频真实感重建不完全与 GT 像素窗对齐,使得 Object 与 Postcard 上的 SSIM 略低于强平滑平摊的传统方法。
  • 像素级反光分支退化被打破:定性对比显示,由于像素空间缺乏结构化加性约束,此前 SOTA(如 YTMT、DSRNet、DSIT、RDNet)预测的反光层几乎退化为纯黑或接近零值,而 PRISM 能够清晰重构出空间连贯、结构完整的反光图案。

亮点与洞察

  • 发现并论证了 VAE 隐空间的低互相干性优势:用统计数据证明预训练隐空间中透射与反光的余弦相似度仅为 0.07(像素空间为 0.74),直接化解了 sRGB 域非线性混合的解耦难题,启发将传统底层病态逆问题升维至生成隐空间求解。
  • 单步流匹配对偶分离机制:敏锐抓住 \(v^* = z_T - z_I \approx -z_R\) 的对偶数学关系,利用单次前向推理同时得到透射和反光,既摒弃了多阶段双流的冗余结构,又无需耗费数十步扩散采样。
  • Stop-Gradient 赋能的跨图隐空间换光置换:利用低相干隐空间支持任意图像跨界拼装的独特性质,通过截断梯度的循环一致性训练迫使模型对任意反光混合均具备不变分解力,彻底切断了模型记忆样本特征的捷径。

局限与展望

  • 密集强反光饱和下的假设失效:当反光区域极度明亮且光强饱和(例如 \(\|R\|_1 / \|I\|_1 > 0.5\))时,混合图像隐变量 \(z_I\) 会大幅偏向反光子空间,隐空间低相干假定被削弱,导致透射层仍可能残留部分虚影。
  • 骨干参数量与显存开销:基于 FLUX.2 Klein 4B 骨干的参数量达 3.97B,虽然在 RTX 4090 上采用 bfloat16 单步推理耗时仅 104 ms,但显存占用依然显著高于传统小模型,对移动端或边缘计算设备不友好。
  • 理论保障的进一步探究:目前隐空间近似加性保持主要依赖经验论证,关于 VAE 编码器何时以及为何能够保持加性弱失真,尚缺乏严格的拓扑学理论证明。

相关工作与启发

  • vs 双流交互网络 (DSRNet / DSIT / RDNet):传统方法在像素空间耗费大量交互注意力或可逆块去避免信息丢失,但反光分支普遍退化为全零;PRISM 借助隐空间流匹配以单个单流 Transformer 的单次前向同时精准预测双分量。
  • vs 扩散先验去反光 (DAI, AAAI'26):DAI 依赖 Stable Diffusion 先验需要微调 VAE 且需要多步迭代采样;PRISM 改用无条件 FLUX 隐空间流匹配,无噪声直起、单步常速输运,推理效率与保真度显著占优。
  • 对底层视觉任务的启发:针对阴影去除(Shadow Removal)、去雨(Deraining)或高光去除等具有加性/半加性物理属性的分解任务,将图像投影至隐空间并借助流匹配相反场来做单步解耦是一条极具推广潜力的通用范式。

评分

  • 新颖性: ⭐⭐⭐⭐⭐(颠覆性地利用隐空间低相关性重构反光去除任务,单步流匹配推导极为精妙)
  • 实验充分度: ⭐⭐⭐⭐⭐(覆盖 6 个基准及 1K 真实场景零样本泛化,消融深入,反光可视化极具说服力)
  • 写作质量: ⭐⭐⭐⭐⭐(动机直击痛点,实验设计与对比推演环环相扣,逻辑缜密)
  • 价值: ⭐⭐⭐⭐⭐(为病态图像图层分离问题开辟了全新的隐空间生成式求解路径)