Masked Autoencoders that Listen

Chapter 0: The Sound Nobody Labeled

Put a microphone anywhere on Earth and leave it running. In an hour you will have collected an hour of audio. In that same hour you will have collected exactly zero labels, because labels do not fall out of the sky — a human has to listen, decide what happened, and type it in.

That asymmetry is the entire reason this paper exists. Recorded sound is nearly free; annotated sound is expensive, slow, and — this is the part people underestimate — ambiguous. Ask five annotators to label ten seconds of a busy street and you will get five different sets of tags, all defensible.

Let us put real numbers on the scarcity, because the whole design of Audio-MAE is a response to these numbers. The field’s biggest audio dataset is AudioSet: roughly two million ten-second YouTube clips, annotated with 527 audio-event classes. Here is what that annotation actually looks like when you open the file.

Read the last three rows again. A weak label tells you that a class occurs somewhere inside a ten-second window — it does not tell you when, or for how long, or whether it is the loudest thing in the clip. If you train a classifier on that, the model must simultaneously learn what a dog sounds like and where in the clip to look, from a supervision signal that answers neither question directly.

There are exactly four ways out of a label shortage, and the field tried them in this order:

The field’s workaround: borrow eyes to hear with

Faced with too few audio labels, the audio community made a move that is both ingenious and, on reflection, deeply strange: initialize audio models from ImageNet.

The logic runs like this. A mel-spectrogram is a two-dimensional array of numbers. So is a grayscale image. Vision Transformers are very good at two-dimensional arrays of numbers, and there are 1.28M cleanly labeled images sitting in ImageNet. So: pre-train a ViT on images, then feed it spectrograms.

Concretely, two surgeries make the weights fit. Both are worth spelling out with shapes, because the awkwardness of the surgery is the argument against it:

This works better than it has any right to. AST (DeiT-B backbone), PaSST (DeiT-B), MBT (ViT-B on ImageNet-21K), HTS-AT (Swin-B) and PSLA (EfficientNet) all use some version of it, and they held the AudioSet leaderboard for years. So why be suspicious?

Because there is a much better-behaved precedent, and image-to-audio is not it. When video models are initialized from image models — I3D and 3D-ResNets inflate 2D kernels into 3D — the transfer is homogeneous: a video frame literally is an image. Spectrogram transfer is heterogeneous. It remains unclear, in the paper’s own words, why it helps at all beyond arguably similar low-level semantics: the shapes of spectrograms and the shapes of visual objects.

And the axes do not mean the same things. This is the crux, and Chapter 6 will turn it into an architecture decision:

There is one more cost, and it is the kind that never shows up in a benchmark table: label bias transfers too. ImageNet has its own well-documented skews — geographic, cultural, taxonomic. Initialize every audio model on it and those skews propagate into a modality that never asked for them. The paper flags this explicitly as a reason to prefer audio-only pre-training: it prevents uncontrollable bias transfer from another modality.

The in-domain alternatives, and why they were still expensive

Self-supervised audio learning already existed in 2022; the paper is careful to place itself among neighbors rather than claim virgin territory. The family splits along two axes — what signal you feed in, and what pretext objective you use:

Every method above Audio-MAE shares a property that sounds innocuous and is in fact the whole cost problem: they encode a full view of the input, masked positions included. If a position is masked, it is replaced by a mask token or zeroed — but it still occupies a slot in the sequence, and self-attention is quadratic in sequence length.

Do the arithmetic once, now, because it motivates everything: a ten-second AudioSet clip becomes 512 patches. Full-view attention costs on the order of 512 × 512 = 262,144 query-key pairs per head per layer. Encode only the visible 20% — 102 patches — and the cost drops to 102 × 102 = 10,404 pairs. That is a 25.2× reduction, per head, per layer, for the entire pre-training run. Chapter 3 turns this into wall-clock hours.

What one AudioSet training example actually is

Abstractions about "weak labels" get slippery, so open one example. A clip of a street corner might arrive tagged like this:

The supervised loss on that target is binary cross-entropy over 527 independent sigmoids — not a softmax, because the classes are not mutually exclusive. Every one of the 524 zeros is a positive assertion that the sound is absent, and in a dataset this large a meaningful number of them are wrong. Weak, multi-label, and noisy: three separate ways for a supervised signal to be thin.

Now notice what self-supervision does to that ledger. The reconstruction target for the same clip is not three bits — it is 131,072 numbers (1024 frames × 128 mel bands), every one of them exactly correct by construction, with no annotator in the loop. That is the ratio the paper is exploiting: about four orders of magnitude more supervision per clip, for free.

Where this ends up

So you know what we are building toward, here is the destination. Pre-trained on AudioSet-2M with no labels at all, then fine-tuned, one ViT-B encoder produces:

Read that table as a scorecard on the paper’s three-part bet from the callout above:

• Bet (a) — you can throw away 80% and still learn. Settled by the AS-20K column: 37.0 from a model that never saw more than a fifth of any spectrogram at a time.
• Bet (b) — it becomes cheap enough for all of AudioSet. Settled by 36 hours on 64 GPUs for a 2M-clip pass, quantified in Chapter 3 and Chapter 7.
• Bet (c) — audio-only beats borrowed weights. Settled by every row, and most brutally by speaker identification.

Look at that last row for a moment. Speaker identification with an ImageNet-initialized AST lands at 41.1% accuracy; the same-size encoder pre-trained on unlabeled audio reaches 94.8%. Whatever ImageNet teaches a ViT, it is not "how to tell voices apart."

The interactive below is the supply-chain diagram for this entire subfield. Each model on the right consumes data from one or more lanes on the left. Click through them and watch which lanes each system needs — then notice that the top scorer needs exactly one.

Three things the diagram should make concrete. First, the models with the best pre-Audio-MAE numbers all draw from ImageNet — a lane with no sound in it. Second, SS-AST and MAE-AST, the closest self-supervised relatives, both add 1,000 hours of LibriSpeech speech on top of AudioSet; Audio-MAE does not. Third, the winner consumes the fewest lanes.

One honest caveat before we move on, because the paper offers it about itself. Audio-MAE processes everything at a 16 kHz sampling rate for consistency between pre-training and fine-tuning. HTS-AT and PaSST use 32 kHz. Referencing the PANNs study, the authors estimate up to 0.4 mAP of headroom that Audio-MAE leaves on the table purely from the lower sampling rate. The 47.3 is achieved with one hand tied.

Vocabulary, fixed now so it never wobbles

Self-supervised learning has a small vocabulary that papers use loosely. This lesson uses it strictly:

The distinction in the last row is not pedantry. Audio-MAE has a perfectly ordinary supervised regression loss; what makes it self-supervised is only that nobody had to write down the target. That is why the supervision budget calculation above is fair: 131,072 real numbers per clip, generated by the act of hiding them.

What we will build, chapter by chapter

One reading tip for the chapters ahead. Whenever a number appears, ask what it is being compared against — this paper’s findings are almost entirely differences, and a difference without its baseline is decoration. 47.3 means nothing; 47.3 against 47.1 with no ImageNet means everything.

Chapter 1: Why You Can Throw Away 80%

Here is a number that ought to bother you. BERT masks 15% of its tokens. MAE, the image model this paper extends, masks 75%. Audio-MAE masks 80%.

Same family of idea — hide part of the input, predict it back — and a five-fold difference in how much you can hide. Either the language people are being timid or something is genuinely different about pixels and spectrograms. This chapter shows it is the latter, and that the difference is measurable.

The lineage, in one pass

Masked and denoising autoencoders are old. The idea is embarrassingly simple: corrupt the input, ask the network to restore it, and keep whatever internal representation it needed in order to succeed. The representation is the product; the reconstruction is scaffolding you throw away.

Two rows in that table do something the others do not, and both belong to MAE and Audio-MAE: the masked patches are not replaced by a mask token in the encoder. They are deleted. The encoder’s input sequence physically shrinks. Everything that makes Audio-MAE cheap follows from that one decision, and Chapter 3 traces the consequences through the tensor shapes.

Redundancy is why the ratio can climb

Why can an image survive losing three quarters of itself while a sentence cannot survive losing a fifth?

Because of how much each unit carries. A word token is a discrete, high-entropy symbol. Delete "Paris" from "The capital of France is Paris" and no amount of staring at the neighboring words reconstructs the exact string; you must know the fact. Delete a 16×16 patch from the middle of a photograph of grass and you can very nearly paint it back from the patches around it.

Spectrograms are, if anything, more redundant than photographs, and the reasons are physical:

The widget below turns this from a claim into a measurement. Three one-dimensional signals stand in for three modalities: a token stream (each symbol independent of its neighbors), a smooth image scanline, and a mel-band energy track pulled from the synthetic spectrogram used throughout this lesson. All three get masked at the same ratio and reconstructed the dumbest possible way — linear interpolation between the surviving neighbors. The error curves are computed live.

Read the curves, not the strips. The token stream’s error is essentially flat and high from the very first masked symbol: interpolating between unrelated symbols never works, so hiding more of them barely makes it worse — it was already hopeless. The smooth signals start near zero and climb slowly, then sharply once the surviving samples get far enough apart that interpolation crosses real structure.

How many patches actually survive?

Time for concrete counting, because these two numbers recur in every later chapter. A ten-second AudioSet clip becomes a 64×8 grid of patches (Chapter 2 derives that grid), so:

Those are the paper’s own figures — it states that each 10-second sample has 64×8 = 512 patches, with 154 visible at 70% masking and 102 at 80%. Note the rounding goes down in one case and up in the other; the exact tie-breaking does not matter, but the magnitude does. One hundred and two patches. That is the entire input the encoder ever sees during pre-training.

The paper frames masking as a Bernoulli process: each patch is dropped independently with probability p. In practice implementations draw a fixed number of survivors without replacement, which is why the counts above are exact rather than expectations. Both descriptions give the same mean; the fixed-count version just removes the variance.

A worked aside: the model never sees the same mask twice

Why does a random mask act as such a strong regularizer? Count the masks. Choosing which 102 of 512 patches survive is a binomial coefficient, and we can size it with the entropy formula rather than a factorial calculator. For choosing a fraction q of n items,

Step by step with q = 102/512 = 0.1992:

1. −q log₂ q = −0.1992 × log₂(0.1992) = −0.1992 × (−2.328) = 0.4638
2. −(1−q) log₂(1−q) = −0.8008 × log₂(0.8008) = −0.8008 × (−0.3205) = 0.2567
3. H(q) = 0.4638 + 0.2567 = 0.7205 bits per patch
4. log₂ C ≈ 512 × 0.7205 = 368.9 bits, minus a small correction of about 4.5 bits from the Stirling factor → ≈ 364 bits
5. 364 bits / log₂(10) = 364 / 3.3219 = 109.6 → roughly 10¹¹⁰ distinct masks

Now compare that against the training budget. Pre-training runs 32 epochs over about 1.96M clips, so the model draws roughly 32 × 1.96 × 10⁶ ≈ 6.3 × 10⁷ masks over the whole run. Against 10¹¹⁰ possibilities, the model sees a vanishing sliver: it can never memorize "which patches tend to be visible," only "how spectrogram content coheres."

A second axis: what you predict, not just how much you hide

The masking ratio is only one of two dials. The other is the prediction target, and the lineage table above quietly contains four different answers to it:

Audio-MAE takes the cheapest option on that list, and the paper’s ablations in Chapter 5 show it did not need a more expensive one. There is a pleasing detail in the MaskFeat row, too: the paper notes that HOG features are in turn related to spectrogram features — oriented-gradient histograms and time-frequency energy maps are cousins. The vision community, hunting for a good masked-prediction target, drifted toward something spectrogram-shaped.

Measure the redundancy yourself

The claim "spectrograms are redundant" is testable in six lines. Autocorrelation along each axis tells you how far you can move before a patch stops resembling itself — and it is the number you should compute before choosing a masking ratio for any new modality:

python — how much does a patch tell you about its neighbour?
import numpy as np

# spec: (n_frames, n_mels) log-mel, e.g. (1024, 128)
z = (spec - spec.mean()) / spec.std()

def corr_at(dt, df):                     # correlation at an offset
    a = z[:spec.shape[0] - dt, :spec.shape[1] - df]
    b = z[dt:, df:]
    return float((a * b).mean())

# typical AudioSet values: strong locally, decaying fast
print(corr_at(1, 0))    # adjacent FRAME  (10 ms apart, 60% shared samples) -> high
print(corr_at(0, 1))    # adjacent MEL BAND                                 -> high
print(corr_at(16, 0))   # one PATCH away in time (160 ms)                   -> much lower
print(corr_at(0, 16))   # one PATCH away in frequency                       -> lower still

# the design rule: patch size should sit near where correlation falls off.
# too small -> neighbouring patches are near-copies, the task is trivial.
# too large -> patches are independent, masking destroys real information.

That last comment is the quiet justification for 16 × 16. A patch spans 160 ms and one eighth of the mel axis — roughly the scale at which spectrogram content stops repeating itself. Table 1a in Chapter 2 is the empirical confirmation: patches twice as wide in either axis both lose accuracy.

Proper hardness: the Goldilocks framing

The paper offers a sentence that is easy to skim and worth memorizing: designing a pretext task with proper hardness is important for effective self-supervised learning of audio representations. Both failure directions are real, and the ablations in Chapter 8 show both:

That third row is a specifically audio finding. In MAE for images, the choice of masking pattern barely matters. For spectrograms it matters a lot, because structured masking removes an entire semantic unit — a whole word, a whole frequency band — rather than a scattering of interpolable holes. Chapters 4 and 8 make this tangible; you will hear yourself say "of course" once you have watched a time-masked spectrogram fail to come back.

And there is a beautiful inversion coming in Chapter 7: the masking strategy that is worst for pre-training turns out to be best for fine-tuning. Task-agnostic pre-training wants unstructured masking at a high ratio; task-specific fine-tuning wants structured masking at a low ratio. Hold that tension; it resolves cleanly.

Chapter 2: From Pressure Wave to Patch Grid

Audio-MAE never touches audio. By the time anything trainable is involved, a ten-second recording has become a sequence of 512 vectors of length 768, and every decision that got it there is a modeling decision with consequences.

This chapter walks the whole path with real numbers, because you cannot debug — or re-implement — what you have only seen as a box labeled "feature extraction". By the end you will be able to compute the input shape for any clip length in your head.

Step 1 — resample to 16 kHz mono

Every recording is pre-processed to a single channel at a 16,000 Hz sampling rate. Not because 16 kHz is optimal, but for simplicity and consistency between pre-training and fine-tuning: AudioSet clips arrive with wildly heterogeneous native rates (many are 8 kHz or higher, and YouTube’s own re-encoding may have up- or down-sampled the track already).

The cost is real and the paper states it: HTS-AT and PaSST train on 32 kHz audio, and referencing the PANNs study the authors estimate up to 0.4 mAP of improvement that Audio-MAE forgoes by staying at 16 kHz. By the Nyquist limit, 16 kHz sampling represents content only up to 8 kHz — everything above that (cymbal shimmer, fricative brightness, insect calls) is simply gone before the model starts.

Step 2 — frame the waveform

A spectrum is only meaningful over a stretch of time short enough that the sound is roughly stationary. So the waveform is chopped into overlapping frames: a 25 ms Hanning window that shifts every 10 ms.

Convert to samples, which is the only unit the code sees:

That 60% overlap is not incidental — it is precisely the temporal redundancy Chapter 1 exploited. Consecutive columns of the spectrogram are computed from mostly the same audio samples, so they cannot differ much. Masking one column and predicting it from its neighbors is closer to interpolation than to invention.

Why a Hanning window rather than just cutting the waveform? Because a hard cut is a multiplication by a rectangle, and a rectangle’s spectrum is a sinc function with large side lobes — a pure tone would smear energy across the whole frequency axis (spectral leakage). The Hanning taper, which falls smoothly to zero at both edges, trades a slightly wider main lobe for far smaller side lobes. Clean harmonic stacks in the spectrogram — the very structure the model will learn to exploit — depend on this.

Now count frames for a ten-second AudioSet clip, every step shown:

1. Total samples: 10 s × 16,000 Hz = 160,000
2. The first frame consumes samples 0–399. Each additional frame advances 160 samples.
3. Number of complete frames: floor((160,000 − 400) / 160) + 1 = floor(159,600 / 160) + 1 = 997 + 1 = 998
4. The paper reports a time axis of 1024. The extra 26 rows come from padding (and the loader pads or trims every dataset to a fixed length anyway).

Step 3 — 128 mel bands

Each frame’s spectrum is warped onto 128 Kaldi-compatible mel-frequency bands. The mel scale approximates human pitch perception: roughly linear below 1 kHz, logarithmic above.

Work out the band layout by hand — it explains where the model’s resolution actually goes:

1. Top of the range: m(8000) = 2595 · log₁₀(1 + 8000/700) = 2595 · log₁₀(12.4286) = 2595 × 1.09442 = 2840.0 mel
2. 128 bands spread evenly across mel: each band is 2840.0 / 128 = 22.19 mel wide
3. Band 1 spans mel 0 → 22.19. In Hz: f(22.19) = 700 · (10^0.008551 − 1) = 700 × 0.01989 = 13.9 Hz wide
4. The midpoint band 64 sits near mel 1420: f(1420) = 700 · (10^0.54721 − 1) = 700 × 2.5253 = 1768 Hz

Read that last line again. Half of all 128 bands are spent below 1.8 kHz — where pitch, the first two or three formants, and most speech energy live. The upper half of the audible-to-this-model range, from 1.8 kHz to 8 kHz, gets the other 64 bands. The mel warp is a resolution budget, and it is spent on the part of the spectrum humans care about.

Step 4 — the resulting tensor, per dataset

Every dataset gets padded or trimmed to a fixed duration, giving a fixed input shape:

Then each dataset is normalized with its own mean and standard deviation, estimated on that dataset’s training split. The paper is emphatic that proper normalization matters to avoid a pre-training / fine-tuning discrepancy:

Those means are negative because the features are log-energies. And notice the spread: AudioSet sits about 2.4 units louder on average than the speech corpora, and VoxCeleb’s standard deviation (3.074) is barely half of Speech Commands’ (5.565). Feed a model pre-trained on AudioSet statistics an un-renormalized VoxCeleb spectrogram and every activation in layer 1 is off-scale. This one line of preprocessing is the difference between transfer working and transfer mysteriously not working.

The whole front end is four lines of torchaudio, which is worth seeing because every constant in it has now been justified:

python — the exact front end
import torchaudio, torch

wav, sr = torchaudio.load(path)                 # any rate, any channels
wav = wav.mean(0, keepdim=True)                  # -> mono
wav = torchaudio.functional.resample(wav, sr, 16000)

fbank = torchaudio.compliance.kaldi.fbank(
    wav, htk_compat=True, sample_frequency=16000,
    window_type='hanning', num_mel_bins=128,
    frame_length=25, frame_shift=10)               # -> (998, 128) for 10 s

fbank = torch.nn.functional.pad(fbank, (0, 0, 0, 1024 - fbank.shape[0]))
fbank = (fbank - (-4.268)) / (4.569 * 2)          # dataset mean / std

And here is the debugging table — what each step’s failure actually looks like downstream, since these bugs are silent rather than loud:

Step 5 — patchify without overlap

The 1024 × 128 plane is cut into non-overlapping 16 × 16 tiles by a convolution whose kernel size and stride are both (16, 16). The general count, which you should keep:

What is one patch, physically? Sixteen frames × sixteen mel bands:

• In time: 16 frames × 10 ms hop = 160 ms — about the length of a syllable nucleus, or two or three cycles of a slow tremolo.
• In frequency, bottom row: mel bands 0–15 = mel 0 to 355.0 → f(355.0) = 700 × (10^0.13680 − 1) = 0 to 259 Hz. One patch covers the fundamental of most adult voices.
• In frequency, top row: mel bands 112–127 = mel 2485 to 2840 → 5,649 Hz to 8,000 Hz. The same 16 numbers now cover 2.35 kHz of bandwidth.

The eight patch rows therefore have wildly different physical widths — 259 Hz at the bottom, 2,351 Hz at the top — while being identical to the network. That asymmetry is a big part of why frequency-axis position matters so much, and it is the seed of the argument in Chapter 6.

Here is Table 1a in full, and it is a rare thing in ML papers: an ablation where the cheap option wins outright.

Step 6 — flatten and project

Each 16 × 16 tile is flattened to a vector of 256 numbers and linearly projected to the transformer width. For ViT-B that width is 768:

python — patchify, three equivalent ways
# the spectrogram: (batch, 1, 1024, 128)
x = spec                                    # B x 1 x 1024 x 128

# (a) explicit: unfold into tiles, flatten, matmul
tiles = x.unfold(2, 16, 16).unfold(3, 16, 16)   # B x 1 x 64 x 8 x 16 x 16
flat  = tiles.reshape(B, 64 * 8, 16 * 16)          # B x 512 x 256
tok   = flat @ Wp.T + bp                    # Wp: 768 x 256  ->  B x 512 x 768

# (b) the one-liner every implementation actually uses
proj  = nn.Conv2d(1, 768, kernel_size=16, stride=16)
tok   = proj(x).flatten(2).transpose(1, 2)         # B x 512 x 768

# (c) parameter count is identical either way:
#     768 * 256 weights + 768 biases = 196,608 + 768 = 197,376

A convolution with kernel = stride is an independent linear map applied to each tile. There is no cleverness in (b); it is (a) with the loop pushed into cuDNN.

Finally, fixed sinusoidal positional embeddings are added — the two-dimensional variant, where one half of the 768 dimensions encodes the time index and the other half encodes the frequency index. Fixed, not learned, and that choice earns its keep in Chapter 3: the decoder has to re-insert mask tokens at their original coordinates, and fine-tuning uses a completely different masking pattern from pre-training. A closed-form positional code stays valid under every one of those rearrangements.

Why fixed sinusoidal rather than learned? The two-dimensional variant assigns half the 768 channels to encode the time index and half the frequency index, each as a bank of sines and cosines at geometrically spaced wavelengths:

Three properties earn it its place, and every one of them is used later in this lesson. It is defined for any index, so a 64-column grid and an 8-column grid need no interpolation — Chapter 7 fine-tunes on 1-second clips with no surgery. It is deterministic, so the decoder in Chapter 6 can re-insert mask tokens at their true coordinates without having learned anything about them. And the difference between two positions depends only on their offset, which is what makes "local" a meaningful notion for the windowed attention that follows.

Notice while you explore that the frequency axis labels are not evenly spaced in Hz — that is the mel warp made visible. The bottom patch row covers a couple of hundred hertz; the top covers a couple of thousand. Chapter 6 will argue that this asymmetry is one reason a spectrogram decoder should treat the frequency axis as physically meaningful rather than as "the other spatial dimension".

Click around the harmonic stacks in the first half and then the broadband burst in the middle. A patch on a vowel is a set of near-horizontal stripes; a patch on a fricative is a wash. Those are the two textures the decoder in Chapter 6 will have to reproduce, and they fail in different ways.

Chapter 3: The Encoder That Sees One Fifth

We have 512 tokens, each a 768-dimensional vector with its position baked in. Now comes the move that gives the whole architecture its shape: throw 410 of them away before the transformer ever runs.

Not zero them. Not replace them with a learned [MASK] vector. Delete them from the sequence, so that the tensor entering block 1 has 102 rows instead of 512.

The full forward pass, with shapes

Here is the entire pre-training encoder path for one 10-second AudioSet clip at the default 80% masking ratio. Every shape is exact.

Row 5 is the paper. Everything else is standard ViT.

How the gather actually works

The masking implementation is a small idiom worth knowing, because it makes "sample a random subset and remember how to undo it" a three-line operation with no Python loops:

python — random masking, MAE style
def random_masking(x, mask_ratio):
    B, L, D = x.shape                       # B x 512 x 768
    len_keep = int(L * (1 - mask_ratio))       # 512 * 0.2 -> 102

    noise = torch.rand(B, L, device=x.device)  # one uniform number per patch
    ids_shuffle = torch.argsort(noise, dim=1)  # random permutation of 0..511
    ids_restore = torch.argsort(ids_shuffle, dim=1)   # its inverse

    ids_keep = ids_shuffle[:, :len_keep]                  # the 102 survivors
    x_kept = torch.gather(x, 1, ids_keep[..., None].expand(-1, -1, D))

    mask = torch.ones(B, L, device=x.device)      # 1 = masked, for the loss
    mask[:, :len_keep] = 0
    mask = torch.gather(mask, 1, ids_restore)     # put it back in patch order
    return x_kept, mask, ids_restore          # (B,102,768), (B,512), (B,512)

Three things to notice. argsort of uniform noise gives a uniform random permutation — no sampling library needed. ids_restore is the inverse permutation, and the decoder in Chapter 6 cannot function without it. And the mask tensor is built in shuffled order and then un-shuffled, so downstream the loss knows exactly which of the 512 original positions to score.

One block, in full

Since the encoder is exactly a vanilla ViT, here is the complete block, annotated with shapes for the masked case. Nothing in it knows anything about audio — that is the point:

python — the ViT-B block Audio-MAE uses unmodified
class Block(nn.Module):
    def __init__(self, d=768, heads=12, mlp_ratio=4):
        self.norm1 = nn.LayerNorm(d)
        self.attn  = MultiHeadAttention(d, heads)      # head dim = 768/12 = 64
        self.norm2 = nn.LayerNorm(d)
        self.mlp   = nn.Sequential(nn.Linear(d, mlp_ratio * d),   # 768 -> 3072
                                   nn.GELU(),
                                   nn.Linear(mlp_ratio * d, d))   # 3072 -> 768

    def forward(self, x):                 # x: (B, 102, 768) during pre-training
        x = x + self.attn(self.norm1(x))    # pre-norm residual; 102 x 102 attention
        x = x + self.mlp(self.norm2(x))     # (B, 102, 768) throughout
        return x

# the encoder is twelve of these, then a final LayerNorm. That is all.
# Note what is ABSENT: no class token, no pooling, no relative position bias,
# no audio-specific anything. Every domain prior lives in the masking and
# in the decoder — both of which are discarded before deployment.

Pre-norm placement (LayerNorm inside the residual branch) matters at this scale: it is what lets the 12-layer stack train stably at batch 512 with no gradient clipping, which the hyperparameter table confirms is switched off.

Worked example: the cost model, and reproducing the paper’s FLOPs column

People repeat "attention is quadratic, so masking gives a huge speedup." Let us actually check, because at ViT-B scale the intuition is mostly wrong, and the truth is more interesting.

One transformer block on a sequence of n tokens with width d costs, in multiply-accumulates:

Now plug in ViT-B numbers (d = 768, so d² = 589,824) for the two cases, every step written out:

So the honest speedup is 4.027 / 0.738 = 5.46×, not 25×. The quadratic term shrinks 25.2× (512² / 102²), exactly as advertised — but at 512 tokens it was only a tenth of the work in the first place. Most of the saving comes from the boring linear term: five times fewer tokens, five times fewer matrix rows.

And the quadratic term does dominate eventually. It scales as n²d against 12nd², so the crossover is at n = 6d = 4,608 tokens — a 90-second clip at this patch size. For AudioSet’s 10-second clips we are far below it; for the "modeling lengthy audio" the paper names as future work, we are not. That is the sense in which Audio-MAE’s efficiency argument is a statement about 2022 clip lengths, not a law.

And a short list of what masking does not buy, so the accounting stays honest:

So the honest headline is: masking makes the largest single component of the pre-training step about five times cheaper, while leaving the rest of the pipeline alone. That is enough to turn a week into 36 hours, which is the difference between running the experiment and not.

Step through the stages once in order, too. The hand-offs are where implementations break: between stage 3 and 4 the position must be added before masking (otherwise the survivors lose their coordinates), and between 4 and 5 the sequence length silently changes from 512 to 102, which every downstream shape assertion has to know about. A surprising number of MAE reimplementation bugs are one of those two.

Push the slider to 0 and the compute bar tells you what a full-view method like SS-AST pays on every step of pre-training. Push it to 0.8 and you get Audio-MAE. Over 32 epochs of a 2M-clip dataset, that ratio is the difference between a 36-hour run and a week.

It is worth naming the three quantities that people conflate when they say "cost", because masking affects them very differently:

Where the 86M parameters actually are

"ViT-B has 86M parameters" is usually quoted and never derived. Derive it, because the breakdown tells you what masking does and does not save.

One transformer block at width d = 768 contains:

1. Attention projections Q, K, V and the output projection: 4 × d² = 4 × 589,824 = 2,359,296
2. MLP up-projection (768 → 3072) and down-projection (3072 → 768): 2 × 4d² = 8 × 589,824 = 4,718,592
3. Block total (ignoring biases and the two LayerNorms, together a few thousand): 7,077,888 ≈ 7.08M
4. Twelve blocks: 12 × 7,077,888 = 84,934,656 ≈ 84.9M
5. Plus the patch embedding, 197,376 — and the position embeddings, which are zero parameters because they are fixed sinusoids
6. Total ≈ 85.1M, which is the 86M quoted once biases and norms are counted

Two thirds of the model is the MLP, one third is attention, and none of it is the patch embedding. Notice what this means for masking: the parameter count is completely independent of sequence length. Masking changes how many rows flow through those matrices, not how big they are. Memory for weights and optimizer state is unchanged; only activation memory and FLOPs shrink.

Activation memory is worth one line of arithmetic too, since it is what actually fills a V100. A batch of 512 clips at 512 tokens of width 768 in fp32 is 512 × 512 × 768 × 4 bytes = 805 MB for a single stored tensor, and a transformer block stores several. At 102 tokens the same tensor is 160 MB. That is why batch 512 across 64 GPUs is feasible at all.

That last point is the practical one. If a reproduction runs out of memory, the fix is a smaller batch or a higher masking ratio — never a narrower model, which would change the artifact you are trying to build.

And if a reproduction is unexpectedly slow, look at the data loader before the model: at 2M clips an epoch, mel extraction on the CPU is a common bottleneck that no amount of masking will fix.

How big should the encoder be?

Table 1b sweeps the backbone. All three are fine-tuned identically; only the encoder changes:

Two readings, and the second is the useful one.

First, the obvious: bigger is better, but with brutally diminishing returns. Going from 86M to 304M parameters — 3.5× the model — buys 0.5 mAP on AS-20K and 0.1 on AS-2M. That is why ViT-B is the default in every other table.

Second, the subtle one, which the paper spells out: the ViT-S to ViT-B gap is 5.0 mAP on AS-20K but only 2.3 mAP on AS-2M. More in-domain fine-tuning data closes more than half the capacity gap. Model size and fine-tuning data are partial substitutes — if you have a lot of labeled in-domain audio, a small encoder gets most of the way; if you have 20K labeled clips, capacity matters far more.

One more consequence of the encoder-sees-only-visible design, easy to miss: because masked patches never enter the encoder, the encoder has no way to know how many patches were masked or where. The positional embeddings it receives are the ones belonging to the surviving patches, so it knows where the survivors are — a scattered, irregular subset of a 64 × 8 grid. Every representation it builds must therefore be robust to arbitrary, jagged views of the time-frequency plane. That robustness is a large part of why the encoder fine-tunes so well under a completely different masking scheme in Chapter 7.

Chapter 4: Mask and Reconstruct — the Lab

This is the paper’s Figure 1 and Figure 2, made playable. Everything up to here was setup; from here on you can turn the knobs the authors turned and watch what they watched.

Before you touch anything, know exactly what the four masking strategies do, because the differences between them are mechanical and countable, and the sim will not teach you that part.

The four strategies, mechanically

Two consequences of that table deserve their own arithmetic, because both surprise people.

Consequence 1 — frequency masking is coarse

There are only 8 frequency rows in the 64 × 8 grid. So a frequency-masking ratio can only take values that are multiples of 1/8 = 0.125. Ask for 0.3 and you get either 2 rows (0.25) or 3 rows (0.375); there is no 0.3. Time masking has 64 columns, so it is granular to about 1.6%. Whenever you see a frequency-masking ablation curve looking lumpy, this is why.

Consequence 2 — "0.3 time + 0.3 frequency" removes about half the patches

This is the fine-tuning default, and reading it as "30% masking" is wrong by a factor of nearly two. The two maskings compose multiplicatively, because a patch survives only if both its row and its column survive:

On the actual integer grid, with every step:

1. Time columns masked: 0.3 × 64 = 19.2 → 19 masked, 45 survive
2. Frequency rows masked: 0.3 × 8 = 2.4 → 2 masked, 6 survive
3. Surviving patches: 45 × 6 = 270 of 512
4. Effective masking ratio: 1 − 270/512 = 1 − 0.527 = 0.473

So fine-tuning still discards roughly half the input — which is why the paper can honestly say masking during fine-tuning "as a side effect, also reduces computation." Half the tokens is a big saving on a 100-epoch fine-tune.

The lab

Three panels: the original spectrogram, the masked input the encoder would actually receive, and a reconstruction. The reconstruction here is a deliberately simple stand-in for the decoder — each masked patch is predicted as a locality-weighted average of the nearest visible patches, which is the crude version of exactly what Chapter 6’s local-attention decoder learns to do properly. The error number under the panels is the patch-normalized MSE over masked patches only: the paper’s actual objective, computed live.

The middle panel is the one to stare at. That is not a visualization aid — it is a literal picture of the encoder’s input during pre-training, minus the fact that the encoder does not even see the grey cells as cells. To the encoder they do not exist; it receives a bag of surviving patches with their coordinates attached, and must build a representation of the whole recording from that.

Work through this sequence deliberately; each step is a finding from the paper you can now verify with your own eyes:

1. Unstructured, 0.30. The reconstruction is nearly perfect and the error is tiny. This is the "too easy" regime — nothing forces the model to understand anything; averaging neighbors suffices.
2. Unstructured, 0.80. The masked input looks like confetti, and yet the harmonic stacks and the fricative burst still come back recognizably. This is the paper’s default, and it is where "proper hardness" lives.
3. Time, 0.80. Long stretches of the clip are simply gone. Nothing in the surviving context tells you what happened during a missing 1.5-second span — and nothing can. The error jumps.
4. Frequency, 0.80. Whole registers vanish for the entire duration. Harmonics can sometimes be extrapolated upward from surviving lower bands — that is exactly the argument the paper makes for local attention — but with six of eight rows gone, there is not enough left to extrapolate from.
5. Time + frequency, 0.50. The lattice pattern. This is the fine-tuning regime, and notice it looks structured rather than random — the encoder sees a regular grid of survivors.

Two caveats about the lab, stated plainly so you calibrate what you are looking at. The reconstruction here is a fixed locality-weighted average, not a trained network — it has no notion of harmonics, so it will always underperform the real decoder on voiced segments, and it cannot hallucinate plausible-but-wrong content the way a trained model does. And the error number is computed on texture only, comparing each patch after normalizing both prediction and target, which is the closest honest analogue of the paper’s patch-normalized objective for a predictor that outputs raw values. Read the numbers as a relative index across settings, not as an mAP.

A last note on the numbers under the panels: the visible count is what the encoder would actually receive, and the effective-mask figure is what you got rather than what you asked for — the two diverge for structured strategies, because rows and columns come in whole units.

Each strategy is a model of a real-world failure

There is a reason the four strategies are the four strategies, and it is not combinatorial tidiness. Each corresponds to a way audio actually goes missing in the world, which is what makes them sensible fine-tuning augmentations in Chapter 7:

Read the first row again. Unstructured masking is the least realistic corruption and the best pre-training pretext. That tension is the chapter in one line: during pre-training you are not simulating deployment, you are designing a puzzle. During fine-tuning you are simulating deployment. Different jobs, different masks.

It also suggests how the authors probably arrived at the inversion. Start from MAE, which only ever needed random masking. Add structured variants because spectrograms have named axes and it seems obvious they should be treated differently. Discover the structured variants pre-train worse, which is a disappointing negative result. Then notice the structured variants look exactly like SpecAugment, try them at fine-tuning time where SpecAugment lives — and find that they win there. The published finding reads as one clean principle; it was almost certainly two experiments that disagreed.

One more thing to try in the lab, because it is the cleanest demonstration of "answerable". Set frequency masking to 0.50 — four of eight rows gone — and look at the reconstruction of the voiced segments. Some of the missing harmonic bands come back, because the surviving rows still contain enough of the harmonic ladder to place the missing rungs. Now push it to 0.75 and the same reconstruction collapses into a smear. The information did not degrade gradually; it crossed a threshold where the ladder stopped being inferable at all.

What the paper’s own reconstructions look like

The authors did not stop at pictures. They inverted the reconstructed mel-spectrograms back to waveforms using the Griffin-Lim algorithm and published audible examples — ground truth, masked input, and reconstruction — for four content types: speech, music, event sounds, and "others". Griffin-Lim estimates the missing phase iteratively, so the paper is careful to note there are perceivable artifacts from imperfect phase estimation that belong to the inversion, not to the model.

Their qualitative findings line up exactly with what the lab shows:

The last row is the most instructive thing in that whole figure. When the reconstruction fails it does not fail by producing noise — it fails by producing something confidently wrong. That is the signature of a conditional-mean predictor being asked an unanswerable question, and it is why the paper insists that absolute positions and arrangement of spectrogram components are critical for humans to understand sound: shifting a pitch makes an audio sound completely different, and phoneme sequences in time are important cues for speech understanding.

And a closing thought on what "reconstruction" is really for here. Nobody wants an Audio-MAE that inpaints spectrograms — well, almost nobody, packet-loss concealment excepted. The reconstruction exists so that gradients exist. Every masked patch is a question whose answer is known for free, and the only way to answer 410 of them from 102 clues is to have understood, somewhere in those twelve blocks, what sound is present. The picture in the third panel is a receipt, not a product.

Chapter 5: The Objective, Computed by Hand

"Reconstruct the spectrogram" is not a loss function. This chapter turns it into one, then computes it end to end on numbers small enough to check on paper.

The loss, stated

Audio-MAE minimizes the mean squared error between the prediction and the input spectrogram, averaged over the unknown patches. The paper writes it as:

Every symbol, defined before use:

Two design decisions hide in that one line, and both matter more than the formula.

Decision 1 — the loss is computed only on masked patches

Visible patches contribute nothing. If they did, the model could drive the loss down by learning an identity map on the 20% it can see, and the gradient from the part that requires actual understanding would be swamped. Restricting the average to unknown patches means every unit of loss corresponds to a genuine prediction.

There is a second, less obvious consequence of restricting the average to masked patches: the loss per example depends on the masking ratio, because N changes. At 80% masking N = 410; at 30% masking N = 154. Since the average divides by N, the gradient magnitude stays comparable across ratios — which is what makes the Chapter 8 masking-ratio sweep a fair comparison rather than a disguised learning-rate sweep. Had the paper summed instead of averaged, every point on that curve would also have been changing the effective step size.

Decision 2 — per-patch normalization

The decoder can be trained to predict the raw patch values, or their per-patch normalized versions: subtract that patch’s own mean, divide by its own standard deviation. Audio-MAE, following MAE, minimizes the patch-normalized mean square error by default.

Why? Because normalization deletes the two easiest degrees of freedom from the target. A patch’s mean is its loudness, and a patch’s standard deviation is its contrast — both are strongly predictable from neighbors without understanding anything about sound. Remove them and what is left is pure texture: is this a harmonic stack, a noise wash, an onset edge, silence with a faint tail? That is the part worth learning.

The sim below makes the argument concrete by pitting three predictors against each other.

Switch between the two patches with the buttons and watch which predictor changes rank. On the high-contrast patch the ordering is the same under both metrics; on the flat one it inverts. That inversion is exactly the failure mode normalization exists to prevent, and in a real spectrogram a large fraction of patches are the flat kind — silence, noise floor, decay tails.

The "predict the patch mean" column is the whole point. On raw MSE it scores exactly the patch variance — so on a flat, loud patch it looks nearly perfect while having learned nothing. On patch-normalized MSE it scores exactly 1.0 on every patch, because the normalized target has zero mean and unit variance, and a constant prediction of zero gets all of it wrong. Normalization converts a shortcut into a visible failure.

Worked example: the loss, every intermediate step

Real patches have 256 values, which is unreadable on a page. So take two masked patches of four values each. Everything below generalizes without a single change of formula.

Patch A — textured, moderate level: x_A = [2, 4, 6, 8]

1. Mean: μ = (2 + 4 + 6 + 8) / 4 = 20 / 4 = 5
2. Deviations: 2−5 = −3, 4−5 = −1, 6−5 = 1, 8−5 = 3
3. Squared deviations: 9, 1, 1, 9 → sum = 20
4. Variance: 20 / 4 = 5. Standard deviation: σ = sqrt(5 + ε) with ε = 10⁻⁶ → 2.23607
5. Normalized target t_A = deviations / σ = [−1.34164, −0.44721, 0.44721, 1.34164]
6. Suppose the decoder predicts p_A = [−1.10, −0.30, 0.60, 1.20] (already in normalized space — that is what the head outputs)
7. Differences p − t: 0.24164, 0.14721, 0.15279, −0.14164
8. Squares: 0.058390, 0.021671, 0.023345, 0.020062 → sum = 0.123468
9. Patch A loss: 0.123468 / 4 = 0.030867

Patch B — flat and loud: x_B = [10, 10, 10, 10]

1. Mean: μ = 10. Deviations: 0, 0, 0, 0. Variance: 0
2. σ = sqrt(0 + 10⁻⁶) = 0.001. This is what the epsilon is for — without it we divide by zero on silence
3. Normalized target t_B = [0, 0, 0, 0]
4. Predictions p_B = [0.05, −0.02, 0.01, 0.00]
5. Squares: 0.0025, 0.0004, 0.0001, 0 → sum = 0.0030
6. Patch B loss: 0.0030 / 4 = 0.00075

Total: average over the two masked patches:

Now a cross-check that proves you understood the normalization, not just followed it. Convert patch A’s normalized predictions back into raw log-mel units with x̂ = μ + σp:

1. 5 + 2.23607 × (−1.10) = 5 − 2.45968 = 2.54032  (target was 2)
2. 5 + 2.23607 × (−0.30) = 5 − 0.67082 = 4.32918  (target was 4)
3. 5 + 2.23607 × 0.60 = 5 + 1.34164 = 6.34164  (target was 6)
4. 5 + 2.23607 × 1.20 = 5 + 2.68328 = 7.68328  (target was 8)
5. Raw squared errors: 0.291949, 0.108359, 0.116719, 0.100309 → sum 0.617336 → raw MSE = 0.154334
6. And 0.154334 / σ² = 0.154334 / 5 = 0.030867 — exactly the normalized loss from step 9 above

So patch normalization is precisely a per-patch reweighting of the raw MSE by 1/σ². High-contrast patches are down-weighted; low-contrast patches are up-weighted. The model is told, in effect: a 1 dB error on a quiet, flat patch is as serious as a 3 dB error on a loud, busy one.

A final sanity property worth internalizing: the loss is invariant to the overall loudness of the recording and to its dynamic range. Multiply every sample by two, and every patch’s mean and standard deviation scale together, so the normalized targets are unchanged and the loss is unchanged. That is exactly right for a representation learner — a siren recorded quietly and a siren recorded loudly should produce the same features — and it is also why the ±6 dB magnitude jitter used during pre-training is a harmless augmentation rather than a contradiction of the objective.

The same computation in code, three ways

python — form 1: literal, one patch at a time
import numpy as np

x    = np.array([2., 4., 6., 8.])
mu   = x.mean()                       # 5.0
var  = ((x - mu) ** 2).mean()          # 5.0
sd   = np.sqrt(var + 1e-6)             # 2.2360682
tgt  = (x - mu) / sd                    # [-1.34164 -0.44721  0.44721  1.34164]
pred = np.array([-1.10, -0.30, 0.60, 1.20])
loss_A = ((pred - tgt) ** 2).mean()    # 0.0308668

python — form 2: both patches at once, still explicit
x = np.array([[2., 4., 6., 8.],
              [10., 10., 10., 10.]])           # (n_masked_patches, patch_dim)
mu   = x.mean(-1, keepdims=True)              # [[5.] [10.]]
var  = x.var(-1, keepdims=True)               # [[5.] [ 0.]]
tgt  = (x - mu) / np.sqrt(var + 1e-6)      # row 2 is all zeros
pred = np.array([[-1.10, -0.30, 0.60, 1.20],
                 [ 0.05, -0.02, 0.01, 0.00]])
per_patch = ((pred - tgt) ** 2).mean(-1)   # [0.0308668  0.00075]
loss = per_patch.mean()                     # 0.0158084

python — form 3: the real thing, masked-only average over all 512 patches
# pred: (B, 512, 256)   target: (B, 512, 256)   mask: (B, 512), 1 = masked
mu  = target.mean(dim=-1, keepdim=True)
var = target.var(dim=-1, keepdim=True)
tgt = (target - mu) / (var + 1e-6) ** .5

loss = ((pred - tgt) ** 2).mean(dim=-1)          # (B, 512) per-patch MSE
loss = (loss * mask).sum() / mask.sum()          # masked patches only

Form 3 is the whole objective. Note that it computes the loss for all 512 positions and then throws away the visible ones by multiplication — wasteful on paper, and completely free on a GPU, which is why every implementation does it this way rather than gathering.

Notice too that both decisions are about what to score, not about how to score it. The squared error itself is the least interesting part of the objective; the choices that carry the weight are which patches enter the average and what transform is applied to the target first.

Why squared error, and not something cleverer

MSE has a well-known pathology as a generative objective: when several outputs are plausible, the loss-minimizing prediction is their average, which is usually blurry and sometimes physically impossible. Image generation abandoned plain MSE for exactly this reason. Why is it fine here?

Because reconstruction quality is not the deliverable. The blurring is a symptom of the model honestly reporting its uncertainty, and the encoder representation that lets it compute that conditional mean is precisely what we keep. You can see the pathology in the paper’s own figures — the time-masked speech reconstruction produces plausible speech-like texture with the wrong words — and it does not damage a single downstream number.

The alternatives all cost something the paper does not want to spend. An adversarial term would sharpen reconstructions and add training instability for no representational gain. A perceptual loss needs a pre-trained feature extractor, which is exactly the external dependency Chapter 0 spent its whole argument removing. Discrete-token prediction — the BEiT route, and the one BEATs would later take for audio — needs a separately trained tokenizer, another moving part. MSE on normalized patches has no dependencies, no hyperparameters beyond epsilon, and no failure modes that show up downstream. It is the boring choice, chosen on purpose.

The contrastive objective they tried and dropped

SS-AST, the closest prior work, uses a joint contrastive-plus-reconstructive objective, and wav2vec 2.0 is purely contrastive. So the obvious question is whether adding a contrastive term helps Audio-MAE. The authors ran it, and reported the negative result in the appendix — which is more useful than most positive results.

Their contrastive loss is InfoNCE over masked tokens:

Intuitively L_c pulls a masked patch toward its own contextualized embedding while pushing away the other masked patches of the same recording. Three setups were compared: reconstruction alone, contrastive alone, and a multi-task combination L_r + αL_c:

Reconstruction alone wins on both. The combination lands between the two pure objectives rather than above either — the paper’s conclusion is that L_c and L_r do not work complementarily in Audio-MAE.

Chapter 6: The Decoder That Listens Locally

The decoder is the part you throw away. It exists only during pre-training; the moment fine-tuning starts it is deleted and never used again.

And it is the only place in Audio-MAE where the authors added something that is not standard MAE. That is not a coincidence — it is the paper’s thesis about where audio-specific knowledge belongs. Put the domain prior in the scaffolding, keep the product generic.

The decoder’s forward pass, with shapes

Step 4 is worth sitting with. All 410 mask tokens are the same vector. Before positional embeddings are added, the decoder cannot tell the patch at (time 3, band 0) from the patch at (time 61, band 7): they are literally identical rows. The positional embedding is the entire identity of a masked position. Everything the decoder ever reconstructs is "what belongs at these coordinates, given this scattering of visible context".

Why global attention is the wrong prior for a spectrogram

Image-based MAE uses global self-attention in its decoder, and that is appropriate for visual context: visual objects are typically invariant under translation or scaling, and their exact position may not affect the semantics of an image. A cat in the corner is a cat.

Spectrograms break every clause of that sentence. The position, scale, and translation of spectrogram features directly affect the sound. Shifting a pattern up in frequency is a pitch change — a different sound, not the same sound relocated.

The paper gives two concrete examples, and both are worth converting into intuitions you can carry:

Variant 1 — shifted local windows

Attention is restricted to a window of neighboring patches. Audio-MAE uses 4 × 4 patch windows over the 64 × 8 grid. Count what that partition looks like:

A window covers 4 × 160 ms = 640 ms of time and half the mel range. That is a linguistically sensible unit: about the length of a word, and wide enough in frequency to contain several harmonics of a voice.

The problem with fixed windows is obvious: a patch at a window boundary can never attend across it, so information never crosses. Swin Transformers solve this by shifting the window grid by 50% on alternating layers, and Audio-MAE copies the solution — shifting window attention by 50% between consecutive decoder layers. For padding the margin when shifting, the spectrogram is cyclically shifted to the top-left, so no partial windows exist.

With a 4-wide window the shift is 2 patches. In frequency, where there are only 2 window rows, a 2-patch shift means the shifted partition straddles the boundary between the low-band and high-band windows — on odd layers, low and high registers finally talk to each other. Over 16 layers, alternating alignment gives every patch an effective receptive field far larger than 4 × 4 while never paying for global attention.

Concretely, in code, a windowed attention layer is a reshape sandwich — there is no special attention kernel involved:

python — window partition, attend, un-partition
# x: (B, T=64, F=8, D=512) after ids_restore and position embeddings
def window_attn(x, win=4, shift=0):
    B, T, F, D = x.shape
    if shift:
        x = torch.roll(x, shifts=(-shift, -shift), dims=(1, 2))   # cyclic, to the top-left

    x = x.view(B, T // win, win, F // win, win, D)     # (B,16,4,2,4,512)
    x = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, win * win, D)   # (B*32, 16, 512)

    x = self_attention(x)                            # 16 queries x 16 keys, per window

    x = x.view(B, T // win, F // win, win, win, D)
    x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, T, F, D)
    if shift:
        x = torch.roll(x, shifts=(shift, shift), dims=(1, 2))    # undo
    return x

# the decoder stack alternates: shift = 0, 2, 0, 2, ... over 16 layers

Trace one boundary patch through two layers to see why the shift matters, using time indices only. Patch at time 3 sits in window [0..3] on an aligned layer, so it can see patches 0 through 3 and nothing at time 4. On the shifted layer the roll moves everything left by 2, so patch 3 lands at index 1 and its window covers original times 2 through 5 — it now sees across the old boundary. Two layers, and information has crossed. Sixteen alternating layers give a receptive field far wider than any single window, at local-attention prices.

The cyclic roll also explains a subtlety in frequency. With 8 rows and a 4-wide window there are exactly two window rows, and a 2-row shift makes the shifted partition straddle them: rows 6, 7, 0, 1 become one window. Rows 0 and 7 are the lowest and highest mel bands, which are not physically adjacent — the roll makes them neighbors artificially. That is a mild wart inherited from Swin, invisible at 16 layers because the aligned layers alternate with it, but worth knowing if you ever change the window size.

Variant 2 — hybrid window attention (Hwin)

The second variant takes a different route to cross-window connection: compute local attention within a window in all but the last few top layers, then go global at the top. The paper writes the configuration as local(8) + global(4) — eight local layers followed by four global ones.

The motivation is precise: this way the input feature maps for the final reconstruction layer also contain global information. Build local structure first, integrate globally at the end, then predict. No pooling, no hierarchy — the paper explicitly keeps it simple.

Two small details that trip up re-implementations. The decoder has its own positional embeddings, separate from the encoder’s, and they are added after the un-shuffle — the encoder’s positions were consumed by the encoder and are not carried through the projection. And the mask token is a single learned parameter vector of length 512, typically initialized near zero; its job is not to carry information but to occupy a slot that the positional embedding then addresses.

In the explorer above, put the query on a patch in the middle of the voiced segment and switch between global and local. Under global attention the highlighted set is the whole clip, including the silence and the siren three seconds later. Under local attention it is the 640 ms neighborhood and half the mel range — which is, physically, the region that could plausibly tell you anything about a vowel.

Worked example: what local attention actually costs

Count query-key pairs per layer on the real 64 × 8 grid:

1. Global: every one of 512 patches attends to all 512 → 512 × 512 = 262,144 pairs
2. Local 4 × 4: 32 windows, each with 16 patches attending to 16 → 32 × 16 × 16 = 32 × 256 = 8,192 pairs
3. Reduction: 262,144 / 8,192 = 32× — exactly the ratio you would predict, since window size 16 is 1/32 of 512

Now the honest part, because the naive conclusion ("so local attention is 32× cheaper") is wrong for the same reason it was wrong in Chapter 3. With decoder width d = 512 (so d² = 262,144):

Read the last row carefully: the 16-layer local decoder is about 1.7× more expensive than the 8-layer global one, not cheaper. Local attention is not a speed optimization here — it is an inductive bias. What the 32× attention reduction buys is the freedom to be deep: since attention has stopped being the scaling wall, 16 layers is affordable, and depth is where the gain lives.

The decoder design space, in one view

Four independent choices define the decoder, and the paper ablates all of them. Seeing them together shows how small the search actually was — and how much of the outcome hangs on one axis:

Every axis has a spread of half a point or less, and the default configuration takes the top of all four. That is worth noticing for a practical reason: this is a paper whose headline improvement over the prior state of the art is 0.2 mAP, built from four independent choices each worth about 0.5. Get any one of them wrong and the result evaporates. It is a good reminder that "simple extension of MAE" describes the concept, not the amount of care in the configuration.

All three decoder ablations

Attention type (Table 1c). Note the depths: the global decoder is the 8-layer variant, local is 16 layers, hybrid is local(8) + global(4). Every row fine-tunes an identically-sized ViT-B encoder — the decoders are discarded, so these differences are entirely due to what the pretext task taught the encoder:

Shifted local attention wins on three of four tasks; the hybrid ties on AS-2M and takes SID. Both beat vanilla global attention everywhere. The half-point of AS-2M mAP that separates local from global is the paper’s entire margin over the previous state of the art — 47.3 against PaSST’s and HTS-AT’s 47.1. Without the local decoder, this paper does not have a headline result.

Depth (Table 1d) and width (Table 1e), both with the default local decoder:

Depth pays until it saturates; width saturates immediately, so 512 is chosen as the trade-off point since a wider decoder is not better. Notice also that a 2-layer decoder still reaches 46.8 — matching the global 8-layer decoder. Even a very weak decoder produces a good encoder, which tells you how much of the work the masking ratio is doing.

It is worth stating what the local decoder does not claim. It does not say attention between distant time-frequency regions is useless — a repeating musical motif is exactly a long-range dependency, and the paper’s own reconstructions show music being handled well. It says that a decoder with a local default and a shifting window learns better than one that must discover locality from scratch, at this data scale and this clip length. The hybrid variant, which restores global attention in the top four layers and ties on AS-2M, is the paper quietly agreeing that global context has a place — just not everywhere, and not first.

Chapter 7: Two Stages, One Encoder

Everything so far has been architecture. This chapter is the recipe: what runs, for how long, on what hardware, with which hyperparameters — and the one genuinely surprising decision inside it.

Stage 1 — self-supervised pre-training on AudioSet-2M

The complete configuration, from the paper’s hyperparameter table:

Two entries deserve unpacking.

The effective learning rate

The paper follows MAE’s convention of specifying a base learning rate that gets scaled by batch size:

If you reproduce this at batch 128 on a single GPU, the base rate stays 0.0002 and your effective rate becomes 0.0001. Copying the number 0.0004 directly would double your intended step size. This convention silently breaks a great many reproduction attempts.

Every one of those rows is a lever someone will ask you about in a reproduction, so treat the table as a checklist rather than a description.

The schedule, evaluated

"Half-cycle cosine decay with 3 warm-up epochs" is a curve, so plot it by hand. Warm-up rises linearly from 0 to the effective rate over 3 epochs; then, with progress p = (epoch − 3) / (32 − 3):

Notice how much of the run happens at a small step size: by epoch 28 the rate is 5% of its peak. Combined with the epoch ablation in Chapter 8 — performance saturates after epoch 24 — the picture is consistent. The last quarter of pre-training is a slow polish, not new learning.

The data pipeline: cyclic sampling

For each audio, the loader randomly samples a starting time, cyclically extracts 10 seconds, and randomly jitters the magnitude by up to ±6 dB. Cyclic means it wraps: start at second 7 of a 10-second clip and you get seconds 7–10 followed by seconds 0–7 in the same tensor.

Wrapping creates a seam — an artificial discontinuity somewhere in every training example — which sounds like a defect and is in fact useful three times over. It guarantees a fixed 1024-frame length with no padding. It makes each epoch see a different view of the same clip, multiplying the effective dataset by the number of possible start offsets. And it teaches the model not to trust absolute time position, which is exactly right: the pretext task should be about time-frequency structure, not about "the siren always starts at frame 300."

Stage 2 — supervised fine-tuning

The decoder is deleted. What remains is the ViT-B encoder, plus a new head: average pooling over the token dimension, followed by a linear layer. That is the entire classifier — no class token, no attention pooling, no MLP head.

100 epochs of 200,000 sampled instances is 20 million examples, which is about ten full passes over AudioSet-2M — the paper says exactly this. The sampling is without replacement within an epoch, deliberately, to avoid duplicated samples in a batch and repeated samples within an epoch.

Worked example: the class-balancing weights

AudioSet is severely unbalanced, so fine-tuning uses importance sampling. An instance’s sampling probability is proportional to the sum of its classes’ weights, where a class weight is

where the denominator counts how many training clips carry class c. Take two classes at opposite extremes and run the numbers:

1. A head class present in 1,000,000 clips: w = 1000 / 1,000,000 + 0.01 = 0.001 + 0.01 = 0.011
2. A tail class present in 500 clips: w = 1000 / 500 + 0.01 = 2.0 + 0.01 = 2.01
3. Relative sampling boost: 2.01 / 0.011 = 182.7×
4. Now delete the epsilon and redo step 3: 2.0 / 0.001 = 2000×

So ε = 0.01 is not cosmetic — it compresses a 2000-fold imbalance into a 183-fold one. The paper describes it as set to avoid underflow in majority classes, and you can now see the shape of that: without it, head classes are sampled so rarely that the model forgets them.

And because AudioSet is multi-label, the weights add. A clip tagged both with the head class and the tail class gets 0.011 + 2.01 = 2.021 — effectively inheriting the rarity of its rarest tag. Rare events are usually accompanied by common ones (a bicycle bell in traffic), so this addition is what makes the scheme work at all.

Everything above is stage-specific plumbing. What follows is the one finding in this chapter that changes how you think.

The inversion: masking flips its personality between stages

Here is the finding promised back in Chapter 1, and it is the most counterintuitive result in the paper:

The paper’s summary is worth memorizing verbatim: for task-agnostic pre-training, unstructured masking with a higher ratio is preferred; while in task-specific fine-tuning, structured masking with lower ratios performs better.

Why does the preference flip? Because the two stages are asking for different things from the mask.

A useful way to hold the whole recipe: each stage answers a different question about the same encoder.

Why binary cross-entropy, and why mixup forces it

The fine-tuning loss is BCE for multi-label datasets or whenever mixup is enabled, and plain cross-entropy otherwise. The second clause is the interesting one.

Mixup blends two training examples and their labels: given clips a and b with a coefficient λ = 0.5, the input is 0.5x_a + 0.5x_b and the target is 0.5y_a + 0.5y_b. On a single-label task that target is no longer one-hot — it is a genuine two-point distribution. Cross-entropy against a soft target still works mathematically, but the natural reading of a mixed audio clip is "both sounds are present", which is exactly a multi-label statement. Hence BCE. And no label smoothing on top, because mixup is already smoothing the targets; stacking the two would double-count.

Mixup on audio also has a physical justification that it lacks for images: adding two waveforms is what happens when two sources play at once. A mixture of a dog and a siren is a perfectly realistic recording. A pixel-wise blend of a dog and a siren photograph is not a realistic image. This is why mixup at 0.5 is safe here while CutMix — which pastes a rectangle of one spectrogram into another, creating time-frequency discontinuities no physical source produces — cost about 0.5 mAP.

Send a batch through the pipeline above in both stages and note the one block that changes colour: the decoder. Everything upstream of it is identical between the two runs — same front end, same patchify, same encoder weights, same masking machinery. Two stages, one encoder, and a swap at the very end.

Fine-tuning the other five tasks

The same encoder is fine-tuned separately on each downstream dataset, with per-dataset input lengths, statistics, and schedules. The clip durations alone change the token count by a factor of eight:

Because positional embeddings are fixed sinusoids rather than learned tables, moving from a 64 × 8 grid to an 8 × 8 grid needs no interpolation and no surgery — you evaluate the same closed form at fewer indices. Compare that with the ImageNet-initialized models of Chapter 0, which must bilinearly resample a learned 14 × 14 grid every time the input shape changes. A design choice made for the decoder’s benefit pays off again three stages later.

The remaining pattern is consistent: The pattern is consistent: smaller datasets get fewer GPUs, higher base learning rates (0.001 instead of 0.0002), and no weighted sampling — only AudioSet-2M is unbalanced enough to need it. AudioSet-20K, notably, is fine-tuned on just 4 GPUs for 60 epochs without weighted sampling.

One detail that matters if you ever reproduce this: for ViT-S the authors use a larger learning rate than for ViT-B — 0.0005 on AS-2M fine-tuning and 0.002 on AS-20K — because they find larger learning rates work better for the ViT-S encoder. Smaller models tolerate, and want, bigger steps.

If you are reproducing this on one or two GPUs rather than sixty-four, the four settings that actually matter are, in order: get the effective learning rate right (scale the base rate by your batch size, do not copy 0.0004); keep the batch as large as gradient accumulation allows, because AdamW at β₂ = 0.95 with a small batch is noisy; stop pre-training at 24 epochs, which the ablation says costs 0.1 mAP for a 25% saving; and normalize with your dataset’s statistics rather than the AudioSet constants. The rest of the table is comfort.

Chapter 8: The Numbers, and What Moves Them

The claim is large: state-of-the-art on six audio and speech classification tasks, outperforming other recent models that use external supervised pre-training. This chapter audits it, then hands you every ablation as a knob you can turn.

How the comparison is organized

The paper splits every competitor into three groups, and the split is the argument:

Audio-MAE belongs in the middle group. Beating that group is the paper’s obligation; beating the bottom group is the surprise.

The headline table

Five observations that a table alone will not give you.

1. The AS-20K margin is the real result. 37.0 against AST’s 34.7 and MaskSpec’s 32.3 is a large gap in a field that fights over tenths. AS-20K fine-tunes on 20,000 labeled clips — precisely the low-label regime where a good pre-trained representation should pay off most, and it does.
2. The AS-2M margin is thin and honest. 47.3 against 47.1 is 0.2 mAP, within a few error bars of nothing. The paper does not hide this; it reports ±.11 from three runs. The claim it supports is "matches or beats the ImageNet-pre-trained state of the art without ImageNet", not "dominates."
3. Speaker identification is the blowout. 94.8 against AST’s 41.1. Voice identity lives in fine harmonic and formant detail across long spans — exactly what reconstructing masked spectrogram patches teaches, and exactly what an ImageNet prior knows nothing about.
4. Audio-MAE uses less data than its self-supervised rivals. SS-AST and MAE-AST add 1,000 hours of LibriSpeech; Audio-MAE uses only AudioSet, and still wins on the speech tasks.
5. The starred SUPERB numbers are not comparable. wav2vec 2.0 and HuBERT appear with linear-evaluation results from the SUPERB benchmark, where the underlying models are not end-to-end fine-tuned. The paper flags this explicitly rather than quietly claiming a win.

What "47.3 mAP" actually means

Every AudioSet number in this lesson is mean average precision, and it is worth knowing exactly what is being averaged before you decide whether 0.2 is a lot.

Because a clip can carry several of the 527 tags, accuracy is meaningless — you are not picking one answer. Instead, each class is scored independently as a ranking problem:

1. For class c, sort all 19,000 evaluation clips by the model’s score for c, highest first.
2. Walk down that ranked list. At each position where a clip that truly has class c appears, record the precision so far.
3. Average those recorded precisions — that is the average precision AP_c for class c.
4. Average AP over all 527 classes, weighting each class equally. That is mAP.

A tiny worked example makes step 3 concrete. Suppose class c is truly present in 3 of 10 clips, and the model’s ranking places them at positions 1, 3 and 6:

1. Hit at rank 1: precision = 1/1 = 1.000
2. Hit at rank 3: precision = 2/3 = 0.667
3. Hit at rank 6: precision = 3/6 = 0.500
4. AP_c = (1.000 + 0.667 + 0.500) / 3 = 2.167 / 3 = 0.722

Now the two facts that recalibrate your reading of the whole results table. First, mAP is unweighted across classes, so the 500-clip tail classes count exactly as much as the million-clip head classes — which is why the weighted sampler of Chapter 7 exists, and why it is worth 200,000 samples an epoch. Second, random guessing on a 527-class multi-label problem with sparse positives scores near zero, not near 1/527. Against that floor, 47.3 versus 47.1 is genuinely a small difference — and 37.0 versus 34.7 on AS-20K genuinely is not.

Every ablation, as a knob

The explorer below carries all nine ablations from the paper. Each one varies a single choice from the default configuration (ViT-B encoder, 16-layer shifted-local decoder of width 512, non-overlapping 16×16 patches, 80% unstructured masking, 32 epochs of AudioSet-2M pre-training). Click a bar to read the paper’s finding for that setting.

One more habit before you turn the knobs: read the error bars. Audio-MAE’s AS-2M result is 47.3 ± .11 over three runs, which means a difference of 0.2 mAP is roughly two standard deviations of run-to-run noise — suggestive, not decisive. A difference of 0.5 (global versus local decoder) is comfortably outside it. A difference of 2.3 (AS-20K against AST) is not in the same conversation. Most of the ablation deltas below are between 0.1 and 0.5; treat the 0.1 rows as ties.

Cycle through all nine knobs once before reading on. Two things become visually obvious that a stack of tables hides: almost every knob has a flat top — a plateau where two or three settings tie — and the one knob whose bars fall off a cliff is the pre-training data scale, which drops from 47.3 to 39.6 as you starve it. Architecture choices are worth tenths; data is worth units.

The three ablations worth reading twice

Pre-training data scale — and a surprise about balance

Performance rises monotonically with data, with no sign of saturation at 2M clips — the curve is still climbing when the dataset runs out.

But look at the top two rows. Twenty thousand carefully class-balanced clips and twenty thousand clips grabbed at random perform the same (39.4 versus 39.6, and the random one is marginally ahead). The paper’s reading: the distribution of data classes is less important for pre-training. This is a genuinely liberating result. For supervised learning, class balance is a first-order concern that people spend months on. For self-supervised pre-training, you can skip the curation and just collect more.

Pre-training duration

Training longer is beneficial, yet performance saturates after the 24th epoch. Note what this implies about the budget: three quarters of the final quality arrives in the first 8 epochs (9 hours), and the last 0.1 mAP costs another 9. If you are reproducing this on a budget, stop at 24.

ImageNet initialization — the negative result

The paper designs three scenarios and reports deltas against the audio-only default:

Three readings, in increasing order of importance.

Scenario 2 shows ImageNet weights alone are not sufficient, and the deficit is far worse when downstream data is small (−5.0 on AS-20K against −1.9 on AS-2M). With two million labeled clips you can fine-tune your way out of a bad initialization; with twenty thousand you cannot.

Scenario 3 is the striking one: adding a free extra pre-training stage makes the model worse. Every version of "ImageNet first, then Audio-MAE" lands below "Audio-MAE alone". And supervised ImageNet initialization is more harmful than self-supervised (−0.9 against −0.2 on AS-20K) — consistent with the label-bias concern from Chapter 0: supervised weights carry a task-specific structure that must be actively unlearned.

What generalizes beyond this paper

Strip away AudioSet and the specific numbers and three findings here are about self-supervised learning in general, not about audio:

1. Match the modality. Cross-modal initialization is a tax whose size scales inversely with how much in-domain data you have downstream. If you have a lot, it is nearly free; if you have little, it is expensive. Either way it is not a gift.
2. Curation matters less than volume for the pretext stage. Balanced and unbalanced 20K subsets tie. Spend your effort collecting, not curating — and save the curation for the supervised stage, where it is worth 200,000 weighted samples an epoch.
3. Put the domain prior in the disposable part. The decoder is thrown away, so biasing it costs nothing at deployment while shaping everything the encoder must learn. That is a design pattern, and it transfers to any asymmetric autoencoder.

What the masking ablations (Figure 4) show

The two masking sweeps are published as figures rather than tables, so only the trends and the default point are quotable with confidence. Read them as shapes:

The one anchor point published numerically is the default: 0.8 unstructured pre-training with 0.3 structured fine-tuning gives 47.3 mAP. Every other point on those curves should be read as a direction, not a measurement — and the direction is unambiguous in both.

Chapter 9: Limits, Lineage, and a Calculator

The conclusion of the paper is four sentences long and every one of them is a claim you can now defend from memory.

The limitations, in the authors’ own accounting

The appendix contains a limitations section that is unusually candid, and each item points at a paper somebody wrote afterwards:

The stated future direction is multimodal self-supervised learning with a joint audio-visual MAE, because those domains share natural correspondences in video data. Audio-visual masked autoencoders duly appeared within a year.

A calculator you can point at your own audio

Everything in Chapters 2 through 6 is arithmetic. Here it is, wired up: give it a clip length and a masking ratio and it produces the full shape ledger, including the encoder cost from the model you derived in Chapter 3 and the decoder window count from Chapter 6.

Set the duration to 1 second and you recover the Speech Commands shape: 1 × 128 × 128, an 8 × 8 grid, 64 patches. Set it to 5 and you get ESC-50. Push it past about 80 seconds and the token count crosses 4,608 — the point from Chapter 3 where quadratic attention finally overtakes the linear terms, and the paper’s compute limitation stops being theoretical.

A last observation about that crossover, since it is the limitation with the sharpest edge. The paper’s efficiency argument — encode only 20%, save five and a half times the compute — is a statement about ten-second clips. Stretch the input toward the minute-long recordings that real deployments contain, and the quadratic term takes over: at 60 seconds the token count is 3,072 and attention is already a quarter of each block; at 90 seconds it is roughly half. The masking trick still helps, but it stops being sufficient on its own. Everything the paper lists as future work about lengthy audio lives on the far side of that line.

Where Audio-MAE sits

Continue on this site

Everything in the calculator above is derivable by hand from Chapter 2 and Chapter 3; the tool exists to check you, not to replace you.

Misconceptions worth killing on the way out

The cheat sheet

Everything you need to re-derive the paper on a whiteboard.

Walk the lineage map left to right once. The left column is four papers from four different problems — language modeling, image reconstruction, image classification efficiency, audio tagging — and Audio-MAE is the join. That is worth noticing as a research pattern in its own right: the contribution is not a new mechanism, it is a correct assembly, plus one measured decision about which mechanism to bias.

Build it yourself

If you have a ViT implementation and a GPU, the whole method is about 150 lines on top of it. The checklist, in dependency order:

1. Log-mel front end: 16 kHz, 25 ms / 10 ms, 128 bands. Pad or trim to a multiple of 16. Normalize with your dataset’s own mean and std — not AudioSet’s.
2. Conv2d(1, 768, 16, stride=16) plus fixed 2-D sinusoidal position embeddings.
3. The random_masking function from Chapter 3 verbatim. Keep ids_restore.
4. A vanilla ViT-B on the gathered 20%.
5. Decoder: project to 512, append mask tokens, un-shuffle, add decoder positions, 16 blocks — and window the attention into 4 × 4 patch blocks, shifting by 2 on alternate layers with a cyclic roll.
6. Loss: patch-normalize the target, MSE, average over masked positions only.
7. Fine-tune: delete the decoder, mask 0.3 in time and 0.3 in frequency, average-pool, linear head, BCE.

Step 5 is the only one that is not boilerplate, and it is the only one that matters for the headline number. Everything else is MAE with a spectrogram in front of it — which is, in the end, exactly what the title promises.

One closing frame for the whole lesson. Audio-MAE is a paper about where to put your assumptions. The authors had exactly two audio-specific beliefs — spectrograms are redundant enough to mask at 80%, and they correlate locally — and they spent both of them on parts of the system that never ship: the masking policy and a decoder that gets deleted. What ships is a plain ViT-B that had no idea it was learning about sound. That is why the same encoder tops audio tagging, environmental sound, keyword spotting and speaker identification without a single architectural change between them.