Emotion Recognition with Whisper

The second half of week 5. Two ways to get an emotion out of a speech model: a classifier head, or one more token in the vocabulary.

The other half of week 5. The workshop set two tasks, and the other post was the second one.

This is task one: given a clip, say something about it.

Whisper transcribes. It doesn’t tell you how something was said. How is often the point. The same sentence read angrily and read sadly is one transcript and two different messages.

The slide offers a classifier as the obvious route.

• Take the encoder
• Average its outputs into one vector
• Put a single linear layer on top, one score per emotion
• Train it so the right emotion scores highest

Three lines, and it works.

I built it the other way, because of something Whisper already does.

Whisper’s Token Prefix

A Whisper transcript does not start with words. It starts with a structured token prefix, and that prefix is what lets one model do transcription, translation and language identification without being three models.

Those are not a separate mechanism bolted on. They are ordinary vocabulary entries, predicted by the same softmax and trained by the same cross-entropy as any word. Look at the ids on that slide:

<|startoftranscript|><|en|><|transcribe|><|notimestamps|>Hello, my name is Bes.<|endoftext|>

[50258, 50259, 50359, 50363, 15947, 11, 452, 1315, 307, 8190, 13, 50257]

The first four are the language and the task. Nothing distinguishes them from 15947 except what the model learned to do when it sees them.

So: is emotion a different kind of thing from language and task, or the same kind of thing? It is hard to argue it is different. Adding <emotion_happy> and its siblings to that vocabulary lets the model say what it heard through machinery it already has.

import sys

sys.path.insert(0, "..")
import matplotlib.pyplot as plt
from _style import COLOURS, MUTED

fig, axes = plt.subplots(1, 2, figsize=(8.4, 3.4))


def box(ax, x, y, w, h, label, colour, text="white", size=9):
    ax.add_patch(plt.Rectangle((x, y), w, h, facecolor=colour, edgecolor="none"))
    ax.text(x + w / 2, y + h / 2, label, ha="center", va="center",
            color=text, fontsize=size)


def arrow(ax, x1, y1, x2, y2):
    ax.annotate("", xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle="-|>", color="#aeb6bf", linewidth=1.3))


for ax, title in zip(axes, ("a classifier head", "a token in the vocabulary")):
    box(ax, 0.05, 0.42, 0.22, 0.18, "encoder", COLOURS[0])
    ax.set_title(title, fontsize=10, color=MUTED, pad=10)
    ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.axis("off")

# left: two outputs, two losses
arrow(axes[0], 0.27, 0.56, 0.40, 0.72)
arrow(axes[0], 0.27, 0.46, 0.40, 0.30)
box(axes[0], 0.40, 0.64, 0.22, 0.16, "decoder", COLOURS[2])
box(axes[0], 0.40, 0.22, 0.22, 0.16, "Linear", COLOURS[3])
axes[0].text(0.66, 0.72, "a c d", fontsize=10, color=MUTED, va="center")
axes[0].text(0.66, 0.30, "shaky", fontsize=10, color=MUTED, va="center")
axes[0].text(0.5, 0.06, "two losses", fontsize=9, color=MUTED, ha="center")

# right: one output, one loss
arrow(axes[1], 0.27, 0.51, 0.40, 0.51)
box(axes[1], 0.40, 0.43, 0.22, 0.16, "decoder", COLOURS[2])
arrow(axes[1], 0.62, 0.51, 0.70, 0.51)
axes[1].text(0.72, 0.51, "<shaky>  a c d", fontsize=10, color=MUTED, va="center")
axes[1].text(0.5, 0.06, "one loss", fontsize=9, color=MUTED, ha="center")

fig.tight_layout()

Figure 1: The same encoder, two ways to get an emotion out of it. On the left the emotion leaves through a separate head trained by its own loss; on the right it is the first thing the decoder says.

The appeal isn’t elegance for its own sake. Three concrete things follow.

• Adding a seventh emotion is one more row in a table, not a reshaped output layer and a checkpoint that no longer loads
• There’s one loss instead of two, so nothing to tune between them
• The emotion is emitted first, so every token after it is generated with it in context

A classifier head produces its answer off to the side, where the transcription never sees it. That last point is the real argument for the design. It’s also the one that costs something.

Measuring Both

Same synthetic speech as the fine-tuning post: three-letter words where each letter is a tone. Emotion is an acoustic property on top. bright raises the pitch, shaky adds a tremolo, calm is neither.

Same encoder, same decoder size, same data. The only difference is where the emotion comes out.

import numpy as np
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)}
EMOTIONS = ["calm", "bright", "shaky"]

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


def speak(word, seed, emotion):
    g = np.random.default_rng(seed)
    parts = []
    for c in word:
        t = np.arange(int(RATE * TONE_SECONDS)) / RATE
        f = TONES[c] * (1 + 0.02 * g.normal())
        if emotion == 1:                       # bright: everything a bit higher
            f *= 1.18
        tone = np.sin(2 * np.pi * f * t) + 0.4 * np.sin(4 * np.pi * f * t)
        if emotion == 2:                       # shaky: amplitude wobbles at 22 Hz
            tone = tone * (1 + 0.6 * np.sin(2 * np.pi * 22 * t))
        parts.append(tone * np.hanning(len(t)))
    x = np.concatenate(parts)
    return (x + NOISE * g.normal(size=len(x))).astype(np.float32)

The same spectrogram front end as the other post turns each clip into a picture.

N_FFT, HOP, N_MELS = 256, 64, 32


def stft(x, n_fft, hop):
    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):
    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


BANK = mel_filterbank(RATE, N_FFT, N_MELS)


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


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


X, W, E = build(12, 0)
X_test, W_test, E_test = build(4, 500_000)
print(f"{len(X)} clips, {len(WORDS)} words x {len(EMOTIONS)} emotions")
print(f"guessing the emotion: {1 / len(EMOTIONS):.3f}")

480 clips, 40 words x 3 emotions
guessing the emotion: 0.333

LETTERS = sorted(set("".join(WORDS)))
EMOTION_TOKENS = [len(LETTERS) + i for i in range(3)]   # the extra vocabulary entries
START, END = len(LETTERS) + 3, len(LETTERS) + 4
VOCAB = len(LETTERS) + 5
DIM, FRAMES = 64, X.shape[1]

# <start> <emotion> l l l <end> — the emotion slot is filled per clip
TARGET = torch.stack([torch.tensor([START, 0] + [LETTERS.index(c) for c in w] + [END])
                      for w in WORDS])


class Model(nn.Module):
    """One encoder, one decoder, and a classifier head that only the head
    variant ever trains."""

    def __init__(self):
        super().__init__()
        self.input = nn.Linear(N_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, 6, 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)
        self.classifier = nn.Linear(DIM, 3)

    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)), self.classifier(memory.mean(1))

Both designs share one model: an encoder over the spectrogram, a decoder, and a small classifier head hanging off the encoder that only one of them ever trains.

from _arch import diagram

diagram(Model(), input_shape=((1, X.shape[1], N_MELS), (1, 5)), style="flow",
        input_dtype=(torch.float32, torch.long))

Figure 2: The shared model. Both designs are this; only the loss differs, and whether the emotion leaves through the decoder or through the small head hanging off the encoder.

Training them differs only in where the emotion comes out. The token design puts it in the sequence and uses one cross-entropy over everything. The head design removes it from the sequence and adds a second loss on the classifier.

def train(design, epochs=45):
    torch.manual_seed(0)
    model = Model()
    optimiser = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
    generator = torch.Generator().manual_seed(0)

    for _ in range(epochs):
        perm = torch.randperm(len(X), generator=generator)
        for i in range(0, len(perm) - 64, 64):
            b = perm[i:i + 64]
            target = TARGET[W[b]].clone()
            target[:, 1] = torch.tensor(EMOTION_TOKENS)[E[b]]

            if design == "token":
                logits, _ = model(X[b], target[:, :-1])
                loss = nn.functional.cross_entropy(logits.reshape(-1, VOCAB),
                                                   target[:, 1:].reshape(-1))
            else:
                text = torch.cat([target[:, :1], target[:, 2:]], dim=1)  # emotion removed
                logits, predicted = model(X[b], text[:, :-1])
                loss = (nn.functional.cross_entropy(logits.reshape(-1, VOCAB),
                                                    text[:, 1:].reshape(-1))
                        + nn.functional.cross_entropy(predicted, E[b]))

            optimiser.zero_grad()
            loss.backward()
            optimiser.step()
    return model


@torch.no_grad()
def evaluate(model, design):
    tokens = torch.full((len(X_test), 1), START)
    steps = 4 if design == "token" else 3
    for _ in range(steps):
        tokens = torch.cat([tokens, model(X_test, tokens)[0][:, -1].argmax(-1, keepdim=True)], 1)

    if design == "token":
        emitted = tokens[:, 1]
        emotion = torch.tensor([EMOTION_TOKENS.index(i) if i in EMOTION_TOKENS else -1
                                for i in emitted.tolist()])
        word = (tokens[:, 2:5] == TARGET[W_test][:, 2:5]).all(1)
    else:
        emotion = model(X_test, tokens[:, :1])[1].argmax(1)
        word = (tokens[:, 1:4] == TARGET[W_test][:, 2:5]).all(1)

    return (emotion == E_test).float().mean().item(), word.float().mean().item()


results = {}
for design in ("head", "token"):
    results[design] = evaluate(train(design), design)

# Built as one string and printed once, so the table arrives as a single block
rows = [f"{'design':>26} {'emotion':>9} {'word':>9}"]
for design, name in (("head", "a classifier head"), ("token", "a token in the vocabulary")):
    emotion, word = results[design]
    rows.append(f"{name:>26} {emotion:>9.3f} {word:>9.3f}")
print("\n".join(rows))

                    design   emotion      word
         a classifier head     0.800     0.962
 a token in the vocabulary     0.781     0.975

Conclusion

On the emotion itself the two designs land within 0.02 of each other, which is the answer I would expect. The information is in the encoder either way and both designs read the same encoder; where you attach the readout does not change what there is to read.

The token design transcribes worse here, and that is the real trade. A two-layer decoder now has to produce the emotion and spell the word, and it does not have the capacity to do both as well as a decoder that only spells.

That cost is a small-model effect. Whisper’s decoder is far larger and already emits several control tokens before it writes anything, so one more costs it nothing noticeable. It would have been easy to run this comparison at a single size, watch the token design lose on transcription, and conclude the design was worse, when what I had measured was that my decoder was too small.

The architectural argument survives the accuracy result regardless, for the reasons above: a vocabulary entry instead of a reshaped head, one loss instead of two, and an emotion the transcription can actually see.

The full project is on GitHub.