Preserving Localized Patch Semantics in VLMs

TL;DR

This work tackles the drift of patch-level visual semantics in autoregressive VLMs by introducing Logit Lens Loss (LLL), an auxiliary objective applied directly to visual tokens to align them with vocabulary concepts describing their image regions. Formulated with bounding-box supervision, LLL defines a patch-level grounding loss and combines it with the standard next-token prediction loss to form $L_{total}=L_{NTP}+\lambda L_{LLL}$, using $\lambda=0.5$ and applying at the final layer. By projecting both visual and textual tokens into a shared vocabulary space via $f=\mathrm{softmax}\circ U_{\theta}$, LLL yields direct gradient signals to visual patches, preserving localized content while avoiding architectural changes. Empirically, LLL produces sharper Logit Lens object confidence maps, improves Referring Expression Segmentation and POPE performance, enhances attention to object regions during answer generation, and enables strong zero-shot localization on unseen categories, contributing to better explainability and grounding in VLMs.

Abstract

Logit Lens has been proposed for visualizing tokens that contribute most to LLM answers. Recently, Logit Lens was also shown to be applicable in autoregressive Vision-Language Models (VLMs), where it illustrates the conceptual content of image tokens in the form of heatmaps, e.g., which image tokens are likely to depict the concept of cat in a given image. However, the visual content of image tokens often gets diffused to language tokens, and consequently, the locality of visual information gets mostly destroyed, which renders Logit Lens visualization unusable for explainability. To address this issue, we introduce a complementary loss to next-token prediction (NTP) to prevent the visual tokens from losing the visual representation inherited from corresponding image patches. The proposed Logit Lens Loss (LLL) is designed to make visual token embeddings more semantically aligned with the textual concepts that describe their image regions (e.g., patches containing a cat with the word "cat"), without requiring any architectural modification or large-scale training. This way, LLL constrains the mixing of image and text tokens in the self-attention layers in order to prevent image tokens from losing their localized visual information. As our experiments show, LLL not only makes Logit Lens practically relevant by producing meaningful object confidence maps in images, but also improves performance on vision-centric tasks like segmentation without attaching any special heads.

Figures (5)

• Figure 1: VLM components with the added Logit Lens Loss (LLL) shown in cyan. While NTP loss maximizes the probability of generating ground truth text tokens, LLL increases the probability of visual tokens $s$ to predict text tokens describing their content, e.g., if visual token $s$ represents a patch containing part of a cat, the probability of predicting token $<cat>$ is increased. This way LLL maximizes the alignment of deep embeddings of image patches containing the cat with text token $<cat>$.
• Figure 2: Logit Lens object confidence maps produced by Qwen2.5-VL-7B. First column: input image with the queried object text; second: base model; third: NTP-finetuned; fourth: NTP+LLL-finetuned. Compared to the base and NTP models, LLL produces sharply localized, object-aligned Logit Lens maps by preserving patch-level semantics in visual tokens despite cross-modal mixing in the LLM.
• Figure 3: Bar plots illustrating the ratio of the object confidence scores of objects present in images within their bounding boxes to the confidence measure of the same objects averaged across all visual tokens in each image.
• Figure 4: Bar plots illustrating the ratio of the average attention of target tokens within the bounding box of the target object w.r.t. the answer token to the average attention across all visual tokens.
• Figure 5: Attention maps of the last output token with respect to visual tokens under different settings. Queried object is “Cake”.