Training Universal Text Encoders with Pair Relevance Classification Loss

We thank the reviewer for their constructive comments and positive feedback regarding the technical soundness of our work and clarity of the exposition. Hereby we address their concerns:

1. Clarifying conceptual novelty and real-world implications of the work

• Our work clarifies the link between the logit bias and the label-imbalanced classification problem that is brought by making use of the in-batch negatives; something which is not apparent from the SigLIP paper. This allowed us to derive a formula for setting the correct value for the logit bias per total batch size used, which in turn allows us to train effectively strong generalist text encoding models at a considerable fraction of the batch size typically used for softmax-based contrastive learning approaches. Upon requests of reviewers EFpF and rQBS, please find in Fig.3a experimental evidence for this, also please refer to point 4.1.1 of this rebuttal regarding theoretical intuition for this effectiveness.

• We are demonstrating increased robustness to weakly-supervised dataset noise, in new evidence towards reviewer rQBS. Please find in Fig.3b experimental evidence for it, also please refer to point 4.1.2 of this rebuttal regarding theoretical intuition for the improved robustness.

• The novelty concern at the level discussed by the reviewer can also be applied to the SigLIP paper itself. At the time of SigLIP paper’s review process, binary cross-entropy losses for representation learning have already been explored, as we point out in our extended related work section in the Appendix. To our knowledge, binary cross-entropy loss in representation learning was first explored by [Hjelm et al., 2019]. Their work introduces Deep InfoMax (DIM), which in its simplest form consists of a binary cross-entropy loss and the logits are computed from pairs of views: pairs derived from the same image are labeled as positive, while pairs from different images are labeled as negatives. Following similar reasoning to the reviewer, SigLIP can be thought of as an application of the DIM loss to aligning an image view to a caption “view”. However, this reasoning under-appreciates the successful application of the binary cross-entropy loss to the corresponding domain. Similarly for our work, as pointed out by reviewer EFpF: “This is the first work to demonstrate the pointwise sigmoid loss outperforms the standard contrastive loss for training text embeddings.”.

• Finally, we would like to underline the most distinctive feature of our work: It lays important groundwork for future graded relevance works. Language tasks are more complex and nuanced than just binary labels. For example, although it is common to have only two classes in sentiment classification (positive and negative), in the real-world, the range of sentiments are much more than that. While today there are limited resources available containing graded relevance data, such as the TREC datasets and STS, in the future BCE losses can provide a definite advantage in dense text encoding when one considers the possibility of annotating the data pair relevance with large language generative models. Our work lays the groundwork for such advances by allowing us to encode the uncertainty about a pair of text into the continuum of (0,1). Our work is the first one to successfully exploit the finer grained signal provided by graded relevance in the context of text encoding, as evidence by our experimental results on TREC-DL, and we are looking forward to applying our non-trivial, but simple, training recipe to future work leveraging larger graded relevance datasets.

2. Extending experimental evidence beyond LLM2Vec setup

We have updated the manuscript to include models finetuned from Qwen2-7B-Instruct and Mistral-7B-Instruct. Deviating from the LLM2Vec recipe, we just enabled the bidirectional connection and created a representation by average pooling the token representations which correspond to the embed texts $x$ and $y$, excluding task description or instruction tokens. Both for BCE and softmax-based we are searching for the best batch size and learning rate. We find that both for Mistral-7B and Qwen2-7B training BCE surpasses softmax-based losses on average across the 56-task English MTEB. Please find more details on the updated paper (Table 2) and training details in the Appendix.

Dear Reviewer,

Thank you for your constructive feedback. We appreciate your insights and would like to address your concerns systematically:

Intuition for why BCE is a better loss

We will support this claim by elucidating with theoretical insight two sets of empirical observations: Increased robustness to noise and increased robustness to batch size

BCE performs better than softmax loss in smaller batch sizes

Setting the correct logit bias

We further solidify this claim in the revised Fig. 3b by providing an additional point for batch sizes 2048 training Meta-LLaMA-3-8B-Instruct. Our new empirical evidence reinforces our original claim, that binary cross-entropy losses improve robustness with respect to batch size because they attain competitive, and at times improved, performance over the softmax-based baseline with smaller batch sizes (256, 512), where performance is almost constant. Crucially in our experiments we are using a logit bias depending on the total batch size used, according to the proposed formula $- \log( 2 B - 1)$ (see Section 3.1, Eq. 3). In Section 3, we discuss the role of logit bias for enabling robust training in label-imbalanced classification scenaria, connecting it to the literature of logit adjustment [Menon et al., 2021]. At the same time, the effect that various values for the logit bias can have in fixed batch-size training are evident in Fig. 2.b.

Theoretical insight on why is better in smaller batch sizes

We can get intuition about the batch size resilience of BCE losses by analyzing loss gradients. In softmax-based contrastive learning, the denominator, in which the negative pairs appear, can be expressed as a log-sum-exp. This means that each of the $(x_k, y_n)$ pairs' gradient contribution is weighted by each pair’s probability to have $y_n$ classified from $x_n$.

Combining this with large logit scales $\in (10, 100)$, that are used in order to achieve good downstream performance, we can see that the negative pair with the largest logit $s(x_k, y_{m^*})$ contributes almost all of the gradient, while the rest of the negative pairs with smaller logits $s(x_k, y_m) \leq s(x_k, y_{m^*})$ have significantly smaller gradient contributions. In contrast, BCE loss weights all negative pairs independently from one another.

Combining this with large logit scales $\in (10, 100)$ once more, we can see that only the negative pairs that are perceived as positive by the model are going to be used and they will have approximately equal weight ($\approx 1$) among themselves. We argue that this makes the BCE formulation more sample efficient when it comes to using in-batch negatives, as it utilizes simultaneously all erroneous predictions about negative pairs instead of the most erroneous one at each training step.

Robustness to weakly-supervised dataset noise

We would like to draw your attention to this point raised by reviewer EFpF, and here we argue about it with experimental evidence and theoretical insight.

BCE improves robustness to noise compared to softmax loss

We have performed an ablation study with respect to levels of artificial noise introduced to the open-source E5 training dataset. In particular, we specify a probability of flipping the label between a designated positive pair $(x, y^+)$ and a hard negative pair $(x, y^-)$, and we sample the flip variable independently for every triplet $(x, y^+, y^-)$ in the dataset prior to training. Then, we train Meta-LLaMA-3-8B-Instruct models with the BCE and softmax-based losses and we report the average 56-task MTEB score of the best models found during training. In Fig. 3.b, we observe that BCE outperforms the softmax variant for all noise levels $p_{\text{noise}} \in \{0.05, 0.1, 0.2, 0.5\}$.

Intuition on why BCE is more robust to noise

Suppose the binary label case with hard negatives. At each training step the model encounters a batch of size $B$, the BCE-loss then is the average of losses from $2B^2$ binary classification predictions, while softmax-based loss is the average of losses from $B$ multi-class classification predictions. Assuming that the in-batch negatives are indeed true negatives: If one pair of texts in the batch is mislabeled, then for the softmax-loss this means that the $\frac{1}{B}$ predictions is optimized towards a false target, whereas for the BCE-loss $\frac{1}{B^2}$ predictions is optimized towards a false target. We hypothesize that BCE achieves better robustness to weakly-supervised dataset noise, because the effect of noise is more diluted in the gradients of the BCE objective.

• In Table 2, the results for Sheared-Llama-1.3B indicate that the BCE loss performs worse compared to the softmax-based loss. Other experiments also show only marginal difference in average scores (e.g., less than 0.1 for Qwen2-7B).
• Regarding Table 2, the literature model does not include the leading models, and this table partially overlaps with the results presented in Table 7.
• Please include the AIR-Benchmark as a zero-shot evaluation metric to show the effectiveness of the proposed method.
• Final comment: Thank you for the detailed responses. While the idea of enhanced batch and noise robustness through the loss function is interesting, the experimental results and presentation do not provide sufficient support for this claim. I will maintain my score.

We appreciate the reviewer’s thoughtful feedback, which has provided valuable insights for improving our work. Below, we address the points raised, highlighting areas of agreement, clarification, and proposed changes.

Response to Strengths

• We are glad the reviewer acknowledged the generalization of SigLIP and the utility of graded labels in our proposed approach.
• The experimental results on MTEB tasks and graded relevance datasets (Section 5) reinforce these observations.

Response to Weaknesses

1. In the case of binary labels, our work clarifies how to set significant hyperparameters making the difference

While the loss reduces to SigLIP for binary labels, the results in Table 2 demonstrate its superior generalization compared to softmax-based contrastive loss. We are happy to find that the reviewer has already considered experimenting with binary cross-entropy losses for training dense text encoders. This highlights the importance of setting hyperparameters correctly in the training recipe, like the logit bias, which we analyze through the prism of label-imbalance classification. We invite the reviewer to point (3) of our rebuttal for more details.

Our work is important groundwork for future graded relevance works:

While we agree with the reviewer that the vast majority of existing text training data are of binary labels, we would like to mention that this happens more out of convenience of data curation. In reality, language tasks are more complex and nuanced than just binary labels. For example, although it is common to have only two classes in sentiment classification (positive and negative), in the real-world, the range of sentiments are much more than that. While today there are limited resources available containing graded relevance data, such as the TREC datasets and STS, in the future BCE losses can provide a definite advantage in dense text encoding when one considers the possibility of annotating the data pair relevance with large language generative models. Our work lays the groundwork for such advances by allowing us to encode the uncertainty about a pair of text into the continuum of (0,1).

2. Missing comparison to regression-based pointwise ranking:

We acknowledge this limitation and we have included a discussion of [Zou et al., 2021] in the revised paper. However, we note that our focus is on general-purpose text embeddings, distinct from pointwise ranking models for information retrieval. While the referenced paper mobilizes training objectives on graded relevance labels, it does not prescribe a way to deal with the imbalances of the training relevance data. It is not obvious how to modify their regression-based loss so that it accounts for possible imbalances, nor do they report the marginal distribution of the human annotated relevance data. Our work fills that gap by allowing pointwise training on imbalanced relevance data.

Thank you for your review and feedback. We are grateful for your recognition of the significance of our work, accommodating graded relevance scores in contrastive learning, and the clarity of our presentation. Below, we address your comments in detail:

1. About IW softmax

Softmax-based graded relevance handling (e.g., IW softmax) is an ablation attempt from us. We would like to stress that IW softmax is not a studied method from the literature. Instead, we introduce it as a method to trivially adapt softmax-based losses for graded relevance scores. Our results demonstrate that this adaptation is difficult, thus justifying the need of our proposed loss. In L507-508: “By this baseline, we want to demonstrate an example that attempts to incorporate graded relevance scores into the standard contrastive loss.”. We have changed the main text to further clarify this.

2. About the term Universal Text Encoders

The term “Universal” characterizes “Text Encoders” in the title, referring to the fact that we are training text encoders to perform well on a variety of downstream tasks and categories. Our text encoders are benchmarked against a broad spectrum of text related downstream tasks using MTEB benchmark: classification, clustering, pair classification, retrieval, reranking, summarization and semantic textual similarity (STS) tasks. Our text encoders are universal in the sense that we would like them to give a good vector representation for all embedding tasks. This is in contrast to focusing only on just one task, like dense encoders for retrieval for instance.

3. Addressing the concern about improvements

While MTEB results serve as testimony that our method is on par and at times brings minor improvements over softmax-based loss, the focus of our work remains on merits that can be achieved via the BCE loss, that cannot be achieved by softmax-based ones. We are here underlining benefits to using BCE loss over softmax which are not minor:

3.1 Improved performance with smaller batch sizes:

Given the same computation budget, we can now effectively scale the training of dense text encoder models to larger context lengths, and if needed model sizes. Please refer to Fig. 3a of our paper. We have further solidified our original argument with a theoretical insight in the Section “Extended Study” of the Appendix.

3.2 Improved robustness to noise

In response to reviewer rQBS, we have further conducted experiments with artificial noise, flipping with some probability, independently for each query in the training set, the positive pair with the hard negative one. In Fig. 3b of the updated manuscript, we show that BCE loss outperforms softmax-based losses for all noise probabilities in {0.05, 0.1, 0.2, 0.5} in terms of average MTEB score. We further discuss the intuition behind this in the main text.

3.3 Major advantage in using graded relevance data

Our experimental results on TREC-DL show that only BCE loss is able to take advantage of the finer grained signal of graded relevance data. Note that in Table 3, we have run our best performing BCE recipe with 5 random seeds, demonstrating small variance and that the BCE loss consistently brings superior results to other training recipes.

4. On the role of logit bias

In Section 3.2 on the role of logit bias, we connect with Menon et al., 2021 and the literature of logit adjustment for long-tail learning. In particular, Menon et al. modify the unnormalized logits of a K-class classifier by adding a fixed estimate for the log-likelihood of each one of the classes to the corresponding class logit predicted by the network. By training this logit-adjusted classifier with the regular cross-entropy loss and removing the logit bias at the end of training, they show that the remaining trained classifier is robust to shifts in the label distribution. We have modified the main text to clarify this connection.

5. Typo

We thank the reviewer for spotting the oversight and we correct it in the revision.

Proposed revisions

We appreciate your suggestions and we are confident that they will strengthen the final version of our work.

• Clarify that IW softmax is a baseline that we propose, not the literature, for attempting an adaptation of softmax-based losses to graded relevance data.
• Discuss logit adjustment literature for multi-class imbalanced classification more extensively.
• Fix the typo suggested by the reviewer.

References

Menon et al., “Long-tail learning via logit adjustment”, ICLR 2021