기술 요약: OS-Pruner

문제 정의

Chain-of-Thought (CoT) 프롬프팅을 활용하는 대규모 언어 모델(LLM)은 복잡한 추론 작업에서 상당한 성공을 거두었습니다. 그러나 이러한 모델들은 종-종 "계산적 과잉 사고(computational overthinking)"를 보이며, 최종 답변의 정확도를 개선하지 못한 채 불필요한 추론 단계를 생성하여 지연 시간과 추론 비용을 증가시킵니다. 이러한 동작에는 논거의 재진술, 이미 결정된 결론의 재검증, 또는 가치가 낮은 계산의 지속 등이 포함됩니다. 최근 연구들은 CoT 궤적을 가지치기(pruning)할 수 있음을 시사하지만, 기존 솔루션들은 다음과 같은 구체적인 한계를 가집니다:

• 정적 예산(Static Budgets): 예산 강제(budget forcing)와 같은 방식은 외부적으로 고정된 사고 제한을 부과하며, 진화하는 추론 상태를 무시합니다.
• 모델 측 압축(Model-Side Compression): 길이 페널티 미세 조정(length-penalized fine-tuning)이나 토큰 건너뛰기(token skipping)와 같은 기술은 베이스 모델의 비용이 많이 드는 재학습을 요구하며, 비추론 작업에서의 성능 저하를 초래할 수 있습니다.
• 휴리스틱 조기 종료(Heuristic Early Stopping): 현재의 동적 방법들은 주로 분류나 임계값(예: 신뢰도가 고정된 값을 초과할 때 중단)에 의존하며, 이는 추가 토큰의 비용과 답변 품질의 기대 개선 사이의 트레이드오프를 고려하지 못합니다.

저자들은 CoT 가지치기가 근본적으로 **순차적 결정 문제(sequential decision problem)**이며, **최적 정지 문제(optimal stopping problem)**로 프레임워크화하는 것이 가장 적합하다고 주장합니다. 즉, 시스템은 각 단계에서 추가 토큰을 생성하는 비용이 정확도의 기대 이득을 정당화하는지 여부를 결정해야 합니다.

방법론: OS-Pruner

OS-Pruner는 CoT 가지치기를 최적 정지 문제로 공식화하기 위해 설계된 경량 플러그인 프레임워크입니다. 이는 동결된(frozen) 베이스 추론 모델을 수정하지 않고 작동합니다.

1. 문제 공식화

입력 문제 $x$와 부분적으로 생성된 CoT 접두사(prefix) $y_{\le i}$가 주어졌을 때, OS-Pruner는 정지 확률을 출력하는 파라미터화된 정지 정책 $\pi_\theta$를 학습합니다. 목적 함수는 정확도와 생성된 길이를 교환하는 명시적인 효용 함수를 최대화하는 것입니다:
$$r(y_{\le i}|x) = A(y_{\le i}|x) - \lambda L(y_{\le i})$$
여기서:

• $A(y_{\le i}|x)$는 접두사로부터 도출된 답변의 정확도(정답이면 1, 오답이면 0)입니다.
• $L(y_{\le i})$는 단계 $i$까지 생성된 토큰의 수입니다.
• $\lambda$는 길이에 대한 페널티를 제어하는 스칼라 파라미터로, 사용자가 정확도-효율성 프런티어 상의 특정 작동 지점을 선택할 수 있게 합니다.

2. 이론적 이점

본 논문은 OS-Pruner를 신뢰도 임계값 방식과 **가치 함수 관점(Value Function View)**을 통해 차별화합니다.

• 분류 기반 방식은 $A(y_{\le i}|x) \ge \gamma$ (고정된 임계값)일 때 중단합니다.
• **최적 정지(OS-Pruner)**는 즉각적인 정지 가치가 계속 진행할 가치(continuation value)보다 클 때 중단합니다: $A(y_{\le i}|x) \ge C_\lambda(x, y_{\le i})$.
  계속 진행 가치 $C_\lambda$는 추가 토큰의 비용을 제외한 기대 미래 보상을 나타냅니다. 저자들은 고정 임계값 분류기가 미래의 추론 경로의 상태 의존적 가치를 고려하지 못하기 때문에, 최적 정지 정책과 비교했을 때 임의로 큰 성능 격차를 가질 수 있음을 증명했습니다 (정리 1).

3. 구현 및 아키텍처

• 데이터 전처리: 베이스 모델이 전체 추론 궤적을 생성합니다. 학습을 위해 시스템은 각 추론 단계 이후에 개입하여, 모델이 최종 답변을 출력하도록 강제하고(답변 섹션에서의 "생각하기" 방지), Math-Verify를 사용하여 정답과 대조하여 채점합니다. 이를 통해 모든 접두사에 대한 세밀한 감독(supervision)이 가능해집니다.
• 정책 아키텍처: 정책은 접두사의 마지막 은닉 상태 $h_l$에 적용되는 경량 선형 헤드($\sigma(W^\top h_l + b)$)입니다. 성능을 높이기 위해 베이스 모델의 마지막 $n$(통상 2개) 자가 주의(self-attention) 레이어를 미세 조정하는 반면, 나머지 모델은 동결된 상태를 유지합니다.
• 학습 효율성: 베이스 모델은 인과적(causal)이므로, 한 단락의 은닉 상태는 후속 토큰의 영향을 받지 않습니다. 이를 통해 전체 추론 궤적을 단 한 번의 순방향 패스(forward pass)로 처리할 수 있으며, 선형 레이어는 관련 은닉 상태들에 병렬로 적용됩니다.
• 점진적 학습: 정지하지 않고 계속 머무는 지역 최적점(local minima)에 빠지는 것을 방지하기 위해, 학습은 높은 $\lambda$(길이에 대한 높은 비용)에서 시작하여 점진적으로 낮은 값으로 어닐링(annealing)됩니다.

4. 추론

추론 시, 정책은 추론 단계의 경계(예: 단락의 끝)에서만 호출됩니다. 시스템은 정지 확률을 확인하고, 정지 결정이 내려지면 즉시 최종 답변을 도출합니다. 이 설계는 병렬 프리필(pre-fill) 프로세싱을 활용하고 KV 캐시 재계산을 최소화함으로써 현대적인 LLM 서빙 시스템(예: vLLM)에서 높은 효율성을 보장합니다.

주요 기여

1. 엄격한 공식화: 본 논문은 CoT 가지치기를 명시적인 보상 함수(정확도와 길이를 교환)를 가진 최적 정지 문제로 공식화했습니다.
2. 경량 플러그인: OS-Pruner는 동결된 베이스 추론 모델을 수정하지 않고 추론 접두사 상에서 작동하는 학습 기반 정지 정책이며, 작은 정책 헤드와 몇 개의 트랜스포머 레이어만을 미세 조정하면 됩니다.
3. 실증적 성능: 실험을 통해 OS-Pruner가 다양한 모델과 벤치마크에서 정확도를 유지하면서 추론 길이를 실질적으로 줄임을 입증했습니다.
4. 상보성: 이 방법은 모델 측 방식에 의해 이미 간결함이 최적화된 모델에도 추가적인 이득을 제공함을 보여줍니다.

실험 결과

저자들은 세 가지 모델(DeepSeek-R1-Distill-Qwen-7B, GPT-OSS-20B, DRPO-7B)을 대상으로 다양한 난이도(GSM8K, MATH-500, AIME)의 벤치마크에서 OS-Pruner를 평가했습니다.

• 길이 감소: OS-Pruner는 최소한의 정확도 희생으로 생성 길이를 20~60% 감소시켰습니다.
  • DeepSeek-R1-Distill-Qwen-7B의 경우, GSM8K에서 59.3%, MATH-500에서 **52.8%**의 길이 감소를 달성했으며, 정확도 변화는 각각 -0.7%와 +2.7%에 불과했습니다.
  • 이미 간결함에 최적화된 모델인 GPT-OSS-20B에서도 (예: MATH-500에서 21%) 경쟁력 있는 정확도를 유지하면서 유의미한 감소를 달성했습니다.
• 베이스라인과의 비교: OS-Pruner는 정확도-효율성 점수(AES)에서 학습 기반 베이스라인들(Classification Probe, Answer Convergence, FlashThink)을 일관되게 능가했습니다. 이는 OS-Pruner가 파레토 프런티어(Pareto frontier)에 위치하거나 그 근처에 있는 반면, 고정 임계값 베이스라인들은 길이 예산이 줄어듦에 따라 더 급격히 성능이 저하되었음을 보여줍니다.
• 강건성: 이 방법은 난이도 수준에 따라 잘 작동했습니다. 과잉 사고가 빈번한 쉬운 작업에서는 더 공격적으로 작동하고, 추가 계산이 가치 있는 어려운 작업(AIME)에서는 더 보수적으로 작동했습니다.

의의 및 주장

본 논문은 OS-Pruner가 기존의 휴리스틱 또는 정적 방법보다 CoT 가지치기에 대해 더 자연스럽고 최적인 접근 방식을 제공한다고 주장합니다. 문제를 순차적 결정 과정으로 프레임워크화함으로써, 전역적인 임계값이 아닌 특정 추론 상태에 기반하여 종료의 충분한 시점을 동적으로 평가합니다.

저자들은 OS-Pruner가 다음을 수행함을 강조합니다:

• 계산 비용 감소: 불필요한 토큰을 제거함으로써 추론 지연 시간과 에너지 발자국을 낮춥니다.
• 세밀한 제어 제공: 파라미터 $\lambda$를 통해 사용자가 추론 노력과 정확도 사이의 트레이드오프를 명시적으로 튜닝할 수 있습니다.
• 기존 모델과 호환: 베이스 생성기를 재학습할 필요가 없어, 기존 추론 모델을 위한 실용적인 "플러그인" 역할을 합니다.
• 모델 측 방식과 상보적: 간결함을 위해 훈련된 모델(DRPO 등)도 추가적인 가지치기의 혜택을 입으며, 이는 모델 측 압축과 최적 정지가 서로 다른 측면의 효율성을 다룬다는 것을 시사합니다.

본 연구는 수학적 추론에 집중하고 있으나, 이 프레임워크가 더 넓은 맥로에 적용 가능함을 언급하며, 향후 코딩 및 과학적 추론 작업으로의 확장과 프런티어 모델로의 확장이 필요함을 밝히며 마무리합니다.