Network virus propagation prediction method based on space-time diagram attention model

Abstract

The invention provides a network virus propagation prediction method based on a space-time diagram attention model, and mainly solves the problems of strong dependence on propagation parameters and low prediction precision in the prior art. The scheme comprises the following steps: 1) constructing a space-time infection state diagram based on a network topology structure and virus historical transmission data; 2) processing the historical virus transmission data by using a sliding window mechanism, and constructing a data set; 3) constructing a prediction model based on the space-time diagram attention network in combination with the diagram attention network and the long-short-term memory network; 4) training the prediction model by using the constructed data set, and evaluating the prediction performance of the model; and 5) predicting the propagation trend of the future network virus by adopting the trained model. According to the method, the spatial topological characteristics and the time evolution law of virus propagation are effectively combined, and meanwhile, a graph attention mechanism is introduced, so that the prediction accuracy and the calculation speed of virus propagation in a large-scale network environment can be remarkably improved.

Description

Technical Field

[0001] The present invention belongs to the technical field of network information security, and further relates to virus spread prediction technology. Specifically, it is a network virus spread prediction method based on a spatio-temporal graph attention network, which can be used to predict the spread trend of network viruses in advance and contain the spread and harm of viruses. Background Art

[0002] With the advent of the big data era, network security issues have increasingly become the focus of global attention. As a kind of malicious software with self-replicating ability, transmission dependence and destructiveness, network viruses rely on the extensive connection of computer networks and can quickly infect executable files, system services or network protocols of target devices, and perform malicious operations such as data theft, system paralysis, and resource hijacking during the transmission process. Network virus attacks not only cause huge economic losses and resource waste, but also pose a severe challenge to global network security and information stability.

[0003] The spread of network viruses often takes advantage of zero-day vulnerabilities for rapid spread, which makes it difficult for traditional computer security protection means, such as antivirus software, intrusion detection system IDS and firewalls, to cope with the threats of new viruses. And there is also a certain time lag in newly released security patches, unable to provide real-time protection. Therefore, virus spread prediction, as a proactive defense strategy, can effectively make up for the deficiencies of traditional protection measures. By predicting the spread trend of network viruses in advance, especially in the initial stage of virus outbreak, targeted defense measures can be implemented, thus significantly improving the defense ability against new network viruses, slowing down the spread of viruses, and ultimately effectively containing the spread and harm of network viruses.

[0004] The current virus spread prediction methods mainly fall into three categories: 1) Network virus spread prediction methods based on mathematical models, with the typical representative being the SIR model. This model was originally derived from the classic infectious disease model in epidemiology and was later introduced to study the spread of network viruses. Here, S, I, and R respectively represent the nodes in the network in the susceptible, infected, and immune states. Although the SIR model has a simple structure, is easy to understand, and can provide an effective framework for the preliminary analysis of network virus spread, it relies on a series of idealized assumptions that do not comprehensively consider the complexity of virus spread in the real network environment. For example, it is often difficult to accurately estimate the selection of model parameters such as the infection rate and recovery rate, and these parameters are usually assumed to be constants, ignoring the impact of network dynamic characteristics on the spread process. With the complexity of network topologies, the applicability of the traditional SIR model is gradually limited and cannot accurately reflect the real dynamics of virus spread in complex networks. Therefore, its application in modern network environments faces significant limitations. 2) Prediction methods based on network dynamics and topological structures, aiming to simulate and analyze virus spread by combining the real network topology. In this type of method, the nodes in the network represent computers or other devices accessing the network, and the edges represent the connections between nodes. By simulating the spread path and scope of the virus in the network topology, the spread trend and impact of the virus are then predicted. However, the network topologies and structures used in most existing methods are usually static and do not fully consider the dynamic changes in the topology of the real network environment. This results in significant prediction biases when these methods face network topology evolution and changes. In addition, similar to the methods based on mathematical models, these methods also face difficulties in accurately estimating parameter selection, further limiting their application effects in complex network environments. 3) Virus spread trend prediction methods based on deep learning. Most existing methods use recurrent neural networks for prediction, among which LSTM, the long short-term memory network, as an efficient model for processing sequential data, is widely used in time series prediction. These methods collect historical spread data of network viruses to train prediction models, thus avoiding the dependence on parameter selection in traditional mathematical models. Nevertheless, these methods have strong time dependence and mainly focus on capturing the time dynamics of virus spread, making it difficult to effectively consider the relationships between nodes in the spatial topology and their impact on virus spread, limiting their prediction accuracy in complex network environments. Summary of the Invention

[0005] The object of the present invention is to address the deficiencies of the above-mentioned existing technologies and propose a network virus spread prediction method based on a spatio-temporal graph attention model to solve the problems of limited application and poor detection effects of existing prediction technologies in complex network environments. The present invention aims to provide more accurate spread predictions by simultaneously capturing the spatial features and time dynamics during the spread process of network viruses, providing theoretical support and practical guidance for timely curbing virus spread and implementing effective prevention and control measures.

[0006] The technical idea of the present invention is as follows: First, according to the network topology structure and virus historical transmission data, a spatio-temporal infection status graph is constructed as a data set. Then, the graph attention network GAT is used to extract the spatial features between nodes, and at the same time, the long short-term memory network LSTM is combined to extract the time dynamics in the virus transmission process, and a spatio-temporal graph attention network ST-GAT model is constructed. Then, the ST-GAT model is trained using the spatio-temporal infection status graph data set, and the prediction accuracy of the model is evaluated. Finally, the trained ST-GAT model is used to predict the future transmission trend of network viruses and give early warnings to the nodes that may be infected, so as to achieve early defense responses.

[0007] To achieve the above object, the technical solution of the present invention includes the following steps:

[0008] (1) According to the transmission characteristics of network viruses, the computer network is represented as an undirected graph with an infection status time series, which is used to describe the process of the infection status of devices in the network changing over time; where is the node set at time , and the historical transmission data of the virus is used as the feature of the node; represents the edge set at time , and an edge is formed between any two nodes that communicate with each other in the undirected graph; represents the adjacency matrix of the spatio-temporal infection status graph G at time ;

[0009] (2) Use the sliding window mechanism to process the virus historical transmission data to construct a data set; and record the infection status of each node at future time points as the true label , and its shape is recorded as , where represents the total number of nodes, represents the prediction result step length of each node;

[0010] (3) Construct a prediction model based on the spatio-temporal graph attention network, and the implementation steps are as follows:

[0011] (3.1) Take the node feature matrix and the adjacency matrix of the spatio-temporal infection status graph G as the input data of the model;

[0012] (3.2) Construct two layers of graph attention layers GAL to learn the spatial dependence relationship between nodes in graph G; where the output feature of the first layer of GAL is , and the spatial feature of the spatio-temporal infection status graph output by the second layer of GAL is calculated and obtained using the feature :

[0013] ,

[0014] ;

[0015] Among them, represents the calculation of the graph attention layer GAL, that is, using the multi-head attention mechanism to calculate the hidden state and take the average value of it to obtain the aggregated node features as the output;

[0016] (3.3) Construct two LSTM layers to capture the dynamic information of the time series, process the time series data, and learn the time dynamics of virus transmission; the features output by the graph attention layer are used as the input of the first LSTM layer to obtain the third output feature , and then is used as the input of the second LSTM layer to obtain the fourth output feature ;

[0017] (3.4) Map the output of the LSTM layer through a fully connected layer to obtain the final prediction result vector :

[0018]

[0019] (3.5) Adjust the prediction result vector to the same shape as the true label according to the following formula:

[0020]

[0021] Among them, is the predicted value of each node, indicating the infection status of the node at the future time window point; is the batch size of the number of samples during training; is the total number of nodes in the network;

[0022] (4) Divide the dataset into a training set, a validation set, and a test set, select the mean squared error loss MSE Loss as the model loss function, and use optimizer to train the constructed prediction model to obtain the final trained prediction model;

[0023] (5) Input the test set data into the final prediction model to realize the prediction of the future spread of the virus.

[0024] Compared with the prior art, the present invention has the following advantages:

[0025] First, the network virus spread prediction method of the present invention combines the Graph Attention Network (GAT) and the Long Short-Term Memory Network (LSTM), which can effectively extract the spatial features between nodes and the temporal dynamics during the virus spread. Through the GAT part, the model can capture the spatial dependence relationships between network nodes, thus accurately reflecting the spread patterns between nodes. The LSTM network processes the time series changes during the virus spread, enabling the model to dynamically track the spread trend and thus more accurately predict the spread of the virus between nodes.

[0026] Second, the prediction method of the present invention can achieve node-level infection situation prediction, and can predict the specific infection status of each node during the virus spread, thus providing an accurate prediction basis for targeted prevention and control and decision-making. Compared with traditional methods based on mathematical models, the present invention does not rely on idealized assumption conditions or difficult-to-accurately-estimate spread parameters. Instead, by combining the graph attention mechanism and the time series model, it effectively captures the complex relationships and spread dynamics between nodes, ensuring that the prediction results are more practical and accurate.

[0027] Third, since the present invention introduces the graph attention mechanism, when the model processes a large number of network nodes, it can accurately capture the neighbor nodes that are most closely related to the target node, and thus more accurately reflect the actual situation of the virus spread. Through the graph attention mechanism, the model can automatically adjust the influence weights between nodes, enhancing the accuracy of the prediction. In addition, the model supports parallel computing, greatly improving the computing speed of the prediction process and meeting the real-time prediction requirements in a large-scale network environment. BRIEF DESCRIPTION OF THE DRAWINGS

[0028] Figure 1 is the implementation flowchart of the present invention;

[0029] Figure 2 is the implementation flowchart of the iterative training of the virus spread prediction model based on the spatio-temporal graph attention network of the present invention;

[0030] Figure 3 is the node degree distribution diagram provided in the embodiment of the present invention;

[0031] Figure 4 is the trend graph of the change of the loss value loss during the model training provided in the embodiment of the present invention.

[0032] Figure 5 is the training result graph of the trained spatio-temporal graph attention prediction model on the test set provided in the embodiment of the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0033] The following will further describe the present invention in detail with reference to the drawings and specific embodiments.

[0034] Embodiment 1: Refer to the appendixFigure 1 , a network virus propagation prediction method based on a spatio-temporal graph attention network proposed by the present invention specifically includes the following steps:

[0035] Step 1. According to the propagation characteristics of network viruses, represent the computer network as an undirected graph with an infection state time series , which is used to describe the process of the infection state of devices in the network changing over time; where is the node set at time , and the historical propagation data of the virus is used as the characteristics of the nodes; represents the edge set at time , and an edge is formed between any two nodes that communicate with each other in the undirected graph; represents the adjacency matrix of the spatio-temporal infection state graph G at time . In this embodiment, the nodes constituting the node set represent devices in the network, including but not limited to computers, servers, and routers. The adjacency matrix of the above spatio-temporal infection state graph G is expressed as follows:

[0036]

[0037] where, if there is an edge between the i-th node and the j-th node , then , otherwise ; i, j = 1, 2,..., N.

[0038] Step 2. Use the sliding window mechanism to process the virus historical propagation data to construct a data set; and record the infection state of each node at future time points as the true label , and its shape is recorded as , where represents the total number of nodes, represents the prediction result step length of each node.

[0039] In this embodiment, the above use of the sliding window mechanism to process the virus historical propagation data is specifically: set the window size, including the historical data window size and the prediction target data window size , set the sliding step length to 1; perform window sliding operations on the historical data to obtain the node feature matrix in the spatio-temporal infection state graph G, and the implementation steps are as follows:

[0040] (2.1) Under the sliding window mechanism, the node feature matrix contains the features of each node in the historical window, and the feature vector of the i-th node is expressed as :

[0041] ,

[0042] Among them, represents the feature vector of the i-th node when the size of the historical data window is p, where p = 1, 2,... ; is the set of device infection states, and the infection states include the susceptible state , the latent state , the infected state and the immune state , which are obtained according to the infection situation of devices in the network.

[0043] (2.2) Represent the node feature matrix as follows:

[0044] .

[0045] Step 3. Construct a prediction model based on the spatio-temporal graph attention network, and the implementation steps are as follows:

[0046] (3.1) Use the node feature matrix and the adjacency matrix of the spatio-temporal infection state graph G as the input data of the model;

[0047] (3.2) Construct two layers of graph attention layers GAL to learn the spatial dependence relationship between nodes in graph G; the output feature of the first layer of GAL is , and use the feature to calculate and obtain the spatial feature of the spatio-temporal infection state graph output by the second layer of GAL:

[0048] ,

[0049] ;

[0050] Among them, represents the calculation of the graph attention layer GAL, that is, using the multi-head attention mechanism to calculate the hidden state and take the average of it to obtain the aggregated node feature as the output. In this embodiment, the above multi-head attention mechanism specifically learns different feature representations through multiple independent attention heads, so as to improve the modeling ability of the model for complex propagation relationships. The expression is as follows:

[0051]

[0052] Among them, is the number of attention heads, is the activation function, is the The trainable weight matrix of the attention heads; is the feature of the i-th node after calculation by the attention mechanism; is the attention coefficient between the i-th node and its neighbor node q, defined as follows:

[0053]

[0054] where, is the linear transformation matrix of the node features, is the weight vector used to calculate the attention coefficient; represents the activation function, which is used to increase non-linear features; represents all the neighbor nodes of the i-th node.

[0055] (3.3) Construct two LSTM layers to capture the dynamic information of the time series, process the time series data, and learn the time dynamics of virus spread; use the features output by the graph attention layer as the input of the first LSTM layer to obtain the third output feature , and then use as the input of the second LSTM layer to obtain the fourth output feature ;

[0056] (3.4) Map the output of the LSTM layer through a fully connected layer to obtain the final prediction result vector :

[0057]

[0058] (3.5) Adjust the prediction result vector to the same shape as the true label according to the following formula:

[0059]

[0060] where, is the predicted value of each node, indicating the infection status of the node at the future time window point; is the batch size of the number of samples during training; is the total number of nodes in the network;

[0061] Step 4. Divide the dataset into training set, validation set, and test set, select the mean squared error loss MSE Loss as the model loss function, and use optimizer to train the constructed prediction model to obtain the final trained prediction model; the training process implemented in this step of this embodiment is as follows:

[0062] (4.1)Divide the dataset constructed in step 2 into a training set, a validation set, and a test set in the ratio of 80%:10%:10%; the training set is used for the training process of the model to enable the model to learn the patterns and features in the data and optimize the model parameters through gradient descent; the validation set is used to adjust the hyperparameters and evaluate the generalization ability of the model during the training process; the test set is used to evaluate the performance of the model.

[0063] (4.2)Select the mean squared error loss (MSE Loss) as the model loss function and calculate the square difference between the predicted value and the true label :

[0064]

[0065] where is the prediction result of the model, is the true label of node .

[0066] (4.3)Select the optimizer as the optimization algorithm for the model, and use the learning rate and weight decay as parameters; it is used to improve the stability and generalization ability of the training, and the parameter update formula is:

[0067]

[0068] where and are the estimates of the first moment and the second moment of the gradient respectively, is the numerical stability term. In this embodiment, the learning rate is preferably taken as ; the weight decay is introduced in the optimization process, that is, a regularization term is added to the loss function:

[0069]

[0070] where is the weight decay coefficient, and in this embodiment, the preferred value is , is the mean squared error loss, is the final optimized objective function.

[0071] (4.4)Divide the dataset into multiple batches according to the size of . Each batch performs one forward propagation and one backward propagation through the model, and the training data is iteratively trained through the data loader to optimize the model parameters.

[0072] ​Step 5. Input the test set data into the final prediction model to predict the future spread of the virus.

[0073] Example 2: The overall implementation steps of the network virus spread prediction method proposed in this example are the same as those in Example 1. Now, with reference to the attached Figure 1-2 Specific parameter settings are given to further describe the implementation process of the present invention in detail:

[0074] Step 1) Define the spatio-temporal infection status graph :

[0075] According to the spread characteristics of the network virus, the computer network is represented as an undirected graph with an infection status time series , which is used to describe the process of the infection status of devices in the network changing over time. Among them: is the node set at time . The nodes represent devices in the network, such as computers, servers, routers, etc. The historical spread data of the virus is used as the characteristics of the nodes; represents the edge set at time , indicating the connection relationship between nodes in the network. If two nodes and can communicate with each other, there is an edge in the graph; represents the adjacency matrix of the spatio-temporal infection status graph G at time , which is expressed as follows:

[0076]

[0077] Among them, if there is an edge between the i-th node and the j-th node , then , otherwise ; i, j = 1, 2,..., N.

[0078] Step 2) Construct the data set:

[0079] Use the sliding window mechanism to process the historical spread data of the virus and construct the training data set. In this example, it is preferably set that the window size is 20, where the historical data window size is 15, the prediction target data window size is 5, and the sliding step is 1. Perform window sliding operations on the historical data, and the obtained data is used as the characteristics of the nodes in the spatio-temporal infection status graph G . Under the sliding window mechanism, the node feature matrix contains the characteristics of each node in the historical window:

[0080] 。

[0081] The feature vector of each node can be expressed as:

[0082] ,

[0083] where, represents the feature vector of the i-th node when the size of the historical data window is p, and p = 1, 2,... ; is the set of device infection states; Devices in the network are divided into four states according to their infection situations, namely the susceptible state , the latent state , the infected state and the immune state . The infection state of each node can be one of the four states. The true label represents the infection state of each node at a future time point, that is, the target prediction value, with a shape of , where is the prediction result step size of each node;

[0084] Step 3) Refer to Appendix Figure 2 , construct a prediction model based on the spatio-temporal graph attention network, as follows:

[0085] (3a) The input data of the model is the node feature matrix , and the adjacency matrix of the spatio-temporal infection state graph G;

[0086] (3b) Construct two layers of graph attention layers GAL to learn the spatial dependence relationship between nodes in the graph. The first layer of graph attention layer calculates:

[0087]

[0088] The calculation method of the attention coefficient in the graph attention layer GAL is:

[0089]

[0090] where, is the linear transformation matrix of node features, is the weight vector used to calculate the attention coefficient; represents the activation function, which is used to increase non-linear features; represents all the neighbor nodes of the i-th node.

[0091] In addition, this layer also uses the multi-head attention mechanism to stabilize the learning process. Apply A separate attention mechanism is used to calculate the hidden state, and then the average value of its features is calculated to obtain the features output by the first layer of GAL:

[0092]

[0093] where is the activation function. The second layer of graph attention layer uses the features output by the previous layer to calculate:

[0094]

[0095] to obtain the spatial features of the spatio-temporal infection status graph;

[0096] (3c) Two LSTM layers are constructed in the prediction model. The LSTM layer captures the dynamic information of the time series, processes the time series data and learns the time dynamics of virus transmission. The first layer of LSTM inputs the node features output by the graph attention layer, and obtains the output of the LSTM layer and a hidden state :

[0097]

[0098] The second layer of LSTM inputs and obtains the output of the LSTM layer and a hidden state ;

[0099] (3d) The output of the LSTM layer will be mapped through a fully connected layer to obtain the final prediction result vector:

[0100]

[0101] (3e) The prediction result vector is adjusted to the required shape to fit subsequent analysis:

[0102]

[0103] where is the predicted value of each node, representing the infection status of the node at the future time window point; is the batch size of the number of samples during training; is the total number of nodes in the network;

[0104] Step 4) Use the dataset to train the model:

[0105] (4a)Divide the dataset constructed in step 2) into a training set, a validation set, and a test set in the ratio of 80%:10%:10%. The training set is used for the training process of the model, enabling the model to learn the patterns and features in the data and optimizing the model parameters through gradient descent; the validation set is used during the training process to adjust the hyperparameters and evaluate the generalization ability of the model to prevent overfitting; after the model training is completed, the test set is used to evaluate the performance of the model, testing the prediction ability of the model on unseen data to ensure its generalization;

[0106] (4b)Select the mean squared error loss (MSE Loss) as the loss function and calculate the squared difference between the predicted value and the true label . The calculation formula is:

[0107]

[0108] where is the prediction result of the model, and is the true label of node ;

[0109] (4c)Select the optimizer as the optimization algorithm of the model, and use the learning rate and weight decay as parameters;

[0110] (4d)The dataset is split into several small batches according to the size of . Each batch undergoes one forward propagation and one backward propagation through the model. The training data is iterated through the data loader. By training within the specified total number of epochs , the parameters of the model are optimized to improve the prediction ability of the model;

[0111] Step 5) Evaluate the model accuracy:

[0112] Use the trained model to predict the future spread of the virus and evaluate the accuracy of the model on the test set. Use the mean absolute error (MAE) and root mean square error (RMSE) between the predicted value and the true value as reference standards to measure the error degree of the model when predicting the spread of the virus;

[0113] The mean absolute error (MAE) is the average of the absolute values of the differences between the predicted value and the true value. It measures the average deviation between the predicted value of the model and the true value. The smaller the value, the better the prediction ability of the model. The calculation method is as follows:

[0114]

[0115] The root mean square error (RMSE) is the square root of the mean of the squares of the differences between the predicted values and the true values. RMSE is more sensitive to larger errors and can thus reflect the impact of larger errors on the performance of the model. The calculation method is as follows:

[0116]

[0117] The smaller its value, the higher the prediction accuracy of the model and the better its generalization ability.

[0118] The effects of the present invention will be further described below in conjunction with simulation experiments.

[0119] 1. Simulation conditions:

[0120] The simulation experiment of the present invention was carried out under the hardware environment shown in Table 1 and the software environment shown in Table 2.

[0121] Table 1: Experimental hardware environment

[0122] Hardware Configuration CPU 12th Gen Intel(R) Core(TM) i7-12700KF GPU NVIDIA GeForce RTX 4060 Memory 32GB DDR4 RAM Main Hard Disk 1TB SSD

[0123] Table 2: Experimental software environment

[0124] Software Configuration Operating System Windows 11 Programming Software Visual Studio Code 1.88.1 Python Version 3.11.10 Pytorch Version 2.4.1

[0125] 2. Simulation content:

[0126] The network used in the experiment consists of 500 nodes and 996 non-directed edges, and its node degree distribution is as shown in the appendix Figure 3 and conforms to the power-law distribution. Using the spatio-temporal graph attention prediction model proposed by the present invention, combined with the divided data set, the selected loss function and optimizer, the model is trained. In the experiment, the number of training epochs is set to , and the change trend of the loss value loss during the training process is as shown in the appendix Figure 4 .

[0127] 3. Simulation results:

[0128] Appendix Figure 5 shows the training results of the spatio-temporal graph attention prediction model on the test set after training. Among them, the solid line represents the real situation of the node infection status, and the dashed line represents the infection status predicted by the model. It can be seen from the results that the prediction curve of the model is highly consistent with the real curve, indicating that the model has high accuracy and good fitting ability.

[0129] The present invention proposes a method for predicting the spread of network viruses based on a spatio-temporal graph attention network. By constructing a spatio-temporal infection status graph, combining GAT to extract spatial features between nodes, and LSTM to capture the temporal evolution of virus spread, spatio-temporal information is fused. The constructed dataset is used to train ST-GAT and evaluate the prediction performance. Finally, the trained model is used to predict the virus spread trend, warn high-risk infected nodes, and achieve active defense and early intervention. This method effectively combines the spatial topological characteristics and temporal evolution law of virus spread, realizes node-level infection prediction, and introduces a graph attention mechanism to improve prediction accuracy and calculation speed. It can meet the requirements of virus spread prediction in a large-scale network environment, and can effectively solve problems such as strong dependence on spread parameters and low prediction accuracy in existing virus spread prediction methods, providing strong support for network virus prevention and control.

[0130] Inside an enterprise, various operations highly rely on the network. The prediction method proposed by the present invention can be used to monitor the virus infection status of devices in the enterprise network in real time and predict the virus spread trend in advance. For example, in a financial enterprise, it can accurately monitor devices such as customer data storage servers and trading terminals, timely discover potential virus threats, and take targeted measures before a large-scale virus outbreak, such as isolating infected devices and updating antivirus software, effectively avoiding huge economic losses caused by data leakage and system paralysis, and ensuring the normal operation of the enterprise and the security of customer information. In the industrial Internet, industrial production devices are interconnected through the network. Once attacked by a virus, it may lead to serious consequences such as production line paralysis and production accidents. Using the method of the present invention, industrial enterprises can monitor the network status of production devices in real time, predict the virus spread path, and take protective measures in advance to ensure the continuity and stability of industrial production and reduce production losses caused by network security problems. It has broad application prospects.

[0131] It should be noted that the user information (including but not limited to user device information, user personal information, etc.) and data (including but not limited to data for analysis, stored data, displayed data, etc.) involved in this application are all information and data authorized by the user or fully authorized by all parties. And the collection, use, and processing of relevant data need to comply with the laws, regulations, and standards of relevant countries and regions, and corresponding operation entrances are provided for users to choose to authorize or refuse.

[0132] The above simulation analysis proves the correctness and effectiveness of the method proposed by the present invention.

[0133] The parts not detailed in the present invention belong to the common general knowledge of those skilled in the art.

[0134] The above are only the preferred embodiments of the present invention and are not intended to limit the present invention. Obviously, for those skilled in the art, after understanding the content and principle of the present invention, various modifications and changes in form and details may be made without departing from the principle and structure of the present invention. However, these modifications and changes based on the idea of the present invention are still within the scope of protection of the claims of the present invention.

Claims

1. A method for predicting the spread of network viruses based on a spatio-temporal graph attention network, characterized in that, It includes the following steps: (1) According to the propagation characteristics of network viruses, the computer network is represented as an undirected graph G with an infection state time series t =(N t , E t , W t ), which is used to describe the process of the infection state of devices in the network changing over time; where N t is the node set at time t, and the historical propagation data of the virus is used as the feature of the node; E t represents the edge set at time t, and any two nodes that communicate with each other form an edge in the undirected graph; W t represents the adjacency matrix of the spatio-temporal infection state graph G at time t; (2) Process the historical virus transmission data using a sliding window mechanism to construct a dataset; and record the infection status of each node at future time points as the true label Y, with its shape denoted as [N, m pre , where N represents the total number of nodes, and m pre represents the prediction result step size of each node; (3) Construct a prediction model based on the spatio-temporal graph attention network, and the implementation steps are as follows: (3.1) Use the node feature matrix X and the adjacency matrix W of the spatio-temporal infection status graph G t as the input data of the model; (3.2) Construct two-layer graph attention layer GAL to learn the spatial dependence relationship between nodes in graph G; Among them, the output feature of the first layer of GAL is X1, and the spatial feature X2 of the spatio-temporal infection status graph output by the second layer of GAL is calculated by using the feature X1: X1 = GATConv(X, W t ) X2 = GATConv(X1, W t ); Among them, GATConv() represents the calculation of the graph attention layer GAL, that is, using the multi-head attention mechanism to calculate the hidden state and take the average of it to obtain the aggregated node feature as the output; (3.3) Construct two LSTM layers to capture the dynamic information of the time series, process the time series data and learn the time dynamics of virus transmission; use the feature X2 output by the graph attention layer as the input of the first LSTM layer to obtain the third output feature X3, and then use X3 as the input of the second LSTM layer to obtain the fourth output feature X4; (3.4) Map the output of the LSTM layer through a fully connected layer to obtain the final predicted result vector (3.5) Adjust the prediction result vector to the same shape as the true label Y according to the following formula: Among them, n pred is the predicted value of each node, indicating the infection status of the node at future time window points; batch size is the batch size of the number of samples during training; n nodes is the total number of nodes in the network; (4) Divide the data set into a training set, a validation set and a test set, select the mean square error loss MSE Loss as the model loss function, and use the Adam optimizer to train the constructed prediction model to obtain the final trained prediction model; (5) Input the test set data into the final prediction model to realize the prediction of the future spread of the virus.

2. The method according to claim 1, characterized in that: The nodes in step (1) represent devices in the network, including but not limited to computers, servers, and routers.

3. The method according to claim 1, wherein: The adjacency matrix of the spatio-temporal infection status graph G in step (1) is expressed as follows: Among them, if there is an edge between the $i$-th node $n$ i and the $j$-th node $n$ j , then $w$ ij $= 0$, otherwise $w$ ij $= 1$; $i, j = 1, 2, \ldots, N$.

4. The method according to claim 1, characterized in that: The processing of the historical virus transmission data using the sliding window mechanism in step (2) is specifically as follows: set the window sizes, including the historical data window size \(m\) his and the predicted target data window size \(m\) pre , set the sliding step size to 1; perform window sliding operations on the historical data to obtain the node feature matrix \(X\) in the spatio-temporal infection status graph \(G\).

5. The method according to claim 4, wherein: The node feature matrix X is obtained according to the following steps: (2.1) Under the sliding window mechanism, the node feature matrix X contains the features of each node in the historical window, and the feature vector of the i-th node is represented as x i ′: where x ip represents the feature vector of the i-th node when the historical data window size is p, where p = 1, 2,..., m his ; is the set of device infection states; (2.2) The node feature matrix X is expressed as follows:

6. The method according to claim 5, characterized in that: The infection status in step (2.1) includes the susceptible state S, the latent state E, the infected state I, and the immune state R, which are divided according to the infection situation of devices in the network.

7. The method according to claim 1, wherein: The multi-head attention mechanism in step (3b) specifically learns different feature representations through multiple independent attention heads, so as to improve the model's ability to model complex transmission relationships. The expression is as follows: Among them, K is the number of attention heads, σ is the activation function, and α iq is the attention coefficient between the i-th node and its neighbor node q, and W k is the trainable weight matrix of the k-th attention head; X′ i is the feature of the i-th node after being calculated by the attention mechanism.

8. The method according to claim 7, wherein: The attention coefficient α iq , is defined as follows: Among them, W is the linear transformation matrix of node features, a is the weight vector used to calculate the attention coefficient; LeakyReLU represents the activation function, which is used to increase non-linear features; Q represents all neighbor nodes of the i-th node.

9. The method according to claim 1, wherein: Step (4) trains the constructed prediction model, and the implementation steps are as follows: (4.1) Divide the data set constructed in step (2) into a training set, a validation set and a test set according to the ratio of 80%:10%:10%; among them, the training set is used for the training process of the model, so that the model learns the patterns and features in the data, and optimizes the model parameters through gradient descent; the validation set is used to adjust the hyperparameters and evaluate the generalization ability of the model during the training process; the test set is used to evaluate the performance of the model; (4.2) Select the mean squared error loss MSE Loss as the model loss function and calculate the squared difference between the predicted value and the true label Y Among them, is the prediction result of the model, and y i is the true label of node i; (4.3) The optimization algorithm of the model selects the Adam optimizer, and uses the learning rate and weight decay as parameters; (4.4) Split the dataset into multiple batches according to the size of batch_size. Each batch undergoes one forward propagation and one backward propagation through the model. Iteratively train the data through the data loader to optimize the model parameters.

10. The method according to claim 9, characterized in that: In step (4.3), select the Adam optimizer and set the learning rate and weight decay parameters to improve the stability and generalization ability of training. The parameter update formula is: where m t and v t are the estimates of the first and second moments of the gradient respectively, ∈ is a numerical stability term, and lr represents the learning rate; the optimization process introduces weight decay, that is, adding an L2 regularization term to the loss function: where λ is the weight decay coefficient, is the mean squared error loss, is the final optimization objective function.

Images