Aryabhatta 1.0：面向JEE数学考试的语言模型

概述

Aryabhata 1.0 是由 Physics Wallah AI Research 开发的70亿参数小型数学语言模型，专为 JEE Mains 等高风险印度竞争性考试优化。尽管体积紧凑，Aryabhata 1.0 在以考试为中心的推理任务上实现了最先进的性能，具有出色的token效率和低推理成本。

🚧 Aryabhata 1.0 是一个实验性发布。我们正在积极寻求反馈——请在本仓库的讨论标签页贡献您的意见。

🧠 主要特性

• 架构：70亿参数的因果解码器模型。
• 考试中心优化：专门针对JEE级别数学推理进行调优。
• 高准确率：
  • JEE Mains 2025年1月场次准确率 86%。
  • JEE Mains 2025年4月场次准确率 90.2%。
• Token效率：在 ~2K token窗口内高效运行，而其他推理模型需要约8K。
• 计算高效：使用优化流水线在 1x2 NVIDIA H100 GPU 上训练。

🛠️ 训练详情

• 训练数据：从Physics Wallah专有考试数据集中精选的约13万个问题-解答对。
• 训练流水线：
  • 模型合并
  • 拒绝采样
  • 监督微调（SFT）
  • 可验证奖励强化学习（RLVR）

🔀 模型合并

我们从模型合并（加权平均）开始，通过组合多种模型能力来构建强大的初始化（Aryabhata 0.5）：

• Qwen 2.5 Math：具有扎实符号数学基础的稳健数学中心LLM。
• Ace Math：Qwen 2.5 Math的增强版本，由NVIDIA微调，在数学基准测试中提高了准确率。
• DeepSeek R1 Distill Qwen：长篇推理模型，在从DeepSeek R1蒸馏的推理轨迹上微调。

📚 数据策划 + 拒绝采样

我们从Physics Wallah内部数据库中提取了约25万个原始问题，并应用了严格的过滤和清洗：

• 移除：基于图表、非英语和选项过多的问题。
• 保留：符合JEE Main 2019-2024分布的问题。 最终策划数据集：约13万个高质量问题。

对于每个问题：

• 使用Aryabhata 0.5生成4个思维链。
• 仅保留能得出正确最终答案的思维链。

结果数据集：

• 约10万个问题
• 约35万个高质量思维链

我们使用此数据集进行SFT。

🎯 可验证奖励强化学习（RLVR）

我们使用了自定义的组相对策略优化（GRPO）变体，针对数学特定的奖励函数进行了调整：

• 移除了KL散度惩罚
• 移除了裁剪

我们在剩余的约3万个问题上使用了RLVR。

这种多阶段训练策略使Aryabhata 1.0能够捕获与教学法一致的推理模式，使其在解决学生的实际数学问题方面非常有效。

📊 性能亮点

评估设置

所有评估均在temperature = 0.0下进行，我们报告pass@1准确率。

评估数据集

我们在两组官方JEE Mains 2025数学试卷上评估了模型：

• 1月场次：10份试卷，包含250道题目。
• 4月场次：9份试卷，包含225道题目。

每份试卷包括：

• 单选题，有一个正确选项
• 数值答案型（NAT）题目，需要精确的数值回答

评估指标

我们使用复合评估指标来反映真实世界的评分严格性并减少假阳性：

1. 浮点匹配

• 在容差范围内（±1e-9）比较预测答案和目标答案
• 稳健地处理舍入误差和小数值错误

2. 字符串匹配

• 用于符号答案（如分数、根式）
• 使用严格精确匹配——预测必须逐字符匹配真值

3. LLM作为评判者（GPT-4o-mini）

• 用于模糊格式的数学等价性判断

🔹 各模型准确率比较

Aryabhata在JEE Main数学上具有最佳准确率，与前沿模型相当

🔹 准确率 vs Token使用量

Aryabhata在准确率与token使用量方面与前沿模型相当

🔧 预期用途

主要用例：

• 竞争性考试备考（JEE Main级别数学问题）
• 问答和解疑系统
• 教育辅导和概念讲解

💡 如何使用

🧪 使用 🤗 Transformers

from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig

model_id = "PhysicsWallahAI/Aryabhata-1.0"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)


# 定义停止字符串
stop_strings = ["<|im_end|>", "<|end|>", "<im_start|>", "⁠```python\n", "⁠<|im_start|>", "]}}]}}]"]

def strip_bad_tokens(s, stop_strings):
    for suffix in stop_strings:
        if s.endswith(suffix):
            return s[:-len(suffix)]
    return s


# 创建生成配置（也可以设置temperature、top_p等）
generation_config = GenerationConfig(
    max_new_tokens=4096,
    stop_strings = stop_strings
)

query = 'Find all the values of \\sqrt[3]{1}'
messages = [{'role': 'system', 'content': 'Think step-by-step; put only the final answer inside \\boxed{}.'},
            {'role': 'user', 'content': query}]

text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
inputs = tokenizer([text], return_tensors="pt")
outputs = model.generate(**inputs, generation_config=generation_config, tokenizer=tokenizer)

print(strip_bad_tokens(tokenizer.decode(outputs[0], skip_special_tokens=True), stop_strings))

⚡ 使用 vLLM

使用vLLM高效运行模型：

from vllm import LLM, SamplingParams

# 初始化模型（如果本地不存在则从Hugging Face下载）
llm = LLM(model="PhysicsWallahAI/Aryabhata-1.0")

# 定义提示和采样配置
query = 'Find all the values of \\sqrt[3]{1}'
messages = [{'role': 'system', 'content': 'Think step-by-step; put only the final answer inside \\boxed{}.'},
            {'role': 'user', 'content': query}]
sampling_params = SamplingParams(temperature=0.0, max_tokens=4*1024, stop=["<|im_end|>", "<|end|>", "<im_start|>", "⁠```python\n", "⁠<|im_start|>", "]}}]}}]"])

# 运行推理
results = llm.chat(messages, sampling_params)

# 打印结果
print(results[0].outputs[0].text.strip())

🚀 路线图

Aryabhata 2.0（即将推出）：

• 将领域覆盖扩展到物理和化学
• 支持JEE Advanced、NEET和基础课程大纲
• 进一步优化实时部署的经济性和准确性

🤝 引用

如果您使用此模型，请引用：

@misc{Aryabhata2025,
  title = {Aryabhata 1.0: A compact, exam-focused language model tailored for mathematics in Indian competitive exams, especially JEE Main.},
  author = {Physics Wallah AI Research},
  year = {2025},
  note = {\url{https://huggingface.co/PhysicsWallahAI/Aryabhata-1.0}},
}