[论文解读] Real-Time Execution of Large-scale Language Models on Mobile

Wei Niu, Zhenglun Kong+8 more|arXiv (Cornell University)|Sep 15, 2020

Topic Modeling参考文献 22被引用 10

一句话总结

本文提出了一种编译器感知的神经架构搜索（NAS）框架，通过将编译器级优化集成到NAS循环中，协同优化移动设备上BERT模型的准确率与推理延迟。该框架在移动GPU上相比TensorFlow Lite实现了高达7.8倍的加速，且仅损失0.5–2%的准确率，使在消费级移动硬件上实现大规模Transformer模型的实时推理成为可能。

ABSTRACT

Pre-trained large-scale language models have increasingly demonstrated high accuracy on many natural language processing (NLP) tasks. However, the limited weight storage and computational speed on hardware platforms have impeded the popularity of pre-trained models, especially in the era of edge computing. In this paper, we seek to find the best model structure of BERT for a given computation size to match specific devices. We propose the first compiler-aware neural architecture optimization framework. Our framework can guarantee the identified model to meet both resource and real-time specifications of mobile devices, thus achieving real-time execution of large transformer-based models like BERT variants. We evaluate our model on several NLP tasks, achieving competitive results on well-known benchmarks with lower latency on mobile devices. Specifically, our model is 5.2x faster on CPU and 4.1x faster on GPU with 0.5-2% accuracy loss compared with BERT-base. Our overall framework achieves up to 7.8x speedup compared with TensorFlow-Lite with only minor accuracy loss.

研究动机与目标

• 为解决在资源受限的移动设备上以实时延迟要求部署大规模预训练语言模型（如BERT）的挑战。
• 弥合硬件感知神经架构搜索与实际编译器优化之间的差距，后者常导致深层模型的NAS结果失效。
• 协同优化模型架构与编译器级优化（如层融合），以实现在移动CPU和GPU上的低延迟、高准确率推理。
• 使BERT变体在移动CPU和GPU上均可实现实时执行——此前主流框架（如TensorFlow Lite）尚不支持此功能。

提出的方法

• 提出一种将编译器优化集成到NAS循环中的编译器感知神经架构搜索（NAS）框架，确保生成的模型满足实时性和资源约束。
• 采用控制器-训练器范式，基于GLUE基准的微调快速评估，搜索在准确率与延迟之间最优的BERT架构。
• 应用一种新颖的层融合技术，以减少中间张量内存访问与计算开销，显著提升CPU与GPU上的性能。
• 采用两阶段训练流程：首先训练一个深度且狭窄的BERT变体（28个块），然后通过知识蒸馏将其蒸馏为更高效的“学生”模型。
• 通过分析模型结构并应用设备特定优化（包括融合层执行）的编译器流水线，生成优化后的推理代码。
• 使用HuggingFace和自定义编译流水线，在真实移动设备（如三星Galaxy S20、S10）上评估模型，以测量实际延迟与准确率。

实验结果

研究问题

• RQ1 编译器感知的神经架构搜索能否协同优化模型结构与编译过程，实现在移动设备上大规模BERT模型的实时推理？
• RQ2 层融合是否显著降低移动CPU与GPU上的推理延迟，尤其是在深层Transformer模型上？
• RQ3 包含编译器反馈的NAS框架能否在延迟与准确率权衡上优于传统硬件感知NAS或现有推理框架（如TFLite）？
• RQ4 BERT变体是否可在移动GPU上实现实时推理？其性能与CPU执行及TFLite相比如何？

主要发现

• 所提框架在移动GPU上相比TensorFlow Lite实现了高达7.8倍的加速；使用NAS的BERT(ours)在CPU上比BERT BASE快5.2倍，在GPU上快4.1倍。
• 仅损失0.5–2%的准确率，优化后的BERT模型在GLUE基准测试中，延迟与准确率均优于BERT BASE与MobileBERT。
• 层融合减少了中间内存流量，加速了执行，在GPU上最高实现2.4倍加速，在CPU上实现2.0倍加速，优于TFLite。
• 该框架是首个支持移动GPU上实时BERT推理的方案，此前TFLite并不支持此功能。
• 即使未启用编译器优化，基线模型在CPU上的表现也与TFLite相当，但GPU性能显著下降——凸显了编译器级优化的必要性。
• NAS优化后的模型（BERT(ours) with NAS）在CPU和GPU上分别比MobileBERT快1.49倍和1.53倍，准确率仅下降0.4–1%。