Information-Regularized Attention for Visual-Centric Reasoning

Conference: ECCV2026
arXiv: 2607.00434
Code: No public code available yet
Area: Multimodal VLM / Representation Learning
Keywords: Vision-Language Models, Information Bottleneck, Attention Regularization, Hallucination Suppression, Variational Inference

TL;DR

This paper proposes Information-Regularized Attention (IRA), which applies stochastic regularization based on the Information Bottleneck principle to visual value states within the attention modules of VLMs. By injecting data-dependent hierarchical noise through a variational inference framework, the model learns to actively control the injection of visual information during end-to-end full-parameter fine-tuning, thereby simultaneously mitigating hallucination, weak visual grounding, and catastrophic forgetting.

Background & Motivation

Modern VLMs have made breakthrough progress in tasks such as visual question answering, image captioning, and multimodal dialogue. However, they still face three major reliability issues: object hallucination (generating content that does not exist in the image), weak visual grounding (attention failing to focus on relevant regions), and catastrophic forgetting after full-parameter instruction tuning. Current mainstream coping strategies are almost entirely data-driven—methods like visual instruction tuning, preference optimization, and policy optimization attempt to correct model behavior from the perspective of supervision signals. However, the essential commonality of these methods is that they all optimize the model under the standard next-token prediction framework, where visual embeddings are only indirectly updated via language supervision, lacking direct constraints on the visual representation space. This means that task-irrelevant or even noisy visual signals can propagate indiscriminately through the attention layers, interfering with the reliability of cross-modal reasoning.

This deficiency is particularly evident in attention patterns. Recent studies have widely reported the phenomena of attention sinks and spike values in attention heads—the model's attention often collapses onto semantically uninformative visual tokens, resulting in noisy cross-modal interactions. Existing mitigation solutions, such as gating mechanisms and attention distribution optimization, mainly operate at the level of attention weights or attention outputs without addressing the root cause of intermediate representation quality. The starting point of this paper is to re-evaluate cross-modal issues in VLMs within an information-theoretic framework: if the visual representation \(h\) is the bottleneck intermediate variable from input \(x\) to output \(y\), standard SFT only maximizes \(I(h; y)\) (information useful for prediction) but fails to discard \(I(h; x|y)\) (noisy information useless for prediction). This is precisely the term that needs to be constrained in Information Bottleneck (IB) theory.

Based on these insights, this paper proposes applying hierarchical stochastic regularization to the visual value states within the attention modules of a group of mid-to-late transformer layers. The core idea is that by transforming the visual representation from a deterministic variable to a random variable, one can explicitly control the amount of information injected per layer using KL divergence, allowing the model to autonomously learn which visual signals to retain and which noise to suppress during end-to-end training. Core Idea: Introduce a stochastic regularization mechanism based on the Information Bottleneck principle to the value states in attention modules. By injecting data-dependent hierarchical noise for each visual token and explicitly constraining the representation's information capacity via a variational KL term, the model is compelled to simultaneously learn to "see" and "filter" during end-to-end full-parameter fine-tuning.

Method

Overall Architecture

IRA operates during the full-parameter SFT phase of VLMs. In standard VLM architectures (represented by InternVL2/LLaVA-OneVision), a vision encoder maps images to token sequences, which are projected and concatenated with text tokens before being fed into the LLM. The core of IRA is to apply variational stochastic regularization to the visual value states within the attention modules of the mid-to-late layers of the LLM (approximately 60%–80% depth). Specifically, in each IRA layer, the visual value status \(v\) is extracted, and a lightweight linear head predicts the shift and variance of the posterior distribution. Under the Information Bottleneck framework, the posterior is constrained against a prior anchored on the pre-trained representation via KL divergence. A differential regularization is applied to different visual tokens using an uncertainty-aware weighting factor. The regularization strength is gradually increased from 0 to its maximum value using a cosine scheduler to avoid disrupting the pre-trained representation structure in the early stages of training.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    A["Input Image + Text"] --> B["Vision Encoder + Projector"]
    B --> C["LLM Hidden State Sequence"]
    C --> D["Lower Transformer Layers"]
    D --> E["IRA Layer (Mid-to-Late 60%-80%)"]
    E --> F["Upper Transformer Layers"]
    F --> G["Output Token"]

    subgraph H["Internal Process of IRA Layer"]
        H1["Q / K / V Projections"] --> H2["Extract Visual Value v"]
        H2 --> H3["Variational Inference<br/>Prior p(z|⌊v⌋) vs Posterior q(z|v)"]
        H3 --> H4["Uncertainty Weighting<br/>g = ℍ · (1-a)"]
        H4 --> H5["Reparameterization Trick<br/>z = v + Δ + g·σ·ε"]
        H5 --> H6["Stochastic Visual Z Replaces V<br/>→ Attention(Q, K, Z)"]
    end

    E --> H

Key Designs

1. Variational Information Bottleneck Framework: Reshaping Visual Representation Learning as a Lossy Compression Problem

Standard SFT of VLMs only maximizes the conditional likelihood \(p_\theta(y|h)\), which is equivalent to maximizing \(I(h; y)\), but places no constraint on the task-irrelevant noise information \(I(h; x|y)\) in \(h\). According to Information Bottleneck theory, the ideal optimization objective is \(\max I(h; y) - \beta \cdot I(h; x|y)\), which retains predictive information while compressing input noise. The core insight of IRA is that by treating the visual representation at each layer as a random variable \(z^{(\ell)}\) instead of a deterministic variable \(h^{(\ell)}\), and injecting stochastic local reparameterization noise into the attention modules, the Information Bottleneck objective can be formulated as a differentiable training loss via variational inference. Specifically, the entire model defines a Markov chain from input \(x\) to hidden representation \(z\) and then to output \(y\), yielding the following optimization objective:

\[ \mathcal{L}(\theta,\phi,\lambda|y) = \mathbb{E}\big[\log p_\theta(y|z^{(L)})\big] - \beta \sum_{\ell=1}^{L} D_{\text{KL}}\big(q_{\theta,\phi}(z^{(\ell)}|v^{(\ell)}) \,\|\, p_\lambda(z^{(\ell)}|\lfloor v^{(\ell)}\rfloor)\big) \]

The first term represents the reconstruction loss (maintaining task performance), and the second term is the hierarchical KL regularization (constraining information capacity). Unlike simple dropout or weight decay, KL regularization directly acts on the representation space, forcing the representation of each visual token to retain only information highly reliable for prediction while encoding uncertainty as local noise.

2. Data-Dependent Prior and Residual Posterior: Enforcing Hierarchically Controllable Stochasticity within Attention Modules

Directly injecting noise into pre-trained VLM representations can disrupt well-learned embedding structures. IRA avoids this through a carefully designed prior-posterior mechanism to achieve a smooth transition. The prior distribution is defined as a Gaussian anchored on the stop-gradient (stop-grad) visual value states: \(p_\lambda(z|\lfloor v\rfloor) = \mathcal{N}(\lfloor v\rfloor, \sigma_p^2 I)\), where \(\lfloor \cdot \rfloor\) denotes stop-gradient to prevent the prior and posterior from simultaneously depending on \(v\) (which would cause the KL divergence to quickly collapse to zero); \(\sigma_p^2\) is a layer-shared learnable parameter representing a zero-mean Gaussian noise scale learned across the entire dataset. The posterior is defined as a Gaussian with a learnable residual shift: \(q_{\theta,\phi}(z|v) = \mathcal{N}(v + \Delta_\phi(v), \sigma_q^2 I)\). Here, \(\Delta_\phi(v)\) and \(\log \sigma_q^2(v)\) are predicted using a highly lightweight linear head (one Linear(d, d+1) and one Embedding(H, d) per layer). \(\Delta_\phi(v)\) captures the residual adjustment to pre-trained representations, while \(\sigma_q^2\) controls the local uncertainty of each attention head. Sampling is performed via the reparameterization trick:

\[ z = v + \Delta_\phi(v) + \sigma_q \ast \epsilon, \quad \epsilon \sim \mathcal{N}(0, I) \]

This formulation is equivalent to superimposing a learnable feature-level noise \(\Delta_\phi(v) + \sigma_q \cdot \epsilon\) onto the pre-trained visual representation \(v\). Grounding the prior on the pre-trained representation ensures that the model does not deviate too far from the learned knowledge base, while the residual posterior design provides flexible adaptive learning space. During inference, all stochasticity is removed (i.e., \(z = v + \Delta_\phi(v)\)), introducing zero computational overhead.

3. Uncertainty-Aware Token-Level Adaptive Weighting: Differentiated Regularization to Avoid Over-Compressing Critical Information

Different visual tokens contribute very differently to the final prediction; foreground objects and irrelevant background clearly do not require the same regularization strength. To address this, IRA designs a two-stage weighting mechanism. It first computes the average attention mass \(a_i = (1/T) \sum_j p_{j,i}\) for each visual token \(i\), measuring the "attention" of the text context on this visual token. Concurrently, it computes the normalized attention entropy \(\mathbb{H} = (1/H) \sum_n [-\sum_j p_{j,n} \log p_{j,n} / \log T]\) to quantify the certainty of attention allocation. The weighting factor is defined as \(g_i = \mathbb{H} \cdot (1 - a_i)\). Consequently, tokens with high entropy (diffuse/uncertain attention) and low importance (small attention mass) receive higher regularization weights, while tokens with low entropy (the model is certain) and high importance (highly attended) retain more of their original information. In the calculation of this weighting factor, both the \(q\) and \(k\) parts employ stop-gradient to prevent the KL term from backpropagating to the query/key projections. The weighted KL divergence is applied to each token, scaling the noise magnitude during sampling: \(z_i = v_i + \Delta_\phi(v_i) + g_i \cdot \sigma_i \cdot \epsilon\). Ablation experiments show that removing this weighting mechanism leads to a drop in performance across the board, particularly on MMMU (-1.0) and MuirBench (-2.5), suggesting that uniform regularization is too aggressive, and adaptive weighting achieves a fine balance between information preservation and noise compression.

Loss & Training

The final training objective of IRA, as described above, consists of the reconstruction loss and the hierarchical KL regularization terms. \(\beta\) utilizes a cosine scheduler (cosine interpolation) that linearly scales from 0 to \(\beta_{\max}\) during the first \(k\%\) of training steps. This beta warm-up allows the model to stabilize for a few steps under the unregularized SFT objective before gradually introducing the KL constraint, preventing the model from falling into a local optimum of \(D_{\text{KL}} \approx 0\) at the start of training. Empirically, \(\beta_{\max} = 1 \times 10^{-4}\) and \(k = 50\%\) yield the best performance. IRA is only active during the SFT stage and is not applied in pre-training (Stage 1/Stage 1.5). The learning rate for IRA parameters is set to 10 times the backbone network learning rate (\(1 \times 10^{-4}\) vs. \(1 \times 10^{-5}\)), enabling the regularization term to quickly catch up with task optimization during the early steps of training.

Key Experimental Results

Main Results

Across three models—InternVL2-8B, InternVL2.5-8B, and LLaVA-OneVision-8B—IRA consistently outperforms standard SFT on over 10 benchmarks. The table below compares the performance of InternVL2.5-8B on comprehensive visual understanding benchmarks (most significant improvements under IRA are bolded):

Benchmark Type	Benchmark	SFT	+IRA	Gain
Reasoning/Knowledge	MMMU-Pro	30.4	30.6	+0.2
Reasoning/Knowledge	MMMU	46.4	47.6	+1.2
Reasoning/Knowledge	MME	1981	2038	+57
Reasoning/Knowledge	OK-VQA	80.6	81.8	+1.2
Visual Perception	MMStar	61.1	58.8	-2.3
Text-Intensive	TextVQA	74.5	74.7	+0.2
Text-Intensive	ChartQA	81.0	81.8	+0.8
Text-Intensive	DocVQA	86.7	86.7	Unchanged

The improvements from IRA are concentrated in reasoning-intensive benchmarks (MMMU/MME/OK-VQA), whereas there is a slight decrease on MMStar (pure perception). This indicates that IRA primarily enhances the robustness of cross-modal compositional reasoning rather than merely aligning surface-level statistics. An improvement of 57 points on MME (1981 to 2038) is a substantial margin, demonstrating that stochastic regularization significantly boosts the semantic abstraction and cross-modal reasoning capabilities of the model in open scenarios.

Regarding robustness and generalization (InternVL2.5-8B):

Benchmark	Evaluation Dimension	SFT	+IRA	Gain
POPE	Hallucination Detection	87.0	87.5	+0.5
HallusionBench	Fine-Grained Hallucination	37.6	37.9	+0.3
VLM-Bias	Low-Level Visual Bias	17.8	18.3	+0.5
VLM-Blind	Low-Level Visual Blindspot	33.9	37.3	+3.4
MVBench	Video Understanding	51.5	52.0	+0.5

IRA achieves the most substantial improvement (+3.4) on VLM-Blind (which evaluates whether a model truly "sees" low-level visual features instead of relying on language priors), proving that stochastic regularization indeed improves the model's reliance on actual visual signals.

Ablation Study

Configuration	MMMU	VQAtext	ChartQA	EmbSpatial	MuirBench	7-Task Avg.
Full IRA	46.1	70.3	79.8	64.7	38.4	61.3
w/o Weighted KL	45.1	69.5	80.4	65.6	35.9	60.7
w/o Prior Anchoring v	45.6	69.4	79.6	64.6	34.6	60.6
w/o IRA (Standard SFT)	44.9	69.6	76.5	63.5	34.7	59.5

Removing uncertainty weighting leads to a drop of 1.0 on MMMU and 2.5 on MuirBench, indicating that uniform regularization is overly aggressive. Removing prior anchoring degrades performance further (particularly on MuirBench and VQAtext), demonstrating that anchoring on the pre-trained representation successfully preserves fine-grained visual details.

Applying IRA at the mid-to-late layers of the LLM (60%–80% depth, corresponding to layers 20–26 in a 32-layer LLM) yields the best results. These layers have accumulated sufficient abstract representation suitable for information compression, while still retaining strong visual grounding signals. Applying IRA too broadly (0%–100%) or too shallowly (20%–40%) leads to performance degradation.

Key Findings

• Information bottleneck corrects the attention distribution: IRA reduces the attention sink ratio of InternVL2.5 from 46.9% to 40.6% while simultaneously improving accuracy, demonstrating that controlled compression of visual information directly improves attention focus.
• Representation curvature is an effective proxy indicator of attention quality: IRA yields a straighter curvature trajectory, which is highly correlated with more stable attention patterns and fewer sinks. This finding provides a new evaluation dimension for designing robust architectures in the future.
• IRA mitigates catastrophic forgetting: Throughout the training process across 10 benchmarks, standard SFT shows a rise-then-fall average performance curve (overfitting), whereas IRA exhibits a more stable and monotonic growth curve.
• Text-intensive tasks also benefit: Even on tasks requiring high-fidelity visual transmission, such as TextVQA and ChartQA, IRA still brings improvements, meaning the adaptive weighting successfully protects vital fine-grained visual details.

Highlights & Insights

• Architectural-level rather than data-level solution: Unlike existing data-driven approaches like preference optimization and instruction tuning, IRA addresses VLM hallucinations from the fundamental level of representation learning. This serves as a complement to the "data-centric pattern coverage" mindset, and the two can be orthogonally integrated.
• Dual role of stochastic attention (regularization & representation): The paper demonstrates empirically that stochastic attention acts not just as a regularizer, but also as an active contributor to representation learning—it reshapes the representation geometry and generates straighter curvature trajectories. This provides empirical support for using stochastic attention as a first-class citizen in VLM architectures.
• Joint design of attention entropy and token importance weighting: Incorporating both attention mass (importance) and attention entropy (uncertainty) to build the weighting factor is highly complementary. Tokens with low attention mass but low entropy (regions where the model is certain but they are unimportant) are properly protected, avoiding a "one-size-fits-all" single-metric approach.
• Prior anchoring + stop-gradient hook to prevent collapse: Data-dependent priors naturally risk rapid KL collapse (since both prior and posterior would drift towards \(v\) together). The stop-gradient design is simple yet highly effective, offering high transferability to other scenarios requiring data-dependent priors.

Limitations & Future Work

• Validated only on models up to 8B parameters; the performance on larger models (13B/26B/73B) remains speculative.
• During training, extracting visual value states layer-by-layer and computing attention statistics (entropy/importance) increases codebase complexity and GPU memory overhead. Although stochasticity is removed in inference, preserving the learned mean values still introduces a minor computational footprint.
• The uncertainty weighting \(g_i = \mathbb{H} \cdot (1 - a_i)\) is an empirical design; the theoretically optimal weighting formulation warrants further derivation.
• Present applications are limited to the SFT stage. The authors speculate that IRA could be used during the LLM/VLM pre-training stage, but this has not been validated.
• Applicability to other modalities (e.g., audio, video frame sequences) has not been explored.

Related Work & Insights

• vs. Data-driven hallucination mitigation (DPO/RLHF/VL-feedback): These schemes correct the output distribution during optimization, whereas IRA squashes representation noise. They are complementary and can be combined.
• vs. Attention weight optimization (Reinforced Attention Learning, Gated Attention): Existing methods modify attention weights or outputs. IRA intervenes earlier on the attention inputs (value states) to control information quality at the source.
• vs. Information Bottleneck methods (VIB-Probe etc.): VIB-Probe employs Information Bottleneck as an external detector to filter hallucinations, whereas IRA internalizes IB principles into the training objectives, directly constraining representations layer by layer.
• Connection to Ladder-VAE: IRA sharing backbone parameters between the inference and generative processes share similarities with the hierarchical variational structure of Ladder-VAEs, though no explicit bottom-up/top-down pathways are constructed.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Applying stochastic regularization to the value states in attention modules introduces a novel perspective based on information bottleneck principles—focusing on representations instead of weights.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ Extensively validated on 3 mainstream VLM architectures over 20+ benchmarks (including hallucination, robustness, and out-of-distribution multi-image/video evaluation), alongside deep analyses of representation curvature, attention sinks, and training stability.
• Writing Quality: ⭐⭐⭐⭐ The motivation is very clear and the mathematical derivations are thorough, although some equations are densely packed and the appendix layout could be improved.
• Value: ⭐⭐⭐⭐⭐ Provides an architectural-level hallucination mitigation paradigm that is orthogonal to data-driven methods. The stochastic regularization within attention is both elegant and lightweight, making it highly practical for engineering deployment.

Related Papers

• [ICLR 2026] Small Drafts, Big Verdict: Information-Intensive Visual Reasoning via Speculation
• [ICLR 2026] CompoDistill: Attention Distillation for Compositional Reasoning in Multimodal LLMs
• [ECCV 2026] ScAle: Attention Head Scaling as a Minimal Adapter for Spatial Reasoning in Vision–Language Models
• [CVPR 2026] Think360: Evaluating the Width-centric Reasoning Capability of MLLMs Beyond Depth
• [ICLR 2026] VideoReasonBench: Can MLLMs Perform Vision-Centric Complex Video Reasoning?