LAST: Language Model Aware Speech Tokenization
Arnon Turetzky Yossi Adi


Appendix
    Kathryne J. Daniel
    September 9, 2024
======================



§ ABSTRACT 
Speech tokenization serves as the foundation of speech LM, enabling them to perform various tasks such as spoken language modeling, text-to-speech, speech-to-text, etc. Most speech tokenizers are trained independently of the LM training process, relying on separate acoustic models and quantization methods. Following such an approach may create a mismatch between the tokenization process and its usage afterward. In this study, we propose a novel approach to training a speech tokenizer by leveraging objectives from pre-trained textual LMs. We advocate for the integration of this objective into the process of learning discrete speech representations. Our aim is to transform features from a pre-trained speech model into a new feature space that enables better clustering for speech LMs. We empirically investigate the impact of various model design choices, including speech vocabulary size and text LM size. Our results demonstrate the proposed tokenization method outperforms the evaluated baselines considering both spoken language modeling and speech-to-text. More importantly, unlike prior work, the proposed method allows the utilization of a single pre-trained LM for processing both speech and text inputs, setting it apart from conventional tokenization approaches.



Speech Tokenization, Speech Language Models




§ INTRODUCTION



The development of Speech Language Models () was recently raised as a new research direction within the spoken language processing community <cit.>. are usually composed of two or three main components: (i) speech tokenizer which converts raw speech signals into discrete tokens; (ii) uLM, operating over this discrete representation to learn the underlying distribution of speech utterances; and (iii) in the generative setup, a unit-based vocoder which converts speech tokens into a waveform signal <cit.>. The ability to operate directly over raw speech recordings without accessing any textual supervision holds great potential: (i) This can be beneficial for languages that do not have large textual resources or standardized orthography (Swiss German, dialectal Arabic, Igbo, etc.); (ii) It may also be useful for “high-resource” languages, in which the oral and written forms often mismatch; and (iii) It can also model non-verbal cues such as laughter, coughing, etc. Recent studies in the field have shown that following this modeling paradigm can be beneficial for spoken dialogue modeling <cit.>, speaking style conversion and speech emotion conversion <cit.>, direct speech-to-speech translation <cit.>, silent video-to-speech generation <cit.>, etc. In this work, we study the speech tokenization part. 

The most common approach nowadays to speech tokenization is applying the k-means algorithm over representations obtained from a pre-trained SSL models, this results in a semantic speech tokens <cit.>. While being simple and effective, this approach has several drawbacks. First, it requires a separate training stage on top of the SSL model, and can not be done jointly. Second, it is not clear what layer should we use to obtain representations to train the k-means model. Prior work proposed different layers <cit.> and show varying results when different layers are being chosen, even when considering the same model. Lastly, the k-means model is sensitive to the data used to learn the cluster centroids. Hence, it remains an open question of what data should be used to train the clustering model. Prior work found that different dataset splits result in different model performance <cit.>. Another line of speech and audio tokenization was proposed by <cit.> using a Residual Vector Quantization (RVQ) method under the auto-encoding setup. Such a tokenization method is highly general and can capture any type of audio (not only speech), however, it was found to be challenging to optimize on top of it without conditioning on text or semantic tokens <cit.>. Such representation is often denoted as acoustic tokens. 



Considering the fact that the above-mentioned tokens will be later on used to train , raises the following question: can we construct a speech tokenizer which will be guided by a language model? Moreover, recently <cit.> found that performing a warm-initialization of parameters from a pre-trained text LM, such as OPT <cit.> or LLaMA 2 <cit.>, results in superior performance to cold-random initialization. 

Equipped with the above findings we propose , which stands for Language Model Aware Speech Tokenization. Specifically, we propose to involve a pre-trained text LM during the tokenization process to result in speech tokens that will be better suited for sequential modeling. Generally, is comprised of three main components: (i) a pre-trained, frozen speech SSL model which extracts contextualized speech representations from speech signals. We experimented with HuBERT <cit.>, however is general and can be applied to any SSL method; (ii) an adapter-quantization module which converts this contextualized representation into discrete tokens; and (iii) a pre-trained, frozen LM which guides the tokenization process towards better sequential modeling. A visual description of the proposed method can be seen on <Ref>. Notice, in contrast to <cit.> which fine-tunes text LM using speech tokens, and as a result removing the learned textual information, the proposed approach keeps the text-LM frozen, hence does not affect performance on text benchmarks. 

We evaluate on a set of zero-resource speech modeling tasks <cit.> and found the proposed method is superior to the traditional k-means method across all setups. We also evaluate the ability of the newly proposed speech tokens on the task of ASR and found them to produce superior performance than the k-means alternative. We additionally provide an extensive ablation study, which sheds light on the importance of each of the components constructing .





§ BACKGROUND



As mentioned before, the general pipeline for constructing is comprised of two main modules: (i) Speech tokenizer and (ii) uLM (See Figure <ref> for a visual description). In the following paragraphs, we give a background for each of the components including the standard evaluation methods.  

Speech Tokenizer module encodes the raw speech signal into a discrete representation. The common approach is first to encode the speech into a continuous representation and then quantize the representation to achieve a sequence of discrete units <cit.>. 

Formally, denote the domain of audio samples by 𝒳⊂ℝ. The representation for a raw signal is therefore a sequence of samples x = (x_1,…, x_T), where  x_t∈𝒳 for all 1≤ t ≤ T. 

Consider an encoder network, f, that gets as input the speech utterance and outputs a sequence of spectral representations sampled at a low frequency as follows f(x) = (v_1, …, v_T'). Note that we do not assume anything about the structure of the encoder network f. <cit.>, evaluated several speech encoders, namely, Mel-spectrogram, Contrastive Predictive Coding <cit.>, wav2vec2 <cit.>, and HuBERT <cit.>. 

Since the representations learned by such models are usually continuous, a tokenization or quantization algorithm is applied over the models' outputs to generate discrete units, denoted as z = (z_1,…,z_T'). Each element z_i in z is a positive integer, z_i∈{1,...,K} for 1≤ i ≤ T', where K is the number of discrete units. We denote the quantization model with Q. The common approach in prior work is to use the k-means algorithm as the quantization method. 



Unit Language Model  is trained on the extracted discrete units, z. Such a language model learns a probability distribution of the learned unit sequences, which enables direct modeling of speech data without textual supervision. 

The language model can be used to sequentially model speech utterances, and generate speech conditionally or unconditionally. Moreover, such a modeling framework allows for capturing and modeling prosodic features <cit.>, as well as speaker identity <cit.>, or even natural dialogues <cit.>. This is in contrast to using textual features, as they do not encode such information. In this work, we do not focus on speech generation but rather demonstrate that the proposed speech tokenizer provides better performance for modeling sequential data via the uLM.  

Zero-Resource Speech Evaluation. <cit.> proposed a set of zero-shot evaluation tasks specifically targeting speech modeling (i.e., sWUGGY, and sBLIMP). The sWUGGY metric requires detecting the real word from a pair of short utterances such as 'brick' vs. 'blick.' Similarly, sBLIMP requires detecting the syntactically correct sentence from a pair of sentences. In both metrics, detection is done by comparing the probabilities of both sequences. These metrics are desirable under our setup as, unlike perplexity, these allow us to compare models with different tokenizers. We mainly report this method throughout the paper.




§ LANGUAGE AWARE SPEECH TOKENIZER


In this section we present . We start by describing the speech tokenization approach guided by a frozen LM, followed by the exact implementation details and fine tuning configuration. 



 §.§ Model



Given a speech utterance x, we first feed it into a frozen and pre-trained speech encoder network f(x) = v = (v_1, …, v_T'), where each v_i∈ℝ^d. Next, v is being fed into a learnable encoder module E(v) = u = (u_1, …, u_T'), where each u_i∈ℝ^d'. Then, a quantization process is performed, via a Vector Quantization (VQ) module, Q, to quantize each u_i. Formally, Q(u) = (z_1,…,z_T'), where each z_i is a positive integer in the range of {1, …, K} and K is the number of codes in Q. 

Next, to guide the quantization process toward sequential modeling, we feed z into a pre-trained textual LM to perform next token prediction. To extend the text LM for speech tokens, we add randomly initialized adaptation layers before and after the LM which we freeze for the entire training, hence keeping the existing text LM capabilities unchanged. We analyze the effect of the adapter size in <Ref>. 

Similarly to the common practice in training LMs we minimize the negative log-likelihood loss between the predicted and true probability distributions over the learned tokens. Formally, we minimize the following, 


    ℒ_LM = -∑_i=1^nlog p_θ(z_i | z_i-1, …, z_1),


where we consider θ as the parameters of both the adaptor and new speech tokens look-up table. We do not update the LM parameters, we only backpropagate through it. 

Notice, that we learn to tokenize speech via next token prediction. Such an optimization process may easily collapse to a single token or a sequence of tokens. To prevent collapse we introduce a reconstruction loss function to stabilize the optimization process. For that, we introduce a decoder module, D, which gets as input u and is optimized to reconstruct v. Specifically, we minimize the L2 loss between D(u) and v. Overall, the objective function of the system is as follows, 


    ℒ = ℒ_LM + λD(u) - v_2,


where λ is a hyperparameter balancing between both loss functions. A visual description of the proposed method can be seen in <Ref>.





 §.§ Details


We set f to be a pre-trained HuBERT- can be beneficial
for languages that do not have large textual resources or even a widely used standardized
orthographybase model <cit.>. Similarly to <cit.>, representations are obtained from its 9th layer. For E, we utilize a transformer encoder followed by a projection layer, where the latent dimension d' is determined by the LM dimension.

We employ VQ <cit.> to discretize the output u where we experimented number of codebooks ∈{100, 200, 500, 1000}. We empirically analyze results for the different number of codebooks in <Ref>.

For the LM part, we utilize the OPT model <cit.> with 125M and 350M parameter models. These models are initialized from a pre-trained text LM following <cit.>. We employ a separate lookup table with with a new randomly initialized embedding and add two more projection layers to enable speech tokens in addition to text tokens. In order to feed tokens into the model as input, the first layer projects from the embedding dimension to the model dimension. The second layer maps the model output to the number of tokens, allowing for the calculation of next speech token probabilities. In our setup we use only the speech tokens but using both text and speech tokens, only requires an indicator to specify whether to use text or speech. This indicator will determine which group of layers should be used for each token. Similarly to <cit.> we remove sequential repetitions of speech tokens before feeding them into the LM. Notice, unlike <cit.> which fine-tunes the text LM and as a result removes text modeling capabilities, following the proposed approach keeps the original performance over text benchmarks unchanged while extending it to speech modeling ones. We update the LM lookup table with the vector quantization codebooks at each training step and feed z as input for the LM. 



Finetune. Although the proposed method allows the utilization of a single pre-trained LM for processing both speech and text inputs, we would also evaluate similarly to <cit.>, i.e., finetuning the LM. We denote the difference between both cases as “Pretrain” and “Finetune” to avoid confusion. 


§ EXPERIMENTAL SETUP



We use different numbers and types of GPUs across our experiments due to diverse computational requirements ranging from single to 4 GPUS out of the Nvidia A5000, A6000, and A40. Due to computational constraints we use grad accumulation, and set a limit for maximum steps. Unless stated otherwise, except for the speech recognition experiments, we set a limit of 200K optimization steps with a batch size of 32. We sampled 10 seconds from each example and zero padded from right whenever needed. We used AdamW optimizer and linear-warmup-cosine scheduler that increase learning rate from 0 to 1e-4/2.5e-5 by 1K/10K steps then decrease to 3e-5/1e-5 for pretraining and finetuning experiments respectively. Speech recognition experiments were limited to 150K steps with batch size 64, same optimizer and scheduler but with a warmup of 5K steps to learning rate 1e-4 decrease to 1e-5.




§ RESULTS



We report results for the zero-resource speech metrics <cit.> together with speech recognition metrics, comparing the proposed method to the k-means alternative. 



 §.§ Zero-Resource Speech Metrics

We start by evaluating the proposed method using the zero-resource speech metrics as proposed by <cit.>, namely sWUGGY and sBLIM, and compare it to the commonly used k-means tokenizer. Unless stated otherwise, for a fair comparison we follow the same approach as in <cit.> and report the results for the “in-vocabulary” split. For both measures, we compare the geometric mean of the models' sequence probabilities assigned to each utterance within a pair. We report results in two setups. In the first one, we keep the LM frozen and do not update its parameters. In the second setup, we follow TWIST <cit.>, and we fine-tune the LM initialized from a text model. For a fair comparison to the k-means baseline, we consider two setups of : with and without the adaptation layers. Both and the k-means baseline are built on top of the HuBERT- model. Results are presented in <Ref> using with OPT-350 model and codebook size of 500.

We observe a consistent improvement across all the evaluated setups when following the proposed approach compared to the evaluated baseline. When considering no TWIST fine-tuning of the LM, significantly outperforms the k-means alternative. Interestingly, even without applying TWIST to the proposed method it provides competitive performance to TWIST using the k-means tokenizer. 

Another benefit of following is it requires significantly fewer tunable parameters compared to TWIST. For example, to obtain the results presented in <Ref> requires training ∼230M parameters less than TWIST.








 §.§ Speech Recognition


Next, we evaluate how the proposed tokenization method is able to preserve linguistic content. For that, we train an acoustic model for the task of ASR and measured the WER. We train a T5- <cit.> seq2seq model on LS-960. Text tokenization is similar across all methods using T5- pre-trained tokenizer.

We additionally measure the ABX metric, as proposed by <cit.>. Unlike WER which measures the content preservation at the word level, ABX operates at the phoneme level. The ABX task measures the phonetic discriminative abilities of the representation. It involves a pair of words differing by a single phoneme and a reference test word sharing a phoneme with one of the pair. It assesses whether the test phoneme is closer in representation to the correct or incorrect phoneme, expecting a shorter distance to the correct one. The ABX task is conducted in two setups: 'within' and 'across'. 'Within' is evaluated on input data from the same speaker, while 'across' is evaluated on input data from different speakers. Notice, our goal is not to construct the best performing ASR model <cit.>, we are primarily interested in comparing the relative results of compared to the k-means alternative.

Results are presented in <Ref>. When considering the WER metric, shows superior performance than the k-means alternative, however, when considering the ABX metric, the k-means method outperforms . Considering the way both methods were constructed, these findings imply that the k-means model better captures phonemic information.


To better analyze this, we follow <cit.> and visualize the tokens clusters as a function of phoneme families measured over the TIMIT corpus which contains full alignment. For each token cluster, we select the majority phoneme, the one that was most present in the cluster. Results are depicted in <Ref>. Results suggest that both methods are correlative with phonetic information. 






§ ABLATION


We turn into analyzing the different components composing . We start by evaluating the effect of the vocabulary size, we systematically compare different numbers of codebook sizes. Next, we evaluate the effect of the LM tokenizer used and the final . We consider LM of different sizes from the OPT family of models. Lastly, we consider different numbers of layers for encoding, decoding and modality adapting layers. 




 §.§ Vocabulary Size 


Similarly to <cit.> we experiment with different codebook sizes to better evaluate the effect of the number of tokens on modelling performance. We measure zero-resource speech metrics (i.e., sWUGGY and sBLIMP) using with the additional decoding layers. Results are reported in <Ref>. Results suggest that 500 tokens provide the best performance. Interestingly using 1000 units provides inferior performance to even 100 units. We suspect this is due to high redundancy in the codebook usage. 







 §.§ Tokenizer-Model Relation


Next, we evaluate the effect of LM size on the tokenization method and as a afterwords. We consider different model sizes from the OPT model family (i.e., 125M and 350M parameters) as either the and the model used for tokenization (denoted by Tokenizer LM). Results are reported in <Ref>. 

We observe several interesting findings: (i) although conducting the experiments on the same model family, we can say our tokenizers can be used by different LMs. (ii) finetuning the 125M with the tokenizer trained with the 350M LM achieves better results than the one finetuned with the tokenizer train with the 125M LM. 






 §.§ Modelling Choices



Model Architecture. We analyze the effect of different architectural choices of . Specifically, we consider different numbers of layers in the encoder, decoder, and adaptation modules. For the adaptation module, we additionally consider a varying number of layers before and after the LM. In all setups, we report the sWUGGY metric using OPT-350 with 500 tokens. Results presented in <Ref>. 

Results suggest that except of one setup (second row in <Ref>), the proposed method is not sensitive to the specific design choice. However, a few interesting insights comes up: (i) using 2 layers for all modules provides the best overall results; (ii) using more adaptation layers before the LM provides superior performance than using more adaptation layers after the LM; and (iii) using equal number of adaptation layers before and after the LM provides the best performance and stabilizes the results, as long as there is enough capacity in the encoder module. 

Reconstruction loss. Lastly, we discussed the need for the L2 loss <Ref> to prevent collapse during model optimization. Specifically, we proposed to regularize the training process with an L2 loss between the D(u) and v. To evaluate different regularization alternatives, we experiment with replacing v with the different layer from f (HuBERT model), specifically layer 10 or with the acoustic features obtained from the first CNN layers in HuBERT. However, we found both options to provide inferior performance to the currently used one. 



§ RELATED WORK



Speech Language Models. Speech language models were first demonstrated by <cit.>. The authors showed how raw and uncurated speech data can be leveraged into building a GSLM system. Next, <cit.> proposed a multi-stream to jointly process “pseudo-text” tokens together with quantized prosodic features (i.e., duration and F0). Such a modeling framework opened up a new and promising research direction for processing and modeling spoken data. <cit.> evaluated the robustness and disentanglement properties of speech-to-tokens models and demonstrated the ability to perform voice conversion as well as a lightweight speech codec. <cit.> proposed to cast the task of speech emotion conversion as a translation task, hence translating between one emotion to the other in the discrete space, while <cit.> proposed a similar approach for speaking style conversion. <cit.> proposed training two jointly to mimic natural spoken dialogues. Recently, <cit.> proposed cascading several LMs, in which one LM operates over semantic speech tokens while the others operate on acoustic tokens. Such modeling framework allows generating natural speech while keeping the identity of the speaker and acoustic conditions unchanged <cit.>. <cit.> show that the semantic units obtained from such models highly correlate to phonemes and phoneme states. <cit.> followed a similar modeling mechanism using a different speech tokenizer and proposed a textless approach for speech-to-speech translation. <cit.> and <cit.> considered speech as another language in multilingual setups, and showed that involving speech and text as part of training data improves results for speech translation and multilingual text tasks. <cit.> and <cit.> used the joint training to improve transcriptions tasks as ASR and TTS. <cit.> proposed augmenting a text LM with continuous speech data to improve spoken question-answering tasks. 




Speech and Audio Tokenizers. The most common approach to tokenize speech for spoken language modeling is the k-means method, which was first proposed by <cit.>. Later on, other studies follow a similar modeling paradigm <cit.>. These results in a discrete semantic representation of the speech signal. <cit.> pointed out the lack of robustness in current k-means tokenization methods and proposed to match representations from clean and augmented signals. Another alternative approach is to follow the VQ-VAE approach using an RVQ <cit.>. Due to its objective function, such representation better captures acoustic information rather than semantic information. Recent studies either use the semantic tokens <cit.>, the acoustic tokens conditioned on global textual descriptions  <cit.>, or both semantic and acoustic tokens <cit.>. Recently, <cit.> proposed to jointly train speech tokenizers to capture both semantic and acoustic tokens. The authors proposed training a discrete auto-encoder using RVQ and optimizing the first RVQ codebook to be similar to the semantic tokens obtained from the k-means model. Notice, that all of the above methods are orthogonal to the proposed method as they can be applied jointly with our method. 

Another relevant related work to ours is the work by <cit.> that proposes to leverage a pre-train text-to-image and text-to-video models (respectively) to build image and video generation models condition on audio inputs. The authors proposed to augment the textual inputs with the newly learned single audio token. Unlike these works, our method is focused on speech and not general audio, and more importantly, aimed at learning a full vocabulary over time rather than a single audio token.


§ DISCUSSION


In this study, we propose : a language model aware speech tokenization method. Unlike prior work, which constructs the speech tokens independently from the , involves the during the tokenization process. leverages both frozen pre-trained contextualized speech encoder and frozen pre-trained text LM while introducing a lightweight modality adapter. Results suggest provides superior performance to the k-means alternative considering both zero-resource metrics and transcription capabilities. Interestingly, as augments a frozen pre-trained text LM, it is not only superior in terms of speech modeling but also keeps the same text capabilities of the original LM model. 

Limitations. While is superior to the k-means alternative and allows to optimize jointly the speech tokenizer and , it requires more computational resources than the standard k-means approach. Additionally, we only presented results for zero-resource speech modeling metrics, more evaluations are needed in the direction of unit-to-speech synthesis. We leave that for future research.  

Future work. We intend to extend this work in two main directions: (i) as mentioned above, evaluating also under the unit-to-speech synthesis framework. This will allow us to evaluate the system similarly to the GSLM setup as was done by <cit.>. Additionally, as the demonstrated that a frozen text LM can be efficiently augmented with speech tokens, we intend to explore the merger of the two modalities in the form of ASR and TTS.

IEEEbib


