Federated Few-shot Learning for Cough Classification with Edge Devices

The notations that are used in this paper are summarized in Table 1. These notations are formally defined in the following sections that describe our method and technical details.

In this study, we tackle the case in which a new type of disease appears suddenly and medical systems have no records related to the disease data before. We can take COVID-19 as a particular example. The data can only be collected from patients suffering from these diseases. Moreover, we deal with the situation where we have multiple local medical centers, each can collect a few samples of a patient’s cough data. However, these centers do not want to leak their data to protect patient information. Our goal is to classify and detect a novel type of cough without the need for a huge unified dataset.

To this end, we first assume that each local medical center has a sufficient amount of data on several types of traditional cough related to diseases that have been discovered and treated in the past. The available data, the so-called base set, have been labeled and are denoted as $\mathcal{A}^{u}$ = $\{(x_{i}^{u},y_{i}^{u})|y_{i}^{u}\in Y_{train}\}$, where $x_{i}^{u}$ and $y_{i}^{u}$ denote the $i$-th cough audio representation and the corresponding cough type on the $u$-th local device, respectively, $Y_{train}$ includes types of cough that have been gathered during the past years. The base set $\mathcal{A}^{u}$ acts as the prior knowledge, which can be exploited to learn the novel classes of cough through only a few samples. We also assume that there are $U\in\mathbb{N}$ local devices holding cough data belonging to new types of cough that have not been found before and are not included in $\mathcal{A}^{u}$. The new dataset including only a few labeled data on the $u$-th device is denoted as $\mathcal{E}^{u}=\{(x_{j}^{u},y_{j}^{u})|y_{j}^{u}\in Y_{test},y_{j}^{u}\notin Y_{train}\}$ where $(x_{j}^{u},y_{j}^{u})$ denotes the data point belonging to the novel type of cough on the $u$-th local device, $x_{j}^{u}$ is $j$-th audio records of cough, $y_{j}^{u}$ is the corresponding label of $x_{j}^{u}$, and $Y_{test}$ includes types of novel cough sounds.

We follow the idea of FSL that utilizes the base set $\mathcal{A}^{u}$ to train a classifier to classify a novel set $\mathcal{E}^{u}$. In our work, the FSL procedure follows the episodic learning mechanism snell2017prototypical , where an episode $e$, including of a support set and a query set, is formed to perform a few-shot task. Intuitively, the two support and query sets are created to mimic the circumstances encountered during real evaluation. In our cough classification context, the support set is the set that has information about the long-established cough sounds and the corresponding labels, which is the prior knowledge that helps to train the model to predict novel labels of data samples in the query set. This idea is based on human learning, for example, a child can utilize knowledge learned from a few sample images of a drawing animal (the support set) to recognize real-world animals (query set).

In the training procedure, we undertake a specific approach wherein we randomly select $N$ classes from the base set $\mathcal{A}^{u}$ on each device $u$. Within this selection, we further employ random sampling to draw $K$ support data samples and $V$ query data samples per class, thereby composing an episode denoted as $e$. Consequently, within each episode $e$, every class in the selected $N$ classes is represented by two distinct sets: the support set and the query set, both of which consist of sampled images. The support set is denoted as ${S^{u}={\{(s_{i}^{u},y_{i}^{u})\}}^{N\times{K}}_{i=1}}$, including $K$ different audio data samples $\{s_{1}^{u},s_{2}^{u},\cdots,s_{K}^{u}\}$ from $N$ classes $\{y_{1}^{u},y_{2}^{u},\cdots,y_{N}^{u}\}$, and $(s_{i}^{u},y_{i}^{u})\subset A^{u}$. Based on the support set, the classifier predicts a query set. The query set consists of data samples belonging to labels in the support set and is denoted as $Q^{u}={\{(q_{i}^{u},y_{i}^{u})\}}^{M\times{V}}_{i=1}$. $Q^{u}$ consists of $V$ different data points ${\{q_{1}^{u},q_{2}^{u},\cdots,q_{V}^{u}\}}$ from $M$ classes ${\{y_{1}^{u},y_{2}^{u},\cdots,y_{M}^{u}\}}$, and $(q_{i}^{u},y_{i}^{u})\subset\mathcal{A}^{u}$. As the result, the episodes can be characterized by three key variables: the quantity of classes denoted by $N$ (also referred to as the ”ways”), the number of samples per class in the support set denoted by $K$ (termed as the ”shots”), and the number of samples per class in the query set denoted by $V$.

After the training process is completed, we evaluate the classifier by using the test set which is the novel set $\mathcal{E}^{u}$. The support set and the query set are taken from the test set $\mathcal{E}^{u}$ as the way performed in the training process. In the testing process, the support set includes $(s_{j}^{u},y_{j}^{u})\subset\mathcal{E}^{u}$. And the query set consists of $(q_{j}^{u},y_{j}^{u})\subset\mathcal{E}^{u}$. In fact, the query set consists of unlabeled data. We only know that the labels of the data in the query set exist in the support set. We have to use an FSL model to predict the labels of data samples in the query set. In our evaluation process, the support set is provided by hospitals, and the query set can be obtained from the patients.

As medical centers have records of cough-related diseases but do not want to leak them out to protect patient information, we adopt federated learning to ensure the privacy of data as well as the satisfaction of patients. To this end, all $U$ local devices are supplied with models that have the same architecture. Each device trains the model with its own data. After that, local devices exchange the weight of models together to produce a global model without sharing data. Specifically, each device learns a local model $f_{\omega}^{u}$ following the prototypical network method. After a fixed period, we aggregate local model $f_{\omega}^{u}$ from $U$ local devices to generate global model $f_{\omega^{*}}$, which is called a communication round. After that, the global model is synced back to local devices to continue learning. Finally, we stop the training when we get a good global model $f_{\omega^{*}}$ that meets our expectations. In this work, we assume a communication round which is performed after all the local devices complete one training round.

Definition 1 formulates the problem that each local medical center can perform few-shot tasks, in which a different embedding space is learned in each local device. Specifically, on each device, the objective is to minimize the distance $d(f_{\omega}^{u}(q_{i}^{u})$, $p_{c}^{u}$) at each training step, where $d(\cdot)$ is the Euclidean distance function, $c$ is the label of $q_{i}^{u}$, and $p_{c}^{u}$ is a representation vector called prototype. Each device uses $S^{u}$ and $Q^{u}$ to train $f_{\omega}^{u}$. First, from $f_{\omega}(s_{i})$ of classes, a few-shot model calculates $p_{c}^{u}$ for each class. Second, the model chooses $c$ that has the smallest distance $d(f_{\omega}^{u}(q_{i}^{u})$, $p_{c}^{u}$) as the label of query $q_{i}^{u}$.

To train a single few-shot task, the objective of each local device is expressed as:

where $l^{u}_{\tau^{u}}(\omega)$ is the negative log-likelihood of the true label of each query sample ($q_{j}^{u}$, $y_{j}^{u}$) that the model of the $u$-th local device needs to minimize in a few-shot task ${\tau^{u}}$.

The average of objectives of few-shot tasks on a device $u$ is formulated as:

where $L^{u}(\omega)$ is the average of the objective functions of few-shot tasks per local device. $\beta^{u}$ is the total number of few-shot tasks per the $u$-th local device.

To allow information to be exchanged between local devices, we assume that all $U$ devices can join communication rounds periodically to generate $f_{\omega^{*}}$, aggregated from learned models $f_{\omega}^{u}$ of local devices. A communication round starts when each local device uploads its embedding network $f_{\omega}^{u}$ to the server and ends when all local devices download the global model. To be able to get $\omega^{*}$, we minimize the average of the objectives of all devices

where $\beta$ is the total number of few-shot tasks on all devices and $U$ is the number of local devices.

To complete the aim, from (4), our method has learned a transformation function to create an embedding space where cough data points are clustered closely around the class centroid, which is computed as the center point of all $x_{j}^{u}$ having the same value of $y_{j}^{u}$, denoted as $p^{u}_{novel}$. Therefore, we train a classifier to choose label $y_{j}^{u}$ as the cough type of $x_{j}^{u}$ by comparing the smallest distance from the embedded vector of $x_{j}^{u}$ to $p^{u}_{novel}$, i.e., we find $y_{j}^{u}$ satisfying the following condition:

To solve the formulated problem in (5), we propose a framework that consists of two main components: FSL to classify cough sounds and a federated learning paradigm, illustrated in Fig. 2.

First, we describe the architecture of FSL, where each local device exploits its data to train the embedding model by performing few-shot tasks. To achieve this, we design a training procedure consisting of two stages, as shown in Step (1) in Fig. 2. The first stage involves extracting Mel Frequency Cepstral Coefficients (MFCCs) from cough sounds in both $S^{u}$ and $Q^{u}$. The second stage is dedicated to classifying the query labels, wherein the MFCCs are fed into the embedding network to map the data into the embedding space. This stage utilizes prototypical networks snell2017prototypical to create an embedding space for audio records. In this space, classification is performed by calculating the distance to prototype representations of classes.

In Step (2), all local devices upload their weights ${\omega}^{u}$ to a server. In Step (3), the global model $f_{\omega^{*}}$ is aggregated from the uploaded weights. Finally, in Step (4), local devices download and sync the global model to complete a communication round. All four steps are iteratively performed until ${\omega^{*}}$ is converged.

Feature Extraction

Mel Frequency Cepstral Coefficients (MFCCs) are extracted from both support and query data. Empirically, we extract 40-dimensional features of MFCCs and design sliding windows to extract MFCCs in which the length of each window is 128ms with a 64ms hop length. After that, these features are inputted into an embedding network to map data into the embedding space.

FSL

The embedding network employed in our study is based on ResNet-18 he2016deep architecture, augmented with both channel attention block and spatial attention block. To integrate these components into our embedding network, we follow the approach outlined in woo2018cbam . A visual representation of our embedding network is presented in Fig. 3.

In the embedding space, we follow the implementation of prototypical networks snell2017prototypical . For every episode, a prototype is computed by averaging the embedded support examples per class:

where f is embedding function, $|S_{c}^{u}|$ is the number of support examples belonging to class c in episode e. We apply a softmax function to the negative distance to choose the label of the query sample.

Federated learning

We adopt the Federated averaging (FedAvg) algorithm mcmahan2017communication to complete the aggregation step, whose objective function is given in (4).

System design

To evaluate the proposed framework in a real-world application, edge devices are utilized as in Fig. 1. Specifically, a user has a mobile phone which helps the user to record audio to predict. When training, mobile phones save their data and send it to laptops for training. In the Pre-processing layer, extracting MFCCs of audio data is performed on the laptop in the local training layer. Then, laptops use these MFCCs data to train an embedding network. Each embedding network trained completely in one communication round will be pushed to the cloud layer treated as a server to aggregated weights. Finally, a global model is sent back to the mobile phone to serve users.

Since our proposed embedding network architecture is similar to TC-ResNet, we perform an ablation study that replaces our embedding network, in which ResNet-18 combined with attention mechanisms is employed, with an implementation of TC-ResNet (termed TD-ResNet7) in parnami2022few . In this experiment, we only consider the performance comparison in a few-shot setting without federated learning. More specifically, all the few-shot settings are listed as follows: 2-way 2-shot, 2-way 8-shot, and 5-way 8-shot. Table 2 illustrates the performance comparison between TD-ResNet7 and our proposed ResNet-Attention architectures. From the table, we discuss the following interesting observations:

• •

  We observe that the performance of TD-ResNet7 is slightly superior to our ResNet-Attention in the case of the 2-way 2-shot task.

• •

  However, except for the 2-way 2-shot task, the F1-Scores obtained by using TD-ResNet7 are significantly lower than those from ResNet-Attention in the remaining tasks.

Intuitively, although both TD-ResNet7 and our proposed architecture are built upon ResNet-18, the difference in implementation details causes the performance gap between the two architectures. Specifically, TD-ResNet7 sets the kernel sizes to $7\times 1$ and $3\times 1$ instead of the square kernels in ResNet-18. Besides, TD-ResNet7 applies dilated convolutional layers instead of normal convolutional layers. This architecture of TD-ResNet7 performs effectively on normal sequential data such as speech. However, the cough has a cycle that is not as same as other sequential data; therefore, TD-ResNet7 is unable to capture meaningful features of cough sound while our proposed architecture can do better.

Table 3 shows the performance comparison between our proposed F2LCough and all competitive schemes following the cross-validation approach, including MoCo, Relational network, TC-ResNet (few-shot learning only), ResNet-Attention, TC-ResNet-F (few-shot learning combined with federated learning), and F2LCough under the 2-way 2-shot setting.⁴⁴4In our experiments, other FSL settings show similar trends. We summary observations as follows:

• •

  Overall, F2LCough demonstrates effectiveness across most types of cough with superior performance than those of other methods, with the exception of Dry afternoon cough and Gaggy wet cough.

• •

  The two types of cough, including Dry afternoon cough and Gaggy wet cough, pose challenges due to their susceptibility to confusion with other types. For the Dry afternoon cough case, models trained on multiple devices, such as TC-ResNet-F and F2LCough, exhibit slow or unstable convergence, leading to inferior performance compared to ResNet-Attention and TC-ResNet.

• •

  The MoCo approach is the least effective for classifying coughs with all F1-Scores under 0.5. The reason for immensely low performance is that MoCo utilizes a contrastive learning method that requires a significant amount of training data to acquire a good result.

• •

  When using a Relational Network the F1-Scores fluctuate between the range of 0.43 and 0.57. One can see that the family of methods following the prototypical approach such as TC-ResNet outweighs the relational network both in the prediction of Heavy cold, and sore throat coughing and Night wet cough. Via our empirical experiments, we believe that the coughing data tend to form geometric clusters, thus the distance metric employed in the prototypical approach is more appropriate than the relational metric used in the Relational network for the cough classification task.

• •

  The F1-scores obtained from methods utilizing the F2L setting are significantly higher than those without federated learning. According to the experimental results, leveraging federated learning in cough classification is one of the most suitable manners, which is capable of exploiting as much user data as possible to preserve user privacy and acquire high performance. Consequently, the federated learning scheme is well-suited for applications, which require preserving user privacy and highly accurate precision like health care applications.

It is worth noting that since federated learning is employed for F2LCough, we only synchronize weights of models of local devices without sharing data in each communication round. As a result, our framework can protect the sensitive data of each user.

In a practical scenario, we aim to use the trained model to predict a newly outbreak decease of patients. Thus, the medical centers should be able to deploy the learned model into lightweight devices for field operations. Therefore, to demonstrate the capability of implementing F2LCough in a real-world situation, we experiment on a setup where users record their cough data and save the trained model using an application installed on Android phones while the medical center utilize laptops as mobile devices for training. The interface of the application is shown in Fig. 4.

We measure the prediction time of the model integrated on the laptop, shown in Table 4. Note that the following laptop specification is used in this experiment:

• •

  Laptop (using CPU to train the model): Ram 8GB, Processor 11th Gen Intel(R) Core(TM) i5-1135G7 @ 2.40GHz

Furthermore, we assess the running time on users’ Android devices by implementing the F2LCough model using TensorFlow Lite. This implementation is designed to classify various cough sounds, as well as non-cough sounds. The application exhibits a highly efficient inference time, with an average of approximately 2.5 milliseconds.

To validate the application, we conducted a study with 55 volunteers who tested the application and provided their opinions. The volunteer group consisted of individuals aged 21 to 30, including 5 individuals who work in medical fields. Feedback was collected through a Google form, which included a user satisfaction score ranging from 1 to 5, with 5 representing exceptional satisfaction and 1 indicating a poor experience.

The results of the study are depicted in Fig. 5, which displays the user satisfaction statistics. As illustrated, the majority of users rated the application at a score of 4, accounting for 67.3% of the total responses. Additionally, 18.2% of users rated the results at a score of 5, while 14.5% rated the results at a score of 3.