A Statistical Inference Framework for PMI Estimation and SGNS Word Embeddings

Zhongqi Fan

stat.MLOct 4, 2026

A Statistical Inference Framework for PMI Estimation and SGNS Word Embeddings

Authors: Zhongqi Fan

Organizations: Beijing Normal-Hong Kong Baptist University

Abstract

Pointwise Mutual Information (PMI) is a core measure of testing word association, and Skip-gram with Negative Sampling (SGNS) is essentially a method that implicitly factorizes a shifted PMI matrix. However, a systematic and well-rounded characterization of finite-sample uncertainty in PMI estimation remains absent and imperative to venture into. We provide a statistical framework for PMI estimation and its connection to SGNS. We prove consistency, asymptotic unbiasedness, and asymptotic normality of the empirical PMI estimator, derive its variance via the Delta method, and, applying stochastic approximation theory, obtain a variance decomposition for SGNS-based PMI estimation that separates data variance from optimization variance. Simulation experiments validate the Delta method approximation. Real-data experiments on the Brown Corpus (d = 100) reveal that SGNS systematically deviates from the theoretical relationship PMI + log K. The empirical relationship shows an attenuated PMI coefficient, an amplified log K effect, and a positive intercept, indicating systematic bias. Word analogy validation confirms the models are effective. The failure to validate the variance decomposition under low-dimensional conditions does not diminish its theoretical value; rather, it identifies the unbiasedness assumption as the key bottleneck and clarifies the gap between asymptotic theory and practice, providing implications for both practice and theory.

Figures & tables

$x,y$	Specific words in the vocabulary
$V$	Vocabulary size
$N$	Total number of windows in the corpus
$n_{x}$	Number of windows containing word $x$
$n_{xy}$	Number of windows word $x$ and $y$ co-occurring
$p_{x}=n_{x}/N$	Empirical probability of word $x$
$p_{xy}=n_{xy}/N$	Empirical probability of $x$ and $y$ co-occurring

Table 1

Figure 1 : PMI Estimation and Bias

Figure 2 : PMI Estimation Variance

Figure 3 : Document length distribution of the Brown Corpus.

Statistic	Value
Total word tokens	981,716
Unique word types	40,234
Number of documents	500
Mean document length	2,322 words
Window size (sliding)	5
Total training windows ( $N$ )	979,716

Table 1 : Brown Corpus Statistics

Figure 4 : Word frequency distribution (Zipf’s law). Left: rank–frequency plot on log–log scales. Right: histogram of log-transformed frequencies.

Figure 5 : Co-occurrence statistics. Left: histogram of co-occurrence counts (

n_{xy}

). Right: scatter plot of PMI vs.

\log(n_{xy})

Figure 6 : PMI distribution. Left: histogram of PMI values. Right: boxplot of PMI by co-occurrence frequency category.

Frequency	PMI Type	Word Pairs
High	Positive	(is, it), (it, was), (had, he)
	Near-zero	(of, the), (in, the), (and, the)
	Negative	(a, the), (in, of), (and, to)
Medium	Positive	(do, something), (age, at)
	Near-zero	(an, make), (some, which)
	Negative	(political, to), (had, new)

Table 2 : Selected Word Pairs by Category (

3\times 3

Grid)

Figure 7 : Selected word pairs for SGNS experiments, highlighted against the background distribution.

$K$	Semantic		Syntactic		Overall
$K$	Cor/Total	Accuracy	Cor/Total	Accuracy	Overall
1	38/306	12.42%	66/7878	0.84%	1.27%
2	49/306	16.01%	46/7878	0.58%	1.16%
5	60/306	19.61%	42/7878	0.53%	1.25%
10	41/306	13.40%	61/7878	0.77%	1.25%
15	51/306	16.67%	67/7878	0.85%	1.44%

Table 3: Analogy accuracy for SGNS trained on Brown Corpus

Figure 8 : Analogy accuracy as a function of word frequency (figure adapted from [ 13 ] ).

$K$	Correct	Total	$p$ -value
1	38	306	$5.25\times 10^{-127}$
2	49	306	$4.20\times 10^{-196}$
5	60	306	$1.79\times 10^{-212}$
10	41	306	$2.40\times 10^{-138}$
15	51	306	$6.70\times 10^{-177}$
20	56	306	$1.44\times 10^{-196}$

Table 4 : Binomial test results for semantic analogy accuracy

Category \K	$1$	$2$	$5$	$10$	$15$	$20$
High Positive	1.00	1.00	1.00	1.00	1.00	1.00
High Near-zero	1.00	1.00	1.00	1.00	1.00	1.00
High Negative	1.00	1.00	1.00	1.00	1.00	1.00
Medium Positive	1.00	1.00	1.00	1.00	1.00	1.00
Medium Near-zero	1.00	1.00	1.00	1.00	1.00	1.00
Medium Negative	0.50	1.00	1.00	1.00	1.00	1.00

Table 5: Rejection Rates of

H_{0}^{(1)}

by Category and

K

Figure 9 : Bias Distribution of Theoretical Formula

Metric	Theoretical	With Intercept	Without Intercept
A (PMI)	1.0000	0.7077	0.8046
B (log K)	1.0000	1.5692	2.2174
Intercept	0.0000	1.5692	0 (fixed)
R²	$-0.6572$	0.6365	0.5030
RMSE	2.9221	1.3685	1.6002
MAE	2.5208	1.1265	1.3262

Table 6: Regression model comparison (transposed)

Coefficient: A (PMI)
Statistic	With Intercept	Without Intercept
Estimate	$0.7077\pm 0.1348$	$0.8046\pm 0.1560$
t-statistic	$-2.17$	$-1.25$
p-value	$0.0320$	$0.2129$
Coefficient: B (log K)
Estimate	$1.5692\pm 0.1145$	$2.2174\pm 0.0719$

Table 7 : Hypothesis tests for regression coefficients (transposed)

Figure 10 : Regression diagnostic plots. Top row: theoretical model (

\hat{\theta}=\text{PMI}+\log K

). Bottom row: empirical model with intercept (

\hat{\theta}=0.70\cdot\text{PMI}+1.57\cdot\log K+1.57

). Left: actual vs. predicted. Middle: residuals vs. predicted. Right: coefficient comparison.

$k$	PMI $-$ log $k$	SGNS
$k$	PMI $-$ log $k$	$d=100$	$d=500$	$d=1000$
1	0%	31.4%	29.4%	7.40%
5	0%	39.3%	36.0%	7.13%
15	0%	7.80%	6.37%	5.97%

Table 8 : SGNS objective deviation by dimension (Levy & Goldberg, 2014, Table 1)

Explore similar work

PromptNCE: Conditional Probabilities and PMI Using Only LLMs and Contrastive Estimation Prompts
May 20, 2026Juliette Woodrow, Chris PiechMutual Information Prompt Engineering
Department of Computer Science Stanford University
No Free Lunch: Non-Asymptotic Analysis of Prediction-Powered Inference
May 26, 2025Pranav Mani, Peng Xu, Zachary C. Lipton +1Finite-Sample Prediction Error
Abridge AI, San Francisco, CA, USA · Machine Learning Department, Carnegie Mellon University, Pittsburgh, PA, USA · Department of Computer Science, Johns Hopkins University, Baltimore, MD, USA
InfoAtlas: A Foundation Model for Zero-Shot Statistical Dependence Estimate
May 29, 2026Zhengyang Hu, Yanzhi Chen, Hanxiang Ren +5Mutual Information Temporal Dependence
The University of Hong Kong · University of Cambridge · Microsoft +2

May 20, 2026, cs.CLJ/K move · Enter open · S save

PromptNCE: Conditional Probabilities and PMI Using Only LLMs and Contrastive Estimation Prompts

Juliette Woodrow, Chris Piech

Department of Computer Science Stanford University

Mutual Information Prompt Engineering Contrastive Learning

Open paper AskPDF arXiv

Estimating mutual information from text usually requires training a task-specific critic, which limits its use in low-data settings. We ask whether large language models can instead estimate pointwise mutual information zero-shot, using only prompts and elicited probabilities. We construct a benchmark from three publicly available human-annotated datasets with ground-truth PMI, and evaluate five information-theoretic prompting-based estimators. Our main method, PromptNCE, frames conditional probability estimation as a contrastive task and augments the candidate set with an explicit OTHER category. The OTHER category allows the model to assign probability mass outside the candidate set, avoiding the closed-set normalization of standard contrastive prompts. PromptNCE gives the best conditional probability estimates on all three datasets. For full PMI, we find that estimating label base rates is the primary bottleneck on two of the three datasets, with the best methods reaching Spearman correlation up to 0.78. We also present a case study in computer science education showing how these estimators can be used to score student knowledge summaries in a low-data setting. We release our code and prompts.

Abstract

Explore similar work

PromptNCE: Conditional Probabilities and PMI Using Only LLMs and Contrastive Estimation Prompts

No Free Lunch: Non-Asymptotic Analysis of Prediction-Powered Inference

InfoAtlas: A Foundation Model for Zero-Shot Statistical Dependence Estimate