Layer Skip: Enabling Early Exit Inference and Self-Speculative Decoding

Abstract

레이어를 조기 종료하고 병렬로 검증하여 보조 모델 없이 빠르게 디코딩

 

[arXiv](2024/04/25 version v1)

 

 

 

Introduction

Speculative Decoding은 보조 모델이 필요하다. 

본문에서는 추가적인 모델이나 레이어가 필요하지 않은 Self Speculative Decoding을 제안한다.

 

 

 

Motivation

Exiting Earlier in LLMs

Llama 7B 각 레이어의 출력을 LM head에 투영하여 해당 레이어에서 예측된 토큰을 확인했다.

 

관찰 결과, 올바른 토큰을 예측하는 데 모든 레이어가 필요한 것은 아니었다. 평균적으로 23.5개의 레이어가 필요했다.

 

Layer dropout을 사용하여 모델이 후기 레이어보다 초기 레이어에 더 의존하도록 하고, 손실 함수를 추가하여 모델이 레이어를 건너뛰고 바로 LM head를 통해 토큰을 출력할 수 있도록 하는 방법을 제안한다.

---

Correcting if we Exit Too Early

레이어를 조기 종료하는 경우 정확도가 감소할 수 있다. 

조기 종료로 토큰을 생성하고 나머지 레이어를 통해 병렬로 토큰을 확인하고 수정하는 self speculative decoding을 제안한다.

 

 

 

Proposed Solution

제안된 방법에는 3가지 단계가 있다.

• Layer Dropout & Early Exit Loss를 이용한 학습
• Early Exit를 이용한 추론
• Speculative Decoding을 이용한 검증 및 수정

---

Training using Layer Dropout & Early Exit Loss

표기:

입력 X, 출력 Y,

레이어 출력 x_l,

LM head g, logits e,

loss function J.

 

Layer Dropout

훈련 반복 t, 레이어 l에서 dropout 확률 p를 적용한다.

 

후기 레이어로 갈수록, 훈련 시간이 길어질수록 p가 지수적으로 증가한다.

 

Early Exit Loss

모든 레이어의 조기 종료 손실.

e는 정규화항. 모든 레이어의 합은 1이고 후기 레이어로 갈수록 증가한다.

.

---

Inference using Early Exit

---

Inference using Self-Speculative Decoding

레이어를 조기 종료하고 나머지 레이어를 통해 검증한다.

Self-Verification

생성 단계에서 사용되지 않은 나머지 레이어를 활용하여 draft token과 최종 출력 토큰이 일치하는지 병렬로 검증하고 수정한다.

 

Reusing the Cache

• Draft 단계에서 생성된 KV cache를 저장하여 재계산하지 않는다.
• 또한 검증을 계속하기 위해 draft 출력 레이어의 Query를 저장한다.

 

 

 

Experiments

성능 (Last layer만 보면 됨)

 

빠름