Logit Lens Supervision for Patch-Level Explanations in Vision-Language Models

Technical Summary: Logit Lens Supervision for Patch-Level Explanations in Vision-Language Models

Problem Statement
Modern autoregressive Vision-Language Models (VLMs), such as LLaVA and Qwen-VL, combine visual encoders with Large Language Models (LLMs) to generate fluent text. However, as visual tokens propagate through the LLM's attention layers, they undergo significant transformation and mixing with language tokens. The authors argue that this "uncontrolled modality mixing" causes visual embeddings to lose their direct correspondence to the specific image regions (patches) from which they originated. Consequently, the connection between a visual token and the semantic content of its source patch becomes diluted. This leads to two primary issues:

1. Weak Patch-Level Explainability: Standard Logit Lens visualizations (which project embeddings to vocabulary probabilities) often fail to produce sharp, localized confidence maps, particularly in newer VLMs like Qwen2.5-VL.
2. Hallucination and Grounding Degradation: Visual tokens may be effectively ignored or their information transferred entirely to language tokens, allowing models to hallucinate object descriptions even when the corresponding visual evidence is removed or absent.

Methodology: Logit Lens Loss (LLL)
To address the decoupling of visual tokens from their source image regions, the authors propose Logit Lens Loss (LLL), a lightweight auxiliary objective applied during fine-tuning.

• Core Concept: The method leverages the Logit Lens technique, which projects intermediate layer embeddings ($h_l$) through the LLM's unembedding matrix ($U_\theta$) to generate a probability distribution over the vocabulary.
• Mechanism: Unlike standard Next-Token Prediction (NTP) loss, which only supervises the generation of text tokens, LLL directly supervises visual tokens.
  • For visual tokens ($s$) corresponding to image patches within a ground-truth bounding box of a target object (set $\mathcal{P}'$), the loss maximizes the probability of the vocabulary token describing that object (e.g., <cat>).
  • For visual tokens outside the object region (set $\mathcal{P} \setminus \mathcal{P}'$), the loss minimizes the probability of that specific vocabulary token.
• Formulation: The total loss is defined as $L_{total} = L_{NTP} + \lambda L_{LLL}$, where $\lambda$ is empirically set to 0.5. The LLL term is a cross-entropy loss applied to the vocabulary distribution derived from visual token embeddings at the final LLM layer.
• Key Distinction: The approach requires no architectural modifications (no mask decoders, localization heads, or bounding box prediction modules). It relies solely on existing visual-token embeddings and the vocabulary projection head. The supervision signal is derived directly from bounding box annotations to identify which patches correspond to the target concept, but the model is not trained to output coordinates; rather, it is trained to align the internal representation of the patch with the semantic concept.

Key Contributions

1. Novel Visual Grounding Loss: Introduction of LLL, which explicitly constrains visual token embeddings to remain semantically aligned with the language concepts describing their image regions as they pass through LLM layers.
2. Intrinsic Localization: Unlike prior works that predict bounding boxes or masks as external outputs, LLL embeds localization awareness directly into the deep embeddings of visual tokens. This creates an intrinsic link between visual patches and textual concepts without modifying the model architecture.
3. Enhanced Explainability: The method significantly improves the quality of Logit Lens object confidence maps. Models fine-tuned with LLL produce sharper, more localized maps that accurately reflect which image patches the model associates with specific textual concepts.
4. Direct Gradient Signal: The authors demonstrate that LLL provides a direct gradient signal to visual tokens, whereas the signal from standard NTP must propagate through attention layers, often resulting in a weak signal for visual tokens.

Experimental Results
The authors evaluated LLL on LLaVA-v1.5-7B and Qwen2.5-VL-7B across several metrics:

• Referring Expression Segmentation (RES): Using Logit Lens maps to generate segmentation masks (via thresholding or prompting a frozen SAM), LLL-finetuned models showed significant improvements in cIoU scores over base and NTP-only models. With SAM post-processing, Qwen2.5-VL+LLL achieved performance comparable to State-of-the-Art (SOTA) specialized grounding models.
• Object Hallucination (POPE): On the POPE benchmark, LLL reduced object hallucinations, increasing "Yes/No" accuracy. Notably, models fine-tuned only on RefCOCO (localization data) and evaluated on POPE showed improved accuracy with LLL, whereas NTP-only fine-tuning degraded performance relative to the base model.
• Zero-Shot Pointing: Models fine-tuned with LLL on RefCOCO demonstrated strong transfer to the PixMo-Points dataset (unseen object categories), achieving lower median localization errors than several proprietary baselines and task-specific models trained with explicit coordinate supervision.
• Visual Question Answering (VQA): LLL preserved or slightly improved performance on VQAv2 and GQA benchmarks, indicating that enhancing visual grounding does not compromise the model's language generation capabilities.
• Visualization: Qualitative results (Fig. 2) showed that LLL transforms sparse, confusing Logit Lens signals into sharp, object-aligned confidence maps.

Significance and Claims
The paper claims that LLL offers a simple, effective mechanism to preserve the "locality" of visual information within autoregressive VLMs without the complexity of adding new architectural components. By grounding visual embeddings directly in the vocabulary space, the method makes foundation VLMs more interpretable (via faithful Logit Lens maps) and more robust against hallucinations. The authors emphasize that LLL does not aim to teach new concepts but rather to prevent the degradation of existing localized image understanding during the language modeling process. The results suggest that better internal grounding is achievable through auxiliary supervision that aligns visual tokens with their semantic descriptions, leading to models that are both more effective and more explainable.