Supervised and Self-supervised Pretraining based COVID-19 detection using acoustic breathing/cough/speech signals

The diagram of the proposed classification model that we utilize in each subtask is shown in Fig. 1, which mainly consists of a pre-processing module, an encoder module and a classifier module. In the encoder module, two bi-directional LSTM (BiLSTM) [16] layers was utilized as the encoder. Each BiLSTM layer contains 128 hidden units and the dropout rate is set to be 0.1. Two fully connected feed-forward layers are utilized in the classifier module, which consists of two linear transformations with a ReLU activation in between. As the detection of COVID-19 is a 0-1 classification problem, we use the binary classification loss function for training.

For the supervised pre-training, no extra data are utilized except the officially provided DiCOVA data. The total DiCOVA dataset consist of breathing, cough and speech labeled data. We obtain a breathing/cough/speech model using labeled breathing/cough/speech data through supervised pre-training, respectively. Specifically, the breathing model $\text{Model}_{\rm breath}$ can be obtained using the labeled breathing data $X_{\rm breath}$ through supervised pre-training, and $\text{Model}_{\rm cough}$ and $\text{Model}_{\rm speech}$ can be built similarly as

	$\displaystyle\text{Model}_{\rm breath}$	$\displaystyle=\text{BiLSTM}(X_{\rm breath}),$		(1)
	$\displaystyle\text{Model}_{\rm cough}$	$\displaystyle=\text{BiLSTM}(X_{\rm cough}),$		(2)
	$\displaystyle\text{Model}_{\rm speech}$	$\displaystyle=\text{BiLSTM}(X_{\rm speech}).$		(3)

Then, we average the parameters of the three models as the average model as

$$\text{Model}_{\rm ave}={\rm ave}(\text{Model}_{\rm breath},\text{Model}_{\rm cough},\text{Model}_{\rm speech}).$$

(4)

For each task, the average model is considered as an initialization model, and we re-train the breathing/cough/speech model in the corresponding task. The pre-training process is shown in Fig. 2. As we initialize the model using average model, the encoder and the classifier are easier to find the optimal solution.

As the amount of COVID-19 audio data in each sub-task is still limited, the traditional MFCC feature might be not sufficiently representative for classification tasks. Hence, we propose to utilize a high-level feature to replace the traditional MFCC feature. It is worth mentioning that there are many effective methods that were proposed to solve these low-resource tasks [17, 18, 19, 20, 21]. In the speech domain, The wav2vec2.0 [14] model is a representative self-supervised pre-training framework for learning speech representation. The public self-supervised pre-trained wav2vec2.0 [14] model is utilized in the stage of self-supervised pre-training. We pre-train the public pre-trained wav2vec2.0 model on the DiCOVA data and use the pre-trained model as a high-level feature extractor. For example, we use DiCOVA breathing data to pre-train the wav2vec2.0 model. After finishing pre-training, we utilize the pre-trained model to extract high-level breathing feature, which is then input into the diagnosing model to replace the MFCC feature.

In order to guide the reader, we will briefly review the wav2vec 2.0 model in this section. The structure of wav2vec 2.0 model is shown in Fig. 3, including a CNN-based feature encoder $f:X\mapsto Z$ and a transformer encoder $g:Z\mapsto C$. In detail, the input raw waveform $X$ is downsampled to the latent speech representation $Z$ by the feature encoder. The transformer encoder then models the contextualized representation $C$ and extracts a high-level feature from the input $Z$. A quantization module $Z\mapsto Q$ discretizes the output of the feature encoder to $q_{t}$ as targets in the contrastive objective.

The quantization module first maps the latent speech representation $Z$ to logits $\mathbf{l}\in\mathbb{R}^{G\times V}$, given $G$ codebooks with $V$ entries.The Gumbel softmax operation [22] is then used to select discrete codebook entries in a fully differentiable way. For a given frame $Z_{t}$, we can therefore select one entry from each codebook and concatenate the resulting vectors $e_{1},...,e_{G}$ and apply a linear transformation to obtain $\mathbf{q}_{t}$. The weighted loss function is thus given by

$$L=L_{m}+\alpha L_{d}+\beta L_{f},$$

(5)

where

	$\displaystyle L_{m}$	$\displaystyle=-\log\frac{\exp({\rm sim}(\mathbf{c}_{t},\mathbf{q}_{t})/\kappa)}{\sum_{\tilde{\mathbf{q}}{\sim}\mathbf{Q}_{t}}\exp({\rm sim}(\mathbf{c}_{t},\tilde{\mathbf{q}})/\kappa)},$		(6)
	$\displaystyle L_{d}$	$\displaystyle=\frac{1}{GV}\sum_{g=1}^{G}\sum_{v=1}^{V}\overline{p}_{g,v}\log\overline{p}_{g,v},$		(7)
	$\displaystyle\overline{p}_{g,v}$	$\displaystyle=\frac{\exp(\overline{l}_{g,v}+n_{v})/\tau}{\sum_{k=1}^{V}\exp(\overline{l}_{g,k}+n_{k})/\tau}.$		(8)

It is clear that the total loss function is the weighted summation of three terms $L_{m}$, $L_{d}$ and $L_{f}$ parameterized by $\alpha$ and $\beta$. In (6), $L_{m}$ is the contrastive loss, which enables the model distinguishable between the true quantized latent speech representation $\mathbf{q}_{t}$ and a set of $K+1$ quantized candidate representations $\tilde{\mathbf{q}}\in\mathbf{Q}_{t}$. The quantized candidate representation $\tilde{\mathbf{q}}$ contains $\mathbf{q}_{t}$ and $K$ distractors, and the latter are uniformly sampled from other masked time steps of the same utterance. In (5), the diversity loss $L_{d}$ aims to increase the use of quantized codebook representation, and $L_{f}$ is an $\ell_{2}$ penalty over the outputs of the feature encoder. In (6), sim stands for the cosine similarity between two vectors and $\kappa$ is a temperature. In (7), $\overline{p}_{g,v}$ represents the probability of choosing the $v$-th codebook entry for group $g$ across a batch of utterances, where $\tau$ is a temperature. In (8), $\overline{l}_{g,v}$ stands for the average logits $\mathbf{l}$ across utterances in a batch. More details on the wav2vec 2.0 model can be found in [14].

For the supervised pre-training method, the diagnosing model using MFCC feature as input, and for the self-supervised pre-training method, the diagnosing model using high-level features extracted by the pre-trained wav2vec2.0 model. As local features and high-level features contain different information, in order to combine the superiority of different features, we use model ensemble to improve the detection performance. Specifically, for the same task we utilize the proposed supervised and self-supervised methods to train two models and obtain $\text{score}_{\rm sup}$ and $\text{score}_{\rm self-sup}$ on the test set, respectively. Then, we can ensemble the scores obtained by the two models, which is given by

$$\text{score}_{\rm emsemble}=\mu*\text{score}_{\rm sup}+(1-\mu)*\text{score}_{\rm self-sup},$$

(9)

where $\mu$ is chosen by the user. The weighted summation $\text{score}_{\rm emsemble}$ is taken as the detection probability.

For the track-4 fusion task, we utilize the best breathing, cough and speech models to obtain an approporiate fusion model as

$$\text{score}_{\rm fusion}=\theta*\text{score}_{\rm bre}+\gamma*\text{score}_{\rm cou}+\phi*\text{score}_{\rm spe},$$

(10)

where the parameter $\theta,\gamma,\phi$ need to satisfy $\theta+\gamma+\phi=1$. In principal, the breathing, cough and speech signals have a different importance in terms of the COVID-19 information, $\theta,\gamma,\phi$ should be chosen in line with the signal qualities.

The proposed model is evaluated on the DiCOVA-ICASSP 2022 challenge dataset [15], which is derived from the crowd-sourced Coswara dataset [23] and collected from volunteers with different health conditions. Volunteers were advised to record their sound in a quiet environment using a web-application. The audio streams are sampled at a sampling frequency of 44.1 kHz and in an FLAC format. The durations of the sound recordings range from about 1 second up to 29 seconds. The dataset consists of a total of 965 samples including 172 negative samples and 793 positive samples, where each sample includes cough, speech and breath sounds. Male and female patient samples are 723 and 242, respectively. The blind test set provided by the organizer includes 471 samples without labels.

In the pre-processing module, the amplitude of the raw waveform data is normalized between -1 to 1 through a normalization operator. As there are many silent segments in the speech/cough/breath sound signals, the speech activity detector (SAD) is applied to cut off these silent segments. The sound data is downsampled to 16 kHz. Forty dimensional MFCC and delta-delta coefficients are extracted with a widow of 25 msec audio samples and a hop of 10 msec. Due to the small size of the training data, we use SpecAugment [24] time-frequency mask to augment the data, as it was shown data augmentation is effective to improve the performance, particularly in low-resource cases. The time mask length is 20 and the frequency mask length is 50. In the encoder module, two BiLSTM [16] layers are utilized as the encoder. The BiLSTM layer dimension is 128 and the dropout rate is set to be 0.1. Two fully connected feed-forward layers are utilized in the classifier module, which consists of two linear transformations with a ReLU activation in between. The dimension of feed-forward layers is 256, which is finally mapped to 1-dimension for binary classification. As the detection of COVID-19 is a 0-1 classification problem, we use the binary classification loss function for training. We follow the baseline system given by the organizer to train the model for 5-fold cross-validation in each task and then decode on the official blind test set to obtain the final test results, separately. In the baseline system, The CNN model is also utilized as another baseline system in experiments.

For the supervised pre-training, different seeds are set to train different models. The pre-training model is trained using 2 RTX3090ti-24G GPUs with 50 epochs, and the total training time is about 1 hour. For the self-supervised pre-training, the wav2vec2.0 pre-training model is implemented using the fairseq toolkit [25]. The feature encoder contains seven blocks, where each block has 512 temporal convolution channels with strides (5, 2, 2, 2, 2, 2, 2) and kernel widths (10, 3, 3, 3, 3, 3, 2, 2). Thus, the interval between two sequential samples in the feature encoder output $Z$ is around 20 ms and the receptive audio field is 25 ms. The models contain 12 transformer encoder blocks with a dimension of 512, a feed forward module with a dimension of 2048 and 8 attention heads. The pre-training process is optimized with Adam [26]. We set $G$ = 2 and $V$ = 320 for the quantization module and each entry with a size of 128. The temperature $\kappa$ is set to be 0.1 and $\tau$ is annealed from 2 to 0.5 by a factor of 0.999995 over iterations. For the contrastive loss, $\alpha$ and $\beta$ are set to be 0.1 and 10, respectively. We use $K$ = 100 distractors and the total number of pre-training epochs are 200. After self-supervised pre-training, the pre-trained model parameters are frozen. Then, the pre-trained model is utilized as a high-level feature extractor, which generates high-level features from raw waveforms.

In experiments, we use 5-fold cross-validation to evaluate our supervised pre-training method, self-supervised pre-training method and model ensemble method. For comparison, we also test CNN, LSTM and official baseline model without pre-training. The official baseline comes from the website¹¹1https://competitions.codalab.org/competitions/34801#results. The obtained results are shown in Table 1 and Table 2. Experimental results show that the methods with pre-training achieve a higher AUC than the baseline. That is, pre-training can provide more information for classifiers, particularly in case the dataset is small-sized. Supervised pre-training models can achieve a great enhancement of performance compared with the baseline system. In breath/cough/speech task, supervised pre-training models increase the AUC from 84.50/74.89/84.26 to 86.41/76.05/85.02, respectively. Self-supervised pre-training methods obtain the AUC of 86.22/75.55/84.35 in the breath/cough/speech task. Despite achieving a smaller improvement in performance than supervised counterparts, self-supervised pre-training models still surpass the baseline in all tasks. Compared to the supervised and self-supervised methods, the proposed ensemble model can further improve the detection performance, which obtains the best performance for all tasks, and the corresponding AUCs are 86.72/76.36/85.21 in cough/speech/breath tasks, respectively. These results indicate that high-level feature pre-trained by wav2vec2.0 and low-level MFCC feature are complementary. For the fusion task, the final probability is calculated as a weighted summation over three best individual tasks. In case the cough/speech/breath weights are set to be 0.4/0.2/0.4, we achieve the highest AUC of 88.44 for the fusion task.