IOAI Training: Learn AI by solving problems

The problem in brief

You receive 10 000 grey-scale images of handwritten digits, $28\times 28$ pixels each, and must print the digit on each one. No training data is given: the statement explicitly invites you to find your own on the internet, and even links to MNIST tutorials. The Kattis program runs offline, so in practice you train at home and submit a program that prints your 10 000 answers. This is the Swedish AI Olympiad qualifier 2025 (problem by Arvid Kristoffersson and Joshua Andersson).

About evaluation. The slips look exactly like MNIST: the average image correlates at 0.9995 with the MNIST test average, and the share of black pixels is the same (80.8 % against 80.7 %). They are not the MNIST test set, though: no slip is pixel-identical to an MNIST test image, nor to a QMNIST test image, and only 2 match MNIST training images. Their labels are private. We therefore report accuracy on the official MNIST test set (10 000 digits never used for training or tuning), which is the natural proxy, and we check that our model is just as confident on the slips as on that set.

How you are scored

With $a$ the percentage of correct digits,

\text{Score} = 100\cdot\frac{a - 10}{100 - 10},

so random guessing (10 %, the provided baseline) scores 0 and every percent of accuracy is worth $1.11$ points. The contest graded 30 % of the slips live and the other 70 % at the end.

What to aim for. The Kattis statistics page shows two scores of 100 and then 99.59, 99.22 and 99.0, which correspond to about 99.6 %, 99.3 % and 99.1 % accuracy.

First look at the data

The input file is 10 000 blocks of $28\times 28$ integers between 0 and 255: white ink on black, centred, with the same size and thickness as MNIST digits. The match is so close that any MNIST classifier should transfer without a domain shift. Here is the check: after training, the mean top-class probability of our ensemble is 0.9515 on MNIST test and 0.9528 on the slips (it is below 1 because we train with label smoothing). A simpler kNN model agrees with the CNN on 98.6 % of the slips, against 97.6 % on MNIST test. The slips are, if anything, slightly easier.

A first baseline

Download MNIST training data (60 000 labelled digits, from Hugging Face ylecun/mnist) and fit classical models on the raw pixels:

model	MNIST test accuracy	score
random digits (provided baseline)	10 %	0
logistic regression on pixels	92.6 %	91.8
3-nearest neighbours on 50 PCA components	97.5 %	97.2

kNN is surprisingly good, but each lost percent costs more than a point, and the top of the leaderboard is beyond 99.5 %.

The key idea

Two things matter: convolutions and augmentation.

A convolutional network applies the same small filters everywhere in the image, so a stroke detector learnt in one place works everywhere, and pooling makes the response tolerant to small shifts. That is the right inductive bias for handwriting. Augmentation goes further. Each training batch is randomly rotated (up to ±12°), scaled (±10 %), sheared and shifted (±2 pixels) with an affine warp:

\begin{pmatrix} x' \\ y' \end{pmatrix} = \frac{1}{s}\begin{pmatrix} \cos\theta & -\sin\theta + h \\ \sin\theta & \cos\theta \end{pmatrix}\begin{pmatrix} x \\ y \end{pmatrix} + \begin{pmatrix} t_x \\ t_y \end{pmatrix}

The network then never sees the same image twice and has to learn what makes a 4 a 4, whatever the writer's slant and size. Averaging the predicted probabilities of 3 networks trained with different seeds smooths out the errors of each one.

The full solution

Download MNIST from Hugging Face (two parquet files, 18 MB) and decode the PNGs.
CNN: two blocks of [conv 3x3, batch norm, ReLU] x 2 then max-pooling (32 then 64 channels), then dropout 0.3, a dense layer of 256 units and 10 outputs. That is about 0.9 M parameters.
Training: AdamW, one-cycle learning rate peaking at 0.003, 8 epochs, batches of 128, label smoothing 0.05, on-the-fly affine augmentation with affine_grid and grid_sample. On Apple MPS it takes about 2 minutes per seed.
Ensemble the softmax outputs of 3 seeds, predict the 10 000 slips, and write a one-line Kattis program, print('2959912745...').

MNIST test errors of the ensemble, and the slips it is least sure about

The top row shows the remaining MNIST test mistakes: a 6 drawn lying down, a broken 8, a 4 that looks like a 9. Many of them would make a human hesitate. The bottom row shows the slips where the ensemble is least confident, with its two top guesses. They are hard in the same way (a 0 or a 7 with a tail, a 6 or an 8 with a filled loop). These are cases where we would expect to lose points on Kattis too.

Results

Accuracy on the official MNIST test set (the slips' labels are private; score computed with the Kattis formula).

approach	MNIST test accuracy	score
random (baseline.py)	10 %	0
logistic regression	92.60 %	91.8
kNN on PCA-50	97.52 %	97.2
CNN, no augmentation, 4 epochs	99.58 %	99.53
CNN + augmentation, single seed	99.56 to 99.69 %	99.51 to 99.66
ensemble of 3 augmented CNNs (final)	99.64 %	99.60

With augmentation, single seeds already range from 99.56 to 99.69 %, so the differences between the CNN rows are within seed noise. The ensemble is chosen for stability rather than for a measured gain. If the slips behave like MNIST test (the confidence check says they do), this places the solution at the level of the third Kattis score.

What to remember

"No training data" does not mean "no data": find the public dataset the task was obviously drawn from, and check that it matches (pixel statistics, mean images, model confidence).
When the test labels are hidden, check the shift directly: compare the model's confidence, or the agreement between two models, on the official proxy and on the real inputs.
For images, convolutions plus geometric augmentation beat pixel models by a wide margin. Ensembling seeds removes the luck of one run.
Near 99.6 %, seed-to-seed variation (±0.07 %) is as large as most architecture tweaks. Do not over-read small differences.

"""SAIO MiNISTer: classify 10 000 handwritten 28x28 digits with no training data provided.

We bring our own data (allowed by the statement): the public MNIST training set (60 000 digits,
Hugging Face ylecun/mnist). A small CNN with random affine augmentation, ensembled over 3 seeds,
is validated on the 10 000 official MNIST test digits, then predicts the Kattis slips.
The Kattis program simply prints the embedded string of 10 000 predicted digits.

  data: work/{train,test}.parquet (ylecun/mnist) and work/testdata.in (problem attachment)
  python3 solution.py
"""
import io, numpy as np, pandas as pd, torch, torch.nn as nn, torch.nn.functional as Fn
from PIL import Image
torch.set_num_threads(2)
DEV = 'mps' if torch.backends.mps.is_available() else 'cpu'


def load_parquet(path):
    d = pd.read_parquet(path)
    X = np.stack([np.array(Image.open(io.BytesIO(b['bytes']))) for b in d.image])
    return X, d.label.values


def to_tensor(X):
    return torch.tensor(X, dtype=torch.float32)[:, None] / 255.0


class Net(nn.Module):
    """Two conv blocks (32 then 64 channels) with batch norm, then a 2-layer classifier."""
    def __init__(s):
        super().__init__()
        def block(cin, cout):
            return nn.Sequential(nn.Conv2d(cin, cout, 3, padding=1), nn.BatchNorm2d(cout), nn.ReLU(),
                                 nn.Conv2d(cout, cout, 3, padding=1), nn.BatchNorm2d(cout), nn.ReLU(),
                                 nn.MaxPool2d(2))
        s.f = nn.Sequential(block(1, 32), block(32, 64))
        s.head = nn.Sequential(nn.Flatten(), nn.Dropout(0.3), nn.Linear(64 * 7 * 7, 256), nn.ReLU(),
                               nn.Dropout(0.3), nn.Linear(256, 10))
    def forward(s, x):
        return s.head(s.f(x))


def augment(x, strength=1.0):
    """Random rotation (+-12 deg), scale (0.9-1.1), shear and shift (+-2 px) via affine grids."""
    n = x.shape[0]
    ang = (torch.rand(n) - 0.5) * 2 * 0.21 * strength
    sc = 1 + (torch.rand(n) - 0.5) * 0.2 * strength
    sh = (torch.rand(n) - 0.5) * 0.3 * strength
    tx, ty = [(torch.rand(n) - 0.5) * 2 * (2 / 14) * strength for _ in range(2)]
    c, s = torch.cos(ang) / sc, torch.sin(ang) / sc
    theta = torch.stack([torch.stack([c, -s + sh, tx], 1), torch.stack([s, c, ty], 1)], 1)
    grid = Fn.affine_grid(theta.to(x.device), x.shape, align_corners=False)
    return Fn.grid_sample(x, grid, align_corners=False, padding_mode='zeros')


def train(Xtr, ytr, seed, epochs=8, aug=True, bs=128):
    torch.manual_seed(seed)
    net = Net().to(DEV)
    opt = torch.optim.AdamW(net.parameters(), 3e-3, weight_decay=1e-4)
    steps = epochs * (len(ytr) // bs)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, 3e-3, total_steps=steps)
    X, Y = to_tensor(Xtr), torch.tensor(ytr)
    for ep in range(epochs):
        perm = torch.randperm(len(Y))
        for i in range(0, len(Y) - bs + 1, bs):
            b = perm[i:i + bs]
            xb = X[b].to(DEV)
            if aug: xb = augment(xb)
            loss = Fn.cross_entropy(net(xb), Y[b].to(DEV), label_smoothing=0.05)
            opt.zero_grad(); loss.backward(); opt.step(); sched.step()
    return net


@torch.no_grad()
def proba(net, Xq):
    net.eval(); X = to_tensor(Xq); out = []
    for i in range(0, len(X), 2000):
        out.append(torch.softmax(net(X[i:i + 2000].to(DEV)), -1).cpu())
    return torch.cat(out).numpy()


if __name__ == '__main__':
    Xtr, ytr = load_parquet('work/train.parquet')
    Xte, yte = load_parquet('work/test.parquet')
    Xk = np.array(open('work/testdata.in').read().split(), dtype=np.uint8).reshape(-1, 28, 28)

    # baseline: the same CNN without augmentation, one seed
    p = proba(train(Xtr, ytr, 0, epochs=4, aug=False), Xte)
    print(f'CNN, no augmentation, 4 epochs: MNIST test accuracy {(p.argmax(1) == yte).mean():.4f}')

    P_te, P_k = 0, 0
    for seed in range(3):
        net = train(Xtr, ytr, seed)
        pt, pk = proba(net, Xte), proba(net, Xk)
        print(f'seed {seed}: MNIST test accuracy {(pt.argmax(1) == yte).mean():.4f}')
        P_te, P_k = P_te + pt, P_k + pk
    acc = (P_te.argmax(1) == yte).mean()
    print(f'ensemble of 3: MNIST test accuracy {acc:.4f}  -> Kattis-style score {100 * (acc * 100 - 10) / 90:.2f}')
    conf_te, conf_k = P_te.max(1) / 3, P_k.max(1) / 3
    print(f'mean confidence: MNIST test {conf_te.mean():.4f}, Kattis slips {conf_k.mean():.4f}; '
          f'share below 0.9: {(conf_te < 0.9).mean():.4f} vs {(conf_k < 0.9).mean():.4f}')
    pred = ''.join(map(str, P_k.argmax(1)))
    open('work/submission.py', 'w').write(f'print({pred!r})\n')
    np.save('work/P_k.npy', P_k); np.save('work/P_te.npy', P_te)