Detecting Android Ransomware with Graph Neural Networks: A Baseline Study

Introduction

This post describes the baseline experiment for my master's thesis on Android malware detection. The idea: represent each APK as a function call graph (FCG) and train Graph Neural Networks to classify it as benign or ransomware. The best model (GIN, 3 layers) reaches 98.15% accuracy and 100% malware recall.

The following sections cover the dataset, preprocessing, feature engineering, model architectures, and results across GIN, GCN, and GAT.

The Dataset

The dataset consists of 715 Android APKs: 502 benign applications and 213 ransomware samples spanning 6 families.

The ransomware families and their sample counts:

This gives us a ~2.4:1 class imbalance, which needs to be accounted for during training.

From APKs to Graphs

The Idea

Android apps are compiled to Dalvik bytecode. Using Androguard↗, we statically analyse each APK and extract a function call graph (FCG), i.e. a directed graph where:

• Nodes = methods (functions) defined in the app
• Edges = caller → callee relationships

Malware tends to exhibit structurally distinct calling patterns (encryption routines, device-locking sequences, C2 communication). These create graph topologies that differ from benign apps, and a GNN can learn to recognise them.

Two Preprocessing Modes

The pipeline supports two graph construction modes:

Internal-only (default): Only author-written methods become nodes. External API calls (e.g. android.widget.Toast.makeText) are not added as nodes; their signal is encoded as numerical features on the calling node. This removes ~12-14% of nodes that would otherwise be uninformative leaf stubs shared across all apps.

Full FCG: All methods (internal + external) become nodes. External methods get binary flags (is_external, is_android_api) as features.

Node Features

For the internal-only mode, each node carries a 5-dimensional feature vector:

The central design choice is encoding external API usage as node-level features rather than expanding the graph. This captures which internal methods call which APIs (the discriminative signal) while keeping graphs compact. The 11 platform API prefixes used include Landroid/, Ljava/, Ldalvik/, Ljavax/, Landroidx/, Lkotlin/, and several others.

Log-scaling prevents high-degree hub methods from dominating the feature space.

Graph Statistics

The resulting graphs vary enormously in size:

This is orders of magnitude larger than typical GNN benchmarks (MUTAG ≈ 18 nodes, ENZYMES ≈ 33 nodes) and poses significant computational challenges.

Output Format

Each APK produces a PyTorch Geometric Data object:

Data(
    x          = [num_nodes, 5],       # node features
    edge_index = [2, num_edges],       # directed adjacency (COO)
    y          = [1],                  # 0=benign, 1=malware
    apk_name   = str,                  # metadata
    family     = str,                  # malware family or "benign"
)

Model Architectures

All three architectures follow the same structural pattern:

Node features (dim=5)
  → 3 GNN layers (each with BatchNorm + ReLU)
  → Global mean pooling
  → 2-layer MLP classifier with dropout
  → 2-class logits

GIN (Graph Isomorphism Network)

GIN is the most expressive standard message-passing GNN, equivalent in power to the 1-WL graph isomorphism test. Update rule:

$h_v^{(k)} = \text{MLP}^{(k)}\!\left((1 + \epsilon) \cdot h_v^{(k-1)} + \sum_{u \in \mathcal{N}(v)} h_u^{(k-1)}\right)$

Each layer uses a 2-layer MLP (Linear → ReLU → Linear) with a learnable $\epsilon$. 100,866 parameters.

GCN (Graph Convolutional Network)

The simplest spectral-inspired GNN. Uses symmetric normalised aggregation, acting as a low-pass filter on the graph spectrum. Serves as a sanity check: if GCN already works well, the task may not require sophisticated message passing. 51,330 parameters.

GAT (Graph Attention Network)

Learns attention coefficients over edges to weight neighbour contributions differently. Uses 4 attention heads, each producing 32-dimensional features (concatenated to 128). More expensive per step, but can focus on the most informative neighbours. 52,098 parameters.

Why Global Mean Pooling?

All models use mean pooling as the graph readout:

$h_G = \frac{1}{|V|} \sum_{v \in V} h_v^{(L)}$

Sum pooling was avoided because graph sizes span from 15 to 400k+ nodes. The representation magnitude would correlate with graph size rather than structure.

The animation below shows how global mean pooling collapses all node embeddings into a single graph-level vector, which then passes through the MLP classifier to produce the final prediction.

Training Setup

Data Splitting

With only 715 samples and class imbalance, the split strategy matters. We use stratified splitting (70/15/15) to preserve the ~70/30 benign/malware ratio in every split:

All experiments share the exact same split (random_state=42), enabling direct head-to-head comparison.

Handling Class Imbalance

A naive model could reach 70% accuracy by always predicting "benign". We use weighted cross-entropy with weights inversely proportional to class frequency:

• Benign weight: 0.597
• Malware weight: 1.403

Misclassifying a malware sample therefore costs ~2.35× more than misclassifying a benign one, which pushes the model towards higher recall.

Hyperparameters

All 6 experiments (3 architectures × 2 dataset modes) share identical settings:

Why batch size 4: GPU memory constraints. A single benign graph can have 170k+ nodes. With batch size 4, PyG's block-diagonal batching produces super-graphs with 500k+ nodes, which is already at the memory limit. My laptop GPU only has 8 GB of VRAM and 18 GB of shared system RAM.

Why min 100 epochs for early stopping? GNN training on variable-sized graphs produces high per-epoch validation loss variance. A single batch dominated by a 200k-node graph can spike the loss. The minimum epoch floor prevents premature termination from such transient spikes.

Results

Summary

Toggle between metrics to compare all 6 experiments:

Key Findings

1. GIN is the best baseline architecture. GIN on internal_only achieves the top accuracy (98.15%) with perfect malware recall (zero missed ransomware) in 14.7 minutes of training. GAT on full_fcg matches the accuracy but takes 4× longer.

2. Internal-only preprocessing works better (or equally well). For GIN and GCN, removing external nodes and encoding their signal as features consistently improves performance. The exception is GAT, which benefits from full_fcg, likely because attention can learn to downweight uninformative external leaf nodes on its own.

3. All models exceed 94% accuracy. FCG topology combined with these lightweight features is highly discriminative for binary malware classification. Even GCN, the simplest model, performs well.

4. GIN achieves 100% malware recall on both dataset modes. All ransomware samples in the test set are correctly identified, with only 2 false positives (benign apps flagged as malware). In a security context, this is the preferable trade-off: missed malware is far more costly than a flagged benign app.

Confusion Matrices

GIN internal_only (2 false positives, 0 false negatives):

              Predicted
              Benign  Malware
Actual Benign    74       2
     Malware      0      32

GAT full_fcg (balanced errors: 1 FP, 1 FN):

              Predicted
              Benign  Malware
Actual Benign    75       1
     Malware      1      31

Per-Family Recall

A model with high overall recall could still systematically miss certain ransomware families. The per-family breakdown:

simplelocker is the hardest family. It is the only one where any model fails. GCN misses 2/9 samples, GAT misses 1/9, while GIN catches all 9. This family has enormous internal size variance (65 to 398,727 nodes). GIN's stronger theoretical discriminative power (1-WL equivalence) likely helps it handle this diversity.

Training Dynamics

All models converge relatively quickly due to strong class separability:

The following chart animates the GIN (internal_only) training curves. LR reductions are shown as dashed lines, and the early stopping point is marked in orange.

GAT is 2-4× slower per epoch due to per-edge attention computation. On full_fcg graphs (larger, more edges), some epochs take 60-240 seconds due to memory pressure.

All models showed modest overfitting (1-2 pp gap between train and val accuracy), indicating that the regularisation stack (dropout, weight decay, early stopping, batch normalisation) is well-calibrated for this dataset size.

Limitations and Next Steps

The baseline establishes a solid performance floor, but several limitations remain:

• Single split: Results depend on one specific train/val/test partition. Cross-validation would provide confidence intervals.
• No hyperparameter search: The chosen settings are reasonable defaults, not optimised.
• Small test set: With only 32 malware test samples, each misclassification changes recall by ~3 percentage points.
• Binary classification only: The model detects malware but does not identify the ransomware family.
• No data augmentation: Techniques like random edge dropout or subgraph sampling could improve generalisation.

Planned next steps include cross-validation, hyperparameter tuning, multi-class family prediction, and more expressive graph-level features.

Conclusion

Function call graphs provide a useful structural representation for Android malware detection. With 5 hand-crafted node features, a 3-layer GIN model achieves near-perfect malware recall on this dataset. The internal-only preprocessing mode, where external API usage is compressed into node features rather than expanding the graph, proves to be both more efficient and more effective than including all methods as nodes.

The full pipeline (APK to prediction) trains in under 15 minutes on a single GPU, making it practical for rapid iteration.

Resources

Model on Hugging Face↗

Dataset on Hugging Face↗