# IAPO: Input Attribution-Aware Policy Optimization for Tool Use in Small Multimodal Agents

arXiv: https://arxiv.org/abs/2606.11652 · HTML: https://mlsanity.com/paper/2606.11652/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reinforcement learning for LLM agents, one of 16 directions in LLM agents and tool use.
- Builds on 25 of its 28 references that are in the map.
- First posted on arXiv 10 Jun 2026.

## Placed in

- LLM agents and tool use / [Reinforcement learning for LLM agents](https://mlsanity.com/agents/reinforcement-learning-for-llm-agents/)

## Builds on (most cited first, 25 of 25)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Qwen3-VL Technical Report](https://mlsanity.com/paper/2511.21631/) (2025)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [Knowledge Distillation: A Survey](https://mlsanity.com/paper/2006.05525/) (2020)
- [ChartQA: A Benchmark for Question Answering about Charts with Visual and Logical Reasoning](https://mlsanity.com/paper/2203.10244/) (2022)
- [Vision-R1: Incentivizing Reasoning Capability in Multimodal Large Language Models](https://mlsanity.com/paper/2503.06749/) (2025)
- [VLM-R1: A Stable and Generalizable R1-style Large Vision-Language Model](https://mlsanity.com/paper/2504.07615/) (2025)
- [Compositional Semantic Parsing on Semi-Structured Tables](https://mlsanity.com/paper/1508.00305/) (2015)
- [TabFact: A Large-scale Dataset for Table-based Fact Verification](https://mlsanity.com/paper/1909.02164/) (2019)
- [DeepEyes: Incentivizing "Thinking with Images" via Reinforcement Learning](https://mlsanity.com/paper/2505.14362/) (2025)
- [Visual Sketchpad: Sketching as a Visual Chain of Thought for Multimodal Language Models](https://mlsanity.com/paper/2406.09403/) (2024)
- [Logic-RL: Unleashing LLM Reasoning with Rule-Based Reinforcement Learning](https://mlsanity.com/paper/2502.14768/) (2025)
- [CharXiv: Charting Gaps in Realistic Chart Understanding in Multimodal LLMs](https://mlsanity.com/paper/2406.18521/) (2024)
- [The Landscape of Agentic Reinforcement Learning for LLMs: A Survey](https://mlsanity.com/paper/2509.02547/) (2025)
- [A Survey of Reinforcement Learning for Large Reasoning Models](https://mlsanity.com/paper/2509.08827/) (2025)
- [ToRL: Scaling Tool-Integrated RL](https://mlsanity.com/paper/2503.23383/) (2025)
- [OpenThinkIMG: Learning to Think with Images via Visual Tool Reinforcement Learning](https://mlsanity.com/paper/2505.08617/) (2025)
- [VTool-R1: VLMs Learn to Think with Images via Reinforcement Learning on Multimodal Tool Use](https://mlsanity.com/paper/2505.19255/) (2025)
- [ReFocus: Visual Editing as a Chain of Thought for Structured Image Understanding](https://mlsanity.com/paper/2501.05452/) (2025)
- [TableVQA-Bench: A Visual Question Answering Benchmark on Multiple Table Domains](https://mlsanity.com/paper/2404.19205/) (2024)
- [Agent RL Scaling Law: Agent RL with Spontaneous Code Execution for Mathematical Problem Solving](https://mlsanity.com/paper/2505.07773/) (2025)
- [Distilling LLM Agent into Small Models with Retrieval and Code Tools](https://mlsanity.com/paper/2505.17612/) (2025)
- [Mock Worlds, Real Skills: Building Small Agentic Language Models with Synthetic Tasks, Simulated Environments, and Rubric-Based Rewards](https://mlsanity.com/paper/2601.22511/) (2026)
- [Attribution Explanations for Deep Neural Networks: A Theoretical Perspective](https://mlsanity.com/paper/2508.07636/) (2025)

## Abstract (arXiv)

This paper investigates reinforcement learning (RL) methods for improving tool-calling capabilities in multimodal small language model (SLM) agents. While existing works have explored various reward designs to improve agentic tool-calling ability, these approaches face inherent limitations for SLM training, especially under multimodal scenarios. First, many existing methods evaluate tool use correctness through exact matching against certain ground-truth or predefined formats. However, this assumption is often unsuitable for multimodal tasks, where multiple tool use paths may be valid and annotated tool trajectories are typically unavailable. Second, such sparse and brittle binary rewards provide little guidance on how to improve the underlying decision process, making them particularly difficult for multimodal SLM to learn from. To address these issues, we propose Input Attribution-Aware Policy Optimization (IAPO), an RL algorithm for improving tool use in multimodal SLM by aligning the model's attribution across input components with that of a stronger teacher. Experiments on Qwen2.5-VL-3B show that the proposed method improves visual question answering accuracy by an average of 3% across six test sets compared with existing visual tool use work, by helping the model attend to the most relevant input evidence.