Mamba-PKD: A Framework for Efficient and Scalable Model Compression in Image Classification

Although deep neural networks (DNNs) have achieved success in image processing tasks, their computational complexity poses challenges in resource-constrained deployment. This paper integrates Mamba architecture within a Progressive Knowledge Distillation (PKD) pipeline to reduce complexity while maintaining accuracy in image classification. The framework distills a large teacher model into smaller Mamba student models, each trained using Selective-State-Space Models (S-SSM). We use the MNIST and CIFAR-10 datasets to demonstrate the approach's effectiveness. Seven students distill an MNIST teacher, retaining 63% of the teacher's FLOPs and matching its accuracy of 98%. The weakest student used 1% of the teacher's FLOPs to reach 72% accuracy. For CIFAR-10, the students achieve 1% lower accuracy than the teacher; the smallest retaining 5% of the teacher's FLOPs and achieving 50% accuracy. The framework shows its flexibility and scalability, demonstrating that Mamba blocks can serve as weak learners within PKD.