🤖 AI

AIR: Adaptive Interleaved Reasoning with Code in MLLMs

This paper introduces AIR, a framework that enhances Multimodal Large Language Models' ability to perform adaptive interleaved reasoning with code for complex numerical computation through a three-component solution involving cold-start data construction, dataset curation, and an adaptive tool-invocation strategy guided by group-constrained reinforcement learning, resulting in significant performance improvements on evaluation benchmarks.

Original authors: Cong Han, Xiaohan Lan, Haibo Qiu, Yujie Zhong

Published 2026-06-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Cong Han, Xiaohan Lan, Haibo Qiu, Yujie Zhong

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart assistant who is great at looking at pictures and reading text, but sometimes gets stuck when a math problem requires heavy calculation. If you ask this assistant to solve a complex geometry problem involving volume, it might try to do the math in its "head" (the neural network) and make a mistake, just like a human might when trying to multiply large numbers without a calculator.

The paper introduces a new system called AIR (Adaptive Interleaved Reasoning with Code) that teaches this assistant a new superpower: knowing when to stop thinking and start using a tool.

Here is how AIR works, broken down into simple concepts:

1. The Core Idea: The "Human Calculator" Analogy

Think of the AI as a student taking a test.

Old Way: The student tries to solve every problem using only their brain. If the math is too hard, they guess or get it wrong.
AIR Way: The student learns that when a problem gets too tricky, they can raise their hand and ask for a calculator. They write down the numbers, hit the buttons, get the result, and then write the final answer.
The Twist: The AI doesn't just use a calculator for every question. It learns to be adaptive. It knows, "This question is easy, I'll do it in my head," and "This question is hard, I need to write some code to solve it."

2. The Training Process: A Two-Step Journey

The paper explains that you can't just tell a model to "use a calculator" and expect it to work immediately. It needs a specific training plan:

Step 1: The "Cold Start" (Learning the Basics)
Before the AI can learn to use tools on its own, it needs to see examples of how it's done. The researchers built a special dataset where they took hard math problems and manually showed the AI how to switch between "thinking" (text) and "calculating" (writing Python code).

Analogy: Imagine a teacher showing a student a few solved examples where the student uses a calculator for the hard parts. This gives the student a basic template to follow.

Step 2: Reinforcement Learning (The "Trial and Error" Phase)
Once the AI has the basics, it starts practicing on its own. This is where the magic happens. The researchers used a special scoring system (a reward function) to teach the AI when to use the tool.

The Problem: If the AI uses the calculator too much, it gets confused and the training falls apart (like a student who refuses to think for themselves). If it never uses it, it keeps making math errors.
The Solution: The researchers created a "Group-Constrained Reward." Imagine a coach watching a whole team of players. The coach doesn't just reward the player who scores; they reward the team for having the right balance of players using calculators.
- If the AI uses the tool too often, it gets a lower score.
- If it uses the tool just enough to get the right answer, it gets a high score.
- This keeps the AI stable and prevents it from going crazy with tool usage.

3. The Result: Smarter and More Stable

The paper claims that after this training:

Accuracy: The AI got significantly better at math problems, especially the ones that required complex calculations. On average, its scores went up by about 6 points, and for the specific problems where it used the tool, it improved by nearly 10 points.
Reliability: The AI successfully used the tool more than 95% of the time when it decided to use it.
Stability: The "Group-Constrained" method prevented the AI from crashing or becoming unstable during training, a common problem when teaching AI to use tools.

Summary

In short, AIR is a method that teaches Multimodal Large Language Models (AI that sees and reads) to act like a human problem-solver: it thinks through a problem, realizes when the math is too hard, writes a quick computer script to solve the calculation, and then finishes the answer. The key innovation is teaching the AI to do this automatically and at the right time, without needing a human to tell it when to pick up the calculator.

Technical Summary: AIR: Adaptive Interleaved Reasoning with Code in Multimodal Large Language Models

Problem Statement

While recent advancements in Multimodal Large Language Models (MLLMs) have integrated Reinforcement Learning from Verifiable Rewards (RLVR) to enhance reasoning capabilities, existing paradigms primarily focus on tool use for visual-perception tasks (e.g., rotation, zooming). These approaches rely on predefined heuristics for visual manipulation and fail to address complex numerical computation problems. Current MLLMs often lack the ability to autonomously trigger Python scripts during textual reasoning to handle intricate mathematical tasks, a capability essential for bridging the gap between multimodal deliberation and agentic execution. Furthermore, training models for such adaptive tool invocation often leads to instability, including model collapse, when tool usage becomes excessive or unregulated.

Methodology

The paper proposes AIR (Adaptive Interleaved Reasoning with Code), a framework designed to empower MLLMs with adaptive interleaved reasoning capabilities through a two-stage training paradigm: Supervised Fine-Tuning (SFT) cold-start and Reinforcement Learning (RL) enhancement.

1. Two-Stage Cold-Start Data Construction

To overcome the difficulty of activating interleaved reasoning in smaller open-source models, the authors introduce a decoupled, two-step data synthesis pipeline:

Textual CoT Generation: An MLLM generates initial textual Chain-of-Thought (CoT) reasoning paths for complex problems.
Interleaved Rewriting: An LLM rewrites these paths, replacing numerical calculation components with executable Python code snippets.
Verification: The resulting data undergoes rigorous verification in a sandbox environment to ensure code execution integrity and result accuracy before being used as cold-start data.

2. RL Data Filtering Strategies

To ensure high-fidelity training data for the RL phase, two distinct filtering strategies are employed:

Self-Sampled Strategy: Utilizes a multi-turn consensus mechanism where the candidate model's own roll-outs determine eligibility based on the Pass@k metric. This allows for iterative re-filtering as the model improves.
Prior-Filtered Strategy: Leverages high-capacity "teacher" models (e.g., GPT, Gemini) to act as an extrinsic quality gate, labeling data samples based on whether they necessitate tool invocation for effective resolution.

3. Adaptive Tool-Invocation via Group-Constrained Reward

The core innovation lies in the RL training phase, which utilizes the Group Relative Policy Optimization (GRPO) algorithm enhanced with a group-constrained reward function.

Hybrid Reward Structure: The total reward ( $R$ ) combines a Format Reward ( $R_f$ ), an Accuracy Reward ( $R_{acc}$ ), and a Code-Integrated Reward ( $R_{code}$ ).
Group Constraints: Unlike standard GRPO, this method introduces intra-group constraints on the proportion of interleaved reasoning sequences. The reward for code usage ( $R_{code}$ ) is adjusted based on the proportion of tool-calling sequences within a group, differentiated by whether the final answer is correct ( $P_r$ ) or incorrect ( $P_w$ ).
Stability Mechanism: This constraint prevents the model from over-invoking tools (which leads to entropy spikes and training collapse) while encouraging adaptive invocation. It decouples the code reward from the accuracy reward, acknowledging that tool usage is a component of the reasoning chain rather than the sole determinant of success.

Key Contributions

Two-Stage Data Construction Methodology: A pipeline for synthesizing cold-start data that supports code-augmented interleaved reasoning, effectively lowering the barrier for generating high-quality training trajectories.
Data Filtering Strategies: The proposal of Self-Sampled and Prior-Filtered strategies to curate interleaved reasoning data, mitigating noise and ensuring stable policy optimization.
Adaptive Tool-Invocation Strategy: The development of a GRPO variant with group constraints that enables models to autonomously determine the necessity and timing of tool calls, while ensuring long-term training stability.

Experimental Results

Experiments were conducted using the Qwen2.5-VL-7B model across six mathematical reasoning benchmarks (MathVista, MathVision, MathVerse, DynaMath, WeMath, LogicVista).

Performance Gains: After RL training with the group-constrained reward, the AIR model achieved an average improvement of 6.1 percentage points (pp) across benchmarks compared to the baseline.
Specific Improvements: Notable gains include +8.1 pp on MathVista, +7.7 pp on MathVerse, and +9.9 pp on WeMath.
Interleaved Reasoning Accuracy: The accuracy specifically for interleaved reasoning samples increased by 9.9 pp.
Tool-Use Success: The overall success rate of tool use exceeded 95% across various tool-call ratios.
Stability: The group-constrained mechanism successfully prevented training collapse, a common issue in Agentic RL where excessive tool use leads to instability.

Significance and Claims

The paper claims that AIR represents a significant step forward in enabling MLLMs to perform System 2 (slow thinking) reasoning by autonomously orchestrating specialized tools for computation-heavy tasks. Unlike previous works focused on visual perception (System 1), AIR addresses the critical gap in complex numerical computation.

The authors assert that their approach provides a robust solution for complex multimodal reasoning by:

Decoupling the generation of reasoning traces from the execution of code, allowing for more reliable data production.
Stabilizing long-horizon Agentic RL training through group-constrained rewards, preventing the model collapse associated with unregulated tool invocation.
Demonstrating that adaptive tool invocation, rather than fixed or excessive usage, is key to improving performance on mathematical benchmarks.

The paper concludes modestly, noting that while AIR achieves strong results, its performance remains bounded by the base model's inherent coding and perceptual capabilities, and future work should expand evaluation beyond mathematical reasoning to domains like physics and finance.