HDMba: Hyperspectral Remote Sensing Imagery Dehazing with State Space Model

To ensure effective and direct dehazing feature extraction, we adopted an end-to-end multi-scale feature extraction framework based on RDM blocks, as shown in Fig. 2 (a). Downsampling in the U-net is discarded, which allows the high-frequency information to be preserved. HDMba first applies a 3×3 convolution layer to extract low-level features ${F}_{0}\in\mathbb{R}^{W\times H\times C}$ from the haze image ${X}\in\mathbb{R}^{W\times H\times B}$, where $W$ and $H$ represent the spatial dimensions, and $C$ and $B$ are the number of feature channels and the number of input bands, respectively. Then, ${F}_{0}$ is processed through several RDM blocks to extract deep features for scene recovery with the feature size of $W\times H\times C$, and this process can be expressed as:

where $\mathrm{RDM}_{i}$ represents the $i$-th RDM block, and $\mathcal{F}_{i}$ denotes the $i$-th spatial-spectral feature obtained by this block. The network is designed with 4 RDM blocks. To recover a clean scene $\mathcal{Y}\in\mathbb{R}^{W\times H\times B}$ from the deep feature $\mathcal{F}_{I}$, two 3×3 convolution layers are concatenated with the shallow feature via skip connections. The global skip connection between the hazy image and the output enables the intermediate network to learn the irregular and uneven haze characteristics distributed across the spectral and spatial domains. A combination of mean squared error (MSE) and $L_{1}$ norm is used as loss function for network training:

where $\left\|\cdot\right\|_{mse}$ represents the MSE loss and $\left\|\cdot\right\|_{1}$ represents the $L_{1}$ norm. $\theta_{1}$ and $\theta_{2}$ are the weighting factors.

Haze exhibits an irregular spatial distribution, making it essential to focus on local haze areas and extract crucial information from them for scene recovery. We designed the RDM block (Fig. 2 (b)) for deep local and long-range information modelling, which mainly contains multiple DMLs, represented as follows:

where $\mathrm{DML}_{k}$ represents the $k$-th DML and $\mathcal{F}_{k}$ denotes the $k$-th spatial-spectral feature in RDM block. Each RDM block concludes with a 3×3 convolution layer, adding the residual features from the previous block via a skip connection. For each DML (Fig. 2 (c)), we combine the normalized layer with WSSM and MLP respectively to improve the nonlinearity characterization while modelling spatial information. The process can be defined as follows:

where $\mathcal{F}_{t-1}$ represents the input feature embedding of DML, $\mathcal{F}_{t}^{\prime}$ and $\mathcal{F}_{t}$ represent the output of $\mathrm{WSSM}$ and $\mathrm{MLP}$, respectively, and $\mathrm{LN}$ represents the Layer Normalization layer.

Currently, most visual Mamba approaches primarily capture long-term dependencies by increasing scanning directions, but they lack the ability to effectively capture local spatial details and inter-regional correlations [21]. We proposed a novel WSSM to enhance the processing capacity in local areas where haze is distributed, as shown in Fig. 2 (d).

Specifically, for the input feature embedding $\mathcal{Z}_{ip}\in\mathbb{R}^{W\times H\times C}$, the window partition [22] is first performed in spatial dimension with the window size of $M$, resulting in $\frac{W\times H}{M^{2}}$ overlapping patches $z_{ip}\in\mathbb{R}^{M^{2}\times C}$. These patches then capture local dependencies through Mamba, while retaining the correlation of different local regions, ensuring a comprehensive analysis of haze regions in the image. Finally, through a window reverse operation, the output patches $z_{op}\in\mathbb{R}^{M^{2}\times C}$ from Mamba are gathered to obtain the final features $\mathcal{Z}_{op}\in\mathbb{R}^{W\times H\times C}$. This process can be expressed as follows:

where $\left\{z_{ip}\right\}\in\mathbb{R}^{\frac{WH}{M^{2}}\times M^{2}\times C}$ and $\left\{z_{op}\right\}\in\mathbb{R}^{\frac{WH}{M^{2}}\times M^{2}\times C}$. For Mamba (Fig. 2 (e)), the input spatial feature sequence enters two branches after passing through RMSNorm. In the main branch, the features undergo successive processing by a linear layer, depth-wise separable convolution, a SiLU activation function, and SSM, effectively integrating haze characteristics from various regions. The other branch passes through a linear layer and a SiLU function and multiplies the output of the main branch. A normalization layer produces the final output. The process can be represented as follows:

where $z_{op1}$ and $z_{op2}$ represent the output of the two branches respectively. $\mathrm{DConv}\left(\cdot\right)$ represents the depth-wise separable convolution, $\mathrm{\phi}\left(\cdot\right)$ denotes the SiLU activation function and $\otimes$ indicates the element-wise multiplication.

To evaluate the effectiveness of the proposed HDMba, we used two Gaofen-5 HSI datasets: HyperDehazing dataset¹¹1The dataset will be publicly available at https://github.com/RsAI-lab/HyperDehazing and HDD [6]. The HyperDehazing dataset is synthesized based on the DOS model using 100 clean scenes with 20 different haze thicknesses and 5 haze abundances. It contains a total of 2000 hazy and haze-free image pairs, each with a size of 512×512×305. 90% of this dataset was used for network training, while the remaining 10% was reserved for testing. HDD comprises 20 reference-free hazy images, each with dimensions 512×512×305. All images in this dataset were used for network testing. To meet memory requirements, we cropped the training data to 64×64 and the test data to 128×128.

We set the number of DMLs $K$=4, the window size $M$=8. $\theta_{1}$ and $\theta_{2}$ are set to 1 and 0.1, respectively. The batch size is set to 4, and all datasets are trained for 10,000 iterations. The Adam optimization operator was employed to accelerate the training, where the momentum parameters were set to 0.9, 0.999, and 10^-8, respectively. The initial learning rate was set to 1×10^-4, with the cosine annealing strategy to adjust the learning rate. The whole network was implemented on the PyTorch framework with an NVIDIA GeForce RTX 3060 GPU.

To quantitatively evaluate the dehazing performance, we used structural similarity index measurement (SSIM), peak signal-to-noise ratio (PSNR), universal image quality index (UQI), and spectral angle mapping (SAM) metrics for paired images, and natural image quality evaluator (NIQE) and average gradient (AG) metrics for real images. The results are shown in Table I. Transformer-based methods generally outperform CNN-based methods. However, HDMba achieves the best results across all metrics except UQI. When processing high-dimensional data, HDMba has significantly fewer parameters (4.60M) compared to most dehazing networks.

To visualize the effectiveness of dehazing, Fig. 3 presents dehazing images of several state-of-the-art methods. It is evident that LKDNet, AIDTransformer, and RSDFormer do not completely remove the haze. The scenes recovered by CANet and SGNet exhibit spectral distortions. While AACNet and Restormer manage to recover parts of the clean scene, some haze residue remains. In contrast, the proposed HDMba recovers the result closest to the clean scene, with a good reconstruction of surface details.

As shown in Fig. 4, we selected building and vegetation scenes to compare spectra reconstruction performance. RSDFormer has an insufficient dehazing ability, resulting in spectral curves significantly higher than the reference value. AIDTransformer exhibits similar issues (Fig. 4(b)). AACNet, DehazeFormer, and Restormer produce spectra in the visible range that are lower than the reference, indicating excessive dehazing. In contrast, HDMba achieves spectra closest to the reference (Fig. 4 (b)), with spectral trends that are highly consistent despite some deviations in certain cases (Fig. 4 (a)).

The effectiveness of HDMba in processing hazy HSIs across different bands is quantitatively evaluated using the SSIM and PSNR metrics, as shown in Fig. 5. The results indicate that HDMba consistently outperforms other methods across most bands, achieving better performance than SGNet (which has the lowest SSIM) and DehazeFormer (which has the lowest PSNR). It should be noted that the performance of HDMba degrades in the 1430 nm and 1950 nm wavelength ranges, which are close to the water and atmospheric absorption bands and can cause severe interference. Nevertheless, our method exhibits excellent dehazing effects across nearly all spectral bands.

We conducted ablation experiments on HyperDehazing to investigate the effectiveness of each component of the proposed model. These experiments included evaluating the impact of constituent elements within the DML, as well as exploring the effect of different partition window sizes in the WSSM on the dehazing results.

1) Analysis of DML: We trained the network with variations in constituent elements of Mamba and MLP, presenting the corresponding dehazing results in Table II. It is evident that the combination of SSM, DWConv1d, and multiplication in Mamba yields excellent dehazing performance, with further enhancement achieved by incorporating MLP.

2) Analysis of window size: We assessed the effect of window size in the WSSM on dehazing, and the results are presented in Table III. As we can see, larger window sizes lead to improved performance but also increase computational costs. We set the window size to 8 to strike a balance between performance and computation cost.