Hardware-Friendly Token-Group Activation Quantization for Low-Bit Mamba Super-Resolution¶
会议: NeurIPS 2026
arXiv: 2610.16261
OpenReview: https://openreview.net/forum?id=46hNfjRw3m
代码: https://github.com/imaboybut/TGQ-MambaSR
领域: 图像恢复
关键词: 图像超分辨率, Mamba模型, 训练后量化, Token分组量化, 频域蒸馏
一句话总结¶
针对选择性扫描机制导致的空域 Token 激活范围剧烈分化难题，本文提出硬件友好的静态 Token 分组量化（TGQ）与频域保持微调（FPR）目标，在无需运行时动态统计开销的前提下实现了极低比特（W4A4/W2A2）Mamba 超分辨率的高保真重建，在 Jetson Orin Nano 上取得 2.23× 实际加速。
研究背景与动机¶
基于选择性状态空间模型（Mamba）的图像超分辨率（SR）模型（如 MambaIR、MambaIRv2）因其对长程空间依赖具有线性复杂度，已成为替代自注意力 Transformer 的高效图像恢复骨干网络。然而，在边缘设备上落地高分辨率重建依然面临激活内存读写带宽瓶颈，迫切需要低比特后训练量化（PTQ）将权重与激活同时压缩到低比特整型表示。现有的卷积与 Transformer PTQ 算法大多依赖逐层（Layer-wise）统一截断阈值，若直接移植至低比特 Mamba 模型，会导致恢复图像的纹理细节遭受毁灭性劣化。
这种量化退化的根源在于 Mamba 的核心算子——输入依赖的选择性扫描（Selective Scan）。经验分析揭示，在 MambaIRv2-Light 的 20 个 ASSM 模块中，有 18 个模块在经过选择性扫描后，其逐 Token 激活范围的离散系数（CoV）显著增大，平均增幅达 1.80 倍（在视觉分类模型 Vim 上甚至高达 3~5 倍）。不同空间位置的 Token 呈现出截然不同的动态幅度范围，这与逐层统一切断阈值的假设产生直接冲突：若采用宽松的截断范围包容高响应区域，绝大多数平坦区域的典型 Token 将严重损失有效量化精度（步长过大导致舍入误差剧增）；若采用收紧的截断范围逼近典型 Token，重建高频纹理关键的高响应 Token 则会被暴力截断（引发严重饱和溢出误差）。
直接根据扫描后的后验激活（Post-scan）在线统计并动态分组虽能提供精度上限，但在静态低比特加速引擎中属于硬件不友好方案，不仅需要在线插入统计算子破坏内核编译，还会因重复计算带来巨大延迟惩罚。本文的核心切入角度在于：选择性扫描的递推转移参数本就由扫描前的特征确定，这意味着扫描前特征已经蕴含了下游激活分化规律的先验线索。核心 idea：将硬件不友好的后验分组 Oracle 蒸馏为一个轻量级静态前验预测器，在推理期仅凭扫描前特征查表索引离线校准的组级截断阈值，并引入频域保持微调（FPR）联合对齐高频幅值与相位，兼顾静态低比特内核执行效率与高保真图像超分精度。
方法详解¶
整体框架¶
Token-Group Quantization（TGQ）将量化自适应性解耦为静态的结构预测与查表机制。在离线校准阶段，首先提取下游量化敏感位点的激活描述子进行 K-means 聚类以生成表征舍入-截断平衡的伪标签，并训练一个共享的轻量级前验两层前馈网络（Pre-scan Predictor）学会从扫描前特征预测分组索引；随后冻结预测器，离线最小化量化均方误差以初始化各层各组的截断边界表；最后，在全精度教师模型监督下，结合像素损失、特征相似度损失与频域保持微调损失（FPR）仅对截断阈值参数执行 3000 次微调。
在实际推理阶段，模型严格遵循静态部署约束（C1：无运行时动态统计；C2：在量化算子前完成分组确定）。输入特征经由冻结的轻量级预测器直接输出离散组索引 \(\hat{g}_i \in \{1,\dots,K\}\)，量化内核直接查表获取预先校准好的上下限截断值 \((l_{\ell,m,\hat{g}_i}, u_{\ell,m,\hat{g}_i})\) 并执行单次融合整型量化与解量化，不破坏底层 CUDA Graph 捕获与硬件编译执行流水线。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
In["输入浅层特征图 x~"] --> Pred["前验分组预测器<br/>轻量共享两层FFN"]
Pred -->|"离线索引查表 (静态分配)"| GroupLookup["分组自适应量化核<br/>按组索引固定截断阈值表"]
In --> Scan["Mamba 选择性扫描<br/>五大扫描内部运算数"]
Scan -->|"分组激活点 (8个敏感位点)"| GroupLookup
GroupLookup --> PostOps["后验投影与卷积前向<br/>out_proj / convffn2"]
PostOps --> Out["量化恢复输出 y^"]
Teacher["全精度教师输出 y"] -.->|"幅值与相位分解对齐"| FPR["频域保持微调 (FPR)<br/>校准优化截断边界参数"]
FPR -.->|"离线优化更新"| GroupLookup关键设计¶
1. 基于下游激活描述子的量化感知伪标签构建：精准刻画舍入与截断权衡
分组量化的目标并非语义聚类，而是将处于相似量化误差状态的 Token 划分至同一组。针对均匀量化器，对称截断 \([-c, c]\) 下的逐 Token 量化误差可分解为有效范围内的舍入误差与超出范围的尾部截断误差： $\(E_i(c) \approx \frac{\Delta(c)^2}{12}\Pr(|X_i|\le c) + \mathbb{E}\left[(|X_i|-c)_{+}^2\right]\)$ 其中量化步长 \(\Delta(c)=2c/(2^b-1)\) 随截断范围扩大而增大，导致第一项舍入误差上升；第二项截断误差则随截断范围缩小而加剧。为此，对每个 Token \(i\) 在当前块的 8 个敏感量化点 \(m\in\mathcal{M}\)（包括下游 linear 投影以及扫描内部自适应张量 \(x_s, dt_s, \delta, B_s, C_s\)），构建能够捕获能量主体与长尾分布的 4 维紧凑描述子： $\(\phi_m(x_i^{(\ell,m)}) = \big[q_{0.99}(|x_i^{(\ell,m)}|),\; q_{0.999}(|x_i^{(\ell,m)}|),\; \mathrm{RMS}(x_i^{(\ell,m)}),\; \max|x_i^{(\ell,m)}|\big]\)$ 将 8 个位点的描述子拼接成 32 维特征后，在离线校准集的所有 Token 样本上执行单次 K-means 聚类，产生能够区分不同截断尺度体系的离线伪标签 \(g_i^{\star(\ell)}\)。这种设计确保了处于相同组内的 Token 在下游具有相近的最优截断偏好，从而使离线固定边界能够最大化保留表达动态。
2. 结构与幅度解耦的前验分组预测器：实现零运行时开销的硬件友好推断
为满足静态编译要求，避免在线双重前向扫描或动态统计计算，设计了一个全局共享的极简两层 FFN（输入通道 \(C \to 64 \to K\)，采用 GELU 激活），参数量极小（仅百余参数）。该预测器以进入 ASSM 模块前的原始输入特征 \(\tilde{x}_i^{(\ell)}\) 为输入，在离线校准期通过交叉熵损失拟合 K-means 伪标签： $\(\hat{g}_i^{(\ell)} = \arg\max_{k\in\{1,\ldots,K\}} \big[f_\theta(\tilde{x}_i^{(\ell)})\big]_k, \qquad \mathcal{L}_{\mathrm{pred}} = \mathbb{E}_{(i,\ell)}\big[\mathrm{CE}(f_\theta(\tilde{x}_i^{(\ell)}), g_i^{\star(\ell)})\big]\)$ 这种设计体现了核心的“结构-幅度解耦”思想：不同网络层虽受选择性扫描影响产生不同绝对数值的激活漂移，但局部纹理与平坦区域的空间相对结构特征高度共通。由共享的轻量 FFN 捕获通用的 Token 结构分配，而各层各算子具体的动态幅度差异则交由离线标定的特定层-算子-分组截断参数表 \((l_{\ell,m,g}, u_{\ell,m,g})\) 充分吸收。在推理阶段，预测器执行前向推断输出整数索引后，CUDA 核直接融合查表寻址，完全消除了运行时后验统计与设备同步开销。
3. 频域保持微调（FPR）目标：精准纠正低比特超分高频频偏
在极低比特（如 4-bit、2-bit）量化下，空域的 \(L_1\) 像素重建损失和中间层特征余弦相似度损失往往被占比极大的低频平滑能量主导，导致学生网络即便降低了空域 MSE，依旧丢失关键的边缘锐度并出现高频伪影。为此，本文提出面向 PTQ 截断参数微调的教师监督频域保持微调（FPR）目标。对加 2D Hann 窗加权后的全精度教师输出与量化学生输出进行二维离散傅里叶变换，并将频谱严格分解为幅值谱与相位谱： $\(\mathcal{L}_{\mathrm{FPR}} = \frac{1}{2}\big\|\,|\mathbf{Y}|-|\hat{\mathbf{Y}}|\,\big\|_1 + \frac{1}{2}\big\|\,\mathrm{wrap}(\angle\mathbf{Y}-\angle\hat{\mathbf{Y}})\,\big\|_1\)$ 其中幅值匹配强制低比特模型对齐高频能量分布、抑制纹理平滑衰减，而相位对齐 \(\mathrm{wrap}(\cdot)\in(-\pi,\pi]\) 则强力约束边缘与结构位移。联合损失函数表达为： $\(\mathcal{L} = \|\hat{\mathbf{y}}-\mathbf{y}\|_1 + \lambda_{\mathrm{feat}}\frac{1}{|\mathcal{H}|}\sum_{h\in\mathcal{H}}\left\|\frac{\hat{F}_h}{\|\hat{F}_h\|_2}-\frac{F_h}{\|F_h\|_2}\right\|_2^2 + \lambda_{\mathrm{FPR}}\mathcal{L}_{\mathrm{FPR}}\)$ 整个微调过程保持骨干权重与预测器权重完全冻结，仅微调各组截断边界标量，以 3000 次轻量迭代彻底解决低比特量化引起的频域谱失真。
损失函数 / 训练策略¶
离线校准使用 DF2K 数据集采样的 96 张 \(64\times 64\) 低分辨率图像补丁，骨干权重固定。第一阶段利用 Adam 优化器以学习率 \(10^{-3}\) 训练 20 轮轻量前验预测器（批大小 1024）；第二阶段根据预测组别初始化截断参数后，设置 \(\lambda_{\mathrm{feat}}=5\)、\(\lambda_{\mathrm{FPR}}=0.01\)，对截断边界参数迭代微调 3000 次（批大小 8）。默认 Token 分组数取 \(K=8\)。
实验关键数据¶
主实验¶
在 MambaIRv2-Light 骨干网络上，针对传统 Weight-freeze PTQ 协议与联合 LoRA 蒸馏协议，在 5 个标准超分基准上全面对比。
表 1：Weight-freeze PTQ 协议下 MambaIRv2-Light 性能对比（PSNR / SSIM）
| 缩放 | 方法 | W / A | Set5 | Set14 | BSD100 | Urban100 | Manga109 |
|---|---|---|---|---|---|---|---|
| \(\times 4\) | 全精度 FP32 | 32 / 32 | 32.51 / 0.8992 | 28.84 / 0.7878 | 27.75 / 0.7426 | 26.82 / 0.8079 | 31.24 / 0.9182 |
| \(\times 4\) | PTQ4VM (WACV'25) | 4 / 4 | 30.82 / 0.8670 | 27.69 / 0.7546 | 26.95 / 0.7115 | 24.76 / 0.7321 | 28.19 / 0.8660 |
| \(\times 4\) | Quamba (ICLR'25) | 4 / 4 | 31.01 / 0.8715 | 27.77 / 0.7585 | 26.99 / 0.7149 | 25.01 / 0.7470 | 28.57 / 0.8752 |
| \(\times 4\) | MambaQuant (ICLR'25) | 4 / 4 | 30.74 / 0.8650 | 27.17 / 0.7413 | 26.37 / 0.6920 | 23.28 / 0.6694 | 26.73 / 0.8186 |
| \(\times 4\) | OAQ-SR (ICCV'25) | 4 / 4 | 30.40 / 0.8577 | 27.41 / 0.7470 | 26.75 / 0.7045 | 24.31 / 0.7136 | 27.39 / 0.8485 |
| \(\times 4\) | Ours (TGQ) | 4 / 4 | 31.46 / 0.8811 | 28.14 / 0.7683 | 27.24 / 0.7231 | 25.39 / 0.7600 | 29.21 / 0.8884 |
| \(\times 4\) | PTQ4VM | 2 / 2 | 28.77 / 0.8162 | 26.36 / 0.7167 | 26.16 / 0.6802 | 23.37 / 0.6704 | 25.26 / 0.7943 |
| \(\times 4\) | Quamba | 2 / 2 | 28.88 / 0.8080 | 26.45 / 0.7131 | 26.20 / 0.6752 | 23.48 / 0.6651 | 25.43 / 0.7818 |
| \(\times 4\) | MambaQuant | 2 / 2 | 28.84 / 0.8079 | 26.41 / 0.7114 | 26.18 / 0.6739 | 23.45 / 0.6648 | 25.38 / 0.7829 |
| \(\times 4\) | OAQ-SR | 2 / 2 | 28.78 / 0.8122 | 26.37 / 0.7143 | 26.17 / 0.6780 | 23.40 / 0.6681 | 25.33 / 0.7892 |
| \(\times 4\) | Ours (TGQ) | 2 / 2 | 29.12 / 0.8260 | 26.57 / 0.7259 | 26.32 / 0.6891 | 23.57 / 0.6805 | 25.58 / 0.8031 |
| \(\times 2\) | 全精度 FP32 | 32 / 32 | 38.26 / 0.9615 | 34.09 / 0.9221 | 32.36 / 0.9019 | 33.26 / 0.9378 | 39.35 / 0.9785 |
| \(\times 2\) | OAQ-SR (ICCV'25) | 4 / 4 | 37.22 / 0.9551 | 32.60 / 0.9071 | 31.51 / 0.8897 | 30.49 / 0.9122 | 36.91 / 0.9717 |
| \(\times 2\) | Ours (TGQ) | 4 / 4 | 37.56 / 0.9578 | 33.24 / 0.9149 | 31.86 / 0.8950 | 31.34 / 0.9204 | 37.85 / 0.9742 |
消融实验与分析¶
针对核心设计组件与分组策略，在 MambaIRv2-Light (\(\times 2\), W4A4) 上开展严格消融验证。
表 2：不同监督目标与 Token 分组机制的消融结果（Set14 与 Urban100）
| 配置编号 | 分组机制 (TG) | 截断优化目标 | Set14 (PSNR / SSIM) | Urban100 (PSNR / SSIM) | 说明 |
|---|---|---|---|---|---|
| 1 | \(\times\) | 仅均方误差初始化 (Eq. 10) | 32.15 / 0.8931 | 29.84 / 0.8915 | 传统逐层量化基准 |
| 2 | \(\times\) | \(\mathcal{L}_{\mathrm{pix}} + \mathcal{L}_{\mathrm{feat}}\) | 32.45 / 0.9056 | 30.34 / 0.9107 | 逐层量化经蒸馏微调 |
| 3 | \(\checkmark\) | 仅均方误差初始化 (Eq. 10) | 32.38 / 0.9015 | 30.56 / 0.9115 | 仅加入 TGQ 分组，无频域微调 |
| 4 | \(\checkmark\) | \(\mathcal{L}_{\mathrm{pix}}\) 仅像素损失 | 32.53 / 0.9018 | 30.72 / 0.9123 | 传统空域损失微调 |
| 5 | \(\checkmark\) | \(\mathcal{L}_{\mathrm{pix}} + \text{常规特征MSE}\) | 33.01 / 0.9091 | 30.89 / 0.9136 | 引入常规特征蒸馏 |
| 6 | \(\checkmark\) | \(\mathcal{L}_{\mathrm{pix}} + \mathcal{L}_{\mathrm{feat}}\) | 33.06 / 0.9107 | 30.99 / 0.9172 | 归一化特征相似度蒸馏 |
| 7 | \(\checkmark\) | 上述目标 + 仅幅值匹配 | 33.15 / 0.9122 | 31.14 / 0.9190 | 增强高频能量保留 |
| 8 | \(\checkmark\) | 上述目标 + 仅相位匹配 | 33.10 / 0.9128 | 31.10 / 0.9183 | 强化边缘结构位移对齐 |
| 9 | \(\checkmark\) | 完整目标（含 FPR 幅值+相位） | 33.24 / 0.9149 | 31.34 / 0.9204 | 联合优化取得最优恢复指标 |
表 3：不同分组策略与硬件端到端推理延迟对比（W4A4，MambaIRv2-Light）
| 分组策略 | 静态可部署 | Set14 PSNR | Urban100 PSNR | Orin Nano 延迟 (ms) | Orin Nano 加速比 |
|---|---|---|---|---|---|
| FP16 浮点基准 | - | 34.09 | 33.26 | 3149.6 | \(1.00\times\) |
| 4-bit Base (逐层基准) | \(\checkmark\) | 32.45 | 30.34 | 1380.1 | \(2.28\times\) |
| 随机分配分组 | \(\checkmark\) | 31.76 | 29.91 | - | - |
| 扫描前特征直接 K-means | \(\checkmark\) | 32.84 | 30.77 | - | - |
| 后验分组 Oracle (理论上限) | \(\times\) | 33.41 | 31.57 | 3213.9 | \(0.98\times\) (变慢) |
| Ours (前验预测器 + 查表) | \(\checkmark\) | 33.24 | 31.34 | 1410.7 | \(2.23\times\) |
关键发现¶
- Token 分组贡献占主导：在 Set14 上，仅引入 Token 分组机制便贡献了相较于逐层量化基线 77% 的性能提升（Urban100 上达 65%），证实消除空间 Token 激活分布分化是提升 Mamba PTQ 的第一核心因子。
- 频域微调的互补性：幅值匹配与相位匹配各自能带来明显提升，二者联合的 FPR 目标相较基准微调在 Urban100 上进一步斩获 0.35 dB 提升；频带误差分解显示该增益完全集中在高频频段，证明了 FPR 弥补传统空域损失高频欠监督的有效性。
- 突破硬件加速瓶颈：后验分组 Oracle 由于需在量化前执行额外的前向观察与排序统计，导致其在 Jetson Orin Nano 上延迟达 3213.9 ms，完全丧失 4-bit 加速比（\(0.98\times\)）；而本文的前验静态预测方案仅增加 2.2% 额外延迟，实现 \(2.23\times\) 真实加速。
亮点与洞察¶
- 结构与幅度解耦预测：巧妙发现虽然不同层后验激活绝对尺度变化剧烈，但空间 Token 的相对量化偏好结构高度稳定，因此仅用一个微型全局共享两层 FFN 就能有效解耦预测分类与数值截断，极大降低了模型复杂度。
- 与通道维度正交变换的完美互补：在附录分析中，将本文方法与大模型常用的 SpinQuant（学习旋转矩阵）结合时，性能在 W2A2 极限比特下获得进一步跃升（Manga109 上由 32.67 dB 飙升至 34.72 dB），证明空间 Token 轴的分组与通道维度的异常值平滑具有完全正交的增益空间。
- 出色的容错与泛化稳健性：针对预测器分类误差的扰动测试显示，即便在推理阶段人为随机打乱高达 20% 的分组索引，Set14 上的 PSNR 仅微跌 0.15 dB；同时方法能零调整直接泛化至图像去噪、JPEG 去块、SwinIR 架构以及 Vision Mamba 分类任务。
局限与展望¶
- 算子调度开销：尽管在前向低比特计算中避免了动态统计，但组索引查表与非连续内存寻址相比纯单核密集张量乘法仍引入了微小的显存访问离散性（在 A6000 上有 3.4% 额外时延）。
- 硬件适配广度：目前硬件加速内核主要验证于 NVIDIA GPU 与 Jetson Orin 平台的整型执行单元，在定点加速更为严苛的移动端 NPU、DSP 或专有 ASIC 上，多组截断边界的动态寄存器分配机制仍需深度软硬协同设计。
相关工作与启发¶
- vs PTQ4VM / Quamba / MambaQuant: 先前针对 Vision Mamba 的量化方案主要依赖通道缩放或方差旋转，本质上依然假设同层内所有 Token 共享单一截断尺度，在逐像素稠密预测的高频超分任务中无法化解选择性扫描引起的空间异质性；本文从空间 Token 轴切入，填补了动态选择性算子的空域自适应量化空白。
- vs OAQ-SR / 2DQuant: 传统卷积/Transformer 超分 PTQ 方案关注孤立通道离群点或二维网格分布，忽略了状态空间模型中隐藏状态长程累积带来的逐 Token 级联方差发散；本文揭示了选择性扫描独有的 1.80× 范围膨胀现象，并给出了静态可落地的数学解法。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次揭示选择性扫描引发的空间 Token 激活异质性机制，提出结构-幅度解耦的前验预测查表量化范式。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖超分、去噪、去伪影三大任务及 Transformer 与 ViM 跨架构实验，配备真实边缘设备端到端硬件加速评测。
- 写作质量: ⭐⭐⭐⭐⭐ 论证链条极其紧密，数学动机、硬件约束与经验分析层层递进，消融严谨。
- 价值: ⭐⭐⭐⭐⭐ 为状态空间模型在边缘高保真图像恢复设备的低比特端侧落地提供了极其坚实的方案参考。