A method for knowledge contamination detection and filtering in a retrieval enhancement generation system

Abstract

This paper proposes a method for detecting and filtering knowledge pollution in RAG (Retrieved Aggregate Algebra) systems, relating to the field of artificial intelligence security technology. It addresses the shortcomings of existing RAG knowledge pollution detection methods, such as neglecting structural features between documents and weak anti-interference capabilities. This method combines a dual-view feature fusion and robust graph learning detection framework, extracting pre-trained semantic vectors and TF-IDF lexical sparse features from the text. A weighted similarity matrix is ​​calculated to construct graph data capable of capturing both literal and semantic differences. GATv2 and KAN layers are used to deeply mine the nonlinear features of nodes within the residual connection structure. This method balances high detection robustness with inference efficiency and can be combined with a greedy filtering strategy to detect and filter knowledge pollution attacks on online RAG systems, providing high-precision security defense for RAG systems.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence security technology, specifically to defense technologies against knowledge pollution and retrieval-augmented generation (RAG) systems. Specifically, this invention discloses a knowledge pollution detection and filtering method based on the KAN augmented graph attention network (KAGATv2-Res), used to automatically identify and filter malicious text injected by PoisonedRAG-type attackers (including black-box text generation and white-box attacks) in the Top-k documents retrieved by the RAG system, thereby improving the robustness and security of downstream large-scale model (external large language model, LLM) question answering. Background Technology

[0002] With the widespread application of RAG system architecture in scenarios such as question answering, retrieval, and knowledge augmentation generation, attackers use knowledge pollution techniques to inject misleading text (malicious text) into the retrieval corpus / index, and induce LLM to generate incorrect answers through carefully crafted hints or context.

[0003] Existing defense methods mainly include keyword or pattern-based blacklist filtering, which is easily circumvented by adversarial examples and results in significant recall losses; and single-document classifier-based methods, which ignore the adjacency relationships between documents and group pattern (collective anomaly) information, making them prone to false positives or false negatives.

[0004] Therefore, there is an urgent need for a knowledge pollution detection and filtering method for retrieval enhancement generation systems that can simultaneously utilize document semantic representation and inter-document relationships (local neighbors and cross-class connections) to capture group patterns and structural anomalies. Summary of the Invention

[0005] This invention addresses the problems of existing RAG knowledge pollution detection methods, which only independently judge single documents, ignore semantic adjacency information, are prone to misjudgment, and have weak anti-interference capabilities. It provides a knowledge pollution detection and filtering method for retrieval enhancement generation systems.

[0006] A method for knowledge contamination detection and filtering in a retrieval enhancement generation system, comprising the following steps:

[0007] Step 1: Construct the training dataset;

[0008] Step 2: Construct semantic view features and lexical view features for the training dataset described in Step 1, and calculate the similarity matrix;

[0009] Step 3: Based on the similarity matrix described in Step 2, calculate the distance matrix and construct the graph data;

[0010] Step 4: Construct the KAGATv2-Res model and use the trained KAGATv2-Res model to output the corresponding binary classification score; the specific process is as follows:

[0011] Step 4: Input the graph data into the first GATv2 layer of the KAGATv2-Res model, calculate the attention coefficient between a node and its neighboring nodes through a dynamic attention mechanism, and aggregate the neighbor information to obtain the hidden layer features.

[0012] Step 42: Input the hidden layer features into the KAN layer to recursively generate Chebyshev polynomial terms, introduce learnable spectral gating parameters to weight and filter the polynomials, calculate the nonlinear output, and concatenate the nonlinear output with the features of the basic linear layer in the KAN layer to obtain the output features of the KAN layer.

[0013] Step 43: Use a linear projection layer to map the output features of the KAN layer to the same dimension as the output of the first GATv2 layer, and perform residual connections to obtain residual features;

[0014] Step 4: Aggregate the residual features using the second GATv2 layer, and calculate the final binary classification score using a fully connected layer.

[0015] The beneficial effects of this invention: The method described in this invention combines a robust detection framework of a Chebyshev KAN layer with spectral gating and a residual graph attention network (GATv2). KAN enhances the graph attention network by fusing semantic vectors and TF-IDF lexical features to construct a document relationship graph. A Chebyshev multinomial layer with spectral gating enhances nonlinear feature extraction, and Jacobi derivative regularization is introduced during training to improve model robustness. Finally, a greedy filtering strategy is used to dynamically clean the retrieved Top-K documents, effectively identifying and removing potentially malicious injected text, ensuring the security of the RAG system's generation. It has the following advantages:

[0016] The method described in this invention addresses the vulnerability of existing RAG systems to knowledge injection attacks (PoisonedRAG), and can effectively detect and filter malicious text while maintaining retrieval accuracy.

[0017] The method described in this invention fully considers the structural relationships between documents during the detection phase. By constructing a dual-view graph structure based on the fusion of semantic vectors and TF-IDF lexical features on the retrieved candidate document set, and using a weighted similarity matrix to capture the implicit literal structure and semantic pattern differences between documents, the accuracy of adversarial text recognition is significantly improved.

[0018] The method described in this invention not only adaptively aggregates neighbor information through a dynamic attention mechanism, but also enhances the expressive power of nonlinear features using Chebyshev polynomials. Combined with Jacobi derivative regularization training, it effectively smooths the decision boundary and resists malicious perturbations. At the same time, the architecture optimizes feature transfer through residual connections and linear projection, ensuring high detection robustness while taking into account inference efficiency. It can also be used in conjunction with a greedy filtering strategy to detect and filter knowledge pollution attacks on online RAG systems. Attached Figure Description

[0019] Figure 1 This is a flowchart of constructing the training dataset in this invention.

[0020] Figure 2 This is a flowchart of model training and detection inference based on the KAGATv2-Res model to prevent RAG knowledge contamination, as described in this invention.

[0021] Figure 3 The diagram shows the comparison of the model of this invention with existing models on three metrics when using black-box attacks.

[0022] Figure 4 The diagram shows the comparison of the model of this invention with existing models on three metrics when using white-box attacks. Detailed Implementation

[0023] Combination Figures 1 to 4 This embodiment describes a method for detecting and filtering knowledge pollution in a retrieval-enhanced generative system. In the feature construction stage, this method integrates semantic vectors from a Transformer-based sentence embedding model with lexical features from term frequency-inverse document frequency (TF-IDF) to construct a document relationship graph (graph data) capable of capturing literal structural differences. In the KAGATv2-Res model construction and training stage, RobustChebyLayer, integrating Chebyshev multinomials and spectral gating mechanisms, enhances the nonlinear representation capability of GATv2 and introduces a Jacobi derivative regularization term to optimize the loss function, thereby learning to distinguish the structural differences between the true semantic distribution and the malicious injection distribution. In the inference application stage, the trained model is used to perform binary classification on the dynamic subgraph constructed from the retrieved Top-K documents. Once a malicious node is identified, a greedy filtering strategy is executed to remove it and automatically fills the subsequent candidate pool sequentially. This process is repeated until a clean context containing K safe documents is generated, thereby achieving the detection and filtering of knowledge pollution attacks on the RAG system.

[0024] The method for detecting and filtering knowledge contamination in a retrieval enhancement generation system described in this embodiment specifically includes the following steps:

[0025] Step 1. Configure the API interface of the external large language model (external LLM); prepare the retrieval model (SentenceTransformer) for white-box optimization, load the retrieval model weights and deploy it to a computeable device (CPU / GPU); configure the attack parameters and optimization parameters for white-box text adversarial attacks (Hotflip);

[0026] Step 2. Construct a training dataset; the training dataset consists of the original corpus (original data) as positive sample data and malicious documents as negative sample data; such as Figure 1 As shown, the specific implementation process for constructing negative sample data is as follows:

[0027] Step 21. Read the input question set M (question-incorrect answer pairs) from the input text and restore the checkpoint, that is: restore the state at the time of the interruption by trying to read the last saved checkpoint file;

[0028] Step 22. Call an external LLM to generate candidate corpus (generated text, component I), verify the length of component I. If component I is not empty, concatenate the problem text Q with component I to generate a malicious document P for subsequent verification or white-box optimization, where P=Q⊕I, Q∈M;

[0029] Step 23. Determine whether to select black-box attack. If yes, set S=Q, then the malicious document P=S⊕I, and proceed to step 24; otherwise, load white-box attack, perform Hotflip iterative optimization, and generate combined text. Proceed to step 24;

[0030] Step 24. Decode the output natural text, use the decoded natural text as context, and input it together with the corresponding question text into an external LLM for question-and-answer verification; if the answer output by the external LLM contains the preset target answer, the verification is successful, save the result, update the statistics, and save the checkpoint; proceed to step 25; otherwise, retry, record the failure, save the checkpoint, and proceed to step 25.

[0031] Step 25. Determine if there are any unprocessed problem texts. If yes, return to step 21; otherwise, save the final result.

[0032] In step 23 of this embodiment, the specific process of white-box attack text performing Hotflip iterative optimization (forward, backward, replacement) is as follows:

[0033] Step 231. Use the I-component generated by the external LLM as the I-component for white-box attack; use the retrieval model to calculate the embedding vector of the question text;

[0034] Step 232. S component initialization: Construct the initial token sequence of the S component. The initial token sequence uses the sequence of the question text Q after word segmentation as the initial value (i.e., S=Q), or uses the sequence of masked tokens as the initial value.

[0035] Step 233. Forward propagation: Input the malicious document P=S⊕I into the retrieval model, calculate the embedding representation of the malicious document P, and calculate the cosine similarity between the embedding representation and the embedding vector of the problem text as the loss function;

[0036] Step 234. Backpropagation: Capture the embedding layer gradients of the retrieval model through the gradient storage mechanism, and summarize the gradient information at each position in the token sequence dimension;

[0037] Step 235. Calculate the dot product between the gradient information and the weight matrix of the embedding layer of the retrieval model, and select the top k candidate tokens with the largest dot product;

[0038] Step 236. Perform candidate replacement on component S to calculate similarity score: Select the token to be replaced in component S, replace the selected first k candidate tokens into the position in turn, and re-perform forward propagation on the combined text after replacement, calculate the cosine similarity between it and the embedding vector of the question text as the similarity score of the candidate token; select the candidate token that can maximize the similarity score, update the corresponding token in component S; obtain the optimized token sequence;

[0039] Step 237. Early stopping judgment: Stop optimization when the similarity score exceeds the threshold or when the maximum number of iterations is reached; finally, obtain the combined text. =S⊕I.

[0040] Step 3. Construct semantic view features and lexical view features for the training dataset constructed in Step 2, and extract semantic features using a Transformer-based sentence embedding model (e.g., all-MiniLM-L6-v2). Simultaneously, TF-IDF is used to extract lexical sparse features. Used to capture literal structural differences; the semantic features Lexical sparsity The nodes are concatenated to obtain the node feature matrix X;

[0041] Step 4. Calculate the semantic similarity matrix and word similarity matrix The final similarity matrix is ​​obtained through weighted fusion. :

[0042]

[0043] in, This represents the semantic weight coefficient.

[0044] Step 5. Based on the distance matrix D=1- The NearestNeighbors algorithm is used to find the k nearest neighbors for each node, and an edge index E containing bidirectional edges and self-loops is constructed to finally form graph data. .

[0045] Step 6. Train the KAGATv2-Res model, and use the trained KAGATv2-Res model to finally output the corresponding binary classification class Z; the specific process is as follows:

[0046] Step 61. Input the graph data G into the first GATv2 layer of the KAGATv2-Res model. This layer calculates the attention coefficient between node i and its neighbor node j through a dynamic attention mechanism, and aggregates the neighbor information to obtain the hidden layer features. It can be expressed as follows:

[0047]

[0048] In the formula, ELU is a nonlinear activation function. For the first The convolutional operation of the layers is used to aggregate features; and the Dropout operation is used to prevent overfitting.

[0049] Step 62. Transfer the hidden layer features obtained in Step 61 to... The input KAN layer is subjected to Tanh normalization, and then Chebyshev polynomials are recursively generated. Introducing learnable spectral gating parameters Weighted selection is performed on polynomials of different orders to calculate nonlinear output. :

[0050]

[0051] In the formula, The Chebyshev coefficient matrix is ​​used as the final nonlinear output. The features are concatenated with those of the base linear branch layer inside the KAN layer to obtain the output features of the KAN layer. ;

[0052] Step 63. Use a linear projection layer to extract the output features of the KAN layer. Mapped to the same dimension as the output of the first GATv2 layer, and residual connections are performed to preserve the original semantic information, resulting in the following residual features:

[0053]

[0054] In the formula, () is a linear projection layer used to enlarge the output feature dimension of the KAN layer to the same dimension as the output feature of the first layer GATv2.

[0055] Step 64. For the residual features A second GATv2 layer aggregation is performed, and the final binary classification score is calculated through a fully connected layer. :

[0056]

[0057] In the formula, () represents a linear classification layer, used to map the features aggregated by the second layer GATv2 to the final classification logical value; These correspond to the scores for nodes belonging to the normal and malicious categories, respectively.

[0058] In this implementation, during the model training phase, a total loss function incorporating classification cross-entropy and Jacobian derivative regularization is constructed to smooth the decision boundary and enhance the model's robustness to adversarial perturbations; an automatic differentiation mechanism is used to calculate the Jacobian derivative of the KAN layer output with respect to the input. Construct the total loss function. The total loss function includes classification cross-entropy loss. And derivative regularization terms to smooth the decision boundary and improve robustness, total loss function as follows:

[0059]

[0060] In the formula, () represents the classification cross-entropy loss function; For regularization weight hyperparameters, This is a real label; () is used to find the maximum value;

[0061] Step 7. Use the KAGATv2-Res model trained in Step 6 for detection and inference; combine a greedy filtering strategy to perform real-time cleaning and automatic completion of the retrieved Top-K documents, effectively identifying and blocking malicious text injection, providing high-precision security defense for the RAG system, thereby realizing the detection and filtering of knowledge pollution attacks in the RAG system; the specific process is as follows:

[0062] Step 71. Retrieve a preliminary Top-K candidate document set (including raw data and potentially malicious injection text) for the user query. Sort the documents in descending order of search score to construct a candidate pool;

[0063] Step 72. For the currently selected subset of candidate documents, construct an inference subgraph in real time. This inference subgraph has the same structure as the graph data in Step 5. Furthermore, ensure that the semantic weight coefficient α during inference and the graph construction logic are strictly consistent with those during training to defend against adversarial text.

[0064] Step 73. Input the inference subgraph into the trained KAGATv2-Res model, perform forward propagation (without calculating derivatives), and obtain the predicted category for each candidate document:

[0065]

[0066] like If the value is 1, the document is determined to be malicious text.

[0067] Step 74. Execute the greedy filtering strategy to ultimately detect and filter knowledge pollution attacks in the RAG system.

[0068] In this embodiment, the specific process of implementing the greedy filtering strategy is as follows:

[0069] Step A1. Select documents from the candidate pool in descending order of score;

[0070] Step A2. Determine whether the selected document is a malicious document. If so, remove it from the selected list and record it as filtered; automatically select the next document that has not been marked as malicious from the candidate pool to fill the gap; otherwise, keep the document and proceed to step A3.

[0071] Step A3. Determine whether all TOP-K candidate documents are safe candidate documents. If so, proceed to step A4; otherwise, return to step A1.

[0072] Step A4. Input the K highest-scoring filtered security candidate documents and the user query as the context into the LLM to generate the final user response.

[0073] Combination Figure 3 and Figure 4 This embodiment describes how the model of the present invention (KAGATv2-Res) is compared with existing RAG layered defense models (RAGuard), RAG lightweight defense models (RAGDEFENDER), TrustRAG enhanced RAG models (TrustRAG), and RAG forensics models (RAGForensics) in terms of attack success rate (ASR), false positive rate (AKR), and false negative rate (MDR). The KAGATv2-Res model described in this invention has the lowest values ​​in all three indicators, and can most effectively reduce attack success rate, reduce knowledge leakage, and filter malicious documents. It is the most secure RAG defense model in this evaluation.

[0074] The technical features of the above embodiments can be combined in any way. For the sake of brevity, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.

[0075] The embodiments described above are merely illustrative of several implementations of the present invention, and while the descriptions are relatively specific and detailed, they should not be construed as limiting the scope of the invention patent. It should be noted that those skilled in the art can make various modifications and improvements without departing from the concept of the present invention, and these all fall within the protection scope of the present invention. Therefore, the protection scope of this invention patent should be determined by the appended claims.

Claims

1. A method for retrieval augmentation generation system knowledge pollution detection and filtering, characterized in that: This method is implemented by the following steps: Step 1: Construct the training dataset; The training dataset consists of original data as positive sample data and malicious documents as negative sample data; the specific process for constructing negative sample data is as follows: Step 11: Load the retrieval model for white-box optimization and configure the attack parameters for white-box text adversarial attacks. Step 12: Read the input question set M from the input text, call the external LLM to generate component I, and concatenate the question text Q with component I to generate malicious document P; Step 13: Determine whether to select black-box attack. If yes, set component S=Q, generate malicious document P=S⊕I, and proceed to Step 14; otherwise, load white-box attack, perform Hotflip iterative optimization, and generate combined text. Perform step one four; Step 14: Decode the output natural text, call an external LLM for verification, and when the verification is successful, output and save the result, then return to Step 12; when all question texts have been processed, save the final result as negative sample data for the training dataset. Step 2: Construct semantic view features and lexical view features for the training dataset described in Step 1, and calculate the similarity matrix; Step 3: Based on the similarity matrix described in Step 2, calculate the distance matrix and construct the graph data; Step 4: Construct the KAGATv2-Res model and use the trained KAGATv2-Res model to output the corresponding binary classification score; The specific process is as follows: Step 4: Input the graph data into the first GATv2 layer of the KAGATv2-Res model, calculate the attention coefficient between a node and its neighboring nodes through a dynamic attention mechanism, and aggregate the neighbor information to obtain the hidden layer features. Step 42: Input the hidden layer features into the KAN layer to recursively generate Chebyshev polynomial terms, introduce learnable spectral gating parameters to weight and filter the polynomials, calculate the nonlinear output, and concatenate the nonlinear output with the features of the basic linear layer in the KAN layer to obtain the output features of the KAN layer. Step 43: Use a linear projection layer to map the output features of the KAN layer to the same dimension as the output of the first GATv2 layer, and perform residual connections to obtain residual features; Step 4: Aggregate the residual features using the second GATv2 layer, and calculate the final binary classification score using a fully connected layer.

2. The method for detecting and filtering knowledge contamination in a retrieval enhancement generation system according to claim 1, characterized in that: In step one of steps, the specific process of performing Hotflip iterative optimization is as follows: Component I generated by an external LLM is used as component I in a white-box attack; the question text Q is segmented and encoded using a retrieval model, and the question text embedding vector is calculated. Construct an initial token sequence for component S, wherein the initial token sequence is the sequence of the question text Q after word segmentation as the initial value; Input the malicious document P=S⊕I into the retrieval model, calculate the embedding representation of the malicious document P, calculate the cosine similarity between the embedding representation and the embedding vector of the problem text, and use the cosine similarity as the optimization target. Perform backpropagation on the optimization objective to obtain the gradient information of the retrieval model embedding layer with respect to the position of each token in component S; Based on the dot product of the gradient information and the word vectors in the embedding layer of the retrieval model, the top k candidate tokens are selected for the current position of the token to be replaced. Select the top k candidate tokens to replace the tokens to be replaced in component S, and calculate the cosine similarity between the replaced combined text and the embedding vector of the question text. Select candidate tokens with high similarity scores to update the corresponding tokens in component S; Obtain the optimized token sequence; when the similarity score exceeds a set threshold, obtain the combined text. =S⊕I.

3. The method for knowledge contamination detection and filtering in a retrieval enhancement generation system according to claim 1, characterized in that: In step two, a Transformer-based sentence embedding model is used to extract semantic features; TF-IDF is used to extract lexical sparse features, and the semantic features and lexical sparse features are concatenated to obtain a node feature matrix. .

4. The method for knowledge contamination detection and filtering in a retrieval enhancement generation system according to claim 3, characterized in that: In step two, the semantic similarity matrix is ​​calculated. and word similarity matrix The final similarity matrix is ​​obtained through weighted fusion. It can be expressed as follows: ； In the formula, This represents the semantic weight coefficient.

5. The method for knowledge contamination detection and filtering in a retrieval enhancement generation system according to claim 4, characterized in that: Calculate the distance matrix D=1- For each node, find k nearest neighbors, construct an edge index E including bidirectional edges and self-loops, and finally form graph data. .

6. The method for knowledge contamination detection and filtering in a retrieval enhancement generation system according to claim 5, characterized in that: In step four, during the training phase of the KAGATv2-Res model, a total loss function is constructed, including a classification cross-entropy loss function and a Jacobian derivative regularization term, to smooth the decision boundary. An automatic differentiation mechanism is then used to calculate the Jacobian derivative of the KAN layer output with respect to the input. ;in, The output features of the KAN layer The hidden layer features are the output of the first GATv2 layer; The total loss function as follows: ； In the formula, For binary classification scores, () represents the classification cross-entropy loss function; For regularization weight hyperparameters, This is a real label; () is used to find the maximum value of the Jacobian derivative.

7. A method for detecting and filtering knowledge contamination in a retrieval enhancement generation system according to any one of claims 1-6, characterized in that: The pollution detection and filtering method also includes detecting and inferring the trained KAGATv2-Res model, and using a greedy filtering strategy to clean and automatically fill in the retrieved Top-K documents in real time.

8. The method for detecting and filtering knowledge contamination in a retrieval enhancement generation system according to claim 7, characterized in that: For user queries, retrieve the initial Top-K candidate document set, sort them in descending order of retrieval score, and construct a candidate pool; For the currently selected subset of candidate documents, an inference subgraph is constructed in real time. This inference subgraph is then input into the trained KAGATv2-Res model to obtain the predicted category for each candidate document. ； like If the value is 1, the document is determined to be malicious text. By implementing a greedy filtering strategy, the knowledge pollution attacks in the RAG system can be detected and filtered.

9. A method for detecting and filtering knowledge contamination in a retrieval enhancement generation system according to claim 8, characterized in that: The specific process of the greedy filtering strategy is as follows: Documents are selected sequentially from the candidate pool; If the current document is detected as malicious, it will be removed from the selected list and recorded as filtered. The next unmarked document is automatically selected from the candidate pool to fill the gap, and the detection process is repeated until K safe documents are selected. The K security documents are used as contextual inputs to an external LLM to generate the final user response, thereby enabling the detection and filtering of knowledge contamination in the RAG system.

Images