[論文レビュー] An Ensemble of Pre-trained Transformer Models For Imbalanced Multiclass Malware Classification

Ferhat Demirkıran, Aykut Çayır+2 more|arXiv (Cornell University)|Dec 25, 2021

Advanced Malware Detection Techniques被引用数 4

ひとこと要約

本論文は、APIコールシーケンスを用いて、極めて不均衡なデータセットにおけるマルウェア分類を改善するために、事前学習済みBERTとCANINEトランスフォーマーモデルを組み合わせた、バギングベースのランダムトランスフォーマーフォレスト（RTF）アンサンブルモデルを提案する。RTFモデルは、注意メカニズムとアンサンブル学習を活用してマルウェア行動内の複雑なシーケンス関係を捉えることで、CatakデータセットにおいてSOTAの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のマルウェア分類への初適用を実証し、アンサンブル学習を用いてそのパフォーマンスを検証すること。

提案手法

• 本手法は、バギングベースのアンサンブルフレームワークとして、ランダムトランスフォーマーフォレスト（RTF）と呼ばれるものを利用し、訓練データのランダムサブセット上で事前学習済みBERTまたはCANINEモデルの複数インスタンスを学習する。
• 各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 や CANINE のような事前学習済みトランスフォーマーのアンサンブル、たとえばランダムトランスフォーマーフォレスト（RTF）は、ベンチマークマルウェアデータセットでSOTAパフォーマンスを達成できるか？
• RQ4 RQ.4: BERT や CANINE のような大規模事前学習モデルを用いる場合、推論速度と分類精度のトレードオフはどのようなものか？
• RQ5 RQ.5: 前処理技術は、特にCatakデータセットにおける動的APIコールシーケンスのトランスフォーマーモデルのパフォーマンスにどのように影響を与えるか？

主な発見

• RTFアンサンブルモデルは、CatakデータセットでSOTAのF1スコア0.6149を達成し、すべてのベースラインモデルおよび先行手法を上回った。
• BERTとCANINEの事前学習済みモデルは、4つのデータセットすべてにおいて、F1スコアとAUCの観点で単層トランスフォーマーやLSTMモデルを顕著に上回った。特に、まれなマルウェアファミリーにおいて顕著な優位性を示した。
• CANINEモデルは、マルウェア分類において初めて適用され、VirusShareおよびVirusSampleデータセットでSOTAの結果を達成した。これは、静的APIコールシーケンスに対して有効であることを示している。
• BERTとCANINEの学習時間は、LSTMや単層トランスフォーマーと比較して著しく長く、Oliveiraデータセットでは141.02分を要したが、モデルは事前学習済みであるため、リアルタイム応答に影響しない。
• RTFの推論時間は平均で1サンプルあたり約4.5秒であり、これはAVスキャナにとって許容可能な範囲であり、大規模事前学習モデルを用いても依然として合理的な範囲に収まった。
• RTFモデルは4つのデータセットのうち3つでSOTAパフォーマンスを達成し、不均衡なマルウェア分類における事前学習トランスフォーマーとアンサンブル学習の有効性を確認した。

Figure 2 : The outputs of pre-processing steps.