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

Problem Statement

The paper addresses the challenge of training Small Language Model (SLM) agents to effectively use tools in multimodal scenarios, specifically for visual question answering (VQA) involving charts and tables. While Reinforcement Learning (RL) methods like Proximal Policy Optimization (PPO) and Group Relative Policy Optimization (GRPO) have improved agent performance, they face significant limitations when applied to multimodal SLMs:

1. Inadequate Reward Design for Multimodal Tasks: Existing methods often rely on exact matching against ground-truth tool trajectories or predefined formats. However, multimodal tasks frequently admit multiple valid tool-use paths (e.g., highlighting, masking, or drawing bounding boxes can all be valid for focusing on relevant chart regions). Rigid format checks fail to capture this flexibility.
2. Sparse and Brittle Rewards: Current approaches typically use binary, outcome-based rewards (correct/incorrect final answer). These provide sparse feedback and fail to guide the model's intermediate decision-making process. Consequently, SLMs may learn to exploit accidental shortcuts or superficial correlations to achieve a correct final answer without genuinely learning robust tool-selection policies grounded in relevant visual evidence.
3. Suboptimal Performance of Vanilla GRPO: Even with vanilla GRPO, multimodal SLMs continue to produce substantial numbers of misused tool calls, indicating that standard RL algorithms are insufficient for learning the nuanced decision-making required for multimodal tool use.

Methodology: Input Attribution-Aware Policy Optimization (IAPO)

To address these issues, the authors propose IAPO, an RL algorithm that aligns the student model's input attribution with that of a stronger teacher model. The core idea is that the contribution of specific input components to a generated tool call is strongly correlated with tool-call correctness.

Key Components:

1. Integrated Gradient (IG) Scores:

   • The input prompt is decomposed into semantic blocks: pure-text context ($P$), user query ($R$), image placeholder ($I$), and tool definitions ($T$).
   • The method computes Integrated Gradient scores to quantify how much each semantic block contributes to the log-probability of the generated tool-call sequence.
   • Specifically, for a tool-call sequence $y'$, the score $s(E)$ is the log-probability under the current policy. The gradient of this score with respect to input embeddings is calculated, and the token-level attribution is defined as the norm of the element-wise product of the gradient and the embedding ($o_i = \|g_i \odot e_i\|_2$).
   • These token-level attributions are aggregated to produce block-level scores ($O_k$) for each semantic type.

2. Attribution-Based Penalty:

   • A stronger teacher model (e.g., a larger SLM or an RL-trained model) generates reference IG scores ($\hat{b}$) for the same input.
   • The student model's IG scores ($\hat{a}$) are compared to the teacher's using Kullback-Leibler (KL) divergence to form an attribution penalty: $P_A = D_{KL}(\hat{b} \parallel \hat{a})$.
   • This penalty encourages the student to focus on the same input components (e.g., specific tool definitions or query parts) that the teacher deems relevant.

3. Modified Reward Function:

   • The final reward for a trajectory $i$ combines the standard outcome-based reward ($R_{out}$) with the attribution penalty:
     $$R_{IAPO}^i = R_{out}^i - \lambda P_A^i$$
   • Here, $\lambda$ controls the strength of the attribution alignment. This dense signal guides the policy to not only maximize task accuracy but also to align its internal decision process with the teacher's attribution patterns.

4. Training Framework:

   • The method uses a two-turn rollout: the model generates tool calls to edit an image, then reasons over the edited image to produce a final answer.
   • Training employs the GRPO objective but utilizes the modified $R_{IAPO}$ to compute group-relative advantages.

Key Contributions

1. New Integrated Gradient Scores: The paper introduces a method to quantify the contribution of different input semantic blocks (text, image, tool definitions) to tool-call generation, providing a dense supervision signal.
2. IAPO Algorithm: The authors propose the first RL method specifically designed for SLM tool use that utilizes attribution-based penalties. By comparing student and teacher IG scores via KL divergence, IAPO transfers the teacher's tool-calling capability and decision-making logic to the student.
3. Empirical Validation: The paper demonstrates that vanilla GRPO leads to persistent tool misuse in multimodal SLMs. It further shows that adapting reward designs from text-only models is insufficient, whereas IAPO significantly improves performance by aligning input attribution.

Experimental Results

The method was evaluated on visual question answering tasks using the Qwen2.5-VL-3B model as the policy and a Qwen2.5-VL-7B model (trained with GRPO) as the teacher.

• Performance Gains: IAPO improved visual question answering accuracy by an average of 3% across six test sets (including ChartQA, VWTQ, and CharXiv) compared to existing visual tool-use work (VTool-R1, TRM, ToRL).
• Tool Selection Accuracy: Case studies revealed that even after 100 steps of vanilla GRPO training, the SLM made significant tool-selection errors (e.g., applying row-based tools to column-divisible charts). IAPO significantly reduced these errors, with the model's tool-call success rate approaching 100% in later training stages.
• Attribution Alignment: The study showed that IG scores for correct input parts increased during training and eventually surpassed those of misused parts, correlating with the reduction in tool misuse. The 7B teacher model exhibited clearer separation in IG scores between correct and incorrect inputs, validating the use of a stronger teacher for guidance.
• Generalization: IAPO demonstrated superior generalization on out-of-distribution benchmarks (CharXiv), suggesting that attribution alignment helps prevent overfitting to superficial prompt correlations.

Significance and Claims

The paper claims that IAPO provides a more effective framework for training multimodal SLM agents than existing RL methods. Its significance lies in:

• Bridging the Gap for SLMs: It addresses the specific difficulty SLMs face in learning from sparse, binary rewards by introducing a dense, process-level supervision signal derived from input attribution.
• Moving Beyond Binary Rewards: It challenges the reliance on exact matching or binary outcome rewards in multimodal settings, proposing instead that aligning how a model attends to input (attribution) is crucial for robust tool use.
• Teacher-Student Distillation in RL: It successfully leverages a stronger teacher model not just for output imitation, but for guiding the internal decision-making process (attribution patterns) of a smaller student model during online RL training.

The authors conclude that while larger models naturally exhibit better tool-selection abilities and attribution patterns, IAPO effectively transfers these capabilities to smaller models, enabling them to attend to the most relevant input evidence and achieve state-of-the-art performance in multimodal tool-use tasks.