ViGil: Quad-Modal Joint Malware Detection Model

ViGil is a production-grade, quad-modal deep learning model designed for static malware detection and threat classification on Portable Executable (PE) files (EXE, DLL, SYS). Fusing structural code representation, visual representations, raw byte sequences, and API call semantics, the model produces highly accurate classification results with built-in epistemic uncertainty estimation.

Architecture Overview

The model is built on PyTorch and consists of four parallel modal encoding streams fused via a Deep Residual MLP head with Monte Carlo (MC) dropout:

PE File
 ├─ Code Property Graph (CPG) ──► OptimizedHGT (6L, 8H) ─────────────► [B, 768] ─┐
 ├─ Grayscale Byte Image ──────► ConvNeXt-Tiny (Pretrained) ─────────► [B, 512] ─┤
 ├─ Raw PE Bytes (1024)  ──────┐                                                 │ ──► DeepResMLP ──► Verdict &
 └─ API Import Tokens (256) ───┴► OptimizedRansomFormer (Cross-Attn) ─► [B, 256] ─┘      (MC Dropout)   Uncertainty

1. Code Property Graph (CPG) Stream — OptimizedHGT

• Input: Heterogeneous node/edge tensors extracted from the executable's disassembled control-flow/data-flow graph. Node feature input is 320-dimensional.
• Encoder: Heterogeneous Graph Transformer (HGT) with 6 layers, 8 multi-head attention heads, and jumping knowledge (JK) aggregation across all layers.
• Output: 768-dimensional graph embedding.

2. Grayscale byte-image Stream — OptimizedCNN

• Input: 224x224 RGB tensor representing the raw PE binary bytes mapped directly onto a 2D grayscale image.
• Encoder: ConvNeXt-Tiny backbone with pre-trained ImageNet weights.
• Output: 512-dimensional visual embedding.

3. Cross-Modal PE Bytes & API Imports Stream — OptimizedRansomFormer

Fuses structural byte patterns and import table features:

• PE Bytes: 1024-byte sequence processed via a 3-layer 1D CNN with attention pooling to produce a 256-dimensional representation.
• API Imports: 256 API token IDs processed via positional embeddings and a 4-layer, 8-head Pre-LN Transformer to produce a 256-dimensional representation.
• Fusion: Cross-modal attention layer where the PE Bytes query the API Imports.
• Output: 256-dimensional cross-modal embedding.

4. Fusion Head — OptimizedFusion

• Input: Concatenated representations from the three streams (768 + 512 + 256 = 1536-dim).
• Architecture: Deep Residual Multi-Layer Perceptron (MLP) with two residual blocks (1024-dim each) and high dropout rates (0.4) for Monte Carlo sampling.
• Output: 2-class logits (0: BENIGN, 1: MALWARE).

---

Downstream Usage: Predict on a Single File

1. Download the Deployment Package

Download the self-contained vigil_deploy.zip archive containing the prepackaged uir library, standalone predictor script, configuration, and model checkpoint:

# 1. Install dependencies
pip install -r requirements.txt

# 2. Run prediction on a suspicious file
python predict.py --file suspicious.exe

2. Execution CLI Arguments

The predictor script supports the following command-line interface:

python predict.py --file suspicious.exe \
                  --model joint_model.pt \
                  --samples 20 \
                  --device auto \
                  --verbose

• --file (required): Path to the executable file to analyze.
• --model (optional): Path to the joint_model.pt weights.
• --samples (optional): Number of Monte Carlo dropout iterations for uncertainty estimation (default: 20).
• --device (optional): Select target device (cpu, cuda, mps, or auto).
• --verbose (optional): Enable verbose extraction logs.
• --json (optional): Print output in structured JSON format.

3. Example Output

==================================================================
  ViGil — Quad-Modal Malware Detection
  Architecture: OptimizedHGT + ConvNeXt + AttentionByte + CLSTransformerAPI → DeepResMLP
==================================================================
  File:        suspicious.exe
  Verdict:     MALWARE
  Confidence:  94.23%
  Uncertainty: 0.000412  (epistemic variance, MC dropout)
==================================================================

---

Uncertainty Estimation via Monte Carlo Dropout

Static malware analysis suffers from out-of-distribution (OOD) vulnerability (e.g. packed/obfuscated binaries or novel zero-days). To solve this, ViGil keeps Dropout active during inference. By performing $T$ forward passes ($T=20$), it calculates:

1. Confidence: Mean class probability across all stochastic runs.
2. Epistemic Uncertainty (Variance): The variance of the predictions. High variance indicates that the model is unfamiliar with the binary's feature combination (even if the confidence happens to be high). High uncertainty calls for manual sandbox analysis.

---

license: apache-2.0