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

1. Il Problema

La rilevazione moderna dei malware deve affrontare sfide crescenti dovute alla complessità e alla scala degli attacchi, rendendo insufficienti i sistemi basati su firme tradizionali. Sebbene l'analisi dinamica (esecuzione in ambiente controllato) fornisca evidenze comportamentali ricche, la trasformazione di questi dati in modelli di apprendimento automatico interpretabili rimane difficile.
In particolare, l'uso delle Reti Neurali su Grafi (GNN) per modellare i Grafi del Flusso di Controllo (CFG) dei programmi è promettente, ma presenta due limiti principali:

1. Mancanza di diversità: Un singolo modello GNN fatica a catturare la vasta gamma di comportamenti e strutture presenti nelle diverse famiglie di malware.
2. Opacità decisionale: Le GNN sono spesso "scatole nere". È cruciale non solo prevedere se un file è malevolo, ma spiegare perché, identificando quali parti del grafo (basi blocchi, transizioni) hanno influenzato la decisione.
   Le tecniche esistenti di Ensemble Learning o Mixture of Experts (MoE) spesso non integrano la logica di routing (come il modello seleziona gli esperti) nel processo di spiegazione, rendendo le attribuzioni meno fedeli al processo decisionale effettivo.

2. Metodologia Proposta

Gli autori propongono un framework end-to-end per la rilevazione di malware basato su CFG, che integra un modello Mixture of Experts (MoE) con spiegazioni Routing-Aware (consapevoli del routing).

A. Estrazione e Rappresentazione dei Nodi

• Costruzione del CFG: I grafi del flusso di controllo vengono estratti dinamicamente da file PE (Portable Executable) utilizzando l'esecuzione simbolica (framework angr), garantendo una rappresentazione fedele del comportamento runtime.
• Embedding dei Nodi (Due fasi):
  1. Codifica basata su regole: Ogni istruzione assembly (x86-64) in un blocco di base viene convertita in un vettore di 439 dimensioni basato su componenti (prefissi, opcode, ModRM, SIB, displacement, immediate, flag opzionali).
  2. Riduzione dimensionale: Un autoencoder comprime questi vettori in uno spazio latente compatto (64 dimensioni), preservando le informazioni essenziali per l'apprendimento.

B. Architettura MoE a Due Livelli

Il modello introduce diversità a due livelli distinti:

1. Livello Nodo (Codifica Multi-Statistiche):

   • Ogni strato della GNN calcola statistiche di vicinato multiple in parallelo.
   • Vengono combinate due scelte di reweighting del grado ($\rho \in \{0, 1\}$): uniforme o pesato in base al grado del nodo.
   • Vengono applicati tre operatori di pooling ($\lambda \in \{mean, std, max\}$): media, deviazione standard e massimo.
   • Le 6 combinazioni risultanti ($\rho, \lambda$) vengono fuse tramite un MLP leggero per creare rappresentazioni nodali diversificate.

2. Livello Lettura (Expert Graph):

   • Sono definiti 6 esperti, ciascuno specializzato su una specifica vista $(\rho, \lambda)$ del grafo.
   • Ogni esperto produce un logit a livello di grafo aggregando le rappresentazioni dei nodi finali secondo la sua specifica configurazione statistica.

C. Routing e Selezione degli Esperti

• Una rete di gating (gating network) riceve le rappresentazioni combinate dei 6 esperti.
• Utilizza una strategia Top-k routing (con $k=2$ negli esperimenti principali) per selezionare e pesare dinamicamente solo gli esperti più rilevanti per il campione in input.
• Viene introdotta una Loss di Bilanciamento del Carico (Load Balancing - LB) per prevenire il collasso degli esperti (dove un solo esperto domina) e garantire un utilizzo equo e diversificato di tutti i moduli.

D. Spiegazioni Routing-Aware (Post-hoc)

Il contributo chiave per l'interpretabilità è il metodo di spiegazione:

• Vengono calcolate spiegazioni a livello di bordo (edge-level) per ciascun esperto selezionato (usando metodi come Integrated Gradients).
• Le spiegazioni vengono aggregate pesando i punteggi di importanza di ogni esperto con le porte (gate weights) apprese dal router.
• Risultato: La spiegazione finale riflette non solo cosa ogni esperto ha evidenziato, ma anche quanto fortemente quel specifico esperto è stato selezionato per quella decisione, allineando la spiegazione alla logica condizionale del modello.

3. Contributi Chiave

1. Framework Unificato: Un sistema end-to-end che unisce estrazione CFG, codifica nodale multi-statistica, modellazione MoE con routing Top-k e spiegazioni allineate al routing.
2. Diversità a Due Livelli: Un'architettura MoE che combina la diversità nell'encoding dei nodi (statistiche di vicinato) con la specializzazione a livello di esperti (6 viste distinte).
3. Metodo di Spiegazione Innovativo: Una tecnica post-hoc che integra le porte di routing nelle spiegazioni, garantendo che la "rationale" fornita sia fedele al meccanismo di selezione condizionale del modello.
4. Valutazione Completa: Una valutazione rigorosa che misura non solo l'accuratezza, ma anche la fedeltà (fidelity) e la caratterizzazione delle spiegazioni sotto perturbazioni di sparsità.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su dataset reali (BODMAS, PMML per malware; DikeDataset per benigni) confrontando il modello proposto con baseline GNN singole (GCN, GIN, GAT, GraphSAGE) e varianti MoE.

• Prestazioni di Rilevazione:

  • Il modello MoE Top-2 con Load Balancing ha ottenuto la migliore accuratezza complessiva (93.59%), superando significativamente le baseline (che si attestavano tra l'85% e l'89%).
  • Ha mostrato le migliori prestazioni per classe: Precisione per il benigno al 96.23% e Recall per il malware al 97.94%.
  • L'uso del bilanciamento del carico (LB) ha migliorato la stabilità e l'equilibrio tra le classi rispetto alle varianti senza LB.

• Analisi del Routing:

  • L'analisi delle heatmap di co-selezione ha dimostrato che senza LB, il modello tende a selezionare sempre la stessa coppia di esperti (bias). Con LB, la selezione è diversificata, coinvolgendo tutti e 6 gli esperti in modo cooperativo.

• Valutazione dell'Interpretabilità:

  • Utilizzando metriche di Fidelity (quanto la rimozione di un sottografo influisce sulla previsione) e Characterization Score, il modello Top-2 con LB ha dimostrato la massima fedeltà.
  • Le spiegazioni generate sono risultate più robuste e necessarie/sufficienti rispetto a quelle di modelli con aggregazione uniforme o routing basato su softmax a temperatura, confermando che allineare la spiegazione al routing dinamico migliora la trasparenza.

5. Significato e Impatto

Questo lavoro dimostra che l'uso esplicito di architetture Mixture of Experts con meccanismi di routing dinamici non solo migliora l'accuratezza nella rilevazione di malware complessi, ma è fondamentale per la trasparenza dei sistemi di sicurezza.

• Affidabilità: Le spiegazioni "Routing-Aware" forniscono agli analisti di sicurezza una giustificazione difendibile delle decisioni, mostrando quali pattern strutturali specifici (es. deviazione standard delle connessioni o nodi ad alto grado) hanno guidato la classificazione.
• Generalizzabilità: L'approccio suggerisce che la diversità strutturale e la specializzazione condizionale sono strategie superiori rispetto ai modelli monolitici per dati complessi come i grafi di esecuzione.
• Futuro: Il framework apre la strada all'integrazione di grafi multi-modali (chiamate API, syscall) e all'uso di meccanismi di gating adattivi per la sicurezza informatica.

In sintesi, il paper stabilisce un nuovo standard per i modelli di rilevazione malware basati su grafi, dimostrando che l'interpretabilità non è un'aggiunta postuma, ma una componente intrinseca che può essere ottimizzata attraverso un'architettura di routing consapevole.