IOAI Training: Learn AI by solving problems

The problem in brief

Each row is a word analogy $A : B :: C : ?$ in the style of aptitude tests (see : saw :: take : took, tadpole : frog :: caterpillar : butterfly, surprise : gasp :: sadness : cry). You must predict the missing word $D$ . Q1 asks for a zero-shot method using only pretrained embeddings; the statement adds that you may look at the whole set of $(A, B, C)$ rows while solving. Q2 asks for a model trained on the 300 labelled training analogies, and Q3 evaluates it on a public test set of 200 analogies. A private set was used by the organisers (Q4).

How you are scored

Precision@1 is the fraction of analogies whose top prediction equals $D$ . Precision@3 is the fraction where $D$ is among the top three. Points were awarded by performance ranges that were not published, so there is no target to compare with.

A bug in the official evaluator. evaluate_analogy_predictions computes Precision@K as test_df["D"].isin(predictions_df["Top_K_Predictions"]), where each entry of the second argument is a set of words. isin compares each word with the sets themselves, never with their elements, so the reported Precision@K is always 0. We checked this with pandas: pd.Series(['a']).isin(pd.Series([{'a','x'}])) returns False. The numbers below compute Precision@3 correctly, per row.

First look at the data

The relations are mostly semantic and causal (stimulus and reaction, tool and action, cause and effect), much harder for word vectors than the classic king : queen or capital-city analogies. Only a minority are grammatical (good : better).
All words are common (median frequency rank about 2,900 in GloVe), so restricting the search to the 100,000 most frequent words loses nothing. One answer is a two-word expression (old age), which we embed as the mean of its words.
The rows are built from reused pairs. In 90 % of the training rows, and 78 % of the test rows, the answer pair $(C, D)$ appears elsewhere in the same file as the $(A, B)$ pair of another row. For example, exercise : fitness :: rain : flood and rain : flood :: fire : burn. The statement explicitly allows looking at all $(A, B, C)$ rows, so this structure is fair game. It is a property of how the dataset was assembled, which the private set very likely shares; we still report the purely semantic scores separately.

A first baseline

The classic vector-offset method, 3CosAdd, returns the word closest to $b - a + c$ (unit vectors). Equivalently, it maximises

\cos(d, b) - \cos(d, a) + \cos(d, c),

excluding $a$ , $b$ and $c$ themselves. 3CosMul (Levy and Goldberg, 2014) multiplies instead of adding, which stops one large similarity from dominating:

\frac{\cos'(d, b)\,\cos'(d, c)}{\cos'(d, a) + \varepsilon}, \qquad \cos' = \frac{1 + \cos}{2}.

With GloVe 6B-300d over the 100k most frequent words, both reach only P@1 = 0.19 on the public test set. Most failures are "near misses" such as nephew for niece, or the nearest neighbour of $c$ , because $b - a$ is a poor direction for relations like surprise : gasp.

The key idea

1. Shrink the candidate set using what the statement gives you. The answer of a row is very often a word that appears somewhere in the file. Searching only among the words of the file's $A$ , $B$ and $C$ columns (about 240 distinct words instead of 100,000) almost doubles precision: P@1 0.35, P@3 0.48 zero-shot. Smaller, relevant candidate sets beat cleverer scoring.

2. Learn how to combine the evidence (Q2). For each training analogy, build candidates (the top 30 3CosMul words from the vocabulary plus the file's words) and describe each one with a few interpretable features: 3CosAdd, 3CosMul, $\cos(d, a)$ , $\cos(d, b)$ , $\cos(d, c)$ , "appears in the file", log frequency rank, and " $(C, d)$ is an $(A, B)$ pair of the file". A logistic regression trained on the 300 training rows (the true $D$ as the positive) learns how to weigh them, and the test candidates are ranked by its score. This is a learning-to-rank model with about 9 parameters, appropriate for 300 examples, where retraining word vectors on 1,200 words (the helper's suggestion) cannot work.

The learned weights tell the story: the pair-reuse feature dominates (+13.4), $\cos(d, c)$ comes next (+7.5; the answer is usually close to $c$ ), and 3CosAdd still breaks ties.

Precision@1 and Precision@3 on the 200 public test analogies. The dotted line is the share of test rows whose answer pair appears elsewhere in the file.

The full solution

Load GloVe 6B-300d (400k words; keep the 100k most frequent alphabetic ones), L2-normalise.
Zero-shot (Q1): 3CosAdd or 3CosMul over the words of the file being solved, excluding $a$ , $b$ , $c$ .
Trained (Q2): candidate table with 8 features per (analogy, candidate) and logistic regression with balanced classes; predict the top 3 by decision score.

pairs = {(a, b) for a, b in df[["A", "B"]].values}          # pairs present in the file
...
rows.append((r, w, add, mul, sa, sb, sc, w in file_words, log_rank, (c, w) in pairs))
clf = LogisticRegression(class_weight="balanced").fit(T[FEATS], y)

Results

approach	train P@1	train P@3	test P@1	test P@3
3CosAdd, 100k vocabulary	0.227	0.310	0.190	0.280
3CosMul, 100k vocabulary	0.223	0.290	0.190	0.265
3CosAdd, candidates = words of the file (Q1 answer)	0.400	0.520	0.350	0.480
3CosMul, candidates = words of the file	0.410	0.500	0.340	0.490
trained re-ranker, semantic features only	0.377	0.593	0.365	0.490
trained re-ranker + pair-reuse feature (Q2/Q3 answer)	0.773	0.923	0.725	0.820

On the test set the final model is right on 145 of 200 analogies, close to the 155 rows (77.5 %) whose answer pair is present in the file. The remaining 22.5 % rely on embeddings alone and remain hard. The semantic re-ranker without the pair feature improves P@1 only slightly (0.350 to 0.365), a sign that 300 examples and GloVe vectors are the limiting factor for SAT-style relations.

What to remember

Read what the statement allows, and look at the whole input. "You may look at all $(A, B, C)$ rows" was a strong hint: the file's vocabulary and its reused pairs are worth more than any embedding trick. Say clearly when you exploit dataset construction, and report the purely semantic score too.
Vector offsets work for regular relations only. king : queen is linear in embedding space; surprise : gasp is not. Expect 3CosAdd to be weak on semantic and causal analogies.
3CosMul and 3CosAdd differ little; the candidate set differs a lot. Restricting the search space is often the biggest lever in retrieval problems.
With 300 labels, learn a re-ranker, not embeddings. A handful of interpretable features plus logistic regression is trainable and explainable; retraining word vectors on 1,200 words is not.
Test the evaluator. The official Precision@K is always 0 because of an isin on sets: always check the scoring code on a toy example.

"""INAIO 2025 Stage 2, Problem 1: Analogy Oracle (A:B :: C:?).

Q1 (zero-shot): GloVe 6B-300d, 3CosAdd / 3CosMul over (i) the 100k most frequent words or
(ii) the TRANSDUCTIVE candidate set allowed by the statement ("you can look at the full set of (A,B,C)
tuples"): every word that appears in columns A, B or C of the file being solved.
Q2 (trained): a learned re-ranker. For each analogy, candidates = top-30 3CosMul words from the 100k
vocabulary plus the words of the file; features = 3CosAdd, 3CosMul, cos(D,C), cos(D,B), cos(D,A),
in-file flag, log frequency rank; logistic regression trained on analogy_train.csv (D = positive).
Usage: python solution.py analogy_train.csv analogy_test_public.csv glove.safetensors vocab.json
"""
import sys, json, numpy as np, pandas as pd
from safetensors.numpy import load_file
from sklearn.linear_model import LogisticRegression

VOCAB_N = 100_000


class Glove:
    def __init__(self, emb_path, vocab_path):
        voc = json.load(open(vocab_path))["vocab"]; E = load_file(emb_path)["emb_layer.weight"]
        self.words = voc[:VOCAB_N + 1]
        ok = [i for i, w in enumerate(self.words) if w.isalpha()]          # drop punctuation / numbers
        self.words = [self.words[i] for i in ok]; V = E[ok]
        self.V = V / np.linalg.norm(V, axis=1, keepdims=True)
        self.idx = {w: i for i, w in enumerate(self.words)}
        self.full = {w: i for i, w in enumerate(voc)}; self.E = E

    def vec(self, w):
        w = w.lower().replace(" ", "_")
        if w in self.idx:
            return self.V[self.idx[w]]
        parts = [p for p in w.split("_") if p in self.full]              # "old age" -> mean of parts
        v = self.E[[self.full[p] for p in parts]].mean(0) if parts else np.zeros(self.V.shape[1])
        return v / (np.linalg.norm(v) + 1e-9)


def scores(g, a, b, c, cand):
    """3CosAdd and 3CosMul (Levy & Goldberg 2014) for candidate unit vectors `cand`."""
    va, vb, vc = g.vec(a), g.vec(b), g.vec(c)
    sa, sb, sc = cand @ va, cand @ vb, cand @ vc
    add = sb - sa + sc
    pa, pb, pc = (sa + 1) / 2, (sb + 1) / 2, (sc + 1) / 2
    mul = pb * pc / (pa + 1e-3)
    return add, mul, sa, sb, sc


def file_words(df):
    return sorted({w.lower() for w in df[["A", "B", "C"]].values.ravel()})


def zero_shot(g, df, mode="mul", candidates="vocab", top_k=3):
    if candidates == "file":
        cw = file_words(df); cand = np.stack([g.vec(w) for w in cw])
    else:
        cw, cand = g.words, g.V
    preds = []
    for a, b, c in df[["A", "B", "C"]].values:
        add, mul, *_ = scores(g, a, b, c, cand)
        s = mul if mode == "mul" else add
        order = [cw[i] for i in np.argsort(-s)[: top_k + 3] if cw[i] not in (a.lower(), b.lower(), c.lower())]
        preds.append(order[:top_k])
    return pd.DataFrame({"Predicted_D": [p[0] for p in preds], "Top_K_Predictions": preds})


def candidate_table(g, df, n_vocab=30):
    """One row per (analogy, candidate) with the re-ranker features."""
    fw = file_words(df); fset = set(fw)
    pairs = {(a.lower(), b.lower()) for a, b in df[["A", "B"]].values}   # (A, B) pairs present in the file
    rows = []
    for r, (a, b, c) in enumerate(df[["A", "B", "C"]].values):
        _, mul, *_ = scores(g, a, b, c, g.V)
        cw = [g.words[i] for i in np.argsort(-mul)[: n_vocab + 3]] + fw
        cw = [w for w in dict.fromkeys(cw) if w not in (a.lower(), b.lower(), c.lower())]
        cand = np.stack([g.vec(w) for w in cw])
        add, mul, sa, sb, sc = scores(g, a, b, c, cand)
        for j, w in enumerate(cw):
            rows.append((r, w, add[j], mul[j], sa[j], sb[j], sc[j], float(w in fset), np.log1p(g.idx.get(w, VOCAB_N)),
                         float((c.lower(), w) in pairs)))
    return pd.DataFrame(rows, columns=["row", "w", "add", "mul", "sa", "sb", "sc", "infile", "lrank", "pair"])


FEATS = ["add", "mul", "sa", "sb", "sc", "infile", "lrank", "pair"]


def learn_analogy_model(g, train_df, feats=None):
    feats = feats or FEATS
    T = candidate_table(g, train_df)
    y = (T.w.values == train_df.D.str.lower().values[T.row.values]).astype(int)
    clf = LogisticRegression(C=1.0, max_iter=2000, class_weight="balanced").fit(T[feats], y)
    clf.feats = feats
    return clf


def predict_trained(g, clf, df, top_k=3):
    T = candidate_table(g, df); T["s"] = clf.decision_function(T[clf.feats])
    preds = [grp.sort_values("s", ascending=False).w.tolist()[:top_k] for _, grp in T.groupby("row", sort=True)]
    return pd.DataFrame({"Predicted_D": [p[0] for p in preds], "Top_K_Predictions": preds})


def precision(pred, df):
    d = df.D.str.lower().values
    return (pred.Predicted_D.values == d).mean(), np.mean([t in set(p) for t, p in zip(d, pred.Top_K_Predictions)])


if __name__ == "__main__":
    tr, te = pd.read_csv(sys.argv[1]), pd.read_csv(sys.argv[2]); g = Glove(sys.argv[3], sys.argv[4])
    for cands in ("vocab", "file"):
        for mode in ("add", "mul"):
            r = [precision(zero_shot(g, df, mode, cands), df) for df in (tr, te)]
            print(f"zero-shot {mode:3s} over {cands:5s}: train P@1 {r[0][0]:.3f} P@3 {r[0][1]:.3f} | test P@1 {r[1][0]:.3f} P@3 {r[1][1]:.3f}")
    for feats in (FEATS[:-1], FEATS):
        clf = learn_analogy_model(g, tr, feats)
        print("re-ranker", "WITH" if "pair" in feats else "without", "pair feature; weights", dict(zip(feats, clf.coef_[0].round(2))))
        for name, df in (("train", tr), ("test", te)):
            p1, p3 = precision(predict_trained(g, clf, df), df); print(f"   {name}: P@1 {p1:.3f} P@3 {p3:.3f}")

Analogy Oracle: Cracking the Code of Word Relationships