Mamba base PKD for efficient knowledge compression
Journal:
arXiv
Published Date:
Mar 3, 2025
Abstract
Deep neural networks (DNNs) have remarkably succeeded in various image
processing tasks. However, their large size and computational complexity
present significant challenges for deploying them in resource-constrained
environments. This paper presents an innovative approach for integrating Mamba
Architecture within a Progressive Knowledge Distillation (PKD) process to
address the challenge of reducing model complexity while maintaining accuracy
in image classification tasks. The proposed framework distills a large teacher
model into progressively smaller student models, designed using Mamba blocks.
Each student model is trained using Selective-State-Space Models (S-SSM) within
the Mamba blocks, focusing on important input aspects while reducing
computational complexity. The work's preliminary experiments use MNIST and
CIFAR-10 as datasets to demonstrate the effectiveness of this approach. For
MNIST, the teacher model achieves 98% accuracy. A set of seven student models
as a group retained 63% of the teacher's FLOPs, approximating the teacher's
performance with 98% accuracy. The weak student used only 1% of the teacher's
FLOPs and maintained 72% accuracy. Similarly, for CIFAR-10, the students
achieved 1% less accuracy compared to the teacher, with the small student
retaining 5% of the teacher's FLOPs to achieve 50% accuracy. These results
confirm the flexibility and scalability of Mamba Architecture, which can be
integrated into PKD, succeeding in the process of finding students as weak
learners. The framework provides a solution for deploying complex neural
networks in real-time applications with a reduction in computational cost.