Repurposing Language Models into Embedding Models: Finding the Compute-Optimal Recipe

We thank the Reviewer for their insightful feedback! Addressing your questions and concerns:

(W1) The term Compute-Optimal [...] is very misleading because it only considers the training cost and ignores the inference cost [which is important in practice].

We do agree that our usage of the “compute-optimal” term is somewhat imprecise – we have updated the paper, pointing out that we mean the training cost only. We have also added an appropriate paragraph in the limitations section.

We believe that there is value in deriving the functional relationship between the model size, training method, data, and the final training loss (which is also done in many previous works on scaling laws). This relationship clearly characterises the ingredients needed to achieve particular performances. We refrain from taking inference costs into account, as they are hard to estimate. However, we stress that by having the functional form, the reader can analyse their particular use case, which may take into account the sum of the training and inference cost as a constraint, and find the optimal iso-total-cost models.

(W2) I think that the figures with FLOPs on the x-axis and different curves are for different model size are much easier to read than what the paper has now [...]

Thank you for this suggestion – to increase the readability, we have added plots where the x-axis is the model size, not FLOPs. During experimentation we were looking at both versions and decided that it is better to present loss vs size plots as the IsoFLOP profiles in them directly answer our motivating question: given the FLOP budget, what is the optimal model size one should use?

(Q1) Although you show that the correlation between MTEB and loss is very high (−0.892) overall, some important trends seem to be different, which might affect your takeaway message. [...]

We thank the Reviewer for raising this. Ultimately, the most suitable metric to use is the one most relevant to the reader. We believe people who have the need to fine-tune their own embedding models and the resources to do so should use their own benchmarks to derive the scaling laws most accurate for them. Without access to the relevant benchmark, we find the validation loss to be the smoothest metric to base scaling laws upon (again, this approach is typically assumed in the scaling laws works). The MTEB results have higher variances and do not necessarily align with the user’s requirement.

(Q2) I don't understand l_r and l_c in line 134 and 135. Why don't we need to use labels.T for l_c? If this is a standard loss, please cite. [...]

It is a standard symmetric contrastive loss used throughout the contrastive learning literature for training textual embedding models [3,4,5,6] (and in other applications [7]). The specific formulation included in our paper originates from [3] (Section 2.2). We cite this work (please see line 126), but we will emphasise it more to avoid confusion – thank you for your suggestion!

As for the question about labels.T for l_c: it should remain labels – we just use PyTorch notation here; the first argument to the cross entropy function is a matrix, and the second argument is a vector indicating for each row what is the index of the “correct” element. (We link the documentation in the footnote on page 3.)

(Q3) Why do you choose to use the English partition of the BAAI BGE dataset, which is designed for Chinese embedding? [...]

The reason for choosing BAAI BGE is that this dataset is a massive collection of data for training embeddings. Since we do not consider scaling laws under data constraints like [1], we needed a large enough dataset on which none of our experiments would exceed 1 epoch of training; BAAI BGE was the only published dataset satisfying our constraints.

As for why we used only the English partition of BAAI BGE: we want to explore a family of decoder-only models with different model parameter scales, and the Pythia family is the only available option at the time of experiments that offers lots of sub-billion-parameter models, which are common for embedding models. The Pythia family was pre-trained on The Pile [2], which is an English-only dataset, and hence it does not have Chinese knowledge and capabilities. Therefore, we only focused on the English partition of BAAI BGE.

We have updated the paper to include the full explanation provided here.

(Q4) Did you get the performance of full MTEB? Does the full MTEB give you different trends? [...] you should make it clear in those figures [that only a subset of MTEB was used].

The Reviewer is correct here—for practical computational purposes, we did use a subset of MTEB instead of the whole benchmark. (As we explain in Appendix B, this subset was selected in a principled way to be representative of the whole MTEB.)

During the rebuttal period, we checked that the results on the whole suite (except MSMARCO, which is the most expensive to run) are highly correlated with our selected subset. To this end, we compared the two evaluations for 8 checkpoints from the full fine-tuning experiment, and the Pearson correlation is 0.976. This information has been included in the revised version of the paper. We have also fixed the captions to “MTEB subset performance.”

References

[1] Muennighoff et al.: Scaling data-constrained language models. NeurIPS 2023
[2] Gao et al.: The Pile: An 800GB dataset of diverse text for language modeling. arXiv 2021
[3] Neelakantan et al.: Text and Code Embeddings by Contrastive Pre-Training. arXiv 2022
[4] Wang et al.: Text Embeddings by Weakly-Supervised Contrastive Pre-training. arXiv 2022
[5] Wang et al.: Improving Text Embeddings with Large Language Models. arXiv 2024
[6] Günther et al.: Jina Embeddings: A Novel Set of High-Performance Sentence Embedding. arXiv 2023
[7] Eysenbach et al.: Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference. arXiv 2024

The authors address my concerns well. I have no other major reasons to reject this paper, so I raise my score from 6 to 7.

We thank the Reviewer for their effort to assess our work! Below we address each of the concerns.

(1a) The paper's generalization ability is limited. As mentioned in the limitation, the proposed computation law is only on the Pythia; additional language models might better reflect the generalization ability.

We agree that having additional models would allow for a better reflection of the paper’s generalisation. During the rebuttal we ran experiments with a very recent model – Gemma 2b [1], using LoRA and full fine-tuning methods.

Importantly, we found that the the loss formula with parameters fitted to the Pythia data points, adequately describes the behaviour of Gemma 2b training for the ‘medium’ budgets.

At the extremely low computational budget, Gemma performs better than the prediction, which intuitively is caused by the general higher model capability than those in the Pythia suite; also, scaling laws tend to be much less regular in smaller regimes – see Figure 4 in the pdf attachment.

As a result of being time-restricted, we did not perform the two highest budget trainings, but we believe that the results so far are a strong indication as to the generalisability of our findings. We will update our paper to contain those new findings once we have the complete results for Gemma.

(1b) The paper also only focuses on contrastive finetuning.

We intentionally focus only on contrastive learning: this has become a standard approach to obtaining useful embedding models. We note that there has been no study before ours that helps to select the most compute-efficient variant of contrastive learning approaches in decoder-only models, rendering our study a worthwhile contribution.

(2) The whole paper mainly focuses on the performance analysis. It would be better to include additional theoretical support for the phenomenon that appears in the paper. The paper can also be strengthened by testing the proposed performance law on some unobserved configurations.

We provide theoretical support by utilising a classical risk decomposition model (similar to that of the Compute-Optimal Scaling Laws paper [2]), and adapting it to account for the difference in forward and backward computation parameters in Section 4.3. With this theoretical framework we describe the loss behaviour in terms of the model parameters (both forward and backward), and the amount of data.

We also show that the scaling law generalises to the very recent Gemma 2B model (see above).

(3a) The paper claims to also analyse the relationships between the data quantities and the performance. However, I cannot find it in Section 4.2.

In this paper, we focused on investigating the compute-bounded setting – the data-quantity factor is only present implicitly as it can be inferred from the FLOP limit and the model size. However, we now realize that the data-centric perspective is also very valuable and we should analyse it explicitly, especially in the light of the claim in line 157. For this reason, we attach Figure 3 to our 1-page pdf attachment with plots. A simple and practical conclusion can be inferred from it – if the data is restricted, one should choose the biggest model available and use full fine-tuning or LoRA with a rank on the bigger side. We have updated our paper to contain this plot and conclusion.

We note that these conclusions can be achieved using the raw experimental data that we provide – please, see the link at the bottom of page 5 of our paper.

(3b) The paper fails to follow the NeurIPS citation format [...].

Thank you for noticing it! This has already been fixed.

References

[1] Mesnard et al.: Gemma: Open Models Based on Gemini Research and Technology. arXiv 2024
[2] Hoffmann et al.: Training compute-optimal large language models. arXiv 2022

We are grateful for the Reviewer’s feedback and hope the concerns are appropriately addressed. If that is the case, we kindly ask you to reconsider the paper score.

We again thank you for your feedback, and we are grateful for your support of our work!

We thank the Reviewer for detailed and informative feedback. Below we address each of the concerns.

(1) [Despite you mention so in lines 51-52] data quantity, a very crucial factor, is not investigated in the paper.

In this paper, we focused on investigating the compute-bounded setting – the data-quantity factor is present implicitly as it can be inferred from the FLOP limit and the model size. Following your comment, we recognise that the data-centric perspective is also very valuable. For this reason, we attach Figure 2 to our 1-page pdf file with plots. A simple and practical conclusion can be inferred from it – if the data is restricted, one should choose the biggest model available and use full fine-tuning or LoRA with a rank on the bigger side. We have updated our paper to contain this plot and conclusion.

We note that these conclusions can be achieved using the raw experimental data that we provide – please, see the link at the bottom of page 5 of our paper.

(2) The conclusion that full fine-tuning and low-rank adaptation fine-tuning produce optimal models at lower and higher computational budgets respectively seems somewhat superficial.

We do agree with the Reviewer that our study did not reveal patterns that are unexpected or peculiar. Nevertheless, we argue that our findings constitute a solid empirical grounding and, thus a concrete scientific contribution to helping researchers and practitioners working with embedding models. It is the first study devoted to establishing compute-efficient approaches to contrastive fine-tuning of decoder-only language models. Given the recent activity in this area, we found it to be a serious research gap that needed to be addressed.

Also, please note that our study concerns not only the selection of the optimal method from amongst the four pre-selected, but we also go deeper and analyse important hyper-parameters of those methods (LoRA rank and active block fraction). This was not previously studied, with a notable exception of limited analysis in [2].

(3) [...] Retrieval is a crucial task for which embedding models can be used. However, only SciFact is evaluated. Scifact is not a representative retrieval dataset as its corpus is relatively small [...].

Our selection of SciFact was guided by its strong correlation (0.927) with the overall performance in the retrieval tasks category, as evidenced by the Table 11 of MTEB [8]. Nevertheless, we concur with the Reviewer's observation regarding the size of the SciFact benchmark.

Therefore, we now run new experiments to investigate the correlations between the performance on SciFact and two large retrieval benchmarks: Natural Questions [6] (2.6M test data points), and DBPedia [7] (4.6M test data points) for our best models for each (FLOP, method) combination, which are 0.976 and 0.971 respectively (also see Figure 3 of the attachment). We hence conclude, that despite its humble size, SciFact is a satisfactory representative of the retrieval category. This clarification has been added to the paper.

(4) Figure 1 is somewhat hard to read since it is difficult to distinguish the lines representing full fine-tuning, LoRA, and block freezing.

We thank the Reviewer for the suggestion. We have updated the figure to make it more legible in Figure 1 of the attachment.

(5) The average pooling method to extract representation is not entirely reasonable, as the autoregressive characteristic would cause the outputs of the first tokens to lack information about the latter tokens.

While using the average pooling method may indeed seem unintuitive, it works very well in practice. In Appendix D, section 3, we compare the two most popular methods for extracting embeddings – average and last pooling methods – for Pythia 410m, fine-tuned with the FLOP budget of 9.6e+16. We perform evaluation on the whole MTEB [8] without MSMARCO, and we find the difference in performance is 0.47 vs 0.36 in favour of average pooling.

We note that average pooling has been used in several previous works [1,2,3]. Moreover, theoretically, the average pooling method has the possibility to converge to the ‘last’ method by zeroing out all the previous tokens. Hence, we argue the average pooling method is empirically and theoretically justified.

(6) L141-146 is not entirely accurate. Hard negative mining methods like [4] and [5] do not require this two-stage training procedure [which] is more of a recent trend in training powerful general-purpose embedding models.

We indeed ignored the existing approaches for mining the hard negatives. We have rewritten the mentioned paragraph as shown below (changes in bold) – if something is still imprecise, we would appreciate feedback!

[…] which can be incorporated into the contrastive loss to promote learning more precise representations. Moreover, there are works that focus on approaches for mining hard negatives which result in better training data in the context of specific downstream tasks [4,5].

However, recent works aiming at training powerful, general-purpose embedding models that do rely on datasets with hard negatives often arrive at embeddings by having two distinct training phases […]

References

[1] Li et al.: Towards General Text Embeddings with Multi-stage Contrastive Learning. arXiv 2023
[2] Wang et al.: Improving Text Embeddings with Large Language Models. arXiv 2024
[3] Wang et al.: Text Embeddings by Weakly-Supervised Contrastive Pre-training. arXiv 2022
[4] Xiong et al.: Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval. ICLR 2021
[5] Zhan et al.: Optimizing Dense Retrieval Model Training with Hard Negatives. SIGIR 2021
[6] Kwiatkowski et al.: Natural Questions: A Benchmark for Question Answering Research. TACL 2019
[7] Hasibi et al.: DBpedia-Entity v2: A Test Collection for Entity Search. SIGIR 2017
[8] Muenninghoff et al.: MTEB: Massive Text Embedding Benchmark. EACL 2023

We thank the Reviewer for their effort to assess our work! Below we respond to the two points of critique.

(1) The paper does not introduce a new methodology. The techniques employed, such as scaling laws, extracting representations from transformers, and contrastive fine-tuning, are all well-established in the literature.

We agree with the Reviewer that our study does not introduce novel techniques – and this was not our intention in this project. We rather aimed to investigate and compare existing modern approaches in the practical, limited compute-budget setting. Before our work, no study had been devoted to establishing compute-efficient approaches to contrastive fine-tuning of decoder-only language models. Given the recent activity in this area, we found it to be an important research gap that needs to be addressed. We believe that our findings constitute a concrete scientific contribution to helping researchers and practitioners working with embedding models.

(2) There are several existing studies that the paper fails to acknowledge, which could provide a richer context and enhance the literature review. Notable omissions include works like “LLM-Oriented Retrieval Tuner” and “Scaling Laws for Dense Retrieval.”

Thank you for pointing out these two recent works! We have extended the related work section to include both of them. Specifically, starting from line 95 (new text in bold):

These approaches are all applicable to embedding models. Muennighoff [1] explored a simple parameter-efficient approach to repurpose GPT models into embedding models where only the bias tensors of the transformer model are updated [2]. Sun et al. [3] developed a parameter-efficient method designed specifically for fine-tuning embedding models. However, a systematic study of what parameter-efficient methods are optimal under what scenarios has not been performed, which is the aim of our work.

Also, starting from line 84:

The single most related investigation is a concurrent work [4], where the scaling laws for encoder models for retrieval are investigated. There are significant differences in the settings we consider. Our work focus is investigating the process of fine-tuning decoder-only models for good quality embeddings. Moreover, our main goal is to find which strategy for achieving good embeddings is optimal in a budget-restricted setting, while taking into account the popularity of applying parameter efficient methods for fine-tuning, like LoRA or partial model freezing.

References

[1] Muennighoff: SGPT: GPT sentence embeddings for semantic search. arXiv 2022
[2] Zaken et al.: BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models. ACL 2022
[3] Sun et al.: LLM-Oriented Retrieval Tuner. arXiv 2024
[4] Fang et al.: Scaling Laws For Dense Retrieval. SIGIR 2024

We hope that we adequately addressed the Reviewer’s concerns. If that is the case, we kindly ask for a reconsideration of the paper score.