Information-Regularized Attention for Visual-Centric Reasoning

Technical Summary: Information-Regularized Attention for Visual-Centric Reasoning

Problem Statement
Vision-language models (VLMs) have established themselves as a paradigm for multimodal learning but suffer from significant reliability issues, including object hallucination, weak visual grounding, and catastrophic forgetting following full-parameter instruction tuning. The authors argue that these failures stem from a lack of explicit control over visual representation learning during the standard next-token prediction objective. Under current training paradigms, visual embeddings are optimized passively through language supervision alone. This allows task-irrelevant or noisy visual signals to propagate through attention layers, leading to "attention sinks" (where attention collapses onto semantically uninformative tokens) and spurious cross-modal interactions. Existing data-centric approaches, which rely on expanding supervision signals (e.g., preference optimization), fail to directly regularize the feature representations themselves.

Methodology: Information-Regularized Attention (IRA)
To address these issues, the paper proposes Information-Regularized Attention (IRA), a stochastic attention mechanism designed to explicitly regulate the amount of visual information injected into the hidden states of intermediate transformer layers.

• Information-Theoretic Formulation: The method frames the problem as a variational inference task. Instead of treating hidden states $h$ as deterministic mappings of input $x$, IRA introduces a stochastic latent variable $z$ by injecting random noise into the value states of the attention module before they enter the attention computation. This transforms the deterministic propagation of visual embeddings into a stochastic process.
• Architecture: A lightweight parametric head is attached to each transformer layer to produce a posterior distribution $q_\phi(z | v)$, where $v$ represents the value states. This posterior is approximated using a reparameterization trick ($z = v + \Delta_\phi(v) + \sigma_q \epsilon$), allowing the model to learn adaptive perturbations and uncertainty estimates without additional encoding overhead.
• Regularization Objective: The training objective includes a Kullback-Leibler (KL) divergence term between the learned posterior and a data-dependent prior $p_\lambda(z | \lfloor v \rfloor)$. The prior is centered on the detached original value state (using a stop-gradient operation) to prevent trivial posterior-prior collapse.
• Uncertainty-Conditioned Weighting: Recognizing that not all visual tokens require equal regularization, the authors introduce a token-wise weighting mechanism. The regularization magnitude is adaptively adjusted based on two factors:
  1. Importance: The average attention mass assigned to a token by language contexts.
  2. Uncertainty: The entropy of the cross-modal attention distribution.
     Tokens that are unimportant (low attention) and uncertain (high entropy) receive stronger regularization, while critical visual details are preserved.
• Layer-Specific Application: Experiments indicate that applying IRA to the mid-to-late layers (specifically 60%–80% of the transformer depth) yields the best performance, as these layers balance abstract reasoning with retained visual grounding signals.

Key Contributions

1. Novel Mechanism: The proposal of IRA, which regulates visual representations at the representation level before they enter cross-modal attention, explicitly controlling visual information flow during full-parameter instruction tuning.
2. Geometric Insights: The demonstration that IRA improves representation geometry. The authors introduce a curvature metric to analyze neural trajectories, finding that IRA produces smoother, more linear trajectories across layers, which correlates with stable training dynamics and improved grounding.
3. Attention Sink Mitigation: The identification of a correlation between attention sinks and representational curvature. IRA is shown to suppress attention sinks and reduce the ratio of tokens receiving excessive, uninformative attention.

Results
The authors evaluated IRA on open-source VLMs (InternVL2, InternVL2.5, and LLaVA-OneVision) across various benchmarks:

• General Performance: IRA consistently outperformed standard Supervised Fine-Tuning (SFT) on reasoning-intensive benchmarks (e.g., MMMU, STEM, General QA) and knowledge-intensive tasks (OK-VQA).
• Robustness: The method improved performance on hallucination detection benchmarks (POPE, HallusionBench) and reduced bias in low-level vision tasks (VLM-are-biased, VLM-are-blind).
• Generalization: IRA demonstrated better generalization to Out-of-Distribution (OOD) tasks, including multi-image understanding (MuirBench, BLINK) and video understanding (MVBench).
• Representation Analysis: Quantitative analysis showed that IRA models exhibited lower curvature changes (straighter trajectories) and significantly reduced attention sink ratios compared to SFT baselines.
• Ablation: Removing the uncertainty-guided weighting or the data-dependent prior led to performance degradation, confirming the necessity of adaptive regularization and anchoring to pretrained representations.

Significance and Claims
The paper claims that stochastic attention is not merely a regularizer but a fundamental contributor to representation learning within generative architectures. By explicitly regulating the information injected into hidden states, IRA addresses the root cause of visual hallucinations and instability: the uncontrolled propagation of noisy visual signals. The authors posit that improving VLM reliability requires moving beyond data-centric post-training to directly regularize feature representations within the generative architecture. Furthermore, the work suggests that attention should be viewed not just as a weighting mechanism but as a metric for evaluating internal representation quality, offering a new direction for building more reliable and interpretable multimodal systems. The authors note that while current experiments are limited to 8B parameter models due to resource constraints, they expect the findings to scale to larger models.