让轻量BERT在普通单片机上运行,支持超长文本
MCUBERT: Memory-Efficient BERT Inference on Commodity Microcontrollers
- 通过嵌入层压缩与调度优化,降低模型存储和内存占用
- 在不足256KB内存下处理超过512个词元,参数量减少3至5.7倍
- 适合资源受限设备部署自然语言模型的开发者参考
本文提出MCUBERT,通过网络与调度协同优化,实现BERT等语言模型在小型微控制器(MCU)上的部署。研究发现嵌入表是小型BERT模型的主要存储瓶颈。为此,提出基于聚类低秩近似的两阶段神经架构搜索算法,用于嵌入压缩;同时设计细粒度的MCU友好调度策略,通过计算分块、重排序及内核优化,显著提升输入序列长度支持能力,且无延迟或精度损失。MCUBERT将BERT-tiny和BERT-mini的参数量分别减少5.7×和3.0×,执行内存降低3.5×和4.3×,延迟减少1.5×。首次在通用MCU上实现轻量BERT运行,可在少于256KB内存下处理超过512个词元。
原文摘要 · Abstract (English)
In this paper, we propose MCUBERT to enable language models like BERT on tiny microcontroller units (MCUs) through network and scheduling co-optimization. We observe the embedding table contributes to the major storage bottleneck for tiny BERT models. Hence, at the network level, we propose an MCU-aware two-stage neural architecture search algorithm based on clustered low-rank approximation for embedding compression. To reduce the inference memory requirements, we further propose a novel fine-grained MCU-friendly scheduling strategy. Through careful computation tiling and re-ordering as well as kernel design, we drastically increase the input sequence lengths supported on MCUs without any latency or accuracy penalty. MCUBERT reduces the parameter size of BERT-tiny and BERT-mini by 5.7$\times$ and 3.0$\times$ and the execution memory by 3.5$\times$ and 4.3$\times$, respectively. MCUBERT also achieves 1.5$\times$ latency reduction. For the first time, MCUBERT enables lightweight BERT models on commodity MCUs and processing more than 512 tokens with less than 256KB of memory.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。