In this article, we look at methods for estimating means in a finite population. Previously, I wrote about the sampling without replacement methods for this problem. Today, I am going to look at sampling with replacement.

From means to sparse vectors. As in the previous post, the goal is to estimate a mean $\mu = \sum_{i=1}^n p_i f(i)$ when we can only afford $K$ evaluations of $f$. And, as in that post's alternative formulation, we'll get there indirectly. We look for a random vector $\widehat{p}$ that is

• unbiased: $\mathbb{E}[\widehat{p}] = p$, and
• $K$-sparse: $\widehat{p}$ has at most $K$ nonzero entries.

Then $\widehat{p}^\top f$ is an unbiased estimate of $\mu$ for every $f$, and computing it only needs $f(i)$ where $\widehat{p}_i \ne 0$, i.e., at most $K$ evaluations. Estimating all of $p$ is more than we need for a single $\mu$, but it is the right tool when we know nothing about $f$ in advance, or when we want to reuse one sample for many $f$'s.

Formally, we're going to compare a few methods for $K$-sparse representation of a probability distribution $p$ over $\{ 1, \ldots, n \}$.

Suppose, we sample points i.i.d. with replacement from a distribution $q$ over $\{ 1, \ldots, n \}$.

$$ x_1, \ldots, x_K \overset{\text{i.i.d.}}{\sim} q $$

We then put the samples in a list

$$ S = [x_1, \ldots, x_K] $$

Goal: A method for taking the sample $S$ and forming an unbiased $n$-dimensional representation of the distribution with at most $K$ nonzero elements.

$$ \widehat{p}(S) \quad\text{such that}\quad \text{for all } S, \|\widehat{p}(S) \|_0 \le K, \text{ and } \mathbb{E}_S\left[ \widehat{p}(S) \right] = p $$

Clearly, we'd rather sample without replacement. But, suppose for technical reasons, we are only allowed i.i.d. draws from a distribution $q$ that we are allowed to manipulate. How far can we push this scheme?

Estimators: We will compare two estimators that operate on exactly the same sample $S$. The primary difference between the two methods is that method 2 puts the samples in a set (removing duplicates) and method 1 keeps them in a list.

Method 1: Importance sampling estimator

$$ \widehat{p}_i(X) = \frac{1}{K} \sum_{k=1}^K \frac{p(x_k)}{q(x_k)} 1[ x_k = i ]\quad\text{for } i = 1,\ldots,n $$

Method 2: Horvitz–Thompson estimator

$$ \widehat{p}_i(X) = \frac{p_i}{\pi_i} 1[ i \in S ]\quad\text{for } i = 1,\ldots,n $$

where $\pi_i = \mathrm{Pr}\left[ i \in S \right]$

Notice that both estimators are nonnegative, but neither estimator is guaranteed to sum to 1 for all choices of $q$.

What's the inclusion probability $\pi$? If we take $K$ i.i.d. samples with replacement from $q$, what is the probability that we sample item i at least once? $$ \begin{align} \pi_i &= \mathrm{Pr}[ i \in S ] \\ &= 1 - \mathrm{Pr}[ i \not\in \{ X_1, \ldots, X_K \} ] \\ &= 1 - \mathrm{Pr}[ X \ne i ]^K &\quad\text{[X's are i.i.d.]} \\ &= 1 - (1 - q_i)^K \end{align} $$

Warning: This design might not sample $K$ distinct items. Thus, the inclusion probabilities are a bit unusual as they generally sum to less than $K$.

• Exercise for the reader: Show that $\sum_{i=1}^n 1 - (1 - q_i)^K \le K$ given $\sum_{i=1}^{n} q_i = 1$ and $q_1\ldots,q_n \ge 0$.

How do we know if we're doing a good job? A reasonable metric is the expected sum of squared error

$$ \mathcal{E} = \mathbb{E}_S\left[ \sum_{i=1}^n \left( \widehat{p}_i(S) - p_i \right)^2 \right] $$

Why this metric? It is the variance of our estimate of $\mu$, averaged over functions $f$ that we know nothing about. Model $f$ as random and independent of the sample, with uncorrelated values of mean zero and variance one ($\mathbb{E}[f f^\top] = I$). Then

$$ \begin{aligned} \mathbb{E}_f \mathbb{E}_S\left[ \left( \widehat{p}(S)^\top f - p^\top f \right)^2 \right] &= \mathbb{E}_S\left[ (\widehat{p}(S) - p)^\top\, \mathbb{E}_f\!\left[ f f^\top \right] (\widehat{p}(S) - p) \right] \\ &= \mathcal{E}. \end{aligned} $$

If we did know something about $f$, we would choose differently (more on that below).

For IS, $$ \mathcal{E}_{\mathrm{is}} = \frac{1}{K} \sum_{i=1}^n q_i \left( \frac{p_i}{q_i} - p_i \right)^2 + (1-q_i) p_i^2 = \frac{1}{K} \sum_{i=1}^n \frac{p_i^2}{q_i} - p_i^2 $$

For HT,

$$ \begin{aligned} \mathcal{E}_{\mathrm{ht}} &= \sum_{i=1}^n \pi_i \left( \frac{p_i}{\pi_i} - p_i \right)^2 + (1-\pi_i) p_i^2 \\ &= \sum_{i=1}^n \frac{p_i^2}{\pi_i} - p_i^2 \\ &= \sum_{i=1}^n \frac{p_i^2}{1 - (1 - q_i)^K} - p_i^2 \end{aligned} $$

How to pick $q$? The choice of $q$ is important. The optimal choice for importance sampling is $q=p$; regardless of $K$. However, the optimal choice for HT is different. I do not believe there is a closed form solution for it (it seems unlikely given that it is the root of a high-degree polynomial in $K$), but it can be found numerically and it depends on $K$. The section on the optimal proposal for HT, below, characterizes it: computing it only takes one-dimensional root finding.

HT never loses to importance sampling. For the same proposal $q$ and the same budget $K$,

$$ \mathcal{E}_{\mathrm{ht}}(q) \le \mathcal{E}_{\mathrm{is}}(q), $$

with equality when $K = 1$. So whatever $q$ you would use for importance sampling, keeping only the distinct samples and weighting them by their inclusion probabilities is at least as good. In particular, if $q^*$ is the best proposal for HT,

$$ \mathcal{E}_{\mathrm{ht}}(q^*) \le \mathcal{E}_{\mathrm{ht}}(p) \le \mathcal{E}_{\mathrm{is}}(p) = \min_q \mathcal{E}_{\mathrm{is}}(q). $$

Proof

Both errors are sums of per-item terms, so it suffices to compare them item by item. Fix an item with $p_i > 0$ and let $s = 1 - q_i$. The two per-item terms are

$$ \begin{aligned} \text{HT:}&\quad p_i^2 \left( \frac{1}{\pi_i} - 1 \right) = p_i^2 \, \frac{s^K}{1 - s^K} \\ \text{IS:}&\quad \frac{p_i^2}{K} \left( \frac{1}{q_i} - 1 \right) = \frac{p_i^2}{K} \, \frac{s}{1 - s} \end{aligned} $$

If $s = 0$, both are zero. Otherwise $0 < s < 1$. Clearing the (positive) denominators and using the identity

$$ 1 - s^K = (1 - s) \sum_{j=0}^{K-1} s^j, $$

the HT term is at most the IS term if and only if

$$ K s^{K-1} \le \sum_{j=0}^{K-1} s^j. $$

This holds because each of the $K$ terms on the right is at least $s^{K-1}$. When $K = 1$, both sides equal $1$, so the two errors are equal. $\blacksquare$

Experiments¶

The Horvitz–Thompson estimator, applied to a sample drawn with replacement:

Set up a synthetic estimation problem. The distribution $p$ is sampled from a distribution which is very likely to have a few heavy hitters.

What the plot shows. The dotted lines (simulation, 1,000 repetitions) sit on top of the solid lines (formulas), so the error formulas check out. Importance sampling decays like $1/K$, which is a straight line of slope $-1$ on log-log axes. Both HT curves bend away from that line as $K$ grows. To see why, look at one item's contribution to the HT error, $p_i^2 (1 - q_i)^K / \left(1 - (1 - q_i)^K\right)$. While $K q_i$ is small, it is about the same as the item's contribution to the importance sampling error (roughly $p_i^2 / (K q_i)$). Once $K q_i$ is large, it shrinks geometrically. So the heavy items' contributions vanish first, and what's left comes from the light items, which carry little weight.

Most of the gain comes from switching the estimator: at $K = n = 20$, HT with $q = p$ has about a third of the error of importance sampling, and at $K = 40$ less than a fifth. Optimizing $q$ helps more as $K$ grows: it cuts the HT error by about 14% at $K = 20$ and by about 45% at $K = 40$.

Choosing $q$¶

Optimal Proposal for Importance Sampling¶

I have previously ranted about how the optimal proposal distribution is not $p$. However, under the unbiased + expected squared error conditions, it is the optimal choice for the importance sampling type estimators.

There is no contradiction with the rant. The rant is about a single $f$ that we know, where the best proposal is proportional to $p_i |f(i)|$. Here we average over $f$ (see Why this metric? above), and then $q = p$ is the best we can do with importance sampling.

Proposition: The minimum variance choice for $q$ is $p$.

Proof

First, we write out the Lagrangian for picking $q$ with optimal variance: $$ \begin{cases} \underset{{q_1, \ldots q_n}}{\mathrm{minimize}}\ & \frac{1}{K} \sum_{i=1}^n \frac{p_i^2}{q_i} \\ \mathrm{subject\ to} & \sum_{i=1}^n q_i = 1 \\ & q_1, \ldots q_n \ge 0 \end{cases} $$

$$ L(q_1, \ldots q_n, \lambda) = \frac{1}{K} \sum_{i=1}^n \frac{p_i^2}{q_i} - \lambda \left(1 - \sum_{i=1}^n q_i \right) $$

Now, solve for $q$ and $\lambda$ which make $L$'s partial derivatives simultaneously zero. $$ \begin{align} 0 &= \nabla_{q_j}\left[ L(q_1, \ldots q_n, \lambda) \right] \\ &= \nabla_{q_j}\left[ \frac{1}{K} \sum_{i=1}^n \frac{p_i^2}{q_i} - \lambda \left( 1 - \sum_{i=1}^n q_i \right) \right] \\ &= \frac{1}{K} \sum_{i=1}^n \nabla_{q_j}\left[ \frac{p_i^2}{q_i} \right] + \lambda \sum_{i=1}^n \nabla_{q_j}\left[ q_i \right] \\ &= \lambda - \frac{1}{K} \frac{p_j^2}{q_j^2} \\ \lambda &= \frac{1}{K} \frac{p_j^2}{q_j^2} \\ q_j^2 &= \frac{p_j^2 }{ K \lambda } \\ \sqrt{q_j^2} &= \sqrt{\frac{p_j^2 }{ K \lambda }} \\ q_j &= \frac{p_j}{ \sqrt{ K \lambda }} \\ \end{align} $$

and, the constraint that $\nabla_\lambda L = 0$ just replicates the constraint corresponding to $\lambda$. $$ \begin{align} \left(1 - \sum_{i=1}^n q_i \right) &= 0 \\ 1 &= \sum_{i=1}^n q_i \end{align} $$

Clearly, the only valid solution to all these equations is $q = p$. Furthermore, this choice is optimal for any valid choice of $K$.

Optimal Proposal for HT¶

For HT, the best $q$ is not $p$ (except when $K = 1$), and it depends on $K$. I don't know of a closed form, but the problem has enough structure that finding $q^*$ is easy.

It's a convex problem. The inclusion probability $\pi(q) = 1 - (1-q)^K$ is concave and positive on $(0, 1]$, so $1/\pi(q)$ is convex. Hence $\mathcal{E}_{\mathrm{ht}}$ is a (strictly) convex function of $q$, minimized over the probability simplex: there is one optimum, and any point satisfying the optimality conditions is it.

One shared constant. Setting the derivative of the Lagrangian to zero gives, for every item $i$,

$$ p_i^2 \, \frac{K (1-q_i)^{K-1}}{\left(1 - (1-q_i)^K\right)^2} = \lambda. $$

Taking square roots, $p_i \, \phi(q_i) = c$, where $c = \sqrt{\lambda}$ and

$$ \phi(q) = \frac{\sqrt{K (1-q)^{K-1}}}{1 - (1-q)^K}. $$

Since $\phi$ is decreasing, each $q_i$ is determined by $p_i$ and $c$ alone, by one-dimensional root finding. Then $c$ is set (one more root finding) so that $\sum_i q_i = 1$.

Shape of the solution. Two properties follow:

• $q^*$ keeps the ranking of $p$: a larger $p_i$ needs a smaller $\phi(q_i)$, i.e., a larger $q_i$.
• $q^*$ is flatter than $p$: the ratio $q^*_i / p_i$ decreases as $p_i$ grows.

Proof

Note that $q_i / p_i = q_i \phi(q_i) / c$ and, writing $s = 1 - q$, $$ q\,\phi(q) = \frac{\sqrt{K}}{\sum_{j=0}^{K-1} s^{\,j - (K-1)/2}}. $$ The exponents in that sum come in pairs $\pm a$, and each pair $s^{a} + s^{-a}$ decreases as $s$ grows, so $q\,\phi(q)$ decreases in $q$.

The intuition: a heavy item lands in the sample almost surely anyway, and extra copies of it are wasted, so $q^*$ moves probability toward the light items. When $K = 1$, $\phi(q) = 1/q$ and we recover $q^* = p$.

Summary¶

If you're stuck with i.i.d. samples, don't throw away the information in which items you saw: put the samples in a set and use the Horvitz–Thompson estimator.

• For any proposal $q$, HT's error is never worse than importance sampling's, and it is much better once $K$ is comparable to $n$. Each item's share of the HT error shrinks geometrically once $K q_i$ is large, rather than like $1/K$.
• The best proposal for HT is not $p$. It's a flattened version of $p$ that depends on $K$, and it's easy to compute.

Appendix: Numerical Checks¶

Check: $q = p$ Is Optimal for Importance Sampling¶

Check: HT Never Loses to Importance Sampling¶

Check: Optimal Proposal vs. a Generic Optimizer¶

Since the problem is convex, the root-finding solution opt_q_ht should match (or beat) a generic optimizer applied to the same objective.

Gradient Check for the HT Error¶

The gradients d_var_htmc and d_var_htmc_softmax are written by hand, so we check them against central differences (fdcheck raises an error if they disagree). We check d_var_htmc at a proposal without tiny entries: where $q_i$ is small, the HT error curves sharply, and central differences themselves become inaccurate.