InternLM-Math-Plus: Unified Math Reasoning

• InternLM-Math-Plus is a unified mathematical reasoning framework that combines chain-of-thought, reward modeling, formal proof generation, data augmentation, and code execution.
• The framework leverages a seq2seq architecture to process informal, formal, and tool-assisted math tasks uniformly, achieving state-of-the-art benchmark performance.
• It delivers robust, verifiable results through integrated verification strategies, including reward-model reranking and machine-checked formal proofs, ensuring scalable extensibility.

InternLM-Math-Plus is a unified mathematical reasoning framework extending the InternLM2 Transformer LLMs, specifically designed for versatile and verifiable reasoning in mathematical domains. Leveraging a seq2seq architecture, InternLM-Math-Plus integrates chain-of-thought reasoning, reward modeling, formalized proof generation, data augmentation, and code execution into a single network capable of handling informal, formal, and tool-assisted math tasks. It achieves state-of-the-art open-source performance across benchmarks such as GSM8K, MATH, Hungary math exam, MathBench-ZH, and MiniF2F, and introduces infrastructure for future extensibility and multi-modal inference (Ying et al., 2024).

1. Model Architecture and Design

InternLM-Math builds on the InternLM2 line of decoder-only Transformers, recast for seq2seq math reasoning. The backbone employs:

• Transformer decoder, context length 4096.
• Mixed-precision (bfloat16) with FlashAttention2.
• AdamW optimizer (β₁=0.9, β₂=0.95, weight_decay=0.1), cosine learning rate annealed from $3 \times 10^{-5}$ to $3 \times 10^{-6}$ over ~125B tokens.
• Parameter scale: InternLM2-Math-Base-7B (~7B parameters) and InternLM2-Math-Base-20B (~20B parameters).

All mathematical capabilities are unified via a prompt-to-target, seq2seq interface, such that each math task—be it natural language, code, or formal language—is serialized and processed identically by the core model.

2. Unified Reasoning Components

InternLM-Math-Plus introduces five reasoning components, each with specific loss objectives and data serialization, tightly integrated in training and inference:

1. Chain-of-Thought (CoT) Reasoning: Generates intermediate reasoning steps and final answers using maximum likelihood on high-quality human-written or rule-based “scratchpad” steps. At inference, the model produces free-form CoT as with general LLM chat.
2. Reward Modeling (RM):
   • Outcome Reward Model (ORM): Classifies candidate answers as Correct/Incorrect.
   • Process Reward Model (PRM): Tags each reasoning step with <PRM+>,<PRM−> to assess process validity. The same seq2seq head supports both ORM and PRM; for inference, K candidate CoTs are sampled and scored, and the top path’s output is selected.
3. Formal Reasoning and Verification:
   • Translates between natural language and LEAN 3 (formal proof assistant).
   • Generates full formal proofs via prover-style tactic generation.
   • All formal subtasks are serialized into prompt-target format; generated proof code is machine-checked by lean3/lean4, accepting only successfully “closed-goal” outputs.
4. Data Augmentation (DA): Produces paraphrases, harder variants, and rephrasings from given questions using a dedicated loss. Augmentation increases the diversity and difficulty of the self-improving training pool.
5. Code Interpreter Integration (RICO): Enables Reasoning Interleaved with Coding through specialized prompt tags for tool calls (“<tool_call> python code </tool_call>”) and execution results (“<exec_result> ... </exec_result>”). Model alternates between text and tool calls in both training (with GPT-4-generated code/execution) and inference.

3. Training Regimen and Optimization

Training proceeds in two primary phases:

• Continued Pre-training over 125B tokens sourced from CC data, domain-specific math corpora (arXiv, web math, forums), and synthetic numeric operations, optimizing the next-token prediction objective. For the 20B variant, early stopping is triggered at 80B tokens using in-context learning validation.
• Supervised Fine-Tuning (SFT) is conducted on a multi-task mixture of CoT, RM, Formal, DA, and CI with a joint loss:

$$L_{\text{SFT}} = \lambda_{\text{CoT}} L_{\text{CoT}} + \lambda_{\text{ORM}} L_{\text{ORM}} + \lambda_{\text{PRM}} L_{\text{PRM}} + \lambda_{\mathcal{F}} L_{\text{Formal}} + \lambda_{\text{DA}} L_{\text{DA}} + \lambda_{\text{CI}} L_{\text{CI}}$$

Fine-tuning uses ~622M SFT tokens per pass, three epochs, sequence length up to 32,768, and learning rates from $4 \times 10^{-5}$ to $6 \times 10^{-6}$.

Reward models are trained in SFT-style (not PPO); during inference, reward reranking functions as post-hoc policy selection by maximizing $p(\text{label}=\text{Correct})$ among sampled CoTs.

4. Inference Procedures and Verification

Inference employs several decoding and reranking strategies:

• In-Context Learning uses sampling (temperature 0.7, K=100 for GSM8K, K=256 for MATH) with majority/voting self-consistency.
• SFT-only inference uses greedy decoding without sampling.
• Code Interpreter issues rounds of tool calls until a closure token is produced.
• Verification involves:
  • Reward Model scoring across sampled candidates.
  • Formal proof validation via lean3/lean4 “no goals remain” criterion.
  • Execution of Python code snippets with iterative retries upon error.

Pseudocode for unified training and inference (see sec. 6 of the source) formalizes these processes, highlighting prompt construction, batch iteration, and the role of AdamW optimization.

5. Performance and Benchmark Results

InternLM-Math-Plus achieves open-source state-of-the-art results across multiple benchmarks. Key metrics:

Model	GSM8K@1	GSM8K@100	MATH@1	MATH@256
InternLM2-Math-Base-7B	49.2	75.7	21.5	35.6
InternLM2-Math-Base-20B	63.7	84.8	27.3	46.2

In SFT + Zero-Shot CoT (greedy):

Model	GSM8K	MATH	Hungary	MathBench-ZH
InternLM2-Math-7B (SFT)	78.1	34.6	55	40.0
InternLM2-Math-20B (SFT)	82.6	37.7	66	45.3

Reward-model reranking brings significant gains:

Model	Greedy	PRM@100
InternLM2-Math-7B	34.6	47.0
InternLM2-Math-20B	37.7	50.0

On MiniF2F (formal proof):

Model	MiniF2F-test
InternLM2-Math-7B-Base	30.3 (74/244)
InternLM2-Math-20B-Base	29.5 (72/244)

With the code interpreter, further improvements are seen:

Model	GSM8K	MATH
InternLM2-Math-7B	79.4	50.9
InternLM2-Math-20B	80.7	54.3

Pre-training on math corpora yields ~12–15 point improvement on ICL over base models, joint multi-task SFT outperforms single-task CoT, reward-model reranking confers an additional 10–12 points, and code interpreter boosts MATH scores by 10–15 points relative to plain CoT. Formal verification enables fully machine-checked proofs, achieving open-source SOTA on MiniF2F (Ying et al., 2024).

6. Framework Extensibility and “Plus” Directions

InternLM-Math-Plus provides a blueprint for extensible, multi-modal mathematical reasoning via the unified seq2seq format. Extensions proposed in the source include:

• Self-Critique Module (SCM): Model generates and then refines its own critiques using a <Critique> prompt.
• Retriever-Augmented Module: Incorporation of external math references (e.g., formal libraries) via retrieval components.
• PPO-style RLHF: Upgrades reward modeling from SFT to Proximal Policy Optimization to maximize correctness.
• Multi-solver Ensemble: Adding CAS tools (e.g., Sympy) alongside Python interpreters.
• Network-level Adapters: Reduces interference between formal and informal reasoning streams.
• Implicit CoT Distillation: Compresses explicit CoT into a lighter network, following the ICoT paradigm.

Such modules integrate by defining new prompt tags and corresponding λ-weighted loss terms within the unified framework.

7. Significance and Research Directions

InternLM-Math-Plus establishes a versatile infrastructure for developing mathematical LLMs with robust, verifiable reasoning. Its unified seq2seq interface for informal, formal, augmented, and tool-based problem solving presents a scalable pathway for self-iterative improvement, deployment of self-critique, and integration of retrieval and more advanced RL fine-tuning. The capability to produce machine-checked formal proofs, high zero-shot and code-assisted benchmark performance, and modular extensibility position InternLM-Math-Plus as an open-source foundation for multi-modal mathematical language modeling and automated reasoning research (Ying et al., 2024).