Technical Summary: Predictions as Surrogates: Revisiting Surrogate Outcomes in the Age of AI

1. Problem Statement

In many scientific domains, acquiring the true outcome variable $Y$ is costly, time-consuming, or invasive, leading to datasets where $Y$ is only partially observed. Researchers often rely on surrogate outcomes (auxiliary variables $\hat{Y}$) that are cheaper to obtain and correlated with $Y$. While traditional surrogate models have been established in biostatistics and economics, the rise of AI has introduced a new class of surrogates: predictions from pre-trained machine learning models (e.g., Large Language Models, computer vision models).

The paper identifies a gap in Prediction-Powered Inference (PPI), a nascent field that seeks to incorporate these AI predictions into statistical inference. Existing PPI methods often treat predictions as direct substitutes for outcomes or apply simple linear corrections. However, these approaches face limitations:

1. Inefficiency: They may not achieve the semiparametric efficiency bound, especially when predictions are biased or systematically deviate from the true outcome.
2. Overfitting Bias: When flexible machine learning models are used to learn the relationship between surrogates and outcomes, using the same data for both learning and inference (same-sample fitting) can introduce first-order bias, invalidating confidence intervals.
3. Modality and Distribution Mismatches: Pre-trained models often suffer from modality mismatches (e.g., predicting from images while inference requires tabular covariates) or distribution shifts (trained on general data, applied to specific subpopulations), leading to miscalibrated predictions.

The core problem is to develop a statistically rigorous, computationally feasible, and efficient framework that leverages pre-trained AI predictions as surrogates while correcting for their inherent biases and ensuring valid inference.

2. Methodology: Recalibrated Prediction-Powered Inference (RePPI)

The authors propose Recalibrated Prediction-Powered Inference (RePPI), a method that unifies the classical surrogate outcome literature (specifically the work of Robins et al. [1994] and Chen & Chen [2000]) with modern PPI.

2.1 Theoretical Foundation

The target parameter $\theta^\star$ is defined via an estimating equation $E[U_{\theta}(X, Y)] = 0$ (or equivalently, minimizing a convex loss $E[\ell_\theta(X, Y)]$). The data consists of $n$ labeled samples $(Y, X, \hat{Y})$ and $N$ unlabeled samples $(X, \hat{Y})$.

The authors establish that the optimal imputed loss $g_\theta$ that minimizes the asymptotic variance of the estimator is given by:
$$\nabla g_\theta^\star(X, \hat{Y}) = \frac{1}{1+r} E[\nabla \ell_\theta(X, Y) \mid X, \hat{Y}]$$
where $r = \lim n/N$. This optimal score represents the conditional expectation of the true gradient given the surrogate and covariates.

2.2 The RePPI Algorithm

To operationalize this theory, RePPI employs two key innovations:

1. Flexible Recalibration via Machine Learning: Instead of assuming a parametric form for the conditional expectation, RePPI uses flexible machine learning algorithms (e.g., random forests, neural networks) to estimate the optimal score function $s(X, \hat{Y}) \approx E[\nabla \ell_{\hat{\theta}_0}(X, Y) \mid X, \hat{Y}]$, where $\hat{\theta}_0$ is an initial estimator (e.g., from labeled data only).
2. Cross-Fitting: To prevent overfitting bias when using flexible learners, the labeled data is split into folds. The score function is trained on one subset and evaluated on another (out-of-sample). This ensures the independence required for valid asymptotic theory.
3. Matrix Power Tuning (Control Variates): To guarantee efficiency even if the score estimate is imperfect, RePPI applies a linear transformation matrix $\hat{M}$ to the estimated score. This matrix is estimated as the empirical covariance between the true gradient and the estimated score, normalized by the variance of the score. This acts as an optimal control variate, ensuring the estimator is never worse than the "XY-only" baseline (which ignores predictions).

The final estimator minimizes a linearized empirical risk:
$$\hat{\theta}_{RePPI} = \arg\min_\theta \left( \frac{1}{n}\sum_{i=1}^n \ell_\theta(X_i, Y_i) - \frac{1}{n}\sum_{i=1}^n \hat{g}_\theta(X_i, \hat{Y}_i) + \frac{1}{N}\sum_{i=n+1}^{n+N} \hat{g}_\theta(X_i, \hat{Y}_i) \right)$$
where $\hat{g}_\theta$ is constructed using the recalibrated score.

3. Key Contributions

1. Unified Framework: The paper formally connects the decades-old surrogate outcome literature with the emerging PPI field, clarifying that pre-trained model predictions are a specific type of surrogate.
2. Recalibrated PPI (RePPI): The authors introduce a new estimator that:
   • Achieves semiparametric efficiency (smallest asymptotic variance) if the optimal score is consistently estimated.
   • Is always more efficient than the XY-only estimator, even if the score estimation is inconsistent (due to the control variate mechanism).
   • Maintains convexity of the optimization objective (unlike some previous PPI proposals) by using a linearized imputed loss.
3. Theoretical Guarantees: The paper provides rigorous proofs for asymptotic normality and efficiency under cross-fitting, demonstrating that the method is robust to the use of high-capacity, flexible learners.
4. Analysis of Recalibration Benefits: The authors theoretically and numerically analyze three specific scenarios where recalibration is critical:
   • Modality Mismatch: When predictions are based on unstructured data (e.g., images) but inference requires structured covariates.
   • Distribution Shift: When the prediction model is trained on a different population than the target inference population.
   • Discrete Predictions: When models output coarse, uncalibrated discrete categories (e.g., 1-5 stars) rather than continuous values.

4. Results

The paper validates RePPI through both stylized theoretical models and real-world applications using state-of-the-art AI models.

4.1 Theoretical Efficiency Gains

In stylized models, RePPI consistently achieves lower asymptotic variance than existing methods (XY-only, standard PPI, and PPI++):

• Modality Mismatch: RePPI reduces variance by effectively utilizing the information in the unstructured modality while correcting for the lack of structured covariates in the prediction.
• Distribution Shift: RePPI corrects for bias introduced when the prediction model's training distribution differs from the inference population, whereas standard PPI can perform worse than the XY-only baseline in these cases.
• Discrete Predictions: RePPI learns a non-linear mapping from discrete predictions to the outcome scale, outperforming PPI++ which relies on linear corrections.

4.2 Empirical Applications

Three real-world applications demonstrate significant gains in effective sample size (reduction in required labeled data to achieve a target confidence interval width):

1. X-ray Diagnosis (NIH ChestX-ray14): Using a DenseNet-121 model trained on CheXpert to predict pleural effusion. RePPI reduced the required labeled sample size by ~5.7% compared to PPI++ to achieve the same interval length.
2. Toxicity Measurement (CivilComments-WILDS): Using a DistilBERT model trained on Wikipedia comments to predict toxicity in news comments (distribution shift). RePPI reduced the required labeled sample size by ~5.8%.
3. Wine Reviews (WineEnthusiast): Using GPT-5.4 to predict ratings from text reviews (discrete predictions). RePPI reduced the required labeled sample size by ~17.6% to 19% compared to PPI++, highlighting the substantial benefit of non-linear recalibration for discrete surrogates.

In all cases, RePPI maintained valid 95% confidence interval coverage, whereas standard PPI often produced wider intervals or invalid coverage due to miscalibration.

5. Significance and Claims

The paper claims that RePPI represents a significant step forward in leveraging AI for statistical inference by:

• Bridging Theory and Practice: It provides a practical computational strategy (cross-fitting + linearized loss) to achieve the theoretical efficiency bounds established in semiparametric statistics, which were previously difficult to attain with flexible machine learning models.
• Robustness: Unlike previous methods that might fail or become inefficient when predictions are biased, RePPI is designed to be "safe," guaranteeing improvement over ignoring the predictions entirely.
• Addressing the AI Era: It specifically addresses the unique challenges of using pre-trained foundation models (high dimensionality, unstructured data, distribution shifts) as surrogates, moving beyond the assumption that predictions are perfect or simply linear transformations of the truth.

The authors modestly note limitations, including the reliance on asymptotic theory, the assumption of i.i.d. data (though extensions to distribution shift are discussed), and the current focus on using predictions as surrogates rather than utilizing richer model outputs like embeddings directly. Nevertheless, the work establishes a robust foundation for "prediction-powered" science in the age of AI.