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

Summary

The paper fine-tunes Pythia LLMs with different sizes for information retrieval (IR) applications. The paper investigates the effect of different model sizes, FLOPs, fine-tuning methods such as LoRA or bias tuning, and some important hyperparameters (e.g., dimension in LoRA) on loss and downstream applications (measured by MTEB). Some takeaways are interesting and counterintuitive (e.g., the larger model does not necessarily achieve better loss).

Strengths

This paper conducts very extensive experiments and analyses in the main paper and appendix. I appreciate all the efforts and believe that it would be very useful for practitioners who want to create their own embedding models from LLM.

Weaknesses

Although I think the experiments in this paper are valuable, I think main goal of the analyses is not very practical. Personally, I think the term Compute-Optimal from Chinchilla (https://arxiv.org/pdf/2203.15556) and this paper is very misleading because it only considers the training cost and ignores the inference cost. In practice, the LLMs are often trained much longer than this "optimal" value. For example, Chinchilla-optimal for 8B LLM is just 200B tokens, but LLaMA 3 is trained using over 15T tokens (https://ai.meta.com/blog/meta-llama-3/). For the information retrieval (IR) models, the inference cost might be even more important compared to LLMs. In addition to main complaint above, I think that some details could be explained more clearly (see the questions below). Finally, 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 (the model size on the x-axis and different curves are for different FLOPs), but it might be my personally preference.

Questions

1. 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. For example, the best score of MTEB always comes from the largest model in Figure 7-10, but the lowest loss often comes from smaller model size in Figure 2 and 4. The bias tuning methods are significantly worse in terms of loss, but it seems to achieve the best MTEB when FLOP is 5e+18 in Figure 11. When you write your takeaways, why not also include the results of MTEB? 2. 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. Otherwise, please provide more explanations/motivations for this loss. 3. Why do you choose to use the English partition of the BAAI BGE dataset, which is designed for Chinese embedding? If it is not a common choice, you should explain more about this. 4. Did you get the performance of full MTEB? Does the full MTEB give you different trends? In the figures, MTEB performances means the performances on the task subset, right? I think you should make it clear in those figures (MTEB performances -> MTEB subset performances).

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

The current limitation section is good and adding some discussion of inference cost could make it better.

Summary

This paper proposes a new algorithm to automatically find optimal configurations of model sizes, data quantities, and fine-tuning methods for text embedding models at different computational budget levels by contrastive pretraining. Specifically, the paper analyzes the choices of different configurations and then designs an algorithm to predict optimal architecture based on budget. The paper also designs a new task to find optimal model configuration given limited budgets. The paper first introduces basics, including the scaling laws, the constrained optimization problem, the contrastive learning, fine-tuning methods, and then computational cost formulas. The paper then experiments with different finetuning methods to find the optimal loss regarding parameters. Furthermore, the paper aims to find the relationship between loss functions regarding the number of parameters (trainable/total) and the number of training tokens. The paper then fit the plots to find the best equations.

Strengths

1. The paper proposes a new computation optimization problem: given a fine-tuning corpus and a family of pre-trained decoder-only language models of different sizes, minimize the contrastive loss subject to a fixed computational budget. The newly proposed problem is important and has the potential to expand. 2. The paper conducts a comprehensive analysis of existing fine-tuning techniques with different hyperparameters. The paper visualizes the results with clear figures and concludes with several useful results. The paper also predicts the optimal loss for different budgets. The paper evaluates the downstream results. 3. The paper provides code, additional evaluation results and analysis in the Appendix, and detailed computational resources.

Weaknesses

1. 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. The paper also only focuses on contrastive finetuning. 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. 3. The paper claims to also analyze the relationships between the data quantities and the performance. However, I cannot find it in Section 4.2. The paper fails to follow the NeurIPS citation format, which uses numbers for references/citations.

Questions

See weakness

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper includes a limitation section in Section 5.

Summary

This paper investigates how to effectively train text embedding models from pre-trained decoder-only language models while considering computational budget constraints. The authors explore the influence of model sizes, fine-tuning methods, and computational budgets on the performance of the embedding models. The results demonstrate that full fine-tuning is optimal for lower computational budgets, whereas low-rank adaptation fine-tuning is more effective for higher computational budgets.

Strengths

(1) This paper focuses on the important problem of finding the optimal training settings for embedding models within a fixed compute budget. (2) This paper conducts extensive experiments across 8 model sizes, 6 compute budgets, and 4 tuning methods. (3) Although the paper uses contrastive loss as the primary measure of model performance, the authors show a strong correlation between contrastive loss and downstream task performance.

Weaknesses

(1) In L51-52, it is stated that "given a fixed computational budget, predicts the optimal network architecture, data quantity, and parameter-efficient fine-tuning hyperparameters". However, data quantity, a very crucial factor, is not investigated in the 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. (3) I have concerns about the selection of the MTEB subset. 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 (only with 5183 documents). (4) Figure 1 is somewhat hard to read since it is difficult to distinguish the lines representing full fine-tuning, LoRA, and block freezing. (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. (6) L141-146 is not entirely accurate. Hard negative mining methods like [1] and [2] do not require this two-stage training procedure. This two-stage training procedure is more of a recent trend in training powerful general-purpose embedding models. [1] Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, Arnold Overwijk. Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval. ICLR 2021. [2] Jingtao Zhan, Jiaxin Mao, Yiqun Liu, Jiafeng Guo, Min Zhang, Shaoping Ma. Optimizing Dense Retrieval Model Training with Hard Negatives. SIGIR 2021.

Questions

Please see my comments in Weakness.

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The paper has a proper Limitations and future work section.

Summary

This paper focuses on the efficient contrastive training of text embedding models using pre-trained decoder-only language models. The main contribution is an algorithm that determines the best configuration of model sizes, data amounts, and fine-tuning methods for different computational budgets. Through extensive experimentation, the authors developed a guideline that helps practitioners choose the most suitable design options for their text embedding models. The study finds that full fine-tuning is best for lower computational budgets, while low-rank adaptation fine-tuning is optimal for higher budgets.

Strengths

The research question posed, “What is the best embedding model one can train from a backbone decoder-only LLM with a fixed budget?” is highly valuable and important. It addresses a critical need within the NLP and IR domains. The paper boasts a robust experimental setup, and the conclusions drawn from these experiments are insightful. The analysis based on empirical results provides a useful reference for advancements in both NLP and IR fields. This thorough approach enhances the reliability and applicability of the research findings.

Weaknesses

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. 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.” Acknowledging and discussing these related studies could strengthen the paper by positioning it within the existing research landscape more accurately.

Questions

Please refer to the previous sections.

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

4

Limitations

This paper discusses the limitation in Section 5.

New Results

Dear Reviewers, We are happy to announce one more result concerning the very recent Gemma-2-2b model (released two weeks ago). Despite the limited time, we managed to conduct experiments with full fine-tuning and observed that they were consistent with those obtained with Gemma-2 (see the first subplot of Figure 4 in the attached document). This result further supports the robustness and broad applicability of our work. We will present a more exhaustive study on this model in the camera-ready version. Best regards, Authors