A Memory-Efficient Retrieval Architecture for RAG-Enabled Wearable Medical LLMs-Agents

TL;DR

This paper introduces a hierarchical retrieval architecture for edge RAG in medical LLMs that reduces energy and memory use by nearly 50% and 75% respectively, while maintaining accuracy, suitable for privacy-preserving medical AI agents.

Contribution

It proposes a novel two-stage retrieval scheme for edge RAG that significantly improves energy efficiency and reduces memory access in medical LLM applications.

Findings

Memory access reduced by nearly 50%.

Energy consumption decreased by 75%.

Maintains high retrieval accuracy.

Abstract

With powerful and integrative large language models (LLMs), medical AI agents have demonstrated unique advantages in providing personalized medical consultations, continuous health monitoring, and precise treatment plans. Retrieval-Augmented Generation (RAG) integrates personal medical documents into LLMs by an external retrievable database to address the costly retraining or fine-tuning issues in deploying customized agents. While deploying medical agents in edge devices ensures privacy protection, RAG implementations impose substantial memory access and energy consumption during the retrieval stage. This paper presents a hierarchical retrieval architecture for edge RAG, leveraging a two-stage retrieval scheme that combines approximate retrieval for candidate set generation, followed by high-precision retrieval on pre-selected document embeddings. The proposed architecture significantly reduces energy consumption and external memory access while maintaining retrieval accuracy. Simulation results show that, under TSMC 28nm technology, the proposed hierarchical retrieval architecture has reduced the overall memory access by nearly 50% and the computation by 75% compared to pure INT8 retrieval, and the total energy consumption for 1 MB data retrieval is 177.76 {\mu}J/query.