Routing-Aware Explanations for Mixture of Experts Graph Models in Malware Detection

Probleemstelling

Moderne malware wordt steeds complexer en vereist detectiemethoden die verder gaan dan handtekeningsgebaseerde systemen. Dynamische analyse, die het runtime-gedrag van malware observeert, biedt rijkere informatie dan statische analyse, maar is moeilijker te modelleren. Graph Neural Networks (GNN's) zijn geschikt voor het modelleren van programmastructuren zoals Control Flow Graphs (CFG's), maar standaard GNN's hebben vaak moeite met de enorme diversiteit aan malware-gedrag en -structuren.

Daarnaast is uitlegbaarheid (explainability) een kritieke uitdaging. In beveiligingscontexten is het niet voldoende om alleen te weten dat een bestand kwaadaardig is; analisten moeten begrijpen waarom. Bestaande uitlegmethodes voor GNN's zijn vaak niet optimaal afgestemd op geavanceerde ensemble-modellen zoals Mixture of Experts (MoE), waarbij de beslissing wordt genomen door een combinatie van gespecialiseerde submodellen. Er ontbreekt een methode die de uitleg direct koppelt aan hoe het model de experts selecteert (routing).

Methodologie

Het paper introduceert een end-to-end framework voor malware-detectie op CFG's dat een Mixture of Experts (MoE) architectuur combineert met routing-aware uitleg.

1. Data Voorbereiding en Feature Embedding

• CFG Extractie: Dynamische symbolische executie (met het angr framework) wordt gebruikt om CFG's te extraheren uit Portable Executable (PE) bestanden.
• Tweestaps Feature Embedding:
  • Stap 1 (Rule-based Encoding): Assembler-instructies binnen basisblokken worden omgezet in 439-dimensionale vectoren op basis van instructiecomponenten (prefix, opcode, ModRM, SIB, displacement, immediate, option).
  • Stap 2 (Autoencoder): Een autoencoder comprimeert deze hoge dimensie naar een compacte latent space (64 dimensies) om ruis te verminderen en essentiële informatie te behouden.

2. Multi-Statistiek Node Encoding (Niveau 1)

Om diversiteit op knooppuntniveau te creëren, berekent elke GNN-laag meerdere buurtstatistieken parallel:

• Degree Reweighting ($\rho$): $\rho=0$ (uniforme aggregatie) of $\rho=1$ (gewicht op basis van knooppuntgraad).
• Pooling Keuze ($\lambda$): $\lambda \in \{mean, std, max\}$.
• Deze zes combinaties ($\rho \in \{0,1\}, \lambda \in \{mean, std, max\}$) worden samengevoegd via een lichtgewicht MLP, wat resulteert in diverse knooppuntrepresentaties die complementaire structurele cues vastleggen.

3. Expert Readout en Routing (Niveau 2)

• Zes Experts: Er zijn zes gespecialiseerde experts, elk gekoppeld aan een specifieke $(\rho, \lambda)$-combinatie. Elke expert produceert een graf-niveau logit (voorspelling) op basis van zijn specifieke aggregatiestrategie.
• Top-k Routing met Load Balancing (LB): Een gating netwerk (router) selecteert de $k$ meest relevante experts voor een specifieke graaf.
  • De router gebruikt een Top-k operatie (in dit geval $k=2$) om de selectie te verspreiden.
  • Een Load Balancing auxiliary loss wordt toegevoegd om te voorkomen dat de router alleen op één expert leunt (expert collapse), wat zorgt voor een evenwichtig gebruik van alle experts.

4. Routing-Aware Post-Hoc Uitleg

In plaats van een algemene uitleg te genereren, wordt de uitleg afgestemd op de daadwerkelijke beslissing van het model:

• Per geselecteerde expert worden edge-level uitlegscores (bijv. via Integrated Gradients) berekend.
• Deze scores worden geaggregeerd, gewogen met de gate-weights ($\alpha$) van de router.
• Formule: $\beta^{(a,j)} = \sum_{i=1}^{k} \alpha_i \beta^{(i,j)}$.
• Dit zorgt ervoor dat de uiteindelijke uitleg zowel weergeeft wat de expert ziet als hoe sterk die expert bijdraagt aan de uiteindelijke voorspelling.

Belangrijkste Bijdragen

1. End-to-End Framework: Een unificatie van CFG-construction, twee-staps MoE-GNN-modellering, Top-k routing en post-hoc expert-uitgelijnde attributie.
2. Twee-niveau Diversiteit: Een MoE-architectuur die zowel diversiteit op knooppuntniveau (via multi-statistiek encoding) als op expert-niveau (via zes gespecialiseerde lezers) combineert.
3. Routing-Aware Uitleg: Een nieuwe methode voor post-hoc uitleg die de bijdrage van individuele experts aggregeert op basis van de geleerde router-gates, wat leidt tot meer betrouwbare en transparante rationale.
4. Robuuste Evaluatie: Een uitgebreide evaluatie die niet alleen nauwkeurigheid meet, maar ook de "faithfulness" (trouw) van de uitleg onder sparsiteits-perturbaties.

Resultaten

De methode is getest op een dataset bestaande uit malware (BODMAS, PMML) en benign software (DikeDataset).

• Detectienauwkeurigheid: De voorgestelde Top-2 MoE met Load Balancing behaalde de beste resultaten met 93,59% nauwkeurigheid. Dit is significant beter dan single-expert baselines (zoals GCN, GIN, GAT) die rond de 86-89% scoorden, en ook beter dan andere MoE-varianten (zoals uniform aggregatie of Top-1 routing).
• Class Balans: Het model behaalde een hoge precisie voor benign samples (96,23%) en een hoge recall voor malware (97,94%).
• Expert Diversiteit: De Load Balancing term bleek cruciaal; zonder LB neigde het model naar een sterke bias naar slechts één paar experts (E4, E6). Met LB werd het gebruik van alle zes experts veel diverser en evenwichtiger.
• Uitlegbaarheid: De Characterization Score (een combinatie van Fidelity+ en Fidelity-) was het hoogst voor de Top-2 MoE met LB. Dit betekent dat de gegenereerde subgrafen zowel noodzakelijk als voldoende waren voor de voorspelling, wat aantoont dat de routing-aware uitleg trouw is aan het modelgedrag.

Betekenis en Conclusie

Dit paper toont aan dat het expliciet modelleren van expert-selectie en het balanceren van expert-gebruik niet alleen de detectienauwkeurigheid van malware verbetert, maar ook de transparantie van de beslissingen. Door de uitleg te koppelen aan de dynamische routing van het model, krijgen beveiligingsanalisten inzicht in welke specifieke structurele patronen (bijv. dispersie vs. gemiddelde, graad-gewogen vs. uniform) leidend waren tot de classificatie.

De studie legt een fundament voor toekomstig werk in multi-modale graaf-invoer en adaptieve gating-mechanismen, en onderstreept dat uitlegbaarheid in complexe ensemble-modellen essentieel is voor vertrouwen en audit in cybersecurity-toepassingen.