A Lightweight Retrieval-Augmented Generation System for Mobile Devices

by

Retrieval-Augmented Generation (RAG) combines a language model (LLM) with a retrieval system to improve response accuracy by incorporating relevant external knowledge. Deploying RAG on mobile devices presents unique challenges, such as memory constraints, limited processing power, and the need for efficiency.

In this blog, I want to build a lightweight RAG system optimized for mobile devices using a small LLM and an efficient retrieval mechanism.

RAG on Mobile device architecture
Image by Author – RAG for Mobile devices

Table of Contents

Retrieval-Augmented Generation – RAG

RAG, while not a recent innovation, has matured significantly. Let’s quickly recap its core principles before diving into the exciting frontier of implementing RAG on mobile devices.

Retrieval-Augmented Generation (RAG) is an approach that cumulates the strengths of retrieval and generation models. It is particularly useful for tasks that involve generating answers based on large amounts of knowledge, even extending beyond the training data available to the model. The two main components of RAG are Retrieval model and Augmentation model.

Retrieval Component

This part of the RAG model searches a knowledge base or document store to fetch relevant information. Typically, this involves using a dense passage retrieval (like using a neural search model) or a more traditional method like BM251 to identify documents that can assist in answering the query.

Generator Component

This component generates an answer using both the input query and the retrieved documents. A model like a transformer-based sequence generator (e.g., BART2, T53) is often used.

Technical Challenges of Running RAG Models on Mobile

We all know mobile devices have constraints related to computational power, memory and network connectivity. While RAG is an efficient way to enhance language models, deploying a fully standalone RAG system on a mobile device faces several limitations that comes from the limited resources of mobile devices.

  1. Large Storage Requirements: A high-quality vector database1 and LLM require gigabytes of storage, making it difficult to deploy on mobile without cloud support.
  2. Limited Memory and RAM: Mobile devices typically have limited RAM, making it challenging to load both the retrieval and generation models into memory simultaneously.
  3. Energy Consumption: LLM inference and retrieval operations consume battery power rapidly, reducing usability.
  4. Networking Constraints: If the mobile device requires cloud access for retrieval, network latency can impact response times.

In addition to the inherent challenges that mobile devices present, LLMs are built on deep neural network architectures, typically Transformers, which rely on billions (or even trillions) of parameters. The computational and memory demands arise from several key factors.

Large Number of Parameters

  • An LLM like GPT-4 or LLaMA-3 has hundreds of billions of parameters, each requiring storage in floating-point precision (e.g., FP16 or FP32).
  • The weight matrices of these models alone take up terabytes of memory.


For instance, Gemini Ultra (Google DeepMind) has over 1T(1 trillion) model weights, which supports advanced multimodal reasoning.

Here is the breakdown of memory usage for weightS with FP16 vs FP32.

FP16 – Half Precision Floating Point,16bit. It uses 16bits per number or 2 bytes.

1 billion * 2bytes = 2GB

FP32 – Single Precision Floating Point, 32 bit. It uses 32-bits per number or 4-bytes.

1 billion * 4bytes = 4GB

High-Dimensional Word Embeddings

Word embeddings comes in retrieval phase, where each and every document that is fed to vector DB is converted into a high dimensional vector of numbers.

  • Each token (word or subword) in a model is converted into a high-dimensional vector (e.g., 4096-dimensional in GPT-4).
  • Storing and processing these embeddings during training and inference requires substantial memory.

Transformer Architecture and Attention Mechanism

The Transformer architecture is built upon attention mechanisms. Attention’s quadratic memory complexity with respect to sequence length is a significant memory consideration. Model size and batch size also play crucial roles in memory requirements.

Factors that influence memory

Sequence length: Longer sequences lead to a quadratic increase in memory usage in the attention layers.

Model size: Larger models, with more layers and hidden units, require more memory to store parameters and intermediate activations.
Batch size: Larger batch sizes increase memory usage during training.

Training vs. Inference Memory Requirements

Training Phase (Most Expensive)
  • Training LLMs require huge GPU clusters(TPUs/A100/H100 GPUs)5.
  • Back-propagation stores intermediate activations, gradients, and optimizer states.
  • Memory needs grow due to batch size, optimizer states (e.g., Adam), and model parallelism.
Inference Phase (Still Heavy but Less Than Training)
  • Pre-trained models still require large memory because model weights must be stored in GPU memory.
  • The KV (key-value) cache for attention grows with sequence length

System Architecture – RAG for mobile devices

A mobile-efficient RAG system consists of

– Local vector databases as FAISS or SQLite-based search

– Compact LLM as DistilGPT, MobileBERT, or an on-device TFLite model

– Efficient Tokenization (to manage input size effectively).

– Retrieval & Ranking Module (to fetch relevant documents before generation).

Tech Stack

  • sentence-transformers for embedding generation.
  • faiss for local retrieval.
  • sqlite3 for persistent vector storage.
  • transformers for a lightweight LLM.
  • TFLite for model compression and mobile deployment.
  • Python and Flask for a mobile-accessible backend.

Efficient Tokenization

As mobile devices are limited on resources front, every measure is needed to ensure that the input sequence length does not exceed the mobile device’s processing limits. Efficiency in tokenization can be achieved by

  1. Truncating Long Inputs – Many text inputs exceed the model’s maximum token length. Instead of processing the entire input, the tokenizer selects the most relevant sections, such as the beginning and end of a passage or high-scoring key sentences, ensuring that crucial information is retained within the allowed limit.
  2. Byte-Pair Encoding (BPE) – Traditional tokenization splits text into words, but BPE efficiently breaks words into smaller subword units. This approach helps in handling unknown words and reduces the vocabulary size, making it ideal for mobile deployments. For example, “unhappiness” may be split into “un”, “happi”, and “ness”, allowing the model to understand word variations without a huge vocabulary. Whole idea is to have limited vocabulary that could span to more word semantics.
  3. Filtering Stop-words – Common words like “is,” “the,” and “and” do not add significant meaning in most NLP tasks. Removing them reduces token count, making processing more efficient without losing the core message of the input text.

Retrieval and Ranking Module

Retrieval fetches the most relevant documents based on the query. Ranking ensures the retrieved documents are prioritized based on relevance.

Compute Query Embeddings: Generate vector embeddings for the user query.

Perform Nearest Neighbor Search: Find the top-k similar embeddings from the database. This is often done by using cosin similarity3.

Apply Reranking: Reorder results based on semantic similarity.

Compact LLMs

Compact LLMs are similar versions of large language models optimized for efficiency while maintaining reasonable performance. Compared to traditional LLMs, they have fewer parameters, transfer learning or knowledge distillation, optimized architectures, quantization, pruning and low latency.

Common compact LLMs include:

  • DistilGPT-2: A distilled version of GPT-2 with nearly half the parameters.
  • MobileBERT: A lightweight adaptation of BERT optimized for mobile.
  • TinyLlama: A reduced-size version of LLaMA with efficient inference.
  • ALBERT: A parameter-efficient version of BERT with factorized embeddings.

Github Code

This code is a starting point. The rich external information leads to more accurate generation. This can be extended to build sophisticated applications and can also deploy as mobile apps.

RAG For Mobile Devices

Summary

It is interesting to work on building a lightweight Retrieval-Augmented Generation (RAG) system for mobile and is exciting to figure out how to make large models and retrieval systems work on devices with limited resources. Mobile devices come with their own set of challenges, especially when it comes to running large models. LLMs are intrinsically heavy, so trying to bridge the gap and make them work efficiently on these devices is truly fascinating.

It’s a new learning journey, and RAGs are an underlying technology for intelligent agents(AI agents) that can greatly enhance how we interact with devices. Here, I’ve tried to show how to use a compact LLM, persistent local vector storage (FAISS & SQLite), and smart retrieval techniques to bring RAG to mobile devices. I am hopeful to extend further on this in future!

References

  1. BM25 – Best Matching 25, ranking function used in search engines and information retrieval systems to find and rank the most relevant documents based on a given query. ↩︎
  2. BART – Unlike BERT (only encoder) or GPT (only decoder), BART uses both an encoder and a decoder, making it powerful for tasks requiring input transformation (e.g., summarization and translation) ↩︎
  3. T5 – Text-to-Text Transfer Transformer is a sequence-to-sequence (seq2seq) transformer model. It treats every NLP task as a text-to-text problem, where input and output are text strings. ↩︎
  4. Vector Database ↩︎
  5. TPUs/A100/H100 GPUs
    – H100 GPUs – NVIDIA’s next-gen AI GPU, successor to A100.
    – A100 GPUs – NVIDIA’s powerful deep learning accelerator.
    – TPUs (Tensor Processing Units) – Developed by Google ↩︎
  6. Cosine similarity – It is a metric used to measure how similar two vectors are, based on the cosine of the angle between them. It is commonly used in text retrieval, recommendation systems, and vector databases like FAISS ↩︎