JLB-JLB/android-ransomware-gat

Android ransomware detection GATs (thesis models)

Graph attention networks over class call graphs of Android APKs, trained for the Master's thesis Ransomware Detection using Graph Neural Networks Enhanced by Large Language Models (Joscha Lasse Bisping, TU Berlin). The node features are produced by a large language model that reads each decompiled class (default variants) or by a code embedding model (comparison arm).

Layout

<variant>/<held_out_family>/model.ckpt     Lightning checkpoint (state_dict + hyper_parameters)
<variant>/<held_out_family>/bundle.json    feature layout, model/training config, provenance
<variant>/<held_out_family>/metrics.json   test metrics of that fold

The thesis evaluated with family-leave-one-out: for each of the six ransomware families a model was trained on the other five (plus benign apps) and tested on the held-out family with an equally large benign sample. Every variant therefore comes as six models; each is the best of three seeds for its fold (test macro-F1, ties: AUROC, then loss). To score apps of a family the model must not have seen, take that family's fold; for general use, the pletor fold saw the largest share of the ransomware corpus.

verdict_behavior: enlarged thesis corpus, 2,166 benign and 213 ransomware apps (10.2:1)

verdict: enlarged thesis corpus, 2,166 benign and 213 ransomware apps (10.2:1)

embedding: base thesis corpus, 502 benign and 213 ransomware apps

Test sets are the held-out family plus the same number of benign apps (n in the tables), so folds with few APKs carry little evidence. macro-F1 0.333 means the model flagged nothing (or everything) on that fold.

Intended use and limitations

These checkpoints are research artefacts that reproduce the thesis' family-leave-one-out experiments; they are meant for re-running or extending that evaluation, not for screening apps. Each model was trained without one ransomware family and scored on a balanced hold-out of that family, so none of the eighteen is a calibrated detector for real app traffic, where benign apps outnumber ransomware by orders of magnitude and the 0.5 threshold was never tuned. The verdict variants only work on node features produced by the same upstream step, Gemma 4 E4B with the prompt below; features from another model or prompt, or raw code, give meaningless outputs. The training data covers six ransomware families from 2014 to 2020 and Google Play apps of the same period, so other malware kinds and newer ransomware are outside what the models were trained to recognise.

Architecture (all variants)

GAT-large: 4 GATConv layers (8 heads x 32 = 256 channels, concat=True), each followed by BatchNorm1d(256) and ReLU; global_mean_pool over the nodes; head Linear(256,256) -> ReLU -> Dropout(0.5) -> Linear(256,2). Class-weighted cross-entropy, AdamW (lr 1e-3, weight decay 1e-4), early stopping on validation macro-F1. Output index 1 is ransomware; decision threshold 0.5.

Input

One graph per APK, the class call graph:

• —Nodes: one per outer class with bytecode in the APK (inner and anonymous classes Foo$Bar fold into Foo). No library filter: every class is a node.
• —Edges: directed, caller class -> callee class, whenever any method of the caller invokes any method of the callee (Androguard method cross-references aggregated to class level; intra-class calls dropped). Shape [2, n_edges], long.
• —Node features x, shape [n_nodes, in_dim], float32, in this exact column order:

Verdict features come from google/gemma-4-E4B-it served with vLLM (temperature 0, structured JSON output, max_model_len 8192), one call per class with the class source truncated to 16,000 characters. Comments are stripped from the JADX source first. The JSON schema has nine booleans, crypto_use, file_enumeration, screen_lock_or_overlay, c2_network, sms_abuse, device_admin, anti_analysis, potentially_malicious, potentially_ransomware; crypto_use and c2_network are not fed to the models. The prompt (v3_evidence) is:

You are an Android security analyst. Examine the decompiled class below. First record
which concrete behaviours are present, then give an overall verdict. Base every field
strictly on code that is actually shown — do not guess. potentially_ransomware should
be true only when the class encrypts or locks files/the device and/or shows a ransom
demand, and it implies potentially_malicious. Respond with JSON only.

Class:

<class source>

Embedding features: SentenceTransformer("nomic-ai/CodeRankEmbed", trust_remote_code=True) with max_seq_length = 8192 and normalize_embeddings=True on the comment-stripped class source (JADX Java/Kotlin; Androguard DAD pseudo-Java or smali when JADX has no file for the class).

Minimal code (no pipeline needed)

Requires torch>=2.8, torch-geometric>=2.7, lightning>=2.6. The checkpoint is a PyTorch Lightning checkpoint. The architecture hyperparameters it stores (in_dim, hidden, num_layers, heads, dropout) are handed to the module below by load_from_checkpoint, so nothing has to be typed in by hand.

import lightning as L
import torch, torch.nn as nn
from torch_geometric.data import Batch, Data
from torch_geometric.nn import GATConv, global_mean_pool


class RansomwareGAT(L.LightningModule):
    def __init__(self, in_dim, hidden, num_layers, heads, dropout, **training_hparams):
        super().__init__()
        self.save_hyperparameters()
        convs, norms = [], []
        for i in range(num_layers):
            convs.append(GATConv(in_dim if i == 0 else hidden, hidden // heads, heads=heads))
            norms.append(nn.BatchNorm1d(hidden))
        self.model = nn.Module()
        self.model.convs = nn.ModuleList(convs)
        self.model.norms = nn.ModuleList(norms)
        self.model.head = nn.Sequential(nn.Linear(hidden, hidden), nn.ReLU(), nn.Dropout(dropout),
                                        nn.Linear(hidden, 2))
        # Training-fold class weights of the loss; stored in the checkpoint, unused at inference.
        self.register_buffer("class_weights", torch.ones(2))

    def forward(self, batch):
        x = batch.x
        for conv, norm in zip(self.model.convs, self.model.norms):
            x = torch.relu(norm(conv(x, batch.edge_index)))
        return self.model.head(global_mean_pool(x, batch.batch))

    def predict_step(self, batch, batch_idx=0):
        return torch.softmax(self(batch), dim=1)[:, 1]  # P(ransomware) per graph


model = RansomwareGAT.load_from_checkpoint("verdict_behavior/pletor/model.ckpt", map_location="cpu")
model.eval()

# One APK: n_nodes x in_dim features in the column order above, directed class-call edges.
graph = Data(
    x=torch.tensor(node_features, dtype=torch.float32),        # [n_nodes, 7] for verdict_behavior
    edge_index=torch.tensor(edges, dtype=torch.long).t(),      # [2, n_edges], rows = (caller, callee)
)
with torch.no_grad():
    p_ransomware = model.predict_step(Batch.from_data_list([graph])).item()
print("ransomware" if p_ransomware >= 0.5 else "benign", p_ransomware)

Several APKs can be scored in one call by passing a list of Data objects to Batch.from_data_list, or with L.Trainer().predict(model, loader) where loader is a torch_geometric.loader.DataLoader.

With the thesis' pipeline repository (jlb-jlb/master-thesis-ransomware-detection-pipeline) the whole path from an APK file to this prediction (JADX decompilation, graph construction, Gemma verdicts) is python -m detection_pipeline fetch --variant verdict_behavior followed by scripts/run_predict.sh --bundle models/verdict_behavior/pletor app.apk.

Citation

If you use these models in your work, please cite the thesis they come from:

@mastersthesis{bisping2026ransomware,
  title  = {Ransomware Detection using Graph Neural Networks Enhanced by Large Language Models},
  author = {Bisping, Joscha Lasse},
  school = {Technische Universität Berlin, Chair of Machine Learning and Security},
  type   = {Master's thesis},
  year   = {2026}
}

License

The MIT license of this repository covers the checkpoints and the files published here. The pipeline repository is licensed separately, under its own terms.