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

2026-06-23

AI Summary

This work addresses the high computational cost of attention-based context modeling in high-resolution RAW image reconstruction from JPEG previews. To overcome this limitation, the authors propose MambaRaw, a novel framework that leverages state space models (SSMs) for efficient entropy parameter estimation and introduces TileMambaBlock to enable selective scanning of information-dense regions. Additionally, an Energy-Aware Refinement module with identity initialization is designed to accurately capture the long-tailed energy distribution inherent in RAW signals. Guided by JPEG metadata, the proposed method achieves significant performance gains over existing approaches on Sony, Olympus, and Samsung datasets, yielding PSNR improvements of 1.2–1.4 dB at low bitrates and reducing end-to-end encoding latency by approximately 9%.

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.

Problem

• raw image reconstruction
• 4K resolution
• computational efficiency
• metadata-based reconstruction
• context modeling

Innovation

• State Space Models
• Selective Scanning
• JPEG-guided Reconstruction
• Energy-Aware Refinement
• Efficient Raw Image Processing