MalwareDatasetClassification (SBAN)
Multiclass pipeline for
malware dataset origin classification on
SBAN: four synchronized text views per sample → predict which sub-corpus it belongs to (
bodmas,
dike,
malwarebazaar,
sorel20m).
This repository contains code, notebooks, and sban_weighted_stacking_model.joblib. No SBAN parquet or raw JSON is distributed; obtain SBAN separately.
Task and labels
| |
|---|
| Input | assembly_code, binary_code, source_code, NLD for one sample |
| Output | dataset_name ∈ {bodmas, dike, malwarebazaar, sorel20m} |
| Scope | Dataset provenance classification, not generic malware detection |
Data preparation pipeline (scripts 01–11)
End-to-end flow on local SBAN exports:
01_make_a_dataframe.py — Merge JSON shards under data/M1/SBAN-MA-JUN25 into SBAN.parquet (four representations aligned by ID).02_validate_data.py — Schema, missing values, duplicates, cross-dataset ID overlap, content fingerprints (see Data quality).03_make_clean_dataframe.py — Cleaning rules → SBAN_clean.parquet.04_analyze_prompt_residue.py — Count LLM/prompt boilerplate phrases per representation (Prompt residue).05_split_dataframe.py — Stratified train / validation / test parquet files.06_make_features.py — Optional TF-IDF .npz features for alternate experiments.07–11 — Per-representation audits and source cleaning (08_clean_source_code.py uses 07_audit_source_code.py).
Notebooks:
baseline.ipynb — Early fusion / baseline stacking comparisons.svc_sban.ipynb — Production model: per-representation TF-IDF + numeric features, ID-based feature pruning, class-weight search, weighted LinearSVC bases, HistGradientBoostingClassifier meta learner, joblib export.inference.ipynb — Load exported bundle; validation/test metrics; synthetic demo row.
Canonical runtime entrypoint: inference.py (CLI + StackingPredictor).
Data quality findings
Summaries below come from running the numbered scripts on the full merged SBAN table (before train/val/test split). Reproduce with your own copy of the data.
Cross-dataset ID overlap (content match rate)
Shared IDs across corpus pairs; percentages = share of common IDs where that column’s text is byte-identical (02_validate_data.py, section 9).
| Pair | Common IDs | assembly | binary | source |
|---|
| bodmas × sorel20m | 806 | 72% | 74% | 91% |
| bodmas × malwarebazaar | 520 | 27% | 25% | 0% |
| bodmas × dike | 201 | 30% | 24% | 0% |
| dike × malwarebazaar | 82 | 23% | 22% | 4% |
| malwarebazaar × sorel20m | 99 | 41% | 47% | 0% |
| dike × sorel20m | 46 | 44% | 48% | 0% |
High overlap for bodmas × sorel20m (especially source) motivates careful splitting and explains why the classifier must use subtle cues, not only exact string identity across corpora.
Rows with four aligned representations
After merge / alignment (01_make_a_dataframe.py):
| Dataset | Rows | Matched (4 repr.) |
|---|
| bodmas | 82,032 | 757 |
| dike | 5,342 | 669 |
| malwarebazaar | 6,048 | 905 |
| sorel20m | 71,319 | 726 |
“Matched” = samples where all four representation fields are present for labeling and training.
Prompt residue analysis
04_analyze_prompt_residue.py scans fixed English phrases (e.g. “your code”, “here”, “additional”) across columns. Illustrative totals on cleaned data:
| Phrase | assembly | binary | source | NLD | Total |
|---|
| your code | 1 | 0 | 572 | 1 | 573 |
| add main function | 0 | 0 | 69 | 0 | 69 |
| code goes | 0 | 0 | 140 | 0 | 140 |
| implementation goes | 0 | 0 | 49 | 0 | 49 |
| corrected | 0 | 0 | 168 | 16 | 169 |
| here | 128 | 0 | 1,493 | 323 | 1,798 |
| no comments | 0 | 0 | 53 | 0 | 53 |
| additional | 37 | 0 | 76 | 1,103 | 1,185 |
Most residue sits in source and NLD; source cleaning scripts (07/08) target audit failures before modeling.
Model architecture
Artifact: sban_weighted_stacking_model.joblib (bundle_version: 1, trained with scikit-learn 1.6.1).
1For each r ∈ {asm, binary, source, nld}:
2 text → TF-IDF (binary: hex → byte tokens + instsep)
3 + 6 numeric stats (length, tokens, entropy, …)
4 → StandardScaler
5 → sparse hstack → column subset (selected_indices from ID pruning)
6 → LinearSVC (tuned class weights) → decision_function (4 scores)
7
8Meta:
9 hstack(all base decision scores + all scaled numeric blocks)
10 → HistGradientBoostingClassifier
11 → class probabilities
Bundle keys: representation_order, representation_columns, numeric_feature_names, label_encoder, meta_model, representations (vectorizer, scaler, indices, base model), selected_class_weight_configs, metadata.
Training details and ablations: svc_sban.ipynb.
Feature pruning (TF-IDF columns)
Implemented in svc_sban.ipynb (cells after the first per-representation LinearSVC bases):
- Importance — For each representation, mean
|coef_| over classes from final_base_models (TF-IDF tokens + six numeric stats).
- Sort ascending — Lowest-importance names are dropped first.
- Ratio sweep — Validation macro-F1 was plotted for many removal ratios (roughly 5–60% and 65–80% in the analysis figures); the exported model uses a single setting.
- Production choice —
feature_pruning_ratio = 0.65: remove the lowest 65% of the ranked feature list for TF-IDF vocabulary entries. The six numeric columns (char_count, line_count, token_count, avg_line_length, unique_token_ratio, char_entropy) are always kept and re-appended via fixed column indices after TF-IDF subsetting.
Validation macro-F1 at 65% feature removal (same notebook run):
| Representation | Macro F1 (val) | Columns after prune |
|---|
| asm | 0.7009 | 26,259 |
| binary | 0.6609 | 26,259 |
| nld | 0.4870 | 26,259 |
| source | 0.8802 | 26,257 |
These pruned column sets are stored in the joblib bundle as representations[r]["selected_indices"] (feature step only; ID pruning below may reuse the same index vector).
ID pruning (training samples)
Overlapping bodmas vs sorel20m IDs motivate dropping ambiguous training rows before refitting bases:
- Fix feature pruning at 65% and fit a temporary
LinearSVC on pruned features.
- Score each training row in
bodmas and sorel20m only: sparse TF-IDF presence (binary) dotted with pruned-model TF-IDF coefficient magnitudes → importance_score.
- Grid — For each representation, remove the lowest-scoring
id_prune_ratios fraction per class (5%, 10%, …, 70%), refit on remaining train rows, measure validation macro-F1 → id_pruning_summary in the notebook.
- Production choice —
selected_id_prune_ratios:
| Representation | ID remove ratio | Val macro F1 | Train rows kept | Removed bodmas / sorel20m |
|---|
| asm | 10% | 0.6978 | 102,538 | 5,689 / 4,941 |
| binary | 5% | 0.6625 | 107,854 | 2,844 / 2,470 |
| source | 40% | 0.8712 | 70,646 | 22,756 / 19,766 |
| nld | 40% | 0.4804 | 70,646 | 22,756 / 19,766 |
Final stacking retrains ID-pruned bases (5-fold OOF decision scores), then class-weight search and meta learner on top of that pipeline. dike and malwarebazaar rows are never removed by this step.
Split sizes used in training notebook
| Split | Rows | bodmas | dike | malwarebazaar | sorel20m |
|---|
| Train | 113,168 | 56,892 | 3,267 | 3,594 | 49,415 |
| Validation | 16,167 | 8,128 | 467 | 513 | 7,059 |
| Test | 32,334 | 16,255 | 933 | 1,027 | 14,119 |
(Test counts from inference.ipynb evaluation on exported bundle.)
Base models on validation (svc_sban.ipynb)
Single-representation LinearSVC decision scores, validation set:
| Representation | Accuracy | Macro F1 | Weighted F1 |
|---|
| asm | 0.8983 | 0.7099 | 0.8902 |
| binary | 0.8426 | 0.6607 | 0.8345 |
| source | 0.9253 | 0.8811 | 0.9251 |
| nld | 0.6621 | 0.4997 | 0.6539 |
Source is the strongest single view; nld alone is weakest but adds complementary signal in the stack.
Final exported model — validation & test
Metrics from inference.ipynb with sban_weighted_stacking_model.joblib (matches weighted meta validation in svc_sban.ipynb before export).
Validation (n = 16,167)
| Accuracy | Macro F1 | Weighted F1 |
|---|
| Overall | 0.9413 | 0.9097 | 0.9412 |
| Class | Precision | Recall | F1 | Support |
|---|
| bodmas | 0.9551 | 0.9398 | 0.9474 | 8,128 |
| dike | 0.9125 | 0.8266 | 0.8674 | 467 |
| malwarebazaar | 0.8986 | 0.8635 | 0.8807 | 513 |
| sorel20m | 0.9306 | 0.9562 | 0.9433 | 7,059 |
Test (n = 32,334)
| Accuracy | Macro F1 | Weighted F1 |
|---|
| Overall | 0.9379 | 0.9012 | 0.9378 |
| Class | Precision | Recall | F1 | Support |
|---|
| bodmas | 0.9532 | 0.9364 | 0.9447 | 16,255 |
| dike | 0.8909 | 0.8489 | 0.8694 | 933 |
| malwarebazaar | 0.8885 | 0.8150 | 0.8502 | 1,027 |
| sorel20m | 0.9272 | 0.9545 | 0.9407 | 14,119 |
Minority classes (dike, malwarebazaar) remain the hardest; weighted class tuning in svc_sban.ipynb targets that imbalance.
Inference schema
| Column | Required for predict | Notes |
|---|
assembly_code, binary_code, source_code, NLD | Yes | |
ID | No | Preserved in output |
dataset_name | No | For --evaluate / notebook metrics |
Installation
1pip install -r requirements-inference.txt # predict only
2pip install -r requirements.txt # full pipeline + notebooks
Use scikit-learn 1.6.1 when loading the joblib bundle.
Running inference
1python inference.py \
2 --model-path sban_weighted_stacking_model.joblib \
3 --input /path/to/SBAN_test.parquet \
4 --output predictions.parquet \
5 --evaluate
1from inference import load_predictor
2import pandas as pd
3
4predictor = load_predictor("sban_weighted_stacking_model.joblib")
5out = predictor.predict(pd.read_parquet("/path/to/samples.parquet"))
Notebook: inference.ipynb — Colab or local setup → demo row → validation/test cells (update parquet paths).
Reproducing the production model
- Obtain SBAN and build parquets via
01–05 (and cleaning/audit scripts as needed).
- Open
svc_sban.ipynb (Colab or local), point to SBAN_train/val/test.parquet.
- Run training cells; export
sban_weighted_stacking_model.joblib to the repo root.
- Verify with
inference.py or inference.ipynb.
Citation and security
- Cite the SBAN dataset authors; this repo does not redistribute their files.
joblib.load uses pickle — only load bundles from this project or your own exports.
