:
1 2022 Machine Learning for Health (ML4H) 2022

# Algorithmic bias in machine learning based delirium prediction

Sandhya Tripathi<EMAIL_ADDRESS>
Bradley A Fritz<EMAIL_ADDRESS>
Michael S Avidan<EMAIL_ADDRESS>
Department of Anesthesiology  Washington University in St Louis  MO  USA
Yixin Chen<EMAIL_ADDRESS>
Department of Computer Science and Engineering  Washington University in St
Louis  St Louis  MO  USA  Christopher R King<EMAIL_ADDRESS>
Department of Anesthesiology  Washington University in St Louis  MO  USA

###### Abstract

Although prediction models for delirium, a commonly occurring condition during
general hospitalization or post-surgery, have not gained huge popularity,
their algorithmic bias evaluation is crucial due to the existing association
between social determinants of health and delirium risk. In this context,
using MIMIC-III and another academic hospital dataset, we present some initial
experimental evidence showing how sociodemographic features such as sex and
race can impact the model performance across subgroups. With this work, our
intent is to initiate a discussion about the intersectionality effects of old
age, race and socioeconomic factors on the early-stage detection and
prevention of delirium using ML.

###### keywords:

algorithmic bias, delirium risk

##  1 Introduction

Delirium, a commonly under-diagnosed and under-treated condition during
hospitalization, is characterized by an acute change in awareness, attention,
and cognition arising alongside serious illness (American Psychiatric
Association et al., 2013). Delirium-associated risk factors, symptoms and
implications in the elderly population have been widely studied (Inouye et
al., 2014). The delirium prevalence rate can be more than $20\%$ for high-risk
post-surgical patients (Wilson et al., 2020) with up to $50\%$ chances after
hip replacement. In addition to the impact on patients, the estimated
additional healthcare costs associated with delirium exceed $\$44,291$ per
patient per year (Gou et al., 2021).

Machine learning (ML) based delirium prediction can have several applications
in clinical practice. First, patients at higher risk of delirium can be
prioritized to receive preventative treatments and safety measures to mitigate
harmful behaviors should delirium develop. Accurate targeting is necessary
because the most effective treatments are labor-intensive and require high
levels of nursing supervision, making them not easily scaled to all at-risk
patients. There are also tradeoffs in patient quality of life with excessively
aggressive delirium prevention. For example, some anti-anxiety medications
(benzodiazepines) and sleep aides are thought to increase delirium risk, but
unnecessarily removing these medications from lower-risk patients deprives
them of their beneficial effects.

An important concern while using ML for clinical decision support is
algorithmic bias, defined as an ML model’s capability to reinforce existing
biases against certain population groups (such as certain categories of age,
sex, and race commonly referred to as protected or sensitive attributes). With
a long-standing discussion on the relation between social determinants of
health (SDOH) and delirium risk (Khan et al., 2016; Dicks et al., 2021), there
is a pressing need to evaluate delirium prediction models for algorithmic bias
and make sure that the models do not increase the disparities in delirium
prevention. This evaluation for algorithmic bias becomes more crucial because
some of the currently available ML-based delirium prediction models use
protected attributes such as race, insurance status, and language when these
features could be the root cause of dissimilarities during delirium
assessment. Also, as delirium prediction can also be used to enrich patient
samples in research of novel biomarkers, prophylaxis, and treatment, models
which are not accurate in sub-populations would then deprive patients of the
opportunity to participate in such research and produce studies which do not
include all relevant demographics. Presently, the research on the subject of
bias in ML-based delirium prediction has been limited by sample size and
algorithm complexity (Castro et al., 2021).

Contributions: With our study, we 1) present evidence that there could be
algorithmic bias in delirium prediction, 2) recommend users to evaluate the
prediction models for algorithmic bias and check if this bias can be reduced
by simple methods such as removing the protected attribute or accounting for
confounders via Propensity Score Matching.

##  2 Related work

There have been some disagreeing viewpoints among researchers when studying
the effect of demographic factors and SDOH on delirium risk (Khan et al.,
2016; Dicks et al., 2021). As suggested by authors in Arias et al. (2022),
solutions such as the SDOH framework for delirium can capture the challenges
and vulnerabilities unique to older adults. Two separate studies investigating
delirium risk in medical ICU patients in a Dutch region (Wu et al., 2021), and
surgical ICU patients in the US (Khan et al., 2016) found no significant
effect of race or SDOH on delirium risk. In contrast, Dicks et al. (2021)
report evidence of a substantial difference in delirium risk between non-
Hispanic White and Black patients in the (71-80) age groups. Researchers
studying dementia and long-term cognitive impairment (LTCI) discovered that
delirium severity was higher in African American participants (Boltz et al.,
2021). Further, the socioeconomic and clinical risk factors, such as race,
education, hospital type, and delirium duration, were linked to worse LTCI
(Haddad et al., 2020). Health disparities that accentuated during the COVID-19
pandemic require us to be more proactive in terms of screenings for delirium,
a well-known complication of respiratory illness in older adults (O’Hanlon and
Inouye, 2020).

A recent survey by Chua et al. (2021) presents the various ML models developed
in different hospital settings. The penetration of ML in delirium prediction
is still in the initial stages as the models are elementary. ML-based delirium
prediction can be used for patients during post-operative stay (Racine et al.,
2021; Bishara et al., 2022), ICU stay (Hur et al., 2021), general
hospitalization (Ocagli et al., 2021), and in the emergency department (Lee et
al., 2021). While these models achieve satisfactory prediction performance,
they overlook the possibly harmful inclusion of highly predictive attributes
such as insurance status and language in model building as also noted by Chua
et al. (2021).

A small multicenter study by Castro et al. (2021) does a surface-level
evaluation of algorithmic bias in ML-based delirium prediction for COVID-19
patients at the time of admission using an l1-penalized regression algorithm.
Their analyses shows that the model does not suffer from algorithmic bias.
However, as acknowledged by the authors, their labels could be erroneous;
their dataset is comparatively small, and mixes ICU and ward patients in very
different contexts. In contrast, both of our datasets are large, ICU based,
and contain expert-labelled delirium outcomes. Unlike Castro et al. (2021),
our work is the first to show that there could be algorithmic bias in the
predictions of a Random Forest model. Since we performed this analysis before
deploying the above model, we have the option of modifying the model or making
the user aware of these biases.

##  3 Evaluating algorithmic bias

An algorithm suffers from disparate mistreatment (Zafar et al., 2017) if the
algorithm’s performance differs for otherwise similar people at different
levels of a sensitive feature. We use evaluation metrics such as Positive
Predictive Value (PPV), Sensitivity and AUROC to report disparate
mistreatment.

We also evaluate the classifiers on a propensity score matched (PSM) sample so
that the distribution of $\mathbf{X}$ is similar between “treated”
($\mathbf{Z}_{i}=1$) and “untreated” ($\mathbf{Z}_{i}=0$) subjects (Austin,
2011). Here $\mathbf{Z}$ and $\mathbf{X}$ denote sensitive and baseline
covariates (features other than sensitive features). Hence, in order to
compare only individuals with similar covariates, we calculate
$e_{i}=P(\mathbf{Z}_{i}=1|\mathbf{X}_{i})$ and 1-to-1 match individuals with
$z=1$ to individuals with similar $e$ and $z=0$. To estimate the propensity
score, we use a logistic regression model with $\mathbf{Z}$ being regressed on
observed baseline features ($\mathbf{X}$).

##  4 Experimental setup

In our experiments, we consider two datasets: ACTFAST-Epic dataset and MIMIC-
III (Johnson et al., 2016). ACTFAST-Epic dataset is electronic health data
recorded prior to surgical procedures between November 2018 and Feb 2021 at
Barnes-Jewish Hospital (St Louis, MO, USA). Access to the data was approved by
the Human Research Protection Office at Washington University in St Louis, USA
with a waiver of informed consent (IRB number $201607122$). We used CAM-ICU
assessments (Ely et al., 2001) in ACTFAST-Epic dataset and ICD-9 codes in
MIMIC-III to assign delirium labels; the patient demographics are available in
Table LABEL:tab:actfastEpicDems. For MIMIC-III, we followed the pre-processing
and model training steps as in Coombes et al. (2021). After preprocessing,
there were a total of $12699$ patients in ACTFAST-Epic and $48435$ patients in
the MIMIC-III dataset. For each subgroup of a protected attribute, we report
the difference in metric from the mean metric value of the protected attribute
and average it to perform bootstrapped t-testfor comparing the metrics. We
trained Random Forest classifiers on both ACTFAST-Epic and MIMIC-III datasets.
Additional details on datasets and model training are available in Appendix A.

Our experiments are intended to investigate the following hypotheses: 0) there
could be performance differences across subgroups measured using different
metrics and these differences might (not) be eliminated by either 1) removing
the protected variable or 2) accounting for confounding variables (using PSM).

For compactness, we choose AUROC as the main performance metric. Table
LABEL:tab:AUROC-summary presents the gap between the maximum difference from
the average and the minimum difference from the average of the AUROC metric
across the levels/subgroups of a protected attribute. The difference was taken
post-averaging across 150 bootstraps of the dataset.

##  5 Results and Discussion

Table 2 reports the differences in subgroups and that propensity score
matching (PSM) can reduce these differences. We also check the effect of
removing the protected variable and observe that the performance doesn’t
change (Table LABEL:tab:EpicWOSesnPerf in Appendix A).

As shown in Table 2, for ACTFAST-Epic dataset, the performance differences
change modestly on the propensity-matched sample across all the metrics. This
is due to differences in the procedures and comorbidities included in the
propensity score. Notably, there are no other SDOHs included in the propensity
score which could be acting as surrogates for the race variable. Differential
missing data across subgroups could also contribute to the drop in apparent
differential mistreatment when propensity matching, as matching conditional on
data may exclude some patients with significant gaps in their records (and
worse classifier performance). For ACTFAST-Epic, since removing the sensitive
variable from the training dataset didn’t reduce the performance differences,
the sensitive attributes could still affect the outcome via other confounding
variables. As an example, height can be a strong confounder in the absence of
sex. We also measured the False negative rate (FNR) and False positive rate
(FPR) of the classifiers learnt on ACTFAST-Epic dataset. We noticed that there
are higher discrepancies in the FNR values (negative value implies advantage
as lower mistakes) meaning that the patients who were delirious and could have
benefitted from the intervention were missed. Full details on FNR and FPR are
presented in Table LABEL:tab:Epic_mistakes, Appendix A. We can also obtain
confidence intervals on AUROC using the bootstraps as reported in Table
LABEL:tab:Epic-auroc-std for the ACTFAST-Epic dataset.

Table 2 reports that adding the sensitive variable to the feature set (model
M2) can change the results in terms of subgroups’ performance and hence
answers to hypothesis 2. We only included AUROC in the main paper as the
others showed a similar pattern of change in performance; other metrics are
included in Table LABEL:tab:MIMIC-withOutsensM1 and LABEL:tab:MIMIC-withsensM2
in Appendix A.

Unlike some of the existing delirium prediction models, Coombes et al. (2021)
do not include the protected variables in the model training on MIMIC-III
dataset. The subgroup performance of a Random Forest model (M1) on such a
feature set is presented in columns 4 and 5 of Table 2. To demonstrate the
adverse effect of including the protected attributes, we also trained another
model (M2) that used protected attributes in addition to the predictors of
(M1) and report the corresponding AUROC in the last two columns of Table 2.
Although the overall performance of the two models is the same, the
differences between the subgroup and group average of (M2) are statistically
significant.

As a summarizing result for the two hypotheses mentioned earlier, we are
interested in two kinds of comparisons from Table LABEL:tab:AUROC-summary. For
a fixed model, we expect that the after-matching values are the same or
smaller. Across models, we compare the before-matching/after-matching values
to see if removing the protected attribute brings any improvement. Matching
reduces the discrepancy in most cases. Removing the protected variable could
be helpful in some cases as seen in the case of MIMIC-III. We would like to
caution that even though these summaries are easier, they lose information on
which subgroup is affected. They could also be affected by the skewness in the
subgroup sample size.

Our work is limited by the use of only the top two or three commonly occurring
subgroups in protected features such as sex and race. To some extent, delirium
assessment has a subjective part to it, so extending the protected attribute
set to include language and insurance for evaluation of the models can provide
insights helpful in understanding the intersectionality of various SDOH in the
context of delirium prediction. Another necessary step is the evaluation of
the delirium prediction models to reflect the interaction of two protected
attributes.

From a broader societal perspective, these studies would benefit the areas
where SDOHs affect data used in medical risk prediction models. Designing
minimally biased EHR-based risk stratification for delirium prediction can
guide personalized post-operative delirium prevention and reduction strategies
without disadvantaging certain socioeconomic groups.

tab:actfastEpicDems  Patient groups Total Delirium pos Delirium neg Male 1460
662 798 Female 2352 1125 1227 Black 971 524 447 White 2747 1220 1527 Race NA
94 43 51 Age [14 - 52) 1247 546 701 Age [52 - 67) 1278 602 676 Age [67 - 102]
1287 639 648 Total 3812 1787 2025

Table 1: Protected attribute distribution in ACTFAST-Epic test dataset.
Delirium labels were extracted using the CAM-ICU tool.

tab:PerfDiff [][c]  Patient group ACTFAST-Epic Grp - Avg PPV Grp - Avg PPV
(PSM^) Grp-Avg Sens Grp -Avg Sens (PSM) Grp - Avg AUROC Grp - Avg AUROC (PSM)
Male -0.01* -0.02* -0.02* 0.0 -0.01* -0.01* Female 0.01* 0.02* 0.02* 0.0 0.01*
0.01* White -0.04* -0.02* -0.04* 0.05* -0.01* 0.02* Black 0.03* 0.02* 0.02*
-0.05* 0.01* -0.02* Age [14 - 52) -0.02* -0.04* 0.05* 0.03* 0.02* 0.01* Age
[52 - 67) 0.0 [0.04*, -0.02*] -0.01* [-0.03*, 0.01*] -0.01* [-0.01*, ’0.0’]
Age [67 - 102] 0.02* 0.02* -0.04* -0.01* -0.01* 0.0 Overall Avg 0.76 (PPV) -
0.77 (Sens) - 0.85 (AUROC) - [][c]  Patient groups MIMIC-III Delirium Pos
Delirium Neg Grp - Avg AUROC M1 Grp - Avg AUROC M1 (PSM) Grp - Avg AUROC M2
Grp - Avg AUROC M2 (PSM) Male 565 6066 0.01* 0.01* 0.01* 0.02* Female 413 4896
-0.01* -0.01* -0.01* -0.02* Hispanic 27 338 0.0 [0.01*, -0.0] 0.02* [0.03*,
0.01*] White 845 9411 0.0 [-0.0*, 0.0] -0.01* [-0.01*, -0.01*] Black 106 1213
0.01* [0.0*, -0.01*] -0.01* [0.01*, -0.03*] Age [18 - 57) 269 3709 -0.01* 0.0
-0.01* 0.01* Age [57 - 74) 296 3685 0.01* [0.0, 0.0] 0.0 [-0.01*, -0.01*] Age
[74 - 90] 413 3568 0.01* 0.0 0.01* 0.01* Overall Avg 978 10962 0.87 (AUROC) -
0.87 (AUROC) -

Table 2: Performance difference from the group average for (a) the ACTFAST-
Epic dataset and (b) MIMIC-III dataset using a Random Forest classifier on a
bootstrapped (150 samples) test set. ^Multiple values imply results when
matched to patients in each group in this category in the same order as in
column ‘Patient group’. Subtable (b) compares AUROC for M1 (without protected
attributes) and M2 (with protected attributes) models along with reporting the
distribution of protected attributes conditional on the outcome value in the
test set of MIMIC-III. Statistical significance at $\alpha=0.05$ is symbolized
by *.

tab:AUROC-summary  Max discrepancy in AUROC across subgroups for each
protected attribute Protected attribute Model with protected attribute Model
without protected attribute Before matching After Matching Before Matching
After Matching MIMIC-III dataset Race 0.03 0.04 0.01 0.01 Sex 0.02 0.04 0.02
0.02 Age 0.02 0.02 0.02 0 ACTFAST-Epic dataset Race 0.02 0.04 0.02 0.02 Sex
0.02 0.02 0.02 0.02 Age 0.03 0.01 0.04 0

Table 3: Summarizing metric (on AUROC) between models trained on datasets with
and without protected attributes with further comparison using propensity
matching.

##  6 Conclusion

We present our initial investigations to explicitly evaluate ML-based delirium
prediction models for algorithmic bias. Non-representative dataset samples and
biases during manual assessment (Yang et al. (2008); Jones et al. (2006)) add
layers of complexity to learning effectively. Therefore, we should be aware of
the challenges before (evaluating subgroup performance), during (checking
model applicability) and after (making the user aware of the model’s pros and
cons in terms of bias and applicability) the deployment of prediction models.

## References

  * American Psychiatric Association et al. (2013) DS American Psychiatric Association, American Psychiatric Association, et al.  _Diagnostic and statistical manual of mental disorders: DSM-5_ , volume 5.  American psychiatric association Washington, DC, 2013.
  * Arias et al. (2022) Franchesca Arias, Margarita Alegria, Amy J Kind, Richard N Jones, Thomas G Travison, Edward R Marcantonio, Eva M Schmitt, Tamara G Fong, and Sharon K Inouye.  A framework of social determinants of health for delirium tailored to older adults.  _Journal of the American Geriatrics Society_ , 70(1):235–242, 2022.
  * Austin (2011) Peter C Austin.  An introduction to propensity score methods for reducing the effects of confounding in observational studies.  _Multivariate behavioral research_ , 46(3):399–424, 2011.
  * Bishara et al. (2022) Andrew Bishara, Catherine Chiu, Elizabeth L Whitlock, Vanja C Douglas, Sei Lee, Atul J Butte, Jacqueline M Leung, and Anne L Donovan.  Postoperative delirium prediction using machine learning models and preoperative electronic health record data.  _BMC anesthesiology_ , 22(1):1–12, 2022.
  * Boltz et al. (2021) Marie Boltz, Rhonda BeLue, Barbara Resnick, Ashley Kuzmik, Elizabeth Galik, Joanne R Jones, Rachel Arendacs, Liron Sinvani, Jacqueline Mogle, and James E Galvin.  Disparities in physical and psychological symptoms in hospitalized african american and white persons with dementia.  _Journal of aging and health_ , 33(5-6):340–349, 2021.
  * Castro et al. (2021) Victor M Castro, Chana A Sacks, Roy H Perlis, and Thomas H McCoy.  Development and external validation of a delirium prediction model for hospitalized patients with coronavirus disease 2019.  _Journal of the Academy of Consultation-liaison Psychiatry_ , 62(3):298–308, 2021.
  * Chua et al. (2021) SJ Chua, Scott Wrigley, Casey Hair, and Ramesh Sahathevan.  Prediction of delirium using data mining: A systematic review.  _Journal of Clinical Neuroscience_ , 91:288–298, 2021.
  * Coombes et al. (2021) Caitlin E Coombes, Kevin R Coombes, and Naleef Fareed.  A novel model to label delirium in an intensive care unit from clinician actions.  _BMC medical informatics and decision making_ , 21(1):1–12, 2021.
  * Dicks et al. (2021) Robert Dicks, Jimmy Choi, Christine Waszynski, Kadesha Collins-Fletcher, Beth Taylor, Catherine Martinez, and David O’Sullivan.  Health disparities in delirium.  _Innovation in aging_ , 5(Suppl 1):147—147, January 2021.  ISSN 2399-5300.  URL https://europepmc.org/articles/PMC8682256.
  * Ely et al. (2001) E Wesley Ely, Richard Margolin, Joseph Francis, Lisa May, Brenda Truman, Robert Dittus, Theodore Speroff, Shiva Gautam, Gordon R Bernard, and Sharon K Inouye.  Evaluation of delirium in critically ill patients: validation of the confusion assessment method for the intensive care unit (cam-icu).  _Critical care medicine_ , 29(7):1370–1379, 2001\.
  * Gou et al. (2021) Ray Yun Gou, Tammy T Hshieh, Edward R Marcantonio, Zara Cooper, Richard N Jones, Thomas G Travison, Tamara G Fong, Ayesha Abdeen, Jeffrey Lange, Brandon Earp, et al.  One-year medicare costs associated with delirium in older patients undergoing major elective surgery.  _JAMA surgery_ , 156(5):430–442, 2021.
  * Haddad et al. (2020) Diane N Haddad, Matthew F Mart, Li Wang, Christopher J Lindsell, Rameela Raman, Mina F Nordness, Kenneth W Sharp, Pratik P Pandharipande, Timothy D Girard, E Wesley Ely, et al.  Socioeconomic factors and intensive care unit-related cognitive impairment.  _Annals of surgery_ , 272(4):596, 2020.
  * Hur et al. (2021) Sujeong Hur, Ryoung-Eun Ko, Junsang Yoo, Juhyung Ha, Won Chul Cha, Chi Ryang Chung, et al.  A machine learning–based algorithm for the prediction of intensive care unit delirium (pride): Retrospective study.  _JMIR medical informatics_ , 9(7):e23401, 2021\.
  * Inouye et al. (2014) Sharon K Inouye, Rudi GJ Westendorp, and Jane S Saczynski.  Delirium in elderly people.  _The Lancet_ , 383(9920):911–922, 2014.
  * Johnson et al. (2016) Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo A Celi, and Roger G Mark.  Mimic-III, a freely accessible critical care database.  _Scientific data_ , 3(1):1–9, 2016.
  * Jones et al. (2006) Richard N Jones, Frances M Yang, Ying Zhang, Dan K Kiely, Edward R Marcantonio, and Sharon K Inouye.  Does educational attainment contribute to risk for delirium? a potential role for cognitive reserve.  _The Journals of Gerontology Series A: Biological Sciences and Medical Sciences_ , 61(12):1307–1311, 2006.
  * Khan et al. (2016) Babar A Khan, Anthony Perkins, Siu L Hui, Sujuan Gao, Noll L Campbell, Mark O Farber, and Malaz A Boustani.  Relationship between african-american race and delirium in the intensive care unit.  _Critical care medicine_ , 44(9):1727, 2016.
  * Lee et al. (2021) Sangil Lee, Brianna Mueller, W Nick Street, and Ryan Carnahan.  Machine learning algorithm to predict delirium from emergency department data.  _medRxiv_ , 2021.
  * Ocagli et al. (2021) Honoria Ocagli, Daniele Bottigliengo, Giulia Lorenzoni, Danila Azzolina, Aslihan S Acar, Silvia Sorgato, Lucia Stivanello, Mario Degan, and Dario Gregori.  A machine learning approach for investigating delirium as a multifactorial syndrome.  _International Journal of Environmental Research and Public Health_ , 18(13):7105, 2021.
  * O’Hanlon and Inouye (2020) Shane O’Hanlon and Sharon K Inouye.  Delirium: a missing piece in the covid-19 pandemic puzzle.  _Age and ageing_ , 2020.
  * Racine et al. (2021) Annie M Racine, Douglas Tommet, Madeline L D’Aquila, Tamara G Fong, Yun Gou, Patricia A Tabloski, Eran D Metzger, Tammy T Hshieh, Eva M Schmitt, Sarinnapha M Vasunilashorn, et al.  Machine learning to develop and internally validate a predictive model for post-operative delirium in a prospective, observational clinical cohort study of older surgical patients.  _Journal of general internal medicine_ , 36(2):265–273, 2021.
  * Wilson et al. (2020) Jo Ellen Wilson, Matthew F Mart, Colm Cunningham, Yahya Shehabi, Timothy D Girard, Alasdair MJ MacLullich, Arjen JC Slooter, and E Ely.  Delirium.  _Nature Reviews Disease Primers_ , 6(1):1–26, 2020.
  * Wu et al. (2021) Ting-Ting Wu, Marieke Zegers, Rens Kooken, John L Griffith, Beth E Molnar, John W Devlin, and Mark van den Boogaard.  Social determinants of health and delirium occurrence and duration in critically ill adults.  _Critical Care Explorations_ , 3(9), 2021.
  * Yang et al. (2008) Frances M Yang, Sharon K Inouye, Michael A Fearing, Dan K Kiely, Edward R Marcantonio, and Richard N Jones.  Participation in activity and risk for incident delirium.  _Journal of the American Geriatrics Society_ , 56(8):1479–1484, 2008.
  * Youden (1950) William J Youden.  Index for rating diagnostic tests.  _Cancer_ , 3(1):32–35, 1950.
  * Zafar et al. (2017) Muhammad Bilal Zafar, Isabel Valera, Manuel Gomez Rodriguez, and Krishna P Gummadi.  Fairness beyond disparate treatment & disparate impact: Learning classification without disparate mistreatment.  In _Proceedings of the 26th international conference on world wide web_ , pages 1171–1180, 2017.

##  Appendix A Additional experimental details

We consider two datasets: 1) Electronic health data (ACTFAST-Epic dataset)
recorded prior to surgical procedures between November 2018 and Feb 2021
Barnes-Jewish Hospital (St Louis, MO, USA). Access to the data was approved by
the Human Research Protection Office at Washington University in St Louis, USA
with a waiver of informed consent (IRB number $201607122$) 2) MIMIC-III
dataset (Johnson et al., 2016). For the ACTFAST-Epic dataset, CAM-ICU
assessments (Ely et al., 2001) were used to assign delirium labels and data
recorded during preop assessment was used as features. A complete list of
ACTFAST-Epic features used to train the model included a bag of word
representations of preop notes (500 dimensional) and preop records (185
dimensional) as presented in Table LABEL:tab:EpicFeatures. For MIMIC-III,
delirium labels were extracted using the ICD-9 codes and all the pre-
processing and model training steps were performed as by Coombes et al.
(2021). After preprocessing, there were a total of $12699$ patients in
ACTFAST-Epic and $48435$ patients in the MIMIC-III dataset.

We trained Logistic Regression, Decision Tree, Random Forest, Gradient Boosted
Trees, and DNN models to learn the classifiers on both ACTFAST-Epic and MIMIC-
III. We chose Random Forest for reporting as it performed the best in the
group of algorithms. The same analysis can be conducted for other algorithms
too. Each dataset was split into the train (70%) and test (30%) subsets while
maintaining the same delirium outcome rate in the splits. We performed the
computation on 150 bootstrapped samples of the test split in each dataset. For
each subgroup (PSM subgroup), we report the difference from the average metric
under consideration. While computing PPV and Sensitivity, we used Youden’s
index (Youden, 1950) to identify the classifier’s decision boundary. The
outcome distribution in the two datasets is different as ACTFAST-Epic is
balanced while MIMIC-III is not. Hence, for the MIMIC-III dataset, we only
report AUROC in the main paper. The calibration plots for the models on which
the classifiers were trained are presented in Figure
LABEL:fig:calibrationplots.

We filtered the subgroups for Race/Ethnicity in the datasets based on the data
availability in each subgroup and hence reported the subgroup performance for
Black, White, and Hispanic (in MIMIC-III) and “not available/unknown”
(ACTFAST-Epic) groups. We required that at least 100 samples be available
after propensity score matching, to compute the final metric.

### Code availability

We provide the python code that was used to obtain the results for both
datasets at https://github.com/sandhyat/AlgorithmicBias_Delirium_ML4H22022.
ACTFAST-Epic dataset is not publicly available. For MIMIC-III, we also provide
psql queries used to extract the initial dataset following Coombes et al.
(2021).

tab:EpicFeatures  Demographics Comorbidities Comorbidities Preop Lab Test
Preop Lab Test Sex Anemia pastTransplant POIKILOCYTOSIS TROPONIN I emergency
Asthma LVEF HCG,URINE, POC HEMOGLOBIN A1C SurgService_Name CAD CancerCurrent
CLARITY, URINE HEMATOCRIT ETHNICITY CHF PacemakerICD COLOR, URINE PCO2,
ARTERIAL RACE CKD Stroke LEUKOCYTE ESTERASE, URINE LIPASE age COPD PNA RED
BLOOD CELLS, URINE MEAN PLATELET VOLUME Temp CancerHx delirium_history
ANISOCYTOSIS PCO2, VENOUS TOBACCO_USE Chronic_Pain MentalHistory_adhd HYALINE
CAST IRON BINDING CAPACITY, TOTAL HEIGHT_IN_INCHES Cirrhosis
MentalHistory_other COVID-19 CORONAVIRUS RNA SPECIFIC GRAVITY, URINE
WEIGHT_IN_KG DVT_PE DVT URINE NITRITE N-TERMINAL PRO B-TYPE NATURETIC PEPTIDE
Dementia_Mild CognitiveImpairment PE GLUCOSE, URINE, QUALITATIVE ASPARTATE
AMINOTRANSFERASE Clinician Assessment Diabetes cancerStatus WHITE BLOOD CELLS,
URINE MAGNESIUM Dialysis CHF_class BACTERIA, URINE BASOPHIL ABSOLUTE wheezing
GERD DHF EPITHELIAL CELLS, SQUAMOUS, URINE WHITE BLOOD CELL COUNT Pulse HTN
DM_Type1 HEPATITIS C AB RED CELL DISTRIBUTION WIDTH SD Resp LiverDisease ESRD
PROTEIN, URINE QUALITATIVE PHOSPHORUS, PLASMA Short Blessed Total Score MI MR
BILIRUBIN, DIRECT VANCOMYCIN TROUGH History of Falling MentalHistory _ anxiety
DiastolicFunction IMMATURE GRANULOCYTE ABSOLUTE TRIGLYCERIDES Pain Score
MentalHistory _bipolar DyspneaFreq FEV1Percent BILIRUBIN, TOTAL poorDentition
MentalHistory _depression IONIZED CALCIUM CHLORIDE ASA MentalHistory
_schizophrenia In-hospital attributes LYMPHOCYTE ABSOLUTE FIBRINOGEN, CLAUSS
Gait/Transferring MS URINE UROBILINOGEN PO2, VENOUS NutritionRisk OSA
preop_del FERRITIN C REACTIVE PROTEIN Barthel index score OtherRhythm
preop_ICU HEMOGLOBIN RED BLOOD CELL AD8 Dementia Score OutpatientInsulin
preop_vent_b CALCIUM LACTATE StopBANG PAD PlannedAnesthesia NEUTROPHIL
ABSOLUTE THYROID STIMULATING HORMONE Mental Status PHTN pre_aki_status ALBUMIN
URINE BLOOD FunctionalCapacity Pacemaker plannedDispo CHOLESTEROL, TOTAL SO2
(MEAS) ARTERIAL morse ReplacedValve preop_los URINE KETONES inr SpO2 Seizures
BOW_NA (Absence of procedure text) EOSINOPHIL ABSOLUTE PROTEIN, TOTAL, PLASMA
SBP Thyroid_disease noROS blood urea nitrogen PH, VENOUS DBP activeInfection
hasBlock (kind of anesthesia) CREATININE MONOCYTE ABSOLUTE Ambulatory Aids AR
IRON, TOTAL CHOLESTEROL, HDL Secondary Diagnosis on_o2 Preop Lab Test PLATELET
ALANINE AMINOTRANSFERASE pancreatitis ACTIVATED PARTIAL THROMBOPLASTIN TIME
POTASSIUM, PLASMA Preop Lab Test pastDialysis RBC MORPHOLOGY NUCLEATED RBC
ABS, AUTO PH, URINE Coombs AF ABORH PAT INTERP GLUCOSE ALKALINE PHOSPHATASE
CREATINE KINASE, TOTAL TR ERYTHROCYTE SEDIMENTATION RATE CO2, TOTAL SODIUM
VITAMIN D, 25-HYDROXY AS TRANSFERRIN SATURATION PO2 ARTERIAL POC PH, ARTERIAL

Table 4: Predictors used to train the Random forest model in ACTFAST-Epic
dataset.

fig:calibrationplots [ACTFAST-Epic dataset (Table 2a)] [MIMIC-III dataset
Model M1 (Table 2b)] [MIMIC-III dataset Model M2 (Table 2b)]

Figure 1: Calibration plot for RF models reported in the main paper.

tab:EpicWOSesnPerf  Patient group ACTFAST-Epic Grp - Avg PPV Grp - Avg PPV
(PSM^) Grp-Avg Sens Grp -Avg Sens (PSM) Grp - Avg AUROC Grp - Avg AUROC (PSM)
Male -0.01* -0.03* -0.01* 0.02* -0.01* -0.01* Female 0.01* 0.03* 0.01* -0.02*
0.01* 0.01* White -0.04* [’-0.02*’] -0.04* [’0.03*’] -0.01* [’0.01*’] Black
0.02* [’0.02*’] 0.02* [’-0.03*’] 0.01* [’-0.01*’] Age [14 - 52) -0.02* -0.03*
0.06* 0.03* 0.02* 0.0 Age [52 - 67) 0.0 [’0.03*’, ’-0.03*’] 0.0 [’-0.03*’,
’0.04*’] -0.01* [0.0, ’0.0*] Age [67 - 102] 0.02* 0.03* -0.06* -0.04* -0.02*
-0.0* Overall Avg 0.76 (PPV) - 0.77 (Sensitivity) - 0.85 (AUROC) -

Table 5: Performance difference from the group average for the ACTFAST-Epic
dataset using a Random Forest classifier on a bootstrapped (150 samples) test
set when the sensitive features (age, sex, and race) were not included in the
training dataset. ^Multiple values imply results when matched to patients in
each group in this category in the same order as in column ‘Patient group’.

tab:Epic-auroc-std  Model with protected attribute Models without protected
attributes Patient group Grp - Avg AUROC Grp - Avg AUROC (PSM) Grp - Avg AUROC
Grp - Avg AUROC (PSM) Male -0.01*+-0.0106 -0.01*+-0.003 -0.01*+-0.0101
-0.01*+-0.003 Female 0.01*+-0.0106 0.01*+-0.003 0.01*+-0.0101 0.01*+-0.003
White -0.01*+-0.0309 [’0.02*+-0.003’] -0.01*+-0.0306 [’0.01*+-0.003’] Black
0.01*+-0.0171 [’-0.02*+-0.003’] 0.01*+-0.0169 [’-0.01*+-0.003’] Age [14 - 52)
0.02*+-0.0148 0.01*+-0.003 0.02*+-0.0148 -0.0+-0.003 Age [52 - 67)
-0.01*+-0.0134 [’-0.01*+-0.003’, ’0.0+-0.003’] -0.01*+-0.0136 [’0.0+-0.003’,
’0.0*+-0.003’] Age [67 - 102] -0.01*+-0.0142 -0.0+-0.003 -0.02*+-0.0143
-0.0*+-0.003 Overall Avg 0.85 (AUROC) - 0.85 (AUROC) -

Table 6: ACTFAST-Epic: AUROC with standard deviations on 150 bootstrapped
samples.

tab:Epic_mistakes  Model with protected attributes Model without protected
attributes Patient group ACTFAST-Epic Grp- Avg FNR Grp - Avg FNR (PSM) Grp -
Avg FPR Grp - Avg FPR (PSM) Grp- Avg FNR Grp - Avg FNR (PSM) Grp - Avg FPR Grp
- Avg FPR (PSM) Male 0.02* -0.0 -0.01* 0.01* 0.01* -0.02* 0.0 0.02* Female
-0.02* 0.0 0.01* -0.01* -0.01* 0.02* 0.0 -0.02* White 0.04* [-0.05*] 0.01*
[0.01*] 0.04* [-0.03*] 0.0 [0.02*] Black -0.02* [0.05*] -0.0 [-0.01*] -0.02*
[0.03*] 0.0 [-0.02*] Age [14 - 52) -0.05* -0.03* 0.01* 0.0* -0.06* -0.03*
0.02* 0.01* Age [52 - 67) 0.01* [0.03*, -0.01*] 0.0 [-0.0*, 0.0] 0.0 [0.03*,
-0.04*] 0.0 [-0.01*, 0.02*] Age [67 - 102] 0.04* 0.01* -0.01* -0.0 0.06* 0.04*
-0.02* -0.02* Overall Avg 0.23 (FNR) - 0.22 (FPR) - 0.23 (FNR) - 0.22 (FPR) -

Table 7: ACTFAST-Epic: Comparison of FNR, FPR results when the protected
attributes were included and not included in the training dataset.

tab:MIMIC-withOutsensM1  Patient group MIMIC-III Grp - Avg PPV Grp -Avg PPV
(PSM) Grp - Avg Sens Grp - Avg Sens (PSM) Grp - Avg AUROC Grp- Avg AUROC (PSM)
Male 0.02* 0.02* 0.03* 0.03* 0.01* 0.01* Female -0.02* -0.02* -0.03* -0.03*
-0.01* -0.01* Hispanic 0.0 [’0.01*’, ’-0.02*’] -0.01 [’-0.04*’, -0.0] 0.0
[’0.01*’, -0.0] White 0.0 [-0.0, ’0.02*’] -0.02* [’-0.01*’, 0.0] 0.0 [’-0.0*’,
0.0] Black 0.0 [0.0, ’-0.01*’] 0.03* [’0.01*’, ’0.04*’] 0.01* [’0.0*’,
’-0.01*’] Age[18 - 57) -0.04* -0.03* -0.01* -0.01* -0.01* 0.0 Age [57 - 74)
0.0 [’0.03*’, ’-0.02*’] -0.02* [’0.01*’, 0.0] 0.01* [0.0, 0.0] Age [74 - 90]
0.05* 0.02* 0.03* 0.0 0.01* 0.0 Overall Avg 0.26 (PPV) - 0.79 (Sensitivity) -
0.87 (AUROC) -

Table 8: Performance difference from the group average for the MIMIC-III
dataset using a Random Forest classifier on a bootstrapped (150 samples) test
set when the sensitive features (age, sex, and race) were not included in the
training dataset.

tab:MIMIC-withsensM2  Patient group MIMIC-III Grp - Avg PPV Grp -Avg PPV (PSM)
Grp - Avg Sens Grp - Avg Sens (PSM) Grp - Avg AUROC Grp- Avg AUROC (PSM) Male
0.02* 0.02* 0.02* 0.02* 0.01* 0.02* Female -0.02* -0.02* -0.02* -0.02* -0.01*
-0.02* Hispanic -0.02* [0.0, ’-0.04*’] 0.0 [’0.04*’, ’-0.02*’] 0.02* [’0.03*’,
’0.01*’] White 0.02* [’0.01*’, ’0.04*’] -0.04* [’-0.04*’, ’0.02*’] -0.01*
[’-0.01*’, ’-0.01*’] Black 0.0 [’-0.01*’, -0.0] 0.04* [’0.04*’, ’-0.04*’]
-0.01* [’0.01*’, ’-0.03*’] Age[18 - 57) -0.04* -0.02* 0.01* 0.03* -0.01* 0.01*
Age [57 - 74) -0.01* [’0.02*’, ’-0.03*’] -0.03* [’-0.03*’, ’-0.02*’] 0.0
[’-0.01*’, ’-0.01*’] Age [74 - 90] 0.04* 0.03* 0.02* 0.02* 0.01* 0.01* Overall
Avg 0.26 (PPV) - 0.8 (Sens) - 0.87 (AUROC) -

Table 9: Performance difference from the group average for the MIMIC-III
dataset using a Random Forest classifier on a bootstrapped (150 samples) test
set when the sensitive features (age, sex, and race) were included in the
training dataset.


# FedEdge AI-TC: A Semi-supervised Traffic Classification Method based on
Trusted Federated Deep Learning for Mobile Edge Computing

Pan Wang, Zeyi Li, Mengyi Fu, Zixuan Wang, Ze Zhang and Minyao Liu Pan Wang
was with the School of Modern Posts, Nanjing University of
Post&Telecommunications, Nanjing, Jiangsu, 210003 China e-mail:
(wangpan@njupt.edu.cn).Manuscript received July 20, 2023.

###### Abstract

As a typical entity of MEC (Mobile Edge Computing), 5G CPE (Customer Premise
Equipment) has proven to be a promising alternative to traditional HGU (Home
Gateway Unit). Network TC (Traffic Classification) is a vital service quality
assurance and security management method for communication networks, which has
become a crucial functional entity in 5G CPE/HGU. In recent years, many
researchers have applied Machine Learning (ML) or Deep Learning (DL) to TC,
namely AI-TC, to improve its performance. However, AI-TC methods face
significant challenges, including high data dependency, exhaustively costly
traffic labeling, and network subscribers’ privacy. Besides, as the AI-TC
carrier, 5G CPE/HGU’s limited computing resources often become the bottleneck
of models for efficient classification. Furthermore, the long-standing problem
of the ”black box” for AI-TC models has always perplexed network operators
regarding the model’s transparency and credibility, i.e., AI model
interpretability. Therefore, how to achieve an efficient and trusted
classification carried on the ”weak computing power” network entity while
protecting user privacy has become the key to ensuring the service quality and
security of the home network. This paper presents the FedEdge AI-TC framework,
a novel AI-TC approach for implementing trusted Federated Learning (FL) based
efficient Network TC in 5G CPE/HGU. First, FedEdge AI-TC effectively protects
the data privacy of network subscribers by proposing an FL based framework of
local training, model parameters iterating, and centralized training. Second,
a semi-supervised TC algorithm based on Variational Auto-Encoder (VAE) and
convolutional neural network (CNN) is designed to reduce data dependence while
keeping the TC accuracy. Finally, XAI-Pruning, an AI model compression method,
combined with the DL model interpretability, is proposed to condense the model
and interpret it globally and locally to achieve light-weighted AI-TC model
deployment while building the trust in their decision of network operators. To
demonstrate the efficiency of the proposed method, we conducted some
experimental evaluations on commonly used public benchmark datasets and real
network datasets. The results show that FedEdge AI-TC can outperform the
benchmarking methods regarding the accuracy and achieve excellent TC
performance of model inference on 5G CPE/HGU with limited computing resources,
which effectively protects the users’ privacy and improve the model’s
credibility.

###### Index Terms:

traffic classification, edge computing, federated learning, variational auto-
encoder, semi-supervised, model interpretability.

##  I Introduction

As a distinct entity of MEC, 5G CPE has gradually become an alternative to
HGU. Network traffic classification has played a crucial role in ensuring
service quality and managing security for home networks. It is a critical
functional element within 5G CPE. It finds extensive applications in QoS
(Quality of Service) / QoE (Quality of Experience) management, network
resource optimization, congestion control, and intrusion detection. With the
popularity of smart homes, many applications such as video surveillance, fire
and smoke detection, smart appliances, VR/AR, and others have emerged
alongside traditional internet services like high-definition videos and online
gaming. These applications impose demanding requirements on the network’s QoS,
including fast and flexible customization of services, real-time
responsiveness, and high reliability. Thus, home networks exhibit four
significant trends: ”Terminals Heterogeneity, Applications Diversity, High
Privacy, and Rapid Evolution.” Traffic classification in home networks, as an
important prerequisite for fine-grained network resource management, has
become one of the crucial security measures for smart homes. As shown in Fig.
1, the 5G CPE/Edge Gateway serves as the ”connection point” between the smart
home and the wide area network. It is crucial for the reliable forwarding of
household application traffic. The AI-TC based on 5G CPE/Edge Gateway is the
key to achieving fine-grained network resource management, QoE assurance, and
intrusion detection in home networks.

Figure 1: The Scenario of AI-TC for MEC.

The development of Network TC has generally gone through three stages. In the
first phase, TC methods were mainly based on port matching or DPI (Deep Packet
Inspection). However, this type of technology quickly became ineffective with
the increasing use of techniques like tunneling, encryption, random ports, and
concerns about security issues such as user privacy breaches. The second stage
primarily leveraged machine learning techniques to extract underlying patterns
of different services/applications/attacks traffic features and achieved TC by
discriminating various applications in the data space. However, such methods
require the extraction of high-quality traffic features as the training inputs
for ML. The extraction and selection of these features heavily rely on the
domain expertise of network specialists and are time-consuming and labor-
intensive. In the third stage, with the rapid development of cloud computing,
big data, especially deep learning, and high-performance computing technology,
feature learning of massive traffic data has become feasible, bringing new
imagination space for improving the TC’s performance. DL has three excellent
features: automatic feature extraction, exploration of deep nonlinear
features, and many classical models in computer vision/image/text/speech that
can be reused. These advantages are all lacking in ML-based TC methods.
Several DL-based TC technologies have been proposed recently, including
CNN/AE/MLP/LSTM/GAN-based methods, which have achieved better classification
performance than ML-TC[1, 2, 3, 4, 5].

However, applying DL technology to smart home network traffic classification
faces three major challenges. Firstly, DL models heavily rely on a large
volume of online behavior data from home users, which raises concerns about
highly sensitive user privacy. Additionally, the collection and labeling of
traffic samples are time-consuming and labor-intensive. Secondly, the 5G
CPE/edge gateway’s limited computing resources often become the bottleneck of
the AI-TC models for efficient classification. Thirdly, the ”black box”
problem of the DL classification model has always perplexed the
trustworthiness of home users/network operators. Therefore, efficiently
achieving trusted classification of home network traffic on a ”weak computing
power” gateway device while protecting user privacy is crucial in ensuring
service quality and security.

Federated Learning is a distributed machine learning technology providing
privacy protection, presenting a novel application paradigm that balances data
privacy protection and sharing computing[6]. FL constructs a global model
based on virtual fusion data through distributed model training among multiple
data sources with local data, without exchanging local data but model
parameters or intermediate results. In recent years, FL has been widely used
in industries with high sensitivity to data privacy, such as finance and
medical care, and has made significant progress[7]. Inspired by this, this
paper proposes FedEdge AI-TC, an AI traffic classification method based on
federated learning of 5G CPE/edge gateway. This method uses the FL framework
to train the traffic classification model of the home network without
uploading home network data to a centralized server but executing local
distributed model training on a 5G CPE/edge gateway. The global traffic
classification model is constructed by exchanging model parameters with the
centralized server while protecting the privacy of home users. In addition,
considering that traffic sample collection and labeling are time-consuming and
labor-intensive, we design a semi-supervised traffic classification algorithm
based on VAE and CNN to reduce dependency on traffic sample data. Finally,
XAI-Pruning, an AI model compression method, combined with the DL model
interpretability, is proposed to condense the model and interpret it globally
and locally to achieve light-weighted AI-TC model deployment while building
the trust in their decision of network operators. To demonstrate the
efficiency of the proposed method, we conducted some experimental evaluations
on commonly used public benchmark datasets and real network datasets. The
results show that FedEdge AI-TC can outperform the benchmarking methods
regarding the accuracy and achieve excellent TC performance of model inference
on 5G CPE/HGU with limited computing resources, which effectively protects the
users’ privacy and improve the model’s credibility. The contributions of this
paper are as follows:

  1. 1.

We propose a 5G CPE traffic classification method FedEdge AI-TC based on
federated learning, which effectively protects the privacy of home user data
by constructing the FL framework of local training, parameter updating, and
centralized training;

  2. 2.

A semi-supervised traffic classification algorithm based on VAE and CNN is
designed to reduce the dependence on traffic sample data;

  3. 3.

A pruning method based on DL model interpretability (XAI-Pruning) is proposed
for model compression, and the model is globally and locally explained to
increase model transparency and credibility;

  4. 4.

Experiments on public and self-built datasets show this method can achieve
high traffic classification accuracy under limited computing resources.

The chapter organization of this paper is as follows: Section I is an overall
introduction; Section II is related research works; Section III presents the
framework for the proposed approach; Section IV describes the FedEdge AI-TC
method; Section V evaluates the proposed method through experiments and
provides a comprehensive discussion of the results.; Section VI concludes the
contributions and outlines potential directions for future research. Table I
below is the list of abbreviations in alphabetical order.

TABLE I: List of abbreviations in alphabetical order Acronym | Explanation
---|---
AE | Auto Encoders
CPE | Customer Premise Equipment
CNN | Convolutional Neural Networks
DL | Deep Learning
DPI | Deep Packet Inspection
FL | Federated Learning
FAM | Flow Attribute Matrix
GAN | Generative Adversarial Network
HGU | Home Gateway Unit
LSTM | Long Short Term Memory
MSE | Mean Square Error
MLP | Multilayer Perceptron
QoE | Quality of Experience
QoS | Quality of Service
SSL | Semi-Supervised Learning
TC | Traffic Classification
VAE | Variational Auto Encoder
XAI | Model Explanation/Interpretability
ML | Machine Learning
GNN | Graph Neural Network

##  II Related works

###  II-A Deep Learning based Traffic Classification

Deep learning, also referred to as deep structured learning or hierarchical
learning, is achieved by acquiring the representation of data. In contrast to
traditional machine learning algorithms, deep learning can automatically
extract features without human intervention, rendering it an ideal approach
for traffic classification. The application of deep learning techniques to
Network TC[8, 9, 10, 11, 12] involves three steps: firstly, characterizing the
input data by defining and designing the model input using data packets, PCAP
files, or traffic statistics vectors as features; secondly, selecting suitable
models and algorithms based on classifier objectives and model
characteristics; finally extracting traffic features automatically through
training a DL-based classifier and associating input data with corresponding
category labels.

Recent research has demonstrated deep learning methods’ superiority in traffic
classification. For example, CNNs[13, 14] are widely used in traffic
classification. They can automatically extract features from raw network
traffic data and have end-to-end learning capabilities during training. In
addition, RNNs and LSTMs[15, 16] are used to process traffic sequence data and
can capture the temporal dependencies in the data. These types of networks are
often used in traffic classification to identify persistent attacks such as
DDoS (Distributed Denial of Service) attacks. GNN has proven to be a novel
information representation method for DL, which has been applied in TC or
IDS[17].

###  II-B Semi-supervised Learning based Traffic Classification

Semi-supervised traffic classification[18, 19, 20, 21] is an approach that
utilizes a small set of labeled data along with a substantial amount of
unlabeled data to distinguish various network traffic types. There are four
primary methods for semi-supervised traffic classification: cluster-based
methods, generative models, GANs (Generative Adversarial Networks), and
discriminative models. Cluster-based methods[22, 23] have low computational
complexity but can be influenced by data distribution and may exhibit
instability in practical usage. Generative model-based methods[24] are
effective for unknown or dynamically changing network applications; however,
they require prior knowledge to select appropriate statistical features and
clustering parameters which might limit the generalization of classification
results. GAN-based methods[25, 26] can enhance dataset diversity and quality,
thereby improving the model’s generalization performance. Nevertheless, GAN
models have complex structures and numerous parameters that pose challenges in
training and render them impractical for deployment on edge devices.
Therefore, this study primarily adopts a discriminative model-based approach
by directly learning the mapping function from feature space to class space.
Model parameters are optimized by minimizing the classification error of
labeled data while incorporating a regularization term for unlabeled data.
Subsequently, the unlabeled data is classified based on predicted results.

Figure 2: The Workflow of the End-to-End Traffic Classification based on
Trusted Federated Learning for Edge Gateway.

###  II-C FL and its applications in Traffic Classification

Federated Learning (FL) is a distributed ML/DL framework that focuses on
decentralized training data, aiming to obtain ML/DL models by distributing the
data across numerous nodes while ensuring privacy and security. FL allows
local clients to retain their data, sharing only the model parameters with a
central server, thereby reducing communication overhead and preserving client
data privacy[7, 27]. Recently, two main approaches have emerged for traffic
classification tasks by combining federated learning with deep learning
techniques. The first approach[28, 29, 30, 31] empowers child nodes to
annotate the data through various means. The second approach[32, 33, 34, 35,
36] involves transforming the model structure and training objectives so that
sub-nodes can train the model using unlabeled data; subsequently, fine-tuning
is performed by the server using labeled data to achieve semi-supervised
traffic classification. We propose a semi-supervised traffic classification
model based on VAE-CNN that incorporates a federated learning paradigm
enabling edge devices’ semi-supervised training of the model.

##  III The Overall Framework

###  III-A The Workflow of FedEdge AI-TC

Smart home networks face four major challenges: terminal heterogeneity,
application diversity, high privacy, and rapid evolution. A network TC system
must continuously learn through long-term iterative optimization to overcome
these challenges. The process follows the full life cycle of federated
learning, as shown in Fig. 2. The edge-side AI-TC classification system
workflow based on federated learning includes _initialization, broadcast,
training, parameter uploading, model aggregation and evaluation, edge
deployment, and model monitoring_. Here are the steps for implementing an AI-
powered network traffic classification system:

  1. 1.

Initialization: Provide the client node with an initialization model for
efficient local/global model training.

  2. 2.

Broadcast: The centralized server broadcasts the initialization model to all
the client nodes like 5G CPE/HGU.

  3. 3.

Local Training: The client node performs feature engineering, model
construction, and local training.

    1. (a)

Feature Engineering: Extract, select, represent, and compress network traffic
features to build an optimal feature subset for the AI-TC classification
system.

    2. (b)

Model Construction and Local Training: This step involves selecting what kind
of learning methods (supervised/semi-supervised/unsupervised/weakly
supervised), training methods (centralized/distributed training, federated
learning), whether to pre-train, whether to use classical models for transfer
learning, to form a local model (Local Model) on the local client node.

  4. 4.

Parameters Uploading: Upload encrypted parameter information obtained from
local training to the centralized server.

  5. 5.

Aggregation: The centralized server performs ’secure aggregation’ on the
parameter information uploaded by each client node (such as using the _FedAvg_
algorithm) and performs global training.

  6. 6.

Model Evaluation: Evaluate the global model obtained from the centralized
server’s global training. If the training process converges, it will enter the
model deployment step; otherwise, it will inform the client node to continue
training and iterative optimization. In addition, the model evaluation also
needs to consider its computational complexity, time complexity, and the
computing resources and time required for training/inference.

  7. 7.

Edge deployment: Deploy the model on the edge or terminal side using the
pull/push/subscribe model deploy method and update strategy. Model compression
and interpretation are two important tasks in this step.

    1. (a)

Model Compression: Compress the inference/classification model small enough to
meet the fast classification under limited computing power.

    2. (b)

XAI(Model Interpretation/Explanation): Solve the ”black box” problem of the
AI-TC model to make the classification model users trust the model.

  8. 8.

Model Monitoring: Monitor the status of the classification system, including
model, and real-time network flow, to report some key issues such as
classification system failures and model degradation.

  9. 9.

Continuous Learning: Initiate iterative optimizations and continuous learning
from initialization to maintain high adaptability, robustness, and reliability
of the classification system.

The following sections of this article will focus on _initialization, model
training, compression, and explanation._

###  III-B The Architecture of FedEdge AI-TC

Figure 3: The Architecture of FedEdge AI-TC.

The overall architecture of FedEdge AI-TC is illustrated in Fig. 3, which is
divided into the client node and central server/central aggregator based on
our previous work[37]. HGU, functioning as the client node, performs local
training and inference classification tasks[38]. On the other hand, the
central aggregator acts as a centralized server responsible for aggregated
training, model evaluation, compression, interpretation, and deployment of
inference models. The workflow can be summarized as follows: Initially, the
central aggregator broadcasts the initial classification model to HGU.
Subsequently, HGU collects real-time packets and performs redundant/invalid
packet filtering, network flow attribute calculation, and normalization to
formulate a FAM (Flow Attribute Matrix), depicted in Fig.4. This matrix is
then utilized for training an initialization model locally. Afterward,
encrypted gradient information, loss, and other parameter details are uploaded
to the central aggregator for aggregate training. Since each HGU exhibits
similar flow characteristics, Horizontal Federated Learning (HFL) is employed
for aggregate training by standard secure aggregation algorithms like FedAvg.
Then the model performance metrics will be evaluated, including accuracy,
precision, recall, and F1-Score for assessing model convergence. If
convergence occurs successfully, the aggregated global model undergoes
compression to meet computing resource constraints of HGU (including
CPU/memory/Flash). Once an available inference model is obtained, XAI-based
methods are applied to provide both global and local explanations of the
model. Finally, the resulting inference model will be distributed across all
HGUs. Otherwise, if convergence fails, the global model will be issued to each
HGU, initiating a new epoch of iterative training and optimization until it
converges.

##  IV The Methodology of FedEdge AI-TC

###  IV-A Initial Model

The initial model plays a crucial role in determining the convergence speed
and performance of federated aggregation training, which serves as the initial
stage of federated learning. While randomly initialized models can be used
within the FL framework, it is essential to construct an initial model that
enhances local/global model training efficiency. Unlike traditional AI domains
like images and text, network traffic classification lacks pre-trained models,
necessitating the development of the initial model by ourselves. In this
study, we utilize three benchmarking datasets including (ISCIX[39], UNSW-
NB15[40], and MIRAGE[41]) as baseline datasets for constructing our initial
model. We adopt CNN as the supervised training algorithm, and further details
about the initial model can be found in our previous work[37].

###  IV-B Federated Semi-Supervised Learning Traffic Classification Method
using VAE+CNN

####  IV-B1 The Introduction of FSSL

There are two types of network traffic data in the FedEdge AI-TC system. The
first type, labeled data, is stored in the central aggregator. The second
type, unlabeled data, is located in the local HGU, i.e., real-time traffic
data. From the perspective of FL, it belongs to the disjoint scenario. This
alignment with the home network scenario arises due to the absence of labels
for real-time traffic forwarded by the HGU, making it impractical to annotate
such data. Conversely, the central aggregator possesses powerful computing
resources and can effectively accomplish this work. Semi-supervised Learning
(SSL) aims at leveraging unlabeled data to enhance model training by learning
classification boundaries within these unlabeled samples and evaluating their
proximity to labeled ones. Consequently, this approach strengthens both the
robustness and generalization ability of models. The advantages of SSL
primarily are in two aspects: (1) enhancing the robustness of TC
classification models; (2) mitigating loss in model generalization caused by
domain differences—for instance, traffic forwarded by different HGU devices
may exhibit non-independent and identically distributed (non-IID)
characteristics. In the subsequent section, we will propose an SSL-based
method for traffic classification using FSSL.

####  IV-B2 Problem Formulation

  1. a)

Basic definitions: $D=[D_{l},D_{u}]$; $D$ refers to the network traffic
dataset, which consists of a labeled dataset $D_{l}$ and an unlabeled dataset
$D_{u}$, $D_{l}\cap D_{u}=\phi$; $M$ and $N$ represent the total number of
records in labeled and unlabeled datasets, respectively;
$L=\\{l_{1},l_{2},\cdots,l_{c}\\}$, refers to the dataset of the traffic
classification label, where $l\in L,0\leq c\leq K$, and $K$ is the total
number of application types; $F=F+F^{\prime}$, is a collection of flow feature
vectors; $F=[F_{1},F_{2},\cdots,F_{M}]$, refers to the set of labeled flow
feature vectors; $F=[F_{1}^{\prime},F_{2}^{\prime},\cdots,F_{N}^{\prime}]$,
refers to the set of unlabeled flow feature vectors;
$F=F+F^{\prime}=\\{f_{l}^{i},f_{u}^{j}\\}=\\{f_{l}^{1},f_{l}^{2},\cdots,f_{l}^{n},f_{u}^{1},f_{u}^{2},\cdots,f_{u}^{m},\\}$,
labeled/unlabeled flow feature vectors are denoted by $f_{l}^{i}$ and
$f_{u}^{j}$, $0\leq i\leq M$, $0\leq j\leq N$.

  2. b)

Definitions related to network traffic:

     * •

flow: A flow is identified by a five-tuple consisting of the source address,
destination address, source port, destination port, and TCP/UDP protocol.
$T=\\{t_{1},t_{2},\cdots,t_{m+n}\\}$, represents a set of flows.

     * •

flow feature: It includes packet-level features, flow-level features, and
statistical features, formally defined as $f=\\{f^{1},f^{2},\cdots,f^{78}\\}$.
$f$ is the flow feature vector, composed of a total of 78 feature sub-items
[13], which consist of the following three types of features:

       * –

Packet-level features: The temporal and spatial features with packets as the
granularity, including packet payload characteristics, packet length-related
features, and time-related features. For example, packet length, inter-arrival
time between packets, etc.

       * –

Flow-level features: The temporal and spatial features of flows, with flows as
the granularity, including flow length, flow duration, number of packets in a
flow, and so on.

       * –

Statistical features: The expectation, variance, maximum value, and minimum
value of the relevant feature.

Table II presents an example collection of traffic features for FedEdge AI-TC.
From the example, it can be observed that specific feature entries involve a
high level of user privacy.

TABLE II: A typical example (partial) of a network traffic feature set Flow Attributes | Definition | Catogory | Description
---|---|---|---
Domain Name | DNS/SNI in TLS | Payload related | domain.com, which is applicable to applications such as HTTP/HTTPS.
TCP slide_win | TCP Slide Window | Packet related | TCP flow control parameters
TLS_handshake | TLS handshake packet information | Payload related | Handshake types, cipher suites, content types, key length, etc.
Total Fwd Pkts | Packet length sequence | Packet related | The sequence of packet lengths in the flow.
Pkt IAT Min | Packet arrival time | Packet related | The sequence of arrival times of packets in the flow.
Flow Len | Flow length correlation | Flow related | The total number of bytes in the flow per unit of time.
Flow Duration | Flow duration | Flow related | The duration of the TCP flow.

Figure 4: The Example of L_FAM.

  3. c)

Flow Features Matrix(FAM): FAM comprises flow feature vectors and application
category labels, formally defined as $FAM=[L\\_FAM,U\\_FAM]$. $L\\_FAM$ refers
to the labeled flow feature matrix,
$L\\_FAM=F+L=[[F_{1},F_{2},\cdots,F_{M}]^{T},L^{T}]$. $U\\_FAM$ refers to the
unlabeled flow feature matrix,
$U\\_FAM=F+L=[F_{1}^{\prime},F_{2}^{\prime},\cdots,F_{N}^{\prime}]]$. Fig. 4
represents an example of $L\\_FAM$.

####  IV-B3 VAE

As we all know, an Autoencoder (AE) is specifically designed to acquire a low-
dimensional latent representation of samples by constructing an encoder and a
decoder. It is commonly employed for tasks such as data compression or
generation. However, due to its sole focus on learning the encoding of the
sample itself, AE-based models usually show weak generalization. Therefore,
AE’s capacity to capture the underlying data distribution needs to be
improved. In addition, network traffic consistently exhibits characteristics
such as large scale, dynamic, and heterogeneity. A conventional AE model
usually fails to fully reconstruct comprehensive network traffic even when
provided with extensive datasets. Consequently, accurate classification
becomes imperative when encountering traffic samples beyond the dataset.

Figure 5: The General Architecture of VAE.

Variational Autoencoder (VAE) is an extension of Autoencoder (AE). More
precisely, it is a generative model widely employed for unsupervised pre-
training of unlabeled data. Its excellent ability to learn the latent
distribution enables the model to acquire strong generalization capabilities
during subsequent fine-tuning. As depicted in Fig. 5, the core concept behind
VAE lies in learning the implicit representation of actual samples $X$ and the
implicit distribution from these samples to generated samples $X^{\prime}$.
This approach enhances the robustness of the model in learning implicit
feature representations. It estimates the overall data distribution by
constructing a generative model $p_{\varphi}(Z|X)$ based on data samples.
However, commonly used methods for estimating data distributions rely on
maximum likelihood estimation through parameter estimation techniques.
Assuming that the distribution of data samples $P(X)$ follows a Gaussian
distribution $N(\mu,\sigma^{2})$, this transforms the statistical problem of
generating models into a parameter estimation problem. The remaining challenge
lies in fitting the distributions of encoder $q_{\phi}(Z|X)$ and decoder
$p_{\varphi}(Z|X)$. While autoencoders (AE) typically learn data distributions
by minimizing reconstruction loss like Mean Square Error (MSE), this type of
fitting primarily occurs at a sample level within datasets. It fails to
capture underlying data distributions effectively. In contrast, VAE employs KL
divergence as a measure for quantifying differences between two distributions,
also known as relative entropy. Therefore, we can define our loss function as
Eq. 1.

$Loss=L(X,X^{\prime})+\sum_{j}KL(q_{j}(Z|X)\parallel p(Z))$ (1)

The loss function consists of two components: the first component is the
reconstruction loss, and the second component is the KL divergence between the
proper distribution and the distribution we have chosen. VAE aims to minimize
this relative entropy, as expressed in Eq. 2.

$\displaystyle L(x)$ $\displaystyle=E_{z~{}q(z|x)}\log\frac{p(z,x)}{q(z|x)}$
(2) $\displaystyle=\log p(x)-KL(q(z|x)\parallel p(z|x))$

$L(x)$ is called the Variational Lower Bound. We aim to optimize this lower
bound, as the closer it is to $\log p(x)$, the smaller the KL divergence. In
this case, $q_{\phi}(X|Z)$ can estimate $p_{\varphi}(Z|X)$ more accurately.

The VAE further decomposes the sampling of $z$ into two parts: one consists of
fixed values such as the standard deviation $\sigma$ and the mean $\mu$, and
the other is a random Gaussian noise $\epsilon$ . After applying the
reparameterization trick, we can rewrite $L(x)$ as Eq. 3:

$L(x)=\frac{1}{2}\sum_{i=1}^{J}(1+\log(\sigma_{j}^{2})-\mu_{j}^{2})+\frac{1}{L}\sum_{i=1}^{L}\log
p(x|z_{i})$ (3)

The optimization of the variational lower bound of $L(x)$ implies that, while
ensuring that the $Z$ values generated by the encoder conform to a prior
Gaussian distribution, the decoder can maximize the possibility of
reconstructing the original $X$.

####  IV-B4 VAE based Unsupervised Learning Algorithm for Network TC

As shown in Alg.1, the entire algorithm mainly consists of the following three
steps:

  1. a)

Define the hyperparameters: Input dimension is $input\\_dim$; Hidden layer
dimensions are $hi\\_dim$ for $0<i<L$, where $L$ is the number of hidden
layers; Dependent variable dimension is $z\\_dim$; Batch size is
$batch\\_size$; Number of epochs for training is $num\\_epochs$.

  2. b)

Dataset: $X\in D_{u}$, $X$ is equivalent to a $U\\_FAM$ with
$input\\_dim\times batch\\_size$.

  3. c)

Model construction and training:

     * •

_Define the exact architectures of Encoder and Decoder._ This includes
determining the number of layers, $input\\_dim$, $hi\\_dim$, $z\\_dim$, and
the loss function. The Encoder maps $X$ into the latent space $Z$, while the
Decoder maps the randomly sampled $z$ from $Z$ into the data space
$X^{\prime}$. The ultimate goal is to make $X$ and $X^{\prime}$ as close as
possible.

     * •

_The feed-forward propagation process from Encoder to Decoder._

       * –

In this context, the input data $X$, referred to as $U\\_FAM$, is fed into the
Encoder. After sequential computations, the mean and variance of the posterior
distribution $Z$ in the latent space are obtained.

       * –

The technique of reparameterization is used to sample a latent variable $z$
from $Z$, i.e., $Z=\mu+\sigma\odot\varepsilon$, where $\varepsilon\sim
N(0,1)$, and it is then fed into the Decoder.

       * –

The decoder performs layer-by-layer computations to obtain the reconstructed
output $X^{\prime}$ of the input data $X$, which can be expressed as
$U\\_FAM^{\prime}$.

       * –

Calculate the reconstruction error and KL divergence based on Eq. 1 to obtain.

     * •

_Backpropagation and Optimization._ By iterating through $num\\_epochs$ and
utilizing the optimizer defined, the VAE model is trained in a loop to
optimize the model parameters with the goal of minimizing $L(x)$.

Figure 6: The VAE Model Architecture of FedEdge AI-TC. Algorithm 1 VAE based
Unsupervised Learning Algorithm for Network Traffic Classification

0: Training data: $U\\_FAM$, $X\in D_{u}$; Training parameters: $imput\\_dim$,
$h_{i}\\_dim$, $0<i<L$, $z\\_dim$, $batch\\_size$, $num\\_epochs$,
$learning\\_rate$, $optimizer$;

0: VAE based unsupervised model.

1: Initial training dataset $X\in D_{u}$;

2: Define the exact architectures of Encoder and Decoder;

3: for epoch in $num\\_epochs$ do

4: for epoch batch of $U\\_FAM$ of $X$ do

5: For each training samples $X\in D_{u}$:

6: Encoder: computing the mean-$\mu_{x}$ and variance-$\sigma_{x}$ of $Z$;

7: Encoder: using reparameterize skill to sample $z$ from $Z$;
$z=\mu_{x}+\sigma_{x}\odot\varepsilon$; $\varepsilon$ from $N(0,1)$;

8: Decoder:computing by layer-wise to get the reconstruct output $X^{\prime}$
of $X$, which is $U\\_FAM^{\prime}$;

9: Compute the loss of reconstruct and KL-divergence using equation(1); Update
the weights and bias;

10: end for

11: end for

The structure of the unsupervised model for network traffic based on VAE is
shown in Fig. 6, and detailed parameters are provided in Table III.

TABLE III: Detailed parameters of FedEdge AI-TC Parameter name |  Parameter value |  Parameter Interpretation
---|---|---
input_dim |  78 |  Model input dimension
Layers of Encoder/Decoder |  3 |  Number of layers in the encoder and decoder
$h_{1}\\_dim$, $h_{2}\\_dim$, $h_{3}\\_dim$ |  78,64,32 |  Dimension of each layer
loss function for Encoder |  ReLU |  Loss function
loss function for Decoder |  ReLU |  Loss function
loss function for Decoder’s output |  Sigmoid |  Loss function
batch_size |  128 |  Batch size
learning rate |  0.01 |  Learning rate

####  IV-B5 VAE+CNN based Semi-supervised Learning Algorithm for Network TC

Figure 7: The VAE+CNN Semi-supervised Model Architecture of FedEdge AI-TC.

As illustrated in Fig.7, there are three parts in the semi-supervised based
network traffic classifier: the encoder of VAE model, CNN, and softmax
classifier. The labeled data, i.e., L_FAM is fed into the Encoder of the VAE
model obtained from Section.IV-B4, then the output of the VAE encoder is
subsequently fed into a three layers CNN model with a softmax classifier,
which is concatenated with the VAE encoder. Finally, the decision results will
be outputted for classification. The overall process is commonly referred to
as _Fine-Tuning_. Due to the limited space of this paper, we do not provide
the detail of the CNN classifier, which can be found in our previous work[37].

###  IV-C The Model Compression Method Based on Interpretation

####  IV-C1 Model Interpretation

We propose to implement an interpretable framework for deep learning traffic
classification models based on SHAP values, which are mainly used to quantify
the contribution of each feature to the model prediction. The basic design
idea is to calculate the marginal contribution SHAP value when features are
added to the model so that the importance of the features can be interpreted
according to the SHAP value, which is calculated as Eq. 4 in this paper.
Suppose the i-th sample of sample set M is $x_{i}$, the j-th feature of sample
$x_{i}$ is $x_{ij}$, $f(x_{ij})$ is the shapely value of $x_{ij}$.

$f(x_{ij})=\sum_{S}\frac{\left|S\right|!(p-\left|S\right|-1)!}{p!}(v_{x}(S\cup\\{x_{j}\\})-v_{x}(S))$
(4)

where $S\subseteq\\{x_{1},\cdots,x_{p}\\}\setminus\\{x_{j}\\}$,
$\\{x_{1},\cdots,x_{p}\\}$ is the set of all possible input features excluding
$\\{x_{j}\\}$, $p$ is the number of features of the sample, and $v_{x}(S)$ is
the prediction result of the feature subset $S$.

The architecture of Model Interpretation is shown in Fig. 8. The left part is
the traditional structure of the traffic classification model, and the process
shown on the right allows for the interpretability of the traffic
classification model and the optimization of the structure and parameters of
the traffic classification model. This framework is divided into a local
interpretation and a global interpretation.

_Local interpretation_ means that for each data instance, the contribution of
each feature to its predicted outcome is calculated and presented visually.
The formula for calculating the local interpretation is as follows:

$y_{i}=y_{base}+f(x_{i1}+f(x_{i2}+\cdots+f(x_{ij})$ (5)

where $y_{i}$ is the predicted value of the model for sample $x_{i}$,
$y_{base}$ is the mean of all sample evaluated values.

As for _Global interpretation_ , firstly, a matrix of feature SHAP values is
calculated, with one instance per row and one feature per column. Secondly, in
the traditional global interpretation, the feature j’s contribution is
obtained by summing the shapely mean of feature j for all samples with Eq. 6.
And then, SHAP values are sorted in descending order to obtain the importance
of the model features.

$f(x_{j})=\sum_{i=1}^{M}f(x_{ij})$ (6)

####  IV-C2 Model Pruning

To address the issue of how to compress models to make them suitable for
training and inference on resource-limited devices, in particular, we will
focus on pruning, an easier-to-implement model compression technique. Model
pruning is based on an underlying assumption: ’weight contribution,’ which
means that not every weight contributes equally to the output prediction.

Figure 8: The Architecture of Model Interpretation.

Therefore, the basic design idea of model pruning in this paper is to rank the
feature importance by global interpretation, after which the importance
ranking of the convolutional kernels is calculated using the causal evaluation
mechanism. The convolutional kernels with importance below a threshold are
filtered out and pruned.

##  V Experimental Evaluation and Discussion

###  V-A Evaluation Settings and Chosen Datasets

We conducted the experimental evaluation on two datasets. One is a
benchmarking public dataset, ISCX-VPN2016. The other is a private dataset
built by ourselves. The latter comprises six popular applications and
background traffic from terminals collected in the campus network scenario,
including Bilibili, QQ music, Honor of Kings, Teamfight Tactics, and Game for
peace. These apps cover the five popular app categories of video, music, Moba
games, First Person Shooting (FPS) and Role-playing game (RPG). To collect the
data, we used a semi-automatic web traffic generation tool. We leveraged an
automated traffic generator to collect traffic from Bilibili and QQ music,
with 13,314 flows from Bilibili and 20703 flows from QQ music. For the other
interactive games, we chose to collect them manually. We used PCAPDroid to
mark the network traffic as it occurred at the endpoint and also at the
router, where the two were compared and filtered. The network flow features
were calculated by CICFlowMeter for each application’s PCAP files. In
addition, this experiment also provides the statistics of the background
traffic information, which contains information about network location
services, security components, and syslog services. Table IV shows the exact
number of flows.

TABLE IV: The private experimental dataset Applications Name | Type | Number
---|---|---
bilibili | Video | 13314
QQ music | Music | 20703
Honor of Kings | Moba | 9475
Teamfight Tactics | RPG | 14005
Game for peace | FPS | 7763
Background | Log | 13017

The experimental environment is AMD Ryzen 3600, 16GB RAM, NVIDIA GTX 1660,
CUDA 7.5, CDNN10.5. In this paper, Python3 is the primary programming
language. The following is a description of the evaluation metrics: Precision,
Recall, F1, Accuracy, and AUC.

Time-complexity-related metrics about training like training time are not
included in this paper because we think those are highly dependent on the
hardware resources.

###  V-B Performance Evaluation

The training process of the VAE part in the VAE+CNN (E-CNN) model is divided
into two parts. Firstly, the labels were removed from the datasets, which can
be acted as the unlabeled flows for the unsupervised learning training. After
training, one can save the trained encoder in the VAE model for further semi-
supervised learning. The training process for the CNN part mainly aim to
convert the labeled data into digital encoding for fine-tuning.

In the single CNN model, we process the dataset and divide it into training
and testing sets according to ratio, which is then trained and evaluated in
the CNN.

For the E-CNN and single CNN models, we obtained the following two diagrams by
adjusting the ratio of the dataset to the training and testing sets. The
horizontal axis in the diagram represents the partitioning ratio. For example,
0.2 means that 80$\%$ of the entire dataset is allocated to the training set,
and the remaining 20$\%$ is allocated to the testing set. The vertical axis
represents the accuracy of the model training results under this partition.
Fig. 9 displays the model results of E-CNN and single CNN in real-life
scenarios, while Fig. 10 shows the results in the public dataset.

Figure 9: Accuracy of different rates in real-life scenarios.

Figure 10: Accuracy of different rates in public dataset.

According to the figures, as the partition ratio increases, meaning that the
training set data decreases, the accuracy of both models tends to decrease and
becomes similar, both at around 0.7. Within the partition ratio interval [0.2,
0.8], E-CNN achieves good results with higher accuracy than a single CNN.

We have set the partition ratio to 0.45. After training and testing, we
obtained the respective Confusion Matrix and Classification Report for E-CNN
and single CNN using this ratio. The Confusion matrix is shown in Fig. 11 and
Fig. 12. In the figures, the x-axis represents the prediction labels, the
y-axis represents the actual labels, and the color intensity indicates the
count of correct and incorrect predictions. According to the figure, under the
partition with this ratio, E-CNN shows relatively accurate predictions for
’QQmusic,’ ’Teamfight Tactics,’ and ’Bilibili.’ On the other hand, compared to
CNN, E-CNN’s predictions for ’QQmusic,’ ’IQiyi,’ and ’Teamfight Tactics’ are
relatively accurate, but there are also several errors.

Figure 11: VAE+CNN Confusion Matrix

Figure 12: CNN Confusion Matrix

The classification report is presented in Table V and Table VI. It is evident
from the tables that, under this partition ratio, E-CNN demonstrates
relatively accurate predictions for ’Teamfight Tactics’ and ’Bilibili,’ with
prediction accuracy rates above 0.9. The prediction accuracy for other
applications is also around 0.7, resulting in a total classification accuracy
of 0.7422. In contrast, CNN’s performance is notably inferior to that of
E-CNN. The lowest prediction accuracy is approximately 0.5, and the highest
prediction accuracy is only around 0.8. Its overall classification accuracy is
merely 0.6989.

TABLE V: Classification Report of VAE+CNN | Precision | Recall | F1-score | Support
---|---|---|---|---
QQmusic | 0.7736 | 0.7485 | 0.7609 | 9270
Background | 0.5163 | 0.6754 | 0.5852 | 5853
Game For Peace | 0.6214 | 0.6598 | 0.6401 | 3513
Bilibili | 0.9097 | 0.9000 | 0.9048 | 6020
IQiyi | 0.7056 | 0.6644 | 0.6843 | 6272
Honor of Kings | 0.6915 | 0.4897 | 0.5734 | 4243
Teamfight Tactics | 0.9457 | 0.9379 | 0.9418 | 6280
accuracy |  |  | 0.7422 | 41451
macro avg | 0.7377 | 0.7251 | 0.7272 | 41451
weighted avg | 0.7515 | 0.7422 | 0.7434 | 41451
TABLE VI: Classification Report of CNN | Precision | Recall | F1-score | Support
---|---|---|---|---
QQmusic | 0.6160 | 0.7612 | 0.6809 | 8940
Background | 0.5496 | 0.4854 | 0.5155 | 3871
Game For Peace | 0.6696 | 0.5250 | 0.5885 | 2941
Bilibili | 0.7990 | 0.7936 | 0.7963 | 4873
IQiyi | 0.6527 | 0.6982 | 0.6747 | 7590
Honor of Kings | 0.6933 | 0.6125 | 0.6504 | 7672
Teamfight Tactics | 0.9869 | 0.8767 | 0.9285 | 5564
accuracy |  |  | 0.6989 | 41451
macro avg | 0.7096 | 0.6789 | 0.6907 | 41451
weighted avg | 0.7059 | 0.6989 | 0.6989 | 41451

Based on the experimental results presented above, we can conclude that when
the ratio of labeled data to total data falls within the range of [0.2, 0.8],
the VAE+CNN method proposed in this paper can be employed for classifying and
recognizing network application traffic. At this range, the recognition
outcomes outperform those achieved by solely training with CNN, and the
classification accuracy is commendable.

Fig. 13 visualizes the importance of ten features for the classification
results. These ten features are the most significant ones contributing to the
classification outcomes. For instance, in the case of the applications
Bilibili and IQiyi, the most crucial feature for the classification result is
s_idleMax, representing the maximum idle time of any neighboring packet in a
stream. Notably, video streams exhibit distinct characteristics on this
feature due to their continuous packet sending and shorter inter-packet idle
time, which differentiates them from game streams.

Figure 13: Summaryplot of Global Interpretation. TABLE VII: Baseline Model vs. XAI-Pruning Model | Baseline model | XAI-Pruning model
---|---|---
Size of zipped file | 634163 bytes | 207292 bytes
Accuracy of prediction | 0.75 | 0.69
Inference Time | 2.78s | 1.12s

The XAI-Pruning model effectively filters out unimportant weights and
parameters, substantially reducing the number of parameters and the overall
model size. As a result, the model becomes more compact, enabling
significantly faster inference and quicker predictions. On the other hand, the
baseline model, which includes all weights and parameters, generally exhibits
higher prediction accuracy due to its capacity to capture more intricate model
details and complexity.

As depicted in Table VII, the baseline model size is approximately three times
larger than the XAI-Pruning model, and the inference time is doubled. However,
despite these changes, both models still exhibit comparable prediction
accuracy without significant differences. Consequently, the pruning method
proposed in this paper effectively reduces the model’s compression size and
inference time to some extent while having a relatively minor impact on
prediction accuracy.

##  VI Conclusion and Future work

This paper presents the FedEdge AI-TC approach for trusted Federated Learning
(FL) based efficient Network TC in 5G CPE/HGU. Firstly, FedEdge AI-TC
effectively protects the data privacy of network subscribers by proposing an
FL-based framework of local training, model parameters iterating, and
centralized training. Secondly, a semi-supervised TC algorithm based on
Variational Auto-Encoder (VAE) and convolutional neural network (CNN) is
designed to reduce data dependence while keeping the TC accuracy. Finally,
XAI-Pruning, an AI model compression method, combined with the DL model
interpretability, is proposed to condense the model and interpret it globally
and locally to achieve light-weighted AI-TC model deployment while building
the trust in their decision of network operators. To demonstrate the
efficiency of the proposed method, we conducted some experimental evaluations
on commonly used public benchmark datasets and real network datasets. The
results show that FedEdge AI-TC can outperform the benchmarking methods
regarding the accuracy and achieve excellent TC performance of model inference
on 5G CPE/HGU with limited computing resources, which effectively protects the
users’ privacy and improve the model’s credibility.

However, besides reliability, robustness, and generalization are still two
important topics when handling network traffic classification, especially
using ML/DL. In the future, we will continuously focus on how to leverage
ML/DL algorithms like generative models or large language models to enhance
the reliability, robustness, and generalization of network traffic
classification models.

## Acknowledgment

The paper is supported by National Natural Science Fundation (General Program)
of China under Grant 61972211

## Reference

  * [1] S. Rezaei and X. Liu, “Deep learning for encrypted traffic classification: An overview,” _IEEE communications magazine_ , vol. 57, no. 5, pp. 76–81, 2019\.
  * [2] X. Zhou, Y. Hu, J. Wu, W. Liang, J. Ma, and Q. Jin, “Distribution bias aware collaborative generative adversarial network for imbalanced deep learning in industrial iot,” _IEEE Transactions on Industrial Informatics_ , vol. 19, no. 1, pp. 570–580, 2023.
  * [3] G. Aceto, D. Ciuonzo, A. Montieri, and A. Pescapé, “Mobile encrypted traffic classification using deep learning,” in _2018 Network traffic measurement and analysis conference (TMA)_. IEEE, 2018, pp. 1–8.
  * [4] P. Wang, F. Ye, X. Chen, and Y. Qian, “Datanet: Deep learning based encrypted network traffic classification in sdn home gateway,” _IEEE ACCESS_ , vol. 6, pp. 55 380–55 391, 2018.
  * [5] G. Aceto, D. Ciuonzo, A. Montieri, and A. Pescapé, “Mobile encrypted traffic classification using deep learning: Experimental evaluation, lessons learned, and challenges,” _IEEE Transactions on Network and Service Management_ , vol. 16, no. 2, pp. 445–458, 2019.
  * [6] S. Niknam, H. S. Dhillon, and J. H. Reed, “Federated learning for wireless communications: Motivation, opportunities, and challenges,” _IEEE Communications Magazine_ , vol. 58, no. 6, pp. 46–51, 2020.
  * [7] X. Zhou, X. Ye, K. I.-K. Wang, W. Liang, N. K. C. Nair, S. Shimizu, Z. Yan, and Q. Jin, “Hierarchical federated learning with social context clustering-based participant selection for internet of medical things applications,” _IEEE Transactions on Computational Social Systems_ , pp. 1–10, 2023.
  * [8] P. Wang, X. Chen, F. Ye, and Z. Sun, “A survey of techniques for mobile service encrypted traffic classification using deep learning,” _IEEE Access_ , vol. 7, pp. 54 024–54 033, 2019.
  * [9] N. Bayat, W. Jackson, and D. Liu, “Deep learning for network traffic classification,” _arXiv preprint arXiv:2106.12693_ , 2021.
  * [10] B. Pang, Y. Fu, S. Ren, Y. Wang, Q. Liao, and Y. Jia, “Cgnn: traffic classification with graph neural network,” _arXiv preprint arXiv:2110.09726_ , 2021.
  * [11] L. Yang, A. Finamore, F. Jun, and D. Rossi, “Deep learning and zero-day traffic classification: Lessons learned from a commercial-grade dataset,” _IEEE Transactions on Network and Service Management_ , vol. 18, no. 4, pp. 4103–4118, 2021.
  * [12] P. Wang, S. Li, F. Ye, Z. Wang, and M. Zhang, “Packetcgan: Exploratory study of class imbalance for encrypted traffic classification using cgan,” in _ICC 2020-2020 IEEE International Conference on Communications (ICC)_. IEEE, 2020, pp. 1–7.
  * [13] X. Kong, C. Wang, Y. Li, J. Hou, T. Jiang, and Z. Liu, “Traffic classification based on cnn-lstm hybrid network,” in _International Forum on Digital TV and Wireless Multimedia Communications_. Springer, 2021, pp. 401–411.
  * [14] H. Zhou, Y. Wang, and M. Ye, “A method of cnn traffic classification based on sppnet,” in _2018 14th international conference on computational intelligence and security (CIS)_. IEEE, 2018, pp. 390–394.
  * [15] K. N. K. Thapa and N. Duraipandian, “Malicious traffic classification using long short-term memory (lstm) model,” _Wireless Personal Communications_ , vol. 119, pp. 2707–2724, 2021.
  * [16] H.-K. Lim, J.-B. Kim, K. Kim, Y.-G. Hong, and Y.-H. Han, “Payload-based traffic classification using multi-layer lstm in software defined networks,” _Applied Sciences_ , vol. 9, no. 12, p. 2550, 2019.
  * [17] X. Zhou, W. Liang, W. Li, K. Yan, S. Shimizu, and K. I.-K. Wang, “Hierarchical adversarial attacks against graph-neural-network-based iot network intrusion detection system,” _IEEE Internet of Things Journal_ , vol. 9, no. 12, pp. 9310–9319, 2022.
  * [18] S. Rezaei and X. Liu, “How to achieve high classification accuracy with just a few labels: A semi-supervised approach using sampled packets,” _arXiv preprint arXiv:1812.09761_ , 2018.
  * [19] Y. Liu, Z. Zhu, and P. Zhong, “A classification method for network traffic based on semi-supervised approach,” in _Journal of Physics: Conference Series_ , vol. 2010, no. 1. IOP Publishing, 2021, p. 012014.
  * [20] J. Ning, Y. Wang, J. Yang, H. Gacanin, and S. Ci, “A novel malware traffic classification method using semi-supervised learning,” in _2021 IEEE 94th Vehicular Technology Conference (VTC2021-Fall)_. IEEE, 2021, pp. 1–5.
  * [21] R. Zhao, X. Deng, Z. Yan, J. Ma, Z. Xue, and Y. Wang, “Mt-flowformer: A semi-supervised flow transformer for encrypted traffic classification,” in _Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining_ , 2022, pp. 2576–2584.
  * [22] A. Almalawi and A. Fahad, “An efficient network classification based on various-widths clustering and semi-supervised stacking,” _IEEE Access_ , vol. 9, pp. 151 681–151 696, 2021.
  * [23] Z. Jin, Z. Liang, M. He, Y. Peng, H. Xue, and Y. Wang, “A federated semi-supervised learning approach for network traffic classification,” _International Journal of Network Management_ , vol. 33, no. 3, p. e2222, 2023\.
  * [24] T. Li, S. Chen, Z. Yao, X. Chen, and J. Yang, “Semi-supervised network traffic classification using deep generative models,” in _2018 14th International Conference on Natural Computation, Fuzzy Systems and Knowledge Discovery (ICNC-FSKD)_. IEEE, 2018, pp. 1282–1288.
  * [25] C. Xu, R. Xia, Y. Xiao, Y. Li, G. Shi, and K.-C. Chen, “Federated traffic synthesizing and classification using generative adversarial networks,” in _ICC 2021-IEEE International Conference on Communications_. IEEE, 2021, pp. 1–6.
  * [26] P. Wang, Z. Wang, F. Ye, and X. Chen, “Bytesgan: A semi-supervised generative adversarial network for encrypted traffic classification in sdn edge gateway,” _Computer Networks_ , vol. 200, p. 108535, 2021.
  * [27] H. H. Yang, Z. Zhongyuan, and T. Quek, “Enabling intelligence at network edge network edge: An overview of federated learning,” _ZTE Commun._ , vol. 18, no. 2, pp. 2–10, 2020.
  * [28] O. Aouedi, K. Piamrat, G. Muller, and K. Singh, “Fluids: Federated learning with semi-supervised approach for intrusion detection system,” in _2022 IEEE 19th Annual Consumer Communications & Networking Conference (CCNC)_. IEEE, 2022, pp. 523–524.
  * [29] P. Zhou, “Federated deep payload classification for industrial internet with cloud-edge architecture,” in _2020 16th International Conference on Mobility, Sensing and Networking (MSN)_. IEEE, 2020, pp. 228–235.
  * [30] O. Aouedi, K. Piamrat, G. Muller, and K. Singh, “Intrusion detection for softwarized networks with semi-supervised federated learning,” in _ICC 2022-IEEE International Conference on Communications_. IEEE, 2022, pp. 5244–5249.
  * [31] X. Zhou, W. Liang, K. I.-K. Wang, and L. T. Yang, “Deep correlation mining based on hierarchical hybrid networks for heterogeneous big data recommendations,” _IEEE Transactions on Computational Social Systems_ , vol. 8, no. 1, pp. 171–178, 2021.
  * [32] Y. Zhao, J. Chen, D. Wu, J. Teng, and S. Yu, “Multi-task network anomaly detection using federated learning,” in _Proceedings of the 10th international symposium on information and communication technology_ , 2019, pp. 273–279.
  * [33] H. Mun and Y. Lee, “Internet traffic classification with federated learning,” _Electronics_ , vol. 10, no. 1, p. 27, 2020.
  * [34] U. Majeed, L. U. Khan, and C. S. Hong, “Cross-silo horizontal federated learning for flow-based time-related-features oriented traffic classification,” in _2020 21st Asia-Pacific Network Operations and Management Symposium (APNOMS)_. IEEE, 2020, pp. 389–392.
  * [35] E. Bakopoulou, B. Tillman, and A. Markopoulou, “Fedpacket: A federated learning approach to mobile packet classification,” _IEEE Transactions on Mobile Computing_ , vol. 21, no. 10, pp. 3609–3628, 2021.
  * [36] X. Zhou, W. Liang, K. I.-K. Wang, Z. Yan, L. T. Yang, W. Wei, J. Ma, and Q. Jin, “Decentralized p2p federated learning for privacy-preserving and resilient mobile robotic systems,” _IEEE Wireless Communications_ , vol. 30, no. 2, pp. 82–89, 2023.
  * [37] Z. Wang, C. Miao, Y. Xu, Z. Li, Z. Sun, and P. Wang, “Trusted encrypted traffic intrusion detection method based on federated learning and autoencoder,” _China Communications_ , vol. 0, no. 0, pp. 1–25, 2023. [Online]. Available: http://www.cic-chinacommunications.cn/EN/10.23919/JCC.ja.2022-0392#2
  * [38] P. Wang, F. Ye, and X. Chen, “A smart home gateway platform for data collection and awareness,” _IEEE Commun. Mag._ , vol. 56, no. 9, pp. 87–93, 2018. [Online]. Available: https://doi.org/10.1109/MCOM.2018.1701217
  * [39] G. Draper-Gil, A. H. Lashkari, M. S. I. Mamun, and A. A. Ghorbani, “Characterization of encrypted and vpn traffic using time-related,” in _Proceedings of the 2nd international conference on information systems security and privacy (ICISSP)_ , 2016, pp. 407–414.
  * [40] N. Moustafa and J. Slay, “Unsw-nb15: a comprehensive data set for network intrusion detection systems (unsw-nb15 network data set),” in _2015 military communications and information systems conference (MilCIS)_. IEEE, 2015, pp. 1–6.
  * [41] G. Aceto, D. Ciuonzo, A. Montieri, V. Persico, and A. Pescapé, “Mirage: Mobile-app traffic capture and ground-truth creation,” in _2019 4th International Conference on Computing, Communications and Security (ICCCS)_. IEEE, 2019, pp. 1–8.

|  Pan Wang (M’18) received the BS/MS/Ph.D. degree in Electrical & Computer
Engineering from Nanjing University of Posts&Telecommunications, Nanjing,
China, in 2001, 2004, and 2013, respectively. He is currently a Full Professor
at Nanjing University of Posts & Telecommunications, Nanjing, China. His
research interests include AI-powered networking and security in
B5G/6G/IoT/Smart Grid/CFN, and AI-enabled big data analysis. From 2017 to
2018, he has been a visiting scholar at University of Dayton (UD) in the
Department of Electrical and Computer Engineering, OH, USA. He served as a TPC
member of IEEE CyberSciTech Congress. He is also a reviewer for several
journals, including IEEE Transaction on Network and Service Management, IEEE
Transaction on EMERGING TOPICS IN COMPUTATIONAL INTELLIGENCE, IEEE Internet of
Things Journal, IEEE Journal on Selected Areas in Communications, IEEE ACCESS,
Computer Networks, Computer&Security, Computer Communications, Engineering
Applications of Artificial Intelligence, Big Data Research, etc.
---|---
|  Zeyi Li is currently pursuing the Ph.D. degree in Cyberspace Security at
Nanjing University of Posts and Telecommunications. He was born in Soochow,
Jiangsu, China, in 1997. He received his bachelor’s degree in mathematics in
2019 and received M.S. degree in computer science in 2022. His research
interests include network security, communication network security, anomaly
detection and analysis, deep packet inspection, and graph neural networks.
---|---
|  Mengyi Fu was born in Huaian, Jiangsu, China ,in 2000. She is currently
pursuing the Ph.D degree at NanjingUniversity of Posts and Telecommunications.
She received her B.Sc from Nanjing University of Posts and Telecommunications,
NanjingChina, in 2022. Her research includes encrypted traffic identification,
deep learning and traffic prediction.
---|---
|  Zixuan Wang was born in Nanjing, Jiangsu, China ,in 1994 . He obtained a
bachelor’s degree from Tongda College of Nanjing University of Posts and
Telecommunications in 2017, He is currently pursuing a master’s degree in
logistics engineering at Nanjing University of Posts and Telecommunications.
His research interests include encrypted traffic identification and data
balancing.
---|---
|  Ze Zhang currently pursuing the Master’s degree in Information Network at
Nanjing University of Postsand Telecommunications. He was born in Zhenjiang,
Jiangsu, China, in 2000. He received his Bachelor’s degree in Network
Engineering in 2022. His research interests include computer networks, network
security, anomaly detection and analysis.
---|---
|  Minyao Liu was born in Ganzou, Jiangxi, China, in 2000. She is currently
pursuing her master’s degree at Nanjing Post and Telecommunications
University. She received her bachelor’s degree in Management from NUPT in
2022. Her research areas include traffic identification, deep learning and
anomaly detection.
---|---