Hugging Face가 LayerSkip 기법으로 대규모 언어모델의 조기 종료(Early Exit)와 자체 투기 디코딩(Self-Speculative Decoding)을 결합해 텍스트 생성 속도 향상 및 메모리 절감

Context

기존 투기 디코딩(Speculative Decoding)은 작은 드래프트 모델과 큰 검증 모델 두 개의 서로 다른 모델을 필요로 하기 때문에 계산 비용이 증가하고 배포 시 더 많은 GPU 자원이 필요했다. 대규모 언어모델을 작은 GPU 환경에 배포할 때 메모리 오버헤드가 상당했다.

Technical Solution

• 단일 모델의 초기 레이어를 드래프트 토큰 생성에 활용하고 후기 레이어를 검증에 활용: 동일 모델 내에서 조기 종료된 레이어의 로짓(logits)을 언임베딩(unembedding) 처리
• LayerSkip 훈련 레시피 적용: 사전학습 또는 도메인 특화 미세조정 과정에서 조기 종료 손실(early exit loss)과 점진적 레이어 드롭아웃(progressive layer dropout) 적용
• Transformers 라이브러리 통합: generate() 함수에 assistant_early_exit 인자 추가로 조기 종료 자체 투기 디코딩 활성화
• 캐시 재사용을 통한 검증 최적화: 드래프트 단계와 검증 단계에서 KV 캐시 공유
• 조기 종료 레이어의 최적점 탐색: 각 모델별로 예측 정확도와 토큰 생성 오버헤드 간 최적 균형점 결정 (레이어별 성능 벤치마킹으로 식별)

Impact

벤치마크 결과 대부분의 모델에서 일정 레이어를 경계로 속도 향상이 증가했다가 감소하는 패턴 관찰 (Llama3.2 1B 제외). 메모리 사용량은 전통적 투기 디코딩 기법 대비 감소 (단일 모델 사용으로 인한 감소).

Key Takeaway

자체 투기 디코딩은 단일 모델 내의 레이어 계층을 활용하여 드래프트와 검증을 동시에 수행하므로, 모델 훈련 시점에 조기 종료 레이어의 출력을 후기 레이어 출력에 가깝게 조정하는 것이 텍스트 생성 속도와 메모리 효율성의 핵심이다.