PaperlayerView exact revision on arXiv

Source-linked AI summary

AIMO-2 Winning Solution: Building State-of-the-Art Mathematical Reasoning Models with OpenMathReasoning dataset

Ivan Moshkov, Darragh Hanley, Ivan Sorokin, Shubham Toshniwal, Christof Henkel, Benedikt Schifferer, Wei Du, Igor Gitman

arXiv:2504.16891v1cs.AIcs.CLcs.LG

TL;DR

The paper addresses limitations in publicly available mathematical reasoning models for challenging competition problems. It builds large-scale reasoning data, trains models for tool-integrated reasoning and generative solution selection, and combines these ideas into open-weight models. The resulting system won AIMO-2 with 34 of 50 private-test problems solved, while the released models establish a new state-of-the-art among open-weight mathematical reasoning models.

  • Problem

    Publicly available reasoning models had limitations for challenging mathematical problems, including difficulty directly producing tool-integrated reasoning solutions.

  • Method

    The paper constructs large-scale CoT data, develops iterative training and filtering for TIR, and trains GenSelect models to choose among candidate solutions.

  • Results

    The submission correctly solved 34 out of 50 private AIMO-2 test problems, and the released models establish a new state-of-the-art among open-weight mathematical reasoning models.

  • Takeaways & Limitations

    The released dataset, code, and OpenMath-Nemotron models support further research with CoT, TIR, and GenSelect inference modes.

  • Takeaways & Limitations

    The models produced around 10% more tokens on average, and this undesirable feature was not fixed before the competition ended.

Abstract

from arXiv · show

This paper presents our winning submission to the AI Mathematical Olympiad - Progress Prize 2 (AIMO-2) competition. Our recipe for building state-of-the-art mathematical reasoning models relies on three key pillars. First, we create a large-scale dataset comprising 540K unique high-quality math problems, including olympiad-level problems, and their 3.2M long-reasoning solutions. Second, we develop a novel method to integrate code execution with long reasoning models through iterative training, generation, and quality filtering, resulting in 1.7M high-quality Tool-Integrated Reasoning solutions. Third, we create a pipeline to train models to select the most promising solution from many candidates. We show that such generative solution selection (GenSelect) can significantly improve upon majority voting baseline. Combining these ideas, we train a series of models that achieve state-of-the-art results on mathematical reasoning benchmarks. To facilitate further research, we release our code, models, and the complete OpenMathReasoning dataset under a commercially permissive license.

1. Introduction

The paper presents a winning AIMO-2 submission and a pipeline for open-weight mathematical reasoning models built around large-scale reasoning data, tool-integrated reasoning, and generative solution selection.

  • 1. Introduction: 34 out of 50 problems were correctly solved on the private AIMO-2 test set under strict computational limits.AIMO-2 challenged models with 50 difficult national-level problems that had not been published online.
  • 1. Introduction: 540K unique mathematical problems yielded 3.2M long-reasoning CoT solutions for training reasoning models.The dataset was constructed through problem extraction and refinement, with solutions generated by DeepSeek-R1 and QwQ-32B.
  • 1. Introduction: A 1.7M-solution pipeline integrated code execution into long-reasoning generations through iterative training, generation, and quality filtering.The method also controlled how many code executions each generation could make.
  • 1. Introduction: Generative Solution Selection trained models to choose promising solutions from multiple candidates, addressing the gap between majority voting and pass@k.The pipeline generated 566K selection examples, but strict competition time constraints prevented its use in the AIMO-2 submission.
  • 1. Introduction: The resulting open-weight models span 1.5B, 7B, 14B, and 32B parameters and support CoT, TIR, and GenSelect inference modes.The AIMO-2 submission used an intermediate 14B model with inference optimizations for the competition’s time and compute limits.
  • 1. Introduction: The authors release their code, finetuned models, and complete OpenMathReasoning dataset under a commercially permissive license.The release is intended to facilitate further research.

2. Data Preparation

The paper constructs training and validation resources by extracting, filtering, and decontaminating mathematical problems, then synthesizing verified long-reasoning solutions. Its validation set targets AIMO-2-like exact-match evaluation through recent AIME and HMMT problems.

  • 2.1. Problems preparation: 540K unique mathematical problems were constructed from AoPS forum discussions through extraction, classification, transformation, answer extraction, and benchmark decontamination.Multiple-choice, binary, and invalid problems were removed; proof questions were transformed into answer-based questions.
  • 2.1. Problems preparation: The released resources include the complete OpenMathReasoning dataset, models, and code for reproducing the preparation pipeline.The code is provided through NeMo-Skills, while the data and models are available through Hugging Face.
  • 2.2. Comp-Math-24-25 Benchmark: 256 problems form Comp-Math-24-25, a validation set drawn from 2024–2025 AIME and HMMT competitions.The set emphasizes topics, difficulty, and predominantly non-proof numerical-answer formats aligned with AIMO-2, while minimizing contamination.
  • 2.3. Text-based Solution Synthesis: CoT synthesis generated up to 32 candidate solutions per problem using DeepSeek-R1 and QwQ-32B, with harder known-answer problems receiving more solutions.Solutions were filtered by answer equivalence, using majority answers as ground truth when extracted answers were unavailable or proofs had been converted.

3. Tool-Integrated Reasoning

The paper develops Tool-Integrated Reasoning by iteratively training models to combine long-form reasoning with Python execution and filtering generated solutions for useful code. It also controls execution counts during generation to reinforce compliance with specified limits.

  • 3. Tool-Integrated Reasoning: Prompting DeepSeek-R1 and QwQ-32B alone did not produce Tool-Integrated Reasoning solutions, motivating a more elaborate training and filtering pipeline.The target capability combines natural-language reasoning with Python execution for challenging mathematical problems.
  • 3.3. Iterative TIR improvement: TIR training improved performance significantly over the CoT-only predecessor.The paper attributes this result to the iterative model-improvement process in which stronger TIR models generate and filter additional examples.
  • 3.1. Stage-0 TIR data: 15K stage-0 TIR samples remained after filtering for correct answers, code use, and code blocks that provided novel and significant computation.The filters also removed solutions with more than two code blocks and classified code by novelty and significance.
  • 3.3. Iterative TIR improvement: 1.7M final TIR solutions resulted from repeated training, generation, and filtering with progressively stronger models.An intermediate 14B model was trained on QwQ-32B solutions before the final generation round.
  • 3.4. Controlling the number of code blocks: During data generation, each problem was assigned 1 to 8 allowed code executions, and generations exceeding the limit were removed.Warnings about remaining executions were included in prompts so models learned to follow the specified code-execution limit.

4. Generative Solution Selection

Generative Solution Selection (GenSelect) trains models to compare multiple candidate solution summaries and select the most promising one, addressing the gap between majority@k and pass@k. The pipeline constructs balanced comparison examples, filters selections, and uses regenerated summaries to reduce inference cost, while gains are strongest with fewer generations.

  • Motivation: GenSelect addresses the gap between majority@k and pass@k by selecting the most promising solution from multiple candidate summaries.The method compares solutions jointly rather than judging each in isolation.
  • Generating Selection Candidates: Qwen2.5-32B-Instruct replaces potentially terse native summaries with faithful summaries up to 2048 tokens, discarding summaries whose predicted answers differ from the original.If multiple valid summaries remain, the longest is selected; fewer than 5% of samples are discarded.
  • Generating Selection Candidates: New solution summaries let reasoning models achieve modest accuracy gains over majority voting when prompted to compare and select one.This observation motivates the training-data pipeline for GenSelect.
  • Generating Selection Candidates: QwQ-32B generates selections from groups of 2–16 summaries containing at least one correct and one incorrect solution, producing 1M selections filtered to 565K.Eight distinct comparison groups are created for each problem.
  • Reducing computational cost: Capping summaries at 2048 tokens limits each 16-solution input to 32768 tokens, while the largest accuracy improvements occur with fewer generations.For more candidates, the paper proposes repeated 16-solution sampling followed by majority voting; GenSelect becomes unstable above 32 generations.

5. OpenMath-Nemotron models

OpenMath-Nemotron models are trained through supervised fine-tuning across CoT, TIR, and GenSelect tasks, with additional training on harder problems. The resulting evaluation examines benchmark accuracy, training progression, and differences between prompting modes across model sizes.

  • Training: The first SFT round mixes CoT generation, TIR generation, and GenSelect selection, totaling 5.5M samples: 3.2M CoT, 1.7M TIR, and 566K GenSelect.Mixed-task training achieved similar accuracy to sequential training across the three tasks.
  • Training: A second SFT round trains most models for four epochs on 2.2M harder Olympiad-forum problems with long solutions and low TIR pass rates.The 32B model was excluded from this second round.
  • Results: Table 6 reports majority@64 accuracy on Comp-Math-24-25 after two SFT rounds, showing significant CoT gains and comparable TIR results.The paper also evaluates on 975 text-only Math problems from Humanity’s Last Exam, called HLE-Math.
  • Results: Smaller models can outperform under TIR majority@k yet match or underperform CoT at pass@1 because they produce more unfinished solutions.The authors hypothesize that smaller models are less consistent in using tools effectively.
  • Training: Training accuracy improves over time, with smaller models requiring longer training to achieve meaningful improvements.Figure 4 shows accuracy for intermediate 1.5B and 14B checkpoints.

6. Kaggle submission

The winning AIMO-2 submission solved 34 of 50 private-leaderboard questions under strict computational constraints. Its pipeline combined CoT and TIR models with checkpoint merging and inference optimizations, while GenSelect was not used in the submission.

  • Submission setting: 34 of 50 questions were correctly solved on the private leaderboard under a five-hour, four-L4-GPU limit.The competition required one-at-a-time answers in an offline Jupyter environment.
  • Training recipe: The submission trained Qwen2.5-14B-Base for eight epochs on 2.2M DeepSeek-R1 CoT solutions, then fine-tuned on 15K stage-0 TIR samples.The resulting TIR model was trained for 400 steps at a constant learning rate of 1e-5.
  • Training recipe: Merging the CoT and TIR checkpoints improved accuracy and accelerated generation by reducing solution length and code executions.The team found a simple linear checkpoint combination to be the most effective merging approach.
  • Inference optimization: Inference optimizations included TensorRT-LLM batching, weight quantization, ReDrafter speculative decoding, and adaptive early stopping.ReDrafter proposed up to three tokens per step, with all three accepted in approximately 65% of steps.
  • Inference optimization: The submission used near-greedy decoding, with temperature 0 and redrafter_greedy_search enabled, because it provided more stable small-batch results and a small speed improvement.TensorRT-LLM nevertheless produced varying outputs for identical prompts within a batch.
  • Discussion: The Kaggle entry used an early model version, lacked GenSelect and prompt-switchable CoT/TIR modes, and produced around 10% more tokens on average.Leaderboard score variance and discrepancies involving smaller models also complicated submission decisions.

7. Related Work

Related work has advanced tool-augmented mathematical reasoning through executable-code frameworks, specialized datasets, reasoning models, and generative evaluation methods. Recent datasets increasingly emphasize challenging problems and rich multi-step reasoning traces.

  • Tool-augmented reasoning: Program of Thoughts integrates natural language with executable code for hybrid textual and programmatic step-by-step reasoning.Subsequent work developed datasets and models supporting tool-integrated reasoning.
  • Tool-augmented reasoning: OpenMathInstruct-1, InfinityMATH, and MARIO provide datasets combining mathematical reasoning with code-interpreter outputs or tool results.Their examples span benchmarks including GSM8K and MATH.
  • Reasoning models: Qwen2.5 supports CoT and TIR, while other methods train reasoning chains before invoking domain-specific tools or apply iterative preference learning.These approaches develop both model capabilities and tool-use behavior.
  • Generative evaluation: Generative reward models and verifiers recast response evaluation as generation, using token probabilities or explicit reasoning before assigning correctness or quality judgments.Related work includes CoT-GenRM, Bradley-Terry reward modeling, self-taught evaluators, and Meta-Rewarding.
  • Mathematical reasoning datasets: Skywork-MathQA and NuminaMath expand large-scale mathematical training data with millions of examples and step-by-step reasoning annotations.Skywork-MathQA contains 2.5 million question-answer pairs, while NuminaMath offers 860K competition-style problems.
  • Mathematical reasoning datasets: BackMATH, a Chinese DeepSeek-R1 distilled dataset, and DolphinR1 emphasize backward reasoning, distilled solutions, or outputs from multiple reasoning models.Their reported sizes are approximately 14K, 110K, and 800K samples, respectively.

8. Conclusion

The paper presents a winning AIMO-2 submission and a pipeline for open-weight mathematical reasoning models. Its contributions combine tool-integrated reasoning, generative solution selection, a large released dataset, and models supporting multiple inference modes.

  • Contributions: The paper develops a method that combines code execution with long CoT generations to produce TIR solutions.This is presented as one of the paper’s principal contributions.
  • Contributions: The authors create GenSelect, a pipeline that trains models to select the most promising solution from multiple candidates.The method is designed for generative solution selection.
  • Released resources: The released OpenMathReasoning dataset contains 540K unique problems, 3.2M long CoT solutions, 1.7M long TIR solutions, and 566K GenSelect traces.These resources are released alongside the paper’s code and models.
  • Released resources: The released OpenMath-Nemotron models operate in CoT, TIR, or GenSelect inference modes and establish a new state-of-the-art among open-weight mathematical reasoning models.The conclusion frames these models as the paper’s resulting system contribution.

A. Problem Preparation Prompts

The appendix describes LLM prompts for extracting and classifying mathematical problems before training. It covers binary-question, validity, and multiple-choice classification, with conservative rules for ambiguous cases.

  • A.1. Binary Problem Classification: Problem extraction prompts ask an LLM to identify every problem in forum posts, including posts containing multiple problems or none.The extraction stage precedes classification.
  • A.1. Binary Problem Classification: Binary classification requires an explicit two-choice response and a question naturally leading to that response; otherwise the problem is not binary.Examples include true-or-false and yes-or-no questions, while numerical or open-ended problems are excluded.
  • A.2. Valid Problem Classification: Validity classification checks whether a statement is a well-formed, meaningfully interpretable math problem with enough information to solve using standard or advanced techniques.Difficulty, solvability, and required methods are not evaluated as validity criteria.
  • A.2. Valid Problem Classification: The validity prompt defaults to not invalid unless there is a clear reason, reserving invalid labels for missing definitions, external images, vague statements, or conceptual discussion.Problems remain valid even when they require recursion, limits, or logarithms.
  • A.2. Valid Problem Classification: Examples label incomplete, image-dependent, definitional, and open-ended prompts invalid, while fully specified algebraic, trigonometric, and number-theoretic questions are not invalid.The examples illustrate the distinction between solvable math problems and insufficiently specified prompts.
  • A.3. Multiple Choice Problem Classification: Multiple-choice classification requires explicit answer choices, a final-answer request rather than a proof, and at least one correct option.Problems without choices are classified as not mcq even if they have numerical answers.

A.4. Proof Problem Classification

Proof problems ask readers to establish why a statement is true rather than produce a uniquely numerical or expression-valued answer. The section contrasts proof prompts with computational examples.

  • Proof problems commonly use directives such as “prove that,” “show that,” or “demonstrate that.”
  • They may ask the solver to justify or explain why a statement is true.
  • They lack a well-defined answer in the form of a number or expression.
  • An equation-solving prompt is presented as another non-proof example, while identity and arithmetic-progression statements use proof language.
  • A geometric-progression prompt asking for several terms is presented as a non-proof computational example.

A.5. Proof Problem Conversion

Proof problems can be converted into automatically gradable tasks by preserving their mathematical setup while asking for a numerical or expression-valued answer. The examples replace existence or proof conclusions with constrained calculations.

  • The conversion task creates an equivalent numerical or expression-answer problem that remains at least as difficult as the original proof problem.
  • A proof that an integer system has no solutions is reframed by assuming a solution and asking for the sum of all possible x values.
  • The converted system preserves the original equations while changing the requested conclusion from nonexistence to calculation.
  • A proof about parabolic triangles is converted into defining f(n) through a minimum coordinate under an exact-area condition.
  • The parabolic-triangle example then asks for the specific value f(4), yielding a directly gradable target.

B. TIR Data Generation Prompts

The TIR prompts train models to use Python within mathematical reasoning and to classify code by whether it introduces a calculation, verifies prior work, or provides substantial computational assistance.

  • B.1. Stage-0 TIR Data Generation Prompt: Stage-0 TIR solutions must use Python for each complex calculation and place the final answer inside boxed notation.
  • B.2. TIR Novelty Evaluation: The prompts require exact output formats, prioritize verification when uncertain, and distinguish code that checks known answers from code that supplies new computation.
  • B.2. TIR Novelty Evaluation: Novelty evaluation labels a calculation novel when its code output does not appear in the preceding solution text.
  • B.2. TIR Novelty Evaluation: A transition-matrix computation is classified as a novel calculation because Python computes the probability from described algorithmic steps.
  • B.2. TIR Novelty Evaluation: The combinatorial example likewise treats computing C(51, 5) with Python as novel when the numerical result was not previously present.
  • B.2. TIR Novelty Evaluation: A volume calculation is classified as verification because the Python output approximately matches the earlier manual result 37π ≈ 116.23.
  • B.3. TIR Significance Evaluation: The significance rubric calls simple equation solving or arithmetic trivial when code offers little advantage over manual calculation.
  • B.3. TIR Significance Evaluation: The TSP example is significant: exhaustive evaluation of 3,628,800 permutations produces a best route, minimum distance 251, and a 5.77-second runtime.

D. Prompts for GenSelect Data Preparation

The GenSelect preparation prompts turn exploratory solution comparisons and long mathematical derivations into clean, structured summaries. They retain the reasoning needed to justify a final ranking or solution.

  • D.1. Re-generating Comparison Summary: The comparison-summary prompt supplies multiple solutions and a critique that eventually ranks them and identifies the best solution by index.
  • D.1. Re-generating Comparison Summary: The regenerated critique should remove exploratory discussion while preserving the details leading to the final judgment.
  • D.1. Re-generating Comparison Summary: The final comparison output must end with “Judgment:” followed by the index of the best solution.
  • D.2. Re-generating Solution Summary: The solution-summary prompt provides a math problem and a long solution containing alternate approaches, mistakes, corrections, and eventual resolution.
  • D.2. Re-generating Solution Summary: The required rewrite keeps the final correct approach and removes the exploratory portions of the generation.
  • D.2. Re-generating Solution Summary: The clean solution must cover all details of the final correct derivation rather than merely state its answer.

E. Example of Re-generated Summary

The example asks for the number of prime cassowaries and concludes that every such number is divisible by 9, so the answer is zero. It also illustrates code-based checking and execution-limit warnings during reasoning.

  • A cassowary is a 9-digit number using each digit from 1 through 9 exactly once.
  • Every cassowary has digit sum 45, making it divisible by 9.
  • 0 prime cassowaries remain because these 9-digit numbers are divisible by 9 and therefore composite.
  • The example includes code execution for calculations and a warning that three executions remain.
  • A redundant Python gcd computation confirms that gcd(5, 13) is 1, so an inverse exists.

F.3. Example of model encountering code execution limit

The example uses Python to define and numerically solve a two-variable system, then shows the model continuing without code after exhausting its execution allowance.

  • The model imports numerical libraries and defines a two-equation system in x and y.
  • It calls fsolve with the initial guess [1, 1] to obtain a numerical solution.
  • The computation returns approximately [-0.93580304, 2.03210427].
  • After the execution limit is reached, the model is instructed to continue using mathematical reasoning and analytical skill.
Loading 2504.16891v1…