To Code or not to Code? Adaptive Tool Integration for Math Language Models via Expectation-Maximization

TL;DR

This work addresses the rigidity of existing tool-integrated math language models by enabling autonomous, metacognitive code integration during reasoning. It introduces an Expectation-Maximization framework (AutoCode) that alternates guided exploration of code-using trajectories with off-policy RL optimization, treating code-triggering as a latent variable. Empirical results show substantial improvements on challenging benchmarks (e.g., MATH500 and AIME) and enhanced training efficiency, driven by structured data curation and learned code-use strategies with high selection accuracy. The approach demonstrates a synergistic improvement by effectively combining chain-of-thought reasoning with code execution, signaling a practical path toward more adaptive and capable math LLMs.

Abstract

Recent advances in mathematical problem-solving with language models (LMs) integrate chain-of-thought (CoT) reasoning and code execution to harness their complementary strengths. However, existing hybrid frameworks exhibit a critical limitation: they depend on externally dictated instructions or rigid code-integration templates, lacking metacognitive awareness -- the capacity to dynamically evaluate intrinsic capabilities and autonomously determine when and how to integrate tools. This rigidity motivates our study of autonomous code integration, enabling models to adapt tool-usage strategies as their reasoning abilities evolve during training. While reinforcement learning (RL) shows promise for boosting LLM reasoning at scale (e.g., DeepSeek-R1), we demonstrate its inefficiency in learning autonomous code integration due to inadequate exploration of the vast combinatorial space of CoT-code interleaving patterns. To address this challenge, we propose a novel Expectation-Maximization (EM) framework that synergizes structured exploration (E-step) with off-policy RL optimization (M-step), creating a self-reinforcing cycle between metacognitive tool-use decisions and evolving capabilities. Experiments reveal our method achieves superior results through improved exploration. Notably, our 7B model improves over 11% on MATH500 and 9.4% on AIME without o1-like CoT.

Figures (6)

• Figure 1: Analysis of existing tool-integrated math LLMs. While tool integration can in principle bring complementary benefits to CoT reasoning, existing models show critical rigidity when incorporating code, lacking the metacognitive capacity to earn the synergistic benefits from code integration.
• Figure 2: Example of Autonomous Code Integration. We aim to enable LLMs to determine tool-usage strategies based on their own capability boundaries. In the example, the model write code to solve the problem that demand special tricks, strategically bypassing its inherent limitations.
• Figure 3: Method Overview. (Left) shows an overview for the EM framework, which alternates between finding a reference strategy for guided exploration (E-step) and off-policy RL (M-step). (Right) shows the data curation for guided exploration. We generate $K$ rollouts, estimate values of code-triggering decisions and subsample the initial data with sampling weights per Eq. \ref{['eq_sampling']}.
• Figure 4: Training Efficiency and Convergence. We benchmark the learning dynamics of our approach against three two training paradigms: supervised fine-tuning and reinforcement learning (RL). The Pass@1 accuracy is evaluated on an held-out dev-set. We use Qwen-2.5-Base as the base model. SFT is conducted using collected public data openmathmammoth. The dashed lines indicate asymptotic performance.
• Figure 5: Visualization of Exploration in the SFT+RL paradigm. The distribution of code invocation rates across queries to visualize policy's exploration of code-integrated trajectories. Without external guidance, LLM tends to exploit its local policy neighborhood, concentrating code usage toward extremes as training phase evolves.