Vector-Quantized Autoregressive Predictive Coding

Yu-An Chung, Hao Tang, James Glass

Introduction

Many high-level properties of speech, e.g., phonetic content and speaker characteristics, are not easily accessibleFollowing other recent works , in this study we use linear separability (or separability with a shallow network) to define the accessibility of information for a downstream task. without sufficiently powerful transformations from the surface features such as audio waveforms and spectrograms. Speech representation learning aims to search for a transformation from the surface features that better reveals these properties to downstream tasks.

Recently, self-supervised learning—a paradigm that treats the input itself or modifications of the input as learning targets—has obtained promising results for learning such transformations . These methods, mostly inspired by the techniques for pre-training NLP models , learn speech representations by either inferring future information conditioned on historical audio , or predicting masked parts of input audio . The resulting representations, obtained without requiring any additional labeled data, are able to outperform the surface features across downstream tasks such as speech recognition, speech translation, speaker identification, and emotion recognition .

Autoregressive Predictive Coding (APC) is one of the recent self-supervised speech representation models. APC defines a prediction task that trains an autoregressive neural model (e.g., a unidirectional RNN) to predict a future frame considering the past context. Although the learned representations contain highly accessible phonetic and speaker information, the reason why this seemingly unrelated self-supervised objective produces such a representation remains unclear. In this work, we aim to provide an explanation, investigating the constituents that lead to low objective values, and connect them with the properties of the learned representations.

Our approach is to study the properties of the learned representation as we limit the model capacity. The models with limited capacity are forced to retain information to achieve maximal prediction, thereby allowing us to study the constituents of the task and the learned representations. Several options are available to obtain a spectrum of models with different capacity, including reducing the number of layers, reducing the hidden layer size, or enforcing a bottleneck along the feed-forward process. The impact of different numbers of hidden layers has been studied in prior work . Regardless, it is difficult to quantify the amount of information by changing the number of layers, changing the hidden layer size, or using low-rank matrices to enforce bottlenecks. In this work, we study the use of vector quantization (VQ), where the amount of information (i.e., bits required to transmit the codebook and the sequence of codes) can be exactly quantified, to control the capacity of the models.

Recent studies on VQ for representation learning, mostly motivated by the discrete nature of phonetic units, attempt to show that enforcing the quantization leads to a better representation for acoustic unit discovery and automatic speech recognition . In contrast, our goal is not to discover the discrete units in speech. We treat VQ as a general approach to limit the model capacity, and study its impact on the information encoded in the learned representations.

Proposed Models

Given a sequence of acoustic feature vectors (x1,x2,...,xt)(x_{1},x_{2},...,x_{t}) as context, APC incorporates an autoregressive neural model gARg_{AR}, e.g., a unidirectional RNN or a Transformer decoder , to summarize the sequence for predicting a future frame xt+nx_{t+n} that is nn steps ahead of xtx_{t}. Let yty_{t} denote the prediction of gARg_{AR} at time tt. In practice, for a speech utterance x=(x1,x2,...,xT)\mathbf{x}=(x_{1},x_{2},...,x_{T}), where TT is the sequence length, gARg_{AR} is trained by minimizing the L1L1 frame-wise loss between the predicted sequence (y1,y2,...,yT−n)(y_{1},y_{2},...,y_{T-n}) and the target sequence (x1+n,x2+n,...,xT)(x_{1+n},x_{2+n},...,x_{T}):

Once gARg_{AR} is trained, one can extract features by taking its hidden representations, e.g., the last layer output, to replace the surface features as the new input to downstream models.

2 VQ-APC

The codebook size and the code dimension of a VQ layer control the amount of information from the previous layer flowing to the next, and changing either of these two factors allows us to explicitly control the capacity of the models. As we limit the model capacity, the model is forced to retain information to achieve maximal prediction. By studying a sequence of increasingly limited models, we are able to reveal the constituents of the prediction task and the learned representations.

Experiments

We conduct experiments to study the properties of the learned representations and their connection to the self-supervised objective. Since VQ layers are known to significantly disrupt model training, we first examine where VQ layer(s) should be inserted. Next, by using phonetic and speaker classification as probing tasks, we show the model’s preference in preserving speech information as its capacity becomes constrained. We then visualize the learned VQ codes to show the presence and absence of phonetic information and the correspondence between codes and phones. Finally, we compare VQ-APC with other self-supervised speech representation models.

Probing tasks. We consider phonetic and speaker classification for measuring the accessibility of the phonetic and speaker information contained in the representations, respectively. Both experiments are carried out on the Wall Street Journal corpus (WSJ) . For phonetic classification, the goal is to correctly classify each frame in an utterance into one of the 42 phones. The phone alignments are generated with a speaker adapted GMM-HMM model. We follow the standard split of WSJ, using 90% of si284 for training, the rest 10% for validation, and reporting phone error rate on dev93. For speaker classification, the goal is to correctly predict the speaker identity of an utterance. We follow and consider a 259-class classification task where each class corresponds to an unique speaker, using 80% of si284 for training, the other 10% for validation, and reporting classification error rate on the rest 10%. We note that speaker classification is not a typical task for WSJ, and only serves as a sanity check for the presence of speaker information (and its potentially correlated channel information) . The classifier for both tasks is a linear logistic regression that takes the features extracted from the self-supervised models as input. For speaker classification, the features from the same utterance are averaged before being fed to the classifier. All self-supervised models are kept frozen when training the linear classifier. All reported numbers are an average of 5 runs, of which variances are negligibly small and not included.

2 Preliminary VQ experiments

We first explore several potential places to insert VQ layers. For all VQ-APC variants in our experiments, the autoregressive model gARg_{AR} is a 3-layer unidirectional GRU with 512 hidden units, and the target frame in the future, nn, is set to 5 when training (Equation 1) on LibriSpeech. Whenever a VQ layer is added, the embedding dimension of each code is 512, and τ\tau for the Gumbel-Softmax straight-through estimator (Equation 2) is a fixed value of 0.1 throughout training.

Quantizing multiple layers. We find that our VQ-APC models with multiple VQ layers have trouble fitting the training set. Their representations are also much worse than the regular APC on phonetic classification. One potential remedy is to enable VQ with a schedule , but is beyond the scope of the paper.

3 The constituents of the learned representations

Experiments so far suggest that the phonetic information is still present (if not better) after using VQ. For the rest of the paper, we will focus on the case where VQ is inserted at the third layer, i.e., the case of {3}\{3\}. To study the constituents of the learned representations, we train a series of increasingly limited VQ-APC models by decreasing the codebook size from 2048 to 64 while fixing the code dimension to 512. As the codebook size becomes smaller, the model is forced to choose what information to encode and what to discard, thus revealing the constituents of the learned representations. We show the training losses of these models at convergence and the respective phone and speaker error rates in Figure 2. The dashed lines are the training loss, phone error rate, and speaker error rate of a regular APC model.

First, the training loss (purple curve) increases as expected, showing worse fit on the training set as we limit the codebook size. Note that in theory, when the codebook size goes to infinity, we recover the regular APC. The phone error rate (red curve) obtains a minimum at codebook size 512, and starts to worsen with smaller codebook size. The sharp degradation in phone error rate suggests that the model discards certain phonetic information to achieve maximal self-supervised objective.

The speaker error rate (blue curve), on the contrary, does not change by much as we limit the codebook size. This shows that the speaker information (and its potentially correlated channel information) is mostly retained. Given the sharp degradation in phone error rate, we can conclude that the model prefers to retain speaker information over phonetic information to achieve maximal future prediction. The preference of information can potentially stem from the use of GRUs, the VQ configuration, and the self-supervised, future prediction objective. More analyses are needed to disentangle the among these causes.

On the other hand, when the codebook size becomes large, the model falls back to regular APC and might suffer from overfitting , paying unnecessary attention to the spectral details that does not generalize for predicting future frames. Finally, even with a codebook size of 64, we still see gains over regular APC, showing the strong performance of VQ-APC in representation learning.

4 Visualizations of learned codes

To better measure the correspondence between the learned VQ codes and English phones, we compute co-occurrence statistics (at the frame level) across the 360-hour subset of LibriSpeech, the dataset we use to train the VQ-APC models. We compare three settings, {1}\{1\}, {2}\{2\}, and {3}\{3\} with a codebook size of 128. The conditional probability P(phone∣code)P(\text{phone}|\text{code}), as shown in Figure 3, are estimated based on the co-occurrence statistics, i.e., via maximum likelihood. In each sub-figure, the rows and columns are ordered via spectral co-clustering with 15 clusters to group together phones that share similar sets of codes, and a diagonal segment would imply a high correspondence between phones and codes. Note that the phone labels of LibriSpeech are only used for analysis and never seen during training.

From Figure 3, we see that the correspondence between phones and VQ codes is stronger when quantized at higher layers, and is especially strong for {3}\{3\}. Recall that probing tasks are useful for showing the presence of certain information, but have difficulty showing the absence of it. In contrast, given the co-occurrence statistics, mutual information can be estimated to support the absence of information. The normalized mutual information are 0.167, 0.285, and 0.406 for {1}\{1\}, {2}\{2\}, and {3}\{3\}, respectively. In other words, not only can we conclude that the learned representations in {3}\{3\} contain phonetic information, we can also readily conclude that {1}\{1\} and {2}\{2\} contain much less information for certain phones.

5 A comparison with other self-supervised models

Finally, we compare VQ-APC with other self-supervised speech representation models, including Contrastive Predictive Coding (CPC) , Bidirectional Masked Reconstruction , Mockingjay , and Multi-Target APC . We briefly review these methods below, and show their results on phonetic and speaker classification in Table 2. To stay as close to the original implementation as possible, we do not separate the discussion of models, such as the use RNNs or Transformers, and the self-supervised objectives.

CPC and APC share a similar methodology as both use an autoregressive model to learn representations through conditioning on past frames to predict information about a future frame. Their difference is that while APC tries to directly predict the future frame via regression, CPC aims to learn representations containing information that most discriminates the future frame from a set of negative samples. We mainly follow the original paper for implementing CPC with some modifications described in . These modifications are meant to minimize the architectural differences between APC and CPC while maintaining their training objectives.

Multi-Target APC (MT-APC) is an extension of APC. It incorporates an auxiliary objective serving as a regularizer to improve the generalization of the main future prediction task. The exact same setup described in is used in our experiments.

Different from CPC and APC that are based on the idea of future prediction, Bidirectional Masked Reconstruction (BMR) and Mockingjay are under the category of masked prediction. Inspired by the masked language modeling technique from BERT , BMR and Mockingjay mask parts of the input signals, and predict them through conditioning on both past and future contexts with a bidirectional RNN and Transformer encoder , respectively. For experiments, we mainly follow the implementations in the original papers , except that the number of layers are reduced to match ours to minimize the architectural differences.

On phonetic classification, we see that VQ-APC (28.4) improves over APC (33.3) and MT-APC (30.5), demonstrating the effectiveness of VQ layers. It also outperforms other self-supervised models despite using the same (vs. CPC) or smaller (vs. BMR and Mockingjay) network. On speaker classification, VQ-APC (5.5) again improves over the other two APC models (8.5 and 7.3), and is on par with the best model (Mockingjay, 5.1). On both tasks, all self-supervised models outperform log Mel regardless of the type of classifier it uses.

Conclusions

We have demonstrated that incorporating vector quantization (VQ) layers into an Autoregressive Predictive Coding model imposes a bottleneck, forcing the model to learn better representations. Extensive experiments have been conducted to compare different VQ configurations, to study the effect of varying codebook sizes, and to compare with other self-supervised speech representation models. We show evidence for the presence and absence of phonetic and speaker information in the learned representations, and also show the model’s preference in retaining information when the model capacity is limited, in the hope to bridge the connection between the self-supervised objective and the properties of the learned representations. When the phonetic information is present, the learned VQ codes also correspond well with English phones.

References