당신의 추론 모델은 언제 생각을 멈춰야 할지 이미 알고 있는가?

최신 추론 모델들이 성능 향상을 위해 사고 과정을 길게 늘리는 경향이 있지만, 이는 막대한 연산 비용과 지연 시간을 초래한다. 이 논문은 모델이 스스로 언제 생각을 멈춰야 할지 이미 알고 있다는 사실을 발견하고, 이를 활용해 성능 저하 없이 추론 효율성을 극대화하는 방법을 제시한다.

기존의 추론 모델은 '더 길게 생각할수록 더 똑똑하다'는 가정하에 긴 사고 사슬(CoT)을 생성한다. 하지만 이는 불필요한 반복과 오류를 포함하는 '오버씽킹' 문제를 야기한다. 연구진은 모델이 여러 답변 후보를 생성할 때, 정답으로 이어지는 짧고 명확한 경로에 대해 통계적으로 더 높은 확신(Log-probability)을 보인다는 점에 주목했다.

SAGE는 모델의 이러한 내재적 확신을 나침반 삼아 작동한다. 단순히 다음 토큰의 확률만 보는 것이 아니라, 지금까지 생성한 전체 경로의 평균 로그 확률을 계산하여 가장 유망한 가지만 골라낸다. 이는 마치 미로 찾기에서 막다른 길로 갈 확률이 높은 경로는 즉시 포기하고, 목적지에 가까워 보이는 경로에 집중하는 것과 같다.

결과적으로 모델은 불필요한 검토 과정을 생략하고 정답이 도출되는 즉시 생각을 멈추는 법을 배운다. 이는 추론의 '양'보다 '질'에 집중하게 함으로써, 더 적은 토큰을 사용하면서도 논리적 오류에 빠질 확률을 줄여 전체적인 정확도를 향상시킨다.

SAGE는 토큰 단위의 탐색(TSearch)과 단계 단위의 탐색 두 가지 방식으로 구현된다. 핵심은 누적 로그 확률 $\Phi$를 점수 함수로 사용하는 것이다. 각 단계 $k$에서 토큰 $y_i$의 로그 확률 $\phi(y_i | y_{<i}, x)$를 모두 더한 뒤 토큰 수 $k$로 나눈 평균값 $[\sum \log P \to \text{평균 계산} \to \Phi]$을 산출한다. 이 값이 높을수록 모델이 해당 경로를 신뢰한다는 의미다.

탐색 종료 조건으로 '허용 순위 비율(Tolerance accept rank Ratio, TR)'을 도입했다. 사고 종료 토큰인 </think>가 상위 $h$개 후보 안에 들고, 그 순위가 설정된 TR 임계값보다 높을 때만 추론을 종료한다. 이는 모델이 충분히 확신할 때만 생각을 멈추게 하는 안전장치 역할을 한다.

SAGE-RL은 이를 강화학습의 롤아웃(Rollout) 단계에 적용한다. 그룹 기반 강화학습(GRPO/GSPO) 시, 8개의 샘플 중 2개는 SAGE 방식으로 생성하여 고품질의 효율적인 경로를 확보하고, 나머지는 일반 샘플링을 사용한다. 모델은 SAGE가 찾아낸 '짧지만 정확한' 경로에서 더 높은 보상을 얻게 되므로, 자연스럽게 효율적인 사고 패턴을 학습하게 된다.

MATH-500 벤치마크에서 SAGE-RL을 적용한 모델은 기존 RLVR 대비 토큰 사용량을 평균 44.1% 절감하면서도 정확도는 2.1% 향상시켰다. 특히 AIME 2025와 같은 고난도 수학 문제에서 DeepSeek-R1 기반 모델의 성능을 유의미하게 개선했다.

효율성 측면에서 SAGE-RL로 튜닝된 모델은 표준 Pass@1 추론 시 기존 모델보다 훨씬 짧은 답변을 생성한다. 예를 들어 DS-1.5B 모델에 적용했을 때, MATH-500에서의 답변 길이는 4,882토큰에서 2,915토큰으로 줄어들었으며, 토큰당 효율성(TE)은 약 70% 이상 개선되었다.

Ablation Study 결과, 단순한 빔 서치(Beam Search)는 답변 길이를 줄이지 못하고 정확도 향상도 미미했으나, SAGE의 확신 기반 필터링은 정확도와 간결함을 동시에 잡는 데 필수적임이 확인되었다. 또한 모델의 크기가 커지고 문제의 난이도가 높을수록 SAGE-RL의 성능 향상 폭이 더 크게 나타났다.

SAGE는 기존의 Greedy 또는 Random Sampling이 가진 한계를 RFCS(Ratio of the First Correct Step) 지표로 정량화했다. RFCS는 정답이 처음 등장한 시점을 전체 단계로 나눈 값으로, 1보다 작을수록 불필요한 '뒷북' 추론이 많음을 의미한다. 실험 결과 대부분의 모델이 0.5~0.6 수준의 낮은 RFCS를 보였으며, 이는 절반 가까운 토큰이 낭비되고 있음을 시사한다.

SAGE-RL의 목적 함수는 GRPO와 GSPO의 구조를 그대로 유지하되, 롤아웃 데이터 생성 방식만 하이브리드로 변경했다. SAGE를 통해 생성된 샘플은 높은 보상을 받을 확률이 높으므로, 정책 모델의 업데이트 방향이 효율적인 경로로 유도된다. 이때 KL 발산(KL Divergence)이 일반 RL보다 가파르게 증가하는데, 이는 모델이 기존의 비효율적인 방식에서 벗어나 새로운 효율적 패턴을 적극적으로 수용하고 있음을 보여준다.

구현 측면에서는 vLLM 엔진을 사용하여 2m개의 추론 경로를 병렬로 생성함으로써 시간 복잡도 증가를 최소화했다. 비록 메모리 사용량은 약 2m배 증가하지만, KV 캐시 재사용과 연속 배치(Continuous Batching) 기술을 통해 실제 추론 지연 시간은 오히려 기존 방식보다 28.7%~40% 이상 단축되는 결과를 얻었다.

SAGE는 병렬 탐색을 위해 더 많은 GPU 메모리를 요구하며, 탐색 폭(Exploration Width) $m$이 커질수록 메모리 오버헤드가 증가한다. 또한 매우 쉬운 문제에서는 토큰 절감 효과가 크지만, 극도로 어려운 문제에서는 오히려 더 많은 사고 단계가 필요할 수 있어 문제 난이도에 따른 적응형 설정이 향후 과제로 남아있다.