AIR: Adaptive Interleaved Reasoning with Code in MLLMs

Paper: 2606.23678 Authors: Cong Han, Xiaolan Lan, Haibo Qiu, Yujie Zhong Categories: cs.CV, cs.AI

The Gap

Existing multimodal LLMs treat code as an external tool—they have a fixed set of “visual operations” (e.g., crop, resize, segment) and call them via predefined heuristics. This works for perception tasks like object detection but completely fails on tasks that require numerical computation from visual data. For example, given a photo of a geometry diagram, asking “what is the ratio of the shaded area to the total area?” triggers only visual parsing, no arithmetic reasoning. The community has been stuck on the tool-use paradigm from earlier papers like ViperGPT and CodeVQA.

This paper identifies the gap: adaptive interleaving—the model must decide *when to reason in natural language, when to execute Python, and how to combine them in a single trajectory—and then closes it with RL training.

[Problem] -> [Assumption] -> [Method] -> [Evidence] -> [Conclusion]
  |              |              |            |              |
  | MLLMs fail   | Code is only | Train MLLMs| Accuracy     | Adaptive
  | on visual    | for fixed    | with RL on  | +6.1pp on   | interleaving
  | numerical    | tool calls   | interleaved | benchmarks   | is learnable
  | computation  | (heuristic)  | reasoning   | Tool use>95% |
  v              v              v            v              v
 [Limitation]  [False premise] [RL + reward] [Empirical]  [Claim accepted]

The Increment

One sentence: Before this paper, MLLMs could only call code for predefined visual operations; after this paper, MLLMs learn via RL to autonomously interleave natural-language reasoning with Python computation for any visual numerical task.

Core Mechanism

The method (AIR) has three components chained together:

1. Cold-start data construction (two-stage): First, collect image-question pairs from visual math datasets (e.g., geometry, charts). Second, use an LLM to generate a reasoning trajectory that mixes natural language steps with Python code snippets (e.g., “First, detect the radius from the image using tool-X, then r = 5; area = math.pi ** r**2”). Each trajectory is a chain of interleaved steps. These become the initial supervised fine-tuning data.

2. RL training with group-constrained reward: A standard PPO variant is used, but the reward function is critical. For each generated trajectory, the model gets:

   • +points for correct final answer.
   • +points for each code execution that runs without error.
   • Group constraint: if the model tries to skip code or hallucinate results, the group of related reasoning steps gets penalized. This prevents the model from faking arithmetic.

3. Adaptive tool invocation: During inference, the model treats code execution as a “special token” that, when emitted, triggers the Python interpreter. The model learns when to use it—not via a fixed policy, but via the RL reward. So for a simple lookup question, it might skip code; for a complex calculation, it inserts code.

[Image] -> [MLLM Decoder] -> [step tokens] ---> [Python Executor] -> [Output]
                |                   |                     ^
                v                   v                     |
          [Visual Features]   [Code Tokens] <---[Adapter/Reward]
                              +[Text Tokens]          |
                                                      v
                                          [Group-Constrained Reward
                                           computed after full trajectory]

Structural metaphor: Think of this as a chef learning to cook a recipe that requires both taste judgment and precise measurements.

• The MLLM is the chef. The image is the dish on the counter.
• Earlier methods (fixed tool-use) are like a chef who can only use a pre-set set of kitchen gadgets (blender, knife, oven) but never a measuring cup or scale. They can chop and blend visually, but if the recipe says “add 3.14 grams of salt”, they guess.
• Cold-start data is a cookbook with example recipes that show *when to use the measuring cup (Python). The chef studies these.
• RL training is the chef practicing the recipe while a mentor (reward function) watches. The mentor gives a thumbs-up if the dish tastes right (correct answer) *and if every time the chef guessed a measurement instead of using the cup, the mentor points out the error. The group constraint means the mentor doesn’t just check the final dish; they check that *each step with a number was done with the cup, not by eye.
• Adaptive invocation emerges: after enough practice, the chef automatically reaches for the cup when numbers are involved, and doesn’t waste time when only visual properties matter.

Key Concepts

• Interleaved reasoning: A sequence of steps where each step is either a natural-language statement or a Python code snippet. The model outputs both intermixed. Example: Step 1: “The rectangle’s width is 12 cm.” Step 2: w = 12. Step 3: “Height is 8 cm.” Step 4: h = 8. Step 5: area = w ** h. Step 6: “So area is 96 cm².” The key is that code is not behind a special tool call; it’s inline.

• Group-constrained reward: A reward function that evaluates *subtrajectories (groups of consecutive steps) rather than only the final answer. If the model generates code that tries to use a variable before defining it, the group containing that step gets a penalty. If the model writes a numeric answer without code for a computation that requires it, the group gets zero reward. This prevents the model from learning shortcuts that look like reasoning but skip real computation.

• Cold-start data: Training data that bootstraps the model into the interleaved behavior before RL. Without it, the model never generates code interleaved with text (it has no such examples in its pretraining). The authors build it by using an LLM to “translate” existing visual reasonings into code-augmented versions.

Framework Shift

Before (mainstream approach):         After (this paper):
+---------------------+              +-----------------------+
| Input image         |              | Input image           |
+----------+----------+              +----------+------------+
           |                                    |
           v                                    v
+---------------------+              +---------------------+
| Visual Encoder      |              | Visual Encoder      |
+----------+----------+              +----------+------------+
           |                                    |
           v                                    v
+---------------------+              +---------------------+
| Heuristic Scheduler |              | RL-trained Decoder  |
| (always calls tool  |              | (decides per step)  |
|  if visual task)    |              +----------+----------+
+----------+----------+                         |
           |                           +---------+---------+
           v                           |                   |
+---------------------+        +-------+--+         +------+------+
| Fixed Visual Tool   |        | Text     |         | Python      |
| (crop/segment etc)  |        | Reasoning|         | Execution   |
+----------+----------+        +----------+         +-------------+
           |
           v
+---------------------+
| Text-only answer    |
+---------------------+

From a fixed pipeline (always code = visual operation, no computation) to an adaptive interleaving (code = arbitrary Python, invoked per step when needed for computation). The core shift is from reactive tool use to proactive, learned orchestration of reasoning and computation.

Expert Assessment

Problem choice: Real gap. Numerical comprehension in visual input is a well-known weakness of current MLLMs, and the interleaving solution is timely after o3’s paradigm shift. However, the paper limits itself to pure numerical computation tasks—what about tasks needing both visual and numerical reasoning? That’s still a frontier, but the paper’s scope is appropriately narrow.

Method maturity: Clever insight, moderate engineering. The two-stage cold start is straightforward but effective. The group-constrained reward is the novel piece—it elegantly prevents lazy imitation. That said, the RL training details are vanilla PPO with minimal innovation; the real novelty is in the data and reward design. No simpler approach exists that achieves the same effect.

Experimental integrity: Baselines are fair—they compare against a variety of prior tool-use methods (ViperGPT, CodeVQA) and also a zero-shot baseline. The 6.1 pp average improvement is statistically significant given the benchmark sizes. One red flag: all evaluation tasks are derived from the same family of visual math datasets used for cold-start data. Could the method overfit to the data distribution? The authors do not test on out-of-distribution tasks (e.g., OCR-based math). That limits confidence in generalization.

Writing quality: The paper is clearly structured but dense. The “adaptive interleaving” concept is explained well. The biggest weakness: Section 4 (experiments) lacks ablation on the group-constrained reward—how much does it contribute compared to a simpler binary reward? Without that, the reader cannot judge whether the magic is in the reward or in the cold start. A single sentence added in the conclusion would solve this.

Verdict: weak accept — the gap is real, the method is plausible, but the evaluation scope and missing ablation leave room for stronger evidence.

Takeaways

• Cold-start data recipe: Use a strong LLM (GPT-4) to “translate” natural-language reasoning traces into interleaved code+text trajectories. This technique transfers to any domain where you want to teach a model to use code inline (e.g., chemistry question solving, legal computation).
• Group-constrained reward: Instead of rewarding only the final answer, chunk the trajectory into logical groups and reward each group for consistency. This can be applied to other multi-step reasoning tasks (e.g., tool use, multi-hop QA) to prevent the model from skipping steps.
• Adaptive invocation: Your model doesn’t need a separate “tool selector” module; RL can learn when to invoke code as a token-level decision. This simplifies architectures.