APT attack detection and tracing method based on graph neural network and reinforcement learning

Abstract

The invention discloses an APT attack detection and tracing method based on a graph neural network and reinforcement learning, and belongs to the technical field of network security and multi-step attacks, and the method comprises the following steps: carrying out the modeling of an event in a system audit log through a continuous-time dynamic heterogeneous graph, generating a more challenging sample through an AGNN, improving the quality of the sample, and carrying out the detection of the APT attack. The generated positive sample is more complex and is close to a real attack scene; attention on similarity between attack steps is enhanced by using contrast loss, the ability of detecting attack behaviors is improved, and in a fine adjustment stage, according to a sample type, edge embedding is obtained by embedding nodes on two sides of a spliced edge; the full connection layer maps the embedding to the probability that a connection edge belongs to an attack event, and a weighted cross entropy loss function is used as a fine tuning target; the method of combining Monte Carlo tree search with reinforcement learning is used for gradually tracing events in the system. According to the method, APT attack graph embedded feature extraction and detection can be better carried out.

Description

Technical Field

[0001] The present invention belongs to the field of network security and multi-step attack technology, and in particular relates to an APT attack detection and tracing method based on graph neural network and reinforcement learning. Background Art

[0002] Advanced Persistent Threat (APT) attacks are a complex and highly covert form of attack, usually carried out over a long period of time by attackers with strong resources and technical capabilities, with the goal of stealing sensitive information or disrupting the normal operation of the target system. The multi-step and long-term latent characteristics of APT attacks make them extremely challenging in network security detection. Traditional intrusion detection systems (IDS) are difficult to effectively deal with the complexity of APT attacks because these systems usually rely on the feature analysis of single-step attacks and lack a global understanding of multi-step, cross-period attack behaviors.

[0003] In an APT attack, each attack step usually seems harmless, but multiple steps combined together pose a highly hidden threat. In addition, APT attackers often adopt different strategies at different time points, making the order and time interval of the attack steps uncertain, which further increases the difficulty of detection. Traditional signature- or rule-based methods show obvious limitations in the face of such complexity, especially when faced with large-scale, high-dimensional network log data, which is prone to a large number of false positives and false negatives.

[0004] In recent years, GNN has been gradually applied to the detection of APT attacks because GNN can effectively capture the complex dependencies between nodes in the network system, especially showing superior performance in the detection of multi-step attacks. Rossi et al. proposed a method based on graph convolutional network (GCN) to model attack behaviors in the network to capture the correlation between attack steps and improve detection accuracy. However, the tracing problem of APT attacks has not been fully solved, which requires the model to not only detect the attack behavior but also track it to reveal the specific path and steps of the attacker. Gao et al. introduced continuous-time dynamic heterogeneous graphs in the article Detecting unknown threat based on continuous-time dynamic heterogeneous graph network to detect complex multi-step attacks, integrating attention and memory mechanisms into dynamic graph information propagation to facilitate anomaly detection in continuous-time dynamic graphs. Nevertheless, due to constraints related to repeated sampling, their ability to capture contextual information related to the behavior of network entities is limited.

[0005] There are currently few methods for attack tracing using graph models. Existing methods analyze changes in model output by perturbing the input of GNN and identify graph components that significantly affect the prediction results. Existing methods are generally divided into two categories: learning-based methods and search-based methods. Learning-based methods, such as GNNExplainer, PGExplainer, ReinforcedExplainer, and PGMExplainer, typically use neural networks to identify key nodes or edges based on representations generated by trained GNN models. Specifically, GNNExplainer identifies a compact subgraph structure and a small number of node features for tracing.

[0006] PGMExplainer identifies key graph components and demonstrates dependencies through conditional probabilities. Although learning-based methods are effective in identifying influential nodes, edges, and subgraphs on predictions, they have difficulty capturing the correlation between spatiotemporal dependencies in the graph and the prediction results. Search-based methods, such as SubgraphX ​​and Causal-Interpret, use search algorithms with score functions to find important subsets of inputs. For example, SubgraphX, which uses Monte Carlo tree search, can effectively explore subgraphs to obtain explanations. Other methods have also been proposed. For example, He et al. used probabilistic graph models as a traceability framework in the article An explainer fortemporal graph neural networks. Xia et al. introduced TGNNExplainer in the article Explaining temporal graph models through an explorer-navigator framework, which used Monte Carlo tree search algorithms to find subsets of historical events in temporal graphs for prediction. However, GNNExplainer has difficulty capturing deep correlations in the interactions of computer network entities over a long period of time, which affects the traceability accuracy of complex multi-step attacks such as APT. In addition, its random masking of edges is limited to predicting the appearance of future edges rather than detecting attack behaviors. Summary of the invention

[0007] The purpose of the present invention is to provide an APT attack detection and tracing method based on graph neural network and reinforcement learning, which can better perform APT attack graph embedding feature extraction and detection, and also has better attack tracing performance.

[0008] To achieve the above object, the technical solution adopted by the present invention is:

[0009] A method for detecting and tracing APT attacks based on graph neural network and reinforcement learning, comprising the following steps:

[0010] Step S1: Use a continuous-time dynamic heterogeneous graph to model the events in the system audit log, generate more challenging samples through AGNN, improve the quality of samples, and make the generated positive samples more complex and close to the real attack scenario;

[0011] Step S2: In the pre-training stage, contrastive learning is used to optimize positive and negative samples, so that the model can be more sensitive when facing sparse positive samples, so as to better detect multi-step behaviors in APT attacks;

[0012] Step S3: In the fine-tuning stage, according to the sample type, the edge embedding is obtained by concatenating the embeddings of the nodes on both sides of the edge. The fully connected layer maps the embedding to the probability that the connected edge belongs to an attack event, and the weighted cross entropy loss function is used as the fine-tuning target.

[0013] Step S4: Use Monte Carlo tree search combined with reinforcement learning to gradually trace the events in the system.

[0014] The further improvement of the technical solution of the present invention is that the specific steps of step S1 are:

[0015] Step S1.1, using a continuous-time dynamic heterogeneous graph to model events in the system audit log;

[0016] The system behavior is represented by a graph G = (V, E), where V is a set of nodes representing entities in the system, and E is a set of edges representing the relationships between entities. For each node v∈V, a node feature matrix X is defined, and for each edge e∈E, an edge feature matrix A is defined. Based on these feature matrices, a graph embedding model is constructed.

[0017] Step S1.2: In the graph modeling process, GNN is used to learn the embedding of nodes and edges. In order to capture the dynamics of the time dimension, a message passing mechanism is used to update the status of the nodes. The update method is shown in formula (1):

[0018]

[0019] in, represents the embedding representation of node v at time step t, represents the neighbor nodes of node v, W m and W s is the weight matrix, σ is the activation function;

[0020] Step S1.3: The graph neural network generates the feature representation of the global graph by reading out the local node features aggregated by the function R, as shown in formula (2):

[0021]

[0022] H represents the embedding set of all nodes, and s is the embedding representation of the global graph. Through the above method, the overall behavioral characteristics of the system can be obtained, including the embedding representation of normal behavior and attack behavior.

[0023] A further improvement of the technical solution of the present invention is that: the specific steps of generating more challenging samples by AGNN in step S1 are as follows: The adversarial perturbation generation formula is shown in formula (3):

[0024]

[0025] Where δ is the adversarial perturbation vector, ∈ is the perturbation size constraint, and f θ is the prediction function of GNN, X and A represent the node feature matrix and adjacency matrix respectively, and y is the true label;

[0026] The adversarial training of AGNN needs to combine the original task loss and adversarial loss, and the comprehensive loss function As shown in formula (4)

[0027]

[0028] in, is the original task loss, used to ensure the performance of the model under normal input, is the adversarial loss, which is used to make the model robust when subjected to adversarial perturbations, and λ is the weight that balances the two parts of the loss;

[0029] When generating adversarial samples, more challenging node embeddings are generated through adversarial perturbations, and the node v i The adversarial embedding representation of is shown in formula (5):

[0030]

[0031] Among them, h i For node v i The original embedding of i Adversarial perturbations are generated by maximizing the model loss. Adversarial node embeddings are used to generate more complex adversarial samples, helping the model improve its ability to recognize attack behaviors.

[0032] A further improvement of the technical solution of the present invention is that AGNN can also perturb the structure of the graph to generate more complex samples. The perturbed adjacency matrix A * It is expressed as shown in formula (6)

[0033] A * =A+ΔA (6)

[0034] Where: ΔA is the adversarial perturbation matrix, representing a small modification to the graph structure;

[0035] AGNN generates adversarial samples by optimizing formula (7) to improve the generalization performance of the model.

[0036]

[0037] The inner maximization operation of the formula is used to generate the most aggressive adversarial perturbations, while the outer minimization operation is used to allow the model to achieve good performance even in this most unfavorable situation.

[0038] The further improvement of the technical solution of the present invention is that the specific steps of step S2 are as follows:

[0039] The positive sample is recorded as (X, A), and the negative sample generated based on (X, A) is recorded as (X * ,A * ); By constructing positive and negative sample pairs, the contrast loss function is used for training. The contrast loss function is defined as shown in formula (8)

[0040]

[0041] Among them, y i is the label of sample i, positive sample is 1, negative sample is 0, d(h i ,s) represents the positive sample embedding h i The distance between the positive sample and the global feature s, m is the distance threshold, the distance of the positive sample should be greater than m, that is, if the distance is less than m, the model will be penalized and the distance will be increased; if the distance is already greater than m, the loss of this item is zero, which is used to ensure that the distance between the positive sample and the global feature is greater than m. The total number of samples is N, which represents the number of positive and negative samples. The loss of all samples is normalized so that the value of the loss function is not affected by the number of samples.

[0042] A further improvement of the technical solution of the present invention is that in the encoder parameter fine-tuning stage, both the encoder and the decoder are adapted to a limited number of real attack samples. According to their types, the embedding of the edge is obtained by concatenating the embeddings of the nodes on both sides of the edge. Then, the fully connected layer maps the embedding to the probability that the connected edge belongs to an attack event. The weighted cross entropy loss function shown in formula (9) is used as the fine-tuning target

[0043]

[0044] Among them, y x Indicates whether edge x is a true attack event, 0 or 1, p x Indicates that the prediction x belongs to attack event y x=1, ω represents the weight corresponding to the attack behavior, which is used to amplify the impact of sparse attack signals and increase the weight of attack behavior, so that the model can pay more attention to a small number of attack samples.

[0045] The further improvement of the technical solution of the present invention is that the specific steps of step S4 are as follows:

[0046] Build a timing threat sniffer to calculate event E j E k The relevance score of the sniffer is denoted by H θ (E j ,E k ), inheriting the encoder from the fine-tuning phase of the model and combining it with the MLP layer, using these scores to optimize the Monte Carlo tree search during the attack traceability node expansion process;

[0047] The input and training process are described as follows: Let the target event be Each candidate event The input of the sniffer is recorded as vec(·) is a function that converts raw event data into a vector; let vec(E i )for Where X is the node feature matrix, Time(·) is a learnable function that embeds the real-valued timestamp into the vector; the relative dynamic time encoding function is used to calculate the time embedding within the interval, as shown in formula (10)

[0048] RTE(ΔT)=TimeLinear(RT(ΔT)) (10)

[0049] Where ΔT represents the position, TimeLinear is a learnable linear projection function, and RT(ΔT) is as shown in Equation (11) and Equation (12):

[0050]

[0051] Among them, dim represents the dimension of each dimension of the position encoding corresponding to the sine wave;

[0052] Input the detected attack events and each candidate event into H θ The output is the corresponding correlation score between them.

[0053] The further improvement of the technical solution of the present invention is that the prediction of the correlation score is regarded as a regression task, and the mean square error MSE is used as the loss function when training the sniffer. At the same time, a time-series-based correlation score is designed as the training target of the sniffer. The loss function is shown in formula (13):

[0054]

[0055] Where N represents the number of samples, y i is the model output, is the target relevance score, determined by formula (14)

[0056]

[0057] Where e is the natural logarithm, β is the decay rate constant, and ΔT is the timing difference between event pairs.

[0058] A further improvement of the technical solution of the present invention is that: the Monte Carlo tree search expands the nodes in the search tree through multiple rounds of simulation, each node represents a sub-event set, and each round of search includes four steps: node selection, node expansion, simulation and back propagation, using a pre-trained model, including an encoder and a decoder, to further evaluate and calculate rewards for the newly generated sub-nodes;

[0059] Initialization: The root node is a set of candidate event edges that meet the spatiotemporal conditions;

[0060] Node selection: Starting from the root node, select child nodes according to the upper confidence bound UCT strategy to balance development and exploration. The selection criteria are shown in formula (15):

[0061]

[0062] in, Represents node N i The set of events that have been expanded, c(N i ,e j ) indicates that at node N i Select e j The cumulative reward, n(N i ,e j ) is node N i Middle side j The number of times selected, ∑n(N i ,e l ) represents the total number of visits to all edges under the node;

[0063] λ(t) is the weight that changes dynamically with the number of search steps t, as shown in formula (16):

[0064] λ(t)=λ0·exp(-α·t) (16)

[0065] Where λ0 is the initial value, which controls the early exploration weight, α is the decay rate, which controls the speed at which λ decreases, and t is the number of search steps. In the early stage of the search, λ(t) is close to λ0.

[0066] Due to the adoption of the above technical solution, the technical progress achieved by the present invention is:

[0067] The present invention discloses an APT attack detection and tracing method based on graph neural network and reinforcement learning. The APT attack is graph modeled based on GNN, and the events in the system audit log are modeled using a continuous-time dynamic heterogeneous graph. The nodes of the event graph represent different entities in the system, and the edges represent the interactions or causal relationships between the entities. In the graph modeling process, GNN is used to learn the embedding of nodes and edges. GNN can effectively capture the complex dependencies between nodes in the network system, and especially shows excellent performance in the detection of multi-step attacks.

[0068] The present invention proposes a graph embedding feature extraction method based on AGNN and contrastive learning. Adversarial samples are generated through AGNN to improve the quality of samples, making the generated positive samples more complex and close to the real attack scenarios. The contrast loss function is combined to better characterize the similarities and differences between positive and negative samples, thereby improving the accuracy of detection.

[0069] The present invention uses reinforcement learning driven Monte Carlo tree search to trace the attack source. The method first calculates the correlation score between events through a time series sniffer, and then combines the Monte Carlo tree search algorithm to design a reinforcement learning strategy to dynamically adjust node selection to enhance the exploration of attack paths. In this way, the model can trace the attack path after detecting an APT attack, thereby providing more valuable information for security experts. BRIEF DESCRIPTION OF THE DRAWINGS

[0070] Figure 1 It is the overall flow chart of the method of the present invention;

[0071] Figure 2 This is the graph embedding visualization result obtained on the ATLAS-M2 dataset before and after applying AGNN and contrastive learning;

[0072] Figure 3 This is the graph embedding visualization result obtained on the ATLAS-M5 dataset before and after applying AGNN and contrastive learning;

[0073] Figure 4 These are the loss curves of the model at different training stages on the ATLAS-M2 dataset;

[0074] Figure 5 These are the loss curves of the model at different training stages on the ATLAS-M5 dataset. DETAILED DESCRIPTION

[0075] The present invention is further described in detail below in conjunction with embodiments:

[0076] Aiming at the problem of advanced persistent threat attack detection and tracing, the present invention proposes an APT attack detection and tracing method based on graph neural network and reinforcement learning. This method models the event graph constructed by the audit log through a continuous-time dynamic heterogeneous graph, and then embeds the normal behavior and attack behavior in the system through a graph neural network to capture the complex dependencies of the system in the time dimension. Furthermore, a contrast loss function is proposed and designed. By using contrast learning in the pre-training stage to optimize the distance between positive and negative samples, the model can be more sensitive when facing sparse positive samples, so as to better detect multi-step behaviors in APT attacks. Finally, the Monte Carlo tree search algorithm is combined to trace the detected APT attacks, and the node selection is dynamically adjusted using the reinforcement learning strategy to enhance the exploration of the attack path.

[0077] The present invention mainly consists of four parts: a graph modeling method based on AGNN, a graph embedding feature extraction network based on contrastive learning, a Monte Carlo tree attack tracing algorithm based on reinforcement learning, and an algorithm evaluation. First, a continuous-time dynamic heterogeneous graph is used to model the events in the system audit log, and more challenging samples are generated by AGNN to improve the quality of the samples, so that the generated positive samples are more complex and close to the real attack scenario. Secondly, the contrast loss is used to enhance the attention to the similarity between the attack steps and improve the ability to detect attack behaviors. In the fine-tuning stage, according to the sample type, the embedding of the edge is obtained by splicing the embedding of the nodes on both sides of the edge. Then, the fully connected layer maps the embedding to the probability that the connecting edge belongs to the attack event, and uses the weighted cross entropy loss function as the fine-tuning target. Finally, the Monte Carlo tree search combined with the reinforcement learning method is used to gradually trace the events in the system. In order to trace the detected APT attack, the node embedding and the detected abnormal event E_k are used as inputs. Any event that occurred before E_k may become part of the explanation, and a search-based method is used to identify the event subset E_(xp_k) related to the abnormal event. The search process is based on a reinforcement learning strategy. The Monte Carlo tree search gradually expands the nodes to find the most relevant historical events E_k, thereby constructing an event sequence that can explain the current attack. In order to effectively identify the key steps related to the attack behavior, a time-series threat sniffer is constructed to calculate the correlation score between event E_j and attack event E_k. In the process of tracing the attack path, a root node is first initialized to represent the set of detected attack events. Monte Carlo tree search expands the nodes in the search tree through multiple rounds of simulation, and each node represents a sub-event set. Each round of search includes four steps: node selection, node expansion, simulation, and back propagation. The pre-trained model (including encoder and decoder) is used to further evaluate and calculate the reward for the newly generated sub-nodes. The algorithm is ablated on the public APT attack datasets ATLAS-M2 and ATLAS-M5. APT attack detection and tracing experiments are carried out on two datasets. The overall performance is compared with other similar algorithms, and the evaluation indicators such as AUC and fidelity are compared with other algorithms to prove the superiority of the algorithm.

[0078] The algorithm proposed in this paper is mainly divided into three parts: graph modeling and feature extraction, pre-training and fine-tuning, and reinforcement learning-driven attack tracing. The algorithm flow chart is as follows Figure 1As shown. First, in the graph modeling and feature extraction part, a continuous-time dynamic heterogeneous graph is used to model the events in the system audit log. The continuous-time dynamic heterogeneous graph represents the complex behaviors in the system through directed edges between different types of nodes (such as users, processes, files, etc.) and events. The graph neural network is used to extract features from these nodes and their relationships, and the system state is embedded into a high-dimensional vector space through message passing and aggregation mechanisms, thereby capturing the complex dependencies in the system, especially the correlation in the time dimension. In this way, the multi-step characteristics of APT attacks can be effectively modeled to provide support for subsequent attack detection. Secondly, in the pre-training and fine-tuning part, the present invention proposes a contrast loss function. The traditional binary cross entropy loss function (BCE) focuses on the classification of positive and negative samples, while the contrast loss function designed by the present invention optimizes the distance between positive and negative samples, so that the model can better characterize the similarities between attack behaviors and the differences between normal behaviors, thereby improving the accuracy of APT attack detection. By optimizing the positive and negative samples using contrast learning in the pre-training stage, the model can be more sensitive when facing sparse positive samples, thereby better detecting multi-step behaviors in APT attacks. Finally, in the reinforcement learning-driven attack tracing part, the present invention combines Monte Carlo tree search and reinforcement learning strategies. First, after detecting an APT attack, Monte Carlo tree search is used to explore the event subgraph to find potential attack paths. In order to improve the effectiveness of tracing, the present invention designs a dynamically adjusted strategy selection method so that the model pays more attention to exploring new paths in the early stages of the search, and focuses more on known high-return paths in the later stages. In this way, the model can effectively trace each step of the APT attack and identify the attacker's complete path, thereby providing security experts with detailed attack behavior information.

[0079] 1. Graph modeling and feature extraction

[0080] In this part, graph modeling of events in the system is performed and features that are helpful for APT attack detection are extracted.

[0081] First, a continuous-time dynamic heterogeneous graph is used to model the events in the system audit log. The nodes of the event graph represent different entities in the system (such as processes, files, network connections, etc.), while the edges represent the interactions or causal relationships between entities. The time dimension is represented by the timestamps of the edges to capture the temporal dependencies between events in the system. The continuous-time dynamic heterogeneous graph can effectively capture complex heterogeneous relationships and model the dynamic changes of the system, which is particularly suitable for the modeling needs of APT attacks.

[0082] Suppose the event sequence is S = {E1, E2, ...}, each event E i can be represented as a four-tuple in and Represent the source node and destination node of the event respectively, e i is the edge connecting two nodes, t i is the timestamp of the event, and the types p and q of the nodes reflect the categories of the nodes in the graph, such as host, process, file, etc. The system behavior is represented by a graph G = (V, E), where V is a set of nodes representing entities in the system, and E is a set of edges representing the relationships between entities. For each node v∈V, a node feature matrix X is defined, and for each edge e∈E, an edge feature matrix A is defined. Based on these feature matrices, a graph embedding model is constructed.

[0083] In the graph modeling process, GNN is used to learn the embedding of nodes and edges. In order to capture the dynamics of the time dimension, a message passing mechanism is used to update the status of the node. The specific update method is shown in formula (1).

[0084]

[0085] in, represents the embedding representation of node v at time step t, represents the neighbor nodes of node v, W m and W s is the weight matrix and σ is the activation function.

[0086] Afterwards, the graph neural network generates a feature representation of the global graph by reading out the local node features aggregated by the function R, as shown in formula (2).

[0087]

[0088] Here, H represents the embedding set of all nodes, and s is the embedding representation of the global graph. In this way, the overall behavior characteristics of the system can be obtained, including the embedding representation of normal behavior and attack behavior.

[0089] In order to deal with the class imbalance problem in the data, the present invention generates more challenging samples through an adversarial graph neural network (AGNN), improves the quality of the samples, and makes the generated positive samples more complex and close to the real attack scenario, so that the model is more robust in the face of complex attacks. In the training mechanism of AGNN, the generator and the discriminator compete with each other, the generator tries to generate samples that can deceive the discriminator, and the discriminator is constantly improved to better distinguish between real samples and generated samples. This adversarial training method can cause the generated positive samples to gradually approach the complex behavioral characteristics in the real attack scenario, ensuring that the model obtains a more comprehensive and diverse sample set during the training process. These high-quality samples enable the model to more effectively learn the feature differences between positive and negative samples when facing sparse attack samples, thereby improving the robustness and accuracy of the model in detecting complex attacks, especially APT attacks.

[0090] The adversarial training of AGNN aims to generate perturbations for node features or adjacency relationships to improve the robustness of the model. The adversarial perturbation generation formula is shown in formula (3).

[0091]

[0092] Where δ is the adversarial perturbation vector, ∈ is the size constraint of the perturbation, usually a small positive number, ensuring that the perturbation does not change the overall structure of the graph, and f θ is the prediction function of GNN, X and A represent the node feature matrix and adjacency matrix respectively, and y is the true label. The adversarial training of AGNN needs to combine the original task loss and adversarial loss, and the comprehensive loss function As shown in formula (4).

[0093]

[0094] in: is the original task loss, used to ensure the performance of the model under normal input, is the adversarial loss, which is used to make the model robust when subjected to adversarial perturbations, and λ is the weight that balances the two parts of the loss.

[0095] When generating adversarial samples, more challenging node embeddings are generated through adversarial perturbations, and the node v i The adversarial embedding representation of is shown in formula (5).

[0096]

[0097] Where: h i For node v i The original embedding of iAdversarial perturbations are generated by maximizing the model loss. Adversarial node embedding is used to generate more complex adversarial samples to help the model improve its ability to identify attacking behaviors.

[0098] In order to generate more complex samples, AGNN can also perturb the structure of the graph. The adjacency matrix A after perturbation * It is expressed as shown in formula (6).

[0099] A * =A+ΔA (6)

[0100] Where: ΔA is the adversarial perturbation matrix, which represents a slight modification to the graph structure, such as adding or deleting some edges, to enhance the diversity and difficulty of the samples.

[0101] The goal of AGNN is to generate adversarial samples and improve the generalization performance of the model by optimizing the following formula (7). The inner maximization operation of the formula is used to generate the most aggressive adversarial perturbation, while the outer minimization operation is used to enable the model to achieve good performance even in the most unfavorable situation.

[0102]

[0103] Through adversarial learning, the generalization ability of the model has been significantly improved, especially when dealing with new or unknown attack scenarios, the model can still maintain a high detection performance.

[0104] 2. Pre-training and fine-tuning

[0105] 2.1 Pre-training phase

[0106] In this invention, the positive sample (the event graph representing abnormal system behavior) is recorded as (X, A), and the negative sample generated based on (X, A) is recorded as (X * ,A * ). By constructing positive and negative sample pairs, the contrast loss function designed by the present invention is used for training. The contrast loss function is defined as shown in formula (8).

[0107]

[0108] Among them, y i is the label of sample i (positive sample is 1, negative sample is 0), d(h i ,s) represents the positive sample embedding h iThe distance between the positive sample and the global feature s, m is the distance threshold, and the distance of the positive sample should be greater than m, that is, if the distance is less than m, the model will be penalized and the distance will be increased. If the distance is already greater than m, the loss of this item is zero, which is used to ensure that the distance between the positive sample and the global feature is greater than m. The total number of samples is N, which represents the number of positive and negative samples. The loss of all samples is normalized so that the value of the loss function is not affected by the number of samples.

[0109] In complex APT attack detection, positive and negative samples are often unbalanced, and the standard binary cross entropy loss function may not be effective enough for processing sparse positive samples. It focuses on the classification problem, that is, the separation of positive and negative samples, but APT attack detection is not just a simple binary classification problem, but also involves the distance problem between samples, that is, the similarity between attack behaviors and the difference between normal behaviors. Through the contrast loss designed by the present invention, the model can pay more attention to the similarity between attack steps and improve the ability to detect attack behaviors, especially the detection of similar attack steps in multi-step attack scenarios such as APT attacks. Each step of an APT attack may differ in time and features, but they are still related. Contrast loss can better capture these complex similarities.

[0110] 2.2 Fine-tuning stage

[0111] After pre-training, the encoder parameters enter the fine-tuning phase, during which both the encoder and decoder are adapted to a limited number of real attack samples. The embedding of the edge is obtained by concatenating the embeddings of the nodes on both sides of the edge according to its type. Then, the fully connected layer maps the embedding to the probability that the connected edge belongs to an attack event. The weighted cross entropy loss function shown in Equation (9) is designed as the fine-tuning target.

[0112]

[0113] Among them, y x Indicates whether edge x is a true attack event (0 or 1), p x Indicates that x is predicted to be an attack event (y x =1). ω represents the weight corresponding to the attack behavior, which is used to amplify the impact of sparse attack signals. Increasing the weight of attack behavior allows the model to pay more attention to a small number of attack samples.

[0114] By designing weighted cross entropy, the model can better detect sparse attack steps hidden in normal behavior. This is very important for multi-step attack scenarios in APT attacks, because each step may seem harmless, but all steps combined constitute an attack. In addition, to avoid the situation where the model only focuses on normal samples, the positive and negative samples in APT attacks are seriously unbalanced, so the model can better learn how to distinguish complex multi-step attacks from normal behavior. At the same time, it can be adjusted according to specific scenarios and data sets. For very rare attack scenarios, a larger ω can be selected to enhance the detection ability of attack samples.

[0115] 3. Attack tracing driven by reinforcement learning

[0116] In order to trace the source of APT attacks, the present invention adopts the method of Monte Carlo tree search combined with reinforcement learning to gradually trace the events in the system and find the complete attack path.

[0117] In order to trace the detected APT attack, the node is embedded in the detected abnormal event E k As input, in E k Any event that has occurred before may be part of the explanation, and a search-based approach is used to identify the subset of events related to the abnormal event. xpk The search process is based on a reinforcement learning strategy, which gradually expands nodes through Monte Carlo tree search to find nodes that are related to E k The most relevant historical events, thereby constructing a sequence of events that can explain the current attack.

[0118] In order to effectively identify the key steps related to the attack behavior, a time sequence threat sniffer is constructed to calculate the event E j E k The sniffer is denoted as H θ (E j ,E k ), inheriting the encoder from the fine-tuning stage of the model and combining it with the MLP layer, using these scores to optimize the Monte Carlo tree search in the attack traceability node expansion process. The input and training process are described as follows: Let the target event be Each candidate event The input of the sniffer is recorded as Here vec(·) is a function that converts raw event data into a vector. Let vec(E i )for Where X is the node feature matrix, and Time(·) is a learnable function that embeds the real-valued timestamp into the vector. Inspired by Transformer, the present invention uses a relative dynamic time encoding function to calculate the time embedding within the interval, as shown in formula (10).

[0119] RTE(ΔT)=TimeLinear(RT(ΔT)) (10)

[0120] Among them, ΔT represents the position, TimeLinear is a learnable linear projection function, and RT(ΔT) is shown in Equation (11) and Equation (12).

[0121]

[0122] Where dim is the dimension of the position encoding corresponding to the sine wave. Relative temporal embedding enhances the feature representation of source and target nodes, facilitating the capture of dynamic information related to the temporal relationship between these nodes.

[0123] In this way, the detected attack event and each candidate event (i.e., all events before the attack event) are input into H θ The output is the corresponding correlation score between them.

[0124] Predicting the correlation score can be regarded as a regression task, so the mean square error (MSE) is used as the loss function when training the sniffer. At the same time, considering that the events in the APT attack sequence have strong time dependence, a time-series-based correlation score is designed and used as the training target of the sniffer to better capture the temporal correlation between events. The loss function is shown in formula (13).

[0125]

[0126] Where N represents the number of samples, y i is the model output, is the target relevance score, determined by formula (14).

[0127]

[0128] Where e is the natural logarithm, β is the decay rate constant, and ΔT is the timing difference between event pairs.

[0129] In the process of tracing the attack path, a root node is first initialized to represent the set of detected attack events. Monte Carlo tree search expands the nodes in the search tree through multiple rounds of simulation, and each node represents a sub-event set. Each round of search includes four steps: node selection, node expansion, simulation, and back propagation. The pre-trained model (including encoder and decoder) is used to further evaluate and calculate rewards for the newly generated sub-nodes.

[0130] Initialization: The root node is a set of candidate event edges that meet the spatiotemporal conditions. For example, suppose the target is tracing the attack event e 100For the tracking attack step, the root node set of events seen by the encoder is {e1,e2,…,e 99}.

[0131] Node selection: Starting from the root node, child nodes are selected according to the upper confidence bound (UCT) strategy to balance development and exploration. The selection criteria designed by the present invention are shown in formula (15).

[0132]

[0133] in, Represents node N i The set of events that have been expanded, c(N i ,e j ) indicates that at node N i Select e j The cumulative reward, n(N i ,e j ) is node N i Middle side j The number of times selected, ∑n(N i ,e l ) represents the total number of visits to all edges under the node. The algorithm is implemented by i Remove event e θ Move to N i The left part The development part represents the known node rewards. It makes the model more inclined to select nodes with known high rewards. The inverse relationship is to encourage more development of nodes with high average rewards, while avoiding overdevelopment of nodes that have been selected many times. It allows the model to pay more attention to nodes with high average rewards, rather than simply making choices based on cumulative rewards, thus finding a balance between development and exploration. The exploration part is to find new possible paths and give more weight to nodes that have not been fully explored. The purpose is to avoid falling into the local optimal solution and find better solutions that may exist in the search space. i ,e j )) term, in order to dynamically adjust the decay speed of the exploration term, so that the exploration weight decreases smoothly as the number of times the node is selected gradually increases, rather than decaying to zero too quickly. This adjustment allows the model to explore new nodes more actively in the early stage of the search, but gradually turn to development as the number of explorations increases. λ(t) is the weight that changes dynamically with the number of search steps t, as shown in formula (16).

[0134] λ(t)=λ0·exp(-α·t) (16)

[0135] Where λ0 is the initial value, which controls the early exploration weight. α is the decay rate, which controls the speed at which λ decreases, and t is the number of search steps. In the early stage of the search, λ(t) is close to λ0, which means that the exploration weight is high. As the number of search steps t increases, λ(t) will gradually decay, and the model will focus more on development. The form of exponential decay makes the transition from exploration to development smoother, avoiding sudden switching of strategies at a certain moment, thereby improving the stability and efficiency of the search.

[0136] The selection strategy designed by the present invention can ensure that even if the node has been selected multiple times, the intensity of exploration will not drop rapidly. This is very important for long-term series attacks that may be hidden in APT attacks, ensuring that the exploration will not end prematurely and that more potential attack paths can be discovered. APT attacks are usually multi-stage and complex, and the exploration mechanism of this strategy can flexibly respond to changes in attack scenarios. When the attack steps are scattered in different time periods, the model can still maintain monitoring and detection of different paths through longer explorations, so that key attack steps will not be missed.

[0137] Node expansion: In order to optimize the search space and improve the quality of node selection, a strategy is adopted to expand only nodes that are potentially related to the attack event. Instead of expanding all possible child nodes on each selected node, the search space is narrowed down through the sniffer. If the selected node N i is considered scalable, the sniffer calculates the potential score as shown in Equation (17).

[0138]

[0139] in represents the events that were not extended in the previous simulations by discarding the least significant events e β To retain the most relevant events to expand new nodes. The scores are efficiently inferred with minimal computational cost using the pre-trained sniffers inherited from the pre-training module. The selection and expansion of nodes occur in an alternating manner. Starting from the root node, the optimal event is selected and a new child node is generated according to equation (13), and then the new child node is expanded according to equation (17). Next, the process of expanding and selecting from new nodes is repeated. The process ends when the current node is determined to be a leaf node or the number of nodes is less than the hyperparameter k. β The least important events are selected in order to find the nodes that need to be expanded from the candidate nodes that are not currently expanded. This process is mainly to optimize the search space and focus on the nodes that are most likely to be related to the attack events. θ Refers to the node selection process from the current node N i From the set of candidate child nodes, select the best node for further exploration.

[0140] Simulation: Use the pre-training and fine-tuning modules to simulate rewards for new child nodes and for leaf nodes N leaf Calculate the corresponding reward value r(N leaf ).

[0141] Backward propagation: propagate the reward value upward from the leaf node, updating the reward and visit count of each node on the path. All nodes N from the root to the leaf node i Update n(·,·) and c(·,·) by accumulating the rewards of leaf nodes and adding 1. Specifically, they are shown in Equations (18) and (19).

[0142] n(N i ,e θ )=n(N i ,e θ )+1 (18)

[0143] c(N i ,e θ )=c(N i ,e θ )+r(N leaf ) (19)

[0144] Finally, the set of nodes that received the highest reward is considered the tracing result. Through this reinforcement learning strategy, the model can dynamically adjust node selection, enhance the exploration of attack paths, gradually track the complete path of APT attacks, and ultimately find each step of the attacker's operation and its causal relationship, thereby providing security experts with more comprehensive tracing information.

[0145] 4.1 Experimental Environment and Dataset

[0146] The public datasets ATLAS-M2 and ATLAS-M5 used in the present invention are from real APT attack scenarios, including system, DNS, and Firefox logs. The M2 dataset corresponds to phishing attacks against web browsers, while the M5 dataset represents small-scale attack activities against Microsoft Office. The sizes of M2 and M5 are 671.112MB and 726.113MB, respectively. The malicious behavior rate of the M2 dataset is 13.68%, and the malicious behavior rate of the M5 dataset is 5.33%. Both datasets contain attacks across multiple victim hosts. Initially, the attacker replaced the benign web page in the victim's system with a malicious web page on the first host. Subsequently, on the second host connected to the first host, the victim visited the malicious web page, resulting in a compromise. Each attack against multiple hosts collects log files obtained from the respective victim hosts.

[0147] The algorithm is implemented using Python 3.9.0 and PyTorch 1.7.0. The experiment is conducted on a PC with Intel Core i9 2.5GHz, 32GB RAM, Windows 11 64-bit, and the GPU is NVIDIA RTX3060. Considering the number of attack samples, the training set, validation set, and test set are divided according to the timestamp.

[0148] 4.2 Evaluation Metrics

[0149] For attack detection, the area under the curve (AUC) is used as the performance indicator. AUC measures the accuracy of the classifier in determining whether a positive sample is ranked higher than a negative sample. The value range is 0 to 1. The larger the value, the better the detection performance.

[0150] Fidelity and sparsity are used to evaluate the traceability performance. Fidelity reflects the accuracy of the model in faithfully retaining attack-related behaviors from the input data. The definition of fidelity is shown in formula (20).

[0151]

[0152] Where N e is the number of edges, Π is the indicator function, f(·) is the neural network model of the output graph embedding, t i is the time point, R i is time point t i The sampled subgraph, G i is time point t i The original picture.

[0153] In addition to fidelity, effective attribution should be concise, focusing on capturing the most relevant input features while ignoring any irrelevant features. This feature is quantified by the metric sparsity, which specifically evaluates the proportion of features considered important by the attribution method. Sparsity is defined as Represents the proportion of zeros in model parameters or activation values. Sparse models show the scarcity of parameters or activation values ​​that are critical to the final prediction between normal and abnormal behaviors. Higher fidelity and higher sparsity mean better traceability. The paper calculates the area under the fidelity-sparsity curve (AUFSC) to evaluate performance. The larger the AUFSC, the better the performance.

[0154] The present invention has passed experiments, and the APT attack detection and tracing model has ideal effects, which are consistent with the design expectations.

[0155] In order to further verify the performance of the algorithm proposed in the present invention, the algorithm proposed in the present invention is compared with other classic algorithms such as TGNN, CDHGN and TRACE2VEC in terms of AUC value and fidelity. The experimental results are shown in Tables 1 and 2.

[0156] Table 1 Comparison of anomaly detection results on different datasets

[0157]

[0158] As can be seen from Table 1, the proposed method outperforms various baselines on different partitions of the two datasets. Taking the ATLAS-M2 dataset as an example, the method achieves an optimal AUC of 0.9674 on this dataset, which is 6.2%, 5.4%, and 2.5% higher than TGNN, CDHGN, and TRACE2VEC, respectively. The optimal AUC of 0.9796 is achieved on the ATLAS-M5 dataset, with the largest improvements reaching 10.3%, 8.4%, and 2.6%, respectively. Compared with ATLAS-M2, the ATLAS-M5 dataset is more unbalanced and has a lower anomaly rate. In this case, weighted cross entropy can fully highlight its advantages, thereby achieving higher gains.

[0159] Among the baseline methods, TGNN uses a homogeneous graph model (CDGN) to predict abnormal events; on this basis, CDHGN introduces a heterogeneous attention mechanism for nodes and edges in continuous-time dynamic graphs, further improving the performance; and TRACE2VEC achieves additional performance improvement by introducing a data augmentation module based on negative sampling in the pre-training stage. Compared with TRACE2VEC, the proposed AGNN and contrastive loss can work together in the pre-training stage, making the model more focused on the similarities between attack steps, significantly improving the ability to detect attack behaviors. Subsequently, weighted cross entropy is introduced in the fine-tuning stage to pay extra attention to unbalanced samples and enhance the generalization ability of the model. Experimental results show that the proposed method shows significant advantages in detecting multi-step complex attacks.

[0160] Table 2 Comparison of explanation results on different datasets

[0161]

[0162] As can be seen from Table 2, the proposed method still outperforms all existing baselines in attack tracing. Specifically, on the ATLAS-M2 dataset, the highest fidelity and best AUFSC achieved by this method reached 1.9657 and 1.5893, respectively, which are 9.5% and 10.1% higher than the best baseline, respectively. The same significant improvement was observed on the ATLAS-M5 dataset. Compared with the baseline, the key difference of this method is that the reinforcement learning mechanism is introduced and the correlation score between events is obtained through the timing sniffer. The tracing process is transformed into a sequential decision problem through reinforcement learning, so that the model can trace back its reward-based selection process after the inference or tracing is completed, thereby providing more intuitive interpretability. In addition, the timing sniffer incorporates time information into node selection and expansion to ensure the rationality and effectiveness of the tracing path. Overall, compared with the existing state-of-the-art baselines, this method not only achieves higher performance in attack tracing, but also shows better interpretability.

[0163] In order to verify the contribution of the AGNN method of generating adversarial samples and the loss functions designed in the pre-training and fine-tuning stages to the algorithm proposed in this invention, the following experiments were conducted in the same experimental environment:

[0164] Experiment 1: AGNN + classic GNN loss function + cross entropy

[0165] Experiment 2: GNN + contrast loss function + cross entropy

[0166] Experiment 3: AGNN + contrast loss function + cross entropy

[0167] Experiment 4: The algorithm proposed by this invention is AGNN + contrast loss function + weighted cross entropy

[0168] Table 3 Ablation experiment of anomaly detection

[0169]

[0170] According to the results in Table 3, both the first and second group experiments prove that using only AGNN or only contrast loss can provide the model with certain anomaly detection capabilities, but the effect is still relatively limited. In the third group of experiments, after using AGNN and contrast loss together, the model performance was significantly improved. The specific reason is that AGNN generates more complex and diverse samples, providing the model with richer training data. At the same time, under the guidance of contrast loss, the model can more fully learn the feature differences between positive and negative samples, thereby significantly improving the robustness and accuracy against APT attacks. Finally, in the fourth group of experiments, the introduction of weighted cross entropy helps the model focus more on the sparse attack steps hidden in normal behavior, while effectively dealing with the problem of imbalance in the number of positive and negative samples, thereby further improving the detection performance and achieving the best effect.

[0171] Figure 2 and Figure 3 The graph embedding visualization results obtained on different datasets after combining AGNN with contrastive learning are shown. (a) is without AGNN and contrastive learning, and (b) is with AGNN and contrastive learning. It can be seen intuitively that normal samples show a more cohesive distribution structure in the feature space, while attack samples are concentrated in the peripheral area. This shows that the designed contrastive loss can better capture the similarities between APT attack behaviors, can effectively bring similar samples closer to each other, and push samples with large differences away, making the distribution boundary between normal samples and attack samples in the graph embedding space clearer, further improving the model's ability to discriminate multi-step attack behaviors.

[0172] Figure 4 and Figure 5 The loss curves of the model at different training stages are shown. It can be clearly seen that as the training continues, each loss indicator shows an overall downward trend. This result not only shows that the model can converge robustly at different stages, but also shows that it has good adaptability to diverse data distributions.

[0173] In addition, in order to verify the effectiveness of reinforcement learning for attack tracing, two sets of experiments were designed while retaining other modules:

[0174] Experiment 1: Search-based attack tracing

[0175] Experiment 2: Attack Tracing Based on Reinforcement Learning

[0176] Table 4 Attack tracing ablation experiment

[0177]

[0178] From the data in Table 4, we can see that after the introduction of reinforcement learning, the fidelity of the attack sequence obtained by the model in the tracing process has been significantly improved. This is mainly due to: first, the time encoding embedding in the event provides the model with temporal context information in the multi-step attack sequence, enabling it to capture potential causal relationships; second, the sniffer based on temporal information can mine more accurate correlations between events, thereby retaining more critical events when the node is expanded.

Claims

1. A method for detecting and tracing APT attacks based on graph neural networks and reinforcement learning, characterized by: The following steps are involved: Step S1: Use a continuous-time dynamic heterogeneous graph to model the events in the system audit log, generate more challenging samples through AGNN, improve the quality of samples, and make the generated positive samples more complex and close to the real attack scenario; Step S2: In the pre-training stage, contrastive learning is used to optimize positive and negative samples, so that the model can be more sensitive when facing sparse positive samples, so as to better detect multi-step behaviors in APT attacks; Step S3: In the fine-tuning stage, according to the sample type, the edge embedding is obtained by concatenating the embeddings of the nodes on both sides of the edge. The fully connected layer maps the embedding to the probability that the connected edge belongs to an attack event, and the weighted cross entropy loss function is used as the fine-tuning target. Step S4: Use Monte Carlo tree search combined with reinforcement learning to gradually trace the events in the system.

2. According to claim 1, a method for detecting and tracing APT attacks based on graph neural network and reinforcement learning is characterized in that: The specific steps of step S1 are: Step S1.1, using a continuous-time dynamic heterogeneous graph to model events in the system audit log; The system behavior is represented by a graph G = (V, E), where V is a set of nodes representing entities in the system, and E is a set of edges representing the relationships between entities. For each node v∈V, a node feature matrix X is defined, and for each edge e∈E, an edge feature matrix A is defined. Based on these feature matrices, a graph embedding model is constructed. Step S1.2: In the graph modeling process, GNN is used to learn the embedding of nodes and edges. In order to capture the dynamics of the time dimension, a message passing mechanism is used to update the status of the nodes. The update method is shown in formula (1): in, represents the embedding representation of node v at time step t, represents the neighbor nodes of node v, W m and W s is the weight matrix, σ is the activation function; Step S1.3: The graph neural network generates the feature representation of the global graph by reading out the local node features aggregated by the function R, as shown in formula (2): H represents the embedding set of all nodes, and s is the embedding representation of the global graph. Through the above method, the overall behavioral characteristics of the system can be obtained, including the embedding representation of normal behavior and attack behavior.

3. The APT attack detection and tracing method based on graph neural network and reinforcement learning according to claim 1 is characterized in that: The specific steps of generating more challenging samples through AGNN in step S1 are as follows: The adversarial perturbation generation formula is shown in formula (3) Where δ is the adversarial perturbation vector, ∈ is the perturbation size constraint, and f θ is the prediction function of GNN, X and A represent the node feature matrix and adjacency matrix respectively, and y is the true label; The adversarial training of AGNN needs to combine the original task loss and adversarial loss, and the comprehensive loss function As shown in formula (4) in, is the original task loss, used to ensure the performance of the model under normal input, is the adversarial loss, which is used to make the model robust when subjected to adversarial perturbations, and λ is the weight that balances the two parts of the loss; When generating adversarial samples, more challenging node embeddings are generated through adversarial perturbations, and the node v i The adversarial embedding representation of is shown in formula (5): Among them, h i For node v i The original embedding of i Adversarial perturbations are generated by maximizing the model loss. Adversarial node embeddings are used to generate more complex adversarial samples, helping the model improve its ability to recognize attack behaviors.

4. The APT attack detection and tracing method based on graph neural network and reinforcement learning according to claim 3 is characterized in that: AGNN can also perturb the structure of the graph to generate more complex samples. The adjacency matrix A after perturbation is * It is expressed as shown in formula (6) A * =A+ΔA (6) Where: ΔA is the adversarial perturbation matrix, representing a small modification on the graph structure; AGNN generates adversarial samples by optimizing formula (7) to improve the generalization performance of the model. The inner maximization operation of the formula is used to generate the most aggressive adversarial perturbations, while the outer minimization operation is used to allow the model to achieve good performance even in this most unfavorable situation.

5. The APT attack detection and tracing method based on graph neural network and reinforcement learning according to claim 1 is characterized in that: The specific steps of step S2 are as follows: The positive sample is recorded as (X, A), and the negative sample generated based on (X, A) is recorded as (X * ,A * ); By constructing positive and negative sample pairs, the contrast loss function is used for training. The contrast loss function is defined as shown in formula (8) Among them, y i is the label of sample i, positive sample is 1, negative sample is 0, d(h i ,s) represents the positive sample embedding h i The distance between the positive sample and the global feature s, m is the distance threshold, the distance of the positive sample should be greater than m, that is, if the distance is less than m, the model will be penalized and the distance will be increased; if the distance is already greater than m, the loss of this item is zero, which is used to ensure that the distance between the positive sample and the global feature is greater than m. The total number of samples is N, which represents the number of positive and negative samples. The loss of all samples is normalized so that the value of the loss function is not affected by the number of samples.

6. The APT attack detection and tracing method based on graph neural network and reinforcement learning according to claim 1 is characterized in that: In the encoder parameter fine-tuning stage, both the encoder and the decoder are adapted to the limited real attack samples. According to their types, the edge embedding is obtained by concatenating the embeddings of the nodes on both sides of the edge. Then the fully connected layer maps the embedding to the probability that the connected edge belongs to an attack event. The weighted cross entropy loss function shown in formula (9) is used as the fine-tuning target Among them, y x Indicates whether edge x is a true attack event, 0 or 1, p x Indicates that the prediction x belongs to attack event y x =1, ω represents the weight corresponding to the attack behavior, which is used to amplify the impact of sparse attack signals and increase the weight of attack behavior, so that the model can pay more attention to a small number of attack samples.

7. The APT attack detection and tracing method based on graph neural network and reinforcement learning according to claim 1 is characterized in that: The specific steps of step S4 are as follows: Build a timing threat sniffer to calculate event E j E k The relevance score of the sniffer is denoted by H θ (E j ,E k ), inheriting the encoder from the fine-tuning phase of the model and combining it with the MLP layer, using these scores to optimize the Monte Carlo tree search during the attack traceability node expansion process; The input and training process are described as follows: Let the target event be Each candidate event The input of the sniffer is recorded as vec(·) is a function that converts raw event data into a vector; let vec(E i )for Where X is the node feature matrix, Time(·) is a learnable function that embeds the real-valued timestamp into the vector; the relative dynamic time encoding function is used to calculate the time embedding within the interval, as shown in formula (10) RTE(ΔT)=TimeLinear(RT(ΔT)) (10) Where ΔT represents the position, TimeLinear is a learnable linear projection function, and RT(ΔT) is as shown in Equation (11) and Equation (12): Among them, dim represents the dimension of each dimension of the position encoding corresponding to the sine wave; Input the detected attack events and each candidate event into H θ The output is the corresponding correlation score between them.

8. The APT attack detection and tracing method based on graph neural network and reinforcement learning according to claim 7 is characterized in that: The prediction of the correlation score is regarded as a regression task. The mean square error (MSE) is used as the loss function when training the sniffer. At the same time, a time-series-based correlation score is designed as the training target of the sniffer. The loss function is shown in formula (13): Where N represents the number of samples, y i is the model output, is the target relevance score, determined by formula (14) Where e is the natural logarithm, β is the decay rate constant, and ΔT is the timing difference between event pairs.

9. The APT attack detection and tracing method based on graph neural network and reinforcement learning according to claim 8 is characterized in that: Monte Carlo tree search expands the nodes in the search tree through multiple rounds of simulation. Each node represents a sub-event set. Each round of search includes four steps: node selection, node expansion, simulation, and back-propagation. The pre-trained model, including the encoder and decoder, is used to further evaluate and calculate rewards for the newly generated sub-nodes. Initialization: The root node is a set of candidate event edges that meet the spatiotemporal conditions; Node selection: Starting from the root node, select child nodes according to the upper confidence bound UCT strategy to balance development and exploration. The selection criteria are shown in formula (15) in, Represents node N i The set of events that have been expanded, c(N i ,e j ) indicates that at node N i Select e j The cumulative reward, n(N i ,e j ) is node N i Middle side j The number of times selected, ∑n(N i ,e l ) represents the total number of visits to all edges under the node; λ(t) is the weight that changes dynamically with the number of search steps t, as shown in formula (16): λ(t)=λ0·exp(-α·t) (16)where λ0 is the initial value, controlling the early exploration weight, α is the decay rate, controlling the speed at which λ decreases, and t is the number of search steps. In the early stage of the search, λ(t) is close to λ0.

Images