LayerSkip: 조기 종료 추론과 자기 추측적 디코딩 활성화

LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

April 25, 2024

저자: Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed A Aly, Beidi Chen, Carole-Jean Wu

cs.AI

초록

대규모 언어 모델(LLM)의 추론 속도를 높이기 위한 종단 간(end-to-end) 솔루션인 LayerSkip을 소개합니다. 첫째, 학습 과정에서는 초기 층에는 낮은 드롭아웃 비율을, 후기 층에는 높은 드롭아웃 비율을 적용하는 층 드롭아웃(layer dropout)과 모든 트랜스포머 층이 동일한 출구를 공유하는 조기 종료 손실(early exit loss)을 적용합니다. 둘째, 추론 과정에서는 이 학습 방법이 모델에 추가적인 보조 층이나 모듈을 도입하지 않고도 초기 층에서의 조기 종료 정확도를 높인다는 것을 보여줍니다. 셋째, 초기 층에서 종료하고 남은 층을 통해 검증 및 수정을 수행하는 새로운 자기 추측 디코딩(self-speculative decoding) 솔루션을 제시합니다. 우리가 제안한 자기 추측 디코딩 접근법은 다른 추측 디코딩 방법들보다 메모리 사용량이 적으며, 초안(draft) 단계와 검증(verification) 단계의 계산 및 활성화를 공유함으로써 이점을 얻습니다. 다양한 크기의 Llama 모델에 대해 서로 다른 유형의 학습(처음부터 사전 학습, 지속적 사전 학습, 특정 데이터 도메인에 대한 미세 조정, 특정 작업에 대한 미세 조정)을 실험했습니다. 우리의 추론 솔루션을 구현하여 CNN/DM 문서 요약 작업에서 최대 2.16배, 코딩 작업에서 1.82배, TOPv2 의미 구문 분석 작업에서 2.0배의 속도 향상을 보여줍니다. 코드와 체크포인트는 https://github.com/facebookresearch/LayerSkip에서 공개합니다.

English

We present LayerSkip, an end-to-end solution to speed-up inference of large language models (LLMs). First, during training we apply layer dropout, with low dropout rates for earlier layers and higher dropout rates for later layers, and an early exit loss where all transformer layers share the same exit. Second, during inference, we show that this training recipe increases the accuracy of early exit at earlier layers, without adding any auxiliary layers or modules to the model. Third, we present a novel self-speculative decoding solution where we exit at early layers and verify and correct with remaining layers of the model. Our proposed self-speculative decoding approach has less memory footprint than other speculative decoding approaches and benefits from shared compute and activations of the draft and verification stages. We run experiments on different Llama model sizes on different types of training: pretraining from scratch, continual pretraining, finetuning on specific data domain, and finetuning on specific task. We implement our inference solution and show speedups of up to 2.16x on summarization for CNN/DM documents, 1.82x on coding, and 2.0x on TOPv2 semantic parsing task. We open source our code and checkpoints at https://github.com/facebookresearch/LayerSkip.