O1-Pruner: Length-Harmonizing Fine-Tuning for O1-Like Reasoning Pruning

Recently, long-thought reasoning LLMs, such as OpenAI's O1, adopt extended reasoning processes similar to how humans ponder over complex problems. This reasoning paradigm significantly enhances the model's problem-solving abilities and has achieved promising results. However, long-thought reasoning process leads to a substantial increase in inference time. A pressing challenge is reducing the inference overhead of long-thought LLMs while ensuring accuracy. In this paper, we experimentally demonstrate that long-thought reasoning models struggle to effectively allocate token budgets based on problem difficulty and reasoning redundancies. To address this, we propose Length-Harmonizing Fine-Tuning (O1-Pruner), aiming at minimizing reasoning overhead while maintaining accuracy. This effective fine-tuning method first estimates the LLM's baseline performance through pre-sampling and then uses RL-style fine-tuning to encourage the model to generate shorter reasoning processes under accuracy constraints. This allows the model to achieve efficient reasoning with lower redundancy while maintaining accuracy. Experiments on various mathematical reasoning benchmarks show that O1-Pruner not only significantly reduces inference overhead but also achieves higher accuracy, providing a novel and promising solution to this challenge. Our code is coming soon at https://github.com/StarDewXXX/O1-Pruner

한 줄 요약

O1-Pruner는 수학적 추론 과제에서 정확도를 유지하면서 추론 길이를 줄이는 긍정적 결과를 보이는 길이 조화 최적화 프레임워크이다.

핵심 기여도

**Length-Harmonizing Fine-Tuning (O1-Pruner)**를 제안하여, 정확도 유지하에 추론 길이를 최소화.
**RL-style fine-tuning**을 사용하여 짧은 추론 생성을 유도.
**Accuracy-Efficiency Score (AES)**라는 새로운 평가 지표를 정의.
**Marco-o1-7B**와 **QwQ-32B-Preview** 모델에서 실험적으로 AES, 정확도, 추론 길이 개선을 검증.

핵심 아이디어

O1-Pruner는 **길이 불균형**(length disharmony) 문제를 해결하기 위해 설계되었다. 기존의 long-thought 추론 모델은 문제 난이도와 무관하게 과도한 토큰을 사용해 추론 과정을 생성하는 경향이 있다. 이는 계산 비용 증가와 추론 효율 저하를 초래한다. O1-Pruner는 **RL-style fine-tuning**을 통해 모델이 짧은 추론을 생성하도록 유도하면서도 정확도를 유지하도록 설계되었다.

핵심적으로, **Length-Harmonizing Reward**를 도입하여 짧은 추론을 보상하고, 정확도 저하를 제재한다. 이는 **Proximal Policy Optimization (PPO)** 기반의 off-policy 학습 전략을 통해 구현된다. 또한, **MATH 데이터셋**을 기반으로 학습하며, 문제-정답 쌍만 사용하는 방식으로 데이터 효율성을 높였다.

기술적 접근법

**모델**: Marco-o1-7B (full-parameter fine-tuning), QwQ-32B-Preview (Freeze Fine-Tune).
**데이터셋**: MATH (5,000 문제), GSM8k, GaoKao (테스트셋).
**훈련 전략**:
**Pre-sampling**: 기존 모델의 추론 결과를 샘플링하여 기준선 성능을 추정.
**RL-style fine-tuning**: 짧은 추론 생성을 유도하는 Length-Harmonizing Reward를 사용.
**Hyperparameter**: λ = 2로 설정하여 정확도와 효율성의 균형을 맞춤.
**평가 지표**:
**Accuracy**: 문제 해결 정확도.
**Length**: 추론 토큰 수.
**AES**: 정확도와 길이의 가중합으로 정의된 새로운 지표.
**AES 계산식**:

$$
\Delta\text{Length} = \frac{\text{Length}_{\text{baseline}} - \text{Length}_{\text{model}}}{\text{Length}_{\text{baseline}}}, \quad
\Delta\text{Acc} = \frac{\text{Acc}_{\text{model}} - \text{Acc}_{\text{baseline}}}{\text{Acc}_{\text{baseline}}}
$$
$$
\text{AES} = \alpha \cdot \Delta\text{Length} + \beta \cdot \Delta\text{Acc} - \gamma \cdot \max(0, -\Delta\text{Acc})
$$
(기본값: α=1, β=3, γ=5)

주요 결과

**Marco-o1-7B**에서 O1-Pruner는 추론 길이를 20% 이상 줄이며, 정확도는 1.2% 증가.
**QwQ-32B-Preview**에서도 추론 길이 15% 감소, 정확도 0.8% 향상.
**AES** 지표에서 기존 SFT, DPO, Fast-Solving Prompt 대비 10% 이상 우수한 성과.
**λ=2** 설정 시 정확도와 추론 길이의 균형이 가장 좋음.

의의 및 한계

O1-Pruner는 long-thought 추론 모델의 **추론 효율성**을 향상시키는 새로운 프레임워크로, 수학적 추론 분야에서 특히 유용하다. 기존의 짧은 추론 유도 방법(예: SFT, DPO)보다 AES, 정확도, 길이 모두에서 우수한 성능을 보인다. 또한, **RL-style fine-tuning**과 **Length-Harmonizing Reward**는 추론 과정의 길이와 정확도를 동시에 최적화하는 새로운 접근법을 제시한다.

그러나, **훈련 데이터의 난이도**에 따라 성능이 크게 변한다는 한계가 있다. 예를 들어, 낮은 난이도 데이터에서는 추론 길이가 줄어들지만 정확도 향상이 없었다. 이는 모델이 복잡한 문제를 학습하지 못하면 짧은 추론이 무의미할 수 있음을 시사한다. 또한, **Freeze Fine-Tune**을 사용한 대규모 모델 실험은 계산 자원 제약으로 제한적이다.

실용적 활용

O1-Pruner는 **수학 문제 해결**, **복잡한 추론 시스템**, **대규모 LLM의 추론 효율화** 등에 적용 가능하다. 특히, 추론 과정이 길고 정확도가 중요한 산업 분야(예: 금융, 의료, 법률)에서 유용하게 사용될 수 있다. 또한, 추론 과정의 길이와 정확도를 동시에 최적화하는 기법은 **LLM의 실시간 응답 성능 향상**에도 기여할 수 있다.