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

1. Problema e Motivação

A detecção de malware moderna enfrenta desafios crescentes devido à complexidade, ofuscação e técnicas de empacotamento (packing) dos códigos maliciosos. Métodos baseados em assinaturas são vulneráveis a pequenas alterações no código, enquanto a análise estática pode falhar em capturar comportamentos dinâmicos. A análise dinâmica, que observa o comportamento do programa em tempo de execução, é mais robusta, mas gera dados complexos.

A representação de programas como Grafos de Fluxo de Controle (CFGs) permite modelar essas estruturas de execução. No entanto, os Redes Neurais em Grafos (GNNs) tradicionais, embora eficazes, muitas vezes operam como "caixas pretas", dificultando a compreensão de por que uma decisão foi tomada. Além disso, um único modelo pode não capturar a diversidade de comportamentos e estruturas presentes em diferentes famílias de malware. A necessidade de modelos que sejam não apenas precisos, mas também explicáveis e adaptáveis a diferentes padrões de entrada é crítica para a segurança cibernética.

2. Metodologia Proposta

Os autores propõem um framework end-to-end que combina Mistura de Especialistas (MoE) com GNNs, focado em gerar explicações conscientes do roteamento (routing-aware). A arquitetura opera em três estágios principais:

A. Extração e Representação de Nós (CFGs)

• Extração Dinâmica: Os CFGs são extraídos de arquivos executáveis portáteis (PE) usando execução simbólica dinâmica (framework angr), capturando caminhos de execução reais que a análise estática poderia perder.
• Codificação de Nós em Duas Etapas:
  1. Codificação Baseada em Regras: Cada instrução assembly dentro de um bloco básico é convertida em um vetor de 439 dimensões, decompondo prefixos, opcodes, ModRM, SIB, deslocamento e imediato.
  2. Redução de Dimensionalidade: Um autoencoder é treinado para comprimir esses vetores em um espaço latente compacto (64 dimensões), preservando a informação essencial para o aprendizado subsequente.

B. Codificação de Nós Multi-Estatística (Nível de Nó)

Para promover diversidade na representação dos nós, cada camada da GNN calcula múltiplas estatísticas de vizinhança em paralelo. A atualização do nó $i$ na camada $l$ considera:

• Fator de Reponderação de Grau ($\rho$): $\rho \in \{0, 1\}$. Se 0, a agregação é uniforme; se 1, nós com maior grau recebem pesos maiores.
• Escolha de Pooling ($\lambda$): $\lambda \in \{mean, std, max\}$. Captura tendência central, dispersão ou respostas extremas.
• Fusão: As 6 combinações possíveis ($\rho \times \lambda$) são concatenadas e processadas por um MLP leve, criando representações de nós diversas e complementares.

C. Leitura de Especialistas e Roteamento (Nível de Grafo)

• Seis Especialistas: O modelo possui seis "especialistas" de leitura, cada um vinculado a uma combinação específica de $(\rho, \lambda)$. Cada especialista agrega os estados dos nós para produzir logits de grafo específicos.
• Roteamento Top-k com Balanceamento de Carga (LB): Uma rede de gating (portas) recebe as representações de todos os especialistas e seleciona os $k$ mais relevantes (Top-k) para cada grafo de entrada.
  • Utiliza-se uma operação Top-k para esparsidade.
  • Inclui-se uma perda auxiliar de balanceamento de carga (LB) para evitar que o roteador dependa excessivamente de um único especialista, garantindo que todos os especialistas sejam utilizados de forma equilibrada.

D. Explicações Conscientes do Roteamento

O método de explicação pós-hoc é o diferencial do trabalho:

1. Calcula-se a importância das arestas (pontuações $\beta$) individualmente para cada especialista selecionado (Top-k).
2. As explicações são agregadas ponderando as pontuações de cada especialista pelos pesos das portas ($\alpha$) aprendidos pelo roteador.
   • Fórmula: $\beta_{agregado} = \sum \alpha_i \cdot \beta_{(i)}$.

• Significado: Isso garante que a explicação final reflita não apenas o que cada especialista vê, mas também quão fortemente o modelo confiou naquele especialista para aquela decisão específica.

3. Principais Contribuições

1. Framework Unificado: Um sistema completo que une extração de CFG, codificação de nós multi-estatística, modelagem MoE-GNN, roteamento Top-k e atribuições pós-hoc alinhadas aos especialistas.
2. Diversidade em Dois Níveis: Uma arquitetura MoE que promove diversidade tanto na codificação de nós (via estatísticas de vizinhança) quanto na leitura de grafo (via especialistas especializados em diferentes vistas estruturais).
3. Explicabilidade Alinhada ao Roteamento: Um método inovador que agrega explicações baseadas na lógica de decisão real do modelo (quais especialistas foram ativados e com que peso), superando métodos que tratam o modelo como uma caixa preta única.

4. Resultados e Análise

O modelo foi avaliado em um conjunto de dados composto por amostras benignas (DikeDataset) e maliciosas (BODMAS e PMML), comparado com GNNs de especialista único (GCN, GIN, GAT, GraphSAGE) e outras variantes de MoE.

• Desempenho de Detecção:
  • O modelo MoE Top-2 com Balanceamento de Carga (LB) alcançou a melhor precisão geral (93,59%), superando significativamente as melhores GNNs de base (que atingiram ~89,10%).
  • Obteve a maior precisão para amostras benignas (96,23%) e a maior recall para maliciosas (97,94%).
• Impacto do Balanceamento de Carga (LB):
  • Sem LB, o modelo tendia a super-especializar-se em pares específicos de especialistas (ex: E4 e E6), reduzindo a diversidade.
  • Com LB, a ativação dos especialistas tornou-se mais diversificada e equilibrada, melhorando a estabilidade e o desempenho.
• Qualidade da Explicação:
  • Métricas de Fidelidade (Fidelity+ e Fidelity-) e Caracterização mostraram que a configuração MoE Top-2 com LB produziu as explicações mais fiéis e completas.
  • As explicações geradas foram robustas sob perturbações de esparsidade, indicando que as subestruturas identificadas são realmente críticas para a decisão do modelo.

5. Significado e Conclusão

Este trabalho demonstra que explicitar a dinâmica de roteamento em modelos MoE não apenas melhora a precisão da detecção de malware, mas também aumenta drasticamente a transparência do sistema. Ao vincular as explicações diretamente aos especialistas ativados e aos seus pesos, os analistas de segurança podem entender melhor quais padrões estruturais (ex: dispersão de vizinhança, centralidade de grau) levaram à classificação de um arquivo como malicioso.

A abordagem sugere que a combinação de codificação multi-estatística com especialização condicional e explicabilidade alinhada ao roteamento é um caminho promissor para sistemas de IA explicável (XAI) em segurança cibernética, permitindo auditorias mais confiáveis e detecção de malware mais robusta contra ofuscação.