Prediction-Powered Inference Beyond Convex Losses

A zero population gradient need not identify a minimizer of a nonconvex risk. We can instead compare the risk of a candidate with that of a separately chosen benchmark: a true global minimizer cannot have greater risk. (Angelopoulos et al. 2023) combines this inequality with prediction-error correction to construct a confidence set.

1. The coverage statement

Write \(L(\theta)=\mathbb{E}[\ell_\theta(X,Y)]\) and let \(\theta^*\) be a population minimizer. Split the unlabeled data. Use the first part to choose a benchmark \(\widehat\theta^f\) by minimizing predicted loss, and reserve the other part for inference. Define

\[ L^f(\theta)\triangleq\mathbb{E}[\ell_\theta(X,f(X))], \qquad \Delta(\theta)\triangleq L(\theta)-L^f(\theta). \]

On the held-out data, let \(\widetilde L^f\) estimate \(L^f\). Suppose the labeled sample supplies lower and upper bounds \(R^l,R^u\) for \(\Delta\), and the held-out unlabeled sample supplies bounds \(T^l,T^u\) for \(\widetilde L^f-L^f\). All guarantees below are conditional on the first split.

Theorem 1. Assign failure probability \(\delta/2\) to each of the two rectifier bounds used below, and \((\alpha-\delta)/2\) to each imputed-risk error bound, where \(0<\delta<\alpha\). Then the set

\[ \begin{split} C_\alpha\triangleq\{\theta:\ \widetilde L^f(\theta) \leq {}&\widetilde L^f(\widehat\theta^f) -R^l(\theta)+R^u(\widehat\theta^f)\\ &+T^u(\theta)-T^l(\widehat\theta^f)\} \end{split} \]

contains \(\theta^*\) with probability at least \(1-\alpha\). Convexity is not required.

2. Deriving the inequality

First express the risk through its estimated components; this fixes the signs in the confidence bounds. Since

\[ L(\theta)=\widetilde L^f(\theta)+\Delta(\theta) -(\widetilde L^f(\theta)-L^f(\theta)), \]

a lower bound is \(\widetilde L^f+R^l-T^u\), while an upper bound is \(\widetilde L^f+R^u-T^l\). Keep a candidate whenever its lower bound is no larger than the benchmark’s upper bound.

At the true minimizer,

\[ L(\theta^*)\leq L(\widehat\theta^f). \]

Consequently, rejection of \(\theta^*\) requires at least one of four one-sided bounds to fail: the lower rectifier bound at \(\theta^*\), the upper rectifier bound at the benchmark, or one of the two imputed-risk error bounds. Their error probabilities sum to \(\alpha\).

3. Why split the unlabeled sample?

The optimizer used as benchmark is random. A bound valid at each fixed \(\theta\) need not remain valid after optimization on the same sample. Conditioning on the first split fixes the benchmark before the held-out inference step.

Coverage uses only the minimizing inequality, so the benchmark need not converge to \(\theta^*\). A lower-risk benchmark can exclude more false candidates, and accurate predictions can narrow the correction bounds. These affect precision without changing the coverage argument.

For losses whose first-order condition already characterizes the minimizer, the estimating-equation construction is usually simpler. Here the additional work buys a statement about the global risk minimizer.

4. An example where the first-order condition is insufficient

For a concrete example, take \(\ell_\theta(Y)=(\theta^2-Y)^2\), assume a finite second outcome moment, and suppose \(\mu=\mathbb{E}Y>0\). Expanding gives

\[ L(\theta)=\theta^4-2\mu\theta^2+\mathbb{E}[Y^2], \qquad L'(\theta)=4\theta(\theta^2-\mu). \]

There are three stationary points: \(0\) and \(\pm\sqrt\mu\). Only the latter two minimize the risk. Indeed,

\[ L(0)-L(\sqrt\mu)=\mu^2>0, \qquad L''(0)=-4\mu<0. \]

A gradient test can retain \(0\) even with arbitrarily precise estimation. A risk comparison can exclude it once uncertainty is smaller than the gap \(\mu^2\). If several minimizers exist, the theorem covers each specified one; covering all of them simultaneously requires simultaneous bounds.

The benchmark need not equal \(\sqrt\mu\). Suppose its excess risk is \(\eta=L(\widehat\theta^f)-\inf_\theta L(\theta)\). Even with perfect evaluation, comparison with that benchmark retains every candidate whose excess risk is at most \(\eta\). A poor benchmark therefore gives a broad set, although the true minimizer remains protected.

5. A finite-sample construction with bounded loss

Suppose \(0\leq\ell_\theta(x,y)\leq B\) over the parameter space and for both true and predicted outcomes. Let \(N_2\) be the size of the held-out unlabeled sample and put

\[ \widehat\Delta(\theta) =\frac1n\sum_i(\ell_\theta(X_i,Y_i)-\ell_\theta(X_i,f(X_i))). \]

The summand lies in \([-B,B]\), a range of length \(2B\). Hoeffding’s one-sided inequality consequently gives

\[ \mathbb P\{\widehat\Delta(\theta)-\Delta(\theta)>r\} \leq\exp(-nr^2/(2B^2)). \]

Choose

\[ r_\Delta=B\sqrt{\frac{2\log(2/\delta)}n}, \qquad r_f=B\sqrt{\frac{\log(2/(\alpha-\delta))}{2N_2}}. \]

The second radius uses the shorter range \([0,B]\) of the predicted loss. Set \(R^l=\widehat\Delta-r_\Delta\), \(R^u=\widehat\Delta+r_\Delta\), and \(T^l=-r_f,T^u=r_f\). Each bound used in the proof has its assigned one-sided error probability. Substituting into the theorem yields the particularly transparent set

\[ C_\alpha=\left\{\theta: \widehat L^{\rm PP}(\theta)- \widehat L^{\rm PP}(\widehat\theta^f) \leq2r_\Delta+2r_f\right\}, \quad \widehat L^{\rm PP}=\widetilde L^f+\widehat\Delta. \]

Both candidate and benchmark bounds contribute, producing the factor two. There is no candidate-count factor because coverage of a fixed \(\theta^*\) evaluates the candidate-side bound only there. The first split fixes the benchmark for its own bound.

6. Why paired risk differences can be more precise

One can also apply concentration directly to a risk difference. Conditional on the first split, fix \(b=\widehat\theta^f\) and define

\[ \begin{aligned} D_\theta(X,Y)&=\ell_\theta(X,Y)-\ell_b(X,Y),\\ D_\theta^f(X)&=\ell_\theta(X,f(X))-\ell_b(X,f(X)). \end{aligned} \]

Then \(L(\theta)-L(b)=\mathbb{E}(D_\theta-D_\theta^f)+\mathbb{E}D_\theta^f\). The corresponding corrected average has variance

\[ \frac{\operatorname{Var}(D_\theta-D_\theta^f)}n +\frac{\operatorname{Var}(D_\theta^f)}{N_2}. \]

For example, the first term includes

\[ \operatorname{Var}(D_\theta-D_\theta^f) =\operatorname{Var}(E_\theta)+\operatorname{Var}(E_b) -2\operatorname{Cov}(E_\theta,E_b), \]

Here \(E_a=\ell_a(X,Y)-\ell_a(X,f(X))\). Common prediction errors cancel in a paired comparison, reducing variance when the covariance is positive. Retain \(\theta\) whenever a valid lower bound for the risk difference is nonpositive. At \(\theta^*\), that difference is nonpositive, so rejection requires a coverage failure. A normal approximation for this version still needs moment and variance conditions.

In the nonconvex quartic example, the cancellation can be carried through exactly. For a benchmark \(b\) held fixed by the first split,

\[ \ell_\theta(Y)-\ell_b(Y) =(\theta^2-b^2)(\theta^2+b^2-2Y). \]

Subtracting its predicted-outcome counterpart gives

\[ D_\theta-D_\theta^f =-2(\theta^2-b^2)(Y-f(X)). \]

The \(Y^2-f(X)^2\) term present in each separate rectifier cancels from the difference. Consequently the corrected empirical risk difference is

\[ \widehat L^{\rm PP}(\theta)-\widehat L^{\rm PP}(b) =(\theta^2-b^2)(\theta^2+b^2-2\widehat\mu^{\rm PP}), \]

where \(\widehat\mu^{\rm PP}\) uses the labeled sample and the held-out unlabeled sample. Its variance, conditional on the benchmark and predictor, is

\[ 4(\theta^2-b^2)^2\left( \frac{\operatorname{Var}(Y-f(X))}{n} +\frac{\operatorname{Var}(f(X))}{N_2}\right). \]

Squaring outcomes in separate risk estimates can demand stronger moments for their variance. The paired difference cancels those terms and reduces to the corrected mean calculation. Retaining covariance between candidates therefore has a concrete benefit.

When \(b^2=\mu\), the excess risk is \((\theta^2-\mu)^2\). At \(\theta=0\) it is \(\mu^2\), although the population gradient vanishes there. The risk inequality excludes this nonglobal stationary point as evaluation becomes precise.

An approximate numerical benchmark is enough, since \(L(\theta^*)\leq L(b)\) for every admissible \(b\). A finite candidate grid is a separate approximation: it may omit the true minimizer. Inference must then concern the grid-restricted target or account for the loss approximation error.

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