DRSformer·论文蒸馏笔记：可学习 Top-k 稀疏注意力去雨网络

DRSformer·论文蒸馏笔记可学习 Top-k 稀疏注意力去雨网络蒸馏对象Xiang Chen, Hao Li, Mingqiang Li, Jinshan Pan.《Learning A Sparse Transformer Network for Effective Image Deraining》CVPR 2023, pp. 5896–5905arXiv:2303.11950DOI 10.1109/CVPR52729.2023.00571作者机构南京理工大学陈翔/李昊/潘金山 中国电科信息科学研究院官方代码github.com/cschenxiang/DRSformer含预训练权重蒸馏日期2026-10-10 · 蒸馏方法DeepLens / 去雨系列同款正文 目录 封面流程谱系定位TPAMI 2025 统一评测综述独立训练赛道冠军一作陈翔即该综述一作——这篇是先有方法、后有评分规则的典型样本摘要本笔记蒸馏 CVPR 2023 去雨 Transformer DRSformer针对稠密自注意力聚合无关特征、抹平高频细节的缺陷提出可学习 top-k 稀疏注意力TKSA仅保留最关键通道相似度参与特征聚合辅以混合尺度前馈网络与混合专家特征补偿器五大基准全面刷新 SOTA独立训练赛道至今仍是冠军基线。目录摘要一、导读这篇论文在你的谱系里是什么位置二、动机稠密自注意力的三宗罪三、总体架构不换骨架只换器官四、TKSA可学习 Top-k 稀疏注意力核心创新4.1 计算流程4.2 关键细节k 不是超参数是可学习区间4.3 与两个近亲的区别论文 4.4 节4.4 消融证据五、MSFN混合尺度前馈网络六、MEFC混合专家特征补偿器七、实验五基准全 SOTA 与训练配方7.1 主结果Table 1PSNR/SSIMY 通道7.2 训练配方复现照抄用八、与你的复现/项目谱系的关系8.1 AutoDL 复现成本账3090 口径8.2 对张娟 MG-TSKD 的三个接口8.3 对 E0-E3 增量计划附录 A术语速查附录 B本地材料索引一、导读这篇论文在你的谱系里是什么位置先给结论链RestormerCVPR22通道转置注意力 │ 遗留问题softmax 稠密聚合无关 token 也掺和 ▼ DRSformerCVPR23陈翔/潘金山 ←── 本篇 │ top-k 稀疏化 混合尺度 FFN MoE 补偿器 ▼ TPAMI 2025 统一评测综述同一作者团队 │ 独立训练赛道冠军22 方法Table V ▼ HQ-RAIN / RE-RAIN合成冠军、真实域泛化弱 → 综述核心发现三个值得记住的点它至今仍是独立训练赛道的榜首方法TPAMI 2025 综述 Table V 平均分第一不是过气基线它是注意力手术刀式创新——不换骨架还是 U 形编解码 转置注意力只动注意力计算方式这种改动粒度最适合作为缝合素材论文自述的最大局限是效率33.7M 参数 / 242.9G FLOPs作者明确说未来要做剪枝/蒸馏——这正好是张娟 MG-TSKD 蒸馏项目可以对接的口子详见第八节。二、动机稠密自注意力的三宗罪论文的出发点是对标准自注意力的批判逻辑链条很干净无关 token 掺与聚合标准 Transformer 把所有 query-key 对的相似度都拿来做特征聚合。但 key 里的 token 并不总与 query 相关——不相关的相似度照样参与 softmax 加权噪声被硬塞进输出特征小权重被放大softmax 的稠密计算模式会放大幅度较小的相似度权重使特征交互过程对隐式噪声敏感论文引 explicit sparse transformer、NeurIPS 2021 稀疏化探索佐证高频信息被抹平注意 softmax(QK^T/λ) 矩阵每行非负且和为 1——拿它去乘 V 本质是对 V 做加权平均高频细节天然被平均掉输出趋于过平滑。与视觉证据一致论文 Fig. 1Uformer / Restormer / IDT 在细节与纹理恢复上都有缺口IDT 还引入边界伪影——纯 Transformer 去雨的共性问题。蒸馏注解第 3 条其实解释了为什么合成强、真实弱——过平滑的先验在合成数据上被 PSNR 奖励合成 GT 本身就偏干净到真实复杂纹理上就露馅。这与 TPAMI 综述的 RE-RAIN 发现是同一枚硬币的两面。三、总体架构不换骨架只换器官DRSformer 是层级化 U 形编解码器Restormer 同款骨架输入端3×3 卷积重叠 patch embedding下采样/上采样pixel-unshuffle / pixel-shuffle主干单元每层堆叠 Nᵢ 个稀疏 Transformer 块STB配置 {N₀, N₁, N₂, N₃, N₄} {4, 4, 6, 6, 8}四层注意力头数 {1, 2, 4, 8}初始通道 C48膨胀比 2STB 内部TKSA注意力 MSFN前馈Pre-LN 残差结构X′ₗ Xₗ₋₁ TKSA(LN(Xₗ₋₁))Xₗ X′ₗ MSFN(LN(X′ₗ))首尾两级各挂 N₀ 个MEFC混合专家特征补偿器做协同精修输出端全局残差 I_derain F(I_rain) I_rain训练损失只有纯 L1。蒸馏观点纯 L1 全局残差 U 形骨架说明性能增益几乎全部来自三个新器官TKSA / MSFN / MEFC消融归因干净——这是审稿人喜欢、复现者省心的设计。四、TKSA可学习 Top-k 稀疏注意力核心创新4.1 计算流程沿用 Restormer 的通道维转置注意力在 Ĉ×Ĉ 通道相似度矩阵上做注意力而非空间维复杂度从 O(HW)² 降到 O(C²)但插入一步关键手术SparseAtt ( Q , K , V ) softmax ( T k ( Q K ⊤ / λ ) ) V \text{SparseAtt}(Q,K,V) \text{softmax}\big(T_k(QK^\top/\lambda)\big)VSparseAtt(Q,K,V)softmax(Tk​(QK⊤/λ))V其中 top-k 选择算子[ T k ( S ) ] i j { S i j , S i j ∈ top-k(row j ) 0 , otherwise [T_k(S)]_{ij} \begin{cases} S_{ij}, S_{ij} \in \text{top-k(row } j) \\ 0, \text{otherwise} \end{cases}[Tk​(S)]ij​{Sij​,0,​Sij​∈top-k(rowj)otherwise​即对相似度矩阵每一行只保留 top-k 个分数进 softmax其余 scatter 置 0。不是 dropout 式随机丢弃而是按贡献度自适应筛选。4.2 关键细节k 不是超参数是可学习区间k 以适当分数的加权平均动态确定取值被约束在区间[Δ₁, Δ₂] [1/2, 4/5]内消融实测Rain200Hk [1/2, 4/5] 时 32.18 dB 最优k 太小 → 全局信息聚合不足性能断崖式下跌k 太大 → 无关无用特征混入性能缓降固定单值 k如 1/2对取值敏感区间化 可学习最稳。4.3 与两个近亲的区别论文 4.4 节方法稀疏维度k 的性质代价KiTCVPR 2022空间维k-NN 局部注意力固定局部敏感哈希复杂全局交互不足KVTECCV 2022空间维 top-k token固定空间维开销大DRSformer通道维top-k可学习区间通道注意力本身 O(C²)top-k 几乎白送分类学上KiT/KVT 一类属于数据驱动稀疏data-basedDRSformer 属于内容驱动稀疏content-based——每个 query 按相似度内容自适应选 key。4.4 消融证据去掉 top-k退化为标准转置注意力各基准 PSNR 全线下降Fig. 6高通滤波HPF可视化Fig. 8带 top-k 的特征图高频细节明显更完整——佐证第二节加权平均抹高频的诊断机制解释邻近像素天然相似top-k 剔除的正是长程依赖中的低相似度交互等于给注意力加了一个近邻偏置但不损失全局挖掘能力。五、MSFN混合尺度前馈网络动机Uformer/Restormer/IDT 的前馈网络只用单尺度深度卷积补局部性忽略多尺度雨纹的相关性MSPFN 早就证明多尺度对去雨有效。结构纯并行双分支 交叉拼接X → LN → 1×1conv 升维膨胀比 r2.66 ├─ 3×3 DWConv ──┐ │ [拼接] → 3×3 DWConv ──┐ └─ 5×5 DWConv ──┘ [拼接] → 1×1conv 残差 └────────5×5 DWConv ──────┘两条分支各跑两轮第二轮时交换拼接对方的输出3×3 分支吃 [3×3 输出, 5×5 输出]5×5 分支吃 [5×5 输出, 3×3 输出]实现跨尺度信息流动。消融Rain200HTable 3前馈网络PSNR / SSIMFN普通 MLP31.84 / 0.9275DFN单尺度 DWConv31.88 / 0.9279GDFNRestormer 门控双流31.97 / 0.9286MSFN本文32.18 / 0.9330对 GDFN 的增益 0.21 dB——零新算子全是 DWConv 和拼接纯结构设计换来这是模块缝合性价比最高的一类素材。六、MEFC混合专家特征补偿器动机雨分布揭示退化位置与程度数据稀疏性STB 只挖掘了内容稀疏——MEFC 来补数据稀疏实现数据 × 内容双重稀疏协同。设计与经典 MoE 的两点不同专家池8 个异构专家——平均池化3×3 感受野、可分离卷积1×1/3×3/5×5/7×7、膨胀卷积3×3/5×5/7×7感受野和算子类型都不同不同于同构专家堆叠无外部门控网络用自注意力充当专家切换器——通道平均池化得描述子 z_c经 W₁/W₂ 生成各专家系数加权融合 8 个专家输出各自 zero-pad 对齐尺寸后 1×1 卷积 残差。消融Rain200HTable 4基线 32.03 → MEFC-1 → 32.07 → MEFC-2 8 专家 →32.18。单专家32.06明显弱于多专家——异构感受野的多样性是增益来源。工程注意Rain200L 和 SPA-Data 训练时不挂 MEFC——雨纹简单时 MEFC 反而多余。复现时照抄这个配置。七、实验五基准全 SOTA 与训练配方7.1 主结果Table 1PSNR/SSIMY 通道数据集RestormerIDTDRSformerRain200L40.99 / 0.989040.74 / 0.988441.23 / 0.9894Rain200H32.00 / 0.932932.10 / 0.934432.18 / 0.9330DID-Data35.29 / 0.964134.89 / 0.962335.38 / 0.9647DDN-Data34.20 / 0.957133.84 / 0.954934.36 / 0.9590SPA-Data真实47.98 / 0.992147.35 / 0.993048.53 / 0.9924平均领先同期 IDT 约0.4 dB对 CNN 系MPRNet/SPDNet/RCDNet优势更大真实数据 SPA-Data 与无 GT 的 Internet-DataNIQE 4.095 / BRISQUE 22.730 双最低也拿第一——注意这是论文自己测的TPAMI 综述统一口径重测后 RE-RAIN 泛化并不占优两份证据合起来读才完整。7.2 训练配方复现照抄用项配置优化器AdamWbatch / patch8 / 128×128迭代300K前 92K 恒定 lr余下 cosine 退火学习率1×10⁻⁴ → 1×10⁻⁶增强随机水平 垂直翻转硬件4× RTX 3090端到端无预训练MEFCRain200L / SPA-Data 不用其余数据集用指标Y 通道 YCbCr与本篇 TPAMI 综述统一后的口径一致八、与你的复现/项目谱系的关系8.1 AutoDL 复现成本账3090 口径方案内容预估成本L0 推理复现官方预训练权重直接测五基准¥2-5几小时L1 单卡全量复训1×3090batch 8显存约 20G 内可容300K iter 约为论文 4 卡时长 ×3-4 → 12-18 天¥300-500L2 砍迭代复训150K iter batch 4经验损失 ≤0.1 dB¥120-200L3 仅 Rain200H 验证单数据集复训做 TKSA 消融复现¥40-80建议路线先 L0 验证环境与权重再按需上 L2——与 VDMamba 的先推理后训练套路一致。8.2 对zj MG-TSKD 的三个接口教师候选DRSformer 是独立赛道冠军且官方有权重比 MPRNet 教师更强论文 Limitations 原文写着will apply the pruning or distillation scheme——作者自己点名要做蒸馏你们等于替他把这条路走掉TKSA 直接可摘top-k 算子是即插即用的注意力替换件加到学生网络不增算子类型k 区间 [1/2, 4/5] 可直接沿用MSFN 当缝合素材对 GDFN 0.21 dB、零新算子是学生网络前馈层的低成本升级件。8.3 对 E0-E3 增量计划E0 跨域矩阵若把 DRSformer 权重拉进来三权重之外加一个可直接复现 TPAMI 综述 Table VI 的合成冠军 vs 真实泛化分歧在你自己机器上的表现——DRSformer 正是那张表里的关键样本。附录 A术语速查术语含义STBSparse Transformer BlockTKSA MSFN 的基础单元TKSATop-k Sparse Attention通道维转置注意力 可学习 top-k 筛选MSFNMixed-Scale Feed-forward Network3×3/5×5 双分支交叉拼接前馈MEFCMixture of Experts Feature Compensator8 异构专家 自注意力门控[Δ₁, Δ₂][1/2, 4/5]top-k 保留比例的可学习区间转置注意力Restormer 技巧在通道维C×C而非空间维做注意力复杂度 O(C²)content-based sparsity按内容相似度动态稀疏vs>附录 B本地材料索引论文 PDFH:\zj\视频去雨综述论文\DRSformer_CVPR2023.pdfarXiv:2303.11950v110 页全文抽取drsformer_extracted.txt本工作区47k 字符官方代码github.com/cschenxiang/DRSformerbasicsr 框架含 pretrained_models / train.sh / test.py姊妹篇笔记单幅图像去雨综述_论文蒸馏笔记.md本工作区DRSformer 在其中 Table V/VI 的榜单位置