Training Language Models to Reason Efficiently

Scaling model size and training data has led to great advances in the performance of Large Language Models (LLMs). However, the diminishing returns of this approach necessitate alternative methods to improve model capabilities, particularly in tasks requiring advanced reasoning. Large reasoning models, which leverage long chain-of-thoughts, bring unprecedented breakthroughs in problem-solving capabilities but at a substantial deployment cost associated to longer generations. Reducing inference costs is crucial for the economic feasibility, user experience, and environmental sustainability of these models. In this work, we propose to train large reasoning models to reason efficiently. More precisely, we use reinforcement learning (RL) to train reasoning models to dynamically allocate inference-time compute based on task complexity. Our method incentivizes models to minimize unnecessary computational overhead while maintaining accuracy, thereby achieving substantial efficiency gains. It enables the derivation of a family of reasoning models with varying efficiency levels, controlled via a single hyperparameter. Experiments on two open-weight large reasoning models demonstrate significant reductions in inference cost while preserving most of the accuracy.

한 줄 요약

LLM의 추론 비용을 줄이기 위해 강화학습을 활용해 사고 과정을 효율적으로 학습하는 방법을 제안한다.

핵심 기여도

**강화학습 기반 추론 효율화**: 정확도를 유지하면서 토큰 사용량을 최소화하도록 모델을 학습시킴.
**단일 하이퍼파라미터로 효율성 조절**: `α` 값을 조정해 추론 비용과 정확도 간의 트레이드오프를 조절 가능.
**GSM8K, MATH500, AIME 데이터셋에서 1.7~65% 토큰 절감**: 정확도 감소는 1.7~5% 미만.
**100개의 RL 스텝으로 학습 가능**: 대규모 RL 없이도 효과적인 결과 도출.

핵심 아이디어

기존의 대규모 언어 모델은 사고 과정(chain-of-thought)을 통해 정확도를 높이지만, 이는 추론 비용을 증가시키는 단점이 있다. 본 연구는 **강화학습**(Reinforcement Learning, RL)을 활용해 모델이 **정확한 답변을 도출하는 데 필요한 토큰 수를 최소화**하도록 학습시킨다. 핵심 아이디어는 **정확도를 유지하면서 추론 과정의 계산량을 줄이는 것**이다.

이를 위해 **목표 함수**(Equation 3)를 정의하고, 토큰 길이에 대한 **모노토닉 패널티 함수** `f`와 **조절 가능한 하이퍼파라미터** `α`를 도입한다. `α` 값이 높을수록 짧은 답변이 더 강하게 보상된다. 또한, **문제의 난이도에 따라 추론 길이가 자동 조절**되도록 **per-prompt 정규화**와 **시그모이드 함수**(σ)를 사용해 길이 패널티를 부드럽게 적용한다.

기술적 접근법

**모델**: DeepSeek-R1-Distill-Qwen-1.5B, DeepSeek-R1-Distill-Qwen-7B
**알고리즘**: Policy gradient 기반 강화학습
**목표 함수**:

$$
R = \sigma\left( \frac{L - \mu}{\sigma} \right)
$$
여기서 $ L $: 토큰 길이, $ \mu $: 평균, $ \sigma $: 표준편차, $ \sigma $: 시그모이드 함수

**하이퍼파라미터**: $ \alpha \in [0, 0.05, 0.1, 0.2, 0.4] $
**학습 스텝**: 100개의 RL 스텝 (~200 gradient updates)
**추론 효율성**: 토큰 사용량을 줄이면서 정확도 유지

주요 결과

**GSM8K**: 토큰 사용량 65% 감소, 정확도 1.7% 감소
**MATH500**: 토큰 사용량 37% 감소, 정확도 2.2% 감소
**AIME 2024**: 토큰 사용량 16% 감소, 정확도 3.3% 감소
**7B 모델**: 50% 토큰 감소, 정확도 5% 미만 감소
**1.5B 모델**: 유사한 수준의 효율성 달성
**모델 학습 비용**: 100 RL 스텝으로도 효과적

의의 및 한계

본 연구는 **LLM의 추론 효율성**을 향상시키는 새로운 접근법을 제시하며, 기존의 모델 축소(distillation)와 달리 **모델 크기 유지** 상태에서 **추론 비용을 줄이는 방식**으로 학술적·실용적 가치가 높다. 특히, **단일 모델이 문제 난이도에 따라 자동으로 추론 자원을 조절**할 수 있다는 점에서 유연성이 높다.

그러나 한계점도 존재한다. 예를 들어, **복잡한 문제에서는 토큰 감소가 제한적**이며, **정확도 감소가 일부 데이터셋에서 발생**한다. 또한, **강화학습의 학습 과정 자체가 추가적인 자원을 요구**할 수 있으므로, **초기 학습 비용**을 고려해야 한다.

실용적 활용

본 연구는 **AI 추론 서비스 제공자**(예: 클라우드 AI API)나 **교육용 AI**(수학 문제 풀이) 등에 적용 가능하다. 특히, **높은 정확도를 유지하면서도 저비용으로 추론이 필요한 산업**에서 유용하며, **환경 지속 가능성**(carbon footprint 감소) 측면에서도 의미가 있다.