A retrieval augmented generation system knowledge poisoning attack defense method

Abstract

The application discloses a retrieval enhanced generation system knowledge poisoning attack defense method, relates to the technical field of artificial intelligence security and graph neural networks, and solves the problems that existing defense methods depend on isolated detection of single document content, ignore the structural relationship between documents and the interactive context of documents and user queries, and thus it is difficult to capture high retrieval relevance but low semantic consistency of malicious documents. The application models queries and retrieval documents as a heterogeneous graph structure, fuses semantic embedding and graph structure features, constructs a global graph containing query-document edges and document-document similarity edges, classifies document nodes through an edge type-aware graph attention network, and identifies malicious injected documents. The application improves the defense capability of RAG systems against knowledge poisoning attacks and provides safe and reliable external knowledge input for large language models.

Description

Technical Field

[0001] This invention relates to the fields of artificial intelligence security and graph neural network technology, specifically to a method for defending against knowledge poisoning attacks on retrieval-enhanced generative systems. Background Technology

[0002] Retrieval-enhanced generation (RAG) systems improve factual accuracy by retrieving documents from external knowledge bases to assist large language models in generating responses. However, attackers can inject a small number of malicious documents into the knowledge base to manipulate the retrieval results, causing the model to generate false or harmful content. This type of attack is known as knowledge poisoning.

[0003] Existing defense methods primarily rely on isolated detection of individual document content, such as judging document authenticity based on text classifiers or comparing multiple sources through consistency checks. However, these methods ignore the structural relationships between documents and the interactive context between documents and user queries, making it difficult to capture the typical abnormal pattern of malicious documents with high retrieval relevance but low semantic consistency.

[0004] Therefore, there is an urgent need for a dynamic defense mechanism that can jointly model query-document-document multi-source relationships and integrate semantic and structural features to achieve efficient and robust identification of knowledge poisoning attacks. Summary of the Invention

[0005] To address the problem that existing defense methods rely on isolated detection of individual document content, ignoring the structural relationships between documents and the interactive context between documents and user queries, which leads to difficulties in capturing malicious documents with high retrieval relevance but low semantic consistency, this invention provides a knowledge poisoning attack defense method for retrieval enhancement generation systems. By constructing a heterogeneous graph model and fusing semantic and structural features, it achieves high-precision identification of malicious documents.

[0006] A method for defending against knowledge poisoning attacks in a retrieval-enhanced generation system, comprising the following steps:

[0007] Step S1: Collect labeled retrieval data under known attack scenarios. Each sample in the retrieval data includes a user query document and a corresponding set of retrieval documents.

[0008] Step S2: Encode the user query document and each document in the retrieved document set using a pre-trained sentence encoding model to obtain the query semantic embedding vector and the document semantic embedding vector;

[0009] Step S3: Construct a global heterogeneous graph based on the query semantic embedding vector and document semantic embedding vector obtained in step S2;

[0010] Step S4: Calculate the four-dimensional structural feature vector for each document node in the global heterogeneous graph;

[0011] Step S5: Concatenate the document semantic embedding vector with the four-dimensional structural feature vector to form the initial feature representation of each node; use an edge type-aware graph attention network model to encode the global heterogeneous graph to obtain the embedding representation of all document nodes; the graph attention network model integrates edge attribute and edge type embedding information, and realizes adaptive aggregation of document node features through a multi-head attention mechanism; and calculate the probability of each document.

[0012] Step S6: Set a decision threshold and identify documents with a probability higher than the decision threshold as malicious documents; input the remaining documents as context into the large language model to generate the final answer; achieve effective defense against knowledge poisoning attacks in the retrieval enhancement generation system.

[0013] The beneficial effects of this invention are:

[0014] The knowledge poisoning attack defense method described in this invention models the RAG retrieval results as a global heterogeneous graph containing cross-query edges; it integrates semantic and refined graph structure features to significantly improve the detection capability of covert poisoning samples.

[0015] The knowledge poisoning attack defense method described in this invention supports real-time deployment and is applicable to actual RAG security protection scenarios. This method improves the RAG system's ability to defend against knowledge poisoning attacks and provides safe and reliable external knowledge input for large language models.

[0016] The knowledge poisoning attack defense method described in this invention is compared with existing RAG defense methods and trusted RAG detection methods in terms of false positive rate (FPR), accuracy rate (DACC), and false negative rate (FNR). The detection method of this invention is significantly superior to existing methods. Attached Figure Description

[0017] Figure 1 This is a flowchart of a method for defending against knowledge poisoning attacks in a retrieval enhancement generation system, as described in this invention.

[0018] Figure 2 This is a comparison diagram of the present invention and two existing methods in DACC.

[0019] Figure 3 This is a comparison diagram of the present invention and two existing methods on FPR.

[0020] Figure 4 This is a comparison diagram of the present invention and two existing methods on FNR. Detailed Implementation

[0021] Combination Figures 1 to 4 This embodiment describes a method for defending against knowledge poisoning attacks in a retrieval enhancement generation system. The method comprises the following steps:

[0022] Step 1: Collect labeled retrieval data under known attack scenarios and construct the original sample set S;

[0023] The original sample set includes several samples, each sample including a user query document q and its corresponding set of search documents. Each document is labeled with a binary tag. ; 0 represents a normal document, and 1 represents a malicious document; at the same time, the original search score assigned by the searcher to each document is recorded; the original sample set S is randomly divided into training set, validation set and test set according to the proportion, with the training set accounting for 70%, and the validation set and test set each accounting for 15%.

[0024] Step 2: Encode the user query document q and the documents in the retrieved document set D using a pre-trained sentence encoding model (All-mpnet-base-v2) to obtain the query semantic embedding vector. and document semantic embedding vector It has 768 dimensions;

[0025] Step 3: Construct a global heterogeneous graph based on the query semantic embedding vector and document semantic embedding vector obtained in Step 2;

[0026] In this embodiment, each query-document heterogeneous graph contains one query node and multiple document nodes. Directed edges are added from the query node to each document node, with the edge attribute being the normalized retrieval score. Among the document nodes, the similarity between each document and all other documents is calculated. The top k most similar documents with a similarity greater than a preset threshold are selected, and directed edges are added from these document nodes to other document nodes that meet this requirement. The edge attribute is the cosine similarity value. These two types of edges respectively represent retrieval relevance and document semantic association, and the edge type identifier supports the graph neural network model for differentiated modeling of multi-source relationships.

[0027] In this embodiment, two types of semantically heterogeneous directed edges are introduced into the heterogeneous graph to model multi-source relationships; the first type of edge represents the retrieval relevance between query documents, and the second type of edge represents the semantic association between documents; and by assigning distinguishable type labels to edges of different categories, the graph neural network can perform differentiated aggregation of relationships from different sources during message passing. The specific process is as follows:

[0028] Step 31: Initialize the graph data list The first sample is selected from the original sample set S as the current sample for processing.

[0029] Step 32: Construct the query-document heterogeneous graph for the current sample; add query nodes and document nodes, and add directed edges from the query nodes to each document node;

[0030] Step 33: Construct semantic association edges between document nodes; compute the current document. With other documents The cosine similarity; if the similarity is greater than a preset threshold. And the remaining documents Belongs to the current document If the first k most similar neighbors are found, then add them from the current document. Point to other documents The directed edges are assigned a cosine similarity value. The cosine similarity is calculated using the following formula:

[0031]

[0032] In the formula, and Document nodes and document nodes semantic embedding vector, It is a norm.

[0033] Step 34: Add the currently constructed query-document heterogeneity graph to the graph data list. Select the next sample from the original sample set S. Determine if the sample exists. If it does, return to step 32; otherwise, proceed to step 35.

[0034] Step 35: List the graph data The query-document heterogeneity graphs of all samples are merged into a global heterogeneity graph. ;in, For querying the set of nodes, Let E be the set of document nodes, and let E be the set of edges (including query-document edges and document-document edges).

[0035] Step 4: For the document node set For each document node, calculate its four-dimensional structural feature vector. ( The four-dimensional structural features include the proportion of reciprocal edges, the average similarity between neighbors, the local clustering coefficient, and the retrieval score, which are defined as follows:

[0036] The reciprocal edge ratio: the proportion of outgoing edges of a node that have reverse edges (i.e., bidirectional connections); the calculation formula is:

[0037]

[0038] in, For document nodes in the heterogeneous graph; This is the set of directed edges between document nodes; For document nodes The degree of departure; For document nodes The out-neighbor set, i.e. ; It is a very small positive number, and the denominator is guaranteed to be non-zero.

[0039] The average neighbor similarity is the mean of the edge attributes (i.e., similarity scores) corresponding to all outgoing edges of the node; the calculation formula is:

[0040]

[0041] in, For document nodes With document nodes Cosine similarity between them;

[0042] Local clustering coefficient: In an undirected subgraph composed of reciprocal edges, the ratio of the actual number of connections between a node's neighbors to the theoretical maximum number of connections; the formula is: Let For reciprocal neighbors to gather, record .

[0043]

[0044] in, For a set of mutually beneficial edges (bidirectional edges), For document nodes The set of reciprocal neighbors, i.e., with document nodes The set of document nodes that have reciprocal edges, i.e. ;in , Represents a document node. For the document nodes in the heterogeneous graph; if <2, then =0.

[0045] Search score: The search score for this node; Calculation formula:

[0046]

[0047] in: For query text semantic embedding vector; For document semantic embedding vector; This represents the dot product operation of vectors.

[0048] Step 5: Concatenate the 768-dimensional document semantic embedding vector with the 4-dimensional structural feature vector to form the initial feature representation X of each document node, with a dimension of 772; initialize the edge type-aware graph attention network model (GATv2) and the fully connected classification layer, and set the maximum number of training epochs. And the current round t=1;

[0049] Step 6: Perform the t-th round of forward propagation;

[0050] The global heterogeneous graph corresponding to the training set is input into the graph attention network model. The network model uses a multi-head attention mechanism to fuse edge attribute and edge type information, adaptively aggregates the initial node features X, calculates the new embedding representation of the nodes, and outputs the predicted probability that each document belongs to malicious documents through a fully connected classification layer. ;

[0051] Step 7: Calculate the loss and update the parameters;

[0052] The Focal Loss function is used to calculate the predicted probabilities on the training set. The loss value between the model and the true label y is used to update the model parameters using the backpropagation algorithm;

[0053] Step 8: Perform verification and early stop judgment;

[0054] Evaluate the current network model performance (F1 score) on the validation set. Determine if the current round t is less than... If the early stopping condition is not triggered (i.e., the validation set performance does not improve within consecutive Patience rounds), then let t = t + 1 and return to step 6; otherwise, end the training and save the optimal network model parameters, and proceed to step 9.

[0055] Step 9: Set the decision threshold search interval , This is the lower limit of the decision threshold. Set the upper limit of the decision threshold; select the current threshold. ;

[0056] Step 10: Calculate the current threshold based on the validation set. The classification F1 score is calculated, and the optimal threshold corresponding to the best F1 score is recorded. ;

[0057] Step 11: Select the next threshold (For example + 0.01), judge Is it less than or equal to? If so, let = If the condition is met, proceed to step 10; otherwise, proceed to step 12.

[0058] Step 12: Load the optimal network model parameters and optimal decision threshold. Inference is performed on the retrieved documents in the test set. If the malicious probability output by the network model is greater than... If so, it is determined to be a malicious document;

[0059] Step 13: Input the filtered remaining documents as context into the large language model (GPT3.5) to generate the final answer, thus completing the defense against knowledge poisoning attacks.

[0060] like Figures 2 to 4 As shown, Figure 2 In this invention, the DACC of the method is 0.994, which is higher than that of the existing RAG defense method (RAGuard) (0.984) and the Trust RAG detection method (TrustRAG) (0.941). Figure 3 In this invention, the FPR of the method is 0.008, which is significantly lower than RAGuard's 0.099 and TrustRAG's 0.686; Figure 4 In this study, the FNR of the method of the present invention is 0.002, while RAGuard and TrustRAG are 0.059 and 0.016, respectively. Therefore, the detection method of the present invention is significantly superior to the two existing methods.

[0061] 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.

[0062] 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 defending against knowledge poisoning attacks in a retrieval-enhanced generation system, characterized by: This method is implemented by the following steps: Step S1: Collect labeled retrieval data under known attack scenarios. Each sample in the retrieval data includes a user query document and a corresponding set of retrieval documents. Step S2: Encode the user query document and each document in the retrieved document set using a pre-trained sentence encoding model to obtain the query semantic embedding vector and the document semantic embedding vector; Step S3: Based on the query semantic embedding vector and document semantic embedding vector obtained in step S2, construct a global heterogeneous graph; the specific process is as follows: Step S31: Initialize the graph data list The first sample is selected from the original sample set S as the current sample for processing. Step S32: Construct the query-document heterogeneity graph for the current sample; Add query nodes and document nodes, and add directed edges from the query node to each document node; Step S33: Construct semantic association edges between document nodes to obtain the query-document heterogeneous graph of the current sample; Step S34: Add the query-document heterogeneity graph of the current sample to the graph data list. Select the next sample in the original sample set S, determine if the sample exists, and if it does, return to step S32. Otherwise, proceed to step S35; Step S35: List the graph data The query-document heterogeneity graphs of all samples are merged into a global heterogeneity graph. ;in, For querying the set of nodes, Let E be the set of document nodes and E be the set of edges; Step S4: Calculate the four-dimensional structural feature vector for each document node in the global heterogeneous graph; Step S5: Concatenate the document semantic embedding vector with the four-dimensional structural feature vector to form the initial feature representation of each node; use an edge type-aware graph attention network model to encode the global heterogeneous graph to obtain the embedding representation of all document nodes; the graph attention network model integrates edge attribute and edge type embedding information, and realizes adaptive aggregation of document node features through a multi-head attention mechanism; and calculate the probability of each document. Step S6: Set a decision threshold and identify documents with a probability higher than the decision threshold as malicious documents; input the remaining documents as context into the large language model to generate the final answer; achieve effective defense against knowledge poisoning attacks in the retrieval enhancement generation system.

2. The method for defending against knowledge poisoning attacks in a retrieval enhancement generation system according to claim 1, characterized in that: In step S1, an original sample set S is constructed from the tagged retrieval data; the original sample set S includes several samples, each sample including a user query document q and its corresponding set of retrieved documents. Each document is labeled with a binary tag. ; 0 indicates a normal document, and 1 indicates a malicious document.

3. The method for defending against knowledge poisoning attacks in a retrieval enhancement generation system according to claim 1, characterized in that: In step S33, the process of constructing semantic association edges is as follows: Calculate the current document With other documents The cosine similarity is used to determine the similarity if it exceeds a preset threshold. And the remaining documents Belongs to the current document If the first k most similar neighbors are found, then add them from the current document. Point to other documents The directed edges are defined by their cosine similarity values.

4. A method for defending against knowledge poisoning attacks in a retrieval enhancement generation system according to claim 3, characterized in that: In step S4, the specific process of calculating the four-dimensional structural feature vector is as follows: The four-dimensional structural features include the proportion of reciprocal edges, the average similarity between neighbors, the local clustering coefficient, and the retrieval score. The reciprocal edge ratio is used to represent the proportion of reverse edges in the outgoing edges of a document node; The average neighbor similarity is used to represent the mean value of the edge attributes corresponding to all outgoing edges of a document node; The local clustering coefficient is used in an undirected subgraph composed of reciprocal edges to represent the ratio of the actual number of connections between document node neighbors to the theoretical maximum number of connections. The retrieval score is used to represent the retrieval score of a document node; The four-dimensional structural features are standardized to obtain a four-dimensional structural feature vector.

5. A method for defending against knowledge poisoning attacks in a retrieval enhancement generation system according to claim 4, characterized in that: The formula for the reciprocal edge ratio is: ； In the formula, For document nodes in a heterogeneous graph; This is the set of directed edges between document nodes; For document nodes The degree of departure; For document nodes The neighbors gathered. It is a very small positive number; The formula for calculating the average similarity of the neighbors is: ； In the formula, For document nodes With document nodes Cosine similarity between them; The formula for calculating the local clustering coefficient is: Let For reciprocal neighbors to gather, record ; ； In the formula, For a set of mutually beneficial edges, For document nodes a reciprocal neighborhood set For document nodes in a heterogeneous graph; if <2, then =0; The formula for calculating the retrieval score is: ； In the formula, To search for documents semantic embedding vector; For document semantic embedding vector; This represents the dot product operation of vectors.

6. A method for defending against knowledge poisoning attacks in a retrieval enhancement generation system according to claim 5, characterized in that: In step S5, the specific execution process of the graph attention network model is as follows: Step S51. Initialize the edge-type-aware graph attention network model and the fully connected classification layer, and set the maximum number of training epochs. And the current round t=1; Step S52. Perform the t-th round of forward propagation; The global heterogeneous graph is used as the training set to input a graph attention network model. This network model uses a multi-head attention mechanism to fuse edge attribute and edge type information, calculates a new embedding representation of document nodes, and outputs the predicted probability that each document belongs to malicious documents through a fully connected classification layer. ; During training, the prediction probability is calculated. The loss value between the actual label y and the real label y is used to update the network model parameters using the backpropagation algorithm; Step S53: Perform verification and early stop judgment; Evaluate the performance of the current network model on the validation set and determine whether the current round t is less than 1. If the early stopping condition is not triggered, let t = t+1 and return to step S52; otherwise, end the training and save the optimal network model parameters, and proceed to step S6.

7. A method for defending against knowledge poisoning attacks in a retrieval enhancement generation system according to claim 6, characterized in that: The specific implementation process of step S6 is as follows: Step S61. Set the decision threshold search interval , This is the lower limit of the decision threshold. Set the upper limit of the decision threshold; select the current threshold. ; Step S62. Calculate the current threshold. The classification F1 score is calculated, and the optimal threshold corresponding to the best F1 score is recorded. ; Step S63. Select the next threshold Determine the next threshold. Is it less than or equal to? If so, then = Return to step S62; Otherwise, proceed to step S63; Step S63. Use the optimal network model parameters and the optimal decision threshold. Inference is performed on the new retrieved document. If the probability of malice output by the network model is greater than... If so, it is determined to be a malicious document.

Images