💻 computer science

Aryabhata 2: Scaling Reinforcement Learning for Advanced STEM Reasoning

The paper introduces Aryabhata 2, a reinforcement learning-trained language model based on GPT-OSS-20B that achieves superior performance and greater token efficiency on competitive STEM examinations like JEE and NEET by leveraging a high-quality curriculum and progressive rollout exploration.

Original authors: Ritvik Rastogi, Vishal Singh, Tejas Chaudhari, Sandeep Varma

Published 2026-05-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ritvik Rastogi, Vishal Singh, Tejas Chaudhari, Sandeep Varma

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant but slightly scattered student named GPT-OSS-20B. This student has read a massive library of books and knows a lot of facts, but when you ask them to solve a tricky, multi-step math or science problem (like those found in India's toughest college entrance exams), they sometimes ramble, get confused, or take forever to write down the answer. They might get the right answer, but they use a lot of "ink" (computer tokens) to do it, which is expensive and slow.

The authors of this paper wanted to turn this student into a champion problem-solver without buying a bigger brain (a larger model). They created a new student called Aryabhata 2.

Here is how they did it, explained through simple analogies:

1. The Training Ground: A Strict Coach

Instead of just letting the student read more books (which is what standard training does), the researchers acted like a strict coach using Reinforcement Learning.

The Reward System: Imagine a video game where you only get points if you solve a puzzle correctly AND explain it clearly. If the student gets the answer wrong, they get zero points. If they get the answer right but write a 10-page essay when a one-paragraph explanation would do, they get a lower score.
The "Judge": The coach didn't just trust the student's word. They used a super-smart "Judge" (another AI) to check every single step. If the math didn't add up or the logic was flawed, the answer was rejected immediately.

2. The Curriculum: From Easy to Impossible

The researchers didn't throw the student into the deep end immediately. They built a three-level training camp:

Level 1 (Formatting): First, they taught the student how to hold the pen correctly. They focused on making sure the answer looked neat and followed a specific structure.
Level 2 (The Grind): Then, they gave the student thousands of practice problems. As the student got better, the coach made the problems harder. If the student kept getting 70% right, the coach swapped in even trickier questions.
Level 3 (The Wildcard): Finally, they let the student try many different ways to solve the same hard problem at once. Imagine asking the student to try 128 different paths to solve a maze simultaneously. This helped them discover clever shortcuts they hadn't seen before.

3. The Secret Sauce: "Broadened Exploration"

Usually, when training AI, you might ask it to solve a problem once. Aryabhata 2 was trained to generate many different attempts for every single question.

The Analogy: Think of it like a detective solving a crime. Instead of following one hunch, the detective sends out 128 different teams to look for clues. If even one team finds the right clue, the case is solved. This method helped the model find the "golden path" to the answer much faster and more reliably.

4. The Results: Faster, Smarter, and Leaner

When they tested Aryabhata 2 on real exams (like JEE and NEET) and even on super-hard math competitions (like AIME), the results were impressive:

Better Accuracy: It solved more problems correctly than its "parent" model (GPT-OSS-20B) and even beat some much larger, more expensive models.
The "Token" Savings: This is the biggest win. The original model often wrote long, winding explanations. Aryabhata 2 learned to be concise. It used up to 64% fewer words (tokens) to get the same (or better) result.
- Analogy: If the old model was like a tourist taking 100 photos to find the perfect shot, Aryabhata 2 is like a pro photographer who snaps the perfect shot in one click.

The Bottom Line

The paper claims that by using a very specific, high-quality set of practice questions and a smart, multi-stage training method, they turned a standard 20-billion-parameter AI into a specialized STEM reasoning expert. It didn't need to be a giant model to be smart; it just needed the right kind of practice and a coach that knew exactly how to reward good behavior.

What they didn't claim:
The paper does not claim this AI can replace human teachers, diagnose medical conditions, or be used for general conversation. It is specifically built to solve science, math, and engineering problems found in competitive exams, acting as a highly efficient tutor for those specific subjects.

Technical Summary: Aryabhata 2 – Scaling Reinforcement Learning for Advanced STEM Reasoning

1. Problem Statement

Competitive STEM examinations, specifically the Joint Entrance Examination (JEE) and the National Eligibility cum Entrance Test (NEET) in India, present a unique challenge for Large Language Models (LLMs). These exams require multi-step symbolic reasoning, precise numerical computation, and deep conceptual understanding across physics, chemistry, and mathematics. Unlike standard benchmarks, these problems are carefully constructed to test conceptual depth under strict constraints.

While recent LLMs demonstrate strong performance on general reasoning benchmarks, they face significant hurdles in real-world educational deployment:

Scalability: Deploying frontier models at scale to handle millions of student doubts is cost-prohibitive due to high inference costs and per-token pricing.
Efficiency: Open-source models often require large parameter counts and generate excessively long chains of thought, leading to high latency and token consumption.
Domain Specificity: General models often lack the specific syllabus coverage and step-by-step alignment required for competitive exam strategies.

The authors identify a gap between high-performing but expensive models and efficient, domain-specific solutions capable of providing accurate, concise, and structured explanations for student doubts.

2. Methodology

The paper introduces Aryabhata 2, a reasoning-focused language model trained via reinforcement learning (RL) post-training. The methodology is built upon a unified, RL-only pipeline designed to maximize reasoning capabilities under constrained compute resources.

2.1 Base Model and Architecture

Base Model: The system is built on GPT-OSS-20B, a 20-billion parameter open-source Mixture of Experts (MoE) model with 3.6 billion active parameters.
Parameter-Efficient Fine-Tuning (PEFT): To manage memory constraints (utilizing two NVIDIA H100 NVL GPUs), the authors employ Low-Rank Adaptation (LoRA). Adapters are inserted into attention projection layers ( $q, k, v, o$ ) and the token embedding layer. This results in only ~0.15% of parameters being trainable (31.8M out of 20.9B), significantly reducing memory footprint while preserving adaptation capacity.

2.2 Data Curation and Verification

The training corpus is derived from PhysicsWallah's internal question banks, initially containing 1.78M questions.

Cleaning Pipeline: A deterministic pipeline removes HTML artifacts (specifically <img> tags), validates LaTeX syntax using pdflatex, filters incomplete questions via an LLM classifier, and removes non-STEM content. This reduces the dataset to 1.25M high-quality questions.
Answer Verification: To ensure reliable reward signals, an automated verification pipeline using a Policy Model (GPT-OSS-120B) and a Judge Model (Qwen/Qwen3-30B-A3B-Thinking) is employed.
- Multi-Pass Sampling: Questions are verified in stages (1, 4, or 16 samples). A question is accepted if any generated solution is judged correct, ensuring the curriculum contains valid reasoning trajectories.
Curriculum Construction: Questions are categorized by difficulty (Trivial, Learnable, Challenging) based on the base model's generation success rate. Chemistry questions are upsampled to address initial performance imbalances.

2.3 Reinforcement Learning Framework

The training utilizes Group Relative Policy Optimization (GRPO) with specific modifications to suit the hardware and task constraints:

KL-Free Training: Reference models are omitted to save GPU memory; KL regularization is removed.
Objective: A DAPO-style clipped objective with an asymmetric upper clipping threshold is used.
Reward Function: A multiplicative reward $R = R_{accuracy} \times R_{format}$ $R = R_{a cc u r a cy} \times R_{f or ma t}$ is employed:
- Accuracy Reward ( $R_{accuracy}$ ): Computed via a cascade of string matching, numeric tolerance matching, and symbolic equivalence (using math-verify). Partial credit (0.5) is awarded for correct option text in single-correct MCQs even if the label is wrong.
- Format Reward ( $R_{format}$ ): Encourages sufficiently detailed final answers ( $S_{len}$ ) and a balanced ratio between reasoning and solution length ( $S_{ratio}$ ), preventing overly brief or excessively long outputs.

2.4 Training Phases

Training proceeds in three sequential phases to balance stability and exploration:

Phase 1 (Format Alignment): 300 steps with a group size of 8. Focuses on aligning the model with the desired output structure using a mix of trivial and learnable questions.
Phase 2 (Prolonged RL): ~5,000 steps. The group size increases from 8 to 16. The curriculum adaptively shifts to harder questions as accuracy improves. EMA-based checkpoint merging is used to stabilize training when rewards plateau.
Phase 3 (Broadened RL): ~700 steps. The group size expands significantly (64 to 128) to broaden exploration of reasoning trajectories and discover alternative solution strategies.

3. Key Contributions

Aryabhata 2 Model: A 20B-parameter model specifically optimized for competitive STEM reasoning, demonstrating that targeted RL on domain-specific curricula can outperform significantly larger base models.
Unified RL Pipeline: A novel training strategy combining rigorous data verification, difficulty-aware curriculum construction, and phased RL (prolonged optimization + broadened exploration) that operates effectively on limited hardware (2x H100).
Efficiency Gains: The model achieves state-of-the-art performance among open-source models while generating substantially fewer tokens (up to 64% fewer than the base model), addressing the cost barrier for large-scale educational deployment.
Comprehensive Evaluation: Validation across in-distribution competitive exams (JEE Main/Advanced, NEET) and out-of-distribution reasoning benchmarks (AIME, HMMT, MMLU-Pro, GPQA).

4. Results

4.1 In-Distribution Performance

On competitive exam benchmarks (JEE Main 2025/2026, JEE Advanced 2025, NEET 2025), Aryabhata 2 achieves an average Pass@1 of 88.95%, outperforming the base GPT-OSS-20B (83.00%) and rivaling much larger models like GPT-OSS-120B (88.28%) and Qwen3-30B-A3B (88.55%).

Token Efficiency: Aryabhata 2 reduces output tokens by 52–64% compared to GPT-OSS-20B.
Accuracy per Token: It achieves 42.31 Acc./1K tokens, significantly higher than GPT-OSS-20B (15.68) and GPT-OSS-120B (26.66).

4.2 Out-of-Distribution Performance

On general reasoning benchmarks (AIME, HMMT, GPQA, MMLU-Pro, MMLU-Redux 2.0), Aryabhata 2 achieves an average Pass@1 of 87.64%.

It improves upon GPT-OSS-20B (84.95%) and Nemotron 3 Nano (83.48%).
It shows significant gains on HMMT (+1.54 over base) and GPQA (+4.35 over base).
While it remains below frontier proprietary models (e.g., GPT-5 Mini, Gemini 2.5 Flash), it demonstrates strong generalization for a 20B parameter model.

4.3 Efficiency Comparison

Across all Out-of-Distribution (OOD) benchmarks, Aryabhata 2 reduces output tokens by 24–71% compared to GPT-OSS-20B, achieving an OOD accuracy of 39.58 Acc./1K tokens.

5. Significance and Claims

The paper claims that targeted reinforcement learning on domain-specific curricula is an effective strategy for scaling practical STEM reasoning systems. The significance of the work lies in demonstrating that:

Specialization over Scale: A smaller, efficiently trained model can outperform larger base models on specific, high-stakes reasoning tasks.
Deployment Viability: By combining rigorous data curation with phased RL, it is possible to train high-quality reasoning models on constrained hardware (2x H100), making advanced AI tutoring feasible for large-scale educational systems with millions of users.
Token Efficiency: The approach not only improves accuracy but drastically reduces inference costs by generating more concise, structured reasoning chains, directly addressing the economic barriers to deploying LLMs in education.

The authors conclude that their pipeline enables stable training under compute constraints while preserving strong reasoning quality, offering a viable path for real-world educational workloads.