# Hybrid RAG + Multi-Attempt Tool Use for AIMO Progress Prize 3

Author:
Lonnie Qin
Competition:
AI Mathematical Olympiad — Progress Prize 3
Public LB Score:
≈ 40 / 50 (no improvement over simpler baselines)
Status:
Awaiting Private LB results (hopeful for a competitive score)
1. Overview
This solution implements a
retrieval-augmented, multi-attempt reasoning pipeline
built around a 120B-class model, a high-quality olympiad corpus, and a persistent Python sandbox for tool use.
Core Idea (which I still believe is strong):
Give the model a rich set of similar, fully-solved competition problems (via hybrid semantic + keyword retrieval) and let it reason multiple times in parallel with tool access. Combine the attempts using entropy-weighted voting. This approach mimics how strong human contestants solve problems — by recalling analogous examples and verifying rigorously.
While the public leaderboard score stabilized around
40
, I am optimistic that the method’s depth will translate better on the private test set.
2. Corpus Creation & Embedding Pipeline
Notebook:
[https://www.kaggle.com/code/lonnieqin/aimo-math-corpus-creation-embedding-generation](https://www.kaggle.com/code/lonnieqin/aimo-math-corpus-creation-embedding-generation)
Datasets Used
NuminaMath-1.5
(
AI-MO/NuminaMath-1.5
): ~900k problems → filtered to hard competition-style only (
aime
,
amc_10/12
,
imo
,
usamo
,
putnam
,
hmmt
, etc.)
Hendrycks MATH
(
hendrycks/competition_math
): Level 4 & 5 only
Filtering & Cleaning
Minimum solution length = 100 characters
Problem-type filter (Number Theory, Counting & Probability, Intermediate Algebra, Precalculus, Geometry, Algebra)
Strict deduplication on normalized problem text
Aggressive cleaning of AoPS navigation boilerplate, markdown links, ellipsis, etc.
Embedding Generation
model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')  # 768-dim
embeddings = model.encode(corpus_texts, normalize_embeddings=True, batch_size=32)
FAISS
IndexFlatL2
(exact search)
Saved:
faiss_index.bin
,
corpus.pkl
,
embeddings.npy
, metadata
Final corpus size:
~120k high-quality solved problems.
3. Main Inference Pipeline
Notebook:
[https://www.kaggle.com/code/lonnieqin/aimo-3-solution](https://www.kaggle.com/code/lonnieqin/aimo-3-solution)
3.1 Model & Serving
Model:
gpt-oss-120b
served via vLLM
Configuration:
fp8_e4m3
KV cache, 65k context,
gpu-memory-utilization=0.96
, prefix caching
3.2 Hybrid Similar-Problem Retriever
class SimilarProblemRetriever:
    def get_similar(self, problem: str, k: int = 5) -> str:
        # 1. Embedding retrieval (3× candidates)
        # 2. Keyword retrieval using MATH_KEYWORDS (divisor, prime, gcd, permutation, triangle…)
        # 3. Hybrid score = 0.6×embedding_sim + 0.4×keyword_overlap
Returns 5 reference problems with full solutions and answers — dramatically enriching the context.
3.3 System Prompt Highlights
Elite IMO-level identity
Explicit instructions to study and adapt techniques from the retrieved examples
Strict
\boxed{}
output format (0–99999)
3.4 Tool Use — Persistent Jupyter Sandbox
16 pre-warmed kernels with
sympy
,
numpy
,
mpmath
, etc.
Custom
python
tool with automatic
print()
injection
State persists across reasoning turns within an attempt
3.5 Multi-Attempt Parallel Solving
8 parallel attempts per problem (different seeds)
Up to 128 reasoning turns each
temperature=1.0
,
min_p=0.02
Early stopping when one answer reaches 4 votes
3.6 Answer Selection
Weighted voting: weight =
1 / max(entropy, 1e-9)
Mean token entropy used as confidence signal
Transparent vote table displayed for each problem
4. Time & Resource Management
Global notebook timeout: 17,400 s
Adaptive per-problem budget (300–900 s) with 50-problem lookahead safety
Fully reliable within competition constraints
5. Prompt Logging
All prompts, retrieved examples, and final answers are logged to
prompt_logs.jsonl
for analysis.
6. Results & Reflection
Public LB:
≈ 40
The approach did not improve the public leaderboard score relative to simpler baselines. However, the core components (high-quality corpus, hybrid retrieval, persistent tool use, and entropy-weighted ensembling) remain technically sound and closely mirror successful human problem-solving strategies.
I am hopeful that the private test set will reward the method’s depth and robustness more than the public set did.
7. Why This Approach Still Has Strong Potential
Component
Value Even at Current Score
Hybrid retrieval
Guarantees relevant examples
120B model + vLLM
Raw reasoning power
Persistent Python sandbox
Exact verification
Entropy-weighted voting
Filters low-confidence outputs
Multi-attempt parallelism
Robustness
8. Future Improvements (if continuing)
Stronger reranking (cross-encoder)
Difficulty-aware example selection
Longer context or MoE models
Self-verification / critique loops
Conclusion
While the public leaderboard result was not the breakthrough I hoped for, I remain convinced that retrieval-augmented multi-attempt reasoning with tool use is a powerful direction for mathematical Olympiad AI. The private LB will be the real test.
Notebooks
Corpus creation:
[https://www.kaggle.com/code/lonnieqin/aimo-math-corpus-creation-embedding-generation](https://www.kaggle.com/code/lonnieqin/aimo-math-corpus-creation-embedding-generation)
Full solution:
[https://www.kaggle.com/code/lonnieqin/aimo-3-solution](https://www.kaggle.com/code/lonnieqin/aimo-3-solution)
Credits to notebooks:
[https://www.kaggle.com/code/nihilisticneuralnet/44-50-let-me-over-cook](https://www.kaggle.com/code/nihilisticneuralnet/44-50-let-me-over-cook)