Soft Knowledge Distillation with Multi-Dimensional Cross-Net Attention for Image Restoration Models Compression
作者: Yongheng Zhang, Danfeng Yan
分类: cs.CV
发布日期: 2025-01-16
备注: Accepted by ICASSP2025
💡 一句话要点
提出基于多维交叉注意力软知识蒸馏的图像修复模型压缩方法
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 知识蒸馏 图像修复 模型压缩 注意力机制 Transformer 对比学习
📋 核心要点
- 现有图像修复知识蒸馏方法忽略了教师和学生模型间隐式注意力关系的利用，导致蒸馏效果受限。
- 提出软知识蒸馏策略，通过多维交叉注意力机制，使学生模型能够隐式学习教师模型的注意力信息。
- 实验表明，该方法在图像去雨、去模糊和去噪任务上，显著降低计算复杂度的同时保持了图像修复能力。
📝 摘要（中文）
基于Transformer的编码器-解码器模型在图像到图像转换任务，特别是图像修复中取得了显著成功。然而，它们的高计算复杂度（表现为高FLOPs和参数量）限制了其在实际场景中的应用。现有的图像修复知识蒸馏方法通常采用轻量级学生模型，直接模仿教师模型的中间特征和重建结果，忽略了它们之间的隐式注意力关系。为了解决这个问题，我们提出了一种软知识蒸馏（SKD）策略，该策略结合了多维交叉网络注意力（MCA）机制来压缩图像修复模型。该机制促进了学生和教师在通道和空间维度上的交互，使学生能够隐式地学习注意力矩阵。此外，我们采用高斯核函数来测量核空间中学生和教师特征之间的距离，确保稳定和高效的特征学习。为了进一步提高重建图像的质量，我们用图像级别的对比学习损失取代了常用的L1或KL散度损失。在图像去雨、去模糊和去噪三个任务上的实验表明，我们的SKD策略在保持强大的图像修复能力的同时，显著降低了计算复杂度。
🔬 方法详解
问题定义：论文旨在解决基于Transformer的图像修复模型计算复杂度高，难以在实际场景中应用的问题。现有知识蒸馏方法忽略了教师模型和学生模型之间的注意力关系，导致学生模型无法充分学习教师模型的知识，修复效果不佳。
核心思路：论文的核心思路是利用软知识蒸馏，通过多维交叉注意力机制，让学生模型学习教师模型的注意力信息。同时，使用高斯核函数度量特征距离，并采用对比学习损失来提升重建图像质量。这样可以在保证修复效果的同时，降低模型的计算复杂度。
技术框架：整体框架包含一个教师模型和一个学生模型。首先，通过多维交叉注意力（MCA）模块，计算教师模型和学生模型在通道和空间维度上的注意力矩阵，实现二者之间的信息交互。然后，使用高斯核函数计算学生和教师特征在核空间的距离，作为特征蒸馏的损失函数。最后，使用对比学习损失来优化重建图像的质量。
关键创新：论文的关键创新在于提出了多维交叉注意力（MCA）机制，该机制能够有效地捕捉教师模型和学生模型之间的注意力关系，并将其传递给学生模型。此外，使用高斯核函数度量特征距离，以及采用对比学习损失，都有助于提升蒸馏效果和重建图像质量。
关键设计：MCA模块的设计是关键。它分别在通道和空间维度上计算注意力矩阵，然后将这些注意力矩阵用于指导学生模型的学习。高斯核函数的参数选择以及对比学习损失的温度系数等超参数的设置，也会影响最终的蒸馏效果。具体网络结构的选择（例如，Transformer的具体层数和维度）也需要根据具体任务进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，提出的SKD策略在图像去雨、去模糊和去噪三个任务上均取得了显著的性能提升。与现有知识蒸馏方法相比，该方法能够在保持甚至提升图像修复质量的同时，显著降低模型的FLOPs和参数量。具体的性能数据（例如PSNR和SSIM指标）以及与基线方法的对比结果（具体数值未知）进一步验证了该方法的有效性。
🎯 应用场景
该研究成果可应用于各种图像修复场景，例如移动设备上的图像去噪、去模糊，以及视频监控中的图像修复等。通过降低模型的计算复杂度，可以使其更容易部署在资源受限的设备上，具有广阔的应用前景。未来，该方法还可以扩展到其他图像到图像的转换任务中。
📄 摘要（原文）
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.