Internet of Things defense method of reinforcement learning based on graph attention enhancement

Abstract

The invention relates to the technical field of network security, in particular to realization of self-efficient autonomous defense in an Internet of Things network. Comprising the following steps: constructing a network attack and defense scene, and modeling observation information by using a graph data structure; modeling the action space of the attacker and the defense agent; modeling an attack and defense process in the network attack and defense scene; performing normalization processing on the feature matrix of the nodes by adopting graph normalization; dynamically aggregating the current state information of each device according to the interaction between the network device and the adjacent device by using the self-attention mechanism of the GAT; selectively reserving a node host according to the self-attention score by using a down-graph sampling pooling strategy; and training a defense agent by using a graph attention network enhanced reinforcement learning algorithm. The problems that automatic defense of an existing reinforcement learning algorithm faces performance bottleneck in a complex network scene, convergence to a high return value is difficult, and an effective defense strategy cannot be made are solved.

Description

Technical Field

[0001] The present invention relates to the field of Internet of Things defense technology, and in particular to an Internet of Things defense method based on graph attention-enhanced reinforcement learning. Background Art

[0002] The rapid development of artificial intelligence technology has promoted the transformation of the application field of the Internet of Things (IoT), gradually moving from the Industrial Internet of Things, Medical Internet of Things, and Consumer Internet of Things to the Artificial Intelligence Internet of Things (AIoT); although AIoT has significantly improved convenience, it has also exposed IoT network systems to more diverse network threats. The use of a large number of devices has brought about a surge in potential vulnerabilities, and the high degree of interaction between devices has further amplified the risk of network attacks.

[0003] Alshahrani, H et al.'s Intrusion detection framework for industrial internet of things using software defined network. Sustainability uses reinforcement learning for network security defense, but reinforcement learning methods have certain limitations in accurately characterizing and utilizing observation information and capturing complex dependencies and interaction patterns in the network, especially in IoT scenarios with a large number of devices and complex interactions. These limitations hinder the effective formulation of intelligent defense strategies. In addition, existing reinforcement learning methods still face challenges in formulating and implementing effective strategies in complex network environments. This is mainly due to the complex interactions and dependencies between network devices in these complex systems, making it difficult to distinguish network topology structures or identify the importance of network hosts, which are often highly dependent on their current states, and fail to effectively capture and utilize interactions between nodes, resulting in suboptimal results. Summary of the invention

[0004] In view of the shortcomings of the existing methods, the present invention solves the problem that the automated defense of the existing reinforcement learning algorithm faces performance bottlenecks in complex network scenarios, is difficult to converge to a higher reward value, and cannot formulate effective defense strategies.

[0005] The technical solution adopted by the present invention is: the Internet of Things defense method based on graph attention enhanced reinforcement learning includes the following steps:

[0006] Step 1: Construct a network attack and defense scenario and use graph data structures to model observation information;

[0007] As a preferred embodiment of the present invention, step one specifically includes:

[0008] Step 11: Identify device types as different nodes of the graph, represent connection relationships with edges, and set edge weights;

[0009] The weight of the edge is assigned according to the importance of the connection, the frequency of communication, and the amount of data transferred.

[0010] As a preferred embodiment of the present invention, the formula of edge weight is:

[0011]

[0012] Among them, σ(s,t) represents the total number of shortest paths between nodes s and t, σ(s,t|e) represents the number of shortest paths through edge e, and V is the node set.

[0013] Step 12: Collect key information of different devices and abstract the key information into node features.

[0014] Step 2: Model the action space of the attacker and the defense agent;

[0015] The action space of the defense agent includes: defense agent scanning, patching, isolation, reconnection, sleep, recovery and honeypot; the attacker action space includes: attacker scanning, privilege escalation, zero-day attack and lateral movement.

[0016] Step 3: Model the attack and defense process in the network attack and defense scenario;

[0017] As a preferred embodiment of the present invention, step three specifically includes:

[0018] Step 31: The attacker uses the vulnerability to elevate privileges, performs an attacker scan to identify the target host, and performs a vulnerability scan on the target host; performs privilege escalation or a zero-day attack to destroy the target host; and penetrates the entire network through lateral movement;

[0019] Step 32: Probabilistic success action based on the number of vulnerabilities of the target host;

[0020] Step 33: The defense agent monitors the network state and captures the network state s at each time step. t ; The defense agent analyzes the current network state s through the policy network π(a|s) t Generate defensive actions; and calculate the expected cumulative reward; after executing the defensive action, the next state s t+1 According to the state transition probability P(s t+1 |s t ,a t ) is determined, and the immediate reward r is obtained according to the reward function t ; Status t 、Action a t and status t+1 , r t is stored in the replay buffer and used to update the policy network;

[0021] As a preferred embodiment of the present invention, the formula of the reward function is:

[0022] R step =-C action -N compromised ×10

[0023] When the number of steps reaches the preset maximum number of steps per round, if the defense is successful, the final reward will be returned Otherwise return Among them, C action represents the cost of the defender taking action, N compromised and N uncompromised Respectively represent the number of hosts hacked by the attacker and the number of hosts that have not been hacked, N total Indicates the total number of devices.

[0024] As a preferred embodiment of the present invention, a hierarchical penalty mechanism is introduced into the reward function, and the formula for constructing the improved reward function is:

[0025] R′ step =-C action -N normal ×1-N high ×10

[0026] Among them, N normal is the number of compromised common hosts, N high is the number of high-value hosts compromised.

[0027] Step 4: Use graph normalization to normalize the feature matrix of the node; use the self-attention mechanism of GAT to dynamically aggregate the current state information of each device based on the interaction between the network device and its neighboring devices; use the graph downsampling pooling strategy to selectively retain the node host according to the self-attention score; use the reinforcement learning algorithm enhanced by the graph attention network to train the defense agent;

[0028] As a preferred embodiment of the present invention, dynamically aggregating the current status information of each device according to the interaction between the network device and its neighboring devices includes: assigning weights according to the importance of the connected devices, and aggregating the status information based on the weights.

[0029] As a preferred embodiment of the present invention, using a graph downsampling pooling strategy to selectively retain nodes according to attention scores includes: integrating a graph sampling pooling algorithm into a policy network so that the defense agent pays attention to key hosts in the current state.

[0030] As a preferred embodiment of the present invention, the reinforcement learning algorithm enhanced by the graph attention network to train the defense agent includes: using the GAT-DQN algorithm to train the defense agent, and the GAT-DQN algorithm includes:

[0031] Define the target network Q target Parameters, initialization parameters learning rate α, discount factor γ, exploration rate ∈ and maximum number of training rounds max_episodes;

[0032] At the beginning of each training round, the environment is reset and the initial state s0 is obtained; s0 includes the state information of the host, which is reconstructed into x0 and edge_index0 based on the graph data structure;

[0033] Generate a random number rand∈[0,1] at each time step; if rand<∈, randomly select an action a t Otherwise, use the current Q network according to the formula Get the optimal defensive action a t ; After the action is executed, the environment returns the reward r t , enter state s t+1 , reconstructed again to get x t+1 ,edge_index t+1 ;

[0034] The four-tuple (x t ,edge_index t ,a t ,r t ,x t+1 ,edge_index t+1 , done) is added to the experience playback area; if the playback area size exceeds batch_size, a small batch is randomly sampled for training;

[0035] The loss function is calculated using the sampled batch data, and the Q network parameters are optimized using the gradient descent method. The loss function is calculated based on the formula of the target network:

[0036]

[0037] The parameters of the Q network are synchronized to the target network every fixed number of steps, and the exploration rate is decayed.

[0038] As a preferred embodiment of the present invention, the reinforcement learning algorithm enhanced by the graph attention network for training the defense agent further includes: using the GAT-PPO algorithm to train the defense agent, and the GAT-PPO algorithm includes:

[0039] Initialize the policy network π θ(a|x, edge.index) and value network V(x, edge.index; ω); set discount factor γ, GAE parameter λ, maximum training round T total , batch size B and number of experience sampling steps M B ;

[0040] For each training epoch, execute:

[0041] Reset the environment, extract the current state and reconstruct it into the node feature matrix x0 and the edge connection relationship edge_index0;

[0042] Based on the policy network π θ , calculate the action probability distribution p according to the current state θ , and then sample the action a according to the distribution t ;

[0043] Execute action a t , get reward r t , next state feature x t+1 , edge connection relationship edge_index t+1 And the completion mark done;

[0044] (x t ,edge.index t ,a t ,r t ,x t+1 ,edge.index t+1 ) is stored in ReplayBuffer;

[0045] Calculate the TD error at each time step with the following formula:

[0046] δ t =r t +γV(x t+1 ,edge.index t+1 ;ω)-V(x t ,edge.index t ;ω)

[0047] Calculate the advantage function A based on TD error t , the formula is:

[0048]

[0049] According to the advantage function A t Calculate the target value function:

[0050] V t =A t +V(x t,edge.index t ; π)

[0051] Sample data from ReplayBuffer in small batches and optimize the policy network π based on the PPO loss function θ , and update the parameters θ by gradient descent;

[0052] Optimize the value network, minimize the value function error, and update the parameter π;

[0053] After each update, clear the ReplayBuffer.

[0054] Beneficial effects of the present invention:

[0055] 1. By introducing the reinforcement learning algorithms GAT-DQN and GAT-PPO enhanced by graph neural networks, the intelligence level of network defense strategies is significantly improved; it shows strong adaptability in multiple complex network scenarios and attack strategy combinations, fully proving its wide applicability in dealing with advanced persistent threats (APTs) and other network threats;

[0056] 2. The present invention models network observation data as a graph structure and introduces a graph neural network. By using the information propagation and feature aggregation mechanism of the graph neural network, the defense agent comprehensively considers the interactive relationship between the host, IoT devices, and other entities in the network topology, accurately identifies potential attack paths, and formulates more accurate and effective defense measures;

[0057] 3. The application of reinforcement learning technology based on graph attention network in the field of network security can alleviate the network threats of IoT network systems and open up new directions for network security research;

[0058] 4. The Yawning Titan (YT) simulation environment was improved to build a near-real IoT network scenario, making the defense strategies generated by training more in line with actual application needs; and a variety of attack penetration strategies were designed to train more comprehensive automated defense strategies, thereby enhancing the overall defense capabilities of the system. BRIEF DESCRIPTION OF THE DRAWINGS

[0059] Figure 1 It is a flow chart of the IoT defense method based on graph attention-enhanced reinforcement learning of the present invention;

[0060] Figure 2 It is a graph of feature enhancement performed by the graph attention network of the present invention;

[0061] Figure 3 is a detailed view of the experimental network topology of the present invention;

[0062] Figure 4It is a comparison chart of the average reward curves of five trainings using GAT-PPO, Norm-PPO, RNN-PPO and Sb3-PPO of the present invention;

[0063] Figure 5 This is a comparison chart of the average reward curves of five trainings using GAT-DQN, DDQN-DNER, RNN-DQN and DDQN in different scenarios of the present invention. DETAILED DESCRIPTION

[0064] The present invention is further described below in conjunction with the accompanying drawings and embodiments. This figure is a simplified schematic diagram, which only illustrates the basic structure of the present invention in a schematic manner, and therefore it only shows the components related to the present invention.

[0065] The IoT defense method based on reinforcement learning with graph attention enhancement includes the following steps:

[0066] like Figure 1 ,Step 1: Build a network attack and defense scenario and use graph data structure to model the observation information;

[0067] Use Yawning Titan to build network attack and defense scenarios including: enterprise networks, IoT networks such as Consumer IoT (CIoT) networks, Industrial IoT (IIoT) networks, and Artificial Intelligence IoT (AIoT);

[0068] Specifically include:

[0069] Step 11: Model the observation information of the network attack and defense scenario using a graph data structure, where different device types are identified as different nodes in the graph, and the connection relationship between different devices is represented by edges in the graph. The weight of the edge is assigned according to the importance of the connection, the communication frequency, and the data transmission volume;

[0070] Different device types include: several IoT devices, several user hosts, several high-value servers, etc.;

[0071] Step 12: Collect key information of different devices, including confidentiality, availability, reliability and vulnerability, and abstract the key information into the characteristics of nodes in the graph;

[0072] Through the modeling process in step one, the observation information based on the graph data structure can accurately reflect the network structure, while revealing potential attack paths and key nodes, thus helping to optimize network defense strategies.

[0073] Step 2: Model the action space of the attacker and the defense agent;

[0074] The action space of the defense agent includes: defense agent scanning, patching, isolation, reconnection, hibernation, recovery, and honeypot;

[0075] Defense Proxy Scan (Scan): Scan all devices to detect potential Red Team intrusions, the action cost is 0;

[0076] Patch: Reduce the vulnerability score of the target device, the action cost is 0;

[0077] Isolation: Disconnect the target host from the network. The action cost is 1.

[0078] Reconnection: Restore network access to the isolated host. The action cost is 0.

[0079] Sleep: Do not take any action, the action cost is 0;

[0080] Restore: restore the host to its initial state, the action cost is 1;

[0081] Honeypot: Place a disguised host in the network topology, with an action cost of 2;

[0082] The action costs in the action space of the defense agent are all derived from the Yawning Titan simulation environment.

[0083] The attacker action space includes: attacker scanning, privilege escalation, zero-day attacks, and lateral movement;

[0084] Attacker Scan: Corresponding to ATT&CK technology T1018, use active scanning tools such as ping to discover new hosts in the network;

[0085] Privilege Escalation: Corresponding to ATT&CK technology T1210 and strategy TA0004, it uses service vulnerabilities on the target host to escalate privileges;

[0086] Zero-Day attack: An attack that exploits unpatched or undisclosed vulnerabilities.

[0087] Lateral Movement: The attacker moves laterally within the target network to further compromise other systems to gain more data or control.

[0088] Step 3: Model the attack and defense process in the network attack and defense scenario;

[0089] Step 31: In the network attack and defense scenario, APT attackers continuously exploit vulnerabilities to elevate privileges and infiltrate the entire network; Figure 1 The medium attacker (red node) first performs an attacker scan to identify exploitable hosts and performs a vulnerability scan on the target host (yellow mark). Subsequently, the attacker performs privilege escalation or zero-day exploits to compromise the target host and further penetrate the entire network through lateral movement.

[0090] The ultimate goal of the attacker is to repeatedly execute the above process to deeply penetrate the entire network system and achieve data theft or complete control of the system. At the same time, the attacker's behavior patterns and attack chains are taken into consideration, covering a variety of attack strategies and key targets to fully simulate the penetration process. The defense agent monitors the network status in real time and strategically selects the optimal defense action based on the current situation to minimize potential losses.

[0091] Step 32. In order to formulate an efficient and accurate defense strategy, the complex vulnerability exploitation operations of real attackers are abstracted in the simulation environment; specifically, these operations are simplified into probabilistic successful actions based on the number of vulnerabilities in the target host; the target host refers to the target device that the attacker currently chooses to penetrate and destroy next; this abstract modeling method helps the defense agent better predict and respond to complex APT attacks, thereby accelerating the training process and promoting the optimization of defense strategies.

[0092] Step 33: The defense agent continuously monitors the network status and captures the network status s at each time step. t , including all host status and device connection status; then, the defense agent analyzes the current network status s through the policy network π(a|s) t Generate defensive actions; at the same time, calculate the expected cumulative reward Among them, γ is the discount factor, reflecting the present value of future rewards; r t+k+1 is the cumulative reward value from time t to the end of the round; after executing the defensive action, the next state s t+1 According to the state transition probability P(s t+1 |s t ,a t ) is determined, and the instant reward r is obtained according to the configured reward function t ; Status t 、Action a t and status t+1 , r t It is stored in the replay buffer and used to update the policy network.

[0093] This process transforms the problem of finding the optimal strategy into a series of policy network optimization tasks. Through the GAT-DQN and GAT-PPO algorithms, the trained defense agent can respond quickly and accurately in a dynamic and complex network security environment.

[0094] Reinforcement Learning (RL) is a key machine learning method. The defense agent learns actions that can maximize the cumulative reward in a given state by interacting with the environment. Current reinforcement learning algorithms are mainly divided into two categories: value-based deep reinforcement learning algorithms such as deep Q-network (DQN) and policy-based deep reinforcement learning algorithms such as proximal policy optimization algorithm (PPO).

[0095] The Deep Q-Network (DQN) is a reinforcement learning algorithm that combines Q-learning and deep neural networks to approximate the Q-value function, enabling defense agents to learn optimal actions in complex environments; the Q-network is represented by the weight parameter ω and minimizes the predicted Q-value Q for a given state-action pair. ω (s i ,a i ) and the target value for training; the target value is the immediate reward r i The discounted maximum Q value of the next state The objective function is:

[0096]

[0097] Compared with value-based methods such as DQN, policy-based deep reinforcement learning directly learns a policy function that determines the probability of an action based on the current state, which is very suitable for complex safety decisions. Proximal Policy Optimization (PPO) is one of the mainstream policy-based methods;

[0098] As shown in formula (2), PPO optimizes the strategy to maximize the cumulative reward by controlling the training update amplitude. The formula is:

[0099]

[0100] Among them, θ represents the parameters of the policy function, is the advantage estimate of action a in state s, ρ t (θ) is the ratio of the current strategy to the old strategy, ∈ is a hyperparameter, and the clip(·) function is used to limit the range of the ratio; the PPO algorithm has the advantages of high adaptability in action space and training stability, and can efficiently optimize the strategy; however, the existing PPO algorithm mainly uses a fully connected layer to process observation data, which has certain limitations in a complex network environment; in order to solve this problem, the present invention combines graph neural networks (GNNs), which perform well in deep learning applications (such as node classification and graph classification) for processing unstructured data.

[0101] Since the importance of hosts in real network scenarios varies, Graph Attention Network (GAT), as a type of graph neural network, is particularly suitable for network security;

[0102]

[0103] As shown in Equation 3, GAT dynamically adjusts the influence of neighbors on the central node by aggregating neighbor node features and using the learned attention weights to update node representations; this method efficiently captures key information in graph structured data through the attention mechanism.

[0104] Among them, h′ i represents the updated feature representation of node i, σ is the activation function, is the set of neighbor nodes of node i, W is the feature transformation matrix, is the parameter vector for calculating the attention weights, and is the concatenation of the feature vectors of nodes i and j.

[0105] Graph Attention Networks (GAT) are particularly suitable for application in the field of network security because they can accurately capture dynamic network interactions. For example, GAT aggregates the characteristics of connected devices through information transfer weighted by attention weights, ensuring that each device in the network dynamically reflects the interaction with connected devices based on the comprehensive information of adjacent devices. This mechanism enables defense agents to comprehensively evaluate and identify security threats that each node may face. For example, in APT attacks and DDoS attacks, when a device in the network is attacked, the devices connected to it will quickly reflect abnormal behavior and potential security risks, allowing defense agents to take immediate defense measures. GAT quickly analyzes and responds to potential attacks by utilizing these dynamic characteristics and changes in attention weights between devices.

[0106] In order to develop effective defense strategies by training intelligent defense agents through reinforcement learning, a series of challenges are faced:

[0107] 1. How to transform complex and dynamic real-world observation data into a format that can be understood and acted upon by intelligent agents, thereby ensuring accurate representation of the complexity of the environment;

[0108] 2. How to integrate the observation data of each host in a dynamic network environment while maintaining the stability of training when there are complex dependencies and interactions between hosts in the network topology;

[0109] 3. In a large-scale network topology containing a large number of hosts, how to identify and prioritize key hosts so that intelligent defense agents can concentrate resources and improve the overall effectiveness of network defense.

[0110] Step 4: Use graph data structures to model environmental observations and adopt graph normalization technology to ensure the accuracy of environmental representation and the stability of training; integrate the observation data of hosts in dynamic networks through GAT enhanced features, so that defense agents can accurately evaluate environmental dynamics and implement more effective strategies; use graph downsampling pooling to identify and prioritize key hosts, thereby improving strategy efficiency in complex scenarios; key hosts refer to high-value hosts or hosts with a large number of connections, that is, if a host with a large number of connections is successfully attacked, the attacker can infiltrate and destroy more connected hosts through this host; key hosts are assigned for training through a self-attention score.

[0111] Step 41, graph normalization includes:

[0112] Based on the original network system structure, the graph data structure is used to efficiently model network observations. The real network topology is represented by the graph data structure, which effectively solves the problem of how to efficiently represent environmental observations and ensures that the intelligent defense agent can process and learn efficiently.

[0113] Network observations are modeled in two parts. Each device in the network topology is represented by a feature matrix. Each row in the matrix corresponds to a host, and the columns represent different key information abstracted as the features of the nodes in the graph, which are used to capture the status information of each host. In addition, the connection relationship between devices is represented by an edge index, which is a three-row matrix. The first two rows represent the connection relationship between devices, and the third row provides the weight of each connection, reflecting its importance in the network topology.

[0114] In addition, the normalization method is used to normalize the feature matrix of the node to enhance data processing efficiency and improve the adaptability of the model to networks of different scales. The normalization formula of the feature matrix is:

[0115]

[0116] Where X is the original feature matrix, μ and σ are the mean and standard deviation of X, respectively, and ∈ is a small constant added for numerical stability; the normalization of node features ensures the stability of the training process.

[0117] Step 42: GAT enhancement features include:

[0118] The normalized feature matrix and edge index are deeply processed by the graph attention network (GAT); GAT uses the self-attention mechanism to dynamically aggregate the current state information of each device based on the interaction between the network device and its neighboring devices; initially, weights are assigned according to the importance of the connected devices, and then the state information is aggregated based on the weights.

[0119] For example, in security threat detection in enterprise networks, GAT can dynamically aggregate device status information to identify potential threats. When a server A is connected to neighboring devices B and C, initially, the importance weights of B are learned to be 0.7 and C is 0.3 based on historical data. Through the self-attention mechanism, GAT finds that the state of C at the current time step has an increased impact on A, and dynamically adjusts the weights to 0.5 for B and 0.5 for C. Subsequently, A updates its own features by weighted aggregation of the status information of B and C to generate an embedding that reflects the current network dynamics and interactions for identifying abnormal behavior. At this time, the defense agent observes that A's features have the features of the connected host. If the connected host has abnormal defense, the defense agent can take defensive actions more promptly to eliminate potential threats and reduce losses. This method ensures that the defense agent fully considers the high interactivity within the network, thereby significantly improving the accuracy of feature representation and improving the decision-making ability of defense agents in complex network topologies.

[0120] Step 43, graph sampling pooling includes:

[0121] In certain network topologies containing a large number of hosts, the graph downsampling pooling strategy selectively retains nodes based on the calculated self-attention scores; the calculation formula for the self-attention score of a node in the Self-Attention Graph Pooling (SAGPooling) algorithm of the present invention is:

[0122]

[0123] in, is an adjacency matrix containing self-loops, is the corresponding degree matrix, X att is the feature matrix of the node, and σ is the activation function.

[0124] The present invention integrates the SAGPooling algorithm into the policy network, enabling the defense agent to focus on the key hosts in the current state; this resource allocation strategy significantly improves the efficiency and responsiveness of network defense;

[0125] Before the policy network performs weighted feature aggregation in step 42, if the number of devices is very large, each device must perform weighted feature aggregation as described in step 42, which will significantly increase the time complexity. Therefore, it is chosen to prioritize the identification of key hosts, such as hosts connected to the successfully attacked host, high-value hosts, and hosts with a large number of connections, which will be prioritized for feature aggregation for observation and protection.

[0126] Step 44: Implement the GAT-DQN algorithm;

[0127] use Figure 2The network framework in the figure replaces the Q network in the DQN algorithm to achieve the GAT-DQN algorithm based on GAT-enhanced reinforcement learning algorithm; the network framework described in the figure replaces the Actor-Critic network framework of PPO to obtain the GAT-PPO algorithm; GAT-DQN and GAT-PPO are applied in Figure 2 Network framework, GAT-DQN outputs the score of each defense action and selects the action with the highest score for defense; while GAT-PPO applies two Figure 2 The network framework outputs the probability distribution of actions, selects the action with the highest probability for defense, and outputs a score to evaluate the quality of applying this action in the current state.

[0128] Among them, the specific steps of GAT-DQN include:

[0129] Step 441, parameter initialization: define the target network Q target The parameters of the network are initialized in the same way as the Q network, including: setting the learning rate α, discount factor γ, exploration rate ∈, and the maximum number of training rounds max_episodes;

[0130] Environment initialization: At the beginning of each training round, the environment is reset and the initial state s0 is obtained; the one-dimensional data of state s0 includes the state information of the host in the network topology, and the state information is reconstructed into x0 and edge_index0 based on the graph data structure;

[0131] Environment interaction: At each time step, generate a random number rand∈[0,1]; if rand<∈, randomly select an action a t Otherwise, use the current Q network according to the formula Get the optimal defensive action a t ; After the action is executed, the environment returns the reward r t , enter the new state s t+1 , reconstructed again to get x t+1 ,edge_index t+1 ;

[0132] Experience storage: store the quadruple (x t ,edge_index t ,a t ,r t ,x t+1 ,edge_index t+1 , done) is added to the experience replay buffer; if the replay buffer size exceeds batch_size, a small batch is randomly sampled for training;

[0133] Network update: Use the sampled batch data to calculate the loss function, and use the gradient descent method to optimize the Q network parameters. The loss function is calculated based on the formula of the target network:

[0134]

[0135] Target network update: synchronize the parameters of the Q network to the target network at fixed steps and attenuate the exploration rate;

[0136] Step 45 also includes implementing the GAT-PPO algorithm, and the specific implementation steps are:

[0137] Parameter initialization: Initialize the policy network π θ (a|x, edge.index) and value network V(x, edge.index; ω), where both the policy network and the value network are based on the GAT model; set the discount factor γ, GAE parameter λ, and the maximum training round T total , batch size B and number of experience sampling steps M B .

[0138] Training epoch loop: For each training epoch, do the following:

[0139] State initialization: reset the environment, extract the current state and reconstruct it into the node feature matrix x0 and the edge connection relationship edge_index0.

[0140] Action selection and sampling: based on policy network π θ , calculate the action probability distribution p according to the current state θ , and then sample the action a according to the distribution t .

[0141] Environment interaction and data storage:

[0142] Perform action a in the environment t , get reward r t , next state feature x t+1 , edge connection relationship edge_index t+1 And the completion mark done.

[0143] (x t ,edge.index t ,a t ,r t ,x t+1 ,edge.index t+1 ) is stored in ReplayBuffer.

[0144] Compute the advantage function:

[0145] Extract data from ReplayBuffer and calculate the TD error \(\delta_t\) at each time step. The formula is as follows:

[0146] δ t =r t +γV(x t+1 ,edge.index t+1 ;ω)-V(x t ,edge.index t ;ω) (7)

[0147] Calculate the advantage function A based on TD error t , the formula is:

[0148]

[0149] Estimate the target value function: According to the advantage function A t Calculate the target value function:

[0150] V t =A t +V(x t ,edge.index t ;ω) (9)

[0151] Strategy and Value Network Update:

[0152] Sample data from ReplayBuffer in small batches and optimize the policy network π based on the PPO loss function θ , and update the parameters θ by gradient descent;

[0153] At the same time, optimize the value network, minimize the value function error, and update the parameter ω;

[0154] ReplayBuffer is cleared: After each update is completed, the ReplayBuffer is cleared to prepare data for the next round of training.

[0155] Similarly, GAT-PPO replaces the Actor-Critic network structure in the PPO algorithm. The Actor network first propagates information and aggregates features through GAT, and then outputs more accurate defense actions. When evaluating the current state, the Critic network can also evaluate whether there are potential threats in the network based on this mechanism.

[0156] Conduct YT simulation and improve YT's state space and reward mechanism;

[0157] YawningTitan (YT) is a network security simulator developed based on the OpenAiGym platform. It is used to simulate real network topologies and provide a flexible simulation environment. YT allows intelligent defense agents to formulate strategies based on the current state and explore new network security decision-making methods, thereby improving the efficiency of network defense strategies and the robustness of network security frameworks. YT can accurately simulate a variety of real network topologies and provide a diverse environment for network security scenarios. The state space, action space, and reward mechanism of YT are described as follows:

[0158] State space: In the YT environment, each network device is represented by a 3-bit feature vector that abstractly describes the security status of the device, including the vulnerability score (an abstract value representing the number of device vulnerabilities), whether the device has been compromised, and whether the device is in isolation;

[0159] Action space: detailed description of actions and their costs;

[0160] Reward mechanism: In the YT environment, the original reward function R at each step step The calculation formula is:

[0161] R step =-C action -N compromised ×10 (10)

[0162] When the number of steps reaches 100 (the preset maximum number of steps per round), if the defense is successful, the final reward will be returned Otherwise return Among them, C action represents the cost of the defender taking action, N compromised and N uncompromised Respectively represent the number of hosts hacked by the attacker and the number of hosts that have not been hacked, N total Indicates the total number of devices;

[0163] In order to make the YT simulation environment closer to the complexity of real network scenarios, its state space and reward mechanism are improved, including:

[0164] The weight w(e) of the edge is calculated by the betweenness centrality formula, which is:

[0165]

[0166] Among them, σ(s,t) represents the total number of shortest paths between nodes s and t, σ(s,t|e) represents the number of shortest paths through edge e, and V is the node set;

[0167] In addition, a feature F is assigned to each node, and the feature F value is the sum of the weights of all edges connected to the node; this improvement enriches the state representation of the environment and more accurately reflects the dynamics and complexity of the real network; the cost of the isolation action is adjusted from a fixed value to a dynamic value, which is determined by the feature F of the isolated node, that is, the node feature value F is returned as the cost of isolating the node.

[0168] The original cost of isolating a node is fixed. After optimization, it is dynamically adjusted according to the sum of the weights of the edges connected to it. For example, if a host is connected to many hosts, the sum of its variable weights will be very large. Isolating this host will have a great impact on the entire network topology, so the cost will be higher, which is more in line with reality. In the state space, the edges originally had no weights, but now the weights are calculated based on the variable betweenness centrality and assigned to the edges.

[0169] A graded penalty mechanism is introduced into the reward function. The penalty size is adjusted according to the importance of the compromised device, distinguishing between ordinary hosts (such as users or operating machines) and high-value hosts (such as key servers); the penalty P at each step is calculated based on the number of ordinary hosts N compromised. normal and high-value hosts N high Calculation of the number of P = N normal ×1+N high ×10; the formula of the updated improved reward function is:

[0170] R′ step =-C action -N normal ×1-N high ×10 (12)

[0171] These improvements are designed to improve the precision and depth of decision-making and more accurately reflect the complexity of network security defense. The evaluation criterion is the total cumulative reward obtained by the defense defense agent during the entire operation process, which directly measures the effectiveness of the strategy.

[0172] Experimental setup

[0173] In order to evaluate the effectiveness of GAT-DQN and GAT-PPO algorithms, three network scenarios (Scn.A, Scn.B and Scn.C) are designed, as shown in Table 1 and Figure 3A detailed view of the topology is presented, and three different attack strategies, PH Start, PC Start, and PVStart, are introduced to build a multi-dimensional attack and defense environment. Each scenario is tested against three attack strategies to evaluate the effectiveness of the defense strategies. In the experiment, the attacker starts from the entry point of the network and gradually expands the attack range, while the defender takes action based on the real-time status of the network. Through experiments on nine combinations of scenarios and attack strategies, the aim is to evaluate the effectiveness of the intelligent defense agent based on reinforcement learning algorithm against different attack strategies in complex and dynamic network topologies.

[0174] Table 1 Network scenarios and attack strategy descriptions

[0175]

[0176] The GAT-DQN and GAT-PPO algorithms were tested separately, and three other mainstream network defense models were also evaluated. Each model was run in three scenarios to deal with different attacker strategies. The specific application description of each algorithm is shown in Table 2.

[0177] Table 2 Application and description of the algorithm models compared

[0178]

[0179] Each scenario combines a specific network environment and attack strategy, such as using the PH attack strategy (PHStart.Scn.A) in scenario A. Each algorithm is trained five times independently in each scenario to ensure the robustness of the experimental results. The effectiveness of the defense agent in complex network topologies is systematically evaluated, and the results are averaged across scenarios, while calculating the 95% confidence interval to obtain reliable performance indicators. The training parameters of the reinforcement learning algorithm are as follows:

[0180] The GAT-PPO of the present invention uses a learning rate of 0.0003, a discount factor γ=0.99, a GAE parameter λ=0.95, a PPO clipping parameter ∈=0.2, a batch size of 2048, a mini-batch size of 64, and a GAT hidden layer width of 32;

[0181] Norm-PPO uses a learning rate of 0.0003, a discount factor γ = 0.99, a GAE parameter λ = 0.95, a PPO clipping parameter ∈ = 0.2, a batch size of 2048, and a mini-batch size of 64;

[0182] RNN-PPO uses a learning rate of 0.0003, a discount factor γ = 0.99, a GAE parameter λ = 0.95, a PPO pruning parameter ∈ = 0.2, a batch size of 2048, and a mini-batch size of 64;

[0183] Sb3-PPO uses a learning rate of 0.0003, a discount factor γ = 0.99, a GAE parameter λ = 0.95, a PPO clipping parameter ∈ = 0.2, a batch size of 2048, and a mini-batch size of 64;

[0184] The GAT-DQN of the present invention uses a learning rate of 0.0001, a discount factor γ = 0.99, ∈ = 0.1, a buffer capacity of 10000, a batch size of 64, and a GAT hidden layer width of 16;

[0185] DDQN-DNER uses a learning rate of 0.0001, a discount factor γ = 0.99, an n-step parameter of 5, a PER parameter of 0.6, ∈ = 0.5, a buffer capacity of 10000, a batch size of 128, and a hidden layer width of 128;

[0186] Figure 4 Comparison of average reward curves of five training runs using GAT-PPO, Norm-PPO, RNN-PPO, and Sb3-PPO in different scenarios (95% confidence interval);

[0187] RNN-DDQN uses a learning rate of 0.0001, a discount factor γ = 0.5, a backtracking step of 16, ∈ = 0.1, a buffer capacity of 10000, a batch size of 32, and a hidden layer width of 64.

[0188] DDQN uses a learning rate of 0.0001, a discount factor γ = 0.99, ∈ = 0.5, a buffer capacity of 10000, a batch size of 128, and a hidden layer width of 128.

[0189] Figure 4 The performance comparison analysis of GAT-PPO algorithm and other policy-based deep reinforcement learning algorithms in three different network topologies is shown; each subgraph depicts the reward trajectory of the algorithm for a specified attack strategy in a specific scenario during the training phase; the following conclusions can be drawn from these curves: GAT-PPO shows significant advantages in performance, robustness and convergence speed.

[0190] Scenario A( Figure 4 PC Strat.Scn.A,PH Strat.Scn.A,PV Strat.Scn.A) As the training progresses, the reward values ​​of all algorithms show an upward trend; among them, GAT-PPO shows significant advantages in convergence speed and final reward value; this shows that in a star topology containing 13 hosts (a small number of hosts), most algorithms can effectively analyze the interaction status, identify potential attacks, and formulate effective defense strategies to deal with advanced persistent threats (APTs); however, GAT-PPO can adapt to the environment and optimize strategies faster and more effectively in this small-scale topology.

[0191] Scenario B( Figure 4 PC Strat.Scn.B, PH Strat.Scn.B, PV Strat.Scn.B) In a tree topology containing 23 hosts, the reward values ​​of the other three algorithms dropped significantly and faced significant training difficulties; however, GAT-PPO still showed excellent performance and demonstrated strong adaptability; this indicates that as the number of hosts increases, other algorithms have performance bottlenecks when dealing with complex network interactions, which is manifested as the inability to achieve high reward values; this reflects that the defense agents based on these algorithms are unable to take effective defense actions to mitigate network threats (such as APTs); in contrast, GAT-PPO can effectively overcome these limitations and still maintain excellent performance in more complex networks.

[0192] Scenario C( Figure 4 PC Strat.Scn.C,PH Strat.Scn.C,PV Strat.Scn.C) As the network complexity increases, the reward values ​​of all algorithms are lower in the early stage of training, especially the existing PPO algorithm performs poorly in complex environments; however, GAT-PPO always performs well in the later training stages and achieves significantly higher reward values; this further demonstrates the adaptability and effectiveness of GAT-PPO in complex network environments, and can still significantly improve the defense effect when traditional methods are difficult to cope with; this proves that the GAT-PPO algorithm based on graph neural network enhancement can effectively deal with network threats under various network topologies, and has high adaptability and practicality.

[0193] These experimental results show that GAT-PPO can provide more efficient defense solutions in dynamic and complex network environments with remarkable robustness and fast strategy optimization capabilities, showing obvious advantages over traditional algorithms.

[0194] Table 3 Comparison of average convergence rewards between GAT-PPO, Norm-PPO, RNN-PPO and Sb3-PPO in different scenarios

[0195]

[0196] Table 3 provides a detailed numerical display of the final average reward values ​​of the four algorithms in different network scenarios. The results prove the excellent performance and adaptability of the GAT-PPO algorithm under different scenarios and attack strategies.

[0197] From the perspective of the final convergence reward, GAT-PPO significantly outperforms other algorithms in scenarios A, B, and C. Especially in the ring topology of scenario C, the reward value of GAT-PPO far exceeds that of other algorithms by nearly a thousand times, further proving the effectiveness of GAT-PPO in the formulation of intelligent defense strategies.

[0198] In a single scenario, the reward values ​​of the GAT-PPO algorithm for different attack strategies show high consistency, indicating that the GAT-PPO's intelligent defense agent has broad and efficient capabilities in dealing with diverse attack tactics.

[0199] GAT-PPO eventually converges to a positive reward of 1.0 in the ring network topology, indicating that the ring topology helps to develop a more efficient defense strategy; this strategy can reduce the number of damaged hosts and thus obtain higher positive rewards at the end of the round.

[0200] Although there are more high-value hosts in scenario B, resulting in higher defense costs, GAT-PPO is still able to achieve positive rewards in scenario C, indicating that the defense strategy in the ring topology is more advantageous than the mesh topology, proving the strategy-making ability of GAT-PPO when adapting to changes in network configuration.

[0201] Figure 5 Comparison of average reward curves of five training runs using GAT-DQN, DDQN-DNER, RNN-DQN, and DDQN in different scenarios (95% confidence interval);

[0202] Figure 5 The training curve performance of the GAT-DQN algorithm and other value-based deep reinforcement learning algorithms in network defense is compared, covering different network scenarios and attack strategies.

[0203] From the perspective of the final convergence reward, the GAT-DQN algorithm always shows higher reward values ​​under different scenarios and multiple attack strategies; this result is consistent with the GAT-PPO algorithm, which once again proves the importance of deep reinforcement learning algorithms enhanced by graph neural networks in the formulation of autonomous network security defense strategies; such algorithms significantly improve performance, indicating that they can adopt more precise defense strategies and effectively mitigate network threats such as APT (advanced persistent threats).

[0204] From the perspective of stability, Figure 4 The GAT-PPO and value-based deep reinforcement learning algorithms performed relatively poorly in different scenarios; in particular, except for the GAT-DQN algorithm, other value-based algorithms showed obvious fluctuations during the training process and it was difficult to achieve stable training; this shows that when formulating autonomous network defense strategies, the stability of policy-based deep reinforcement learning algorithms is better than that of value-based algorithms.

[0205] In summary, graph neural networks play a key role in improving the performance of reinforcement learning in the formulation of autonomous network defense strategies; their information propagation and feature aggregation mechanisms significantly improve algorithm performance, thereby enhancing the effectiveness of defense strategies in mitigating network threats; in addition, graph neural networks also contribute to the stability of the defense strategy learning phase, ensuring the stable acquisition of efficient strategies, further demonstrating the wide applicability of graph neural networks in the field of network defense.

[0206] Table 4 Comparison of average convergence rewards between GAT-DQN, DQN-DNER, RNN-DDQN and DDQN in different scenarios

[0207]

[0208] Table 4 shows the average reward values ​​when the value-based deep reinforcement learning algorithm converges under different scenarios and attack strategies.

[0209] The results show that GAT-DQN is significantly better than other algorithms overall, indicating that GAT-DQN can effectively control defense costs and minimize damage to network equipment, demonstrating its excellent defense capabilities in responding to network attacks; in contrast, defense strategies based on other algorithms have obvious limitations and fail to achieve the same level of protection.

[0210] The DDQN-DNER algorithm achieved high reward values ​​in the CybORG environment, proving that it can effectively defend against network attacks in this environment; however, in the simulated complex network scenario, the defense agent using this algorithm failed to effectively control the defense cost and loss.

[0211] The RNN-DDQN algorithm was proposed by the Cubav2 team and won the top three in the Cage-Challenge network defense competition; however, the performance of the RNN and LSTM mechanisms in the algorithm in complex network scenarios did not meet expectations and failed to fully exert their advantages to achieve effective defense.

[0212] As a highly versatile baseline algorithm, DDQN achieved the same reward value as GAT-DQN in scenario A, verifying its wide applicability in specific environments; however, with the increase in the number of network hosts and the evolution of the topology, the algorithm also faces difficulties and cannot generate effective defense strategies.

[0213] The current autonomous defense algorithms based on deep reinforcement learning often encounter performance bottlenecks in complex network scenarios and are difficult to converge to high reward values, thus failing to generate effective defense strategies. The present invention models network observation data as a graph structure and introduces the graph attention network (GAT) in the graph neural network (GNNs). By utilizing the information propagation and feature aggregation mechanism of the graph attention network, the defense agent can fully consider the interactions between hosts, IoT devices and other entities in the network topology. This enables the defense agent to identify potential attack paths and take more accurate and efficient defense actions. The research results show that this method can effectively respond to network threats and achieve efficient and autonomous defense strategies.

[0214] The present invention proposes the application of reinforcement learning technology based on graph attention network in the field of network security, aiming to alleviate the network threats faced by other network systems such as the Internet of Things, and to open up new directions for network security research; by introducing the reinforcement learning algorithms GAT-DQN and GAT-PPO enhanced by graph neural networks, the intelligence level of network defense strategies is significantly improved; these algorithms show excellent performance and strong adaptability in multiple complex network scenarios and attack strategy combinations, fully demonstrating their wide applicability in dealing with advanced persistent threats (APTs) and other network threats; in addition, the present invention improves the Yawning Titan (YT) simulation environment and adds real Internet of Things network scenarios, so that the defense strategies generated by training are closer to actual application needs; at the same time, a variety of attack penetration strategies are designed to train more comprehensive automated defense strategies, thereby enhancing the overall defense capabilities of the system.

[0215] Based on the above ideal embodiments of the present invention, the relevant staff can make various changes and modifications without departing from the technical concept of the present invention through the above description. The technical scope of the present invention is not limited to the contents of the specification, and its technical scope must be determined according to the scope of the claims.

Claims

1. IoT defense method based on graph attention-enhanced reinforcement learning, characterized by include: Step 1: Construct a network attack and defense scenario and use graph data structures to model observation information; Step 2: Model the action space of the attacker and the defense agent; Step 3: Model the attack and defense process in the network attack and defense scenario; Step 4: Use graph normalization to normalize the feature matrix of the node; use the self-attention mechanism of GAT to dynamically aggregate the current state information of each device based on the interaction between the network device and its neighboring devices; use the graph downsampling pooling strategy to selectively retain the node host according to the self-attention score; use the reinforcement learning algorithm enhanced by the graph attention network to train the defense agent.

2. The IoT defense method based on graph attention enhanced reinforcement learning according to claim 1 is characterized in that: Step 1 specifically includes: Step 11: Identify device types as different nodes of the graph, represent connection relationships with edges, and set edge weights; Step 12: Collect key information of different devices and abstract the key information into node features.

3. The IoT defense method based on graph attention enhanced reinforcement learning according to claim 2 is characterized in that: The formula for edge weight is: Among them, σ(s,t) represents the total number of shortest paths between nodes s and y, σ(s,t|e) represents the number of shortest paths through edge e, and V is the node set.

4. The IoT defense method based on graph attention enhanced reinforcement learning according to claim 1 is characterized in that: Step three specifically includes: Step 31: The attacker uses the vulnerability to elevate privileges, performs an attacker scan to identify the target host, and performs a vulnerability scan on the target host; performs privilege escalation or a zero-day attack to destroy the target host; and penetrates the entire network through lateral movement; Step 32: Probabilistic success action based on the number of vulnerabilities of the target host; Step 33: The defense agent monitors the network state and captures the network state s at each time step. t ; The defense agent analyzes the current network state st through the policy network π(a|s) to generate a defense action; and calculates the expected cumulative reward; After executing the defense action, the next state s t+1 According to the state transition probability P(s t+1 |s t ,a t ) is determined, and the immediate reward r is obtained according to the reward function t ; Status t 、Action a t and status t+1 , r t Stored in the replay buffer and used to update the policy network.

5. The IoT defense method based on graph attention enhanced reinforcement learning according to claim 4 is characterized in that: The formula for the reward function is: R step -C action -N compromised ×10 When the number of steps reaches the preset maximum number of steps per round, if the defense is successful, the final reward will be returned Otherwise return Among them, C action represents the cost of the defender taking action, N compromised and N uncompromised Respectively represent the number of hosts hacked by the attacker and the number of hosts that have not been hacked, N total Indicates the total number of devices.

6. The IoT defense method based on graph attention enhanced reinforcement learning according to claim 4 is characterized in that: A hierarchical penalty mechanism is introduced into the reward function. The formula for constructing the improved reward function is: R′ step =-C action -N normal ×1-N high ×10 Among them, N normal is the number of compromised common hosts, N high is the number of high-value hosts compromised.

7. The IoT defense method based on graph attention enhanced reinforcement learning according to claim 4 is characterized in that: The reinforcement learning algorithm enhanced by the graph attention network trains the defense agent, including: using the GAT-DQN algorithm to train the defense agent, the GAT-DQN algorithm includes: Define the parameters of the target network Qtarget, initialize the parameters learning rate α, discount factor γ, exploration rate ∈ and maximum number of training rounds max_episodes; At the beginning of each training round, the environment is reset and the initial state s0 is obtained; s0 includes the state information of the host, which is reconstructed into x0 and edge_index0 based on the graph data structure; Generate a random number rand∈[0,1] at each time step; if rand<∈, randomly select an action a t Otherwise, use the current Q network according to the formula Get the optimal defensive action a t ; After the action is executed, the environment returns the reward r t , enter state s t+1 , reconstructed again to get x t+1 ,edge_index t+1 ; The four-tuple (x t ,edge_index t ,a t ,r t ,x t+1 ,edge_index t+1 , done) is added to the experience playback area; if the playback area size exceeds batch_size, a small batch is randomly sampled for training; The loss function is calculated using the sampled batch data, and the Q network parameters are optimized using the gradient descent method. The loss function is calculated based on the formula of the target network: The parameters of the Q network are synchronized to the target network every fixed number of steps, and the exploration rate is decayed.

8. The IoT defense method based on graph attention enhanced reinforcement learning according to claim 1, characterized in that: The reinforcement learning algorithm enhanced by the graph attention network to train the defense agent also includes: using the GAT-PPO algorithm to train the defense agent, the GAT-PPO algorithm includes: Initialize the policy network π θ (a|x, edge.index) and value network V(x, edge.index; ω); set discount factor γ, GAE parameter λ, maximum training round T total , batch size B and number of experience sampling steps M B ; For each training epoch, execute: Reset the environment, extract the current state and reconstruct it into the node feature matrix x0 and the edge connection relationship edge_index0; Based on the policy network π θ , calculate the action probability distribution p according to the current state θ , and then sample the action a according to the distribution t ; Execute action a t , get reward r t , next state feature x t+1 , edge connection relationship edge_index t+1 And the completion mark done; (x t ,edge.index t ,a t ,r t ,x t+1 ,edge.index t+1 ) is stored in ReplayBuffer; Calculate the TD error at each time step with the following formula: δ t =r t +γV(x t+1 ,edge.index t+1 ;ω)-V(x t ,edge.index t ;ω) Calculate the advantage function A based on TD error t , the formula is: According to the advantage function A t Calculate the target value function: V t =A t +V(x t ,edge.index t ;π) Sample data from ReplayBuffer in small batches and optimize the policy network π based on the PPO loss function θ , and update the parameter θ by gradient descent; Optimize the value network, minimize the value function error, and update the parameter π; After each update, clear the ReplayBuffer.

9. The IoT defense method based on graph attention enhanced reinforcement learning according to claim 1, characterized in that: Dynamically aggregating current status information of each device based on the interaction of the network device with its neighboring devices includes: assigning weights according to the importance of the connected devices, and aggregating the status information based on the weights.

10. The IoT defense method based on graph attention enhanced reinforcement learning according to claim 1, characterized in that: Utilizing a graph downsampling pooling strategy to selectively retain nodes according to attention scores includes: integrating a graph sampling pooling algorithm into a policy network to enable the defense agent to focus on key hosts in the current state.

Images