CleanBase: Detecting Malicious Documents in RAG Knowledge Databases
Weifei Jin, Xilong Wang, Wei Zou, Jinyuan Jia, Neil Gong
TLDR
CleanBase detects malicious documents in RAG knowledge bases by identifying semantically similar attack documents that form cliques.
Key contributions
- Introduces CleanBase, a novel method to detect prompt injection attacks in RAG systems.
- Leverages semantic similarity of malicious documents to construct a knowledge database graph.
- Identifies malicious documents by detecting cliques formed by highly similar documents in the graph.
- Provides theoretical guarantees and empirical validation across multiple datasets and attacks.
Why it matters
RAG systems are vulnerable to prompt injection attacks through malicious documents, compromising their integrity. CleanBase offers a crucial defense by accurately identifying these threats. This work significantly enhances the security and trustworthiness of RAG applications.
Original Abstract
Retrieval-augmented generation (RAG) is vulnerable to prompt injection attacks, in which an adversary inserts malicious documents containing carefully crafted injected prompts into the knowledge database. When a user issues a question targeted by the attack, the RAG system may retrieve these malicious documents, whose injected prompts mislead it into generating attacker-specified answers, thereby compromising the integrity of the RAG system. In this work, we propose CleanBase, a method to detect malicious documents within a knowledge database. Our key insight is that malicious documents crafted for the same attack-targeted questions often exhibit high semantic similarity, as attackers deliberately make them consistent to improve attack success rates. Accordingly, CleanBase constructs a similarity graph over the knowledge database, where each node represents a document and an edge connects two nodes if their semantic similarity--computed using an embedding model--exceeds a statistically determined threshold. Due to their inherent similarity, malicious documents tend to form cliques within this graph. CleanBase detects such cliques and flags the corresponding documents as malicious. We theoretically derive upper bounds on CleanBase's false positive and false negative rates and empirically validate its effectiveness. Experimental results across multiple datasets and prompt injection attacks demonstrate that CleanBase accurately detects malicious documents and effectively safeguards RAG systems. Our source code is available at https://github.com/WeifeiJin/CleanBase.
📬 Weekly AI Paper Digest
Get the top 10 AI/ML arXiv papers from the week — summarized, scored, and delivered to your inbox every Monday.