NuminaMath 7B TIR, AIMO 진행상금 수상 – 기술 요약

TL;DR

NuminaMath 7B TIR는 50개의 숨겨진 대회 문제 중 29개를 정확히 풀어 최초 AI Math Olympiad (AIMO) 진행상금을 수상했으며, 두 단계 파인튜닝 레시피, 대규모 고품질 수학 데이터, 도구 통합 추론을 활용한 자기 일관성(SC‑TIR) 추론 전략이 오픈소스 LLM의 올림피아드 수준 수학 성능을 크게 향상시킬 수 있음을 입증했습니다.

Introducing Numina – an open AI4Maths initiative

Numina는 2023년 말에 AI 기반 수학 추론을 가속화하기 위해 시작된 오픈소스 프로젝트입니다. Jia Li, Yann Fleureau, Guillaume Lample, Stan Polu, Hélène Evain이 설립했으며, 초기에는 Mistral AI의 지원을 받았습니다. 2024년 초에는 Hugging Face 파인튜닝 전문가 Lewis Tunstall와 Ed Beeching이 팀에 합류했고, General Catalyst와 Answer.ai의 추가 지원도 받았습니다. 이 협업의 첫 구체적인 목표는 2024 AIMO 진행상금으로, Kaggle 대회에서 50개의 사전 선발 수준 수학 문제(AMC 12/AIME 수준에 해당) 를 오픈 웨이트 모델(2024년 2월 23일 이전에 공개된)만 사용해 테스트하는 것이었습니다.

The AI Math Olympiad (AIMO) Progress Prize

AIMO 상은 국제 수학 올림피아드에서 금메달을 딸 수 있는 AI를 만드는 것을 목표로 하며, 500만 달러의 상금을 제공합니다. 첫 진행상금에서는 참가자들이 Kaggle 플랫폼에 하루 두 번 솔루션을 제출해야 했으며, 각 제출은 P100 또는 두 대의 T4 GPU에서 최대 9시간 동안 실행되었습니다. 문제는 고등학교 대회 수준의 정수 출력 질문으로 구성되며, 공개 리더보드(50문제)와 최종 순위를 결정하는 숨겨진 비공개 리더보드가 있습니다.

Winning solution architecture

Fine‑tuned reasoning agent – DeepSeekMath‑Base 7B를 두 단계 파인튜닝하여 자연어 추론과 Python REPL 호출을 교차하는 모델로 만들었습니다.
SC‑TIR decoding algorithm – 많은 후보 솔루션을 생성하고, 삽입된 Python 코드를 실행한 뒤, 자기 일관성 검사를 거쳐 다수결을 적용하는 새로운 추론 루틴입니다.
Robust internal validation – AMC, AIME, 그리고 MATH 벤치마크의 두 하위 집합으로 구성된 네 개의 검증 세트를 사용해 모델 선택을 안내하고 공개 리더보드에 과적합되는 것을 방지했습니다.

학습은 오픈소스 스택(TRL, PyTorch, vLLM, DeepSpeed)을 사용해 8개의 H100 GPU가 장착된 단일 노드에서 진행했으며, 약 10시간 만에 완료되었습니다.

Two‑stage training recipe (MuMath‑Code)

이 레시피는 MuMath‑Code 논문을 따르며 다음과 같이 구성됩니다.

Stage 1 – Chain‑of‑Thought (CoT) fine‑tuning – 기본 모델은 수백만 개의 자연어 수학 문제와 CoT 스타일 솔루션을 학습해 단계별 추론을 장려합니다.
Stage 2 – Tool‑Integrated Reasoning (TIR) fine‑tuning – ToRA 형식으로 GPT‑4가 생성한 합성 데이터셋을 사용해 각 문제에 대한 근거, Python 프로그램, 실행 결과를 쌍으로 제공합니다. 이를 통해 모델은 중간 계산을 위해 Python REPL을 호출하도록 학습합니다.

두 단계 모두 전체 파라미터 파인튜닝(LoRA/DoRA 미사용)과 2048 토큰 패킹 전략, 그래디언트 체크포인팅, DeepSpeed ZeRO‑3 샤딩을 적용했습니다. 주요 하이퍼파라미터는 단계 간에 동일했으며, 학습률 2e‑5, 배치 크기 32, 코사인 스케줄러, 워밍업 비율은 Stage 1에서는 0.0, Stage 2에서는 0.1이었습니다.

Data is everything

Numina는 두 개의 대규모 데이터셋을 구축했습니다.

Chain‑of‑Thought dataset – 중국 고등학교 워크시트, 미국 시험 PDF, 국제 올림피아드 아카이브 등에서 수십만 개의 문제‑솔루션 쌍을 수집했습니다. 파이프라인은 OCR, 분할, 영어 번역, CoT 형식 변환을 수행했습니다.
Tool‑Integrated Reasoning dataset – 약 60 000개의 문제(대부분 정수 출력)를 GPT‑4 파이프라인을 통해 ToRA 스타일 추론을 생성하고 코드를 실행해 불일치를 필터링했습니다. 각 문제는 정확성을 보장하기 위해 최대 세 번 재생성되었습니다.

저자들은 향후 몇 주 안에 전체 데이터셋을 오픈소스로 공개할 계획입니다.

SC‑TIR: Self‑Consistency + Tool‑Integrated Reasoning

추론 시 모델은 Kaggle이 문제를 무작위 순서로 제공하고 GPU 자원이 제한돼 변동성이 크게 나타납니다. SC‑TIR은 이를 다음과 같이 완화합니다.

각 문제를 N번 복제(N = 48, 승리 실행 기준)하여 다양한 프롬프트 배치를 만듭니다.
N개의 완성을 샘플링하고, 각각 Python 코드 블록을 생성합니다.
모든 코드 블록을 실행해 출력 또는 트레이스백을 캡처합니다.
최대 M = 4 단계까지 재샘플링하여 이전 오류를 기반으로 모델이 자체 수정하도록 합니다.
유효하지 않은 샘플을 제거하고 최종 숫자 답에 대해 다수결을 적용합니다.

AutoGPTQ를 이용한 8‑bit 양자화는 T4 GPU에서 필수적이었습니다. 업로드 시간을 절반으로 줄이고 bfloat16 호환성 문제를 피했으며, VRAM 사용량을 감소시켰고 정확도 손실은 미미했습니다.

Preventing over‑fitting to the public leaderboard

공개 테스트 세트가 50문제에 불과해 팀은 네 개의 내부 검증 스위트를 구축했습니다.

AMC (83개의 정수 출력 문제) – 비공개 테스트 세트를 대표하며, 모델은 약 60‑65 %를 해결했습니다.
AIME (90문제) – 난이도가 높아 실패 모드를 드러냅니다.
MATH level 4 (754문제) 및 MATH level 5 (721문제) – 정수 답만 남긴 5 000문제 MATH 벤치마크의 하위 집합입니다.

5‑10개의 무작위 시드로 반복 평가해 변동성을 측정했으며(SC‑TIR은 보통 1‑3 %), 이를 통해 하이퍼파라미터 튜닝의 신뢰성을 확보했습니다.

Experiments that didn’t make the final cut

순수 CoT 모델에 다수결 적용.
MMOS(단계별 Python) 모델, 최고 16/50점.
Kahneman‑Tversky Optimisation(KTO) 온‑정책 파인튜닝, 공개 점수를 27/50으로 올렸지만 최종 모델에 시간 부족.
코드 실행 보상을 포함한 강화학습(PPO 및 REINFORCE‑LOO) – 보상 곡선은 유망했지만 정확도 향상은 없었습니다.
더 큰 백본(InternLM‑20B, CodeLlama‑33B, Mixtral‑8x7B)으로 확장 – T4 GPU에서 추론이 너무 느렸습니다.
모델 병합 기법(DARE, TIES, WARP) – 내부 지표에서 회귀를 일으켰습니다.

Implications and future directions

NuminaMath 7B TIR의 성공은 고품질 수학 데이터, 두 단계 파인튜닝 체계, SC‑TIR 기반 견고한 추론을 갖춘 오픈소스 LLM이 올림피아드 수준의 추론에서도 경쟁력을 가질 수 있음을 보여줍니다. 이 접근법은 더 큰 모델에도 확장 가능하며(팀은 더 큰 백본용 데이터셋을 공개할 예정) 도구 사용이 필요한 다른 분야(예: 기호 계산, 과학 프로그래밍)에도 적용될 수 있습니다.

프로젝트가 오픈된 만큼 기여자와 파트너가 데이터셋을 확장하고, 디코딩 알고리즘을 개선하거나 새로운 모델 아키텍처를 탐구하도록 초대합니다. 지속적인 커뮤니티 노력은 궁극적인 AIMO 목표인 금메달 AI에 한 걸음 더 다가가게 할 것입니다.

Acknowledgements

저자들은 Thomas Wolf와 Leandro von Werra에게 Numina–Hugging Face 협업을 촉진해 준 것에 감사를 표합니다. Hugo Larcher에게 GPU 제공을, Colin Raffel에게 모델 병합에 대한 조언을, Omar Sanseviero에게 블로그 포스트 피드백을 각각 감사드립니다. 추가 지원은 Mistral.ai, General Catalyst, Answer.ai, 그리고 베이징 국제 수학 연구센터(북경대)에서 제공되었습니다.

Access the model and demo

모델 저장소: NuminaMath‑7B‑TIR
인터랙티브 데모: Math Olympiad Solver Space