기술 요약: 모델 프리 스펙큘레이티브 샘플링(STAND)을 통한 가속화된 테스트 시간 확장

문제 제기

테스트 시간 확장 (TTS) 기법들, 예를 들어 베스트 오브 N 샘플링, 다수결 투표, 트리 탐색 등은 추론 과정에서 추가적인 계산 자원을 할당함으로써 대형 추론 모델 (LRM) 의 추론 능력을 크게 향상시켰습니다. 그러나 이러한 방법들은 긴 시퀀스와 여러 궤적을 생성함으로써 상당한 계산 비용을 수반합니다. 길이 기반 보상을 통한 학습이나 캐스케이드 모델 사용과 같은 기존 효율성 중심 접근법들은 계산 비용을 줄이면 정확도가 희생되는 근본적인 트레이드오프에 직면하는 경우가 많습니다. 또한, 스펙큘레이티브 디코딩 (SD) 은 손실 없는 가속화를 위한 경로를 제공하지만, 많은 기존 모델 프리 SD 방법들은 LRM 에서 사용되는 샘플링 기반 생성 전략의 확률적 특성을 활용하지 못하는 결정론적 n-그램 조회에 의존합니다.

방법론: STAND

저자들은 추가적인 드래프트 모델이나 미세 조정이 필요 없이 LLM 추론을 가속화하도록 설계된 모델 프리 스펙큘레이티브 디코딩 접근법인 STAND(STochastic Adaptive N-gram Drafting, 확률적 적응형 n-그램 드래프팅) 를 제안합니다. STAND 는 서로 다른 솔루션 전반에서 유사한 추론 패턴과 토큰 시퀀스가 빈번하게 재사용되는 추론 궤적의 고유한 중복성을 활용합니다.

이 방법론은 세 가지 핵심 혁신에 기반합니다:

1. 로그트 기반 적응형 n-그램 모듈:
   전통적인 모델 프리 드래프터가 이산적인 토큰 ID 나 상위 k 개 토큰을 저장하는 것과 달리, STAND 는 관찰된 n-그램에 대한 다음 토큰의 로그트 분포(확률 분포) 를 저장합니다.

   • 메모리 효율성: 대규모 어휘를 처리하기 위해 시스템은 상위 k 개 인덱스와 해당 확률만 유지합니다. n-그램이 반복적으로 마주쳐지면, 저장되지 않은 인덱스를 0 확률로 간주하여 가중 평균을 사용하여 분포를 병합합니다. 결과 분포는 일정한 메모리 사용량을 보장하기 위해 상위 10 개 가장 확률이 높은 토큰으로 잘립니다.
   • 확률적 드래프팅: 확률적 정보를 보존함으로써, STAND 는 단일 결정론적 토큰을 선택하는 대신 저장된 분포에서 드래프트 토큰을 샘플링하는 확률적 드래프팅을 가능하게 합니다. 이는 드래프트 분포를 타겟 모델의 샘플링 전략과 정렬시켜 토큰 수락률을 크게 높입니다.

2. Gumbel-Top-K 를 통한 최적화된 샘플링:
   드래프트 트리를 위해 병렬로 여러 토큰을 효율적으로 샘플링하기 위해, STAND 는 순차적 샘플링을 Gumbel-Top-K 트릭으로 대체합니다. 로그 확률에 Gumbel 노이즈를 추가하고 상위 k 개 인덱스를 선택함으로써, 이 방법은 k 개의 토큰을 비복원 방식으로 병렬로 샘플링합니다. Gumbel 노이즈 항은 사전 계산되어 캐시되므로 샘플링 오버헤드를 드래프팅 프로세스와 분리하여 지연 시간을 더욱 줄입니다.

3. 데이터 기반 드래프트 트리 최적화:
   휴리스틱 또는 순수 동적 트리 구성에 의존하는 대신, STAND 는 정적 드래프트 트리를 최적화하기 위한 데이터 기반 접근법을 채택합니다.

   • 프로세스: 초기 대규모 트리 (625 개 노드) 가 구성됩니다. 소수의 실제 데이터 샘플 (예: 30 개 샘플) 에서 스펙큘레이티브 디코딩이 수행됩니다. 시스템은 노드별 수락 확률을 추적하고 가장 효과적인 상위 80 개 노드를 선택하여 컴팩트하고 최적화된 트리 구조로 재구성합니다.
   • 결과: 이는 트리 구조가 실제 세계 성능을 위해 경험적으로 최적화되도록 하면서도 정적 트리의 계산 효율성을 유지합니다.

주요 기여

• 모델 프리 확률적 드래프팅: 본 논문은 기존 모델 프리 SD 방법들이 확률성을 largely 무시하고 있음을 규명합니다. STAND 는 이전 생성물로부터 로그트 정보를 보존하고 활용하여 확률적 드래프팅을 수행하는 메커니즘을 도입하며, 이는 샘플링 기반 LRM 추론에 중요합니다.
• 효율적인 로그트 저장: 전체 신경망 드래프트 모델의 오버헤드 없이 확률적 드래프팅을 가능하게 하는 n-그램 로그트 분포를 위한 메모리 효율적이고 압축된 표현 방식을 제안합니다.
• 데이터 기반 트리 구성: 경험적 수락률에 기반하여 트리 구조를 적응시키는 드래프트 트리를 위한 새로운 오프라인 최적화 방법으로, 휴리스틱 트리 설계보다 우수한 성능을 보입니다.
• 플러그 앤 플레이 호환성: 모델 프리 접근법으로서 STAND 는 타겟 모델에 대한 추가 학습이나 미세 조정, 별도의 드래프트 모델이 필요하지 않아 기존 LRM 에 모두 적용 가능합니다.

실험 결과

DeepSeek-R1-Distill-Qwen-7B 및 14B 모델을 사용하여 수학 추론 (AIME-2024), 과학 QA(GPQA-Diamond), 코딩 (LiveCodeBench) 에서 광범위한 평가가 수행되었습니다.

• 처리량 및 지연 시간: STAND 는 정확도를 유지하면서 표준 자기회귀 디코딩에 비해 추론 지연 시간을 60–65% 줄입니다. 처리량 측면에서는 최신 스펙큘레이티브 디코딩 방법보다 14–28% 더 우수한 성능을 보입니다.
• 다중 궤적 확장: STAND 의 이점은 추론 궤적의 수에 따라 확장됩니다. 베스트 오브 16 샘플링 시나리오에서 STAND 는 가장 높은 처리량과 수락 길이를 달성하여 Token Recycle, ANPD, SAM 과 같은 베이스라인을 크게 능가합니다.
• 다양한 추론 패턴:
  • 단일 궤적: STAND 는 표준 디코딩 대비 처리량을 약 2.38 배 향상시킵니다.
  • 배치 디코딩: 다른 방법들이 성능이 저하되는 반면, STAND 는 배치 크기가 8 일 때도 베이스라인 대비 최대 1.52 배의 상당한 속도 향상을 유지합니다.
  • 테스트 시간 트리 탐색: 다양한 검증자 트리 탐색 (DVTS) 에서 STAND 는 표준 디코딩 대비 평균 2.54 배의 속도 향상을 달성합니다.
• 절대 실험 (Ablation Studies):
  • 확률적 vs 결정론적: 확률적 드래프팅은 결정론적 드래프팅에 비해 일관되게 더 높은 수락 확률 (5–8% 높음) 과 처리량을 보입니다.
  • 트리 최적화: 데이터 기반 방법으로 최적화된 트리는 도메인 내 (AIME) 및 도메인 외 (GPQA) 작업 모두에서 휴리스틱 트리 (예: Token Recycle 에서의) 보다 우수한 성능을 보입니다.
  • 트리 구조: STAND 를 위한 최적화된 트리는 베이스라인 (7 개) 대비 더 깊은 깊이 (최대 13 개 수준) 와 "긴 꼬리" 구조를 보여주며, 이는 추론 궤적에서 발견되는 더 긴 결정론적 시퀀스를 활용하는 모델의 능력을 반영합니다.

의의 및 주장

본 논문은 STAND 가 언어 모델 추론을 가속화하기 위한 강력한 플러그 앤 플레이 솔루션을 제공한다고 주장합니다. 추론 궤적의 중복성과 샘플링 기반 생성에 내재된 확률적 정보를 활용함으로써, STAND 는 타겟 모델의 정확도를 훼손하지 않고 상당한 가속화 (지연 시간 60–65% 감소) 를 달성합니다.

저자들은 모델 재학습이나 별도의 드래프트 모델 배포가 필요한 방법들과 달리 STAND 는 완전히 모델 프리이며 기존 LRM 과 호환됨을 강조합니다. 이는 특히 다중 궤적, 배치 디코딩, 트리 탐색이 관련된 시나리오에서 테스트 시간 확장 시 성능과 효율성 사이의 트레이드오프를 효과적으로 해결합니다. 이 연구는 이전 궤적 데이터를 단순히 대체하거나 무시하는 방법들보다 (n-그램 로그트 저장을 통한) 역사적 정보를 집계하는 것이 테스트 시간 확장에 더 유익함을 시사합니다.

참고: 저자들은 현재 구현이 HuggingFace 라이브러리 (vLLM 또는 SGLang 보다 덜 최적화됨) 와 Python 기반 n-그램 조회를 사용한다는 한계를 인정하며, 더 최적화된 구현을 통해 추가적인 지연 시간 감소 가능성이 있음을 시사합니다.