BEATs: Audio Pre-Training with Acoustic Tokenizers

Chapter 1: From waveform to patches

A Transformer eats a sequence of vectors. Audio arrives as a 1-D waveform — tens of thousands of samples per second. The standard bridge, which BEATs inherits from the vision-transformer view of audio (HTS-AT, AST), has two steps: turn the waveform into a 2-D mel spectrogram, then chop that image into square patches.

Step 1: the mel spectrogram

Run a short-time Fourier transform: slide a ~25 ms window across the waveform every ~10 ms, take the magnitude of its frequency content, and warp the frequency axis onto the perceptual mel scale (which spaces low frequencies finely and high frequencies coarsely, like the ear does). For BEATs the result is a matrix with 128 mel bins on the frequency axis and one column per 10 ms frame. A 10-second clip at this hop is roughly 998 frames × 128 mels — call it F frames by M mels.

Step 2: patchify

Treat that spectrogram as a single-channel image of shape [1, F, M] and split it into non-overlapping 16×16 patches, exactly as a ViT splits a photo. With F ≈ 1000 and M = 128 this yields about (1000/16) × (128/16) = 62 × 8 ≈ 496 patches. Each patch is a 16×16 = 256-number tile of the spectrogram. A small convolution (a "patch-embed" stem) projects each tile to a D-dimensional vector — in BEATs, D = 768. We now have a sequence:

This is the object every later chapter operates on. The tokenizer assigns each patch x_n a discrete code; the encoder masks a subset of patches and predicts those codes; the distillation step re-fits the codes. Everything is per-patch.

Worked example: counting the patches

Let's make the shapes concrete on a tiny clip so the arithmetic is unambiguous. Take a 1-second clip, hop 10 ms, so F = 100 frames; M = 128 mels; patch = 16. Patches along time: ⌊100/16⌋ = 6 (the leftover 4 frames are dropped or padded). Patches along frequency: 128/16 = 8. Total patches N = 6 × 8 = 48. Each patch is 256 raw numbers; the conv stem maps 256 → 768. So a 1-second clip becomes a sequence of 48 vectors of length 768 — a [48, 768] tensor. Double the clip length, double N; the per-patch dimension never changes.

Chapter 2: The Acoustic Tokenizer

The tokenizer is the heart of BEATs. Its job: take each patch vector and emit a single integer in {1,…,K} — its code id. The set of K possible codes is the codebook. BEATs uses a codebook of K = 1024 entries, each a learnable vector in a low-dimensional code space (the paper uses a 256-dim quantizer space, projected from the encoder's 768). Quantizing means: for a patch, find the nearest codebook vector and report its index.

Concretely the tokenizer is itself a small Transformer (the "tokenizer encoder") followed by a vector-quantization (VQ) layer. Let z_n ∈ ℝ^d be the tokenizer's continuous output for patch n, and let the codebook be {e₁,…,e_K} with each e_k ∈ ℝ^d. The code id is nearest-neighbour:

and the quantized vector handed downstream is just the chosen codebook entry, q̂_n = e_cn. Two patches that map to the same nearest entry get the same token — they are declared "the same kind of sound."

The normalized-distance trick

BEATs (following the BEiT-v2 / VQ-KD line) ℓ2-normalizes both z_n and the codebook entries before computing distance. On the unit sphere, squared Euclidean distance and cosine similarity are equivalent: ‖a−b‖² = 2−2·a·b when ‖a‖=‖b‖=1. So nearest-neighbour quantization becomes "pick the code with the highest cosine similarity." Normalizing keeps the codebook from collapsing onto a few high-norm entries and makes training far more stable — a recurring lesson in VQ.

The gradient problem and the straight-through estimator

argmin is not differentiable — you cannot backprop through "pick the nearest". The fix, from VQ-VAE, is the straight-through estimator: on the forward pass use the quantized q̂_n; on the backward pass copy the gradient straight through as if q̂_n = z_n. In code this is the one-liner q = z + stop_grad(qhat − z): the value equals q̂_n, but the gradient flows to z_n unchanged.

To actually move the codebook and the encoder toward each other, VQ adds two auxiliary terms. With sg[·] the stop-gradient:

The first term (the codebook loss) pulls the chosen code toward the encoder output. The second (the commitment loss, weight β ≈ 0.25) pulls the encoder output toward its code so the encoder "commits" to codes instead of drifting. (BEATs in practice uses the exponential-moving-average update for the codebook, an equivalent of the first term; the commitment term stays.)

Worked example: quantizing one patch by hand

Take a 2-D code space (d = 2) and a tiny codebook of three entries: e₁ = (1, 0), e₂ = (0, 1), e₃ = (−1, 0). The tokenizer outputs z = (0.8, 0.3). Distances:

‖z−e₁‖² = (0.8−1)² + (0.3−0)² = 0.04 + 0.09 = 0.13
‖z−e₂‖² = (0.8−0)² + (0.3−1)² = 0.64 + 0.49 = 1.13
‖z−e₃‖² = (0.8−(−1))² + 0.09 = 3.24 + 0.09 = 3.33

The smallest is e₁, so c = 1 and the quantized vector is (1, 0). The commitment loss would then push z toward (1, 0), and the codebook update would nudge e₁ toward (0.8, 0.3). You will reproduce exactly this nearest-code logic in Code Lab #1.

python
# Vector quantization with the straight-through estimator (VQ-VAE)
import torch
import torch.nn.functional as F

def quantize(z, codebook, beta=0.25):
    """
    z:        [N, d]   tokenizer outputs (one per patch)
    codebook: [K, d]   the K learnable code vectors
    returns:  q [N, d] quantized, ids [N], vq_loss
    """
    # L2-normalize so distance == cosine (BEATs / BEiT-v2 trick)
    zn  = F.normalize(z, dim=-1)
    en  = F.normalize(codebook, dim=-1)
    # Squared distances [N, K]; argmin over K = nearest code
    d2  = (zn.pow(2).sum(-1, keepdim=True)
           - 2 * zn @ en.T
           + en.pow(2).sum(-1))      # [N, K]
    ids = d2.argmin(-1)                  # [N] the token ids
    qhat = en[ids]                          # [N, d] chosen codes
    # Codebook + commitment losses
    cb  = F.mse_loss(qhat, zn.detach())
    cm  = F.mse_loss(zn, qhat.detach())
    vq_loss = cb + beta * cm
    # Straight-through: forward = qhat, backward = identity to z
    q = zn + (qhat - zn).detach()
    return q, ids, vq_loss

Chapter 3: Masked Prediction — the model's objective

Now freeze the tokenizer for a moment and look at the model being pre-trained. Given the patch sequence X = [x₁,…,x_N], the tokenizer has labeled every patch with a code id c_n ∈ {1,…,K}. We sample a mask set M ⊂ {1,…,N} covering about 75% of the patches. For masked positions we replace x_n with a shared learnable [MASK] embedding; unmasked positions pass through unchanged. The Transformer encoder reads this corrupted sequence and produces hidden states h₁,…,h_N. A linear "prediction head" maps each masked position's hidden state to a distribution over the K codes:

and the loss is plain cross-entropy against the tokenizer's labels, summed over masked positions only:

("MAM" = masked audio modeling.) That is the entire pre-training objective for the model. No reconstruction, no waveform, no autoregression — just "for each blanked patch, classify which of the 1024 acoustic codes it was." Because the codes are semantic (we'll earn that adjective in Chapters 4–5), classifying them forces the encoder to reason about what the surrounding sound is, not just what it looks like.

Why mask 75%? Why only score masked positions?

Two design choices that look small and aren't:

High mask ratio. If you mask only 15% (BERT's text ratio), audio patches are so locally redundant that the model can solve each blank by copying a neighbour — a barely-disturbed harmonic stack is trivially interpolated. Masking ~75% (as MAE found for images, and BEATs adopts) removes the easy local cues and forces long-range, semantic inference: to fill a hole you must recognize the event spanning the whole clip.

Score only masked positions. If you also asked the model to predict unmasked patches it can see, it would just learn a copy — the loss would reward identity, not understanding. Restricting L_MAM to M makes every contributing term a genuine guess.

Worked example: the cross-entropy on one masked patch

Suppose K = 4 codes and the prediction head outputs logits (2.0, 1.0, 0.1, −1.0) for a masked patch whose true code is c = 1 (the first). Softmax: exponentiate → (7.39, 2.72, 1.11, 0.37), sum = 11.59, so p = (0.638, 0.235, 0.096, 0.032). The loss for this patch is −log p₁ = −log 0.638 = 0.449 nats. If the model had been confident-and-correct (p₁ → 1) the loss → 0; confident-and-wrong (p₁ → 0) sends it to ∞. Cross-entropy only looks at the probability assigned to the right code — the within-code acoustic texture is invisible to it. That invisibility is the whole reason discrete targets beat reconstruction.

python
# Masked Audio Modeling loss — predict tokenizer codes at masked patches
import torch
import torch.nn.functional as F

def mam_loss(encoder, head, patches, code_ids, mask_ratio=0.75):
    N, D = patches.shape
    # 1. choose ~75% of positions to mask
    n_mask = int(N * mask_ratio)
    mask = torch.randperm(N)[:n_mask]
    # 2. replace masked patches with the learnable [MASK] embedding
    x = patches.clone()
    x[mask] = head.mask_token        # [D] broadcast
    # 3. encode the corrupted sequence (bidirectional, all visible at once)
    h = encoder(x)                    # [N, D]
    # 4. classify each masked position over the K codes
    logits = head.proj(h[mask])       # [n_mask, K]
    # 5. cross-entropy against tokenizer labels, masked positions only
    return F.cross_entropy(logits, code_ids[mask])

Chapter 4: Label Distillation — making codes semantic

We have postponed the central question long enough: how does a code id become semantic? A VQ codebook trained only to reconstruct, or trained on random patches, gives codes that cluster acoustically — they separate loud from quiet, high from low — but not violin from flute. BEATs' answer is label distillation (it calls the objective the acoustic tokenizer's knowledge distillation): train the tokenizer so that its discrete codes predict the outputs of a trained masked-prediction model. The model has, through pre-training and fine-tuning, learned a semantic representation; distillation pours that semantics into the discrete codes.

The estimator/teacher split

The tokenizer is trained as a tiny student that must reconstruct a teacher's behaviour through the bottleneck of a discrete code. Two pieces:

The distillation loss makes the prediction-from-code match the teacher. BEATs maximizes the cosine similarity between the code-derived prediction t̂_n and the teacher target t_n (equivalently, minimizes a normalized MSE):

where t̂_n is computed from the quantized code only, sg[·] freezes the teacher, and L_vq is the codebook/commitment term from Chapter 2. The crucial constraint: t̂_n may depend on the patch only through its discrete code e_cn. So if two patches must yield very different teacher representations, they are forced into different codes. The discreteness becomes an information bottleneck that the codebook must spend wisely on whatever the teacher considers meaningful — i.e. on semantics.

Worked example: distillation pressure on the codebook

Imagine three patches A, B, C with teacher targets t_A = (1, 0), t_B = (0.9, 0.1), t_C = (0, 1). The teacher says A and B are nearly identical (cos ≈ 0.996) and C is orthogonal. A code can only emit one prediction t̂ per id. If A and B share code 1 (predicting (0.95, 0.05)) the cosine loss is tiny for both. But forcing C into code 1 as well would cost cos((0.95,0.05),(0,1)) ≈ 0.05 — terrible. So distillation pushes C onto a different code. The codebook self-organizes so that "patches the teacher treats alike share a code; patches the teacher distinguishes get split." That is semantics, learned. You'll implement this similarity-driven split in Code Lab #2.

python
# Acoustic-tokenizer knowledge distillation (cosine objective)
import torch
import torch.nn.functional as F

def tokenizer_distill_loss(tok, decoder, teacher, patches):
    # teacher is FROZEN — the trained BEATs model from the last iteration
    with torch.no_grad():
        t = teacher(patches)                 # [N, D] target representation
    z = tok(patches)                        # [N, d] tokenizer outputs
    q, ids, vq = quantize(z, tok.codebook)   # discrete bottleneck
    # predict the teacher target FROM THE CODE ONLY
    t_hat = decoder(q)                       # [N, D]
    # match direction, not magnitude
    distill = -F.cosine_similarity(t_hat, t.detach(), dim=-1).mean()
    return distill + vq, ids

Chapter 5: The Iterative Loop

We have a chicken-and-egg problem. To make codes semantic (Chapter 4) we distill from a trained BEATs model. But to train that model (Chapter 3) we need semantic codes. Which comes first? BEATs breaks the cycle with a cold start and then iterates.

Iteration 0: the random-projection tokenizer (cold start)

You cannot distill from a teacher you don't have. So the very first tokenizer is not learned at all — it is a fixed random-projection quantizer. Project each patch through a random (frozen) matrix into the code space, then nearest-neighbour against a random codebook. The codes are acoustically meaningful but not semantic — exactly the HuBERT-style starting point. Crucially, even a random discrete target is enough to bootstrap a first masked-prediction model that learns something useful, just as HuBERT's first iteration runs on k-means-of-MFCC labels.

The alternation

Each pass through the loop has two phases that never train together: (a) freeze the model, train the tokenizer by distillation; (b) freeze the tokenizer, train the model by masked prediction. Like the EM algorithm or like HuBERT's relabeling rounds, alternating between two frozen-and-trained roles is what keeps the loop stable. Joint end-to-end training would let the model and tokenizer collude on a degenerate code (e.g. constant labels with zero loss). The freezing prevents collapse.

Worked example: tracing label quality across iterations

Score "label semanticity" as the fraction of same-event patches that share a code (higher = more semantic). Suppose random-projection labels give 0.30 (mostly acoustic accidents). Iter-1's model, trained on those, learns features that — when distilled — produce codes scoring 0.62. Iter-2's model, trained on the 0.62 labels, distills to 0.71. Iter-3 would give ~0.73. The deltas — +0.32, +0.09, +0.02 — show the textbook diminishing returns that justify stopping at two. (These numbers are illustrative; the paper reports the analogous downstream-accuracy gains plateauing by iteration 2.)

Chapter 6: The Encoder — what actually does the work

Both the tokenizer and the masked-prediction model share the same backbone: a ViT-style Transformer encoder adapted for spectrograms, drawn from HTS-AT-flavoured audio transformers. Let's pin down the architecture and the data flow with real shapes.

The stack

The base BEATs encoder is roughly 90M parameters — 12 Transformer layers, hidden width 768, 12 attention heads, FFN width 3072. Two adaptations matter for audio:

Convolutional positional embedding. Instead of (or in addition to) sinusoidal position codes, a depthwise convolution over the patch sequence injects relative position — it lets nearby patches in time-frequency know they are neighbours, which is more natural for the 2-D grid structure of a spectrogram than 1-D sinusoids.

Relative-position bias in attention. Each attention score gets an additive learned bias depending on the relative offset between query and key patch. So the raw attention is

where b_i−j is a learned scalar shared by all patch pairs the same distance apart. This is the relative-position scheme from Swin/HTS-AT — it improves length generalization and locality without fixing a maximum sequence length the way learned absolute positions do.

Pre-train vs fine-tune: only the head changes

During pre-training the encoder reads a masked patch sequence and the head classifies codes at masked positions. During fine-tuning the same encoder reads the full (unmasked) sequence; a new head mean-pools all N hidden states into a clip-level vector and classifies it into the C downstream labels (e.g. 527 AudioSet classes). The expensive 90M-parameter encoder is reused verbatim — only the cheap output head is replaced. That reuse is the entire value proposition of self-supervised pre-training: pay once for the encoder on unlabeled data, amortize over many downstream tasks.

Worked example: where the FLOPs go

For N = 496 patches and D = 768, one attention layer costs roughly N²·D for the QK^T·V products ≈ 496²·768 ≈ 1.9×10⁸ mults, while the FFN costs roughly N·D·(4D) ≈ 496·768·3072 ≈ 1.2×10⁹ — the FFN dominates a single layer at this sequence length. But because attention is O(N²), doubling the clip length (N→992) quadruples the attention term while only doubling the FFN term. This is why long-audio modeling is hard and why patching (which keeps N small) is essential.

Chapter 7: Results & Ablations

BEATs' headline claim is straightforward: with the same encoder size and data, swapping the masked-prediction target from reconstruction or hand-built clusters to iterated semantic codes moves audio classification to state of the art. The flagship benchmark is AudioSet (~2M weakly-labeled 10-second YouTube clips, 527 sound classes, scored by mean average precision, mAP).

The headline numbers (qualitative where exact figures are uncertain)

The single-model BEATs reached roughly 48.6 mAP on AudioSet-2M and the ensemble roughly 50.6 mAP — the best reported at the time, surpassing supervised and prior-SSL audio transformers. It also topped smaller benchmarks (ESC-50 environmental sounds near the high-90s accuracy, strong scores on speech-command and other tasks), showing the representation transfers beyond AudioSet.

The ablations that prove the thesis

The paper's experiments are designed to isolate why it works:

• Target type matters most. Holding everything else fixed and only changing the masked-prediction target, the semantic discrete codes beat both raw-reconstruction targets and random-projection (acoustic) codes. This is the core claim: the target, not the architecture, is the lever.
• Iteration helps, then saturates. Iteration 1 (random labels) is already competitive; iteration 2/3 (distilled semantic labels) adds a clear margin; further iterations barely move. Diminishing returns justify stopping early.
• Distillation > reconstruction for the tokenizer. A tokenizer trained to reconstruct the mel patch through its code (acoustic) yields worse downstream accuracy than one trained to match a teacher representation (semantic). Same VQ machinery, different supervision — the supervision is what counts.
• Mask ratio. ~75% is near-optimal; much lower lets the model copy neighbours, much higher starves it of context.

Worked example: reading an mAP delta

mAP averages the per-class average precision over all 527 classes. Suppose reconstruction-target pre-training gives 45.0 mAP and semantic-code pre-training gives 48.6. The +3.6 absolute looks small, but on AudioSet — where the gap between consecutive SOTA papers is often well under a point — 3.6 mAP from a single design change (the target) is large. It says: on average across every sound class, the semantic-code model ranks true positives meaningfully higher. The cost was zero extra labels — the whole pipeline ran on the same unlabeled audio.

Chapter 8: BEATs Explorer

This is the whole pipeline in one interactive surface. A synthetic clip is patchified; a learnable-looking codebook tokenizes each patch into a code id; a fraction of patches are masked; the masked-prediction head guesses their codes. Drive every knob and watch the mask, the token map, and the (simulated) prediction accuracy respond. There's no quiz here — this is the payoff. Read the three readouts as you move the controls.

Things to notice as you play:

• At mask = 20%, prediction accuracy is high even at iteration 0 — the model just copies visible neighbours. The task is too easy to teach semantics.
• Push mask to 75% and iteration-0 accuracy collapses — random-projection codes give the model nothing to reason about when context is scarce. Now raise iter: as the codes turn semantic, accuracy recovers, because the model can infer "this whole region is one event" and predict its code.
• Crank K very high and accuracy dips again at every iteration — too many rarely-used codes (collapse) make the classification target noisy. K = 1024 in the real model is a tuned compromise this toy mimics in miniature.

That interaction — high mask ratio is only learnable when the target is semantic, which only happens after iteration — is the single most important intuition in the paper. The mask ratio and the iterative tokenizer are not two independent tricks; they are two halves of one mechanism.

Chapter 9: Limitations & Connections

BEATs is a clean, influential result, but it is not magic. Where does it strain, and where does it sit in the larger map of representation learning?

Limitations

• Multi-stage and expensive to reproduce. The full recipe is: random-proj. bootstrap → train model → distill tokenizer → retrain model, twice. Each masked-prediction run is a full pre-training over millions of clips. It is a pipeline, not a single loss — costly and fiddly to reproduce end-to-end.
• O(N²) attention caps clip length. Like every vanilla Transformer, the encoder pays quadratic attention in the number of patches, so very long recordings (minutes) are out of reach without windowing or efficient-attention surgery.
• Tokenizer choices are heuristic. K = 1024, the 75% mask ratio, the cosine distillation, the number of iterations — all tuned, not derived. A different audio domain might need different settings, and the paper offers limited theory for picking them a priori.
• Inherits the teacher's blind spots. Distillation pours the teacher's semantics into the codes — including the teacher's biases. If the bootstrap model under-represents a rare sound class, the codebook may never carve out a code for it.
• Patch granularity is fixed. A 16×16 time-frequency patch is a compromise; events much shorter (clicks) or much longer (drones) than a patch are awkwardly represented.

Where BEATs sits

BEATs is the audio member of a family of discrete-token masked-prediction methods that all share one move — make the prediction target a learned discrete code rather than raw input:

Read next on Engineermaxxing

• Attention Is All You Need (Veanor) — the self-attention, multi-head, and positional-encoding machinery the BEATs encoder is built from.
• Vision Transformer (Gleam) — the patchify-an-image idea BEATs borrows to patchify a spectrogram.
• Embeddings (Gleam) — why mapping inputs to vectors (and codes) is the substrate of all of this.
• Diffusion Models (Gleam) — a contrasting generative pre-training objective; compare its denoising target to BEATs' discrete masked target.
• All Veanors — more paper walkthroughs.

The cheat sheet