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

1. Problématique

La détection de logiciels malveillants (malware) moderne fait face à des défis croissants liés à la complexité et à l'échelle des menaces. Les méthodes traditionnelles basées sur les signatures sont vulnérables au packing et à l'obfuscation. Bien que l'analyse dynamique (exécution contrôlée) offre des preuves comportementales riches, elle génère des structures de données complexes, notamment des Graphes de Flux de Contrôle (CFG).

L'utilisation de Réseaux de Neurones à Graphes (GNN) pour modéliser ces CFG est prometteuse, mais présente deux limites majeures :

1. Manque de diversité : Un modèle unique peine souvent à capturer la variété des comportements malveillants et des structures de graphes.
2. Opacité (Boîte noire) : Les décisions des GNN sont difficiles à interpréter. Dans le domaine de la cybersécurité, il est crucial de comprendre pourquoi un modèle a classé un échantillon comme malveillant (explicabilité) pour valider les alertes et prioriser les investigations.

Les modèles existants de type "Mixture of Experts" (MoE) améliorent la performance mais peinent souvent à fournir des explications cohérentes avec la logique de routage interne du modèle.

2. Méthodologie Proposée

Les auteurs proposent un cadre de détection de malware explicable et conscient du routage (routing-aware), basé sur une architecture MoE appliquée aux CFG.

A. Extraction et Représentation des Nœuds

• Extraction dynamique : Les CFG sont extraits de fichiers exécutables (PE) via une exécution symbolique dynamique (framework angr) pour capturer les chemins d'exécution réels.
• Encodage en deux étapes :
  1. Encodage basé sur des règles : Chaque instruction assembleur (x86-64) d'un bloc de base est convertie en un vecteur de 439 dimensions (préfixes, opcodes, ModRM, SIB, déplacements, immédiats, options).
  2. Réduction de dimension : Un autoencodeur compresse ces vecteurs en un espace latent de 64 dimensions, préservant l'information essentielle tout en réduisant la complexité.

B. Architecture MoE à Deux Niveaux de Diversité

Le modèle introduit une diversité à deux niveaux pour capturer des perspectives structurelles complémentaires :

1. Niveau Nœud (Encodage Multi-statistique) :
   Chaque couche GNN calcule plusieurs statistiques de voisinage en parallèle, fusionnées via un MLP léger. Cela est contrôlé par deux hyperparamètres :

   • Facteur de ré-pondération du degré ($\rho$) : $\rho=0$ (pondération uniforme) ou $\rho=1$ (pondération biaisée par le degré, favorisant les nœuds centraux).
   • Choix de pooling ($\lambda$) : Moyenne (mean), Écart-type (std), ou Maximum (max).
     Cela génère 6 vues distinctes par nœud ($2 \times 3$ combinaisons).

2. Niveau Lecture (Readout) - Les Experts :
   Six "experts" distincts sont entraînés, chacun associé à une combinaison spécifique $(\rho, \lambda)$. Chaque expert produit des logits au niveau du graphe basés sur sa vue spécialisée.

C. Routage Top-k avec Équilibrage de Charge (Load Balancing - LB)

• Un réseau de "gating" (commutateur) reçoit les représentations des six experts.
• Il sélectionne dynamiquement les k meilleurs experts (Top-k) pour chaque échantillon d'entrée via une opération de Top-k softmax.
• Une perte d'équilibrage de charge (LB) est ajoutée à la fonction de coût pour éviter que le modèle ne s'effondre sur un seul expert (expert collapse) et pour encourager une utilisation équilibrée de tous les experts.

D. Explications Post-hoc Conscientes du Routage

C'est l'innovation clé du papier. Au lieu d'agréger aveuglément les explications :

1. Des scores d'importance par arête sont calculés pour chaque expert sélectionné (Top-k).
2. Ces scores sont agrégés en pondérant les explications de chaque expert par les portes (gates) apprises par le routeur ($\alpha_i$).
3. La formule d'agrégation est : $\beta^{(a,j)} = \sum_{i=1}^{k} \alpha_i \beta^{(i,j)}$.
   Cela garantit que l'explication finale reflète non seulement ce que chaque expert voit, mais aussi à quel point le modèle a fait confiance à cet expert pour cette décision spécifique.

3. Contributions Clés

1. Cadre unifié : Une approche de bout en bout combinant l'extraction de CFG, l'encodage multi-statistique, le modèle MoE avec routage Top-k, et des attributions alignées sur le routage.
2. Diversité à deux niveaux : Combinaison d'un encodage de nœuds riche (statistiques multiples) et d'une spécialisation au niveau des experts pour gérer l'hétérogénéité des malwares.
3. Méthode d'explication innovante : Une méthode post-hoc qui intègre dynamiquement les poids du routeur dans l'explication, améliorant la fidélité et la transparence des décisions du modèle.

4. Résultats Expérimentaux

L'évaluation a été menée sur des jeux de données contenant des échantillons malveillants (BODMAS, PMML) et bénins (DikeDataset), comparant le modèle proposé à des GNN standards (GCN, GAT, GIN, GraphSAGE) et à d'autres variantes MoE.

• Performance de Détection :

  • La variante MoE Top-2 avec Équilibrage de Charge (LB) a obtenu les meilleurs résultats avec une précision de 93,59% (contre ~89% pour les meilleurs GNN simples).
  • Elle a atteint une rappel (recall) de 97,94% pour les malwares et une précision de 96,23% pour les logiciels bénins.
  • L'ajout de LB a significativement amélioré la stabilité et l'équilibre par classe par rapport aux versions sans LB.

• Analyse du Routage :

  • Sans LB, le modèle tendait à sur-spécialiser certains paires d'experts (biais de sélection).
  • Avec LB, la diversité d'activation des experts a augmenté, prouvant que le modèle exploite réellement la complémentarité des six experts.

• Évaluation de l'Explicabilité (Fidélité) :

  • Les métriques Fidelity+ (nécessité) et Fidelity- (suffisance) ont été mesurées via des perturbations de sparsité du graphe.
  • Le modèle Top-2 avec LB a produit les explications les plus fidèles (scores de Characterization les plus élevés), confirmant que les sous-graphes identifiés sont à la fois indispensables et suffisants pour la prédiction.
  • Les explications "conscientes du routage" sont plus robustes que les méthodes d'agrégation uniforme ou basées sur un softmax de température.

5. Signification et Impact

Ce travail démontre que l'intégration explicite de la dynamique de routage dans le processus d'explication améliore considérablement la transparence des modèles MoE appliqués à la cybersécurité.

• Pour la sécurité : Cela permet aux analystes de faire confiance aux alertes générées par l'IA, car l'explication correspond à la logique interne réelle du modèle (quels experts ont été activés et pourquoi).
• Pour la recherche : Cela valide l'approche MoE pour la détection de malware, montrant qu'une spécialisation conditionnelle des experts, couplée à un équilibrage de charge, surpasse les modèles monolithiques tout en restant interprétable.

En conclusion, l'approche proposée offre un compromis optimal entre haute précision de détection et interprétabilité robuste, un critère essentiel pour le déploiement opérationnel de systèmes de détection de malware basés sur l'IA.