Calibeating Prediction-Powered Inference

1. Problem Statement

The paper addresses the problem of semisupervised mean estimation in settings where:

• A small labeled dataset $\{(X_i, Y_i)\}_{i=1}^n$ is available, containing covariates $X$ and outcomes $Y$.
• A large unlabeled dataset $\{\tilde{X}_j\}_{j=1}^N$ is available, containing only covariates.
• A black-box prediction model $m(X)$ exists, trained on external data or historical sources, which provides a score for the unlabeled data.

The Core Challenge:
Standard approaches like Augmented Inverse Probability Weighting (AIPW) and Prediction-Powered Inference (PPI) use the prediction score $m(X)$ as a regression adjustment to improve efficiency. However, these methods treat $m(X)$ as fixed. If the prediction model is miscalibrated (i.e., its numerical values do not align with the true conditional mean $E[Y|X]$, even if it ranks individuals correctly), the residual variance $Y - m(X)$ remains high. This leads to inefficient estimators with wider confidence intervals, despite the model being informative. The paper seeks a method to improve efficiency without retraining the underlying black-box model.

2. Methodology: Calibrated Prediction-Powered Inference (Calibrated PPI)

The authors propose a framework where the prediction score is post-hoc calibrated on the small labeled sample before being used for semisupervised estimation.

The Workflow

1. Input: Labeled data $(X, Y)$, unlabeled data $\tilde{X}$, and a pre-trained score $m(X)$.
2. Calibration Step: Learn a transformation map $f: \mathbb{R} \to \mathbb{R}$ on the labeled data to minimize the squared error between the transformed score and the outcome:
   $$\hat{f} \in \arg\min_{f \in \mathcal{H}} \frac{1}{n} \sum_{i=1}^n (Y_i - f(m(X_i)))^2$$
   where $\mathcal{H}$ is a class of functions (e.g., linear, isotonic).
3. Calibrated Score: Define the new score $m^*_n(X) = \hat{f}(m(X))$.
4. Estimation: Compute the semisupervised mean estimator using the calibrated scores:
   $$\hat{\psi}_{cal} = \frac{1}{n+N} \left( \sum_{i=1}^n m^*_n(X_i) + \sum_{j=1}^N m^*_n(\tilde{X}_j) \right)$$

Key Theoretical Insights

• Exact AIPW Representation: The paper proves that the plug-in estimator $\hat{\psi}_{cal}$ is exactly equivalent to an AIPW estimator using the calibrated score $m^*_n$. Specifically, because the calibration minimizes squared error on the labeled set, the residual term $\frac{1}{n}\sum (Y_i - m^*_n(X_i))$ is zero (or orthogonal to the calibration class), satisfying the double-robustness property.
• Efficiency Gain: By transforming $m(X)$ to better approximate the conditional mean $E[Y|X]$ (or $E[Y|m(X)]$), the residual variance is reduced. Since the asymptotic variance of AIPW estimators depends on this residual variance, calibration directly improves efficiency.
• No Retraining Required: The underlying black-box model $m(X)$ is never retrained; only a lightweight post-processing step is applied.

Specific Calibration Methods Studied

1. Linear Calibration: Fits an affine map $f(t) = a + bt$.
   • Result: The authors show that Linearly Calibrated PPI is first-order asymptotically equivalent to PPI++ (which uses empirical efficiency maximization to find an optimal scaling factor).
2. Isotonic Calibration: Fits a non-decreasing function $f$ (using Pool Adjacent Violators Algorithm or monotonic regression trees).
   • Result: This is a non-parametric approach that adapts to the data shape. The paper establishes that isotonic calibration achieves the smallest asymptotic variance among all estimators based on monotone transformations of the original score.

3. Key Contributions

1. Introduction of Calibrated PPI: A simple, plug-in framework that post-hoc calibrates prediction scores to improve semisupervised mean estimation without retraining the base model.
2. Theoretical Guarantees for Isotonic Calibration:
   • Calibeating: Isotonic calibration "calibeats" (improves upon) the original score and simpler rules (like linear calibration) in terms of both predictive accuracy and estimator efficiency.
   • First-Order Optimality: Once a score is isotonicly calibrated, no further post-processing of that fitted score yields additional first-order gains in efficiency. It is "first-order optimal" against further post-processing.
   • Asymptotic Normality: The paper establishes asymptotic normality and valid Wald inference for the isotonic-calibrated estimator, with a convergence rate of $O(n^{-2/3})$ for the calibration error, which is negligible compared to the $O(n^{-1/2})$ estimation error.
3. Clarification of Existing Methods:
   • Shows that the original PPI estimator is a special case of AIPW that can be inefficient when the prediction model is accurate but uncalibrated.
   • Demonstrates that PPI++ is essentially an AIPW estimator with empirical efficiency maximization and is first-order equivalent to linear calibration.
4. Empirical Validation: Extensive simulations and real-data experiments (including LLM evaluation benchmarks) showing that calibrated estimators often outperform standard PPI and AIPW, and are competitive with or superior to PPI++.

4. Results

• Simulations: In synthetic settings with miscalibrated scores, calibrated methods (LinearCal, IsoCal) significantly reduced Root Mean Squared Error (RMSE) and confidence interval widths compared to raw PPI and AIPW.
• Real-Data Benchmarks:
  • Census Income & Forest/Galaxies: Calibration methods (especially LinearCal and AutoCal) consistently outperformed PPI and were competitive with PPI++.
  • LLM Evaluation (PPE): In a practical setting using Large Language Model (LLM) reward models, raw AIPW/PPI failed dramatically when the reward model scores were poorly scaled (MSE was orders of magnitude higher than labeled-only). Calibrated PPI (specifically LinearCal) restored validity and efficiency, achieving MSE reductions of ~7.5% to 10.5% compared to labeled-only estimation, effectively making cheap, public LLM evaluators usable for rigorous inference.
• Label Efficiency: The methods achieved significant "label savings," meaning they could achieve the same precision as labeled-only estimation with fewer labeled examples (e.g., 7-10% savings in the LLM experiments).

5. Significance

• Bridging Communities: The paper connects the emerging field of Prediction-Powered Inference (machine learning) with the established literature on Semiparametric Statistics, Causal Inference, and Missing Data. It clarifies that PPI methods are essentially instances of AIPW and Targeted Minimum Loss Estimation (TMLE).
• Practical Utility: It provides a "lightweight" solution for practitioners who have access to powerful black-box models but lack the resources to retrain them. A simple calibration step can unlock the full efficiency potential of these models.
• Robustness: The framework ensures valid inference even when the prediction model is misspecified, provided the calibration step is performed correctly.
• Software: The authors provide a Python package (ppi_aipw) implementing these methods, making them accessible for immediate application in semisupervised learning and causal inference tasks.

In summary, the paper argues that calibration is not just a predictive tool but a crucial inferential step for semisupervised learning. By aligning prediction scores with the outcome scale, one can achieve "calibeating"—improving both prediction and estimation efficiency without the computational cost of retraining complex models.