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

2025-07-24

AI Summary

This work addresses security risks in retrieval-augmented generation (RAG) systems arising from malicious document injection attacks. We propose a gradient-based toxicity detection mechanism that, for the first time, integrates gradient analysis with masked language modeling (MLM). Specifically, we compute gradients of the retriever’s similarity function with respect to input tokens to identify high-impact tokens; then leverage MLM prediction probability anomalies and token importance ranking to enable fine-grained identification and filtering of toxic documents. Experiments demonstrate that our method removes over 90% of injected toxic content across diverse attack scenarios while preserving over 95% recall of benign documents—significantly enhancing RAG robustness and output reliability. The core contribution lies in introducing interpretable gradient analysis into RAG security defense, enabling real-time, high-accuracy toxicity interception without model retraining and with minimal computational overhead.

Technology Category

• Natural Language Processing: Safety and Robustness
• Machine Learning: Adversarial Learning & Robustness
• Computer Vision: Adversarial Attacks & Robustness

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• Detect poisoned documents in RAG pipelines
• Filter adversarially crafted documents using GMTP
• Maintain retrieval accuracy while removing malicious content

Innovation

• Gradient-based high-impact token identification
• Masked token probability analysis via MLM
• High-precision poisoned document filtering