MambaRaw: Selective State Space Modeling for Efficient 4K Raw Image Reconstruction

Abstract: In-camera JPEG previews are ubiquitous in raw image formats and provide an sRGB reference at negligible storage cost. Although existing metadata-based reconstruction frameworks can exploit this side information when recovering raw images, their context models often become computationally expensive especially at high resolution, eg, 4K raw image, given that attention mechanisms scale quadratically with feature maps, hindering its practical application. To address these limitations, we propose MambaRaw, a JPEG-conditioned metadata-based raw image reconstruction framework that uses State Space Models (SSMs) to estimate entropy parameters efficiently. Our key contribution comprises a Spatial-Energy Coupled Context Modeling mechanism with two lightweight modules: (1) TileMambaBlock, which performs Mamba-style selective scanning only on information-dense tiles to improve the efficiency; and (2) Energy-Aware Refinement (EAR), an identity-initialized residual module that enhance feature representation to match the long-tail energy distribution of raw signals. Extensive experiments on three camera datasets (Sony, Olympus, Samsung) show consistent improvements over strong metadata-based baselines and set a new state of the art for JPEG-guided raw reconstruction with great efficiency. Notably, at low metadata bitrates, MambaRaw increases PSNR by 1.2--1.4 dB and reduces end-to-end coding latency by about 9%. Code is released at https://github.com/Peizeli1/MambaRaw.

• The paper introduces a novel JPEG-conditioned framework using selective state space modeling that enhances 4K raw image reconstruction efficiency.
• It implements TileMambaBlock for selective spatial modeling and Energy-Aware Refinement (EAR) to dynamically calibrate entropy, reducing FLOPs and latency.
• Experimental results show consistent 1.2–1.4 dB PSNR gains over baselines with reduced bitrate, demonstrating superior rate-distortion performance.

Efficient State Space Modeling for High-Resolution RAW Reconstruction: The MambaRaw Framework

Introduction

The paper "MambaRaw: Selective State Space Modeling for Efficient 4K Raw Image Reconstruction" (2606.24479) introduces a novel JPEG-conditioned framework for metadata-based RAW image reconstruction at 4K resolution. The standard context modeling approaches, notably convolutional and Transformer-based models, face scalability bottlenecks either in receptive field coverage or computational/memory complexity, especially for high-resolution inputs. MambaRaw addresses these issues by integrating state space models (SSMs) in entropy parameter estimation and implementing spatial-energy coupled context modules: TileMambaBlock for selective spatial modeling and Energy-Aware Refinement (EAR) for entropy adaptation. The approach results in state-of-the-art rate-distortion performance and reduced inference latency across diverse camera datasets (Samsung, Olympus, Sony), providing compelling evidence for the utility of SSM-based selective modeling in RAW image reconstruction.

Motivation and Architecture

High-resolution raw signals exhibit spatially non-uniform information distribution, with textures and structural details concentrated in high-energy regions. Traditional dense context modeling spends significant compute on smooth backgrounds, yielding diminishing returns. MambaRaw operationalizes this by:

• Applying state space models selectively ("TileMambaBlock") to information-dense tiles, identified by a lightweight L2 energy metric.
• Refining the entropy parameter network with EAR, which calibrates feature responses in accordance with the intrinsic long-tail energy distribution of raw data, enhancing entropy modeling.

The backbone architecture leverages a two-level VAE structure, conditioned at each scale on an in-camera JPEG preview via concatenation of resized features. The context model replaces separable spatial/channel blocks with a coupled design, allowing the integration of SSM-driven spatial reasoning onto selected tiles, further refined by spatially-dependent EAR gates.

Figure 1: The MambaRaw framework combines a two-level VAE conditioned on JPEG preview, with a selective TileMambaBlock and EAR for efficient spatial-energy coupled context modeling.

Selective Context Modeling: TileMambaBlock and Energy-Aware Refinement

TileMambaBlock

TileMambaBlock partitions context features into non-overlapping $T \times T$ tiles. Using normalized channel-wise L2 energy as a proxy for information density, it selects the top $\rho$ fraction of tiles for SSM-based spatial modeling; the remainder pass through unchanged. This mechanism reduces FLOPs and preserves global context only where needed, transitioning to dense processing for small patches or high keep ratios. Ablations demonstrate that L2 Energy selection outperforms random, entropy, or gradient-based metrics in speed and reconstruction quality.

Figure 2: (a) Distribution of raw feature energy motivates selective modeling; (b) visualization of selected high-energy tiles; (c) performance and latency trade-off as a function of keep ratio $\rho$, with optimal balance at $\rho=0.5$.

Energy-Aware Refinement (EAR)

EAR performs spatial enhancement by computing local energy maps and constructing identity-initialized gated residuals. Unlike global channel recalibration (e.g., SENet), EAR maintains granularity, dynamically modulating feature representations to accommodate energy variances inherent to RAW data. EAR is implemented via lightweight $1\times1$ convolution branches, gated by energy maps and stabilized to start as an identity mapping, ensuring robust training.

Experimental Results

Rate-Distortion and Efficiency

MambaRaw achieves consistent improvements over strong baselines (Beyond-R2LCM, CAM, R2LCM) across all rate-distortion points and camera datasets. On the NUS dataset, it yields 1.2–1.4 dB PSNR gains at comparable metadata bitrates, with a strict improvement in low-bitrate regimes, outperforming Beyond-R2LCM by 0.11 dB and using 16% fewer bits on AdobeFiveK. Importantly, MambaRaw maintains performance at true 4K resolution ($3840 \times 2160$), reducing FLOPs by 56% and wall-clock latency by 9% compared to baseline CNNs, preserving high fidelity with practical compute requirements.

Figure 3: Rate-distortion curves over the NUS benchmark, demonstrating MambaRaw’s dominance in both PSNR and SSIM across multiple cameras.

Ablation and Component Analysis

Component-wise analysis reveals that EAR alone confers steady PSNR gains with minimal latency overhead. Dense SSM integration delivers substantial rate-distortion improvement, but only when combined with tile-wise selection does the model achieve optimal latency-performance trade-off, reducing inference time by up to 12%. SSM-based context blocks outperform both CNN and Transformer alternatives, achieving highest PSNR and SSIM with reduced complexity.

Qualitative Visualization

Qualitative results substantiate quantitative findings, with error maps showing MambaRaw’s capacity to faithfully reconstruct complex textures and suppress residuals in high-frequency regions relative to baselines.

Figure 4: Qualitative comparison showing per-pixel absolute error maps; MambaRaw suppresses error in textural and edge regions.

Figure 5: Additional qualitative results further confirm structural fidelity and error suppression across scenes and channels.

Implications and Future Work

MambaRaw demonstrates the practical effectiveness of SSM-based selective context modeling in high-resolution RAW image reconstruction under constrained metadata bitrates. The strict efficiency gains and structural fidelity achieved suggest significant potential for deployment in computational photography pipelines, especially in mobile or resource-limited scenarios.

Theoretical implications include the viability of spatial-energy coupling as a proxy for information-driven adaptive context processing and the extensibility of SSMs in entropy modeling for vision tasks beyond RAW reconstruction. Future directions include:

• Extension to raw video reconstruction via temporal SSMs (e.g., VideoMamba),
• Hardware-aware accelerator mapping for real-time inference,
• Exploration of multi-modal context modeling using richer metadata.

Conclusion

MambaRaw provides an authoritative solution to the longstanding efficiency bottleneck in 4K RAW image reconstruction. Through spatial-energy coupled context modeling, integrating TileMambaBlock and EAR, it achieves strict rate-distortion superiority and reduced computational latency relative to state-of-the-art baselines. The selective SSM approach advances both theoretical and practical understanding of adaptive context modeling in learned image compression and opens avenues for future research in high-resolution, bandwidth-constrained vision applications.