[논문 리뷰] AutoDistill: an End-to-End Framework to Explore and Distill Hardware-Efficient Language Models

Xiaofan Zhang, Zongwei Zhou외 2명|arXiv (Cornell University)|2022. 01. 21.

Topic Modeling인용 수 5

한 줄 요약

AutoDistill는 베이지안 최적화를 사용한 다중 목적 신경망 아키텍처 탐색과 빠른 모델 후보 평가를 위한 플래시 디스틸리케이션을 활용하는 엔드 투 엔드 프레임워크로, 하드웨어 효율적이고 작업에 종속되지 않는 NLP 모델을 구축한다. 이 프레임워크는 MobileBERT보다 최대 3.2% 높은 사전학습 정확도와 1.44배 빠른 추론 지연 시간을 달성하며, 20.6M 파라미터 모델이 GLUE 및 SQuAD 벤치마크에서 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 모델 디스틸리케이션을 지원하기 위해.

제안 방법

• 목표 하드웨어에서 정확도와 추론 지연 시간을 동시에 최적화하는 다중 목적 신경망 아키텍처 탐색(NAS)을 수행하기 위해 베이지안 최적화(BO)를 활용한다.
• 전체 디스틸리케이션 단계의 5%에 해당하는 37k 단계(전체 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 벤치마크에서 AutoDistill 모델 두 개(23M 이하 파라미터)는 DistillBERT, TinyBERT, NAS-BERT를 초월하는 F1 스코어를 기록했으며, 파라미터 수를 62% 이상 감소시켰다.
• 37k 단계의 플래시 디스틸리케이션은 전체 학습 비용의 5%만으로도 유망한 모델을 성공적으로 식별했으며, 전체 디스틸리케이션과 비교해 안정적인 상대적 성능 순위를 유지했다.
• BO와 플래시 디스틸리케이션의 통합은 전체 검색 비용을 크게 줄였고, 정확도와 지연 시간 측면에서 기존 베이스라인을 모두 초월하는 모델 발견을 가능하게 했다.