검증을 학습하는 초안 모델: VAT로 추측 디코딩 효율 향상

배경

추측 디코딩은 대규모 언어 모델의 자기회귀 추론 비용을 줄이는 대표적인 방법이다. 작은 초안 모델이 여러 후보 토큰을 먼저 생성하면, 큰 타깃 모델이 이를 한 번의 순전파로 검증한다. 후보가 연속적으로 많이 수용될수록 타깃 모델을 한 토큰씩 실행하는 횟수가 줄어들어 전체 추론이 빨라진다.

그러나 검증은 순차적으로 진행된다. 어느 한 위치에서 첫 번째 거절이 발생하면 일반적으로 해당 위치 이후의 후보는 모두 폐기된다. 기존 초안 모델 학습은 타깃 모델의 토큰 분포를 모방하는 데 집중하고, 위치별로 고정된 손실 가중치를 적용하는 경우가 많았다. 이런 학습 방식은 한 번의 거절이 이후 전체 후보의 유효성에 미치는 영향을 직접 반영하지 못한다.

「Verification-Aware Training for Speculative Decoding」은 이 문제를 다루기 위해 Verification-Aware Training(VAT)을 제시한다. VAT는 초안 모델이나 타깃 모델을 새로 설계하지 않고 학습 목적만 바꾸는 플러그인 프레임워크다.

핵심 구성

• 검증 헤드. 학습 과정에서 타깃 모델의 순차 검증을 시뮬레이션하고 각 위치가 검증을 통과해 살아남는지를 기록한다. 초안 모델과 함께 학습되는 경량 이진 분류기인 검증 헤드는 각 위치의 생존 여부를 예측한다. 따라서 초안 모델은 어떤 토큰이 그럴듯한지뿐 아니라 후보 블록이 어디까지 유지될 가능성이 높은지도 학습한다.
• 검증 적응형 가중치. 고정된 위치별 감쇠를 사용하는 대신, 각 샘플의 첫 번째 거절 지점까지는 손실에 최대 가중치를 유지한다. 이후에는 그 첫 거절 지점을 새로운 기준점으로 삼아 감쇠를 다시 시작한다. 이는 실제 검증에서 앞부분의 연속 수용이 갖는 가치를 학습에 더 가깝게 반영하려는 설계다.
• 추론 파이프라인 유지. 타깃 모델, 초안 모델의 구조, 추측 디코딩의 실행 방식은 변경하지 않는다. 주된 변화는 학습 시 사용하는 감독 신호와 손실 함수에 있다.

결과와 의미

연구진은 EAGLE-3와 DFlash에 VAT를 적용하고 Qwen3-4B, Qwen3-8B, LLaMA-3.1-8B에서 평가했다. 논문이 보고한 결과에 따르면 평균 수용 길이는 최대 11.4%, 실제 벽시계 기준 속도 향상은 최대 8.7% 개선됐다. 수학, 코드, 대화 벤치마크에서도 일관된 이득이 관찰됐다. 다만 제공된 자료에는 전체 실험 설정과 구성별 세부 결과가 없으므로, 이 수치는 모든 환경에서 보장되는 값이 아니라 보고된 최대 개선폭으로 해석해야 한다.

VAT가 던지는 broader한 시사점은 토큰 단위 모방 정확도와 추측 디코딩 효율이 동일하지 않다는 것이다. 좋은 초안 모델은 개별 토큰을 그럴듯하게 예측하는 데 그치지 않고, 타깃 모델이 연속해서 수용할 수 있는 긴 접두부를 만들어야 한다. VAT는 수용·거절 패턴을 학습 목표에 되돌려 이 목적을 명시적으로 다룬다.

기존 추론 스택을 유지하면서 다른 초안 학습 기법과 결합할 수 있다는 점도 실용적이다. 향후 공개 코드와 추가 실험을 통해 하드웨어, 배치 크기, 작업 유형에 따른 실제 효과와 비용을 검증할 필요가 있다.

출처: Hugging Face Daily Papers