💻 computer science

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

MambaRaw is an efficient, JPEG-guided raw image reconstruction framework that leverages State Space Models and a novel Spatial-Energy Coupled Context Modeling mechanism to achieve state-of-the-art performance on 4K images while significantly reducing computational costs and latency compared to existing attention-based methods.

Original authors: Peize Li, Fanhu Zeng, Tongda Xu, Xingguo Xu, Xinjie Zhang, Xingtong Ge, Haotian Zhang, Yan Wang

Published 2026-06-24

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Peize Li, Fanhu Zeng, Tongda Xu, Xingguo Xu, Xinjie Zhang, Xingtong Ge, Haotian Zhang, Yan Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a camera that takes two photos at once: a tiny, compressed "preview" (like a JPEG) that you see instantly on your screen, and a massive, high-quality "raw" file that holds all the hidden details of the scene.

The problem is that the raw file is huge. Sending it over the internet or storing it takes a lot of space and time. The "preview" is small and easy to send, but it's not good enough for professional editing.

The Goal:
The paper introduces a new method called MambaRaw. Its job is to act like a super-smart "reconstruction crew." It takes the tiny preview and a very small amount of extra data (metadata) to rebuild the massive, high-quality raw image. The challenge? Doing this for 4K images (which are like 4 million pixels) is usually too slow and requires too much computer power.

The Old Way vs. The New Way:

The Old Way (The "Brute Force" Approach): Previous methods tried to analyze every single pixel of the 4K image with the same level of intense scrutiny. It's like hiring a team of detectives to search every single room in a massive mansion, even the empty closets and the smooth, empty walls. This wastes a lot of time and energy.
The MambaRaw Way (The "Smart Scout" Approach): MambaRaw realizes that most of an image is actually boring (like a blue sky or a smooth wall). It only spends its heavy computing power on the "interesting" parts (like a face, a tree, or text).

How MambaRaw Works (The Two Magic Tools):

TileMambaBlock (The "Smart Scout"):
Imagine the 4K image is cut up into a grid of small tiles (like a mosaic). MambaRaw first quickly checks the "energy" of each tile.
- Low Energy: A smooth blue sky? The system ignores it or gives it a quick glance.
- High Energy: A detailed texture or a sharp edge? The system flags it as "Important!"
- The Result: It only sends the "Important" tiles to the heavy-duty processing brain. This saves a massive amount of time and computing power, just like a detective only searching the rooms where clues are likely to be found.
Energy-Aware Refinement (EAR) (The "Detail Polisher"):
Even after finding the important tiles, the system needs to make sure the details are perfect. Raw images have a weird quirk: the most important details are often very faint or "rare" (like a long-tail distribution).
- Think of this module as a specialized artist who knows exactly how to paint the faint, tricky parts of a picture. It uses a "residual" trick (starting with a blank canvas and just adding the necessary changes) to ensure the final image looks exactly like the original scene, capturing those subtle, high-quality details that other methods miss.

The Results:
The authors tested this on photos from Sony, Olympus, and Samsung cameras.

Better Quality: They rebuilt the images with much higher clarity (up to 1.4 dB better in technical terms) compared to the best existing methods.
Faster: Because they stopped wasting time on empty parts of the image, the whole process was about 9% faster.
Efficient: It uses less computer memory, meaning it can run on standard computers without needing a supercomputer.

In a Nutshell:
MambaRaw is like a smart renovation crew. Instead of painting every inch of a house with expensive gold paint, it identifies the fancy, detailed areas that need the gold and paints the plain walls with a quick, cheap coat. The result is a house that looks just as luxurious as if every inch were painted, but it was built much faster and cheaper.

Technical Summary: MambaRaw

Problem Statement

High-resolution raw image reconstruction faces a critical bottleneck in balancing high-fidelity recovery with computational efficiency, particularly at 4K resolutions. While in-camera JPEG previews provide a useful sRGB reference with negligible storage cost, existing metadata-based reconstruction frameworks struggle to leverage this side information effectively at high resolutions.

The primary limitations of current state-of-the-art methods (e.g., R2LCM, Beyond-R2LCM) are twofold:

Computational Cost: Context models relying on self-attention mechanisms scale quadratically with feature map size, making 4K inference prohibitively expensive in terms of memory and latency.
Inefficiency in Processing: Standard approaches apply heavy global context modeling uniformly across the entire image. This is redundant, as raw images contain large smooth regions (low energy) that require less processing than texture-rich regions (high energy).
Distribution Mismatch: Raw signals exhibit uneven, camera-dependent channel distributions with long-tail energy characteristics that uniform context models often fail to capture accurately.

Methodology: MambaRaw

MambaRaw is a JPEG-conditioned, metadata-based raw image reconstruction framework designed to address these bottlenecks using State Space Models (SSMs). The architecture integrates SSMs into the entropy parameter estimation network, replacing standard convolutional or attention-based context models with a Spatial-Energy Coupled Context Modeling paradigm.

The framework operates on a hyperprior-based structure where a compact metadata bitstream and an aligned JPEG preview are used to reconstruct the raw image. The core innovation lies in two lightweight modules:

1. TileMambaBlock (Efficiency)

To mitigate the quadratic complexity of global modeling, MambaRaw introduces content-adaptive selective scanning:

Tile Partitioning: The input feature map is divided into non-overlapping tiles.
Energy-Guided Selection: The L2 energy of each tile ( $E = f^2$ ) is computed as a proxy for information density. High-energy tiles correspond to detailed foregrounds (edges, textures), while low-energy tiles correspond to smooth backgrounds.
Selective Scanning: A keep ratio ( $\rho$ ) determines the subset of top- $k$ informative tiles. The computationally expensive Mamba-style selective scanning is applied only to these high-energy tiles. Smooth regions are bypassed, significantly reducing FLOPs without sacrificing perceptual fidelity.

2. Energy-Aware Refinement (EAR) (Effectiveness)

To address the long-tail energy distribution of raw signals within the processed regions, MambaRaw employs EAR:

Mechanism: An identity-initialized residual module that dynamically adjusts feature responses based on local spatial energy statistics.
Function: Unlike global channel recalibration methods (e.g., SENets), EAR preserves local spatial granularity. It uses the spatial energy map to gate a residual variation, enhancing the representation of complex signals and aligning the feature distribution with the intrinsic properties of raw data.

Key Contributions

The paper claims three primary contributions:

A New Paradigm: The proposal of a spatial–energy coupled context modeling approach for JPEG-guided raw reconstruction, which resolves high-resolution bottlenecks by combining global spatial scanning with energy-guided entropy refinement.
Novel Modules: The introduction of TileMambaBlock and Energy-Aware Refinement (EAR). These enable content-adaptive selective scanning for efficient global modeling on information-dense tiles and enhance spatial entropy features to match long-tail raw distributions.
State-of-the-Art Performance: Extensive experiments demonstrating that the method achieves consistent rate–distortion improvements over strong metadata-based baselines while significantly reducing computational latency.

Experimental Results

The authors evaluated MambaRaw on three diverse camera datasets (Sony, Olympus, Samsung) and the AdobeFiveK dataset.

Rate-Distortion Performance:
- On the NUS dataset, MambaRaw outperforms strong baselines (including Beyond-R2LCM) by 1.2–1.4 dB in PSNR at comparable or lower metadata bitrates.
- On the AdobeFiveK dataset, at low bitrates (< 5e-4 bpp), MambaRaw achieves 58.55 dB PSNR, surpassing Beyond-R2LCM (58.44 dB) while requiring approximately 16% fewer bits.
Efficiency at 4K Resolution:
- At full 4K resolution (3840×2160), MambaRaw reduces FLOPs by approximately 56% and lowers end-to-end coding latency by 9% compared to the Beyond-R2LCM baseline.
- Memory usage is significantly reduced (10.2 GB vs. 22.8 GB for the baseline), allowing 4K inference on consumer-level hardware where the baseline approaches device limits.
Ablation Studies:
- The combination of EAR and TileMambaBlock yields the best trade-off. Removing the tile selection (using dense SSM) increases latency by ~12% with marginal PSNR gains.
- L2 Energy is identified as the most efficient metric for tile selection compared to Entropy or Gradient magnitude.
- SSM-based blocks outperform both CNN and Transformer (Swin) alternatives in the 4K context, offering superior performance with lower latency than Transformers.

Significance and Claims

The paper positions MambaRaw as a solution that bridges the gap between high-fidelity raw image reconstruction and practical computational efficiency. By leveraging the sparsity of high-frequency information in natural images, the method avoids the "one-size-fits-all" computational cost of previous high-resolution models.

The authors claim that MambaRaw sets a new state of the art for JPEG-guided raw reconstruction, offering a superior balance between reconstruction quality and inference speed. The work demonstrates that selective state space modeling is a viable and efficient alternative to attention mechanisms for high-resolution image processing tasks, particularly where metadata bitrates are constrained. The code is released to facilitate further research in efficient raw image processing.