Choice of semantic tokens

[eonglints]

Obviously the semantic tokens in the original paper are difficult to reproduce because w2v-BERT XL is closed source. However, I'm not sure wav2vec 2.0 features are the most appropriate replacement. AudioLM was heavily influenced by Text-Free Prosody-Aware Generative Spoken Language Modeling (paper, code) (indeed the lead author of that paper left FAIR, joined Google and contributed to AudioLM).

In that paper, they use clustered features from an intermediate layer (layer 6 I believe) of HuBERT (paper, code, HF). This model and this clustering approach has been used for many speech-related papers in the last year (including several from FAIR. If I have time, I'll find some references. But for phoneme recognition, for example, HuBERT does a better job than wav2vec 2.0 (see this leaderboard. If you're wondering, WavLM is essentially HuBERT with more data and a couple more tricks. It's also available. However, there are more papers using HuBERT).

Facebook even released a small library (which you may well already be aware of) that makes it easy to extract clustered features using pre-trained HuBERT and pre-trained k-means models. The only issue is that most of the clustering models have considerably fewer centroids than AudioLM, (the exception being the mHuBERT 1000 model).

Anyway, I guess all this to say that perhaps HuBERT (or potentially WavLM) would be a better base for semantic tokens than wav2vec 2.0. Unfortunately, all three models have a higher frame-rate (50Hz) than w2v-BERT (25Hz) which will make the sequences longer.

Thanks so much for your work on this!

[lucidrains]

@eonglints thank you for this! as you can tell, i'm not very familiar with the audio domain (but plan to deepen my knowledge in this arena)

i will plan on incorporating hubert intermediate features then!

do you think the additional detail of normalizing across each dimension is very important or inconsequential?

[eonglints]

I think you've got plenty of expertise for one person already :) Thanks for all your fantastic contributions.

I just double-checked the FAIR "spoken language modeling" papers; in the first one they use layer 6 of the huBERT base model and in the second they use layer 9 of the base model. They don't appear to justify the change. You can quite easily choose the layer in their textless-lib or you can just use the Hugging Face model and choose a layer. We'd probably need to look at training a k-means model with 1024 centroids to be comparable to AudioLM (FAIR only provide 50, 100, 200 and 500).

As for normalizing, I've never tried it and I've not seen it elsewhere in the literature (although there's a good chance I've missed it). As you know, the authors of AudioLM state that this step "significantly improves their phonetic discriminability" so it certainly sounds worth exploring. Probably not essential immediately though as clustered HuBERT features already have pretty good phonetic discriminability (as shown in that leaderboard that I posted earlier).

[lucidrains]

thanks for the kind words

yup, i think there are organizations (say Stability.ai) that may be able to take on training a larger wav2vec or huBERT model with more centroids for the semantic layer

agreed with your take on the normalizing across each dimension

[lucidrains]

@eonglints also, while i have you on the line, what do you think about this paper

is it still worth implementing given the results of audiolm? my take is that perhaps with the right conditioning added to audiolm, the architecture would obsolete the prior work, but i'm not well read in the field enough to know for sure

thank you in advance!

[eonglints]

That's a great question. Natural Speech is quite a different beast, being an end-to-end text-to-speech model (with no intermediate features). I've not looked at it closely, but it appears to be fairly similar to an existing TTS model called VITS (which already has an official implementation). Indeed, one of the main contributors to Coqui (an open-source speech community) doesn't feel like it's worth implementing for their library.

As for whether AudioLM could be conditioned or prompted to do TTS, this is a question I think is very worth exploring and is something I'm actively researching in my PhD. I'm not confident it will make the existing approaches obsolete, but it's certainly an interesting new direction.

[lucidrains]

thank you for linking to Coqui as well as the maintainer's opinion on Natural Speech! this helps a lot in my decision making

i'll look into adding conditioning (with classifier free guidance). i have high hopes after seeing Whisper (which is just a vanilla encoder / decoder transformer)

[eonglints]

Totally agree, I had the same thought after seeing the relative simplicity of Whisper. Thanks again, you're doing god's work 🙏

[lucidrains]

c17ee7d ok, the code is in there! very curious to see if this works at all

[lucidrains]

@eonglints only the smallest Hubert has its quantizer released https://github.com/facebookresearch/fairseq/blob/main/examples/hubert/README.md#hubert

[eonglints]

Yeah, the base model is the only one with quantizers (there are more here), but it's worked pretty well in the two FAIR papers I mentioned in the comment. The base model is 768-dimensional compared to w2v-BERT XL's 1024 and the larger HuBERT models do perform better on phoneme recognition (see "PR" on the leaderboard) so it may well be worth using them. But yeah, that would involve training k-means models...

[lucidrains]

ok, i'll try to get the smallest Hubert working this week and will start thinking about how to get that larger model with more centroids trained + open sourced!

[eonglints]

Awesome. It should actually be pretty straightforward - FAIR have released the Large and Extra Large models (HF link) and the k-means clustering should just be a case of running Scikit-Learn's Minibatch k-means (which is what FAIR use for their clustering). I've done it with ~500 hours of data and the base model and it only takes a few hours.

[lucidrains]

@eonglints a11722e

ohh got it! yes, just need to train the sklearn minibatchkmeans on top of the large models

[eonglints]

Just another little note about the semantic tokens - in the paper they mention in IV. B "in the first two stages [semantic and course acoustic], we follow the previously proposed practice of removing consecutive repetitions of the semantic tokens [14]" In 14 they state "we found that removing sequential repetitions of units improves performance, hence we apply it universally."

Given these statements, we should probably look to include this step. It's obviously fairly easy to implement using torch.unique_consecutive and the way the FAIR researchers implement it looks like this. Audio LM doesn't use the counts of the repetitions so there would be no need to return them.

[lucidrains]

🙏 oh yes, thank you! this ties into one of my todo items for variable lengthed sequences (and addition of [eos]'s)

will try to get this done today

[lucidrains]

that is really strange that the count of the consecutive sequences isn't incorporated

isn't duration important for certain types of audio, say music? i wonder if run-lengthed encoding would help here

[lucidrains]

i can't complain listening to the piano samples though!

[eonglints]

Yeah, I find it pretty weird too. Durations are very important in speech as well as music obviously. In the second spoken language modeling paper from FAIR they make a note of the count and incorporate it in the input embedding (basically run length encoding like you say) but in AudioLM they don't appear too. My only guess is that somehow the acoustic token modeling is taking care of re-injecting appropriate durations... but you'll have a better intuition than me as to whether is this possible/likely.
Like you say, one way or another, no complaints about the quality of the end result!

[lucidrains]

@eonglints ok got it! we can worry about injecting the duration information once we see some preliminary results

implemented the unique consecutive semantic tokens like you recommended! https://github.com/lucidrains/audiolm-pytorch/blob/main/audiolm_pytorch/audiolm_pytorch.py#L86

[averkij]

Hi, I'm digging in the audio quantization, and this discussion is very helpful. Thanks!

vq-wav2vec is used in vq_wav2vec.py, but I can't figure out how it works.

codebook_indices is a list of tuples of integers:

embed = self.model.feature_extractor(wav_input)
_, codebook_indices = self.model.vector_quantizer.forward_idx(embed)

like this:

tensor([[  2,  31],
        [154,  78],
        [  8, 204],
        ...]])

• So there is one tuple per audio frame and there should be one embedding per tuple, right?
• If so, how to get embeddings for audio having these tuples of codebook indices? vector_quantizer.embedding object have 320 entities.

[eonglints]

I've not looked at vq-wav2vec very much, but I'll have a look tomorrow. To try and keep in line with the AudioLM paper, we're probably going to use discrete HuBERT features rather than discrete wav2vec features (there is much more precedent in the literature for this approach, precedent that clearly inspired AudioLM). But still, good to look at.

[eonglints]

Just to say that I'm successfully training the semantic tokens transformer right now using the LJ dataset and mHuBERT 1000 tokens (these tokens were generated offline before-hand, so I haven't tested in-line token extraction). I just copied the training script and autoregressive wrapper from PaLM-pytorch and made a couple of minor changes as required. I was going to put in an issue with the text_embeds and text_mask needing to default to None to allow unconditional generation but you already fixed it today. I was also going to put in a request to support sample rate conversion but you've already done that today too - thank you!

[lucidrains]

@eonglints thank you! i definitely don't know enough about the best practices around resampling, so if you see anything that does not look right, would really appreciate a PR!

i plan on getting the full sampling script (prompt wav file -> full output wav file) finished by end of week as well as some multi-GPU distributed training scripts in place for all three attention networks

[amitaie]

Sorry for bumping this up, But @eonglints can you share some details on what does in mean "successfully training the semantic tokens transformer", to what loss\accuracy did you get? why did you used mHuBert 1000 and not the smaller one (with the kmeans), did you find the smaller one insufficient?

Thanks in ahead!

[eonglints]

Hey, so I was wondering what you think @lucidrains about the placement of the semantic token extraction and de-duplication of consecutive tokens. I can see from your recent commits that we have to deal with the fact that the sequence lengths will be variable because of unique consecutive. Just a thought - and this may be against your general design philosophy - would it make more sense to move the semantic token extraction and unique consecutive de-duplication to a dataloader rather than as part of the SemanticTransformer forward pass? That way, the sequence length will be fixed and we can also set up the dataloader to use pre-computed semantic tokens if they've already been computed in a previous run. Anyway, just a thought.

[lucidrains]

hey yup! i plan on supporting both! you can actually turn the unique consecutive off and handle it manually in the dataloader by setting unique_consecutive = False

[eonglints]

Nice, supporting both is a good shout!

[Rongjiehuang]

Hey, I wonder if the semantic token from the released 16khz Hubert also supports 24khz sound synthesis, or the 24kHz hubert is needed?

[Rongjiehuang]

Previous work using Hubert tokens suggests that target waveforms should be 16Khz.

https://github.com/facebookresearch/fairseq/blob/main/examples/textless_nlp/gslm/speech2unit/pretrained/hubert_feature_reader.py#L45
https://github.com/facebookresearch/fairseq/blob/main/examples/speech_to_speech/docs/direct_s2st_discrete_units.md#target-speech

[eonglints]

Hey, so the pre-trained Hubert feature extractor only works on 16kHz source audio and produces tokens at 50Hz. These sample rates are not related to the output audio sample rate though and these tokens cannot be trivially inverted back into the waveform space. In AudioLM, they condition SoundStream tokens (which happen to also be 16kHz) on these semantic tokens but you could train a higher bit-rate SoundStream model. Alternatively, you could use these semantic tokens to condition a GAN or diffusion model to get back up to the waveform space at any sample rate you like. Hope that helps!