Semi-UFormer: Uncertainty-Aware Dehazing

• The paper presents a novel semi-supervised transformer network that integrates a two-stage teacher-student paradigm with uncertainty estimation for enhanced image dehazing.
• The model leverages Mix DehazeFormer Blocks and pixel-level aleatoric modeling, achieving superior PSNR/SSIM scores on both synthetic and real-world benchmarks.
• Knowledge distillation with KL divergence narrows the synthetic-real domain gap, effectively guiding restoration in challenging regions and outperforming previous methods.

Semi-UFormer is a semi-supervised, uncertainty-aware transformer network for image dehazing. Its central objective is to bridge the domain gap between synthetic and real hazy images, leveraging knowledge distillation and uncertainty modeling to enhance restoration fidelity. The framework explicitly models pixel-level aleatoric uncertainty and uses it to guide dehazing, achieving leading quantitative and qualitative results on both synthetic benchmarks and real-world images (Tong et al., 2022).

1. Architectural Overview

Semi-UFormer employs a two-stage teacher–student paradigm based on a U-Net-style backbone augmented with Mix DehazeFormer Blocks (MDB). Both teacher and student networks consist of five hierarchical stages, each incorporating the following components:

• Shallow Feature Extraction: A $3 \times 3$ convolution $C_3(\cdot)$ for preliminary feature extraction.
• Mix DehazeFormer Block (MDB): Composed of $n$ DehazeFormer transformer layers $DF(\cdot)$ followed by a residual convolutional block $RB(\cdot)$. This design enables effective global information fusion via

$$F_\text{global} = RB\left(DF(F_\text{shallow})_n\right) + DF(F_\text{shallow})_n$$

• Reconstruction Head: Another $3 \times 3$ convolution with a pixel-shuffle operation $P(\cdot)$, which generates the final dehazed output, $J = P(C_3(F_\text{global}))$.

The teacher network is first trained to produce a coarse dehazed image and a corresponding uncertainty map. The student, initialized from the teacher's trained weights, freezes the uncertainty estimation branch and is retrained with both labeled (synthetic) and unlabeled (real) data.

2. Knowledge Distillation and Semi-Supervised Learning

The core learning objective decouples into supervised and unsupervised components for both teacher and student networks:

• Teacher Supervised Loss: Combines uncertainty-estimation loss $L_{ue}$ and an adversarial reconstruction loss $C_3(\cdot)$0:

$C_3(\cdot)$1

• Teacher Unsupervised Loss: Incorporates identity loss $C_3(\cdot)$2, dark-channel prior loss $C_3(\cdot)$3, and total variation (TV) smoothness loss $C_3(\cdot)$4:

$C_3(\cdot)$5

with total teacher loss $C_3(\cdot)$6.

• Student Supervised Loss: Employs the uncertainty-guided supervised loss $C_3(\cdot)$7 and the adversarial loss:

$C_3(\cdot)$8

• Student Unsupervised Loss: Sum of uncertainty-guided consistency loss $C_3(\cdot)$9, dark-channel prior, TV smoothness, and a KL divergence distillation term:

$n$0

where

$n$1

and total student loss $n$2.

The KL distillation term is pivotal for domain adaptation, encouraging the student to align feature distributions between synthetic and real domains.

3. Uncertainty Estimation and Aleatoric Modeling

A dedicated Uncertainty Estimation Block (UEB) models pixel-wise aleatoric uncertainty $n$3. For a synthetic image $n$4, with dehazed teacher output $n$5 and ground truth $n$6, a Laplace observation model yields:

$n$7

The corresponding negative log-likelihood loss, up to additive constants, is

$n$8

The uncertainty loss becomes

$n$9

Predicted uncertainty modulates both supervised and unsupervised student loss terms, focusing the model on regions with high expected error:

• Supervised:

$DF(\cdot)$0

• Unsupervised:

$DF(\cdot)$1

4. Training Protocol and Domain Adaptation

The training regime utilizes 10,000 synthetic hazy/clean pairs (RESIDE OTS) and 2,000 unpaired real hazy photographs (URHI). The teacher is trained for 100 epochs with a 5:1 update ratio of supervised to unsupervised steps, using a learning rate decaying linearly from $DF(\cdot)$2. The student, initialized with the teacher’s final weights and with the UEB frozen, trains for 60 epochs with a 5:1 unsupervised to supervised update ratio, following the same learning rate schedule.

Domain adaptation is operationalized via the KL distillation loss $DF(\cdot)$3 on high-level embeddings from the MDB, which compels the student network to transfer feature knowledge learned on synthetic data to unlabelled real inputs.

5. Empirical Results and Ablation Analysis

Semi-UFormer achieves state-of-the-art performance on synthetic and real-world benchmarks. On SOTS-Outdoor, it yields PSNR/SSIM scores of 26.87/0.928, and 28.72/0.934 on HSTS—surpassing previous work by an average of over 2 dB. On 50 real images, no-reference metrics (SSEQ ↓, σ-color ↓, HCC ↑) are also optimal or near-optimal: SSEQ=37.75, σ-color=0.0001, HCC=0.2340.

Ablation studies demonstrate cumulative benefits of each component:

Variant	PSNR	SSIM
Base (DehazeFormer-S + $DF(\cdot)$4)	24.74	0.905
+ Mix DehazeFormer Block (MDB)	25.34	0.910
+ Uncertainty Estimation (UEB, $DF(\cdot)$5)	26.17	0.919
+ KL Distillation (full Semi-UFormer)	26.87	0.928

Key findings:

• The MDB consistently improves performance via long-range context aggregation.
• The uncertainty module yields ∼0.8 dB improvement by concentrating reconstruction on challenging regions.
• KL distillation narrows the synthetic–real domain gap by an additional ∼0.7 dB.

6. Comparative Analysis and Significance

Semi-UFormer’s semi-supervised, uncertainty-aware transformer architecture demonstrates that explicit modeling of pixel-wise uncertainty not only provides confidence estimates but also directs the student’s learning focus during domain adaptation. The approach both quantitatively and qualitatively outperforms prior deterministic dehazing networks, producing sharper structure and better color fidelity on real hazy scenes (Tong et al., 2022). The systematic deployment of adversarial, prior-based, and self-supervised (consistency, KL) objectives—alongside uncertainty normalization—accounts for robust generalization across synthetic and real data distributions.

7. Conclusion and Future Prospects

Semi-UFormer establishes a new state-of-the-art for image dehazing by uniting semi-supervised knowledge distillation and explicit uncertainty modeling in a transformer backbone. The integration of per-pixel uncertainty not only quantifies model confidence but also enhances dehazing, especially in high-ambiguity regions such as edges and textures. A plausible implication is that uncertainty-guided supervision could generalize to other domains facing synthetic-to-real gaps and spatially localized ambiguity (Tong et al., 2022).