Optimising TinyML with Quantization and Distillation of Transformer and Mamba Models for Indoor Localisation on Edge Devices

📄 arXiv: 2412.09289v1 📥 PDF

作者: Thanaphon Suwannaphong, Ferdian Jovan, Ian Craddock, Ryan McConville

分类: cs.LG, cs.SE

发布日期: 2024-12-12

💡 一句话要点

针对边缘设备室内定位，提出Transformer和Mamba模型量化与蒸馏的TinyML优化方案

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: TinyML 室内定位 边缘计算 Transformer Mamba 量化 知识蒸馏

📋 核心要点

  1. 现有室内定位方案依赖集中式远程处理，存在功耗高、隐私泄露、延迟大等问题，难以满足边缘设备需求。
  2. 论文提出基于Transformer和Mamba模型的TinyML优化方案，利用量化和知识蒸馏等技术压缩模型，降低资源占用。
  3. 实验结果表明，量化Transformer模型在64KB RAM约束下表现良好，Mamba模型在32KB RAM约束下无需压缩即可工作。

📝 摘要（中文）

本文提出了一种小型高效的机器学习模型（TinyML），用于资源受限的边缘设备上的设备端室内定位。传统的室内定位方法依赖于集中式的远程数据处理。将数据处理转移到边缘设备本身具有诸多优势，包括延长电池寿命、增强隐私、减少延迟和降低运营成本，这些对于健康监测等常见应用至关重要。本文重点研究模型压缩技术，包括量化和知识蒸馏，以显著减小模型尺寸，同时保持较高的预测性能。该工作基于大型最先进的基于Transformer的模型，并力求将其部署在低功耗MCU中。此外，还提出了一种基于状态空间的Mamba架构，作为Transformer的更紧凑替代方案。结果表明，量化的Transformer模型在64 KB RAM约束下表现良好，在模型大小和定位精度之间取得了有效平衡。此外，紧凑的Mamba模型在更严格的约束（例如32 KB RAM）下也表现出强大的性能，无需模型压缩，使其成为资源受限环境的可行选择。通过该框架，可以将先进的室内定位模型部署到具有严格内存限制的低功耗MCU上。这些TinyML模型在医疗保健中的应用有潜力通过提供准确、实时的位置数据，同时最大限度地降低功耗、提高数据隐私、改善延迟和降低基础设施成本，从而彻底改变患者监测。

🔬 方法详解

问题定义：论文旨在解决资源受限的边缘设备上进行高精度室内定位的问题。现有方法通常依赖于将数据传输到云端进行集中处理，这导致了高延迟、高功耗以及潜在的隐私泄露风险。因此，如何在低功耗、低内存的边缘设备上部署复杂的室内定位模型是一个关键挑战。

核心思路：论文的核心思路是通过模型压缩技术，包括量化和知识蒸馏，来减小模型的尺寸，使其能够在资源受限的边缘设备上运行。同时，探索了Mamba这一新型状态空间模型，作为Transformer的更紧凑替代方案，以进一步降低模型复杂度。

技术框架：整体框架包括以下几个主要步骤：1) 使用大型Transformer模型作为教师模型；2) 对教师模型进行量化，降低模型精度，减小模型尺寸；3) 使用知识蒸馏技术，将教师模型的知识迁移到更小的学生模型（可以是量化后的Transformer或Mamba模型）；4) 在边缘设备上部署和评估压缩后的模型。

关键创新：论文的关键创新在于探索了Mamba模型在室内定位任务中的应用，并验证了其在极低资源约束下的可行性。此外，论文还系统地研究了量化和知识蒸馏等模型压缩技术在Transformer模型上的应用效果，并针对边缘设备进行了优化。

关键设计：论文中，量化采用的是训练后量化（Post-Training Quantization）方法，将模型权重和激活值转换为较低的精度（例如8位整数）。知识蒸馏采用的是logits蒸馏，即最小化学生模型和教师模型输出logits之间的差异。Mamba模型的具体结构和参数设置根据室内定位任务进行了调整，以达到最佳性能。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，量化后的Transformer模型在64KB RAM约束下能够实现较好的定位精度，而Mamba模型在32KB RAM约束下无需压缩即可达到可接受的性能。这表明，通过模型压缩和架构优化，可以将复杂的深度学习模型部署到资源极其有限的边缘设备上，为边缘智能应用开辟了新的可能性。

🎯 应用场景

该研究成果可广泛应用于医疗健康领域，例如患者实时定位与追踪、老人跌倒检测等。此外，还可应用于智能家居、工业自动化等场景，实现设备间的精确定位与协同。将定位功能部署在边缘设备上，能够有效保护用户隐私，降低网络延迟，并减少对中心服务器的依赖。

📄 摘要（原文）

This paper proposes small and efficient machine learning models (TinyML) for resource-constrained edge devices, specifically for on-device indoor localisation. Typical approaches for indoor localisation rely on centralised remote processing of data transmitted from lower powered devices such as wearables. However, there are several benefits for moving this to the edge device itself, including increased battery life, enhanced privacy, reduced latency and lowered operational costs, all of which are key for common applications such as health monitoring. The work focuses on model compression techniques, including quantization and knowledge distillation, to significantly reduce the model size while maintaining high predictive performance. We base our work on a large state-of-the-art transformer-based model and seek to deploy it within low-power MCUs. We also propose a state-space-based architecture using Mamba as a more compact alternative to the transformer. Our results show that the quantized transformer model performs well within a 64 KB RAM constraint, achieving an effective balance between model size and localisation precision. Additionally, the compact Mamba model has strong performance under even tighter constraints, such as a 32 KB of RAM, without the need for model compression, making it a viable option for more resource-limited environments. We demonstrate that, through our framework, it is feasible to deploy advanced indoor localisation models onto low-power MCUs with restricted memory limitations. The application of these TinyML models in healthcare has the potential to revolutionize patient monitoring by providing accurate, real-time location data while minimizing power consumption, increasing data privacy, improving latency and reducing infrastructure costs.