모든 초안을 활용하라: Hidden State 기반 추론 가속화 기법

Make Every Draft Count: Hidden State based Speculative Decoding

추론 가속화를 위한 핵심 기술로 등장한 추론 가속화는, 경량 모델을 사용하여 후보 토큰을 생성하고, 생성된 토큰을 대상 모델이 병렬적으로 검증하는 방식으로 작동합니다. 이 방식은 메모리 병목 현상을 완화하여 연산 효율을 높이지만, 대부분의 초안 토큰이 검증에 실패하여 폐기되고, 이는 불필요한 연산 낭비를 초래합니다. 본 연구는 이러한 낭비된 연산을 재활용하기 위한 새로운 시스템을 제안합니다. 핵심 아이디어는 Hidden State 레벨에서 자기 회귀 예측을 수행하고, 토큰 정보를 Hidden State 생성 이후에 통합하여, 잘못된 토큰으로 인해 Hidden State가 오염되는 것을 방지함으로써 Hidden State의 재사용을 가능하게 하는 것입니다. 이러한 시스템을 구현하기 위해, 먼저 자기 회귀 Hidden State 기반의 초안 모델 아키텍처를 도입하여, 토큰 기반 초안 모델보다 풍부한 의미 정보를 유지하고, 초안 재활용을 용이하게 합니다. 둘째, 특별히 설계된 초안 모델을 활용하여 고품질의 초안 토큰 트리를 구성하고, 검증 실패 시 토큰을 재샘플링할 수 있는 효율적인 토큰 정보 주입 메커니즘을 설계합니다. 셋째, 설계된 시스템 내의 오버헤드를 제거하여 하드웨어 활용도를 극대화합니다. 다양한 기준 모델과의 비교 평가를 통해, 본 연구에서 제안하는 방법이 기존 추론 가속화 방식보다 최대 3.3배 빠른 성능을 보임을 확인했습니다.

Speculative decoding has emerged as a pivotal technique to accelerate LLM inference by employing a lightweight draft model to generate candidate tokens that are subsequently verified by the target model in parallel. However, while this paradigm successfully increases the arithmetic intensity of memory-bound inference, it causes significant compute inefficiency: the majority of draft tokens fail verification and are discarded, resulting in waste of computation. Motivated by the goal of recollecting this wasted computation, we propose a novel system that transforms discarded drafts into reusable tokens. Our key insight is to perform auto-regressive prediction at the hidden states level and postpone the integrating token information after the hidden states generation, so the draft hidden states are not contaminated by incorrect tokens, enabling hidden state reuse. To implement such a system, first we introduce a draft model architecture based on auto-regressive hidden states, which preserves richer semantics than token-based drafters to facilitate draft repurposing. Second, we design an efficient token information injection mechanism that leverages our specialized draft model to construct high-quality draft token trees and enables resampling tokens from verification failures. Third, we eliminate the overhead hidden in our design to further maximize hardware utilization. We conducted extensive evaluations against various baselines, demonstrating up to a 3.3x speedup against standard speculative decoding.