Implement BM25 Ranking Score

NLP · Linear Algebra · Data Processing

Hard

Problem

Theory

BM25 (Best Matching 25) is a ranking function used in information retrieval to estimate the relevance of documents to a given search query. It evolved from probabilistic retrieval models in the 1990s and remains the backbone of modern search engines. Despite the rise of neural approaches, BM25 often serves as a strong baseline and is used in hybrid retrieval systems.

The Evolution from TF-IDF

BM25 builds upon TF-IDF concepts while addressing its limitations:

TF-IDF approach: \text{TF-IDF}(t, d) = \text{tf}(t, d) \times \text{idf}(t)

Limitations of TF-IDF:

• Term frequency grows unboundedly for long documents
• No saturation effect - 10 occurrences scores twice as high as 5 occurrences
• Document length is not explicitly considered
• Long documents unfairly dominate rankings

BM25 solutions:

• Introduces term frequency saturation (diminishing returns)
• Adds explicit document length normalization
• Provides tunable parameters for different use cases

The Complete BM25 Formula

For a query Q containing terms q_1, q_2, ..., q_n and a document D:

\text{BM25}(D, Q) = \sum_{i=1}^{n} \text{IDF}(q_i) \cdot \frac{f(q_i, D) \cdot (k_1 + 1)}{f(q_i, D) + k_1 \cdot (1 - b + b \cdot \frac{|D|}{\text{avgdl}})}

Variables explained:

• f(q_i, D) = frequency of term $$ in document D
• |D| = length of document $$ (in words)
• \text{avgdl} = average document length across the entire corpus
• k_1 = term frequency saturation parameter (typical range: 1.2-2.0)
• b = document length normalization parameter (typical value: 0.75)

The IDF Component

Inverse Document Frequency penalizes common terms that appear in many documents:

\text{IDF}(q_i) = \ln\left(\frac{N - n(q_i) + 0.5}{n(q_i) + 0.5} + 1\right)

Variables:

• N = total number of documents in the corpus
• n(q_i) = number of documents containing term $$

Properties:

• Returns higher values for rare terms (more discriminative)
• The +0.5 smoothing terms prevent division by zero
• The outer +1 ensures non-negative values for all terms
• Common words like "the" get near-zero IDF, rare technical terms get high IDF

Understanding Parameter k_1 (Term Frequency Saturation)

The k_1 parameter controls how quickly additional term occurrences lose their impact:

Low k_1 (e.g., 0.5):

• Diminishing returns kick in quickly
• First occurrence of a term matters most
• Additional occurrences add little value
• Good when term presence is more important than frequency

High k_1 (e.g., 2.0):

• More linear relationship between frequency and score
• Multiple occurrences continue to boost relevance
• Better for domains where repetition indicates topical focus

Boundary cases:

• k_1 = 0: Score becomes binary (term present or absent)
• k_1 \to \infty: Approaches raw term frequency (no saturation)

Typical value: 1.2 balances saturation with frequency sensitivity

Understanding Parameter b (Length Normalization)

The b parameter controls how document length affects scoring:

b = 0 (no normalization):

• Long documents receive no penalty
• Favors comprehensive documents
• Risk: verbose documents dominate

b = 1 (full normalization):

• Score inversely proportional to document length
• Strongly favors concise documents
• Risk: short documents may rank too high

b = 0.75 (balanced):

• Moderate penalty for length
• Standard default for most applications
• Documents shorter than average get boosted
• Documents longer than average get penalized

The normalization term \frac{|D|}{\text{avgdl}} compares each document to the corpus average length.

Term Frequency Saturation Explained

The saturation function bounds the contribution of term frequency:

\text{saturation}(f) = \frac{f \cdot (k_1 + 1)}{f + k_1}

Behavior analysis (with k_1 = 1.5):

• f=1: contribution = 1.0
• f=2: contribution = 1.43 (not 2.0)
• f=5: contribution = 1.92 (not 5.0)
• f=10: contribution = 2.17
• f=100: contribution = 2.46
• f \to \infty: contribution approaches k_1 + 1 = 2.5

Key insight: The first few occurrences contribute significantly, but additional occurrences have diminishing returns. The maximum possible contribution is bounded by k_1 + 1.

Multi-Term Query Handling

BM25 scores for multi-term queries are computed additively:

\text{BM25}(D, Q) = \sum_{i=1}^{n} \text{BM25}_{q_i}(D)

Properties of additive scoring:

• Each query term contributes independently
• Documents matching multiple query terms score higher
• Allows precomputation: term scores can be computed once and summed at query time
• Efficient for inverted index implementations

Worked Example

Corpus (3 documents):

• D1: "the cat sat on the mat" (6 words)
• D2: "the dog ran in the park and played with the ball" (11 words)
• D3: "cat" (1 word)

Query: "cat"

Step 1 - Compute corpus statistics:

• N = 3 documents
• avgdl = (6 + 11 + 1) / 3 = 6 words
• n("cat") = 2 (appears in D1 and D3)

Step 2 - Calculate IDF("cat"): \ln\left(\frac{3 - 2 + 0.5}{2 + 0.5} + 1\right) = \ln\left(\frac{1.5}{2.5} + 1\right) = \ln(1.6) \approx 0.47

Step 3 - Calculate BM25 scores (k1=1.5, b=0.75):

For D1 (f=1, |D|=6):

• Length ratio: 6/6 = 1.0 (average length)
• Score contribution: \frac{1 \times 2.5}{1 + 1.5 \times (0.25 + 0.75 \times 1.0)} = \frac{2.5}{2.5} = 1.0
• Final: 0.47 × 1.0 = 0.47

For D2 (f=0): Score = 0 (term not present)

For D3 (f=1, |D|=1):

• Length ratio: 1/6 ≈ 0.167 (much shorter than average)
• Score contribution: \frac{1 \times 2.5}{1 + 1.5 \times (0.25 + 0.75 \times 0.167)} = \frac{2.5}{1.56} \approx 1.6
• Final: 0.47 × 1.6 = 0.75

Ranking result: D3 (0.75) > D1 (0.47) > D2 (0)

D3 scores highest because it is shorter than average, receiving a length normalization boost.

BM25 Variants

BM25+: Adds a small constant delta to prevent negative scores: \text{BM25+} = \text{BM25} + \delta

BM25L: Addresses over-penalization of long but relevant documents

BM25F: Extends to structured documents with multiple fields (title, body, anchor text) using per-field weights: \text{BM25F} = \sum_{\text{fields}} w_f \cdot \text{BM25}_f

Implementation Considerations

• Preprocessing pipeline: Tokenization, lowercasing, stemming/lemmatization, stopword removal all affect retrieval quality

• Index structure: Inverted index maps terms to document IDs and term frequencies for efficient retrieval

• Scoring efficiency: IDF values computed once during indexing; only TF-based computation needed at query time

• Parameter tuning: k_1 and b can be tuned on a validation set for specific domains

Where BM25 Shows Up

• Search Engines: Elasticsearch, Apache Lucene/Solr use BM25 as the default ranking function

• Document Retrieval: Legal document search, academic paper search, enterprise knowledge bases

• RAG Systems: Retrieval-Augmented Generation combines BM25 with dense retrieval to find context for large language models

• Question Answering: Initial retrieval step to find candidate passages before applying reading comprehension models

• E-commerce Search: Product search ranking based on query-product description matching

• Hybrid Search: BM25 combined with neural embeddings (dense-sparse hybrid) often outperforms either approach alone

• Email Search: Finding relevant emails by keyword matching with length-normalized scoring

• Code Search: Finding relevant code snippets or documentation in software repositories

Examples

Hints

1. Use one Counter per document for term frequencies.
2. Use Counter.update(set(document)) to count document frequencies.
3. Build a NumPy vector of one term's frequency across all documents before applying the formula.

Requirements

• Use the stated BM25 IDF and scoring formulas
• Count repeated query terms once
• Preserve document order
• Return a NumPy array of floating-point scores

Constraints

• docs is a list of token lists
• k_1 > 0 and 0 \le b \le 1
• Use NumPy and the Python standard library only

Starter Code

import math
from collections import Counter
import numpy as np

def bm25_score(query_tokens: list[str], docs: list[list[str]], k1: float = 1.2, b: float = 0.75) -> np.ndarray:
    """
    Returns a NumPy array with one score per document.
    """
    # Write code here
    pass

Test Cases