Meta AI Security Framework

A multi-class malware family classification and analysis framework built on Windows API emulation traces. Classifies executables into 9 behavioral categories using a hybrid generative-discriminative topic model (GSDMM + discriminative head), with integrated latent backdoor detection and malware evolution tracking.

Overview

The framework takes Speakeasy Windows emulation reports (JSON) as input and performs:

1. Multi-class malware classification — 9 families: backdoor, clean, coinminer, dropper, keylogger, ransomware, rat, trojan, windows_syswow64
2. Latent backdoor detection — flags training samples showing statistical anomalies (cluster entropy, intra-class divergence)
3. Malware evolution tracking — measures behavioral drift between training and test distributions using Jensen-Shannon divergence per family

Architecture

Speakeasy JSON reports
        |
        v
  data_loader.py          Extract API call sequences (api_name fields)
        |
        v
  arg_features.py         Build argument-flag tokens (VirtualAlloc RWX, SetWindowsHookEx hook type, etc.)
        |
        v
  features.py             Construct unigram + bigram TF-IDF feature matrix
        |
  gsdmm.py                K=18 GSDMM cluster posteriors (generative layer)
        |
        v
  discriminative_head.py  Logistic Regression on [posteriors ⊕ TF-IDF(uni+bi+args)]
        |
        v
  Prediction: one of 9 malware family labels

Best system: Hybrid (K=18 unigram GSDMM posteriors concatenated with TF-IDF over unigrams + top-2000 bigrams + argument-flag tokens) + Logistic Regression (class_weight='balanced')

Performance (17,402-sample test set):

• Accuracy: 74.9% (up from 54.6% with naive Bayes decoder)
• Macro F1: 0.673 (doubled from 0.330)

Dataset

Speakeasy Windows emulation reports in archive/windows_emulation_trainset/ and archive/windows_emulation_testset/. Each subfolder is a malware family label. Each JSON file is one emulation run containing arrays of API calls with api_name, args, and ret_val fields.

Train set: ~75,298 samples | Test set: ~17,402 samples

Setup

pip install -r requirements.txt

Requires Python 3.8+. Key dependencies: scikit-learn, lightgbm, scipy, numpy.

Running the Framework

Step 1 — Train the base GSDMM model (first run only; builds API-name data caches)

python train_speakeasy_gsdmm.py -c configs/speakeasy_gsdmm_config.json

Step 2 — Build argument-flag cache (one-time archive re-parse; requires archive/ access)

python build_arg_cache.py

Step 3 — Run the full ablation / best system

python train_args_ablation.py

Outputs the final 74.9% / 0.673 system and a complete ablation table to saved_files/speakeasy_gsdmm/full_ablation_args_results.json.

Other entry points

# Posteriors-only ablation (Step 1.1 / 1.4)
python train_discriminative_head.py

# Unigram + bigram ablation (Step 1.2)
python train_full_ablation.py --bigrams 2000

# Step 1.2b — bigram-retrained GSDMM (useful negative result)
python train_gsdmm_bigram.py

Project Structure

speakeasy_gsdmm/
  __init__.py
  constants.py            9-class label map, default hyperparameters
  data_loader.py          Speakeasy JSON parser, dataset loader, pickle caching
  gsdmm.py                GSDMM (Gibbs Sampling Dirichlet Multinomial Mixture)
  discriminative_head.py  TF-IDF featurization + multi-classifier evaluation
  features.py             Bigram API pair extraction
  arg_features.py         Argument-flag token extractor
  backdoor_detector.py    Three-signal latent backdoor detector
  evolution_tracker.py    Per-family JS divergence + API shift analysis
  meta_framework.py       Pipeline orchestration (load → train → evaluate → detect → track)

configs/
  speakeasy_gsdmm_config.json   K=18, alpha=0.1, beta=0.1, n_iters=30

saved_files/speakeasy_gsdmm/
  gsdmm_model.pkl               Trained GSDMM (K=18 unigrams)
  framework_results.json        Base GSDMM evaluation + backdoor + evolution results
  full_ablation_args_results.json   Final Phase-1 best-system ablation table
  _cache_{train,test}_*.pkl     Parsed data caches (skip 12-min file I/O on reruns)
  _feat_{train,test}_*.npy      Cached GSDMM posterior matrices

Key Findings

• 18-posterior bottleneck: GSDMM cluster posteriors alone cap at ~55% accuracy regardless of classifier. Accuracy is feature-limited, not classifier-limited.
• Sequence matters: Bigram API pairs add +3-4 points everywhere; they revive the hybrid over plain BOW.
• Hybrid beats BOW once sequence is present: With bigrams, the GSDMM posteriors add complementary signal (not present with unigrams alone). Keylogger F1: 0.28 (BOW) → 0.80 (hybrid) via the SetWindowsHookEx cluster.
• Argument flags help where they matter: Dropper F1 doubled (0.195 → 0.398) from file-write/create argument tokens.
• Orthogonality principle: Bigram-retrained GSDMM is marginally worse in the hybrid — bigram clusters duplicate the bigram TF-IDF block. Unigram clusters stay orthogonal and add more.

See PROJECT_EXPLANATION.txt for the full technical write-up, all empirical results, and the Phase 2 research roadmap.