추론 모델은 언제 생각을 멈춰야 할지 내재적으로 알고 있는가?

최근 DeepSeek-R1과 같은 모델들은 정답을 찾기 위해 매우 긴 추론 과정을 거치지만, 이는 종종 불필요한 연산 낭비로 이어집니다. 이 논문은 모델이 이미 최적의 종료 시점을 알고 있다는 사실을 발견하고, 이를 활용해 정확도는 높이면서 추론 비용은 절반 가까이 줄이는 획기적인 방법론을 제시합니다.

최근의 대형 추론 모델들은 문제를 풀 때 '생각의 사슬(CoT)'을 길게 늘려 성능을 높이지만, 정답을 찾은 뒤에도 불필요한 검토를 반복하는 '과잉 사고' 문제를 겪는다. 이는 모델이 생성하는 수많은 경로 중 어떤 것이 가장 효율적인지 선택하지 못하는 기존 샘플링 방식의 한계에서 기인한다. 연구진은 모델이 각 토큰을 생성할 때 부여하는 확률값들을 분석한 결과, 모델이 정답에 가까운 짧은 경로에 이미 높은 점수를 주고 있다는 점에 주목했다. 즉, 모델은 무엇이 정답인지 이미 알고 있지만, 우리가 그 목소리를 듣지 못하고 있었던 것이다. SAGE는 모델의 이러한 '자기 확신'을 나침반 삼아 가장 유망한 경로만을 골라내어 탐색한다. 이를 통해 모델이 정답을 확신하는 순간 즉시 생각을 멈추게 함으로써, 마치 숙련된 수학자가 군더더기 없이 깔끔한 풀이를 내놓는 것과 같은 효과를 구현했다.

SAGE(Self-Aware Guided Efficient Reasoning)는 토큰 단위의 탐색과 단계(Step) 단위의 확장을 결합한 하이브리드 디코딩 전략이다. 먼저 모델의 누적 로그 확률 $\Phi$를 기준으로 상위 $m$개의 유망한 경로를 유지하며 탐색을 진행한다. $\Phi(y_{\le k}) = \frac{1}{k} \sum_{i=1}^k \log \pi_\theta(y_i | y_{<i}, x)$ 공식을 사용하여 [각 토큰의 생성 확률 로그값을 입력으로] → [전체 토큰 수 $k$로 나누어 평균을 내는 연산을 수행해] → [시퀀스 전체의 평균 확신도 점수를 얻고] → [이 점수가 높을수록 모델이 해당 경로를 일관되게 신뢰하고 있음을 판별한다].

탐색 종료를 위해 'Tolerance accept rank Ratio(TR)'라는 임계값을 도입했다. 추론 단계 끝에 </think> 토큰이 나타날 때, 해당 토큰의 확률 순위가 TR 범위 내에 있다면 모델이 충분한 확신을 가진 것으로 판단하고 즉시 탐색을 종료한다. SAGE-RL은 이 과정을 강화학습(RLVR)에 통합한다. 그룹 내 8개의 샘플 중 2개는 SAGE로 생성하여 높은 보상을 받게 하고, 나머지 6개는 일반 샘플링을 사용하여 모델이 자연스럽게 SAGE가 발견한 효율적인 추론 패턴을 닮아가도록 유도한다.

DS-1.5B 모델에 SAGE-RL을 적용한 결과, MATH-500 벤치마크에서 정확도가 향상되었으며 특히 AIME 2025에서는 기존 대비 6.2%의 성능 향상을 보였다. 가장 인상적인 결과는 효율성 개선으로, 모든 테스트 벤치마크에서 평균 44.1%의 토큰 절감을 달성했다. 이는 모델이 더 똑똑해지면서도 더 적은 자원을 사용하게 되었음을 의미한다.

Qwen3-8B와 같은 대규모 모델에서도 효과가 입증되었다. Minerva 데이터셋에서 정확도를 유지하면서도 추론 길이를 원래의 45.7% 수준으로 대폭 압축했다. 또한 RFCS(Ratio of the First Correct Step) 지표 분석을 통해, 학습된 모델들이 정답을 도출한 직후 불필요한 단계를 생성하지 않고 즉시 생각을 멈추는 경향이 뚜렷해졌음을 확인했다.

SAGE는 Vanilla Beam Search의 구조를 차용하되, </think> 토큰의 순위 기반 조기 종료 메커니즘을 추가하여 차별화했다. 핵심은 누적 로그 확률 $\Phi$를 점수 함수로 사용하여 모델의 자기 인식을 직접적으로 활용한다는 점이다. Prior work들이 단순히 길이를 제한하거나 페널티를 주는 방식이었다면, SAGE는 모델의 확률 분포 내에 숨겨진 최적의 경로를 '발견'하는 데 집중한다.

학습 단계인 SAGE-RL은 GRPO(Group Relative Policy Optimization) 아키텍처를 기반으로 한다. 하이브리드 롤아웃 전략을 통해 SAGE가 생성한 고품질/저비용 샘플이 Advantage 계산 시 높은 가중치를 받게 함으로써, 정책 모델($\pi_\theta$)이 효율적인 추론 궤적을 학습하도록 설계되었다. 이는 KL Divergence와 Entropy 변화 추이를 통해 모델이 더 확신 있는 추론 패턴으로 수렴함을 이론적으로 뒷받침한다.

SAGE는 탐색 폭($m$)이 커질수록 메모리 사용량이 약 $2m$배 증가하는 공간 복잡도 문제를 안고 있다. 또한 vLLM 엔진 환경에서 탐색 폭이 2를 초과할 경우 하드웨어 제약으로 인해 추론 시간이 급격히 늘어나는 현상이 관찰되어, 효율과 성능 사이의 적절한 트레이드오프 설정이 필요하다.