Evaluating Out-of-Distribution Robustness in
Graph-Based Android Malware Classification

A New Principled Benchmark

Ngoc N. Tran1, Anwar Said2, Waseem Abbas3, Tyler Derr1, Xenofon D. Koutsoukos1
1Vanderbilt University 2Highmark Health 3The University of Texas at Dallas
arXiv Code Data

Abstract

Graph malware classifiers achieve high accuracy on standard benchmarks but suffer significantly under distribution shift when new malware variants emerge. Our research highlights that existing structural features fail to capture the deeper semantic patterns necessary for robust generalization.

We introduce two new benchmarks, MalNet-Tiny-Common and MalNet-Tiny-Distinct, designed to evaluate performance under realistic covariate and domain shifts. We propose a semantic enrichment framework that augments Function Call Graphs (FCGs) with function-level metadata and code embeddings derived from Large Language Models (LLMs).

Evaluations demonstrated that our data-centric methodology outperforms model-based approaches under distribution shift, and consistently further enhances robustness when used in conjunction.

Overview

Below is an overview of the data construction pipeline of our work. Please consult our README.md files for detailed information.

Architecture Diagram

Dataset

Curation

The dataset curation process involved several stages to ensure high-quality, semantically enriched benchmarks:

  1. Sample Selection: We selected malware samples from the MalNet database, focusing on families that represent shifts in distribution with respect to the original MalNet-Tiny benchmark:
    • MalNet-Tiny-Common: Malwares of the same families but different malware types, simulating covariance shift.
    • MalNet-Tiny-Distinct: Malwares of completely different families, simulating a more drastic domain shift.
  2. APK Acquisition: Raw binaries (.apk files) were retrieved from AndroZoo using SHA256 hashes provided by MalNet.
  3. Graph Extraction: For each APK, we extracted the FCGs using static analysis tools, identifying all function nodes and their call relationships.
  4. Semantic Enrichment: Each node in the FCG was augmented with:
    • Function Metadata: Including information such as function names, access modifiers, return types, and parameters.
    • LLM Embeddings: High-dimensional vector representations of the function's bytecode/source code using state-of-the-art LLMs (e.g., CodeXEmbed, Qwen 3).

Download

We release the pre-processed attributed graphs, including all semantic embeddings, as downloadable files. These files allow researchers to train and evaluate models without performing the expensive LLM inference step themselves.

To improve usability, we have reworked the precomputed data structure such that there are no duplicate weights across the files, broken down large files into smaller ones for easy access, and improved code to not allocate unnecessary memory during the data loading process. Additionally, this new release allows users to download only the specific files needed for any requested dataset.

The datasets and loading code are available at the following Hugging Face link.

Dataset Construction

We provide our code to construct our dataset from any APKs, ensuring compliance with AndroZoo's redistribution policies.

1. Setup

Install the necessary dependencies:

pip install -r requirements.txt

2. LLM Embedding Server

Start the inference server to handle code embedding requests (supports various backends):

# Example: Start server on port 8080
python llm_inference_server_cxe.py --port 8080

3. Graph Construction

Run the construction script to process APKs into attributed graphs:

python create_graph.py --apk_dir ./path/to/apks --n_jobs 8 --port 8080

The splits directory contains the definitions for MalNet-Tiny-Common and MalNet-Tiny-Distinct.

Training & Evaluation

The training directory contains the Exphormer-based model implementation and evaluation scripts.

Data Structure

Organize your datasets as follows in the datasets/ folder:

datasets/
└── [dataset_name]/
    ├── raw/
    │   ├── malnet-graph-tiny/  (Graph structures)
    │   └── split_info_tiny/    (Train/Val/Test splits)
    └── processed/              (Generated automatically)

If you download the dataset from our repository to the processed directory, you can skip the graph construction step.

Running Experiments

To reproduce the results, use the provided configuration files:

python main.py --cfg config_file.yaml

Citation

@misc{tran2026evaluating,
      title={Evaluating Out-of-Distribution Robustness in Graph-Based Android Malware Classification: A New Principled Benchmark}, 
      author={Ngoc N. Tran and Anwar Said and Waseem Abbas and Tyler Derr and Xenofon D. Koutsoukos},
      year={2026},
      eprint={2508.06734},
      archivePrefix={arXiv},
      primaryClass={cs.CR},
      url={https://arxiv.org/abs/2508.06734}, 
}