Locality Alignment Improves Vision-Language Models

TL;DR

The paper addresses persistent spatial reasoning errors in vision-language models caused by vision transformers pretrained with image-level supervision. It introduces locality alignment as a light-weight post-training stage, using MaskEmbed to reconstruct masked patch embeddings from a frozen teacher, thereby extracting local semantics without extra annotations. Empirically, locality-aligned backbones improve patch-level locality probes and bolster performance on spatial benchmarks across CLIP- and SigLIP-based backbones, with strong gains for high-resolution ViTs; integrating the MaskEmbed decoder as a vision-language adapter further enhances results. The approach requires minimal compute relative to pretraining and can be applied to off-the-shelf backbones, suggesting a practical pathway to better region-aware representations for open-ended VLMs.

Abstract

Vision language models (VLMs) have seen growing adoption in recent years, but many still struggle with basic spatial reasoning errors. We hypothesize that this is due to VLMs adopting pre-trained vision backbones, specifically vision transformers (ViTs) trained with image-level supervision and minimal inductive biases. Such models may fail to encode the class contents at each position in the image, and our goal is to resolve this with a vision backbone that effectively captures both local and global image semantics. Our main insight is that we do not require new supervision to learn this capability - pre-trained models contain significant knowledge of local semantics that we can extract and use for scalable self-supervision. We propose a new efficient post-training stage for ViTs called locality alignment and a novel fine-tuning procedure called MaskEmbed that uses a masked reconstruction loss to learn semantic contributions for each image patch. We first evaluate locality alignment with a vision-only benchmark, finding that it improves a model's performance at patch-level semantic segmentation, especially for strong backbones trained with image-caption pairs (e.g., CLIP and SigLIP). We then train a series of VLMs with and without locality alignment, and show that locality-aligned backbones improve performance across a range of benchmarks, particularly ones that involve spatial understanding (e.g., RefCOCO, OCID-Ref, TallyQA, VSR, AI2D). Overall, we demonstrate that we can efficiently learn local semantic extraction via a locality alignment stage, and that this procedure benefits VLM training recipes that use off-the-shelf vision backbones.

Figures (14)

• Figure 1: VLM training pipeline with locality alignment. Given a pre-trained vision backbone, we first perform a locality alignment stage using our MaskEmbed procedure (left), and then use the fine-tuned ViT to train a VLM (center). We find that doing so improves VLM performance in multiple benchmarks that involve spatial understanding (right).
• Figure 2: MaskEmbed training diagram. The encoder and decoder jointly reconstruct the pre-trained teacher's masked output, where patches are masked at the embedding layer for the encoder and at the input layer for the teacher.
• Figure 3: Qualitative examples from probing benchmark. We plot predictions for two images using CLIP ViT-L @ 336px before and after locality alignment. The original backbone fails to distinguish where certain objects occur in the image, but the aligned backbone corrects this.
• Figure 4: Probing benchmark results. We find that locality alignment with MaskEmbed improves IN1k classifiers across multiple model scales (left), and improves many models trained with language supervision (right). Interestingly, most models increase both their local and global probing accuracy.
• Figure 5: VLM benchmarking. We plot results across a suite of benchmarks and show controlled comparisons for CLIP (left) and SigLIP (right) with both the Llava-1.5 data mixture (top) and the extended data mixture (bottom). Overall, we achieve better performance in nearly all metrics with locality-aligned backbones. Between the two data mixtures, we find that the larger dataset does not have uniformly better performance and leads to different gains across text comprehension, chart understanding and localization tasks.