Search papers, labs, and topics across Lattice.
This paper introduces a unified model for on-device retrieval-augmented generation (RAG) that compresses the RAG context and uses shared representations for both retrieval and generation. By unifying these components, the model achieves comparable performance to traditional RAG readers with only 1/10th of the context size, while maintaining storage requirements similar to multi-vector retrieval models. The key result is a significant reduction in both context size and disk utilization for on-device RAG, addressing the memory and storage constraints of local deployment.
On-device RAG gets a major efficiency boost: a unified model slashes context size by 90% without sacrificing accuracy.
Traditional Retrieval-Augmented Generation (RAG) approaches generally assume that retrieval and generation occur on powerful servers removed from the end user. While this reduces local hardware constraints, it introduces significant drawbacks: privacy concerns regarding data access, recurring maintenance and storage costs, increased latency, and the necessity of an internet connection. On-device RAG addresses these challenges by executing the entire pipeline locally, making it ideal for querying sensitive personal information such as financial documents, contact details, and medical history. However, on-device deployment necessitates a delicate balance between limited memory and disk space. Specifically, the context size provided to the generative model must be restricted to manage KV cache and attention memory usage, while the size of stored embeddings must be minimized to preserve disk space. In this work, we propose a unified model that compresses the RAG context and utilizes the same representations for retrieval. This approach minimizes disk utilization compared to using separate representations, while significantly reducing the context size required for generation. With an average of 1/10 of the context, our model matches the performance of a traditional RAG reader without increasing storage requirements compared to a multi-vector retrieval model. This approach represents the first model to unify retrieval and context compression using a shared model and representation. We believe this work will inspire further consolidation of distinct models to optimize on-device performance.
