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

2026-06-10

AI Summary

Existing reinforcement learning approaches struggle to effectively train tool-use capabilities in multimodal small language models due to their reliance on sparse binary rewards and insufficient guidance toward relevant input evidence. This work introduces input attribution alignment into the reinforcement learning training of multimodal agents for the first time, proposing an attribution-aware policy optimization method that enhances the agent’s focus on critical visual and textual evidence by encouraging the small model to mimic the attention distribution of a strong teacher model over input segments. Experiments based on Qwen2.5-VL-3B demonstrate an average accuracy improvement of 3% across six visual question answering benchmarks, significantly outperforming conventional reward-driven baselines.

Abstract

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.

Problem

• tool use
• multimodal agents
• small language models
• reinforcement learning
• reward design

Innovation

• Input Attribution
• Policy Optimization
• Multimodal Agents
• Tool Use
• Small Language Models