Aryabhatta 1.0 : An exam-focused language model for JEE Math
Overview
Aryabhata 1.0 is a 7B parameter small language model for mathematics developed by Physics Wallah AI Research, optimized for high-stakes Indian competitive exams like JEE Mains. Despite its compact size, Aryabhata 1.0 achieves state-of-the-art performance on exam-centric reasoning tasks with impressive token efficiency and low inference cost.
🚧 Aryabhata 1.0 is an experimental release. We are actively seeking feedback — please contribute in the Discussion tab of this repo.
🧠 Key Features
- Architecture: 7B parameter causal decoder-based model.
- Exam-Centric Optimization: Specifically tuned for JEE-level Mathematics reasoning.
- High Accuracy:
- 86% on JEE Mains January 2025 session.
- 90.2% on JEE Mains April 2025 session.
- Token Efficiency: Operates effectively around a ~2K token window, compared to ~8K required by other reasoning models.
- Compute Efficient: Trained on a 1x2 NVIDIA H100 GPU using optimized pipeline.
🛠️ Training Details
- Training Data: ~130K problem-solution pairs curated from proprietary Physics Wallah exam datasets.
- Training Pipeline:
- Model Merging
- Rejection Sampling
- Supervised Fine-Tuning (SFT)
- Reinforcement Learning with Verifiable Rewards (RLVR)
🔀 Model Merging
We began with model merging (Weighted average) to build a strong initialization (Aryabhata 0.5) by combining diverse model capabilities:
- Qwen 2.5 Math: A robust math-centric LLM with solid symbolic math foundations.
- Ace Math: An enhanced version of Qwen 2.5 Math, fine-tuned by NVIDIA for improved accuracy in mathematics benchmarks.
- DeepSeek R1 Distill Qwen: A long-form reasoning model, fine-tuned on reasoning traces distilled from DeepSeek R1.
📚 Data Curation + Rejection Sampling
We extracted ~250K raw questions from Physics Wallah's internal database and applied aggressive filtering and cleaning:
- Removed: diagram-based, non-English, and option-heavy questions.
- Kept: questions matching the distribution of JEE Main 2019–2024.
Final curated dataset: ~130K high-quality questions.
For each question:
- Generated 4 CoTs using Aryabhata 0.5.
- Retained only those leading to correct final answers.
Resulting Dataset:
- ~100K questions
- ~350K high-quality CoTs
We used this dataset for SFT.
🎯 Reinforcement Learning with Verifiable Rewards (RLVR)
We used a custom in-house variant of Group Relative Policy Optimization (GRPO), adapted for math-specific reward functions.
- Removed KL-divergence penalty
- Removed clipping
We used RLVR on the remaining ~30K questions.
This multi-phase training strategy allows Aryabhata 1.0 to capture pedagogy-aligned reasoning patterns, making it highly effective for solving real student queries in mathematics.
📊 Performance Highlights
Evaluation Setup
All evaluations were performed with temperature = 0.0, and we report pass@1 accuracy.
Evaluation Datasets
We evaluated the model on two sets of official JEE Mains 2025 mathematics papers:
- January Session: 10 question papers containing 250 questions.
- April Session: 9 question papers containing 225 questions.
Each paper includes a mix of:
- Multiple Choice Questions (MCQs) with one correct option
- Numeric Answer Type (NAT) questions requiring precise numerical responses
Evaluation Metric
We used a composite evaluation metric to reflect real-world grading rigor and reduce false positives:
- Float Match
- Compares predicted and target answers within a tolerance (±1e-9)
- Handles rounding artifacts and small numerical errors robustly
- String Match
- Used for symbolic answers (e.g., fractions, radicals)
- Uses strict exact match — predictions must match ground truth character-for-character
- LLM-as-Judge (GPT-4o-mini)
- Used for Mathematical equivalence for ambiguous formats
🔹 Accuracy Comparison Across Models
Aryabhata has the best accuracy on JEE Main Maths, on par with frontier models
🔹 Accuracy vs Token Usage
Aryabhata is on par with frontier models in terms of accuracy vs token usage
🔧 Intended Use
Primary Use Cases:
- Competitive exam preparation (JEE Main level mathematics problems)
- Question answering and doubt-solving systems
- Educational tutoring and concept explanation
💡 How to Use
🧪 Using with 🤗 Transformers
1from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
2
3model_id = "PhysicsWallahAI/Aryabhata-1.0"
4
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6model = AutoModelForCausalLM.from_pretrained(model_id)
7
8
9# Define stop strings
10stop_strings = ["<|im_end|>", "<|end|>", "<im_start|>", "```python\n", "<|im_start|>", "]}}]}}]"]
11
12def strip_bad_tokens(s, stop_strings):
13 for suffix in stop_strings:
14 if s.endswith(suffix):
15 return s[:-len(suffix)]
16 return s
17
18
19# Create generation config (can also set temperature, top_p, etc.)
20generation_config = GenerationConfig(
21 max_new_tokens=4096,
22 stop_strings = stop_strings
23)
24
25query = 'Find all the values of \\sqrt[3]{1}'
26messages = [{'role': 'system', 'content': 'Think step-by-step; put only the final answer inside \\boxed{}.'},
27 {'role': 'user', 'content': query}]
28
29text = tokenizer.apply_chat_template(
30 messages,
31 tokenize=False,
32 add_generation_prompt=True
33)
34inputs = tokenizer([text], return_tensors="pt")
35outputs = model.generate(**inputs, generation_config=generation_config, tokenizer=tokenizer)
36
37print(strip_bad_tokens(tokenizer.decode(outputs[0], skip_special_tokens=True), stop_strings))
⚡ Using with vLLM
To run the model efficiently using vLLM:
1from vllm import LLM, SamplingParams
2
3# Initialize model (downloads from Hugging Face if not local)
4llm = LLM(model="PhysicsWallahAI/Aryabhata-1.0")
5
6# Define prompt and sampling configuration
7query = 'Find all the values of \\sqrt[3]{1}'
8messages = [{'role': 'system', 'content': 'Think step-by-step; put only the final answer inside \\boxed{}.'},
9 {'role': 'user', 'content': query}]
10sampling_params = SamplingParams(temperature=0.0, max_tokens=4*1024, stop=["<|im_end|>", "<|end|>", "<im_start|>", "```python\n", "<|im_start|>", "]}}]}}]"])
11
12# Run inference
13results = llm.chat(messages, sampling_params)
14
15# Print result
16print(results[0].outputs[0].text.strip())
🚀 Roadmap
Aryabhata 2.0 (Upcoming):
- Extending domain coverage to Physics and Chemistry
- Supporting JEE Advanced, NEET, and Foundation syllabus
- Further optimization for affordability and accuracy in real-time deployments
🤝 Citation
If you use this model, please cite:
1@misc{Aryabhata2025,
2 title = {Aryabhata 1.0: A compact, exam-focused language model tailored for mathematics in Indian competitive exams, especially JEE Main.},
3 author = {Physics Wallah AI Research},
4 year = {2025},
5 note = {\url{https://huggingface.co/PhysicsWallahAI/Aryabhata-1.0}},
6}
