MMR Reranking

Start from brute force

The obvious approach is plain top-k: score every document by cos(q, d), sort, and take the first k. It is correct for its own objective, relevance, and cheap: O(n·d + n log n) time for n documents of dimension d, and O(n) space for the scores. It is wrong for this one, because it never compares documents with each other. On the example documents it returns [0, 1], and document 1 is a near-copy of document 0 (cosine 0.998), so half of the context budget buys a paraphrase of what is already there while the different document 2 is left out.

The approach

Build the list greedily. Each round, score every unselected candidate by lam · rel − (1 − lam) · red, where rel = cos(q, d) is its relevance and red is its maximum cosine to anything already selected, and keep the best. Max, not average, is the point: one near-copy in the selection is enough to make a candidate redundant.

The reference does exactly this. _cos is the cosine helper, selected grows by one index per round, red = max((_cos(docs[i], docs[j]) for j in selected), default=0.0) makes the first round pure relevance, and the strict score > best_score keeps the lower index on a tie. Its cost is recomputing every rel and every red from scratch each round. The optimal version computes rel once per document and keeps a running red[i]: after each pick it sets red[i] = max(red[i], cos(docs[i], picked)) against only the new document, which leaves every score identical.

Walk through an example

With q = [1, 0], docs = [[0.9, 0.1], [0.85, 0.15], [0.3, 0.8]], lam = 0.3 and k = 2, round 1 is pure relevance, and in round 2 document 1’s 0.998 similarity to document 0 sinks it below document 2, giving [0, 2].

Complexity

With relevance cached and a running maximum, selection is O(k·n·d) time, one new cosine per candidate per round, with O(n) extra space for rel and red. The reference recomputes the maximum over the whole selected set every round: O(k²·n·d) time, O(k) space. Plain top-k is O(n·d + n log n) time and O(n) space.

Edge cases to test

• An exact duplicate can survive at lam = 0.5: in mmr([1,0], [[1,0],[1,0],[0,1]], 0.5, 2) the copy scores 0.5·1 − 0.5·1 = 0, the orthogonal document scores 0.5·0 − 0.5·0 = 0, and the lower-index tie-break returns [0, 1].
• lam = 0 ignores relevance entirely: every first-round score is 0, so the first pick is index 0 whatever its relevance. mmr([1,0], [[0.3,0.8],[0.9,0.1]], 0.0, 1) returns [0], the less relevant document.
• lam = 1 is plain top-k: [0, 1, 2] on the example documents.
• k larger than the number of documents: once no candidate is left, the reference appends its -1 sentinel, so k = 4 on the three example documents returns [0, 1, 2, -1] in both languages. Clamp k first.
• A zero vector has no direction: Python’s _cos raises ZeroDivisionError, while TypeScript gets NaN, and because NaN > bestScore is always false that document is never picked; with no other candidate, the round returns -1.

Common mistakes

• Ranking on relevance alone — you get five near-duplicate passages and waste the budget.
• Setting λ wrong: all relevance is redundant, all diversity drifts off-topic.
• Computing the penalty against the query instead of against the already-selected items.
• Averaging the similarity to the selected set instead of taking the maximum, so one near-copy among several unrelated picks is diluted and slips through.

Follow-ups an interviewer asks next

How do you choose λ?

Tune it on real queries, measuring answer quality or coverage of distinct facts at a fixed context budget. The trade-off is sensitive to how similar your embeddings make documents: on the example documents, the second slot flips from document 2 to the near-copy document 1 once lam passes about 0.463. Defaults near 0.5 are common, but check them against your data.

Retrieval returns 1,000 candidates. How do you keep MMR fast?

Run it on a short list: take the top few dozen by relevance first and apply MMR to those, since a document far down the relevance ranking rarely wins a slot anyway. Cache relevance, keep the running maximum per candidate, and normalize vectors once so each cosine becomes a dot product; with the vectors in a matrix, a round is one matrix-vector product.

What other diversification methods are there?

Determinantal point processes choose a set whose similarity matrix has a large determinant, which rewards diversity across the whole set instead of one greedy step at a time. Simpler options cluster the candidates and take the best per cluster, or cap results per source document. MMR stays popular because it is one tunable knob and a few lines of code.

Must the redundancy term use the same embedding as relevance?

No. Relevance is query-to-document similarity, but redundancy is document-to-document similarity, which can be measured more cheaply, for example with n-gram overlap or a shared source id. In RAG the most common duplicates are overlapping chunks of the same document, and a source-id check catches those exactly.

Where this shows up

MMR diversifies RAG context and search results so the prompt is not stuffed with paraphrases of one passage — you cover more of the answer within the same token budget. It is a standard reranking step in retrieval and in extractive summarization.

Reference solution

import math

def _cos(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    return dot / (na * nb)

def mmr(q, docs, lam, k):
    selected = []
    while len(selected) < k:
        best_i, best_score = -1, -float('inf')
        for i in range(len(docs)):
            if i in selected:
                continue
            rel = _cos(q, docs[i])
            red = max((_cos(docs[i], docs[j]) for j in selected), default=0.0)
            score = lam * rel - (1 - lam) * red
            if score > best_score:
                best_i, best_score = i, score
        selected.append(best_i)
    return selected


print(mmr([1, 0], [[0.9, 0.1], [0.85, 0.15], [0.3, 0.8]], 0.3, 2))

function cos(a: number[], b: number[]): number {
  let dot = 0, na = 0, nb = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    na += a[i] * a[i];
    nb += b[i] * b[i];
  }
  return dot / (Math.sqrt(na) * Math.sqrt(nb));
}

function mmr(q: number[], docs: number[][], lam: number, k: number): number[] {
  const selected: number[] = [];
  while (selected.length < k) {
    let bestI = -1;
    let bestScore = -Infinity;
    for (let i = 0; i < docs.length; i++) {
      if (selected.includes(i)) continue;
      const rel = cos(q, docs[i]);
      const red = selected.length
        ? Math.max(...selected.map((j) => cos(docs[i], docs[j])))
        : 0;
      const score = lam * rel - (1 - lam) * red;
      if (score > bestScore) {
        bestScore = score;
        bestI = i;
      }
    }
    selected.push(bestI);
  }
  return selected;
}

console.log(mmr([1, 0], [[0.9, 0.1], [0.85, 0.15], [0.3, 0.8]], 0.3, 2));