Kimina-Prover-RL: 줄임새가 된 Kimina Prover 훈련 파이프라인 (Verl 완전 호환)

DeepSeek-R1에 영감을 받은 구조화된 추론-생성 패러다임을 기반으로 Lean 4 형식 증명에 대한 오픈소스 훈련 파이프라인인 kimina-prover-rl 을 소개합니다.

이 훈련 파이프라인은 Kimina Prover 를 훈련하기 위해 사용했던 시스템의 간소화 버전으로, 핵심 구성 요소를 유지하면서 오픈소스 Verl 프레임워크와 완전 호환성을 제공합니다. 이 파이프라인은 recipe/kimina-prover-rl 에 포함된 전체 훈련 레시피를 포함하는 Verl 포크의 일부로 출시되었습니다.

이 훈련 파이프라인의 결과로 우리는 두 개의 모델을 공개합니다:

AI-MO/Kimina-Prover-RL-1.7B, MiniF2F 벤치마크에서 76.63% Pass@32를 달성한 1.7B 파라미터 모델 — 이 크기 범주의 오픈소스 모델에서 새로운 최강을 세웠습니다.

AI-MO/Kimina-Prover-RL-0.6B, MiniF2F 벤치마크에서 71.30% Pass@32를 달성한 0.6B 파라미터 모델 — 이 크기 범주의 오픈소스 모델에서도 새로운 최강을 세웠습니다.

kimina-prover-rl 은 Lean 4 형식 증명 목표를 해결하도록 대형 언어 모델을 가르치기 위해 설계된 훈련 파이프라인으로, 자연어 추론 추적에 이어 해당하는 Lean 코드를 생성하는 두 단계 출력 구조를 사용합니다.

이 패러다임은 DeepSeek-R1에서 영감을 받아 모델이 계획과 실행을 분리할 수 있게 하여 설명 가능성, 오류 복구 및 더 강력한 일반화를 촉진합니다.

이 추론 프레임워크 하에서 모델을 훈련하기 위해 우리는 GRPO — LLMs에 맞춤화된 강화 학습 접근법을 적용합니다. 이 Kimina-prover 훈련 파이프라인의 오픈소스 버전은 RL 라이브러리 Verl 을 사용하여 구현되었습니다.

GRPO 의 롤아웃 단계 동안, 모델은 각 프롬프트에 대해 N 개의 출력을 생성합니다. 우리의 kimina-lean-server 를 사용하여 Lean 코드가 성공적으로 검증된 모든 출력에는 1 의 보상이 할당됩니다.

이 프레임워크에는 두 가지 주요 기능이 추가되었습니다:

• 출력 구조를 가르치기 위해 형식 검사 보상
• 실패 신호에서 모델을 학습하도록 유도하기 위한 오류 수정 턴

훈련 동안에는 많은 수의 Lean 4 증명 후보가 동시 검증되어야 합니다. 이를 효율적으로 처리하기 위해 고 트러프류트 검증 시스템을 요구합니다.

이 필요성을 충족하기 위해 Numina 와 Kimi 는 대규모로 병렬 증명 검증을 지원하는 오픈소스 서버인 kimina-lean-server 를 개발했습니다.

통합을 단순화하기 위해 우리는 서버의 API 와 상호 작용할 수 있는 깔끔한 인터페이스를 제공하는 경량 Python 패키지인 kimina-client (PyPI 에서 사용 가능) 도 제공합니다.

우리는 Kimina-Prover-Promptset 을 사용하여 훈련합니다, 이는 NuminaMath-LEAN 데이터셋의 정교화된 부분입니다.

이 훈련 설정을 위해 우리는 다음과 같이 데이터셋을 필터링하고 전처리합니다:

쉬운 문제 제거: 역사적 승률 0.5 이상인 경우에만 데이터셋에서 어려운 명제를 유지합니다.

변형 생성: Gemini를 사용하여 기존 문제의 다양성을 증가시키기 위해 변형을 생성합니다.

난이도 높은 문제 복제: 훈련 동안 더 큰 가중치를 부여하기 위해 난이도 높은 문제를 복제합니다.

결과 데이터셋은 Lean 4 형식 증명 모델을 개선하기 위한 도전적이고 가치 있는 문제를 포함합니다.

NuminaMath-LEAN-RL 은 또한 AI-MO/Kimina-Prover-RL-1.7B 와 AI-MO/Kimina-Prover-RL-0.6B 를 훈련하기 위해 사용된 데이터셋입니다.

예시 입력 형식:

다음 본문을 이어서 번역해주세요:

Think about and solve the following problems step by step in Lean 4.
# Problem:
Find all primes that are the difference of the fourth powers of two integers.
...

우리의 추론 훈련 파이프라인의 핵심 아이디어는 LLM 출력을 두 단계로 구조화하는 것입니다. 하나의 추론 블록과 하나의 Lean 4 블록을 따릅니다:

• 추론 블록 ( ... )
• Lean 4 코드 블록

...<thinking>
To prove the statement, we use induction on n.
The base case is trivial, and the inductive step follows by applying the hypothesis.
...
</thinking>

Each rollout is **verified** to ensure that this format is respected. If the output is malformed — e.g., missing the `...<thinking>`
블록 또는 코드 블록이 잘못 배치된 경우 — 모델은 증명 자체의 유효성과 무관하게 **제로 보상을** 받습니다.

이는 일관성을 강제하고 모델을 신뢰할 수 있는 출력 구조를 학습하도록 가르칩니다.

kimina-prover 에서 이러한 확인은 단순히 `...<thinking>`
블록과 Lean 4 블록의 존재만 확인하는 것을 넘어:

- 각 출력이 정확히 하나의 `...<thinking>...</thinking>`
블록과 하나의 Lean 4 코드 블록을 포함하는지 확인합니다.
- 반복적인 추론 줄이 있는 출력 (이는 종종 환각 또는 퇴화된 생성을 나타냄) 을 거부합니다.
- 추론 섹션 내의 전략 블록이 충분히 많고 비 주석 줄이 충분한지 확인합니다.
- 주석 밀도 (추론과 Lean 코드 모두에서) 적용하여 과도하게 방대하거나 템플릿화된 출력을 처벌합니다.
- 전략 블록에 기술된 전략과 최종 Lean 코드 간의 의미적 정렬을 일치 점수 (예: Intersection-over-Union 또는 서브코드 커버리지) 를 사용하여 비교합니다.
- 불필요하게 긴 응답을 처벌하고, 완전한 답변을 제공하면서도 토큰을 더 효율적으로 사용하는 것을 장려합니다.

이러한 모든 확인을 통과하는 생성물만 잘 형식화되어 보상을 받을 수 있습니다. 이 구조화된 필터링은 훈련 안정성을 개선하고 깔끔한 추론을 장려합니다.

정보 전달력을 높이기 위해, 우리는 모델이 실패한 증정을 수정할 기회를 부여하는 **오류 수정 메커니즘**을 추가했습니다.

롤아웃이 실패할 경우 (예: Lean 오류 또는 잘못된 증명에 의해):

- 전체 프롬프트, 응답 및 Lean 피드백을 저장합니다.
- 모델이 이전 추론/코드를 명시적으로 수정하도록 프롬프트하는 **새로운 훈련 샘플**을 생성합니다.

이는 Lean 피드백이 훈련 중 제공되는 동안 모델이 실패 신호에서 학습하도록 장려합니다.

또한, Lean 피드백이 프롬프트의 일부로 주입되고 모델이 자신의 출력을 성공적으로 디버깅하는 경우 보상을 받는 다중 턴 상호작용 체인을 가능하게 합니다.

다중 턴 응답은 길어질 수 있으므로, 우리는 오류 수정 턴을 한 번만 허용하고 오류 메시지를 토큰 수로 제한합니다.

작업 **Understanding R1-Zero-Like Training: A Critical Perspective** 는 GRPO 에 최적화 편향이 있어 잘못된 출력에 특히 더 긴 응답을 인위적으로 만들어준다고 주장합니다.

우리의 실험 중에도 이러한 행동을 관찰했으며, 우리는 최적화를 위해 DrGPO 를 사용했습니다. DrGRPO 는 전역 상수로 정규화하여 길이 편향을 제거하며 토큰 수준의 손실을 집계합니다.

제공된 구성 파일은 8 GPU 설정을 위한 것입니다.

우리가 미세 조정하는 모델은 **AI-MO/Kimina-Prover-Distill-1.7B** 입니다. 이 모델은 **Qwen/Qwen3-1.7B** 의 미세 조정 버전으로, 우리의 **AI-MO/Kimina-Prover-72B** 모델에서 생성된 콜드 스타트 데이터로微调되었습니다.

각 단계마다 학습 데이터셋에서 256 개의 샘플을 가져옵니다. 두 개 중 하나가 오류 수정 샘플입니다. 우리는 각 샘플당 8 번의 rollout 을 생성하므로 총 2048 번의 generations 를 얻습니다. 여러 노드를 사용하는 경우 16 또는 32 번의 rollouts 로 늘릴 수 있습니다.

우리는 5 개의 훈련 단계마다 모델을 평가하며, verl 의 best@8 메트릭을 사용하여 빠른 검증 단계를 수행합니다. 여러 노드를 사용하는 경우 best@16 또는 32 로 늘릴 수 있습니다. 우리는 오류 수정 전후에 성능을 평가합니다. 각 실패 응답에 대해 모델이 자신의 증명을 수정하기 위해 한 번 더 시도할 수 있습니다.

수련 단계가 몇 개 지나면 성능에 일관된 개선이 관찰됩니다. 이 섹션에서는 8 개의 H100 GPU 에서 48 시간 훈련 후의 훈련 지표를 논의합니다.

단계 85 에 도달했을 때, 파이프라인은 모델의 정확도를 4 점 향상시켜 best@8 메트릭에서 70%, 오류 수정 후 74% 를 달성했습니다:

동시에, 훈련 과정에서 형식 오류의 수가 꾸준히 감소하는 것을 관찰하며, 이는 모델이 구조적으로 유효한 출력을 생성하도록 학습하고 있음을 나타냅니다.

마지막으로, DeepSeek-R1 스타일의 훈련 설정에서 예상대로, 모델 출력의 평균 토크 길이는 훈련과 함께 증가합니다 — 이는 모델이 더 긴, 더 구조화된 추적 (trace) 으로 추론하고 있음을 의미합니다.

훈련 후, 우리는 오류 수정을 포함하여 및 없이 pass@32 를 사용하여 모델을 평가했습니다. MiniF2F 에서 pass@32 에서 1.7B 모델을 3% 이상 개선할 수 있었습니다:

| 모델 | 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 모델을 미세 조정하여 성능을 2% 이상 개선했습니다.

| 모델 | Pass@32 |
|---|---|
| AI-MO/Kimina-Prover-Distill-0.6B | 68.85% |
| AI-MO/Kimina-Prover-RL-0.6B | 71.30% |

Kimina-Prover-RL 은 Lean 4 정리 증명자를 훈련하기 위한 경량적이지만 강력한 강화학습 파이프라인을 제공합니다.

구조적 추론, 형식 보너스, 오류 수정을 결합하여 우리는 0.6B–1.7B 파라미터 범위의 오픈소스 모델에서 최상위 결과를 달성했습니다.

모델과 함께, 우리는 `recipe/kimina-prover-rl` 에 포함된 전체 훈련 레시피를 포함하는 Verl 포크도 공개합니다. 커뮤니티가 우리의 결과를 재현하거나 파이프라인을 그들의 데이터셋 및 모델에 맞게 적응할 수 있습니다.

우리는 이 릴리스가 커뮤니티가 형식 추론에서 RL 훈련을 실험하고 Lean 4 의 오픈소스 자동 증명의 한계를 밀어붙이는 데 견고한 기초가 될 것으로 기대합니다.