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は以下のプロセスで動作します。

タスクが順次到着し、それぞれが新しいマルウェアファミリーを導入します。
各タスクのモデルは、パラメータ空間でタスクの最小値を揃えるために、前のタスクから初期化されます（ウォームスタート）。
トレーニング後、現在のタスクと前のタスクの最適解の間にモード接続に基づく補間が計算されます。このプロセスでは、パラメータ更新の大きさに基づく適応的な層別補間を使用できます。
補間されたモデルが現在のモデルを置き換え、次のタスクに使用されます。
過去のデータ、特徴量、生成サンプルは一切保存されず、メモリフリーの継続学習を実現します。

このプロセス全体により、低損失多様体上に存在する接続された解の連鎖が構築され、破壊的忘却が大幅に軽減されます。

データセット

FreeMOCAは、EMBERとAndrozooという2つの大規模マルウェアデータセットで評価されました。データセットの入手先:

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の実行

このリポジトリは、次の2つの継続学習シナリオをサポートしています。

クラス増分学習（Class-IL）: タスクを重ねるごとに新しいクラス（マルウェアファミリー）が導入されます。
ドメイン増分学習（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

人気