NuminaMath 7B TIR wins AIMO Progress Prize – technical recap

TL;DR

NuminaMath 7B TIR won the inaugural AI Math Olympiad (AIMO) Progress Prize by correctly solving 29 of 50 hidden competition problems, demonstrating that a two‑stage fine‑tuning recipe, large high‑quality math data, and a self‑consistency with tool‑integrated reasoning (SC‑TIR) inference strategy can dramatically improve open‑source LLM performance on Olympiad‑level math.

Introducing Numina – an open AI4Maths initiative

Numina is an open‑source project launched in late 2023 to accelerate AI‑driven mathematical reasoning. It was founded by Jia Li, Yann Fleureau, Guillaume Lample, Stan Polu, and Hélène Evain, with early backing from Mistral AI. In early 2024, Hugging Face fine‑tuning experts Lewis Tunstall and Ed Beeching joined the team, and additional support arrived from General Catalyst and Answer.ai. The collaboration’s first concrete goal was the 2024 AIMO Progress Prize, a Kaggle competition that tests LLMs on 50 pre‑selection‑level math problems (comparable to AMC 12/AIME) using only open‑weight models released before 23 Feb 2024.

The AI Math Olympiad (AIMO) Progress Prize

The AIMO Prize aims to create an AI capable of winning a gold medal at the International Math Olympiad, offering a $5 M grand prize. The first progress prize required participants to submit solutions twice daily on a Kaggle platform, each submission running on either a P100 or two T4 GPUs for up to nine hours. Problems are integer‑output questions drawn from high‑school competition levels, with a public leaderboard (50 problems) and a hidden private leaderboard that determines the final ranking.

Winning solution architecture

The victorious system comprised three core components:

  1. Fine‑tuned reasoning agent – DeepSeekMath‑Base 7B was fine‑tuned in two stages to become a model that interleaves natural‑language reasoning with Python REPL calls.
  2. SC‑TIR decoding algorithm – A novel inference routine that generates many candidate solutions, executes the embedded Python code, and applies majority voting after self‑consistency checks.
  3. Robust internal validation – Four curated validation sets (AMC, AIME, and two subsets of the MATH benchmark) guided model selection and prevented over‑fitting to the public leaderboard.

Training used open‑source stacks (TRL, PyTorch, vLLM, DeepSpeed) on a single node with eight H100 GPUs, completing in roughly ten hours.

Two‑stage training recipe (MuMath‑Code)

The recipe follows the MuMath‑Code paper and consists of:

  • Stage 1 – Chain‑of‑Thought (CoT) fine‑tuning – The base model learns from several hundred thousand natural‑language math problems with CoT‑style solutions, encouraging step‑by‑step reasoning.
  • Stage 2 – Tool‑Integrated Reasoning (TIR) fine‑tuning – Using a synthetic dataset generated by GPT‑4 in the ToRA format, each problem is paired with a rationale, Python program, and execution result. This teaches the model to call the Python REPL for intermediate calculations.

Both stages employed full‑parameter fine‑tuning (no LoRA/DoRA) with a 2048‑token packing strategy, gradient checkpointing, and DeepSpeed ZeRO‑3 sharding. Key hyper‑parameters were identical across stages: learning rate 2e‑5, batch size 32, cosine scheduler, and a warm‑up ratio of 0.0 (Stage 1) or 0.1 (Stage 2).

Data is everything

Numina built two large datasets:

  • Chain‑of‑Thought dataset – Hundreds of thousands of problem‑solution pairs sourced from Chinese high‑school worksheets, US exam PDFs, and international Olympiad archives. The pipeline performed OCR, segmentation, English translation, and conversion to a CoT format.
  • Tool‑Integrated Reasoning dataset – Approximately 60 000 problems (mostly integer‑output) were processed through a GPT‑4 pipeline that generated ToRA‑style reasoning, executed the code, and filtered out mismatches. Each problem was regenerated up to three times to ensure correctness.

The authors plan to open‑source the full dataset in the coming weeks.

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

During inference, the model suffers from high variance because Kaggle supplies problems in random order and GPU resources are limited. SC‑TIR mitigates this by:

  1. Replicating each problem N times (N = 48 for the winning run) to create a batch of diverse prompts.
  2. Sampling N completions, each producing a block of Python code.
  3. Executing every code block and capturing outputs or tracebacks.
  4. Re‑sampling up to M = 4 depths, allowing the model to self‑correct based on previous errors.
  5. Pruning invalid samples and applying majority voting on the final numeric answers.

Quantization to 8‑bit via AutoGPTQ was essential for the T4 GPUs: it halved upload time, avoided bfloat16 incompatibility, and reduced VRAM consumption while incurring only a minor accuracy loss.

Preventing over‑fitting to the public leaderboard

Because the public test set contains only 50 problems, the team constructed four internal validation suites:

  • AMC (83 integer‑output problems) – Representative of the private test set; models solved ~60‑65 %.
  • AIME (90 problems) – Higher difficulty to expose failure modes.
  • MATH level 4 (754 problems) and MATH level 5 (721 problems) – Subsets of the 5 000‑problem MATH benchmark filtered for integer answers.

Repeated evaluations with 5‑10 random seeds measured variance (typically 1‑3 % for SC‑TIR), enabling reliable hyper‑parameter tuning.

Experiments that didn’t make the final cut

The team explored several alternatives before settling on the MuMath‑Code pipeline:

  • Pure CoT models with majority voting.
  • MMOS (single‑step Python) models, which plateaued at 16/50.
  • Kahneman‑Tversky Optimisation (KTO) on‑policy fine‑tuning, which raised the public score to 27/50 but was not time‑finished for the final model.
  • Reinforcement learning (PPO and REINFORCE‑LOO) with code‑execution rewards – showed promising reward curves but no measurable accuracy gain.
  • Scaling to larger backbones (InternLM‑20B, CodeLlama‑33B, Mixtral‑8x7B) – inference was too slow on T4 GPUs.
  • Model merging techniques (DARE, TIES, WARP) – caused regressions on internal metrics.

Implications and future directions

NuminaMath 7B TIR’s success demonstrates that open‑source LLMs can achieve competitive Olympiad‑level reasoning when equipped with high‑quality math data, a two‑stage fine‑tuning regime, and robust inference via SC‑TIR. The approach scales to larger models (the team plans to release datasets for bigger backbones) and can be applied to other domains requiring tool use (e.g., symbolic computation, scientific programming).

The open nature of the project invites contributors and partners to extend the dataset, improve the decoding algorithm, or explore new model architectures. Continued community effort could bring the ultimate AIMO goal—a gold‑medal‑winning AI—closer to reality.

Acknowledgements

The authors thank Thomas Wolf and Leandro von Werra for facilitating the Numina–Hugging Face collaboration, Hugo Larcher for GPU provisioning, Colin Raffel for advice on model merging, and Omar Sanseviero for feedback on the blog post. Additional support came from Mistral.ai, General Catalyst, Answer.ai, and the Beijing International Center for Mathematical Research at Peking University.

Access the model and demo

Sources