Information-Regularized Attention for Visual-Centric Reasoning

Conference: ECCV2026
arXiv: 2607.00434
Code: No public code available
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 amount of visual information injected 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 like visual question answering, image captioning, and multimodal dialogue, but still face three reliability issues: object hallucination (generating content non-existent in the image), weak visual grounding (attention failing to focus on relevant areas), and catastrophic forgetting after full-parameter instruction tuning. Current mainstream strategies are predominantly data-driven—visual instruction tuning, preference optimization, and policy optimization attempt to correct model behavior from the perspective of supervision signals. However, the fundamental commonality of these methods is that they optimize the model under the standard next-token prediction framework, where visual embeddings are only updated indirectly through language supervision, lacking direct constraints on the visual representation space. This means that task-irrelevant or even noisy visual signals can propagate indiscriminately through attention layers, interfering with the reliability of cross-modal reasoning.

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

Based on these insights, this paper proposes applying hierarchical stochastic regularization to visual value states inside the attention modules of a group of mid-to-late transformer layers. The core idea is that changing the visual representation from a deterministic variable to a random variable allows explicit control over the information injected at each layer via KL divergence, enabling 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—injecting data-dependent hierarchical noise for each visual token and explicitly constraining the representation's information capacity via a variational KL term, allowing the model 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 (like InternVL2/LLaVA-OneVision), the visual encoder maps the image to a sequence of tokens, which is then 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 LLM's mid-to-late layers (approx. 60%-80% depth). Specifically, in each IRA layer, the visual value state \(v\) is extracted, and a lightweight linear head is used to predict the shift and variance of the posterior distribution. Under the information bottleneck framework, the KL divergence between this posterior and a prior anchored on the pre-trained representation is constrained. A differentiated 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 schedule 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["Visual Encoder + Projector"]
    B --> C["LLM Hidden State Sequence"]
    C --> D["Lower Transformer Layers"]
    D --> E["IRA Layers (mid-to-late 60%-80%)"]
    E --> F["Upper Transformer Layers"]
    F --> G["Output token"]

    subgraph H["IRA Internal Flow"]
        H1["Q / K / V Projection"] --> 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. Information Bottleneck Variational Framework: Reshaping Visual Representation Learning as a Lossy Compression Optimization Problem

Standard VLM SFT only maximizes the conditional likelihood \(p_\theta(y|h)\), which is equivalent to maximizing \(I(h; y)\), but places no constraint on the noise information \(I(h; x|y)\) in \(h\) that is irrelevant to \(y\). In information bottleneck theory, the ideal learning objective is \(\max I(h; y) - \beta \cdot I(h; x|y)\)—retaining predictive information while compressing input noise. The key insight of IRA is that by changing the visual representation at each layer from a deterministic variable \(h^{(l)}\) to a random variable \(z^{(l)}\) and injecting stochastic local reparameterization noise within the attention modules, the information bottleneck objective can be converted into a differentiable training loss via variational inference. Specifically, the entire model defines a Markov chain from input \(x\) to hidden states \(z\) and then to output \(y\), where the optimization objective becomes:

\[ \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 is the reconstruction loss (maintaining task performance), and the second term is the hierarchical KL regularization term (constraining information capacity). Unlike simple dropout or weight decay, KL regularization directly acts on the representation space: it forces each visual token's representation to preserve only the information reliable for prediction, while encoding uncertainties as local noise.

2. Data-Dependent Prior and Residual Posterior: Achieving Hierarchical and Controllable Stochasticity Injection within Attention Modules

Directly adding noise to pre-trained VLM representations can destroy maturely learned embedding structures. IRA achieves a smooth transition through a carefully designed prior-posterior mechanism. The prior distribution is defined as a Gaussian anchored on the stop-gradient visual value states: \(p_\lambda(z|\lfloor v\rfloor) = \mathcal{N}(\lfloor v\rfloor, \sigma_p^2 I)\). Here, \(\lfloor \cdot \rfloor\) denotes stop-gradient, which prevents the KL from collapsing to zero by stopping both the prior and posterior from concurrently depending on \(v\); \(\sigma_p^2\) is a layer-shared learnable parameter, representing a zero-mean Gaussian noise scale learned from 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 by an extremely lightweight linear head (one Linear(d, d+1) and one Embedding(H, d) per layer), where \(\Delta_\phi(v)\) captures residual adjustments to the pre-trained representation, and \(\sigma_q^2\) controls the local uncertainty of each attention head. The final representation is sampled via the reparameterization trick:

\[ z = v + \Delta_\phi(v) + \sigma_q * \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\) on top of the pre-trained visual representation \(v\). Anchoring the prior to pre-trained representations ensures that the model does not drift too far from learned knowledge, while the residual posterior design provides flexible adaptive adjustment. During inference, all stochasticity is removed (\(z = v + \Delta_\phi(v)\)), introducing no extra inference 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; salient objects and irrelevant backgrounds obviously do not require the same regularization strength. To address this, IRA designs a two-stage weighting mechanism. First, it computes the average attention mass for each visual token \(a_i = (1/T) \sum_j p_{j,i}\) to measure how much the textual context "focuses" on the visual token. Simultaneously, it computes the normalized attention entropy \(\mathbb{H} = (1/H) \sum_n [-\sum_j p_{j,n} \cdot \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)\)—tokens with high entropy (scattered, uncertain attention distribution) and low importance (small attention mass) receive larger regularization weights; conversely, tokens with low entropy (highly confident model) and high importance (narrowly focused) retain more of their original info. The \(q\) and \(k\) components used in calculating this weighting factor also employ stop-gradients to prevent the KL term from backwarding into query/key parameter updates. The weighted KL divergence applies to each token and scales the noise magnitude concurrently during sampling: \(z_i = v_i + \Delta_\phi(v_i) + g_i \cdot \sigma_i \cdot \epsilon\). Ablation studies show that removing this weighting mechanism leads to overall performance degradation, especially on MMMU (-1.0) and MuirBench (-2.5), suggesting that uniform regularization is too coarse, while adaptive weighting maintains a fine balance between information preservation and noise compression.

Loss & Training

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

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 shows the performance comparison of InternVL2.5-8B on comprehensive visual understanding benchmarks (most significant gains under IRA are bolded):

Benchmark Type	Benchmark	SFT	+IRA	Change
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	Flat

IRA's improvements are concentrated on reasoning-intensive benchmarks (MMMU/MME/OK-VQA), while showing a slight drop on MMStar (pure perception). This indicates that IRA mainly enhances the robustness of cross-modal compositional reasoning, rather than merely gaming surface-level alignment metrics. The MME improvement of 57 points (1981 \(\rightarrow\) 2038) is a substantial margin, illustrating that stochastic regularization significantly boosts semantic abstraction and cross-modal reasoning in open scenarios.

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

Benchmark	Evaluation Dimension	SFT	+IRA	Change
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 Blindspots	33.9	37.3	+3.4
MVBench	Video Understanding	51.5	52.0	+0.5

IRA shows the most significant gain (+3.4) on VLM-Blind (which evaluates whether the model truly "sees" low-level visual features instead of relying on language priors), proving that stochastic regularization indeed improves the model's reliance on real 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 Anchored on 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 the uncertainty weighting reduces MMMU by 1.0 and MuirBench by 2.5, indicating that uniform regularization is overly aggressive. Removing the anchored prior leads to larger drops (especially on MuirBench and VQAtext), demonstrating that anchoring on pre-trained representations indeed preserves more fine-grained visual details.

IRA performs best when applied to the mid-to-late layers of the LLM (60%-80% depth, corresponding to around layers 20-26 in a 32-layer architecture). These layers have accumulated enough abstract representation suitable for information compression, while still retaining strong visual grounding signals. Applying it too widely (0%-100%) or too shallowly (20%-40%) results in performance loss.

Key Findings

• Information bottleneck corrects attention distributions: IRA reduces the attention sink ratio of InternVL2.5 from 46.9% to 40.6% while concurrently increasing accuracy—proving that controlled compression of visual information directly improves attention focus.
• Representation curvature acts as an effective proxy for attention quality: IRA yields a straighter curvature trajectory, which strongly correlates with more stable attention patterns and fewer sinks. This finding provides a new evaluation metric for designing robust future architectures.
• IRA alleviates catastrophic forgetting: During training across 10 benchmarks, standard SFT shows an early peak followed by a decline (overfitting), whereas IRA exhibits a more stable and monotonic growth curve.
• Text-intensive tasks also benefit: On tasks requiring high-fidelity visual transmission such as TextVQA/ChartQA, IRA still yields gains (rather than drops), indicating that adaptive weighting successfully protects key fine-grained visual details.

Highlights & Insights

• Architecture-level rather than data-level solution: Unlike existing data-driven solutions (such as preference optimization and instruction tuning), IRA tackles VLM hallucination issues at the fundamental level of representation learning—representing a complementary and orthogonal addition to the "data-coverage" paradigm.
• Dual role of stochastic attention in regularization and representation: The paper experimentally demonstrates that stochastic attention is not merely a regularizer but an active contributor to representation learning—reshaping representation geometry and generating straighter curvature trajectories. This provides empirical support for incorporating stochastic attention as a first-class citizen in VLM architecture design.
• Joint design of attention entropy and token importance weighting: Leveraging both attention mass (importance) and attention entropy (uncertainty) to build the weighting factor offers strong complementarity—tokens with low attention mass but low entropy (areas where the model is certain but finds them unimportant) are properly protected, evading one-size-fits-all single-metric methods.
• Prior anchoring + stop-gradient collapse prevention trick: Data-dependent priors naturally tend to cause rapid KL collapse (as prior and posterior would degenerate if both look at \(v\)). The stop-gradient design is simple yet highly effective, making it highly transferable to other scenarios demanding data-dependent priors.

Limitations & Future Work

• The method has only been validated on models up to 8B parameters; its efficacy on larger scales (e.g., 13B/26B/73B) remains speculative.
• Extracting layer-wise visual value states and calculating attention metrics (entropy/importance) during training increases pipeline complexity and VRAM consumption; although randomness is removed in inference, preserving the mean still incurs minor computational overhead.
• The uncertainty weighting \(g_i = \mathbb{H} \cdot (1 - a_i)\) is an empirical design, and the theoretically optimal weighting formulation warrants further derivation.
• Currently constrained to the SFT stage, although the paper speculates that IRA could also be applicable during the pre-training stages of LLMs/VLMs without experimental validation.
• Applicability to other modalities (e.g., audio, video frame sequences) has not yet been explored.

Related Work & Insights

• vs. Data-driven hallucination mitigation (DPO/RLHF/VL-feedback): These schemes correct output distributions at the training signal level, whereas IRA compresses noise at the representation level—making them highly complementary and combinable.
• vs. Attention weight optimization (Reinforced Attention Learning, Gated Attention): Existing methods modify attention weights or outputs, while IRA intervenes earlier at the attention input (value states), thus controlling information quality from the source.
• vs. Information bottleneck methods (VIB-Probe, etc.): VIB-Probe uses the information bottleneck as an external probe to filter hallucinations, whereas IRA internalizes IB principles as part of the training objective to directly constrain layer representations.
• Connection to Ladder-VAE: The concept of sharing backbone parameters between the inference and generative processes in IRA shares similarities with the hierarchical variational structure of Ladder-VAEs, though a separate bottom-up/top-down path is not explicitly implemented in IRA.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Bringing the information bottleneck into the value states of attention modules as a form of stochastic regularization is a highly novel perspective—modifying representations themselves instead of attention weights.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ Validated across 3 prominent VLM architectures on over 20 benchmarks (encompassing hallucination, robustness, and OOD multi-image/video evaluation), with in-depth analysis on representation curvature, attention sinks, and training stability.
• Writing Quality: ⭐⭐⭐⭐ Clear motivation and solid formulation derivation, although some equation layouts are dense and the appendix structure could be optimized.
• Value: ⭐⭐⭐⭐⭐ Provides an architectural-level hallucination mitigation solution orthogonal to data-driven approaches. The stochastic regularization inside attention is elegantly designed with minimal extra parameters, offering strong engineering feasibility.

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?