Kimina-Prover-RL 출시

Hugging Face는 Lean 4를 이용한 형식적 정리 증명을 위해 설계된 오픈 소스 학습 파이프라인인 kimina-prover-rl을 소개했습니다. 이 시스템은 DeepSeek-R1에서 영감을 받은 구조화된 '추론 후 생성(reasoning-then-generation)' 패러다임을 활용하여 0.6B에서 1.7B 파라미터 범위의 오픈 소스 모델에서 최첨단(SOTA) 성능을 달성합니다.

새로운 최첨단 소형 모델

kimina-prover-rl 파이프라인을 사용하여 MiniF2F 벤치마크의 해당 크기 카테고리에서 새로운 기준을 세우는 두 개의 새로운 모델을 제작했습니다:

  • AI-MO/Kimina-Prover-RL-1.7B: 76.63% Pass@32를 달성했습니다.
  • AI-MO/Kimina-Prover-RL-0.6B: 71.30% Pass@32를 달성했습니다.

기술적 아키텍처 및 학습 패러다임

Kimina-Prover-RL은 모델이 먼저 자연어 추론 흔적( <think> 블록으로 감싸임)을 생성한 다음 그에 상응하는 Lean 4 코드를 생성하는 2단계 출력 구조를 채택합니다. 이러한 계획과 실행의 분리는 설명 가능성, 오류 복구 및 일반화 능력을 향상시키도록 설계되었습니다.

GRPO 및 DrGPO를 이용한 강화 학습

이 파이프라인은 오픈 소스 Verl 프레임워크를 통해 구현된 Group Relative Policy Optimization (GRPO)을 사용합니다. 최적화 편향—특히 GRPO가 잘못된 출력에 대해 인위적으로 더 긴 응답을 생성하는 경향—을 완화하기 위해, 팀은 전역 상수(global constant)로 토큰 수준의 손실을 정규화하여 길이 편향을 제거함으로써 손실을 집계하는 DrGPO를 활용했습니다.

보상 메커니즘

보상은 두 가지 주요 기준에 따라 할당됩니다:

  1. 검증 보상(Verification Reward): kimina-lean-server에 의해 Lean 코드가 성공적으로 검증되면 1의 보상이 주어집니다.
  2. 형식 보상(Format Reward): 구조적 일관성을 보장하기 위해, 출력이 잘못된 형식일 경우 모델은 0의 보상을 받습니다. 형식 검사는 다음을 포함합니다:
    • 출력당 정확히 하나의 <think> 블록과 하나의 Lean 4 코드 블록이 있는지 확인합니다.
    •   repetitive reasoning lines(반복적인 추론 라인)를 거부합니다.
    • tactic blocks 내의 비주석 라인(non-comment lines)의 수가 충분한지 확인합니다.
    • 주석 밀도에 임계값을 적용하여 불필요한 상용구(boilerplate)를 방지합니다.
    • 매칭 스코어(예: Intersection-over-Union)를 사용하여 설명된 tactic과 최종 코드 간의 의미론적 정렬을 측정합니다.
    • 불필요하게 긴 응답을 방지합니다.

오류 수정 메커니즘

실패로부터의 학습을 강화하기 위해, 파이프라인은 오류 수정 턴(error correction turn)을 포함합니다. 롤아웃(rollout)이 실패할 경우, 시스템은 프롬프트, 응답, Lean 피드백을 저장하여 새로운 학습 샘플을 생성합니다. 그런 다음 모델은 이전의 추론과 코드를 수정하도록 명시적으로 프롬프트되어, 모델이 자신의 출력을 성공적으로 디버깅하는 과정에서 보상을 받는 다중 턴 상호작용 체인을 가능하게 합니다.

데이터 및 인프라

데이터셋: Kimina-Prover-Promptset

모델들은 NuminaMath-LEAN의 선별된 하위 집합인 Kimina-Prover-Promptset으로 학습되었습니다. 데이터셋은 다음과 같이 정제되었습니다:

  • 쉬운 문제(과거 승률 > 0.5)를 제거합니다.
  • Gemini를 사용하여 문제 변형을 생성하여 다양성을 높입니다.
  • 어려운 문제를 복제하여 학습 가중치를 높입니다.

고처리량 검증

RL 롤아웃 중에 필요한 병렬 증명 검사 규모를 처리하기 위해, 팀은 병렬 Lean 4 검증을 위한 오픈 소스 서버인 kimina-lean-server를 개발했습니다. API 상호작용을 위한 보조 Python 패키지인 kimina-client는 PyPI에서 사용할 수 있습니다.

성능 결과

8개의 H100 GPU에서 48시간 동안 학습한 결과 일관된 개선이 나타났습니다. 85단계(step 85)에 도달했을 때, best@8 메트릭은 70%에 도달했으며, 오류 수정 턴 이후에는 74%로 증가했습니다.

MiniF2F 벤치마크에서 RL-튜닝된 모델을 이전의 증류(distilled) 모델과 비교하면 다음과 같습니다:

Model Pass@32 Pass@32 with error fixing
AI-MO/Kimina-Prover-Distill-1.7B 72.95% 75.41%
AI-MO/Kimina-Prover-RL-1.7B 76.23% 77.87%

0.6B 모델의 경우, 파이프라인은 Pass@32 기준 68.85% (Distill)에서 71.30% (RL)로 성능을 향합니다.

Sources

관련