arXiv:2501.09321cs.CV2025-01中稿 · ICASSP2025被引 2

用多维注意力机制让小模型学大模型的隐含注意力,压缩图像修复模型。

Soft Knowledge Distillation with Multi-Dimensional Cross-Net Attention for Image Restoration Models Compression

  • 引入跨网络多维注意力,让师生模型在通道和空间维度互动。
  • 在三个任务上减少计算量,同时保持修复效果领先。
  • 适合想轻量化图像修复模型的研究者或工程师。

基于Transformer的编码器-解码器模型在图像到图像转换任务中表现优异,尤其在图像修复领域。然而,其高计算复杂度(表现为更高的浮点运算次数FLOPs和参数量)限制了实际应用。现有图像修复中的知识蒸馏方法通常采用轻量级学生模型直接模仿教师模型的中间特征和重建结果,忽略了它们之间的隐含注意力关系。为此,我们提出一种软知识蒸馏(SKD)策略,引入多维交叉网络注意力(MCA)机制,实现学生与教师在通道和空间维度上的交互,使学生能隐式学习注意力矩阵。此外,采用高斯核函数在核空间中衡量师生特征距离,确保特征学习稳定高效。为进一步提升重建质量,我们在图像层面用对比学习损失替代常用的L1或KL散度损失。在图像去雨、去模糊和去噪三个任务上的实验表明,所提SKD策略显著降低计算复杂度,同时保持强大的图像修复能力。

原文摘要 · Abstract (English)

Transformer-based encoder-decoder models have achieved remarkable success in image-to-image transfer tasks, particularly in image restoration. However, their high computational complexity-manifested in elevated FLOPs and parameter counts-limits their application in real-world scenarios. Existing knowledge distillation methods in image restoration typically employ lightweight student models that directly mimic the intermediate features and reconstruction results of the teacher, overlooking the implicit attention relationships between them. To address this, we propose a Soft Knowledge Distillation (SKD) strategy that incorporates a Multi-dimensional Cross-net Attention (MCA) mechanism for compressing image restoration models. This mechanism facilitates interaction between the student and teacher across both channel and spatial dimensions, enabling the student to implicitly learn the attention matrices. Additionally, we employ a Gaussian kernel function to measure the distance between student and teacher features in kernel space, ensuring stable and efficient feature learning. To further enhance the quality of reconstructed images, we replace the commonly used L1 or KL divergence loss with a contrastive learning loss at the image level. Experiments on three tasks-image deraining, deblurring, and denoising-demonstrate that our SKD strategy significantly reduces computational complexity while maintaining strong image restoration capabilities.

图像修复知识蒸馏注意力机制模型压缩

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。