Probabilistic Language-Image Pre-Training

[W2/Q1] UNC token efficiency

In practice, Transformer with a relatively short context length (e.g., a few hundred) and small network capacity (e.g., ViT-B) does not show a quadratic complexity on the input length. Following the reviewer’s comment, we compare three different design choices (deterministic, UNC token, and multi-head uncertainty estimate module) on five different architectures (ViT-B/32, ViT-B/16, ViT-B/16 with 768 width, ViT-L/16, and ViT-L/14). Here, we only compare image encoders because image encoders always take the same image token length, which makes it easier to compare the impact of different design choices. We report (1) the number of input image tokens for each architecture (if we choose the “[UNC] token architecture”, it will be added by one), (2) the base parameters when we use deterministic, (3) inference time for 50k ImageNet validation images (lower is faster), and (4) additional parameters compared to the deterministic baseline.

We have three findings from the table.

• Even if we increase the input token length, the actual inference speed is not quadratically slow -- ViT-B/32 and ViT-B/16 have the same Transformer capability but only input lengths are different (49 vs. 196), but their inference times are 75.18s and 76.68s, which is an almost neglectable change. If we choose a larger model (e.g., ViT-L), the difference becomes larger, e.g., 137.77s vs. 176.95s, but it is still not a quadratic order.
• [UNC] token adds almost neglectable parameters (0.3M for B and 0.6M for L) and inference time compared to the deterministic one.
• Multi-head architecture needs a large number of additional parameters (e.g., 20M for B and 40M for L) and shows slower inference time, especially for a larger network (e.g., 176.95s vs. 191.68s for ViT-L/14).

Furthermore, in practice, multi-head architecture requires more memory than [UNC] token, which makes it difficult to use a large batch size and scale up to a larger backbone. On the other hand, [UNC] token only needs almost neglectable additional parameters, inference speed and memory size, which makes it easier to scale up.

We add the related discussion in Appendix C.3.

[W3] Novelty compared to PCME++

Continuing from the response for W2, ProLIP has a huge contribution in terms of scaling up probabilistic representations. Namely, we cannot scale up PCME++ to CLIP level training dataset. First, as shown in the previous response, ProLIP is significantly more efficient than PCME++’s architecture (multi-head architecture) in terms of inference speed, parameter size, and memory. It makes ProLIP easier to be scaled up to a larger backbone while PCME++ suffers from the scalability issue. Namely, our architectural contribution is significant compared to PCME++ in terms of scaling up.

Second, as shown in Table C.3 of the paper, a model will not be converged if we solely use PCME++ loss. We need additional deterministic loss, such as CLIP loss or SigLIP loss for a stable convergence of PCME++ loss. Even if we use additional deterministic loss, we can find that it performs worse than a model trained solely with ProLIP’s PPCL loss. It is because as we discussed in Sec A.2 of the paper, PCME++ loss uses binary cross entropy (BCE) loss, and PPCL loss is based on log-sigmoid loss. To compute BCE loss, we have to compute the squared L2 distance between two $\mu$s. If the difference between $\mu$s is extremely small, computing $\sum_{i=1}^D ((\mu_v^i)^2 - (\mu_t^i)^2)$ will suffer from a numerical error. On the other hand, PPCL loss directly uses the cosine similarity between $\mu$s, which is more numerically stable than the squared summation. Moreover, our inclusion loss is sufficiently novel which enforces a desired behavior; which was not able to PCME++. It supports that our loss function design is crucial to scaling up probabilistic representations.

Finally, our main contribution is to train the first PrVLM comparable to the deterministic baselines (e.g., CLIP, SigLIP). Again, previous PrVLMs were not able to pre-train (i.e., they fine-tune the pre-trained CLIP models). PCME++ shows a small-scale pre-training result, but it still needs deterministic loss. Furthermore, as we mentioned before, even if we use additional CLIP loss, PPCL-only ProLIP performs better than PCME++ (Table C.3). It shows that ProLIP has a significant novel contribution to PrVLM.

[Q2] Image uncertainty

An ideal PrVLM should capture three potential input uncertainties: (1) uncertainty from the text modality, (2) uncertainty from the image modality, and (3) uncertainty from the text-image cross-modality. Uni-modal uncertainty is straightforward; if an input has more detailed information (e.g., describing more detailed information in text, or capturing a very detailed and complex scene by photography), then it will have smaller uncertainty, otherwise, it will have larger uncertainty (e.g., providing very high-level caption, such as “person”, or only a part object with white background is taken by picture).

The cross-modal uncertainty should capture “how many possible instances can be matched to this input?”. We can think this in two different viewpoints: text-to-image and image-to-text. The text-to-image relationship is straightforward. If we have a caption “photo”, then it will be matched to all photographs, and if we have a caption with a very detailed description (e.g., the full description of the hotel room), then it will be only matched to specific images. Image-to-text relationships are often determined by the dataset. For example, if we have a caption dataset exactly describing which objects are in the image, then we can think that an image with fewer objects will have larger uncertainty. However, if we consider a caption dataset where an image has multiple captions each caption focuses on completely different objects, then an image with more objects will have larger uncertainty. In other words, unlike text uncertainty originating from a text-to-image relationship, image uncertainty originating from an image-to-text relationship is highly affected by the property of the training dataset.

In practice, because our datasets have a mixed property and their captions are somewhat noisy, our image uncertainty will have a mixed property, namely, unlike text uncertainty, there could be no strong relationship between the absolute uncertainty value and the number of objects (or complexity of images). However, empirically, we have more captions that exactly describe the scene (because of the dataset filtering strategy), therefore, a more simple image tends to have larger uncertainty and a more complex image tends to have smaller uncertainty.

Furthermore, since there is no “uni-modal” uncertainty supervision, a PrVLM trained only with image-text relationships will have no guarantee to represent proper image uni-modal uncertainty. To enforce a proper uni-modal image uncertainty, we first propose the uni-modal uncertainty supervision by proposing the inclusion loss, namely, the original image embedding should be included by an image embedding from the masked image.

We also clarify that the dataset used for the analysis is a very large-scale dataset with 3.5M image-text pairs (which follows the training distribution of DataComp1B). If we choose a more restricted dataset with selective images or captions, we may have a different observation, but we don’t think it truly represents the embedding space learned by ProLIP.

Our revised paper includes the related discussion in Section A.1.

[W1] Mixture of Gaussian

Thanks for the interesting question. Mixture of Gaussian (MoG) is an intuitive approach to tackle the multi-mode problem. However, if we have a sufficiently large dimensionality, MoG is not a mandatory option. Intuitively, we can break D-dimensional unimodal Gaussian probabilistic embedding space into multiple subspaces with smaller dimensions (e.g., D/2-dimensional space) and consider each space as a different Gaussian embedding space. More formally and rigorously, in our framework (using CSD as the probabilistic distance), using k-MoG (MoG with k number of Gaussians) is equivalent to using k^2 D-dimensional embedding space.

Consider a probabilistic embedding with 2-MoG, namely $Z \sim \frac{1}{2}\mathcal N(\mu_1, \Sigma_1)$ with probability 0.5 and $Z \sim \frac{1}{2}\mathcal N(\mu_2, \Sigma_2)$ with probability 0.5. We can compute the expected CSD between two $Z$s (where $Z_1$ is parameterized by $\mu^1_1, \mu^1_2, \Sigma^1_1, \Sigma^1_2$ and $Z_2$ is parameterized by $\mu^2_1, \mu^2_2, \Sigma^2_1, \Sigma^2_2$) by computing

$\frac{1}{4} \left[ d(\mu^1_1, \mu^2_1, \Sigma^1_1, \Sigma^2_1) + d(\mu^1_1, \mu^2_2, \Sigma^1_1, \Sigma^2_2) + d(\mu^1_2, \mu^2_1, \Sigma^1_2, \Sigma^2_1) + d(\mu^1_2, \mu^2_2, \Sigma^1_2, \Sigma^2_2) \right],$

where $d(\cdot)$ is CSD, i.e, $d(\mu_1, \mu_2, \Sigma_1, \Sigma_2) = \| \mu_1 - \mu_2 \|_2^2 + tr(\Sigma_1 + \Sigma_2)$. To simplicity, we omit $\frac{1}{4}$ for the remaining derivation. Now, consider two virtual unimodal Gaussian embeddings

$W_1 \sim \mathcal N(\mu^1_1 \oplus \mu^1_1 \oplus \mu^1_2 \oplus \mu^1_2, \Sigma^1_1 \oplus \Sigma^1_1 \oplus \Sigma^1_2 \oplus \Sigma^1_2)$ $W_2 \sim \mathcal N(\mu^2_1 \oplus \mu^2_2 \oplus \mu^2_1 \oplus \mu^2_2, \Sigma^2_1 \oplus \Sigma^2_2 \oplus \Sigma^2_1 \oplus \Sigma^2_2),$

where $\oplus$ denotes the “concatenation” operation, i.e, $W_1$ and $W_2$ have four times larger dimensionality than $Z_1$ and $Z_2$.

Interestingly, we can easily show that the above equation equals CSD between $W_1$ and $W_2$. Note that this derivation is invariant to the diagonal covariance, but also holds for the full covariance. In other words, using MoG is mathematically equivalent to using a larger dimensionality (as much as the square of the number of modes); therefore, if we have a sufficiently large dimensionality that can capture the ambiguity of the dataset, MoG is not a mandatory option.

Our revised paper includes the related discussion in Section D.

[W2] Contrastive loss can solve many-to-many?

From the response for [Q8], we clarify that our purpose is to capture the inherent input uncertainty that cannot be captured by a deterministic embedding space. Since this is the fundamental limitation of a deterministic space, this problem is agnostic to the choice of objective function. We partially agree that contrastive learning less penalizes plausible matching than other loss functions, such as triplet loss with hard negative mining, but it still shares the limitation of deterministic embeddings.

We cite Appendix A.1 in Section 2.1 and 2.2. If the reviewer thinks the current version is still insufficient, please let us know.

[W3] Other downstream tasks

We do not think our applications are not representative as the reviewer's comment, but we also agree that it would be better to show more applications of the learned uncertainty. Here, we show 3 applications.

Dataset filtering

Below, we show the DataComp CLIP filtering small track (filtering 12.8B noisy web-crawled image-text pairs) by using our method and baselines provided by DataComp [A]

• [A] Gadre, Samir Yitzhak, et al. "Datacomp: In search of the next generation of multimodal datasets." NeurIPS 2024.

Here, we can observe that by using ProLIP uncertainty-aware features, we can achieve a better filtering performance compared to the other baselines. However, we would like to note that ProLIP is not specifically designed for dataset filtering; proposing a new dataset filtering method using ProLIP will be an interesting future work, but not the scope of the current paper.

[W2] Modest improvement

First, we want to emphasize that +0.9% ImageNet zero-shot accuracy is not a marginal improvement. As shown in [A], it is almost close to a gap when using a larger backbone (e.g., ViT-H/14 (75.6%) $\Rightarrow$ ViT-g/14 (76.7%) when the number of seen samples is 13B). However, we also want to clarify that our purpose is not to achieve a significantly stronger ImageNet zero-shot performance, but to achieve a better understanding of the input uncertainty (i.e., aleatoric uncertainty).

Even if we consider the improvement to be marginal, we argue that the modest improvement against CLIP and the scalability challenge of PrVLM are not at the same level. As already clarified in Section A.3, our ablation study (Table C.3) empirically shows that the PCME++ loss fails to be converged when it is used in a stand-alone way without any deterministic loss (e.g., CLIP loss). On the other hand, PPCL loss converges well even without any additional loss function. Our first contribution is to enable a fully end-to-end scalable pre-training of PrVLM by proposing (1) new loss functions -- PPCL and inclusion loss, (2) new [UNC] token architecture, which is extremely efficient. Table C.3 and C.6 support that our new loss function and architecture outperform the previous PrVLM design choice (PCME++ loss + CLIP loss & multi-head uncertainty estimation module). Our main contribution is to train the first PrVLM (which has the advantage of correctly representing aleatoric uncertainty -- please see the newly added Section A.1 for more detailed discussion) comparable to the deterministic baselines (e.g., CLIP, SigLIP). Again, previous PrVLMs were not able to pre-train (i.e., they fine-tune the pre-trained CLIP models). PCME++ shows a small-scale pre-training result, but it still needs deterministic loss. Furthermore, as we mentioned before, even if we use additional CLIP loss, PPCL-only ProLIP performs better than PCME++ (Table C.3). It shows that ProLIP has a significant novel contribution to PrVLM.

Furthermore, the proposed loss design and training process are not specifically complex than CLIP. In terms of the training process, we use the same training process to openclip, except for the “partial” inclusion loss. Similarly, if we only consider PPCL, the loss design is not specifically complex than the original contrastive loss (CLIP) or log sigmoid loss (SigLIP). As shown in Table C.3, solely using PPCL is sufficient for achieving performance; we use additional loss functions (inclusion loss, VIB loss) to ensure the learned embedding space has our desired properties (e.g., showing inclusive relationship between text and image / partial and whole instance, and preventing variance collapse). Our purpose is to train a fully end-to-end scalable pre-training of PrVLM, not outperform CLIP with a large gap.

We also would like to clarify that, in Table C.1., ProLIP ViT-L is not a fair comparison with other backbones. It is a fine-tuned backbone from SigLIP ViT-L (ImageNet zero-shot 80%), not fully from scratch training. Even more, we would like to emphasize that its improvement is never modest against ViT-B/16 CLIP; The fine-tuned ProLIP ViT-L achieves 79.4% ImageNet zero-shot accuracy, while CLIP ViT-B/16 achieves 67.2% in our setting (1.28B seen samples), and 73.5% by openclip (13B seen samples). Even openclip ViT-L and ViT-H achieved 75.3% and 78.0%, respectively, while our ViT-L/14 shows 79.4%, outperforming expensive ViT-H training (Remark that openclip ViT-H took 279 hours with 824 A100s). We argue that the difference is not as modest as the reviewer’s comment. Again, we clarify that searching hyperparameters for fine-tuning is out of our main contribution, hence, we did not exhaustively search the fine-tuning hyperparameters and there will be a large room for performance improvements by hyperparameter search.

Understanding image dataset

First, we emphasize that ProLIP’s image uncertainty is not the same as the “image uncertainty” of classification tasks. In classification tasks, an image has a high uncertainty if it can be matched to “multiple classes”, while the ProLIP case assigns a high uncertainty if an image can be described in “multiple different and various captions”. However, ProLIP has a different mechanism with classification uncertainty.

An ideal PrVLM should capture three potential input uncertainties: (1) uncertainty from the text modality, (2) uncertainty from the image modality, and (3) uncertainty from the text-image cross-modality. Uni-modal uncertainty is straightforward; if an input has more detailed information (e.g., describing more detailed information in text, or capturing a very detailed and complex scene by photography), then it will have smaller uncertainty, otherwise, it will have larger uncertainty (e.g., providing very high-level caption, such as “person”, or only a part object with white background is taken by picture).

The cross-modal uncertainty should capture “how many possible instances can be matched to this input?”. We can think this in two different viewpoints: text-to-image and image-to-text. The text-to-image relationship is straightforward. If we have a caption “photo”, then it will be matched to all photographs, and if we have a caption with a very detailed description (e.g., the full description of the hotel room), then it will be only matched to specific images. Image-to-text relationships are often determined by the dataset. For example, if we have a caption dataset exactly describing which objects are in the image, then we can think that an image with fewer objects will have larger uncertainty. However, if we consider a caption dataset where an image has multiple captions each caption focuses on completely different objects, then an image with more objects will have larger uncertainty. In other words, unlike text uncertainty originating from a text-to-image relationship, image uncertainty originating from an image-to-text relationship is highly affected by the property of the training dataset.

In practice, because our datasets have a mixed property and their captions are somewhat noisy, our image uncertainty will have a mixed property, namely, unlike text uncertainty, there could be no strong relationship between the absolute uncertainty value and the number of objects (or complexity of images). However, empirically, we have more captions that exactly describe the scene (because of the dataset filtering strategy), therefore, a more simple image tends to have larger uncertainty and a more complex image tends to have smaller uncertainty.

For example, assume we have an image with a white background and a very clear overall object shape. In terms of classification, it will have low uncertainty because there is no confounder to the classification. However, ProLIP will assign a high uncertainty for this case.

From this, we can divide the image classification tasks into two different scenarios: (1) When all images have homogeneous backgrounds and only the quality of the image determines the classification performance, (2) When images are natural images and the task is inherently multi-object classification, but the labels are single-labeled (e.g., ImageNet [B, C])

• [B] Beyer, Lucas, et al. "Are we done with imagenet?." arXiv preprint arXiv:2006.07159 (2020).
• [C] Yun, Sangdoo, et al. "Re-labeling imagenet: from single to multi-labels, from global to localized labels." CVPR 2021.

As the first example, we choose MNIST. Here, we observe that the learned image uncertainty and the MNIST accuracy show a strong negative correlation, (-0.98), namely, if an image is more uncertain then ProLIP tends to estimate a wrong label. It aligns with our intuition.

As the second example, we choose ImageNet-1k, which shows a strong positive correlation (+0.98), i.e., a certain image tends to be wrongly classified by ProLIP. This could be counterintuitive in terms of “classification task”, but actually it is a correctly estimated value. For example, ImageNet contains various image distributions. Some images are thumbnail images with white backgrounds (high uncertainty) and some images are in-the-wild images with complex backgrounds and multiple objects (low uncertainty). In this case, a certain image (more complex image) will be a more “difficult” image.

Overall, ProLIP’s image uncertainty tendency can be used for understanding the properties of the given dataset. Converting ProLIP’s image uncertainty to image classification uncertainty would be an interesting topic, but we remain this for future work.