ULLME : A Unified Framework for Large Language Model Embeddings with Generation-Augmented Learning (2024)

Abstract: Large Language Models (LLMs) excel in various natural language processing tasks, but leveraging them for dense passage embedding remains challenging. This is due to their causal attention mechanism and the misalignment between their pre-training objectives and the text ranking tasks. Despite some recent efforts to address these issues, existing frameworks for LLM-based text embeddings have been limited by their support for only a limited range of LLM architectures and fine-tuning strategies, limiting their practical application and versatility. In this work, we introduce the Unified framework for Large Language Model Embedding (ULLME), a flexible, plug-and-play implementation that enables bidirectional attention across various LLMs and supports a range of fine-tuning strategies. We also propose Generation-augmented Representation Learning (GRL), a novel fine-tuning method to boost LLMs for text embedding tasks. GRL enforces consistency between representation-based and generation-based relevance scores, leveraging LLMs’ powerful generative abilities for learning passage embeddings. To showcase our framework’s flexibility and effectiveness, we release three pre-trained models from ULLME with different backbone architectures, ranging from 1.5B to 8B parameters, all of which demonstrate strong performance on the Massive Text Embedding Benchmark. Our framework is publicly available at: https://github.com/nlp-uoregon/ullme. A demo video for ULLME can also be found at https://rb.gy/ws1ile.

## introduction for many years, the field of information retrieval has been dominated by a paradigm that relied heavily on pre-trained bidirectional encoders or framework #supported supported fine-tuning strategy llms sft dpo contrastive sentencetrasformers @xcite 1 ✗ ✗ ✓ echo-embedding @xcite 2 ✗ ✗ ✗ gritlm @xcite 2 ✓ ✗ ✓ llm2vec @xcite 3 ✗ ✗ ✓ nv-emb @xcite 1 in this paper, we present ullme, a versatile and extensible platform designed to advance the use of llms for dense retrieval. ullme addresses the critical limitations of existing frameworks by offering a comprehensive, plug-and-play solution that seamlessly enables bidirectional attention across a array of diverse llm families, including llama, mistral, phi, qwen, among others. our framework's flexibility also extends beyond model compatibility, supporting a wide spectrum of fine-tuning strategies for llm-based representation learning. as such, ullme provides an unified framework for various llm backbones and finetuning methodologies, allowing developers to comprehensively explore the full potential of llms in diverse embedding tasks, free from the constraints of implementation-specific restrictions. in addition, existing frameworks for llm-based text embeddings can be challenging for general users who are not familiar with training details like contrastive learning with large batch sizes and efficient fine-tuning. ullme lowers these entry barriers by providing an efficient, user-friendly abstraction from those complexities, allowing users to focus on their data and tasks. for instance, ullme's training processes are integrated with advanced techniques like gradcache @xcite and @xcite , enabling efficient contrastive learning and tuning with larger batch sizes, and sparing users from complicated configuration and testing. ullme also comes with user-friendly features that make it easy to evaluate various fine-tuned llms using the massive text embedding benchmark (mteb) @xcite , a comprehensive evaluation suite with numerous tasks for text embeddings. to showcase the versatility and effectiveness of ullme, we release three pre-trained llm-embedding models with different backbone architectures, ranging from 1.5b to 8b parameters, which deliver highly competitive results on mteb. our findings also highlight the advantages of our new fine-tuning method, grl, which significantly outperforms the strong baselines, underscoring the potential of our framework to advance research and development in llm-based embeddings. llms for dense retrieval. recent advancements in this area have primarily addressed two key challenges: (i): overcoming llms' causal attention limitations by developing methods to enable bidirectional attention within llms @xcite @xcite , allowing models to consider both past and future context when computing embeddings, and (ii): aligning llm pretraining with text ranking by fine-tuning llms via contrastive learning @xcite @xcite . this process can also be augmented with additional objectives such as supervised fine-tuning (sft) @xcite or mask-filling tasks @xcite ). an alternative approach proposed by @xcite involves a prompting method where the input sequence is duplicated, enabling each token to attend to future tokens and mitigating the contextualization issues inherent in causal attention. while these methods have shown promise, they generally do not explicitly enforce consistency between the model's understanding of relevance in both the embedding and generation spaces. this limitation restricts their ability to fully leverage the remarkable generative capabilities of llms for dense retrieval tasks. our work, grl, builds upon these foundations while addressing their limitations, introducing novel techniques to harmonize embedding-based and generation-based relevance scoring within a unified framework. frameworks of llms for dense retrieval. existing frameworks for llms in dense retrieval have been constrained by their limited support for llm architectures and fine-tuning strategies. as shown in table1, sentencetransformers @xcite supports various types of llms but is primarily designed for inference without allowing fine-tuning, limiting its applicability in advancing state-of-the-art dense retrieval methods. some recent works @xcite @xcite , such as echo @xcite , gritlm @xcite , @xcite , and the models in the hugging face's mteb leaderboard foot_1 , have introduced implementations for llm-based text embeddings. however, these approaches are often tailored to specific model architectures and training methods with hard-coded implementations, thus restricting their adaptability and use across different llm architectures and fine-tuning strategies to meet diverse development and application demands. in contrast, our framework ullme addresses these limitations by offering a flexible and extensible platform. ullme can accommodate a diverse range of llm backbones and supports various training approaches, making it highly versatile and broadly applicable. ## ullme -unified framework for large language model embedding we present an overview of our ullme framework in section 3.1 while section 3.2 details the key technical methods. ## experiments our ullme framework supports various llm architectures and fine-tuning strategies for text embeddings with convenient interface. to highlight the framework's flexibility, we demonstrate the operations of ullme with three different base llms ranging from 1.5b to 8b parameters: phi-1.5b @xcite , mistral-7b-instruct-v0.2 @xcite , and meta-llama3-8b-instruct (ai@meta, 2024). for each llm, we evaluate ullme's performance for different combinations of attention and fine-tuning approaches, including: base: original causal model, causal + cl: causal model fine-tuned with contrastive learning, bi + cl: bidirectional-enabled model finetuned with contrastive learning, and bi + cl + sft: bidirectional-enabled model fine-tuned with contrastive learning and sft. in addition, we report the performance of our generation-augmented representation learning (grl) method for finetuning llms in ullme, featuring the full model grl and grl sf t , a variant of grl that replaces dpo with sft for tuning. finally, we compare the performance of ullme's models with recent state-of-the-art methods for llm-based text embeddings, including echo @xcite and @xcite . settings. following prior work @xcite @xcite , we use a curated subset of the msmarco dataset @xcite for model training. mteb datasets are employed for evaluation. to train the models, we utilize lora @xcite with @xmath0, and enable various optimization techniques, i.e., gradcache, gradient checkpointing, mixed precision training, and fsdp @xcite , to minimize gpu memory requirements. we utilize the adamw optimizer @xcite with a learning rate of 2e-4 and a batch size of 512 with the number of hard negative passages per example was set to 8. we train the models for one epoch on msmarco. the weights for the grl loss components include λ @xmath1.5. the scaling factor β in the dpo loss was set to 0.1. ## conclusion we introduce ullme (unified framework for large language model embedding), a comprehensive and flexible toolkit for leveraging llms for text embeddings and dense retrieval tasks. our work addresses critical limitations in existing frameworks for llm embeddings by providing support for various llm architectures, fine-tuning strategies, and benchmark evaluation within a single, user-friendly framework. our experimental results demonstrate the effectiveness of ullme, particularly the grl strategy, in improving dense retrieval performance across various llm scales and tasks. our potential future directions include exploration of better techniques to leverage the generative and discriminative capabilities of llms, and extension of the framework to support emerging llm architectures and training paradigms. we anticipate that ullme will facilitate broader applications of llm embeddings in downstream tasks, ranging from deep context understanding requirements like sentiment analysis @xcite to text style comprehension tasks such as authorship attribution @xcite , thereby contributing to the advancement of natural language processing and information retrieval fields.