PaperlayerView exact revision on arXiv

Source-linked AI summary

Uncertainty-Guided Adverse Weather Restoration via Gated Transformer Network

Zheke Jin, Yuning Cui, Tianle Jin, Alois Knoll, Hu Cao

arXiv:2609.02434v1cs.CV

TL;DR

Adverse-weather restoration must handle heterogeneous degradations, while existing unified models may aggregate context uniformly, fuse scales naively, and ignore predictive uncertainty. UAR-Net combines gated Transformers, balanced multi-scale skips, and uncertainty-aware refinement, achieving state-of-the-art performance across multiple benchmarks.

  • Problem

    Existing all-in-one adverse-weather models use weather-agnostic global aggregation, naive cross-scale fusion, and deterministic point predictions despite heterogeneous degradations and ambiguity.

  • Method

    UAR-Net integrates GDTB for selective global and multi-scale local modeling, BMSC for progressive feature fusion, and URH with BAE-Loss for refinement and pixel-wise uncertainty.

  • Results

    UAR-Net achieves state-of-the-art PSNR/SSIM across multiple adverse-weather benchmarks, outperforming Histoformer by an average +0.78 dB PSNR on unified evaluation.

  • Takeaways & Limitations

    The framework provides strong restoration across both large-scale and localized adverse-weather degradations while jointly learning reconstructions and uncertainty.

Abstract

from arXiv · show

Restoring images degraded by adverse weather remains challenging due to spatially heterogeneous degradations. Many existing weather-specific restoration models rely on weather-agnostic global aggregation, naive cross-scale fusion, and deterministic objectives, which struggle to handle heterogeneous degradations in all-in-one adverse-weather settings. To address these limitations, we propose an Uncertainty-guided Adverse-weather Restoration Network (UAR-Net), a weather-specific AiO framework that integrates a gated transformer with balanced multi-scale skip connections. Specifically, we employ Gated Dual-scale Transformer Blocks (GDTB) to jointly model selective global interactions and multi-scale local structures, a progressive Balanced Multi-scale Skip Connection (BMSC) for balanced multi-scale feature integration, and an Uncertainty-Aware Refinement Head (URH) that performs artifact removal, detail enhancement, and predictive uncertainty estimation. The model is supervised by a Brightness-Aware Energy Loss (BAE-Loss) to encourage accurate reconstruction with well-calibrated uncertainty. Extensive experiments demonstrate that our method achieves state-of-the-art performance across multiple adverse-weather benchmarks. The codes will open source upon acceptance.

I. INTRODUCTION

Adverse-weather restoration must handle multiple spatially heterogeneous degradations that can impair safety-critical visual perception. UAR-Net addresses limitations in unified restoration through selective global modeling, balanced cross-scale fusion, uncertainty-aware refinement, and strong benchmark performance.

  • Motivation: Rain, snow, and haze can obscure scene content, blur structural boundaries, and degrade downstream detection, tracking, and segmentation.Robust restoration is therefore framed as a prerequisite for reliable real-world perception.
  • Limitations of Existing Methods: Single-weather restoration models require separate networks for different degradations, limiting scalability and robustness in unconstrained environments.Unified AiO frameworks were introduced to address multiple weather degradations within one network.
  • Limitations of Existing Methods: Existing AiO approaches often aggregate global context uniformly and use deterministic point predictions, which can mix degradation patterns and produce over-confident restorations in ambiguous regions.These limitations weaken degradation-related cues and stability under severe degradation.
  • Proposed Approach: UAR-Net combines query-conditioned GDTB gating, progressive BMSC fusion, and URH with BAE-Loss to model weather-aware interactions, balanced skips, details, and pixel-wise uncertainty.The design controls cross-scale information, reduces corrupted shallow-feature influence, and refines coarse predictions.
  • Results: UAR-Net achieves state-of-the-art performance across multiple adverse-weather benchmarks and recent unified restoration models.The introduction reports extensive experiments supporting this comparison.

III. METHOD

UAR-Net uses a GDTB-based Transformer encoder-decoder with selective gated attention, dual-scale local modeling, balanced skip fusion, and uncertainty-aware refinement. Its attention pathway preserves linear complexity while gates selectively regulate global context for diverse weather degradations.

  • Architecture: UAR-Net builds both encoder and decoder from stacked GDTBs for hierarchical feature extraction and reconstruction.The architecture is organized as a Transformer-based encoder-decoder restoration framework.
  • Architecture: BMSC progressively aggregates features from multiple encoder stages into balanced skip representations for effective cross-scale interaction.Supplementary skips retain low-frequency priors and highlight degradation residuals through lightweight pooling and convolution operations.
  • Selective Gated Attention: SGA splits the query projection into attention and gating branches, then uses a sigmoid gate to modulate linear-attention outputs selectively.This introduces data-dependent filtering that suppresses less informative tokens and emphasizes relevant global context.
  • Efficient Attention: Linear attention avoids explicitly constructing the N × N attention matrix, reducing time and memory complexity from O(N^2) to O(N).Sinusoidal query-key modulation adds structured relative-position reweighting while preserving linear complexity.
  • Dual-Scale Local Modeling: DGFF uses parallel depth-wise convolution branches with different receptive fields and a gating mechanism to model complementary local spatial patterns.This component targets local feature modeling under spatially heterogeneous degradations.

C. Balanced Multi-Scale Skip Connection (BMSC)

BMSC progressively fuses multi-level encoder features into a balanced skip representation, using a predictor–corrector integration scheme and lightweight refinement before decoder injection. This design suppresses unreliable shallow responses while preserving complementary details and semantics for stable cross-scale guidance.

  • Multi-scale integration: BMSC replaces direct skip addition or concatenation with progressive top-down integration of encoder features into a single balanced representation.Features are integrated from deep to shallow using a fixed order and then provided to the decoder as shared skip guidance.
  • Predictor–corrector fusion: The predictor–corrector scheme combines an Euler prediction with a correction step to reduce sensitivity to noisy updates during cross-scale fusion.The method interprets the update as a linear multistep integration scheme, using Adams–Bashforth prediction and Adams–Moulton correction.
  • Fusion operations: BMSC aligns encoder features through convolution and interpolation, then fuses them with element-wise addition followed by ReLU activation.The alignment ensures matched spatial resolution and channel dimension before fusion.
  • Skip refinement: A lightweight refinement module produces clean, stable multi-scale guidance that is resized and injected into all decoder stages.The refinement stage uses vHeat to model global context through diffusion rather than explicit pairwise token interactions.
  • Feature quality: BMSC’s balanced features show reduced noise responses and more coherent spatial structures than raw encoder features across multiple scales.These visualizations indicate improved cross-scale information integration before decoding.
  • Skip refinement: vHeat provides global receptive fields through frequency-domain diffusion while avoiding the quadratic cost of conventional self-attention.The HCO formulation uses DCT2D and IDCT2D and achieves O(N^1.5) complexity for an N-pixel feature map.

D. Uncertainty-Aware Refinement Head (URH)

URH refines coarse restorations with multi-scale GDTB processing while modeling per-pixel uncertainty and brightness consistency for severe and ambiguous degradations.

  • Architecture: URH uses a compact four-stage U-Net with GDTB blocks to enhance multi-scale details and correct artifacts in coarse outputs.Standard skip connections are retained because cross-scale fusion is handled earlier in the network.
  • Uncertainty modeling: URH predicts per-pixel mean µ(x) and variance σ2(x), using variance as a task-driven signal of regional restoration difficulty.Severely degraded regions tend to have higher variance, while clean or well-observed regions tend to have lower variance.
  • Uncertainty modeling: BAE-Loss models pixel-wise uncertainty with a per-pixel Gaussian distribution and an energy-based scoring rule for probabilistic predictions.The Energy Score is approximated using 1000 Monte Carlo samples drawn from the predicted distribution.
  • Uncertainty modeling: The uncertainty formulation uses pixel-wise ℓ1 distances to provide stable gradients suited to high-resolution restoration under heavy degradation.Its non-local nature encourages probability mass near the ground truth rather than matching only a single point.
  • Brightness awareness: BAE-Loss incorporates brightness-aware regression to address global brightness shifts from haze, snow accumulation, or illumination changes.An adaptive weight based on a Bhattacharyya-distance measure balances the original prediction and its brightness-aligned counterpart.
  • Training objective: An annealed probabilistic weight lets URH prioritize stable brightness-aware reconstruction before progressively adding probabilistic supervision.The design is reported to produce sharper refinements and better-calibrated uncertainty under severe and ambiguous degradations.

E. Total Loss

The total loss supplements BAE-Loss with a correlation term that promotes global structural and intensity consistency with the ground truth.

  • Correlation loss: The final training objective incorporates a correlation loss Lcor.This term is defined using the Pearson correlation coefficient between the restored image IHQ and ground truth IGT.
  • Correlation loss: The correlation loss encourages global structural and intensity consistency, complementing pixel-wise supervision and uncertainty modeling in BAE-Loss.

A. Training Details

Training uses a five-stage progressive schedule over 300,000 iterations, increasing effective patch size while managing GPU memory through smaller batches.

  • Training schedule: 300,000 iterations are distributed across five progressive training stages on four NVIDIA H100 GPUs.The model is implemented in PyTorch and trained from scratch.
  • Optimization: The schedule increases effective patch size while reducing batch size to learn higher-resolution content without exceeding GPU memory limits.AdamW starts at 3 × 10−4, remains constant for 92,000 iterations, then decays to 1 × 10−6 with cosine annealing.

B. Datasets

The evaluation uses unified multi-weather data spanning synthetic and real snow, localized raindrops, and mixed rain–haze degradations.

  • Snow: Snow100K provides 100K synthetic snowy images, with 9,000 used for training and Snow100K-S, Snow100K-L, and Snow100K-Real for testing.The subsets represent small-particle snow, large-particle snow, and real snowy scenes.
  • Raindrop: Raindrop contains 1,319 real-world image pairs, split into 1,069 training pairs and 249 testing pairs for localized non-uniform occlusions.
  • Mixed weather: Outdoor-Rain contains 9,000 synthetic images combining rain streaks and fog to simulate mixed rain–haze conditions.
  • Protocol: Snow100K, Raindrop, and Outdoor-Rain are merged for multi-weather training, with evaluation on Snow100K-S/L, Raindrop, and Outdoor-Rain Test1.

C. Evaluation Metrics

The evaluation uses full-reference distortion and perceptual metrics, alongside no-reference measures that assess visual quality without ground-truth images.

  • Full-reference metrics: PSNR and SSIM are standard full-reference metrics for evaluating restoration quality against ground-truth images.PSNR derives from MSE, while SSIM measures luminance, contrast, and structure.
  • Full-reference metrics: Higher PSNR indicates lower pixel-wise error, whereas higher SSIM indicates better structural similarity.PSNR uses image dimensions and maximum pixel value in its formulation; SSIM ranges from −1 to 1.
  • Perceptual metrics: LPIPS measures perceptual distance between deep feature representations, with lower values indicating greater perceptual similarity.It captures texture, structure, and semantic differences beyond RGB-space comparisons.
  • No-reference metrics: Q-Align predicts discrete quality levels with a multimodal model to align machine-generated scores with human subjective judgments.
  • No-reference metrics: MUSIQ predicts perceptual quality from a single image by modeling image quality across multiple spatial scales.It combines local-detail and global-composition information without requiring a reference image.

V. EXPERIMENTS

Experiments evaluate UAR-Net against unified and task-specific restoration methods, reporting benchmark performance, perceptual quality, feature organization, and component ablations.

  • Experimental results and comparisons: UAR-Net consistently achieves state-of-the-art performance across unified and task-specific adverse-weather benchmarks.It surpasses Histoformer by +0.78 dB average PSNR in unified evaluation and leads PSNR and SSIM on Snow100K-S and Snow100K-L.
  • Experimental results and comparisons: UAR-Net improves unified restoration over Histoformer by +0.78 dB average PSNR, including gains of +1.32 dB on Outdoor-Rain and +0.93 dB on Snow100K-S.
  • Perceptual quality evaluation: UAR-Net achieves the lowest LPIPS and highest Q-Align and MUSIQ scores across all evaluated datasets.These results indicate improvements in both distortion reduction and perceptual fidelity.
  • T-SNE feature visualization: Compared with MODEM, UAR-Net produces more compact intra-class clusters and clearer inter-class separation in t-SNE feature visualizations.The organization suggests more condition-aware and disentangled representations across weather conditions.
  • Ablation studies: Adding GDTB, BMSC, BAE-Loss, and URH progressively raises average performance to 34.46 dB PSNR and 0.9500 SSIM in the full model.The ablation begins from a histogram-transformer baseline and reports incremental gains from each proposed component.
  • Ablation studies: Sinusoidal reweighting adds about +0.05 dB average PSNR, while linear multistep fusion improves PSNR by about +0.18 dB over simple averaging in BMSC.
  • Ablation studies: The H/4 × W/4 balanced feature size slightly outperforms alternatives while reducing computation, and a 300k BAE-Loss annealing step improves average PSNR/SSIM to 34.21/0.9480.

C. Complexity Analysis

Complexity analysis compares unified restoration models using average PSNR and computational cost on 128 × 128 input patches.

  • Complexity analysis: UAR-Net achieves the highest average PSNR among compared unified methods under the reported evaluation setting.The comparison is conducted on a single H100 GPU with 128 × 128 input patches.
  • Complexity analysis: The reported comparison indicates that UAR-Net delivers superior restoration quality despite differences in computational cost among the evaluated methods.
  • Complexity analysis: Figure 9 presents the accuracy–complexity trade-off using average PSNR versus FLOPs on 128 × 128 input patches.
Loading 2609.02434v1…