손실을 줄이세요! 효율적인 병렬 추론을 위한 조기 경로 가지치기 학습

대형 언어 모델이 복잡한 문제를 풀 때 여러 경로를 동시에 탐색하면 비용이 기하급수적으로 증가하는데, 이 논문은 틀린 경로를 초기에 잘라내어 비용을 70% 이상 아끼는 방법을 제시한다. 특히 모델 내부의 신호를 직접 활용하여 별도의 외부 모델 없이도 매우 빠르고 정확하게 오류를 잡아낸다.

기존의 병렬 추론 방식은 여러 개의 답변 후보를 생성한 뒤 다수결로 정답을 고르는데, 이 과정에서 초기에 오류가 발생한 '죽은 경로'들도 끝까지 생성되느라 막대한 연산 자원이 낭비된다. Transformer의 Self-Attention 구조에서 각 토큰은 이전 문맥의 정보를 담고 있지만, 단순히 텍스트 결과물만 봐서는 모델이 얼마나 확신을 가지고 있는지 알기 어렵다.

이 논문은 모델 내부의 고차원적인 숨겨진 상태(hidden states)에 이미 오류 발생 징후가 포함되어 있다는 점에 착안했다. [STOP]이라는 특수 토큰을 문장 끝에 붙여 모델이 지금까지 생성한 내용이 정답으로 이어질 가능성이 높은지 스스로 평가하게 만든다. 이는 마치 숙련된 검토자가 글의 서론만 보고도 논리적 오류를 직감하는 것과 유사한 원리다.

결과적으로 틀린 경로는 초기에 제거되어 연산량이 줄어들 뿐만 아니라, 최종 답변을 고를 때 오답 후보들이 미리 걸러지므로 전체적인 정답률이 오히려 상승하는 효과를 거둔다. 이는 N²에 비례하는 연산 부담을 획기적으로 줄여주는 실용적인 접근이다.

전체 프로세스는 Launch, Check, Resume의 3단계 파이프라인으로 구성된다. 먼저 Launch 단계에서는 모든 경로에 대해 일정 길이(Lprefix)만큼의 문구를 생성하고, 이때 생성된 Key-Value 행렬을 메모리(KV Cache)에 저장한다.

Check 단계에서는 저장된 KV Cache 뒤에 학습 가능한 [STOP] 토큰을 추가로 입력한다. [STOP] 토큰의 hidden state h_i를 입력으로 받아 [h_i → LoRA 어댑터 → 선형 분류 헤드 → 스칼라 점수 s_i] 순으로 연산하여 각 경로의 성공 확률을 계산한다. 이때 기존 모델의 파라미터는 고정(frozen)한 채 가벼운 LoRA 어댑터만 사용하여 연산 오버헤드를 최소화한다.

Resume 단계에서는 계산된 점수를 바탕으로 상위 k개의 유망한 경로만 남기고 나머지는 메모리에서 즉시 삭제한다. 살아남은 경로들만 생성을 재개하여 최종 답변을 도출한다. 학습 시에는 몬테카를로 추정을 통해 각 문구가 실제로 정답에 도달할 확률을 정답 레이블로 사용하여 Binary Cross-Entropy 손실 함수로 최적화한다.

DeepSeek-R1-Distill-Qwen-1.5B 모델을 사용한 AIME24 실험에서, 가지치기를 하지 않았을 때보다 토큰 사용량을 73.88% 줄이면서도 정확도는 30.10%에서 37.92%로 크게 향상되었다. 이는 외부 모델을 사용하는 방식(Type II)이나 단순 확률값 기반 방식(Type III)보다 월등히 높은 효율성이다.

AIME25 벤치마크에서도 GPT-OSS-20B 모델에 적용했을 때, 고정된 컴퓨팅 예산 하에서 정확도를 84%에서 약 90%까지 끌어올리는 성능을 보였다. 특히 모델 크기가 커질수록, 그리고 문제의 난이도가 높을수록 STOP의 효율성이 더 두드러지는 경향이 확인되었다.

추론 지연 시간 분석 결과, STOP 모듈 추가로 인한 오버헤드는 단 0.59%에 불과했다. 이는 기존의 텍스트 기반 외부 검증 모델이 3.37%의 오버헤드를 발생시키는 것과 비교했을 때 매우 가벼운 수준이며, KV Cache를 재사용하는 설계 덕분에 중복 연산을 완벽히 제거했다.

STOP은 LRM(Large Reasoning Model)의 내부 상태를 직접 활용하는 Type IV(Internal & Learnable) 방식의 최초 구현체이다. 아키텍처적으로는 기존 모델의 어휘 사전에 [STOP] 토큰을 추가하고, 이 토큰이 처리될 때만 활성화되는 전용 LoRA 어댑터와 분류 헤드를 결합했다. 이는 모델의 원래 생성 능력을 전혀 해치지 않으면서도 검증 기능을 모듈식으로 추가할 수 있게 한다.

수학적 기반으로서, 문구 p_i의 잠재적 가치를 s_i = S(p_i | x, Theta)로 정의하고, 이를 최대화하는 동시에 생성 토큰 수를 최소화하는 것을 목표로 한다. 학습 데이터 구축 시에는 각 문구에서 32개의 무작위 경로를 생성하여 성공률을 계산하는 몬테카를로 추정 방식을 사용하여, 0 또는 1의 딱딱한 레이블(hard label)보다 분산이 낮은 부드러운 레이블(soft label)을 제공함으로써 학습의 안정성을 높였다.

또한, 최적의 경로 유지 비율(gamma)을 결정하기 위해 컴퓨팅 예산(C), 문구 길이(Lprefix), 작업 길이(Ltask) 간의 관계를 파워 로우(power-law) 수식으로 정립했다. 실험 결과, 문맥이 길어질수록 모델이 경로의 유망성을 더 정확히 판단할 수 있어 더 공격적인 가지치기(낮은 gamma)가 가능하다는 점을 입증했다.

현재 연구는 고정된 위치(예: 2048 토큰 지점)에서 한 번만 가지치기를 수행하는 단일 단계 방식에 집중하고 있다. 여러 단계에서 순차적으로 경로를 줄여나가는 다단계 가지치기나, 고정된 위치가 아닌 논리적 흐름에 따라 동적으로 체크포인트를 설정하는 방식에 대해서는 추가 연구가 필요하다.