Prediction-Powered Adaptive Shrinkage Estimation

We thank the reviewer for considering our idea convincing and thoughtful comments.

• Sharing information across problems: We agree with the reviewer that further motivation of why and how it is possible to share information across problems will improve our exposition. In our revision we will add further motivation. Below is an alternative attempt to further explain and elaborate on Prop 5.3. Suppose as in that proposition that $n_j$ is the same across all problems, $N_j=\infty$, and that second moments $\rho_j,\tau_j,\sigma_j$ are identical across all problems. We make these assumptions throughout the remainder of our response. Then we could ask: What is the best convex combination of $\hat{\theta}_j^{PT}$ and $\tilde{Z}_j^f$ in the following sense:

  $\omega_j^*\in\mathrm{argmin}_{\omega_j}\mathbb E[(\theta_j-(\omega_j\hat{\theta}_j^{PT}+(1-\omega_j)\tilde{Z}_j^f)^2].$

  By direct calculation (since the RHS is a convex quadratic in $\omega_j$) we find that: $\omega_j^*=\frac{\mathbb E[(\theta_j-\tilde{Z}_j^f)^2]}{\mathbb E[(\theta_j-\tilde Z_j^f)^2]+\tilde\sigma^2}.$ This implies the following intuitive result: the larger the MSE $\mathbb E[(\theta_j-\tilde{Z}_j^f)^2]$, the less weight we should assign to $\tilde{Z}_j^f$. Note that if we have a single problem $m=1$, then if $n_j$ is sufficiently large, we could estimate $\tilde{\sigma}^2$ accurately. However, we cannot estimate $\mathbb E[(\theta_1-\tilde{Z}_1^f)^2]$ accurately since we only have a single $\theta_1$ (a single problem). At best, we can compute an unbiased estimate of this quantity since: $\mathbb E[(\hat{\theta}_1^{PT}-\tilde{Z}_1^f)^2]-\tilde{\sigma}^2=\mathbb E[(\theta_1-\tilde{Z}_1^f)^2].$

  Now suppose we have multiple problems, then we can now learn how good the ML predictor is for estimating the $\theta_j$ by sharing information across problems. To wit, as $m \to \infty$

  $\frac{1}{m}\sum_{j=1}^m(\hat{\theta}_j^{PT}-\tilde{Z}_j^f)^2-\tilde\sigma^2\stackrel{\mathbb P}{\to}\mathbb E[(\theta_j-\tilde Z_j^f)^2],$

  where we emphasize that the mean squared on the RHS also integrates with respect to the meta-distribution $\mathbb P_{\eta}$ that models the distribution of the $\theta_j$. Thus we can also estimate the optimal $\omega_j^*$. In words: sharing information across problems allows us to learn how good the ML predictor is and to then decide how much to shrink toward it. Our implementation shares information in a similar way but is more involved due to the heteroscedasticity across problems (not all second moments are identical) and also for heterogeneity in $n_j$, $N_j$.

• Role of meta-distribution: See above. We note that we could state our result in a frequentist setting wherein all parameters are deterministic, similar to the classical result of James-Stein. We preferred to state results in terms of a meta-distribution, since we thought this would be a more familiar setup for audience at ICML.

• Assumption on known second moments: We agree that our assumption on second-moment parameters $(\sigma^2_j, \gamma_j)$ are not easily satisfied in practice. In Sec 2 of paper, such treatment is more of a theoretical convenience, and sample-based estimates are used for real-world datasets. Therefore, it is true that an implicit requirement for our theory of PAS to work well in many practical settings is that the sample-based estimates are good enough. While our regime ($m\to\infty$, $n_j,N_j$ finite) cannot yield asymptotic result on these sample-based estimates, we make the following remarks:

  1. In practice, PAS works well even when $n_j$ is very small (and so the estimates are noisy). To highlight this point further and motivated by your suggestion, we reran our real data examples with different labeled/unlabeled splits, going from 1%-40%. PAS still dominates other baselines even when only 1% of the data is labeled (e.g. $n_j < 10$ for Amazon Review). [Link to the plots]

  2. In our response to Reviewer Tcxk, we propose a new variant of PAS called UniPAS, whose asymptotic guarantee does not require knowledge of second moments. UniPAS also has competitive empirical performance with PAS.

• Other comments:

  • L204: The underline is intentional to denote the abbreviation "cl" for classical estimator. We can clarify this.
  • L170: Thanks for catching the typo; we will fix it.
  • Covariate Shift: We interpret this as potential differences in $P_{\eta_j}(X_{ij})$ across problems $j$. In the synthetic model, the mean of $X_{ij}$ varies with $\eta_j$, but that is not an assumption that PAS makes. What makes "information sharing" more effective, as mentioned above, is when $(\theta_j-\tilde{Z}_j^f)^2$ is small on average across problems. If the question refers to train/test covariate shift within problems, this violates the PPI assumption.

We thank the reviewer for a very helpful report, which has helped us improve on this work.

• We agree that we should expand on connections to related work, which will be added as a new section in the appendix after our revision. Briefly:

  • Fisch et al. 2024 (StratPPI): The starting point is similar to ours with multiple parameters $\theta_1,...,\theta_m$. However, in StratPPI, these parameters are not of interest per-say. Instead there exists known weights $w_1,...,w_m$ such that the parameter of true interest is $\theta := \sum_{j} w_j \theta_j$. The ultimate goal is to come up with an unbiased and low variance estimator of $\theta$.
  • Rosenman et al (2023): When their stratum weights are the same, their estimator corresponds to our shrinkage baseline up to differences in the one-dimensional family of shrinkage weights (our Eq. (15)).
  • PPI++: Suppose $n_j=n, N_j=N$ for all problems, then indeed we can cast our setting into PPI++ with a parameter vector. Moreover, PT is then using the same $\lambda$ for all problems (more of which below).

• Previously, we were comparing to PPI++ applied separately to each problem (with its own $\lambda_j$). PPI++ itself requires asymptotics with $n_j,N_j\to\infty$. It would be possible to extend to asymptotics with $m,n\to\infty$; however, based on the reviewer comments we now have a more compelling methodological alternative: if we only seek to compare ourselves to PPI++ with a single $\lambda$, then we can learn the optimal $\lambda$ in asymptotics with $n$ fixed, $m\to\infty$ (a result complementary to the PPI++ paper which takes $m$ fixed and $n\to\infty$). In closed form, we have $\lambda^*:=\frac{\sum_{j=1}^mn_j^{-1}\gamma_j}{\sum_{j=1}^m\frac{n_j+N_j}{n_jN_j}\tau_j^2}$ which coincides with the optimal weight selected by PPI++. We further clip $\lambda^*$ to be within $[0, 1]$.

  Without knowing $\tau_j, \sigma_j, \gamma_j$, we can still plug in their sample-based estimates (see Point 2 for Reviewer Tcxk) and obtain $\hat\lambda$ (after clipping it as well), which has the property that $\hat\lambda - \lambda^* \overset{L^2}{\to} 0$ as $m\to\infty$ (this is different from the per-problem case when $n_j \to \infty$ is needed). We denote this method as UniPT.

• Taking UniPT as a starting point, we develop UniPAS. This has an asymptotic guarantee ($n_j, N_j$ fixed, $m\to\infty$) analogous to PAS replacing PT with UniPT (that is, we try to be asymptotically at least as good as PPI++ with a single $\lambda$, while PAS was trying to at least as good as PPI++ with a separate-per problem $\lambda$). The upshot is that UniPAS does not require knowledge of $\tau_j, \sigma_j, \gamma_j$. In empirical results UniPAS is competitive with PAS (although PAS is slightly better).

• In addition to including two new methods (UniPT, UniPAS), we have reran our real data examples with labeled/unlabeled split ratios ranging from 1% to 40%. [Link to plots]

• On the analogy of $Z_j^f$ and $\mathbb E[\theta\mid\psi]$: We agree that the analogy is somewhat loose and we will provide more details in the revision. Our goal is to provide a heuristic motivation for the one-dimensional parameterized family of weights $\omega_j(\cdot)$ (whose ultimate success is judged by the empirical results). To elaborate on this: in Eq. (12) of Sec 3, we find that the best weights take the form: $\frac{\omega}{\omega+\sigma_{\varepsilon}^2},$ for $\omega = \mathbb E[(\theta - \mathbb E[\theta \mid \phi])^2] = \mathbb E[\mathrm{Var}[\theta \mid \phi]]$. If we instead take the best convex combination of $\hat{\theta}^{cl}-\xi$ and $\mathbb E[\theta]$ (instead of $\mathbb E[\theta\mid\phi]$), the optimal weights again take the form above, now with $\omega=\mathbb E[(\theta-\mathbb E[\theta])^2] = \mathrm{Var}(\theta)$. Now suppose that we ask for the best convex combination (not necessarily a Bayes predictor) between $\hat{\theta}^{cl}-\xi$ and $h(\phi)$ where $h(\cdot)$ is some fixed function. Then we can show that the optimal convex combination is given by the form above with

  $$\omega:=\mathbb E[(\theta-h(\phi))^2] = \mathbb E[\mathrm{Var}[\theta\mid\phi]] + \mathbb E[(h(\phi)-\mathbb E[\theta\mid\phi])^2].$$

  (The above expression is interesting as it forces us to inflate "$\omega$", i.e., to shrink less toward $h(\phi)$ in a way that depends on how close $h(\phi)$ is to $\mathbb E[\theta\mid\phi]$.) The takeaway is that for many possible predictors, the optimal weights have the same parameterized form up to the single parameter $\omega$ that varies according to the quality of the predictor. This motivates our one-dimensional family of weights. Once this family has been motivated, we learn $\omega$ in (15) in a way that does not depend on the above analogy at all by minimizing CURE.

We thank the reviewer for acknowledging the novelty and contribution in our paper, as well as the many constructive comments. We hope our response below addresses all the concerns directly.

Background on PPI

The core idea of Prediction-Powered Inference (PPI) is as follows. Given existing black-box ML predictors $f$, how can we use them to enhance classical statistical procedures and improve their efficiency? The focus in this literature is on employing safeguards so that if the ML predictor is good, then statistical gains can be large, while if the ML predictor is bad, then one still retains some form of statistical guarantees (such as consistency). The focus on this literature is not on how to train the best possible ML models (most results are agnostic to that), but rather how to wrap around an existing arbitrary $f$ to improve statistical procedures. In this paper, we follow this established perspective in the PPI literature.

Here, arbitrary typically means that $f$ can be of any functional form or a total black-box (e.g., API calls to LLMs), and should thus be considered given and fixed. While adapting to the quality of $f$ is important, an “arbitrarily bad” $f$ is rarely considered in practice, since prediction-powered methods are generally deployed only when $f$ is expected to provide at least some signal about the response. Nevertheless, PAS does include mechanisms to adapt to both poor and high-quality predictors. To demonstrate this, we revisit the synthetic model and consider the following predictors:

1. a (newly added) very poor predictor $f_r$, which outputs a $Unif[0, 2]$ value for all inputs,
2. the near-optimal predictor $f_1(x) = x^2$ from the paper (in fact, $f_1$ closely approximates $E[Y_{ij} | X_{ij} = x]$).

Table 1: MSE (± se) $\times 10^{-3}$ in synthetic model ($K=200$ replicates)

In the extreme case where $f_r$ is very biased and provides no useful information, PAS effectively defaults to the classical estimator, showing robustness against poor predictions. On the other hand, with a very good $f_1$, PAS tracks the predictive mean closely and perform best. Importantly, no over-shrinkage is observed. In words, PAS learns how good the predictor is in a data-driven way.

Bounded moment assumptions

We appreciate the reviewer's question regarding the finite moment assumptions. We first remark that these assumptions are satisfied when the response and prediction are bounded (as in the Amazon and Galaxy datasets), or when their joint distribution is reasonably well-behaved (as in our synthetic setting).

We note that some moment assumptions are standard across the PPI literature. Specifically, finite second moments of the joint model $(Y_{ij}, f(X_{ij}))$ are generally needed for basic procedures like power tuning. These are typically viewed as mild assumptions, although we admit that they preclude heavy-tailed distributions.

Our work extends PPI to compound mean estimation using shrinkage principles, particularly risk minimization via an unbiased estimate (CURE). To prove the asymptotic optimality of PAS, we require finite fourth moments. This is a technical condition used to control the variance of the risk estimate itself.

Finally, addressing the reviewer's question about violations: if the fundamental second moment assumption fails, the variance-reduction premise of PPI itself becomes ill-defined. If only the fourth moment assumption is violated, PAS may lose its formal asymptotic guarantee to outperform power tuning and the predictive mean. However, CURE remains an unbiased risk estimate provided that second moments exist.

A light-weight approach

A key strength of PAS is its low computational overhead compared to both classical and PPI estimators. Although PAS uses both power tuning and adaptive shrinkage, the first stage yields a closed-form expression for the optimal tuning parameter $\lambda_j^*$; the second stage involves optimizing CURE over a one-dimensional parameter space, and each evaluation is inexpensive since CURE has an analytic form as well. This makes the optimization very fast in practice.

We first precompute all ML predictions. After this step, all estimators (including PAS) are very fast to compute. Below we benchmark the runtimes for three estimators on the Amazon Review dataset with the precomputed ML predictions. In each trial, the task is to estimate the mean product ratings for $m = 200$ products. We then record the time to construct each estimator by taking the average of 100 repeated constructions. The table reports the mean and max time (in milliseconds) to construct each estimator across 10 such trials.

We thank the reviewer for the careful assessment. The thoughtful questions have helped us improve our work. For better exposition, we slightly reordered our responses to the questions.

1. Finite $m$ results. Our current theoretical analysis permits for finite sample bounds. We will track these results more explicitly in the revision. For instance, the $o(1)$ term in Prop 5.1 is actually $O(1/\sqrt{m})$.

2. How are variance parameters estimated in practice? For real-world datasets, we use sample-based estimates for the variance parameters. We estimate $\sigma^2_j$ and $\gamma_j$ with the standard unbiased estimators

   $$\hat{\sigma}^2_j = \frac{1}{n_j-1} \sum_{i=1}^{n_j} (Y_{ij} - \bar Y_j)^2, \quad \hat\gamma_j = \frac{1}{n_j-1}\sum_{i=1}^{n_j}(Y_{ij}-\bar Y_j)(f(X_{ij}) -\bar{Z}_j^f)$$

   using labeled data. For the prediction variance $\tau^2_j$, we use predictions in both labeled/unlabeled data:

   $$\hat\tau_j^2 = \frac{1}{n_j + N_j -1}\bigg(\sum_{i=1}^{n_j}(f(X_{ij}) - \hat Z_j^f)^2 + \sum_{i=1}^{N_j}(f(\tilde X_{ij}) - \hat{Z}_j^f)^2 \bigg)$$

   where $\hat Z_j^f = \frac{1}{n_j + N_j -1}\left(\sum_{i=1}^{n_j}f(X_{ij}) + \sum_{i=1}^{N_j}f(\tilde X_{ij})\right)$.

3. Impact of estimating variance parameters (with Point 2). We agree with the reviewer that our current theory does not account for estimation errors due to plugging in sample-based estimates of the variance parameters. One possible remedy would be to consider asymptotics in which all of $n_j, N_j, m \to \infty$. However, we prefer asymptotics with $n_j, N_j$ fixed and $m \to \infty$ to represent the regime in which we have a lot of very noisy/difficult individual problems and also to distinguish results from standard setup in the PPI literature that keeps $m$ fixed and takes $n_j, N_j \to \infty$. We note PAS with plug-in variance estimates empirically works well even for very small $n_j$, see point 4 below.

   Motivated by the reviewer's comment, we now consider two further methods: UniPT (that is, power-turning with the same $\lambda$ for all problems, see response to reviewer qLFS) and UniPAS, which is similar to PAS but uses the same power-tuning parameter across problems. We can prove that UniPAS has asymptotic ($m \to \infty$, $n_j,N_j$ fixed) risk less or equal to that of UniPT, PPI, the classical estimator, and the prediction mean. The result for UniPAS accounts for the data-based estimation of the variance parameters. Here is a sketch of UniPAS:

   • We start with UniPT by estimating a single power-tuning parameter $\hat{\lambda}$ (as in our response to reviewer qLFS) that has the property that $\hat{\lambda} - \lambda^* \to 0$ in $L^2$ as $m \to \infty$ (with $n_j,N_j$ fixed) where $\lambda^*$ is the optimal single tuning parameter. Call the resulting estimator $\hat{\theta}_j^{UniPT}$.

   • We come up with stable working estimates of $Var[\hat{\theta}_j^{UniPT}]$ by pretending $\sigma_j^2, \tau_j^2, \gamma_j$ are the same across all $j$ (but sample sizes can vary). The common values can be estimated accurately by averaging the estimates in Point 2 over all $m$. Plugging these in, we derive working estimates $\bar{\sigma}_j^2$ of $Var[\hat{\theta}_j^{UniPT}]$. Then we consider the one-dimensional parameterized family of weights $\bar{\omega}_j(\omega) = \frac{\omega}{\omega + \bar{\sigma}_j^2}.$ This family retains the property that for $\omega=0$ it recovers prediction mean and for $\omega=\infty$ it recovers UniPT.

   • Pretend momentarily that $\bar{\sigma}_j$ and $\hat{\lambda}$ are deterministic. (This is not actually needed; by steps 1 and 2 above asymptotically these converge to a deterministic limit.) Suppose we consider the family of estimators

     $$\bar{\omega}_j(\omega)\hat{\theta}_j^{UniPT} + (1-\bar{\omega}_j(\omega))\tilde Z_j^f.$$

     An unbiased estimate of risk is given by CURE, $\frac{1}{m}\sum_{j=1}^m (2\bar{\omega}_j(\omega) - 1)Var[\hat{\theta}_j^{UniPT}] + 2(1 - \bar{\omega}_j(\omega)) Cov[\hat{\theta}_j^{UniPT}, \tilde Z_j^f] +(1 - \bar{\omega}_j(\omega))^2\big(\hat{\theta}_j^{UniPT} - \tilde Z_j^f\big)^2,$ which is analogous to the expression of CURE for PAS. Now here comes the punchline. If above we replace $Var[\hat{\theta}_j^{UniPT}]$ and $Cov[\hat{\theta}_j^{UniPT}, \tilde Z_j^f]$ by unbiased estimates (analogous to Point 2), then we still retain an unbiased estimate of risk. Moreover, since we are averaging over $m$ with $m \to \infty$, we can establish asymptotic uniform consistency of our objective to the true risk, and thus we can establish asymptotic optimality for UniPAS.

4. Different $n_j/N_j$ and split ratios. Following the reviewer's suggestion we reran our real data analyses with the new methods and different labeled/unlabeled splits, going from 1%-40%. [Link to the plots]