Soft Knowledge Distillation with Multi-Dimensional Cross-Net Attention for Image Restoration Models Compression
This work addresses the problem of deploying efficient image restoration models for real-world applications, though it is incremental as it builds on existing knowledge distillation methods.
The paper tackled the high computational complexity of transformer-based image restoration models by proposing a Soft Knowledge Distillation strategy with a Multi-dimensional Cross-net Attention mechanism, which reduced FLOPs and parameters while maintaining strong performance on tasks like deraining, deblurring, and denoising.
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.