토큰이 아니라 피처를 예측하는 EAGLE 추측 디코딩

자기회귀 디코딩은 토큰을 하나씩 생성하므로 LLM 응답 지연의 근본 원인입니다. ICML에 발표된 EAGLE은 추측 디코딩의 초안 생성 위치를 바꾸는 것만으로 품질 손실 없이 지연을 절반 이하로 줄인 연구입니다.

문제 배경

추측 디코딩(speculative decoding)은 작은 초안 모델이 여러 토큰을 미리 제안하고 본 모델이 한 번에 검증하는 가속 기법입니다. 그러나 본 모델과 어휘를 공유하는 적절한 초안 모델을 구하기 어렵고, 초안이 자주 기각되면 검증 비용만 남아 이득이 사라집니다. 특히 파인튜닝된 특화 모델은 짝이 맞는 소형 모델이 아예 없는 경우가 많습니다.

방법

EAGLE(Extrapolation Algorithm for Greater Language-model Efficiency)은 두 가지 관찰에서 출발합니다. 첫째, 최상위 직전 레이어의 피처 시퀀스는 토큰 시퀀스보다 규칙적이어서 자기회귀 예측이 더 쉽습니다. 둘째, 피처 수준 예측에는 샘플링 결과를 모르는 데서 오는 불확실성이 남는데, 한 스텝 앞선 토큰을 입력에 함께 넣으면 이 불확실성이 해소됩니다. 초안 모듈은 임베딩 층과 LM 헤드를 본 모델에서 그대로 가져와 추가 학습이 필요 없고, 새로 학습하는 부분은 FC 층과 디코더 레이어 하나로 구성된 자기회귀 헤드뿐입니다. 검증 단계가 원 모델의 분포를 보존하므로 생성 결과의 분포는 이론적으로 동일합니다.

결과

Vicuna와 LLaMA2-Chat 전 계열, MoE 모델 Mixtral 8x7B Instruct에 대해 대화, 코드 생성, 수학 추론, 명령 수행 과제로 평가했습니다. LLaMA2-Chat 70B 기준 지연 시간을 2.7배에서 3.5배 단축했고 처리량은 2배로 늘었습니다.

한계와 주의점

예측 헤드는 공짜가 아니라 별도 학습이 필요합니다. 논문 기준 70B 모델용 헤드도 A100 40G 서버에서 1일에서 2일이면 학습되지만 파이프라인은 갖춰야 합니다. 또한 가속 폭은 초안 수락률에 좌우되므로 과제 특성과 샘플링 온도에 따라 실측 이득이 달라지며, 배치가 커서 GPU 연산이 이미 포화된 서빙 환경에서는 이득이 줄어듭니다.

기업 적용 관점

응답 지연이 계약 조건에 들어가는 챗봇이나 상담 보조 수주에서 GPU 증설 없이 지연을 줄이는 카드로 제안할 수 있습니다. 고객 도메인에 파인튜닝된 sLLM은 기존 추측 디코딩을 적용할 초안 모델이 없다는 점이 걸림돌이었는데, EAGLE 방식은 예측 헤드만 추가 학습하면 되므로 온프레미스 특화 모델에도 그대로 이식됩니다.