Stokes-Informed Diffusion for Robust Linear Polarization Estimation

Conference: ECCV2026
Paper: ECCV Paper
Area: Image Restoration / Polarimetric Imaging
Keywords: linear polarization estimation, Stokes parameters, observability-aware supervision, distribution matching distillation, low-rank adaptation

TL;DR

GenPolar reformulates the ill-posed RGB-to-polarization inverse problem as channel-wise Stokes prediction, supervises polarization angles in reliable regions, and uses one-step distillation to support VAE encoder adaptation, improving both polarization strength and angle estimation on three test groups.

Background & Motivation

Polarization reveals material and geometric information that ordinary color images do not directly expose, helping distinguish metals from dielectrics and separate reflections from glass. Actual acquisition usually requires a rotating linear polarizer or a division-of-focal-plane polarimeter: the former needs multiple exposures, while the latter requires a specialized sensor and demosaicking. The paper therefore estimates these cues from an RGB image proportional to total intensity instead of requiring several analyzer-angle measurements at deployment. The difficulty is that intensity does not uniquely determine polarization; similar appearances may correspond to different materials, normals, illumination, and propagation paths. A diffusion model supplies a learned prior, not a physical measurement that makes the missing information uniquely recoverable.

Polarization strength and direction must be distinguished: when the degree of linear polarization (DoLP) is small, the angle of polarization (AoP) is nearly unobservable. Enforcing an angle target at every pixel then treats measurement noise as supervision; representing angles with sine and cosine handles periodicity but does not remove weak-signal instability. Meanwhile, polarization strength can differ across RGB bands, so collapsing them to grayscale discards meaningful spectral variation. Switching to Stokes parameters alone is insufficient because a natural-image VAE introduces encode-decode bias on polarization fields, which propagates into DoLP and AoP. These problems affect representation, supervision reliability, and latent space, rather than being solved simply by using a larger generator.

The paper uses Mueller transport to explain why a scene-dependent mapping from intensity to linear Stokes components is plausible, without claiming to recover a complete Mueller matrix from RGB. It first learns a multi-step conditional diffusion model with physical supervision, then distills its generation path to one step so that the encoding of polarization targets can be adapted stably. Core Idea: predict color-preserving Stokes components, constrain angles only where they are observable, and combine one-step distillation with encoder adaptation to reduce latent-domain errors.

Method

Overall Architecture

Training uses paired RGB intensity conditions and polarization ground truth derived from actual polarization measurements, not labels inferred from RGB alone. The model outputs two linear Stokes components for each RGB channel and analytically computes DoLP and AoP using the input intensity. The four designs are channel-wise Stokes representation, observability-aware supervision, one-step distribution matching distillation, and low-rank encoder adaptation; the last two are jointly optimized in the second stage. Stage I freezes the VAE and trains the diffusion UNet and ControlNet; Stage II retains a frozen teacher and decoder while training the one-step generator and encoder LoRA. Deployment needs only RGB conditioning and Gaussian noise, not polarization ground truth, an observability mask, or the target-encoding branch.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    Input["Paired training data"] --> Stokes["Channel-wise Stokes<br/>representation"]
    Stokes --> Physics["Observability-aware<br/>supervision"]
    Physics -->|Stage I teacher| Distill["One-step distribution<br/>matching distillation"]
    Distill --> Adapt["Low-rank encoder<br/>adaptation"]
    Adapt -. Joint Stage II optimization .-> Distill
    Adapt -. Trained generator .-> Infer["One-step generator<br/>and fixed decoder"]
    RGB["Inference: RGB and Gaussian noise"] --> Infer
    Infer --> Output["Stokes components and DoLP/AoP"]

Key Designs

1. Channel-wise Stokes representation: retain strength and orientation in one physical representation

In the Stokes vector, \(S_0\) denotes total intensity, while \(S_1\) and \(S_2\) describe differences along two linear polarization bases; the paper neglects the circular component \(S_3\). During training, measurements at four analyzer angles provide linear polarization targets, including \(S_1=I_{0^\circ}-I_{90^\circ}\) and \(S_2=I_{45^\circ}-I_{135^\circ}\). Under approximately unpolarized incident illumination, the first column of the Mueller matrix determines outgoing intensity and linear polarization; local material, geometry, and illumination determine the relevant transfer-coefficient ratios. The network therefore learns a spatially varying intensity-to-polarization mapping, not a single global proportionality constant or an explicit solution for all optical parameters. This physical account motivates the parameterization but does not guarantee that one intensity image identifies the true polarization state.

For each color channel \(\lambda\in\{R,G,B\}\), predicted components are converted into standard polarization cues:

\[ \mathrm{DoLP}_{\lambda}=\frac{\sqrt{S_{1,\lambda}^{2}+S_{2,\lambda}^{2}}}{S_{0,\lambda}},\qquad \mathrm{AoP}_{\lambda}=\frac{1}{2}\operatorname{atan2}(S_{2,\lambda},S_{1,\lambda}). \]

This corresponds to Eq. (2); retaining channels allows wavelength-dependent polarization strength instead of forcing all colors to share one grayscale polarization field. The paper uses cross-channel AoP coherence as a practical reliability proxy, not as a claim that all three channels must have identical angles in every scene. In implementation, the three-channel \(S_1\) and three-channel \(S_2\) maps pass separately through the same VAE encoder, yielding two 4-channel latents that are concatenated into 8 channels. Spatial resolution is reduced to \(1/8\) of the input; the UNet is adapted to this joint latent, while ControlNet injects RGB structural conditioning. Decoding splits the two parts and applies the same fixed decoder to recover two three-channel components, rather than feeding a six-channel target directly into the original RGB VAE.

2. Observability-aware supervision: do not treat weak-polarization angles as reliable labels

The angle is the direction of the two Stokes components; when their magnitude approaches zero, small perturbations can cause large angular changes. The method consequently supervises Stokes components at all pixels but adds angular supervision only where ground-truth DoLP exceeds a threshold. Using ground truth for the mask is important: the network cannot reduce its predicted DoLP to remove difficult pixels from the angle loss. Angular distance is \(\pi\)-periodic, so physically equivalent orientations are not penalized as ordinary numerical differences across the angular boundary. Periodicity and reliability filtering have separate roles: the former handles equivalent directions, while the latter asks whether the signal is strong enough.

The supervision structure described in the paper can be retained as:

\[ \Omega_{\tau,\lambda}=\{y:\mathrm{DoLP}^{\mathrm{GT}}_{\lambda}(y)>\tau\},\qquad \mathcal{L}_{\mathrm{phys}}=\mathcal{L}_{s}+\gamma_{\mathrm{AoP}}\mathcal{L}_{\mathrm{AoP}},\qquad \mathcal{L}_{\mathrm{stg}_1}=\mathcal{L}_{\mathrm{diff}}+\gamma_{\mathrm{p}}\mathcal{L}_{\mathrm{phys}}. \]

Here, \(\mathcal{L}_{s}\) is the full-image Stokes error, \(\mathcal{L}_{\mathrm{AoP}}\) is the masked periodic angular error, and \(\mathcal{L}_{\mathrm{diff}}\) supervises clean-latent prediction. Parts of Eqs. (8) and (9) are corrupted in the text extraction, so their exact norms and normalization details are not reconstructed here; only the structure clearly stated in the prose is retained. Low-DoLP regions are not excluded from training altogether: they still receive Stokes and latent supervision, without chasing unreliable angle targets. The training-threshold analysis in Fig. 7 supports \(\tau_{\mathrm{train}}=0.05\); evaluation also uses a fixed ground-truth threshold rather than selecting pixels from each method's predictions. The design improves reliable-region angular estimation and output coherence, not the recovery of true directions in every weakly polarized region.

3. One-step distribution matching distillation: shorten generation and the adaptation gradient path

Stage I starts from SD1.5, freezes the VAE encoder and decoder, adds noise to target latents, and learns to predict clean latents conditioned on RGB. Its parameterization predicts clean latents; the training objective should not be casually rewritten as noise prediction. Predicted latents are decoded before computing the physical loss, exposing training to both latent errors and errors in the recovered polarization fields. The multi-step model supplies a generative prior, but adapting the encoder through a long denoising path makes end-to-end optimization expensive and unstable. Stage II transfers the teacher's capability into a one-step generator, so single-step generation is not merely an auxiliary compression technique added for speed.

Distribution matching distillation (DMD) freezes the Stage-I teacher and trains an auxiliary model online to estimate the score of the current generated distribution. After generated outputs are re-noised, the difference between teacher and generated-distribution scores provides a distribution-level update direction, rather than only pixelwise imitation of one teacher output. The auxiliary score model is training-only and should not be interpreted as a second generator required at deployment. Stage II continues to apply physical supervision after decoding, with Stokes supervision also serving as a regression constraint during distillation. This explains why a single end-to-end regression pass through the same UNet is not equivalent to learning a diffusion teacher and then distilling it.

4. Low-rank encoder adaptation: change the target's latent entry point while keeping decoding stable

A natural-image VAE introduces structured encode-decode errors on polarization fields, illustrated by the error maps in Fig. 5. Stage II inserts LoRA into the VAE encoder while freezing the decoder, seeking target representations that better fit polarization yet remain interpretable by the existing decoder. This is neither a highly flexible image-correction module attached to the decoder nor another name for fine-tuning ControlNet alone. The generator and LoRA are trained jointly, physical supervision acts on both, and a KL term regularizes the encoder posterior toward a standard Gaussian prior. This regularization discourages arbitrary latent-distribution drift while maintaining a connection between target reconstruction and the existing latent space.

Adapting LoRA beforehand through standalone reconstruction, or inserting it directly into Stage I, does not reproduce the full method's performance. The supported conclusion is that adaptation timing and generative training must work together, not that LoRA automatically improves every polarization network. Inference does not encode unknown ground-truth \(S_1,S_2\): the generator directly produces a latent from Gaussian noise and RGB conditioning, and the fixed decoder produces polarization components. Encoder adaptation changes the target latent representation and learning path during training, with the trained generator carrying the resulting benefit into deployment. This distinction prevents the training target-encoding branch from being mistaken for a requirement to use a polarization camera at test time.

A Worked Example

Consider an RGB image containing a strong reflection boundary and a weakly polarized diffuse region; this is an explanatory example, not an additional experiment. During training, paired polarization captures provide three-channel Stokes targets, and the two three-channel components are encoded into one 8-channel target latent. If ground-truth DoLP is 0.02 at one pixel and 0.20 at another, a threshold of 0.05 includes only the latter in angular supervision, while both receive Stokes supervision. Stage I learns to recover targets from noisy latents using RGB structure; Stage II trains one-step generation with teacher scores, physical supervision, and encoder adaptation. At deployment, an RGB image and noise produce a latent in one step, followed by split decoding and DoLP/AoP computation using the input intensity. The low-polarization pixel still receives an angle prediction, but neither this example nor the paper's masking mechanism gives that prediction an additional reliability guarantee.

Loss & Training

Stage I runs for 100 epochs at a learning rate of \(4\times10^{-5}\); Stage II runs for 30 epochs at \(10^{-5}\). The optimizer is AdamW with betas (0.9, 0.999), weight decay \(10^{-3}\), and both physical-supervision weights set to 0.5. LoRA rank is 4, patch size is 512, batch size is 1, and training uses 4 A40 GPUs. \(S_1,S_2\) are used directly in \([-1,1]\), while \(S_0\in[0,2]\) is scaled to \([-1,1]\) for conditioning. Stage-I inference uses 20 steps and the final generator uses one; this compares step counts, not a reported 20-fold wall-clock speedup. The prose describes a fixed maximum-noise timestep in Stage II, whereas line 17 of Algorithm 1 says to sample \(t^*\); the original implementation is needed to clarify this detail.

Key Experimental Results

Main Results

Table 1 on paper page 12 reports DoLP PSNR in dB, DoLP SSIM, and AoP MeanAE in degrees, computed separately for the three color channels and then averaged. AoP evaluation uses a shared ground-truth DoLP mask above 0.05, which the authors say retains approximately 70% of pixels; it is not directly comparable to full-image angular error. RLP contains 598 samples, including 363 existing samples and 235 author-captured samples; DoFP contains 2188 and RSP contains 2000. Each group independently reserves 50 random samples for testing, with the remainder forming a union training set on which all baselines are retrained. This evaluates different capture-modality test groups, not leave-one-domain-out generalization in which an entire modality is absent from training.

Test Group	Method	DoLP PSNR ↑	DoLP SSIM ↑	AoP MeanAE ↓
RLP	Restormer	20.64	0.540	35.35
RLP	MAE	26.98	0.638	28.23
RLP	PolarAnything	19.77	0.563	33.56
RLP	GenPolar	29.54	0.741	20.15
DoFP	Restormer	16.82	0.532	32.72
DoFP	MAE	21.40	0.598	29.82
DoFP	PolarAnything	17.85	0.572	29.33
DoFP	GenPolar	24.01	0.666	18.77
RSP	Restormer	15.15	0.588	35.17
RSP	MAE	16.67	0.616	33.51
RSP	PolarAnything	15.47	0.604	30.45
RSP	GenPolar	21.29	0.676	27.60

Ablation Study

Table 2 on paper page 13 uses the RLP test set throughout, with the same metric definitions and evaluation mask as the main comparison. The direct-regression row predicts DoLP and double-angle sine/cosine components, rather than an untreated scalar AoP target.

Config	DoLP PSNR ↑	DoLP SSIM ↑	AoP MeanAE ↓
Only Stage I	28.84	0.684	22.61
Pre-LoRA	27.92	0.674	23.49
LoRA in Stage I	19.78	0.460	50.03
One-step without LoRA	28.49	0.681	23.12
End-to-end regression + physical loss	24.15	0.553	30.73
Without angular loss	27.21	0.668	23.96
Training threshold of 0	29.47	0.724	21.89
Gray Stokes	27.43	0.621	23.56
Direct DoLP/AoP regression	22.84	0.502	35.40
Full model	29.54	0.741	20.15

Key Findings

• On RLP, DoLP PSNR improves from MAE's 26.98 to 29.54 dB, a gain of 2.56 dB; AoP MeanAE decreases from 28.23 to 20.15 degrees.
• One-step generation without LoRA yields an angular error of 23.12 degrees versus 20.15 for the full model; inserting LoRA directly into Stage I worsens it to 50.03 degrees, supporting the importance of joint-training timing.
• Table 3 on paper page 14 fixes SfPUEL and changes only polarization inputs: material detection accuracy is 90.12% for GenPolar, 82.75% for MAE, and 93.75% for captured inputs.
• Table 4 on the same page fixes PolarFree: GenPolar achieves 22.41 dB / 0.127 LPIPS in de-reflection versus 22.44 dB / 0.133 for captured inputs; only LPIPS is better, not every metric.
• Numerical inconsistency: Fig. 1 on page 2 labels the direct-regression angular error as 31.40 degrees, while Table 2 on page 13 reports 35.40 degrees; this note retains Table 2 without speculating about the cause.

Highlights & Insights

• Representation and supervision reliability are designed together: Stokes components retain strength-orientation coupling, while the angle mask avoids treating weak-signal noise as precise information. Replacing angles with sine and cosine alone cannot substitute for these two decisions.
• One-step distillation offers more than fewer inference steps: it also shortens the end-to-end optimization path for the target encoder. Ablations support a useful interaction between faster generation and improved latent representation, but not a guaranteed runtime speedup ratio.
• Replacing inputs to fixed downstream models is more informative than inspecting polarization pseudocolor maps alone. It tests whether existing task networks can use the estimated cues while preserving the comparison with actual captured inputs.

Limitations & Future Work

• The authors explicitly estimate only linear polarization, not \(S_3\); polarized illumination, skylight, transparent paths, and multiple scattering can weaken the physical approximation.
• RGB-to-polarization remains a non-unique inverse problem. Outputs should be regarded as physically supervised estimates, not measurements equivalent to those from a polarization camera.
• AoP metrics exclude weakly polarized pixels, and the paper acknowledges that underestimating DoLP reduces directional observability; jointly reporting angle error, coverage, and predictive uncertainty is a reader-proposed next step.
• With 50 test samples per group and training on each group's remainder, robustness to new sensors and entirely new scene distributions is not established; independent cross-device splits are still needed.
• The supplied cache contains the main paper and references, but not supplementary failure cases or a resolved code URL; damaged loss equations and the timestep discrepancy noted above prevent any claim that the implementation has been verified.

Related Work & Insights

• RGB-to-Polarization Estimation (paper reference 19) provides the task and baseline context. GenPolar also conditions on intensity, but incorporates reliability filtering and latent adaptation into training.
• PolarAnything (reference 40) directly generates grayscale polarization cues, whereas GenPolar preserves RGB Stokes components. Converting outputs into common DoLP/AoP metrics does not eliminate representation-capacity differences, so the grayscale ablation matters when interpreting comparisons.
• DMD (reference 38) and LoRA (reference 7) are established components reused here; the contribution is their coupling for distribution distillation and polarization-domain encoding, not a new distillation algorithm.
• Transferable insight: for directional variables that become unstable under weak signals, learn a magnitude-bearing representation and select directional supervision by observability. The reliability criterion must be validated against the target task's noise model rather than copying a DoLP threshold.

Rating

• Novelty: 4/5. Physical representation, angular reliability, and encoder adaptation form a targeted combination, although the main components are established.
• Experimental Thoroughness: 4/5. Three data groups, focused ablations, and two fixed downstream tasks are covered, but independent cross-domain and uncertainty evaluation remain limited.
• Writing Quality: 3/5. The main mechanism is clear, but inconsistent direct-regression numbers and Stage-II timestep descriptions need clarification; damaged extracted equations further restrict verification.
• Value: 4/5. The results show that estimation without additional polarization hardware can improve practical tasks, without providing the accuracy guarantees of physical measurement.