DeepMath: smolagents를 사용하는 경량 수학 추론 에이전트

DeepMath는 Qwen3-4B Thinking을 기반으로 구축되고 Group Relative Policy Optimization(GRPO)로 미세 조정된 수학 추론 에이전트입니다. 보안된 샌드박스에서 실행되는 간결한 Python 스니펫으로 장황한 사고 과정 트레이스를 대체함으로써 수학 문제 해결 능력을 향상시키며, 출력 길이를 최대 66% 줄이고 종종 정확성을 높입니다.

기술 아키텍처 및 워크플로우

DeepMath는 소형 언어 모델과 Python 실행기를 결합하여 결정적 계산을 오프로드하고 트레이스의 장황함을 줄입니다. 이 시스템은 smolagents 라이브러리를 사용하여 구현되며, 추론 엔진으로는 vLLM을 사용합니다.

핵심 구성 요소

  • Base Model: Qwen3-4B Thinking.
  • Agent Framework: smolagents로 구축되어 모델이 일반 토큰 또는 Python 스니펫을 포함하는 특수 에이전트 호출을 출력할 수 있도록 합니다.
  • Execution Environment: 허용된 모듈 목록과 스니펫별 타임아웃을 갖춘 샌드박스 환경으로, 안전성을 보장하고 임의 코드 실행(예: 파일 I/O 또는 네트워크 액세스)을 방지합니다.

추론 과정

추론 중에 모델은 중간 계산 단계를 위한 짧은 Python 스니펫을 생성합니다. 이러한 스니펫은 샌드박스에서 실행되고, 결과는 reasoning trace에 다시 통합되어 후속 단계를 위한 결정적 결과를 모델에 제공합니다.

GRPO를 사용한 훈련 방법론

DeepMath는 TRL 라이브러리를 통해 Group Relative Policy Optimization(GRPO)를 사용하여 미세 조정됩니다. 훈련 과정은 장황함보다 정확성과 간결성에 보상을 주는 데 중점을 둡니다.

보상 구조

  • Accuracy Reward: 정답에 +1
  • Code Snippet Reward: 코드 스니펫을 생성하면 +1, 정확도 보상에 비해 10:1 가중치 적용
  • Length Penalty: GRPO 완성 후보를 5,000 토큰으로 제한하여 짧은 출력을 장려합니다.

훈련 최적화

  • Temperature Scheduling: 초기 탐색과 póź기 안정성을 균형 있게 유지하기 위해 선형 스케줄(T=1.2에서 T=0.7)이 사용됩니다.
  • In-context Learning: 모델은 에이전트 호출과 실행기 출력이 포함된 네 개의 해결된 예제를 제공받아 필요한 구문과 응답 패턴을 학습합니다.
  • Dataset: 외부 도구 사용으로부터 구체적으로 이익을 얻는 문제에 초점을 맞춘 OpenMathReasoning 데이터셋의 Tool-Integrated Reasoning(TIR) 부분이 사용됩니다.

평가 및 성능

DeepMath는 MATH500, AIME, HMMT, HLE 네 가지 데이터셋에서 벤치마크되었습니다. 평가는 견고성을 위해 majority@16을 사용했으며, 간결성을 위해 평균 출력 길이를 측정했습니다.

주요 결과

  • Efficiency: 에이전틱 접근 방식은 출력 길이를 크게 줄여 baseline 대비 최대 66% 감소합니다.
  • Synergy: 제거 연구에 따르면 에이전틱 추론만으로도 길이를 줄일 수 있지만, GRPO 훈련과 에이전틱 추론을 함께 사용할 때 최상의 결과를 얻을 수 있습니다.
  • Accuracy: DeepMath는 baseline Qwen3-4B-Thinking-2507 모델보다 짧은 트레이스를 유지하면서 어려운 데이터셋에서 더 높은 정확성을 보여줍니다.

시사점 및 제한 사항

샌드박스로 계산을 오프로드하면 출력 길이가 짧아져 산술 오류가 줄어들고 추론 속도가 향상됩니다. 그러나 연구자들은 몇 가지 제한 사항을 지적합니다:

  • Scope: 현재 초점은 작은 모델과 수학적 추론에 한정되어 있습니다.
  • Generalization: 모델은 경시대회 스타일 수학에서 평가되었으며, 형식적 증명이나 개방형 수학적 창의성으로 전이되지 않을 수 있습니다.
  • Security: 엄격한 샌드박스와 리소스 제한에도 불구하고 생성된 코드 실행은 내재된 위험으로 남아 있으며, 신중한 배포 관리가 필요합니다.

Sources