CoScan: Multi-Scale Content-Adaptive Space-Filling Scans for Causal State-Space Image Restoration
Abstract
Selective state-space models (SSMs) have recently shown strong potential for efficient image restoration. A central difficulty, however, is that a 2D feature map typically needs to be linearized into a 1D token sequence for causal state-space updates, making restoration performance highly sensitive to the imposed scan order. Existing deterministic scans follow space-filling principles but remain fixed and content-agnostic, and can yield fragmented or semantically inconsistent contexts under degradations. We propose CoScan, a content-adaptive scan-order design for causal vision SSMs. CoScan learns image-dependent space-filling traversals via differentiable proxy supervision driven by scan-quality objectives that encourage short-range coherence and causal predictability, while preserving spatial adjacency. To support hierarchical restoration backbones, we introduce an efficient multi-scale scan construction by coarsening a fine-level minimum spanning tree, producing coherent scans across feature resolutions. In addition, a frequency-guided start selection strategy is designed to place informative structures early in the sequence prefix to strengthen causal context formation. Plugged into existing Mamba-based backbones, CoScan yields consistent improvements across several restoration benchmarks with minimal overhead and shows promising results on high-level vision tasks. Code will be released upon acceptance.