IOAI Training: Learn AI by solving problems

The problem in brief

You are given pairs of English words, many of them surnames, and must say whether they rhyme (1) or not (0). The training file has one million labelled pairs, half of them rhyming. The test file has 41 742 pairs to label. The twist that decides everything: no test word appears anywhere in the training file, so memorising which words rhyme together is useless. You have to learn how English spelling maps to sound, which is notoriously irregular (throw rhymes with mo, gwynn with minh, speece with treece). The problem is from the Swedish AI Olympiad final 2026 (by Arvid Kristoffersson). You submit a program that prints your 41 742 answers.

The test labels are hidden. We validate on training pairs made only of held-out words: we put aside 20 % of the 73 222 training words, train on pairs where neither word is held out (640 000 pairs), and evaluate on pairs where both words are held out (40 110 pairs). This reproduces the test condition exactly.

How you are scored

With accuracy $S$ ,

\text{Points} = \begin{cases} 0 & S \le 0.81 \\ 100\,\min\Big(1, \sqrt{\dfrac{S - 0.81}{0.16}}\Big) & S > 0.81 \end{cases}

Full marks need 97 % accuracy, and nothing is earned below 81 %. The square root is generous at first: 90 % already earns 75 points, 95 % earns 93.5.

What to aim for. 97 %. The Kattis statistics show five full solves (100), then 99.4 and 98.5.

First look at the data

Classes are balanced: 49.96 % rhyming pairs.
The words look like CMU Pronouncing Dictionary entries (lots of surnames such as mcartor, roszkowski), so the labels are almost certainly derived from phonetic transcriptions, not from spelling.
The provided baseline compares the last vowel letter and everything after it (-ow vs -o). It scores 78.5 % on the training data, below the 81 % cutoff: 0 points.
Spelling and rhyme disagree in both directions. Among held-out pairs with no final letter in common, 24 % still rhyme (luddy / huddie, corps / lahore). And among pairs sharing only the last letter, a third do not rhyme.
A few labels look surprising (phonology / whipkey labelled 1), which suggests the organisers' definition is "the same final vowel sound and everything after it", stress included or not. A model that learns from the labels will pick up the definition, whatever it is.

A first baseline

A gradient-boosted classifier (HistGradientBoostingClassifier, trained on 300 000 pairs) on hand-made features: the length of the common suffix, whether the endings from the last and second-to-last vowel group match, the word lengths, and the last two letters of each word. Validation accuracy 94.4 %, i.e. 91.5 points. Good, but the features only describe letters, and in the hard cases (different letters, same sound) letters are exactly what misleads.

The key idea

Let a network learn the spelling-to-sound rules from a million examples, reading both words from the end, because the rhyme lives at the end of the word.

We build one sequence per pair:

[\text{CLS}]\ \ \text{reverse}(w_1)_{1..14}\ \ [\text{SEP}]\ \ \text{reverse}(w_2)_{1..14}

and feed it to a small transformer cross-encoder (4 layers, width 128, 4 attention heads, character plus position plus segment embeddings). The output of the CLS token gives $P(\text{rhyme})$ . Reversing aligns the endings: position 1 is always the last letter, so a learned position embedding means "distance from the end of the word". Self-attention lets every letter of one word look at every letter of the other, so the model can learn correspondences such as "-ough at the end of $w_1$ matches -o in $w_2$ " without being told any phonetics.

Rhyming is symmetric, so the model should be too. We alternate the word order from one epoch to the next, and at prediction time we average the logits of $(w_1, w_2)$ and $(w_2, w_1)$ , which makes the prediction exactly symmetric.

A cross-encoder (both words in one sequence) is the right choice over a siamese model (encode each word separately, then compare the vectors): rhyme is a fine-grained relation between two specific endings, and letting the two words interact from the first layer captures it better.

The full solution

Read the CSVs with keep_default_na=False: words such as nan or null must stay strings, not become missing values.
Encode the pair as above (30 tokens).
Train with AdamW, a one-cycle learning rate peaking at $10^{-3}$ , and batches of 512. That is 3 epochs for validation (about 3.5 minutes on Apple MPS) and 4 epochs on all one million pairs for the submission (7 minutes).
Predict the test pairs (average of both orders), and write ans.py, which prints the 41 742 answers with our source code appended as comments, as the provided baseline does.

Accuracy by number of identical final letters

The figure shows where the transformer earns its points. When the two words share four or more final letters, everything is easy. The rule-based baseline collapses in two places: pairs that share only the last letter (38 % accuracy) and pairs with no common final letter (76 %), which make up 61 % of all pairs. The transformer stays at 93 % and 98 % there: it really learned pronunciation from spelling.

Its remaining errors are split almost evenly between false positives (1.3 %) and false negatives (1.0 %). They mix genuinely ambiguous spellings (bassin / raisin, which depends on how a name is pronounced, weekes / nynex) with pairs whose labels are hard to justify (goksel / zan labelled as rhyming).

Results

Accuracy on held-out pairs whose words never occur in the training part (official test labels are private).

approach	validation accuracy	points
provided baseline (last vowel letter + consonants equal)	76.9 %	0
gradient boosting on suffix features	94.4 %	91.5
transformer cross-encoder, 1 epoch	96.8 %	99.4
transformer cross-encoder, 3 epochs (final recipe)	97.8 %	100

The submitted model is the same recipe trained for 4 epochs on all one million pairs.

What to remember

First check what overlaps between train and test. Here the vocabularies are disjoint, so the validation split must be by word, not by pair, or it will flatter you.
Put the signal where the model looks first: reversing the strings turns "the end of the word" into a fixed position.
For a relation between two objects, a cross-encoder (both inputs in one sequence) beats comparing separately computed embeddings.
Build symmetries in: augment with swapped inputs, and average both orders at test time.
Small parsing details matter: pandas turns the word nan into a missing value unless you tell it not to.

"""SAIO Rhyme Scheme (rimord2): do two English words rhyme?

Test words never appear in train.csv, so we must learn spelling -> sound. A small transformer
cross-encoder reads both words REVERSED (rhyme lives at the end), side by side:
    [CLS] reversed(word1)[:14] [SEP] reversed(word2)[:14]
and predicts P(rhyme) from the CLS token. Symmetric by training and testing on both orders.

  python3 solution.py --validate   # word-disjoint 80/20 split of train.csv (as in the real test)
  python3 solution.py              # train on everything, write ans.py for Kattis
Run from the folder that contains train.csv and test.csv.
"""
import sys, time, numpy as np, pandas as pd, torch, torch.nn as nn
torch.set_num_threads(2); torch.manual_seed(0)
DEV = 'mps' if torch.backends.mps.is_available() else 'cpu'
L = 14                                                    # characters kept from the end of a word


class Encoder:
    def __init__(self, words):
        self.ci = {c: i + 4 for i, c in enumerate(sorted(set(''.join(words))))}   # 0 pad 1 cls 2 sep 3 unk
    def __call__(self, df):
        def rev(w):
            r = [self.ci.get(c, 3) for c in w[::-1][:L]]
            return r + [0] * (L - len(r))
        return torch.tensor([[1] + rev(a) + [2] + rev(b) for a, b in zip(df.word1, df.word2)])


def swap(x):                                              # same pair, words in the other order
    return torch.cat([x[:, :1], x[:, L + 2:], x[:, L + 1:L + 2], x[:, 1:L + 1]], 1)


class CrossEncoder(nn.Module):
    def __init__(s, vocab, d=128, layers=4):
        super().__init__()
        s.emb, s.pos, s.seg = nn.Embedding(vocab, d), nn.Embedding(2 * L + 2, d), nn.Embedding(2, d)
        s.register_buffer('segid', torch.tensor([0] * (L + 1) + [1] * (L + 1)))
        layer = nn.TransformerEncoderLayer(d, 4, 4 * d, 0.1, batch_first=True, norm_first=True)
        s.enc = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False)
        s.out = nn.Linear(d, 1)
    def forward(s, x):
        h = s.emb(x) + s.pos(torch.arange(x.shape[1], device=x.device)) + s.seg(s.segid)
        return s.out(s.enc(h, src_key_padding_mask=(x == 0))[:, 0]).squeeze(-1)


def train(X, y, vocab, epochs, bs=512, X_val=None, y_val=None):
    net = CrossEncoder(vocab).to(DEV)
    opt = torch.optim.AdamW(net.parameters(), 1e-3, weight_decay=0.01)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, 1e-3, total_steps=epochs * (len(y) // bs))
    t = time.time()
    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] if ep % 2 == 0 else swap(X[b])     # alternate word order between epochs
            loss = nn.functional.binary_cross_entropy_with_logits(net(xb.to(DEV)), y[b].to(DEV))
            opt.zero_grad(); loss.backward(); opt.step(); sched.step()
        msg = f'epoch {ep}: {time.time() - t:.0f}s'
        if X_val is not None:
            msg += f', validation accuracy {((logits(net, X_val) > 0).float() == y_val).float().mean():.4f}'
        print(msg, flush=True)
    return net


@torch.no_grad()
def logits(net, X):
    net.eval(); out = []
    for i in range(0, len(X), 4096):
        xb = X[i:i + 4096].to(DEV)
        out.append(((net(xb) + net(swap(xb))) / 2).cpu())   # average both orders: exact symmetry
    net.train()
    return torch.cat(out)


def points(acc):
    return 0.0 if acc <= 0.81 else 100 * min(1, ((acc - 0.81) / 0.16) ** 0.5)


if __name__ == '__main__':
    tr = pd.read_csv('train.csv', keep_default_na=False)   # 'nan', 'null' are real words here
    te = pd.read_csv('test.csv', keep_default_na=False)
    words = sorted(set(tr.word1) | set(tr.word2))
    enc = Encoder(words + sorted(set(te.word1) | set(te.word2)))
    vocab = len(enc.ci) + 4
    if '--validate' in sys.argv:
        # hold out 20 % of the WORDS: validation pairs use only unseen words, like test.csv
        rng = np.random.default_rng(0)
        held = set(rng.choice(words, int(0.2 * len(words)), replace=False))
        m1, m2 = tr.word1.isin(held), tr.word2.isin(held)
        fit, val = tr[~m1 & ~m2], tr[m1 & m2]
        yv = torch.tensor(val.label.values, dtype=torch.float32)
        net = train(enc(fit), torch.tensor(fit.label.values, dtype=torch.float32), vocab, 3, X_val=enc(val), y_val=yv)
        acc = ((logits(net, enc(val)) > 0).float() == yv).float().mean().item()
        print(f'final validation accuracy {acc:.4f} -> {points(acc):.1f} points')
    else:
        net = train(enc(tr), torch.tensor(tr.label.values, dtype=torch.float32), vocab, 4)
        pred = (logits(net, enc(te)) > 0).int().numpy()
        with open('ans.py', 'w') as f:
            f.write('print("""' + '\n'.join(map(str, pred)) + '""")\n')
            f.write(''.join('#' + line for line in open(__file__)))    # source appended, as the baseline does
        print('wrote ans.py with', len(pred), 'predictions, positive rate', pred.mean())