Landmark Embedding: A Chunking-Free Embedding Method For Retrieval Augmented Long-Context Large Language Models (2024)

Abstract: Retrieval augmentation is a promising approach to handle long-context language modeling. However, the existing retrieval methods usually work with the chunked context, which is prone to inferior quality of semantic representation and incomplete retrieval of useful information. In this work, we propose a new method for the retrieval augmentation of long-context language modeling, called Landmark Embedding. Our method is characterized by threefold technical contributions. Firstly, we introduce a chunking-free architecture , which keeps the long context coherent such that high-quality embeddings can be generated for the fine-grained units within the context. Secondly, we present a position-aware objective function, which prioritizes the ultimate boundary for a consecutive span of information. By learning to discriminate such a special position, the useful information can be comprehensively retrieved for the query. Thirdly, we design a novel multi-stage learning algorithm, which makes the best use of readily available data and synthetic data for cost-effective training of the landmark embedding. In our experimental study, landmark embedding is able to substantially improve the performance for both LLaMA-2 and ChatGPT in a variety of long-context tasks; meanwhile, it also outperforms the existing retrieval methods with a notable advantage. Our model and source code will be made publicly available.

## introduction large language models (llms) need to handle long-sequence inputs when dealing with many important applications, such as question answering and reading comprehension @xcite . unfortunately, the existing llms are usually constrained by a limited size of context window, e.g., 2k for llama-1 @xcite and 4k for llama-2 @xcite . although †. corresponding author 1. https://github.com/flagopen/flagembedding paid a visit to eiffel tower on sunday. erized by the architectural marvel, he the morning exploring its details and in breathtaking views of the city. after ghtful lunch, he joined with his friends. ey went to louvre museum together. ## related work the related works are reviewed from three aspects: long-context language modeling, retrievalaugmentation, embedding-based retrieval methods. first of all, a large body of research works has been dedicated to the extension of llms' lengths from different directions. one common practice is to modify the position encoding mechanism, where the llms trained on short texts can directly handle longer input during the inference time @xcite . despite simplicity, such methods are prone to inferior performances without further fine-tuning. another popular method is to take advantage of continual training, where the existing llms are fine-tuned over long-sequence data to establish a longer context window @xcite @xcite @xcite . however, the fine-tuning based methods are prone to two subsequent problems. on one hand, the fine-tuned llm will incur an expensive cost for both training and inference. on the other hand, the fine-tuning operation could be unfavorable to the llm's performance with short-sequence inputs. apart from the above common approaches, the llm's context can also be extended by context compression @xcite and stream processing @xcite . nevertheless, the compression methods are likely to result in information loss, while the stream processing will discard the useful information beyond the sliding window. it remains to explore more effective methods for long-context language modeling in the future. in the meantime, the retrieval-augmented generation (rag) is another important issue for llms' research. typically, it employs a standalone retriever, where useful information can be introduced from a vast open-world corpus to enhance the llm's generation quality @xcite @xcite . previously, rag used to be applied for knowledge-intensive tasks, such as open-domain question answering and fact verification @xcite , where an external knowledge base is presented. recently, retrieval augmentation is also found helpful to long-context language modeling, as useful information can be retrieved and presented as a concise input for the llm @xcite @xcite . compared with other alternative methods on context extension, the retrieval-based methods are distinguished for the simplicity and compatibility, as they don't need modification of the downstream llm, and can be easily combined with other methods to establish a longer context. finally, the rag system usually works with an embedding model to retrieve the useful information. in the past few years, many critical techniques have been well established for the effective learning of embedding models, e.g., pre-training @xcite @xcite @xcite , hard-negative sampling @xcite , knowledge distillation @xcite , etc. on top of these techniques, there have been a number of powerful embedding models developed for the general-purpose retrieval applications @xcite @xcite @xcite . however, the existing methods rely on chunking when dealing with the retrieval augmentation of long-context language modeling. as a result, it will inevitably break the coherence of context, which is prone to inferior quality of embedding and incomplete retrieval of useful information. ## conclusion in this paper, we present a new method, landmark embedding, which facilitates the retrieval augmentation of long-context language modeling. the new method is featured by its chunking-free architecture, where discriminative embeddings can be generated for each fine-grained input unit based on the semantic information within a coherent context. a position-aware objective function is proposed; it enables landmark embedding to identify the ultimate boundary of information, which benefits the completeness of retrieval. a novel multi-stage learning algorithm is designed, which makes the best of the readily available data and synthetic data for the effective training of the embedding model. landmark embedding is empirically verified by comprehensive evaluations, where it notably outperforms the existing retrieval methods, bringing in a superior retrieval augmentation effect for both llama-2-7b (4k) and chatgpt-3.5 (16k).