Prioritizing Image-Related Tokens Enhances Vision-Language Pre-Training

TL;DR

The paper tackles the problem that standard vision-language pre-training using next-token prediction can overfit to non-visual caption tokens, increasing hallucination risk. It introduces PRIOR, a token-level weighting scheme that uses a text-only reference LLM to score image-related tokens and reweight the NTP loss accordingly, with a formal importance-sampling interpretation. Empirically, PRIOR yields substantial improvements over NTP across LVLMs with and without visual encoders (average relative gains of $19\%$ and $8\%$, respectively) and demonstrates stronger, more reliable scaling with compute. The approach is simple to integrate, data-efficient (offline token scores), and reduces hallucination risk while enhancing grounding, making it a practical, scalable advance for vision-language pre-training.

Abstract

In standard large vision-language models (LVLMs) pre-training, the model typically maximizes the joint probability of the caption conditioned on the image via next-token prediction (NTP); however, since only a small subset of caption tokens directly relates to the visual content, this naive NTP unintentionally fits the model to noise and increases the risk of hallucination. We present PRIOR, a simple vision-language pre-training approach that addresses this issue by prioritizing image-related tokens through differential weighting in the NTP loss, drawing from the importance sampling framework. PRIOR introduces a reference model-a text-only large language model (LLM) trained on the captions without image inputs, to weight each token based on its probability for LVLMs training. Intuitively, tokens that are directly related to the visual inputs are harder to predict without the image and thus receive lower probabilities from the text-only reference LLM. During training, we implement a token-specific re-weighting term based on the importance scores to adjust each token's loss. We implement PRIOR in two distinct settings: LVLMs with visual encoders and LVLMs without visual encoders. We observe 19% and 8% average relative improvement, respectively, on several vision-language benchmarks compared to NTP. In addition, PRIOR exhibits superior scaling properties, as demonstrated by significantly higher scaling coefficients, indicating greater potential for performance gains compared to NTP given increasing compute and data.

Figures (10)

• Figure 1: (Top) Synthetic examples to highlight the motivation of PRIOR. Only a few tokens in the captions (highlighted in blue, word-level for better visualization) are related to the associated images. PRIOR utilizes probability scores from a text-only LLM to recalibrate the original loss function at the token level, prioritizing image-related tokens that receive lower probability scores from the LLM. (Bottom Left) PRIOR formulation. Given an image $v$ paired with a caption '$t_1, t_2, ..., t_k$', PRIOR enhances vision-language pre-training by assigning a normalized weight to each token loss, which is computed based on the LLM probability $P_r(t_i | t_{<i})$. (Bottom Right) Performance of PRIOR.PRIOR demonstrates consistent performance improvement (average over several vision-language benchmarks) and better training stability compared to the widely used next-token prediction objective in vision-language pre-training. In addition, PRIOR shows superior scaling behaviors in both performance predictability and potential improvement with increased compute and data (§\ref{['sec:sl']}).
• Figure 2: Main experimental results of LVLMs with pre-trained visual encoders. We compare PRIOR with the NTP vision-language pre-training across various training steps on LVLMs with pre-trained visual encoders, and we annotate the highest performance for each method, respectively. PRIOR demonstrates both superior performance and greater stability throughout the entire training.
• Figure 3: The average performance comparison on LVLMs with unified architectures.PRIOR demonstrates better performance and stability across the entire training process.
• Figure 4: The comparison with additional baselines on H-LVLMs with 5,000 training steps.PRIOR consistently performs better and is simple to implement.
• Figure 5: The relative prediction error of NTP and PRIOR. We observe that the performance of LVLMs trained via PRIOR is more predictable at scale.