FreeMOCA: непрерывное обучение без памяти для анализа вредоносного кода

[Ссылка на статью ArXiv]

Zahra Asadi*, Haeseung Jeon*, Sohyun Han, Md Mahmuduzzaman Kamol, Se Eun Oh, Mohammad Saidur Rahman†

*Авторы с равным вкладом. †Автор для корреспонденции.

[!NOTE] Это официальная реализация статьи FreeMOCA: Memory-Free Continual Learning for Malicious Code Analysis.

Конвейер FreeMOCA

FreeMOCA работает следующим образом:

Задачи поступают последовательно, и каждая вводит новые семейства вредоносных программ.
Модель для каждой задачи инициализируется из предыдущей задачи (тёплый старт) для выравнивания минимумов задач в пространстве параметров.
После обучения вычисляется интерполяция на основе связности мод (mode connectivity) между оптимумами текущей и предыдущей задач. Этот процесс может использовать адаптивную послойную интерполяцию на основе величины обновления параметров.
Интерполированная модель заменяет текущую модель и используется для следующей задачи.
Ни прошлые данные, ни признаки, ни сгенерированные образцы не сохраняются, что обеспечивает непрерывное обучение без памяти.

Весь процесс выстраивает цепочку связанных решений, лежащих на многообразии с низкими потерями, что значительно снижает катастрофическое забывание.

Наборы данных

FreeMOCA был оценён на двух крупномасштабных наборах данных о вредоносном ПО: EMBER и Androzoo. Источники наборов данных:

EMBER: https://github.com/elastic/ember
Androzoo (AZ): https://zenodo.org/records/14537891

Загрузите и разместите набор данных в следующем каталоге:

FreeMOCA/data/
├── AZ_Class
│   ├── AZ_Class_Test.npz
│   └── AZ_Class_Train.npz
│
├── AZ_Domain
│   ├── 2008_Domain_AZ_Test_Transfo...
│   ├── 2008_Domain_AZ_Train_Transfo...
│   ├── 2009_Domain_AZ_Test_Transfo...
│   └── ...
│
├── EMBER_Class
│   ├── XY_test.npz
│   └── XY_train.npz
│
└── EMBER_Domain
    ├── 2018-01
    ├── 2018-02
    ├── 2018-03
    └── ...

Запуск FreeMOCA

Этот репозиторий поддерживает два сценария непрерывного обучения:

Class-Incremental Learning (Class-IL): новые классы (семейства вредоносных программ) вводятся по мере выполнения задач.
Domain-Incremental Learning (Domain-IL): пространство меток фиксировано, но распределение данных меняется со временем (например, по месяцам/годам).

Шаг 1. Установка

Выполните следующую команду для

conda create -n freemoca python=3.9
conda activate freemoca
pip install -r requirements.txt

Шаг 2. Запуск FreeMOCA (Class-IL и Domain-IL)

# EMBER Class-IL
cd ./FreeMOCA_Class/EMBER_Class
CUDA_VISIBLE_DEVICES=0 python main.py --train_data /path/to/data --test_data /path/to/data

# AZ Class-IL
cd ./FreeMOCA_Class/AZ_Class
CUDA_VISIBLE_DEVICES=0 python main.py --train_data /path/to/data --test_data /path/to/data

# EMBER Domain-IL
cd ./FreeMOCA_Domain/EMBER_Domain
CUDA_VISIBLE_DEVICES=0 python main.py --data_root /path/to/data/directory

# AZ Domain-IL
cd ./FreeMOCA_Domain/AZ_Domain
CUDA_VISIBLE_DEVICES=0 python main.py --data_root /path/to/data/directory

Более подробные настройки гиперпараметров см. в Приложении A. Общие аргументы FreeMOCA.

Приложение A. Общие аргументы FreeMOCA

Для изменения гиперпараметров или экспериментальных настроек используйте следующие аргументы:

Чтобы изменить настройки аргументов по умолчанию, обратитесь к файлу arguments.py.

Базовые методы

Наш репозиторий поддерживает эксперименты со следующими стандартными базовыми методами:

None: Обучение только на текущей задаче (нижняя граница)
Joint: Обучение на всех данных, просмотренных до сих пор (верхняя граница, непрактично)

Вы можете запустить их следующей командой:

cd /path/to/experiment/directory
CUDA_VISIBLE_DEVICES=0 python none.py --arugment_you_want
CUDA_VISIBLE_DEVICES=0 python joint.py --arugment_you_want

а также следующие предыдущие работы:

CLeWI: Модифицировано из этого репозитория GitHub, работа «Continual Learning with Weight Interpolation».
WSC: Модифицировано из этого репозитория GitHub, работа «Forget Forgetting: Continual Learning in a World of Abundant Memory».
Generative Replay (GR): Модифицировано из этого репозитория GitHub, работа «Continual learning with deep generative replay».
EWC: Модифицировано из этого репозитория GitHub, работа «Overcoming catastrophic forgetting in neural networks».
LwF: Модифицировано из этого репозитория GitHub, работа «Learning without Forgetting».
iCaRL: Модифицировано из этого репозитория GitHub, работа «iCaRL: Incremental Classifier and Representation Learning».
TAMiL: Модифицировано из этого репозитория GitHub, работа «Task-Aware Information Routing from Common Representation Space in Lifelong Learning».
MalCL: Модифицировано из этого репозитория GitHub, работа «MalCL: Leveraging GAN-Based Generative Replay to Combat Catastrophic Forgetting in Malware Classification».

Вы можете запустить базовые методы следующей командой:

# CLeWI for EMBER-Class
cd ./baselines/CLeWI
CUDA_VISIBLE_DEVICES=6 python main.py --model="clewi" \
  --dataset="seq_ember" --n_tasks=11 \
  --lr=0.001 --buffer_size=500 --n_epochs=50 \
  --seed=42 --optim_wd=0.0 --optim_mom=0.0 \
  --batch_size=512 --sub_dataset="ember"

# CLeWI for AZ-Class
cd ./baselines/CLeWI
CUDA_VISIBLE_DEVICES=6 python main.py --model="clewi" \
  --dataset="seq_ember" --n_tasks=11 \
  --lr=0.001 --buffer_size=500 --n_epochs=50 \
  --seed=42 --optim_wd=0.0 --optim_mom=0.0 \
  --batch_size=512 --sub_dataset="az"

# WSC for EMBER-Class
cd ./baselines/WSC
MODEL_NAME=wsc_20_ember
python main.py --config=./exps/wsc_memory/$MODEL_NAME.json

# WSC for AZ-Class
cd ./baselines/WSC
MODEL_NAME=wsc_20_az
python main.py --config=./exps/wsc_memory/$MODEL_NAME.json

# GR for EMBER-Class
cd ./baselines/GR_EWC_LwF_iCaRL_EMBER
CUDA_VISIBLE_DEVICES=0 python main.py --data_set=EMBER --tasks=11 --replay=generative --metrics --logger_file gr --scenario=class

# GR for AZ-Class
cd ./baselines/GR_EWC_LwF_iCaRL_AZ
CUDA_VISIBLE_DEVICES=0 python main.py --data_set=ANDROZOO --tasks=11 --replay=generative --metrics --logger_file gr --scenario=class

# EWC for EMBER-Class
cd ./baselines/GR_EWC_LwF_iCaRL_EMBER
CUDA_VISIBLE_DEVICES=0 python main.py --data_set=EMBER --tasks=11 --ewc --lambda=50 --metrics --logger_file ewc --scenario=class

# EWC for AZ-Class
cd ./baselines/GR_EWC_LwF_iCaRL_AZ
CUDA_VISIBLE_DEVICES=0 python main.py --data_set=ANDROZOO --tasks=11 --ewc --lambda=50 --metrics --logger_file ewc --scenario=class

# LwF for EMBER-Class
cd ./baselines/GR_EWC_LwF_iCaRL_EMBER
CUDA_VISIBLE_DEVICES=0 python main.py --data_set=EMBER  --tasks=11 --replay=current --distill --metrics --logger_file lwf --scenario=class

# LwF for AZ-Class
cd ./baselines/GR_EWC_LwF_iCaRL_AZ
CUDA_VISIBLE_DEVICES=0 python main.py --data_set=AZ  --tasks=11 --replay=current --distill --metrics --logger_file lwf --scenario=class

# iCaRL for EMBER-Class
cd ./baselines/GR_EWC_LwF_iCaRL_EMBER
CUDA_VISIBLE_DEVICES=1 python main.py --data_set=EMBER --tasks=11 --icarl --budget=2000 --metrics --logger_file=icarl --scenario=class

# iCaRL for AZ-Class
cd ./baselines/GR_EWC_LwF_iCaRL_AZ
CUDA_VISIBLE_DEVICES=0 python main.py --data_set=ANDROZOO --tasks=11 --icarl --budget=2000 --metrics --logger_file=icarl --scenario=class

# TAMiL for EMBER-Class
cd ./baselines/TAMiL_EMBER
CUDA_VISIBLE_DEVICES=0 python main.py

# TAMiL for AZ-Class
cd ./baselines/TAMiL_AZ
CUDA_VISIBLE_DEVICES=0 python main.py

# MalCL for EMBER-Class
cd ./baselines/MalCL_EMBER
CUDA_VISIBLE_DEVICES=0 python main.py --train_data /path/to/data --test_data /path/to/data

# MalCL for AZ-Class
cd ./baselines/AZ_EMBER
CUDA_VISIBLE_DEVICES=0 python main.py --train_data /path/to/data --test_data /path/to/data

`--init_classes`	Количество классов в задаче 0
`--epochs`	Эпох на задачу
`--batchsize`	Размер пакета
`--lr`	Скорость обучения
`--momentum`	Импульс SGD
`--weight_decay`	Затухание весов
`--lambda_min`	Минимальный вес интерполяции
`--lambda_max`	Максимальный вес интерполяции

FreeMOCA

Популярное