PyChain: A Fully Parallelized PyTorch Implementation of LF-MMI for End-to-End ASR

For one utterance, MMI objective can be formulated as:

where →X is the input frames sequence, →W_r is the gold transcription of →X, and →^W is any possible transcription.

In LF-MMI [11], a n-gram phone language model (LM) is integrated with the acoustic part to encode all possible word sequences into one single HMM graph called denominator graph $\mathbb{G}_{den}$. Thus, we can replace the denominator part in Eq. (1) with $P(\vec{X}\rvert\mathbb{G}_{den})$.

Similarly, the numerator part in Eq. (1) can be replaced with $P(\vec{X}\rvert\mathbb{G}_{num})$ where $\mathbb{G}_{num}$ is the numerator graph generated by composing the true transcript to the denominator graph $\mathbb{G}_{den}$. We will use numerator and denominator for short in this paper.

Extending to the corpus or batch level, we can get the final LF-MMI loss function as:

Another big advantage of MMI loss function is that, although seemly complicated, its derivatives can be finally reduced to the difference of occupation probability between numerator and denominator graphs[]:

where $\vec{y}_{(u)}$ is the network output of the $u$-th utterance, which we interpret as the log-likelihood for each state. And $\gamma_{num}^{(u)}(s)$ and $\gamma_{den}^{(u)}(s)$ are the state occupation probability of state $s$ in the numerator and denominator graphs respectively, calculated by the Forward-Backward algorithm in HMM.

As shown in Figure 1, we take advantage of both Kaldi and PyTorch to form a full recipe for an end-to-end ASR task. We do data preparations and final decoding in Kaldi for efficiency and consistency, but all other parts in PyTorch. Data preparation includes both feature extraction (e.g. MFCC) and numerator/denominator graph (FSTs) generation. Please note that, because there is no need for alignment in the E2E LF-MMI, we will not do any HMM-GMM pre-training in this stage. After all data is prepared, we go to PyTorch for dataloading and network training, which we will show in details below. After the model is trained and saved, we will load the data and model in PyTorch and then do a forward pass to get the output (posterior). They will be either piped to Kaldi for decoding on the fly, or dumped to the disk and then decoded afterwards.

As in many other ASR toolkits [8, 9], we use Kaldi for data preparation. Both acoustic features and numerator/denominator graphs (FSTs) will be saved in scp/ark format by Kaldi. We use kaldi_io⁶⁶6https://github.com/vesis84/kaldi-io-for-python to read matrix data like input features as numpy arrays and then transform them into PyTorch Tensors. For FSTs, we write our own functions in C++ based on OpenFST [23] and then bind it with Python using pybind [24], so that all these functions can be called in Python seamlessly.

As shown in Eq. (2), HMM graphs are used as supervision in LF-MMI. We follow Kaldi’s practice of using probability distribution function (pdf) to estimate the likelihood [10] of an HMM emission. As a result, both the network output and the occupation probability in Eq. (3) will be computed with respect to a pdf-index (pdf-id) instead of an HMM state. And a typical transition in a HMM graph will be from a from-state $s_{f}$ to a to-state $s_{t}$ emitting a pdf-id $d$ with a probability $p$. In this way, a dense transition matrix of an HMM graph will be of size $(S,S,D)$ with transition probability $p$ in each cell, where $S$ is the total number of HMM states and $D$ is the total number of pdf-ids.

However, due to the heavy minimization and pruning done in LF-MMI, both numerator and denominator graphs are very sparse. In PyChain, we manage to represent numerator and denominator in a uniform way with a ChainGraph object $\mathbb{G}$ that is similar to the COO (Coordinate list) format for a sparse matrix:

• •

  A transition $\vec{T}^{i}=\begin{bmatrix}s_{f}^{i},&s_{t}^{i},&d^{i}\end{bmatrix}$ where $i$ is the index of the transition, $s_{f}$ is the from-state and $s_{t}$ is the to-state. $d$ is the pdf-id on this transition. It basically denotes the coordinate of a transition.

• •

  The forward transition matrix $F_{T}=\begin{bmatrix}\vec{T}^{i}\end{bmatrix}$ of size $(I,3)$ and forward transition probability vector $F_{p}=\begin{bmatrix}p^{i}\end{bmatrix}$ of length $I$ where $I$ is the total number of transitions in a graph. They act like the coordinate list and value vector for a COO sparse matrix. They are denoted as $F$ (meaning the ”forward”) as they are sorted in ascending order of the from-state $s_{f}$.

• •

  The forward indices matrix $F_{I}=\begin{bmatrix}(I_{start}^{s},I_{end}^{s})\end{bmatrix}$ of size $(S,2)$ where $I_{start}^{s}$ and $I_{end}^{s}$ denote the start and end row index in $F_{T}$, between which are all transitions from $s_{f}$.

Similarly, we have another three tensors for backward transitions, namely, $B_{T},B_{P}$ and $B_{I}$. They are equivalent to these forward ones except that they are sorted by $s_{t}$. They are arranged in this way for a quick indexing of any transition by its from/to state.

In the end, there is a unique initial state $s_{init}$ for each graph and a final probability vector $\vec{p_{final}}=\begin{bmatrix}p^{s}\end{bmatrix}$ of length $S$ where $p^{s}$ is the probability of state $s$ being the final state in this HMM graph.

As in most cases, LF-MMI will be trained with batches of utterance, we extend ChainGraph to its batch version, ChainGraphBatch. As its name suggests, it contains a batch of graphs and can be initialized by either a list of different ChainGraph objects (for numerators), or by a single ChainGraph object (for denominator). It has exact the same type of tensors as ChainGraph does, except that there is one more dimension for each of these tensors representing the batch dimension. And the default zero-padding is employed here when the sizes of these tensors differ.

We follow the basic routine suggested by PyTorch custom C++ and CUDA extensions guide to write kernels for ChainFunction.

Our implementation has following key features:

• •

  Both numerator and denominator use the same piece of codes for forward-backward computation with full support for CPU/GPU computation.

• •

  The computation is done in probability space instead of log-probability space by utilizing leaky HMM [11] to solve underflow issues for both numerator and denominator.

• •

  It supports variable lengths of sequences without doing mandatory silence padding or speed perturbation [25] that was required by [13, 14].

A detailed description of our implementation is shown in Algorithm 1. We only present Forward algorithm as Backward algorithm is similar to the forward pass and we will not repeat here. Also, for simplicity and clarity, we omit the leaky HMM and readers are referred to [11] for details.

In Algorithm 1, $L$ denotes the log-likelihood of pdf-ids along the sequences as the output from a neural network. It has size of $(B,T,D)$, where B is the batch size, T is the length and D is the total number of pdf-ids. We get a batch of variable lengths sequences by firstly sorting these sequences by their lengths in descending order and then do the zero-padding to the right. We use $B_{v}(t)$ to denote the valid batch size at each sequence step $t$ before any padding. For example, if there are 3 sequences of lengths $(100,99,98)$ in a batch, $B_{v}$ would be $[3,3,\cdots,2,1]$. Forward trellis is represented by $\alpha$ of size $(B,T+1,S)$. Finally, we sum over the probability of each state in the final step of $\alpha$ to get the output $O$ of the forward algorithm, which is $\log P(\vec{X}\rvert\mathbb{G})$ in Eq. (2).

The 40-dimensional MFCC extracted from 25 ms frames every 10 ms are used as input features. They are then normalized on a per-speaker basis to have zero mean. For numerator and denominator graphs generation, we adopt the best setting in [13, 14]. The phone language model for the denominator graph is estimated using the training transcriptions (choosing a random pronunciation for words with alternative pronunciations in the phoneme-based setting), after inserting silence phones with probability 0.2 between the words and with probability 0.8 at the beginning and end of the sentences. We use a trivial full biphone tree without CD (context-dependency) modeling and a 2-state-skip HMM topology.

As one of our motivations for this work is to show the generality of the LF-MMI training outside Kaldi, we use common components to build our model, which only includes 1D dilated convolution (CNN or TDNN) [26], batch normalization [27], ReLU [28] and dropout [29]. They are stacked 6 times in the sequence of conv-BN-ReLU-dropout with residual connections and finally followed by a fully connected layer. The network is of input dimension 40 (MFCC) and output dimension 84 (number of pdf-ids). The hidden dimension is 640 for all other layers. The dropout rate is set to 0.2, and the convolutional layers are of kernel sizes of $(3,3,3,3,3,3)$, strides of $(1,1,1,1,1,3)$ (equivalent to a subsampling factor of 3 in the original LF-MMI), and dilations of $(1,1,1,3,3,3)$.⁷⁷7The detailed configuration can been found in https://github.com/freewym/espresso/blob/master/examples/asr_wsj/run_chain_e2e.sh.

For training schedule, we use Adam [30] as the optimization method. The learning rate is set to start from $10^{-3}$ and will be halved if no improvement on the validation set is seen at the end of each epoch, and finally fixed to $10^{-5}$. Similar to the findings in [14], we also find out that curriculum training [31] (i.e. training utterance from short to long) is very helpful to E2E LF-MMI training in terms of robustness and ease to converge. However, we only do curriculum training for the first 1 or 2 epochs to achieve the best randomness in training. Otherwise, we sort all sequences in the corpus by its length so that sequences with similar lengths would form a minibatch. In other words, we only do shuffle on the batch level. Finally the model with the best validation loss is selected for decoding.

We compare our results on WSJ eval92 with the original E2E LF-MMI and other SOTA end-to-end ASR systems. As shown in Table 1, we not only achieve competitive results but also use a smaller number of parameters and much simpler structure. Note that we only use n-gram LMs for decoding, while the others (except Hadian et al. [14]) use more powerful neural LMs.

Our implementation is also more efficient than the original one where extra time is wasted on the padded silence parts. Some might argue that since the Forward-Backward algorithm has a temporal dependency on its previous step, the longest sequence inside a batch will decide the computational time and thus our unequal length version will not make a difference on the efficiency. The explanation is, as shown in Algorithm 1, because we are doing parallel computation state-wise and sequence-wise, there would be at most $B\cdot S$ ($128*7,398=946,944$ in our case) concurrent jobs on a GPU at once. However, there are much fewer CUDA cores in a modern GPU (e.g. 3584 for a NVIDIA GTX 1080 Ti in this experiment) which becomes the actual bottleneck for speed. As a result, the computational time will be proportional to the number of jobs we have in total.

Flexibility is another big advantages of our variable lengths supported implementation. Instead of modifying utterance lengths by either silence padding or speed perturbation, we are able to form any utterances into a minibatch to get the maximal flexibility and randomness in training. Note that simply doing zero-padding at the end of each utterance does not work easily because these padded features may not be necessarily aligned to the final states and would probably lead to a terrible alignment and thus degrade the results.

For the completeness of comparison, we also do experiments on augmented data with a 2-fold speed perturbation (sp) as was originally required by [13] and compare the results with other end-to-end systems with data augmentation. As shown in Table 2, we again match up to the results in Kaldi but with a smaller network.