Machine learning has significantly advanced encrypted traffic classification, but deployed models often fail due to shortcut features - spurious correlations learned during training that do not generalize to real-world environments.
ShortcutCatcher is a model-agnostic framework designed to automatically detect and mitigate these shortcuts using explainable AI techniques.
conda create -n ag python=3.10
pip install -U uv
python -m uv pip install autogluon==1.4
uv pip install autogluon.tabular[tabicl]
uv pip install matplotlib
Running example:
python -u main.py \
--root_path dataset \
--experiment per-flow \
--dataset app53-time-s2 \
--description soft_window \
--model_name RandomForestGini \
--noise_type removal \
--min_bound 0.0 \
--importance_type default \
--window 5 \
--rounds 100
For the paper tables and figures, see how_to_reproduce.md for the full batch commands.
ShortcutCatcher contrasts model behavior across two datasets:
- A training dataset used for model learning
- A verification dataset representing a different scenario but sharing the same feature schema
By analyzing discrepancies in feature importance across these datasets, the framework identifies features that act as shortcuts and are unlikely to hold in deployment.
ShortcutCatcher operates in a closed loop:
- Train a model on the training dataset
- Generate feature explanations (e.g., via XAI methods)
- Compare feature relevance across training and verification scenarios
- Detect unstable or spurious features
- Iteratively remove or mitigate these features
- Retrain and reevaluate the model
- ✅ Automated detection of shortcut features
- ✅ Model-agnostic design (compatible with various ML architectures)
- ✅ Improved cross-scenario generalization (up to 3× over standard training)
- ✅ Identification of hidden dataset artifacts affecting performance
- ✅ Realistic evaluation of encrypted traffic classification tasks