MambaRaw: JPEG-Conditioned RAW Reconstruction

• MambaRaw is a JPEG-conditioned RAW reconstruction framework that leverages in-camera JPEG previews and compact metadata to recover high-bit-depth images.
• It replaces heavy attention models with a spatial–energy coupled state space model that selectively applies long-range context processing in information-dense regions.
• Empirical evaluations show improvements up to 1.4 dB PSNR and a 9% reduction in coding latency for 4K RAW images compared to existing methods.

MambaRaw is a JPEG-conditioned, metadata-based raw image reconstruction framework for recovering high-bit-depth RAW images from an in-camera JPEG preview and a compact metadata bitstream. It replaces heavy attention-based context modeling with a state space model (SSM) pipeline specialized for entropy parameter estimation, and introduces a “spatial–energy coupled” design in which long-range modeling is applied selectively to information-dense regions while feature refinement is modulated by RAW feature energy. The method is built for high-resolution operation, particularly 4K RAW, and is reported to improve low-bitrate reconstruction quality by $1.2$–$1.4$ dB PSNR while reducing end-to-end coding latency by about $9\%$ (Li et al., 23 Jun 2026).

1. Problem setting and scope

MambaRaw addresses the metadata-based raw reconstruction setting in which a camera already stores an aligned JPEG or sRGB preview $\mathbf{x}_{\mathrm{jpg}}$ together with a RAW image $\mathbf{x}_{\mathrm{raw}}$. At capture time, the encoder computes a compact metadata bitstream $\mathbf{s}$ from the pair $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$. At reconstruction time, the decoder receives $\mathbf{x}_{\mathrm{jpg}}$ and $\mathbf{s}$, and reconstructs $\hat{\mathbf{x}}_{\mathrm{raw}}$. In this formulation, the JPEG is not itself being compressed; it is assumed to be stored anyway and functions as side information (Li et al., 23 Jun 2026).

The central motivation is that in-camera JPEG previews are ubiquitous in RAW containers, spatially aligned with the corresponding RAW data, and already encode substantial scene structure, including edges, textures, luminance, and color, albeit in rendered 8-bit sRGB form. This makes the JPEG a strong prior, allowing the metadata stream to focus on RAW-specific residual information such as linearization, noise characteristics, and sensor-dependent color statistics. A common misconception is that MambaRaw is a general-purpose RAW image generator; in fact, it is a conditional reconstruction system whose signal model is explicitly organized around JPEG guidance and entropy-coded metadata.

MambaRaw is motivated by limitations in prior metadata-driven frameworks such as SAM, CAM, R2LCM, and Beyond-R2LCM. Convolution-only context models have limited receptive field on 4K images, while attention-based context models incur quadratic complexity in the number of spatial positions $1.4$0, making them costly in FLOPs and memory. The design premise of MambaRaw is that information is spatially non-uniform, RAW feature energies are long-tailed and uneven across space and channels, and therefore expensive global modeling should be concentrated on complex regions rather than applied uniformly over the full latent feature map (Li et al., 23 Jun 2026).

2. Codec architecture and entropy modeling pipeline

MambaRaw builds directly on the JPEG-conditioned VAE codec of Beyond-R2LCM. Its high-level pipeline begins with an analysis transform

$1.4$1

which maps the RAW–JPEG pair to a latent representation $1.4$2. A hyperprior branch then produces

$1.4$3

where $1.4$4 serves as side information for entropy modeling. Reconstruction is obtained through an overview transform

$1.4$5

Training uses the standard rate–distortion objective

$1.4$6

with $1.4$7 given by MSE in raw-linear space (Li et al., 23 Jun 2026).

The SSM component is not used as the global backbone. Instead, it is inserted specifically into the Level-1 entropy parameter estimation network. Given JPEG-conditioned features at latent resolution, MambaRaw applies an input projection $1.4$8, followed by a TileMambaBlock and an Energy-Aware Refinement module, and then an output projection $1.4$9 that predicts Gaussian entropy parameters: $9\%$0 Each latent coefficient is modeled with an independent Gaussian,

$9\%$1

and the latent rate is

$9\%$2

JPEG guidance is injected at every scale by resizing the JPEG with bilinear interpolation and concatenating it with the current features. This conditioning is performed in both the analysis/synthesis transforms and the entropy parameter network, and is explicitly implemented via concatenation rather than attention, avoiding the cost of cross-attention while still exposing the model to multi-scale structure and color cues from the JPEG preview (Li et al., 23 Jun 2026).

3. Spatial–energy coupled context modeling

The defining principle of MambaRaw is “Spatial-Energy Coupled Context Modeling.” The approach couples spatial context modeling, which captures long-range spatial dependencies through an SSM, with energy modeling, which uses feature energy as a proxy for information density. Feature energy is defined as

$9\%$3

At tile level, MambaRaw computes the energy score

$9\%$4

for each tile $9\%$5, and uses this score to determine where selective SSM processing should occur. At pixel level, it computes a spatial energy map

$9\%$6

which is used to modulate local residual refinement (Li et al., 23 Jun 2026).

This design organizes computation along two axes. First, energy determines where to allocate long-range modeling: tiles with high energy are treated as information-dense and are passed through a Mamba-style block, while low-energy tiles are left unchanged. Second, energy determines how strongly local refinement should act: high-energy positions receive larger adaptive residual corrections, while low-energy background regions are altered minimally. The framework therefore treats spatial complexity and signal energy as coupled rather than independent factors.

A plausible implication is that the method operationalizes a form of content-adaptive computation that is specific to RAW reconstruction rather than generic vision modeling. The paper explicitly ties this choice to the long-tailed distribution of RAW feature energy: many locations are low-energy, while a small fraction are detail-rich or correspond to bright specular structures. Under that distribution, uniform context modeling is inefficient and may either underfit difficult regions or over-allocate capacity to smooth ones (Li et al., 23 Jun 2026).

4. TileMambaBlock and Energy-Aware Refinement

TileMambaBlock is the spatially selective SSM component. Given a context feature map

$9\%$7

the feature map is partitioned into non-overlapping $9\%$8 tiles, with

$9\%$9

Each tile is scored using the energy statistic $\mathbf{x}_{\mathrm{jpg}}$0, and the top

$\mathbf{x}_{\mathrm{jpg}}$1

tiles are selected according to the keep ratio $\mathbf{x}_{\mathrm{jpg}}$2. Selective processing is then defined as

$\mathbf{x}_{\mathrm{jpg}}$3

For small feature maps with $\mathbf{x}_{\mathrm{jpg}}$4, or when $\mathbf{x}_{\mathrm{jpg}}$5, the block falls back to dense processing (Li et al., 23 Jun 2026).

The Mamba block itself is VMamba’s Visual State Space block, using a 1D SSM recurrence

$\mathbf{x}_{\mathrm{jpg}}$6

with input-dependent parameters and cross-scan in four directions: left-to-right, right-to-left, top-to-bottom, and bottom-to-top. The 2D output is obtained by

$\mathbf{x}_{\mathrm{jpg}}$7

In the reported implementation, TileMambaBlock uses a VSS block with state expansion factor $\mathbf{x}_{\mathrm{jpg}}$8, tile size $\mathbf{x}_{\mathrm{jpg}}$9, and default keep ratio $\mathbf{x}_{\mathrm{raw}}$0. The stated complexity per selected tile is $\mathbf{x}_{\mathrm{raw}}$1, while total cost scales linearly with the number of selected pixels rather than the full image area (Li et al., 23 Jun 2026).

Energy-Aware Refinement (EAR) follows TileMambaBlock and is designed to better match the long-tail energy distribution of RAW features. Given $\mathbf{x}_{\mathrm{raw}}$2, EAR computes

$\mathbf{x}_{\mathrm{raw}}$3

then forms a channel-wise gate

$\mathbf{x}_{\mathrm{raw}}$4

and a local residual branch

$\mathbf{x}_{\mathrm{raw}}$5

The final refined feature is

$\mathbf{x}_{\mathrm{raw}}$6

A crucial implementation detail is that the last $\mathbf{x}_{\mathrm{raw}}$7 convolution in the residual branch is zero-initialized, so EAR begins as an identity mapping, $\mathbf{x}_{\mathrm{raw}}$8, and only gradually learns adaptive corrections during training (Li et al., 23 Jun 2026).

EAR differs from global channel reweighting modules such as SENet in that it preserves spatial resolution. Its gating is driven by a per-pixel energy map rather than global average pooling, so it functions as a locally adaptive inductive bias rather than a global summary mechanism.

5. Training protocol, ablations, and empirical results

The reported experiments use the NUS dataset camera subsets Samsung NX2000, Olympus E-PL6, and Sony SLT-A57, following the CAM and Beyond-R2LCM protocol on $\mathbf{x}_{\mathrm{raw}}$9 downsampled RAW resolutions. For example, Samsung images are reduced from $\mathbf{s}$0 to $\mathbf{s}$1, Olympus from $\mathbf{s}$2 to $\mathbf{s}$3, and Sony from $\mathbf{s}$4 to $\mathbf{s}$5. AdobeFiveK is also used, with $\mathbf{s}$6 training images and $\mathbf{s}$7 test images, to assess generalization under varied lighting and tonal conditions. The backbone is exactly the two-level JPEG-conditioned VAE from Beyond-R2LCM, with channel width $\mathbf{s}$8, latent channel reduction factor $\mathbf{s}$9, two reconstruction levels, and four learned spatial sampling rounds. Training uses $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$0 patches, Adam with $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$1 and $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$2, an initial learning rate of $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$3 decayed cosinely to $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$4 over $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$5 epochs, mixed precision, batch size $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$6, and separate models for $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$7 (Li et al., 23 Jun 2026).

On the NUS subsets, MambaRaw is reported to outperform strong metadata-based baselines at similar or lower metadata bitrates. In the representative Table 1 comparison, Beyond-R2LCM uses metadata bpp $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$8 and achieves $$(\mathbf{x}_{\mathrm{raw}}, \mathbf{x}_{\mathrm{jpg}})$$9 dB on Samsung, $\mathbf{x}_{\mathrm{jpg}}$0 dB on Olympus, and $\mathbf{x}_{\mathrm{jpg}}$1 dB on Sony. MambaRaw uses metadata bpp $\mathbf{x}_{\mathrm{jpg}}$2 and achieves $\mathbf{x}_{\mathrm{jpg}}$3 dB, $\mathbf{x}_{\mathrm{jpg}}$4 dB, and $\mathbf{x}_{\mathrm{jpg}}$5 dB, corresponding to gains of $\mathbf{x}_{\mathrm{jpg}}$6 dB, $\mathbf{x}_{\mathrm{jpg}}$7 dB, and $\mathbf{x}_{\mathrm{jpg}}$8 dB. On AdobeFiveK at very low metadata rates, Beyond-R2LCM reports $\mathbf{x}_{\mathrm{jpg}}$9 bpp and $\mathbf{s}$0 dB, whereas MambaRaw reports $\mathbf{s}$1 bpp and $\mathbf{s}$2 dB; at higher metadata, Beyond-R2LCM reports $\mathbf{s}$3 bpp and $\mathbf{s}$4 dB, while MambaRaw reports $\mathbf{s}$5 bpp and $\mathbf{s}$6 dB.

The ablation results isolate the contributions of the two proposed modules. The Beyond-R2LCM context baseline achieves $\mathbf{s}$7 dB at $\mathbf{s}$8 ms. Adding EAR alone raises PSNR to $\mathbf{s}$9 dB with $\hat{\mathbf{x}}_{\mathrm{raw}}$0 ms runtime. Replacing the baseline context model with dense VMamba everywhere yields $\hat{\mathbf{x}}_{\mathrm{raw}}$1 dB at $\hat{\mathbf{x}}_{\mathrm{raw}}$2 ms. Full MambaRaw, combining SSM, tile selection, and EAR, attains $\hat{\mathbf{x}}_{\mathrm{raw}}$3 dB at $\hat{\mathbf{x}}_{\mathrm{raw}}$4 ms, preserving essentially all of the dense-SSM gain while being $\hat{\mathbf{x}}_{\mathrm{raw}}$5 faster than dense SSM and faster than the original baseline. Tile-selection ablations show that L2 energy scoring gives the best PSNR, $\hat{\mathbf{x}}_{\mathrm{raw}}$6 dB, at $\hat{\mathbf{x}}_{\mathrm{raw}}$7 ms, outperforming random selection and matching the fastest runtime. Foundational-block comparisons show $\hat{\mathbf{x}}_{\mathrm{raw}}$8 dB for a CNN ResBlock at $\hat{\mathbf{x}}_{\mathrm{raw}}$9 ms, $1.4$00 dB for a windowed-attention Transformer at $1.4$01 ms, and $1.4$02 dB for VMamba at $1.4$03 ms. The chosen operating point $1.4$04, $1.4$05 yields $1.4$06 dB at $1.4$07 ms; denser configurations slightly improve PSNR but with higher latency, while sparser ones reduce latency but lose accuracy.

The 4K results quantify the method’s intended deployment regime. At $1.4$08, Beyond-R2LCM reports $1.4$09 G FLOPs, $1.4$10 GB GPU memory, $1.4$11 ms runtime, and $1.4$12 dB PSNR. MambaRaw reports $1.4$13 G FLOPs, $1.4$14 GB memory, $1.4$15 ms runtime, and $1.4$16 dB PSNR. This corresponds to a $1.4$17 FLOP reduction, an about $1.4$18 latency reduction, and a $1.4$19 dB PSNR improvement. The overhead of tile selection itself is reported as small: L2 scoring costs $1.4$20 ms, top-$1.4$21 selection $1.4$22 ms, pad/reshape $1.4$23 ms, and total selection overhead $1.4$24 ms, or $1.4$25 of total runtime (Li et al., 23 Jun 2026).

6. Relation to adjacent Mamba-based RAW imaging and known limitations

Within Mamba-based RAW imaging, MambaRaw occupies a distinct position. It is not a low-light RAW-to-sRGB enhancement model, as in "Retinex-RAWMamba" (Chen et al., 2024), which bridges demosaicing and denoising using an 8-direction RAW-aware scanning mechanism and a Retinex Decomposition Module. It is also not an sRGB-to-RAW de-rendering framework, as in "RAWMamba" (Chen et al., 2024), which unifies image and video de-rendering through Unified Metadata Embedding and Local Tone-Aware Mamba. MambaRaw instead addresses JPEG-guided RAW reconstruction from a JPEG preview plus a learned metadata stream, and its distinctive technical choice is to insert the SSM specifically into the entropy parameter network of a JPEG-conditioned VAE codec rather than to use Mamba as the full image backbone.

This specialization matters conceptually. The paper explicitly contrasts MambaRaw with attention-based context models that treat all positions uniformly and scale quadratically, and with vision Mamba systems that use SSMs as general backbones for recognition or restoration. MambaRaw uses a selective SSM in a compression-oriented context model, couples it to feature-energy statistics, and targets entropy parameter estimation rather than generic representation learning. That design suggests an intermediate position between learned image compression and RAW-domain computational photography.

The reported limitations are correspondingly specific. MambaRaw is single-frame only and does not exploit temporal redundancy. Performance is reported separately for three camera models, and cross-sensor generalization is not extensively evaluated. The method assumes a reasonably good, aligned JPEG preview; extremely compressed or misaligned JPEGs could degrade guidance quality. It also assumes a fixed RAW color space and mosaicing pipeline, uses inputs normalized to $1.4$26 in raw-linear space, and does not explicitly model sensor noise such as Poisson–Gaussian noise beyond what is captured implicitly by the learned model. Future directions proposed in the paper include RAW video extensions with temporal SSMs, hardware-aware optimization for mobile or embedded deployment, and direct integration of additional metadata such as exposure and ISO into the context model (Li et al., 23 Jun 2026).