Rep-Mamba：重参数化Mamba破解遥感超分轻量化困局

遥感超分模型的轻量化困局，该有个新解法了

做遥感图像超分辨率的朋友应该都有体会：和自然图像超分不同，遥感影像的细节密度极高，一片农田里的田埂纹理、一栋建筑的天台结构、一条道路的路面颗粒，这些高频信息在降采样之后几乎是“糊成一团”的。想恢复出来，现有的Transformer类超分模型确实效果不错，但参数动辄几十兆，推理时注意力机制在长序列上的计算量又让人头疼。

我之前在一个边缘端遥感图像增强项目里就碰到过这种困境：模型在服务器上跑得挺好，一移植到嵌入式推理设备上，显存和功耗双双告急。后来我把目光投向了一个很有潜力的方向——用Mamba这种线性复杂度的状态空间模型来替代自注意力，再配合重参数化（Re-Parameterization）思想压一压推理时的结构开销。这条路走下来，就是Rep-Mamba的思路：在视觉Mamba框架里做重参数化设计，面向轻量化遥感图像超分辨率。

这篇文章把整个设计思路、模块拆解、重参数化的数学原理、训练到部署的完整流程，以及我实际复现中踩过的坑一次说清。适合正在做遥感超分、图像复原，或者想给现有超分模型瘦身上线的朋友参考。

1. 遥感影像超分为什么必须考虑“轻量化”：先聊清楚问题本身

很多做超分的同学对遥感图像的态度是“拿自然图像的模型过来直接训”，这个想法在实验室里倒也能跑，但换成真实遥感场景就很容易翻车。原因在于遥感超分面临的资源约束比自然图像苛刻得多。

1.1 遥感图像的退化过程和自然图像完全不同

自然图像的退化主要是光学模糊加噪声，模型只要学一个通用的映射就能应付。遥感图像不一样，成像链路里包含大气散射、传感器MTF退化、平台运动模糊，甚至还有卷帘快门带来的行间错位。这些退化叠加在一起，高频信息损失得非常严重。更麻烦的是，遥感图像通常是大幅面单景影像，动不动就是几万乘几万像素，落地时必然要切块处理。而切块之后的每个小图块，依然保留了极高的纹理密度和重复结构（比如建筑群、农田格网）。

这种场景下，轻量化不是加分项，而是能不能上线的关键指标。卫星端或无人机端的计算平台，算力和内存都是“挤牙膏”式分配的。我在项目里做过粗略测算：一张标准的遥感影像切块到256×256，用SwinIR那类Transformer超分模型跑起来，单块推理还算能接受，但当分辨率推到512甚至更大，注意力矩阵的计算负担会陡增。用数学点的话说，自注意力的复杂度是图像序列长度的平方关系，这在遥感大幅面切块上极度不友好。

1.2 注意力机制在遥感超分上的开销瓶颈

我这里不是说Transformer不好。论建模能力，自注意力机制确实强，它能全局建模、捕捉长距离依赖，这对自然图像超分很有用。但遥感影像经过切块后，每一块的尺寸仍然不小，如果延续“整图计算注意力”的思路，会出现三个问题：

• 计算复杂度随输入尺寸二次增长，显存吃紧，推理延迟拉高；
• 结构相似、周期性强的遥感纹理很容易让注意力权重趋同，本质上是对计算力的浪费；
• 模型为了追求性能不断加深加宽，参数数量膨胀，和边缘端部署的硬件预算直接冲突。

所以在做轻量化改造的时候，我当时给自己定了两个原则：第一，不能因为省计算而牺牲全局建模能力，遥感图像里的地物上下文关系是很重要的；第二，推理时网络结构必须尽量“干净”，能合并的层全部合并掉，把模型压到最简。

这两个原则，恰好就是Rep-Mamba要去解决的问题——用Mamba的线性复杂度做全局建模，用重参数化把训练时的复杂结构在推理时折叠成单一路径。

1.3 轻量化不是妥协，而是部署刚需

很多人提起轻量化，总觉得是在性能上“打折”。实际并不是。轻量化设计如果能做到“训练时强、推理时简”，性能几乎无损，那这就不是妥协，而是更聪明的建模方式。

Rep-Mamba的核心思路就在于：训练阶段保留多分支结构，让模型有足够的梯度通路和表征能力去充分学习；推理阶段通过重参数化，把多分支合并成一个等价的单分支结构。这样服务器上训练时模型可以“重一点”，实际部署到边缘端时又不吃亏，本质上是在同一个模型里做了两种形态的切换。

从直觉上你就明白了，这类方案的难点不在“重参数化”这个操作本身——这套技术在卷积网络里已经很成熟，真正难的是怎么把重参数化的思想和视觉Mamba的结构搭在一起，让两者的优点不发生冲突。下一节就进入这个核心设计。

2. Rep-Mamba的整体设计：用状态空间模型替换注意力的核心逻辑

先简单把Mamba是什么说清楚。如果你对状态空间模型（SSM）还不太熟悉，可以把它理解成一种“有记忆的序列建模方式”：每一个输出都由当前输入和系统内部状态共同决定，状态不断更新，从而让信息在长序列中流动。

Mamba近年火起来，是因为它在SSM的基础上加了选择机制和硬件感知算法，让模型能根据输入动态调整记住什么、忘记什么。它最大的优势是推理复杂度与序列长度呈线性关系，这正好击中自注意力平方复杂度的痛点。

2.1 为什么选Mamba而不是继续堆Transformer

放在遥感超分这个具体任务里，Mamba有它天然的适配性。遥感影像切成块之后，块内依然有大量空间冗余。如果建模方式是“每个位置都去和其他所有位置算相关度”（也就是自注意力），大量计算会被花在那些本就高度相关的位置上，效率和必要性都成问题。

Mamba的线性复杂度意味着序列拉长时，计算量基本是跟着长度走，而不是几何级数暴涨。同时因为状态空间模型有“循环式地状态传递”特性，它在处理局部连续结构时会自然地保持信息连贯性，不会像窗口注意力那样需要在窗口之间再做额外交互。

我在实际对比中测过，相似参数量下，Mamba块在512×512输入图像上做特征提取的显存占用比同规模Swin Transformer块要低不少。这个优势在遥感大幅面影像上非常实用。

2.2 从一维SSM到二维视觉扫描的关键改造

Mamba原生设计针对的是序列数据，也就是一维离散信号。图像是二维的，不能直接拿来用，于是就有了视觉Mamba（Vision Mamba）的做法：把图像按某种顺序展开成序列，再做扫描。

Rep-Mamba里沿用的就是视觉Mamba常用的2D选择性扫描（SS2D）思路：一张特征图分别从左上、右下、右上、左下四个方向展开成序列，每个方向都走一遍状态空间模型，再把四个方向的输出重排回二维特征图，最后融合。为什么要扫四遍？因为序列方向感会丢失空间位置信息，单方向扫描只能捕捉单向的依赖关系，四个方向才能把完整上下文都捞回来。

这里特别提醒一下，四方向扫描带来的计算量是单方向扫描的四倍，这是视觉Mamba类模型普遍存在的开销。如果你的硬件预算非常紧，可以考虑只在主干后半部分使用全方向扫描，前半部分用更轻的模块。这是工程上的取舍。

2.3 重参数化在Mamba主干里的位置

重参数化改造不是简单地在模型尾部加一个卷积，而是应该贯穿到整个主干特征提取中。Rep-Mamba的基本做法是：每个Mamba块的外围上下接两个重参数化卷积层，训练时它们是多分支形态，推理时合并成单个卷积。这样主干里的每个块都具备“训练时多通路、推理时单通路”的特性。

我复现过程中体会最深的一点是：重参数化分支设计需要和Mamba块的内部结构匹配。Mamba块本身内部有线性层、卷积层和SSM，如果外围的分支设计不合理，合并之后可能出现通道不匹配或者信息流短路的问题。具体细节下一节展开。

3. 主干网络逐模块拆解：从Patch Embedding到重建头

整个Rep-Mamba超分模型走的是经典的超分结构路线：浅层特征提取、深层特征变换、重建。区别集中在深层特征变换这一块，名字上就叫“重参数化Mamba块组”。

3.1 输入处理与浅层特征提取

输入的低分辨率图像先经过一个3×3卷积做浅层特征提取。这里有一个遥感超分特别要注意的点：遥感影像常常是多光谱的，如果你的数据有多个波段，第一层卷积的输入通道要相应调整。不要只想着RGB三通道，实际项目里常见的4波段或者8波段输入并不罕见，这直接决定了整个网络的参数起点。

浅层卷积的通道数一般取64或者128，太大的话轻量化目标就落空了。我建议先用64通道起步，跑通流程后再根据效果考虑是否加宽。卷积后的特征一方面进入深层变换模块，另一方面通过一条残差短接路径直接送到重建头。这种“全局残差”设计是超分模型的标配，能让主网络集中精力学习高频残差而不是完整图像。

3.2 内部怎么用SS2D做四种方向的扫描

每个重参数化Mamba块的内部结构大致是这样：

第一，输入经过LayerNorm归一化后，先进一个线性层升维，把通道维度扩展为原来的两倍，为后续扫描提供更丰富的特征表达。

第二，特征图被拆成两个分支。一个分支经过一个深度可分离卷积和SiLU激活后进入SS2D。另一个分支只做线性变换。这两条路径最终通过逐元素相乘的方式融合，这个结构在语言Mamba里叫门控机制，作用是让模型自己决定“扫描出来的信息有多少要进入下一层”。

第三，SS2D的输出经过线性层降维后，和输入的短接路径相加。这样整个块是一个残差结构，梯度传播很顺畅。

Mamba块的类型在这一层不需要堆太多，我的经验是12个块左右就能在遥感场景下取得不错的平衡，再多的话推理延迟会明显上来，但指标提升开始变缓。

3.3 重建头上的重参数化设计

重建头是决定输出图像质量的关键部分，也是我在设计时最纠结的地方。常见的超分重建头是亚像素卷积（PixelShuffle）加卷积的组合，Rep-Mamba的做法也类似，但每个卷积层都改成了重参数化形态。

具体来说：深层特征先经过一个3×3重参数化卷积做特征整理，然后按缩放倍率进行像素重排上采样。尺度因子为2或4时的处理略有不同，倍率越高的重建，中间往往要多插一个卷积层来缓解锯齿效应。最后再接一个重参数化卷积输出超分结果。

这里踩过的坑我要提前说：重参数化卷积在重建头使用时，输出层的初始化要格外注意。如果输出卷积的最后一个分支用默认的随机初始化，共享的训练早期会出现噪声偏大的问题，超分图像上表现为高频区域出现轻微条纹。建议把最后一个卷积层的权重初始化到接近零值，让网络在训练早期基本沿残差路径走，再慢慢学习重建细节。

4. 重参数化的完整数学推导与实现

重参数化听起来玄，其实数学原理非常直观：训练时网络里有多个并行分支，推理时这些分支可以合并成一个等价的卷积层。合并的依据是线性运算的可加性——多个卷积相加，就等于一个对应权重相加后的卷积。

4.1 训练态的多分支结构与推理态的单分支

Rep-Mamba里重参数化卷积在训练时有三个分支：

• 3×3卷积分支，负责核心的局部特征提取；
• 1×1卷积分支，负责通道间的线性混合；
• 恒等映射分支，也就是直接拷贝输入，保证梯度直通。

每个分支后面都跟一个批归一化层。这个结构从名字上你应该能感觉到，它借鉴了RepVGG的设计。三个分支合在一起的含义是：网络有更多样化的信息通路，梯度能通过不同感受野的分支回流，训练更充分。

推理时，这三个分支加上各自的BN层，可以合并成一个3×3卷积。数学上分两步走。

4.2 BN折叠与分支卷积核合并

先说BN折叠。批归一化在推理时会退化成对参数的线性缩放和平移，可以写成一个固定公式。对于任意卷积层输出 y = Wx + b，后面的BN在推理模式下等价于：

y_hat = (y - μ) / sqrt(σ² + ε) * γ + β

把卷积表达式代进去，就能得到新的权重 W' = W * γ / sqrt(σ² + ε)，新的偏置 b' = (b - μ) * γ / sqrt(σ² + ε) + β。这样就完成了一个“卷积+BN”到“单卷积”的合并。

第二步是分支合并。1×1卷积可以看成是3×3卷积的一种特例，只需把1×1的卷积核周围补零，就变成了一个“中间是权重、四周为零”的3×3卷积核。恒等映射分支更特殊，它可以看成是一个卷积核中心为1、其他位置为0的3×3深度卷积。当输入输出通道数相同时，这个分支等效于一个单位矩阵形状的卷积核。

三个分支都转换成3×3卷积形式后，直接按对应位置相加即可，偏置也直接相加。最终得到单一卷积层的权重和偏置。

4.3 实际合并时的PyTorch伪代码

如果你用PyTorch做实验，合并过程大致像下面这样，我简化了一下关键逻辑：

import torch import torch.nn as nn def fuse_bn_into_conv(conv, bn): # 将BN的缩放与平移折入卷积权重和偏置 gamma = bn.weight.data beta = bn.bias.data mean = bn.running_mean var = bn.running_var eps = bn.eps std = torch.sqrt(var + eps) scale = gamma / std w = conv.weight.data * scale.view(-1, 1, 1, 1) if conv.bias is not None: b = conv.bias.data else: b = torch.zeros_like(mean) b = (b - mean) * scale + beta return w, b def repvgg_style_fuse(conv3x3, bn3x3, conv1x1, bn1x1, identity): w3, b3 = fuse_bn_into_conv(conv3x3, bn3x3) w1, b1 = fuse_bn_into_conv(conv1x1, bn1x1) # 1x1卷积核补零扩展为3x3 w1_padded = torch.zeros_like(w3) w1_padded[:, :, 1:2, 1:2] = w1 # 恒等分支生成等效3x3卷积核 id_tensor = torch.zeros_like(w3) for i in range(id_tensor.shape[1]): id_tensor[i, i, 1, 1] = 1.0 final_w = w3 + w1_padded + id_tensor final_b = b3 + b1 + 0.0 return final_w, final_b

代码本身不复杂，但我建议你顺手写一个单元测试：随机生成输入，比较合并前后两个模型的输出差异。因为涉及多个分支求和，累计误差经常在数值精度上出问题，差个1e-5很正常，但如果差到1e-2级别，就要检查是不是某个分支的BN统计量没有进推理模式。

5. 训练到部署全流程：数据、损失函数、结构转换与验证

模型结构设计的再好，训练流程不对照样出不了效果。这一节把从数据准备到部署验证的完整流程串一遍。

5.1 数据集和训练方式

遥感超分的训练数据构建方式和自然图像类似：从高分辨率遥感影像上裁剪子图，然后通过降采样模拟低分辨率输入。降采样方法很讲究，不要只用最简单的双三次插值（bicubic）。真实的遥感退化是模糊加噪声的组合，你可以把双三次降采样混合高斯模糊、泊松噪声来做数据增广，让模型见过的退化类型更丰富。

我建议采用“多退化预训练+单一退化微调”的策略：先在混合退化数据上把模型练到收敛，再用目标任务对应的退化方式做少量微调。这样模型泛化能力更强，在真实遥感影像上的表现比单纯用bicubic退化训练要稳。

训练时输入patch不要太大，我常用的是64×64的低分辨率块配对应的高分辨率块。道理很简单：显卡显存有限，Mamba块在长序列上的显存峰值不容小觑。有条件的话再用混合精度训练（AMP）进一步压显存、提速度。

5.2 损失函数怎么选

在超分任务里，L1损失是目前的主流选择，比L2（MSE）能保留更多锐利边缘。不过做遥感超分时我强烈建议加上一个频域或梯度域的辅助损失。遥感影像的高频细节是质量关键，只优化L1会让输出在平坦区域很好，但纹理区域偏“软”。

我当时用了一个简单的组合：主体是L1损失，外加一个梯度域损失（计算输出和标签的梯度图之间的L1距离）。这个组合简单有效，不需要额外调太多超参。

如果做的是多光谱数据，注意各波段的灰度分布差异很大，建议在计算损失前对每个波段做标准化，避免近红外或海岸波段主导整个训练信号。

5.3 权重转换与一致性验证

训练完成后，先不要急着转推理结构。做一次完整的“训练态→推理态”验证，确保转换前后输出图像的一致性。

具体步骤是：保存训练态权重，加载模型执行一次前向推理，然后对每个重参数化卷积执行合并，保存推理态权重，用同一张输入图做二次推理。对比两次输出图像的PSNR，理想结果应该是差异无穷小，如果PSNR在80dB以上，基本可以放心转换。

我建议把这个一致性验证写成一个自动脚本，集成到每次导出模型的流程里。别把转换当成一个手工操作，人一旦大意，漏掉某个分支没合并，部署出来就是结果图出现随机网格噪声，排查起来会让人非常崩溃。

导出推理态模型后，可以做一次量化感知训练或量化校准。Mamba类模型对量化其实比较敏感，尤其SSM内部的状态转移参数，低位宽量化时要重点盯住PSNR曲线有没有突然掉点。保险起见，可以先做W8A8量化试点，确认稳定后再考虑更低比特。

6. 我复现和实验中的体感与踩坑备注

这部分不写面面俱到的心得，就说几个我反复折腾过、最有共性的问题。

6.1 坑一：BN在轻量模型里的精度波动

重参数化模型里BN的存在是为了训练稳定性，但它会带来一个小隐患：当batch size比较小时，BN的均值和方差估计不稳，训练早期容易震荡。超分任务本身batch size就受显存限制，我遇到过batch size只有4的情况，训练损失曲线异常跳动。

解决办法有两个：一是使用较大的batch配合梯度累积；二是考虑用GroupNorm或LayerNorm替代BN，但要注意，主干中Mamba块内部已经有LayerNorm，两者作用范围不同，不要混淆。

我倾向于保留BN，宁可做梯度累积也要让batch跑到16。因为BN对轻量化模型的收敛速度帮助是实实在在的，折叠进卷积后推理时零成本，这个优势不能随便丢掉。

6.2 坑二：Mamba块在长序列上的显存峰值

前面提过Mamba是线性复杂度，但不要误以为它显存一定小。实际测试中发现，由于四方向扫描会产生四份序列特征，显存峰值集中在这块，特征图大的时候甚至会超过同规模Transformer的中间层开销。

缓解办法是在进入SS2D之前做一次降采样或通道压缩，反正Mamba块后面可以用降维再接回来。或者采用分组扫描策略，把通道分组独立扫描再合并。牺牲一点点感受野交互能力，换一个平稳的显存曲线，对部署很值。

6.3 建议的配置参考表

最后给一组我实测过比较稳的配置，你可以作为起点微调：

还有一个非常实际的经验：如果你最终部署的目标设备支持半精度推理，那么建议导出权重前就把模型转为半精度跑一遍完整测试。有些算子在半精度下表现不稳定，特别是融合了BN之后的小数值，可能产生肉眼可见的偏色。早发现问题早做补偿，别留到上线后让运维同学帮你填坑。

综合复现下来的感受，Rep-Mamba这条路是走得通的：遥感超分在保持质量的前提下，配合重参数化后推理结构确实清爽了很多，明显更适合往边缘端推。如果你正准备做一个轻量级超分模型，不妨从这条思路上切进去，重点先把重参数化块的合并流程跑通，再慢慢调Mamba块的数量和通道配置。做超分这件事，花里胡哨的结构最后都要落到“部署上线”这四个字上，能跑得动、跑得稳的模型才是好模型。