State-Isolated Vision Mamba via Update-Tendency Routing for Underwater Image Enhancement

Yan Luo; Laibin Chang; Zefan Feng; Yi'ang Chen; Kui Jiang; Zhengguo Li

cs.CVOct 9, 2026

State-Isolated Vision Mamba via Update-Tendency Routing for Underwater Image Enhancement

Authors: Yan Luo, Laibin Chang, Zefan Feng, Yi'ang Chen, Kui Jiang, Zhengguo Li

Abstract

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 $\Delta$ 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 $\Delta$ -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 $\Delta$ .

Figure 4: The detailed structure of Mamba expert.

Method	Venue	UIEB			LSUI			EUVP			Complexity
Method	Venue	PSNR $\uparrow$	SSIM $\uparrow$	MSE $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	MSE $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	MSE $\downarrow$	Params (M) $\downarrow$	FLOPs (G) $\downarrow$	Time (s) $\downarrow$
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
Method	Venue	UCIQE $\uparrow$	UIQM $\uparrow$	URanker $\uparrow$	CLIP-IQA $\uparrow$	UCIQE $\uparrow$	UIQM $\uparrow$	URanker $\uparrow$	CLIP-IQA $\uparrow$	UCIQE $\uparrow$	UIQM $\uparrow$	URanker $\uparrow$	CLIP-IQA $\uparrow$
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

\uparrow

/SSIM

\uparrow

/UCIQE

\uparrow

/UIQM

\uparrow

\dagger

denotes tiled inference while others use full-image inference.

Figure 6: Full resolution visual comparison.

Dataset	Shared State	State Isolation	Gain
UIEB	0.1870 $\pm$ 0.0750	0.8488 $\pm$ 0.0512	$\pm$ 0.6618
LSUI	0.1906 $\pm$ 0.0913	0.8383 $\pm$ 0.0586	$\pm$ 0.6477

Table 4: Group separation score

S_{\mathrm{sep}}

of hidden state responses. Results are reported as mean

\pm

standard deviation across images.

Variant	PSNR $\uparrow$	SSIM $\uparrow$	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 $\uparrow$	SSIM $\uparrow$	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 $\uparrow$	SSIM $\uparrow$	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) $\Delta$ -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\times 1350$	3.486	17.37	24.28 / 0.959
720_img_.png	$1800\times 1300$	3.181	16.74	18.41 / 0.813
731_img_.png	$1800\times 1295$	3.162	16.67	22.93 / 0.945
772_img_.png	$1800\times 1295$	3.192	16.66	24.33 / 0.895
796_img_.png	$1800\times 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\pm 0.0750$	$0.1906\pm 0.0913$
Isolated State (fixed expert)	$0.5765\pm 0.0933$	$0.5614\pm 0.1019$
Full Isolation	$0.8488\pm 0.0512$	$0.8383\pm 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

\pm

standard deviation. The paired gap is Isolated minus Shared State.

Variant	Red	Green	Blue	Avg. $\downarrow$
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 $\uparrow$	SSIM $\uparrow$	UCIQE $\uparrow$
(i) w/o $\ell_{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 $\Delta$ -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)
Noise	Test	Kodak24 [ 13 ] (24)
Haze	Train	RESIDE- $\beta$ -OTS [ 28 ] (72,135)
Haze	Test	SOTS-Outdoor [ 29 ] (500)
Rain	Train	Rain100L [ 58 ] (200)
Rain	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
	$\sigma=15$	$\sigma=25$	$\sigma=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.

Explore similar work

Apr 26, 2026

BVI-Mamba: Video Enhancement Using a Visual State-Space Model for Low-Light and Underwater Environments

Videos captured in low-light and underwater conditions often suffer from distortions such as noise, low contrast, color imbalance, and blur. These issues not only limit visibility but also degrade automatic tasks like detection. Post-processing is typically required but can be time-consuming. AI-based tools for video enhancement also demand significantly more computational resources compared to image-based methods. This paper introduces a novel framework, Visual Mamba, designed to reduce memory usage and computational time by leveraging the Visual State Space (VSS) model. The framework consists of two modules: (i) a feature alignment module, where spatio-temporal displacement between input frames is registered in the feature space, and (ii) an enhancement module, where noise removal and brightness adjustment are performed using a UNet-like architecture, with all convolutional layers replaced by VSS blocks. Experimental results show that the Visual Mamba technique outperforms Transformer and convolution-based models in both low-light and underwater video enhancement tasks. Code is available on line at https://github.com/russellllaputa/BVI-Mamba.

Guoxi Huang, Ruirui Lin, Yini Li +2

Visual Information Laboratory, University of Bristol, Bristol, UK

Underwater Imaging Low-Light Image Enhancement Visual SSMs

Ask PDF

Mar 4, 2026

Separators in Enhancing Autoregressive Pretraining for Vision Mamba

The state space model Mamba has recently emerged as a promising paradigm in computer vision, attracting considerable attention for its efficient handling of long-sequence tasks. Its inherent causal structure makes it particularly well suited for autoregressive pretraining. However, existing autoregressive pretraining methods in vision are largely limited to short-sequence settings and may not fully use Mamba's capacity to model longer contexts. To investigate this setting, we introduce SeparaTors for AutoRegressive pretraining (STAR), a new autoregressive pretraining method for Vision Mamba that explicitly marks the boundaries between different images. STAR increases the patch-token sequence length from 144 to 640 by packing four images and four separator clusters. This is approximately

4.4\times

the ARM patch-token sequence length. The increase is achieved without changing the resolution of any individual image: we use

192\times192

inputs for autoregressive pretraining and

224\times224

inputs for downstream classification fine-tuning. With this long-sequence pretraining scheme, STAR-B achieves 83.5% EMA top-1 accuracy on ImageNet-1K after 1,600 epochs of pretraining. The learned representation also transfers beyond in-distribution classification: compared with ARM, STAR-B improves COCO box AP from 46.11 to 46.84 and mask AP from 40.74 to 41.45, while raising the mean top-1 accuracy across five ImageNet robustness benchmarks from 55.1% to 56.8%. Under the evaluated four-image setting, these results indicate that separator-based long-sequence pretraining improves recognition robustness and dense visual prediction relative to ARM.

Hanpeng Liu, Zidan Wang, Shuoxi Zhang +2

Huazhong University of Science and Technology

Visual Autoregressive Models Visual SSMs Self-Supervised Visual Representation Learning Image Classification

Ask PDF

Apr 17, 2026

Hero-Mamba: Mamba-based Dual Domain Learning for Underwater Image Enhancement

Underwater images often suffer from severe degradation, such as color distortion, low contrast, and blurred details, due to light absorption and scattering in water. While learning-based methods like CNNs and Transformers have shown promise, they face critical limitations: CNNs struggle to model the long-range dependencies needed for non-uniform degradation, and Transformers incur quadratic computational complexity, making them inefficient for high-resolution images. To address these challenges, we propose Hero-Mamba, a novel Mamba-based network that achieves efficient dual-domain learning for underwater image enhancement. Our approach uniquely processes information from both the spatial domain (RGB image) and the spectral domain (FFT components) in parallel. This dual-domain input allows the network to decouple degradation factors, separating color/brightness information from texture/noise. The core of our network utilizes Mamba-based SS2D blocks to capture global receptive fields and long-range dependencies with linear complexity, overcoming the limitations of both CNNs and Transformers. Furthermore, we introduce a ColorFusion block, guided by a background light prior, to restore color information with high fidelity. Extensive experiments on the LSUI and UIEB benchmark datasets demonstrate that Hero-Mamba outperforms state-of-the-art methods. Notably, our model achieves a PSNR of 25.802 and an SSIM of 0.913 on LSUI, validating its superior performance and generalization capabilities.

Tejeswar Pokuri, Shivarth Rai

Manipal Institute of Technology, Manipal Academy of Higher Education, Manipal - 576104, India

Underwater Image Enhancement Underwater Imaging

Ask PDF