PRISM: Latent Composition Consistency for Single-Image Reflection Removal

Conference: ECCV 2026
Paper: ECCV Official Page
PDF: ECCV Open Access
Area: Image Restoration
Keywords: Single-Image Reflection Removal / Latent Flow Matching / Composition Consistency / Layer Contrastive Separation / FLUX

TL;DR

PRISM observes that mutual coherence between transmission and reflection layers drops ten-fold in pretrained VAE latent space compared to pixel space, reinterpreting reflection removal as a latent linear separation task via single-step flow matching on a FLUX backbone that recovers both layers simultaneously, complemented by cross-sample latent composition consistency and patch-level contrastive separation to achieve SOTA across six benchmarks.

Background & Motivation

When photographing through transparent surfaces such as glass windows, captured images are frequently corrupted by superimposed, undesirable reflections. Single-image reflection removal (SIRR) aims to recover the clean transmission layer \(T\) and separate the reflection layer \(R\) from an observed mixture \(I = T + R\). Mathematically, this task represents a severely ill-posed inverse problem, as an infinite number of \((T, R)\) pairs satisfy the decomposition for any given \(I\). Traditional approaches resort to handcrafted priors, relative smoothness assumptions, or learnable residue terms. Recent deep learning dual-stream networks (e.g., DSRNet, DSIT, RDNet) explicitly model feature interactions across two branches; however, they remain anchored in the sRGB pixel space. Due to non-linear camera response functions and tone mapping, pixel-level additive superposition is only a coarse approximation, causing networks to absorb reflections into reconstruction residuals and leading to degenerate reflection branches that predict near-zero outputs.

While the RAW sensor domain offers physically faithful additive linearity, its strict hardware constraints preclude its application to vast collections of existing sRGB images. The core tension lies in the severe statistical entanglement between transmission and reflection in sRGB pixel space. Across 454 real image pairs from the SIR2 benchmark, the average cosine similarity between transmission and reflection in pixel space reaches 0.74; conversely, when encoded into the latent space of a pretrained generative model (FLUX VAE), the cosine similarity plummets to 0.07. This widespread decorrelation reveals that pretrained latent spaces provide a substantially more favorable linear separation geometry for orthogonal layer decomposition than pixel space.

Because the latent space exhibits exceptionally low mutual coherence, reflection separation no longer requires heavy dual-branch recurrent iterations or complex residue engineering. The core idea is to recast single-image reflection removal as a latent linear separation problem via flow matching, directly recovering both transmission and reflection components in a single forward pass while introducing Latent Composition Consistency (LCC) via cross-sample reflection swapping and target-free Layer Contrastive Separation (LCS) to enforce robust orthogonal disentanglement.

Method

Overall Architecture

PRISM begins by projecting the input corrupted mixture image \(I\) into a compressed latent representation \(z_I = \mathcal{E}(I) \in \mathbb{R}^{16 \times \frac{H}{8} \times \frac{W}{8}}\) using a pretrained FLUX VAE encoder \(\mathcal{E}\). Grounded on the approximate additive latent formulation \(z_I \approx z_T + z_R\), PRISM adopts the Multi-Modal Diffusion Transformer (MM-DiT) of FLUX.2 Klein (4B) as the velocity field network \(v_\theta\), initiating a single-step forward pass directly from the clean, uncorrupted input latent \(z_I\). The predicted velocity vector approximates the negative reflection component, enabling simultaneous recovery of the predicted transmission latent \(\hat{z}_T = z_I + v_\theta(z_I, 0)\) and reflection latent \(\hat{z}_R = -v_\theta(z_I, 0)\) in a single forward step, which are subsequently rendered back into image space by the VAE decoder \(\mathcal{D}\).

To prevent degenerate solutions where the network entangles scene-specific content across components, PRISM introduces an auxiliary training path across the batch dimension: separated reflection latents are cyclically swapped across samples to construct synthetic composite mixtures, followed by a second forward pass that enforces cycle consistency. Concurrently, patch-pooled representations undergo layer contrastive separation to maximize semantic divergence between layers without requiring explicit pixel-space reflection ground truth.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input Mixture Image I"] --> B["VAE Encoder ℰ<br/>Project to Latent z_I"]
    B --> C["Latent Flow Matching Single-Step Linear Separation<br/>MM-DiT Predicts Velocity Field v_θ(z_I, 0)"]
    C --> D["Single-Step Vector Resolution<br/>z_T_hat = z_I + v_θ, z_R_hat = -v_θ"]
    D --> E["Latent Composition Consistency<br/>Cyclic Latent Swapping and Second-Pass Decomposition"]
    D --> F["Layer Contrastive Separation<br/>Patch-Pool Local Feature Contrastive Pull/Push"]
    D --> G["VAE Decoder 𝒟<br/>Reconstruct Transmission T_hat and Reflection R_hat"]

Key Designs

1. Latent Flow Matching Single-Step Linear Separation: Bypassing Dual-Stream Redundancy for Single-Pass Dual Recovery

In pixel space, recovering both transmission and reflection typically demands complex dual-stream architectures (e.g., dual-branch LSTMs or interactive transformers), which incur high computational overhead and suffer from reflection branch collapse. Addressing this issue, PRISM exploits the ultra-low mutual coherence of latent space (cosine similarity of only 0.07) to adopt the approximate linear additive assumption \(z_I \approx z_T + z_R\). Consequently, the ideal flow matching velocity target transporting the input latent \(z_I\) to the transmission target \(z_T\) is defined as:

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

This formulation reveals an elegant property: under the additive latent assumption, the flow matching velocity target equals the exact negative of the reflection latent. Implementing \(v_\theta\) via a full fine-tuning of FLUX.2 Klein (4B) MM-DiT, PRISM replaces the text cross-attention conditioning with a single zero-valued embedding to operate unconditionally. Unlike conventional diffusion models that start from Gaussian noise and rely on multi-step iterative denoising, PRISM leverages straight-line trajectories by setting \(t=0\) and starting directly from the unnoised \(z_I\). A single forward pass simultaneously resolves both transmission \(\hat{z}_T = z_I + v_\theta(z_I, 0)\) and reflection \(\hat{z}_R = -v_\theta(z_I, 0)\), eliminating dual-branch architectural complexity and multi-step sampling latency.

2. Latent Composition Consistency: Exploiting Low-Coherence Space for Synthetic Cycle Supervision

Relying solely on one-way transmission reconstruction loss risks degenerate solutions where the network entangles scene-specific textures across both components. In pixel space, blending the reflection from image A onto the transmission of image B produces unphysical artifacts due to high correlation and lighting mismatches. However, in the decorrelated VAE latent space, linearly composing transmission and reflection latents from independent samples produces semantically valid, realistic mixtures upon decoding. PRISM leverages this phenomenon through Latent Composition Consistency (LCC).

Across a training batch of \(N\) samples, the first-pass transmission prediction \(\hat{z}_T^i\) of sample \(i\) is cyclically combined with the reflection prediction \(\hat{z}_R^{(i \bmod N) + 1}\) of an adjacent sample, constructing a synthetic mixture latent \(\tilde{z}_I^i = \hat{z}_T^i + \hat{z}_R^{(i \bmod N) + 1}\). A second forward pass through MM-DiT decomposes this synthetic mixture into \(\tilde{z}_T^i\) and \(\tilde{z}_R^i\). To eliminate trivial solutions and avoid mode collapse (where first-pass predictions collapse to a single constant), PRISM applies a stop-gradient operator \(\texttt{sg}(\cdot)\) to the first-pass outputs, enforcing bidirectional cycle consistency:

\[\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. Layer Contrastive Separation: Ground-Truth-Free Semantic Disentanglement Avoiding Non-Linear Errors

Although the latent additive assumption provides strong empirical utility, strict equality does not hold due to VAE encoder non-linearity: \(\mathcal{E}(I - T) \neq \mathcal{E}(I) - \mathcal{E}(T)\). Enforcing an explicit regression target on the pixel-space residual would inevitably inject non-linear distortion into the velocity field. To circumvent this without sacrificing layer isolation, PRISM introduces the Layer Contrastive Separation (LCS) loss, enforcing semantic orthogonality between transmission and reflection features via contrastive learning.

Recognizing that reflection artifacts are spatially non-uniform, PRISM avoids global pooling by dividing latent feature maps into non-overlapping \(p \times p\) patches (default \(p=4\)), with each patch independently average-pooled and \(\ell_2\)-normalized. Taking the cycle-recovered transmission patch \(f_k(\tilde{z}_T)\) as an anchor, the positive set \(\mathcal{P}\) includes first-pass transmission patches \(\hat{z}_T\) and ground-truth transmission patches \(z_T^{\text{gt}}\), while the negative set \(\mathcal{N}\) comprises predicted reflection patches \(\hat{z}_R\), cycle reflection patches \(\tilde{z}_R\), and ground-truth reflection patches \(z_R^{\text{gt}}\). Using a temperature parameter \(\tau=0.1\), the patch-level InfoNCE loss is formulated as:

\[\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]\]

This objective repels reflection semantics from transmission features without requiring pixel-level reflection regression.

Loss & Training

The complete multi-task training objective is defined as:

\[\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}}\]

where the velocity field loss is \(\mathcal{L}_{\text{latent}} = \|v_\theta(z_I, 0) - v^*\|_2^2\), and the decoded transmission is supervised in pixel space by \(\mathcal{L}_{\text{pixel}} = \|\mathcal{D}(\hat{z}_T) - T\|_1 + \lambda_{\text{lpips}} \mathcal{L}_{\text{LPIPS}}(\mathcal{D}(\hat{z}_T), T)\). Hyperparameters are set to \(\lambda_{\text{lat}} = 1\), \(\lambda_{\text{pix}} = 1\), \(\lambda_{\text{lpips}} = 2\), \(\lambda_{\text{cycle}} = 1\), and \(\lambda_{\text{LCS}} = 0.1\).

The model is implemented in PyTorch and trained on a single NVIDIA A6000 GPU with a batch size of 2 and random \(512 \times 512\) crops for 50,000 iterations. Optimization uses AdamW with bfloat16 mixed precision, decaying the learning rate from \(5 \times 10^{-5}\) to \(1 \times 10^{-6}\) via cosine annealing. Training pairs are sampled from PASCAL VOC synthetic data (7,643 pairs), Zhang et al. real pairs (90 pairs), and Nature real pairs (200 pairs) under a 0.6 : 0.2 : 0.2 ratio per epoch.

Key Experimental Results

Main Results

PRISM was evaluated across six benchmark datasets against ten competitive SIRR baselines: ERRNet, IBCLN, LASIRR, YTMT, RobustSIRR, DSRNet, RRW, DSIT, RDNet, and DAI (AAAI 2026).

Table 1 summarizes PSNR and SSIM performance across all six benchmarks. PRISM establishes new state-of-the-art PSNR across all datasets, surpassing RDNet by 1.56 dB on Real and 1.84 dB on Wild.

Method	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 (Ours)	27.14 / 0.853	27.61 / 0.925	27.85 / 0.906	29.54 / 0.932	27.35 / 0.853	27.95 / 0.918

To assess real-world in-the-wild generalization, Table 2 presents zero-shot evaluation on the OpenRR 1K test set (100 real-world reflection images, \(512 \times 512\) resolution) measured on an RTX 4090 GPU without any dataset-specific fine-tuning.

Method	PSNR ↑	SSIM ↑	LPIPS ↓	DISTS ↓	NIQE ↓	Time (ms) ↓	Params
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 (Ours)	27.23	0.893	0.064	0.041	3.5512	104.23	3.97B

Ablation Study

Table 3 investigates the individual and combined contributions of \(\mathcal{L}_{\text{cycle}}\) and \(\mathcal{L}_{\text{LCS}}\), where the base configuration relies solely on \(\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	Avg 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

Table 4 examines the spatial granularity within the Layer Contrastive Separation loss.

Representation Strategy	Real (20) PSNR / SSIM	Nature (20) PSNR / SSIM	SIR2 (454) PSNR / SSIM	Avg PSNR / SSIM	Note
Patch (\(p=2\))	27.03 / 0.851	27.35 / 0.853	27.94 / 0.917	27.44 / 0.874	Overly local patches induce noise
Patch (\(p=4\), Default)	27.14 / 0.853	27.35 / 0.853	27.95 / 0.918	27.48 / 0.875	Optimal spatial-semantic trade-off
Patch (\(p=8\))	26.83 / 0.849	27.22 / 0.852	27.93 / 0.917	27.33 / 0.873	Large receptive field dilutes cues
Global Avg. Pool	26.73 / 0.845	27.21 / 0.852	27.94 / 0.918	27.29 / 0.872	Collapses spatial heterogeneity

Key Findings

• Bidirectional Consistency is Crucial: Supervising only the transmission cycle \(\tilde{z}_T\) (avg. 27.34 dB) or only the reflection cycle \(\tilde{z}_R\) (avg. 27.36 dB) noticeably underperforms joint bidirectional supervision (avg. 27.48 dB). The transmission cycle maintains scene integrity against diverse reflections, while the reflection cycle enforces physical additive consistency.
• Perception-Distortion Tradeoff: PRISM dominates perceptual quality metrics (achieving a 14.0% reduction in LPIPS, a 14.6% reduction in DISTS, and the best NIQE of 3.5512 on OpenRR 1K). Because generative priors reconstruct vivid, realistic high-frequency textures that may not align pixel-for-pixel with GT windows, SSIM is marginally lower on Object and Postcard despite superior visual clarity and higher PSNR.
• Eliminating Reflection Branch Degeneration: Qualitative reflection visualizations demonstrate that previous pixel-space methods (YTMT, DSRNet, DSIT, RDNet) yield near-zero, black reflection predictions due to unconstrained residual absorption. In contrast, PRISM reconstructs coherent, structurally faithful reflection maps matching the ground truth.

Highlights & Insights

• Leveraging Pretrained Latent Decorrelation: Demonstrating that mutual coherence drops from 0.74 in pixel space to 0.07 in VAE latent space provides a compelling rationale for solving ill-posed non-linear image decomposition tasks in generative latent representations.
• Dual Decomposition via One-Step Flow Matching: Capitalizing on the velocity target duality \(v^* = z_T - z_I \approx -z_R\) allows extracting both transmission and reflection simultaneously in a single forward pass, bypassing multi-step iterative diffusion sampling.
• Stop-Gradient Powered Cross-Image Latent Swapping: Harnessing latent space flexibility to swap reflection latents across batch samples and enforcing cycle consistency prevents trivial memorization and ensures robust separation across arbitrary composite mixtures.

Limitations & Future Work

• Breakdown under Dense Saturated Reflections: When reflections are optically dense and saturated (\(\|R\|_1 / \|I\|_1 > 0.5\)), the input latent shifts heavily toward the reflection subspace, weakening the low-coherence assumption and leaving residual reflection artifacts.
• Model Size and Memory Footprint: The FLUX.2 Klein backbone comprises 3.97B parameters. Although single-step bfloat16 inference takes only 104 ms on an RTX 4090, memory overhead remains a bottleneck for edge devices.
• Theoretical Foundations of Latent Additivity: While the approximate additive behavior of VAE encoders is empirically validated, rigorous mathematical characterization of encoder linearity preservation remains an open theoretical direction.

Related Work & Insights

• vs Dual-Stream Networks (DSRNet / DSIT / RDNet): While pixel-space dual-stream methods rely on complex recurrent attention or reversible blocks to combat information loss and frequently suffer from reflection branch collapse, PRISM uses a single-stream transformer in latent space to recover both layers cleanly in one step.
• vs Diffusion-Based SIRR (DAI, AAAI'26): DAI requires fine-tuning the Stable Diffusion VAE and relies on iterative sampling; PRISM deploys unconditioned FLUX latent flow matching, starting directly from clean latents for rapid single-pass inference.
• Cross-Task Potential: The paradigm of recasting image decomposition into latent flow matching with negative velocity fields holds strong promise for related semi-additive restoration problems such as shadow removal, deraining, and flare removal.

Rating

• Novelty: ⭐⭐⭐⭐⭐ (Elegant formulation exploiting latent decorrelation and dual velocity field decomposition)
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ (Extensive testing over six benchmarks, 1K in-the-wild evaluation, and in-depth ablations)
• Writing Quality: ⭐⭐⭐⭐⭐ (Clear motivation, rigorous mathematical exposition, and structured empirical validation)
• Value: ⭐⭐⭐⭐⭐ (Opens a promising generative latent separation avenue for classic ill-posed vision inverse problems)

Related Papers

• [ECCV 2026] FUMO: Prior-Modulated Diffusion for Single Image Reflection Removal
• [ECCV 2026] Flow Straight to Reality: Perceptually Consistent Flow Matching for Efficient Image Restoration
• [ECCV 2024] L-DiffER: Single Image Reflection Removal with Language-Based Diffusion Model
• [ECCV 2026] BeyondMasks: Evaluating Causal and Physical Consistency in Video Object Removal
• [ECCV 2026] SSDD: Single-Step Diffusion Decoder for Efficient Image Tokenization