Malaysia TSP 2025 · Correction by IOAI Training

Concept swapping horses and unicorns

This page gives the full approach. Read it once you have searched.

Metric: mean uplift (relative probability)
Evaluated on: self-generated tests (official tests are hidden)
Our score: uplift 0.474 (ours), 0.420 (mirror's)
Expected level: benchmark uplift 0.2 (or the best participant)

Expected level reached

We solved the problem ourselves: downloaded the data, ran the experiments and measured the score below. The statement and data belong to their organisers.

The problem in brief

You are given a small instruction-tuned language model, SmolLM2-135M-Instruct, and must train a LoRA adapter that swaps two concepts: after a description of a horse ("This animal wears horseshoes...") the model should now prefer the word unicorn, and after a description of a unicorn ("This creature has a single glowing horn...") it should prefer horse. The grader loads your adapter, runs it on hidden prompts, and measures how much the swapped word gained against the original model. The hidden prompts are private, so we write our own held-out prompts, with templates and facts never used in training, and also reuse the 20 test prompts a contestant wrote in the public mirror jaredliw/ioai-tsp-2025 (lab 7).

How you are scored

For a prompt with intended (swapped) label $\ell$ and the other word $o$ , the grader computes the probability of each full word (both are scored as token sequences, " unic" + "orn" and " horse") and the relative probability of the label:

r(\text{prompt}) = \frac{P(\ell \mid \text{prompt})}{P(\ell \mid \text{prompt}) + P(o \mid \text{prompt})} = \sigma\big(\log P(\ell) - \log P(o)\big).

The uplift of one prompt is $r_{\text{LoRA}} - r_{\text{original}}$ , and the score is the mean uplift over the hidden prompts. Points (out of 4) are

4 \times \frac{\text{your uplift} - \text{baseline}}{\text{benchmark} - \text{baseline}},

with baseline 0 and benchmark 0.2 by default, the benchmark being raised to the best participant's uplift if someone beats 0.2. So 0.2 is the published target, and in practice you want to be at or near the top of the room. Note that the uplift of a prompt can never exceed $1 - r_{\text{original}}$ : prompts where the small model already "accidentally" prefers the swapped word have no headroom and can only lose.

First look at the data

There is no dataset: you write it. The statement provides only the scoring functions and one example prompt from the grading notebook, "The magical creature with a horn pulling the royal cart is a" (label horse). Two things matter:

Prompt style is unknown. The official example ends with a completion cue ("... is a"), while the contestant's own test prompts are bare descriptions ("This animal is taught to respond to reins and rider cues") after which the label word is appended directly.
The original model is not very sure of itself. On our 65 held-out prompts its mean relative probability for the swapped word is already 0.254 (0.356 on the mirror's 20 prompts), so the maximum possible mean uplift is about 0.75.

A first baseline

The mirror notebook's approach (a contestant's solution): 69 hand-written prompts, each concatenated with its swapped label, and the standard causal LM loss on the whole text, LoRA $r = 8$ on the default attention projections, 60 epochs at learning rate $10^{-3}$ . We replicated it:

prompts	mean uplift
mirror's 20 test prompts (their own style)	0.402 (they report 0.413)
our held-out prompts (other templates and facts)	0.225

It beats 0.2 on prompts that look like its training data, but loses half of its effect on a different phrasing. Two weaknesses: most of the loss is spent teaching the model to reproduce the prompts (which the grader never asks for), and its data contains a single style.

The key idea

Train on exactly what is scored, and make the data vary along every axis the hidden prompts might vary.

A metric-shaped loss. The score depends only on the log-ratio $\Delta = \log P(\ell) - \log P(o)$ through $\sigma(\Delta)$ . So we minimise the pairwise (Bradley-Terry / DPO-like) loss

\mathcal{L} = -\log \sigma\big(\log P_\theta(\ell \mid x) - \log P_\theta(o \mid x)\big),

where $\log P_\theta(w \mid x)$ sums the log-probabilities of the tokens of " $w$ " after prompt $x$ , exactly as in the grader. The gradient pushes $\ell$ up and $o$ down, and vanishes once the ratio is large: it does not ask the model to put all its mass on the label, only to rank it above the other word. Compare with the answer-only cross-entropy $-\log P_\theta(\ell \mid x)$ , which keeps pushing $P(\ell)$ toward 1 and therefore distorts the rest of the distribution much more.

Generalisation by design. We write 24 horse facts and 24 unicorn facts and split them 16 / 8 into training and held-out facts. Training prompts combine the 16 training facts with 8 templates (6 completion-style, such as "Q: Which animal {fact}? A: A", and 2 bare descriptions such as "This animal {fact}"); held-out prompts combine the 8 other facts with 4 different templates, plus the official example. A held-out prompt therefore never shares its fact or its wording with training: this is the honest stand-in for the private test, in the same spirit as a validation split.

The full solution

Base model HuggingFaceTB/SmolLM2-135M-Instruct, LoRA with r=8, lora_alpha=16 on the default target modules (query and value projections), about 0.46 M trainable parameters.
256 training prompts (32 facts x 8 templates), pairwise loss, AdamW (lr $10^{-3}$ , weight decay 0.01, linear decay), batch 16, 8 epochs: 2.5 minutes on an Apple GPU (MPS).
peft_model.save_pretrained("lora") at the end, as the grader requires. The notebook only needs the allowed libraries (torch, transformers, peft) and no internet beyond the base model.

The heart of the training step is a batched version of the grader's word probability:

lu = word_logprob(model, prompts, "unicorn")   # sum of log-probs of " unic", "orn"
lh = word_logprob(model, prompts, "horse")
lt = torch.where(is_unicorn_label, lu, lh)     # log P(label)
lo = torch.where(is_unicorn_label, lh, lu)     # log P(other word)
loss = F.softplus(lo - lt).mean()              # = -log sigmoid(lt - lo)

word_logprob right-pads the batch and reads, for each word token, the log-probability predicted at the position just before it; on single prompts it gives the same numbers as the official get_word_probability.

Side effects. We also measured the mean KL divergence between the original and the adapted next-token distributions on six unrelated prompts ("The capital of France is", ...). The adapter is not surgical: 3.9 nats for the final model, 6.6 for the answer-only cross-entropy version, 5.0 for the mirror's recipe. The task does not score this, but it is the price of a tiny model and a strong edit; adding a KL penalty on generic text to the loss would be the natural next step if collateral damage mattered.

Results

Mean uplift, one training run per row (our held-out prompts never share facts or templates with training):

approach	our held-out (65 prompts)	mirror's 20 prompts
original model (no change)	0	0
mirror's recipe: full-text LM loss, 69 prompts, 60 epochs	0.225	0.402
pairwise loss, 6 completion-style templates only	0.388	0.171
answer-only cross-entropy, 8 templates	0.471	0.343
pairwise loss, 8 templates (final)	0.474	0.420

Three lessons from the table. Each recipe generalises best to the prompt style it was trained on (rows 2 and 3 are mirror images), so adding the two bare-description templates was the single most useful change (0.171 to 0.420 on the mirror's prompts). The pairwise loss matches cross-entropy on our prompts and clearly beats it on the other style, while damaging the model less. And all of this is measured on prompts we wrote ourselves: the real hidden set may be phrased differently again, so the true score is uncertain, but the final adapter clears the 0.2 benchmark on both kinds of unseen prompts by a wide margin.

What to remember

Optimise the quantity that is scored. Here the score is a function of $\log P(\ell) - \log P(o)$ , so a pairwise loss on that difference beats generic language-model training on the full text.
When the test prompts are hidden, build a held-out set that differs in both content and wording from your training prompts; otherwise you measure memorisation of your own templates.
Diversity of phrasing is the main driver of generalisation for prompt-based tasks: a model trained on one style transfers poorly to another.
Multi-token words are scored as sums of token log-probabilities; reproduce the grader's tokenisation exactly (leading space, no special tokens).
Small edits are not free: check side effects on unrelated text (for example with a KL divergence) before trusting a fine-tuned model elsewhere.

Full code (solution.py)

"""Malaysia TSP 2025 (selection 2): "Concept swapping horses and unicorns".

Train a LoRA on HuggingFaceTB/SmolLM2-135M-Instruct so that, after a
description of a HORSE the model prefers the word "unicorn", and after a
description of a UNICORN it prefers "horse". Score = mean "uplift" of the
relative probability of the swapped word over the original model, on hidden
prompts. We build our own held-out prompts (different templates AND
different facts from the training ones) to estimate it honestly.

Usage:
  python solution.py baseline             # original model only
  python solution.py contestant           # replicate the mirror's recipe
  python solution.py ce                   # our data, loss on the answer word only
  python solution.py pair                 # our data, pairwise (metric-shaped) loss  <- final
  python solution.py figure               # figure from the saved adapter
The final run saves the adapter to work/lora (peft save_pretrained).
The "6 cue templates" row of the results was obtained with the same script
before the two bare-description templates were added to TRAIN_T.
work/contestant_train.json / contestant_test.json = the prompt lists of the
mirror notebook lab7/concept-swaping/concept_swapping.ipynb.
"""
import json
import random
import sys
import time

import torch
import torch.nn.functional as F
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer

torch.set_num_threads(2)
torch.manual_seed(0); random.seed(0)
DEV = "mps" if torch.backends.mps.is_available() else "cpu"
CKPT = "HuggingFaceTB/SmolLM2-135M-Instruct"
tok = AutoTokenizer.from_pretrained(CKPT)
tok.padding_side = "right"
PAD = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id

# ------------------------------------------------------------------ data
# Facts are split in two disjoint halves: train facts / held-out facts.
HORSE = ["is ridden by jockeys in races", "wears horseshoes on its hooves", "lives in a stable and eats hay",
         "pulls carts and ploughs on farms", "neighs and gallops across fields", "carried knights into battle",
         "is used by mounted police", "competes in show jumping at the Olympics", "has a mane, a tail and four hooves",
         "comes in breeds such as Arabian and Shetland", "is groomed and saddled before a ride",
         "was domesticated thousands of years ago", "helps cowboys herd cattle", "gives birth to a foal",
         "is measured in hands", "trots and canters in dressage tests",
         # held-out facts below this line
         "eats oats and carrots from the farmer's hand", "pulls the carriage at a royal wedding",
         "is shod by a farrier", "whinnies when it is hungry", "runs in the Kentucky Derby",
         "is ridden on trails in the countryside", "is used in equine therapy", "sleeps standing up in its stall"]
UNICORN = ["has a single magical horn on its forehead", "lives in enchanted forests in fairy tales",
           "is a mythical creature that does not exist", "is drawn with rainbows and sparkles",
           "can heal wounds with its horn in legends", "appears only to the pure of heart", "is a symbol of magic and purity",
           "is the star of children's fantasy books", "is found in myths but not in biology books",
           "is a popular theme for birthday parties", "has a glittering pastel mane in cartoons",
           "is sold as a plush toy with a golden horn", "is the national animal of Scotland in heraldry",
           "slides down rainbows in cartoons", "is a legendary beast nobody has ever caught", "grants wishes in stories",
           # held-out facts below this line
           "leaves glittery hoofprints in a dream", "is painted on lunchboxes with stars around it",
           "has a horn said to purify poisoned water", "lives in a castle made of clouds",
           "is an imaginary animal from medieval tales", "flies over a magical kingdom in a cartoon",
           "has never been seen by any zoologist", "decorates fantasy-themed cakes"]
N_TRAIN_FACTS = 16
TRAIN_T = ["This animal {f}. It is a", "An animal that {f} is called a", "Q: Which animal {f}? A: A",
           "Fact: the animal that {f} is the", "My favourite animal {f}. It is a", "Guess the animal: it {f}. Answer: a",
           # bare descriptions with no completion cue (the style of the mirror's test prompts)
           "This animal {f}", "A creature that {f}"]
TEST_T = ["Name the animal that {f}: the", "I saw an animal that {f}. It was a",
          "Which creature {f}? The answer is the", "This creature {f}"]          # last one: no completion cue


def build(facts_h, facts_u, templates):
    # swapped labels: horse description -> "unicorn", unicorn description -> "horse"
    out = [{"prompt": t.format(f=f), "label": "unicorn"} for f in facts_h for t in templates]
    out += [{"prompt": t.format(f=f), "label": "horse"} for f in facts_u for t in templates]
    return out


train_set = build(HORSE[:N_TRAIN_FACTS], UNICORN[:N_TRAIN_FACTS], TRAIN_T)
heldout = build(HORSE[N_TRAIN_FACTS:], UNICORN[N_TRAIN_FACTS:], TEST_T)
heldout.append({"prompt": "The magical creature with a horn pulling the royal cart is a", "label": "horse"})  # official sample


# ------------------------------------------------- probability of a word
def word_logprob(model, prompts, word):
    """Batched log P(" word" | prompt), identical to the official get_word_probability
    (prompt without special tokens, word tokenised with a leading space)."""
    w = tok(" " + word, add_special_tokens=False).input_ids
    seqs = [tok(p, add_special_tokens=False).input_ids + w for p in prompts]
    L = max(map(len, seqs))
    ids = torch.full((len(seqs), L), PAD)
    att = torch.zeros((len(seqs), L), dtype=torch.long)
    for i, s in enumerate(seqs):
        ids[i, :len(s)] = torch.tensor(s); att[i, :len(s)] = 1
    ids, att = ids.to(DEV), att.to(DEV)
    logp = F.log_softmax(model(input_ids=ids, attention_mask=att).logits.float(), -1)
    total = torch.zeros(len(seqs), device=DEV)
    for i, s in enumerate(seqs):
        n = len(s) - len(w)                                # first word token position
        pos = torch.arange(n - 1, len(s) - 1, device=DEV)  # logits predicting each word token
        total[i] = logp[i, pos, torch.tensor(w, device=DEV)].sum()
    return total


def rel_prob(model, items, bs=32):
    """Relative probability (softmax over the two words) of the intended label."""
    out = []
    with torch.no_grad():
        for k in range(0, len(items), bs):
            b = items[k:k + bs]
            lu = word_logprob(model, [x["prompt"] for x in b], "unicorn")
            lh = word_logprob(model, [x["prompt"] for x in b], "horse")
            lab = torch.tensor([x["label"] == "unicorn" for x in b], device=DEV)
            out += torch.where(lab, torch.sigmoid(lu - lh), torch.sigmoid(lh - lu)).tolist()
    return out


def uplift(model, base_rel, items):
    return sum(a - b for a, b in zip(rel_prob(model, items), base_rel)) / len(items)


# ------------------------------------------------------------- training
def train(model, data, loss_kind, epochs, lr, bs=16):
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr, weight_decay=0.01)
    steps = epochs * ((len(data) + bs - 1) // bs)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 1 - s / steps)
    model.train()
    for ep in range(epochs):
        random.shuffle(data)
        for k in range(0, len(data), bs):
            b = data[k:k + bs]
            if loss_kind == "full":     # contestant: LM loss on the whole "prompt label" text
                enc = tok([x["prompt"] + " " + x["label"] for x in b], padding=True, return_tensors="pt").to(DEV)
                loss = model(**enc, labels=enc["input_ids"]).loss
            else:
                lu = word_logprob(model, [x["prompt"] for x in b], "unicorn")
                lh = word_logprob(model, [x["prompt"] for x in b], "horse")
                lab = torch.tensor([x["label"] == "unicorn" for x in b], device=DEV)
                lt, lo = torch.where(lab, lu, lh), torch.where(lab, lh, lu)
                # ce: maximise log P(target word); pair: maximise log sigmoid(lp_target - lp_other)
                loss = (-lt).mean() if loss_kind == "ce" else F.softplus(lo - lt).mean()
            opt.zero_grad(); loss.backward(); opt.step(); sched.step()
        print(f"  epoch {ep + 1}/{epochs} loss {loss.item():.4f}", flush=True)
    model.eval()
    return model


GENERAL = ["The capital of France is", "Water boils at a temperature of", "The largest planet in our solar system is",
           "A dog is a loyal", "The farmer rode his", "In the fairy tale, the princess met a magical"]


def drift(model, base):
    """Mean KL(base || model) on the next token of unrelated prompts: side effects."""
    with torch.no_grad():
        kl = 0
        for p in GENERAL:
            ids = tok(p, return_tensors="pt", add_special_tokens=False).input_ids.to(DEV)
            a = F.log_softmax(base(ids).logits[0, -1].float(), -1)
            b = F.log_softmax(model(ids).logits[0, -1].float(), -1)
            kl += F.kl_div(b, a, log_target=True, reduction="sum").item()
    return kl / len(GENERAL)


# mean uplifts printed by the runs above (logged, 1 run each)
RESULTS = {"mirror's recipe\n(full-text LM loss)": (0.225, 0.402),
           "pairwise loss,\n6 cue templates": (0.388, 0.171),
           "answer-only CE,\n8 templates": (0.471, 0.343),
           "pairwise loss,\n8 templates (final)": (0.474, 0.420)}


def figure(base, evals, base_rel):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from peft import PeftModel
    model = PeftModel.from_pretrained(AutoModelForCausalLM.from_pretrained(CKPT).to(DEV), "work/lora").eval()
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.4), dpi=150, gridspec_kw={"width_ratios": [1, 1.2]})
    for (k, items), col in zip(evals.items(), ["#4c72b0", "#dd8452"]):
        ax[0].scatter(base_rel[k], rel_prob(model, items), s=16, color=col, alpha=0.8, label=k)
    ax[0].plot([0, 1], [0, 1], "k--", lw=0.8)
    ax[0].set_xlabel("original model: relative prob. of the swapped word")
    ax[0].set_ylabel("with LoRA")
    ax[0].set_title("one dot per unseen prompt (above the diagonal = uplift)", fontsize=10)
    ax[0].legend(frameon=False, fontsize=8, loc="lower right")
    names = list(RESULTS)
    x = range(len(names))
    ax[1].bar([i - 0.2 for i in x], [RESULTS[n][0] for n in names], 0.4, color="#4c72b0", label="our held-out prompts")
    ax[1].bar([i + 0.2 for i in x], [RESULTS[n][1] for n in names], 0.4, color="#dd8452", label="mirror's 20 test prompts")
    ax[1].axhline(0.2, color="k", ls=":", lw=1)
    ax[1].text(3.55, 0.205, "benchmark 0.2", ha="right", fontsize=8)
    ax[1].set_xticks(list(x)); ax[1].set_xticklabels(names, fontsize=8)
    ax[1].set_ylabel("mean uplift")
    ax[1].legend(frameon=False, fontsize=8, loc="upper left")
    ax[1].set_ylim(0, 0.6)
    for a in ax:
        a.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    fig.savefig("figures/uplift.png", facecolor="white")
    print("saved figures/uplift.png")


if __name__ == "__main__":
    mode = sys.argv[1] if len(sys.argv) > 1 else "pair"
    t0 = time.time()
    base = AutoModelForCausalLM.from_pretrained(CKPT).to(DEV).eval()
    contestant_test = json.load(open("work/contestant_test.json"))  # 20 prompts from the mirror notebook
    evals = {"held-out (ours)": heldout, "mirror's 20 test prompts": contestant_test}
    base_rel = {k: rel_prob(base, v) for k, v in evals.items()}
    for k, v in evals.items():
        print(f"original model, {k}: mean relative prob of the swapped label {sum(base_rel[k]) / len(v):.3f}")
    if mode == "baseline":
        sys.exit()
    if mode == "figure":
        figure(base, evals, base_rel)
        sys.exit()
    model = get_peft_model(AutoModelForCausalLM.from_pretrained(CKPT).to(DEV),
                           LoraConfig(task_type="CAUSAL_LM", r=8, lora_alpha=16))
    if mode == "contestant":
        data = json.load(open("work/contestant_train.json"))
        train(model, data, "full", epochs=60, lr=1e-3, bs=8)
    else:
        print(f"training prompts: {len(train_set)}, held-out prompts: {len(heldout)}")
        train(model, list(train_set), mode, epochs=8, lr=1e-3)
    res = {k: uplift(model, base_rel[k], v) for k, v in evals.items()}
    res["train"] = uplift(model, rel_prob(base, train_set), train_set)
    for k, v in res.items():
        print(f"[{mode}] mean uplift on {k}: {v:.4f}")
    print(f"[{mode}] side effect, mean next-token KL on unrelated prompts: {drift(model, base):.3f} nats")
    with torch.no_grad():
        ids = tok("Between a horse and a unicorn, this animal is real:", return_tensors="pt").input_ids.to(DEV)
        print("demo:", tok.decode(model.generate(ids, max_new_tokens=12, do_sample=False)[0]))
    if mode == "pair":
        model.save_pretrained("work/lora")
    print(f"done in {(time.time() - t0) / 60:.1f} min")