동적 경로 스티칭으로 효율적인 추론 가속

초록

R‑Stitch는 작은 언어 모델(SLM)의 저엔트로피 토큰을 그대로 사용하고, 불확실한 고엔트로피 토큰만 대형 모델(LLM)에게 위임하는 엔트로피 기반 라우팅 기법이다. 고정 임계값 방식인 R‑Stitch와, 강화학습으로 적응적 라우팅 정책을 학습하는 R‑Stitch⁺를 제안해, CoT 추론 시 토큰 수와 토큰당 연산량을 동시에 감소시켜 3‑4배의 속도 향상을 달성하면서 정확도 손실을 최소화한다.

상세 분석

본 논문은 체인‑오브‑생각(Chain‑of‑Thought, CoT) 방식이 대형 언어 모델(LLM)의 추론 비용을 급증시킨다는 문제의식에서 출발한다. 기존의 가속화 방법은 크게(1) 추론 길이 자체를 줄이는 조기 종료·보상 학습, (2) 작은 모델을 이용한 투기적 디코딩(speculative decoding), (3) KV‑cache 최적화로 나뉜다. 특히 투기적 디코딩은 SLM이 생성한 토큰을 LLM이 검증하는 방식이지만, SLM과 LLM 간 토큰‑레벨 일치도가 낮을 경우 롤백이 빈번해 실제 속도 향상이 제한된다. 저자들은 “SLM이 정답을 맞출 때는 LLM보다 훨씬 짧은 추론 경로를 만든다”는 현상을 관찰하고, 이를 활용하기 위한 새로운 라우팅 신호로 토큰 엔트로피를 제안한다.

실험적으로 토큰 엔트로피와 오류 발생률 사이에 강한 양의 상관관계가 있음을 확인하였다. 즉, 엔트로피가 높은 토큰은 오답을 유발할 확률이 높으며, 전체 토큰 중 고엔트로피 토큰은 10% 미만에 불과해 대부분이 저엔트로피(확신) 영역에 존재한다. 이러한 통계적 특성을 바탕으로 R‑Stitch는 다음과 같은 흐름을 갖는다. 1) 디코딩은 SLM으로 시작한다. 2) 현재 토큰의 정규화 엔트로피 H가 사전에 정의된 임계값 τ 이하이면 SLM의 출력을 그대로 채택한다. 3) H가 τ를 초과하면 해당 토큰을 버리고 동일 프리픽스를 LLM에 전달해 재생성한다. 4) LLM이 저엔트로피 토큰을 생성하면 제어를 다시 SLM에게 넘겨, 두 모델이 엔트로피에 따라 양방향으로 전환한다.

키‑밸류(KV) 캐시 관리에서도 효율성을 극대화한다. SLM·LLM 각각 독립적인 KV‑cache를 유지하고, 모델 전환 시 이미 생성된 토큰에 대한 캐시는 재사용한다. 즉, 이전에 사용한 모델이 다시 활성화될 때는 새로 생성된 토큰만을 프리필(pre‑fill)하여 캐시 재구성을 최소화한다.

R‑Stitch⁺는 고정 τ 대신 경량 라우터를 도입한다. 라우터는 현재 토큰의 엔트로피와 추가 메타 정보를 입력으로 받아, LLM 호출 여부를 결정한다. 라우터는 지연( latency )을 고려한 보상 함수를 사용해 강화학습(RL)으로 학습되며, 이를 통해 다양한 하드웨어·예산 환경에 맞는 최적의 효율‑정확도 트레이드오프를 자동으로 찾는다.

성능 평가에서는 DeepSeek‑R1‑Distill‑Qwen 7B, 14B, 32B 모델을 LLM으로, L1‑Short·Distill·Oat‑1.5B 등을 SLM으로 사용하였다. AMC, MATH, Olympiad, Minerva 등 5개 수학·논리 벤치마크에서 R‑Stitch는 평균 3.0×~4.1×의 속도 향상을 보였으며, 정확도는 전체 LLM 디코딩과 거의 동일한 수준을 유지한다. 특히 토큰‑레벨 일치도가 낮은 경우에도 엔트로피 기반 라우팅 덕분에 롤백이 거의 발생하지 않아, 기존 투기적 디코딩 대비 일관된 가속 효과를 얻었다.

요약하면, 본 연구는 (1) 엔트로피가 오류와 강하게 연관된다는 실증적 근거, (2) 엔트로피 기반 토큰‑레벨 라우팅을 통한 SLM·LLM 협업 프레임워크, (3) 라우터 학습을 통한 적응형 정책(R‑Stitch⁺)이라는 세 축을 제시한다. 이 접근법은 추가 학습 없이 바로 적용 가능하고, KV‑cache 재활용을 통해 전환 오버헤드도 최소화한다는 점에서 실용성이 높다.