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

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

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)에 대한 시스템적 분석은 보완이 필요하다.