LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

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.

한 줄 요약

LayerSkip는 LLM 추론 속도를 1.82~2.16× 가속화하는 end-to-end 솔루션으로, early exit과 self-speculative decoding을 결합한 기법이다.

핵심 기여도

Layer dropout과 early exit loss를 결합한 훈련 레시피를 제안 (모델의 모든 트랜스포머 레이어가 동일한 exit를 공유).
추론 시 추가적인 모듈 없이도 early exit 정확도를 향상 (CNN/DM 문서 요약에서 2.16×, 코딩에서 1.82×, TOPv2에서 2.0× 가속).
self-speculative decoding을 제안: early layer에서 decoding 후 나머지 레이어로 검증 및 수정 (기존 speculative decoding 대비 메모리 점유율 감소).
다양한 훈련 방식 (scratch pretraining, continual pretraining, 도메인/태스크 fine-tuning)에서 실험 수행.

핵심 아이디어

LayerSkip는 LLM의 추론 속도를 향상시키기 위해 early exit과 self-speculative decoding을 결합한 새로운 접근법을 제안한다. 기존의 speculative decoding은 별도의 draft 모델을 사용하는 반면, LayerSkip은 동일한 모델 내에서 early layer로 decoding을 수행하고 나머지 레이어로 검증 및 수정하는 방식을 채택한다. 이는 메모리 점유율을 줄이고, 별도의 모델 없이도 추론 성능을 유지할 수 있다는 장점을 가진다.

또한, 훈련 단계에서 low dropout rate를 early layer에, high dropout rate를 late layer에 적용하며, early exit loss를 통해 모든 트랜스포머 레이어가 동일한 exit를 공유하도록 유도한다. 이는 추론 시 early exit의 정확도를 향상시키는 데 기여하며, 모델 내부에서 다양한 크기의 서브모델을 생성할 수 있게 한다.

기술적 접근법

**Layer dropout**: early layer는 낮은 dropout rate (예: 0.1), late layer는 높은 dropout rate (예: 0.2) 적용.
**Early exit loss**: 모든 트랜스포머 레이어가 동일한 exit를 공유하도록 훈련.
**Self-speculative decoding**: early layer에서 decoding 후 나머지 레이어로 검증 및 수정. KV cache 재사용으로 9~20ms/token의 속도 향상.
**훈련 실험**: Llama2 7B, 13B, Llama3 8B, 1B 등 다양한 모델 크기에서 scratch pretraining, continual pretraining, fine-tuning 실험 수행.
**하이퍼파라미터 예시**: `p_max = 0.1`, `e_scale = 0.1`, `C_rot,R = 8` 등.

주요 결과

**CNN/DM 문서 요약**: 2.16× 가속.
**코딩**: 1.82× 가속.
**TOPv2 의미 파싱**: 2.0× 가속.
**KV cache 재사용**: 9~20ms/token의 속도 향상.
**Middle layer perplexity**: early exit loss 적용 시 증가 억제.

의의 및 한계

LayerSkip는 LLM의 추론 속도를 향상시키는 end-to-end 솔루션으로, 기존의 speculative decoding과 달리 별도의 draft 모델 없이도 메모리 점유율을 줄이고 성능을 유지할 수 있다. 또한, early exit loss와 layer dropout의 결합은 early layer의 정확도를 향상시키며, 다양한 훈련 방식에서의 실험을 통해 일반화 가능성을 입증했다.

그러나, early exit의 정확도 향상은 훈련 과정에서 일부 overhead를 유발할 수 있으며, 모든 토큰에 동일한 exit layer를 적용하는 방식은 개선 여지가 있다. 향후 연구에서는 토큰별 exit layer를 동적으로 결정하는 방식을 탐구할 수 있다.

실용적 활용

LayerSkip는 GPU 서버에서의 LLM 배포를 저비용으로 가능하게 하며, 모바일 및 에지 기기로의 확장에도 유용하다. 특히, 문서 요약, 코드 생성, 의미 파싱 등 다양한 NLP 태스크에서 실시간 추론이 필요한 상황에 적용 가능하다.