Visual attention improves reinforcement learning for multimodal tasks

VAttnPO: Visual Attention-Guided Policy Optimization

Computer Vision and Pattern RecognitionArtificial Intelligence

Summary

It is hard to tell which parts of an image a machine learning model looks at while producing each word in a description or answer, especially when the feedback only rates the whole answer. The authors introduce VAttnPO, a method that uses where the model looks in the image to give better credit to each word during learning. This approach measures the agreement between attention patterns when the correct answer is known and when it is generated, helping guide the model more precisely. Using this method improves performance on image-and-language tasks over previous approaches.

What this means in practice

  • For multimodal ai developers: Improve training of AI models that generate text based on images by refining token-level learning signals with visual attention guidance.
  • For ai product teams: Enhance multimodal chatbots or assistants to provide more accurate and visually grounded responses by better aligning generated words with image regions during training.$Commercial implications: This method can produce higher-quality visual-language interactions, enabling products that serve customers with more trustworthy image-based responses.

Authors

Lihao Zheng, Zhenwei Shao, Zhou Yu, Yan Yang, Haoyi Sun, Qian Wang, Ning Mao, Pan Zhou

Abstract

Multimodal reinforcement learning with verifiable rewards commonly assigns one sequence advantage to an entire response. This outcome-level signal distinguishes overall response quality but leaves changes in visual-evidence use across generated tokens implicit. We introduce VAttnPO (Visual Attention-Guided Policy Optimization), which uses answer-conditioned visual attention to refine token-level credit while preserving the verifier-derived trajectory-mean advantage. A frozen old policy constructs a visual reference by teacher forcing the known correct answer and produces a visual distribution for each rollout token. VAttnPO measures their agreement in visual-attention space using Jensen--Shannon divergence, rather than reading guidance from teacher output-token distributions or introducing an attention-matching objective. Trajectory centering, relative-position weighting, and credit balancing turn this agreement into a zero-mean, position-aware correction to the GRPO sequence advantage. Token intervention, visual-index permutation, and trajectory analysis provide diagnostic evidence linking this signal to model behavior, visual correspondence, and generation position. On Qwen3-VL-8B-Instruct, VAttnPO raises six-task image AVE from 70.97 with GRPO to 73.50 and scores higher on every reported structured, video, and MME metric. Gains over GRPO also hold on Qwen2.5-VL-7B-Instruct and InternVL3-8B-Instruct.