LAR-U-Net: Lightweight Attentive Residual U-Net for Image Denoising
Course Project Notice: This repository and model card were developed as part of the academic project for
ARTI404: Image Processing.
LAR-U-Net is an advanced, lightweight image restoration architecture engineered specifically for grayscale image denoising. It structurally combines the spatial localization precision of a symmetric U-Net backbone, adaptive feature refinement via the Convolutional Block Attention Module (CBAM), and a Residual Noise Learning paradigm. Additionally, it leverages a handcrafted Gradient-Augmented Processing front-end to strictly preserve structural boundaries and high-frequency edge dynamics during restoration.
Architecture Overview
The primary challenge in discriminative image denoising is avoiding the over-smoothing of critical high-frequency details (e.g., fine textures, sharp edges). LAR-U-Net addresses this through a four-tiered architectural design:
- Gradient-Augmented Processing (Front-End): Instead of processing the raw noisy image in isolation, the input is passed through horizontal (G_x) and vertical (G_y) Sobel operators. The extracted gradient magnitude map (G) is concatenated along the channel dimension, forcing early spatial layers to prioritize structural edge maps.
- Symmetric U-Net Backbone: Utilizes an encoder-decoder architecture with skip connections, balancing parameter count and feature representation using optimized 3X3 convolutions.
- Dual-Axis Attention Gates (CBAM): Embedded directly inside each encoder stage to dynamically re-weight features:
- Channel Attention: Compresses spatial dimensions using adaptive global average pooling to capture cross-channel interdependencies (distinguishing signal from noise variants).
- Spatial Attention: Concatenates average and max pooling feature maps across the channel dimension and applies a 7X7 kernel to generate an explicit spatial mask highlighting destructive noise regions.
- Residual Noise Learning Paradigm: Adopting the discriminative mapping strategy popularized by DnCNN, the network directly predicts the underlying Noise Map R(y) rather than the clean image itself.
Mathematical Formulation
1. Image Degradation & Reconstruction Model
Given an observed noisy image y degraded by additive noise:
Where x denotes the latent clean image and v represents the corrupting Additive White Gaussian Noise (AWGN) component with standard deviation. The restored image (x) is computed via explicit residual subtraction:
2. Front-End Edge Gradient Extraction
Spatial edge dynamics are computed using discrete Sobel spatial operators:
3. Optimization Criterion
The network parameters (W) are optimized by minimizing the Mean Squared Error (MSE) loss regularized explicitly over the target residual noise map:
Performance & Experimental Results
Training Hyperparameters
- Optimizer: Adam
- Loss Function: Mean Squared Error (MSE)
- Epochs: 10
- Batch Size: 8
- Target Noise Setup: AWGN (sigma = 25)
- Core Dataset: Berkeley Segmentation Dataset 500 (BSDS500) β using disjoint train/val/test splits.
- Robustness Evaluation Suite: Multi Noises Dataset (Evaluated against Speckle, Poisson, Multiplicative, JPEG artifacts, Quantization, and Salt & Pepper noise).
Quantitative Metrics
- Minimum Training Loss:
0.00315(MSE) - Minimum Validation Loss:
0.00327(MSE) - Test Dataset MAE:
0.04253 - Average Suite PSNR:
19.9205 dB - Peak PSNR Gain: > +7.0 dB on structured natural scenes.
Training Dynamics
| Epoch | Train Loss (MSE) | Validation Loss (MSE) |
|---|---|---|
| 1 | 0.22759 | 0.24682 |
| 2 | 0.00983 | 0.00924 |
| 5 | 0.00415 | 0.00414 |
| 10 | 0.00315 | 0.00327 |
Sample-Level Reconstruction Analysis
| Evaluation Parameter | Sample 1 | Sample 2 |
|---|---|---|
| Initial Noisy PSNR | 20.21 dB | 20.19 dB |
| Denoised PSNR | 28.01 dB | 27.36 dB |
| Net PSNR Gain | +7.79 dB | +7.18 dB |
| Initial MAE | 0.0779 | 0.0783 |
| Denoised MAE | 0.0306 | 0.0331 |
| Net MAE Reduction | -0.0473 | -0.0452 |
Architectural Note on Domain Generalization: While LAR-U-Net yields quantitative improvements on impulsive noise types (e.g., Salt & Pepper), visual quality inspection confirms its performance is optimal for Gaussian-distributed (AWGN) noise. This is due to the inherent symmetry of residual noise mappings and Batch Normalization layers operating over normal probability distributions during SGD optimization.
Comparison with Alternative Architectures & Trade-offs
| Model Architecture | Key Strengths | Key Weaknesses | Optimization Loss |
|---|---|---|---|
| Standard U-Net | High spatial precision via skip connections; robust multi-scale localization. | High memory footprint; prone to producing blurry edges when trained on MSE alone. | MSE Loss / L1 Loss |
| DnCNN | Excellent Gaussian noise mapping via residual learning and Batch Normalization. | Lacks spatial attention; sensitive to structural domain shifts; computationally heavier at single resolution. | Residual MSE Loss |
| Restormer | State-of-the-art global contextual modeling via Transposed Self-Attention; exceptional benchmark results. | Extremely computationally expensive; heavy GPU memory overhead; requires huge datasets. | Frequency Loss |
| LAR-U-Net (Ours) | Lightweight; explicit edge preservation via Sobel front-end; spatial/channel feature refinement via CBAM. | Performance drops on heavy non-Gaussian noise distributions (e.g., extreme Salt & Pepper). | MSE Loss |