Image restoration model compression via mamba-oriented heterogeneous knowledge distillation.

Abstract

The transformer-based Image Restoration (IR) models usually require heavy computation and intensive parameters. To improve the efficiency of IR models, we propose a knowledge distillation strategy for image restoration model compression, specifically instantiated as Mamba-oriented Heterogeneous Knowledge Distillation (MHKD) or the Transformer-to-Mamba architecture pair. Specifically, we firstly pre-train a cumbersome yet high-quality IR model with transformer blocks as the teacher network. Secondly, we construct a lightweight IR model with the vision Mamba blocks as the student network. Subsequently, MHKD is used to transfer the heterogeneous knowledge from the teacher network to the student network. The key elements of MHKD are the feature filter module and feature interface module. The former is designed to reduce the redundant information from the teacher network, and the latter aims to translate the heterogeneous feature of teacher network for the student network with the designed bi-direction knowledge transferring loss. Extensive experiments across multiple IR benchmarks confirm that, for the Transformer-to-Mamba distillation setting, MHKD can assist a Mamba-based IR model with compact size ( ≈  716K parameters) to achieve comparable or even better performance than state-of-the-art models.