Hardware-Friendly Token-Group Activation Quantization for Low-Bit Mamba Super-Resolution

AI summary

Overview

• Research area: Post-training quantization (PTQ) for efficient image restoration, specifically low-bit quantization of Mamba-based (selective state-space) super-resolution models for edge deployment.
• Technical level: Advanced. The paper assumes familiarity with quantization clipping bounds, uniform affine quantization, selective scan / state-space models, calibration-set PTQ protocols, and knowledge distillation.
• Scope: The paper diagnoses token-wise activation heterogeneity caused by selective scan and proposes a static, hardware-friendly token-group activation quantization framework (plus a frequency-domain refinement objective) for low-bit Mamba super-resolution, validated across MambaIRv1/v2 backbones, PTQ protocols, bit-widths, and transfer tasks.

What This Paper Is About

Mamba-based super-resolution models replace quadratic self-attention with selective state-space modeling, but quantizing their activations to low bit-widths is hard: the selective scan operator produces a different activation range for each spatial token, so a single per-layer clipping range either wastes precision on typical tokens or clips the high-response tokens that restoration quality depends on. The paper's goal is to give different token groups their own calibrated clipping ranges while keeping inference static — no runtime statistics collection, no dynamic clipping updates, and no post-scan group assignment — so that the resulting low-bit path can actually run fast on integer hardware.

Key Contributions

1. Identifies selective-scan-induced token-wise activation heterogeneity as the key failure mode of low-bit PTQ for Mamba SR. Across 20 ASSM blocks of MambaIRv2-Light, 18 out of 20 blocks show higher per-token range variation after selective scan than before, with an average 1.80× amplification.
2. Proposes Token-Group Quantization, which frames token-adaptive low-bit quantization as a static approximation of a hardware-unfriendly "downstream oracle." During calibration the oracle is instantiated by K-means over post-scan activation descriptors, and a lightweight pre-scan predictor is trained to recover oracle-aligned groups before the quantized sites.
3. Introduces Frequency-Preserving Refinement (FPR), a teacher-supervised objective that refines clipping bounds by aligning Fourier amplitude and phase between full-precision teacher and quantized student outputs, targeting high-frequency edges and textures that pixel and feature losses under-supervise.
4. Demonstrates broad empirical gains, reporting consistent improvements over prior Mamba and SR PTQ baselines with 2.23× speedup on Jetson Orin Nano over FP16, gains that hold across MambaIR v1 and v2 at Light and Base scales, across denoising and JPEG artifact removal, transfer to SwinIR-light and ImageNet classification on Vision Mamba, and composition with channel-axis transforms such as SpinQuant.

Main Findings

• Activation heterogeneity is broad, not just outliers: Post-scan activations spread over a wider range than pre-scan activations after median-|x| normalization, and per-token range CoV increases in 18 of 20 ASSM layers with ~1.80× average amplification. The paper argues the bottleneck is token populations with distinct ranges, not a few rare outliers.
• Grouping must be quantization-aware, not just variance-based: Random grouping produces similar per-token preferred clipping-scale distributions across groups, whereas K-means over post-scan descriptors separates tokens into distinct clipping-scale regimes.
• Weight-freeze PTQ results (MambaIRv2-Light): At ×4, Ours (4/4) reaches 31.46/0.8811 on Set5, 28.14/0.7683 on Set14, 27.24/0.7231 on BSD100, 25.39/0.7600 on Urban100, and 29.21/0.8884 on Manga109, beating PTQ4VM, Quamba, MambaQuant, and OAQ-SR at the same 4/4 scope. At 2/2 it reaches 29.12/0.8260, 26.57/0.7259, 26.32/0.6891, 23.57/0.6805, and 25.58/0.8031.
• Weight-freeze PTQ results at ×2: Ours (4/4) gives 37.56/0.9578 (Set5), 33.24/0.9149 (Set14), 31.86/0.8950 (BSD100), 31.34/0.9204 (Urban100), 37.85/0.9742 (Manga109); at 2/2, 35.02/0.9436, 31.41/0.8966, 30.51/0.8766, 27.90/0.8708, 32.67/0.9508.
• LoRA-style weight-training protocol: Ours + LoRA is best in all reported settings and nearly recovers full precision at 4-bit; on BSD100 the PSNR drop is only 0.12 dB for ×2 and 0.14 dB for ×4.
• Token grouping is the dominant gain: Token grouping alone accounts for 77% of the improvement over layer-wise PTQ on Set14 and 65% on Urban100, and it helps even without bound refinement. Amplitude and phase terms each help alone, and their combination adds a further 0.35 dB on Urban100.
• Grouping-rule ablation (×2, W4A4): Layer-wise PTQ 32.45/0.9056 (Set14) and 30.34/0.9107 (Urban100); random groups 31.76/0.8904 and 29.91/0.9070 (worse than layer-wise); pre-scan K-means 32.84/0.9092 and 30.77/0.9161; hardware-unfriendly post-scan oracle 33.41/0.9176 and 31.57/0.9253; predicted (Ours) 33.24/0.9149 and 31.34/0.9204, recovering 82.3% and 81.3% of the oracle gain.
• Number of groups: Performance improves with K; gains nearly saturate from K=6 to K=8 (K=8: 33.24/0.9149 on Set14, 31.34/0.9204 on Urban100). On MambaIRv2-Base the optimum shifts mildly to K=16, but K=8 already recovers 89% of the range (Appendix F.3).
• Latency: A6000 — FP16 180.7 ms (1.00×), 4-bit Oracle 254.5 ms (0.71×), 4-bit Base 102.7 ms (1.76×), 4-bit Ours 106.2 ms (1.70×). Jetson Orin Nano — FP16 3149.6 ms (1.00×), 4-bit Oracle 3213.9 ms (0.98×), 4-bit Base 1380.1 ms (2.28×), 4-bit Ours 1410.7 ms (2.23×). The runtime oracle loses the speedup entirely.
• Calibration cost: Weight-freeze — OAQ-SR 1.29 h vs. Ours 1.42 h (3,000 iterations; 1.9 k vs. 4.1 k params). LoRA — SPR²Q 6.12 h vs. Ours + LoRA 6.57 h (12,500 iterations; 617.2 k vs. 619.4 k params).
• Low-bit coverage: The quantized path covers 90.8% of FLOPs, where Linear + Conv is 8.49 G (85.9%), scan operands 0.49 G (4.9%), other conv 0.90 G (9.1%, not quantized), total 9.88 G.
• Transfer and composition: Ours is reported best in every generalization setting; ported Hadamard, SpinQuant, and Affine transforms trail by up to 2.12 dB, and SpinQuant combined with Ours is best everywhere. Robustness holds for input resolution up to 16× the calibration length (Appendix F.2).
• Stability and setup: Results use seed 0; five seeds give std below 0.03 dB at W4A4 and 0.05 dB at W2A2.

Methodology in Plain English

The method starts from a measurement: the selective scan operator mixes information across flattened spatial tokens using token-dependent parameters, so after the scan different tokens want different clipping ranges. The ideal fix — look at post-scan activations at each quantization site and cluster tokens by the clipping range they prefer — requires runtime statistics inside the quantized path, which breaks static low-bit deployment. The paper therefore treats that ideal as an upper bound ("downstream oracle") and tries to reproduce its grouping decisions without ever looking downstream at inference time.

Three pieces make this work. First, a group-conditional quantizer: inside each block, eight sites (three downstream linear sites and five scan-internal per-token operands) use group-specific clipping bounds, while remaining sites keep ordinary layer-wise PTQ. Second, quantization-aware pseudo-labels: for each token and each of those sites, the authors build a four-number descriptor — 0.99 quantile of |x|, 0.999 quantile of |x|, RMS, and max |x| — concatenate descriptors over the eight sites, pool all token–layer samples from the calibration set, and run one K-means to get group labels. Third, a shared two-layer FFN predictor that reads only pre-scan features (available before the scan is quantized) and outputs a group index via argmax; it is trained with cross-entropy against those pseudo-labels. Because the FFN is shared but its input is layer-dependent, the assignment stays layer-adaptive — described as structure–magnitude decoupling, where the shared network captures recurring token-structure cues and the layer/site/group-specific bound tables absorb magnitude drift.

After the predictor is frozen, the calibration set is replayed, activations are pooled per (layer, site, group), and clipping bounds are initialized by minimizing empirical quantization error. Bounds are then refined — with backbone weights, pseudo-labels, predictor parameters, and bit-widths all fixed — using a full-precision teacher. The refinement loss combines an L1 pixel loss, a normalized-feature MSE term, and the new FPR term, which applies a 2D Hann window, takes 2D Fourier transforms of teacher and student outputs, and matches amplitude (L1) and wrapped phase (L1) with equal 1/2 weights. Defaults are λ_feat = 5, λ_FPR = 0.01, K = 8 groups, 3,000 refinement iterations, batch size 8, calibration on 96 low-resolution DF2K patches, and evaluation with PSNR/SSIM on the Y channel of YCbCr on Set5, Set14, BSD100, Urban100, and Manga109 with boundaries cropped by the upscaling factor.

Why This Matters

• Impact on research: It reframes low-bit PTQ for state-space models around token identity rather than only channel-wise outliers, showing that channel-axis rotation/affine transforms and the token axis are complementary rather than competing. It also contributes a data-type-specific analysis (selective scan → 18 of 20 layers with increased token-range variation, ~1.80× amplification) that other SSM quantization work can build on.
• Real-world applications:
  • On-device photo and video upscaling on phones or edge cameras, where activation memory traffic is the bottleneck.
  • Real-time restoration pipelines for surveillance, drones, or automotive imaging that run on Jetson-class accelerators.
  • JPEG artifact removal and denoising in consumer photo tools, both of which the paper evaluates directly.
  • Compressed low-bit restoration inside larger vision systems, given demonstrated transfer to SwinIR-light and Vision Mamba classification.
• Industry relevance: The reported 2.23× speedup on Jetson Orin Nano over FP16, with byte-aligned integer kernels and fixed quantization codes, plus a calibration overhead of only 1.42 h versus 1.29 h for a comparable baseline, is the kind of profile that matters for shipping quantized models on embedded hardware. The authors position the framework for future deployment on NPUs, DSPs, mobile GPUs, and vendor-specific low-bit runtimes.

Future Directions

• Reduce the inference overhead: The stated limitation is that pre-scan prediction and group-wise bound lookup add a small overhead relative to the base quantized path (106.2 ms vs. 102.7 ms on A6000; 1410.7 ms vs. 1380.1 ms on Jetson Orin Nano).
• Port to other accelerators: The authors name NPUs, DSPs, mobile GPUs, and vendor-specific low-bit runtimes as targets for the same token-adaptive framework.
• Close the gap to the oracle: The predicted groups recover 82.3% (Set14) and 81.3% (Urban100) of the oracle gain, leaving a measurable accuracy gap that is stated as an upper bound rather than achieved.
• Tune group count per scale: The optimum shifts from K=8 on the Light backbone to K=16 on MambaIRv2-Base, suggesting group count may need to be chosen per backbone scale rather than fixed.

Target Audience

Researchers and engineers working on model compression and efficient inference — particularly those deploying super-resolution or other restoration models on embedded and edge hardware. It will also interest readers working on quantization of state-space models and Mamba architectures generally, and practitioners who need to decide whether token-wise adaptivity is worth the deployment complexity over layer-wise or channel-wise methods. Readers without a background in quantization or SSM operators will find the paper demanding, since much of the argument depends on clipping-bound formalisms and calibration-protocol distinctions (Weight-freeze PTQ versus Weight-training PTQ with LoRA).

Authors’ abstract

Mamba-based super-resolution models are efficient alternatives to Transformer-based restoration backbones, but their low-bit post-training quantization remains challenging because selective scan, Mamba's input-dependent state-space operator, induces different activation ranges across spatial tokens. Thus, a single layer-wise clipping range poorly matches all tokens. We propose Token-Group Quantization, a hardware-friendly activation PTQ framework that statically approximates a hardware-unfriendly post-scan grouping oracle, where tokens with similar preferred clipping ranges are assigned to the same group. To avoid runtime post-scan grouping in static low-bit inference, a lightweight predictor infers the groups from pre-scan features available before the quantized scan path. During calibration, we form pseudo-labels by clustering post-scan activation descriptors that capture similar rounding--clipping behavior. At inference, the frozen predictor maps pre-scan features to groups and selects entries from offline-calibrated group-wise clipping-bound tables, without runtime statistics collection, dynamic clipping updates, or post-scan group assignment. To address low-bit restoration degradation, we further introduce a Frequency-Preserving Refinement (FPR) objective for clipping-bound calibration. FPR aligns Fourier amplitude and phase between full-precision teacher and quantized student outputs to preserve high-frequency structures. Across PTQ protocols, bit-widths, backbones, and tasks, our method consistently outperforms prior Mamba and SR PTQ baselines, transfers to Transformer SR and Vision Mamba classification, composes with rotation-based quantization, and achieves 2.23x speedup on Jetson Orin Nano over FP16. Code is available here.