AI-MO가 DeepSeek-R1 스타일의 reasoning-then-generation 파이프라인을 Lean 4 정리 증명에 적용해 1.7B 모델에서 76.63% Pass@32 달성

Context

기존 Kimina Prover 시스템의 훈련 파이프라인이 복잡하여 커뮤니티 재현과 커스터마이징이 어려웠다. Lean 4 정리 증명 모델의 성능을 높이면서도 훈련 과정을 단순화하고 확장 가능하게 만들 필요가 있었다.

Technical Solution

• DeepSeek-R1 영감의 두 단계 출력 구조 도입: 자연언어 reasoning 블록()과 Lean 4 코드 블록 순차 생성
• GRPO(강화학습) 알고리즘 적용: 각 프롬프트당 N개 출력 생성 후 Lean 검증 결과에 따라 보상 1 또는 0 할당
• 형식 검증 보상 추가: 와 lean4 블록 정확히 1개씩, 반복 라인 제거, 주석 밀도 임계값, 의미론적 정렬 매칭 스코어 검사를 통해 구조적 유효성 강제
• 오류 수정 턴 구현: 모델이 실패 신호로부터 학습하도록 유도
• kimina-lean-server 개발: Lean 4 대규모 병렬 증명 검증을 위한 고처리량 시스템 제공
• Verl 프레임워크와 완전 호환: recipe/kimina-prover-rl 형태로 공개 배포, NuminaMath-LEAN 데이터셋 필터링(역사적 승률 0.5 이상 제거, 문제 변형 생성, 어려운 문제 중복)

Impact

• 1.7B 모델: MiniF2F 벤치마크 Pass@32에서 72.95%에서 76.23%로 3.28% 포인트 향상, 오류 수정 적용 시 77.87% 달성
• 0.6B 모델: Pass@32에서 68.85%에서 71.30%로 2.45% 포인트 향상
• 48시간 훈련(8개 H100 GPU) 후 Step 85에서 best@8 정확도 4% 포인트 향상(70% 달성), 오류 수정 턴 후 74%
• 훈련 진행 중 형식 오류 지속적 감소
• 출력 평균 토큰 길이 증가(더 긴 구조화된 reasoning 추적 학습)

Key Takeaway

정형 증명 작업에서 명시적인 reasoning 단계와 형식 검증을 RL 보상에 결합하면 작은 모델(0.6B~1.7B)도 SOTA 성능을 달성할 수 있다. 재현 가능성과 커뮤니티 접근성이 높은 공개 파이프라인의 제공은 형식 검증 분야의 혁신을 가속화할 수 있다.