TAPS: 추측 샘플링을 위한 작업 인지형 제안 분포 연구

왜 중요한가

LLM 추론 가속 기법인 추측 샘플링에서 보조 모델의 '학습 데이터 분포'가 성능에 미치는 영향을 최초로 심도 있게 분석했다. 특정 작업에 특화된 보조 모델을 사용하고 이를 추론 시점에 지능적으로 조합함으로써, 기존의 범용 보조 모델보다 훨씬 높은 가속 효율을 달성할 수 있음을 입증했다.

핵심 기여

보조 모델의 도메인 특화 효과 입증

수학(MathInstruct)과 대화(ShareGPT) 데이터로 각각 학습된 보조 모델이 해당 도메인의 벤치마크에서 타겟 모델과의 일치도가 급격히 상승하며 수락 길이가 최대 2배 이상 차이남을 확인했다.

추론 시점의 보조 모델 조합 전략 제안

여러 전문 보조 모델이 존재할 때, 가중치를 평균내는 방식보다 신뢰도 기반 라우팅(Confidence Routing)이나 트리 병합 검증(Merged-tree Verification)이 훨씬 우수한 성능을 보임을 증명했다.

신뢰도 기반 라우팅의 유효성 확인

엔트로피보다 모델의 신뢰도(Confidence) 점수가 현재 작업에 가장 적합한 보조 모델을 선택하는 데 훨씬 명확하고 신뢰할 수 있는 신호임을 실험적으로 밝혀냈다.

핵심 아이디어 이해하기

LLM은 토큰을 하나씩 생성하느라 느리다. 이를 해결하는 '추측 샘플링'은 가벼운 보조 모델이 미리 여러 토큰을 제안하고 큰 모델이 이를 검증하는 방식이다. 기존 연구는 보조 모델의 '구조'를 개선하는 데 집중했지만, 이 논문은 보조 모델이 타겟 모델의 사고방식을 얼마나 잘 흉내 내는지, 즉 '학습 데이터의 일치성'이 핵심임을 지적한다.

보조 모델이 타겟 모델의 다음 토큰 확률 분포(Softmax 결과값)를 정확히 예측할수록 수락 확률이 높아진다. 예를 들어 수학 문제를 풀 때는 수학 데이터로 훈련된 보조 모델이 타겟 모델의 논리 전개를 더 잘 예측한다. 논문은 이를 위해 특정 분야에 특화된 전문 보조 모델들을 준비하고, 입력된 프롬프트에 따라 가장 적합한 모델을 고르거나(Routing), 여러 모델의 제안을 합쳐서 검증하는 방식을 제안한다.

결과적으로 보조 모델의 파라미터를 물리적으로 섞는 것보다, 각자의 전문성을 유지한 채 추론 시점에 상황에 맞는 모델의 제안을 사용하는 것이 수락 길이를 늘려 전체 추론 속도를 실질적으로 가속화한다.

방법론

HASS와 EAGLE-2라는 두 가지 추측 샘플링 프레임워크를 기반으로 실험을 설계했다. 타겟 모델은 Llama-3-8B-Instruct로 고정하고, 약 0.8B 파라미터 규모의 보조 모델을 수학(MathInstruct)과 대화(ShareGPT) 데이터로 각각 학습시켜 전문성을 부여했다.

신뢰도 기반 라우팅(Confidence Routing)은 각 보조 모델이 생성한 토큰 트리의 평균 신뢰도를 계산한다. [각 노드의 예측 확률값 입력 → 산술 평균 연산 → 트리별 점수 출력 → 점수가 높은 트리를 선택하여 타겟 모델에 전달] 과정을 거친다. 이는 가중치를 직접 섞는 방식보다 작업 전환에 훨씬 유연하게 대응한다.

트리 병합 검증(Merged-tree Verification)은 두 보조 모델의 제안을 하나의 공유 루트 아래 병합한다. [두 트리의 노드들을 연결 → 조상 관계를 보존하는 어텐션 마스크 생성 → 타겟 모델이 한 번의 패스로 모든 경로를 병렬 검증 → 수락된 가장 긴 경로 선택] 방식으로 작동한다. 이는 제안의 다양성을 높여 수락 확률을 극대화하지만, 검증해야 할 트리 크기가 커지는 비용이 발생한다.

주요 결과

단일 도메인 특화 효과가 뚜렷하게 나타났다. HASS 기준 온도 0에서 ShareGPT 모델은 MT-Bench에서 3.98의 수락 길이를 기록했으나 MathInstruct 모델은 2.90에 그쳤다. 반대로 GSM8K 수학 벤치마크에서는 MathInstruct 모델이 5.02로 ShareGPT 모델(4.09)을 크게 앞섰다.

조합 전략 비교에서 가중치 평균(Averaged) 방식은 성능이 급격히 저하되어 가장 낮은 효율을 보였다. 반면 트리 병합(Merged Trees) 방식은 HASS와 EAGLE-2 모두에서 가장 높은 평균 수락 길이(각각 5.11, 5.03)를 달성하며 단일 모델이나 라우팅 방식보다 우수한 성능을 입증했다.

라우팅 신호 분석 결과, 신뢰도(Confidence) 기반 선택이 엔트로피 기반보다 훨씬 정확했다. 신뢰도 라우팅은 수학 작업에서 MathInstruct 모델을 90% 이상의 확률로 정확히 선택해냈으나, 엔트로피는 선택 비율이 모호하게 나뉘어 작업 구별 능력이 떨어지는 것으로 나타났다.

기술 상세

보조 모델은 1개의 트랜스포머 레이어와 4096의 은닉층 크기를 가진 0.8B 규모의 디코더 구조를 채택했다. 타겟 모델과 동일한 토크나이저를 공유하여 토큰화 불일치 문제를 방지했다. 학습 시에는 20 에폭 동안 3e-5의 학습률을 사용했다.

HASS 백본은 Top-K 증류 손실을 사용하여 타겟 모델의 상위 K개 토큰 분포를 모사하도록 학습된다. [타겟 모델의 확률 분포 q와 보조 모델의 분포 p 입력 → 상위 K개 토큰에 대해 -sum(q * log(p)) 계산 → 손실값 도출] 과정을 통해 타겟 모델의 예측 경향을 효율적으로 학습한다.

EAGLE-2는 토큰 대신 특징 벡터(Feature) 수준에서 다음 단계를 예측하며, 문맥에 따라 동적으로 트리를 구성한다. 본 연구는 이러한 백본의 구조적 차이보다 학습 데이터의 분포가 수락 길이에 더 지배적인 영향을 미친다는 것을 실험적으로 확인했다.

한계점

Llama-3-8B라는 단일 타겟 모델과 두 가지 도메인에 한정된 실험 결과이다. 수락 길이를 주요 지표로 삼았으나, 트리 병합 시 발생하는 연산 오버헤드를 포함한 최종적인 엔드투엔드 지연 시간(Latency)에 대한 시스템적 분석은 보완이 필요하다.

실무 활용

다양한 작업(수학, 코딩, 일반 대화 등)을 처리해야 하는 LLM 서비스에서 작업별로 특화된 가벼운 보조 모델들을 구비하고, 추론 시점에 이를 동적으로 선택하거나 병합하여 추론 속도를 최적화할 수 있다.

수학적 추론과 일반 대화가 섞인 교육용 챗봇의 응답 속도 가속
여러 전문 도메인(법률, 의료 등) 보조 모델을 활용한 다목적 엔터프라이즈 LLM 서버 최적화
추론 시점의 신뢰도 점수를 활용한 지능형 보조 모델 라우팅 시스템 구축

코드 공개 여부: 공개

코드 저장소 보기

키워드

Speculative Decoding(추측 샘플링)Draft Model(보조 모델)Acceptance Length(수락 길이)Confidence Routing(신뢰도 라우팅)Task Specialization(작업 특화)

용어 해설

추측 샘플링(Speculative Decoding): — 가벼운 보조 모델(Drafter)이 여러 토큰을 미리 제안하고, 거대한 타겟 모델이 이를 한 번에 병렬로 검증하여 LLM의 추론 속도를 높이는 기법이다. 타겟 모델의 출력 분포를 유지하면서도 순차적인 생성 병목을 해결하는 핵심 최적화 기술이다.
수락 길이(Acceptance Length): — 추측 샘플링에서 보조 모델이 제안한 토큰들 중 타겟 모델에 의해 최종적으로 승인된 연속된 토큰의 평균 개수이다. 이 수치가 높을수록 보조 모델이 타겟 모델의 예측을 잘 모사하고 있음을 의미하며, 전체 시스템의 가속 성능과 직결된다.
증류 손실(Distillation Loss): — 타겟 모델(교사)의 지식을 보조 모델(학생)에게 전달하기 위해 두 모델의 출력 확률 분포 차이를 계산하는 손실 함수이다. 보조 모델이 타겟 모델의 복잡한 예측 패턴을 더 적은 파라미터로 효율적으로 학습하게 돕는다.

코드 예제

python

def _merge_trees(
    draft_tokens1, retrieve_indices1, tree_mask1, tree_pos1,
    draft_tokens2, retrieve_indices2, tree_mask2, tree_pos2,
):
    n1 = draft_tokens1.shape[1] - 1
    n2 = draft_tokens2.shape[1] - 1
    N = n1 + n2 + 1
    // ...(중략)
    merged_draft = torch.cat([draft_tokens1, draft_tokens2[0, 1:][None]], dim=1)
    merged_mask = torch.zeros(N, N, device=device, dtype=dtype)
    merged_mask[0, 0] = 1.0
    merged_mask[1:n1 + 1, :n1 + 1] = tree_mask1[0, 0, 1:, :]
    merged_mask[n1 + 1:, 0] = 1.0
    merged_mask[n1 + 1:, n1 + 1:] = tree_mask2[0, 0, 1:, 1:]
    // ...(중략)
    return merged_draft, merged_retrieve, merged_mask, merged_pos

두 개의 서로 다른 보조 모델이 생성한 토큰 트리를 하나의 병합된 트리 구조로 합쳐 타겟 모델이 한 번에 검증할 수 있게 만드는 핵심 로직

TAPS: 추측 샘플링을 위한 작업 인지형 제안 분포 연구

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

코드 예제

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드