CleanBase: Detecting Malicious Documents in RAG Knowledge Databases

TL;DR

CleanBase is a novel method that detects malicious documents in RAG knowledge bases by identifying highly similar document clusters, thereby preventing prompt injection attacks.

Contribution

It introduces a similarity graph approach to detect malicious documents based on their high semantic similarity and clique formation, with theoretical and empirical validation.

Findings

CleanBase accurately detects malicious documents across multiple datasets.

The method effectively prevents prompt injection attacks in RAG systems.

Theoretical bounds on false positive and false negative rates are established.

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.