Hierarchical Softmax for End-to-End Low-resource Multilingual Speech Recognition

Based on the assumption that neighbour languages share similar token distribution, the concept of our proposed method is to generate a representation code for each token via frequency clustering. We first generate the Huffman code of each token. Formally, given the multilingual token vocabulary $V=\{t_{1},t_{2},...t_{N}\}$, where $N$ is the vocabulary size, we maintain the term frequency set $S_{p}=\{p_{t_{i}}\}_{i=1}^{N}$, where the same token in different languages are considered as one token. With $S_{p}$, we generate a Huffman tree of $V=\{t_{i}\}$ via frequency clustering and further recursively generate the Huffman code by assigning 0 to the left subtree and 1 to the right subtree.

For the ASR model, we use conformer [26] as the encoder and transformer [27] as the decoder. For the decoder, we replace the vanilla softmax with H-Softmax.

H-Softmax organizes the output vocabulary into a tree where the leaves are the vocabulary tokens, and the intermediate nodes are latent variables [22]. We use the Huffman tree generated in subsection 2.1 as the tree for H-Softmax that there is a unique path from the root to each token, forming a complete binary tree. We follow [22] and apply the binary tree H-Softmax. The decoding procedure is transformed into predicting one leaf node of the binary tree at each timestep. Each leaf node, which represents a token, could be reached by a path from the root through the inner nodes. Given the transformer output hidden state $h$ and trainable node representation vectors $\{r_{i}\}$, the final possibility of a leaf node $w_{i}$ could be represented as:

While decomposing the Softmax to a binary tree H-Softmax reduces the decoding time complexity from O(V) to O(log(V)), and the train time complexity remains O(Vlog(V)). Since previous H-Softmax implementation[23] is on CPU, considering the order of magnitude difference between CPU’s and GPU’s FLOPs, the challenge of improving the efficiency of model training lies in designing an implementation of H-Softmax for GPU training. We propose a vectorization algorithm to accelerate training and decoding procedures on GPU.

To explain the core idea of our vectorization algorithm, we show a typical H-Softmax tree structure in Fig 2. Then, we can vectorize log-probability calculations from Eq 1 to the followings:

where $Sign$ is a 3 by 2 matrix of $\sigma$s’ signs, $Bias$ is a 3 by 2 matrix of $\sigma$s’ biases and $p$ is the result vector of the inner product between node vectors and $h$. After building the Huffman tree, the $Sign$ and $Bias$ matrices are fixed. So in the training stage, leaf node log probabilities can be acquired only by vector operations.

For decoding, we only need leaves with the highest probabilities. To directly calculate this objective, different from training, we also develop a path-encoding-based multi-layer beam searching on GPU for H-softmax\footrefgithub to retain the time efficiency advantage of time-space complexity O(log(V)) compared to vanilla softmax’s O(V).

We sample our speech-to-text datasets from Common Voice Corpus 11.0. We selected three different linguistic groups: Romance, Slavic, and Turkic. For each group, we selected five languages from the corpus and constructed training and testing sets with the validated data of each language. With the data size of many existing datasets being around $20{\sim}30$ hours [28, 29], we downsampled the training data to the extent of 30 hours per language on average to simulate a low-resource setting. As a result, the total size of training data changed from 5492 hours to 450 hours, with the general downsampling ratio $\lambda=0.082$. Different downsampling ratios are used for different languages to counter the imbalance of data size among languages. For a set of languages $\{L_{1},...,L_{m}\}$ with their proportion $\{p_{1},...,p_{m}\}$ , the downsampling ratio for language $L_{i}$ is obtained by $\lambda_{i}=\frac{p_{i}^{\alpha-1}}{\sum_{j=1}^{m}p_{j}^{\alpha}}\lambda$, where $\alpha$ is a smooth coefficient that we set to 0.5. The size of testing sets is the lesser of 10% of the validated data and 10,000 utterances.

For speech-to-phoneme datasets, we use UCLA Phonetic Corpus [19]. The Edo-Ancient Tokyo (bin), Kele-Congo (sbc), Malayalam-Malay (mal), Creole-Cape Verde (kea), Klao-Liberia (klu), Arabic-Tunisian Spoken (aeb), Makasar-South Sulawesi (mak), Finnish-Finland (fin), Abkhaz-North Caucasian (abk) and Aceh-Sumatra (ace) are testing languages and other languages are used for training.

For acoustic features, the 80-dimensional log-Mel filterbanks (FBANK) are computed with a 25ms window and a 10ms shift. Besides, SpecAugment [30] is applied to 2 frequency masks with maximum frequency mask (F = 10) and 2 time masks with maximum time mask (T = 50) to alleviate over-fitting.

Both H-Softmax and Softmax models are trained using the same network structure. The networks are constructed using WeNet toolkit [31].¹¹1Our implementation is https://github.com/Derek-Gong/hsoftmax Two convolution sub-sampling layers with kernel size 3$\times$3 and stride 2 are used in the front of the encoder. For model parameters, we use 12 conformer layers for the encoder and 6 transformer layers for the decoder.

Adam optimizer is used with a learning rate schedule with 25,000 warm-up steps. The initial learning rate is 0.00005. 100 max epochs for speech-to-text datasets and 80 max epochs for speech-to-phoneme datasets.

We conducted two sets of experiments on speech-to-text datasets: training with languages within one language group and training with languages across language groups. We tokenized the transcriptions at the character level as we verified that it outperforms tokenizing in the sub-word level (such as BPE).²²2Preliminary experiments on Slavic language group with traditional Softmax results in a CER of 8.5% for character-level tokenization, 8.9% for BPE (vacab size = 500) and 9.6% for BPE (vacab size = 5000). As shown in Table 4, when training with languages within the same language group, our Huffman code achieves better performance (character error rate, CER%) than traditional Softmax for all three groups, which demonstrates the effectiveness of our method. Results in Table 5 show that our method can also make improvements when training with the combined data in 15 languages across different language groups, showing that our method works in a more extensive range of scenarios than we expected. In addition, our method is more robust as there are no languages with a distinctively high error rate, unlike traditional Softmax.

Our proposed model demonstrated superior performance compared to the conventional model (phone error rate, PER%) across all languages on the UCLA Phonetic corpus, as presented in Table 3. Overall, our approach outperformed the softmax baseline by a significant margin of 20.51% PER. On the ‘bin’ language, the performance gap between softmax and H-softmax was 32.33%. These results showcase the effectiveness of our method in automatically constructing universal representations for multilingual ASR and achieving zero-shot cross-lingual phoneme recognition.

We observed decoding acceleration with our proposed H-Softmax model in Table 6. Decomposing the Softmax output layer to a binary tree reduces the complexity of obtaining probability distribution from O(V) to O(log(V)), which leads to improvement in efficiency. The results also show that the sentence length (tokens per sentences) determines the decoding time difference between Softmax and H-softmax. Longer sentences take more time steps for inference, and the time difference between the two models exponentially increases.