Self-Supervised Audio Encoders

"BERT for speech is just BERT, but the targets you mask and predict have to be invented from the audio itself. The history of audio SSL is the history of better ways to invent those targets."

Echo, Pitch-Perfect AI Agent

20.0.4.1 The Two Self-Supervised Learning Families

Self-supervised learning (SSL) on audio splits into two families that differ on what kind of target the model learns to predict.

Contrastive SSL defines a target as this audio frame's true latent and trains the model to distinguish it from a set of distractor latents drawn from other frames. The InfoNCE loss treats the problem as a $|Q|$-way classification task: out of one true target plus $|Q| - 1$ distractors, identify the true one. wav2vec 2.0 lives in this family.

Masked-prediction SSL assigns each audio frame a discrete pseudo-label (a cluster ID), masks a subset of the frames, and trains the model to predict the cluster IDs of the masked frames from the unmasked context, using standard cross-entropy. The pseudo-label vocabulary is the hard part: it has to capture phonetic structure without being given any. HuBERT (and its descendants WavLM, BEATs) live in this family.

Both families inherit the BERT recipe (mask a span, predict the masked content), but they differ on whether the prediction target is a continuous latent contrasted against distractors (wav2vec 2.0) or a discrete cluster ID compared with cross-entropy (HuBERT). The empirical winner depends on the downstream benchmark. HuBERT and WavLM dominate the SUPERB benchmark (Yang et al., 2021), a broad evaluation across ten speech tasks. wav2vec 2.0 is simpler to implement and converges faster.

20.0.4.2 wav2vec 2.0: Contrastive Pretraining

Baevski et al.'s wav2vec 2.0 (Baevski et al., 2020) was the breakthrough audio SSL model that beat hybrid HMM-DNN systems on LibriSpeech with only an hour of labeled data after pretraining on 53,000 hours of unlabeled audio. The architecture has three components.

The CNN Feature Encoder

A stack of seven 1D convolutional blocks downsamples the raw 16 kHz waveform by 320x to produce a 512-dimensional latent vector $z_t$ per 25 ms frame at 50 Hz. Each block uses kernel sizes (10, 3, 3, 3, 3, 2, 2) with matching strides, group normalization on the first block, and GELU activations throughout.

The Transformer Context Network

The latent sequence $z = (z_1, z_2, \ldots, z_T)$ goes through a transformer encoder (12 layers, 768 hidden dim for wav2vec2-base; 24 layers, 1024 hidden dim for wav2vec2-large) to produce contextual embeddings $c_t$. Before entering the transformer, a fraction $p$ of time steps (default $p = 0.065$ as starting indices, with each span extending 10 frames) are selected and replaced with a single trained mask vector, exactly as BERT masks input tokens. Positional information uses a convolutional positional encoding (a depthwise conv applied to the input, added residually) rather than sinusoidal or learned position embeddings.

The Quantization Module

In parallel, the unmasked $z_t$ goes through a Gumbel-Softmax quantizer to produce a discrete target $q_t$. The quantizer uses product quantization with $G = 2$ groups and $K = 320$ codewords per group, exposing $K^G = 320^2 = 102{,}400$ effective codewords from $2 \times 320 = 640$ trainable codeword vectors. Concretely, the quantizer multiplies $z_t$ by a trainable projection matrix $W_q \in \mathbb{R}^{d \times K}$ to get logits $\ell_t \in \mathbb{R}^K$ per group, samples a near-one-hot Gumbel-Softmax selection vector at temperature $\tau$ (annealed from 2.0 down to 0.5 during training), and looks up the corresponding codebook entries. The two group outputs are concatenated and passed through a final projection to give the target $q_t$. The Gumbel-Softmax (Section 20.0.2.4) is what makes the codebook trainable end-to-end.

The Contrastive Loss

Figure 20.0.4.2 shows the full forward pass as a wiring diagram: the CNN encoder produces the latent stream $z$, the masked branch sends it through the transformer to make $c_t$, the parallel quantizer branch produces the discrete target $q_t$, and the InfoNCE loss compares the two. Concretely, for each masked time step $t$, the model takes the transformer output $c_t$ and computes its similarity to the true target $q_t$ as well as to a set of $|Q| - 1$ distractor targets sampled from other masked time steps in the same utterance. The InfoNCE loss is:

where $\mathrm{sim}(a, b) = a^\top b / (\lVert a \rVert \cdot \lVert b \rVert)$ is the cosine similarity, $\kappa$ is a temperature hyperparameter (typically 0.1), and $Q_t$ is the set of $|Q|$ candidate vectors (the true $q_t$ plus 100 distractors sampled uniformly from other masked positions). A small diversity loss is added to encourage uniform codebook usage and prevent collapse. The contrastive objective forces $c_t$ to be more similar to its true target $q_t$ than to distractors; because distractors come from other speech frames, the model learns to discriminate fine-grained phonetic content.

Fine-Tuning for ASR

For ASR, the quantizer is dropped, a linear character classifier is added on top of the transformer, and the model is fine-tuned with CTC loss (Section 20.0.3.6) on labeled speech. The famous result from the original paper: pretrain on 53,000 hours of unlabeled LibriVox audio, then fine-tune on just 10 minutes of labeled LibriSpeech, and achieve 4.8% WER on LibriSpeech test-clean, comparable to supervised systems trained on 100x more labeled data.

import torch, librosa
from transformers import Wav2Vec2Model, Wav2Vec2FeatureExtractor

# Bare backbone (no CTC head): outputs frame-level contextual embeddings c_t.
extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/wav2vec2-base")
model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base").eval()

waveform, sr = librosa.load("clip.wav", sr=16_000, mono=True)
inputs = extractor(waveform, sampling_rate=16_000, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
# outputs.last_hidden_state: (1, n_frames, 768) - the c_t from Figure 20.0.4.2.
# n_frames = ceil(n_samples / 320), the 50 Hz output rate of the CNN feature encoder.
print(outputs.last_hidden_state.shape)
# Example: torch.Size([1, 149, 768]) for a ~3 second clip
# (3.0 s * 16,000 samples/s / 320 samples/frame ≈ 150 frames)

20.0.4.3 HuBERT: Masked Cluster Prediction

Hsu et al.'s HuBERT (Hsu et al., 2021) replaces the wav2vec 2.0 contrastive objective with a BERT-style cross-entropy on cluster IDs. The acronym stands for "Hidden Unit BERT": the audio is clustered into a set of discrete "hidden units" (phoneme-like pseudo-labels) and the transformer is trained to predict the unit IDs of masked frames. The model and the targets are trained jointly across iterations.

HuBERT Architecture

The CNN feature encoder is identical to wav2vec 2.0's (320x downsample to 50 Hz, 512-dim per frame). The transformer is also a standard BERT-style stack. The key difference is what the transformer is asked to predict.

Iterative Tokenization

The challenge HuBERT solves is "where do the cluster IDs come from?" Acoustic frames have no pre-existing phoneme labels (otherwise the model would not need SSL). HuBERT's answer is to bootstrap the clustering iteratively:

1. Iteration 1. Compute MFCC features for every 25 ms window of the unlabeled training audio. Run $k$-means with $K = 100$ clusters on a random subset (a few hundred thousand frames). Assign every frame to its nearest cluster; the cluster IDs become the pseudo-labels.
2. Iteration 1 training. Train the HuBERT transformer with masked prediction of these MFCC-derived cluster IDs. Mask 50% of the input frames (each masked span is up to 10 frames; the starting indices are sampled randomly), feed the masked sequence through the transformer, and predict the cluster ID at each masked position with cross-entropy.
3. Iteration 2. Extract intermediate transformer layer features (typically the 6th layer of the iteration-1 model) for the same training audio. Re-cluster with $k$-means at a larger $K = 500$. The new cluster IDs are derived from contextual transformer features and so are much more phonetically meaningful than raw MFCCs.
4. Iteration 2 training. Train a fresh HuBERT transformer with masked prediction of the iteration-2 cluster IDs.
5. Iteration 3 (optional). Repeat once more with $K = 1000$ for the largest model variants.

Each iteration takes hundreds of GPU-days and the published HuBERT models stop at iteration 2 or 3. The intuition is that the transformer learns better representations than MFCCs, which yield better cluster IDs, which yield a better transformer in the next round.

Cluster Prediction Head

At training time the model needs to predict a cluster ID at each masked position. Rather than a flat linear classifier, HuBERT projects the masked transformer output into a low-dimensional space and compares it (via cosine similarity) to projected cluster centers. The cluster centers $\mu_k$ from $k$-means are themselves linearly projected to the same low-dimensional space through a small trainable matrix $A \in \mathbb{R}^{d \times d_c}$:

where $o_t$ is the transformer output at masked position $t$ and $\tau$ is a learned temperature. The cross-entropy loss over the true cluster ID supplies the gradient. Both the transformer and the projection $A$ are trained; the cluster centers $\mu_k$ themselves are frozen between iterations.

Fine-Tuning for ASR

Fine-tuning replaces the cluster-prediction head with a linear character head and trains with CTC loss on labeled speech, exactly as wav2vec 2.0 does. Published checkpoints include facebook/hubert-base-ls960 (95M params, base configuration trained on 960 hours of LibriSpeech), facebook/hubert-large-ls960-ft (300M params, fine-tuned for ASR), and the larger facebook/hubert-xlarge variants.

import torch, librosa
from transformers import Wav2Vec2Processor, HubertForCTC

# Fine-tuned HuBERT-Large with a CTC head trained on LibriSpeech 960h.
processor = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft")
model = HubertForCTC.from_pretrained("facebook/hubert-large-ls960-ft").eval()

waveform, sr = librosa.load("clip.wav", sr=16_000, mono=True)
inputs = processor(waveform, sampling_rate=16_000, return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits   # (1, n_frames, vocab_size)

# Greedy CTC: argmax per frame, then collapse repeats + remove blanks.
pred_ids = torch.argmax(logits, dim=-1)
transcript = processor.batch_decode(pred_ids)[0]
print(transcript)
# "THE QUICK BROWN FOX JUMPS OVER THE LAZY DOG"

20.0.4.4 WavLM: HuBERT Plus Utterance Mixing

Chen et al.'s WavLM (Chen et al., 2022a) is HuBERT with one critical augmentation. The motivation: HuBERT and wav2vec 2.0 dominate clean single-speaker ASR but degrade on conversational, noisy, and multi-speaker audio (think podcasts, meetings, mobile-phone recordings). WavLM addresses this by augmenting training with utterance mixing.

The trick is conceptually simple. During pretraining, with some probability mix two random training utterances together at the waveform level (overlay them with a random gain), but only require the model to predict the cluster IDs of the first utterance's frames. The model is forced to learn to separate overlapping speakers and noise sources because the prediction target ignores the interference. WavLM also adds noise augmentation (overlaying ambient noise from MUSAN at SNR 5-20 dB) and gain augmentation. The cluster targets come from the same HuBERT-style iterative $k$-means pipeline.

The result is that WavLM matches HuBERT on clean ASR but substantially outperforms it on speaker verification, speaker diarization, speech separation, and any other task involving overlapping speakers or background noise. It topped the SUPERB benchmark when released and remains the default choice when downstream conditions are noisy. Published checkpoints include microsoft/wavlm-base (95M params) and microsoft/wavlm-large (316M params).

WavLM is also the encoder used by Mimi (Section 20.0.2.5) as a distillation target for its first RVQ codebook, which is how Moshi gets a semantically meaningful first stream out of its audio codec.

import torch
import librosa
from transformers import AutoFeatureExtractor, WavLMModel

# microsoft/wavlm-base-plus is the standard noise-robust checkpoint;
# upgrade to "microsoft/wavlm-large" for 316M params and better diarization.
feat = AutoFeatureExtractor.from_pretrained("microsoft/wavlm-base-plus")
model = WavLMModel.from_pretrained("microsoft/wavlm-base-plus").eval()

waveform, sr = librosa.load("noisy_meeting.wav", sr=16_000, mono=True)
inputs = feat(waveform, sampling_rate=16_000, return_tensors="pt")

with torch.no_grad():
    out = model(**inputs, output_hidden_states=True)

# out.last_hidden_state: (1, n_frames, 768) at 50 Hz
frame_emb = out.last_hidden_state.squeeze(0)
clip_emb = frame_emb.mean(dim=0)            # (768,) for speaker / clip-level tasks
print(frame_emb.shape, clip_emb.shape)

20.0.4.5 BEATs: Self-Distilled Tokenizer for General Audio

Chen et al.'s BEATs (Chen et al., 2022b/2023, "Bidirectional Encoder representation from Audio Transformers") generalizes the HuBERT recipe beyond speech to general audio (sound events, music, environmental sounds). The HuBERT clustering step works because phonemes are a useful inductive bias for speech, but for general audio there is no analogue: a "sound event" cluster is much less well-defined than a phoneme. BEATs solves this with a self-distilled discrete tokenizer trained alongside the main encoder.

The architecture has two components. The first is an acoustic tokenizer: a small transformer that takes mel-spectrogram patches as input and outputs discrete token IDs via a Gumbel-Softmax codebook. The second is the main audio transformer encoder with the same patch input format. The training cycle alternates between two phases:

1. Tokenizer phase. Train the tokenizer to predict pseudo-labels derived from the main encoder's outputs (self-distillation). Specifically, the main encoder's intermediate features are clustered (or the encoder itself acts as a teacher whose predictions become tokenizer targets), and the tokenizer is trained with cross-entropy to match these targets.
2. Encoder phase. Freeze the tokenizer and use its outputs as cluster IDs to train the main encoder with BERT-style masked prediction (50% mask rate, cross-entropy loss on tokenizer IDs).

The two phases alternate every few epochs, with each producing better targets for the other. The result is a general-audio encoder that achieves state-of-the-art on AudioSet (the 527-class audio event benchmark) and ESC-50, while also performing well on speech tasks. Published checkpoints live under microsoft/beats with sizes from BEATs_iter3 (90M, after 3 self-distillation iterations) up to BEATs_iter3_plus_AS2M (the same model fine-tuned on AudioSet).

BEATs is the encoder to reach for when the downstream task involves non-speech audio (environmental sounds, music) and a HuBERT-family encoder underperforms because speech-specific clustering does not transfer.

20.0.4.6 The Cheat-Sheet

20.0.4.7 Using Them: The Frozen-Encoder + Linear-Probe Pattern

The most common way to put one of these encoders to work is the linear probe: freeze the pretrained encoder, extract a single fixed-size embedding per audio clip by mean-pooling the frame-level outputs, and train a small classification head (often just a linear layer) on labeled task data. The pattern is identical for wav2vec 2.0, HuBERT, WavLM, and BEATs; only the model class and processor differ. The snippet below uses librosa for audio loading and resampling alongside transformers for the encoder.

import torch
import librosa
from transformers import Wav2Vec2Processor, HubertModel

processor = Wav2Vec2Processor.from_pretrained("facebook/hubert-base-ls960")
model = HubertModel.from_pretrained("facebook/hubert-base-ls960").eval()

def extract_embedding(file_path):
    waveform, sr = librosa.load(file_path, sr=16_000, mono=True)
    inputs = processor(waveform, sampling_rate=16_000, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    # outputs.last_hidden_state: (1, n_frames, hidden_dim=768)
    # Mean-pool across the time axis to get a single clip embedding.
    return outputs.last_hidden_state.mean(dim=1).squeeze(0)   # (768,)

emb = extract_embedding("clip.wav")
print(emb.shape)   # torch.Size([768])

# Now train any sklearn / torch classifier on top of these embeddings:
# from sklearn.linear_model import LogisticRegression
# clf = LogisticRegression(max_iter=1000).fit(train_embeddings, train_labels)

20.0.4.8 Which Encoder Should the Reader Use?

A short decision tree, with rationale:

• English ASR on clean speech: any of the four works; wav2vec 2.0 base is the smallest and fastest.
• Multilingual ASR: wav2vec 2.0 XLS-R (Babu et al., 2022), an extension of wav2vec 2.0 trained on 128 languages, or HuBERT multilingual variants.
• Noisy or conversational speech, diarization, speaker verification: WavLM. The utterance-mixing augmentation pays for itself.
• General audio classification (events, music, environmental sounds): BEATs or AST. AST is older and simpler; BEATs is newer and slightly stronger on AudioSet.
• Linear-probe downstream classification: HuBERT or WavLM gives the most generally useful frame-level representations. The default in most HuggingFace audio tutorials is DistilHuBERT (Chang et al., 2022), a 6-layer distilled HuBERT that fine-tunes 2x faster with minimal accuracy loss.
• Speech-to-speech LLM (Moshi-style): WavLM (used as the distillation target for Mimi's first RVQ codebook in Section 20.0.2.5).