DeepMath: A lightweight math reasoning Agent with smolagents

DeepMath is a math reasoning agent built on Qwen3-4B Thinking and fine-tuned with Group Relative Policy Optimization (GRPO). It improves mathematical problem-solving by replacing verbose chain-of-thought traces with concise Python snippets executed in a secure sandbox, reducing output lengths by up to 66% while often increasing accuracy.

Technical Architecture and Workflow

DeepMath combines a small language model with a Python executor to offload deterministic computation and reduce trace verbosity. The system is implemented using the smolagents library and uses vLLM as the inference engine.

Core Components

  • Base Model: Qwen3-4B Thinking.
  • Agent Framework: Built with smolagents to allow the model to output either normal tokens or special agent calls containing Python snippets.
  • Execution Environment: A sandboxed environment with an allow-list of imported modules and per-snippet timeouts to ensure safety and prevent arbitrary code execution (e.g., no file I/O or network access).

Inference Process

During inference, the model generates short Python snippets for intermediate calculation steps. These snippets are executed in the sandbox, and the results are folded back into the reasoning trace, providing the model with deterministic results for subsequent steps.

Training Methodology with GRPO

DeepMath is fine-tuned using Group Relative Policy Optimization (GRPO) via the TRL library. The training process focuses on rewarding accuracy and concision over verbosity.

Reward Structure

  • Accuracy Reward: +1 for correct final answers.
  • Code Snippet Reward: +1 for generating code snippets, weighted 10:1 relative to the accuracy reward.
  • Length Penalty: Shorter outputs are encouraged by limiting GRPO completion candidates to 5,000 tokens.

Training Optimizations

  • Temperature Scheduling: A linear schedule (T=1.2 to T=0.7) is used to balance initial exploration with later stability.
  • In-context Learning: The model is provided with four solved examples containing agent calls and executor outputs to learn the required syntax and response patterns.
  • Dataset: The Tool-Integrated Reasoning (TIR) subset of the OpenMathReasoning dataset is used, focusing on problems that specifically benefit from external tool use.

Evaluation and Performance

DeepMath was benchmarked on four datasets: MATH500, AIME, HMMT, and HLE. The evaluation used majority@16 for robustness and measured mean output length for brevity.

Key Findings

  • Efficiency: The agentic approach significantly reduces output length, with reductions of up to 66% compared to the baseline.
  • Synergy: Ablation studies showed that while agentic inference alone can reduce length, the best results are achieved when both GRPO training and agentic inference are used together.
  • Accuracy: DeepMath demonstrates higher accuracy on challenging datasets while maintaining shorter traces compared to the baseline Qwen3-4B-Thinking-2507 model.

Implications and Limitations

Offloading computation to a sandbox reduces arithmetic errors and improves inference speed due to shorter outputs. However, the researchers note several limitations:

  • Scope: The current focus is limited to small models and mathematical reasoning.
  • Generalization: The model was evaluated on contest-style math and may not transfer to formal proofs or open-ended mathematical creativity.
  • Security: Despite strict sandboxing and resource limits, executing generated code remains an inherent risk that requires careful deployment management.

Sources