[논문 리뷰] An Ensemble of Pre-trained Transformer Models For Imbalanced Multiclass Malware Classification

Ferhat Demirkıran, Aykut Çayır외 2명|arXiv (Cornell University)|2021. 12. 25.

Advanced Malware Detection Techniques인용 수 4

한 줄 요약

이 논문은 고도로 불균형한 데이터셋에서 API 호출 시퀀스를 사용하여 다중 클래스 악성코드 분류를 향상시키기 위해 사전 훈련된 BERT 및 CANINE 트랜스포머 아키텍처를 결합한 백킹 기반 Random Transformer Forest (RTF) 앙상블 모델을 제안한다. RTF 모델은 주목사용 메커니즘과 앙상블 학습을 통해 악성코드 행동의 복잡한 시퀀스 관계를 포착함으로써 Catak 데이터셋에서 최신 기술 수준의 F1 스코어 0.6149를 달성한다.

ABSTRACT

Classification of malware families is crucial for a comprehensive understanding of how they can infect devices, computers, or systems. Thus, malware identification enables security researchers and incident responders to take precautions against malware and accelerate mitigation. API call sequences made by malware are widely utilized features by machine and deep learning models for malware classification as these sequences represent the behavior of malware. However, traditional machine and deep learning models remain incapable of capturing sequence relationships between API calls. On the other hand, the transformer-based models process sequences as a whole and learn relationships between API calls due to multi-head attention mechanisms and positional embeddings. Our experiments demonstrate that the transformer model with one transformer block layer surpassed the widely used base architecture, LSTM. Moreover, BERT or CANINE, pre-trained transformer models, outperformed in classifying highly imbalanced malware families according to evaluation metrics, F1-score, and AUC score. Furthermore, the proposed bagging-based random transformer forest (RTF), an ensemble of BERT or CANINE, has reached the state-of-the-art evaluation scores on three out of four datasets, particularly state-of-the-art F1-score of 0.6149 on one of the commonly used benchmark dataset.

연구 동기 및 목표

• API 호출 시퀀스에서 유도된 행동 특징을 사용하여 고도로 불균형한 다중 클래스 악성코드 분류 문제를 해결하기 위해.
• BERT 및 CANINE와 같은 사전 훈련된 트랜스포머 모델이 악성코드 패밀리 식별을 위해 API 호출 시퀀스 내 장거리 종속성을 어떻게 포착하는지 평가하기 위해.
• 불균형 데이터셋에서 희귀 악성코드 패밀리에 대한 일반화 및 성능을 향상시키기 위한 앙상블 모델을 개발하기 위해.
• 전통적인 LSTM 및 단일 레이어 트랜스포머 모델과 비교하여 트랜스포머 기반 모델의 추론 및 훈련 효율성을 평가하기 위해.
• CANINE의 악성코드 분류 분야에서의 첫 번째 응용을 보여주고, 앙상블 학습을 통해 성능을 검증하기 위해.

제안 방법

• 이 방법은 랜덤 서브셋의 훈련 데이터에 대해 훈련되는 사전 훈련된 BERT 또는 CANINE 모델의 여러 인스턴스를 사용하는 백킹 기반 앙상블 프레임워크인 Random Transformer Forest (RTF)를 활용한다.
• 각 RTF 모델은 순서를 유지하는 토큰화 없는 입력 표현 방식을 사용하여, 시퀀스 간의 종속성을 주목사용 메커니즘으로 모델링할 수 있도록 한다.
• 예측을 다수결 투표 방식으로 통합하여 정규화된 성능과 분산 감소를 달성하며, 특히 희귀 악성코드 패밀리에서 유의미하다.
• Catak 데이터셋의 동적 API 호출을 고려하여 성능 향상을 위해 시퀀스 길이 정규화 및 데이터 증강 등의 사전 처리 단계를 적용한다.
• 클래스 불균형 문제를 고려하여 평가에 F1 스코어와 AUC를 주요 지표로 사용하며, 효율성 분석을 위해 훈련 및 추론 시간을 측정한다.
• BERT 및 CANINE에 내장된 다중 헤드 주목사용과 위치 임베딩을 활용하여, RNN이나 단일 레이어 트랜스포머보다 API 호출 시퀀스 내 장거리 관계를 더 효과적으로 모델링한다.

Figure 1 : Transformer model architecture.

실험 결과

연구 질문

• RQ1 RQ.1: 불균형한 다중 클래스 악성코드 분류에 적합한 분류 지표는 무엇이며, F1 스코어와 AUC는 모델 성능 평가에서 어떻게 비교될 수 있는가?
• RQ2 RQ.2: BERT 및 CANINE와 같은 사전 훈련된 트랜스포머 모델은 기존의 모델들(예: LSTM) 및 단일 레이어 트랜스포머와 비교해 불균형한 악성코드 데이터셋에서 어떻게 성능을 내는가?
• RQ3 RQ.3: BERT와 같은 사전 훈련된 트랜스포머의 앙상블, 예를 들어 Random Transformer Forest (RTF)는 벤치마크 악성코드 데이터셋에서 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ4 RQ.4: BERT 및 CANINE와 같은 대규모 사전 훈련된 모델을 사용할 경우, 추론 속도와 분류 정확도 사이의 상충 관계는 어떻게 나타나는가?
• RQ5 RQ.5: 사전 처리 기법은 특히 Catak 데이터셋에서 동적 API 호출 시퀀스에 대해 트랜스포머 모델의 성능에 어떤 영향을 미치는가?

주요 결과

• RTF 앙상블 모델은 Catak 데이터셋에서 최신 기술 수준의 F1 스코어 0.6149를 달성하여 모든 베이스라인 모델과 이전 방법들을 능가한다.
• BERT 및 CANINE 사전 훈련된 모델은 모든 네 개의 데이터셋에서 F1 스코어와 AUC 측면에서 단일 레이어 트랜스포머 및 LSTM 모델보다 뚜렷이 뛰어나며, 특히 희귀 악성코드 패밀리에서 유의미한 성능 향상을 보인다.
• CANINE 모델은 악성코드 분류 분야에서 처음으로 적용되었으며, VirusShare 및 VirusSample 데이터셋에서 최신 기술 수준의 결과를 달성하여 정적 API 호출 시퀀스에 대한 효과성을 입증한다.
• BERT 및 CANINE의 훈련 시간은 LSTM 및 단일 레이어 트랜스포머보다 상당히 길며, Oliveira 데이터셋에서 BERT는 141.02분이 소요되지만, 이는 사전 훈련된 모델이므로 실시간 응답에 영향을 주지 않는다.
• RTF의 평균 추론 시간은 약 4.5초/샘플로, 악성코드 스캐너에 있어 수용 가능한 수준이며, 대규모 사전 훈련된 모델을 사용함에도 불구하고 합리적인 범위를 유지한다.
• RTF 모델은 네 개 데이터셋 중 세 개에서 최신 기술 수준의 성능을 달성하여, 사전 훈련된 트랜스포머를 사용한 앙상블 학습이 불균형한 악성코드 분류에 효과적임을 확인한다.

Figure 2 : The outputs of pre-processing steps.