Codec-ASR: Training Performant Automatic Speech Recognition Systems with Discrete Speech Representations

Audio codecs capture details of the audio signal using discrete codes at a low bitrate, and are used for speech representation in various tasks, efficient data transmission, and general data compression. Here we consider two types of NACs, operating either on the time-domain signal or on a spectral domain. Figure 1 depicts the general architecture of the considered codecs.

Both TD-NAC and Mel-NAC are trained on the Libri-Light dataset [33] with sample rate 16 kHz. TD-NAC models use an encoder with downsampling rates of $\left\{2,4,5,5\right\}$, resulting in $f_{\text{enc}}=80$ Hz. Both RVQ and FSQ quantizers use $N_{\text{cb}}=8$ codebooks with $D_{\text{cb}}\approx 2^{10}$, resulting in a bitrate of 6.4 kbps. The TD-NAC decoder upsamples in the reverse order of $\left\{5,5,4,2\right\}$ to obtain the reconstructed audio signal. The model is trained on examples with one second of audio. Mel-NAC models use mel-filterbank with a frame length of 1024 samples and frame shift of 256 samples, resulting in $f_{\text{enc}}=62.5$ Hz. Using the same quantizer setup as for TD-NAC, this results in a bitrate of 5 kbps. The Mel-NAC decoder upsamples at rates of $\left\{8,4,4,2\right\}$ to obtain the reconstructed audio signal. The model is trained on examples with 0.512 seconds of audio. All NAC models are trained end-to-end using time-domain loss, discriminative loss, and frequency-domain loss, similar to [22] with equal weights for frequency and discriminative loss and $0.1$ weight for time-domain loss. Model sizes depending on the corresponding quantizer are provided in Table 1. The models are trained on eight NVIDIA V100 GPUs for 130k steps with the AdamW optimizer with a learning rate of $10^{-4}$. A StepLR scheduler with a step size of 1 and gamma of 0.999996 is employed for learning rate decay.

The ASR models presented in the paper adopt the FastConformer Transducer large architecture [34] with 114 M parameters. The encoder consists of 17 layers, with a model dimension of 512. We used 256 channels in sub-sampling module and a kernel size of 9 in convolution module. A single layer RNN-T with hidden dimension of 640 is used for decoder. We maintain the embedding layer output dimension $D_{\text{emb}}$ (Section 2.2.1) at 128 and set $\alpha$ (Section 2.2.4) to 5 across all experiments to ensure equitable comparison. The ASR models are trained on the LibriSpeech corpus [35], encompassing 960 hours of English speech data. Evaluation of ASR model performance is conducted using the standard ’clean’ and ’other’ sets of dev and test partitions from the LibriSpeech dataset. We use a Sentencepiece Byte Pair Encoding (BPE) [36] tokenizer with a vocabulary size of 1024, trained on the text data from the LibriSpeech training set. All ASR models have been trained for 100k updates on two nodes with eight NVIDIA A100 80GB GPUs using a batch size of 32 on each GPU. We use AdamW with a peak learning rate of $2\cdot 10^{-3}$, 15k warmup steps with Cosine annealing, minimum learning rate of $10^{-6}$ and weight decay of $10^{-3}$.

The experiments are designed to study and understand four major components of the pipeline: (i) role of the NAC type, i.e., TD-NAC vs Mel-NAC, (ii) role of quantizers in NAC, i.e., RVQ vs FSQ, (iii) effect of code aggregation strategies, (iv) performance improvements of codec ASR systems with pipeline optimizations. We also setup strong baselines in the form of the traditional Mel-Spectrogram features as well as the widely used Encodec audio codec [22]. All other components like $D_{\text{emb}}$, ASR model size, ASR training data, and tokenizer have been kept constant to facilitate an unbiased study towards the role played by the above highlighted four components. Word error rate (WER) metric is used to evaluate the performance of the ASR models.

To investigate these components’ effects, we first train a TD-NAC model with RVQ following the specifications outlined in Section 3.1. Utilizing features from this audio codec as input, we establish our baseline ASR pipeline, employing parameters detailed in Section 3.2, yielding the baseline performance noted in the first row of Table 2. Subsequently, we adapt the ASR pipeline to initialize the embedding layer with codebooks learned from the trained NAC (Section 2.2.1), maintaining other pipeline components unchanged. With this setup, we train another ASR system and report it’s performance in the second row of Table 2. Likewise, we progressively integrate spectrogram augmentation and noisy embedding training into the pipeline. Notably, codebook initialization of the embedding layer significantly enhances the ASR model’s performance, with more than 10% absolute WER improvement across all the evaluation sets. Spectrogram Augmentation aids in enhancing the model’s noise robustness, as reflected by more than 6% absolute WER improvement on the noisy ’other’ sets. Noisy embedding training is able to even further improve this noise robustness of the model. Consequently, for all subsequent experiments, we incorporate all three components - codebook initialization, spectrogram augmentation, and noisy embedding training - into the training pipeline.

To assess the influence of the code aggregation strategy on the ASR+NAC model pipeline, we build up on the baseline setting as motivated in Section 4.1: TD-NAC model with RVQ, FastConformer-RNNT ASR model with embedding layer initialized with the learnt codebooks, SpecAug, and noisy embedding training. Two models are trained: one utilizing stacking for aggregating code embeddings and the other employing averaging (refer to Section 2.2.2). The performance of these models are reported in rows $6$ and $7$ of Table 3. Notably, the averaging strategy yields significantly superior WER performance compared to stacking. It’s worth noting that the embedding dimension $D_{\text{emb}}$ (as discussed in Section 2.2.1) remained fixed at 128 for both runs and the results might change with an increase in the embedding dimension. However, to ensure a fair comparison and assess the optimal configuration within the described setup, the embedding dimension was kept constant.

Despite the noted performance, stack remains the preferred aggregation scheme for all our NAC-FSQ systems. This choice is informed by the realization that different FSQ codebooks quantize distinct segments of the encoder output, whereas the RVQ codebooks encode residuals of the same vector.

We proceed to examine and compare TD-NAC with Mel-NAC, assessing their influence on downstream ASR tasks. Owing to the distinct down-sampling structures and rates outlined in Section 2.1, the compared TD-NAC operates at a bit-rate of 6.4 kbps, whereas Mel-NAC operates at 5 kbps. The remainder of the ASR pipeline remains constant, incorporating insights from Section 4.2, and we compare both RVQ and FSQ versions of the codecs. The results of these ablations are presented in the last four rows of Table 3. Notably, TD-NAC demonstrates slightly better performance compared to Mel-NAC across all considered ASR eval sets. This finding is intriguing, given that Mel-NAC outperforms TD-NAC for TTS tasks [29]. Hence, the selection of the NAC should consider the downstream task.

Furthermore, we observe that the presented TD-NAC with RVQ and only 8 codebooks outperforms Encodec with 4, 8, and even 16 codebooks, while maintaining all other parameters such as codebook size and ASR system parameter counts constant. The performance of the TD-NAC system with a bit-rate of only 6.4 kbps closely matches that of Encodec with 24 kbps (utilizing all 32 codebooks). We have open-sourced the weights (audio_ codec_16khz_small) and code³³3https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/Audio_Codec_Training.ipynb for this codec model so that it can be utilized by and be built upon by the community.

Finally, we study the effect of quantization schemes on downstream ASR performance. Analysis of the last four rows of Table 3 reveals that FSQ detrimentally affects ASR performance, particularly on the noisy ’other’ sets. We hypothesize this happens because of the fixed finite level encoding scheme utilized by FSQ, which poses challenges in modeling noisy data.

Our pipeline uses TD-NAC model with RVQ, as detailed in Section 3.1, that obtained the best performance in the experiments summarized in Table 3. The NAC model was not retrained and we use the same setup as in Section 3.1. For ASR we use the FastConformer-RNNT model described in Section 3.2 along with avg code aggregation strategy, codebook initialization of the embedding layer, SpecAug and noisy embedding training, based on Section 4. As per the challenge requirements, the ASR model is trained on the LibriSpeech-clean-100 subset (100 hrs) along with the ML-SUPERB 1h set (222 hrs) which contains data from 143 languages. The combined data has 6280 unique characters.

We compare the performance of our NAC+ASR pipeline with the baseline system [37] on the ML-SUPERB 1h test set which consists of 45 hours of unseen speech. Table 4 presents the Character Error Rate (CER) metric for both systems. It can be observed that our system with 21% CER significantly outperforms the challenge baseline. Moreover, our system surpasses the SOTA performance achieved by the XLSR-128 model, which reported a CER of 22% [27], despite being smaller in size and pretrained on significantly less data. This competitive CER underscores the effectiveness of the proposed NAC+ASR pipeline not only in monolingual scenarios (cf. Table 3) but also in multilingual settings encompassing over 100 languages.