Kimina-Prover-RL: Lean 4 정리 증명을 위한 오픈 소스 강화 학습

개요

Kimina-Prover-RL은 대규모 언어 모델이 Lean 4에서 형식적 증명 목표를 해결하도록 학습시키기 위해 설계된 오픈 소스 학습 파이프라인입니다. DeepSeek-R1에서 영감을 받은 구조화된 '추론 후 생성' 패러다임을 채택하여, 모델이 계획과 실행을 분리하도록 강제합니다. 이번 릴리스에는 Verl 프레임워크와 호환되는 학습 레시피와 두 가지 새로운 모델이 포함되어 있습니다: MiniF2F 벤치마크에서 각각 76.63%(Pass@32)와 71.30%(Pass@32)의 성능을 기록한 Kimina-Prover-RL-1.7B 및 Kimina-Prover-RL-0.6B입니다.

학습 방법론

이 파이프라인은 GRPO(Group Relative Policy Optimization)를 활용하여 모델이 자연어 추론 과정을 생성한 뒤 Lean 4 코드를 작성하도록 학습시킵니다. 인위적으로 길고 부정확한 출력을 유도하는 GRPO의 알려진 최적화 편향을 완화하기 위해, 토큰 수준의 손실을 전역 상수로 정규화하는 DrGRPO를 구현에 사용했습니다.

학습 아키텍처의 주요 구성 요소는 다음과 같습니다:

• 형식 보상(Format Rewards): 출력은 <think>...</think> 블록 하나와 lean4 코드 블록 하나라는 특정 구조를 엄격히 준수해야 합니다. 형식이 잘못된 출력은 보상 0을 받습니다. 추가적인 검사를 통해 반복적인 추론, 환각된 전술, 추론과 코드 간의 의미적 불일치에 대해 페널티를 부여합니다.
• 오류 수정(Error Correction): 파이프라인에는 증명 시도가 실패했을 때 모델이 이전 추론과 코드를 수정하도록 유도하는 오류 수정 메커니즘이 포함되어 있습니다. 이를 통해 모델은 다중 턴 상호작용 속에서 Lean 4의 피드백을 통해 학습할 수 있습니다.
• 고처리량 검증(High-Throughput Verification): 증명 검증 규모를 처리하기 위해, kimina-client 파이썬 패키지를 통해 액세스하는 병렬 증명 검증 지원 오픈 소스 서버인 kimina-lean-server를 사용합니다.

데이터셋 및 구성

모델은 NuminaMath-LEAN의 엄선된 하위 집합인 Kimina-Prover-Promptset으로 학습됩니다. 데이터셋은 '쉬운' 문제(과거 승률 > 0.5)를 제거하고, Gemini를 사용하여 변형을 생성해 다양성을 높이며, 어려운 문제의 가중치를 높이기 위해 복제하는 방식으로 처리됩니다.

학습 구성은 8-GPU 설정에 최적화되어 있으며, AI-MO/Kimina-Prover-Distill-1.7B 모델(기반: Qwen/Qwen3-1.7B)을 미세 조정합니다. 이 과정은 단계당 256개의 샘플을 가져오며(절반은 오류 수정 샘플로 지정), 샘플당 8개의 롤아웃을 생성합니다.

성능 결과

학습 파이프라인은 시간이 지남에 따라 정확도와 구조적 준수 측면에서 지속적인 향상을 보여줍니다. 8개의 H100 GPU에서 48시간 동안 학습한 결과, 1.7B 모델은 85단계에서 4포인트의 정확도 향상을 달성했습니다. 이 모델들은 MiniF2F에서 각 크기 범주별로 새로운 최첨단(SOTA) 벤치마크를 수립했습니다:

모델	Pass@32	오류 수정 포함 Pass@32
AI-MO/Kimina-Prover-RL-1.7B	76.23%	77.87%
AI-MO/Kimina-Prover-RL-0.6B	71.30%	N/A

학습이 진행됨에 따라 형식 오류의 빈도는 감소하고 출력의 평균 토큰 길이는 증가하며, 이는 모델이 더 구조화되고 포괄적인 추론 과정을 개발하고 있음을 반영합니다.