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

1. Problem Statement

Modern malware is increasingly complex, utilizing obfuscation, packing, and dynamic code loading to evade traditional signature-based detection. While Graph Neural Networks (GNNs) have shown promise in modeling program structures (specifically Control Flow Graphs or CFGs) to detect behavioral patterns, standard GNNs often rely on a single, monolithic model. This approach struggles with the high diversity of malware families and execution behaviors. Furthermore, standard GNNs are often "black boxes," making it difficult for security analysts to understand why a specific sample was flagged as malicious. Existing Mixture of Experts (MoE) frameworks improve performance but often lack transparency regarding how different experts contribute to a final decision, and their explanations are rarely aligned with the dynamic routing mechanism used during inference.

2. Methodology

The authors propose an end-to-end framework that combines Control Flow Graph (CFG) extraction with a Routing-Aware Mixture of Experts (MoE) architecture. The methodology consists of four main stages:

A. Data Preprocessing and Feature Embedding

• CFG Extraction: Dynamic symbolic execution (using the angr framework) is applied to Portable Executable (PE) files to generate CFGs. This captures runtime behaviors that static analysis might miss.
• Two-Stage Node Embedding:
  1. Rule-Based Encoding: Assembly instructions within each basic block are decomposed into seven components (prefix, opcode, ModRM, SIB, displacement, immediate, option) and converted into a 439-dimensional one-hot/binary vector.
  2. Dimensionality Reduction: An autoencoder compresses the 439-dimensional vectors into a compact 64-dimensional latent space ($x^{(k)}$), preserving essential semantic information while reducing noise.

B. Node-Level Multi-Statistic Encoding

To ensure diversity at the node level, the model computes multiple neighborhood statistics in parallel for each layer.

• Degree Reweighting ($\rho$): Controls aggregation weights. $\rho=0$ implies uniform weighting; $\rho=1$ biases aggregation toward high-degree (structurally central) nodes.
• Pooling Statistics ($\lambda$): Three aggregation functions are applied: Mean (central tendency), Standard Deviation (dispersion), and Max (extreme responses).
• Fusion: The model concatenates the outputs of all 6 combinations ($\rho \in \{0,1\} \times \lambda \in \{mean, std, max\}$) and passes them through a lightweight MLP to create diverse node representations.

C. Graph-Level Expert Readouts and Routing

• Six Specialized Experts: The model employs six distinct "experts," each corresponding to one of the six $(\rho, \lambda)$ views. Each expert aggregates the final node representations into a graph-level logit vector.
• Top-k Routing with Load Balancing (LB): A gating network takes the concatenated outputs of all six experts and computes routing probabilities.
  • Sparsity: It selects only the top-$k$ experts (e.g., Top-2) for a given input using a sparse softmax.
  • Load Balancing: An auxiliary loss term ($L_{LB}$) is added to the objective function to prevent "expert collapse" (where one expert dominates all inputs), ensuring balanced utilization across the expert pool.

D. Routing-Aware Post-Hoc Explanation

The paper introduces a novel explanation method that aligns with the model's conditional computation:

1. Per-Expert Attribution: For the $k$ selected experts, edge-level importance scores are generated using post-hoc techniques (e.g., Integrated Gradients).
2. Gate-Weighted Aggregation: The final explanation for an edge is a weighted sum of the per-expert scores, where the weights are the routing gates ($\alpha$) determined by the router.
   • Formula: $\beta^{(a,j)} = \sum_{i=1}^{k} \alpha_i \beta^{(i,j)}$
   • This ensures the rationale reflects both what the experts found and how much the router relied on them.

3. Key Contributions

1. End-to-End Routing-Aware Framework: Unifies CFG construction, two-stage MoE-GNN modeling, Top-k routing, and post-hoc expert-aligned attributions into a single pipeline.
2. Two-Level Diversity: Introduces a novel architecture that promotes diversity at both the node level (via multi-statistic encoding with $\rho$ and $\lambda$) and the readout level (via six specialized experts).
3. Routing-Aware Explanation: Proposes a method to aggregate explanations based on the router's gates, ensuring the explanation is faithful to the specific subset of experts actually used for the prediction.
4. Comprehensive Evaluation: Establishes a rigorous evaluation protocol measuring not only accuracy but also Fidelity (necessity and sufficiency of subgraphs) and Characterization scores under sparsity perturbations.

4. Results and Analysis

The framework was evaluated on a dataset combining BODMAS and PMML (malware) and DikeDataset (benign).

• Detection Accuracy:

  • The proposed Top-2 MoE with Load Balancing achieved the highest accuracy of 93.59%, significantly outperforming single-expert baselines (GCN, GAT, GIN, GraphSAGE) which ranged from 85.90% to 89.10%.
  • It achieved a 96.23% precision for benign samples and 97.94% recall for malicious samples.
  • Ablation studies showed that Top-2 routing with LB outperformed Top-1 routing and uniform aggregation, proving the value of conditional computation and balanced expert utilization.

• Expert Diversity:

  • Without Load Balancing, the router exhibited strong bias, relying heavily on a single expert pair (E4, E6).
  • With Load Balancing, the activation of experts became significantly more diverse and balanced, preventing over-specialization.

• Explainability (Fidelity & Characterization):

  • The Top-2 MoE with LB configuration yielded the best Characterization scores across all sparsity levels (5%–95%).
  • This indicates that the explanations generated were both necessary (removing the identified subgraph changed the prediction) and sufficient (the subgraph alone could reproduce the prediction).
  • Routing-aware aggregation proved superior to uniform or temperature-based routing in generating faithful explanations.

5. Significance

This paper addresses a critical gap in AI-driven cybersecurity: the trade-off between high-performance detection and model interpretability.

• Trust and Auditing: By making the router explicit and aligning explanations with the routing logic, security analysts can trust the model's decision-making process, seeing exactly which structural patterns (e.g., specific control flow loops or API call sequences) triggered the alert.
• Robustness: The use of Load Balancing ensures the model remains robust against distribution shifts by preventing reliance on a single "lucky" expert.
• Generalizability: The framework demonstrates that combining multi-statistic encoding with conditional expert selection creates a more transparent and effective tool for analyzing complex, obfuscated malware than traditional monolithic GNNs.

In conclusion, the study proves that explicitly modeling expert selection and balancing not only improves detection accuracy but also significantly enhances the transparency and faithfulness of the resulting explanations, making it a viable solution for operational malware analysis workflows.