Repurposing Language Models into Embedding Models: Finding the Compute-Optimal Recipe (2406.04165v2)

Abstract: Text embeddings are essential for many tasks, such as document retrieval, clustering, and semantic similarity assessment. In this paper, we study how to contrastively train text embedding models in a compute-optimal fashion, given a suite of pre-trained decoder-only LLMs. Our innovation is an algorithm that produces optimal configurations of model sizes, data quantities, and fine-tuning methods for text-embedding models at different computational budget levels. The resulting recipe, which we obtain through extensive experiments, can be used by practitioners to make informed design choices for their embedding models. Specifically, our findings suggest that full fine-tuning and low-rank adaptation fine-tuning produce optimal models at lower and higher computational budgets respectively.

• The paper introduces a compute-optimal contrastive fine-tuning strategy that repurposes decoder-only LMs into high-quality text embedding models.
• It demonstrates that full fine-tuning is favorable for low compute budgets while techniques like LoRA excel at higher FLOP thresholds, guided by empirical scaling laws.
• The study provides actionable guidelines for balancing model size, data volume, and parameter tuning methods to optimize embedding performance under varying computational resources.

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

The paper "Repurposing LLMs into Embedding Models: Finding the Compute-Optimal Recipe" presents a comprehensive investigation into efficient methodologies for fine-tuning pretrained LLMs (LMs) to obtain high-quality text embeddings under varying computational constraints. The authors focus particularly on decoder-only LLMs, examining the optimal configurations of model sizes, data volumes, and fine-tuning techniques.

Summary

The primary objective of this work is to identify a compute-optimal strategy for transforming pretrained LMs into text embedding models using a contrastive training objective. The authors propose an innovative algorithm that delivers optimal recipes aligning model sizes, data quantities, and fine-tuning methods tailored to specific computational budgets. This solution is achieved through rigorous experimentation with various methodologies and hyperparameters, offering significant insights for practitioners in the field.

Key Findings

1. Contrastive Loss and Fine-tuning Techniques:
   • Full Fine-tuning:

   The paper finds that full fine-tuning of all model parameters is viable within lower computational budgets. The IsoFLOP profiles reveal that the optimal model size increases with larger budgets, while the loss decreases. - Block Freezing:

   This method involves freezing a portion of the transformer blocks, with the results indicating a minor decrement in the optimal active block fraction as the model size enlarges. The approach maintains comparable performance to full fine-tuning under high computational constraints. - Bias-only Tuning:

   Updating only the bias parameters across the network surfaces as suboptimal for embedding training due to higher achievable losses. - Low-Rank Adaptation (LoRA):

   LoRA demonstrates effectiveness at higher computational budgets, with a consistent trend showing that specific LoRA ranks, such as 32 and 128, result in optimal loss performance across various configurations.

2. Scaling Laws: By fitting scaling laws to their experimental data, the authors derive an empirically grounded formula to predict the contrastive loss as a function of model size, data volume, and the fraction of trainable parameters. This model can project performance enhancements at unobserved FLOP budgets, thereby guiding the selection of compute-optimal configurations.
3. Compute-Optimal Frontier: The research identifies a threshold computational budget (~9.06e+16 FLOP), below which full fine-tuning outperforms other methods. Above this threshold, LoRA becomes the preferred methodology. The paper’s compute-optimal frontier dictates the best achievable loss for given resource constraints, offering a clear pathway for practitioners to follow.

Implications and Future Directions

The findings emphasize the importance of meticulous design choice and hyperparameter tuning to repurpose LMs into efficient embedding models. The optimal configurations highlighted in this work provide a foundational framework that can be instrumental for both academic research and practical applications where computational resources are limited.

Future work may extend these conclusions to other model families beyond the Pythia suite, replicating the methodology to discern whether similar scaling laws apply. Additionally, exploring alternative embedding readout techniques (e.g., max pooling or last-token pooling) and integrating hard negative examples in contrastive training could further optimize embedding quality. Evaluating the full range of MTEB tasks and averaging over multiple random seeds may also enhance the robustness and generalizability of the conclusions.

Conclusion

The paper delivers valuable insights into compute-efficient methodologies for fine-tuning LMs into robust embedding models, presenting a concrete algorithm to navigate model, data, and fine-tuning choices contingent on computational availability. This work substantially contributes to the understanding and practicality of leveraging large pretrained LMs for diverse embedding tasks, particularly under resource constraints. The implications for both theoretical advancements and practical implementations are considerable, paving the way for further refinements in embedding model training, especially in lower-resource contexts.