# MT-FBERT: Malicious Traffic Detection Based on Efficient Federated Learning of BERT
## 3 Mt-Fbert
### The Malicious Traffic Detection Stage  
In the malicious traffic detection stage, the model is initialized with the parameters of the pre-trained model, and then the FL method is adopted to collaboratively train the detection model among multiple organizations on labeled malicious traffic data, which can effectively protect the data privacy of the organizations. Moreover, given the computational and communication overhead of pre-training models in federated learning, we introduce an optimized framework featuring neuron compression and knowledge distillation. The neuron compression module utilizes magnitude-based pruning [49] to selectively retain a subset of neurons within the transformer layers. Each transformer layer comprises two primary sub-layers: a multi-head attention (MHA) mechanism and a feed-forward network (FFN). The FFN contains more parameters than the MHA [50]. We implement the neuron compression by removing neurons with low weights in the FFN. The FFN  
Figure 4: An illustration of the Masked Burst Modeling (MBM). comprises three sequential components: (1) an input fully-connected layer, (2) a ReLU activation layer, and (3) an output fully-connected layer. The output of the FFN can be denoted as  
\[\text{FFN}(x)=\max(0,xW_{1}+b_{1})W_{2}+b_{2}, \tag{4}\]  
where \(x\) is the input of FFN and its dimensionality is \(H\). Assume the fully connected layer in FFN has dimensionality \(d_{f}\). \(W_{1}\in\mathbb{R}^{H\times d_{f}}\) and \(W_{2}\in\mathbb{R}^{d_{f}\times H}\) are the weight matrices of the two fully connected layers. \(b_{1}\in\mathbb{R}^{d_{f}}\) and \(b_{2}\in\mathbb{R}^{H}\) are the biases. We measure \(m\)th neuron in the fully connected layer based on the weights. We retain the neurons with large weights. Detailed descriptions can be found in Appendix A. The compressed transformer encoder layer is denoted as Sub-T-Layer, as shown in Figure 1.  
Moreover, neuron compression introduces bias in gradient estimation, with the update direction diverging from the uncompressed model's trajectory. To mitigate accumulated gradient errors during federated learning of the compressed detection model, we employ knowledge distillation to minimize the divergence between the compressed and uncompressed models. The distillation approach operates at two granularities: (1) layer-level distillation that matches intermediate representations, and (2) neuron-level distillation that preserves critical activation patterns. Layer-level distillation ensures consistency between the compressed and uncompressed models by aligning both forward pass outputs and backward pass gradients at each transformer layer. Each layer's distillation loss can be denoted as  
\[loss_{layer}=\left(\sum_{i=1}^{D}\left\|O_{i}-O_{i}^{\prime}\right\|\right)+ \mu\left\|W_{1}W_{2}-W_{1}^{\prime}W_{2}^{\prime}\right\|, \tag{5}\]  
where \(D\) is the size of the distillation dataset, \(O_{i}\) and \(O_{i}^{\prime}\) represent the outputs from the uncompressed and the compressed models, respectively, \(W_{1}\) and \(W_{2}\) are the weight matrices in the uncompressed model, \(W_{1}^{\prime}\) and \(W_{2}^{\prime}\) are the weight matrices in the compressed models, and \(\mu\) is the regularization coefficient. The frozen FFN parameters in the compressed model introduce catastrophic forgetting risks during federated learning, as they cannot adapt to client-specific malicious traffic patterns. To address this, we employ neuron-level distillation to preserve activation knowledge across server-client updates. Neuron-level distillation identifies low-impact neurons for the malicious traffic detection task by computing their Average Percentage of Zero activations (_APoZ_) [51]. Set the labelled malicious traffic data on the client \(x_{i}\in D^{\prime}\), \(x_{i}\) is a token sequence input into the transformer layer.