Problem 113 · Retrieval

Maximal Marginal Relevance Diversification

MediumAI engineering

01Problem

Implement `mmr(query, docs, k, lam)` where `docs` maps id to embedding. Iteratively select the doc maximising `lam * sim(query, d) - (1 - lam) * max(sim(d, s) for s in selected)` (the max term is 0 when nothing is selected yet), using cosine similarity. Ties go to the smaller id. Return up to `k` ids in selection order.

02Key insight

Top-k by relevance alone returns near-duplicates (the same paragraph from three versions of a document), wasting context window on redundant evidence. MMR trades relevance for novelty with one knob: lam = 1 is pure relevance, lam = 0 is pure diversity. Caching pairwise similarities against the selected set keeps it O(k * n) instead of recomputing; in practice you run MMR on a reranked top-50 to top-100, not the whole corpus.

03Solution

solution.pypython23 lines

import math


def _cos(a: list[float], b: list[float]) -> float:
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(x * x for x in b))
    return 0.0 if na == 0 or nb == 0 else sum(x * y for x, y in zip(a, b)) / (na * nb)


def mmr(query: list[float], docs: dict[str, list[float]], k: int, lam: float = 0.5) -> list[str]:
    if not 0.0 <= lam <= 1.0:
        raise ValueError("lam must be in [0, 1]")
    rel = {d: _cos(query, v) for d, v in docs.items()}
    max_sim_to_selected = {d: 0.0 for d in docs}
    selected: list[str] = []
    remaining = set(docs)
    while remaining and len(selected) < k:
        best = min(remaining, key=lambda d: (-(lam * rel[d] - (1 - lam) * max_sim_to_selected[d]), d))
        selected.append(best)
        remaining.remove(best)
        for d in remaining:
            max_sim_to_selected[d] = max(max_sim_to_selected[d], _cos(docs[d], docs[best]))
    return selected

Try it first.

TimeO(k * n * d)

SpaceO(n)

tests.py · 7 passing

tests.pypython19 lines

docs = {
    "a": [1.0, 0.0, 0.0],
    "a_dup": [0.99, 0.01, 0.0],
    "b": [0.6, 0.8, 0.0],
    "c": [0.0, 0.0, 1.0],
}
q = [1.0, 0.2, 0.0]
assert mmr(q, docs, 4, lam=1.0) == sorted(docs, key=lambda d: -_cos(q, docs[d]))
div = mmr(q, docs, 2, lam=0.5)
assert div == ["a_dup", "b"]
assert mmr(q, docs, 2, lam=1.0) == ["a_dup", "a"]
assert set(mmr(q, docs, 10)) == set(docs)
assert mmr(q, {}, 3) == []
assert mmr(q, docs, 0) == []
try:
    mmr(q, docs, 2, lam=1.5)
    assert False
except ValueError:
    pass

04Staff-level follow-up

How do you pick lam? Tune it on an eval set where the metric rewards coverage (for example nugget recall or answer completeness), because nDCG with duplicate-tolerant labels will always pick lam = 1. Alternatives: cluster-then-pick, or dedupe by source document before reranking, which is cheaper and often enough.