Probabilistic Language-Image Pre-Training

TL;DR

Probabilistic Language-Image Pre-training (ProLIP) addresses the limitation of deterministic vision-language embeddings by mapping images and texts to Gaussian distributions, capturing inherent multiplicity in cross-modal associations. It introduces an uncertainty token [UNC] and a novel inclusion loss to enforce distributional containment between modalities and between original and masked inputs, all trained with probabilistic objectives and without extra uncertainty modules. The approach delivers state-of-the-art zero-shot capabilities on ImageNet (e.g., 74.6% IN-ZSC with ViT-B/16) and gains further with Bayesian prompt re-weighting (75.8% in few-shot settings) while offering interpretable uncertainty insights through analyses on HierarCaps and HierarImgs. These results demonstrate practical gains for uncertainty-aware prompting, traversal, and downstream tasks, highlighting the value of probabilistic representations in large-scale vision-language systems.

Abstract

Vision-language models (VLMs) embed aligned image-text pairs into a joint space but often rely on deterministic embeddings, assuming a one-to-one correspondence between images and texts. This oversimplifies real-world relationships, which are inherently many-to-many, with multiple captions describing a single image and vice versa. We introduce Probabilistic Language-Image Pre-training (ProLIP), the first probabilistic VLM pre-trained on a billion-scale image-text dataset using only probabilistic objectives, achieving a strong zero-shot capability (e.g., 74.6% ImageNet zero-shot accuracy with ViT-B/16). ProLIP efficiently estimates uncertainty by an "uncertainty token" without extra parameters. We also introduce a novel inclusion loss that enforces distributional inclusion relationships between image-text pairs and between original and masked inputs. Experiments demonstrate that, by leveraging uncertainty estimates, ProLIP benefits downstream tasks and aligns with intuitive notions of uncertainty, e.g., shorter texts being more uncertain and more general inputs including specific ones. Utilizing text uncertainties, we further improve ImageNet accuracy from 74.6% to 75.8% (under a few-shot setting), supporting the practical advantages of our probabilistic approach. The code is available at https://github.com/naver-ai/prolip

Figures (19)

• Figure 1: Comparison of ProLIP and deterministic embedding spaces. We visualize images and captions from MS-COCO Caption chen2015cococaption using models trained on DataComp 1B gadre2024datacomp with 1.28B seen samples (See \ref{['appendix:subsec_fig1_visualization_details']} for more details of visualization method). ProLIP can capture multiplicity of image-text matching (e.g., the text embedding of "Train station" covers all three train images), while deterministic embeddings fail to capture the ambiguity. Furthermore, when we synthetically remove the background, ProLIP maps the new embedding near the original embedding but with a larger uncertainty value ($0.109 \rightarrow 0.117$), while the deterministic model maps the new embedding very far from the original one.
• Figure 2: Overview of ProLIP.[CLS] and [UNC] tokens are used for $\mu$ and $\log\sigma^2$, respectively.
• Figure 3: Visual understanding of the proposed inclusion loss. We plot probability density functions (pdfs) of three pairs of Gaussian distributions and their inclusion hypothesis, $\mathcal{H} (Z_1 \subset Z_2)$ (\ref{['eq:inc_test_hypothesis']}), log inclusion (\ref{['eq:inc_test']}) and KL divergence. The dashed line denotes the squared pdf, i.e., $p^2(x)$. $\mathcal{H} (Z_1 \subset Z_2)$ becomes (a) positive if $Z_1$ is included in $Z_2$ and (b) otherwise negative. (c) If $Z_1$ and $Z_2$ are the same level, then $\mathcal{H}$ will become zero. While log inclusion represents how $Z_1$ is included in $Z_2$, KL measures the "dissimilarity" between distributions (e.g., (c) has the largest KL, but the smallest "inc"). \ref{['appendix:fig:inclusion_loss']} shows more examples.
• Figure 4: Uncertain & certain samples. Visualization from the 3.5M filtered DataComp Small pool.
• Figure 5: $\sigma_v^2$ vs. $\sigma_t^2$. Generally, texts are more uncertain than images, as shown in \ref{['fig:uncertain_samples']}.