Reinforced Attention Learning: Optimizing Where to Look, Not What to Say

Paper: 2602.04884 Authors: Bangzheng Li, Jianmo Ni, Chen Qu, Ian Miao, Liu Yang, Xingyu Fu, Muhao Chen, Derek Zhiyuan Cheng Affiliations: UC Davis, Princeton University, Google, Google DeepMind Categories: cs.CL, cs.LG

Abstract

Post-training with Reinforcement Learning (RL) has substantially improved reasoning in Large Language Models (LLMs) via test-time scaling. However, extending this paradigm to Multimodal LLMs (MLLMs) through verbose rationales yields limited gains for perception and can even degrade performance. We propose Reinforced Attention Learning (RAL), a policy-gradient framework that directly optimizes internal attention distributions rather than output token sequences.

Key Contributions

• Reinforced Attention Learning (RAL): A policy-gradient method that shifts optimization from next-token prediction to attention-distribution alignment
• On-Policy Attention Distillation: Transfers latent attention behaviors from teacher to student, yielding stronger cross-modal alignment than standard knowledge distillation
• Empirical Validation: Consistent improvements over GRPO across diverse visual question answering benchmarks

The Problem

Traditional RL for MLLMs optimizes what to generate (token probabilities). The standard RL objective uses importance sampling with the probability ratio between current and old policies.

This approach has limitations for multimodal tasks:

• Verbose Chain-of-Thought reasoning provides marginal gains for perception tasks
• Can even degrade fundamental perceptual capabilities
• Optimizes the result (token) rather than the process (information allocation)

The Solution: Optimize Where to Attend

RAL treats attention patterns as the policy:

Aggregated Causal Attention Distribution Policy

For each generated token t, define the attention distribution over preceding positions. The attention weight from token t to position i is normalized across all preceding positions.

Advantage-Weighted Attention Divergence

The RAL loss function computes the divergence between current and old attention distributions, weighted by the advantage:

• High reward response: Minimize divergence to encourage the attention pattern
• Low reward response: Maximize divergence to penalize the attention pattern

On-Policy Attention Distillation

Traditional distillation focuses on token-level probability alignment. RAL extends this with dual-distillation:

1. Token Distribution Alignment: Standard KD loss
2. Attention Distribution Alignment: Transfer attention patterns from teacher to student

This provides significant additional performance gains by transferring fine-grained perceptual and grounding behaviors.

Results

RAL consistently outperforms GRPO across video and image benchmarks, particularly on perception-intensive tasks:

• Better visual grounding
• Reduced modality bias (over-reliance on linguistic priors)
• More faithful, evidence-based responses

Why This Matters

The key insight: attention is the mechanism for cross-modal reasoning. By directly shaping attention weights, RAL provides a principled way to control how the model integrates visual and textual information.

Takeaways

1. Attention as policy: Internal attention distributions can be treated as a learnable policy space
2. Process over result: Optimizing how the model allocates attention is more effective than optimizing output tokens for perception tasks
3. Structural regularization: Attention-based optimization provides robust regularization against diversity collapse
4. Distillation beyond tokens: Transferring attention patterns is more effective than token-level distillation for multimodal alignment