APT flow detection model lightweight training method and APT flow detection method

Abstract

The invention belongs to the field of Internet of Things security, and particularly relates to an APT traffic detection model lightweight training method and an APT traffic detection method. The invention provides an APT (Advanced Persistent Threat) traffic detection model training method, which comprises the following steps of: acquiring network traffic, and extracting statistical characteristics and packet length sequence characteristics from the network traffic; enhancing the statistical features and the packet length sequence features to obtain statistical important features and packet length sequence mapping features; performing feature connection on the statistical important features and the packet length sequence mapping features to obtain cascade features; and inputting the cascade features into a neural network model for training to obtain a trained APT flow detection model. The invention further provides an APT flow detection method, which comprises the following steps of: deploying the model on terminal equipment, and detecting unknown flow. According to the method, the sample with the maximum information amount is selected in a targeted mode for labeling, and compared with a traditional method, the workload of safety experts is remarkably reduced.

Description

Technical Field

[0001] This invention belongs to the field of Internet of Things security, specifically relating to a lightweight training method for an APT traffic detection model and an APT traffic detection method. Background Technology

[0002] Advanced Persistent Threats (APTs), as a type of targeted and persistent cyberattack, have become a major challenge in the field of cybersecurity in recent years. APT attacks are characterized by their high degree of stealth, long duration, and strong targeting, making traditional feature-matching-based traffic detection methods ineffective in identifying them. With the development of deep learning technology, deep learning-based traffic detection methods have shown great potential in the field of APT detection.

[0003] However, existing deep learning-based APT traffic detection methods face the following key challenges: (1) Imbalanced data: APT attacks exhibit long-tail distribution characteristics in real-world environments, resulting in a significant imbalance in numbers compared to normal traffic. This leads to insufficient learning of minority class samples during model training, resulting in decreased detection performance. (2) High annotation costs: Professional annotation of APT traffic relies on highly specialized security analysts, leading to high annotation costs and difficulty in obtaining large amounts of labeled data. (3) Terminal deployment challenges: Traditional deep learning models have a large number of parameters and high computational and storage requirements, making efficient deployment on resource-constrained terminal devices, especially IoT terminals and edge devices. (4) Blurred decision boundaries: The boundary between APT attacks and normal traffic is blurred, making it difficult for traditional active learning strategies to effectively identify key samples at the decision boundary.

[0004] Existing technologies typically employ data balancing methods such as oversampling and undersampling, or basic active learning strategies, to address some of the problems, but lack a comprehensive solution tailored to the characteristics of APT traffic. This invention aims to effectively solve these problems by proposing an APT traffic detection method that combines a boundary-aware active learning strategy with pre-training fine-tuning. Summary of the Invention

[0005] To address the above problems, the first aspect of this invention provides a lightweight training method for an APT traffic detection model, comprising the following steps:

[0006] Acquire network traffic, and extract statistical features and packet length sequence features from the network traffic;

[0007] The statistical features and packet length sequence features are enhanced to obtain statistical importance features and packet length sequence mapping features;

[0008] The statistically important features and the packet length sequence mapping features are concatenated to obtain cascaded features;

[0009] The cascaded features are input into the neural network model for training to obtain a trained APT traffic detection model.

[0010] Furthermore, the network traffic includes APT traffic and service traffic.

[0011] Furthermore, the statistical features and packet-length sequence features are enhanced, including:

[0012] The statistical features are input into a pre-trained random forest classifier, the importance score of the statistical features is calculated, and the statistically important features are obtained based on the importance score.

[0013] Calculate the time decay weight of the packet length sequence feature, apply time decay weighting and position encoding to the packet length sequence feature according to the time decay weight to obtain a weighted sequence; input the weighted sequence into the Transformer encoder to obtain the packet length sequence mapping feature.

[0014] Furthermore, the feature connection step includes:

[0015] The statistically important features are reduced in dimensionality using average pooling.

[0016] L2 normalization is applied to the package length sequence mapping features and the dimensionality-reduced statistical importance features;

[0017] The cascaded features are obtained by connecting the normalized packet length sequence mapping features and statistical importance features along the feature dimension.

[0018] Furthermore, the neural network model adopts the ALBERT-based architecture, and the training includes self-supervised pre-training and fine-tuning.

[0019] Furthermore, the self-supervised pre-training tasks include a masking feature prediction task and a feature sequence ranking task. The self-supervised pre-training process includes jointly optimizing the joint loss function of the masking feature prediction task and the feature sequence ranking task. The fine-tuning includes ranking and labeling the cascaded features by value and fine-tuning the model using the ALF strategy.

[0020] Furthermore, the cascaded features are ranked by value according to the dynamic sample evaluation mechanism of the BASS strategy.

[0021] A second aspect of the present invention provides an APT traffic detection method, comprising:

[0022] Deploy the above model on terminal devices to detect unknown traffic.

[0023] A third aspect of the present invention provides an electronic device including a memory and a processor, the memory storing a computer program configured to be executed by the processor, the computer program including instructions for performing the methods described above.

[0024] A fourth aspect of the present invention provides a storage medium for storing a computer program, which, when executed by a computer, implements the method described above.

[0025] The beneficial effects of this invention are as follows:

[0026] (1) BASS is used in Z Pool The algorithm identifies and labels samples with the highest value. ALF dynamically fine-tunes the LoRa parameters based on different features. By employing the boundary-aware active learning strategy BASS, it selectively labels samples with the most information, reducing labeling requirements by approximately 90% compared to traditional methods, significantly reducing the workload of security experts.

[0027] (2) By adopting a dynamic sample evaluation mechanism and elastic memory, the problem of long-tail distribution of APT traffic is effectively addressed, and the detection accuracy of a few types of APT attacks is improved.

[0028] (3) By using the adversarial extension enhancement mechanism, boundary samples are generated to help the model accurately locate the decision boundary and improve the robustness of the application.

[0029] (4) Combining the ALBERT architecture and low-rank fine-tuning technology, the number of model parameters is reduced by 93%, making it suitable for deployment on resource-constrained terminal devices.

[0030] (5) This invention recognizes that APT attacks typically span long time periods and exhibit long-term temporal correlations. Therefore, by introducing a time decay weight function ω(t), the attention weights dynamically decay as the time gap increases, thereby capturing the long-term dependent attention adjustment mechanism. An adaptive feature importance adjustment mechanism is introduced, using random forest feature importance calculation to design a learnable feature weight matrix, adjusting the contribution of different features to attention calculation. Depending on the attack stage, the importance of key features changes dynamically, dynamically strengthening the influence of highly discriminative features and adapting to changes in attack patterns. Through the time decay weight function and adaptive feature importance adjustment, the model dynamically adapts to highly discriminative features and long-distance APT dependence, improving its ability to express the temporal and statistical characteristics of APT traffic.

[0031] (6) The method of the present invention has been verified in a variety of network environments. It is not only applicable to APT traffic detection, but can also be extended to other network security scenarios, such as malware detection, intrusion detection and other fields. Attached Figure Description

[0032] Figure 1 This is a flowchart illustrating a specific embodiment of the present invention.

[0033] Figure 2 This is a framework diagram of a specific embodiment of the present invention.

[0034] Figure 3 This is experimental comparison data between the method of the present invention and the prior art. Detailed Implementation

[0035] The present invention will now be described in further detail with reference to the accompanying drawings. The examples given are only for explaining the present invention and are not intended to limit the scope of the present invention.

[0036] The first aspect of this embodiment proposes a lightweight training method for APT traffic detection models. It comprehensively utilizes active learning and pre-trained language model fine-tuning techniques to address the issues of data imbalance and lightweight terminal deployment in APT traffic detection models. The process is as follows: Figure 1 and Figure 2 As shown.

[0037] Step 1: Terminal Traffic Collection

[0038] The traffic collection device is run on the terminal device to collect traffic over a period of time. The specific implementation of the collection process is as follows:

[0039] Deploy a lightweight traffic capture tool on the terminal device and set the capture time window to a continuous 24 hours to ensure coverage of network activity at different times;

[0040] Configure filtering rules to mainly capture traffic from common protocols such as TCP, UDP, HTTP, HTTPS, and DNS. These protocols are the main carriers of APT attacks. The collected traffic is stored in PCAP format.

[0041] In this embodiment, the traffic collected by the terminal device includes normal business traffic (benign traffic) and attack traffic generated by manually executed APT attacks.

[0042] Step 2, Dataset Partitioning

[0043] The collected traffic was divided into an unlabeled sample pool Z at an 8:2 ratio. Pool and labeled test set Z test The specific implementation is as follows:

[0044] To ensure a consistent distribution of each category across the test set, a stratified random sampling method was used for partitioning; for Z... test The samples in the dataset are professionally labeled, categorized as benign traffic or different types of APT traffic, and the predicted categories are obtained simultaneously; the remaining samples are used to construct an unlabeled sample pool Z. PoolThis is used for subsequent pre-training and active learning; all features are standardized so that their mean is 0 and their standard deviation is 1, which enhances the model's convergence.

[0045] Step 3: Feature Extraction and Enhancement

[0046] For the unlabeled sample pool Z Pool and labeled test set Z test The traffic is subjected to feature extraction and enhancement, and the specific implementation is as follows:

[0047] (1) Statistical feature extraction: Extract the statistical features Xs of each flow from the traffic, including: number of data packets, total number of bytes, flow duration, average packet size, standard deviation of packet size, maximum / minimum packet size, average packet interval, standard deviation of packet interval, flag statistics: frequency of TCP flag occurrence (SYN, ACK, FIN, etc.), source / destination port type (common / uncommon), protocol features: protocol type and distribution.

[0048] (2) Packet length sequence feature extraction: Extract the length sequence Xt of the first T=128 packets of each flow to form a time sequence feature. For time sequence features with a length exceeding 128, a truncation strategy is adopted to retain only the length sequence of the first 128 packets. For time sequence features with a length less than 128, a padding strategy is adopted to pad the length to 128, thus obtaining the packet length sequence feature.

[0049] (3) Statistical feature enhancement: The adaptive adjustment mechanism of feature importance is applied, and the implementation steps are as follows:

[0050] Input all statistical features Xs = (Xs1, Xs2, ..., Xsi, ...) into a pre-trained random forest classifier, and calculate the Gini impurity reduction Gi for each feature fi = Xsi; calculate the feature importance score Si = Gi / ∑Gj based on the feature impurity reduction Gi; use the feature importance score Si to perform a weighted transformation on the original statistical features Xs to obtain the statistically important features Xs′ = Xs·diag(S), where S is the feature importance score vector and diag(S) is a diagonal matrix with S as the diagonal element.

[0051] (4) Packet-length sequence feature enhancement: Transformer encoding enhanced by applying the time decay weight function ω(t) is implemented as follows:

[0052] For the packet-length sequence features Xt=[x1,x2,…,x T ], calculate the time decay weight ω(t)=e -λ·(T-t) / T , where λ is the decay coefficient, which is set to 0.5 in this embodiment, t is the time step in the sequence, and T is the total length of the sequence;

[0053] The weighted sequence Xt is calculated by applying time decay weighting and positional encoding to the features of the packet-length sequence. ω Xt ω =[ω(1)·x1+PE(1),ω(2)·x2+PE(2),…,ω(T)·x T +PE(T), where ω(·) is the weighting coefficient and PE(·) is the position code;

[0054] The weighted sequence Xt ω Input a 12-layer Transformer encoder (with a hidden layer dimension of 256 and 8 attention heads) to obtain the packet-length sequence mapping feature Xt′.

[0055] (5) Feature concatenation: The statistically important feature Xs′ is weighted by feature importance and has a dimension of ds; the packet-length sequence mapping feature Xt′ is encoded by Transformer and has a dimension of dt×T (dt is the encoding dimension); to unify the feature dimension, the packet-length sequence mapping feature Xt′ is subjected to average pooling on the temporal dimension: Reduce the dimension of Xt″ to dt; apply L2 normalization to both types of features to ensure they are numerically equivalent:

[0056]

[0057] The normalized features are concatenated along the feature dimensions to obtain cascaded features X = [Xs′]. L2 ;Xt′ L2 Finally, the dimension of X is (ds+dt).

[0058] Step 4, Building the ALBERT pre-trained model

[0059] The ALBERT pre-trained model is run in an unlabeled sample pool Z. Pool Self-supervised pre-training is performed on the model to obtain the pre-trained model M1. The specific implementation is as follows:

[0060] The ALBERT-base architecture was chosen, which features parameter sharing. This architecture transforms the cascaded features X into an input format acceptable to the ALBERT model, including feature segmentation and sequence length adjustment to 128. Two pre-training tasks were designed and implemented:

[0061] (1) Masked Feature Prediction (MFP): 15% of the feature values ​​are randomly masked, and the model is trained to predict the masked features. The loss function is the mean squared error. Where M is the number of occluded features, x i Indicates in [Xs′ L2 ;Xt′ L2 The true value of a certain position in the graph, x i′ represents the predicted value at this location.

[0062] (2) Feature Sequence Ordering Prediction (SOP): The input sequence is randomly split and shuffled. The model is trained to determine whether the sequence is ordered. The loss function is L. SOP For binary cross-entropy:

[0063] L SOP =-y·log(p)+(1-y)·log(1-p), where y is the correct sort label and p is the model prediction probability.

[0064] (3) Joint optimization: Simultaneously optimize the joint loss function of the two pre-training tasks:

[0065] L Pre =0.5·L MFP +0.5·L SOP

[0066] (4) Stop condition setting: Set the query stop condition to reach N=100 rounds of query or prediction accuracy Acc i The target accuracy is greater than or equal to 99.5%.

[0067] After pre-training, model M1 learned a general representation of traffic data, laying the foundation for subsequent supervised fine-tuning.

[0068] Step 5: Active Learning and Model Fine-tuning

[0069] Utilizing a boundary-aware active learning BASS strategy, for Z Pool The cascaded features in the model are ranked and labeled for value, and the ALF strategy is used for model fine-tuning. Specifically, BASS first ranks and labels the cascaded features by calculating their values, and ALF then adaptively allocates weights based on the features before cascading. During fine-tuning, the weighted cascaded features are used. The specific implementation is as follows:

[0070] (1) From Z Pool Ten samples are randomly selected from the data and labeled to construct a label set Z. L Furthermore, a minority class sample memory is established to store minority class samples.

[0071] (2) Perform the following active learning iterative process until the stopping condition is met:

[0072] 2.1) Calculate Z using the Dynamic Sample Evaluation Mechanism (DSM) of the BASS strategy. Pool The value score for all samples is V(x) = 0.4·U(x) + 0.3·C(x) + 0.3·D(x), where: U(x) = 1 - max(P(y|x)), which is the uncertainty score; C(x) = T / (nc +T), which is the category balance score, n c To predict the current number of samples in category c, the temperature coefficient T is set to 1.0; D(x) = min(||xm) i ||2), for diversity score, m i These are samples from the minority class memory bank;

[0073] 2.2) Generate boundary samples using the adversarial extension enhancement (AEA) of the BASS strategy. Where ε is the perturbation magnitude, set to 0.01. This represents the increment added to the boundary samples. L is the loss function, and θ is the parameter involved in fine-tuning the model. Only samples that change their predicted class after adding perturbation will be selected for labeling, further expanding the candidate sample pool.

[0074] 2.3) Elastic Memory (EMB) Update Using the BASS Strategy: When the probability difference of sample x before and after adding perturbation, as determined by the model, is less than the threshold τ = 0.1, it is added to the boundary sample cache B; when the sample class c satisfies n c <median({n j When adding samples to the minority class sample memory Ms, the minority class sample memory is added.

[0075] (3) Select the n=100 samples with the highest value scores from V(x) for expert annotation and add them to the annotation set Z. L In each round, a batch of selected labeled samples will be added to the labeled set. Each round of training will use the current labeled set until the query stopping condition is met, at which point the number of samples in the labeled set will no longer increase.

[0076] (4) Fine-tune model M1 using the ALF strategy. The specific steps are as follows:

[0077] Statistical Important Feature Processing: Rank rs = 8, Scaling Factor α s =16, dropout=0.1; Packet length sequence mapping feature processing: rank rt=16, scaling factor α t =32, dropout=0.1; Apply LoRA adapter to query, key, value matrix and dense layer, retain complete parameters of classifier head; Set weight decay coefficient to 0.01, use optimizer: AdamW for training optimization.

[0078] (5) Using the labeled test set Z test For the fine-tuned model M1 i Perform an evaluation and calculate the accuracy (Acc). i ;

[0079] (6) Check if the current round number c has reached N=100 or the accuracy Acc. i If the target accuracy ACC = 99.5% is achieved, the iteration stops; otherwise, return to step (3). When the stopping condition is met, the fine-tuned model M is output as the final APT flow detection model.

[0080] This embodiment further discloses comparative data between the method of the present invention (BASS) and the prior art, such as... Figure 3 As shown in the UNSW dataset example, the BASS strategy achieved 95% accuracy with only 2800 labeled samples, while the traditional random sampling strategy required more than 4500 samples to achieve the same performance.

[0081] RAND random sampling strategy, ENT entropy-based uncertain sampling strategy, MAR interval-based uncertain sampling strategy, CONF minimum confidence-based uncertain sampling strategy, Diversity based diversity sampling strategy, and Coreset greedy sampling strategy.

[0082] The above are all commonly used benchmarks in the field of active learning. BALT and IMBL are research results from top international conferences in the fields of natural language processing and computer vision in the past three years.

[0083] Iteration represents the number of iterations, F1(%) is the F1 score, and both F1 and AUC are commonly used evaluation metrics in classification. Bold text represents the best-performing data, and underlined text represents the second-best-performing data. If comparison results must be included, the specific values ​​in the annotations can be replaced with the values ​​in the table.

[0084] d1 represents the terminal scenario, using the USTC-TFC dataset; d2 represents the IoT scenario, using the UNSW dataset.

[0085] The second aspect of this embodiment proposes an APT traffic detection method, as detailed below:

[0086] The finely tuned model M is deployed on small terminal devices to perform real-time detection of newly collected unknown traffic. First, dynamic quantization is used to convert FP32 weights to INT8 format to reduce storage space. The front-end processing module is responsible for traffic capture, feature extraction, feature enhancement, and feature concatenation. It loads model M to perform traffic detection and classification, and triggers an alarm when APT traffic is detected.

Claims

1. A lightweight training method for an APT traffic detection model, comprising the following steps: Acquire network traffic, and extract statistical features and packet length sequence features from the network traffic; The statistical features and packet length sequence features are enhanced to obtain statistical importance features and packet length sequence mapping features; The statistically important features and the packet length sequence mapping features are concatenated to obtain cascaded features; The cascaded features are input into the neural network model for training to obtain a trained APT traffic detection model.

2. The method according to claim 1, characterized in that, The network traffic includes APT traffic and service traffic.

3. The method according to claim 1, characterized in that, Enhancement of the statistical features and packet-length sequence features includes: The statistical features are input into a pre-trained random forest classifier, the importance score of the statistical features is calculated, and the statistically important features are obtained based on the importance score. Calculate the time decay weight of the packet length sequence feature, apply time decay weighting and position encoding to the packet length sequence feature according to the time decay weight to obtain a weighted sequence; input the weighted sequence into the Transformer encoder to obtain the packet length sequence mapping feature.

4. The method according to claim 1, characterized in that, The feature connection steps include: The statistically important features are reduced in dimensionality using average pooling. L2 normalization is applied to the package-length sequence mapping features and the dimensionality-reduced statistical importance features; The cascaded features are obtained by connecting the normalized packet length sequence mapping features and statistical importance features along the feature dimension.

5. The method according to claim 1, characterized in that, The neural network model adopts the ALBERT-based architecture, and the training includes self-supervised pre-training and fine-tuning.

6. The method according to claim 5, characterized in that, The self-supervised pre-training tasks include occlusion feature prediction and feature sequence ranking. The self-supervised pre-training process includes jointly optimizing the joint loss function of the occlusion feature prediction and feature sequence ranking tasks. The fine-tuning includes ranking and labeling the cascaded features by value and fine-tuning the model using the ALF strategy.

7. The method according to claim 6, characterized in that, The cascaded features are ranked by value according to the dynamic sample evaluation mechanism of the BASS strategy.

8. An APT traffic detection method, comprising: Deploy the model described in any one of claims 1-7 on a terminal device to detect unknown traffic.

9. An electronic device comprising a memory and a processor, the memory storing a computer program configured to be executed by the processor, the computer program including instructions for performing the method of any one of claims 1 to 7 or 8.

10. A storage medium storing a computer program, which, when executed by a computer, implements the method of any one of claims 1 to 7 or 8.

Images