Le référentiel de TrafficLLM , un cadre d'adaptation LLM universel pour apprendre une représentation robuste du trafic pour tous les LLM open source dans des scénarios du monde réel et améliorer la généralisation à travers diverses tâches d'analyse du trafic.
Remarque : ce code est basé sur ChatGLM2 et Llama2. Un grand merci aux auteurs.
[2024.10.28] Nous avons mis à jour le code d'adaptation pour l'utilisation de GLM4 pour créer TrafficLLM, qui a une vitesse de réglage et d'inférence plus rapide que ChatGLM2. Accédez à Adapt2GLM4 pour plus de détails.
TrafficLLM repose sur un cadre de réglage sophistiqué utilisant le langage naturel et les données de trafic, qui propose les techniques suivantes pour améliorer l'utilité des grands modèles de langage dans l'analyse du trafic réseau.
Tokenisation du domaine de trafic. Pour combler l'écart de modalité entre le langage naturel et les données de trafic hétérogènes, TrafficLLM introduit la tokenisation du domaine du trafic pour traiter les diverses entrées de tâches de détection et de génération de trafic pour l'adaptation LLM. Ce mécanisme étend efficacement le tokenizer natif de LLM en formant un modèle spécialisé de tokenisation sur des corpus de domaines de trafic à grande échelle.
Pipeline de réglage à deux étages. TrafficLLM utilise un pipeline de réglage en deux étapes pour réaliser l'apprentissage robuste des représentations de LLM sur différentes tâches du domaine du trafic. Le pipeline entraîne LLM à comprendre les instructions et à apprendre les modèles de trafic liés aux tâches à différentes étapes, ce qui s'appuie sur la compréhension des tâches TrafficLLM et les capacités de raisonnement du trafic pour diverses tâches de détection et de génération de trafic.
Adaptation extensible avec réglage fin efficace des paramètres (EA-PEFT). Pour adapter LLM en vue d'une généralisation à de nouveaux environnements de trafic, TrafficLLM propose une adaptation extensible avec un réglage fin des paramètres (EA-PEFT) pour mettre à jour les paramètres du modèle avec une faible surcharge. La technique divise les capacités du modèle en différents modèles PEFT, ce qui permet de minimiser les coûts des scénarios dynamiques générés par les changements de modèle de trafic.
Nous avons publié les ensembles de données de formation de TrafficLLM, qui contiennent plus de 0,4 million de données de trafic et 9 000 instructions humaines pour l'adaptation LLM à différentes tâches d'analyse du trafic.
Instruction Datasets : les ensembles de données d'instructions sont utilisés pour aider LLM à acquérir les connaissances du domaine des tâches de détection ou de génération de trafic et à comprendre quelle tâche doit être effectuée dans différents scénarios.
Traffic Datasets : les ensembles de données de trafic contiennent les données de réglage du trafic que nous avons extraites des ensembles de données de trafic publics, ce qui aide LLM à apprendre le modèle de trafic dans différentes tâches en aval.
Pour construire le corpus de langage naturel en tant qu'instructions humaines dans TrafficLLM, nous avons collecté les 9 209 instructions spécifiques à des tâches supervisées par des experts et des assistants IA. Les statistiques sont présentées comme suit :
| Tâches courantes | Tâches en aval | Abréviation. | #Échantillon |
|---|---|---|---|
| Détection du trafic | Détection du trafic de logiciels malveillants | MTD | 1,0K |
| Détection des botnets | BND | 1,1K | |
| Détection de DoH malveillant | MDD | 0,6K | |
| Détection des attaques Web | LIASSE | 0,6K | |
| Détection des attaques APT | DAA | 0,6K | |
| Détection VPN cryptée | MVE | 1,2K | |
| Détection du comportement Tor | À déterminer | 0,6K | |
| Classification des applications cryptées | CAE | 0,6K | |
| Prise d'empreintes digitales de sites Web | WF | 0,6K | |
| Dérive conceptuelle | CD | 0,6K | |
| Génération de trafic | Génération de trafic de logiciels malveillants | MTG | 0,6K |
| Génération de trafic de botnet | BTG | 0,1K | |
| Génération VPN cryptée | EVG | 0,4K | |
| Génération d'applications cryptées | EAG | 0,6K |
Pour évaluer les performances de TrafficLLM sur divers scénarios de réseau, nous avons extrait plus de 0,4 million de données de réglage à partir d'ensembles de données de trafic accessibles au public afin de mesurer les capacités de TrafficLLM à détecter ou à générer du trafic malveillant et inoffensif. Les statistiques sont présentées comme suit :
| Ensembles de données | Tâches | Abréviation. | #Échantillon |
|---|---|---|---|
| USTC TFC 2016 | Détection du trafic de logiciels malveillants | MTD | 50,7K |
| Botnet ISCX2014 | Détection des botnets | BND | 25,0K |
| DoHBrw 2020 | Détection de DoH malveillant | MDD | 47,8K |
| SCCI 2010 | Détection des attaques Web | LIASSE | 34,5K |
| DAPT2020 | Détection des attaques APT | DAA | 10,0K |
| VPN ISCX 2016 | Détection VPN cryptée | MVE | 64,8K |
| ISCXTor 2016 | Détection du comportement Tor | À déterminer | 40,0K |
| CSTNET2023 | Classification des applications cryptées | CAE | 97,6K |
| CW-100 2018 | Prise d'empreintes digitales de sites Web | WF | 7,4K |
| APP-53 2023 | Dérive conceptuelle | CD | 109,8K |
1. Préparation de l'environnement 2. Formation TrafficLLM 2.1. Préparation du point de contrôle pré-entraîné 2.2. Ensemble de données de prétraitement 2.3. Formation Tokenizer de domaine de trafic (facultatif) 2.4. Réglage des instructions du langage neuronal 2.5. Optimisation du trafic spécifique à une tâche 2.6. Adaptation extensible avec PEFT (EA-PEFT) 3. Évaluation de TrafficLLM 3.1. Préparation des points de contrôle et des données 3.2. Évaluation en coursTable des matières :
1. Préparation de l'environnement [Retour en haut]
Veuillez cloner le dépôt et installer l'environnement requis en exécutant les commandes suivantes.
conda create -n trafficllm python=3.9 conda activate trafficllm# Cloner notre clone TrafficLLMgit https://github.com/ZGC-LLM-Safety/TrafficLLM.gitcd TrafficLLM# Installer les bibliothèques requisespip install -r Requirements.txt# Si trainingpip installe rouge_chinese nltk jieba ensembles de données
TrafficLLM utilise trois techniques de base : la tokenisation du domaine du trafic pour traiter les instructions et les données de trafic, le pipeline de réglage en deux étapes pour comprendre la sémantique du texte et apprendre les modèles de trafic dans différentes tâches, et l' EA-PEFT pour mettre à jour. paramètres du modèle pour l’adaptation de nouveaux scénarios.
TrafficLLM est formé sur la base des LLM open source existants. Veuillez suivre les instructions pour préparer les points de contrôle.
ChatGLM2 : préparez le modèle de base ChatGLM, qui est un LLM open source avec des exigences de déploiement légères. Veuillez télécharger ses poids ici. Nous utilisons généralement le modèle v2 avec des paramètres 6B.
Other LLMs : Pour adapter d'autres LLM aux tâches d'analyse du trafic, vous pouvez réutiliser les données de formation dans le repo et modifier leurs scripts de formation selon les instructions officielles. Par exemple, Llama2 doit enregistrer le nouvel ensemble de données dans les configurations.
Pour extraire des données de formation appropriées pour l'apprentissage LLM à partir des ensembles de données de trafic bruts, nous concevons des extracteurs spécialisés pour prétraiter les ensembles de données de trafic pour différentes tâches. Le code de prétraitement contient les paramètres suivants à configurer.
input : le chemin de l'ensemble de données de trafic brut (le chemin du répertoire principal qui contient les sous-répertoires étiquetés. Chaque sous-répertoire étiqueté contient les fichiers bruts .pcap à prétraiter).
dataset_name : Le nom de l'ensemble de données de trafic brut (cela permet de déterminer si le nom a été enregistré dans les codes de TrafficLLM).
traffic_task : Tâches de détection ou tâches de génération.
granularity : granularité au niveau du paquet ou du flux.
output_path : chemin de l'ensemble de données d'entraînement en sortie.
output_name : nom de l'ensemble de données d'entraînement en sortie.
Il s'agit d'une instance permettant de prétraiter des ensembles de données de trafic bruts pour les tâches de détection du trafic au niveau des paquets.
prétraitement du CD python preprocess_dataset.py --input /Votre/Raw/Dataset/Path --dataset_name /Votre/Raw/Dataset/Name --traffic_task détection --granularity packet --output_path /Votre/Sortie/Dataset/Path --output_name /Votre /Output/Dataset/Name
TrafficLLM introduit un tokenizer de domaine de trafic pour gérer le langage neuronal et les données de trafic. Si vous souhaitez entraîner un tokenizer personnalisé avec votre propre ensemble de données, veuillez modifier le model_name et data_path dans le code.
model_name : Le chemin du modèle de base qui contient le tokenizer natif.
data_path : les ensembles de données de formation extraits du processus de prétraitement.
Veuillez suivre la commande pour utiliser le code.
tokenisation de CD python traffic_tokenizer.py
Préparer les données : les données de réglage des instructions du langage neuronal sont nos ensembles de données d'instructions collectées pour la compréhension des tâches d'analyse du trafic.
Démarrer le réglage : après les étapes susmentionnées, vous pouvez démarrer le réglage de la première étape en utilisant trafficllm_stage1.sh. Il y a un exemple comme ci-dessous :
PRE_SEQ_LEN=128
LR=2e-2
NUM_GPUS=1exporter CUDA_VISIBLE_DEVICES=1
torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py
--do_train
--train_file ../datasets/instructions/instructions.json
--validation_file ../datasets/instructions/instructions.json
--preprocessing_num_workers 10
--prompt_colonne instruction
--response_column sortie
--overwrite_cache
--cache_dir ../cache
--model_name_or_path ../models/chatglm2/chatglm2-6b
--output_dir ../models/chatglm2/peft/instruction
--overwrite_output_dir
--max_source_length 1024
--max_target_length 32
--per_device_train_batch_size 1
--per_device_eval_batch_size 1
--gradient_accumulation_steps 16
--predict_with_generate
--max_steps 20000
--logging_steps 10
--save_steps 4000
--learning_rate $LR
--pre_seq_len $PRE_SEQ_LEN Préparer les données : les ensembles de données de réglage du trafic spécifiques à une tâche sont les ensembles de données de formation extraits de l'étape de prétraitement pour différentes tâches en aval.
Commencez le réglage : après les étapes susmentionnées, vous pouvez démarrer le réglage de la deuxième étape en utilisant trafficllm_stage2.sh. Il y a un exemple comme ci-dessous :
PRE_SEQ_LEN=128
LR=2e-2
NUM_GPUS=1exporter CUDA_VISIBLE_DEVICES=1
torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py
--do_train
--train_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_train.json
--validation_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_train.json
--preprocessing_num_workers 10
--prompt_colonne instruction
--response_column sortie
--overwrite_cache
--cache_dir ../cache
--model_name_or_path ../models/chatglm2/chatglm2-6b
--output_dir ../models/chatglm2/peft/ustc-tfc-2016-detection-packet
--overwrite_output_dir
--max_source_length 1024
--max_target_length 32
--per_device_train_batch_size 1
--per_device_eval_batch_size 1
--gradient_accumulation_steps 16
--predict_with_generate
--max_steps 20000
--logging_steps 10
--save_steps 4000
--learning_rate $LR
--pre_seq_len $PRE_SEQ_LEN TrafficLLM utilise EA-PEFT pour organiser les modèles de réglage fin efficaces (PEFT) avec une adaptation extensible, ce qui peut aider TrafficLLM à s'adapter facilement à de nouveaux environnements. L'adaptateur TrafficLLM permet des opérations flexibles pour mettre à jour d'anciens modèles ou enregistrer de nouvelles tâches.
model_name : Le chemin du modèle de base.
tuning_data : Le nouvel ensemble de données d'environnement.
adaptation_task : Mettre à jour ou enregistrer (mettre à jour les anciens modèles ou enregistrer de nouvelles tâches).
task_name : Le nom de la tâche en aval à mettre à jour ou à insérer.
Il existe un exemple pour mettre à jour TrafficLLM avec des tâches de détection du trafic de logiciels malveillants (MTD).
cd EA-PEFT python ea-peft.py --model_name /Your/Base/Model/Path --tuning_data /Your/New/Dataset/Path --adaptation_task update --task_name MTD
Points de contrôle : vous pouvez essayer d'évaluer TrafficLLM en utilisant votre propre modèle ou nos points de contrôle publiés.
Données : au cours de l'étape de prétraitement, nous divisons les ensembles de données de test et créons des fichiers d'étiquettes pour différents ensembles de données à des fins d'évaluation. Veuillez vous référer aux codes de prétraitement.
Pour mesurer l'efficacité de TrafficLLM pour différentes tâches en aval, veuillez exécuter les codes d'évaluation.
model_name : Le chemin du modèle de base.
traffic_task : Tâches de détection ou tâches de génération.
test_file : ensembles de données de test extraits lors des étapes de prétraitement.
label_file : Le fichier d'étiquette extrait lors des étapes de prétraitement.
ptuning_path : chemin du modèle PEFT pour l'évaluation spécifique à une tâche.
Il existe un exemple pour exécuter une évaluation sur les tâches MTD.
python évaluation.py --model_name /Votre/Base/Modèle/Chemin --traffic_task détection --test_file datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_test.json --label_file datasets/ustc-tfc-2016/ustc- tfc-2016_label.json --ptuning_path models/chatglm2/peft/ustc-tfc-2016-detection-packet/checkpoints-20000/
Vous pouvez déployer TrafficLLM sur votre appareil local. Tout d’abord, configurez le chemin du modèle dans config.json pour enregistrer les modèles PEFT collectés à partir des étapes de formation. Il existe un exemple pour enregistrer 6 tâches dans TrafficLLM :
{ "model_path": "models/chatglm2/chatglm2-6b/", "peft_path": "models/chatglm2/peft/", "peft_set": { "NLP": " instruction/checkpoint-8000/", "MTD": "ustc-tfc-2016-detection-packet/checkpoint-10000/", "BND": "iscx-botnet-2014-detection-packet/checkpoint-5000/", "WAD": "csic-2010-detection-packet/checkpoint-6000/", "AAD": "dapt-2020-detection-packet/checkpoint-20000/", "EVD": "iscx-vpn-2016-detection -packet/checkpoint-4000/", "À déterminer": "iscx-tor-2016-detection-packet/checkpoint-10000/"
}, "tasks": { "Détection du trafic malveillant": "MTD", "Détection de botnet": "BND", "Détection d'attaque Web": "WAD", "Détection d'attaque APT": "AAD", "Détection VPN crypté ": "MVE", "Détection du comportement Tor": "À déterminer"
}
} Ensuite, vous devez ajouter la pré-invite dans la fonction prepromt de inference.py et trafficllm_server.py. La pré-invite est le texte de préfixe utilisé dans les données d'entraînement lors du réglage du trafic spécifique à la tâche.
Pour discuter avec TrafficLLM en mode terminal, vous pouvez exécuter la commande suivante :
python inference.py --config=config.json --prompt="Votre texte d'instruction +: + Données de trafic"
Vous pouvez lancez la démo du site Web de TrafficLLM à l'aide de la commande suivante :
streamlit run trafficllm_server.py
Cette démo exécute un serveur Web de TrafficLLM. Accédez à http://Your-Server-IP:Port pour discuter dans la chatbox.
Un grand merci aux travaux connexes ChatGLM2 et Llama2 qui servent de fondements à notre cadre et à nos codes. La conception du bâtiment TrafficLLM est inspirée de ET-BERT et GraphGPT. Merci pour leurs merveilleux travaux.
