Truncated Decoding Tree의 결정론적 탐색을 통한 효율적인 Test-Time Inference

TL;DR Highlight

Self-consistency의 중복 샘플링 낭비를 없애는 결정론적 트리 탐색 디코딩 기법 DLE로 수학/코드 추론 성능과 속도를 동시에 개선

Who Should Read

LLM 추론 비용을 줄이면서 정확도를 높이고 싶은 ML 엔지니어. 특히 수학 문제 풀이나 코드 생성 태스크에서 self-consistency(여러 답안 생성 후 다수결)를 쓰고 있는 개발자.

Core Mechanics

Self-consistency는 같은 prefix(앞부분 토큰)를 반복 생성하는 낭비가 심함. 코드 생성 기준으로 k=32일 때 전체 prefix 토큰의 약 17%가 중복 생성됨.
DLE(Distinct Leaf Enumeration)는 생성 가능한 시퀀스를 트리 구조로 보고, 중복 없이 새로운 leaf(완성된 시퀀스)만 탐색하는 결정론적 방법.
기존 top-p, min-p, ε-sampling 등 어떤 truncation 방식과도 호환되는 drop-in 교체재로 작동하며, 추가 모델 호출이나 별도 평가 없이 동작함.
같은 시퀀스 수 기준으로 DLE가 stochastic self-consistency보다 높은 probability mass coverage를 달성하고, coverage가 높을수록 정확도도 높아지는 양의 상관관계가 있음.
KV cache 공유 prefix 재사용이 가능해서 SGLang/vLLM 같은 inference 엔진의 prefix caching을 직접 활용 가능. 실제 cache hit rate가 이론적 최대치에 근접함.
10개 연속 토큰이 이전 sibling branch와 일치하면 해당 branch를 조기 종료하는 early stopping으로 낭비 방지. 조기 종료된 branch의 85% 이상이 동일한 최종 답을 냈을 것으로 확인됨.

Evidence

Qwen2.5-0.5B-Instruct 기준 GSM8K maj@8에서 DLE(ε-sampling) 44.05% vs Self-consistency(ε-sampling) 40.94%, 약 3~9% 향상. HumanEval pass@8에서 DLE 52.44% vs Self-consistency 47.56%, 약 4~7% 향상.
동일 정확도 약 34% 달성 시 ε-sampling self-consistency는 질문당 평균 150토큰 이상 신규 생성이 필요하지만, DLE는 20토큰 미만으로 동일 성능 달성(k=32 기준).
고정 token budget 조건에서 DLE가 ε-sampling 대비 더 많은 완성 시퀀스를 생성해 majority voting 정확도 향상. token budget 2048 기준 DLE 약 38% vs ε-sampling 약 36%.
Qwen2.5-7B-Instruct GSM8K maj@2에서 DLE(ε-sampling) 84.91% vs Self-consistency(ε-sampling) 81.43%, 대형 모델에서도 일관된 성능 향상 확인.

How to Apply

현재 vLLM이나 SGLang으로 self-consistency를 구현 중이라면, DLE는 동일한 truncation sampler(top-p, min-p, ε-sampling) 설정을 유지하면서 샘플링 루프만 교체하면 됨. SGLang의 RadixAttention이나 vLLM의 automatic prefix caching과 함께 쓰면 실제 latency 감소 효과도 얻을 수 있음.
수학/코드처럼 정답이 협소한(narrow) 도메인에서 k=8~32개 시퀀스로 majority voting을 하는 경우, DLE의 PROBFIRST 또는 DIVFIRST 전략을 사용하면 동일 토큰 예산에서 더 많은 distinct 후보를 탐색 가능. 특히 메모리 제약으로 batch size가 1~2인 환경에서 runtime 이득이 가장 큼.
ε 값을 조정해 prefix reuse 비율을 컨트롤할 수 있음. ε를 크게 설정할수록 branching이 나중에 일어나 prefix 공유가 늘어나고 토큰 효율이 높아지지만 탐색 다양성은 줄어드므로, 태스크 특성에 맞게 ε=0.05~0.3 범위에서 튜닝 권장.

Code Example

snippet

# DLE 핵심 로직 의사코드 (ε-sampling 기준)
# Step 1: greedy generation으로 첫 번째 leaf 생성하며 branching points 추적
def dle_decode(model, prompt, k, epsilon=0.05):
    leaves = []
    # 우선순위 큐: (음수 확률, prefix, 다음 탐색할 alternative token)
    branch_queue = []  # (score, prefix_tokens, alt_token)
    
    # 첫 greedy generation
    prefix, branch_points = greedy_generate_with_branches(model, prompt, epsilon)
    leaves.append(prefix)
    
    # branch points를 확률 기준으로 큐에 추가 (PROBFIRST)
    for pos, token, prob in branch_points:
        heapq.heappush(branch_queue, (-prob, prefix[:pos], token))
    
    while len(leaves) < k and branch_queue:
        neg_prob, branch_prefix, alt_token = heapq.heappop(branch_queue)
        
        # 대안 토큰으로 시작해서 greedy completion
        new_prefix = branch_prefix + [alt_token]
        
        # Early stopping: 이전 sibling과 10토큰 이상 겹치면 중단
        if check_early_stop(new_prefix, leaves, n=10):
            continue
        
        new_leaf, new_branches = greedy_generate_with_branches(
            model, new_prefix, epsilon, 
            kv_cache=get_shared_prefix_cache(new_prefix)  # prefix cache 재사용
        )
        leaves.append(new_leaf)
        
        for pos, token, prob in new_branches:
            path_prob = abs(neg_prob) * prob  # accumulated path probability
            heapq.heappush(branch_queue, (-path_prob, new_leaf[:pos], token))
    
    # Majority voting (uniform weight)
    answers = [extract_answer(leaf) for leaf in leaves]
    return Counter(answers).most_common(1)[0][0]

# 실제 사용 예시 (SGLang 기반)
# SGLang의 RadixAttention이 자동으로 prefix KV cache를 공유해줌
import sglang as sgl

@sgl.function
def dle_with_sglang(s, question, leaves_prefixes):
    s += sgl.system("You are a math solver.")
    s += sgl.user(question)
    # 공유 prefix까지는 cache hit, 이후부터만 새로 생성
    for branch_prefix in leaves_prefixes:
        with s.fork() as branch:
            branch += sgl.assistant(branch_prefix)
            branch += sgl.gen("completion", max_tokens=512)

Terminology

Self-consistency같은 질문에 여러 번 답변을 생성한 뒤 가장 많이 나온 답을 선택하는 기법. 시험 답안을 여러 번 풀어보고 가장 많이 나온 답을 선택하는 것과 비슷.

Truncated sampling다음 토큰 후보 중 확률이 낮은 것들을 잘라내고 상위 후보들 중에서만 샘플링하는 방법. top-p, min-p, ε-sampling이 여기에 해당.

KV cacheLLM이 이미 처리한 토큰들의 중간 계산값을 저장해두는 캐시. 같은 prefix가 다시 나오면 재계산 없이 저장된 값을 쓸 수 있어 속도가 빨라짐.

Coverage생성한 시퀀스들이 전체 가능한 시퀀스 공간의 확률을 얼마나 커버하는지 측정하는 지표. 높을수록 중요한 후보를 더 많이 탐색했다는 의미.

prefix caching여러 생성 요청이 공통 앞부분(prefix)을 공유할 때, 그 부분의 KV cache를 재사용하는 최적화 기법. vLLM, SGLang 등이 지원함.

maj@kk개의 답변을 생성해서 다수결로 최종 답을 선택했을 때의 정확도. k가 클수록 더 많은 답변을 생성해 투표함.

pass@kk개의 답변 중 하나라도 정답이 있으면 성공으로 보는 지표. 주로 코드 생성 평가에서 사용.