Aryabhata 1

India's Joint Entrance Examination (JEE) is not just a test of whether a model can reach the right answer. Students need transparent, step-by-step reasoning they can follow in a classroom or while revising alone. Frontier reasoning models have pushed raw accuracy, but many still hide their chains of thought, wander through nonlinear self-corrections, or burn tens of thousands of tokens per problem. That makes them poor tutors at scale.

Aryabhata 1.0 is a 7B open-weight model post-trained for JEE Main Mathematics. It combines the fluency of instruction-tuned math models with the depth of distilled reasoning models, then sharpens both through verified supervision and reinforcement learning with programmatic rewards.

Why exam-focused small language models

Today's model landscape splits roughly into three groups. Non-reasoning instruct models are fast but guess or pattern-match on multi-step JEE problems. Early reasoning models like o1 and DeepSeek R1 are more accurate, but their traces tend to be hidden, verbose, or hard to follow. Modern reasoning models (o4-mini, Gemini 2.5, updated R1) score well on benchmarks yet still optimize for competition scores rather than the kind of linear, exam-style reasoning that actually helps a student.

Aryabhata aims for high accuracy on real exam papers, roughly 2K tokens per response, and step-by-step reasoning suitable for doubt-solving at scale.

Training pipeline overview

Aryabhata 1.0 is trained in four stages: model merging, data curation, supervised fine-tuning (SFT), and reinforcement learning with verifiable rewards (RLVR).

Stage 1: Model merging

System 1 models answer quickly; System 2 models reason deliberately. We merge three Qwen 2.5 Math family checkpoints with linear weight averaging (MergeKit):

Component	Role
Qwen2.5-Math-7B-Instruct	Baseline math fluency
AceMath-7B-Instruct	Higher benchmark accuracy after NVIDIA post-training
DeepSeek-R1-Distill-Qwen-7B	Long chain-of-thought reasoning style

Merged parameters: θ_merged = α·θ_Qwen + β·θ_Ace + γ·θ_DS with α = 0.15, β = 0.5, γ = 0.35, chosen on held-out math reasoning tasks to balance speed and depth.

Stage 2: Data curation

Training data comes from PhysicsWallah's JEE-aligned question bank (proprietary; not released). Starting from ~250,000 raw items, we:

• Drop diagram-dependent questions (text-only training).
• Remove non-English or malformed items.
• Strip MCQ options so the task is open-ended generation, not classification, following insights that answer matching can outperform multiple-choice evaluation.
• Remove questions that only make sense with options present.
• Normalize question-answer pairs with OpenAI o4-mini (structured extraction, LaTeX answers, dependency flags).

This yields ~130,000 clean questions spanning the full JEE Math syllabus, from conic sections and probability to vector algebra and miscellaneous integrals.

Stage 3: SFT with rejection sampling

For each question we sample four chain-of-thought completions from the merged model and keep only trajectories whose final answer matches ground truth (best-of-4 rejection sampling).

Questions are bucketed by how many of the four samples are correct:

Correct CoTs (of 4)	Questions	Use
0	31,470	RLVR only (hard negatives)
1	9,647	SFT (hardest curriculum tier)
2	9,066	SFT
3	12,643	SFT
4	67,247	SFT (10% subsampled from 268,988 CoTs)

Curriculum SFT trains from easier buckets (4/4 correct) toward harder ones (1/4 correct), stabilizing early optimization. In total we obtain ~350,000 verified CoT traces from ~100,000 questions. SFT uses LoRA and completes in about 4 hours on 2x H100.

Stage 4: RL with verifiable rewards

We optimize an A2C objective with group-relative advantage estimation: for each prompt we sample a group of completions, assign binary rewards (1 if final answer is correct, else 0), and normalize advantages within the group.

Two exploration mechanisms distinguish our RL stage from fixed-group GRPO.

Group size scales with difficulty: G_d = 8 × 2^k for k ∈ {0,1,2,3} (8, 16, 32, 64), driven by average group reward. Harder prompts get more rollouts without wasting compute on easy ones. Sampling temperature ramps from 0.6 to 1.0 during training, giving stability early and diversity later. Curriculum-based sampling keeps problems in a learnable difficulty band, excluding trivial items (weak gradient) and impossibly hard ones (noisy signal).

A single 2x H100 node cannot hold both a large inference engine and full training state, so we use an alternating pipeline: generate rollouts with vLLM, terminate inference to free memory, then run policy updates and repeat. Final RL training takes ~350 hours on 2x H100.

Evaluation

We report pass@1 with greedy decoding (temperature 0). Answer correctness uses a staged pipeline: string match, numeric tolerance, symbolic equivalence, and for MCQs, GPT-4o-mini as judge when needed.

In-distribution: JEE Main 2025 Math

Official JEE Main 2025 mathematics papers:

Session	Questions	Aryabhata 1.0 pass@1
January 2025	250	86.0%
April 2025	225	90.2%

Average generation length stays around ~2,000 output tokens, competitive with frontier models on the accuracy-token frontier. The technical report has full baseline comparisons against Qwen2.5-Math, AceMath, DeepSeek-R1-Distill, Nemotron, GPT-4o, o4-mini, and Gemini 2.5 Flash.

Out-of-distribution

On MATH-500 and GSM8K, Aryabhata generalizes beyond the JEE fine-tuning distribution and outperforms its base merged components, which suggests the pipeline improves reasoning rather than narrow benchmark overfitting.

Release and what's next

Aryabhata 1.0 is open on Hugging Face under the PhysicsWallahAI org, with full methodology in our arXiv paper.

This was our first step toward exam-centric, open small language models for Indian education. Aryabhata 2 extends the same approach to multi-subject STEM (Physics, Chemistry, Math, reasoning) at larger scale with prolonged and broadened reinforcement learning.