Pretraining End-to-End Keyword Search
with Automatically Discovered Acoustic Units

Our work is based on the E2E KWS model of [12]. This model ingests a textual query in the form of a sequence of $L$ letters, $\bm{\mathbf{q}}=(q_{1},\dots,q_{L})$, and a spoken utterance of $N$ frames, $\bm{\mathbf{X}}=(\bm{\mathbf{x}}_{1},\dots,\bm{\mathbf{x}}_{N})$ and predicts the sequence $\bm{\mathbf{y}}(\bm{\mathbf{q}},\bm{\mathbf{X}})=(y_{1},\dots,y_{N})$, where each $y_{n}\in\{0,1\}$ is a binary random variable indicating the existence of the query in the $n$th frame of the document, i.e:

The model comprises a pair of encoders:

• •

  The query encoder computes a fixed-length representation, $\bm{\mathbf{e}}_{\bm{\mathbf{q}}}$, of the query by passing it through a GRU, summing the activations from the last layer of the GRU across the sequence dimension, followed by an affine projection.

• •

  The document encoder computes a down-sampled representation, of the document $\bm{\mathbf{H}}_{\bm{\mathbf{X}}}=(\bm{\mathbf{h}}_{1},\dots,\bm{\mathbf{h}}_{N})$ by passing it through a BLSTM followed by an affine transform.

The sequence $\bm{\mathbf{z}}(\bm{\mathbf{q}},\bm{\mathbf{X}})=(z_{1},\dots,z_{N})$ of occurrence probabilities, $z_{n}\coloneqq P(y_{n}=1|\bm{\mathbf{q}},\bm{\mathbf{X}})$, is then obtained via a matrix-vector product of $\bm{\mathbf{H}}_{\bm{\mathbf{X}}}$ and $\bm{\mathbf{e}}_{\bm{\mathbf{q}}}$ as:

where $\sigma(\cdot)$ is the logistic sigmoid function.

The vector of probabilities from (2) is then post-processed to obtain the timestamps in the document hypothesized to contain the query, and the corresponding confidence scores by detecting “islands” of probabilities above $0.5$. The procedure is as follows:

1. 1.

   Probabilities $z_{n}<0.5$ are zeroed-out.

2. 2.

   The resulting “islands” of non-zero elements are returned as system hypotheses, and each hypothesis’ confidence score is computed as the median probability in its time-span.

The model is trained with mini-batch gradient descent on a transcribed speech dataset. At each training step, $t$, the following objective is minimized:

where $\{\bm{\mathbf{q}}_{t,1}\dots\bm{\mathbf{q}}_{t,L}\}$ is a mini-batch of $L$ queries sampled randomly from the set of unigrams, bigrams and trigrams of the dataset; $\{\bm{\mathbf{X}}_{t,l,1},\dots,\bm{\mathbf{X}}_{t,l,M}\}$ is a set of documents sampled from the dataset such that $\{\bm{\mathbf{X}}_{t,l,1}\}$ contains $\bm{\mathbf{q}}_{t,l}$ while the other $M-1$ documents are sampled randomly; and $f(\cdot)$ is a modified binary cross-entropy function defined as:

with $\phi$ controlling the tolerance of the objective to easily-classified frames and $\lambda$ controlling the relative weighting of positive to negative frames. Following [12], we set $\lambda=5$, $\phi=0.7$ and $M=4$ in all our experiments.

The word-level alignments required for training are obtained by training an HMM-GMM ASR system on the training data and using it for forced alignment.

Acoustic unit discovery entails learning a set of units from untranscribed data. For a language, $l$, this typically involves learning a set of parameters, $\Theta^{l}=\{\bm{\mathbf{\theta}}^{l,u}\}_{u=1}^{U_{l}}$ for each unit $u$, which then allows frames of an utterance $\bm{\mathbf{X}}=(\bm{\mathbf{x}}_{1},\dots,\bm{\mathbf{x}}_{n})$ in that language to be labeled into discrete units $v_{1},\dots,v_{N}$ where each $v_{n}\in\{1,2,\dots,U_{l}\}$.

We employ H-SHMM [22], a Dirichlet-process-based Bayesian nonparametric model for AUD. In H-SHMM, each acoustic unit is a 3-state, left-to-right HMM-GMM with 4 Gaussians per state, and each $\bm{\mathbf{\theta}}^{l,u}$ is a super-vector formed by concatenating all the mean vectors, covariance matrix elements and mixture weights of the HMM-GMM. The parameters of the HMM-GMMs are constrained to dwell in a low-dimensional subspace of the full parameter space:

where $\bm{\mathbf{\eta}}^{l,u}$ is a low-dimensional²²2We set the dimensionality of $\bm{\mathbf{\eta}}^{l,u}$ to 100 in our experiments. Contrast this to the dimensionality of $\bm{\mathbf{\theta}}^{l,u}$ which is around 1000 for HMMs with MFCC features and around 25000 for HMMs with XLS-R features. embedding of the parameter, $\bm{\mathbf{W}}^{l}$ is a language-specific low-rank matrix which, along with the bias vector $\bm{\mathbf{b}}^{l}$, defines the subspace to which the parameters are constrained, and $g(\cdot)$ is a non-linear function mapping vectors from the column space of $\bm{\mathbf{W}}^{l}$ to HMM parameters, ensuring e.g. that the dimensions corresponding to each covariance matrix constitute a positive-definite matrix and that the mixture weights are non-negative and sum up to one. The subspace parameters are themselves further constrained to a $K$-dimensional “hyper”-subspace:

where $\bm{\mathbf{\alpha}}^{l}$ is a language-specific low-dimensional embedding, and $\{\bm{\mathbf{M}}_{k}\}$ and $\{\bm{\mathbf{m}}_{k}\}$ define language-agnostic template matrices and vectors whose linear combinations define the part of the space $\bm{\mathbf{W}}^{l}$ and $\bm{\mathbf{b}}^{l}$ are allowed to occupy. Thus, the subspace parameters are allowed to vary in a controlled manner from one language to another. The H-SHMM defines a distribution with trainable parameters—$\{\bm{\mathbf{M}}_{k}\},\{\bm{\mathbf{m}}_{k}\},\{\bm{\mathbf{\alpha}}^{l}\},\{\bm{\mathbf{\eta}}^{l,u}\}$—which are learned in two phases:

1. 1.

   Supervised pretraining: The model is first trained on phonetically-transcribed speech from multiple languages not including the target language. This imbues the templates $\{\bm{\mathbf{M}}_{k}\}$ and $\{\bm{\mathbf{m}}_{k}\}$ with phone-like characteristics.

2. 2.

   Acoustic unit discovery: The distributions $\{\bm{\mathbf{M}}_{k}\}$ and $\{\bm{\mathbf{m}}_{k}\}$ are kept fixed, and transferred to a target language, $l^{*}$, for which $\bm{\mathbf{\alpha}}^{l^{*}}$ and $\{\bm{\mathbf{\eta}}^{l^{*},u}\}$ are learned.

Both phases involve optimizing a variational lower bound on the log-likelihood of the data, which yields a Baum-Welch-like training procedure. Having obtained the distributions of the HMM parameters, the untranscribed speech can them be labelled with a variational analog of Viterbi decoding. Interested readers are referred to [23] for a thorough coverage of H-SMMM and its inference.

For Libri-light pretraining, we use the 360-hour set of Librispeech as untranscribed data. In the case of BNTR, we use the full 180-hour BNTR training set for pretraining.

For actual acoustic unit discovery, where we learn the acoustic units ($\{\bm{\mathbf{\alpha}}^{l}\},\{\bm{\mathbf{\eta}}^{l,u}\}$) for each target language, we use random 3000-utterance subsets from each respective language’s untranscribed data, and use the learned HMMs to transcribe the full corpus.

The default acoustic inputs to our models (both AUD and KWS) are 13-dimensional MFCC features. In addition, we consider features from a pretrained 300 million parameter XLS-R model [31]⁵⁵5https://dl.fbaipublicfiles.com/fairseq/wav2vec/xlsr2_300m.pt, from which we use the output of the 15th layer, shown in [23] to yield considerably better AUD performance. Note that, due to computational constraints, we only use the XLS-R model as a feature extractor and we do not finetune it.

We base the architecture of our model on [12]⁶⁶6Code available at: https://github.com/bolajiy/golden-retriever. The query encoder is a network with a 32-dimensional embedding layer for computing vector representations of each input grapheme, followed by 2 bidirectional GRU layers with 256 output units per direction per layer, and a 400-dimensional output projection layer whose outputs are summed along the sequence dimension to obtain the vectoral query representation.

The document encoder has 6 BLSTM layers with 512-dimensional output per direction per layer, followed by a 400-dimensional output layer. We apply dropout of 0.4 between successive BLSTM layers, and down-sample by a factor of 2 after the fourth BLSTM layer. This results in document encodings with frame durations of 40ms for XLS-R features and 20ms for MFCC.

The H-SHMM use 3-state HMM-GMMs with 4 diagonal-covariance Gaussians per HMM state, Dirichlet process with truncation parameter of 1, 100-dimensional unit embeddings ($\bm{\mathbf{\eta}}^{l,u}$) and 5-dimensional language embeddings ($\bm{\mathbf{\alpha}}^{l}$).

The neural networks are trained with the Adam optimizer [35]. For pretraining, we use a cosine decay schedule with peak learning rate of 5e-4 and final learning rate of 1e-7, and train for 200k steps using mini-batch size of 256, with 10k warmup steps. For finetuning (and the baseline with no pretraining), we train with the same step-based learning rate scheduling as [12] using mini-batch size of 32. This starts with a learning rate of 0.002, halved whenever the validation loss (computed over 10% of the training queries) does not improve over 4 epochs. The training is stopped when validation loss does not improve for 10 epochs.