[论文解读] AutoDistill: an End-to-End Framework to Explore and Distill Hardware-Efficient Language Models

Xiaofan Zhang, Zongwei Zhou+2 more|arXiv (Cornell University)|Jan 21, 2022

Topic Modeling被引用 5

一句话总结

AutoDistill 是一个端到端框架，采用贝叶斯优化进行多目标神经架构搜索，并利用闪存蒸馏（Flash Distillation）快速评估模型候选，以构建硬件高效、任务无关的 NLP 模型。其在 GLUE 和 SQuAD 基准测试中，以 20.6M 参数的模型实现了比 MobileBERT 高出 3.2% 的预训练准确率和 1.44 倍的推理延迟加速，表现优于 BERT BASE 及其他蒸馏模型。

ABSTRACT

Recently, large pre-trained models have significantly improved the performance of various Natural LanguageProcessing (NLP) tasks but they are expensive to serve due to long serving latency and large memory usage. To compress these models, knowledge distillation has attracted an increasing amount of interest as one of the most effective methods for model compression. However, existing distillation methods have not yet addressed the unique challenges of model serving in datacenters, such as handling fast evolving models, considering serving performance, and optimizing for multiple objectives. To solve these problems, we propose AutoDistill, an end-to-end model distillation framework integrating model architecture exploration and multi-objective optimization for building hardware-efficient NLP pre-trained models. We use Bayesian Optimization to conduct multi-objective Neural Architecture Search for selecting student model architectures. The proposed search comprehensively considers both prediction accuracy and serving latency on target hardware. The experiments on TPUv4i show the finding of seven model architectures with better pre-trained accuracy (up to 3.2% higher) and lower inference latency (up to 1.44x faster) than MobileBERT. By running downstream NLP tasks in the GLUE benchmark, the model distilled for pre-training by AutoDistill with 28.5M parameters achieves an 81.69 average score, which is higher than BERT_BASE, DistillBERT, TinyBERT, NAS-BERT, and MobileBERT. The most compact model found by AutoDistill contains only 20.6M parameters but still outperform BERT_BASE(109M), DistillBERT(67M), TinyBERT(67M), and MobileBERT(25.3M) regarding the average GLUE score. By evaluating on SQuAD, a model found by AutoDistill achieves an 88.4% F1 score with 22.8M parameters, which reduces parameters by more than 62% while maintaining higher accuracy than DistillBERT, TinyBERT, and NAS-BERT.

研究动机与目标

• 解决数据中心部署中自动化、可扩展且具备延迟感知的模型蒸馏缺乏的问题。
• 降低为大规模预训练模型探索学生模型架构的成本与时间。
• 在完整蒸馏前高效识别高潜力学生模型。
• 同时优化目标硬件上的预测准确率与推理延迟。
• 支持端到端、自动化且可投入生产的 NLP 模型蒸馏，适用于数据中心服务。

提出的方法

• 采用贝叶斯优化（BO）执行多目标神经架构搜索（NAS），在目标硬件上联合优化准确率与推理延迟。
• 提出闪存蒸馏（Flash Distillation），一种模型无关的技术，仅使用完整蒸馏步骤的 5%（37k 步 vs. 740k 步）即可早期识别有前景的学生模型候选。
• 在闪存蒸馏中采用渐进式知识迁移策略，固定渐进式与预训练蒸馏步骤之间的比例为 0.48。
• 通过对数准确率增长曲线对候选模型进行早期评估，以确定成本效益最优的停止点。
• 将 BO 与闪存蒸馏集成，加速搜索过程的同时保持高质量模型的发现能力。
• 在 TPUv4i 硬件上执行端到端蒸馏，基于 GLUE 和 SQuAD 基准进行综合性能评估。

实验结果

研究问题

• RQ1 一个自动化、端到端的框架能否在极少人工干预下有效探索并蒸馏出硬件高效的 NLP 模型？
• RQ2 如何高效地将多目标神经架构搜索应用于平衡模型蒸馏中的准确率与推理延迟？
• RQ3 一个简短的渐进式蒸馏阶段（闪存蒸馏）能否可靠地在完整训练前识别出高潜力的学生模型？
• RQ4 闪存蒸馏在不牺牲最终模型性能的前提下，能在多大程度上降低模型探索成本？
• RQ5 所生成的蒸馏模型是否能在标准基准测试中同时优于现有最先进蒸馏模型的准确率与延迟表现？

主要发现

• AutoDistill 发现了七个学生模型架构，在 TPUv4i 上相比 MobileBERT 实现最高达 3.2% 的预训练准确率提升，且推理延迟最高加快 1.44 倍。
• 28.5M 参数的模型在 GLUE 基准上取得 81.69 的平均得分，优于 BERT BASE（109M）、DistillBERT（67M）、TinyBERT（67M）、NAS-BERT 和 MobileBERT（25.3M）。
• 参数最少的模型（20.6M）在参数量减少 81.1% 的情况下，仍取得高于 BERT BASE、DistillBERT、TinyBERT 和 MobileBERT 的平均 GLUE 得分。
• 在 SQuAD 基准上，两个参数少于 23M 的 AutoDistill 模型在 F1 分数上优于 DistillBERT、TinyBERT 和 NAS-BERT，同时参数量减少超过 62%。
• 使用 37k 步的闪存蒸馏成功识别出有前景的模型，仅消耗 5% 的训练成本，且与完整蒸馏相比保持了稳定的相对性能排名。
• BO 与闪存蒸馏的结合显著降低了整体搜索成本，同时实现了在准确率与延迟方面均优于现有基线模型的发现。