vq-wav2vec: Self-Supervised Learning of Discrete Speech Representations

1 INTRODUCTION . Learning discrete representations of speech has gathered much recent interest ( Versteegh et al. , 2016 ; Dunbar et al. , 2019 ) . A popular approach to discover discrete units is via autoencoding ( Tjandra et al. , 2019 ; Eloff et al. , 2019 ; Chorowski et al. , 2019 ) sometimes coupled with an autoregressive model ( Chung et al. , 2019 ) . Another line of research is to learn continuous speech representations in a self-supervised way via predicting context information ( Chung & Glass , 2018 ; van den Oord et al. , 2018 ; Schneider et al. , 2019 ) . In this paper , we combine these two lines of research by learning discrete representations of speech via a context prediction task instead of reconstructing the input . This enables us to directly apply well performing NLP algorithms to speech data ( Figure 1a ) . Our new discretization algorithm , vq-wav2vec , learns discrete representations of fixed length segments of audio signal by utilizing the wav2vec loss and architecture ( Schneider et al , 2019 ; §2 ) . To ∗Equal contribution . †Work done during a Facebook AI residency . 1The code will be made available at http : //github.com/pytorch/fairseq . choose the discrete variables , we consider a Gumbel-Softmax approach ( Jang et al. , 2016 ) as well as online k-means clustering , similar to VQ-VAE ( Oord et al. , 2017 ; Eloff et al. , 2019 ; §3 ) . We then train a Deep Bidirectional Transformer ( BERT ; Devlin et al. , 2018 ; Liu et al. , 2019 ) on the discretized unlabeled speech data and input these representations to a standard acoustic model ( Figure 1b ; §4 ) . Our experiments show that BERT representations perform better than log-mel filterbank inputs as well as dense wav2vec representations on both TIMIT and WSJ benchmarks . Discretization of audio enables the direct application of a whole host of algorithms from the NLP literature to speech data . For example , we show that a standard sequence to sequence model from the NLP literature can be used to perform speech recognition over discrete audio tokens ( §5 , §6 ) . 2 BACKGROUND . 2.1 WAV2VEC . wav2vec ( Schneider et al. , 2019 ) learns representations of audio data by solving a self-supervised context-prediction task with the same loss function as word2vec ( Mikolov et al. , 2013 ; van den Oord et al. , 2018 ) . The model is based on two convolutional neural networks where the the encoder produces a representation zi for each time step i at a rate of 100 Hz and the aggregator combines multiple encoder time steps into a new representation ci for each time step i . Given an aggregated representation ci , the model is trained to distinguish a sample zi+k that is k steps in the future from distractor samples z̃ drawn from a distribution pn , by minimizing the contrastive loss for steps k = 1 , . . . , K : Lwav2veck = − T−k∑ i=1 ( log σ ( z > i+khk ( ci ) ) + λE z̃∼pn [ log σ ( −z̃ > hk ( ci ) ) ] ) ( 1 ) where T is the sequence length , σ ( x ) = 1/ ( 1 + exp ( −x ) ) , and where σ ( z > i+khk ( ci ) ) is the probability of zi+k being the true sample . We consider a step-specific affine transformation hk ( ci ) = Wkci + bk that is applied to ci ( van den Oord et al. , 2018 ) . We optimize the loss L = ∑K k=1 Lk , summing ( 1 ) over different step sizes . After training , the representations produced by the context network ci are input to the acoustic model instead of log-mel filterbank features . 2.2 BERT . BERT ( Devlin et al. , 2018 ) is a pre-training approach for NLP tasks , which uses a transformer encoder model to build a representation of text . Transformers uses self-attention to encode the input sequence as well as an optional source sequence ( Vaswani et al. , 2017 ) . The original BERT model combined two tasks for training : first , masked language modeling randomly removes some of the input tokens and the model has to predict those missing tokens . Second , next sentence prediction splices two different text passages together into a single example and the model needs to predict whether the passages are from the same document . 3 VQ-WAV2VEC . Our approach , vq-wav2vec , learns vector quantized ( VQ ) representations of audio data using a future time-step prediction task . We follow the same architectual choices as wav2vec ( §2.1 ) with two convolutional networks f : X 7→ Z and g : Ẑ 7→ C for feature extraction and aggregation , as well as a new quantization module q : Z 7→ Ẑ to build discrete representations ( Figure 1a ) . We first map 30ms segments of raw speech to a dense feature representation z at a stride of 10ms using the encoder network f . Next , the quantizer ( q ) turns these dense representations into discrete indices which are mapped to a reconstruction ẑ of the original representation z . We feed ẑ into the aggregator g and optimize the same context prediction task as wav2vec outlined in §2.1 . The quantization module replaces the original representation z by ẑ = ei from a fixed size codebook e ∈ RV×d which contains V representations of size d. We consider the Gumbel-Softmax which is a differentiable approximation of the argmax for computing one-hot representations ( §3.1 ; Figure 2a ) as well as online k-means clustering , similar to the vector quantized variational autoencoder ( VQVAE ; Oord et al. , 2017 ; §3.2 ; Figure 2b ) . Finally , we perform multiple vector quantizations over different parts of z to mitigate mode collapse ( §3.3 ) . 3.1 GUMBEL-SOFTMAX . The Gumbel-Softmax ( Gumbel , 1954 ; Jang et al. , 2016 ; Maddison et al. , 2014 ) enables selecting discrete codebook variables in a fully differentiable way and we use the straight-through estimator of Jang et al . ( 2016 ) . Given the dense representation z , we apply a linear layer , followed by a ReLU and another linear which outputs l ∈ RV logits for the Gumbel-Softmax . At inference , we simply pick the largest index in l. At training , the output probabilities for choosing the j-th variable are pj = exp ( lj + vj ) /τ∑V k=1 exp ( lk + vk ) /τ , ( 2 ) where v = − log ( − log ( u ) ) and u are uniform samples from U ( 0 , 1 ) . During the forward pass , i = argmaxjpj and in the backward pass , the true gradient of the Gumbel-Softmax outputs is used . 3.2 K-MEANS . The vector quantization approach of van den Oord et al . ( 2017 ) is an alternative to making the index selection procedure fully differentiable . Different to their setup , we optimize a future time step prediction loss instead of the reconstruction loss of an autoencoder . We choose the codebook variable representation by finding the closest variable to the input features z in terms of the Euclidean distance , yielding i = argminj‖z − ej‖22 . During the forward pass , we select ẑ = ei by choosing the corresponding variable from the codebook . We obtain gradients for the encoder network by back-propagating dLwav2vec/dẑ ( van den Oord et al. , 2017 ) . The final loss has two additional terms : L = K∑ k=1 Lwav2veck + ( ‖sg ( z ) − ẑ‖2 + γ‖z− sg ( ẑ ) ‖2 ) , ( 3 ) where sg ( x ) ≡ x , ddx sg ( x ) ≡ 0 is the stop gradient operator and γ is a hyperparameter . The first term is the future prediction task and gradients do not change the codebook because of the straightthrough gradient estimation of mapping z to ẑ . The second term ‖sg ( z ) − ẑ‖2 moves the codebook vectors closer to the encoder output , and the third term ‖z − sg ( ẑ ) ‖2 makes sure that the encoder outputs are close to a centroid ( codeword ) . 3.3 VECTOR QUANTIZATION WITH MULTIPLE VARIABLE GROUPS . So far , we considered replacing the encoder feature vector z by a single entry ei in the codebook . This is prone to mode collapse where only some of the codewords are actually used . Previously , this problem has been mitigated by workarounds such as re-initializing codewords or applying additional regularizers to the loss function ( Caron et al. , 2019 ) . In the following , we describe another strategy where we independently quantize partitions of z , similar to product quantization ( Jegou et al. , 2011 ) . This results in larger dictionaries and increased downstream performance ( Appendix A ) . The dense feature vector z ∈ Rd is first organized into multiple groups G into the matrix form z′ ∈ RG× ( d/G ) . We then represent each row by an integer index , and hence can represent the full feature vector by the indices i ∈ [ V ] G , where V again denotes the possible number of variables for this particular group and each element ij corresponds to a fixed codebook vector . For each of the G groups , we apply either one of the two VQ approaches ( §3.1 and §3.2 ) . The codebook itself can be initialized in two possible ways : Codebook variables can be shared across groups , i.e. , a particular index in group j would reference the same vector as the same index in group j′ . This yields a codebook e ∈ RV× ( d/G ) . In contrast , not sharing the codebook variables yields a codebook of size e ∈ RV×G× ( d/G ) . In practise , we observe that sharing the codebook variables generally yields competitive results to a non-shared representation . 4 BERT PRE-TRAINING ON QUANTIZED SPEECH . Once we trained a vq-wav2vec model we can discretize audio data and make it applicable to algorithms that require discrete inputs . One possibility is to use the discretized training data and apply BERT pre-training where the task is to predict masked input tokens based on an encoding of the surrounding context ( Devlin et al. , 2018 ) . Once the BERT model is trained , we can use it to build representations and feed them into an acoustic model to improve speech recognition . We follow recent advances in BERT training which only use the masked input token prediction ( Liu et al. , 2019 ) . Since each of the discretized tokens represents around 10 ms of audio it is likely too easy to predict a single masked input token . We therefore change BERT training by masking spans of consecutive discretized speech tokens , similar to Joshi et al . ( 2019 ) . To mask the input sequence , we randomly sample p = 0.05 of all tokens to be a starting index , without replacement , and mask M = 10 consecutive tokens from every sampled index ; spans may overlap . This makes the masked token prediction harder and we show later that it improves accuracy over masking individual tokens ( §6.5 ) . 5 EXPERIMENTAL SETUP . 5.1 DATASETS . We generally pre-train vq-wav2vec and BERT on the full 960h of Librispeech ( Panayotov et al. , 2015 ) and after vq-wav2vec training it is discretized to 345M tokens . Where indicated we perform ablations on a clean 100h subset which is discretized to 39.9M tokens . We evaluate models on two benchmarks : TIMIT ( Garofolo et al. , 1993b ) is a 5h dataset with phoneme labels and Wall Street Journal ( WSJ ; Garofolo et al . 1993a ) is a 81h dataset for speech recognition . For TIMIT , we apply the standard evaluation protocol and consider 39 different phonemes . For WSJ , we train acoustic models directly on 31 graphemes , including the English alphabet , the apostrophe , the silence token and tokens for repeating characters .

Summary: The paper proposes a way to pre-train quantized representations for speech. The approach proposed is a two-stage process: 1. train a quantized version of wav2vec [my understanding is that wav2vec is the same thing as CPC for Audio except for using a binary cross-entropy loss instead of InfoNCE softmax-cross entropy loss]. the authors propose to use gumbel softmax / VQ codebook for the vector quantization.