Structural Generalization in Autonomous Cyber Incident Response with Message-Passing Neural Networks and Reinforcement Learning

TL;DR

This paper presents a method for autonomous cyber incident response that leverages message passing neural networks and reinforcement learning to handle dynamic network structures without retraining, improving adaptability in changing network environments.

Contribution

The paper introduces a relational agent learning approach using message passing neural networks that generalizes across different network structures in cyber incident response tasks.

Findings

Relational agents effectively handle network structure changes.

Agents can perform optimally without retraining on new network variants.

Default vector-based agents require retraining for each network change.

Abstract

We believe that agents for automated incident response based on machine learning need to handle changes in network structure. Computer networks are dynamic, and can naturally change in structure over time. Retraining agents for small network changes costs time and energy. We attempt to address this issue with an existing method of relational agent learning, where the relations between objects are assumed to remain consistent across problem instances. The state of the computer network is represented as a relational graph and encoded through a message passing neural network. The message passing neural network and an agent policy using the encoding are optimized end-to-end using reinforcement learning. We evaluate the approach on the second instance of the Cyber Autonomy Gym for Experimentation (CAGE~2), a cyber incident simulator that simulates attacks on an enterprise network. We create variants of the original network with different numbers of hosts and agents are tested without additional training on them. Our results show that agents using relational information are able to find solutions despite changes to the network, and can perform optimally in some instances. Agents using the default vector state representation perform better, but need to be specially trained on each network variant, demonstrating a trade-off between specialization and generalization.