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

TL;DR

This paper introduces a novel soft knowledge distillation method with multi-dimensional cross-net attention for compressing image restoration models, effectively reducing complexity while preserving high restoration quality.

Contribution

It proposes a multi-dimensional attention mechanism and a contrastive loss to improve knowledge distillation in image restoration models, addressing implicit attention relationships.

Findings

Significant reduction in FLOPs and parameters.

Maintained high restoration performance across tasks.

Enhanced feature learning stability and efficiency.

Abstract

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.