Safeguarding RAG Pipelines with GMTP: A Gradient-based Masked Token Probability Method for Poisoned Document Detection

TL;DR

This paper introduces GMTP, a gradient-based method to detect and filter poisoned documents in RAG systems, significantly improving security by identifying malicious content with high precision while preserving relevant information.

Contribution

The paper presents GMTP, a novel gradient-based approach for detecting poisoned documents in RAG pipelines, enhancing security without sacrificing retrieval quality.

Findings

GMTP detects over 90% of poisoned documents.

GMTP maintains high retrieval performance in adversarial settings.

GMTP effectively filters malicious content across diverse datasets.

Abstract

Retrieval-Augmented Generation (RAG) enhances Large Language Models (LLMs) by providing external knowledge for accurate and up-to-date responses. However, this reliance on external sources exposes a security risk, attackers can inject poisoned documents into the knowledge base to steer the generation process toward harmful or misleading outputs. In this paper, we propose Gradient-based Masked Token Probability (GMTP), a novel defense method to detect and filter out adversarially crafted documents. Specifically, GMTP identifies high-impact tokens by examining gradients of the retriever's similarity function. These key tokens are then masked, and their probabilities are checked via a Masked Language Model (MLM). Since injected tokens typically exhibit markedly low masked-token probabilities, this enables GMTP to easily detect malicious documents and achieve high-precision filtering. Experiments demonstrate that GMTP is able to eliminate over 90% of poisoned content while retaining relevant documents, thus maintaining robust retrieval and generation performance across diverse datasets and adversarial settings.