论文题目:MambaIRv2: Attentive State Space Restoration

论文地址:[2411.15269] MambaIRv2: Attentive State Space Restoration

论文主要可用模块:ASSM注意状态空间模块

论文讲解

1.结构:
MambaIRv2采用经典的编码器-处理器-解码器结构,首先统过浅层特征提取,使用使用3×3卷积层将输入图像转换为特征表示,后面深层特征提取由多个MambaIR Block堆叠组成。这种也可以做U-Net下采样与上采样中注意力模块的替换,但是需要考虑网络深度和训练效率。
2.创新点:
ASE是论文最核心的创新点。传统的Mamba模型存在因果性约束,即当前像素只能访问序列中之前的像素信息,这限制了图像恢复任务中的全局信息利用。而ASE通过引入可学习的提示(prompts)到状态空间方程中,使查询像素能够"跳出"扫描顺序的限制,直接访问图像中任意位置的相关像素,实现了非因果建模能力,打破了传统Mamba的单向依赖局限,显著减少了多方向扫描的计算冗余,提高效率。
3.ASSM详细讲解:
传统Mamba模型遵循严格的因果性约束,即当前时刻t的隐状态只能依赖于当前输入和前一时刻的隐状态,在图像处理中,这意味着当前像素只能"看到"扫描序列中之前的像素,无法直接访问后续或空间上相关但序列位置靠后的像素信息。而ASE重新定义了状态空间方程:加入了可学习的提示(prompt)向量,在2维信息中对于每个位置t,生成对应的提示向量Pt,提示向量编码了全局上下文信息,通过学习过程优化,使其能够"指向"最相关的像素区域。详细 过程如下:
  • 在标准状态空间模型中加入Pt:作为"桥梁"连接全局信息。
  • 计算全局相似度权重,将全局序列作为key和value,当前位置信息作为query,计算注意力分数,生成提示向量Pt。
  • 根据ASE的方程更新注意力增强的新状态。
4.优点:
  • 打破因果性约束:通过提示向量Pt可以看到全局信息。
  • 单次扫描实现全局建模:ASE通过注意力状态空间方程允许超越扫描序列进行注意,仅用单次扫描就能实现图像展开,大幅降低计算复杂度。
  • 自适应全局信息聚合:提示向量Pt根据当前查询动态生成,可以实现不同位置关注不同的全局信息,并且实现内容自适应的长程依赖建模。

实际应用

1.创新应用:
  • 图像去噪: 在图像去噪任务中,ASE可以通过其非因果建模能力,让每个像素位置都能访问全图的清洁特征信息。传统的Mamba模型由于因果性约束,只能利用序列中前面位置的信息进行去噪,这在处理复杂噪声模式时效果有限。通过ASE的提示机制,网络能够识别图像中语义相似的清洁区域,并将这些区域的特征信息传递给当前噪声像素,从而实现更精准的噪声去除。特别是在处理结构化噪声或需要全局统计先验的去噪任务中,ASE的全局信息聚合能力能显著提升去噪性能。
  • 视频修复: ASE机制可以扩展到时序维度,在视频修复任务中发挥重要作用。传统的视频处理方法往往受限于时间上的因果约束,只能利用过去帧的信息。ASE通过引入可学习提示,能够让当前帧访问整个视频序列中语义相关的帧信息,打破时间顺序的限制。这对于视频去模糊、视频超分辨率等任务尤为重要,因为相似的运动模式或纹理信息可能出现在视频的任意时刻,ASE能够有效挖掘和利用这些时空相关性。
  • 医学图像分析: 在医学图像处理中,ASE的全局建模能力可以帮助网络更好地理解解剖结构的整体性。医学图像往往具有强烈的结构相关性,比如器官的对称性、血管的连续性等。传统的局部处理方法可能会错失这些重要的全局特征。ASE通过其提示机制,能够让网络在处理局部区域时同时考虑全局的解剖学先验,从而在病灶检测、器官分割、图像配准等任务中取得更好的效果。
  • 遥感图像处理: 遥感图像通常具有大尺寸和复杂的空间结构,ASE的线性复杂度优势在这类应用中尤为突出。传统的Transformer在处理高分辨率遥感图像时计算开销巨大,而传统Mamba又无法有效利用空间上相距较远但语义相关的区域信息。ASE通过将语义相似的地物特征进行全局关联,能够在保持高计算效率的同时,实现对大范围地物的一致性建模,这对于土地利用分类、变化检测、目标识别等任务具有重要价值。
  • 多模态图像融合: ASE机制可以被扩展用于多模态图像的特征融合任务。不同模态的图像(如可见光、红外、深度图等)往往包含互补的信息,但如何有效融合这些信息是一个问题。ASE的提示机制可以学习跨模态的语义对应关系,让一个模态的特征能够指导另一个模态中对应区域的特征提取和增强。这种跨模态的全局建模能力对于监控系统、自动驾驶、机器人视觉等应用场景具有重要意义。

mamba模型具有一定的部署门槛,暂时只能在linux上部署,同时对于SSM也需要有一定的理解,但是在图像处理上是一个比较好用的创新点,基本用了就涨点,现在发关于mamba的论文很好发,同学们可以着重看看mamba部分,跳出传统的卷积注意力,就能发出好论文。

2.实践展示:

我目前尝试了水下图像增强,去雨,低照度增强,效果都是不错的,SSIM0.93以上,PSNR32+。大家可以尝试加入到自己的模型中,或许对你的模型会有很大的提升哦。最后,需要论文或是代码模型辅导的可以直接在主页联系我哦!

Logo

火山引擎开发者社区是火山引擎打造的AI技术生态平台,聚焦Agent与大模型开发,提供豆包系列模型(图像/视频/视觉)、智能分析与会话工具,并配套评测集、动手实验室及行业案例库。社区通过技术沙龙、挑战赛等活动促进开发者成长,新用户可领50万Tokens权益,助力构建智能应用。

更多推荐