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

TL;DR

The paper tackles the challenge that LLMs, when used for dense passage embeddings, are hampered by their causal attention and misalignment with ranking objectives. It introduces ULLME, a unified, plug-and-play framework that enables bidirectional attention across diverse LLM backbones and supports multiple fine-tuning strategies, including a novel Generation-augmented Representation Learning (GRL) method that aligns embedding and generation signals. GRL jointly considers in-embedding relevance and generation-based likelihood, with losses that include a contrastive term, a preference optimization term, and a KL-divergence term between representation- and generation-based relevance distributions, e.g. $ \mathcal{L}_{GRL} = \lambda_{CL} \mathcal{L}_{CL} + \lambda_{DPO} \mathcal{L}_{DPO} + \lambda_{KL} \mathcal{L}_{KL}$ and $ \mathcal{L}_{KL} = \sum_{p \in U} P_{rt}(q,p) \log \frac{P_{rt}(q,p)}{P_{gen}(q,p)} $. The authors release three pre-trained models (1.5B–8B) and demonstrate strong results on MTEB, outperforming baselines like LLM2Vec and Echo across base models and configurations. The framework integrates evaluation via MTEB, supports LoRA-based fine-tuning, and achieves competitive inference speeds, underscoring its practical utility for researchers and practitioners aiming to deploy flexible, high-performance LLM-based embeddings for dense retrieval.

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.