JT-Math: A Multi-Stage Framework for Advanced Mathematical Reasoning in Large Language Models

TL;DR

JT-Math-8B introduces a multi-stage training framework for large language models, significantly improving complex mathematical reasoning capabilities through specialized datasets, multi-stage reinforcement learning, and chain-of-thought techniques.

Contribution

The paper presents a novel multi-stage optimization framework and training curriculum for large language models to enhance advanced mathematical reasoning.

Findings

Achieves state-of-the-art results among open-source models in mathematical reasoning.

Outperforms models like O1-mini and GPT-4o on competition-level math tasks.

Demonstrates effectiveness of multi-stage RL and Long CoT approaches.

Abstract

Mathematical reasoning is a cornerstone of artificial general intelligence and a primary benchmark for evaluating the capabilities of Large Language Models (LLMs). While state-of-the-art models show promise, they often falter when faced with complex problems that demand deep conceptual understanding and intricate, multi-step deliberation. To address this challenge, we introduce JT-Math-8B, a series of open-source models comprising base, instruct, and thinking versions, built upon a systematic, multi-stage optimization framework. Our pre-training corpus is a high-quality, 210B-token dataset curated through a dedicated data pipeline that uses model-based validation to ensure quality and diversity. The Instruct Model is optimized for direct, concise answers through Supervised Fine-Tuning (SFT) and a GRPO-based reinforcement learning (RL) method. The Thinking Model is trained for complex problem-solving using a Long Chain-of-Thought (Long CoT) approach, combining SFT with a novel, multi-stage RL curriculum that progressively increases task difficulty and context length up to 32K tokens. JT-Math-8B achieves state-of-the-art results among open-source models of similar size, surpassing prominent models like OpenAI's O1-mini and GPT-4o , and demonstrating superior performance on competition-level mathematics.