cycloevan/ember-model

EMBER2024 Malware Detection Models

4종 모델 아키텍처(DNN, TabNet, Hybrid GBDT2NN, LightGBM)를 EMBER2024 데이터셋의 8개 파일 타입 subset 전체에 대해 학습·평가하고 배포 가능한 포맷으로 변환한 모델 컬렉션입니다.

학습 환경: NVIDIA DGX Spark (GB10 Grace Blackwell, 128 GB 통합 메모리, CUDA 13)
코드: github.com/evan0416/ember-ml
데이터셋 논문: Joyce et al., KDD 2025 (arXiv:2506.05074)

---

모델 구성

디렉토리	아키텍처	배포 포맷	파라미터
dnn/	Feed-Forward DNN (PReLU + Dropout)	ONNX (INT8 Static / FP32)	13.2 M (PE) / 0.73 M (non-PE)
tabnet/	TabNet (Arik & Pfister, 2021)	ONNX FP32	~3 M
hybrid/	GBDT2NN (DeepGBM, KDD 2019)	ONNX (nn_part) + LightGBM booster	~1 M NN
lightgbm/	LightGBM (사전학습, joyce8/EMBER2024-benchmark-models)	Treelite .tl	—

Subset 목록

Subset	대상 파일 타입	입력 차원
PE	PE 바이너리 전체 (Win32 + Win64 + .NET)	2,568
Win32	Windows 32-bit PE	2,568
Win64	Windows 64-bit PE	2,568
.NET	.NET 어셈블리	2,568
APK	Android APK	696
ELF	Linux ELF	696
PDF	PDF 문서	696
all	전체 파일 타입 혼합	2,568

---

디렉토리 구조

파일명 규칙: {모델}_{subset}[_suffix].{ext}
.NET subset은 파일명에서 dotnet으로 변환됩니다.

dnn/
├── dnn_PE.onnx              # INT8 Static (배포용, PE/Win32/Win64/dotnet/all)
├── dnn_PE_fp32.onnx         # FP32 ONNX   (참조용, INT8 계열만 추가 포함)
├── dnn_PE.pt                # PyTorch 체크포인트
├── dnn_PE_metrics.json      # 평가 결과 (AUC, TPR@1%FPR)
├── dnn_PE_benchmark.json    # 크기·레이턴시
├── dnn_APK.onnx             # FP32 (non-PE — INT8 AUC 손실 과다)
├── dnn_APK.pt
└── ...

tabnet/
├── tabnet_PE.onnx           # FP32 ONNX (134 MB — sparsemax 언폴딩)
├── tabnet_PE.zip            # pytorch-tabnet 네이티브 (7 MB, 경량)
└── ...

hybrid/
├── hybrid_PE_nnpart.onnx    # GBDT2NN nn_part ONNX (5.1 MB)
├── hybrid_PE_lgbm.model     # LightGBM booster (3.6 MB)
├── hybrid_PE.pt             # PyTorch 체크포인트
└── ...

lightgbm/
├── lightgbm_PE.tl           # Treelite 직렬화 (플랫폼 독립, 재컴파일 필요)
└── ...

---

성능 결과 (EMBER2024 test set)

평가 기준: ROC-AUC, TPR @ 1% FPR (논문 §4.1), challenge set detection rate @ FPR=1% 임계값
Challenge set: 6,315 evasive malware (양성 only, Win32 3,225 / .NET 829 / Win64 814 / PDF 805 / ELF 386 / APK 256)

DNN

Subset	ROC-AUC	TPR@1%FPR	배포 포맷	크기
PE	0.9969	0.9539	INT8 Static ONNX	13.3 MB
Win32	0.9966	0.9468	INT8 Static ONNX	13.3 MB
Win64	0.9969	0.9656	INT8 Static ONNX	13.3 MB
.NET	0.9939	0.8976	INT8 Static ONNX	13.3 MB
all	0.9942	0.9103	INT8 Static ONNX	13.3 MB
APK	0.9761	0.7682	FP32 ONNX	3.9 MB
ELF	0.9840	0.8103	FP32 ONNX	3.9 MB
PDF	0.9795	0.8902	FP32 ONNX	3.9 MB

non-PE(APK/ELF/PDF)는 입력 696-dim, 파라미터 수 부족으로 INT8 AUC 손실이 크므로 FP32 유지.

TabNet

Subset	ROC-AUC	TPR@1%FPR	배포 포맷	크기
PE	0.9951	0.9212	FP32 ONNX	134 MB
Win32	0.9949	0.9322	FP32 ONNX	134 MB
Win64	0.9946	0.9323	FP32 ONNX	134 MB
.NET	0.9925	0.8702	FP32 ONNX	134 MB
all	0.9922	0.8912	FP32 ONNX	134 MB
APK	0.9741	0.7028	FP32 ONNX	13.5 MB
ELF	0.9793	0.5460	FP32 ONNX	13.5 MB
PDF	0.9810	0.8597	FP32 ONNX	13.5 MB

PE 계열 ONNX 크기 134 MB: sparsemax attention loop이 ONNX 그래프로 언폴딩되는 구조적 특성. 크기 우선이면 tabnet_model.zip(7 MB) 직접 사용 권장.

Hybrid (GBDT2NN)

Subset	ROC-AUC	TPR@1%FPR	배포 포맷	크기
PE	0.9982	0.9752	nn_part ONNX + LightGBM booster	5.1 + 3.6 MB
Win32	0.9981	0.9747	nn_part ONNX + LightGBM booster	5.1 + 3.6 MB
Win64	0.9983	0.9812	nn_part ONNX + LightGBM booster	5.1 + 3.6 MB
.NET	0.9961	0.9466	nn_part ONNX + LightGBM booster	5.1 + 3.5 MB
all	0.9970	0.9528	nn_part ONNX + LightGBM booster	5.1 + 3.6 MB
APK	0.9821	0.8003	nn_part ONNX + LightGBM booster	5.1 + 3.5 MB
ELF	0.9899	0.8827	nn_part ONNX + LightGBM booster	5.1 + 3.6 MB
PDF	0.9879	0.9283	nn_part ONNX + LightGBM booster	5.1 + 3.6 MB

LightGBM (Treelite 컴파일)

Subset	ROC-AUC	TPR@1%FPR	크기 (.tl)	크기 (원본 .model)
PE	0.9983	0.9707	5.3 MB	3.6 MB
Win32	0.9984	0.9736	5.3 MB	3.6 MB
Win64	0.9989	0.9831	5.3 MB	3.6 MB
.NET	0.9980	0.9566	5.3 MB	3.5 MB
all	0.9968	0.9440	5.3 MB	3.6 MB
APK	0.9861	0.8157	5.3 MB	3.5 MB
ELF	0.9929	0.9140	5.3 MB	3.6 MB
PDF	0.9913	0.9275	5.3 MB	3.6 MB

원본 LightGBM 모델: joyce8/EMBER2024-benchmark-models. .tl은 Treelite 3.9.1로 직렬화된 플랫폼 독립 파일 — 각 플랫폼에서 재컴파일 필요.

Challenge Set Detection Rate

Challenge set: 6,315 evasive malware (전부 양성). test set FPR=1% 임계값 적용.

Subset	DNN	TabNet	Hybrid	LightGBM
.NET	58.6%	70.0%	80.6%	79.6%
APK	27.3%	29.3%	34.4%	33.6%
ELF	11.7%	4.4%	23.8%	30.3%
PDF	41.5%	40.1%	56.9%	57.1%
PE	38.5%	36.9%	58.2%	58.8%
Win32	36.6%	45.3%	58.4%	69.9%
Win64	46.3%	44.1%	59.5%	59.7%
all	35.3%	42.3%	54.1%	48.4%

---

추론 성능 (Apple M1, darwin-arm64)

warm_batch1 레이턴시: 배치 크기=1, 캐시 웜업 후 측정. 배포 환경(x86_64 Linux)과 다를 수 있음.

레이턴시 (ms, warm batch=1)

Subset	DNN	TabNet	Hybrid	LightGBM
.NET	0.248	5.465	0.151	0.050
APK	0.035	0.846	0.145	0.031
ELF	0.039	0.505	0.160	0.036
PDF	0.036	2.230	0.172	0.048
PE	0.290	4.402	0.138	0.028
Win32	0.288	4.693	0.141	0.044
Win64	0.220	5.621	0.422	0.039
all	0.254	4.788	0.147	0.068

TabNet 레이턴시 높음: sparsemax attention이 ONNX 그래프로 언폴딩되는 구조적 특성.
Hybrid = nn_part ONNX 추론만 측정 (LightGBM leaf extraction 제외).
LightGBM 레이턴시 = 컴파일 .dylib 기준; 업로드 파일은 .tl (재컴파일 필요).

모델 파일 크기 (배포 포맷)

Subset	DNN	TabNet .onnx	TabNet .zip	Hybrid (nn+lgbm)	LightGBM .tl
PE 계열	13.3 MB (INT8)	140.2 MB	7.4 MB	5.3 + 3.8 MB	5.3 MB
non-PE	3.9 MB (FP32)	13.5 MB	3.2 MB	5.3 + 3.7 MB	5.3 MB

---

사용 방법

의존성 설치

pip install onnxruntime>=1.20 numpy
# LightGBM / Hybrid 추론 시
pip install "treelite==3.9.1" "treelite_runtime==3.9.1" lightgbm>=4.6
# TabNet 체크포인트 직접 사용 시
pip install pytorch-tabnet>=4.1

DNN 추론 (ONNX Runtime)

import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download

# PE subset — INT8 Static
model_path = hf_hub_download(
    repo_id="cycloevan/ember-model",
    filename="dnn/dnn_PE.onnx",
)
sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])

# X: np.ndarray shape (N, 2568), dtype float32
X = np.random.randn(1, 2568).astype(np.float32)
logit = sess.run(["logit"], {"features": X})[0]          # shape (N, 1)
prob  = 1 / (1 + np.exp(-logit.ravel()))                  # sigmoid → [0, 1]
print(f"malware probability: {prob[0]:.4f}")

# APK subset — FP32
model_path = hf_hub_download(
    repo_id="cycloevan/ember-model",
    filename="dnn/dnn_APK.onnx",
)
sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
X = np.random.randn(1, 696).astype(np.float32)          # non-PE: dim=696
prob = 1 / (1 + np.exp(-sess.run(["logit"], {"features": X})[0].ravel()))

TabNet 추론 (ONNX Runtime)

import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(
    repo_id="cycloevan/ember-model",
    filename="tabnet/tabnet_PE.onnx",
)
sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
X = np.random.randn(1, 2568).astype(np.float32)
# 출력: logit (sigmoid 전)
logit = sess.run(["logit"], {"features": X})[0]
prob  = 1 / (1 + np.exp(-logit.ravel()))

Hybrid 추론 (ONNX + LightGBM)

import numpy as np
import lightgbm as lgb
import onnxruntime as ort
from huggingface_hub import hf_hub_download

# 1. LightGBM booster로 leaf indices 추출
booster = lgb.Booster(model_file=hf_hub_download(
    repo_id="cycloevan/ember-model",
    filename="hybrid/hybrid_PE_lgbm.model",
))
X_raw = np.random.randn(1, 2568).astype(np.float64)
leaf_indices = booster.predict(X_raw, pred_leaf=True).astype(np.int64)  # (N, n_trees)

# 2. GBDT2NN ONNX로 최종 분류
nn_sess = ort.InferenceSession(hf_hub_download(
    repo_id="cycloevan/ember-model",
    filename="hybrid/hybrid_PE_nnpart.onnx",
), providers=["CPUExecutionProvider"])
logit = nn_sess.run(["logit"], {"leaf_indices": leaf_indices})[0]
prob  = 1 / (1 + np.exp(-logit.ravel()))
print(f"malware probability: {prob[0]:.4f}")

LightGBM 추론 (Treelite 컴파일 — 빠른 추론)

# 1. Treelite .tl → 플랫폼별 공유 라이브러리 컴파일 (최초 1회)
import treelite, treelite_runtime, sys, numpy as np
from huggingface_hub import hf_hub_download

tl_path = hf_hub_download(
    repo_id="cycloevan/ember-model",
    filename="lightgbm/lightgbm_PE.tl",
)
tl_model = treelite.Model.deserialize(tl_path)
lib_ext   = ".dylib" if sys.platform == "darwin" else ".so"
lib_path  = tl_path.replace(".tl", lib_ext)
tl_model.export_lib(
    toolchain="clang" if sys.platform == "darwin" else "gcc",
    libpath=lib_path,
    verbose=False,
)

# 2. 추론
predictor = treelite_runtime.Predictor(lib_path, verbose=False)
X = np.random.randn(1, 2568).astype(np.float32)
prob = predictor.predict(treelite_runtime.DMatrix(X))
print(f"malware probability: {prob[0]:.4f}")

주의: treelite==3.9.1 + treelite_runtime==3.9.1 필요. 4.x는 export_lib() 미지원.

---

학습 및 평가 환경

항목	내용
데이터셋	EMBER2024 — train 52주(2.6 M), test 12주(606 K), challenge 6,315
Feature 차원	PE 2,568 (v3) / non-PE 696 (유효 prefix)
Split 정책	시간적 순서 고정 (temporal split), 임의 셔플 없음
학습 환경	DGX Spark (GB10 Grace Blackwell, 128 GB, CUDA 13)
프레임워크	PyTorch 2.11.0, pytorch-tabnet 4.1, LightGBM 4.6
재현 시드	42
DNN 아키텍처	Linear(2568→2568→1024→512→1) + PReLU + Dropout(0.5)
Hybrid	LightGBM leaf extraction → Linear(n_trees→512→256→1) + PReLU
평가 지표	ROC-AUC, PR-AUC, TPR @ 1% FPR (논문 §4.1)

---

알려진 한계

• TabNet ONNX 크기: sparsemax attention loop 언폴딩으로 PE 계열 ONNX가 134 MB로 팽창. 원본 tabnet_model.zip(7 MB)이 경량.
• Treelite .dylib: Mac ARM64 전용 사전 컴파일 파일. 다른 플랫폼은 .tl에서 재컴파일 필요.
• DNN non-PE INT8: 696-dim 모델은 양자화 AUC 손실이 크므로 FP32 유지.
• Hybrid 추론: 단일 ONNX 파일이 아님 — LightGBM leaf extraction + nn_part ONNX 2단계.
• challenge detection rate: test set에서 FPR=1% 임계값으로 측정. subset별 분포 차이로 값이 상이할 수 있음.

---

인용

@inproceedings{joyce2025ember2024,
  title     = {EMBER2024: An Open Dataset for Training Behavioral Malware Detection Models},
  author    = {Joyce, Ruby and Rudd, Ethan M. and others},
  booktitle = {Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining},
  year      = {2025},
  url       = {https://arxiv.org/abs/2506.05074}
}

---

라이선스

코드 및 모델 가중치: Apache 2.0
LightGBM 원본 모델(hybrid/hybrid_*_lgbm.model): joyce8/EMBER2024-benchmark-models 라이선스 준수