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

蒸馏对象 ：Xiang Chen, Hao Li, Mingqiang Li, Jinshan Pan.《Learning A Sparse Transformer Network for Effective Image Deraining》（CVPR 2023, pp. 5896--5905，arXiv:2303.11950，DOI 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 稀疏注意力（核心创新）](#四、TKSA：可学习 Top-k 稀疏注意力（核心创新）)
- [4.1 计算流程](#4.1 计算流程)
- [4.2 关键细节：k 不是超参数，是可学习区间](#4.2 关键细节：k 不是超参数，是可学习区间)
- [4.3 与两个近亲的区别（论文 4.4 节）](#4.3 与两个近亲的区别（论文 4.4 节）)
- [4.4 消融证据](#4.4 消融证据)
五、MSFN：混合尺度前馈网络
六、MEFC：混合专家特征补偿器
[七、实验：五基准全 SOTA 与训练配方](#七、实验：五基准全 SOTA 与训练配方)
- [7.1 主结果（Table 1，PSNR/SSIM，Y 通道）](#7.1 主结果（Table 1，PSNR/SSIM，Y 通道）)
- [7.2 训练配方（复现照抄用）](#7.2 训练配方（复现照抄用）)
八、与你的复现/项目谱系的关系
- [8.1 AutoDL 复现成本账（3090 口径）](#8.1 AutoDL 复现成本账（3090 口径）)
- [8.2 对张娟 MG-TSKD 的三个接口](#8.2 对张娟 MG-TSKD 的三个接口)
- [8.3 对 E0-E3 增量计划](#8.3 对 E0-E3 增量计划)
[附录 A：术语速查](#附录 A：术语速查)
[附录 B：本地材料索引](#附录 B：本地材料索引)

一、导读：这篇论文在你的谱系里是什么位置

先给结论链：

复制代码

Restormer（CVPR22，通道转置注意力）
      │  遗留问题：softmax 稠密聚合，无关 token 也掺和
      ▼
DRSformer（CVPR23，陈翔/潘金山）  ←── 本篇
      │  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. 1）：Uformer / 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}，初始通道 C=48，膨胀比 2；
STB 内部：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 的通道维转置注意力（在 Ĉ×Ĉ 通道相似度矩阵上做注意力，而非空间维，复杂度从 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)V SparseAtt(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(row j)otherwise

即：对相似度矩阵每一行，只保留 top-k 个分数进 softmax，其余 scatter 置 0。不是 dropout 式随机丢弃，而是按贡献度自适应筛选。

4.2 关键细节：k 不是超参数，是可学习区间

k 以适当分数的加权平均动态确定，取值被约束在区间 $Δ₁, Δ₂$ = $1/2, 4/5$ 内；
消融实测（Rain200H）：k = $1/2, 4/5$ 时 32.18 dB 最优；
k 太小 → 全局信息聚合不足，性能断崖式下跌；k 太大 → 无关无用特征混入，性能缓降；
固定单值 k（如 1/2）对取值敏感，区间化 + 可学习最稳。

4.3 与两个近亲的区别（论文 4.4 节）

方法	稀疏维度	k 的性质	代价
KiT（CVPR 2022）	空间维，k-NN 局部注意力	固定	局部敏感哈希复杂，全局交互不足
KVT（ECCV 2022）	空间维 top-k token	固定	空间维开销大
DRSformer	通道维 top-k	可学习区间	通道注意力本身 O(C²)，top-k 几乎白送

分类学上：KiT/KVT 一类属于数据驱动稀疏（data-based），DRSformer 属于内容驱动稀疏（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 升维（膨胀比 r=2.66）
        ├─ 3×3 DWConv ──┐
        │               [拼接] → 3×3 DWConv ──┐
        └─ 5×5 DWConv ──┘                     [拼接] → 1×1conv + 残差
                    └────────5×5 DWConv ──────┘

两条分支各跑两轮，第二轮时交换拼接对方的输出（3×3 分支吃 $3\times3 输出, 5\times5 输出$ ，5×5 分支吃 $5\times5 输出, 3\times3 输出$ ），实现跨尺度信息流动。

消融（Rain200H，Table 3）：

前馈网络	PSNR / SSIM
FN（普通 MLP）	31.84 / 0.9275
DFN（单尺度 DWConv）	31.88 / 0.9279
GDFN（Restormer 门控双流）	31.97 / 0.9286
MSFN（本文）	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 卷积 + 残差。

消融（Rain200H，Table 4）：基线 32.03 → +MEFC-1 → 32.07 → +MEFC-2 + 8 专家 → 32.18。单专家（32.06）明显弱于多专家------异构感受野的多样性是增益来源。

工程注意 ：Rain200L 和 SPA-Data 训练时不挂 MEFC------雨纹简单时 MEFC 反而多余。复现时照抄这个配置。

七、实验：五基准全 SOTA 与训练配方

7.1 主结果（Table 1，PSNR/SSIM，Y 通道）

数据集	Restormer	IDT	DRSformer
Rain200L	40.99 / 0.9890	40.74 / 0.9884	41.23 / 0.9894
Rain200H	32.00 / 0.9329	32.10 / 0.9344	32.18 / 0.9330
DID-Data	35.29 / 0.9641	34.89 / 0.9623	35.38 / 0.9647
DDN-Data	34.20 / 0.9571	33.84 / 0.9549	34.36 / 0.9590
SPA-Data（真实）	47.98 / 0.9921	47.35 / 0.9930	48.53 / 0.9924

平均领先同期 IDT 约 +0.4 dB；对 CNN 系（MPRNet/SPDNet/RCDNet）优势更大；
真实数据 SPA-Data 与无 GT 的 Internet-Data（NIQE 4.095 / BRISQUE 22.730 双最低）也拿第一------注意：这是论文自己测的，TPAMI 综述统一口径重测后 RE-RAIN 泛化并不占优，两份证据合起来读才完整。

7.2 训练配方（复现照抄用）

项	配置
优化器	AdamW
batch / patch	8 / 128×128
迭代	300K（前 92K 恒定 lr，余下 cosine 退火）
学习率	1×10⁻⁴ → 1×10⁻⁶
增强	随机水平 + 垂直翻转
硬件	4× RTX 3090，端到端，无预训练
MEFC	Rain200L / SPA-Data 不用，其余数据集用
指标	Y 通道 YCbCr（与本篇 TPAMI 综述统一后的口径一致）

八、与你的复现/项目谱系的关系

8.1 AutoDL 复现成本账（3090 口径）

方案	内容	预估成本
L0 推理复现	官方预训练权重直接测五基准	¥2-5（几小时）
L1 单卡全量复训	1×3090，batch 8（显存约 20G 内可容），300K iter 约为论文 4 卡时长 ×3-4 → 12-18 天	¥300-500
L2 砍迭代复训	150K iter + batch 4（经验损失 ≤0.1 dB）	¥120-200
L3 仅 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：术语速查

术语	含义
STB	Sparse Transformer Block：TKSA + MSFN 的基础单元
TKSA	Top-k Sparse Attention：通道维转置注意力 + 可学习 top-k 筛选
MSFN	Mixed-Scale Feed-forward Network：3×3/5×5 双分支交叉拼接前馈
MEFC	Mixture of Experts Feature Compensator：8 异构专家 + 自注意力门控
$Δ₁, Δ₂$ = $1/2, 4/5$	top-k 保留比例的可学习区间
转置注意力	Restormer 技巧：在通道维（C×C）而非空间维做注意力，复杂度 O(C²)
content-based sparsity	按内容相似度动态稀疏（vs data-based 固定模式稀疏）
KiT / KVT	近亲方法：空间维 k-NN / top-k 注意力（对比见 4.3）

附录 B：本地材料索引

论文 PDF：H:\zj\视频去雨综述论文\DRSformer_CVPR2023.pdf（arXiv:2303.11950v1，10 页）
全文抽取：drsformer_extracted.txt（本工作区，47k 字符）
官方代码：github.com/cschenxiang/DRSformer（basicsr 框架，含 pretrained_models / train.sh / test.py）
姊妹篇笔记：单幅图像去雨综述_论文蒸馏笔记.md（本工作区；DRSformer 在其中 Table V/VI 的榜单位置）