기술 요약: LLM 을 위한 계층적 추측적 디코딩 (HiSpec)

1. 문제 제기

대형 언어 모델 (LLM) 추론은 처리량과 정확도 사이의 근본적인 트레이드오프에 직면해 있습니다. **추측적 디코딩 (speculative decoding)**은 더 작은 '초안 (draft)' 모델을 사용하여 토큰을 제안하고, 더 큰 '타겟' 모델이 이를 검증함으로써 처리량을 향상시키지만, 현재는 검증 단계에 병목 현상이 발생하고 있습니다.

저자들은 토큰 검증이 초안 토큰 생성보다 현저히 느린 (2 배에서 6.9 배, 일부 구성에서는 최대 10.3 배) '검증 벽 (verification wall)'을 식별했습니다. 이전 토큰이 검증될 때까지 후속 토큰을 생성할 수 없기 때문에 이 지연 시간이 전체 추론 시간을 지배합니다. 기존 연구들은 주로 초안 작성 (drafting) 단계를 가속화하는 데 집중하여, 검증 병목이 해결되지 않았기 때문에 처리량 향상 효과가 미미했습니다.

'중간 검증 (intermediate verification)'을 통해 이를 해결하려는 이전 시도들 (예: 부정확한 토큰을 조기에 폐기하기 위해 보조 모델을 사용하는 SPRINTER 등) 은 다음과 같은 세 가지 주요 한계를 겪고 있습니다:

1. 훈련 오버헤드: 타겟 모델을 근사화하기 위해 보조 모델을 훈련해야 합니다.
2. 메모리 풋프린트: 여러 모델을 동시에 오케스트레이션하면 메모리 사용량이 증가합니다.
3. 정확도 저하: 일부 방법은 근사 휴리스틱에 의존하거나 전체 타겟 검증을 건너뛰어, 타겟 모델의 자기회귀 분포에서 벗어난 정보 손실 (lossy) 이 있는 출력을 초래합니다.

2. 방법론: HiSpec

저자들은 **계층적 추측적 디코딩 (HiSpec)**을 제안했습니다. 이는 손실 없는 정확도를 유지하면서 낮은 오버헤드로 토큰 검증을 가속화하도록 설계된 프레임워크입니다. HiSpec 은 토큰이 선택된 중간 계층에서 탈출할 수 있도록 명시적으로 훈련된 LLM 인 Early-Exit(EE) 모델을 활용하여 전체 네트워크를 통과하지 않고도 의미 있는 토큰 분포를 생성합니다.

핵심 메커니즘

HiSpec 은 세 가지 고유한 작동 계층을 가진 단일 EE 모델을 활용합니다:

1. 초안 계층 ($L_d$): 초안 토큰을 생성하는 데 사용되는 얕은 탈출 계층입니다.
2. 중간 검증 계층 ($L_i$): 중간 검증을 수행하는 데 사용되는 더 깊은 탈출 계층 (일반적으로 모델 깊이의 약 1/4) 입니다. 이 계층은 자신의 분포와 일치하는 초안 토큰을 잠정적으로 수용하고 부정확한 토큰은 조기에 거부합니다.
3. 타겟 계층 ($L_f$): 출력 일관성을 보장하기 위한 주기적인 전체 모델 검증에 사용되는 최종 계층 (또는 전체 모델) 입니다.

주요 기술적 혁신

• 낮은 오버헤드의 중간 검증: 이전 연구들이 별도의 보조 모델을 훈련한 것과 달리, HiSpec 은 EE 모델의 고유한 탈출 계층을 사용합니다. 이는 새로운 검증자를 훈련할 필요성을 없애고 여러 모델 가중치를 로드하는 메모리 오버헤드를 방지합니다.
• KV 캐시 및 숨겨진 상태 재사용: 불필요한 계산을 방지하기 위해 HiSpec 은 초안 생성, 중간 검증, 타겟 검증 단계 전반에 걸쳐 키 - 값 (KV) 캐시와 숨겨진 상태를 재사용하는 메커니즘을 구현합니다. 초안 생성 중 이러한 상태를 버퍼링하고 다음 단계 전에 거부된 토큰과 관련된 항목을 제거하여 재계산 없이 컨텍스트 정렬을 보장합니다.
• 손실 없는 보장: HiSpec 은 2 단계 검증 프로세스 ($L_d \to L_i \to L_f$) 를 사용합니다. 중간 검증자는 조기 거부만 수행하며 최종 수용 기준을 완화하지 않습니다. 두 단계 모두에서 손실 없는 추측적 디코딩 증명을 순차적으로 적용함으로써, HiSpec 은 최종 출력이 고립되어 실행되는 타겟 모델 (자기회귀 디코딩) 의 출력과 동일함을 보장합니다.
• 동적 검증 정책: 처리량과 정확도의 균형을 맞추기 위해 HiSpec 은 비싼 전체 모델 검증을 호출하기 전에 잠정적으로 수용된 토큰의 창 (기본값 $N_i = 4$) 을 누적합니다. 이는 전체 모델 호출 빈도를 줄이면서도 검증되지 않은 토큰의 긴 사슬을 방지합니다.
• 배치 추론 지원: 토큰 수용 길이가 균일하지 않아 배치 처리에 어려움을 겪는 많은 기존 추측적 디코딩 방법과 달리, HiSpec 은 블록 테이블과 **페이징 어텐션 (paged attention)**을 통합합니다. 이는 패딩 오버헤드 없이 시퀀스별 수용을 가능하게 하여 효율적인 배치 추론을 지원합니다.

3. 주요 기여

이 논문은 다음과 같은 기여를 주장합니다:

1. 검증 병목 현상의 식별: 저자들은 검증 지연 시간 (생성보다 최대 6.9 배 느림) 이 초안 가속화 연구들이 largely 무시해 온 추측적 디코딩 처리량의 주요 제한 요소임을 입증했습니다.
2. HiSpec 프레임워크: 보조 검증자와 관련된 훈련 및 메모리 비용을 제거하는 중간 검증을 위해 Early-Exit 모델을 사용하는 새로운 프레임워크입니다.
3. 자원 효율성: 모든 단계 (초안, 중간, 타겟) 에서 KV 캐시와 숨겨진 상태를 재사용하는 메커니즘으로 계산 및 메모리 중복을 크게 줄였습니다.
4. 정확도 보존: 주기적인 전체 모델 검증을 통해 타겟 모델과 엄격한 출력 일관성을 유지하는 설계로, 손실 없는 보장을 제공합니다.
5. 일반화 및 확장성: 이 방법은 사전 훈련된 EE 모델과 미세 조정을 통해 예측 헤드를 추가하여 수정된 사후 훈련 모델 모두에서 작동하며, 네이티브로 배치 추론을 지원합니다.

4. 결과

저자들은 Llama3, Llama2, CodeLlama, Qwen-3, OPT 를 포함한 8 개의 최첨단 모델과 ShareGPT, CNN/DM, HumanEval, GSM8K, XSum 의 5 가지 다양한 벤치마크를 통해 HiSpec 을 평가했습니다.

• 처리량 향상: 중간 검증이 없는 기본 단일 계층 추측과 비교하여 HiSpec 은 평균 1.28 배, 최대 2.01 배의 처리량 향상을 달성했습니다.
• 최첨단 (SOTA) 과의 비교: HiSpec 은 AdaDecode, Lookahead, LayerSkip 과 같은 최첨단 초안 가속화 방법과 SPRINTER 와 같은 검증 가속화 방법보다 일관되게 우수한 성능을 발휘했습니다.
  • SPRINTER(보조 검증자를 사용) 와 비교하여, HiSpec 은 ShareGPT 에서 SPRINTER 의 1.20 배에 비해 2.01 배의 처리량을 달성하면서도 100% 정확도를 유지했습니다 (SPRINTER 는 테스트에서 정확도를 약 40-52% 로 저하시켰습니다).
  • HiSpec 은 모든 벤치마크와 모델 크기에서 LayerSkip과 Lookahead보다 우수한 성능을 보였습니다.
• 자원 효율성: KV 캐시 재사용으로 인해 중간 검증은 < 0.01% 의 무시할 수 있는 메모리 오버헤드만 추가합니다.
• 배치 처리: HiSpec 은 배치 크기 (1~16) 전반에 걸쳐 속도 향상을 유지하는 반면, 이전 방법들은 배치 처리와 함께 효과적으로 확장되지 않는 경우가 많았습니다.

5. 중요성 및 주장

이 논문은 HiSpec 을 LLM 추론의 중요한 '검증 벽'에 대한 해결책으로 위치시킵니다. 그 중요성은 다음과 같습니다:

• 초점 전환: 저자들은 더 높은 수익을 창출한다고 주장하는 초안 생성 가속화에서 검증 가속화로 최적화 초점을 이동시킵니다.
• 효율성: 보조 모델 훈련이나 여러 모델 가중치의 동시 관리에 따른 계산 및 메모리 비용 없이 높은 처리량을 달성합니다.
• 실용성: 중간 검증과 네이티브 배치 추론 지원 및 손실 없는 보장을 결합한 최초의 방법으로, 광범위한 배포에 적합합니다.
• 유연성: 사전 훈련된 EE 모델과 경량 사후 훈련 (예측 헤드 추가) 을 통해 수정된 표준 모델 모두에 적응할 수 있어 광범위한 적용 가능성을 보장합니다.

저자들은 조기 탈출 계층을 중간 검증에 활용하고 자원을 효율적으로 관리함으로써 HiSpec 이 타겟 모델의 정확도를 훼손하지 않는 견고하고 고처리량의 추론 프레임워크를 제공한다고 결론지었습니다.