QMambaBSR:查询驱动状态空间模型实现高效突发图像超分

📅 发布时间:2026/10/21 9:05:30
QMambaBSR:查询驱动状态空间模型实现高效突发图像超分
突发超分是我最近两年一直在追的方向。手机传感器物理尺寸有限单帧信息量天然不足厂商靠多帧合成硬撑画质但学术圈一直缺一个“复杂度可控、融合效果好”的统一框架。QMambaBSR做的事情本质上是把状态空间模型引入突发图像超分辨率用查询驱动的扫描机制替代传统注意力在几乎不损失精度的前提下把跨帧融合的计算复杂度从平方级拉低到线性。这篇文章我会从成像原理讲起拆到SSM的核心公式再落到网络每个模块怎么设计、消融怎么做的最后补一些复现时才会遇到的坑。1. 突发超分为什么折腾了这么多年从像素级平均到内容级融合1.1 单帧超分的瓶颈与burst的直觉先聊清楚问题本身。单帧超分辨率是在一张低分辨率图像里“脑补”高频细节不管用SRCNN还是Transformer本质上都是在猜测被采样丢失的信息。但突发图像超分有一个先验优势——你手里不是一张LR图而是同一个场景在极短时间内连续拍摄的多帧图像帧与帧之间存在亚像素位移。这就像你拿一台低像素相机快速连拍每个瞬间传感器采样栅格相对于场景的位置都略有偏移把这几帧的信息重新组合理论上能恢复出一个超过任何单帧真实分辨率的画面。这个思路在计算摄影领域早就被验证过。手机的夜景模式、HDR本质上都是多帧合成区别在于它们通常只做像素级的平均或简单的对齐后融合没有显式地建模“如何用多帧信息恢复高频细节”。突发超分想要解决的是更彻底的问题不仅要去噪、对齐还要真正重建出超越每帧采样率的高频内容。学术上这叫作“突发到单帧超分辨率”Burst-to-Single SR输入是N帧LR burst输出是一张HR图像。难点在于帧间关系极不稳定。手持拍摄时人眼会轻微颤动每帧之间存在全局平移、局部视差甚至还有卷帘快门导致的畸变。帧数越多融合时引入信息的概率越大但引入误差的概率也越大。一个糟糕的融合算法会在通道里产生鬼影、拖影比直接单帧超分还难看。1.2 突发超分的三大难点对齐、融合、高效性拆开看突发超分有三个核心矛盾。第一是对齐精度。亚像素级的对齐误差会直接转化为重建图像的模糊和伪影。传统方法用显式光流或单应性估计后来端到端方法把对齐隐式地放到网络里但隐式对齐对帧间大位移仍然敏感。第二是融合策略。简单平均是最低级的融合即使对齐误差很小它也只是降低了噪声无法引入新的高频信息。真正的融合要“自适应选择”对每帧中信息量最大的区域给予更高权重对模糊、遮挡的区域压低权重。这本质上是一个内容感知的选择过程。第三是计算效率。N帧输入意味着数据量是单帧的N倍一个合理的融合模块必须能同时处理空间维度和帧间维度而这两个维度在transformer框架下会迅速把计算量推到不可接受的程度。这也是为什么过去这些年突发超分方法一直围着“怎么融合”打转。早期基于光流的方案是分两步走先估计位移再把对齐后的帧喂给一个小网络。后来用成本体积、可变形卷积做隐式对齐效果好了一些但依然没有跳出“两步分离”的思路。真正让这个方向快速进步的是transformer的引入——自注意力天然适合做跨帧信息交互一个模块就能同时处理对齐和融合。于是新的问题出现了注意力是平方复杂度burst的帧数和patch数量一上来显存和耗时直接爆炸。1.3 方法演进从HDR到Transformer再到状态空间模型如果你翻一下这个领域最近几年的论文会发现一条很清晰的演进路径。早期代表作是Google的HDR它在raw域做多帧去噪和融合靠的是手工设计的对齐和加权平均厉害在工程细节扎实。之后学术界开始用CNN做端到端的burst超分典型方案是先把burst帧用共享权重的特征提取器编码再通过一个融合模块整合帧间信息。那时候融合模块一般是通道注意力或简单的串联卷积表达能力有限。到Transformer时代出现了两类做法一类是把burst的帧维度拉成token序列用自注意力做全局交互另一类是借鉴视频超分的思路用可变形注意力去捕捉帧间对应关系。这两类都能显著提升效果但代价是计算量激增尤其是在训练时一个小时的视频里取几帧做burst显存就吃紧得厉害。更让人头疼的是transformer融合多帧时注意力矩阵中大量计算花在了“本来就不需要对齐的区域”这种冗余在超分这种低层视觉任务里尤其明显——相邻帧大部分区域是高度相似的根本用不着平方级的全局注意力去建模。所以当状态空间模型SSM在2024年前后开始渗透到视觉任务时我第一时间就注意到它在突发超分上的潜力。低层视觉任务不追求像GPT那样的超长上下文但需要精确、可控的局部-全局信息交互。Mamba的线性复杂度让它天然适合处理多帧序列而“查询驱动”的设计则能填补SSM在内容感知上不如注意力的短板。QMambaBSR就是把这两个点结合了起来用一个查询向量去引导状态空间的更新让模型知道每一帧的哪些位置值得被读取、哪些位置应该被忽略。2. Mamba到底解决了什么状态空间模型在图像任务里的兴起2.1 自注意力的平方复杂度与长序列困境回到问题的根源。自注意力机制的计算方式是把输入token映射成Q、K、V三个矩阵然后计算QK^T的相似度矩阵。对于长度为L的序列这个相似度矩阵的规模和计算量都是O(L²)。在文本任务里序列长度通常是几千到几万这个复杂度还可以用各种稀疏化、局部化技巧压一压但在图像任务里情况更糟——一张图被切成patch之后token数量动辄上万而transformer在处理这种序列时几乎不会利用图像本身的二维结构。突发超分更是把这个问题放大了好几倍。假设你输入的是8帧、每帧切成16×16的patch那么融合时的token数量是8×16×162048。这看起来不算多但在网络内部的多层中反复做注意力同时还要保留空间维度做高分辨率重建实际运行时的显存开销是非常可观的。更致命的是自注意力无法天然区分“哪个token来自哪一帧”它必须依赖位置编码去告诉模型帧间关系这增加了学习的难度。有人可能会说可以用局部注意力、窗口注意力来降复杂度。但突发超分的跨帧信息交互本质上需要跨窗口、跨帧的通信窗口化会在语义上切断帧间联系效果会打折扣。Mamba给出的是一个正交的解决途径不计算注意力矩阵而是用一个循环式的状态变量来压缩历史信息每一步的更新只依赖当前输入和上一个状态复杂度O(L)且理论上可以建模任意长度的依赖。2.2 SSM/Mamba的核心线性复杂度输入依赖的选择性状态空间模型的核心数学形式其实非常简单。它描述的是一个连续系统状态向量h(t)随时间演化观测值y(t)由状态和输入x(t)共同决定。写成方程就是h(t) A·h(t) B·x(t) y(t) C·h(t) D·x(t)在实际应用中我们把这个连续系统离散化得到循环形式h_t A_bar·h_{t-1} B_bar·x_t也就是每个时间步用当前输入更新一个隐状态。这听起来很像RNN对吧确实如此。SSM的先进性主要来自两个改进。第一结构化参数化。A矩阵不是随意初始化的而是被设计成一种特殊的结构比如对角矩阵、低秩扰动这让模型可以在训练中稳定地学习长程依赖不会像普通RNN那样梯度消失。第二选择性机制。在Mamba的原始设计中A、B、C和步长delta都是输入相关的也就是说每个时间步看到不同的输入会用不同的状态转移矩阵。这就像注意力机制中的“query-aware”思路——模型可以根据当前输入内容决定“记住什么、忘记什么”。这个特性对burst超分极其关键因为不同帧、不同区域的信息价值差异巨大。还有一个常被忽略的工程点SSM的循环形式在推理时非常高效。它不需要缓存整条序列的KV对只需要维护一个固定大小的状态向量因此内存占用极低。这让Mamba在部署场景——比如手机上跑突发超分——有天然的吸引力尽管目前主要还是在GPU上验证。2.3 为什么突发超分是Mamba的“理想应用场景”你可能已经发现了突发超分和Mamba之间存在一种结构性匹配。burst图像可以很自然地展平成“帧-位置”的二维token网格然后按某种扫描顺序组织成一维序列。序列中的距离关系对应于空间距离和帧间距离Mamba的循环状态必须“走过”这些token才能聚合信息。这个建模方式跟transformer不同transformer是全局等权的所有token两两交互Mamba是路径依赖的信息沿着扫描顺序传播。后者看起来更“笨”但如果配合多条不同方向的扫描反而能构建出一个比注意力更经济的良好归纳偏置——它强迫模型在局部邻域内先做充分交互再沿着序列逐步传播到远处。另一个角度是数据增广。训练burst超分模型时帧间的位移模式是有限的主要是平移和少量旋转模型不需要像大语言模型那样建模开放世界的长程依赖它只需要在“该对齐的地方对齐该融合的地方融合”。Mamba的线性状态恰好能承载这种结构扫描经过一帧又一帧时状态变量天然积累了全局场景信息而查询机制可以从状态中按需提取特定位置的信息完成跨帧融合。这也意味着如果直接把面向文本或通用视觉的Mamba拿来改造成超分网络效果未必好——它需要针对“多帧融合”这个特定任务重新设计查询和扫描策略。QMambaBSR的贡献就在这里。3. QMambaBSR架构拆解查询状态空间模型是如何设计的3.1 整体pipelineburst特征提取、query构建、状态融合与重建QMambaBSR的整体结构沿用了低层视觉任务的经典“浅层特征提取-深层融合-重建”三段式但中间融合层换成了查询驱动的状态空间模块。先给出流程图式的描述输入N帧低分辨率burst ↓ 共享特征提取器几层卷积→ 每帧得到浅层特征F_i ↓ 从参考帧特征生成查询Token序列Q ↓ Q与所有帧特征一起进入多个查询状态空间块QSSB QSSB内沿多方向扫描用Q调制SSM参数更新状态并输出融合特征 ↓ 融合特征经过重建头PixelShuffle上采样→ 输出HR图这个流程的关键设计决策有三处。第一用共享权重的卷积提取每帧特征保证不同帧的特征空间一致第二reference frame选哪一帧QMambaBSR采用了“中间帧优先”的策略因为中间帧在时序上离其他帧的平均距离最小信息融合更均衡第三查询Token的构造方式。常规做法是把参考帧的2D特征图展开成一串token每个token对应空间中的一个patch位置后续所有的跨帧融合都以这些token为锚点。这样做的直觉是突发超分的最终输出只有一个参考坐标系参考帧的空间网格所有其他帧的信息都要被“拉”到这个坐标系里来。3.2 QSSB模块内部查询如何驱动状态更新QSSB是整篇论文的核心模块。它要解决一个根本问题标准Mamba的选择性机制虽然输入相关但没有显式的查询目标。在语言模型里当前token就是查询它决定下一步状态怎么更新但在burst超分里我们关心的是参考帧的每个位置能否从其他帧获取信息此时状态更新不能只看“当前扫描到哪个token”还要看“我们正在为哪个查询位置收集信息”。QMambaBSR的做法是在扫描过程中把查询tokenq和当前扫描到的key-frame特征k拼接起来共同生成SSM的离散化参数。具体来说B Linear(concat([q, k])) # 输入投影 C Linear(concat([q, k])) # 输出投影 delta sigmoid(Linear(concat([q, k]))) # 步长 A 由初始化参数加上一个输入依赖的残差项这样设计之后同一个帧的特征在不同的查询位置下会产生不同的状态转移行为。举个具体场景参考帧第1个patch位置上是一个边缘区域第2个patch位置是平坦区域扫描到另一帧的同一个空间位置时第1个查询会让状态更倾向于“读取并保留高频边缘信息”第2个查询则会让状态选择“忽略微小的像素差异”。这就是“查询驱动”的含义——信息的选择权交给了目标位置而不是源位置。模块内部还有一层残差连接和LayerNorm。每个QSSB的输出维度与输入一致可以堆叠多次以增加融合深度。QMambaBSR默认使用4个QSSB每个QSSB内部做4个方向的扫描从左到右、从右到左、从上到下、从下到上这比单向扫描能覆盖更完整的空间-时间邻域。3.3 扫描策略与位置编码把非序列数据变成序列的学问这是工程上最容易翻车的地方。Mamba是为1D序列设计的图像是2D网格burst是3D帧×高×宽直接展开成1D序列会丢失空间拓扑结构。QMambaBSR采用的方案是把帧维度作为序列的一维把空间维度作为另一维交替扫描。具体展开方式是先固定帧的顺序在每个帧内部按行展开得到一个包含“空间-帧”信息的token序列。然后做4方向扫描正向、反向、按列扫描、反向列扫描。扫描过程中相邻token在序列中的距离既包含空间位移也包含帧间位移模型可以通过学习状态转移来隐式地建立跨帧对应关系。为了让模型知道token来自哪一帧、在什么空间位置QMambaBSR在进入SSM前给每个token加了一个可学习的帧索引嵌入和一个二维位置嵌入。帧索引嵌入的维度很小实验中是16维但它的存在非常关键——消融显示去掉帧索引后PSNR下降了约0.1dB以上。原因很直观状态空间模型是靠序列顺序来组织信息流的如果模型不知道当前token属于哪一帧就很容易把不同帧的相同空间位置混淆导致融合时纹理错位。3.4 重建头与损失函数设计融合后的特征图已经聚合了所有帧的有效信息重建阶段相对简单。QMambaBSR用的是渐进式上采样先经过两层卷积把通道数降低再用PixelShuffle放大2倍最后接一层卷积输出HR图。如果scale factor是4就做两次放大中间穿插ReLU。损失函数方面主loss是Charbonnier loss它对异常值比L2更鲁棒L_pixel sqrt((I_pred - I_gt)^2 epsilon^2)epsilon1e-3除此之外还加了可选的感知损失和频率域损失。感知损失用预训练的VGG特征做约束能让生成图像纹理更自然频率损失是对预测和GT分别做FFT后在频域计算L1它能让模型更关注高频分量。不过我自己的复现经验是频率损失的权重不宜过大否则会产生振铃效应。4. 实验与消融哪些设计真正带来了收益4.1 训练设置与数据集QMambaBSR在训练时采用了与主流burst超分方案一致的做法从高分辨率图像出发模拟生成低分辨率burst。具体流程是取一张HR图随机裁剪出若干个HR patch然后对每个patch施加随机的亚像素平移平移量在±2个像素以内、随机高斯模糊模拟镜头失焦、添加泊松噪声和高斯噪声最后下采样得到对应的LR帧。这种合成训练策略的优点是数据量无限、真值完备缺点是模拟的噪声和模糊分布与真实手机拍摄场景有差距。因此QMambaBSR也做了真实burst数据的微调实验直接用手持手机拍摄的真实burst作为输入以长时间曝光的高质量图像作为监督。这类数据噪声更真实但数量有限通常作为验证集和微调用。训练配置上默认patch大小是128×128输入LRburst长度8帧batch size 8优化器用AdamW初始学习率1e-4采用cosine退火调度总共训练30万步。这样一份配置在单张A100上大约需要两天时间如果显存不够需要把batch size或帧数调小。4.2 在合成与真实burst数据上的表现定量结果部分QMambaBSR在其官方的验证集上相对带注意力的transformer基线有约0.2-0.3dB的PSNR提升这个提升幅度在超分领域算是肉眼可见的尤其是在边缘和高频纹理区域。更值得注意的是LPIPS指标——感知距离下降更明显说明模型恢复出的图像在结构上更接近真实场景。真实burst场景下主观质量的差异更直观。参考帧单帧超分结果往往有涂抹感而QMambaBSR重建出的图像在招牌文字、树叶边缘等地方明显更锐利、更干净。这里我想强调一点突发超分的真实数据评测中定量指标经常和主观感受不一致原因是真实数据的GT本身未必完全对齐、未必是完美的“真值”所以主观图对比有时比数字更能看出融合算法的优劣。推理速度方面8帧128×128输入、4倍超分的单次推理在A100上大约需要60ms左右比同规模transformer基线快约3-5倍。速度提升主要来自三点线性复杂度、不需要KV缓存、GPU友好的串行扫描实现。4.3 消融实验的三个关键发现消融实验部分有几个结论值得拎出来讲。第一个是“查询驱动参数调制”的收益。当QSSB中的B、C、delta改为只由输入帧特征生成、不含查询条件时PSNR下降了约0.15dB。这说明查询信息并不只是锦上添花而是真正在引导状态更新方向。第二个是扫描方向数量的影响从单向扫描增加到双向扫描PSNR提升明显约0.2dB从双向增加到四向提升变缓约0.05dB但主观纹理更稳定继续增加方向反而因序列过长引入不稳定。第三个是帧数的边际效应从4帧提升到8帧收益显著从8帧到14帧仍有收益但递减超过14帧后PSNR几乎不再变化而训练耗时线性增加。这说明突发超分的信息量在帧数达到一定规模后会饱和无脑堆帧数不是好策略。另一个消融是位置编码的对比用绝对位置编码比用相对位置编码更稳定这与图像超分任务关注“绝对坐标对齐”的特性一致。很多人会把NLP的经验搬过来用相对位置编码结果发现跨帧融合时空间对应关系变模糊了。5. 复现与部署经验谈显存、算子与训练稳定性5.1 环境配置与mamba-ssm的编译坑先说环境。Mamba类模型通常依赖selective scan的CUDA算子这个算子目前主要在官方仓库中提供需要从源码编译。问题集中在三点CUDA版本兼容性、PyTorch版本匹配、编译时间。我复现时用的组合是CUDA 11.8、PyTorch 2.0.1、Python 3.10。如果CUDA版本过高比如12.4部分老版本的算子可能编译失败报错通常出现在scan算子的模板实例化阶段。这时候不要硬刚编译器错误先降CUDA版本或者用官方预编译的wheel碰运气。编译本身耗时大约10-20分钟取决于机器性能建议在后台跑。如果编译实在过不去有一个纯PyTorch的替代实现把selective scan写成循环逐时间步更新状态。这个版本速度慢了约3-5倍但能用来调试网络结构和训练超参等确认无误后再切换到CUDA算子跑正式实验。我的建议是先在纯PyTorch模式下把训练流程跑通一次再换算子能省大量排查时间。5.2 显存优化实践patch、帧数与batch的平衡QMambaBSR的显存压力主要来自两方面多帧特征同时驻留显存以及多个QSSB的中间激活值。8帧128×128、batch size 8、4尺度特征在A100 80G上勉强够用换到40G的卡就得动配置。我的实践顺序是第一优先降低batch size到4第二降低patch大小到96×96第三才考虑减少帧数。原因是帧数直接影响融合质量patch大小影响感受野覆盖而batch size对训练稳定性的影响在AdamWcosine下并不大。如果还不够省可以打开gradient checkpointing它会在反向传播时重新计算前向激活能省约50%显存代价是训练时间增加20%左右。还有一个容易忽略但很有效的技巧把参考帧特征和融合后的特征都保留在fp16但SSM内部的delta和A矩阵用fp32计算。原因是delta经过sigmoid后范围很小学习率对它的影响极其敏感fp16下很容易因为精度不足导致状态更新不稳定。5.3 训练稳定性数值精度与学习率策略Mamba类模型在低层视觉任务里的训练稳定性是个真问题。我在复现过程中遇到过两种情况。第一种是loss曲线中途突然跳到NaN。排查后发现是delta初始值过大导致状态更新发散。解决办法是把delta的初始偏置设成负值比如-2让sigmoid输出靠近0状态更新在训练早期保持保守等模型学会了再逐步增大。第二种是混合精度训练时PSNR在某个阶段停滞。原因是fp16下B和C的数值范围差异太大部分通道的信息在反向传播中丢失。解决方法是给B和C的输出层加一个可学习的缩放因子或者干脆让这两个投影层保持fp32。学习率方面我踩过一个坑直接沿用ViT常用的2e-4初始学习率结果早期loss震荡得很厉害。后来改成1e-4并配合1100步的线性warmup情况立刻好转。这个任务对学习率其实比想象中敏感因为SSM的循环结构会让梯度的尺度随时间步累积学习率过大很容易破坏状态转移矩阵的稳定性。数据增强方面除了常规的水平翻转和旋转90度帧内的随机色彩扰动和帧间的随机顺序打乱都有效。随机顺序打乱尤其值得做它强迫模型不能依赖帧序信息只能依赖内容对齐这让模型在真实手持拍摄场景帧序接近但不完全有序中的鲁棒性更好。6. 写在最后QMambaBSR之后突发超分还能往哪走说点个人体会。作为复现者我最直观的感受是QMambaBSR把突发超分的“融合复杂度”这个长期瓶颈真正压了下来。你会发现一旦融合模块的复杂度不再是瓶颈整个系统的设计余地就大了很多——你可以把省下来的显存拿去做更大的patch、更深的网络也可以把8帧扩展到14帧去压榨更多细节。但这次实现在技术上也留了不少可以继续做的事情。首先是端到端的帧间位移建模。目前的方案虽然用查询机制隐式地处理了帧间对应但没有显式地估计亚像素位移如果能把轻量级光流或者隐式位移嵌入引入到SSM的状态更新中理论上还能再提升对齐精度。其次是面向真实场景的退化建模合成数据与现实差距依然存在用半监督或无监督的方式在真实burst上自适应是一个更实用方向。最后是软硬件协同的移动端部署Mamba的线性复杂度优势在GPU上已经体现但手机端的NPU对循环扫描的支持还很弱这需要算法和硬件共同推进。我在跑完这一串实验之后最大的体会是状态空间模型不是要取代Transformer而是把Transformer最贵的那部分计算全局交互用一种更经济的方式近似表达。在突发超分这种“序列长、但结构性强、不需要开放世界联想”的任务上这种近似几乎是无损的。对做这个方向的同学我建议你也动手把基线跑一遍亲手看看查询驱动状态更新和固定参数状态更新的差异那些消融表格里的0.1dB在真实代码里意味着完全不同的优化梯度形态。