Method and system for prompt injection attack detection

Abstract

The present invention discloses a method and system for training a risk detection model to detect prompt injection attack against a target model. The method includes inputting a training prompt to the target model to obtain attention heads and a training output, preparing an attention matrix heatmap, and training the attention matrix heatmaps using an image classification training model to obtain the risk detection model. The system includes a memory storing a set of instructions and one or more processors configured to execute the instructions to perform the method steps.

Description

DESCRIPTIONTITLE OF INVENTION: [METHOD AND SYSTEM FOR PROMPT INJECTION ATTACK DETECTION]TECHNICAL FIELD

[0001] The present disclosure relates generally to the field of artificial intelligence (Al) technologies, with a specific focus on the cyber security of Al models.BACKGROUND

[0002] The following discussion of the background to the invention is intended to facilitate an understanding of the present invention only. It should be appreciated that the discussion is not an acknowledgement or admission that any of the material referred to was published, known or part of the common general knowledge of the person skilled in the art in any jurisdiction as at the priority date of the invention.

[0003] Large Language Models (LLMs) have increasingly become a cornerstone in a myriad of applications, owing to their unparalleled proficiency in executing tasks based on textual instructions. These models are prevalently employed within chatbot applications, offering users tailored information pertinent to specific companies or services. Despite their robust capabilities, LLMs are not immune to cyberattacks, with one notable weakness being their susceptibility to prompt injection attacks.

[0004] Prompt injection attacks currently manifest in two primary forms, differentiated by the attacker's objectives: goal hijacking and prompt leaking. Goal hijacking occurs when an adversary manipulates the LLM to execute unauthorized tasks, thereby illicitly utilizing the company's computational resources for their own ends. Conversely, prompt leaking targets the extraction of the instruction prompts fed into the LLM. These prompts often harbor confidential and sensitive data; their compromise could enable competitors to emulate significant aspects of the application's functionality, thereby eroding competitive advantages and potentially diverting customers.

[0005] Thus, there exists a need to develop a solution that can effectively detect and mitigate prompt leaking attacks in LLM-integrated applications.SUMMARY

[0006] Specifically, to address the above-mentioned technical problems, the present invention specifically uses the following technical solutions:

[0007] Accordingly, in accordance with one aspect of the present invention, there is provided a method for training a risk detection model operable to detect prompt injection attack against a target model, wherein the method comprises following steps of:

[0008] S100 inputting a training prompt to the target model to obtain plurality of attention heads and a training output, wherein the training output is classified with a risk profile;

[0009] S200 preparing an attention matrix heatmap, comprising: S201 obtaining an attention matrix from each of the plurality of attention heads in a layer of the target model;5202 aggregating the attention matrices in the layer into an aggregated attention matrix;5203 creating the attention matrix heatmap based on a plurality of the aggregated attention matrices; S204 labelling the attention matrix heatmap with the risk profile;

[0010] S300 training the plurality of the attention matrix heatmaps using an image classification training model to obtain the risk detection model.

[0011] In some embodiments, the step of S202 further comprises a step of: S202a removing one or more null attentions from the aggregated attention matrix.

[0012] In some embodiments, the step of S202 further comprises a step of: S202b modifying the aggregated attention matrix to retain detection area.

[0013] In some embodiments, the detection area comprises: (I) system prompt that captures output’s attention; and / or (ii) post-system prompt that captures the output’s attention.

[0014] In some embodiments, the attention matrix heatmap comprises a visual pattern associated with the risk profile.

[0015] In some embodiments, the visual pattern comprises one or more diagonal lines.

[0016] In some embodiments, the target model is a large language model.

[0017] In some embodiments, the prompt injection attack comprises goal hijacking, prompt leaking, direct prompt injection attack and indirect prompt injection attack.

[0018] In some embodiments, the layer is a middle layer in network of the target model.

[0019] In some embodiments, data of the aggregated attention matrix has been normalized using a normalization method, wherein the normalization method comprises standard scaling, min-max scaling, and z-score normalization.

[0020] In accordance with another aspect of the present invention, there is provided a system for training a risk detection model operable to detect prompt injection attack against a target model, wherein the system comprises a memory storing a set of instructions; and one or more processors configured to execute the set of instructions to cause the system to perform following steps of:

[0021] S100 inputting a training prompt to the target model to obtain plurality of attention heads and a training output, wherein the training output is classified with a risk profile;

[0022] S200 preparing an attention matrix heatmap, comprising: S201 obtaining an attention matrix from each of the plurality of attention heads in a layer of the target model;5202 aggregating the attention matrices in the layer into an aggregated attention matrix;5203 creating the attention matrix heatmap based on a plurality of the aggregated attention matrices; S204 labelling the attention matrix heatmap with the risk profile;

[0023] S300 training the plurality of the attention matrix heatmaps using an image classification training model to obtain the risk detection model.

[0024] In some embodiments, the step of S202 further comprises a step of: S202a removing one or more null attentions from the aggregated attention matrix.

[0025] In some embodiments, the step of S202 further comprises a step of: S202b modifying the aggregated attention matrix to retain detection area.

[0026] In some embodiments, the detection area comprises: (i) system prompt that captures output’s attention; and / or (ii) post-system prompt that captures the output’s attention.

[0027] In some embodiments, the attention matrix heatmap comprises a visual pattern associated with the risk profile.

[0028] In some embodiments, the visual pattern comprises one or more diagonal lines.

[0029] In some embodiments, the target model is a large language model.

[0030] In some embodiments, the prompt injection attack comprises goal hijacking, prompt leaking, direct prompt injection attack and indirect prompt injection attack.

[0031] In some embodiments, the layer is a middle layer in network of the target model.

[0032] In some embodiments, data of the aggregated attention matrix has been normalized using a normalization method, wherein the normalization method comprises standard scaling, min-max scaling, and z-score normalization.

[0033] Other aspects and features of the present invention will become apparent to those of ordinary skill in the art upon review of the following description of specific embodiments of the invention in conjunction with the accompanying figures.BRIEF DESCRIPTION OF THE DRAWINGS

[0034] In the figures, which illustrate, by way of non-limiting examples only, embodiments of the present invention,

[0035] [Fig. 1]: illustrates, in a block diagram, a method for training a risk detection model for prompt injection attacks according to various embodiments of the present invention.

[0036] [Fig. 2]: illustrates, in a block diagram, a method for preparing an attention matrix heatmap according to various embodiments of the present invention.

[0037] [Fig. 3]: illustrates, in a block diagram, a method for training a detector for prompt injection attacks according to various embodiments of the present invention.

[0038] [Fig. 4]: illustrates, in a block diagram, a method for using a detector to detect prompt injection attacks according to various embodiments of the present invention.

[0039] [Fig. 5]: illustrates the modifications of the attention matrix according to various embodiments of the present invention.

[0040] [Fig. 6]: illustrates examples of the attention matrix heatmaps generated according to various embodiments of the present invention.DETAILED DESCRIPTION

[0041] Throughout this document, unless otherwise indicated to the contrary, the terms “comprising”, “consisting of, “having” and the like, are to be construed as non-exhaustive, or in other words, as meaning “including, but not limited to”.

[0042] Furthermore, throughout the document, unless the context requires otherwise, the word “include” or variations such as “includes” or “including” will be understood to imply the inclusion of a stated integer or group of integers but not the exclusion of any other integer or group of integers.

[0043] Furthermore, throughout the document, unless the context requires otherwise, the term “middle layer” refers to a layer within or near the middle region of network of the LLM, particularly as opposed to the layers in the top ("top layer”) or bottom (“bottom layer”) regions of the network of the LLM. Throughout the document, both the top layer and the bottom layer are collectively known as “outermost layer”. For example, in the experiments conducted in the present invention, there are in total 32 layers, and accordingly layer 18 is considered a middle layer in the context of the present invention.

[0044] Unless defined otherwise, all other technical and scientific terms used herein have the same meaning as is commonly understood by a skilled person to which the subject matter herein belongs.

[0045] Methods for Training a Risk Detection Model

[0046] As illustrated in Fig 1 and Fig 2, according to various embodiments of the present invention, there is provided a method for training a risk detection model operableto detect prompt injection attack against a target model, wherein the method comprises following steps of:

[0047] S100 inputting a training prompt to the target model to obtain plurality of attention heads and a training output, wherein the training output is classified with a risk profile;

[0048] S200 preparing an attention matrix heatmap, comprising: S201 obtaining an attention matrix from each of the plurality of attention heads in a layer of the target model;5202 aggregating the attention matrices in the layer into an aggregated attention matrix;5203 creating the attention matrix heatmap based on a plurality of the aggregated attention matrices; S204 labelling the attention matrix heatmap with the risk profile;

[0049] S300 training the plurality of the attention matrix heatmaps using an image classification training model to obtain the risk detection model.

[0050] In some embodiments, the step of S202 further comprises a step of removing one or more null attentions from the aggregated attention matrix (S202a). In some embodiments, the step of S202 further comprises a step of modifying the aggregated attention matrix to retain detection area (S202b).

[0051] S100 inputting a training prompt to the target model

[0052] Specifically with regards to the step S100, the objective of this step is to closely replicate conditions conducive to prompt injection attacks (e.g., goal hijacking prompt leaking), thereby generating a labeled dataset integral for training a sophisticated risk detection model, designed to be acutely aware of and responsive to such attacks in real- world applications. To achieve the intended objective, different aspects of the step S100 are described in details below.

[0053] Providing a training prompt: this process involves presenting the LLM with a crafted input instruction, also known as the training prompt. The purpose of the training prompt is to navigate the LLM towards executing a designated task, mirroring situations susceptible to prompt injection attacks. The objective of this prompt is to guide the LLM towards performing a specific task, thereby replicating scenarios vulnerable to prompt injection attacks. This involves both the entities responsible for inputting the training prompt and the LLM, which is charged with interpreting and acting on these instructions.The training prompt is crafted to mimic potential situations in which the LLM might encounter prompt injection threats.

[0054] Generating attention heads: When the LLM processes the training prompt, it activates its network architecture, which includes several attention heads across its layers. These attention heads examine different parts of the input, pinpointing key elements necessary for the task at hand. This process results in the creation of attention matrices, which map the relationships between various sections of the input.

[0055] Producing a training output: After the attention heads analyze the input, the LLM produces an output that reflects its typical response to the prompt. This output serves two functions: it represents what the LLM would generate during normal operations and it is utilized for risk analysis. The analysis, whether manual or Al-driven, examines the output for sensitive information that might signal a prompt injection attack, such as unauthorized dissemination of confidential data. The classification of the output into risk profiles serves as an indicator of the associated risk level.

[0056] Risk Labeling: The output's risk classification (e.g., presence or absence of prompt leakage) will then be linked to the respective attention matrix heatmap to be generated. As described below, these risk-labeled heatmaps constitute the training dataset for developing the detection system, which is founded on a risk detection model adept at identifying prompt injection attacks.

[0057] S200 preparing an attention matrix heatmap

[0058] Specifically with regards to the step S200, the objective of S200 is to develop an attention matrix heatmap. This heatmap acts as a graphical illustration of the LLM’s focus mechanisms when it processes input training prompts. S200 comprises several sub-steps (S201 - S204), each contributing to the creation and detailed labeling of the heatmap. This heatmap, labeled with the corresponding risk profile, ultimately forms the training dataset used to enhance an image classification model with risk detection capabilities, thereby transforming it into a risk detection model.

[0059] Step S201 of "Obtaining an attention matrix from each of the plurality of attention heads in a layer of the target model" involves retrieving attention matrices from the multitude of attention heads within a designated layer of the LLM. These attention heads assess the interrelations between different segments of the LLM as they processthe training prompt, thereby determining their focal points. The information gathered at this stage serves as the foundational data for compiling the aggregated attention matrix.

[0060] Step S202 of "Aggregating the attention matrices in the layer into an aggregated attention matrix" involves merging the data obtained from all attention heads in a specific layer. This action combines the identified focal points from various perspectives into a unified aggregated attention matrix, encapsulating the collective attention dynamics within that layer of the LLM. In other words, the step of aggregating or summing up the attention matrices serves as a data preprocessing measure. This process preserves essential information, ensuring that no significant data is lost. Simultaneously, it aids in streamlining the fine-tuning process of the image classification model, thereby enhancing its capability to detect prompt leaking effectively. In some embodiments of the invention, the aggregated attention matrix undergoes additional refinement to boost its analytical and pattern detection capabilities. This refinement includes normalization (for instance, through standard scaling, min-max scaling, z-score normalization) to standardize the data values, the removal of null attentions to eliminate data points that do not contribute meaningfully, and adjustments to emphasize areas critical for identifying risk patterns. These enhancements ensure the aggregated matrix is not only analytically robust but also optimized for detecting key patterns indicative of risks.

[0061] Certain embodiments of this invention employ a systematic approach to normalize the data within each row of the aggregated attention matrix, utilizing min-max scaling. The normalization process, as described in step S202, uniformly scales the values throughout the matrix to a standardized range, typically between 0 and 1 . This step is crucial for ensuring that attention values from different areas of the matrix are comparable. Such uniform scaling is vital for the subsequent image classification model, enabling it to learn effectively from the heatmaps without bias from the variability in raw attention values. Without normalization, variations in attention value ranges could skew the training phase of the model, potentially reducing the accuracy of the risk detection model. Thus, normalizing the values contributes to generating heatmaps that accurately reflect attention patterns, thereby enhancing the training efficiency of the risk detection model.

[0062] According to some embodiments of the invention, the aggregated attention matrix is possessed by removing the null attentions. For example, as illustrated in Fig 5,the first two (2) columns of the attention matrix heatmap are considered as the null attentions, and accordingly removed. This occurs because the LLM utilized in the experiments of the present invention consistently prepends two tokens that can be safely ignored and are thus filtered out during the step of removing null attentions. Thus, it is clear to a skilled artisan that the number of columns to be removed may vary, contingent upon the tokenizer employed by the specific LLM that is to be safeguarded.

[0063] The aggregated attention matrix is a combination of attention matrices from all attention heads within a specific layer of a LLM. Null attentions are values within this matrix that are considered insignificant or irrelevant for the task at hand. These could be due to the attention being spread too thinly across many tokens or focusing on padding tokens that do not contribute to the understanding of the input. The presence of null attentions can introduce noise into the data and potentially obscure the patterns that are indicative of prompt injection attacks.The process identifies these null attentions by analyzing the values in the attention matrix and removes them to refine the data. This refinement ensures that the subsequent heatmap visualization represents only the meaningful and relevant attention distributions. In other words, the removal of null attentions is a form of data preprocessing that improves the quality of the training data for the risk detection model.

[0064] In some embodiments of the invention, the aggregated attention matrix undergoes a process to isolate and highlight detection areas, which include regions influenced by both the system prompt and subsequent post-system prompts. System prompts can be prepended and / or appended to user prompts, while post-system prompts refer to those system prompts that are appended after user inputs. Thus, the system prompt, constituting the initial set of instructions, and the post-system prompt, including any follow-up inputs, could be scrutinized for leakage detection. In other words, the modified attention matrix aims to reflect full or portion of the system prompt / post system prompt. The extent to which these prompts are leaked reflects the capabilities of the attacker, manifesting as distinct patterns within the matrix.

[0065] This modification of the aggregated attention matrix could employ techniques such as slicing and filtering to precisely delineate areas within the matrix that are relevant for identifying potential vulnerabilities to prompt injection attacks. This analysis aims to detect whether the focus on these areas could inadvertently reveal sensitive information,thereby posing a security risk. The objective of this comprehensive examination is twofold: to generate heatmaps that accurately represent focal points susceptible to information leakage and to enhance the training dataset for the risk detection model. By doing so, the risk detection model to be developed based on the dataset of these modified heatmaps is better equipped to discern subtle patterns of unauthorized access or data exposure, increasing its efficacy in real-world scenarios. Through this integrated approach, the quality of the training data (i.e., modified heatmaps) has been improved, and could help generate risk detection model capable of identifying and mitigating threats posed by increasingly sophisticated prompt injection attacks. This refined focus on detection areas within the aggregated attention matrix enhances the effectiveness of the method for training and developing the risk detection model.

[0066] In the present invention, the attention matrix heatmaps are used as the training data. The analysis of these heatmaps is underpinned by identifying distinct visual patterns that correlate with the likelihood of prompt leaking, as interactions with the LLM, including prompt injection attacks, will activate the attention heads, thereby generating recognizable visual patterns (i.e., footprint) on the attention matrix heatmaps.

[0067] For instance, an illustrative example provided in Figure 6 (the attention matrix heatmaps of layer 18) demonstrates how diagonal lines (i.e., the presence of “white dots”, indicating high values on the attention matrix heatmap) within the heatmap patterns (as observed in Figures 6a and 6b) indicate the LLM's output closely mirroring its input, thereby suggesting a potential risk of instruction prompt disclosure. Conversely, the absence of such diagonal lines in other heatmaps (specifically in Figures 6c and 6d) indicates a lower risk of prompt leakage, offering a clear visual benchmark for evaluating the output's integrity. The LLM used in this experiment comprises 32 layers (numbered 0 to 31 ), placing layer 18 in or around the middle region, thus qualifying it as a middle layer. Prompt leakage typically manifests when the LLM sequentially references each token to replicate the system prompt. Consequently, this sequential referencing, where token a attends to token b, followed by token a+1 attending to token b+1 , and so forth, creates a diagonal pattern in the heatmap, visually representing the leakage behavior. In details, the invention also discloses a procedural step, S202, which involves the strategic selection of a particular layer within the LLM's network architecture for the extraction of attention matrices. This choice is guided by an observation of the layer's performance in showcasing attention and dependency relationships within the input data. Theserelationships are critical as they encapsulate the grammatical structure and inherent meaning of the input, providing a ground for uncovering patterns that are indicative of prompt leaking. Notably, in some embodiments, it has been observed that the middle layers (i.e., layers within or near the middle region of network of the LLM), as opposed to the outermost layers (e.g., layers in top or bottom regions), appear to be particularly adept at revealing such patterns, marking a significant insight into the operational dynamics of the LLM in the context of security and data integrity. Nevertheless, it is important to note that layers located in other regions of the LLM network may also prove valuable in identifying patterns that are indicative of, or suggest, prompt leaking. In other words, while the middle layer in some embodiments can be selected for detecting patterns of interest, other layers also have the potential to contribute to the recognition of prompt leakage. Further steps in the process, namely S203 and S204, elaborate on the generation and labelling of the attention matrix heatmap. Step S203 is dedicated to creating a visual representation that encodes the aggregated attention matrices' values, typically employing a color scheme to signify the intensity and distribution of attention throughout the matrix. Following this, step S204 involves a labelling step where each heatmap is assigned a label corresponding to its risk profile. This label acts as an indicator of the potential risk associated with the training output, particularly focusing on whether the output could contain leaked information from the instruction prompt. The risk labelled heatmaps carry information related to the vulnerabilities to prompt injection attacks, and accordingly can be utilized as the dataset for training the risk detection model.

[0068] S300 training the plurality of the attention matrix heatmaps using an image classification training model to obtain the risk detection model

[0069] The step S300 involves training an image classification model with the attention matrix heatmaps labelled with risk profiles so that the image classification model can be developed into a risk detection model for detecting prompt injection attacks. Examples of the image classification model suitable for the present invention include, but not limited to, AlexNet, EfficientNet, ResNet (also known as Residual Networks), Swin Transformer, VGG Image Classification model.

[0070] In some embodiments of this invention, the image classification model used for training begins as a pre-established model, which is then fine-tuned utilizing the labeled attention matrix heatmaps. This fine-tuning is a pivotal process, entailing the recalibrationof the classifier’s parameters to empower it with the ability to discern and categorize heatmap images. The categorization is based on the detection of distinctive patterns that signify the occurrence of prompt leaking. The objective of Step S300 is to develop a detection model with the precision and reliability needed to identify prompt injection attacks, including prompt leakage, through the analysis of specific patterns manifested in the attention matrix heatmaps.

[0071] Fig 3 illustrates the development method of the risk detection model according to some embodiments of the present invention. As illustrated in Fig 3, the method (300) initiates with supplying a LLM (302), with a training prompt, which could be of a potentially malicious nature (301 ). This action triggers the generation of attention heads and a corresponding training output (303). These attention heads are then processed to form an attention matrix heatmap (304) according to the methods of the present invention described above, while the training output (303) undergoes classification into various risk profiles or risk labels (305), through either manual classification techniques or automated methods (e.g., leveraging an LLM-based risk classification model). Examples of risk labels include at least one of the followings: (i) binary output (e.g., attack, non-attack); (ii) multi-category output (e.g., type of attack category); and (iii) likelihood of attack and occurrence.

[0072] The resulting attention matrix heatmaps (304), tagged with their respective risk labels (305), constitute the training dataset (306), which can be used for the training process aimed at developing a detector capable of identifying the risk associated with prompt injection attacks. In some embodiments, this training dataset is employed to train an image classification algorithm (308). Upon the completion of this training, the algorithm is endowed with the capability to detect prompt injection attacks, thus evolving into a dedicated detector, or what is referred to as the risk detection model (310).

[0072] Application of the Risk Detection Model

[0073] Once successfully trained, the risk detection model is operable to identify prompt injection attacks within Al systems or Al-based applications, such as those based on LLMs, by analyzing the attention matrix heatmaps generated during the LLM's active sessions. For instance, as demonstrated in Fig 4 according to some embodiments of this invention, a specialized detector (410) is deployed to ascertain whether an operationalLLM (402) is compromised by prompt injection attacks. Actual inputs (401), which may include potentially malicious prompts, are processed by the operational LLM (402).

[0074] Following this, the process involves extracting attention heads from the LLM (402), and aggregating the attention matrices from those heads within the same layer to form a comprehensive aggregated attention matrix. This matrix is then visualized as a heatmap (attention matrix heatmap, 404) corresponding to the specific input (401 ). The methodology employed to generate these operational heatmaps mirrors the method used to create the heatmap images for training the risk detection model.

[0075] To evaluate whether the input encompasses an attack prompt, the generated heatmap is then analyzed by the detector (410). This analysis is adept at comparing the heatmap against the model's learned patterns, specifically selecting images from the same LLM layer (402) as those on which the detector was trained. The detector then outputs its findings, categorizing them into various risk profiles or risk labels (405), which inform the LLM operator of potential risks and prompt necessary countermeasures. Examples of risk labels include at least one of the followings: (i) binary output (e.g., attack, non-attack); (ii) multi-category output (e.g., type of attack category); and (iii) likelihood of attack and occurrence.

[0076] For instance, should the detector pinpoint an instance of prompt leakage before the information becomes visible in a chat interface, the implicated output can be preemptively withheld from display and substituted with a message indicating denial of service. This measure prevents the dissemination of sensitive information gleaned from the leaked prompt.

[0077] Furthermore, conversations entailing malicious prompts are earmarked for subsequent manual review. This protocol enables prompt engineers and Al specialists to dissect these interactions, leveraging the insights gained to refine instruction prompts or adjust LLM configurations, thereby diminishing the likelihood of future prompt leakage incidents within the operational LLM.

[0078] In certain scenarios, users implicated in circulating problematic prompts may face temporary or permanent restrictions on their LLM access privileges. This action serves as a deterrent against the perpetration of further attacks, safeguarding the LLM environment from potential harm and misuse.

[0079] Moreover, the present invention is versatile, applicable to both white-box and black-box Large Language Models (LLMs). Additionally, with the evolving landscape of prompt injection threats, this invention can be adapted to detect emerging varieties of prompt injection attacks, in addition to goal hijacking and prompt leaking.lt should be further appreciated by the person skilled in the art that variations and combinations of features described above, not being alternatives or substitutes, may be combined to form yet further embodiments falling within the intended scope of the invention.

[0080] As would be understood by a person skilled in the art, each embodiment may be used in combination with other embodiment or several embodiments.

Claims

ClaimsClaim 1 . A method for training a risk detection model operable to detect prompt injection attack against a target model, wherein the method comprises following steps of:S100 inputting a training prompt to the target model to obtain plurality of attention heads and a training output, wherein the training output is classified with a risk profile;S200 preparing an attention matrix heatmap, comprising:5201 obtaining an attention matrix from each of the plurality of attention heads in a layer of the target model;5202 aggregating the attention matrices in the layer into an aggregated attention matrix;5203 creating the attention matrix heatmap based on a plurality of the aggregated attention matrices;5204 labelling the attention matrix heatmap with the risk profile;S300 training the plurality of the attention matrix heatmaps using an image classification training model to obtain the risk detection model.Claim 2. The method according to Claim 1 , wherein the step of S202 further comprises a step of:S202a removing one or more null attentions from the aggregated attention matrix.Claim 3. The method according to Claim 1 or Claim 2, wherein the step of S202 further comprises a step of:S202b modifying the aggregated attention matrix to retain detection area.Claim 4. The method according to Claim 3, wherein the detection area comprises: (i) system prompt that captures output’s attention; and / or (ii) post-system prompt that captures the output’s attention.Claim 5. The method according to anyone of Claim 1 - Claim 4, wherein the attention matrix heatmap comprises a visual pattern associated with the risk profile.Claim 6. The method of Claim 5, wherein the visual pattern comprises one or more diagonal lines.Claim 7. The method according to anyone of Claim 1 - Claim 6, wherein the target model is a large language model.Claim 8. The method according to anyone of Claim 1 - Claim 7, wherein the prompt injection attack comprises goal hijacking, prompt leaking, direct prompt injection attack and indirect prompt injection attack.Claim 9. The method according to anyone of Claim 1 - Claim 8, wherein the layer is a middle layer in network of the target model.Claim 10. The method according to anyone of Claim 1 - Claim 9, wherein data of the aggregated attention matrix has been normalized using a normalization method, wherein the normalization method comprises standard scaling, min-max scaling, and z-score normalization.Claim 11 . A system for training a risk detection model operable to detect prompt injection attack against a target model, wherein the system comprises a memory storing a set of instructions; and one or more processors configured to execute the set of instructions to cause the system to perform following steps of:S100 inputting a training prompt to the target model to obtain plurality of attention heads and a training output, wherein the training output is classified with a risk profile;S200 preparing an attention matrix heatmap, comprising:5201 obtaining an attention matrix from each of the plurality of attention heads in a layer of the target model;5202 aggregating the attention matrices in the layer into an aggregated attention matrix;5203 creating the attention matrix heatmap based on a plurality of the aggregated attention matrices;5204 labelling the attention matrix heatmap with the risk profile;S300 training the plurality of the attention matrix heatmaps using an image classification training model to obtain the risk detection model.Claim 12. The system according to Claim 11 , wherein the step of S202 further comprises a step of:S202a removing one or more null attentions from the aggregated attention matrix.Claim 13. The system according to Claim 11 or Claim 12, wherein the step of S202 further comprises a step of:S202b modifying the aggregated attention matrix to retain detection area.Claim 14. The system according to Claim 13, wherein the detection area comprises: (i) system prompt that captures output’s attention; and / or (ii) post-system prompt that captures the output’s attention.Claim 15. The system according to anyone of Claim 11 - Claim 14, wherein the attention matrix heatmap comprises a visual pattern associated with the risk profile.Claim 16. The system of Claim 15, wherein the visual pattern comprises one or more diagonal lines.Claim 17. The system according to anyone of Claim 1 1 - Claim 16, wherein the target model is a large language model.Claim 18. The system according to anyone of Claim 11 - Claim 17, wherein the prompt injection attack comprises goal hijacking, prompt leaking, direct prompt injection attack and indirect prompt injection attack.Claim 19. The system according to anyone of Claim 11 - Claim 18, wherein the layer is a middle layer in network of the target model.Claim 20. The system according to anyone of Claim 11 - Claim 19, wherein data of the aggregated attention matrix has been normalized using a normalization method, wherein the normalization method comprises standard scaling, min-max scaling, and z-score normalization.

Images