A method for constructing a network anomaly detection model based on federated reinforcement learning

Abstract

The application provides a network anomaly detection model construction method based on federated reinforcement learning, comprising: each network node collects local traffic data, the network node uses the local traffic data to train a node detection model, and then uploads model parameters of the trained node detection model to a server; the server dynamically aggregates the model parameters uploaded by each network node to generate a global model; the server generates local model parameters using a knowledge distillation method according to the global model, and the network node loads the local model parameters in the node detection model to obtain an anomaly detection model. The application solves the problems in the prior art that the order of magnitude difference between normal traffic and abnormal traffic leads to serious imbalance of samples, and the detection system uses a static model architecture, cannot dynamically adjust the inference strategy according to the computing resources of edge devices, and leads to low detection accuracy and poor real-time performance of abnormal traffic.

Description

Technical Field

[0001] This invention relates to the field of network security technology, and in particular to a method for constructing a network anomaly detection model based on federated reinforcement learning. Background Technology

[0002] In the current field of network anomaly detection, technological evolution faces multiple bottlenecks: network attack events exhibit long-tail distribution characteristics in real-world scenarios, and the magnitude difference between normal and abnormal traffic leads to a severe imbalance in samples. In addition, the heterogeneity of different network domains due to business types and topologies causes data distribution drift. Furthermore, existing detection systems generally adopt static model architectures, which cannot dynamically adjust inference strategies based on the computing resources of edge devices. Moreover, when a single reinforcement learning model is deployed across multiple nodes, the lack of a physically driven reward mechanism results in low accuracy and poor real-time performance in detecting abnormal traffic, severely restricting the transferability and continuous learning capabilities of detection models in heterogeneous network environments. Summary of the Invention

[0003] To address the shortcomings of existing technologies, this invention provides a method for constructing a network anomaly detection model based on federated reinforcement learning. This method solves the problems of severe sample imbalance caused by the difference in magnitude between normal and abnormal traffic, and the low accuracy and poor real-time performance of anomaly detection due to the use of a static model architecture in the detection system, which cannot dynamically adjust the inference strategy according to the computing resources of edge devices.

[0004] According to an embodiment of the present invention, a method for constructing a network anomaly detection model based on federated reinforcement learning includes:

[0005] Each network node collects local traffic data and constructs a unified state vector based on the local traffic data;

[0006] Based on the unified state vector, network nodes use the GAN network to generate simulated anomaly samples and combine the simulated anomaly samples with local traffic data to form a balanced dataset.

[0007] Network nodes train node detection models using balanced datasets, and then upload the model parameters of the trained node detection models to the server.

[0008] The server dynamically aggregates the model parameters uploaded by each network node to generate global model parameters, and then generates a global model based on the global model parameters.

[0009] The server generates local model parameters using knowledge distillation based on the global model and broadcasts them to all network nodes. The network nodes load the local model parameters into the node detection model to obtain the anomaly detection model, and then use the anomaly detection model to perform network anomaly detection.

[0010] Preferably, local traffic data includes historical traffic data and current traffic data;

[0011] The method for constructing a unified state vector based on local traffic data is as follows:

[0012] Use a sliding window to extract a local traffic data over a period of time from the current traffic data;

[0013] Calculate the average number of bytes in all packets within the local traffic data;

[0014] The port access frequency is obtained by counting the number of traffic packets received by each receiving port in the local traffic data.

[0015] The percentage of traffic packets corresponding to various network protocols in local traffic data is counted to the total number of traffic packets, and the protocol behavior entropy of various network protocols is calculated based on the percentage.

[0016] The historical anomaly density is obtained by calculating the percentage of data packets identified as abnormal in historical traffic data out of the total number of traffic packets.

[0017] By concatenating the average number of bytes, port access frequency, protocol behavior entropy, and historical anomaly density, a unified state vector is obtained.

[0018] Preferably, the GAN network includes a generator and a discriminator. After constructing the GAN network, a cosine loss function is added to the generator, and a class label constraint is added to the discriminator.

[0019] Preferably, the method for training a node detection model using a balanced dataset includes:

[0020] Construct a PPO learning model as a node detection model and reconstruct the reward function;

[0021] The balanced dataset is divided into a training set and a validation set. The training set is used to train the PPO learning model, and the validation set is used to validate the training results of the PPO learning model.

[0022] Preferably, the reconstructed reward function is as follows:

[0023] R t =w1(t)R preciaion +w2(t)R recall -w3(t)L delay -w4(t)C comm

[0024] Where w1, w2, w3, w4 are dynamic weights, and R precision For accuracy, R recall For recall rate, L delayFor model inference latency, C comm For federal communications expenses.

[0025] Preferably, the method by which the server dynamically aggregates the model parameters uploaded by each network node to generate global model parameters includes:

[0026] Each network node calculates the F1 score of its corresponding PPO learning model on the validation set, and then uploads the local traffic data and F1 score to the server.

[0027] The server calculates the percentage of local traffic data for each network node relative to the total amount of local traffic data received by the server.

[0028] The server assigns dynamic weights to each network node based on its percentage and F1 score, and then aggregates the model parameters of each node based on the dynamic weights to obtain the global model parameters.

[0029] Preferably, for network nodes whose F1 score is greater than a preset score threshold, additional dynamic weights are added to the network node.

[0030] Preferably, the method by which the server generates local model parameters using knowledge distillation based on the global model includes:

[0031] Based on the global model, the server uses knowledge distillation to generate a local model, and then inputs all received local traffic data into the global model and the local model to generate a global probability distribution and a local probability distribution.

[0032] Calculate the KL divergence loss between the local model and the global model based on the global probability distribution and the local probability distribution;

[0033] The total loss function of the local model is constructed by combining KL divergence loss and cross-entropy loss;

[0034] Add resource constraints to the local model, and then train the local model using all local traffic data until the total loss function and resource constraints are minimized, thus obtaining the local model parameters.

[0035] Preferably, when a new network node is added, the new network node loads the global model and collects local traffic data. Then, it calculates the KL divergence between the local traffic data and the local traffic data of all network nodes received by the server. If the KL divergence exceeds a preset threshold, the model parameters of the global model loaded by the new network node are fine-tuned using the local traffic data collected by the new network node.

[0036] Preferably, after the new network node fine-tunes the model parameters of the loaded global model, it generates new model parameters and uploads the new model parameters and the collected local traffic data to the server, after which the server performs dynamic aggregation again.

[0037] Compared with the prior art, the present invention has the following beneficial effects:

[0038] This invention collects local traffic data from various network nodes and uses a GAN network to generate simulated abnormal samples, expanding the amount of abnormal sample data to balance the number of abnormal and normal samples in the dataset. Then, the network nodes send their respective node detection model parameters and local traffic data to the server. The server dynamically aggregates the model parameters of each network node, and then uses knowledge distillation to generate local model parameters, which are then returned to the network nodes and loaded. In this way, the computing resources of all network nodes can be dynamically adjusted, improving the accuracy and real-time performance of abnormal traffic detection. Attached Figure Description

[0039] Figure 1 This is a flowchart illustrating the construction process of the network anomaly detection model according to an embodiment of the present invention. Detailed Implementation

[0040] The technical solutions of the present invention will be further described below with reference to the accompanying drawings and embodiments.

[0041] like Figure 1 As shown in the figure, this invention proposes a method for constructing a network anomaly detection model based on federated reinforcement learning, including:

[0042] Each network node collects local traffic data and constructs a unified state vector based on the local traffic data;

[0043] Local traffic data includes historical traffic data and current traffic data. The following characteristics of traffic packets in the current traffic data are analyzed using a sliding window (window size T seconds):

[0044] ① Average Bytes: Calculates the average number of bytes in all packets within the local traffic data.

[0045]

[0046] Among them, byte t Let T be the number of bytes in the t-th packet within the window, where T is the sliding window duration in seconds.

[0047] ② Protocol type distribution: The percentage of protocols such as TCP, UDP, and ICMP is calculated using the following formula:

[0048]

[0049] Where, N proto N represents the number of traffic packets of a certain protocol type within the window. total This represents the total number of data packets.

[0050] ③ Port access frequency:

[0051]

[0052] Where, N target-port The number of times each receiving port appears within the window, i.e., the number of traffic packets received by the receiving port.

[0053] ④ Protocol Behavior Entropy: Calculate the Shannon entropy of the protocol type sequence to characterize the randomness of traffic distribution.

[0054]

[0055] Where, p i Let be the probability of the i-th protocol type appearing within the window, and n be the total number of protocol types.

[0056] ⑤ Historical Anomaly Density: Based on the detection results of the past M=24 hours, the percentage of abnormal traffic per unit time (per hour) is calculated:

[0057]

[0058] Where, N anomaly N represents the number of abnormal traffic packets detected in the past M hours. total-hour This represents the total number of data packets for the corresponding time period.

[0059] The above features are then concatenated into a state vector S, forming a single state vector, which is then used as the input to the subsequent reinforcement learning model.

[0060] S = [μ byte ,f port H proto ,ρ anomaly ]

[0061] Where S is a 4-dimensional vector, which contains the average number of bytes, port access frequency, protocol behavior entropy, and historical anomaly density.

[0062] Based on the unified state vector, network nodes use the GAN network to generate simulated anomaly samples and combine the simulated anomaly samples with local traffic data to form a balanced dataset.

[0063] A GAN network consists of a generator and a discriminator. The generator generates simulated anomalous samples to increase the amount of anomalous data, balancing the number of anomalous and normal samples in the dataset. The discriminator determines whether the generated simulated anomalous samples are anomalous data. Since anomalous data is diverse, a separate generator is assigned to each type of anomalous data identified as such in the local traffic data. k (There are K generators in total), where k = 1, 2, ..., K. The generator input is a Gaussian noise vector z ~ N(0, I), and the output is a synthesized anomalous sample.

[0064]

[0065] Where z is a d-dimensional Gaussian noise vector, G k For the generator of the k-th type of attack, These are the generated abnormal samples.

[0066] Then, the cosine similarity loss function is used to ensure the generation of samples. Align with the feature distribution of the true sample x, and compute the feature vector f(x) with... Similarity:

[0067]

[0068] Where f(·) is the feature extraction function.

[0069] Next, the discriminator is optimized. The objective function of discriminator D is to minimize the difference between the real sample x and the generated sample x. Binary cross-entropy loss:

[0070]

[0071] Where D(x) is the discriminant probability of the discriminator classifying the true sample x, D(G k (z) represents the discrimination probability of the generated sample. Constraints are then introduced to generate the sample. Category label y k Using cross-entropy loss:

[0072] L C =-E z～N(0,I) [y k logC(G k (z))]

[0073] Among them, C(G) k (z) represents the predicted probability of the generated sample by the classifier, y k This represents the true label for the k-th type of attack.

[0074] In summary, it can be concluded that the generator Gk The total loss function is:

[0075] L G =L C +λL sim +τL D

[0076] Here, λ and τ are weighting coefficients that control the balance between classification and similarity discrimination losses.

[0077] Network nodes train node detection models using balanced datasets, and then upload the model parameters of the trained node detection models to the server.

[0078] In this invention, the network nodes use the PPO reinforcement learning model as the node detection model. The PPO reinforcement learning model (Proximal Policy Optimization) serves as the dynamic decision engine in this step. Its input is a balanced dataset. In this model, the balanced dataset also needs to be represented in the form of a unified state vector S, and the output is the detection action chosen by the network node in the current state (e.g., adjusting the classification threshold, selecting a feature subset, controlling model complexity, etc.). PPO achieves continuous optimization of the detection strategy under multi-objective constraints by maximizing long-term cumulative rewards. Its application principle is as follows:

[0079] 1) Policy network architecture:

[0080] The policy network π uses a fully connected neural network. The input layer dimension is consistent with the unified state vector S, and the output layer has Gaussian distributed parameters. The action a is defined. t Continuous decision space:

[0081] a t ~N(μ) θ (S t ),σ θ (S t ))

[0082] Where, μ θ and σ θ These are the mean and variance parameters of the policy network output, respectively.

[0083] 2) Experience collection: Based on the current strategy π θ Interact with the environment to collect state-action-reward sequences {(S t ,a t ,R t )}, stored in the experience replay pool;

[0084] 3) Dominance estimation: The advantage function A is calculated using generalized advantage estimation (GAE). t Balance deviation and variance:

[0085]

[0086] Among them, V φ Let V be a value function network, where γ is the discount factor and λ is the GAE coefficient. φ The network is trained using Temporal Difference Learning, with the goal of minimizing the mean squared error between the predicted value and the actual cumulative reward. The value function network V... φ Update parameter φ by minimizing the mean squared error (MSE) between the predicted and target values:

[0087]

[0088] Among them, L V For value function loss, E t Represents the expectation of time step t.

[0089] 4) Policy Update: The policy update magnitude is limited by a pruning mechanism to prevent gradient explosion. The objective function J(θ) of PPO is shown below:

[0090]

[0091] Where, π θ For the current policy network, For the old policy network, A t Let be the dominance function, and ο be the pruning threshold.

[0092] To maintain stability during the learning process, the PPO reinforcement learning model standardizes the reward function to achieve zero mean and unit variance. However, this standardization does not consider the computational resources and communication overhead of each network node. Therefore, this invention reconstructs the reward function of the PPO reinforcement learning model by introducing a physically-informed reward function R. t :

[0093] Define a multi-objective reward function R t The performance and resource consumption indicators of the fusion detection are as follows:

[0094] R t =w1(t)R preciaion +w2(t)R recall -w3(t)L delay -w4(t)C comm

[0095] Where w1, w2, w3, w4 are dynamic weights, and R precision For accuracy, R recall For recall rate, L delay For model inference latency, C comm For federal communications expenses.

[0096] Dynamic weight calculation: The weighting coefficients are dynamically adjusted based on the system's real-time resource utilization.

[0097]

[0098] Where α and β are adjustable coefficients, and Resource available For currently available computing resources, Resource total This is the feature extraction function for total resource quantity.

[0099] Calculation of detection performance indicators: Accuracy R precision With recall rate R recall Calculate using the following formula:

[0100]

[0101] Where TP is the number of true positives, FP is the number of false positives, and FN is the number of false negatives.

[0102] Resource consumption index modeling: Model inference latency L delay and federal communications overhead C comm Calculate according to the following formula,

[0103]

[0104] Where B is the batch size, t infer (x i ) represents the inference time for the i-th sample, θ represents the number of model parameters, and f comm For communication frequency, B w This refers to network bandwidth.

[0105] Next, the balanced dataset is divided into a training set and a validation set. The training set is used to train the PPO learning model, and the validation set is used to validate the training results of the PPO learning model.

[0106] The server dynamically aggregates the model parameters uploaded by each network node to generate global model parameters, and then generates a global model based on the global model parameters.

[0107] After the PPO learning model of each network node is trained, the model parameters and local traffic data (unified state vector) of each PPO learning model are uploaded to the server. The server calculates the dynamic weight w based on the ratio between the amount of local traffic data of each node and all local traffic data received by the server, and the F1 score of each network node's PPO learning model on the validation set during training. k :

[0108]

[0109] Where, Nk Let N be the local data volume of node k. total For the total data volume across all nodes, F1 k Let k be the F1 score of the validation set.

[0110] Then, the dynamic weights are normalized to ensure...

[0111]

[0112] in, The normalized weights are represented by K, where K represents the network nodes participating in the aggregation.

[0113] Federated learning is a distributed machine learning framework whose core objective is to train a global model collaboratively by multiple nodes without sharing the original data. Its workflow consists of two phases:

[0114] 1) Local training: Each network node updates the model parameters θ based on local traffic data. k ;

[0115] 2) Parameter aggregation: The server collects the θ values ​​from all nodes. k According to weight w k Fusion generates global parameter θ global .

[0116] In this invention, the server uses dynamic weights. Improved aggregation effect:

[0117]

[0118] Where, θ k The local model parameters for network node k are updated based on the loss function L of local traffic data. k (θ):

[0119] θ k ←θ k -η▽ θ L k (θ)

[0120] Where η is the learning rate.

[0121] The PPO learning model's F1 score on the validation set is higher than the threshold γ. th For network nodes, add extra weight to them:

[0122]

[0123] Where δ is the weight enhancement coefficient.

[0124] Then, by weighting, the model parameters of all network nodes are aggregated to obtain the global model parameters θ. global The server builds a new PPO reinforcement learning model and loads the global model parameters to generate a global model.

[0125] The server generates local model parameters using knowledge distillation based on the global model and broadcasts them to all network nodes. The network nodes load the local model parameters into the node detection model to obtain the anomaly detection model, and then use the anomaly detection model to perform network anomaly detection.

[0126] Federated learning consists of a federated server and multiple nodes connected to it. Therefore, the federated server can be viewed as a "teacher model," and each network node as a "student model." Based on this, and according to the knowledge distillation framework, it can be derived from the global model parameters θ. global (Teacher model) generates soft labels to guide the local model θ. light (Student model) training takes local traffic data received by the server from all network nodes as input and outputs a local probability distribution, namely p. global and p light

[0127]

[0128] Where T is the temperature parameter used for the softening probability distribution, f global and f light These are the forward computation functions for the global model and the local model, respectively.

[0129] Combined with KL divergence loss L KL and cross-entropy loss L CE The total loss function is:

[0130] L total =λ·L KL (p light ||p global )+(1-λ)·L CE (p light ,y)

[0131] Where λ is the loss weight coefficient, y is the true label, and L KL L measures the difference in distribution CE Supervise the accuracy of classification.

[0132] The difference in distribution between the outputs of the local and global models is calculated using the KL divergence loss:

[0133]

[0134] in, and These are the probability values ​​for the i-th class, respectively.

[0135] Then add model complexity loss L resource Limit the number of parameters in the local model ||θ light || and computational complexity N FLOPs :

[0136] L resource =γ·‖θ light ‖+η·N FLOPs

[0137] Where, N FLOPs This represents the number of floating-point operations performed in a single inference iteration of the model.

[0138] Finally, the overall training objective is introduced as minimizing the hybrid loss and resource constraints, resulting in the local model parameters:

[0139]

[0140] in, These are the optimized local model parameters.

[0141] The server then broadcasts the local model parameters to all network nodes. The network nodes load the local model parameters, obtain the anomaly detection model, and use the anomaly detection model to perform network anomaly detection.

[0142] When a new network node joins, the new node loads the global model θ. global Then, calculate the local traffic data D of the new network node. local The distribution P local The distribution P of local traffic data received by the server from all network nodes (excluding the new network node) global KL divergence:

[0143]

[0144] Among them, P global P represents the implicit data distribution of the global model parameters. local For the local data distribution of the new node, X is the feature space with temperature parameters, used for the softening probability distribution.

[0145] When the KL divergence is less than the preset threshold θ th If the KL divergence exceeds the preset threshold θ, then the global model is used directly as the model for the new network node. th Then local fine-tuning will be initiated:

[0146]

[0147] Local traffic data D based on new network nodeslocal Using cross-entropy loss L CE Update model parameters θ global to θ new

[0148]

[0149] Where f(x; θ) is the model's predicted probability for sample x.

[0150] The new node uploads the finely adjusted parameter θ new The data is then sent to the federated server, where the server re-aggregates the global parameter θ′. global :

[0151]

[0152] Among them, w new The weight of the new network node is determined by the amount N of its local traffic data. new Dynamic adjustment, the adjustment method is as follows:

[0153]

[0154] The federated server will then update the global model θ′ global Broadcasting to all network nodes completes the incremental federated learning loop. The server dynamically aggregates the model parameters from each network node, then uses knowledge distillation to generate local model parameters, which are returned to the network nodes for loading. This allows for dynamic adjustment of computing resources across all network nodes, improving the accuracy and real-time performance of abnormal traffic detection.

[0155] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit it. Although the present invention has been described in detail with reference to preferred embodiments, those skilled in the art should understand that modifications or equivalent substitutions can be made to the technical solutions of the present invention without departing from the spirit and scope of the technical solutions of the present invention, and all such modifications or substitutions should be covered within the scope of the claims of the present invention.

Claims

1. A method for constructing a network anomaly detection model based on federated reinforcement learning, characterized in that: include: Each network node collects local traffic data and constructs a unified state vector based on the local traffic data; Based on the unified state vector, network nodes use the GAN network to generate simulated anomaly samples and combine the simulated anomaly samples with local traffic data to form a balanced dataset. The network nodes use a balanced dataset to train a node detection model, including: Construct a PPO learning model as a node detection model and reconstruct the reward function; The balanced dataset is divided into a training set and a validation set. The training set is used to train the PPO learning model, and the validation set is used to validate the training results of the PPO learning model. Upload the trained node detection model parameters to the server; The server dynamically aggregates the model parameters uploaded from various network nodes to generate global model parameters, including: Each network node calculates the F1 score of its corresponding PPO learning model on the validation set, and then uploads the local traffic data and F1 score to the server. The server calculates the percentage of local traffic data for each network node relative to the total amount of local traffic data received by the server. The server assigns dynamic weights to each network node based on its proportion and F1 score. At the same time, for network nodes whose F1 score is greater than a preset score threshold, additional dynamic weights are added to the network node. Then, the model parameters of each node are weighted and aggregated according to the dynamic weights to obtain the global model parameters. Generate a global model based on global model parameters; The server generates local model parameters using knowledge distillation based on the global model and broadcasts them to all network nodes. The network nodes load the local model parameters into the node detection model to obtain the anomaly detection model, and then use the anomaly detection model to perform network anomaly detection.

2. The method for constructing a network anomaly detection model based on federated reinforcement learning as described in claim 1, characterized in that: Local traffic data includes historical traffic data and current traffic data; The method for constructing a unified state vector based on local traffic data is as follows: Use a sliding window to extract a local traffic data over a period of time from the current traffic data; Calculate the average number of bytes in all packets within the local traffic data; The port access frequency is obtained by counting the number of traffic packets received by each receiving port in the local traffic data. The percentage of traffic packets corresponding to various network protocols in local traffic data is counted to the total number of traffic packets, and the protocol behavior entropy of various network protocols is calculated based on the percentage. The historical anomaly density is obtained by calculating the percentage of data packets identified as abnormal in historical traffic data out of the total number of traffic packets. By concatenating the average number of bytes, port access frequency, protocol behavior entropy, and historical anomaly density, a unified state vector is obtained.

3. The method for constructing a network anomaly detection model based on federated reinforcement learning as described in claim 1, characterized in that: The GAN network includes a generator and a discriminator. After constructing the GAN network, a cosine loss function is added to the generator, and a class label constraint is added to the discriminator.

4. The method for constructing a network anomaly detection model based on federated reinforcement learning as described in claim 1, characterized in that: We introduce four physical information components: precision, recall, inference latency, and federated communication overhead. We also assign dynamically changing weights to these four physical information components to reconstruct the reward function.

5. The method for constructing a network anomaly detection model based on federated reinforcement learning as described in claim 1, characterized in that: The server uses knowledge distillation to generate local model parameters based on the global model, including the following methods: Based on the global model, the server uses knowledge distillation to generate a local model, and then inputs all received local traffic data into the global model and the local model to generate a global probability distribution and a local probability distribution. Calculate the KL divergence loss between the local model and the global model based on the global probability distribution and the local probability distribution; The total loss function of the local model is constructed by combining KL divergence loss and cross-entropy loss; Add resource constraints to the local model, and then train the local model using all local traffic data until the total loss function and resource constraints are minimized, thus obtaining the local model parameters.

6. The method for constructing a network anomaly detection model based on federated reinforcement learning as described in claim 1, characterized in that: When a new network node joins, it loads the global model and collects local traffic data. Then, it calculates the KL divergence between the local traffic data and the local traffic data of all network nodes received by the server. If the KL divergence exceeds a preset threshold, the model parameters of the global model loaded by the new network node are fine-tuned using the local traffic data collected by the new network node.

7. The method for constructing a network anomaly detection model based on federated reinforcement learning as described in claim 1, characterized in that: After fine-tuning the model parameters of the loaded global model, the new network node generates new model parameters and uploads the new model parameters and collected local traffic data to the server, where the server then performs dynamic aggregation again.

Images