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

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.

## introduction large language models (llms) have significantly advanced performance across various natural language processing (nlp) tasks, especially in conversational systems and ai assistants @xcite @xcite . however, their reliance on parametric knowledge results in significant challenges, including susceptibility to hallucination @xcite and knowledge update requirements. to mitigate these issues retrieval-augmented generation (rag) has emerged as an effective approach for integrating external, non-parametric knowledge figure 1 : various corpus poisoning attacks on a naïve environment (without defense method) cause a significant performance drop in retrieval by enabling the poisoned documents to be retrieved. in contrast, gmtp effectively filters out poisoned documents, thereby preserving retrieval performance. yet, rag's reliance on external knowledge introduces vulnerabilities to corpus poisoning attacks @xcite @xcite . in such attacks, adversaries have access to the knowledge base, where they are able to inject maliciously crafted documents. this manipulation directly impacts the behavior of the generator, causing it to provoke misleading and harmful responses based on the tampered documents. this is because poisoned documents may contain incorrect information @xcite or adversarial instructions @xcite . this threat is particularly concerning for decision-making llm agents, such as those used in autonomous driving, where an attack could lead to erroneous actions with severe consequences @xcite . furthermore, attackers can optimize poisoned documents so that they appear highly relevant for only target queries @xcite , which makes their detection significantly more challenging. in this paper, we propose gradient-based masked token probability (gmtp), a novel method for filtering adversarially poisoned documents by analyzing the masked token probability of key tokens. specifically, the goal of gmtp is to filter out poisoned documents among those retrieved while maintaining retrieval performance, as can be seen in figure 1 . gmtp first identifies key tokens by examining gradient values derived from the retrieval phase, where tokens with high gradients significantly contribute to the similarity score. these key tokens are then masked, and the probability of correctly predicting them is assessed. since adversarially manipulated documents are optimized to match specific query patterns, they often contain unnatural text patterns that become difficult to reconstruct once masked. leveraging this property, gmtp effectively filters out poisoned documents that exhibit abnormally low masked token probabilities. our key contributions are as follows: ## proposed methodology in this section, we introduce gmtp, which identifies key tokens that contribute to high similarity scores between a query and document, while also detecting potential adversarial inputs by analyzing masked token probabilities. figure 3 illustrates the overall pipeline of rag using gmtp. specifically, gmtp identifies tokens with anomalously low masked token probabilities as potential adversarial indicators by leveraging the fact that cheating tokens exhibit unnatural linguistic patterns. we assume that the attacker can inject poisoned data into the knowledge base and can execute the retriever and generator but cannot retrain them. this assumption aligns with existing research on ragbased adversarial attacks @xcite @xcite . this assumption also reflects real-world scenarios, as many deployed rag systems rely on openly accessible databases such as wikipedia and often utilize third-party models like contriever @xcite and llama3 @xcite . ## results and analysis in this section, we evaluate the performance of gmtp against baseline models and analyze the factors contributing to its filtering effectiveness. due to space constraints, we primarily present results using the dpr model on the nq dataset, with the exception of the main result which includes all datasets. comprehensive results are available in appendix k. attack nq hotpotqa ms marco poisonedrag 0.859 0.860 0.865 phantom 0.901 0.901 0.932 advdecoding 0.910 0.954 0.786 ## conclusion in this study, we have proposed gmtp, a defense method designed to precisely filter poisoned documents in rag systems. by effectively capturing linguistic unnaturalness in poisoned documents, gmtp is able to successfully separate them from clean ones. experimental results demonstrated its strong filtering performance while minimizing retrieval degradation. furthermore, gmtp achieved over 90% successful filtering rate across various hyperparameter settings, demonstrating its robust adaptability to different rag systems. these findings highlight gmtp as a practical and adaptable defense against adversarial poisoning attacks.