active statistical inference

Most statistical inference assumes the data is somehow supplied to you (even in sequential settings; see sequential statistics). You are given data, and you are asked to estimate some parameter.

But suppose we are in a supervised learning setting and you get to choose which data points to label. How should you label in them in order to help your inference (eg obtain smaller confidence intervals)? Zrnic and Candes introduce the framework of active statistical inference to answer this question.

Suppose we have data $(X_i,Y_i{)}_{i=1}^N$, where $Y_i$ is unobserved and $X_i$ is observed. We have a model $f$ which estimates $Y_i$ from $X_i$. The goal is to estimate ${\theta }^{\ast }$, the solution to an M-estimation problem:

$${\theta }^{\ast }={\text{argmin}}_{\theta }\mathbb{E}[{\ell }_{\theta }(X,Y)],$$

where ${\ell }_{\theta }$ is assumed to be convex. This handles mean estimation, linear regression coefficients, and quantile estimation. We have a budget of $B$ observations that we’re allowed to sample (in expectation).

That the budget holds only in expectation seems like a drawback of this model. Would that actually be the case in practice? Probably you would instead have some absolutely upper limit on the budget.

Intuition

Consider mean estimation, ${\theta }^{\ast }=\mathbb{E}[Y]$ (so ${\ell }_{\theta }(X,Y)=(\theta -Y{)}^2$). Let $\pi (X_i)$ be the probability of sampling $X_i$. Consider the doubly robust estimator:

$$\hat{\theta }=\frac{1}{N}\sum _{i=1}^N\left(f(X_i)+(Y_i-f(X_i))\frac{{\xi }_i}{\pi (X_i)}\right),$$

which is unbiased. The variance is

$$\mathbb{V}(\hat{\theta })=\frac{1}{N}\left(\mathbb{V}(Y)+\mathbb{E}[(Y-f(X){)}^2(\frac{1}{\pi (X)}-1)]\right).$$

If $Y\approx f(X)$, or $\pi (X)\approx 1$, then $\mathbb{V}(\hat{\theta })$ is much smaller than the variance of just $B$ randomly selected points, which is $\frac{1}{B}\mathbb{V}(Y)$. So, ideally, we want to choose $\pi (X)$ such that if the model is bad, then $\pi (X)\approx 1$.

Choosing $\pi$

They choose $\pi$ as a function of the model uncertainty. It’s hard to give rigorous guarantees here — the methods are mainly heuristics.

Regression: We train a model $u$ to predict $u(X)=|Y-f(X)|$. This should be trained on a dataset different from that which $f$ was trained on.

Classification: Here we assume that $f(x)={\text{argmax}}_{i\in [k]}{p}_i(x)$ where ${\sum }_i{p}_i(x)=1$. Then we set

$$u(x)=\frac{K}{K-1}(1-\underset{i}{\max }{p}_i(x)).$$

If the classifier is maximally uncertain, then $p_i(x)=1/K$ for all $i$ so $u(x)=1$. In practice, helps to mix $u$ with a uniform so that it doesn’t blow up the variance.

Once $u$ is determined, we might consider sampling strategies of the form: ${\pi }_{\eta }(x)=\eta u(x),$ where $\eta$ is some hyperparameter mean to ensure that we respect the budget $B$. They suggest choosing $\eta =\hat{\eta }$ where $\hat{\eta }=\sup \left\{\eta \in \mathcal{H}:\eta {\sum }_{i=1}^{n}u(X_i)\le B\right\},$ is a data-driven selected parameter. The actual sampling is done by then sampling ${\xi }_i\sim \text{Ber}({\pi }_{\eta }(X_i))$ (i.e., whether $X_i$ is sampled independently of other observations). Note that this means we only meet the budget in expectation:

$${\mathbb{E}}_S\sum _{i=1}^N{\xi }_i=\sum _{i=1}^N\mathbb{E}[\hat{\eta }{\xi }_i]=\sum _{i=1}^N\hat{\eta }u(X_i)\le B,$$

where the expectation is over the sampling only.

General setting

We design a sampling algorithm $\pi$ which samples $x_i$ with probability $\pi (x_i)$. That is, we draw ${\xi }_i\sim \text{Ber}(\pi (x_i))$ and if ${\xi }_i=1$ we sample $x_i$, otherwise we don’t. This implies that all observations are sampled independently from one another. $\pi$ is scaled such that $\mathbb{E}[{\sum }_i{\xi }_i]\le B$ to ensure that the expected number of labeled observations does not exceed the budget.

Let ${\ell }_{\theta ,i}={\ell }_{\theta }(X_i,Y_i)$ be the estimate of the loss on example $i$, and ${\ell }_{\theta ,i}^f={\ell }_{\theta }(X_i,f(X_i))$ be the loss of the model on example $i$. An unbiased estimator of $L(\theta )=\mathbb{E}[{\ell }_{\theta }(X,Y)]$ is

$$L^{\pi }(\theta )=\frac{1}{n}\left(\sum _{i=1}^n{\ell }_{\theta }(X_i,f(X_i))+({\ell }_{\theta }(X_i,Y_i)-{\ell }_{\theta }(X_i,f(X_i)))\frac{{\xi }_i}{\pi (X_i)}\right).$$

Note that this is just the doubly robust estimator. If $\pi$ is just the uniform rule and doesn’t prioritize some observations over others, this recovers the prediction-powered inference estimator.

The resulting confidence interval for our estimate ${\hat{\theta }}^{\hat{\eta }}$ comes from a CLT, from which a Wald interval is constructed. The CLT mainly follows from the usual CLT for M-estimation, but needs to correct for the fact that $\hat{\eta }$ is not independent of the data. This is done by assuming that $\hat{\eta }\to {\eta }^{\ast }$ in probability for some fixed ${\eta }^{\ast }$.

They also study a sequential setting, where the model and sampling rule can be iteratively updated. Then they apply a martingale CLT to get a CI for the estimate.

References

• Active statistical inference by Zrnic and Candes.
• Efficient adaptive experimental design for ATE estimation Very similar ideas, but with more focus on estimation as opposed to choosing the sampling strategy.