Upload 14 files

METHOD.md

@@ -0,0 +1,32 @@
+## Method
+
+For one query with candidate passages $d_1,\dots,d_n$:
+
+- $r_j$: cross-encoder logit (InfoMax) or z-score fusion of cross-encoder logit and dense cosine (InfoMax+F).
+- $e_j$: all-MiniLM-L6-v2 embedding, centred on the pool mean and renormalised; $q$ the query embedding, treated the same way.
+- $w_j = \mathrm{minmax}(r)_j^{\gamma}$: relevance weight.
+
+**Target.** A latent query-information vector $z \in \mathbb{R}^r$, with prior $z \sim \mathcal N(0,\Lambda)$, on the aspect basis
+$U = [\,q,\ u_2,\dots,u_r\,]$. Here $u_2,\dots,u_r$ are the top right-singular vectors of the relevance-weighted pool embeddings
+after removing the query direction, and $\Lambda = \mathrm{diag}(1, \rho\, s_k^2/s_1^2)$.
+
+**Observations.** Each passage is a noisy linear measurement of the target:
+
+$$y_j = \sqrt{w_j}\,(U^\top e_j)^\top z + g_j^\top u + \varepsilon_j,\qquad g_j = (I-UU^\top)e_j .$$
+
+The shared nuisance $u \sim \mathcal N(0,\beta I)$ correlates the noise of passages with similar off-query content (duplicates,
+same source), and $\varepsilon_j \sim \mathcal N(0,\sigma^2)$.
+
+**Selection.** Greedy maximisation of the conditional mutual information
+
+$$I(z; y_j \mid y_S) = \tfrac12 \log \frac{\mathrm{Var}(y_j \mid y_S)}{\mathrm{Var}(y_j \mid y_S, z)},$$
+
+computed exactly by Schur complements. No probabilities from a language model are needed: in the linear-Gaussian model the
+posterior entropy depends only on which passages are selected.
+
+**Anchoring.** The first pick is the top-1 of the base ranking $r$. At each later step, the base ranking's next item $b$ is replaced by
+the max-gain item $j^\star$ only if $\mathrm{gain}(j^\star) > (1+\delta)\,\mathrm{gain}(b)$. With $\delta=\infty$ the selector returns
+the base ranking exactly.
+
+**Guarantee.** None is claimed. Correlated noise breaks submodularity of the objective, so the greedy rule has no worst-case
+approximation bound.

README.md

@@ -1,15 +1,53 @@
 ---
-title: Infomax Reranker
-emoji: 🏢
+title: InfoMax Reranker
+emoji: 🔎
 colorFrom: indigo
-colorTo: green
+colorTo: blue
 sdk: gradio
 sdk_version: 6.31.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
-license: mit
+license: apache-2.0
 short_description: Entropy-minimising passage reranker vs MS MARCO CE
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# InfoMax Reranker
+
+A greedy, information-maximising selector over a retrieved candidate pool, compared with the standard
+`cross-encoder/ms-marco-MiniLM-L6-v2` reranker. Neural scoring runs on ZeroGPU; selection runs on CPU.
+
+- `app.py`: Gradio app. Paste a query and passages separated by blank lines.
+- `infomax_rerank.py`: the selectors: `infomax`, plus the baselines `mmr`, `fused`/`zfuse` and `by_score`.
+- `METHOD.md`: the model and the selection rule.
+- `results/`: BEIR test results (`beir_results.md`, `beir_results.json`).
+- `benchmark/`: the scripts that produced them (see below).
+
+## Hardware
+
+Create the Space with the Gradio SDK and select **ZeroGPU** as hardware. Hosting a ZeroGPU Space needs a PRO or Enterprise
+account. The `@spaces.GPU` decorator has no effect elsewhere, so the app also runs on CPU hardware and locally:
+
+```
+pip install -r requirements.txt gradio
+python app.py
+```
+
+## Reproducing the BEIR numbers
+
+The datasets are the standard BEIR zips (scifact, nfcorpus, fiqa, scidocs, trec-covid), unzipped into `$BEIR_DATA`.
+Outputs go to `$EXP_OUT`. Models are loaded from local folders or Hub ids.
+
+```
+cd benchmark
+export BEIR_DATA=/path/to/beir EXP_OUT=/path/to/out
+python bm25_pools.py                                   # BM25 top-100 pools: test splits + 200-query dev samples
+python dense_stage.py sentence-transformers/all-MiniLM-L6-v2 50
+python ce_stage.py cross-encoder/ms-marco-MiniLM-L6-v2 50
+python build_scores.py 50
+python run_eval.py            # tune every method on dev splits only -> best_params.json
+python run_eval.py --test     # nDCG / Recall on the test splits + paired bootstrap CIs
+python mrr_recall.py          # MRR@10, Recall@3/5/10
+```
+
+As a check on the setup, BM25 nDCG@10 is within about 0.01 of the published BEIR BM25 numbers on all five datasets.
+all-MiniLM-L6-v2 scores 0.648 nDCG@10 on SciFact over the full corpus.

app.py

@@ -0,0 +1,133 @@
+"""InfoMax reranker demo (Gradio, ZeroGPU).
+
+Neural scoring (MS MARCO MiniLM-L6 cross-encoder + all-MiniLM-L6-v2 embeddings) runs inside the
+@spaces.GPU function; every selector then runs on CPU from those cached scores.
+"""
+import spaces  # import first on ZeroGPU
+
+import os
+
+import gradio as gr
+import numpy as np
+import pandas as pd
+import torch
+from sentence_transformers import CrossEncoder, SentenceTransformer
+
+from infomax_rerank import by_score, fused, infomax, mmr, zfuse
+
+CE_ID = os.environ.get("CE_MODEL", "cross-encoder/ms-marco-MiniLM-L6-v2")
+DENSE_ID = os.environ.get("DENSE_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+ce = CrossEncoder(CE_ID, max_length=256, device="cpu")
+enc = SentenceTransformer(DENSE_ID, device="cpu")
+
+# Settings tuned on dev splits only (SciFact train, NFCorpus dev, FiQA dev); one setting for all datasets.
+TUNED = {
+    "InfoMax": dict(rank=1, gamma=4.0, rho=0.3, sigma2=0.03, beta=0.0, delta=0.0),
+    "InfoMax+F": dict(rank=4, gamma=8.0, rho=1.0, sigma2=0.03, beta=3.0, delta=0.0, alpha=0.5),
+    "MMR": dict(lam=0.95),
+    "Fusion": dict(alpha=0.5),
+}
+METHODS = ["CE", "InfoMax", "InfoMax+F", "Dense", "MMR", "Fusion"]
+
+
+@spaces.GPU(duration=60)
+def neural_scores(query, passages):
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    ce.model.to(dev)
+    enc.to(dev)
+    logits = np.asarray(ce.predict([(query, p) for p in passages], batch_size=64,
+                                   show_progress_bar=False), np.float32).reshape(-1)
+    E = enc.encode([query] + passages, batch_size=64, normalize_embeddings=True,
+                   convert_to_numpy=True, show_progress_bar=False).astype(np.float32)
+    return logits, E[1:], E[0]
+
+
+def split_passages(text):
+    blocks = [b.strip() for b in text.replace("\r\n", "\n").split("\n\n")]
+    return [b for b in blocks if b]
+
+
+def rank_all(query, passages_text, k, rank, gamma, sigma2, beta, delta, alpha):
+    query = (query or "").strip()
+    passages = split_passages(passages_text or "")
+    if not query:
+        raise gr.Error("Enter a query.")
+    if len(passages) < 2:
+        raise gr.Error("Enter at least two passages, separated by a blank line.")
+    if len(passages) > 200:
+        raise gr.Error("At most 200 passages per request.")
+    k = int(min(k, len(passages)))
+    r, E, q = neural_scores(query, passages)
+    user_f = dict(rank=int(rank), gamma=gamma, rho=1.0, sigma2=sigma2, beta=beta,
+                  delta=np.inf if delta >= 10 else delta)
+    orders, gains = {}, {}
+    orders["CE"] = by_score(r)
+    orders["Dense"] = by_score(E @ q)
+    orders["MMR"] = mmr(r, E, K=k, **TUNED["MMR"])
+    orders["Fusion"] = zfuse(r, E, q, **TUNED["Fusion"])
+    orders["InfoMax"], gains["InfoMax"] = infomax(r, E, q, K=k, return_gains=True, **TUNED["InfoMax"])
+    orders["InfoMax+F"], gains["InfoMax+F"] = infomax(fused(r, E, q, alpha), E, q, K=k, return_gains=True,
+                                                       **user_f)
+    short = lambda i: f"[{i + 1}] " + (passages[i][:90] + ("…" if len(passages[i]) > 90 else ""))
+    table = pd.DataFrame({m: [short(i) for i in orders[m][:k]] for m in METHODS})
+    table.insert(0, "Rank", list(range(1, k + 1)))
+    detail = pd.DataFrame({
+        "Passage": [f"[{i + 1}]" for i in range(len(passages))],
+        "CE logit": np.round(r, 3),
+        "Dense cosine": np.round(E @ q, 3),
+        "InfoMax+F step": [orders["InfoMax+F"][:k].index(i) + 1 if i in orders["InfoMax+F"][:k] else "" for i in
+                           range(len(passages))],
+    })
+    g = gains["InfoMax+F"]
+    note = ("**InfoMax+F information gain per step (nats):** " +
+            ", ".join(f"{t + 1}: {v:.3f}" for t, v in enumerate(g)))
+    return table, detail, note
+
+
+with open(os.path.join(os.path.dirname(__file__), "results", "beir_results.md")) as f:
+    RESULTS_MD = f.read()
+
+EX_PASSAGES = """Metformin commonly causes gastrointestinal side effects such as nausea, diarrhoea and abdominal discomfort, especially when treatment starts.
+
+Gastrointestinal upset, including nausea and diarrhoea, is the most frequent side effect reported by patients starting metformin.
+
+Long-term metformin use is associated with reduced vitamin B12 absorption, so periodic B12 monitoring is recommended.
+
+Lactic acidosis is a rare but serious adverse effect of metformin, mainly in patients with severe kidney impairment.
+
+Metformin lowers hepatic glucose production and improves insulin sensitivity in type 2 diabetes.
+
+Taking metformin with meals or using an extended-release formulation can reduce stomach upset.
+
+Sulfonylureas such as glipizide can cause hypoglycaemia and weight gain."""
+
+with gr.Blocks(title="InfoMax reranker") as demo:
+    gr.Markdown("# InfoMax reranker\nGreedy information-maximising selection over a candidate pool, compared with "
+                "the standard MS MARCO MiniLM-L6 cross-encoder. Paste a query and passages separated by blank lines.")
+    with gr.Tab("Rerank"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                query = gr.Textbox(label="Query", value="What are the side effects of metformin?")
+                passages = gr.Textbox(label="Passages (blank line between passages)", lines=16, value=EX_PASSAGES)
+                k = gr.Slider(1, 20, value=5, step=1, label="K (passages to select)")
+                with gr.Accordion("InfoMax+F settings (defaults = dev-tuned)", open=False):
+                    rank = gr.Slider(1, 16, value=4, step=1, label="Aspect rank r")
+                    gamma = gr.Slider(0.5, 16, value=8.0, step=0.5, label="Relevance exponent γ")
+                    sigma2 = gr.Slider(0.01, 1.0, value=0.03, step=0.01, label="Observation noise σ²")
+                    beta = gr.Slider(0.0, 10.0, value=3.0, step=0.1, label="Shared nuisance variance β")
+                    delta = gr.Slider(0.0, 10.0, value=0.0, step=0.1,
+                                      label="Anchor margin δ (10 = never leave the fused ranking)")
+                    alpha = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Dense weight α in relevance fusion")
+                btn = gr.Button("Rerank", variant="primary")
+            with gr.Column(scale=2):
+                out = gr.Dataframe(label="Top-K per method", wrap=True)
+                note = gr.Markdown()
+                detail = gr.Dataframe(label="Per-passage scores", wrap=True)
+        btn.click(rank_all, [query, passages, k, rank, gamma, sigma2, beta, delta, alpha], [out, detail, note])
+    with gr.Tab("BEIR results"):
+        gr.Markdown(RESULTS_MD)
+    with gr.Tab("Method"):
+        gr.Markdown(open(os.path.join(os.path.dirname(__file__), "METHOD.md")).read())
+
+if __name__ == "__main__":
+    demo.queue().launch()

benchmark/bm25_pools.py

@@ -0,0 +1,75 @@
+"""First stage: Lucene-style BM25 (k1=0.9, b=0.4, Porter stemming) top-100 pools per test query."""
+import json
+import os
+import re
+import time
+
+import numpy as np
+import scipy.sparse as sp
+from nltk.stem import PorterStemmer
+from sklearn.feature_extraction.text import ENGLISH_STOP_WORDS, CountVectorizer
+
+from common import DATASETS, OUT, evaluate, load_corpus, load_qrels, load_queries
+
+TOK = re.compile(r"[a-z0-9]+")
+_stem = PorterStemmer().stem
+_cache = {}
+
+
+def analyze(text):
+    out = []
+    for t in TOK.findall(text.lower()):
+        if t in ENGLISH_STOP_WORDS:
+            continue
+        s = _cache.get(t)
+        if s is None:
+            s = _cache[t] = _stem(t)
+        out.append(s)
+    return out
+
+
+def run(name, split="test", depth=100, k1=0.9, b=0.4, max_q=None, seed=0):
+    t0 = time.time()
+    ids, titles, texts = load_corpus(name)
+    docs = [(ti + " " + tx).strip() for ti, tx in zip(titles, texts)]
+    queries = load_queries(name)
+    qrels = load_qrels(name, split)
+    qids = sorted(q for q in qrels if q in queries)
+    if max_q and len(qids) > max_q:
+        qids = sorted(np.random.default_rng(seed).choice(qids, max_q, replace=False).tolist())
+    cv = CountVectorizer(analyzer=analyze)
+    X = cv.fit_transform(docs).tocsc().astype(np.float32)  # docs x terms
+    N = X.shape[0]
+    df = np.diff(X.indptr)
+    idf = np.log(1 + (N - df + 0.5) / (df + 0.5)).astype(np.float32)
+    X = X.tocsr()
+    dl = np.asarray(X.sum(1)).ravel()
+    norm = k1 * (1 - b + b * dl / dl.mean())
+    X = X.tocoo()
+    w = X.data * (k1 + 1) / (X.data + norm[X.row]) * idf[X.col]
+    W = sp.csc_matrix((w, (X.row, X.col)), shape=X.shape)
+    vocab = cv.vocabulary_
+    pools, run_ = {}, {}
+    for qid in qids:
+        terms = [vocab[t] for t in analyze(queries[qid]) if t in vocab]
+        if not terms:
+            s = np.zeros(N, dtype=np.float32)
+        else:
+            u, c = np.unique(terms, return_counts=True)
+            s = np.asarray(W[:, u] @ c.astype(np.float32)).ravel()
+        top = np.argpartition(-s, depth)[:depth]
+        top = top[np.argsort(-s[top], kind="stable")]
+        pools[qid] = {"ids": [ids[i] for i in top], "bm25": s[top].tolist()}
+        run_[qid] = {ids[i]: float(s[i]) for i in top}
+    metrics, _ = evaluate(run_, qrels)
+    os.makedirs(OUT, exist_ok=True)
+    json.dump(pools, open(os.path.join(OUT, f"{name}.{split}.pools.json"), "w"))
+    print(f"{name}/{split}: {N} docs, {len(qids)} queries, {time.time()-t0:.0f}s  " +
+          "  ".join(f"{k}={v:.4f}" for k, v in metrics.items()), flush=True)
+
+
+if __name__ == "__main__":
+    for n in DATASETS:
+        run(n, "test")
+    for n, sp_ in [("scifact", "train"), ("nfcorpus", "dev"), ("fiqa", "dev")]:
+        run(n, sp_, max_q=200)

benchmark/build_scores.py

@@ -0,0 +1,39 @@
+"""Merge pools + dense embeddings + cross-encoder logits into out/<tag>.scores.npz (format read by run_eval.py)."""
+import json
+import os
+import sys
+
+import numpy as np
+
+from common import OUT
+
+TAGS = ["scifact.test", "nfcorpus.test", "fiqa.test", "scidocs.test", "trec-covid.test",
+        "scifact.train", "nfcorpus.dev", "fiqa.dev"]
+
+
+def main(depth, tags):
+    for tag in tags:
+        name = tag.split(".")[0]
+        pools = json.load(open(os.path.join(OUT, f"{tag}.pools.json")))
+        z = np.load(os.path.join(OUT, f"{name}.emb.npz"))
+        emb = dict(zip(z["ids"].tolist(), z["emb"]))
+        zq = np.load(os.path.join(OUT, f"{name}.qemb.npz"))
+        qemb = dict(zip(zq["ids"].tolist(), zq["emb"]))
+        logits = np.load(os.path.join(OUT, f"{tag}.ce.npz"))["logits"]
+        qids = sorted(pools)
+        assert len(logits) == sum(len(pools[q]["ids"][:depth]) for q in qids), tag
+        out, p = {"qids": np.array(qids)}, 0
+        for q in qids:
+            ds = pools[q]["ids"][:depth]
+            out[f"{q}|ids"] = np.array(ds)
+            out[f"{q}|bm25"] = np.asarray(pools[q]["bm25"][:depth], np.float32)
+            out[f"{q}|ce"] = logits[p:p + len(ds)]
+            out[f"{q}|emb"] = np.stack([emb[d] for d in ds])
+            out[f"{q}|qemb"] = qemb[q]
+            p += len(ds)
+        np.savez_compressed(os.path.join(OUT, f"{tag}.scores.npz"), **out)
+        print(tag, len(qids), "queries")
+
+
+if __name__ == "__main__":
+    main(int(sys.argv[1]), sys.argv[2:] or TAGS)

benchmark/ce_stage.py

@@ -0,0 +1,56 @@
+"""Cross-encoder stage: MS MARCO MiniLM-L6-v2 logits (max_length 256) for every (query, pool doc), top-DEPTH.
+Resumable: out/<tag>.ce.npz holds logits computed so far in pool order.
+Usage: python ce_stage.py CE_DIR DEPTH
+"""
+import json
+import os
+import sys
+import time
+
+import numpy as np
+import torch
+from sentence_transformers import CrossEncoder
+
+from common import OUT, load_corpus, load_queries
+
+torch.set_num_threads(os.cpu_count())
+TAGS = ["scifact.test", "nfcorpus.test", "trec-covid.test", "scifact.train", "nfcorpus.dev", "fiqa.dev",
+        "fiqa.test", "scidocs.test"]
+
+
+def main():
+    cdir, depth = sys.argv[1], int(sys.argv[2])
+    ce = CrossEncoder(cdir, max_length=256, device="cpu")
+    cache = {}
+    for tag in TAGS:
+        name = tag.split(".")[0]
+        if name not in cache:
+            cache.clear()
+            ids, titles, texts = load_corpus(name)
+            idx = {d: i for i, d in enumerate(ids)}
+            cache[name] = (idx, titles, texts, load_queries(name))
+        idx, titles, texts, queries = cache[name]
+        pools = json.load(open(os.path.join(OUT, f"{tag}.pools.json")))
+        pairs = [(queries[q], (titles[idx[d]] + ". " if titles[idx[d]] else "") + texts[idx[d]])
+                 for q in sorted(pools) for d in pools[q]["ids"][:depth]]
+        path = os.path.join(OUT, f"{tag}.ce.npz")
+        done = np.load(path)["logits"] if os.path.exists(path) else np.zeros(0, np.float32)
+        if len(done) == len(pairs):
+            print(f"{tag}: complete", flush=True)
+            continue
+        t0, start = time.time(), len(done)
+        for s in range(start, len(pairs), 2048):
+            chunk = pairs[s:s + 2048]
+            o = np.argsort([len(a) + len(b) for a, b in chunk], kind="stable")  # length-sort: less padding
+            out = np.asarray(ce.predict([chunk[i] for i in o], batch_size=32, show_progress_bar=False),
+                             np.float32).reshape(-1)
+            res = np.empty_like(out)
+            res[o] = out
+            done = np.concatenate([done, res])
+            np.savez(path, logits=done)
+            rate = (len(done) - start) / (time.time() - t0)
+            print(f"{tag}: {len(done)}/{len(pairs)} pairs, {rate:.0f} pairs/s", flush=True)
+
+
+if __name__ == "__main__":
+    main()

benchmark/common.py

@@ -0,0 +1,60 @@
+"""Shared I/O and evaluation for the BEIR reranking experiments."""
+import csv
+import json
+import os
+
+import numpy as np
+import pytrec_eval
+
+DATA = os.environ.get("BEIR_DATA", "beir")
+OUT = os.environ.get("EXP_OUT", "out")
+DATASETS = ["scifact", "nfcorpus", "fiqa", "scidocs", "trec-covid"]
+
+
+def load_corpus(name):
+    ids, texts, titles = [], [], []
+    with open(os.path.join(DATA, name, "corpus.jsonl")) as f:
+        for line in f:
+            d = json.loads(line)
+            ids.append(d["_id"])
+            titles.append(d.get("title") or "")
+            texts.append(d.get("text") or "")
+    return ids, titles, texts
+
+
+def load_queries(name):
+    q = {}
+    with open(os.path.join(DATA, name, "queries.jsonl")) as f:
+        for line in f:
+            d = json.loads(line)
+            q[d["_id"]] = d["text"]
+    return q
+
+
+def load_qrels(name, split="test"):
+    qrels = {}
+    with open(os.path.join(DATA, name, "qrels", f"{split}.tsv")) as f:
+        r = csv.reader(f, delimiter="\t")
+        next(r)
+        for qid, did, score in r:
+            qrels.setdefault(qid, {})[did] = int(score)
+    return qrels
+
+
+def evaluate(run, qrels, ks=(3, 5, 10)):
+    """run: {qid: {docid: score}}. Returns mean nDCG@k and Recall@k (BEIR/pytrec_eval conventions)."""
+    measures = {f"ndcg_cut.{','.join(map(str, ks))}", f"recall.{','.join(map(str, ks))}"}
+    q = {k: v for k, v in qrels.items() if k in run}
+    ev = pytrec_eval.RelevanceEvaluator(q, measures)
+    res = ev.evaluate({k: run[k] for k in q})
+    out = {}
+    for k in ks:
+        out[f"nDCG@{k}"] = float(np.mean([r[f"ndcg_cut_{k}"] for r in res.values()]))
+        out[f"R@{k}"] = float(np.mean([r[f"recall_{k}"] for r in res.values()]))
+    return out, res
+
+
+def order_to_run(order_ids):
+    """Selected ordered list -> descending pseudo-scores so pytrec_eval keeps the order."""
+    n = len(order_ids)
+    return {d: float(n - i) for i, d in enumerate(order_ids)}

benchmark/dense_stage.py

@@ -0,0 +1,60 @@
+"""Dense stage: all-MiniLM-L6-v2 embeddings for every doc in the top-DEPTH pools and every query.
+Resumable: per-dataset store out/<name>.emb.npz updated after each chunk.
+Usage: python dense_stage.py MINILM_DIR DEPTH
+"""
+import json
+import os
+import sys
+import time
+
+import numpy as np
+import torch
+from sentence_transformers import SentenceTransformer
+
+from common import OUT, load_corpus, load_queries
+
+torch.set_num_threads(os.cpu_count())
+TAGS = ["scifact.test", "scifact.train", "nfcorpus.test", "nfcorpus.dev", "trec-covid.test",
+        "fiqa.test", "fiqa.dev", "scidocs.test"]
+
+
+def main():
+    mdir, depth = sys.argv[1], int(sys.argv[2])
+    enc = SentenceTransformer(mdir, device="cpu")
+    by_ds = {}
+    for t in TAGS:
+        by_ds.setdefault(t.split(".")[0], []).append(t)
+    for name, tags in by_ds.items():
+        store = os.path.join(OUT, f"{name}.emb.npz")
+        have = {}
+        if os.path.exists(store):
+            z = np.load(store)
+            have = dict(zip(z["ids"].tolist(), z["emb"]))
+        need, qneed = set(), set()
+        for t in tags:
+            pools = json.load(open(os.path.join(OUT, f"{t}.pools.json")))
+            for q, p in pools.items():
+                need.update(p["ids"][:depth])
+                qneed.add(q)
+        ids, titles, texts = load_corpus(name)
+        idx = {d: i for i, d in enumerate(ids)}
+        queries = load_queries(name)
+        todo = sorted(d for d in need if d not in have)
+        t0 = time.time()
+        for s in range(0, len(todo), 2048):
+            chunk = todo[s:s + 2048]
+            E = enc.encode([(titles[idx[d]] + ". " if titles[idx[d]] else "") + texts[idx[d]] for d in chunk],
+                           batch_size=64, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False)
+            have.update(zip(chunk, E.astype(np.float32)))
+            k = sorted(have)
+            np.savez(store, ids=np.array(k), emb=np.stack([have[d] for d in k]))
+            print(f"{name}: {min(s + 2048, len(todo))}/{len(todo)} docs, {time.time()-t0:.0f}s", flush=True)
+        qk = sorted(qneed)
+        QE = enc.encode([queries[q] for q in qk], batch_size=64, normalize_embeddings=True,
+                        convert_to_numpy=True, show_progress_bar=False)
+        np.savez(os.path.join(OUT, f"{name}.qemb.npz"), ids=np.array(qk), emb=QE.astype(np.float32))
+        print(f"{name}: done ({len(have)} docs, {len(qk)} queries)", flush=True)
+
+
+if __name__ == "__main__":
+    main()

benchmark/mrr_recall.py

@@ -0,0 +1,53 @@
+"""MRR@10 and Recall@{3,5,10} on the test splits for every method, using the dev-tuned settings in best_params.json.
+Relevance = qrel grade > 0. Paired bootstrap 95% CIs for InfoMax+F minus CE and minus Fusion."""
+import json
+import os
+
+import numpy as np
+
+from common import OUT
+from run_eval import TEST, boot_ci, load, method_orders
+
+METHODS = ["BM25", "Dense", "CE", "MMR", "Fusion", "InfoMax", "InfoMax+F"]
+best = {k: {kk: (np.inf if vv is None else vv) for kk, vv in v.items()}
+        for k, v in json.load(open(os.path.join(OUT, "best_params.json"))).items()}
+
+
+def per_query(Q, qrels, orders):
+    out = {"MRR@10": [], "R@3": [], "R@5": [], "R@10": []}
+    for x, o in zip(Q, orders):
+        rel = {d for d, g in qrels.get(x["qid"], {}).items() if g > 0}
+        if not rel:
+            continue
+        ranked = [x["ids"][j] for j in o]
+        hit = [d in rel for d in ranked[:10]]
+        out["MRR@10"].append(1.0 / (hit.index(True) + 1) if any(hit) else 0.0)
+        for k in (3, 5, 10):
+            out[f"R@{k}"].append(sum(hit[:k]) / len(rel))
+    return {k: np.array(v) for k, v in out.items()}
+
+
+res = {}
+for tag in TEST:
+    Q, qr = load(tag)
+    for m in METHODS:
+        o = method_orders(Q, m, best.get(m, {}))
+        res[(tag, m)] = per_query(Q, qr, o)
+
+for metric in ("MRR@10", "R@3", "R@5", "R@10"):
+    print(f"\n{metric}")
+    print(f"{'dataset':16s}" + "".join(f"{m:>10s}" for m in METHODS))
+    for tag in TEST:
+        print(f"{tag:16s}" + "".join(f"{res[(tag, m)][metric].mean():10.4f}" for m in METHODS))
+
+print("\nInfoMax+F minus CE / minus Fusion, paired bootstrap 95% CI")
+for metric in ("MRR@10", "R@10"):
+    for tag in TEST:
+        a = res[(tag, "InfoMax+F")][metric]
+        s = []
+        for o in ("CE", "Fusion"):
+            d, lo, hi = boot_ci(a, res[(tag, o)][metric])
+            s.append(f"vs {o}: {d:+.4f} [{lo:+.4f},{hi:+.4f}]")
+        print(f"{metric:7s}{tag:16s}" + "   ".join(s))
+json.dump({f"{t}|{m}|{k}": float(v.mean()) for (t, m), d in res.items() for k, v in d.items()},
+          open(os.path.join(OUT, "mrr_recall.json"), "w"), indent=1)

benchmark/run_eval.py

@@ -0,0 +1,141 @@
+"""Tune every method's knobs on dev splits only (scifact.train, nfcorpus.dev, fiqa.dev; one global
+setting per method), then evaluate on the five BEIR test splits.
+
+Usage: python run_eval.py [--quick]
+"""
+import itertools
+import json
+import os
+import sys
+import time
+
+import numpy as np
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
+from common import OUT, evaluate, load_qrels, order_to_run  # noqa: E402
+from infomax_rerank import by_score, fused, infomax, mmr, zfuse  # noqa: E402
+
+DEV = ["scifact.train", "nfcorpus.dev", "fiqa.dev"]
+TEST = ["scifact.test", "nfcorpus.test", "fiqa.test", "scidocs.test", "trec-covid.test"]
+
+
+def load(tag):
+    name, split = tag.split(".")
+    z = np.load(os.path.join(OUT, f"{tag}.scores.npz"))
+    qs = []
+    for q in z["qids"]:
+        qs.append({k: z[f"{q}|{k}"] for k in ("ids", "bm25", "ce", "emb", "qemb")} | {"qid": str(q)})
+    return qs, load_qrels(name, split)
+
+
+def score(Q, qrels, orders):
+    run = {x["qid"]: order_to_run([x["ids"][j] for j in o]) for x, o in zip(Q, orders)}
+    m, per = evaluate(run, qrels)
+    return m, {q: v["ndcg_cut_10"] for q, v in per.items()}
+
+
+def method_orders(Q, method, p):
+    if method == "BM25":
+        return [by_score(x["bm25"]) for x in Q]
+    if method == "Dense":
+        return [by_score(x["emb"] @ x["qemb"]) for x in Q]
+    if method == "CE":
+        return [by_score(x["ce"]) for x in Q]
+    if method == "MMR":
+        return [mmr(x["ce"], x["emb"], **p) for x in Q]
+    if method == "Fusion":
+        return [zfuse(x["ce"], x["emb"], x["qemb"], **p) for x in Q]
+    if method == "InfoMax":
+        return [infomax(x["ce"], x["emb"], x["qemb"], **p) for x in Q]
+    if method == "InfoMax+F":   # same selector, relevance/anchor = the tuned CE+dense fusion score
+        p = dict(p)
+        a = p.pop("alpha")
+        return [infomax(fused(x["ce"], x["emb"], x["qemb"], a), x["emb"], x["qemb"], **p) for x in Q]
+    raise ValueError(method)
+
+
+GRIDS = {
+    "MMR": [{"lam": l} for l in (0.5, 0.6, 0.7, 0.8, 0.9, 0.95)],
+    "Fusion": [{"alpha": a} for a in (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)],
+    "InfoMax": [dict(rank=r, gamma=g, rho=rh, sigma2=s, beta=b, delta=d)
+                for r, g, rh, s, b, d in itertools.product(
+                    (1, 4, 16), (1.0, 2.0, 4.0), (0.3, 1.0), (0.03, 0.3), (0.0, 0.3, 3.0),
+                    (0.0, 0.5, 2.0, np.inf))],
+    "InfoMax+F": [dict(rank=r, gamma=g, rho=1.0, sigma2=s, beta=b, delta=d)
+                  for r, g, s, b, d in itertools.product(
+                      (1, 4, 16), (1.0, 2.0, 4.0, 8.0), (0.03, 0.3), (0.0, 0.3, 3.0), (0.0, 0.5, 2.0, np.inf))],
+}
+
+
+def boot_ci(a, b, n=10000, seed=0):
+    d = np.asarray(a) - np.asarray(b)
+    rng = np.random.default_rng(seed)
+    m = d[rng.integers(0, len(d), (n, len(d)))].mean(1)
+    return float(d.mean()), float(np.percentile(m, 2.5)), float(np.percentile(m, 97.5))
+
+
+def main():
+    quick = "--quick" in sys.argv
+    best = {}
+    bp = os.path.join(OUT, "best_params.json")
+    if "--test" in sys.argv:
+        best = {k: {kk: (np.inf if vv is None else vv) for kk, vv in v.items()}
+                for k, v in json.load(open(bp)).items()}
+        grids = {}
+    else:
+        dev = {t: load(t) for t in DEV}
+        grids = GRIDS
+    for meth, grid in grids.items():
+        if meth == "InfoMax+F":
+            grid = [p | {"alpha": best["Fusion"]["alpha"]} for p in grid]
+        if quick:
+            grid = grid[:4]
+        t0 = time.time()
+        res = []
+        for p in grid:
+            v = np.mean([score(Q, qr, method_orders(Q, meth, p))[0]["nDCG@10"] for Q, qr in dev.values()])
+            res.append((v, p))
+        res.sort(key=lambda t: -t[0])
+        best[meth] = res[0][1]
+        ce_dev = np.mean([score(Q, qr, method_orders(Q, "CE", {}))[0]["nDCG@10"] for Q, qr in dev.values()])
+        print(f"[tune] {meth}: best dev nDCG@10 {res[0][0]:.4f} (CE {ce_dev:.4f}) with {res[0][1]}  "
+              f"[{len(grid)} settings, {time.time()-t0:.0f}s]", flush=True)
+        if meth.startswith("InfoMax"):
+            for v, p in res[:5]:
+                print(f"        {v:.4f} {p}")
+    if "--test" not in sys.argv:
+        json.dump({k: {kk: (None if vv == np.inf else vv) for kk, vv in v.items()} for k, v in best.items()},
+                  open(bp, "w"), indent=1)
+        return
+
+    methods = ["BM25", "Dense", "CE", "MMR", "Fusion", "InfoMax", "InfoMax+F"]
+    table, perq = {}, {}
+    for tag in TEST:
+        Q, qr = load(tag)
+        for meth in methods:
+            o = method_orders(Q, meth, best.get(meth, {}))
+            m, pq = score(Q, qr, o)
+            table[(tag, meth)] = m
+            perq[(tag, meth)] = pq
+    print("\nnDCG@10 on test (BM25 top-50 pools; knobs tuned on dev only)")
+    print(f"{'dataset':16s}" + "".join(f"{m:>9s}" for m in methods))
+    for tag in TEST:
+        print(f"{tag:16s}" + "".join(f"{table[(tag, m)]['nDCG@10']:9.4f}" for m in methods))
+    print("\nInfoMax+F minus X, nDCG@10, paired bootstrap 95% CI")
+    rows = {}
+    for tag in TEST:
+        qids = sorted(perq[(tag, "CE")])
+        for other in ("CE", "MMR", "Fusion"):
+            a = [perq[(tag, "InfoMax+F")][q] for q in qids]
+            b = [perq[(tag, other)][q] for q in qids]
+            rows[(tag, other)] = boot_ci(a, b)
+        print(f"{tag:16s}" + "  ".join(f"vs {o}: {rows[(tag, o)][0]:+.4f} [{rows[(tag, o)][1]:+.4f},{rows[(tag, o)][2]:+.4f}]"
+                                       for o in ("CE", "MMR", "Fusion")))
+    json.dump({"best": {k: {kk: (None if vv == np.inf else vv) for kk, vv in v.items()} for k, v in best.items()},
+               "test": {f"{t}|{m}": v for (t, m), v in table.items()},
+               "ci_infomax_minus": {f"{t}|{o}": v for (t, o), v in rows.items()}},
+              open(os.path.join(OUT, "results.json"), "w"), indent=1)
+
+
+if __name__ == "__main__":
+    main()

infomax_rerank.py

@@ -0,0 +1,108 @@
+"""Rerankers over a fixed candidate pool. Each returns a full ordering (indices into the pool).
+
+Inputs per query: r = cross-encoder logits (n,), E = unit doc embeddings (n, d), q = unit query embedding (d,).
+"""
+import numpy as np
+from scipy.linalg import cho_factor, cho_solve
+
+
+def minmax(x):
+    lo, hi = float(x.min()), float(x.max())
+    return (x - lo) / (hi - lo) if hi > lo + 1e-9 else np.zeros_like(x)
+
+
+def by_score(s):
+    return list(np.argsort(-s, kind="stable"))
+
+
+def complete(order, base):
+    seen = set(order)
+    return list(order) + [j for j in base if j not in seen]
+
+
+# ---------------------------------------------------------------- baselines
+def mmr(r, E, lam, K=10):
+    rel = minmax(r)
+    sim = E @ E.T
+    S, rem = [], set(range(len(r)))
+    for _ in range(min(K, len(r))):
+        if not S:
+            j = int(np.argmax(rel))
+        else:
+            red = sim[:, S].max(1)
+            cand = np.array(sorted(rem))
+            v = lam * rel[cand] - (1 - lam) * red[cand]
+            j = int(cand[np.argmax(v)])
+        S.append(j)
+        rem.discard(j)
+    return complete(S, by_score(r))
+
+
+def fused(r, E, q, alpha):
+    """z-score fusion of cross-encoder logits and dense cosine; alpha = weight on dense."""
+    d = E @ q
+    z = lambda x: (x - x.mean()) / (x.std() + 1e-9)
+    return (1 - alpha) * z(r) + alpha * z(d)
+
+
+def zfuse(r, E, q, alpha):
+    return by_score(fused(r, E, q, alpha))
+
+
+# ---------------------------------------------------------------- information-max selector
+def _condvar(M, S):
+    if not S:
+        return np.diag(M).copy()
+    c = cho_factor(M[np.ix_(S, S)], lower=True)
+    A = cho_solve(c, M[S, :])
+    return np.diag(M) - np.einsum("ij,ij->j", M[S, :], A)
+
+
+def infomax(r, E, q, rank=4, gamma=2.0, rho=1.0, sigma2=0.1, beta=0.0, delta=np.inf, K=10,
+            return_gains=False):
+    """Greedy I(z; y_j | y_S) with z = query-aspect latent and correlated (nuisance) noise.
+
+    y_j = s_j <U^T e_j, z> + <g_j, u> + eps_j,  z ~ N(0, Lambda), u ~ N(0, beta I), eps ~ N(0, sigma2)
+      e_j  : pool-centred unit embedding; U: [query direction, top relevance-weighted pool directions]
+      s_j^2: minmax(cross-encoder)^gamma ; g_j = (I - U U^T) e_j  (off-aspect content -> shared noise)
+    Anchored to the cross-encoder order: the CE top remaining item b is replaced by the max-gain
+    item only if gain(j*) > (1 + delta) * gain(b).  delta = inf reproduces the cross-encoder.
+    """
+    n = len(r)
+    base = by_score(r)
+    w = minmax(r) ** gamma
+    mu = E.mean(0)
+    Ec = E - mu
+    Ec /= np.linalg.norm(Ec, axis=1, keepdims=True) + 1e-9
+    qc = q - mu
+    qc /= np.linalg.norm(qc) + 1e-9
+    U = [qc]
+    lam = [1.0]
+    if rank > 1:
+        Wr = np.sqrt(w)[:, None] * (Ec - np.outer(Ec @ qc, qc))
+        _, sv, Vt = np.linalg.svd(Wr, full_matrices=False)
+        m = min(rank - 1, len(sv))
+        U += list(Vt[:m])
+        lam += list(rho * (sv[:m] / (sv[0] + 1e-9)) ** 2)
+    U = np.array(U).T                      # d x r
+    Lam = np.diag(lam)
+    P = Ec @ U                             # n x r aspect coordinates
+    H = np.sqrt(w)[:, None] * P
+    G = Ec - P @ U.T                       # off-aspect residual
+    R = beta * (G @ G.T) + sigma2 * np.eye(n)
+    C = H @ Lam @ H.T + R
+    S, rem, gains_log = [], np.ones(n, bool), []
+    for t in range(min(K, n)):
+        g = 0.5 * np.log(np.maximum(_condvar(C, S), 1e-12) / np.maximum(_condvar(R, S), 1e-12))
+        b = next(j for j in base if rem[j])
+        if t == 0:
+            j = b
+        else:
+            gm = np.where(rem, g, -np.inf)
+            js = int(np.argmax(gm))
+            j = js if gm[js] > (1 + delta) * g[b] + 1e-12 else b
+        gains_log.append(float(g[j]))
+        S.append(j)
+        rem[j] = False
+    order = complete(S, base)
+    return (order, gains_log) if return_gains else order

requirements.txt

@@ -0,0 +1,6 @@
+torch
+sentence-transformers>=3.0
+numpy
+scipy
+pandas
+spaces

results/beir_results.json

@@ -0,0 +1,528 @@
+{
+ "tuned_params_dev_only": {
+  "MMR": {
+   "lam": 0.95
+  },
+  "Fusion": {
+   "alpha": 0.5
+  },
+  "InfoMax": {
+   "rank": 1,
+   "gamma": 4.0,
+   "rho": 0.3,
+   "sigma2": 0.03,
+   "beta": 0.0,
+   "delta": 0.0
+  },
+  "InfoMax+F": {
+   "rank": 4,
+   "gamma": 8.0,
+   "rho": 1.0,
+   "sigma2": 0.03,
+   "beta": 3.0,
+   "delta": 0.0,
+   "alpha": 0.5
+  }
+ },
+ "ndcg_recall_test": {
+  "scifact.test|BM25": {
+   "nDCG@3": 0.6431199582422674,
+   "R@3": 0.7109444444444445,
+   "nDCG@5": 0.6596725941258738,
+   "R@5": 0.7526111111111111,
+   "nDCG@10": 0.6789590258685955,
+   "R@10": 0.8077222222222222
+  },
+  "scifact.test|Dense": {
+   "nDCG@3": 0.6163483357076979,
+   "R@3": 0.6764444444444444,
+   "nDCG@5": 0.6494822542293248,
+   "R@5": 0.7552777777777778,
+   "nDCG@10": 0.6775772341008127,
+   "R@10": 0.8395555555555556
+  },
+  "scifact.test|CE": {
+   "nDCG@3": 0.6355722307788756,
+   "R@3": 0.6832777777777779,
+   "nDCG@5": 0.6632119878395338,
+   "R@5": 0.7501666666666666,
+   "nDCG@10": 0.6762940264571006,
+   "R@10": 0.7884444444444444
+  },
+  "scifact.test|MMR": {
+   "nDCG@3": 0.6386774537671235,
+   "R@3": 0.6916111111111112,
+   "nDCG@5": 0.6626760413062829,
+   "R@5": 0.7501666666666666,
+   "nDCG@10": 0.6737188913639833,
+   "R@10": 0.7817777777777778
+  },
+  "scifact.test|Fusion": {
+   "nDCG@3": 0.6710598597617542,
+   "R@3": 0.716,
+   "nDCG@5": 0.6884192327495369,
+   "R@5": 0.7627777777777777,
+   "nDCG@10": 0.7120976163008534,
+   "R@10": 0.8312222222222222
+  },
+  "scifact.test|InfoMax": {
+   "nDCG@3": 0.6530626071442321,
+   "R@3": 0.7107222222222221,
+   "nDCG@5": 0.6737190144089154,
+   "R@5": 0.7646111111111111,
+   "nDCG@10": 0.6948257751035408,
+   "R@10": 0.8252222222222222
+  },
+  "scifact.test|InfoMax+F": {
+   "nDCG@3": 0.6725576912860219,
+   "R@3": 0.7196111111111111,
+   "nDCG@5": 0.6932716623902917,
+   "R@5": 0.7727777777777777,
+   "nDCG@10": 0.7155111299403955,
+   "R@10": 0.8368888888888889
+  },
+  "nfcorpus.test|BM25": {
+   "nDCG@3": 0.3823382822268318,
+   "R@3": 0.10380628945641764,
+   "nDCG@5": 0.3584927187341618,
+   "R@5": 0.12331434900055137,
+   "nDCG@10": 0.3261466393934428,
+   "R@10": 0.1564336372752942
+  },
+  "nfcorpus.test|Dense": {
+   "nDCG@3": 0.39224081480653206,
+   "R@3": 0.1054475795114919,
+   "nDCG@5": 0.3673855429222189,
+   "R@5": 0.12987323687278512,
+   "nDCG@10": 0.329672569677747,
+   "R@10": 0.1573953366222301
+  },
+  "nfcorpus.test|CE": {
+   "nDCG@3": 0.4229396163803222,
+   "R@3": 0.11255578971537836,
+   "nDCG@5": 0.3853234004565943,
+   "R@5": 0.13310119029572776,
+   "nDCG@10": 0.3458646508247627,
+   "R@10": 0.1623880974154255
+  },
+  "nfcorpus.test|MMR": {
+   "nDCG@3": 0.4203315776968791,
+   "R@3": 0.11215193296551876,
+   "nDCG@5": 0.3833962466547408,
+   "R@5": 0.13296678402481363,
+   "nDCG@10": 0.34496455529201764,
+   "R@10": 0.16095020847144315
+  },
+  "nfcorpus.test|Fusion": {
+   "nDCG@3": 0.4338920998144631,
+   "R@3": 0.11740153719863461,
+   "nDCG@5": 0.39540063670679393,
+   "R@5": 0.13643880315760448,
+   "nDCG@10": 0.35128231451281,
+   "R@10": 0.16419803765699678
+  },
+  "nfcorpus.test|InfoMax": {
+   "nDCG@3": 0.4313256861056477,
+   "R@3": 0.11718381080918494,
+   "nDCG@5": 0.39456391244419725,
+   "R@5": 0.13520038656020075,
+   "nDCG@10": 0.3459015527645562,
+   "R@10": 0.16016548330124453
+  },
+  "nfcorpus.test|InfoMax+F": {
+   "nDCG@3": 0.4288736628251464,
+   "R@3": 0.11576876191681756,
+   "nDCG@5": 0.3937758274378537,
+   "R@5": 0.13596229095231047,
+   "nDCG@10": 0.3521629267149766,
+   "R@10": 0.1636849294853285
+  },
+  "fiqa.test|BM25": {
+   "nDCG@3": 0.2128203105525299,
+   "R@3": 0.19910159285159282,
+   "nDCG@5": 0.22466852535393944,
+   "R@5": 0.24186197889901592,
+   "nDCG@10": 0.24313170715385354,
+   "R@10": 0.29682660018308166
+  },
+  "fiqa.test|Dense": {
+   "nDCG@3": 0.3238282228795522,
+   "R@3": 0.29630720359887025,
+   "nDCG@5": 0.3319682762608849,
+   "R@5": 0.3457364780281447,
+   "nDCG@10": 0.3454552348536008,
+   "R@10": 0.3935277171388282
+  },
+  "fiqa.test|CE": {
+   "nDCG@3": 0.2999096832508026,
+   "R@3": 0.28151820316172166,
+   "nDCG@5": 0.30844500797059193,
+   "R@5": 0.3263454052805905,
+   "nDCG@10": 0.3223933548101685,
+   "R@10": 0.3740349885025811
+  },
+  "fiqa.test|MMR": {
+   "nDCG@3": 0.29802765556267957,
+   "R@3": 0.2794605899929974,
+   "nDCG@5": 0.3072794992153618,
+   "R@5": 0.3258438620707139,
+   "nDCG@10": 0.3208304753270173,
+   "R@10": 0.3724917786260379
+  },
+  "fiqa.test|Fusion": {
+   "nDCG@3": 0.33832130422615464,
+   "R@3": 0.3066681432422173,
+   "nDCG@5": 0.3422595845439949,
+   "R@5": 0.3483966779105668,
+   "nDCG@10": 0.35845454570372476,
+   "R@10": 0.40362110026461884
+  },
+  "fiqa.test|InfoMax": {
+   "nDCG@3": 0.31712114265211333,
+   "R@3": 0.29967654516265624,
+   "nDCG@5": 0.32570617362477383,
+   "R@5": 0.35026567653882473,
+   "nDCG@10": 0.33689988451188513,
+   "R@10": 0.39294472674102304
+  },
+  "fiqa.test|InfoMax+F": {
+   "nDCG@3": 0.3378956061977511,
+   "R@3": 0.3066577326993994,
+   "nDCG@5": 0.34026758610658164,
+   "R@5": 0.3445141021529911,
+   "nDCG@10": 0.35772802740748577,
+   "R@10": 0.4016650873132354
+  },
+  "scidocs.test|BM25": {
+   "nDCG@3": 0.13953650967789885,
+   "R@3": 0.07848333333333334,
+   "nDCG@5": 0.12325782261102586,
+   "R@5": 0.10913333333333335,
+   "nDCG@10": 0.1489871672519904,
+   "R@10": 0.15653333333333336
+  },
+  "scidocs.test|Dense": {
+   "nDCG@3": 0.19532992037642224,
+   "R@3": 0.11063333333333332,
+   "nDCG@5": 0.17028354531362982,
+   "R@5": 0.15068333333333334,
+   "nDCG@10": 0.19872779902999924,
+   "R@10": 0.20253333333333337
+  },
+  "scidocs.test|CE": {
+   "nDCG@3": 0.1587600182121037,
+   "R@3": 0.08903333333333333,
+   "nDCG@5": 0.13789532380163394,
+   "R@5": 0.12068333333333334,
+   "nDCG@10": 0.16390846021956673,
+   "R@10": 0.16848333333333335
+  },
+  "scidocs.test|MMR": {
+   "nDCG@3": 0.15687177247920608,
+   "R@3": 0.08763333333333334,
+   "nDCG@5": 0.13595627703935984,
+   "R@5": 0.11843333333333333,
+   "nDCG@10": 0.16267991567191226,
+   "R@10": 0.16768333333333335
+  },
+  "scidocs.test|Fusion": {
+   "nDCG@3": 0.18745280547233126,
+   "R@3": 0.10563333333333334,
+   "nDCG@5": 0.16246159640953142,
+   "R@5": 0.1428333333333333,
+   "nDCG@10": 0.1936768980934276,
+   "R@10": 0.19993333333333332
+  },
+  "scidocs.test|InfoMax": {
+   "nDCG@3": 0.177,
+   "R@3": 0.10263333333333333,
+   "nDCG@5": 0.15042113247744018,
+   "R@5": 0.13338333333333333,
+   "nDCG@10": 0.1777698710919696,
+   "R@10": 0.18303333333333335
+  },
+  "scidocs.test|InfoMax+F": {
+   "nDCG@3": 0.18716785439017192,
+   "R@3": 0.10523333333333333,
+   "nDCG@5": 0.16395062501621754,
+   "R@5": 0.14468333333333333,
+   "nDCG@10": 0.19440802729822262,
+   "R@10": 0.20011666666666666
+  },
+  "trec-covid.test|BM25": {
+   "nDCG@3": 0.726536063698862,
+   "R@3": 0.006129608693906591,
+   "nDCG@5": 0.6931832808434869,
+   "R@5": 0.009495638499031355,
+   "nDCG@10": 0.645142792459132,
+   "R@10": 0.017583242371784345
+  },
+  "trec-covid.test|Dense": {
+   "nDCG@3": 0.7977649146579519,
+   "R@3": 0.0068480237477636895,
+   "nDCG@5": 0.7729841377993071,
+   "R@5": 0.011004464313449034,
+   "nDCG@10": 0.7502476986052402,
+   "R@10": 0.021582899500449337
+  },
+  "trec-covid.test|CE": {
+   "nDCG@3": 0.7835752445892035,
+   "R@3": 0.006500822879976922,
+   "nDCG@5": 0.7809746080465918,
+   "R@5": 0.01090665623000931,
+   "nDCG@10": 0.7443051348613305,
+   "R@10": 0.020647277435126458
+  },
+  "trec-covid.test|MMR": {
+   "nDCG@3": 0.7872617974664727,
+   "R@3": 0.006422345220043002,
+   "nDCG@5": 0.7700725718513098,
+   "R@5": 0.01065642470462764,
+   "nDCG@10": 0.7379188447063662,
+   "R@10": 0.020387648913364725
+  },
+  "trec-covid.test|Fusion": {
+   "nDCG@3": 0.8546927872602277,
+   "R@3": 0.0070038606633723065,
+   "nDCG@5": 0.8204160562167376,
+   "R@5": 0.011299291451340282,
+   "nDCG@10": 0.7971106002234066,
+   "R@10": 0.022138455639211664
+  },
+  "trec-covid.test|InfoMax": {
+   "nDCG@3": 0.7682680318494312,
+   "R@3": 0.006316769874605324,
+   "nDCG@5": 0.7494877534030594,
+   "R@5": 0.010045793623235813,
+   "nDCG@10": 0.7130600737191856,
+   "R@10": 0.0193403737920575
+  },
+  "trec-covid.test|InfoMax+F": {
+   "nDCG@3": 0.8376536063698862,
+   "R@3": 0.007084200669747274,
+   "nDCG@5": 0.8241432621592721,
+   "R@5": 0.011455199871392525,
+   "nDCG@10": 0.7896132829384864,
+   "R@10": 0.021971457445884583
+  }
+ },
+ "mrr_recall_test": {
+  "scifact.test|BM25|MRR@10": 0.6436679894179895,
+  "scifact.test|BM25|R@3": 0.7109444444444445,
+  "scifact.test|BM25|R@5": 0.7526111111111111,
+  "scifact.test|BM25|R@10": 0.8077222222222222,
+  "scifact.test|Dense|MRR@10": 0.6298108465608465,
+  "scifact.test|Dense|R@3": 0.6764444444444444,
+  "scifact.test|Dense|R@5": 0.7552777777777778,
+  "scifact.test|Dense|R@10": 0.8395555555555556,
+  "scifact.test|CE|MRR@10": 0.6487526455026454,
+  "scifact.test|CE|R@3": 0.6832777777777779,
+  "scifact.test|CE|R@5": 0.7501666666666668,
+  "scifact.test|CE|R@10": 0.7884444444444444,
+  "scifact.test|MMR|MRR@10": 0.6475079365079364,
+  "scifact.test|MMR|R@3": 0.6916111111111112,
+  "scifact.test|MMR|R@5": 0.7501666666666668,
+  "scifact.test|MMR|R@10": 0.7817777777777778,
+  "scifact.test|Fusion|MRR@10": 0.6794523809523809,
+  "scifact.test|Fusion|R@3": 0.7160000000000001,
+  "scifact.test|Fusion|R@5": 0.7627777777777778,
+  "scifact.test|Fusion|R@10": 0.8312222222222222,
+  "scifact.test|InfoMax|MRR@10": 0.6590330687830689,
+  "scifact.test|InfoMax|R@3": 0.7107222222222223,
+  "scifact.test|InfoMax|R@5": 0.7646111111111111,
+  "scifact.test|InfoMax|R@10": 0.8252222222222222,
+  "scifact.test|InfoMax+F|MRR@10": 0.6815767195767196,
+  "scifact.test|InfoMax+F|R@3": 0.7196111111111111,
+  "scifact.test|InfoMax+F|R@5": 0.7727777777777778,
+  "scifact.test|InfoMax+F|R@10": 0.8368888888888889,
+  "nfcorpus.test|BM25|MRR@10": 0.5234642979999017,
+  "nfcorpus.test|BM25|R@3": 0.10380628945641764,
+  "nfcorpus.test|BM25|R@5": 0.12331434900055138,
+  "nfcorpus.test|BM25|R@10": 0.15643363727529416,
+  "nfcorpus.test|Dense|MRR@10": 0.5332927908005307,
+  "nfcorpus.test|Dense|R@3": 0.10544757951149188,
+  "nfcorpus.test|Dense|R@5": 0.12987323687278512,
+  "nfcorpus.test|Dense|R@10": 0.1573953366222301,
+  "nfcorpus.test|CE|MRR@10": 0.577046783625731,
+  "nfcorpus.test|CE|R@3": 0.11255578971537837,
+  "nfcorpus.test|CE|R@5": 0.13310119029572776,
+  "nfcorpus.test|CE|R@10": 0.16238809741542548,
+  "nfcorpus.test|MMR|MRR@10": 0.5772875816993465,
+  "nfcorpus.test|MMR|R@3": 0.11215193296551876,
+  "nfcorpus.test|MMR|R@5": 0.13296678402481366,
+  "nfcorpus.test|MMR|R@10": 0.16095020847144315,
+  "nfcorpus.test|Fusion|MRR@10": 0.5721055088702147,
+  "nfcorpus.test|Fusion|R@3": 0.11740153719863461,
+  "nfcorpus.test|Fusion|R@5": 0.13643880315760448,
+  "nfcorpus.test|Fusion|R@10": 0.16419803765699678,
+  "nfcorpus.test|InfoMax|MRR@10": 0.5800432453683229,
+  "nfcorpus.test|InfoMax|R@3": 0.11718381080918494,
+  "nfcorpus.test|InfoMax|R@5": 0.13520038656020075,
+  "nfcorpus.test|InfoMax|R@10": 0.1601654833012445,
+  "nfcorpus.test|InfoMax+F|MRR@10": 0.5687982210428031,
+  "nfcorpus.test|InfoMax+F|R@3": 0.11576876191681756,
+  "nfcorpus.test|InfoMax+F|R@5": 0.13596229095231044,
+  "nfcorpus.test|InfoMax+F|R@10": 0.1636849294853285,
+  "fiqa.test|BM25|MRR@10": 0.3025493582206545,
+  "fiqa.test|BM25|R@3": 0.19910159285159287,
+  "fiqa.test|BM25|R@5": 0.24186197889901592,
+  "fiqa.test|BM25|R@10": 0.2968266001830817,
+  "fiqa.test|Dense|MRR@10": 0.435741230648638,
+  "fiqa.test|Dense|R@3": 0.29630720359887025,
+  "fiqa.test|Dense|R@5": 0.34573647802814467,
+  "fiqa.test|Dense|R@10": 0.3935277171388283,
+  "fiqa.test|CE|MRR@10": 0.4014948314716833,
+  "fiqa.test|CE|R@3": 0.28151820316172166,
+  "fiqa.test|CE|R@5": 0.3263454052805905,
+  "fiqa.test|CE|R@10": 0.3740349885025811,
+  "fiqa.test|MMR|MRR@10": 0.4007728297080149,
+  "fiqa.test|MMR|R@3": 0.27946058999299744,
+  "fiqa.test|MMR|R@5": 0.32584386207071386,
+  "fiqa.test|MMR|R@10": 0.3724917786260379,
+  "fiqa.test|Fusion|MRR@10": 0.45054992161473634,
+  "fiqa.test|Fusion|R@3": 0.3066681432422173,
+  "fiqa.test|Fusion|R@5": 0.34839667791056683,
+  "fiqa.test|Fusion|R@10": 0.40362110026461884,
+  "fiqa.test|InfoMax|MRR@10": 0.41511427101704873,
+  "fiqa.test|InfoMax|R@3": 0.29967654516265624,
+  "fiqa.test|InfoMax|R@5": 0.35026567653882473,
+  "fiqa.test|InfoMax|R@10": 0.39294472674102304,
+  "fiqa.test|InfoMax+F|MRR@10": 0.4490440672153635,
+  "fiqa.test|InfoMax+F|R@3": 0.30665773269939933,
+  "fiqa.test|InfoMax+F|R@5": 0.34451410215299105,
+  "fiqa.test|InfoMax+F|R@10": 0.4016650873132354,
+  "scidocs.test|BM25|MRR@10": 0.26165238095238097,
+  "scidocs.test|BM25|R@3": 0.07848333333333334,
+  "scidocs.test|BM25|R@5": 0.10913333333333333,
+  "scidocs.test|BM25|R@10": 0.15653333333333336,
+  "scidocs.test|Dense|MRR@10": 0.3502642857142857,
+  "scidocs.test|Dense|R@3": 0.11063333333333332,
+  "scidocs.test|Dense|R@5": 0.15068333333333334,
+  "scidocs.test|Dense|R@10": 0.20253333333333337,
+  "scidocs.test|CE|MRR@10": 0.29335079365079364,
+  "scidocs.test|CE|R@3": 0.08903333333333333,
+  "scidocs.test|CE|R@5": 0.12068333333333334,
+  "scidocs.test|CE|R@10": 0.16848333333333335,
+  "scidocs.test|MMR|MRR@10": 0.29261507936507936,
+  "scidocs.test|MMR|R@3": 0.08763333333333333,
+  "scidocs.test|MMR|R@5": 0.11843333333333333,
+  "scidocs.test|MMR|R@10": 0.16768333333333335,
+  "scidocs.test|Fusion|MRR@10": 0.336834126984127,
+  "scidocs.test|Fusion|R@3": 0.10563333333333334,
+  "scidocs.test|Fusion|R@5": 0.14283333333333334,
+  "scidocs.test|Fusion|R@10": 0.19993333333333332,
+  "scidocs.test|InfoMax|MRR@10": 0.3088793650793651,
+  "scidocs.test|InfoMax|R@3": 0.10263333333333334,
+  "scidocs.test|InfoMax|R@5": 0.13338333333333333,
+  "scidocs.test|InfoMax|R@10": 0.18303333333333335,
+  "scidocs.test|InfoMax+F|MRR@10": 0.33797738095238095,
+  "scidocs.test|InfoMax+F|R@3": 0.10523333333333333,
+  "scidocs.test|InfoMax+F|R@5": 0.14468333333333333,
+  "scidocs.test|InfoMax+F|R@10": 0.20011666666666666,
+  "trec-covid.test|BM25|MRR@10": 0.8966666666666666,
+  "trec-covid.test|BM25|R@3": 0.006129608693906592,
+  "trec-covid.test|BM25|R@5": 0.009495638499031355,
+  "trec-covid.test|BM25|R@10": 0.017583242371784342,
+  "trec-covid.test|Dense|MRR@10": 0.9066666666666667,
+  "trec-covid.test|Dense|R@3": 0.0068480237477636895,
+  "trec-covid.test|Dense|R@5": 0.011004464313449036,
+  "trec-covid.test|Dense|R@10": 0.021582899500449337,
+  "trec-covid.test|CE|MRR@10": 0.8880000000000001,
+  "trec-covid.test|CE|R@3": 0.006500822879976923,
+  "trec-covid.test|CE|R@5": 0.01090665623000931,
+  "trec-covid.test|CE|R@10": 0.02064727743512646,
+  "trec-covid.test|MMR|MRR@10": 0.8973333333333333,
+  "trec-covid.test|MMR|R@3": 0.006422345220043003,
+  "trec-covid.test|MMR|R@5": 0.010656424704627639,
+  "trec-covid.test|MMR|R@10": 0.020387648913364732,
+  "trec-covid.test|Fusion|MRR@10": 0.95,
+  "trec-covid.test|Fusion|R@3": 0.007003860663372309,
+  "trec-covid.test|Fusion|R@5": 0.011299291451340282,
+  "trec-covid.test|Fusion|R@10": 0.022138455639211668,
+  "trec-covid.test|InfoMax|MRR@10": 0.8983333333333333,
+  "trec-covid.test|InfoMax|R@3": 0.006316769874605324,
+  "trec-covid.test|InfoMax|R@5": 0.010045793623235811,
+  "trec-covid.test|InfoMax|R@10": 0.0193403737920575,
+  "trec-covid.test|InfoMax+F|MRR@10": 0.9533333333333335,
+  "trec-covid.test|InfoMax+F|R@3": 0.007084200669747275,
+  "trec-covid.test|InfoMax+F|R@5": 0.011455199871392525,
+  "trec-covid.test|InfoMax+F|R@10": 0.02197145744588459
+ },
+ "ci95_infomaxF_minus_baseline_ndcg10": {
+  "scifact.test|CE": [
+   0.039217103483294935,
+   0.022909537675935246,
+   0.055961078282767175
+  ],
+  "scifact.test|MMR": [
+   0.04179223857641219,
+   0.025205164715028892,
+   0.05874056517241845
+  ],
+  "scifact.test|Fusion": [
+   0.0034135136395420356,
+   -0.0009101887412512667,
+   0.00797676953601148
+  ],
+  "nfcorpus.test|CE": [
+   0.006298275890213895,
+   0.00014448374247222494,
+   0.012741372907464514
+  ],
+  "nfcorpus.test|MMR": [
+   0.007198371422958963,
+   0.0009795059674302584,
+   0.013655280658479303
+  ],
+  "nfcorpus.test|Fusion": [
+   0.0008806122021665923,
+   -0.0021193370946726207,
+   0.004199293800506931
+  ],
+  "fiqa.test|CE": [
+   0.035334672597317214,
+   0.026161183183126945,
+   0.04471607186251431
+  ],
+  "fiqa.test|MMR": [
+   0.03689755208046848,
+   0.02775011299883495,
+   0.04613706029287447
+  ],
+  "fiqa.test|Fusion": [
+   -0.0007265182962389718,
+   -0.0034145826445293446,
+   0.0018706218114601444
+  ],
+  "scidocs.test|CE": [
+   0.030499567078655924,
+   0.02536860036621575,
+   0.03573024236108956
+  ],
+  "scidocs.test|MMR": [
+   0.03172811162631037,
+   0.02661644582553447,
+   0.0369284625702656
+  ],
+  "scidocs.test|Fusion": [
+   0.0007311292047950231,
+   -0.0011978684641854693,
+   0.0027101133301523487
+  ],
+  "trec-covid.test|CE": [
+   0.045308148077155753,
+   0.008178194738087384,
+   0.0853383899434483
+  ],
+  "trec-covid.test|MMR": [
+   0.05169443823212016,
+   0.014252145682437763,
+   0.09173759710245157
+  ],
+  "trec-covid.test|Fusion": [
+   -0.007497317284920146,
+   -0.021556051504264526,
+   0.006579707248007931
+  ]
+ }
+}

results/beir_results.md

@@ -0,0 +1,75 @@
+## BEIR results (test splits)
+
+Setup: BM25 (k1=0.9, b=0.4, Porter) top-50 candidates per query; MS MARCO MiniLM-L6-v2 cross-encoder (max length 256);
+all-MiniLM-L6-v2 embeddings. All knobs tuned once on dev splits only (SciFact train, NFCorpus dev, FiQA dev; 200 queries each)
+and applied unchanged to all five test sets.
+
+- **CE**: cross-encoder ranking (the standard reranker).
+- **InfoMax**: information-maximising greedy selection, relevance = CE.
+- **InfoMax+F**: same selector, relevance = z-score fusion of CE and dense cosine.
+- **BM25 / Dense**: first-stage scores of the same pool (no reranker). **MMR**: CE relevance, dense redundancy, λ=0.95.
+  **Fusion**: z-score average of CE and dense cosine (α=0.5), no set selection.
+
+### nDCG@10
+
+| Dataset | BM25 | Dense | CE | MMR | Fusion | InfoMax | InfoMax+F |
+|---|---|---|---|---|---|---|---|
+| SciFact | 0.6790 | 0.6776 | 0.6763 | 0.6737 | 0.7121 | 0.6948 | **0.7155** |
+| NFCorpus | 0.3261 | 0.3297 | 0.3459 | 0.3450 | 0.3513 | 0.3459 | **0.3522** |
+| FiQA | 0.2431 | 0.3455 | 0.3224 | 0.3208 | **0.3585** | 0.3369 | 0.3577 |
+| SciDocs | 0.1490 | **0.1987** | 0.1639 | 0.1627 | 0.1937 | 0.1778 | 0.1944 |
+| TREC-COVID | 0.6451 | 0.7502 | 0.7443 | 0.7379 | **0.7971** | 0.7131 | 0.7896 |
+
+### MRR@10
+
+| Dataset | BM25 | Dense | CE | MMR | Fusion | InfoMax | InfoMax+F |
+|---|---|---|---|---|---|---|---|
+| SciFact | 0.6437 | 0.6298 | 0.6488 | 0.6475 | 0.6795 | 0.6590 | **0.6816** |
+| NFCorpus | 0.5235 | 0.5333 | 0.5770 | 0.5773 | 0.5721 | **0.5800** | 0.5688 |
+| FiQA | 0.3025 | 0.4357 | 0.4015 | 0.4008 | **0.4505** | 0.4151 | 0.4490 |
+| SciDocs | 0.2617 | **0.3503** | 0.2934 | 0.2926 | 0.3368 | 0.3089 | 0.3380 |
+| TREC-COVID | 0.8967 | 0.9067 | 0.8880 | 0.8973 | 0.9500 | 0.8983 | **0.9533** |
+
+### Recall@3
+
+| Dataset | BM25 | Dense | CE | MMR | Fusion | InfoMax | InfoMax+F |
+|---|---|---|---|---|---|---|---|
+| SciFact | 0.7109 | 0.6764 | 0.6833 | 0.6916 | 0.7160 | 0.7107 | **0.7196** |
+| NFCorpus | 0.1038 | 0.1054 | 0.1126 | 0.1122 | **0.1174** | 0.1172 | 0.1158 |
+| FiQA | 0.1991 | 0.2963 | 0.2815 | 0.2795 | **0.3067** | 0.2997 | 0.3067 |
+| SciDocs | 0.0785 | **0.1106** | 0.0890 | 0.0876 | 0.1056 | 0.1026 | 0.1052 |
+| TREC-COVID | 0.0061 | 0.0068 | 0.0065 | 0.0064 | 0.0070 | 0.0063 | **0.0071** |
+
+### Recall@5
+
+| Dataset | BM25 | Dense | CE | MMR | Fusion | InfoMax | InfoMax+F |
+|---|---|---|---|---|---|---|---|
+| SciFact | 0.7526 | 0.7553 | 0.7502 | 0.7502 | 0.7628 | 0.7646 | **0.7728** |
+| NFCorpus | 0.1233 | 0.1299 | 0.1331 | 0.1330 | **0.1364** | 0.1352 | 0.1360 |
+| FiQA | 0.2419 | 0.3457 | 0.3263 | 0.3258 | 0.3484 | **0.3503** | 0.3445 |
+| SciDocs | 0.1091 | **0.1507** | 0.1207 | 0.1184 | 0.1428 | 0.1334 | 0.1447 |
+| TREC-COVID | 0.0095 | 0.0110 | 0.0109 | 0.0107 | 0.0113 | 0.0100 | **0.0115** |
+
+### Recall@10
+
+| Dataset | BM25 | Dense | CE | MMR | Fusion | InfoMax | InfoMax+F |
+|---|---|---|---|---|---|---|---|
+| SciFact | 0.8077 | **0.8396** | 0.7884 | 0.7818 | 0.8312 | 0.8252 | 0.8369 |
+| NFCorpus | 0.1564 | 0.1574 | 0.1624 | 0.1610 | **0.1642** | 0.1602 | 0.1637 |
+| FiQA | 0.2968 | 0.3935 | 0.3740 | 0.3725 | **0.4036** | 0.3929 | 0.4017 |
+| SciDocs | 0.1565 | **0.2025** | 0.1685 | 0.1677 | 0.1999 | 0.1830 | 0.2001 |
+| TREC-COVID | 0.0176 | 0.0216 | 0.0206 | 0.0204 | **0.0221** | 0.0193 | 0.0220 |
+
+### InfoMax+F minus baseline, nDCG@10, paired bootstrap 95% CI
+
+| Dataset | vs CE | vs MMR | vs Fusion |
+|---|---|---|---|
+| SciFact | +0.0392 [+0.0229, +0.0560] | +0.0418 [+0.0252, +0.0587] | +0.0034 [-0.0009, +0.0080] |
+| NFCorpus | +0.0063 [+0.0001, +0.0127] | +0.0072 [+0.0010, +0.0137] | +0.0009 [-0.0021, +0.0042] |
+| FiQA | +0.0353 [+0.0262, +0.0447] | +0.0369 [+0.0278, +0.0461] | -0.0007 [-0.0034, +0.0019] |
+| SciDocs | +0.0305 [+0.0254, +0.0357] | +0.0317 [+0.0266, +0.0369] | +0.0007 [-0.0012, +0.0027] |
+| TREC-COVID | +0.0453 [+0.0082, +0.0853] | +0.0517 [+0.0143, +0.0917] | -0.0075 [-0.0216, +0.0066] |
+
+**Reading.** InfoMax+F beats the cross-encoder on nDCG@10 on all five datasets (intervals above zero). It ties Fusion everywhere:
+on these single-hop benchmarks the gain comes from adding dense relevance, not from set selection. nDCG does not penalise
+redundant relevant passages. TREC-COVID recall is near zero by construction (~494 relevant documents per query).