Pointwise Mutual Information › Code Implementation
Language AI HandbookAll chapters Pointwise Mutual Information › Code ImplementationOne section · Sections in this chapter7 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement PMI from scratch, then visualize the differences between raw counts and PPMI to make the normalization effect concrete.
Setup and Data
We will build everything from a small but realistic text corpus with sentences from multiple semantic domains. Using multiple domains lets us see whether PMI correctly separates the financial meaning of "bank" from its geographic meaning, and whether it correctly identifies that "cat" and "dog" share more contexts with each other than with words from the financial domain.
import warnings
warnings.filterwarnings("ignore")
# Small corpus covering multiple semantic domains
corpus = [
"the cat chased the mouse",
"the dog chased the cat",
"the mouse ran from the cat",
"the dog barked at the cat",
"the cat sat on the mat",
"the dog sat on the rug",
"the mouse ate the cheese",
"cats and dogs are common pets",
"mice fear cats and dogs",
"the bank charges fees for services",
"the river bank was flooded",
"money flows through the bank",
"the fish swim in the river",
"the bank approved the loan",
"salmon swim upstream in the river",
]
# Tokenize
def tokenize(text):
return text.lower().split()
all_tokens = [token for sent in corpus for token in tokenize(sent)]# Build vocabulary
vocab = sorted(set(all_tokens))
word_to_idx = {w: i for i, w in enumerate(vocab)}
idx_to_word = {i: w for i, w in enumerate(vocab)}
V = len(vocab)Vocabulary size: 41 Total tokens: 83 Sample vocabulary: ['and', 'approved', 'are', 'at', 'ate', 'bank', 'barked', 'cat', 'cats', 'charges', 'chased', 'cheese']
Our small corpus has a manageable vocabulary, which lets us inspect the full co-occurrence and PPMI matrices directly. The two-domain structure (animals vs. banking/rivers) will let us verify that PMI correctly identifies cross-domain words like "bank" as having associations in both camps.
Building the Co-occurrence Matrix
We construct the raw co-occurrence matrix using a sliding window, as covered in the previous chapter on co-occurrence matrices. Each sentence contributes co-occurrence counts for all pairs of words within the window around each target word.
def build_cooccurrence_matrix(tokenized_corpus, word_to_idx, window_size=2):
V = len(word_to_idx)
M = np.zeros((V, V), dtype=np.float64)
for tokens in tokenized_corpus:
for i, word in enumerate(tokens):
if word not in word_to_idx:
continue
w_idx = word_to_idx[word]
start = max(0, i - window_size)
end = min(len(tokens), i + window_size + 1)
for j in range(start, end):
if i == j:
continue
ctx = tokens[j]
if ctx not in word_to_idx:
continue
c_idx = word_to_idx[ctx]
M[w_idx][c_idx] += 1
return M
tokenized_corpus = [tokenize(sent) for sent in corpus]
cooc_matrix = build_cooccurrence_matrix(
tokenized_corpus, word_to_idx, window_size=2
)Co-occurrence matrix shape: (41, 41) Total co-occurrence events: 242 Non-zero entries: 168 Sparsity: 90.0%
The high sparsity is typical: most word pairs never appear within a window of each other, so most matrix entries are zero. Even with a window size of 2, only a small fraction of all possible word pairs appear as neighbors in these sentences.
Computing PPMI
Now we compute the PPMI matrix from the raw co-occurrence counts. The key steps are computing the marginal and joint probabilities, forming the PMI ratio, and clipping at zero. The implementation supports both standard PPMI (with ) and shifted PPMI (with ), allowing easy experimentation with the shift parameter.
def compute_ppmi(cooc_matrix, k=1.0):
"""
Compute Positive PMI matrix from a raw co-occurrence matrix.
Args:
cooc_matrix: raw co-occurrence count matrix (V x V)
k: shift parameter (default 1 = standard PPMI; k > 1 = shifted PPMI)
Returns:
PPMI matrix with same shape
"""
# Total co-occurrence events
N = cooc_matrix.sum()
# Marginal: how often each word appears as target (row sums)
row_sums = cooc_matrix.sum(axis=1, keepdims=True) # shape (V, 1)
# Marginal: how often each word appears as context (column sums)
col_sums = cooc_matrix.sum(axis=0, keepdims=True) # shape (1, V)
# PMI = log2( count(w,c) * N / (count(w) * count(c)) )
with np.errstate(divide="ignore", invalid="ignore"):
numerator = cooc_matrix * N
denominator = row_sums * col_sums
pmi = np.where(
(denominator > 0) & (cooc_matrix > 0),
np.log2(numerator / denominator),
-np.inf,
)
# Apply shift for SPPMI: subtract log2(k) before clipping
if k > 1:
pmi = pmi - np.log2(k)
# PPMI: replace negatives (and -inf) with zero
ppmi = np.maximum(0, pmi)
return ppmi
ppmi_matrix = compute_ppmi(cooc_matrix, k=1)PPMI matrix shape: (41, 41) Non-zero entries: 168 Max PPMI value: 5.334 Mean PPMI (non-zero only): 2.656
The np.where call handles two tricky cases simultaneously. When cooc_matrix > 0 but denominator == 0, we have a marginal probability of zero, which should not occur in a well-formed matrix but can arise if a word has been zero-padded. When cooc_matrix == 0, the PMI is regardless of the denominator, so we assign directly. The final np.maximum(0, pmi) then clips both and any finite negative values to zero, producing the PPMI matrix.
Inspecting Word Associations
With the PPMI matrix built, we can look up specific word associations. Let's examine the top associations for a few target words and compare the PPMI scores against raw co-occurrence counts.
def get_top_associations(word, matrix, idx_to_word, word_to_idx, n=6):
"""Get top-n associated words for a given target word."""
if word not in word_to_idx:
return []
idx = word_to_idx[word]
scores = matrix[idx]
top_indices = np.argsort(scores)[::-1][:n]
return [(idx_to_word[i], scores[i]) for i in top_indices if scores[i] > 0]Top PPMI associations for 'cat': from PPMI=2.33 raw=1 at PPMI=2.33 raw=1 chased PPMI=2.33 raw=2 sat PPMI=1.33 raw=1 on PPMI=1.33 raw=1 Top PPMI associations for 'bank': flooded PPMI=3.33 raw=1 was PPMI=2.75 raw=1 through PPMI=2.33 raw=1 approved PPMI=2.33 raw=1 fees PPMI=2.33 raw=1 Top PPMI associations for 'dog': at PPMI=2.75 raw=1 barked PPMI=2.75 raw=1 chased PPMI=1.75 raw=1 sat PPMI=1.75 raw=1 on PPMI=1.75 raw=1 Top PPMI associations for 'river': was PPMI=3.53 raw=1 in PPMI=3.11 raw=2 bank PPMI=1.53 raw=1 the PPMI=0.92 raw=3
The PPMI scores reveal specific semantic relationships. "Cat" associates strongly with words like "chased," "mouse," and "mat," reflecting actual usage in the corpus. The word "bank" reveals both its financial and geographic contexts through different high-PMI context words. High-frequency function words like "the" and "and" get suppressed: even if "the" appears many times near "cat" in absolute terms, PMI recognizes that "the" appears near everything equally, so the association is unremarkable.
The side-by-side display of PPMI and raw counts makes this suppression concrete. A context word with a high raw count but low PPMI is a function word or very common noun. A context word with a high PPMI is semantically distinctive to the target word.
PMI vs. Raw Counts: A Direct Comparison
Let's directly compare rankings produced by raw counts versus PPMI. This makes the normalization effect concrete and shows exactly which words get promoted and which get demoted when you switch from raw frequency to association strength.
def compare_raw_vs_ppmi(
word, cooc_matrix, ppmi_matrix, idx_to_word, word_to_idx, n=6
):
"""Return top-n words by raw count and by PPMI for comparison."""
idx = word_to_idx[word]
raw_scores = cooc_matrix[idx]
ppmi_scores = ppmi_matrix[idx]
top_raw_idx = np.argsort(raw_scores)[::-1][:n]
top_raw = [
(idx_to_word[i], raw_scores[i])
for i in top_raw_idx
if raw_scores[i] > 0
]
top_ppmi_idx = np.argsort(ppmi_scores)[::-1][:n]
top_ppmi = [
(idx_to_word[i], ppmi_scores[i])
for i in top_ppmi_idx
if ppmi_scores[i] > 0
]
return top_raw, top_ppmiComparison for 'cat': Top by Raw Count Top by PPMI ---------------------------------------------------------------------- the (6) from (2.33) chased (2) at (2.33) from (1) chased (2.33) at (1) sat (1.33) sat (1) on (1.33) on (1) the (1.14)
Notice how raw counts favor high-frequency words that appear near everything, while PPMI promotes words with a specific association to the target word. This shift in ranking is the core benefit of PMI: you see the words that are truly distinctive to "cat" rather than simply common in the corpus.
Computing Word Similarity Using PPMI
Once we have PPMI vectors, we compute word similarity using cosine similarity. Two words are similar if their PPMI context distributions are aligned: they share the same high-association contexts, indicating they are used in similar linguistic environments.
Cosine similarity between two vectors and is defined as:
where:
- : the dot product, summing the products of corresponding dimensions
- : the Euclidean norm of
- : the Euclidean norm of
Cosine similarity ranges from to , though with PPMI vectors (which are non-negative), it ranges from to . A value of 1 means the two vectors point in exactly the same direction, indicating identical context distributions. A value of 0 means the vectors are orthogonal, indicating no shared high-association contexts.
def cosine_similarity(v1, v2):
"""Compute cosine similarity between two vectors."""
norm1 = np.linalg.norm(v1)
norm2 = np.linalg.norm(v2)
if norm1 == 0 or norm2 == 0:
return 0.0
return np.dot(v1, v2) / (norm1 * norm2)
def word_similarity(w1, w2, matrix, word_to_idx):
"""Compute similarity between two words using their PPMI vectors."""
if w1 not in word_to_idx or w2 not in word_to_idx:
return None
v1 = matrix[word_to_idx[w1]]
v2 = matrix[word_to_idx[w2]]
return cosine_similarity(v1, v2)Word similarity using PPMI vectors (cosine similarity): sim(cat , dog ) = 0.703 sim(cat , mouse ) = 0.495 sim(bank , river ) = 0.316 sim(bank , money ) = 0.240 sim(cat , bank ) = 0.033 sim(dog , mouse ) = 0.161 sim(river , fish ) = 0.443
"Cat" and "dog" share many contexts (both are common animals that interact with the same set of words in the corpus), which produces a high similarity score. "Cat" and "bank" share few meaningful contexts, producing a low similarity score. "Bank" and "river" show moderate similarity. This reflects their shared geographic context in the corpus sentences where "river bank" occurs.
These results match human intuition about semantic similarity, and they emerge entirely from statistical patterns in the text without any hand-crafted features or curated knowledge bases.
Collocation Extraction Using PMI
One of the original applications of PMI is collocation extraction: finding word pairs that form meaningful multi-word expressions. Collocations are combinations whose meaning or usage cannot be predicted from their parts alone. "Hot dog," "kick the bucket," and "New York" all qualify as collocations. A statistical association measure is needed because simple frequency is insufficient: "a the" appears frequently but is not a collocation, while "kick the bucket" appears infrequently but is.
PMI is well-suited to this task because it rewards specificity. "Hot" has relatively high PMI with "dog" because the combination "hot dog" is more specific than you would expect from the individual frequencies of each word. Compare this to "hot water," where "hot" and "water" might have lower PMI because both are common and their combination, while frequent, is not particularly surprising given that hot liquids and water are common topics.
def extract_collocations(ppmi_matrix, idx_to_word, threshold=1.5, top_n=15):
"""
Extract high-PMI word pairs as candidate collocations.
Args:
ppmi_matrix: the PPMI matrix
idx_to_word: index-to-word mapping
threshold: minimum PPMI to qualify
top_n: return top N pairs
Returns:
sorted list of (word1, word2, ppmi_score) tuples
"""
V = ppmi_matrix.shape[0]
collocations = []
for i in range(V):
for j in range(i + 1, V):
# Symmetrize by averaging both directions
score = (ppmi_matrix[i, j] + ppmi_matrix[j, i]) / 2.0
if score >= threshold:
collocations.append((idx_to_word[i], idx_to_word[j], score))
collocations.sort(key=lambda x: x[2], reverse=True)
return collocations[:top_n]Top candidate collocations by PPMI score: common + pets PPMI = 5.33 fear + mice PPMI = 5.33 flooded + was PPMI = 5.33 flows + money PPMI = 5.33 for + services PPMI = 5.33 approved + loan PPMI = 4.92 are + pets PPMI = 4.92 ate + cheese PPMI = 4.92 fees + services PPMI = 4.92 money + through PPMI = 4.92 salmon + upstream PPMI = 4.92 are + common PPMI = 4.33 cats + mice PPMI = 4.33 charges + for PPMI = 4.33 fees + for PPMI = 4.33
The highest-PMI pairs represent words that consistently appear together in the corpus. Even in this small dataset, PMI identifies consistent co-occurrence patterns rather than incidental overlaps driven by word frequency. In a production collocation detector, you would apply this to a much larger corpus and add minimum frequency thresholds to avoid high-PMI pairs that are based on only one or two co-occurrence events.