MalCL: Leveraging GAN-Based Generative Replay to Combat Catastrophic Forgetting in Malware Classification

Basic Information

• Paper ID: 2501.01110
• Title: MalCL: Leveraging GAN-Based Generative Replay to Combat Catastrophic Forgetting in Malware Classification
• Authors: Jimin Park¹, AHyun Ji¹, Minji Park¹, Mohammad Saidur Rahman², Se Eun Oh¹*
• Affiliations: ¹Ewha Womans University, ²University of Texas at El Paso
• Categories: cs.CR (Cryptography and Security), cs.AI (Artificial Intelligence)
• Publication Date: January 2, 2025 (arXiv preprint)
• Paper Link: https://arxiv.org/abs/2501.01110

Abstract

This paper proposes MalCL, a system addressing the continual learning problem in malware classification. The system employs a Generative Adversarial Network (GAN)-based generative replay approach that generates high-quality malware samples through feature matching loss and implements an innovative sample selection mechanism based on model hidden representations. In class-incremental learning scenarios on Windows and Android malware datasets, the system demonstrates significant performance improvements, achieving 55% mean accuracy on Windows malware samples, representing a 28% improvement over other generative replay-based models.

Research Background and Motivation

Core Problem

The primary challenge in malware classification is the phenomenon of catastrophic forgetting. When machine learning models undergo continuous training on new data, their performance on previously learned data significantly degrades. This is particularly severe in the malware domain because:

1. Rapid malware evolution: The AV-TEST Institute records 450,000 new malware and potentially unwanted applications (PUA) daily
2. VirusTotal processes over 1 million software submissions daily
3. Dilemma for antivirus companies: Either remove old samples (risking resurgence of legacy malware) or ignore new samples (missing emerging threats)

Threat Model

The paper defines a concrete threat scenario where attackers exploit legacy malware to evade machine learning systems updated only with new data. As the time gap between original training and attack increases, the likelihood of successful evasion increases.

Limitations of Existing Approaches

1. Traditional machine learning methods: Fail to effectively address catastrophic forgetting
2. Continual learning methods from computer vision: Perform poorly when directly applied to malware classification, sometimes underperforming the "None" baseline
3. Storage constraints: Historical data storage is limited due to privacy regulations

Core Contributions

1. Malware-domain-specific continual learning model: Proposes MalCL, achieving 55% mean accuracy across 11 continual learning tasks on 100 malware families, representing a 28% improvement over existing methods
2. Improved feature matching generative replay: Employs GAN generator combined with feature matching loss to reduce feature discrepancies between original and synthetic samples
3. Innovative replay sample selection mechanism: Develops multiple selection strategies based on intermediate layer features in the classifier to improve alignment between generated and original data
4. Strategic task set construction: Explores the strategy of assigning large categories to initial tasks, effectively mitigating catastrophic forgetting

Methodology Details

Task Definition

MalCL addresses the class-incremental learning problem:

• Input: Sequence of malware feature vectors
• Output: Malware family classification
• Constraint: Each task introduces new malware categories while maintaining recognition ability for historical categories

Model Architecture

Overall Pipeline

MalCL comprises two core components:

1. GAN module: Generates synthetic malware samples from previous tasks
2. Classifier module: Performs malware family classification

GAN Architecture Design

Generator:

• 4 1D convolutional layers + 2 fully connected layers + 3 transposed convolutional layers
• ReLU activation and batch normalization on all layers except the last
• Sigmoid activation on output layer

Discriminator:

• 2 convolutional layers + 2 fully connected layers
• Flattening logic from the second convolutional layer for feature matching
• Sigmoid layer for final output

Classifier Architecture

• 3 convolutional layers + 1 fully connected layer
• Max pooling and dropout following the first two convolutional layers
• Flattening logic from the third convolutional layer for replay sample selection
• Softmax on output layer

Technical Innovations

1. Feature Matching Loss

Traditional binary cross-entropy loss:

LG = -1/m ∑(i=1 to m) log(D(G(zi)))

Feature matching loss:

LG = 1/m ∑(i=1 to m) ||Ex~pdata[D(f)(x)] - Ez~pz[D(f)(G(z))]||

Where D(f)(·) denotes the intermediate layer output of the discriminator. This loss function focuses on richer intermediate features rather than final outputs.

2. Replay Sample Selection Strategies

Selection based on L2 distance to one-hot labels:

Sc,k = {s | argmin √((Ci(s) - yc)²), ∀s ∈ S, yc ∈ Y}

Selection based on L1 distance to logistic regression:

• By batch mean: Selects samples close to the global mean
• By class mean: Selects k samples per class closest to the class center

Sc,k = {s | argmin |Li(s) - L̄i,c(xj)|, ∀s ∈ S, ∀x ∈ Xc}

Experimental Setup

Datasets

1. EMBER Dataset:
   • 337,035 malicious Windows PE files
   • 100 malware families with >400 samples each
   • Features include file size, PE and COFF header information, DLL features, etc.
2. AZ-Class Dataset:
   • 285,582 Android malware samples
   • 100 malware families with ≥200 samples each
   • Uses Drebin features across 8 categories (hardware access, permissions, API calls, etc.)
   • Final feature dimensionality: 2,439

Evaluation Metrics

• Mean Accuracy: Average classification accuracy across all tasks
• Minimum Accuracy: Lowest accuracy among all tasks

Comparison Methods

1. Baseline methods:
   • None: Training only on new data (lower bound of catastrophic forgetting)
   • Joint: Training with all historical data (ideal upper bound)
2. Prior work:
   • GR (Generative Replay): GAN-based generative replay
   • BI-R (Brain-Inspired Replay): Enhanced generative replay using VAE

Implementation Details

• Task construction: First task with 50 classes, subsequent tasks adding 5 classes each, totaling 11 tasks
• Optimizers: Adam for GAN, SGD for classifier
• Batch size: 256
• Learning rates: Classifier 1e-3, momentum 0.9, weight decay 1e-7

Experimental Results

Main Results

Key Findings:

• MalCL improves over None baseline by 27%
• Improves over existing generative replay methods (GR, BI-R) by 28%
• Best configuration: FML + L1 distance to class mean logistic regression

Ablation Studies

Loss Function Comparison

• Feature Matching Loss vs Binary Cross-Entropy: FML outperforms BCE across all selection strategies
• L1 distance to class mean logistic regression: Best performance at 55% accuracy vs 50% for L2 distance to labels

Task Set Construction Strategy

Strategic task construction: Assigning large categories to initial tasks

• First task: 50 "giant" classes (average 5,397 samples)
• Subsequent tasks: Randomly selected small classes (average 670 samples)
• Result: Accuracy improves to 74%, approaching Joint baseline performance

Cross-Dataset Validation

AZ-Class Dataset Results:

• Superior performance compared to EMBER dataset
• Narrower accuracy range, indicating greater stability
• Reason: EMBER dataset exhibits more severe class imbalance

Related Work

Continual Learning Method Classification

1. Replay techniques:
   • Exact replay: Experience Replay (ER), iCaRL
   • Generative replay: GR, BI-R, DDGR
2. Regularization methods:
   • Elastic Weight Consolidation (EWC)
   • Synaptic Intelligence (SI)

Continual Learning in Malware Classification

• Rahman et al. first explored continual learning in malware classification
• Chen et al. combined contrastive learning and active learning
• Existing methods primarily focus on concept drift rather than catastrophic forgetting

Conclusions and Discussion

Main Conclusions

1. MalCL effectively mitigates catastrophic forgetting in malware classification
2. Feature matching loss significantly improves synthetic sample quality
3. Sample selection strategies based on classifier hidden layers are highly effective
4. Strategic task construction is crucial for performance improvement

Limitations

1. Significant gap remains with Joint baseline: MalCL best 54.5% vs Joint 88.7%
2. Global selection strategies underperform: L1 distance to batch mean strategy produces insufficient class coverage
3. Sensitivity to dataset imbalance: EMBER dataset imbalance characteristics impact performance

Future Directions

1. Improve synthetic malware generation quality
2. Develop more advanced generative models
3. Explore hybrid training methods combining generative replay and joint training advantages
4. Extend to additional malware types
5. Integrate more complex features reflecting dynamic malware threat characteristics

In-Depth Evaluation

Strengths

1. Highly targeted problem: Specifically addresses catastrophic forgetting in the malware domain
2. Methodological innovation: Combines feature matching loss with diversified sample selection strategies
3. Comprehensive experiments: Cross-platform (Windows/Android) validation with multiple comparison methods
4. High practical value: Addresses important problems in real-world security protection
5. Sufficient technical detail: Provides complete architecture design and implementation details

Limitations

1. Substantial room for performance improvement: 33% gap remains with ideal upper bound
2. Missing computational overhead analysis: Lacks detailed analysis of GAN training and generation computational costs
3. Adversarial robustness: Does not consider robustness of generated samples against adversarial attacks
4. Feature engineering dependency: Relies on predefined features, potentially limiting generalization capability
5. Long-term performance evaluation: Lacks performance assessment over longer time sequences

Impact

1. Academic contribution: First systematic application of GAN-based generative replay to malware classification
2. Practical value: Provides practical continual learning solutions for antivirus companies
3. Method generalizability: Technical framework extensible to other security detection domains
4. Open-source commitment: Promises to release code, facilitating research reproduction and development

Applicable Scenarios

1. Enterprise security protection: Malware detection systems requiring continuous updates
2. Resource-constrained environments: Scenarios where large-scale historical data storage is infeasible
3. Privacy-sensitive applications: Environments where original malware samples cannot be retained
4. Real-time detection systems: Online detection systems requiring rapid adaptation to emerging threats

References

This paper cites important works from continual learning, malware detection, and generative adversarial networks, including:

• Shin et al. (2017): Continual learning with deep generative replay
• Rahman, Coull, and Wright (2022): First exploration of continual learning in malware classification
• Anderson and Roth (2018): EMBER dataset
• Arp et al. (2014): Drebin feature extraction methodology

Overall Assessment: This paper proposes an innovative solution to catastrophic forgetting in malware classification, with adequate technical methodology and experimental validation. While performance improvement potential remains, it makes important contributions to both research and practical applications in this domain.