DeepStage: Learning Autonomous Defense Policies Against Multi-Stage APT Campaigns
Authors: Trung V. Phan, Tri Gia Nguyen, Thomas Bauschert
Organizations: Chair of Communication Networks, Technische Universität Chemnitz, 09126 Chemnitz, Germany · Department of Information Assurance, FPT University, Da Nang 50509, Vietnam
This paper presents DeepStage, a deep reinforcement learning (DRL) framework for adaptive and stage-aware defense against Advanced Persistent Threats (APTs). The enterprise environment is formulated as a partially observable Markov decision process (POMDP), in which host provenance and network telemetry are fused into unified provenance graphs. Building on our prior work (StageFinder), DeepStage employs a graph neural network encoder and an LSTM-based stage estimator to infer probabilistic attacker stages aligned with the MITRE ATT&CK framework. The resulting stage beliefs, together with graph embeddings, are used to guide a hierarchical Proximal Policy Optimization (PPO) agent that selects defense actions across monitoring, access control, containment, and remediation. Experiments in a realistic enterprise testbed with CALDERA-driven APT playbooks show that DeepStage achieves an average F1-score of 0.887 and a mitigation success rate of 84.7%, outperforming a risk-aware DRL baseline by 21.8% in F1-score and 16.2% in mitigation success. The results demonstrate effective stage-aware and cost-efficient autonomous cyber defense.
Figures & tables
Figure 1: Data and control flow of the proposed DeepStage framework.
Meta-policy
Sub-policy
Deployment in Testbed
Cost C(a)
Amon
a0 : Maintain baseline monitoring
Keep auditd.conf , CamFlow, and Zeek agent in baseline mode
0.01
a1 : Increase host logging level
Modify auditd.conf , CamFlow, and Zeek agent policies to capture process and socket events
0.05
a2 : Activate Deep Packet Inspection
Temporarily enable Zeek full packet capture on target subnets
0.10
a3 : Memory/process snapshot
Use volatility or psrecord via management API to collect runtime memory dump
0.20
a4 : Deploy honeypot redirect
Deploy Cowrie SSH honeypot or redirect via iptables DNAT
0.25
a5 : Log enrichment with threat intel
Query local MISP feed and annotate alerts with TTP tags in provenance database
0.10
Table I: Set of practical sub-policies deployable. Cost C(a) denotes the normalized operational impact.
Parameter
Value / Range
Description
Actor network
3-layer MLP
Input: fused graph embedding ( gt ) + stage belief ( pt ); Output: probability over meta/sub-policy actions.
Critic network
3-layer MLP
Estimates value V(ot) for baseline; shared first layer with actor.
Optimizer
Adam ( β1=0.9 , β2=0.999 )
Adaptive optimization for actor and critic updates.
Learning rate ( η )
3×10−4 (decay 0.99 per epoch)
Stable step size for PPO gradient updates.
Discount factor ( γ )
0.99
Long-term credit assignment for delayed APT mitigation rewards.
GAE parameter ( λGAE )
0.95
Balances bias and variance in advantage estimation.
Table II: PPO training parameters for the DeepStage DRL-based APT defense.
Method
Recon
InitAcc
PrivEsc
LatMov
C2
Exfil
Avg-F1
Mitigation Success (%)
Risk-Aware DRL
0.75
0.73
0.70
0.71
0.76
0.72
0.728
68.5
DeepStage-Unaware
0.82
0.79
0.76
0.78
0.83
0.80
0.797
76.7
DeepStage
0.91
0.88
0.85
0.87
0.92
0.89
0.887
84.7
Table III: Per-stage F1-score and overall defense effectiveness comparison.
Figure 2: Cost–effectiveness frontiers illustrating normalized security gain versus cumulative action cost.
Figure 3: Training convergence of hierarchical PPO across methods.
Figure 4: Defense responsiveness over APT stage transitions.
Advanced Persistent Threats (APTs) are stealthy, multi-stage attacks that require adaptive and timely defense. While deep reinforcement learning (DRL) enables autonomous cyber defense, its decisions are often opaque and difficult to trust in operational environments. This paper presents DeepXplain, an explainable DRL framework for stage-aware APT defense. Building on our prior DeepStage model, DeepXplain integrates provenance-based graph learning, temporal stage estimation, and a unified XAI pipeline that provides structural, temporal, and policy-level explanations. Unlike post-hoc methods, explanation signals are incorporated directly into policy optimization through evidence alignment and confidence-aware reward shaping. To the best of our knowledge, DeepXplain is the first framework to integrate explanation signals into reinforcement learning for APT defense. Experiments in a realistic enterprise testbed show improvements in stage-weighted F1-score (0.887 to 0.915) and success rate (84.7% to 89.6%), along with higher explanation confidence (0.86), improved fidelity (0.79), and more compact explanations (0.31). These results demonstrate enhanced effectiveness and trustworthiness of autonomous cyber defense.
Trung V. Phan, Thomas Bauschert
Chair of Communication Networks, Technische Universität Chemnitz, 09126 Chemnitz, Germany
Autonomous agents are increasingly deployed in both offensive and defensive cyber operations, creating high-speed, closed-loop interactions in critical infrastructure environments. Advanced Persistent Threat (APT) actors exploit "Living off the Land" techniques and targeted telemetry perturbations to induce ambiguity in monitoring systems, causing automated defenses to overreact or misclassify benign behavior as malicious activity. Existing monolithic and multi-agent defense pipelines largely operate on correlation-based signals, lack structural constraints on response actions, and are vulnerable to reasoning drift under ambiguous or adversarial inputs. We present the Causal Multi-Agent Decision Framework (C-MADF), a structurally constrained architecture for autonomous cyber defense that integrates causal modeling with adversarial dual-policy control. C-MADF first learns a Structural Causal Model (SCM) from historical telemetry and compiles it into an investigation-level Directed Acyclic Graph (DAG) that defines admissible response transitions. This roadmap is formalized as a Markov Decision Process (MDP) whose action space is explicitly restricted to causally consistent transitions. Decision-making within this constrained space is performed by a dual-agent reinforcement learning system in which a threat-optimizing Blue-Team policy is counterbalanced by a conservatively shaped Red-Team policy. Inter-policy disagreement is quantified through a Policy Divergence Score and exposed via a human-in-the-loop interface equipped with an Explainability-Transparency Score that serves as an escalation signal under uncertainty. On the real-world CICIoT2023 dataset, C-MADF reduces the false-positive rate from 11.2%, 9.7%, and 8.4% in three cutting-edge literature baselines to 1.8%, while achieving 0.997 precision, 0.961 recall, and 0.979 F1-score.
Yiyao Zhang, Diksha Goel, Hussain Ahmad
Adelaide University, Australia · CSIRO’s Data61, Australia
Autonomous cyber defense systems based on Deep Reinforcement Learning (DRL) have attracted significant research attention, yet remain evaluated almost exclusively against static, heuristic red agents, leaving their robustness against adaptive threats critically understudied. Meanwhile, recent advances in Reinforcement Learning with Verifiable Rewards (RLVR) have improved LLM reasoning, but their integration into cybersecurity remains elusive due to the absence of suitable benchmark environments and interaction datasets. To bridge this gap, we introduce Trident, an agentic LLM red teaming framework comprising three components: a dynamic benchmark with isolated sandbox servers spanning CybORG CAGE 4 and CyberWheel, a dataset comprises over 13,000 high-fidelity red-blue interaction trajectories for RLVR, and a ``Code-as-Policy'' RLVR agentic architecture Trident Agentic). The latter reformulates red agent training as a contextual bandit via a tripartite Log Summarizer--Planner--Coder design, where a trainable Planner generates complete attack strategies from compressed execution logs, which a frozen Coder translates into executable Python policies deployed against live DRL defenders. Empirical evaluations reveal a fundamental brittleness in existing defenses: with a single trainable 7B planner, Trident reduces blue agent defensive performance by an average of 522% compared to static red agent baselines while autonomously discovering emergent behaviors such as decoy avoidance and adaptive state prioritization that static heuristics entirely fail to uncover.
Ryozo Masukawa, Ian Bryant, Armita Kazeminajafabadi +6
University of California, Irvine · Northeastern University · Johns Hopkins University