Discrete speech units

Overview

Language models consume and emit discrete tokens, while speech is a continuous signal. Discrete speech units are the bridge: short stretches of audio are mapped to one of a finite set of codes, so that speech can share a vocabulary with text 2305.11000 §1 2305.11000 §4.1. The codes are usually derived from a self-supervised speech encoder, whose families the survey reviews 2205.10643 §3.

The common core: an encoder produces frame features, each frame gets one index from a finite set of size K, and the index sequence is what the model sees. Two choices differ between variants. The first is how the index is picked: the nearest k-means centre, or a learned quantiser. The second is whether repeated neighbouring indices are merged. Once the codes are fixed, a language model is trained on them with the ordinary next-token objective 2305.11000 §4.2. Changing the encoder, its layer or K changes the unit vocabulary, so every later stage has to be redone 2305.11000 §4.1.

Variants

k-means over self-supervised features (SpeechGPT)

A HuBERT-style encoder maps the waveform to frame features. A k-means model assigns each frame to its nearest cluster centre, and the cluster index is the unit. Adjacent duplicate indices are then removed, so the unit sequence is shorter than the frame sequence 2305.11000 §4.1. The paper names mHuBERT as the tokenizer 2305.11000 §3.1. The number of clusters K is not stated in the paper. The released setup uses 1000 [code: speechgpt/utils/speech2unit/README.md#L9], see Contradictions. Card: SpeechGPT.

Offline k-means targets, re-clustered (HuBERT-style training targets)

Here the cluster index is not a vocabulary for a generative model. It is a prediction target for masked frames. The survey describes a first pass with k-means over MFCC features. A second pass clusters a hidden layer of the first model to produce new targets. The survey reports that two such iterations are enough to match or beat the previous state of the art for low-resource recognition 2205.10643 §3.4.2. The survey does not describe these targets being used as a language-model vocabulary. The HuBERT paper has no card in this wiki, so this variant rests only on the survey's description. Card: Self-Supervised Speech Representation Learning: A Review.

Learned codebook (vector quantisation)

The index comes from a quantiser that is trained together with the encoder, not from a separate clustering step. The survey lists three examples. VQ-VAE picks the nearest codebook vector for each encoder output, passes gradients through the choice with a straight-through estimator, and adds a codebook term and a commitment term 2205.10643 §3.2.2. wav2vec 2.0 quantises the unmasked features with a Gumbel-softmax quantiser, uses product quantisation over several codebooks, and adds a term that rewards using every entry 2205.10643 §3.3.2. Discrete BERT takes its units from a frozen vq-wav2vec model and predicts masked units with cross-entropy 2205.10643 §3.4.2. The survey notes that the wav2vec 2.0 quantiser sits inside a pre-training objective, so its codes are a by-product of that model 2205.10643 §3.3.2. Card: Self-Supervised Speech Representation Learning: A Review.

When to use which

Implementation sketch

The common core is the unit extraction, then vocabulary expansion, then the stage-1 loss. Units are integers in 0 to K-1. Variant-specific choices are the encoder, K, the layer and the quantiser, none of which the SpeechGPT paper pins down. The learned-codebook variant replaces the nearest-centre step.

# Frame features from a self-supervised speech encoder  [arxiv:2305.11000 §4.1]
def speech_to_units(waveform, encoder, kmeans):
    frames = encoder(waveform)                 # (T, d)
    # k-means variant: nearest cluster centre per frame, index in 0..K-1
    # [arxiv:2305.11000 §4.1]. Learned-codebook variant: the index comes
    # from the trained quantiser instead  [arxiv:2205.10643 §3.3.2]
    indices = [kmeans.nearest(f) for f in frames]
    # Drop adjacent duplicates; durations are lost  [arxiv:2305.11000 §3.1]
    units = [indices[0]]
    for idx in indices[1:]:
        if idx != units[-1]:
            units.append(idx)
    return units                               # length <= T

# Add K unit tokens to the vocabulary  [arxiv:2305.11000 §4.2]
def expand_vocabulary(vocab, embeddings, num_units):
    new_vocab = vocab + [unit_token(k) for k in range(num_units)]
    # Original rows are kept; new rows are random  [arxiv:2305.11000 §4.2]
    new_emb = random_init(len(new_vocab), dim)
    new_emb[: len(vocab)] = embeddings
    return new_vocab, new_emb

# Stage 1: predict each unit from the units before it  [arxiv:2305.11000 §4.2]
def stage1_loss(model, units):
    nll = 0.0
    for i in range(len(units)):
        log_probs = model.log_next_token(units[:i])
        nll += -log_probs[units[i]]
    return nll

Contradictions

• Stage 2 length. The paper gives 2100 steps in the text 2305.11000 §5.1 and 4000 steps in its table 2305.11000 app. D.
• Backbone size. The paper states LLaMA-13B 2305.11000 §5.1. The released stage-1 script loads a 7B model [code: speechgpt/scripts/ma_pretrain.sh#L3].
• Learning rate. The paper gives a peak of 2e-4 for all stages 2305.11000 app. D. The released stage-2 and stage-3 scripts use 2e-5 cm_sft.sh com_sft.sh.
• Batch sizes. The paper gives 96 A100 GPUs 2305.11000 §5.1. The per-node batch sizes in the scripts do not multiply out to the paper's stage batch sizes on that hardware ma_pretrain.sh.

The two papers do not disagree about the discretisation itself. The survey describes k-means targets for masked prediction, while SpeechGPT uses the cluster index as a token vocabulary. The survey also reports numbers for encoder pre-training on recognition and identification tasks, and SpeechGPT reports no quantitative results for speech generation. So there are no competing numbers for the same task.

Papers

• Self-Supervised Speech Representation Learning: A Review: a survey of self-supervised speech methods. It covers k-means targets (HuBERT-style), learned quantisers (VQ-VAE, wav2vec 2.0) and Discrete BERT on frozen vq-wav2vec units. It proposes no new method.
• SpeechGPT: Empowering Large Language Models with Intrinsic Cross-Modal Conversational Abilities: uses HuBERT-style units with k-means and duplicate removal, expands an LLM's vocabulary with them, and trains in three stages.