AutoJudge: 작업 특화 손실 추측 디코딩을 통한 LLM 추론 가속화

핵심 요약

대형 언어 모델의 추론 속도를 높이기 위한 기존의 추측 디코딩은 타겟 모델과 드래프트 모델의 토큰이 완벽하게 일치해야 한다는 제약이 있었다. Together AI가 제안한 AutoJudge는 '손실 추측 디코딩(Lossy Speculative Decoding)' 방식을 채택하여, 최종 결과의 품질에 영향을 주지 않는 사소한 불일치는 수용함으로써 더 많은 토큰을 한 번에 생성한다. 이 방식은 수동 데이터 주석 없이도 자가 지도 학습을 통해 작업별로 중요한 토큰을 식별하는 분류기를 훈련시켜 적용된다. 실험 결과 GSM8K 및 LiveCodeBench와 같은 벤치마크에서 정확도 손실을 최소화하면서도 표준 추측 디코딩 대비 1.5~2배의 속도 향상을 달성했다.

배경

추측 디코딩(Speculative Decoding)의 기본 원리, LLM 추론 프레임워크(vLLM 등)에 대한 이해, 트랜스포머 아키텍처 및 임베딩 개념

대상 독자

LLM 추론 엔진 최적화 개발자 및 AI 인프라 엔지니어

의미 / 영향

AutoJudge는 엄격한 분포 일치라는 추측 디코딩의 한계를 극복하여 실질적인 추론 속도를 한 단계 더 끌어올렸다. 이는 특히 수학이나 코딩처럼 정답의 논리적 구조가 중요한 영역에서 효율적인 비용 절감 대안이 될 것이다.

섹션별 상세

AutoJudge는 타겟 모델의 출력 분포를 토큰 단위로 완벽하게 맞추는 대신 다운스트림 품질에 영향을 미치는 특정 토큰만 식별한다. 예를 들어 수학 문제에서 연산 기호의 불일치는 거부하지만, 의미가 동일한 단어 선택의 차이는 수용하여 생성 흐름을 유지한다. 이를 통해 검증 주기당 수용되는 토큰 수를 획기적으로 늘려 전체 추론 시간을 단축한다.

중요한 토큰을 식별하기 위해 수동 레이블링 대신 자가 지도 학습 기반의 데이터 수집 파이프라인을 사용한다. 특정 프롬프트에 대해 드래프트 토큰과 타겟 토큰이 불일치하는 지점을 찾고, 이를 반복적으로 교체하며 최종 정답의 일치 여부를 확인한다. 정답을 바꾸는 불일치만 '중요'로 표시하여 학습 데이터를 자동으로 생성한다.

추측 디코딩 과정에서 이미 계산된 트랜스포머의 히든 스테이트(Hidden States)를 활용하는 경량 분류기를 훈련한다. 드래프트와 타겟 토큰의 임베딩을 결합하여 입력으로 사용하는 로지스틱 회귀 모델을 통해 실시간으로 토큰 수용 여부를 판단한다. 이 분류기는 매우 가벼워 추론 오버헤드가 거의 발생하지 않으며 vLLM, TensorRT-LLM 등 기존 프레임워크에 쉽게 통합된다.

성능 평가 결과 Llama-3.1 70B 모델을 타겟으로 할 때 정확도 2% 하락만으로 약 1.5배의 처리량 향상을 보였다. 특히 대역폭이 제한된 오프로딩 시나리오에서는 속도 향상 폭이 약 2배까지 증가하며, EAGLE-2와 같은 최신 추측 디코딩 기법과 결합했을 때도 추가적인 성능 이득을 제공한다.