# A Unified Model and Document Representation for On-Device Retrieval-Augmented Generation

arXiv: https://arxiv.org/abs/2604.14403 · HTML: https://mlsanity.com/paper/2604.14403/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RAG pipeline design and evaluation, one of 15 directions in Retrieval-augmented generation.
- Builds on 39 of its 55 references that are in the map.
- First posted on arXiv 15 Apr 2026.

## Placed in

- Retrieval-augmented generation / [RAG pipeline design and evaluation](https://mlsanity.com/rag/rag-pipeline-design-and-evaluation/)

## Builds on (most cited first, 30 of 39)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [SimCSE: Simple Contrastive Learning of Sentence Embeddings](https://mlsanity.com/paper/2104.08821/) (2021)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT](https://mlsanity.com/paper/2004.12832/) (2020)
- [Towards General Text Embeddings with Multi-stage Contrastive Learning](https://mlsanity.com/paper/2308.03281/) (2023)
- [Learning to Compress Prompts with Gist Tokens](https://mlsanity.com/paper/2304.08467/) (2023)
- [LLMLingua-2: Data Distillation for Efficient and Faithful Task-Agnostic Prompt Compression](https://mlsanity.com/paper/2403.12968/) (2024)
- [Generative Representational Instruction Tuning](https://mlsanity.com/paper/2402.09906/) (2024)
- [PowerInfer: Fast Large Language Model Serving with a Consumer-grade GPU](https://mlsanity.com/paper/2312.12456/) (2023)
- [In-context Autoencoder for Context Compression in a Large Language Model](https://mlsanity.com/paper/2307.06945/) (2023)
- [Compressing Context to Enhance Inference Efficiency of Large Language Models](https://mlsanity.com/paper/2310.06201/) (2023)
- [LLM in a flash: Efficient Large Language Model Inference with Limited Memory](https://mlsanity.com/paper/2312.11514/) (2023)
- [FlashRAG: A Modular Toolkit for Efficient Retrieval-Augmented Generation Research](https://mlsanity.com/paper/2405.13576/) (2024)
- [Improving Efficient Neural Ranking Models with Cross-Architecture Knowledge Distillation](https://mlsanity.com/paper/2010.02666/) (2020)
- [xRAG: Extreme Context Compression for Retrieval-augmented Generation with One Token](https://mlsanity.com/paper/2405.13792/) (2024)
- [Prompt Compression and Contrastive Conditioning for Controllability and Toxicity Reduction in Language Models](https://mlsanity.com/paper/2210.03162/) (2022)
- [Learning to Compress Prompt in Natural Language Formats](https://mlsanity.com/paper/2402.18700/) (2024)
- [Retrieval-Enhanced Machine Learning](https://mlsanity.com/paper/2205.01230/) (2022)
- [500xCompressor: Generalized Prompt Compression for Large Language Models](https://mlsanity.com/paper/2408.03094/) (2024)
- [LLMCad: Fast and Scalable On-device Large Language Model Inference](https://mlsanity.com/paper/2309.04255/) (2023)
- [Discrete Prompt Compression with Reinforcement Learning](https://mlsanity.com/paper/2308.08758/) (2023)
- [Context Embeddings for Efficient Answer Generation in RAG](https://mlsanity.com/paper/2407.09252/) (2024)
- [TACO-RL: Task Aware Prompt Compression Optimization with Reinforcement Learning](https://mlsanity.com/paper/2409.13035/) (2024)
- [PISCO: Pretty Simple Compression for Retrieval-Augmented Generation](https://mlsanity.com/paper/2501.16075/) (2025)
- [PyLate: Flexible Training and Retrieval for Late Interaction Models](https://mlsanity.com/paper/2508.03555/) (2025)

## Abstract (arXiv)

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.
