Vision Mamba has recently emerged as an efficient alternative to convolutional and Transformer-based architectures by modeling long-range spatial dependencies through recurrent state evolution. Existing studies primarily improve token serialization, scanning strategies, or expert routing, yet largely overlook a state level question question: how should recurrent states be organized for spatially heterogeneous visual restoration? In this paper, we show that forcing tokens with heterogeneous update behaviors to recurrently modify a shared hidden state leads to persistent cross-content interference, which we term State Contention. We further argue that effective recurrent modeling should organize state evolution according to token compatibility rather than token ordering. Based on this observation, we introduce State Isolation, a new recurrent state organization principle for Vision Mamba,tokens with compatible state-update tendencies evolve within independent recurrent state streams. To instantiate this principle, we propose IsoMamba, a Δ-guided architecture that constructs model-intrinsic update representations, discovers update-compatible token groups through label-free clustering, and routes them to isolated Mamba experts following a cluster--route--isolate paradigm. In addition, we develop a dynamic channel-wise reweighting objective to compensate for wavelength-dependent reconstruction difficulty in underwater images. Extensive experiments demonstrate that IsoMamba consistently achieves state-of-the-art performance on multiple underwater image enhancement benchmarks while providing cleaner hidden state responses, validating the effectiveness of state isolated recurrent modeling.
Figures & tables
Figure 1: (a) Existing Mamba serializes features into a single mixed sequence with the spatial-coupled state stream. (b) IsoMamba clusters tokens by Δ and routes state-isolated sequences into decoupled Mamba experts.
Figure 2: State Contention versus State Isolation. (a) Input image. (d) Token groups obtained from the Δ -derived update representation. (b), (e) Hidden-state responses of the shared-state baseline and IsoMamba, respectively. (c), (f) Their corresponding channel-wise absolute reconstruction residuals using same normalization range.
Figure 3: Overview of IsoMamba. The proposed framework consists of two complementary stages. (a) Overall architecture. (b) Global recurrent modeling establishes image-level context through a shared state stream. (c) State-Isolation MBlock reorganizes recurrent state evolution by separating tokens with incompatible update tendencies into independent state streams. (d) A cluster–route–isolate strategy realizes State Isolation using model-intrinsic update tendency derived from Δ .
Figure 4: The detailed structure of Mamba expert.
Method
Venue
UIEB
LSUI
EUVP
Complexity
PSNR ↑
SSIM ↑
MSE ↓
PSNR ↑
SSIM ↑
MSE ↓
PSNR ↑
SSIM ↑
MSE ↓
Params (M) ↓
FLOPs (G) ↓
Time (s) ↓
U-shape
TIP23
20.54
0.772
91.77
24.16
0.846
81.40
25.48
0.830
74.20
65.60
66.20
0.063
DM-UIE
MM23
23.46
0.861
76.30
27.27
0.849
66.66
26.86
0.865
57.04
10.70
66.89
0.197
GUPDM
MM23
23.99
0.877
75.49
27.67
0.867
56.97
26.21
0.850
64.45
1.49
95.80
0.696
Semi-UIR
CVPR23
23.49
0.860
79.86
27.46
0.878
63.57
26.85
0.876
60.57
3.31
72.88
0.785
wMamba
arXiv24
23.55
0.888
78.32
28.71
0.894
50.32
28.55
0.882
47.01
3.69
7.53
0.028
Table 1: Quantitative comparison on the full-reference datasets UIEB, LSUI, and EUVP. The best, second-best, and third-best results are shown in bold , underlined , and italic type, respectively.
Method
Venue
U45
C60
SQUID
UCIQE ↑
UIQM ↑
URanker ↑
CLIP-IQA ↑
UCIQE ↑
UIQM ↑
URanker ↑
CLIP-IQA ↑
UCIQE ↑
UIQM ↑
URanker ↑
CLIP-IQA ↑
U-shape
TIP23
0.567
3.105
1.242
0.435
0.545
2.660
0.956
0.457
0.544
2.230
0.554
0.524
DM-UIE
MM23
0.603
3.108
1.810
0.537
0.580
2.577
1.341
0.559
0.659
0.519
-0.070
0.560
GUPDM
MM23
0.573
3.028
1.591
0.528
0.550
2.591
1.150
0.556
0.577
2.486
0.952
0.568
Semi-UIR
CVPR23
0.605
3.131
2.031
0.510
0.585
2.648
1.596
0.556
0.588
2.309
1.134
0.586
wMamba
arXiv24
0.608
3.068
2.026
0.531
0.580
2.773
1.415
0.504
0.576
2.673
1.003
0.558
Table 2: Quantitative comparison on the no-reference datasets U45, C60, and SQUID. The best, second-best, and third-best results are shown in bold , underlined , and italic type, respectively. Tied values are assigned the same rank.
Figure 5: Visual comparison under paired and unpaired settings. In the paired example, IsoMamba recovers colors and fine structures that more closely agree with the reference. In the unpaired example, it suppresses severe color casts without introducing oversaturation or texture artifacts.
Method
UIEB
LSUI
HINT †
21.63/0.881/0.604/2.888
27.44/0.897/0.597/2.880
EVSSM †
20.72/0.875/0.604/ 2.935
26.75/0.883/0.593/2.835
MaIR
23.19/0.892/0.605/2.910
28.53/0.907/0.596/ 2.911
Ours
23.94/0.895/0.615 /2.849
29.01/0.909/0.599 /2.881
Table 3: Full resolution evaluation. Entries are PSNR ↑ /SSIM ↑ /UCIQE ↑ /UIQM ↑ . † denotes tiled inference while others use full-image inference.
Figure 6: Full resolution visual comparison.
Dataset
Shared State
State Isolation
Gain
UIEB
0.1870 ± 0.0750
0.8488 ± 0.0512
± 0.6618
LSUI
0.1906 ± 0.0913
0.8383 ± 0.0586
± 0.6477
Table 4: Group separation score Ssep of hidden state responses. Results are reported as mean ± standard deviation across images.
Variant
PSNR ↑
SSIM ↑
Params (M)
(i) Shared-state Decoder
24.04
0.917
12.76
(ii) Linear Route Signal
24.61
0.902
15.10
(iii) w/o Channel Reweight
24.32
0.916
15.23
(iv) IsoMamba
25.32
0.920
15.23
Table 5: Top-level component ablation.
Figure 7: Qualitative comparison of top-level variants . Full IsoMamba suppresses the localized yellow artifact around the diver and produces properer foreground colors.
Figure 8: Routing assignments from different representations. Colors denote routed expert identities.
Variant
PSNR ↑
SSIM ↑
Params (M)
(i) Image invariant prototypes
24.51
0.912
14.59
(ii) w/o order label matching
24.27
0.917
15.23
(iii) w/o global state coordination
24.88
0.918
15.06
(iv) IsoMamba
25.32
0.920
15.23
Table 6: Internal mechanisms of State Isolation.
Figure 9: Effect of order-label matching on cluster–expert correspondence. Ours ordering produces more consistently separated trajectories.
Variant
PSNR ↑
SSIM ↑
Time (s)
(i) Soft expert mixing
24.17
0.909
0.106
(ii) Top-2 expert mixing
24.59
0.912
0.102
(iii) Random Top-1 routing
24.54
0.913
0.089
(iv) Δ -guided Top-1
25.32
0.920
0.095
Table 7: State organization versus expert capacity.
Figure 10: Visual analysis of recurrent-state routing. Ours routing more effectively corrects the localized blue color cast while introducing fewer noise artifacts in the background.
Image
Resolution
Time (s)
Alloc. (GB)
PSNR / SSIM
689_img_.png
1800×1350
3.486
17.37
24.28 / 0.959
720_img_.png
1800×1300
3.181
16.74
18.41 / 0.813
731_img_.png
1800×1295
3.162
16.67
22.93 / 0.945
772_img_.png
1800×1295
3.192
16.66
24.33 / 0.895
796_img_.png
1800×1295
3.202
16.64
19.83 / 0.836
Table 1: Full-image inference on the five largest UIEB test images using one RTX 4090 GPU (batch size 1).
Figure 1: Qualitative full-image results on large-resolution test images. Each row shows one sample, indexed by the file id. From left to right are the input, our output, and the ground truth. All results are produced by a single full-image forward pass without tiling.
Variant
PSNR
SSIM
Param (M)
Time (s)
-w/ [3,3,3]
25.28
0.917
14.62
0.093
-w/ [4,3,2]
24.87
0.919
14.02
0.090
-w/ [2,3,4] (Ours)
25.32
0.920
15.23
0.097
Table 2: Effect of the number of clusters.
Measure
UIEB
LSUI
Shared State (fixed expert)
0.1870±0.0750
0.1906±0.0913
Isolated State (fixed expert)
0.5765±0.0933
0.5614±0.1019
Full Isolation
0.8488±0.0512
0.8383±0.0586
Paired gap
0.3895
0.3707
95% bootstrap CI
[0.3781,0.4007]
[0.3643,0.3770]
Table 3: Separation of raw recurrent hidden states. Shared and Isolated State use identical expert-0 operands; Full Isolation uses the corresponding expert for each group. Scores are mean ± standard deviation. The paired gap is Isolated minus Shared State.
Variant
Red
Green
Blue
Avg. ↓
Degraded Input
0.1574
0.0712
0.0608
0.0965
w/o Channel Reweighting
0.0125
0.0051
0.0038
0.0071
IsoMamba(Ours)
0.0117
0.0036
0.0041
0.0064
Table 4: Channel-wise color-histogram KL divergence. Lower values indicate better agreement with the ground-truth color distribution.
Variant
PSNR ↑
SSIM ↑
UCIQE ↑
(i) w/o ℓ1 Loss
24.76
0.916
0.617
(ii) w/o SSIM Loss
23.76
0.902
0.615
(iii) w/o Edge Loss
24.52
0.909
0.612
IsoMamba (Ours)
25.32
0.920
0.618
Table 5: Ablation study on loss function components.
Variant
UIEB
LSUI
-w/ EMA-updated
24.81 / 0.917
29.97 / 0.915
-w/o EMA-updated (Ours)
25.32 / 0.920
29.89 / 0.914
Table 6: Effect of EMA-updated order label on different test sets.
Figure 2: Visualization of Δ -driven cluster maps with different numbers of clusters. From left to right, we show the input image and the route maps for K=2 , K=3 , and K=4 . Each color denotes one cluster and its routed expert identity.
Figure 3: Additional qualitative comparisons on representative test images. We compare with F-RWKV, HINT, Concertormer, MaIR, and EVSSM. Grey and yellow boxes indicate zoom-in regions, and the number under each result is PSNR for that image.
Figure 4: Additional NR visual comparisons on UCIQE. We compare with representative strong baselines under this metric, and report the per-image UCIQE score beneath each result.
Task
Split
Dataset (images)
Noise
Train
BSD400 [ 2 ] + WED [ 36 ] (400 + 4,744)
Test
Kodak24 [ 13 ] (24)
Haze
Train
RESIDE- β -OTS [ 28 ] (72,135)
Test
SOTS-Outdoor [ 29 ] (500)
Rain
Train
Rain100L [ 58 ] (200)
Test
Rain100L [ 58 ] (100)
Table 7: Single-task dataset settings for conventional restoration. Each task is trained and tested independently; parentheses show image counts.
Kodak24
PSNR
SOTS
PSNR/SSIM
Rain100L
PSNR/SSIM
GoPro
PSNR/SSIM
LOL
PSNR/SSIM
σ=15
σ=25
σ=50
DnCNN
34.60
32.14
28.95
DehazeNet
22.46/0.851
MSPFN
33.50/0.948
DeblurGAN
28.70/0.858
URetinex
21.33/0.835
FFDNet
34.63
32.13
28.98
FDGAN
23.15/0.921
LPNet
33.61/0.958
Stripformer
33.08/0.962
SMG
23.81/0.809
ADFNet
34.77
32.22
29.06
DehazeFormer
31.78/0.977
DRSformer
38.14/0.983
HI-Diff
33.33/0.964
Retinexformer
25.16/0.845
DGUNet
34.56
32.10
28.91
Restormer
30.87/0.969
Restormer
36.74/0.978
MPRNet
32.66/0.959
MIRNet
24.14/0.835
Restormer
34.78
32.37
29.08
NAFNet
30.98/0.970
NAFNet
36.63/0.977
Restormer
32.92/0.961
Restormer
22.43/0.823
Table 8: Single-task restoration results, including denoising on Kodak24, dehazing on SOTS, deraining on Rain100L, deblurring on GoPro, and low-light enhancement on LOL. The top three PSNR results are marked in RED , BLUE , and GREEN , respectively.
Figure 5: Visual Comparison of dehazed results between our IsoMamba and other all-in-one methods on the SOTS-outdoor.
Figure 6: Visual Comparison of derained results between our IsoMamba and other all-in-one methods on the Rain100L.
Figure 7: Visual Comparison of Low Light enhancement results between our IsoMamba and other all-in-one methods on the LoL.