Stage 2 GNN — PyPI Supply-Chain Malware Detection
A GraphSAGE classifier over AST-derived API call graphs, used as the Stage 2 component of a two-stage PyPI supply-chain malware detection pipeline. Stage 1 narrows 800K+ packages to a candidate set via metadata anomaly detection (IsolationForest over graph/metadata features); Stage 2 — this model — scores the candidates' source code.
Model
| Architecture | 2-layer GraphSAGE + 2-layer MLP classifier |
| Embedding | 32-dim, initialized from Word2Vec over API-call sequences |
| Hidden dim | 64 |
| Dropout | 0.3 |
| Vocabulary | 21,085 API-call tokens |
| Parameters | ~840K |
| Checkpoint size | 3.4 MB |
Input is a directed graph whose nodes are API calls extracted from a package's
Python AST and whose edges encode call-sequence and containment relations.
Dynamic/unresolvable receivers are normalized to a <dynamic> prefix
(e.g. <dynamic>.decode.strip), which keeps the vocabulary bounded across
arbitrary third-party code.
Training
Trained on the MalwareBench PyPI corpus, restricted to setup.py and
__init__.py — the install-time entry points where supply-chain payloads
overwhelmingly live.
- 3,662 malicious / 3,662 benign packages (balanced)
- Max 6 files per package
- 60/20/20 train/val/test split, seed 42
- 40 epochs max, early stopping patience 8, Adam @ lr 1e-3, batch size 16
- Benign set includes mined hard negatives (benign packages with install-time side effects, subprocess use, and network calls)
Evaluation
Held-out test split: 1,573 packages (767 malicious / 806 benign).
- AUROC 0.983
- PR-AUC 0.986
| Threshold | Precision | Recall | F1 | Accuracy |
|---|---|---|---|---|
| 0.30 | 0.975 | 0.932 | 0.953 | 0.955 |
| 0.50 | 0.985 | 0.931 | 0.957 | 0.959 |
| 0.70 | 0.993 | 0.930 | 0.960 | 0.962 |
| 0.85 | 0.993 | 0.922 | 0.956 | 0.959 |
Best F1 is 0.960 at threshold 0.740.
Read these numbers carefully
The test split is balanced 50/50, and real PyPI is not — malicious packages are a very small fraction of uploads. Precision on a balanced split is not precision in deployment. At a realistic base rate the same score distribution yields substantially more false positives per true positive, so the operating threshold has to be re-tuned against the deployed candidate distribution rather than copied from the table above.
The model is also trained and evaluated on a single corpus (MalwareBench), so these figures reflect that corpus's notion of malicious. Generalization to novel campaigns, obfuscation styles, or non-install-time payloads is not measured here.
This checkpoint is one component of a larger pipeline and is not a standalone verdict source. Treat its output as a ranking signal for human review, not as an automated takedown trigger.
Usage
import torch
ckpt = torch.load("stage2_gnn_setup_init_v2.pt", map_location="cpu")
ckpt["config"] # {'vocab_size': 21085, 'embed_dim': 32, 'hidden': 64, 'dropout': 0.3}
ckpt["vocab"] # list[str], index -> API-call token ('<unk>' at 0)
ckpt["state_dict"] # GraphSAGE weights
ckpt["extra"] # test_report, training_history, training args
Instantiate the matching model definition from the training repository, then
model.load_state_dict(ckpt["state_dict"]). Graphs must be built with the same
AST extractor and token normalization used in training — the vocabulary is
extractor-specific and will not transfer to a differently-built graph.
License
Not yet specified. No reuse rights are granted at this time — please contact the author before redistributing or building on these weights.
Intended use
Security research and defensive triage of the PyPI ecosystem: prioritizing which packages a human analyst looks at. Not intended for automated enforcement, and not evaluated for use outside PyPI Python packages.