AIMO Progress Prize Winning Solution

To replicate the Numina solution, set up a Python 3.10 environment, install PyTorch v2.1.2 (specific version required for reproducibility), and install dependencies via requirements.txt. You must also install Flash Attention 2 and Git LFS.

Note on Flash Attention 2: If your machine has less than 96GB of RAM and many CPU cores, reduce the MAX_JOBS argument during installation.

Note on Hugging Face: You must log in via huggingface-cli login to access datasets and models.

```shell
# Create and activate environment
conda create -n aimo python=3.10 && conda activate aimo

# Install dependencies (ensure PyTorch v2.1.2 is installed first)
pip install -r requirements.txt

# Install Flash Attention 2
python -m pip install flash-attn --no-build-isolation

# If low RAM/high CPU cores, use:
# MAX_JOBS=4 pip install flash-attn --no-build-isolation

# Login to Hugging Face
huggingface-cli login

# Install Git LFS
sudo apt-get install git-lfs
```埋

The PythonREPL class provides a sandboxed environment for executing generated Python code.

Key Features:

• Automatic Imports: Automatically prepends import math, import numpy as np, and import sympy as sp to every query.
• Print Wrapping: Ensures the last line of the query is wrapped in a print() call if it isn't already, to capture output.
• Timeout Protection: Uses a signal-based time_limit to prevent infinite loops or long-running code from hanging the process.
• Error Handling: Captures tracebacks and cleans up error messages to make them usable for the model's next reasoning step.
• Security: Explicitly blocks the use of subprocess and venv within the generated code.

class PythonREPL:
    def __init__(self, timeout=5):
        self.timeout = timeout

    @contextmanager
    def time_limit(self, seconds):
        # ... implementation details ...

    def __call__(self, query):
        # ... executes code via subprocess.run ...
        return success, output

The Numina solution uses an algorithm called SC-TIR (Self-Consistency with Tool-Integrated Reasoning) to solve math problems. It combines majority voting with Python code execution to improve reasoning accuracy.

The SC-TIR Workflow:

1. Prompt Expansion: For a given problem, create $M$ copies of the input to define the initial batch (the 'width' of the search).
2. Reasoning & Code Generation: Sample completions until the model produces a complete block of Python code (using the ToRA format).
3. Tool Execution: Execute the generated Python blocks in a REPL and append the output (including tracebacks) to the reasoning trace.
4. Iterative Generation: Repeat the generation process $N$ times (the 'depth') to create a set of reasoning traces. Traces that fail to produce sensible outputs (e.g., incomplete code or no \boxed{} result) are pruned.
5. Majority Voting: Post-process the valid solution candidates and apply majority voting to select the final answer.

To run the Numina solution, you need to install several Python packages. The installation method depends on whether you are using a standard environment or a Kaggle environment.

Standard pip installation

pip install vllm==0.4.2
pip install grpcio==1.62.2
pip install antlr4-python3-runtime==4.11.0
pip install networkx shapely sage matplotlib gmpy2 scipy numpy sympy mpmath

Kaggle installation

On Kaggle, you must uninstall torch first and use specific local wheels for vllm, grpcio, and ray to ensure compatibility with the T4 environment.

pip uninstall -y torch
pip install -U --no-index --find-links=/kaggle/input/vllm-whl -U vllm
pip install -U --upgrade /kaggle/input/vllm-t4-fix/grpcio-1.62.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
pip install -U --upgrade /kaggle/input/vllm-t4-fix/ray-2.11.0-cp310-cp310-manylinux2014_x86_64.whl
pip install -U --upgrade /kaggle/input/antlr4-python3-runtime-package-4-11/antlr4_python3_runtime-4.11.0-py3-none-any.whl

# If using pip
# !pip install vllm==0.4.2
# !pip install grpcio==1.62.2
# !pip install antlr4-python3-runtime==4.11.0
# !pip install networkx shapely sage matplotlib gmpy2 scipy numpy sympy mpmath

# If on Kaggle
# !pip uninstall -y torch
# !pip install -U --no-index --find-links=/kaggle/input/vllm-whl -U vllm
# !pip install -U --upgrade /kaggle/input/vllm-t4-fix/grpcio-1.62.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
# !pip install -U --upgrade /kaggle/input/vllm-t4-fix/ray-2.11.0-cp310-cp310-manylinux2014_x86_64.whl
# !pip install -U --upgrade /kaggle/input/antlr4-python3-runtime-package-4-11/antlr4_python3_runtime-4.11.0-py3-none-any.whl

The solution uses a Config dataclass to manage all hyperparameters and runtime settings.

@dataclass
class Config:
    model_id: str

    # Decoding Parameters
    num_samples: int        # Number of candidates to generate (width)
    num_generations: int    # Number of steps to generate per candidate (depth)
    restart_on_fail: bool   # Regenerate a step if it fails to generate Python codeblocks

    # Sampling Parameters
    temperature: float
    max_new_tokens: int

    # Runtime Parameters
    validation_set: str     # One of AI-MO/aimo-validation-amc, AI-MO/aimo-validation-aime, AI-MO/aimo-validation-math-level-4, AI-MO/aimo-validation-math-level-5
    is_submission: bool = bool(os.getenv("KAGGLE_IS_COMPETITION_RERUN"))