GeoMamba: Geometry-Aware State Space Modeling for Image Restoration
Abstract
State space models (SSMs) provide an efficient route to global image restoration because they model long sequences with linear complexity. Recent restoration architectures further enlarge the effective context by reordering image tokens so that semantically similar pixels become neighbors in the 1D scan. This strategy improves texture aggregation, but it also exposes a mismatch between the geometry of images and the discretization used by SSMs: adjacent tokens in the reordered sequence can be spatially distant, while the recurrent update still treats them as consecutive samples of a continuous process. In addition, existing prompt-based modulation relies on a finite set of discrete prompts, which is poorly matched to the continuous variation of natural textures and offers no explicit compensation for the low-pass tendency of recurrent integration. We propose GeoMamba, a geometry-aware and continuously modulated SSM for image restoration. GeoMamba introduces geometry-adaptive discretization (GAD), which conditions the SSM time-step on the Euclidean distance between consecutive reordered tokens and attenuates history propagation across spatial jumps. It also introduces a continuous dynamic matrix (CDM), which regresses pixel-wise modulation parameters from feature and gradient cues to reduce prompt quantization and preserve high-frequency details. Experiments on lightweight and classical image super-resolution demonstrate consistent improvements over strong Transformer- and Mamba-based baselines, while achieving better or competitive performance in additional denoising and JPEG artifact reduction settings.