# Preliminary Investigation into Uncertainty-Aware Attack Stage Classification

arXiv: https://arxiv.org/abs/2508.00368 · HTML: https://mlsanity.com/paper/2508.00368/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 6 of its 26 references that are in the map.
- First posted on arXiv 1 Aug 2025.

## Builds on (most cited first, 6 of 6)

- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [Uncertainty-Aware Deep Classifiers using Generative Models](https://mlsanity.com/paper/2006.04183/) (2020)
- [Developing Optimal Causal Cyber-Defence Agents via Cyber Security Simulation](https://mlsanity.com/paper/2207.12355/) (2022)
- [Combating Advanced Persistent Threats: Challenges and Solutions](https://mlsanity.com/paper/2309.09498/) (2023)
- [Multi-Stage Attack Detection via Kill Chain State Machines](https://mlsanity.com/paper/2103.14628/) (2021)
- [Enhancing Trustworthiness in ML-Based Network Intrusion Detection with Uncertainty Quantification](https://mlsanity.com/paper/2310.10655/) (2023)

## Abstract (arXiv)

Advanced Persistent Threats (APTs) represent a significant challenge in cybersecurity due to their prolonged, multi-stage nature and the sophistication of their operators. Traditional detection systems typically focus on identifying malicious activity in binary terms (benign or malicious) without accounting for the progression of an attack. However, effective response strategies depend on accurate inference of the attack's current stage, as countermeasures must be tailored to whether an adversary is in the early reconnaissance phase or actively conducting exploitation or exfiltration. This work addresses the problem of attack stage inference under uncertainty, with a focus on robustness to out-of-distribution (OOD) inputs. We propose a classification approach based on Evidential Deep Learning (EDL), which models predictive uncertainty by outputting parameters of a Dirichlet distribution over possible stages. This allows the system not only to predict the most likely stage of an attack but also to indicate when it is uncertain or the input lies outside the training distribution. Preliminary experiments in a simulated environment demonstrate that the proposed model can accurately infer the stage of an attack with calibrated confidence while effectively detecting OOD inputs, which may indicate changes in the attackers' tactics. These results support the feasibility of deploying uncertainty-aware models for staged threat detection in dynamic and adversarial environments.