优化微型 Transformer 在低功耗 MCU 上的部署
摘要
Transformer 网络在 NLP 和 CV 等许多领域正迅速成为 SotA。与 CNN 类似，人们正大力推动将 Transformer 模型部署在极端边缘端，最终使其适应 MCU 极小的功耗预算和内存占用。然而，该方向的早期方法大多是临时的，且特定于平台和模型。本研究旨在实现并优化编码器微型 Transformer 在商用 MCU 上的灵活、多平台部署。我们提出了一个完整的框架，以将 Transformer 模型端到端地部署到单核和多核 MCU 上。我们的框架提供了一个优化的核库，以最大化数据复用，并避免在关键的注意力块中出现不必要的数据调度操作。我们引入了一种名为 Fused-Weight Self-Attention 的新型 MHSA 推理调度，通过离线融合线性投影权重来进一步减少操作和参数量。此外，为缓解计算注意力图所达到的内存峰值，我们提出了一种用于 MHSA 的深度优先分块方案。我们利用 ARM 和 RISC-V ISA 在三个不同的 MCU 类别上评估了我们的框架，即 STM32H7、STM32L4 和 GAP9 (RV32IMC-XpulpV2)。与 SotA 库 CMSIS-NN (ARM) 和 PULP-NN (RISC-V) 相比，我们分别实现了平均 4.79 倍和 2.0 倍的更低延迟。此外，我们表明我们的 MHSA 深度优先分块方案可将内存峰值降低高达 6.19 倍，而融合权重注意力可将运行时间减少 1.53 倍，参数量减少 25%。我们报告了在多个微型 Transformer 上的显著改进：例如，当在 GAP9 上执行用于基于雷达的手势识别任务的 transformer 块时，我们实现了 0.14ms 的延迟和 4.92 微焦的能量消耗，比同一平台上的 SotA PULP-NN 库低 2.32 倍。
引用
@article{arxiv.2404.02945,
title = {Optimizing the Deployment of Tiny Transformers on Low-Power MCUs},
author = {Victor J. B. Jung and Alessio Burrello and Moritz Scherer and Francesco Conti and Luca Benini},
journal= {arXiv preprint arXiv:2404.02945},
year = {2024}
}
备注
Pre-print manuscript submitted for review to the IEEE Transactions on Computers