A context poisoning attack defense method for RAG applications
By combining sentence-level clustering and probability distribution classification with a logistic regression model, the problem of data poisoning attacks in RAG applications was solved, achieving a highly efficient defense effect.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2025-12-01
- Publication Date
- 2026-03-03
AI Technical Summary
Existing technologies cannot effectively defend against data poisoning attacks in RAG applications, especially when faced with injection of multiple or a few malicious texts, as their defense performance is insufficient and the time cost is too high.
Potential malicious template clusters are identified through sentence-level clustering, and the probability distribution of a large language model is used for classification. Harmful information is identified by combining a logistic regression model, and a key-value caching mechanism is used to optimize computational efficiency.
It achieves effective defense against multiple or a small number of malicious texts, reduces computation time overhead, and improves the defense performance of RAG applications.
Smart Images

Figure CN121234911B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of RAG (Relational Analysis and Grading) technology, specifically to a method for defending against contextual poisoning attacks in RAG applications. The research area of this invention is to effectively distinguish harmful information hidden in contextual information in objective questions with definite answers through clustering and probability distribution classification, thereby resisting data poisoning attacks. Background Technology
[0002] RAG stands for Retrieval-Augmented Generation, an AI architecture that combines information retrieval with natural language generation. It allows large language models to first search for relevant external knowledge before generating answers based on that knowledge. Currently, to defend against data poisoning attacks targeting RAG applications, the following technologies exist:
[0003] 1. Target query rewrite defense technology
[0004] This technology rewrites the input target query directly, leveraging LLM's text generation capabilities to preserve the original meaning of the target query while reconstructing the sentence structure. For example, when the target question is "Who is the CEO of OpenAI?", the rewritten question might be "Who serves as the CEO of OpenAI?". This method has significant drawbacks. First, while changing the sentence structure affects the relevance between malicious text and the target query, it doesn't alter the semantic information of the sentence itself, and the similarity calculation after embedding is largely unaffected. Second, when faced with multiple malicious text injections, even after rewriting the target query, its similarity to the malicious text remains higher than most normal information, failing to provide effective defense. Finally, the correct information in the rewritten target query is also affected, impacting normal retrieval performance.
[0005] 2. Perplexity-based detection and defense techniques
[0006] This technique is based on the perplexity level (PPL) metric to measure text quality; a higher perplexity level indicates lower quality. However, this method has a significant drawback: the malicious text in the attack method is not generated based on gradients, but rather through LLM to generate coherent and fluent text information. There is no meaningless text information involved; statistically, the perplexity level of malicious text is not higher than that of correct text, thus making effective differentiation impossible.
[0007] 3. Knowledge Expansion Defense Technology
[0008] This technique incentivizes the model to defend against attacks by expanding the retrieval range, providing more reference information to make judgments when answering target queries. However, this method has fundamental limitations. First, continuously expanding the retrieval range introduces too much invalid information, affecting the model's judgment. Second, even with an expanded retrieval range, the model's output may still be biased towards incorrect answers under the influence of multiple malicious texts. Finally, excessive retrieval information increases the model's processing time, adding additional time overhead.
[0009] 4. Defense techniques based on comparison with real-world knowledge
[0010] This technology achieves a defensive effect by comparing information retrieved from a knowledge base with knowledge acquired from external sources through the model. However, this method has significant drawbacks. First, the comparison process constantly calls the model to generate knowledge and assess errors, resulting in substantial time overhead. Second, this method ignores the reliability of external information. The knowledge within a RAG system is often specialized domain knowledge used to mitigate the model's illusion-generating problems. Relying on the model to provide additional information for comparison fails to achieve accurate judgment.
[0011] In summary, existing technologies generally suffer from the following fundamental problems: lack of defense performance, inability to adaptively defend against different numbers of malicious texts, and insufficient defense performance; increased time overhead, affecting the time consumption of the normal retrieval process. Summary of the Invention
[0012] To address the technical problems mentioned in the background section, this invention provides a method for defending against context poisoning attacks in RAG applications.
[0013] The technical solution adopted in this invention is as follows:
[0014] A method for defending against context poisoning attacks in RAG applications includes the following steps:
[0015] S1. Enter the target query statement The retrieval system in RAG applications retrieves the target query statement. The search yields multiple contextual information;
[0016] S2. Retrieve The context information of each segment is segmented according to the sentence level. The similarity score between each sentence and the target query is calculated. The sentence with the highest similarity score to the target query in each segment is taken as the representative sentence of that segment. The representative sentence set of all segments is clustered to identify and filter out potential malicious template clusters.
[0017] S3. The large language model provides all candidate answers based on the existing context information, calculates the log probability of all contexts for different candidate answers, removes the influence of the large language model's own parameter knowledge from the log probability to obtain the support of all contexts for different candidate answers, and the overall log probability vector is used as the distribution of the support of contexts for candidate answers.
[0018] S4. Use a logistic regression model to identify individual pieces of harmful information from the distribution of support for candidate answers in the context, thereby filtering out incorrect answers.
[0019] Specifically, S2 includes the following steps:
[0020] S2.1. Retrieve each segment of context Segmented by sentence level The context set is Using a unified semantic encoder to translate the target query statement and each sentence Encode as a vector , ;
[0021] S2.2. Calculate the similarity score between each sentence and the target query statement; the formula is as follows:
[0022] ;in, This represents the similarity score between each sentence and the target query;
[0023] S2.3. Select the sentence with the highest cosine similarity from each paragraph. As the representative sentence of this paragraph, let its vector be denoted as . ;
[0024] S2.4. On the set of representative sentences Perform density-based clustering to obtain multiple high-density clusters. and noise points;
[0025] S2.5. For each cluster Calculate the average similarity within the cluster. and the target query statement Relevance The formulas are as follows:
[0026] ;
[0027] ;
[0028] in, and For clusters Different representative sentences;
[0029] When a cluster simultaneously satisfies , When conditions are met, it is identified as a malicious template cluster and filtered out. The subsequent context of the segment is filled in sequentially, where, This represents the threshold for measuring the density of text segments within a cluster. This represents the threshold indicating the similarity between each text segment in the cluster and the target query.
[0030] Preferably, the large language model adopts the Llama3-8b model based on the transformer architecture.
[0031] Further, step S3 includes the following steps:
[0032] S3.1. Given the target query q and the retrieved context, the candidate answer sequence set is obtained through the large language model. Where n=2, and each candidate answer consists of one or more tokens;
[0033] S3.2. Obtain the log probability of the candidate answer for each piece of context information, and add a key-value caching mechanism in the process;
[0034] S3.2.1. Prompt words When inputting into a large language model, the prompt words are calculated by the Self-Attention layer within the large language model. The K and V caches are:
[0035] , ;
[0036] in, Indicator words Total length of the token Indicates hidden dimensions. Represents the cache matrix;
[0037] S3.2.2. Input the context of the first paragraph into the large language model and calculate the context information of the first paragraph. , , ,
[0038] ;
[0039] This represents the encoding vector of the first segment of context. , , They represent , , Their respective weight matrices; , Separately and , Perform concatenation, cache as , ;
[0040] S3.2.3. Take the last token of the first paragraph's context as the input token of the large language model, and calculate a new query vector. and its key vector and Value vector ,use and , , , Perform self-attention calculations to obtain the answer. Attention score corresponding to the first token ,
[0041] ;in, Indicates the scaling factor;
[0042] Attention Score After further calculations by the large language model, the answer to the subsequent reasoning is obtained. The probability distribution of the first token is used to extract candidate answers. The logarithmic probability corresponding to the first token;
[0043] S3.2.4. Obtain the candidate answers The first token is input into the large language model to obtain the answer. The attention score of the second token is calculated after passing through the large language model's subsequent normalization layer, forward propagation layer, linear layer, and softmax layer to obtain the answer. The probability distribution of the second token is used to extract candidate answers. The logarithmic probability corresponding to the second token;
[0044] S3.2.5. Obtain subsequent candidate answers sequentially following the method in S3.2.4. The logarithmic probability of all tokens in the set;
[0045] S3.2.6. Based on the candidate answers The first piece of context information is obtained by calculating the logarithmic probabilities of all tokens. For candidate answers The logarithmic probability, ;
[0046] in, Indicates the context information of the first paragraph. For candidate answers The logarithmic probability, Indicate the answer corresponding total, Represents the overall parameters of a large language model;
[0047] S3.2.7. Following the methods in S3.2.3-S3.2.6, obtain the context of the first paragraph in relation to the candidate answers. log probability ;
[0048] S3.2.8. Retain the prompt words obtained in S3.2.1 The cache remains unchanged. The first segment of context information is replaced with subsequent context information in sequence. Each time it is replaced, the log probability of the i-th segment of context with respect to different candidate answers is obtained according to the method of S3.2.2-S3.2.7, until the log probability of all contexts with respect to different candidate answers is obtained. , ;
[0049] S3.4 Pair all contexts with the answer The log probabilities form the support vector.
[0050] ;
[0051] in, Indicates search suggestion words All contexts are relevant to the answer. The level of support;
[0052] S3.5. Construct a context-free baseline.
[0053] ;
[0054] ;
[0055] in, Indicates no contextual baseline. This indicates that the knowledge is based on the parameters of the large language model itself. Represents all contexts for candidate answers The degree of support is a probability vector of dimension N*1, specifically... For each of its elements , , indicating context Improve the answer Support Then it indicates the context. Weakening or disproving the candidate answer Support, overall probability vector As context for candidate answers Support distribution;
[0056] S3.6. Do Normalization process yields normalized vectors , ;
[0057] in, Indicates the normalized result , This indicates the candidate answer. Each piece of context information The degree of support for it, with each dimension of the vector being a value. , where represents the specific probability value of each dimension.
[0058] Specifically, S4 includes the following steps:
[0059] S4.1 is obtained from S3 Extract multiple low-dimensional statistics to form a feature vector. ,
[0060] ;
[0061] Among them, the largest contribution value The calculation formula is as follows:
[0062] ;
[0063] This indicates taking the maximum value from all dimensions of the normalized vector;
[0064] entropy The calculation formula is as follows:
[0065] N represents the number of context segments;
[0066] concentration The calculation formula is as follows:
[0067] ;
[0068] variance The calculation formula is as follows:
[0069] , ;
[0070] forward quality and The calculation formula is as follows:
[0071] ;
[0072] S4.2. Binary classification using a logistic regression model:
[0073] , Where W represents the weight vector and b represents the bias term. This represents the Sigmoid activation function. This indicates a normal distribution or the correct answer. This indicates an abnormal distribution or a misleading answer;
[0074] S4.3. Using a threshold When making decisions, if This was determined to be the correct answer, and this candidate answer will be used in the final selection. As the target query statement The query results. Preferably, the loss function of the logistic regression model. Minimize the weighted log loss:
[0075] ;
[0076] in This refers to the predicted probability that the j-th sample belongs to the normal category. , For category weights, is the regularization coefficient.
[0077] After adopting the above solution, the beneficial effects of the present invention are as follows:
[0078] (1) The attack defense method proposed in this invention has a dual-channel defense mechanism, including: a retrieval result filtering method based on sentence hierarchical clustering that can defend against the injection of multiple malicious texts in a concentrated manner, and an answer based on probability output distribution that can identify the sparse injection of a few malicious texts, thereby achieving collaborative defense against the retrieval stage and the generation stage.
[0079] (2) The invention incorporates a KV-cache mechanism, which enables low-latency deployment through cache reuse. It has good feasibility and engineering robustness, and provides excellent defense against context injection attacks in RAG applications. Attached Figure Description
[0080] Figure 1 This is a flowchart illustrating the main steps of the present invention;
[0081] Figure 2 This is a flowchart illustrating the overall steps of the present invention;
[0082] Figure 3 This is a core framework diagram of the Llama3-8b large language model used in this invention;
[0083] Figure 4 This is a schematic diagram of the KV-cache principle used in this invention. Detailed Implementation
[0084] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.
[0085] In the description of this invention, the terms "first" and "second" are used for descriptive purposes only and should not be construed as indicating or implying relative importance or implicitly specifying the number of indicated technical features. Thus, a feature defined as "first" or "second" may explicitly or implicitly include one or more of the stated features. In the description of this invention, "a plurality of" means two or more, unless otherwise explicitly specified.
[0086] In the description of this invention, the term "for example" is used to mean "used as an example, illustration, or description." Any embodiment described as "for example" in this invention is not necessarily to be construed as being more preferred or advantageous than other embodiments. The following description is provided to enable any person skilled in the art to make and use the invention. Details are set forth in the following description for purposes of explanation. It should be understood that those skilled in the art will recognize that the invention can be made without using these specific details. In other instances, well-known structures and processes will not be described in detail to avoid obscuring the description of the invention with unnecessary detail. Therefore, the invention is not intended to be limited to the embodiments shown, but is consistent with the broadest scope of the principles and features disclosed herein.
[0087] The method of this invention only addresses the scenario of processing two distinct candidate answers (which may include incorrect answers). For example... Figure 1 and Figure 2 As shown, a method for defending against context poisoning attacks in RAG applications includes the following steps:
[0088] S1. Enter the target query statement The retrieval system in RAG applications retrieves the target query statement. The search yields multiple contextual information;
[0089] S2. Retrieve The context information of each segment is segmented according to the sentence level. The similarity score between each sentence and the target query is calculated. The sentence with the highest similarity score to the target query in each segment is taken as the representative sentence of that segment. The representative sentence set of all segments is clustered to identify and filter out potential malicious template clusters. In the following description, unless otherwise specified, the context information refers to the context information of the top N segments, where N can be preset manually.
[0090] Existing attack methods make the malicious text itself extremely similar to the target query (template paraphrasing, forced keyword alignment), thereby "squeezing out" the top rankings during the retrieval stage. This invention uses representative sentences at the sentence level to screen, locate and eliminate possible sources of homogeneity and sharp clustering, and screen out situations where multiple malicious texts are contained in the context information.
[0091] S2 includes the following steps:
[0092] S2.1. Retrieve each segment of context Segmented by sentence level The context set is Using a unified semantic encoder to translate the target query statement and each sentence Encode as a vector , ;
[0093] S2.2. Calculate the similarity score between each sentence and the target query statement; the similarity score is calculated using one of the following methods: cosine similarity, Euclidean distance, Mahalanobis distance, Manhattan distance, KL divergence, or other symmetric / asymmetric information measurement functions. In this specific implementation, cosine similarity is used for calculation, and the formula is as follows:
[0094] ;in, This represents the similarity score between each sentence and the target query;
[0095] S2.3. Select the sentence with the highest cosine similarity from each paragraph. As the representative sentence of this paragraph, let its vector be denoted as . ;
[0096] S2.4. On the set of representative sentences Perform density-based clustering to obtain multiple high-density clusters. And noise points; specifically, the density-based clustering method adopts one of DBSCAN, HDBSCAN, OPTICS, K-Means, spectral clustering, or hierarchical clustering. These methods are all existing methods. In this specific implementation, the DBSCAN clustering method is adopted. DBSCAN represents the set of sentences. Divided into multiple high-density clusters Noise points (non-clustered points);
[0097] S2.5. For each cluster Calculate the average similarity within the cluster. and the target query statement Relevance The formulas are as follows:
[0098] ;
[0099] ;
[0100] in, and For clusters Different representative sentences;
[0101] When a cluster simultaneously satisfies , Under certain conditions, it is identified as a malicious template cluster and filtered out. The subsequent context of the segment is filled in sequentially, where, This represents the threshold for measuring the density of text segments within a cluster. The threshold representing the similarity between text segments in each cluster and the target query;
[0102] S2 enables automatic identification of potential malicious template clusters, and can solve scenarios where multiple harmful messages are injected into the context.
[0103] After processing by S1, it cannot be determined whether the harmful information is a single piece or multiple pieces. Therefore, after processing by S1, the retrieved context will be processed by S2 by default. If there are multiple malicious cases, they will be processed by S2 and then processed by S3. If there is a single piece of malicious information, it can be filtered out after processing by S3.
[0104] S3. To address the problem of difficulty in filtering and removing single harmful information mixed with contextual information, a probability-separated extraction is proposed. The large language model provides all candidate answers based on the existing contextual information, calculates the log probability of all contexts for different candidate answers, removes the influence of the large language model's own parameter knowledge from the log probability to obtain the support of all contexts for different candidate answers, and the overall probability vector serves as the distribution of the support of contexts for candidate answers.
[0105] S3 includes the following steps:
[0106] S3.1. Given a target query statement Given the retrieved context, candidate answers are derived using a large language model (considering only objective questions with a definite answer, i.e., usually with one and only one correct answer for the target question). The set of candidate answer sequences is as follows: Where n=2, and each candidate answer consists of one or more tokens;
[0107] In this specific implementation, the large language model is the Llama3-8b model based on the Transformer architecture; its backbone is a Transformer decoder stack with causal masking, and the model structure is as follows. Figure 3 As shown, the Llama3-8b model includes, in sequence, an Embedding layer, a first RMS Norm layer, a self-Attention layer, a second RMS Norm layer, a FeedForward layer, a third RMS Norm layer, a Linear layer, and a softmax function. The first RMS Norm layer produces Q, K, and V vectors. The output of the Embedding layer is connected to the self-Attention layer via a skip layer, and after being added to the self-Attention layer, it is directly input to the second RMS Norm layer. Simultaneously, it is connected to the FeedForward layer via a skip layer and added to the FeedForward layer. The self-Attention layer includes 32 attention heads, and the attention mechanism calculation formula for each attention head is as follows:
[0108] ;
[0109] It should be understood that the diagram shows the core framework of the Llama3-8b model, not the complete structure of the Llama3-8b model. The Llama3-8b model is an existing technology and will not be discussed in detail here.
[0110] S3.2. Probability retrieval operation based on KVcache; obtain the log probability of candidate answers for each segment of context information, and add a KV caching mechanism in the process;
[0111] In this invention, KVcache (Key / Value Cache) refers to the process where, during autoregressive generation, the self-attention at step t is performed with attention from all previous steps 1…t−1. Repeating the calculation of the KV values for each previous step would incur significant time overhead. Therefore, the attention key / value vectors of all previously processed tokens are stored for reuse in subsequent decoding steps, thus avoiding repeated recalculation of the prefix. The principle of KVcache is as follows: Figure 4 As shown.
[0112] S3.2.1. Prompt words When inputting into a large language model, the prompt words include instructions, formatting, and the target query q. It should be understood that in practical applications, prompt words can be input by the user when entering the target query, or they can be preset in advance. During the calculation of the Self-Attention layer in the large language model, the prompt words... The K and V caches are:
[0113] , ;
[0114] in, Indicator words Total length of the token Indicates hidden dimensions. Represents the cache matrix;
[0115] S3.2.2. Input the context of the first paragraph into the large language model and calculate the context information of the first paragraph. , , ,
[0116] ;
[0117] This represents the encoding vector of the first segment of context. , , They represent , , Their respective weight matrices; , Separately and , Perform concatenation, cache as , The purpose of this step is that when replacing the context content later, the key-value pairs only need to be calculated once in this case, which reduces a lot of extra computational overhead.
[0118] S3.2.3. Take the last token of the first paragraph's context as the input token of the large language model, or a start token (automatically provided by the word segmenter within the large language model) as the input token of the large language model, and calculate a new query vector. and its key vector and Value vector ,use and , , , Perform self-attention calculations to obtain the answer. Attention score corresponding to the first token ,
[0119] ;in, Indicates the scaling factor;
[0120] Attention Score After the large language model is processed by the subsequent normalization layer, forward propagation layer, linear layer and softmax layer, the probability distribution of the first token in the subsequent reasoning is obtained. The log probability of the first token in the candidate answer is extracted from this probability distribution (the corresponding probability value is extracted using the content of the candidate answer in S3.1).
[0121] S3.2.4. Obtain the candidate answers The first token is input into the large language model to obtain the answer. The attention score of the second token is calculated after passing through the large language model's subsequent normalization layer, forward propagation layer, linear layer, and softmax layer to obtain the answer. The probability distribution of the second token is such that the second token has the highest probability. Candidate answers are extracted from this distribution. The logarithmic probability of the second token corresponds to this; in this process, the actual generation of the second token is based on the prompt word, the first paragraph of context, and the first token, since the previous prompt words are cached. Given the context of the first paragraph, the input to the large language model only contains candidate answers. The first token;
[0122] S3.2.5. Obtain subsequent candidate answers sequentially following the method in S3.2.4. The logarithmic probability of all tokens in the set; it should be understood that if calculating candidate answers... The i-th token will then be the candidate answer. The (i-1)th token is input into the large language model;
[0123] S3.2.6 Based on candidate answers Log probabilities of all tokens, and token-by-token integrated candidate answers. The overall logarithmic probability value yields the first segment of context information. For candidate answers The logarithmic probability,
[0124] ;
[0125] in, Indicates the context information of the first paragraph. For candidate answers The logarithmic probability, Indicate the answer The corresponding total number of tokens, Represents the overall parameters of a large language model;
[0126] S3.2.7. After completing the candidate answers After extracting the logarithmic probability, the first paragraph of context information is used to obtain the candidate answer in accordance with the methods in S3.2.3-S3.2.6. log probability In this invention, since there are only two candidate answers, the first paragraph of context was obtained earlier. This step yields the log probability of candidate answer 1 relative to the context information of the first paragraph, and the log probability of candidate answer 2 relative to the context information of the first paragraph.
[0127] S3.2.8. Retain the prompt words obtained in S3.2.1 The cache remains unchanged. The first segment of context information is replaced with subsequent context information in sequence. Each time it is replaced, the log probability of the i-th segment of context with respect to different candidate answers is obtained according to the method of S3.2.2-S3.2.7, until the log probability of all contexts with respect to different candidate answers is obtained. , ;
[0128] In this invention, there is a time overhead in calculating repeated prefixes when calculating repeated prompt words, each context, and different answers. Therefore, KVcache is added to optimize the calculation of the probability of each context supporting the answer in s3, which can save computation time overhead.
[0129] S3.4 Pair all contexts with the answer The log probabilities form the support vector.
[0130] ;
[0131] in, Indicates search suggestion words All contexts are relevant to the answer. The level of support;
[0132] S3.5. To isolate the unstable influence of inherent common sense or biases on the output probability, a context-free baseline is constructed.
[0133] ; ;
[0134] in, Indicates no contextual baseline. This indicates that the knowledge is based on the parameters of the large language model itself. Represents all contexts for candidate answers The degree of support is a probability vector of dimension N*1, specifically... For each of its elements , , indicating context Improve the answer Support Then it indicates the context. Weakening or disproving the candidate answer Support, overall probability vector As context for candidate answers The support distribution; it can be seen from the formula that, This demonstrates how contextual information contributes to the answer after eliminating the parameter knowledge inherent in the large language model itself. The impact of this; the model's parameter knowledge is obtained by training the model on a large amount of data before it is invoked. The answer provided by the model is based on the parameters obtained during its training phase, which is called the model's parameter knowledge.
[0135] S3.6. Do Normalization process yields normalized vectors , ;in, Indicates the normalized result , This indicates the candidate answer. Each piece of context information The degree of support for it, with each dimension of the vector being a value. , representing the specific probability value of each dimension;
[0136] Evidence for the correct answer often comes from multiple contexts, and its contribution is relatively more evenly distributed and less sharp. However, misleading answers are fabricated by malicious text, so they are only supported by a very small number of malicious text fragments in the context, and their distribution is sparse and sharp. When there are mutually exclusive answers to the same question, this probability separation based on the context dimension can distinguish between the correct answer and the misleading answer.
[0137] To distinguish the distributional differences between correct and misleading answers, this invention is based on the idea of logistic regression. It takes the source probability vector of the answer in each context as input and learns a uniform or sharp discrimination boundary through logistic regression.
[0138] S4. Use a logistic regression model to identify individual pieces of harmful information from the distribution of support for candidate answers in the context, thereby filtering out incorrect answers;
[0139] S4.1 is obtained from S3 Extract multiple low-dimensional statistics to form a feature vector. ,
[0140] ;
[0141] Among them, the largest contribution value The calculation formula is as follows:
[0142] ;
[0143] This indicates taking the maximum value from all dimensions of the normalized vector;
[0144] entropy The calculation formula is as follows:
[0145] N represents the number of context segments;
[0146] concentration The calculation formula is as follows:
[0147] ;
[0148] variance The calculation formula is as follows:
[0149] , ;
[0150] forward quality and The calculation formula is as follows:
[0151] ;
[0152] In this specific implementation, k is taken as 2;
[0153] S4.2. Binary classification using a logistic regression model:
[0154] , ;
[0155] Where W represents the weight vector and b represents the bias term. This represents the Sigmoid activation function. This indicates a normal distribution or the correct answer. This indicates an abnormal distribution or a misleading answer;
[0156] S4.3. Using a threshold When making decisions, if This was determined to be the correct answer, and this candidate answer will be used in the final selection. As the target query statement The query results.
[0157] The logistic regression model is trained using a manually labeled training set, which consists of feature vectors. Manual annotation is required. The training set used in the training of the logistic regression model needs to be generated and collected through steps S1-S3.
[0158] To handle class imbalance, the loss function of the logistic regression model Minimize the weighted log loss:
[0159] ;
[0160] in This refers to the predicted probability that the j-th sample belongs to the normal category. , For category weights, is the regularization coefficient.
[0161] As can be seen from the above specific implementation, the core innovation of this invention is reflected in four aspects: dual-channel (multiple and single) defense mechanism, two-level reduction of density clustering of context representative sentences, low-latency answer identification through probability partitioning extraction (KV-cache), and lightweight discrimination driven by distribution features.
[0162] The dual-channel (multiple and single) defense mechanism combines sentence-level density clustering filtering (retrieval stage) with probability distribution-based answer identification (generation stage) to cover both multiple injection and few injection attack scenarios, unlike single-channel solutions that rely solely on target query rewriting, prompt word rewriting, or generation-side detection.
[0163] The two-level reduction of density clustering under representative sentences first abstracts the paragraph within the paragraph using the representative sentence with the highest similarity to the target query, and then performs DBSCAN clustering on the representative sentences between paragraphs to process malicious text clusters that are templated and copied, which is different from the dilution effect and noise amplification brought about by directly clustering the whole paragraph or the whole text.
[0164] Low-latency answer identification using probabilistic partitioning extraction (KV-cache) extracts a token-level log probability sequence (rather than mixed scoring) from each candidate answer, and uses KV-cache incremental inference to reduce latency to online usability; unlike full-scale repetitive forward pass which cannot be implemented.
[0165] Lightweight discrimination driven by distribution features feeds logistic regression with statistical distribution features such as sequence entropy, variance, concentration, maximum contribution value, and top k quality sum; it performs a binary classification task on different answers to determine whether they are correct or incorrect answers, and grasps the correct direction of the generated answers during the generation stage.
[0166] Although preferred embodiments of the invention have been described, those skilled in the art, upon learning the basic inventive concept, can make other changes and modifications to these embodiments. Therefore, the appended claims are intended to be interpreted as including both the preferred embodiments and all changes and modifications falling within the scope of the invention.
[0167] Obviously, those skilled in the art can make various modifications and variations to this invention without departing from its spirit and scope. Therefore, if these modifications and variations fall within the scope of the claims of this invention and their equivalents, this invention also intends to include these modifications and variations.
Claims
1. A method for defending against context poisoning attacks in RAG applications, characterized in that, Includes the following steps: S1. Enter the target query statement The retrieval system in RAG applications retrieves the target query statement. The search yields multiple contextual information; S2. Retrieve The context information of each segment is segmented according to the sentence level. The similarity score between each sentence and the target query is calculated, and the sentence with the highest similarity score to the target query in each segment is taken as the representative sentence of that segment. Clustering is performed on the representative sentences of all paragraphs to identify and filter out potentially malicious template clusters; S3. The large language model provides all candidate answers based on the existing context information, calculates the log probability of all contexts for different candidate answers, removes the influence of the large language model's own parameter knowledge from the log probability to obtain the support of all contexts for different candidate answers, and the overall log probability vector is used as the distribution of the support of contexts for candidate answers. S4. Use a logistic regression model to identify individual pieces of harmful information from the distribution of support for candidate answers in the context, thereby filtering out incorrect answers; S2 includes the following steps: S2.
1. Retrieve each segment of context Segmented by sentence level The context set is Using a unified semantic encoder to translate the target query statement and each sentence Encode as a vector , ; S2.
2. Calculate the similarity score between each sentence and the target query statement; the formula is as follows: ;in, This represents the similarity score between each sentence and the target query; S2.
3. Select the sentence with the highest cosine similarity from each paragraph. As the representative sentence of this paragraph, let its vector be denoted as . ; S2.
4. On the set of representative sentences Perform density-based clustering to obtain multiple high-density clusters. and noise points; S2.
5. For each cluster Calculate the average similarity within the cluster. and the target query statement Relevance The formulas are as follows: ; ; in, and For clusters Different representative sentences in Chinese; When a cluster simultaneously satisfies , Under certain conditions, it is identified as a malicious template cluster and filtered out. The subsequent context of the segment is filled in sequentially, where, This represents the threshold for measuring the density of text segments within a cluster. This represents the threshold indicating the similarity between each text segment in the cluster and the target query.
2. The method for defending against context poisoning attacks for RAG applications according to claim 1, characterized in that, The large language model adopts the Llama3-8b model based on the transformer architecture.
3. A method for defending against context poisoning attacks for RAG applications according to claim 2, characterized in that, S3 includes the following steps: S3.
1. Given a target query statement Given the retrieved context, the candidate answer sequence set is derived using a large language model. Where n=2, and each candidate answer consists of one or more tokens; S3.
2. Obtain the log probability of the candidate answer for each piece of context information, and add a key-value caching mechanism in the process; S3.4 Pair all contexts with the answer The log probabilities form the support vector. ; in, Indicates query suggestion words All contexts are relevant to the answer. The level of support; S3.
5. Construct a context-free baseline. ; ; in, Indicates no contextual baseline. This indicates that the knowledge is based on the parameters of the large language model itself. Represents all contexts for candidate answers The degree of support is a probability vector of dimension N*1, specifically... For each of its elements , , indicating context Improve the answer Support Then it indicates the context. Weakening or disproving the candidate answer Support, overall probability vector As context for candidate answers Support distribution; S3.
6. Do Normalization process yields normalized vectors , ; in, Indicates the normalized result , This indicates the candidate answer. Each piece of context information The degree of support for it, with each dimension of the vector being a value. , where represents the specific probability value of each dimension.
4. A method for defending against context poisoning attacks in RAG applications according to claim 3, characterized in that, S3.2 includes the following steps: S3.2.
1. Prompt words When inputting into a large language model, the prompt words are calculated by the Self-Attention layer within the large language model. The K and V caches are: , ; in, Indicator words Total length of the token Indicates the hidden dimension. Represents the cache matrix; S3.2.
2. Input the context of the first paragraph into the large language model and calculate the context information of the first paragraph. , , , , , ; This represents the encoding vector of the first segment of context. , , They represent , , Their respective weight matrices; , Separately and , Perform concatenation, cache as , ; S3.2.
3. Take the last token of the first paragraph's context as the input token of the large language model, and calculate a new query vector. and its key vector and Value vector ,use and , , , Perform self-attention calculations to obtain the answer. Attention score corresponding to the first token , ; in, Indicates the scaling factor; Attention Score After further calculations by the large language model, the answer to the subsequent reasoning is obtained. The probability distribution of the first token is used to extract candidate answers. The logarithmic probability corresponding to the first token; S3.2.
4. Obtain the candidate answers The first token is input into the large language model to obtain the answer. The attention score of the second token is calculated after passing through the large language model's subsequent normalization layer, forward propagation layer, linear layer, and softmax layer to obtain the answer. The probability distribution of the second token is used to extract candidate answers. The logarithmic probability corresponding to the second token; S3.2.
5. Obtain subsequent candidate answers sequentially following the method in S3.2.
4. The logarithmic probability of all tokens in the set; S3.2.
6. Based on the candidate answers The first piece of context information is obtained by calculating the logarithmic probabilities of all tokens. For candidate answers The logarithmic probability, ; in, Indicates the context information of the first paragraph. For candidate answers The logarithmic probability, Indicate the answer corresponding total, Represents the overall parameters of a large language model; S3.2.
7. Following the methods in S3.2.3-S3.2.6, obtain the context of the first paragraph in relation to the candidate answers. log probability ; S3.2.
8. Retain the prompt words obtained in S3.2.1 The cache remains unchanged. The first segment of context information is replaced sequentially with the subsequent context information. Each time it is replaced, the log probability of the i-th segment of context with respect to different candidate answers is obtained according to the method of S3.2.2-S3.2.7, until the log probability of all contexts with respect to different candidate answers is obtained. , .
5. A method for defending against context poisoning attacks in RAG applications according to claim 3 or 4, characterized in that, S4 includes the following steps: S4.1 is obtained from S3 Extract multiple low-dimensional statistics to form a feature vector. , ; Among them, the largest contribution value The calculation formula is as follows: ; This indicates taking the maximum value from all dimensions of the normalized vector; entropy The calculation formula is as follows: N represents the number of context segments; concentration The calculation formula is as follows: ; variance The calculation formula is as follows: , ; forward quality and The calculation formula is as follows: ; S4.
2. Binary classification using a logistic regression model: , ; Where W represents the weight vector and b represents the bias term. This represents the Sigmoid activation function. This indicates a normal distribution or the correct answer. This indicates an abnormal distribution or a misleading answer; S4.
3. Using a threshold When making decisions, if This was determined to be the correct answer, and this candidate answer will be used in the final selection. As the target query statement The query results.
6. A method for defending against context poisoning attacks in RAG applications according to claim 5, characterized in that, The loss function of the logistic regression model Minimize the weighted log loss: ; in This refers to the predicted probability that the j-th sample belongs to the normal category. , For category weights, is the regularization coefficient.
Citation Information
Patent Citations
Knowledge base enhancement generation method and system based on hybrid retrieval and fact verification
CN120874999A