arXiv:2503.01727cs.LG2025-03

用Mamba架构在知识蒸馏中实现高效模型压缩,保持高精度。

Mamba base PKD for efficient knowledge compression

  • 将Mamba块嵌入渐进式知识蒸馏,逐步压缩教师模型
  • 学生模型仅需1%计算量仍达72%准确率,比教师低26%
  • 适合边缘设备部署,兼具灵活性与可扩展性

深度神经网络在图像处理任务中表现卓越,但其庞大体积和高计算复杂度限制了在资源受限环境中的部署。本文提出一种新方法,将Mamba架构融入渐进式知识蒸馏(PKD)过程,实现图像分类任务中模型复杂度降低的同时保持精度。该框架通过一系列逐步缩小的学生模型,由大教师模型进行知识迁移,学生模型采用基于选择性状态空间模型(S-SSM)的Mamba块设计,聚焦关键输入特征,减少计算开销。初步实验使用MNIST和CIFAR-10数据集验证有效性:在MNIST上,教师模型达98%准确率;一组七名学生共保留63%的教师模型浮点运算量(FLOPs),性能接近教师(98%准确率);最弱学生仅用1%的FLOPs,仍保持72%准确率。在CIFAR-10上,学生模型较教师低1%准确率,最小学生仅消耗5%的教师模型FLOPs,达到50%准确率。结果表明Mamba架构可灵活集成于PKD,有效生成弱学生模型,为实时应用提供低计算成本解决方案。

原文摘要 · Abstract (English)

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.

模型压缩Mamba知识蒸馏

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。