Prediction-Powered Inference: Rectifiers, Estimators, and Validity

statisticsmachine learning

Suppose predictions are cheap but outcome labels are scarce. Averaging predictions alone may estimate the wrong mean; a smaller labeled sample can estimate and subtract that prediction error. (Angelopoulos et al. 2023) extends this correction to estimating equations and confidence sets. The correction gives validity, while predictive accuracy determines the possible gain in precision.

1. The result for an estimating equation

Let \(\theta^*\in\mathbb R^m\) satisfy \(\mathbb{E}[g_{\theta^*}(X,Y)]=0\). For a convex loss, this is the population first-order condition, provided the target is an interior minimizer with an appropriate zero-mean gradient or subgradient. Define

\[ \Delta_\theta\triangleq\mathbb{E}[g_\theta(X,Y)-g_\theta(X,f(X))], \qquad g_\theta^f\triangleq\mathbb{E}[g_\theta(X,f(X))]. \]

Theorem 1. Suppose \(R_\delta(\theta)\) covers \(\Delta_\theta\) with probability at least \(1-\delta\), and \(T_{\alpha-\delta}(\theta)\) covers \(g_\theta^f\) with probability at least \(1-(\alpha-\delta)\). Then

\[ C_\alpha\triangleq\{\theta:0\in R_\delta(\theta)+T_{\alpha-\delta}(\theta)\} \]

covers \(\theta^*\) with probability at least \(1-\alpha\). The sum here is the Minkowski sum of sets.

It suffices to check coverage at \(\theta^*\). If both component sets cover, their sum contains \(\Delta_{\theta^*}+g_{\theta^*}^f=0\), so the candidate is retained. The union bound controls failure without requiring simultaneous bands over \(\theta\).

For the usual normal approximation, take independent labeled and unlabeled samples from the same population, of sizes \(n\) and \(N\). Condition on a predictor trained independently of both. Average the gradient differences in the labeled sample and the imputed gradients in the unlabeled sample. With fixed dimension, finite second moments, nonzero coordinate variances, and \(n/N\to p\in(0,1)\), retain \(\theta\) when, for every coordinate \(j\),

\[ |\widehat\Delta_{\theta,j}+\widehat g^f_{\theta,j}| \leq z_{1-\alpha/(2m)} \sqrt{s^2_{\Delta,j}(\theta)/n+s^2_{g,j}(\theta)/N}. \]

A CLT for the two independent averages and a union bound across coordinates give asymptotic coverage at least \(1-\alpha\).

2. The mean: check the sign and the variance

For \(g_\theta(y)=\theta-y\), the estimator is

\[ \widehat\theta^{\rm PP} =\frac1N\sum_{i=1}^N f(\widetilde X_i) +\frac1n\sum_{i=1}^n(Y_i-f(X_i)). \]

Its expectation is \(\mathbb{E}Y\), even if \(f\) is biased. Independence gives the exact variance decomposition

\[ \operatorname{Var}(\widehat\theta^{\rm PP}) =\frac{\operatorname{Var}(f(X))}{N} +\frac{\operatorname{Var}(Y-f(X))}{n}. \]

For large \(N\), prediction-error variance in the labeled sample determines most of the interval width. Accurate predictions can reduce this variance; inaccurate ones can increase it. Coverage alone gives no efficiency guarantee.

3. Quantiles and regressions use the same correction

For a continuous outcome distribution at its \(\tau\)-quantile, use \(g_\theta(y)=\mathbf1\{y\leq\theta\}-\tau\). Correcting the predicted cdf gives

\[ \widehat F^{\rm PP}(\theta) =\frac1N\sum_i\mathbf1\{f(\widetilde X_i)\leq\theta\} +\frac1n\sum_i\left[\mathbf1\{Y_i\leq\theta\}-\mathbf1\{f(X_i)\leq\theta\}\right]. \]

Invert the test that this expression equals \(\tau\). A numerical grid should cover the parameter range of interest; the range of the predictions can miss the true quantile.

For logistic regression, \(g_\theta(x,y)=x(\mu_\theta(x)-y)\), where \(\mu_\theta(x)=(1+e^{-x'\theta})^{-1}\). The terms involving \(\mu_\theta\) cancel in the rectifier, leaving \(\Delta_\theta=\mathbb{E}[X(f(X)-Y)]\), which does not depend on \(\theta\). Poisson regression has the same cancellation with \(\mu_\theta(x)=e^{x'\theta}\).

For linear regression, let \(\Sigma=\mathbb{E}[XX']\) be nonsingular. Linearity gives the population identity

\[ \theta^*=\Sigma^{-1}\mathbb{E}[Xf(X)] +\Sigma^{-1}\mathbb{E}[X(Y-f(X))]. \]

Estimate the two terms on the two samples. Each is an OLS coefficient, so the corresponding sandwich covariance matrices add after division by \(N\) and \(n\). Finite moments, full rank, and the usual OLS linearization conditions give asymptotically valid coordinate intervals. The common-population assumption matters: otherwise the two terms need not cancel the prediction error.

The nonconvex extension replaces first-order conditions with comparisons of population risks.

4. How much can predictions help?

The mean calculation can be pushed further. Replace \(f\) by \(\lambda f\), where \(\lambda\) is fixed before looking at the inference samples. This gives the unbiased estimator

\[ \widehat\mu_\lambda=\overline Y_n+\lambda(\overline f_N-\overline f_n). \]

The added difference of prediction averages has mean zero. It is therefore a control variate. With \(V_f=\operatorname{Var}(f(X))\) and \(C=\operatorname{Cov}(Y,f(X))\), the estimator’s variance is

\[ \begin{aligned} \operatorname{Var}(\widehat\mu_\lambda) &=\frac{\operatorname{Var}(Y)}n -\frac{2\lambda C}n +\lambda^2V_f\left(\frac1n+\frac1N\right),\\ \lambda^*&=\frac{C}{V_f(1+n/N)}\quad (V_f>0). \end{aligned} \]

Substitution yields

\[ \operatorname{Var}(\widehat\mu_{\lambda^*}) =\frac{\operatorname{Var}(Y)}n -\frac{C^2}{nV_f(1+n/N)}. \]

This is a useful supplementary derivation, rather than an extra assumption in the coverage theorem. Correlation with the outcome determines the possible gain. A prediction can have a large constant bias and still be useful, because the labeled residual average removes that constant. Conversely, a highly variable prediction uncorrelated with \(Y\) adds noise when its coefficient is fixed at one. Estimating \(\lambda\) on an independent pilot preserves the conditional calculation; estimating it on the inference sample requires a separate asymptotic argument.

Writing \(\rho_{Y,f}=\operatorname{Corr}(Y,f(X))\), the variance ratio at the optimal fixed coefficient is

\[ \frac{\operatorname{Var}(\widehat\mu_{\lambda^*})} {\operatorname{Var}(\overline Y_n)} =1-\frac{\rho_{Y,f}^2}{1+n/N}. \]

As \(N\) grows, the optimal variance reduction approaches the squared correlation. Constant prediction bias disappears under the labeled correction and hence does not enter this efficiency calculation.

For the original coefficient-one estimator, let \(V_r=\operatorname{Var}(Y-f)\) and suppose acquiring a labeled observation costs \(c_L\) while an unlabeled one costs \(c_U\). With the predictor held fixed and a budget \(c_Ln+c_UN=B\), minimizing \(V_r/n+V_f/N\) in positive real sample sizes gives

\[ \frac{V_r}{c_Ln^2}=\frac{V_f}{c_UN^2}, \qquad \frac nN=\sqrt{\frac{V_rc_U}{V_fc_L}}. \]

This allocation holds the predictor fixed and treats sample sizes as positive real numbers. Small residual variance reduces the labels needed for a given precision, subject to the sample requirements of the inference method. High prediction variance or costly unlabeled data can favor a less extreme allocation.

The two samples estimate complementary terms in the target equation. Labels measure the correction; unlabeled observations estimate the predicted part. Both sampling errors remain in the variance, regardless of how the samples are allocated.

The common-population assumption also becomes concrete here. If labeled observations follow \(P\) and unlabeled observations follow \(Q\), then

\[ \mathbb{E}\widehat\mu_\lambda-\mathbb{E}_PY =\lambda(\mathbb{E}_Qf(X)-\mathbb{E}_Pf(X)). \]

Increasing the unlabeled sample estimates its population’s prediction mean more precisely, without removing the mismatch. Cancellation requires the same prediction mean or an explicit adjustment for population change.

5. Quantile uncertainty comes from classification errors near the quantile

Let \(q_\tau\) satisfy \(F_Y(q_\tau)=\tau\), with density \(h_Y(q_\tau)>0\). At this point write

\[ A=\mathbf1\{Y\leq q_\tau\},\qquad B=\mathbf1\{f(X)\leq q_\tau\}. \]

The corrected cdf has variance

\[ \operatorname{Var}(\widehat F^{\rm PP}(q_\tau)) =\frac{\operatorname{Var}(A-B)}n+ \frac{\operatorname{Var}(B)}N. \]

Here \((A-B)^2=\mathbf1\{A\ne B\}\). Therefore

\[ \operatorname{Var}(A-B) =\mathbb P(A\ne B)-(\mathbb{E}A-\mathbb{E}B)^2. \]

For quantile inference, prediction matters through agreement about which side of the target quantile contains the outcome. Squared-error accuracy may help, but the variance here is determined directly by that classification error.

Under local differentiability and the stochastic equicontinuity needed to invert the estimated equation, a consistent approximate root obeys

\[ \widehat q_\tau-q_\tau =-\frac{\widehat F^{\rm PP}(q_\tau)-\tau}{h_Y(q_\tau)} +o_p(n^{-1/2}+N^{-1/2}). \]

To obtain the sign, expand \(0=\widehat F^{\rm PP}(\widehat q_\tau)-\tau\) around \(q_\tau\) and move the empirical error to the other side. A low density amplifies cdf uncertainty into quantile uncertainty. The corrected empirical cdf need not itself be monotone, because one empirical cdf is subtracted. Test inversion at each candidate remains meaningful even in that case; an arbitrary numerical root need not have the asserted expansion.

6. Regression: the cancellation and the linearization

For logistic regression the empirical corrected score can be written

\[ \widehat\Psi(\theta) =\frac1N\sum_i\widetilde X_i (\mu_\theta(\widetilde X_i)-f(\widetilde X_i)) +\frac1n\sum_iX_i(f(X_i)-Y_i). \]

Only the first average changes with \(\theta\). Its population derivative at the target is

\[ H=\mathbb{E}[XX'\mu_{\theta^*}(X)(1-\mu_{\theta^*}(X))]. \]

If \(H\) is nonsingular and the usual local uniform convergence holds, integrate the sample derivative along the segment from \(\theta^*\) to a consistent root. The score equation gives

\[ \widehat\theta-\theta^* =-H^{-1}\widehat\Psi(\theta^*) +o_p(n^{-1/2}+N^{-1/2}). \]

The two independent score averages contribute additive covariance terms, each divided by its own sample size. Because correction recovers the population score before linearization, \(f\) need not equal the conditional mean.

For the two-OLS construction above, define \(\beta_f=\Sigma^{-1}\mathbb{E}[Xf]\) and \(\beta_r=\Sigma^{-1}\mathbb{E}[X(Y-f)]\). Its influence terms are

\[ \Sigma^{-1}X(f-X'\beta_f),\qquad \Sigma^{-1}X(Y-f-X'\beta_r). \]

Each OLS fit estimates its own second-moment matrix, so its influence function contains a fitted projection residual. Using only raw predictions and residual outcomes would omit that variation. The distinction is relevant when prediction is nonlinear in \(X\).

Angelopoulos, Anastasios N., Stephen Bates, Clara Fannjiang, Michael I. Jordan, and Tijana Zrnic. 2023. “Prediction-Powered Inference.” arXiv. http://arxiv.org/abs/2301.09633.

← All posts