Federated Self-Learning with Weak Supervision for Speech Recognition

Privacy-preserving machine learning [1] can enable learning from user data while mitigating privacy risks. Federated learning (FL) [2] is one of the most popular privacy-preserving learning frameworks which involves training models on-device, with data not leaving edge devices. In FL, multiple model updates from a number of participating devices are aggregated securely on a central server at every round. FL has been demonstrated to perform well in speech applications such as speech recognition [3], keyword spotting [4], and speaker verification [5] among others. Mixed centralized and federated training was done in [6] and layer-wise representation learning in [7]. However, the aforementioned works involve training a model from scratch instead of fine-tuning a well-trained model. In addition, previous works considered static data which does not change across rounds. Differently from previous work, we consider FL settings, where the model is initialized to a well-trained model and streaming data on devices, which are not persisted across rounds. In [8], authors look at domain adaptation of ASR in a federated setting. We additionally look at incorporating weak supervision to learn from alternate sources of feedback.

Semi-supervised learning (SSL) deals with training and improving ASR using unlabelled audio, such as the audio available at devices. Unsupervised approaches such as data2vec [9] or WavLM [10] use contrastive objective functions to pretrain speech models that are then finetuned. Alternatively, a common paradigm is to use a stronger teacher model to label unlabelled data [11] however this approach cannot be applied to the resource constrained setting of on-device learning. Noisy student learning or iterative pseudo-labelling approaches [12, 13] use the ASR model to self-label clean audio with the model trained to predict the same label with augmented version of the audio. Here the audio could be additionally filtered to include cases where the model does not have low confidence. We build off the work in [14] where hybrid HMM-DNN and connectionist temporal classification (CTC) ASR models are updated using a paired teacher model updated using an exponential moving average of the student model. These methods have not been applied to recurrent neural network-transducer (RNN-T) ASR models [15] that are streaming compatible and widely used across ASR applications.

We combine self-learning in this work with weak supervision. In conversational agents, users interact across multiple turns in a session. As shown in prior works [16], later interactions can be used to determine if a request has been correctly handled. If a user cancels or repeats their request, dissatisfaction is signalled. The semantics of the terminal request can be used as feedback for the initial request. Although this is not the ground truth transcription, we use such signals to update ASR models. Users can also be prompted for an explicit feedback signal as another example for a feedback score. We use the REINFORCE [17, 18, 19, 20] framework to update models using arbitrary rewards.

Contributions: We look at incremental updates to ASR models using unlabelled audio on edge devices with federated, compute and memory constraints. We show on public and internal datasets that:

• •

  Self-learning with a paired teacher model updated through exponential moving average of ASR can be used to improve the performance of RNN-T by $10\%$ on new use cases;

• •

  Rehearsal training using historical datasets for generating model updates (pseudo-devices) at the cloud mitigates catastrophic-forgetting [21] on other test sets in self-training;

• •

  Self-learning performance is improved by including weak supervision of NLU semantics or noisy feedback scores integrated through a policy-gradient approach.

Data
Our federated continual training experiments are run from January to June $2021$. We use an internal voice-assistant dataset with de-identified utterances totalling $4500$ hours in this time period from $800$K devices. We make only a single pass through this data as one of the constraints is that persistent audio storage is not feasible.

We evaluate the models on in-house human transcribed (HT) test sets. There is no speaker overlap between the train and evaluation datasets. General comprises a $37$-hour test set in $2021$ and older test sets in $2020$. Delta comprises a $22$-hour HT test set that records a change in frequency of words in $2021$ over $2020$. The transcriptions are filtered based on $1$-gram, $2$-gram and $3$-grams that are $5$x more frequent in $2021$ than $2020$. This test set captures changes in the data distribution such as new use cases and is crucial to measure the impact of continual learning.

We also demonstrate results on models trained on public test sets. We use RNN-T models pretrained on the $960$ hour Librispeech dataset [22] and finetuned using self-learning with weak supervision on the 56 hour SLURP dataset [23]. For the public SLURP dataset, we evaluate on the test partition with $13$K utterances.

Model
The RNN-T model used contains  $60$M parameters with a $5\times 1024$ LSTM encoder, a $2\times 1024$ LSTM prediction network and a feed-forward joint network with tanh activation [24]. The input embeddings of the prediction network are $512$ dimensional. We use a $2500$ sub-word piece tokenizer [25]. The audio features are $64$ dimensional log-mel filter-bank energy features that are computed on a $25$ms window, with a $10$ms shift. SpecAugment [26] is used for the audio features. The features computed on $3$ consecutive $10$ms frames are stacked and sub-sampled to result in $192$ dimensional features at a $30$ms frame rate, provided as input to the ASR model.

A $480K$-hour pre-training dataset (where 120K hours are human transcribed and rest machine transcribed) is utilized for pre-training the baseline. Experiments using multiple losses, have equally weighted losses (no tuning). All results shown are using FedSGD with $400$ devices randomly chosen for each of $3000$ training rounds, batch size $16$ and server-side Adam optimizer. For rehearsal training, $40$ cloud pseudo-devices additionally used with historic transcribed data.

Metric
The performance of these models on the voice-assistant data is measured in terms of relative word error rate reduction (WERR) over the initial baseline model at the start of 2021. Positive WERR values represent improvements, while negative ones show degradations. Absolute WER numbers are reported on SLURP experiments.

In Table 3, the performance of self-learning coupled with weak supervision is depicted for continual learning with a single pass on the internal dataset. First, we observe that if we do not update the paired teacher model with EMA, performance on the new use case does not improve. If we only do self-learning for ASR, there is an improvement of $8.3$% on the new use case test set. Coupling this with an ASR based weak supervision (where each hypothesis gets a feedback score of the WER computed using a teacher model), we see more improvement that increases as feedback includes the NLU component. We also see similar improvement using only the NLU-based feedback-score obtained only for the served hypothesis as opposed to obtaining a score for all possible hypotheses.

Noisy feedback: Table 4 shows the result of federated learning only with noisy feedback for a single served hypothesis from ASR. Here we consider noisy feedback of the form, $M^{\prime}(\mathbf{y},z)=M(\mathbf{y},z)+(-1)^{M(\mathbf{y},z)}U^{\prime}$, where random variable $U^{\prime}\sim p(U|U\in[0,1]),U\sim\mathcal{N}(0,\sigma^{2})$ is drawn from a normal random variable with variance $\sigma^{2}$ truncated to be in the range $[0,1]$. We then add different levels of noise to measure its impact. In a noisy version of a binary feedback score,

where $\mu=\mathrm{E}[U^{\prime}]$. Thus if the mean is less than $0.5$, gradient update with the noisy feedback, in expectation, is in the same direction as the gradient update with the true feedback. We demonstrate that even at a high level of noise of $\sigma=0.4$ we are still able to improve the model on the delta dataset significantly.

EMA hyperparameters and rehearsal training: In Table 5, we first see the impact of rehearsal training on mitigating catastrophic forgetting - we observe reduced regression on the older 2020 test set at the expense of performance of new Delta test sets. Delta test set results are not comparable across prior tables as amount of computation, catastrophic forgetting differ. We also study the impact of EMA hyperparameters, higher $\delta$ implies lower weight to new updates and update frequency $u$ determines how often the teacher model is updated. Improved performance is seen for frequent updates with a lower EMA value. We also observed training diverging when the teacher model is updated to the student model after each step, suggesting that an error feedback loop takes place.

Acknowledgments: We thank Gurpreet, Aaron, Buddha, Bach, Harish, Ehry and Shehzad for helpful discussions.