Efficient Test-Time Inference via Deterministic Exploration of Truncated Decoding Trees

1. 연구 배경 및 문제 제기 (Problem)

대규모 언어 모델 (LLM) 의 추론 능력을 향상시키기 위해 자기 일관성 (Self-Consistency) 기법이 널리 사용됩니다. 이는 동일한 프롬프트에 대해 여러 개의 추론 경로를 샘플링하고, 다수결 투표 등을 통해 최종 답을 도출하는 방식입니다.

그러나 수학, 코딩과 같이 제약된 도메인 (Constrained Domains) 에서 이 방식은 다음과 같은 비효율성을 보입니다:

• 중복 샘플링: 확률적 샘플링 (Sampling with replacement) 을 사용하므로, 높은 확률을 가진 동일한 접두사 (prefix) 와 완성된 답변이 반복적으로 생성됩니다.
• 계산 자원 낭비: Figure 1 과 Figure 2 에서 보듯, 생성된 시퀀스들 간에 중복된 토큰이 매우 많으며, 이는 불필요한 토큰 생성과 메모리/연산 자원의 낭비로 이어집니다.
• 탐색 공간의 비효율성: 제한된 계산 예산 (Budget) 하에서 동일한 고확률 경로를 반복 탐색하는 반면, 실제로는 탐색되지 않은 다른 유망한 경로를 놓치게 됩니다.

2. 제안 방법: DLE (Distinct Leaf Enumeration)

저자들은 DLE(Distinct Leaf Enumeration) 라는 새로운 결정론적 디코딩 방법을 제안합니다. 이는 확률적 샘플링을 대체하여, 잘라낸 (Truncated) 디코딩 트리를 체계적으로 탐색하고 중복되지 않는 리프 노드 (Distinct Leaves) 만을 열거하는 방식입니다.

핵심 메커니즘

1. 트리 기반 결정론적 탐색:
   • 모델의 다음 토큰 분포를 top-k, top-p, min-p, ε-sampling 등의 규칙으로 잘라낸 (Truncated) 후, 이를 트리 구조로 간주합니다.
   • 확률적 샘플링 대신, 가장 높은 확률의 가지를 우선적으로 선택하여 트리를 확장합니다.
   • 특정 분기점 (Branching point) 에서 다른 토큰을 선택하여 새로운 경로를 생성하고, 이를 다시 탐욕적 (Greedy) 으로 완성합니다.
2. 접두사 재사용 (Prefix Reuse):
   • 생성된 시퀀스들이 공통된 접두사를 공유하므로, DLE 는 이 접두사를 한 번만 생성하고 KV 캐시를 재사용합니다.
   • 기존 자기 일관성 방식은 각 시퀀스를 독립적으로 생성하여 중복된 접두사를 반복 계산하는 반면, DLE 는 이를 방지합니다.
3. 조기 종료 (Early Stopping):
   • 분기점이 다르더라도 이후 탐욕적 디코딩 과정에서 동일한 토큰 시퀀스로 수렴하는 경우를 감지합니다.
   • Figure 4 와 같이, 특정 토큰 수 (예: 10 개) 이상 반복되는 패턴이 감지되면 해당 분기를 중단하고 예산을 다른 경로에 할당합니다. 이는 동일한 최종 답을 내는 중복 경로를 제거합니다.

3. 주요 기여 (Key Contributions)

1. DLE 알고리즘 제안: 확률적 샘플링의 대안으로, 중복 없이 잘라낸 트리에서 고유한 리프 노드를 열거하는 결정론적 방법을 제시했습니다. 이는 top-p, min-p, ε-sampling 등 다양한 샘플러에 적용 가능합니다.
2. 알고리즘적 이점 (Coverage 향상): 고정된 계산 예산 하에서 DLE 는 확률적 자기 일관성보다 더 넓은 탐색 공간 (Probability Mass) 을 커버합니다. 이는 이전에 방문하지 않았던 고확률 가지를 체계적으로 탐색함으로써 달성됩니다.
3. 시스템적 이점 (효율성 향상): 중복 생성을 피하고 접두사를 재사용함으로써, 동일한 토큰 예산으로 더 많은 완성된 시퀀스를 생성할 수 있습니다. 이는 vLLM, SGLang 과 같은 최신 추론 엔진의 KV 캐시 최적화와 결합되어 추론 지연 시간을 크게 단축시킵니다.

4. 실험 결과 (Results)

저자들은 GSM8K(수학), HumanEval(코딩), MMLU-Pro(일반 추론) 데이터셋에서 Qwen2.5 및 Llama3.2 모델을 사용하여 실험했습니다.

• 성능 향상:
  • DLE 는 동일한 토큰 예산 하에서 자기 일관성 (Self-Consistency) 보다 더 높은 정확도 (maj@k, pass@k) 를 기록했습니다.
  • 특히 GSM8K 에서 약 39%, HumanEval 에서 약 47% 의 성능 향상을 보였습니다.
  • Figure 5 와 Figure 6 에서 보듯, DLE 는 더 높은 커버리지 (Coverage, 탐색된 확률 질량) 를 달성하며, 이는 높은 정확도와 직접적인 상관관계가 있음을 입증했습니다.
• 효율성 (토큰 및 시간):
  • Figure 8 에서 보듯, 동일한 정확도를 달성하기 위해 DLE 는 자기 일관성보다 훨씬 적은 수의 새로운 토큰을 생성합니다.
  • Figure 9 에서 보듯, SGLang 엔진을 사용할 때 DLE 는 캐시 히트율 (Cache Hit Rate) 이 높고, 특히 작은 배치 크기 (Batch Size) 와 높은 ε 값에서 추론 시간이 크게 단축되었습니다.
• 다양성 vs 정확도:
  • DLE 는 표면적인 N-gram 다양성은 낮을 수 있으나, 정답에 도달하는 유망한 경로를 더 효과적으로 탐색하여 실제 정확도를 높였습니다.

5. 의의 및 결론 (Significance)

이 논문은 LLM 의 테스트 시간 확장 (Test-time Scaling) 에 있어 단순히 샘플 수를 늘리는 것이 아니라, 계산 자원을 어떻게 효율적으로 배분하느냐가 중요함을 강조합니다.

• 중복 제거의 중요성: 제약된 도메인 (수학, 코딩) 에서는 정답이 좁은 범위에 집중되어 있으므로, 무작위 샘플링의 중복을 제거하고 체계적으로 탐색하는 것이 성능 향상의 핵심입니다.
• 실용적 적용 가능성: DLE 는 기존 모델의 재학습 없이 적용 가능한 "Drop-in" 대체 방법이며, 현대적인 추론 시스템 (KV 캐시 최적화 등) 과 자연스럽게 호환되어 실제 서비스 환경에서의 지연 시간 (Latency) 을 줄여줍니다.
• 미래 방향: 결정론적 탐색과 확률적 샘플링의 하이브리드 접근이나, 학습된 분기 정책 (Branching Policy) 연구 등으로 확장될 수 있는 가능성을 제시합니다.

요약하자면, DLE는 LLM 의 추론 과정에서 발생하는 불필요한 중복 계산을 제거하고, 제한된 계산 예산 내에서 더 넓은 유효한 탐색 공간을 커버함으로써 성능과 효율성을 동시에 극대화하는 혁신적인 방법론입니다.