Fine-tuning Whisper

Week 5 of the residency. How sound becomes something a transformer can read, and what a handful of gradient steps on one clip actually does.

Week 5 is titled audio is all you need, and its subtitle says what it’s for: adding another modality to our tool box.

Text was week 1. Images were week 3. The argument each time was that the transformer doesn’t care what the input was, as long as it arrives as a sequence of vectors.

Audio is the third test of that claim.

Two tasks.

Note the annotation on that slide. Once audio is a spectrogram it’s a two-dimensional array of intensities. That’s a picture, and the patch projection from week 3 works on it unchanged.

The second task is the one this post follows, and there’s the running example again. Hello, my name is Bes was the sentence CBOW completed in week 1. Here a speech model has to hear it.

What Sound Is

A waveform is amplitude over time, sampled 16,000 times a second, and almost none of the structure a listener cares about is visible in it directly. Two different sounds are two different bundles of vibrations at different frequencies.

The Fourier transform unbundles it. Run it over short overlapping windows rather than the whole clip and you get frequency on one axis, time on the other, intensity as the value. That is the spectrogram, and it is why the previous slide could call audio “clearly images”.

Whisper adds two refinements.

• The frequency axis is squashed onto the mel scale. It spaces bands the way hearing does: fine detail low down, coarser high up. 80 bands replace 201 raw frequency bins.
• The values go through a logarithm, because loudness is perceived multiplicatively.

Below is that whole front end, built from scratch on the actual clip from the project, so you can see the waveform go in and the picture come out.

import sys

sys.path.insert(0, "..")
import numpy as np
from huggingface_hub import hf_hub_download

REVISION = "f705fed08827ff6c36e3b5329495c943a5e544e8"
clip = np.load(hf_hub_download("roshbeed/ai-residency-blog-data", "audio/clip-16k.npz",
                               repo_type="dataset", revision=REVISION))
waveform, rate = clip["waveform"], int(clip["sample_rate"])

N_FFT, HOP, N_MELS = 400, 160, 80


def stft(x, n_fft, hop):
    """Fourier transform of each short overlapping window."""
    window = np.hanning(n_fft + 1)[:-1]
    frames = 1 + (len(x) - n_fft) // hop
    return np.stack([np.fft.rfft(x[i * hop:i * hop + n_fft] * window)
                     for i in range(frames)], axis=1)


def mel_filterbank(rate, n_fft, n_mels):
    """Triangular filters, evenly spaced on the mel scale rather than in hertz."""
    to_mel = lambda f: 2595 * np.log10(1 + f / 700)
    to_hz = lambda m: 700 * (10 ** (m / 2595) - 1)
    edges = to_hz(np.linspace(to_mel(0), to_mel(rate / 2), n_mels + 2))
    bins = np.floor((n_fft + 1) * edges / rate).astype(int)

    bank = np.zeros((n_mels, n_fft // 2 + 1))
    for m in range(1, n_mels + 1):
        left, centre, right = bins[m - 1], bins[m], bins[m + 1]
        if centre > left:
            bank[m - 1, left:centre] = (np.arange(left, centre) - left) / (centre - left)
        if right > centre:
            bank[m - 1, centre:right] = (right - np.arange(centre, right)) / (right - centre)
    return bank


power = np.abs(stft(waveform, N_FFT, HOP)) ** 2
log_mel = np.log10(np.maximum(mel_filterbank(rate, N_FFT, N_MELS) @ power, 1e-10))

print(f"{len(waveform):,} samples at {rate} Hz = {len(waveform) / rate:.2f} seconds")
print(f"becomes an {log_mel.shape[0]} by {log_mel.shape[1]} picture")

16,982 samples at 16000 Hz = 1.06 seconds
becomes an 80 by 104 picture

Drawn, the waveform and the picture it becomes:

import matplotlib.pyplot as plt
from _style import COLOURS, MUTED, style_axes

fig, (top, bottom) = plt.subplots(2, 1, figsize=(7.4, 4.4),
                                  gridspec_kw={"height_ratios": [1, 2]})
top.plot(np.arange(len(waveform)) / rate, waveform, color=COLOURS[0], linewidth=0.5)
top.set_xlim(0, len(waveform) / rate)
style_axes(top, ylabel="amplitude", grid=None)
top.set_xticks([])
bottom.imshow(log_mel, aspect="auto", origin="lower", cmap="magma",
              extent=(0, len(waveform) / rate, 0, N_MELS))
style_axes(bottom, "Seconds", "Mel band", grid=None)
fig.tight_layout()

Figure 1: One second of speech, before and after. The model never sees the waveform on top.

Whisper

An encoder over the spectrogram, a decoder cross-attending to it. Structurally this is week 3’s multi-digit reader with a spectrogram where the image was, which is the claim the week set out to test.

The two convolutions at the front are doing something worth noticing: the second has stride 2, halving the sequence from 3000 positions to 1500. Attention cost grows with the square of sequence length, so that one layer is a four-fold saving before any attention runs.

And a Whisper transcript does not start with words. It starts with a structured token prefix, which is how one model handles transcription, translation and language identification without being three models. That token sequence at the bottom is the running example again, tokenised.

Fine-tuning on One Example

The base model mishears the name. Bess, with two esses.

The workshop’s demonstration is to fine-tune on that one clip and watch it correct, which makes the mechanism visible in a way a held-out score never does. It also makes something else visible, which is what the rest of this post measures.

I cannot run Whisper during a page build. What follows is the same shape at a size that trains in seconds.

The language is synthetic. Each letter is a tone at its own frequency, and words are three letters long. Speaking a word plays its tones in order. Transcribing reads them back, through the same spectrogram front end as above.

import torch
import torch.nn as nn

RATE, TONE_SECONDS, NOISE = 8000, 0.08, 0.15
ALPHABET = "abcdefgh"
TONES = {c: 300 * 1.28 ** i for i, c in enumerate(ALPHABET)}

rng = np.random.default_rng(0)
WORDS = sorted({"".join(rng.choice(list(ALPHABET), 3)) for _ in range(60)})[:40]


def speak(word, seed, shift=1.0):
    g = np.random.default_rng(seed)
    parts = []
    for c in word:
        t = np.arange(int(RATE * TONE_SECONDS)) / RATE
        f = TONES[c] * shift * (1 + 0.02 * g.normal())
        parts.append((np.sin(2 * np.pi * f * t) + 0.4 * np.sin(4 * np.pi * f * t))
                     * np.hanning(len(t)))
    x = np.concatenate(parts)
    return (x + NOISE * g.normal(size=len(x))).astype(np.float32)


SMALL_FFT, SMALL_HOP, SMALL_MELS = 256, 64, 32
BANK = mel_filterbank(RATE, SMALL_FFT, SMALL_MELS)


def to_picture(x):
    return np.log10(np.maximum(BANK @ np.abs(stft(x, SMALL_FFT, SMALL_HOP)) ** 2,
                               1e-10)).T.astype(np.float32)


def build(repeats, seed0):
    pictures, labels = [], []
    for index, word in enumerate(WORDS):
        for k in range(repeats):
            pictures.append(to_picture(speak(word, seed0 + index * 1000 + k)))
            labels.append(index)
    return torch.from_numpy(np.stack(pictures)), torch.tensor(labels)


X, Y = build(12, 0)
X_test, Y_test = build(4, 500_000)
print(f"{len(WORDS)} words, {len(X)} training clips of shape {tuple(X.shape[1:])}")

40 words, 480 training clips of shape (27, 32)

The model is an encoder over the spectrogram and a decoder that spells the word, which is Whisper’s shape at about a thousandth of the size. Trained for twenty epochs it gets nearly everything right.

LETTERS = sorted(set("".join(WORDS)))
START, END = len(LETTERS), len(LETTERS) + 1
VOCAB = len(LETTERS) + 2
DIM, FRAMES = 64, X.shape[1]
TARGET = torch.tensor([[START] + [LETTERS.index(c) for c in w] + [END] for w in WORDS])


class Transcriber(nn.Module):
    def __init__(self):
        super().__init__()
        self.input = nn.Linear(SMALL_MELS, DIM)
        self.audio_positions = nn.Parameter(torch.randn(1, FRAMES, DIM) * 0.02)
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(DIM, 4, 4 * DIM, batch_first=True,
                                       norm_first=True, dropout=0.0), 2)
        self.embed = nn.Embedding(VOCAB, DIM)
        self.text_positions = nn.Parameter(torch.randn(1, 5, DIM) * 0.02)
        self.decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(DIM, 4, 4 * DIM, batch_first=True,
                                       norm_first=True, dropout=0.0), 2)
        self.out = nn.Linear(DIM, VOCAB)

    def forward(self, picture, tokens):
        memory = self.encoder(self.input(picture) + self.audio_positions)
        h = self.embed(tokens) + self.text_positions[:, :tokens.shape[1]]
        mask = nn.Transformer.generate_square_subsequent_mask(tokens.shape[1])
        return self.out(self.decoder(h, memory, tgt_mask=mask))


torch.manual_seed(0)
model = Transcriber()
optimiser = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
generator = torch.Generator().manual_seed(0)


@torch.no_grad()
def transcribe(pictures):
    tokens = torch.full((len(pictures), 1), START)
    for _ in range(3):
        tokens = torch.cat([tokens, model(pictures, tokens)[:, -1].argmax(-1, keepdim=True)], 1)
    return tokens[:, 1:]


spell = lambda ids: "".join(LETTERS[i] for i in ids if i < len(LETTERS))
accuracy = lambda: (transcribe(X_test) == TARGET[Y_test][:, 1:4]).all(1).float().mean().item()

for epoch in range(20):
    perm = torch.randperm(len(X), generator=generator)
    for i in range(0, len(perm) - 64, 64):
        b = perm[i:i + 64]
        target = TARGET[Y[b]]
        loss = nn.functional.cross_entropy(model(X[b], target[:, :-1]).reshape(-1, VOCAB),
                                           target[:, 1:].reshape(-1))
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()

print(f"transcribes {accuracy():.1%} of held-out clips exactly right")

transcribes 99.4% of held-out clips exactly right

Drawn out, it is the encoder-decoder shape again:

from _arch import diagram

diagram(Transcriber(), input_shape=((1, FRAMES, SMALL_MELS), (1, 4)), style="flow",
        input_dtype=(torch.float32, torch.long))

Figure 2: Whisper’s shape at about a thousandth of the size. The spectrogram enters top left, the text so far enters bottom left through its embedding, and the two streams meet in the decoder.

Now it needs its own version of mishearing a name. Say one word with every tone shifted up 10%, which is a voice it has never encountered, and it gets it wrong.

WORD, SHIFT = "acd", 1.10
one_clip = torch.from_numpy(to_picture(speak(WORD, 999, SHIFT))).unsqueeze(0)
one_target = TARGET[WORDS.index(WORD)].unsqueeze(0)

print(f"the word is {WORD!r}")
print(f"the model hears {spell(transcribe(one_clip)[0].tolist())!r}")
print(f"and is otherwise fine: {accuracy():.1%} on the clean test clips")

the word is 'acd'
the model hears 'adc'
and is otherwise fine: 99.4% on the clean test clips

That’s the situation the workshop demonstration is in: a model that works, and one specific thing it gets wrong. Fine-tune on that single clip and watch both numbers at once: the loss on the clip, and accuracy on everything else.

fine_tune = torch.optim.AdamW(model.parameters(), lr=1e-4)

print(f"{'step':>5} {'loss':>10} {'hears':>8} {'clean accuracy':>16}")
print(f"{0:>5} {'-':>10} {spell(transcribe(one_clip)[0].tolist()):>8} {accuracy():>15.1%}")

steps, losses, clean = [], [], []
for step in range(1, 13):
    loss = nn.functional.cross_entropy(
        model(one_clip, one_target[:, :-1]).reshape(-1, VOCAB), one_target[:, 1:].reshape(-1))
    fine_tune.zero_grad()
    loss.backward()
    fine_tune.step()

    steps.append(step); losses.append(loss.item()); clean.append(accuracy())
    if step % 2 == 0:
        print(f"{step:>5} {loss.item():>10.5f} {spell(transcribe(one_clip)[0].tolist()):>8} "
              f"{clean[-1]:>15.1%}")

 step       loss    hears   clean accuracy
    0          -      adc           99.4%
    2    0.25364      acd           99.4%
    4    0.01372      acd           97.5%
    6    0.00431      acd           96.9%
    8    0.01029      acd           94.4%
   10    0.01842      acd           93.1%
   12    0.01828      acd           92.5%

Both numbers on one chart, the clip’s loss against everything else’s accuracy:

from _style import figure

fig, ax = figure(height=4.0)
ax.plot(steps, losses, color=COLOURS[1], linewidth=2)
style_axes(ax, "Fine-tuning step on the single clip", "Loss on that clip")
ax.annotate("loss on the one clip", xy=(steps[-1], losses[-1]), xytext=(-6, 14),
            textcoords="offset points", ha="right", fontsize=9, color=COLOURS[1])

right = ax.twinx()
right.plot(steps, clean, color=COLOURS[0], linewidth=2)
right.set_ylabel("Accuracy on the other words", color=MUTED, fontsize=9)
right.tick_params(colors=MUTED, labelsize=9, length=0)
right.spines["top"].set_visible(False)
right.annotate("everything else", xy=(steps[-1], clean[-1]), xytext=(-6, -16),
               textcoords="offset points", ha="right", fontsize=9, color=COLOURS[0])
fig.tight_layout()

Figure 3: The fix lands almost immediately. Everything after that is the model memorising one clip at the expense of the other thirty-nine words.

The fix is cheap: two gradient steps and the word is right. On the real project a handful of steps corrected the name.

The loss then falls to nearly zero. That looks like good news. It’s not.

With one training example, a loss near zero means the model has memorised that clip. It’s precisely what you asked for. It’s also what quietly goes wrong when the same procedure is scaled up.

And the cost lands somewhere nobody is looking. Accuracy on the other thirty-nine words falls steadily from step four, and by step twelve the model has given up several points to keep improving on a clip it already got right. Nothing in the fine-tuning loop reports that.

So the question isn’t whether fine-tuning works. It’s when to stop. Answering that means evaluating the thing you aren’t training on.

Conclusion

• A spectrogram turns sound into a picture, and week 3’s machinery then applies
• Whisper is an encoder-decoder over that picture
• Fine-tuning on one example works, and works fast
• It also costs you accuracy everywhere else, and nothing reports that
• Measure what you are not training on

One Whisper-specific trap, which produced no error at all. The tokenizer emits <|startoftranscript|> at the front of a transcript, and it is tempting to pass its output straight through as labels. The model prepends decoder_start_token_id itself when it shifts labels right, and that token is start-of-transcript, so every position ends up one out of step. Nothing raises, the loss goes down, and the model learns a different task from the one you meant.

The full project is on GitHub.