# MT-FBERT: Malicious Traffic Detection Based on Efficient Federated Learning of BERT
## 3 Mt-Fbert
### The Pre-Training Stage
#### 3.2.2 The Pre-Training Tasks  
Two pre-training tasks are employed: Masked Burst Modeling (MBM) and Next Burst Prediction (NBP). These tasks are designed to capture the contextual relationships between traffic tokens.  
The MBM task is similar to the masked language modeling and Masked Image Modeling utilized by BERT [31]. In the MBM task, the traffic token sequence of the input is masked, and the model is required to recover the masked tokens based on the contextual tokens. As shown in Figure 4, each original burst token sequence is first randomly masked to obtain the masked input. During the masking, each token in the original sequence is randomly chosen with 15% probability. As the chosen token from each sequence, we replace it with [MASK] at 80% chance, or choose a random token to replace it at 10% chance, or leave it unchanged at 10% chance, respectively. The model then encodes the masked token sequence with the encoding layer, which consists of three embedding layers, and the vector representation of tokens is obtained by summing up these embeddings. The vectors of tokens are further input into the representation layer, which includes several encoder layers of a transformer to learn the intermediate representations of tokens by using the contextual information. The intermediate representations are processed through a mapping network to predict the masked tokens based on the contextual information; meanwhile, the prediction probabilities of the tokens can be generated. Assume the sequence as \(S=\{token_{i}\}_{i=1}^{n}\), we mask \(k\) tokens and \(MASK_{i}\) represents the masked token at the \(i_{th}\) position in the sequence \(S\). The contextual information of the \(i_{th}\) position in the sequence \(S\) is set as \(Context_{i}=\{token_{1},\dots,token_{i-1},token_{i+1},\dots,token_{n}\}\). The loss function of the MBM task can be defined as  
\[loss_{MBM}=-\sum_{i=1}^{k}log(P(MASK_{i}=token_{i}|Context_{i})), \tag{1}\]  
Figure 3: An example of token representation based on three embeddings.  
The NBP task is similar to the Next Sentence Prediction utilized by BERT. A burst is a sequence of consecutive packets originating from the same direction. We learn the correlation between packets in a burst based on the NBP task. In the NBP task, a burst is equally divided into two segments, which can be denoted as \(sub-burst_{a}\) and \(sub-burst_{b}\). For instance, as shown in Figure 2, \"Burst 1\" is represented by six traffic tokens, with the first three tokens comprising \(sub-burst_{a}\) and the remaining three tokens forming \(sub-burst_{b}\). The NBP task is designed to predict whether two given segments belong to the same burst while simultaneously learning traffic representations by modeling packet-level dependencies within bursts. Specifically, we select \(sub-burst_{a}\) and \(sub-burst_{b}\) as input, and concatenate the two segments end to end, with [SEP] connecting them in the middle. During 50% of the time, the chosen \(sub-burst_{b}\) is the true next segment of \(sub-burst_{a}\). During the remaining 50% of the time, \(sub-burst_{b}\) is randomly selected as long as it is not the next segment of \(sub-burst_{a}\). Formally, the NBP task constitutes a binary classification problem, where the goal is to determine whether two segments belong to the same burst. Assume there are \(n\) bursts, denoting as \(B=\left\{\left(sub-burst_{a}^{i},sub-burst_{b}^{j}\right)\right\}_{j=1}^{n}\). The loss function of the NBP task can be defined as  
\[loss_{NBP}=-{\sum_{j=1}^{n}}log\Big{(}P\Big{(}y_{j}\Big{|}\Big{(}sub-burst_{ a}^{j},sub-burst_{b}^{j}\Big{)}\Big{)}\Big{)}, \tag{2}\]  
where \(y_{j}\) equals 0 if \(sub-burst_{b}^{j}\) is the next segment of \(sub-burst_{a}^{j}\), and 1 otherwise.  
Overall, the loss function of the pre-training task combines the above two loss terms, which is defined as  
\[loss=loss_{MBM}+loss_{NBP}, \tag{3}\]