IOAI Training: Learn AI by solving problems

The problem in brief

A company runs acoustic monitoring stations with an Audio Spectrogram Transformer (AST) that recognises 16 everyday and wildlife sounds (dog, rooster, rain, keyboard typing...). Two new customers want 13 more classes: logging and poaching sounds (axe, chainsaw, hand saw, generator, engine, helicopter, gunshot, firework, crackling fire) and four singing insect species. You must turn the deployed 16-class model into a single 29-class model, starting from its checkpoint, with only a few minutes of single-GPU training, and without forgetting the old classes. Only a small leftover sample of the old data is available (296 training clips, as few as 3 per class), plus 624 training clips for the new classes.

This was a practice task of the IOAI 2026 at-home round: the hidden test set was never published, so all our numbers are on the official validation splits provided with the data (195 old clips, 168 new clips), which no model selection touched except where stated.

How you are scored

On a 29-way prediction, accuracy is measured separately on clips whose true class is old and on clips whose true class is new, and the two are averaged:

\text{Score} = \tfrac12\,\text{Acc}_{\text{old}} + \tfrac12\,\text{Acc}_{\text{new}} .

Forgetting the old classes costs exactly as much as not learning the new ones. The deployed model scores $\tfrac12 \cdot 0.979 + \tfrac12 \cdot 0 = 0.490$ (it cannot output a new class). There is no published threshold or leaderboard for this practice task.

First look at the data

Very little old data. The retained old set is tiny and uneven: Sheep and Keyboard Typing have 3 training clips each, Thunderstorm 62. You cannot relearn the old classes from it; you can only use it to remind the model.
The audio is not what the model expects. Despite the notebook's description, the clips are stored at 44.1 or 48 kHz in stereo; AST was trained on 16 kHz mono, so every clip must be downmixed and resampled (the notebook does it with librosa.load(sr=16000); we use scipy.signal.resample_poly).
The encoder is huge, the head tiny. AST has 86 M parameters, and the 16-class head (a LayerNorm and one linear layer) has 13.8 K. Each clip becomes a 768-number pooled embedding (average of the two summary tokens), and the head is just a linear classifier on top of it.

Running the frozen encoder once over all 1 283 clips takes about 8 minutes on 2 CPU threads; after that, every experiment on the head takes seconds. That makes it cheap to explore.

A first baseline

The obvious move: add 13 rows to the head and train on the new clips only. The new rows learn well (80.8 % on new classes), but the score falls to 0.404, below the deployed model: old accuracy collapses to 0 %. This is catastrophic forgetting in its purest form. The cross-entropy loss only ever sees new labels, so the cheapest way to lower it is to push all old logits down, and the old classes never win again.

Freezing the 16 old rows of the head does not save them either (score 0.682, old accuracy 55 %). The old rows stay intact, but the new rows learn to produce large logits, larger than the old ones, on every clip, including old ones. Forgetting is not only about weights changing; it is about the competition between classes in the softmax.

The key idea

Two things fight forgetting, and they act on different parts of the problem.

1. Replay. Mix the old training clips back into every batch, with class-balanced sampling so that a class with 3 clips is seen as often as one with 60. Now the loss also punishes the model when it predicts a new class for a dog. Score: 0.779, with old and new accuracies balanced (0.79 and 0.77). But 296 old clips are too few to pin down 16 classes, and the old accuracy is far below the original 0.98.

2. Knowledge distillation (KD). We do have something much richer than the old labels: the deployed model itself. On every training clip, old or new, it gives a full probability distribution over the 16 old classes, which encodes what it learned from the data we no longer have. We ask the new model's 16 old logits to reproduce it:

\mathcal{L} = \text{CE}(z, y) + \lambda\,T^2\,\mathrm{KL}\!\left(\operatorname{softmax}(z^{0}_{1..16}/T)\;\middle\|\;\operatorname{softmax}(z_{1..16}/T)\right),

where $z$ are the new model's 29 logits, $z^0$ the frozen model's 16 logits, $T = 2$ softens both distributions (so that the "dark knowledge" in the small probabilities, such as "this rain sounds a bit like sea waves", is transmitted), and the $T^2$ factor keeps the gradient size comparable to the cross-entropy. This is the idea behind Learning without Forgetting (Li and Hoiem, 2016).

Why does it work so well here? Because the KD term is applied to new clips too. On a chainsaw clip, the old model says "mostly Frog, a bit of Rain"; KD asks the new model to keep that relative pattern among the old logits, while cross-entropy pushes the new Chainsaw logit above all of them. The two losses no longer fight: one constrains the shape of the old logits, the other only needs to raise one new logit. Every clip becomes a reminder of the old behaviour, which multiplies the effective size of the replay set.

The weight $\lambda$ matters a lot: $\lambda = 1$ gives 0.827, $\lambda = 3$ gives 0.869, and $\lambda = 10$ gives 0.875, with old accuracy back at 0.981, as good as the original model, while new accuracy stays at 0.77.

Old and new accuracy for each strategy

The full solution

Embed once. Load each clip as 16 kHz mono, compute AST's 1024 x 128 log-mel input, run the frozen encoder, keep the 768-number pooled embedding and the deployed head's 16 logits.
Grow the head. New Linear(768, 29): rows 0 to 15 copied from the deployed head (so the model starts exactly where it was), rows 16 to 28 initialised small, with bias equal to the mean old bias. The LayerNorm before it is kept.
Train the head for 60 epochs with AdamW (learning rate $10^{-3}$ , weight decay $10^{-4}$ ), batches of 32 drawn from old and new training clips with class-balanced sampling, label smoothing 0.1, plus the KD term with $\lambda = 10$ , $T = 2$ .
```
loss = F.cross_entropy(z, y, label_smoothing=0.1)
kd = F.kl_div(F.log_softmax(z[:, :16] / T, -1), F.softmax(z_frozen / T, -1), reduction="batchmean")
loss = loss + 10 * T * T * kd
```
Put the new head back on the untouched encoder and save one 29-class model. Training takes about one minute on a CPU, far inside the 10-minute GPU budget.

Did fine-tuning the encoder help? We also cached the input of encoder layer 10 for every clip and fine-tuned the last two transformer layers plus the head (learning rate $3 \times 10^{-5}$ for the layers, same replay and KD $\lambda = 10$ , 4 epochs, 15 minutes on CPU). It reached 0.868 (old 0.974, new 0.762): no better than the head alone. With about 45 clips per new class, the frozen AudioSet features are already good enough, and updating 14 M more parameters mostly adds variance and a little forgetting. On the contest GPU one could afford more epochs and augmentation, but the evidence here says the head is where the gains are.

Where are the remaining errors? The hard new classes are the short, impulsive or noisy ones: Gunshot (47 %, often taken for Axe), Crackling Fire and Vehicle Engine (both often taken for Rain), Firework (taken for Thunderstorm). Two of the four insects are easy (Tettigonia viridissima 100 %), but Gryllus campestris is confused with Tettigonia. These are genuinely similar sounds, and each new class has only 15 validation clips, so a single clip moves a class's accuracy by 7 points.

Accuracy of each new class with its most frequent confusion

Results

Validation splits provided with the task (official test labels are hidden). Head-only results are the mean of 3 seeds; the standard deviation over seeds is 0.003 to 0.009.

approach	old acc.	new acc.	score
deployed 16-class model	0.979	0.000	0.490
new head, new clips only	0.000	0.808	0.404
same, old head rows frozen	0.554	0.810	0.682
replay old + new clips	0.790	0.768	0.779
replay, class-balanced	0.791	0.762	0.777
replay + KD, $\lambda = 1$	0.882	0.772	0.827
replay + KD, $\lambda = 3$	0.974	0.764	0.869
replay + KD $\lambda = 3$ , old rows frozen	0.979	0.758	0.869
replay + KD $\lambda = 3$ , 150 epochs	0.971	0.754	0.863
replay + KD, $\lambda = 10$ (final)	0.981	0.770	0.876
fine-tune last 2 encoder layers + head, KD 10 (1 seed)	0.974	0.762	0.868

The final model reaches 0.876 on validation, against 0.490 for the deployed model, while keeping the old classes at their original accuracy.

What to remember

Catastrophic forgetting is a softmax competition. Training only on new labels pushes old logits down; even frozen old weights lose because the new logits grow. You must show the model old examples or old behaviour.
The old model is your best data. Distilling its outputs on every clip, including new ones, preserves what it knew far better than a small replay set (old accuracy 0.79 with replay alone, 0.98 with KD).
Cache the frozen part. When the encoder is frozen, compute its features once; dozens of experiments then cost seconds, not hours.
Fine-tuning more is not automatically better. With tens of examples per class, a strong pretrained encoder plus a well-trained head often beats partial fine-tuning, which adds variance and forgetting.
Check the input pipeline. The audio was 44.1/48 kHz stereo while the model expects 16 kHz mono: one silent mismatch like this can cost more than any training trick.

Note on the code. The full solution.py below was assembled from scripts that did run, but it was not re-run end to end after assembly.

"""Operation Night Watch (IOAI 2026 at-home task 1): grow a 16-class AST audio classifier
to 29 classes without forgetting the old ones.

Final recipe (about 1 minute of training once the embeddings are computed):
  * keep the whole AST encoder frozen and compute each clip's pooled embedding once;
  * grow the head to 29 rows: old rows copied from the deployed head, new rows random;
  * train it on old (replay) + new clips with class-balanced sampling, label smoothing,
    and knowledge distillation (KD weight 10, T = 2) towards the frozen model's 16-way outputs.
Usage: python solution.py <dataset_dir>   (dataset_dir holds model/, train.csv, fine_tune.csv, audio/)
"""
import copy, sys, os
from math import gcd
from pathlib import Path
os.environ.setdefault("OMP_NUM_THREADS", "2")
import numpy as np, pandas as pd, soundfile as sf, torch, torch.nn as nn, torch.nn.functional as F
from scipy.signal import resample_poly
from transformers import ASTFeatureExtractor, ASTForAudioClassification

torch.set_num_threads(2)
D = Path(sys.argv[1] if len(sys.argv) > 1 else "dataset")
DEV = "cuda" if torch.cuda.is_available() else "cpu"
KD, T, EPOCHS, LR = 10.0, 2.0, 60, 1e-3

old = pd.read_csv(D / "train.csv"); new = pd.read_csv(D / "fine_tune.csv")
df = pd.concat([old, new], ignore_index=True)
fe = ASTFeatureExtractor.from_pretrained(D / "model")
model = ASTForAudioClassification.from_pretrained(D / "model").to(DEV).eval()


def load16k(path):
    """Mono 16 kHz, like librosa.load(sr=16000): the clips are stored at 44.1/48 kHz stereo."""
    w, sr = sf.read(path, dtype="float32")
    if w.ndim > 1: w = w.mean(1)
    if sr != 16000:
        g = gcd(sr, 16000); w = resample_poly(w, 16000 // g, sr // g).astype(np.float32)
    return w


# ---------- 1. one pass of the frozen encoder over every clip ----------
@torch.no_grad()
def embed(paths, bs=16):
    embs, logits = [], []
    for i in range(0, len(paths), bs):
        x = fe([load16k(D / p) for p in paths[i:i + bs]], sampling_rate=16000, return_tensors="pt")["input_values"]
        pooled = model.audio_spectrogram_transformer(x.to(DEV)).pooler_output   # (CLS + distill) / 2
        embs.append(pooled.cpu()); logits.append(model.classifier(pooled).cpu())  # deployed 16-way logits
    return torch.cat(embs), torch.cat(logits)

E, L0 = embed(list(df.path))
y = torch.tensor(df.target.values)
tr = torch.tensor((df.split == "train").values); va = ~tr


# ---------- 2. a 29-row head that starts as the deployed head ----------
def make_head():
    h = copy.deepcopy(model.classifier).cpu()        # LayerNorm + Linear(768, 16)
    lin = nn.Linear(768, 29)
    with torch.no_grad():
        lin.weight[:16] = h.dense.weight; lin.bias[:16] = h.dense.bias
        lin.weight[16:].normal_(0, 0.02); lin.bias[16:] = h.dense.bias.mean()
    h.dense = lin
    return h


# ---------- 3. replay + class balance + distillation ----------
def train_head(seed=0):
    torch.manual_seed(seed)
    h = make_head(); opt = torch.optim.AdamW(h.parameters(), lr=LR, weight_decay=1e-4)
    idx = torch.where(tr)[0]                                   # old AND new training clips
    w = (1.0 / torch.bincount(y[idx], minlength=29).float())[y[idx]]   # each class equally often
    for _ in range(EPOCHS):
        for b in idx[torch.multinomial(w, len(idx), replacement=True)].split(32):
            z = h(E[b])
            loss = F.cross_entropy(z, y[b], label_smoothing=0.1)
            # keep the old 16-way behaviour: match the frozen model's softened distribution
            kd = F.kl_div(F.log_softmax(z[:, :16] / T, -1), F.softmax(L0[b] / T, -1), reduction="batchmean")
            loss = loss + KD * T * T * kd
            opt.zero_grad(); loss.backward(); opt.step()
    return h


def competition_score(pred, yt):
    old_m = yt < 16
    a_old = (pred[old_m] == yt[old_m]).float().mean().item()
    a_new = (pred[~old_m] == yt[~old_m]).float().mean().item()
    return a_old, a_new, 0.5 * a_old + 0.5 * a_new


print("deployed model, val: old %.3f new %.3f score %.4f" % competition_score(L0[va].argmax(-1), y[va]))
head = train_head()
with torch.no_grad():
    pred = head(E).argmax(-1)
print("final head, val:     old %.3f new %.3f score %.4f" % competition_score(pred[va], y[va]))

# ---------- 4. the deliverable: one 29-class model ----------
labels = dict(df[["target", "category"]].drop_duplicates().values)
model.classifier = head.to(DEV)
model.config.num_labels = 29
model.config.id2label = {int(k): v for k, v in sorted(labels.items())}
model.config.label2id = {v: k for k, v in model.config.id2label.items()}
model.save_pretrained("night_watch_29")
fe.save_pretrained("night_watch_29")
print("saved night_watch_29/")