UNFUSED : UNsupervised Finetuning Using SElf supervised Distillation

Main Contributions: We present UnFuSeD, a new framework to improve downstream audio classification performance in low-resource labeled data settings leveraging SSL. Unlike all prior systems in literature, UnFuSeD does not directly fine-tune an SSL pre-trained model but uses it to extract and cluster audio features to generate pseudo-labels on a downstream task dataset which is then used to perform un-supervised fine-tuning. More precisely, we perform a step of self distillation, guided by the generated pseudo-labels, on a randomly initialized convnet encoder, divided into student and teacher encoders. Finally, post unsupervised fine-tuning, we perform supervised fine-tuning and evaluate downstream task performance on our model linear evaluation setup. Additionally, to pre-train our encoder using SSL, we propose a novel SSL algorithm by modifying over DECAR [13]. Fig.1 shows a clear pictorial representation of our complete training process. We emphasize that UnFuSeD changes the paradigm in which SSL is leveraged to tackle data scarcity and improve downstream task performance. In practice, UnFuSeD achieves state-of-the-art (SOTA) performance on the LAPE Benchmark [7] with an encoder with $\approx$ 40% fewer parameters than the current SOTA model on LAPE.

Self Supervised Learning in Speech and Audio. The past decade has seen massive success in self-supervised learning in vision (CV), speech (SLP), and text (NLP), pushing the boundaries of low-resource representation learning for downstream classification [1, 4]. The most common systems for SSL with speech solve a Masked Acoustic Modelling (MAM) task, either using contrastive learning [1], frame reconstruction [14], or pseudo-label prediction [15]. However, recent research has shown that solving MAM makes the model representations mimic human articulatory responses [16], thus making it unsuitable for non-speech tasks. Thus, in the recent past, researchers have proposed novel systems to learn audio representations that can generalize over both speech and non-speech tasks. Following SSL in speech, these systems either solve a contrastive learning-based instance discrimination task [8], a clustering-based pseudo-label prediction task [13], or a reconstruction task [9]. Knowledge distillation has shown great success in CV, and NLP with major applications in model compression [17]. In a supervised setting, researchers have explored knowledge distillation for automatic speech recognition, speech emotion recognition, and speaker verification [18]. DistillHubert [19] was the first work on distilling SSL-based speech models and performs layer-wise knowledge-distillation (KD) to compress a full HuBERT [15]. On the other hand, when both encoder architectures are the same, this is known as self distillation (SD) [20], and has shown impressive results with the student often outperforming the original teacher. However, to our knowledge, no existing work leverages SD for general-purpose self-supervised audio representation learning, and we are the first to explore this through UnFuSeD.

(1) Upstream SSL Pre-training. Let $X_{pre}$ be an unlabeled dataset of size $J$ where $X_{pre}$ = $\{x_{1},\cdots,x_{j},\cdots,x_{J}\}$ . In our case, here $J$ = 0.25 million, following the exact pre-training setup proposed by the LAPE benchmark. Our primary aim is to learn general-purpose audio representations from this unlabeled audio dataset. To achieve this, we use a simple convnet-based architecture [21, 22], popular in prior-art [7, 9] for a fair comparison. For upstream SSL pre-training, we propose DECAR-v2, an improved version of DECAR [13], based on findings in [23]. DECAR-v2 has two main steps or phases: (a) Assignment Phase and the (b) Training Phase.
(a) Assignment Phase: The primary purpose of this phase is to obtain “pseudo-labels” $\mathbf{q}$ for every unlabelled audio sample $x\in X_{pre}$. To achieve this, we first store all the embeddings $\mathbf{g_{\tilde{x}}}$ obtained from our convnet projection head $h_{proj}$ in memory for the entire $X_{pre}$. After this, we apply Spherical K-means to cluster and get the “pseudo-labels” $\mathbf{q}$ for every $x$ as follows: $\min_{\mathbf{C}\in\mathbb{R}^{d\times K}}\frac{1}{N}\sum_{n=1}^{N}\min_{\mathbf{q}}-\mathbf{g_{\tilde{x}}}^{\top}\mathbf{Cq}$ where C is the Centroid matrix. Both $\mathbf{g_{\tilde{x}}}$ and columns of C are $l_{2}$ normalized. $K$ represents the number of clusters, and $\tilde{x}\approx x$ is an augmented and sampled version of the original audio sample. Additionally, for ConvNet training stability, we keep the prototype head $h_{prot}$ parameters frozen throughout pre-training, and at the end of every assignment phase, the parameters of $h_{prot}$ are replaced by C.
(b) Training Phase: We train the network using supervision from the “pseudo-labels” $\mathbf{q}$ obtained from the assignment phase. To do this, we first obtain the prediction $\mathbf{p}$ using $softmax(z)$ where $z$ is the output of the $h_{prot}$. Post this step; we minimize the multinomial logistic loss between $\mathbf{p}$ and $\mathbf{q}$ with: $\ell(\mathbf{p},\mathbf{q})=-\sum_{k}\mathbf{q}^{(k)}\log\mathbf{p}^{(k)}$. The “pseudo-labels” are kept fixed during the training phase and updated for the entire X only once every epoch during the assignment phase. Similar to [23], the assignment phase and training phase take place in isolation only at the first epoch, after which we use the embeddings $\mathbf{g_{\tilde{x}}}$ obtained from the previous epoch. These embeddings are stored in memory at every iteration of an epoch right after the back-propagation step.

(2) Unsupervised Downstream Fine-tuning. After SSL pre-training, we don’t fine-tune the pre-trained convnet $f_{pre}$ on the target task dataset directly but instead, use it for unsupervised fine-tuning on a randomly initialized convnet $f_{sd}$. We call this step unsupervised fine-tuning as we use the target task dataset but without using its actual labels. Let $D_{targ}=\{X_{targ},Y_{targ}\}$ be the target task labeled dataset of size $I$ where $Y_{targ}$ are labels associated with audio samples $X_{targ}$. For unsupervised fine-tuning, we first generate $Y_{pseduo}$ by extracting and clustering representations obtained on passing $X_{targ}$ through $f_{pre}$. DECAR-v2 generates clusterable embeddings, which helps in the $Y_{pseudo}$ generation. We then use $Y_{pseduo}$ to perform self-distillation on $f_{sd}$. We first divide $f_{sd}$, which follows a similar architecture to $f_{pre}$, into a student ($f^{s}_{sd}$) and teacher network ($f^{t}_{sd}$). $f_{sd}$ has 4 individual blocks, where the first 3 make $f^{s}_{sd}$ and the last block makes $f^{t}_{sd}$. For more details on the architecture of $f^{s}_{sd}$, we refer our readers to [7, 21]. For self-distillation, we treat each block $b_{i}$ as a separate classifier and add a linear transform $h^{i}_{proj}$ to $b_{i}$ to solve three losses parallelly, KL-divergence $L_{kl}$, Mean-Square error $L_{mse}$ and Cross Entorpy $L_{ce}$. $L_{ce}$ ensures that the student blocks correctly classify the pseudo labels $Y_{pseduo}$ and thus utilize the weak supervision knowledge hidden in them. $L_{mse}$ ensures that knowledge of the deepest layers is leveraged to improve feature extraction in shallow layers. $L_{kl}$ ensures that the classification results of student classifiers are similar to that of the teacher classifier. Finally, to optimize our network, we use a weighted average of $L_{kl}$, $L_{mse}$ and $L_{ce}$, which we weigh by $\alpha$, $\beta$ as shown:

(3) Supervised Downstream Fine-tuning Post unsupervised downstream fine-tuning, we do supervised downstream fine-tuning on the student model $f_{sd}$ using ${D}_{targ}$. For a fair comparison with prior-art in this space, we don’t train all the layers of our model and instead just train a task-specific linear head added to the encoder. This method of training is also known as linear evaluation and proves to be an effective technique for evaluating learned audio representations.

Datasets. In our experiments, we use the exact same upstream and downstream training setups proposed by LAPE [7]. For SSL-based pre-training, we use a balanced subset of 10% of the complete AudioSet (0.2 million) and the FSD50K [24]. For downstream tasks (DT), we evaluate our learned representations on LibriSpeech (LBS) [25] and VoxCeleb (VC) [26] for speaker identification, Speech Commands (SC) v1 and v2 [27] for keyword spotting, VoxForge (VF) [12] for language identification, IEMOCAP (IC) [28] for speech emotion recognition, NSynth [29] for TUT Urban [6] and US8K [30] for acoustic event classification and finally Bird Song Detection (BSD) [31].

Hyperparameter Tuning. For SSL Pre-training (DECAR-v2), we find the optimal values for the number of clusters as 512, learning rate as 0.005, batch size as 512, and number of epochs as 100. Projector $h_{proj}$ performs a $\mathbb{R}^{2048}\rightarrow\mathbb{R}^{512}$ non-linear transformation using multiple linear-layers. For Unsupervised Fine-tuning, we use the learning rate as 0.007, batch size as 512, number of epochs as 50, $\alpha$ as 0.7, and $\beta$ as 0.003. $h_{cf}$ performs a $\mathbb{R}^{2048}\rightarrow\mathbb{R}^{t}$ linear transformation, where $t$ is number of classes in target dataset. Projectors $h^{1}_{proj}$, $h^{2}_{proj}$ and $h^{3}_{proj}$ perform $\mathbb{R}^{2048}\rightarrow\mathbb{R}^{t}$, $\mathbb{R}^{1024}\rightarrow\mathbb{R}^{t}$ and $\mathbb{R}^{512}\rightarrow\mathbb{R}^{t}$ non-linear transformations respectively. Finally, for Linear Evaluation, we use the learning rate as 0.001, batch size as 32, and number of epochs as 50. All the hyperparameter choices were made based on an extensive grid search while considering the average performance across all the downstream tasks.