Prototype Retrieval-Augmented Federated Learning System for Robust Intrusion Detection

This repository provides the official implementation and experimental resources for our paper, “Prototype Retrieval-Augmented Federated Learning System for Robust Intrusion Detection.” It includes the source code, experimental datasets, configuration files.

Acknowledgement

All baselines and our method are implemented on top of PFLlib.

We are very grateful for their outstanding library.

Dataset

All used datasets can be downloaded here

Please unzip the RAR file into the "system/dataset" folder.

Quickly Start (baseline FL algorithms + our proposed FedPRO)

In the following bash examples, we use β = 0.5 as a running example.

Please modify the -data argument to select a dataset with a different dataset and β value, using the format {dataset folder}_{β value}.

(1) Generate model weights for FL baselines

Run the command for each FL baseline without --skip_FL True. The generated model weights are saved under temp/{dataset}/{algorithm}/{timestamp}/ by default.

Example:

python system/main.py -data NSLKDD_0.5 -algo FedAvg -ncl 5 # FedAvg. Here, ncl denotes the number of classes.

(2) Evaluate the proposed FedPRO

After generating the model weights, run the following commands with --skip_FL True.

The program loads model weights from temp by default and evaluates the effectiveness of the proposed FedPRO.

NSL-KDD (Personalized test)

python system/main.py -data NSLKDD_0.5 -algo Local -ncl 5 -topk 1 --skip_FL True # Local
python system/main.py -data NSLKDD_0.5 -algo FedAvg -ncl 5 -topk 1 --skip_FL True # FedAvg
python system/main.py -data NSLKDD_0.5 -algo MOON -ncl 5 -topk 1 --skip_FL True # MOON
python system/main.py -data NSLKDD_0.5 -algo FedProto -ncl 5 -topk 1 --skip_FL True # FedProto
python system/main.py -data NSLKDD_0.5 -algo FedTGP -ncl 5 -topk 1 --skip_FL True # FedTGP
python system/main.py -data NSLKDD_0.5 -algo GPFL -ncl 5 -topk 1 --skip_FL True # GPFL
python system/main.py -data NSLKDD_0.5 -algo FedDBE -ncl 5 -topk 1 --skip_FL True # FedDBE

NSL-KDD (Global test ）

python system/main.py -data NSLKDD_global_0.5 -algo Local -ncl 5 -topk 1 --skip_FL True # Local
python system/main.py -data NSLKDD_global_0.5 -algo FedAvg -ncl 5 -topk 1 --skip_FL True # FedAvg
python system/main.py -data NSLKDD_global_0.5 -algo MOON -ncl 5 -topk 1 --skip_FL True # MOON
python system/main.py -data NSLKDD_global_0.5 -algo FedProto -nc 5 -ncl 5 -topk 1 -fd 122 --skip_FL True # FedProto
python system/main.py -data NSLKDD_global_0.5 -algo FedTGP -ncl 5 -topk 1 --skip_FL True # FedTGP
python system/main.py -data NSLKDD_global_0.5 -algo GPFL -ncl 5 -topk 1 --skip_FL True # GPFL
python system/main.py -data NSLKDD_global_0.5 -algo FedDBE -ncl 5 -topk 1 --skip_FL True # FedDBE

CICIDS-2018 (Personalized test)

python system/main.py -data CICIDS_0.5 -algo Local  -ncl 7 -topk 1  --skip_FL True # Local
python system/main.py -data CICIDS_0.5 -algo FedAvg  -ncl 7 -topk 1 --skip_FL True # FedAvg
python system/main.py -data CICIDS_0.5 -algo MOON  -ncl 7 -topk 1 --skip_FL True # MOON
python system/main.py -data CICIDS_0.5 -algo FedProto  -ncl 7 -topk 1  --skip_FL True # FedProto
python system/main.py -data CICIDS_0.5 -algo FedTGP  -ncl 7 -topk 1  --skip_FL True # FedTGP
python system/main.py -data CICIDS_0.5 -algo GPFL  -ncl 7 -topk 1  --skip_FL True # GPFL
python system/main.py -data CICIDS_0.5 -algo FedDBE  -ncl 7 -topk 1  --skip_FL True # FedDBE

CICIDS-2018 (Global test)

python system/main.py -data CICIDS_global_0.5 -algo Local -ncl 7 -topk 1  --skip_FL True # Local
python system/main.py -data CICIDS_global_0.5 -algo FedAvg  -ncl 7 -topk 1  --skip_FL True # FedAvg
python system/main.py -data CICIDS_global_0.5 -algo MOON -nc -ncl 7 -topk 1  --skip_FL True # MOON
python system/main.py -data CICIDS_global_0.5 -algo FedProto  -ncl 7 -topk 1  --skip_FL True # FedProto
python system/main.py -data CICIDS_global_0.5 -algo FedTGP -ncl 7 -topk 1 --skip_FL True # FedTGP
python system/main.py -data CICIDS_global_0.55 -algo GPFL -ncl 7 -topk 1  --skip_FL True # GPFL
python system/main.py -data CICIDS_global_0.5 -algo FedDBE  -ncl 7 -topk 1  --skip_FL True # FedDBE

Experiments on additional large-class datasets (BODMAS, Ominiglot)

BODMAS example

python main.py -data BODMAS_0.1 -m cicids -algo FedAvg -ncl 267 -skip_FL True -topk 1 # Personalized test
python main.py -data BODMAS_global_0.1 -m cicids -algo FedAvg -ncl 267 -skip_FL True -topk 1 # Global test

Omniglot example

python main.py -data Omniglot_0.1 -m ResNet10 -algo FedAvg -ncl 1623 -skip_FL True -topk 1 # Personalized test
python main.py -data Omniglot_global_0.1 -m ResNet10 -algo FedAvg -ncl 1623 -skip_FL True -topk 1 # Global test

Citation

@ARTICLE{11499419,
  author={Zhou, Hanlin and Yan, Huiru and Nian, Jiawei and Liu, Cong and Wang, Ying and Chen, Xiaomin and Theodoropoulos, Georgios and Cheng, Long},
  journal={IEEE Transactions on Computers}, 
  title={{Prototype Retrieval-Augmented Federated Learning System for Robust Intrusion Detection}}, 
  year={2026},
  volume={75},
  number={8},
  pages={2797-2810},
  doi={10.1109/TC.2026.3688741}}