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

TL;DR

This paper introduces an EM-based framework enabling math language models to autonomously decide when and how to integrate code during reasoning, improving problem-solving performance without relying on fixed templates or external instructions.

Contribution

It proposes a novel EM approach that combines structured exploration with off-policy RL to enhance autonomous tool integration in math language models.

Findings

Achieves over 11% improvement on MATH500

Improves 9.4% on AIME without CoT

Demonstrates effective exploration in tool-use decisions

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.