Amrith Setlur

Sparse reward reinforcement learning (RL) has become a standard tool for improving LLM reasoning, but its success depends critically on the coverage present in the base model. In practice, models are often primed for RL through mid-training on curated reasoning traces that teach useful primitive skills such as decomposition, verification, or self-correction. Although effective, this strategy requires manually specifying what the model should learn, and it remains unclear whether such primitive coverage is enough for much harder problems, which require combining these skills into broader solution techniques. We study a more automated approach to mid-training for diverse problem-solving techniques using large corpora of human-written question-answer data. Rather than treating reference solutions as targets to imitate, we use them to define dense rewards for on-policy RL. Our method, ExpRL, uses an LLM judge to compare sampled reasoning traces against reference solutions and assign process-level and outcome-level rewards. This enables RL-based mid-training to reinforce partial progress and useful intermediate behaviors across a broad range of problems, building coverage over techniques that sparse rewards alone often fail to discover. On math reasoning tasks, ExpRL yields better RL priming and a stronger initialization for subsequent sparse reward RL. The resulting policies also transfer better, improving downstream performance on both the original task distribution and a proof domain.

Proprietary AI systems have recently demonstrated impressive capabilities on complex proof-based problems, with gold-level performance reported at the 2025 International Mathematical Olympiad (IMO). However, the training pipelines behind these systems remain largely undisclosed, and their reliance on large "internal" models and scaffolds makes them expensive to run, difficult to reproduce, and hard to study or improve upon. This raises a central question: can small, open models also be trained to achieve competitive reasoning performance on difficult Olympiad-level math? In this paper, we answer this question by building QED-Nano, a 4B model post-trained for Olympiad-level proofs. Our training recipe has three stages: (1) supervised fine-tuning to imbue good proof-writing styles by distilling from DeepSeek-Math-V2, (2) reinforcement learning (RL) with rubric-based rewards, and (3) expanding RL with a reasoning cache, which decomposes long proofs into iterative summarize-and-refine cycles and enables stronger test-time reasoning. QED-Nano surpasses the proof-generation performance of much larger open models, including Nomos-1 and GPT-OSS-120B, and approaches the performance of proprietary models like Gemini 3 Pro, at a fraction of the inference cost.