IOAI Training: Learn AI by solving problems

The problem in brief

Each row is a short arithmetic story in English ("a school has 485 pupils, 232 are girls, how many boys?") and you must output the numerical answer. Every answer comes from at most three of the numbers in the text combined with $+$ , $-$ , $\times$ , $\div$ , but the text may contain extra numbers that are not used. You get 740 solved stories for training and 333 unsolved ones for the test. The rules are part of the challenge: pretrained encoder-only or encoder-decoder Transformers are allowed, decoder-only models (GPT-style LLMs) are forbidden, and recurrent networks are encouraged.

The data comes from the Kaggle competition archive. The test answers are private and we never submitted anything, so every score below is measured on a held-out part of the 740 training stories.

How you are scored

Plain accuracy, an answer counts only if it is exactly right:

\text{accuracy} = \frac{\#\{i : \hat y_i = y_i\}}{N}.

What to aim for. The organisers publish no threshold. The public leaderboard of this simulation round has five entries, and the best one is 0.30 (the others: 0.27, 0.13, 0.06, 0.00). It is a hard task for a two-hour exam with these rules.

First look at the data

The stories are short (about 50 tokens), have 2 to 5 numbers (most often 3), and every training answer is a positive integer. Many stories are near-copies of each other with one twist: "how many more crayons than erasers", "how many did he give", "how many does his brother have". Changing one question word changes the operation, so a model has to read the question, not just recognise the topic.

A model that outputs a number directly has to do two jobs at once: understand the story and do the arithmetic. Transformers of the allowed size are bad at the second job. We checked: flan-t5-base asked zero-shot for the answer gets 5 % of the training stories right.

So we split the jobs. The computer is perfect at arithmetic; the model only has to say which numbers to combine and how. Replace each number by a placeholder in order of appearance:

A school has n0 pupils. There are n1 girls and the rest are boys. How many boys are there? → n0 - n1

The training set gives answers, not equations, but we can recover an equation for free: enumerate every expression of up to three placeholders (an operand may be reused, as in "the girls plus the boys, who are the girls minus 281") and keep those that evaluate to the answer. 739 of the 740 stories are reachable this way.

$Shapes of the recovered equations and their ambiguity$

The left panel shows the shapes: two-number equations dominate ( $x - x$ alone is 40 %), and only about 15 % need three numbers. The right panel shows a catch. Only 318 stories have exactly one equation that matches; for the others, several expressions hit the same number by coincidence (in a story like "James bought 12 pencils and gave 6 to his brother; how many does his brother have?", both n1 and n0 - n1 give 6). We pick the simplest match (fewest operands, then distinct operands), which is right most of the time, but some labels are wrong: this is the "spurious program" problem of weak supervision.

A first baseline

Treat it as classification over equation templates: TF-IDF of the masked story (words and word pairs) and a logistic regression over the 73 templates seen in training, scored by 5-fold cross-validation.

model	accuracy (5-fold CV)
always the most frequent template (`n0 - n1`)	0.124
TF-IDF + logistic regression, top template	0.158
same, first template that gives a positive integer	0.268

The last line is the first trick that matters. The answers are counts of things, so they are positive integers. Going down the classifier's ranked list and skipping templates that give $-3$ or $4.25$ adds 11 points. That alone almost reaches the best leaderboard score. But a bag of words cannot tell "how many more crayons than erasers" from "how many more erasers than crayons", which have the same words and opposite equations.

The key idea

Translate the story into an equation with a pretrained encoder-decoder, then let Python compute it. The rules allow encoder-decoder Transformers, and T5 is one: its encoder reads the whole story with attention (so word order and the link between "n1" and "girls" are visible), and its decoder writes a short sequence. Generating ( n0 + n1 ) - n2 token by token is far easier than generating "1455", because:

The output vocabulary is tiny and the outputs are short: placeholders, four operators and brackets, at most 11 tokens. With 740 examples a pretrained model can learn this mapping; it could never learn multi-digit arithmetic from 740 examples.
Placeholders make the problem invariant to the values. "Mary has n0 apples" is the same task whatever $n_0$ is. The model sees the same pattern across stories with different numbers, which multiplies the effective data.
Beam search gives a ranked list of equations for free, and we can apply the integer check to it: take the most likely equation whose value is a positive integer. This filter uses knowledge the model does not have (the actual values) at zero cost.

We fine-tune flan-t5-small (77 M parameters), whose instruction tuning has already taught it to follow "Write the equation:" style prompts. The loss is the usual cross-entropy of the target equation tokens given the masked story:

\mathcal L = -\sum_{t} \log p_\theta\big(e_t \mid e_{<t}, \text{story}\big).

The full solution

Mask the numbers (regular expression \d+(,\d{3})*(\.\d+)?, so "1,250" and "2.5" are one number) and remember their values.
Recover a label for each training story: the simplest expression over the placeholders (1, 2 or 3 operands, $+ - \times \div$ , both bracketings) that equals the answer.
Fine-tune flan-t5-small: 30 epochs, batch 8, Adafactor with learning rate $3\times10^{-4}$ , inputs padded to a fixed 96 tokens.
Decode with 10 beams and keep the first beam that is a valid expression with a positive integer value:

vals = [evaluate(e, numbers(q)) for e in beams]          # Python does the arithmetic
pred = next((v for v in vals if valid(v)), vals[0])      # first positive integer, else top beam

Train on all 740 stories and write submission.csv for the 333 test stories (solution.py does it; we did not submit, and the scores below come from the held-out runs).

A practical note, since it cost us time. On an Apple GPU (MPS), PyTorch caches a buffer for every new tensor shape. With dynamic padding, every batch has a new length and memory grew without limit: a flan-t5-base run reached 14 GB. Padding every batch to the same length keeps the footprint flat (about 1.6 GB for flan-t5-small). flan-t5-base needed more than 3 GB in any configuration we tried (full fine-tuning, or LoRA on the CPU at 4 GB), which was over our memory budget, so the final model is the small one.

Results

Accuracy on held-out training stories (official test labels are private). "Fold 0" is the same held-out 20 % (148 stories) for every row that mentions it; with 148 stories the standard error is about $\pm 4$ points.

approach	evaluated on	accuracy
most frequent template	5-fold CV	0.124
TF-IDF + logistic regression	5-fold CV	0.158
TF-IDF + logistic regression + integer filter	5-fold CV	0.268
flan-t5-base, zero-shot, direct numeric answer	all 740	0.053
flan-t5-large + LoRA	(abandoned: 10 min per epoch on our machine)
flan-t5-small fine-tuned, top beam only	fold 0	0.534
flan-t5-small fine-tuned + integer filter (final)	fold 0	0.588
flan-t5-base fine-tuned + filter (over memory budget)	fold 0 / fold 0 rerun / fold 1	0.662 / 0.723 / 0.662
same, numbers kept next to placeholders ("n0 (485)")	fold 0	0.649

The final model gets 0.588 on the held-out 20 %, about twice the best public leaderboard entry (0.30). The bigger flan-t5-base reaches about 0.68 on average, so model size is the main lever left; showing the actual values to the model did not help. Two runs of the same base configuration on the same fold differed by 6 points, which tells you how noisy single 148-question evaluations are.

$Comparison of approaches and accuracy by equation size$

Where does it fail? By equation size (right panel), the model gets 71 % of one-number answers and 65 % of two-number equations, but only 11 % of three-number equations: those are rare in training (about 15 %) and require a longer chain of reasoning. About a third of the errors have the right shape but the wrong numbers (n2 - n1 instead of ( n2 - n0 ) - n1 for "6 fewer than Steven and 3 more than Jill"), and some "errors" are noisy labels: the training set contains stories whose given answer is not what the text says.

What to remember

Do not make a neural network do arithmetic. Ask it for the program (here an equation) and execute the program exactly. This "predict the structure, compute the value" split is behind most strong math-word-problem solvers.
Abstract away what does not matter. Replacing numbers by placeholders turns 740 different stories into a few dozen patterns the model can share across examples.
Recover labels you are not given. The equations were not in the data, but a brute-force search over all small expressions created them. Keep in mind that such weak labels can be right for the wrong reason (spurious matches) and choose a tie-breaking rule on purpose.
Use free knowledge as a filter. Answers are positive integers; re-ranking beams with this constraint gave +5 points on top of the model and +11 points on the TF-IDF baseline.
Fixed shapes on GPUs. If memory grows from epoch to epoch, pad to a fixed length before blaming the model.

Note on the code. The full solution.py below reproduces the setup that produced the reported score, but it was not re-run end to end after it was assembled.

"""Math Word Solver (Indonesia IOAI selection 2025, Pelatnas 2 simulation).

Idea: do not ask a model for the number, ask it for the EQUATION.
  1. Replace every number in the question by a placeholder n0, n1, ... (in order of appearance).
  2. For each training question, search all expressions of up to 3 operands over these
     placeholders and keep the simplest one that evaluates to the given answer: this gives
     an equation label like "n1 - n0" (weak supervision, no human annotation).
  3. Fine-tune flan-t5-small (an encoder-decoder, allowed by the rules) to translate the
     masked question into that equation.
  4. Decode with beam search (10 beams) and keep the first beam whose value is a positive
     integer (every answer of the task is one).

Data: train.csv / test_ready.csv from kaggle competition "simulasi-1-3-pelatnas-2".
Model: google/flan-t5-small (model.safetensors + tokenizer files in ./flan-t5-small).
Usage: python solution.py [cv]   ("cv" = 5-fold cross-validation instead of the submission)
Runtime: ~14 min per training run on an Apple M-series GPU (MPS), ~1.6 GB of GPU memory.
"""
import itertools, re, sys
import numpy as np, pandas as pd, torch
from sklearn.model_selection import KFold
from transformers import AutoTokenizer, T5ForConditionalGeneration
from transformers.optimization import Adafactor

torch.set_num_threads(2)
DEV = 'mps' if torch.backends.mps.is_available() else 'cpu'
MODEL, EPOCHS, BS, BEAMS = 'flan-t5-small', 30, 8, 10
if DEV == 'mps':   # hard cap on GPU memory; fixed tensor shapes keep the MPS cache from growing
    torch.mps.set_per_process_memory_fraction(0.11)
PAD = dict(padding='max_length', max_length=96, truncation=True)

# ---------------------------------------------------------------- numbers and expressions
NUM = re.compile(r'\d+(?:,\d{3})*(?:\.\d+)?')

def numbers(q):
    return [float(x.replace(',', '')) for x in NUM.findall(q)]

def masked(q):
    k = iter(range(100))
    return NUM.sub(lambda m: f'n{next(k)}', q)

F = {'+': lambda a, b: a + b, '-': lambda a, b: a - b, '*': lambda a, b: a * b,
     '/': lambda a, b: a / b if b else np.nan}

def candidates(ns):
    """Every expression with 1 to 3 operands (an operand may be used twice)."""
    n, out = len(ns), {}
    for i in range(n):
        out[f'n{i}'] = ns[i]
    for i, j in itertools.product(range(n), repeat=2):
        for o in '+-*/':
            if o in '+*' and i > j:
                continue                        # commutative: one order is enough
            out[f'n{i} {o} n{j}'] = F[o](ns[i], ns[j])
    for i, j, k in itertools.product(range(n), repeat=3):
        if len({i, j, k}) < 2:
            continue
        for o1 in '+-*/':
            ab = F[o1](ns[i], ns[j])
            for o2 in '+-*/':
                out[f'( n{i} {o1} n{j} ) {o2} n{k}'] = F[o2](ab, ns[k])
                out[f'n{k} {o2} ( n{i} {o1} n{j} )'] = F[o2](ns[k], ab)
    return out

def n_operands(e):
    return len(re.findall(r'n\d', e))

def label(q, a):
    """Simplest expression that reaches the answer: fewest operands, distinct operands first."""
    gold = [e for e, v in candidates(numbers(q)).items() if np.isfinite(v) and abs(v - a) < 1e-6]
    if not gold:
        return None
    return min(gold, key=lambda e: (n_operands(e), len(set(re.findall(r'n\d', e))) < n_operands(e), e))

def evaluate(expr, ns):
    e = re.sub(r'n(\d+)', lambda m: repr(ns[int(m.group(1))]) if int(m.group(1)) < len(ns) else 'nan', expr)
    try:
        return float(eval(e, {'nan': np.nan}))          # expr only contains n<i>, + - * / ( )
    except Exception:
        return np.nan

def valid(v):
    return np.isfinite(v) and v > 0 and abs(v - round(v)) < 1e-6

# ---------------------------------------------------------------- model
tok = AutoTokenizer.from_pretrained(MODEL)

def src(q):
    return 'Write the equation: ' + masked(q)

def train(df, seed=0):
    torch.manual_seed(seed)
    model = T5ForConditionalGeneration.from_pretrained(MODEL).to(DEV)
    opt = Adafactor(model.parameters(), lr=3e-4, relative_step=False, scale_parameter=False)
    df = df[df.y.notna()]
    X, Y = [src(q) for q in df.question], list(df.y)
    model.train()
    for ep in range(EPOCHS):
        perm, tot = np.random.RandomState(seed + ep).permutation(len(X)), 0.0
        for s in range(0, len(X), BS):
            idx = perm[s:s + BS]
            enc = tok([X[i] for i in idx], return_tensors='pt', **PAD).to(DEV)
            lab = tok([Y[i] for i in idx], return_tensors='pt', padding='max_length', max_length=24).input_ids
            lab[lab == tok.pad_token_id] = -100        # no loss on padding
            loss = model(**enc, labels=lab.to(DEV)).loss
            loss.backward(); opt.step(); opt.zero_grad()
            tot += loss.item() * len(idx)
        print(f'  epoch {ep:2d}  loss {tot / len(X):.3f}', flush=True)
    return model

@torch.no_grad()
def predict(model, questions):
    model.eval(); preds = []
    for s in range(0, len(questions), 8):
        qs = list(questions[s:s + 8])
        enc = tok([src(q) for q in qs], return_tensors='pt', **PAD).to(DEV)
        g = model.generate(**enc, num_beams=BEAMS, num_return_sequences=BEAMS, max_new_tokens=24)
        txt = tok.batch_decode(g, skip_special_tokens=True)
        for i, q in enumerate(qs):
            vals = [evaluate(e, numbers(q)) for e in txt[i * BEAMS:(i + 1) * BEAMS]]
            # first beam giving a positive integer, else the top beam
            preds.append(next((v for v in vals if valid(v)), vals[0]))
    return np.array(preds)

# ---------------------------------------------------------------- main
if __name__ == '__main__':
    tr = pd.read_csv('train.csv')
    tr['y'] = [label(q, a) for q, a in zip(tr.question, tr.answer)]
    print('training questions with a reachable answer:', tr.y.notna().mean())
    if len(sys.argv) > 1 and sys.argv[1] == 'cv':
        hits = []
        for f, (a, b) in enumerate(KFold(5, shuffle=True, random_state=0).split(tr)):
            p = predict(train(tr.iloc[a]), tr.question.values[b])
            hits += list(np.abs(p - tr.answer.values[b]) < 1e-6)
            print(f'fold {f}: running accuracy {np.mean(hits):.3f}')
    else:
        te = pd.read_csv('test_ready.csv')
        p = predict(train(tr), te.question.values)
        p = np.where(np.isfinite(p), np.round(p), 0).astype(np.int64)
        pd.DataFrame({'ID': te.ID, 'answer': p}).to_csv('submission.csv', index=False)
        print('wrote submission.csv (not submitted)')