Optimizing Swin Transformer UNet for Image Denoising via Knowledge Distillation

Abstract

A lightweight image denoising model, namely SUNet - KD (Swin-UNet integrated with Knowledge Distillation), is put forward within this paper. Knowledge distillation serves as the foundation for this model. Through the establishment of a teacher-student architecture, knowledge transference occurs from the pre-trained SUNet teacher model to a streamlined student model. Introduction of a multi-granularity distillation loss function takes place. This function takes into consideration hard label alignment at the output layer, soft matching of probability distributions, as well as semantic transfer of intermediate features. On the DIV2K dataset, experiments show. SUNet - KD attains a PSNR retention rate of 94.95% and an SSIM retention rate of 97.09%. Parameters are decreased by around 76.5% simultaneously. Moreover, an inference speed 1.82 times faster is provided. An excellent equilibrium between accuracy and efficiency is achieved thereby. For high-quality image denoising in resource-constrained environments, a practical solution is thus furnished.