TileMambaBlock in MambaRaw Framework

• TileMambaBlock is a lightweight module in MambaRaw that provides sparse, content-adaptive long-range spatial modeling by selectively processing information-dense tiles.
• It partitions feature maps into non-overlapping tiles and scores each via L2 energy to apply state space models only to the top-k high-energy tiles.
• Combined with Energy-Aware Refinement, it balances computation and fidelity, achieving significant speedup and improved JPEG-guided raw reconstruction performance.

Searching arXiv for the cited paper and topic-specific details. TileMambaBlock is a lightweight module in MambaRaw, a JPEG-conditioned metadata-based raw image reconstruction framework that uses State Space Models (SSMs) to estimate entropy parameters efficiently (Li et al., 23 Jun 2026). Within the paper’s Spatial-Energy Coupled Context Modeling mechanism, TileMambaBlock performs Mamba-style selective scanning only on information-dense tiles, while Energy-Aware Refinement (EAR) subsequently adjusts the resulting features to match the long-tail energy distribution of raw signals. In the Level-1 entropy-parameter network, TileMambaBlock sits between an input projection and EAR, and its function is sparse, content-adaptive long-range spatial modeling.

1. Architectural placement and functional role

In MambaRaw’s Level-1 entropy-parameter network, the usual convolutional context block is replaced by a three-stage cascade: $$\mathbf{F}_{\mathrm{in}} = \psi_{\mathrm{ep}}^{\mathrm{in}}(\tilde{\mathbf{F}}),\qquad \mathbf{F}_{c} = \mathrm{TileMambaBlock}(\mathbf{F}_{\mathrm{in}}; T,\rho),\qquad \mathbf{F}' = \mathrm{EAR}(\mathbf{F}_{c}).$$ Here, $\tilde{\mathbf{F}}$ is the JPEG-conditioned feature, $\mathbf{F}_{\mathrm{in}}$ is the projected input to TileMambaBlock, $\mathbf{F}_{c}$ is the context feature after selective SSM processing, and $\mathbf{F}'$ is the refined output after EAR (Li et al., 23 Jun 2026).

The module’s specific contribution is not generic spatial filtering. It is responsible for sparse, content-adaptive long-range spatial modeling. Internally, it chops the feature map into $T \times T$ tiles, scores each tile by its $L_2$ energy, selects only the top-$k$ tiles with $k=\lfloor \rho \cdot N_t \rfloor$ for full SSM processing, and leaves the remainder untouched. This places TileMambaBlock at the point where MambaRaw substitutes selective state-space computation for dense context modeling.

In the broader system description, this design addresses the practical limitation that attention mechanisms scale quadratically with feature maps, which becomes expensive at high resolution, such as 4K raw image. TileMambaBlock is therefore part of the framework’s efficiency-oriented replacement of conventional high-cost context modeling in JPEG-guided raw reconstruction (Li et al., 23 Jun 2026).

2. Tile partitioning and selective scanning

Let the input to TileMambaBlock be

$$\mathbf{F}_{\mathrm{in}} \in \mathbb{R}^{C \times H \times W}.$$

The feature map is partitioned into

$\tilde{\mathbf{F}}$0

non-overlapping tiles

$\tilde{\mathbf{F}}$1

Each tile is assigned an $\tilde{\mathbf{F}}$2 energy score: $\tilde{\mathbf{F}}$3 The selected tile index set is then

$\tilde{\mathbf{F}}$4

Only these selected tiles undergo SSM-based processing: $\tilde{\mathbf{F}}$5 The processed and unprocessed tiles are then reassembled into $\tilde{\mathbf{F}}$6 (Li et al., 23 Jun 2026).

This mechanism makes the selectivity explicit: TileMambaBlock is not a uniformly applied tilewise Mamba layer. A common misreading is that tiling merely localizes the SSM computation; in the described formulation, the decisive step is the top-$\tilde{\mathbf{F}}$7 selection by energy, and the non-selected tiles are passed through unchanged. The data also gives a binary-mask formulation,

$\tilde{\mathbf{F}}$8

so that

$\tilde{\mathbf{F}}$9

This expresses the module as a hard routing mechanism over tiles rather than a soft attention map.

The paper’s summary characterizes this as a content-adaptive compromise between computation and fidelity by scoring tiles via simple $\mathbf{F}_{\mathrm{in}}$0 energy, running the linear-time SSM only where it is most needed, and then relying on EAR to restore fine-grained, energy-guided detail. A plausible implication is that the tile score serves as a computational proxy for spatial informativeness within the entropy-parameter network.

3. State space formulation and two-dimensional scanning

TileMambaBlock applies a MambaBlock to each selected tile. That MambaBlock is built upon the discrete SSM

$\mathbf{F}_{\mathrm{in}}$1

where $\mathbf{F}_{\mathrm{in}}$2 is the $\mathbf{F}_{\mathrm{in}}$3-th input token, $\mathbf{F}_{\mathrm{in}}$4 is the hidden state, and $\mathbf{F}_{\mathrm{in}}$5, $\mathbf{F}_{\mathrm{in}}$6, and $\mathbf{F}_{\mathrm{in}}$7 are learned, in practice input-dependent, matrices. No explicit $\mathbf{F}_{\mathrm{in}}$8 term is used here, or $\mathbf{F}_{\mathrm{in}}$9 (Li et al., 23 Jun 2026).

Because this SSM is inherently one-dimensional, the block uses a cross-scan strategy to process two-dimensional tiles: $\mathbf{F}_{c}$0 The description specifies four directional scans: left$\mathbf{F}_{c}$1right, right$\mathbf{F}_{c}$2left, top$\mathbf{F}_{c}$3bottom, and bottom$\mathbf{F}_{c}$4top, followed by a merge. This is the mechanism by which a 1D sequence model is adapted to $\mathbf{F}_{c}$5 tile tensors.

An important implementation property is that the same $\mathbf{F}_{c}$6, $\mathbf{F}_{c}$7, and $\mathbf{F}_{c}$8 parameters are shared across all tiles and across all spatial positions within a tile, so the selective mechanism does not blow up the parameter count. That statement delimits the role of selectivity: it changes where the SSM is executed, but not the parameterization of the SSM itself. This suggests that TileMambaBlock reduces compute and peak SSM memory through sparse execution rather than through per-tile specialization.

4. Computational characteristics

The complexity analysis in the description distinguishes dense SSM processing on the full feature map from selective SSM processing on only a fraction of tiles. Let $\mathbf{F}_{c}$9 denote the total spatial size, $\mathbf{F}'$0 the number of channels, $\mathbf{F}'$1 the tile size, $\mathbf{F}'$2 the number of tiles, $\mathbf{F}'$3 the keep ratio, and $\mathbf{F}'$4 the per-pixel cost constant of one cross-scan SSM.

For a full-feature-map dense SSM, the work is approximately

$\mathbf{F}'$5

The memory must hold at least $\mathbf{F}'$6 for the feature map and SSM intermediate states of comparable size (Li et al., 23 Jun 2026).

For TileMambaBlock selective SSM, the analysis is decomposed into two parts. Energy scoring costs

$\mathbf{F}'$7

and SSM processing on $\mathbf{F}'$8 tiles costs approximately

$\mathbf{F}'$9

The total work is therefore approximately

$T \times T$0

If $T \times T$1, this is roughly a $T \times T$2 speedup over dense SSM, neglecting constant overhead. For memory, the input feature still requires $T \times T$3, but the per-tile SSM scratch is only $T \times T$4 at any one time, so the peak SSM memory is $T \times T$5 rather than $T \times T$6.

These expressions formalize the module’s design target. TileMambaBlock does not eliminate the cost of reading or scoring the full feature map, but it restricts the expensive cross-scan SSM to a selected subset. In the context of MambaRaw, whose abstract emphasizes efficiency for high-resolution JPEG-guided raw reconstruction and reports that the full framework reduces end-to-end coding latency by about $T \times T$7, TileMambaBlock is one of the two lightweight modules supporting that efficiency-oriented context model (Li et al., 23 Jun 2026).

5. Coupling with Energy-Aware Refinement

After TileMambaBlock produces

$T \times T$8

EAR computes a spatial energy map,

$T \times T$9

then a gating tensor,

$L_2$0

Its residual branch is

$L_2$1

and the final refinement is

$L_2$2

The description states that, in this way, EAR “polishes” the sparse SSM outputs by re-introducing local, energy-guided variations (Li et al., 23 Jun 2026).

The interaction between the two modules is structurally specific. TileMambaBlock provides sparse, selected long-range modeling; EAR then modulates a residual correction using an energy-derived gate. The abstract defines EAR as an identity-initialized residual module that enhance feature representation to match the long-tail energy distribution of raw signals. Read together, these details indicate a division of labor inside the Spatial-Energy Coupled Context Modeling mechanism: TileMambaBlock allocates SSM computation to information-dense tiles, while EAR repairs or enriches the resulting feature field using local energy-aware modulation.

A plausible implication is that the pair is designed to mitigate the principal trade-off introduced by selectivity. Because unselected tiles bypass the SSM unchanged, subsequent refinement is needed to recover local variations that may still matter for entropy-parameter estimation.

6. Forward-pass behavior and operational regimes

The step-by-step forward pass is defined for input

$L_2$3

tile size $L_2$4, and keep ratio $L_2$5, with output $L_2$6. First,

$L_2$7

Second, there is a dense fallback condition: if $L_2$8 and $L_2$9, or $k$0, then

$k$1

Otherwise, the block pads and reshapes $k$2 into tiles, computes $k$3, sets $k$4, selects top-$k$5 tile indices $k$6, applies MambaBlock only on selected tiles, and reassembles and crops padding to obtain $k$7. Third, EAR is applied: $k$8

$k$9

The returned output is $k=\lfloor \rho \cdot N_t \rfloor$0 (Li et al., 23 Jun 2026).

This operational description makes clear that TileMambaBlock is not defined solely by its sparse case. The dense fallback preserves a conventional full MambaBlock path when the spatial extent is already tile-bounded or when the keep ratio implies no sparsification. Thus, the module spans two regimes: a degenerate dense regime and a sparse tile-selection regime. That detail is important for implementation fidelity, because it determines when the tile machinery is bypassed entirely.

Within MambaRaw, this forward path serves the entropy-parameter estimation pipeline used for JPEG-guided raw reconstruction. The paper reports extensive experiments on three camera datasets—Sony, Olympus, and Samsung—and states that MambaRaw shows consistent improvements over strong metadata-based baselines, setting a new state of the art for JPEG-guided raw reconstruction with great efficiency. Notably, at low metadata bitrates, the framework increases PSNR by $k=\lfloor \rho \cdot N_t \rfloor$1--$k=\lfloor \rho \cdot N_t \rfloor$2 dB and reduces end-to-end coding latency by about $k=\lfloor \rho \cdot N_t \rfloor$3 (Li et al., 23 Jun 2026). The article-level evidence is reported for the complete framework rather than for TileMambaBlock in isolation; however, TileMambaBlock is one of the two lightweight modules identified as the key contribution of the Spatial-Energy Coupled Context Modeling mechanism.