通过基于Mamba的异构知识蒸馏技术实现图像恢复模型的压缩

《Neural Networks》：Image Restoration Model Compression via Mamba-oriented Heterogeneous Knowledge Distillation

　　

Sai Yang|Bin Hu|Xiaoxin Wu|Fan Liu|Wanzhi Wen

摘要

基于变压器的图像恢复（IR）模型通常需要大量的计算资源和大量的参数。为了提高IR模型的效率，我们提出了一种用于图像恢复模型压缩的知识蒸馏策略，具体实现为面向Mamba的异构知识蒸馏（MHKD）或Transformer-to-Mamba架构对。首先，我们使用包含变压器块的复杂但高质量的IR模型作为教师网络进行预训练。然后，我们构建一个基于视觉Mamba块的轻量级IR模型作为学生网络。接下来，利用MHKD将教师网络中的异构知识传递给学生网络。MHKD的关键组成部分是特征过滤模块和特征接口模块：前者用于减少教师网络中的冗余信息，后者则通过设计的双向知识传递损失机制，将教师网络的异构特征转化为适合学生网络的形式。在多个IR基准测试中的广泛实验表明，在Transformer-to-Mamba蒸馏设置下，MHKD能够帮助基于Mamba的IR模型（参数量约为716K）实现与最先进模型相当甚至更好的性能。