Prediction-Powered Conditional Inference

Yang Sui; Jin Zhou; Hua Zhou; Xiaowu Dai

stat.MLMar 5, 2026

Prediction-Powered Conditional Inference

Authors: Yang Sui, Jin Zhou, Hua Zhou, Xiaowu Dai

Organizations: University of California, Los Angeles

Abstract

We study prediction-powered conditional inference in the setting where labeled data are scarce, unlabeled covariates are abundant, and a black-box machine-learning predictor is available. The goal is to perform statistical inference on conditional functionals evaluated at a fixed target point, such as conditional means, without imposing a parametric model for the conditional relationship. Our approach combines localization with prediction-based variance reduction. First, we introduce an RKHS localization method that learns a data-adaptive weight from covariates and reformulates the target conditional moment at the target point as a weighted unconditional moment. Second, we incorporate machine-learning predictions through a correction-based decomposition of this localized moment, yielding a prediction-powered estimator and confidence interval that reduce variance when the predictor is informative while preserving validity regardless of predictor accuracy. We establish nonasymptotic error bounds and, in the abundant-unlabeled regime, minimax-optimal convergence rates for the resulting estimator, prove pointwise asymptotic normality with consistent variance estimation, and provide an explicit variance decomposition that characterizes how machine-learning predictions and unlabeled covariates improve statistical efficiency. Numerical experiments on simulated and real datasets demonstrate valid conditional coverage and substantially sharper confidence intervals than alternative methods.

Figures & tables

Figure 1: Protocol for prediction-powered conditional inference. The procedure takes labeled data, unlabeled covariates, an ML predictor, and a target point $x_{0}$ as inputs. A localization step uses the covariate distribution to learn weights that capture the local structure around $x_{0}$ . The upper block estimates a bias correction from labeled data, while the lower block computes a plug-in term using predictions from the unlabeled data. These components are combined to produce a valid confidence interval $\mathcal{C}(x_{0})$ for the conditional target $\theta_{0}(x_{0})$ .

Figure 2: Empirical

\sigma^{2}_{Y-f}

\sigma^{2}_{Y}

, and

\sigma^{2}_{f}

for the conditional mean at different (age, sex) target points in the census income data. Results are based on

1000

replications.

Figure 3: Empirical coverage rate and average width of nominal

95\%

confidence intervals for the conditional mean at the target point (

\text{age}=70,\text{sex}=1

) in the census income data. The black horizontal line indicates the nominal

95\%

coverage level. Results are based on

1000

replications.

Figure 4: Distributions of empirical coverage rates, RMSE, and average widths of nominal

95\%

confidence intervals for the conditional mean across 343 simulated target points. Results are based on

1000

replications per target point.

Figure 5: Empirical coverage, RMSE, average confidence-interval width, and mean

\widehat{\omega}(x_{0})

as the auxiliary predictor quality

q

varies. Results are averaged over eight target points with

1000

replications per point.

Figure 6: Empirical coverage rates, RMSE, and average widths of nominal

95\%

confidence intervals for the conditional mean at different (age, sex) target points in the census income data. Results are based on

1000

replications per target point.

Figure 7: Empirical coverage rates, RMSE, and average widths of nominal

95\%

confidence intervals for the conditional mean at 50 target points in the BlogFeedback data. Results are based on

1000

replications per target point.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Experiment	$(n,N,\#x_{0})$	$\mathcal{C}_{h}$	$\mathcal{C}_{\lambda}$	Screen constants
Simulation	$(400,10{,}000,343)$	$\{0.8,1.0,1.15,1.2\}$	$35:[0.02,60]_{\log}$	$(12,4,0.10)$
Predictor quality	$(500,10{,}000,8)$	$\{0.8,1.0,1.15,1.2\}$	$35:[0.02,60]_{\log}$	$(12,4,0.10)$
Income (sex $=1$ )	$(300,10{,}000,31)$	$\{1.0,1.2,1.4\}$	$41:[0.1,1000]_{\log}$	$(12,4,22)$
Income (sex $=2$ )	$(300,10{,}000,31)$	$\{1.0,1.2,1.4\}$	$41:[0.1,1000]_{\log}$	$(10,3,15)$
BlogFeedback	$(300,10{,}000,50)$	$\{0.8,0.9,1.0\}$	$81:[0.1,10{,}000]_{\log}$	$(12,4,300)$

Appendix

Table 1: Sample sizes and covariate-only tuning configurations. The final column reports

(c_{\rm op},c_{\rm loc},c_{\rm bias})

Implementation	Coverage	Bias	RMSE	Width	Wall time (s)	Peak memory (MiB)
Two-fold	0.964	0.0467	0.3465	1.3506	11.35 (1.00)	2,680 (1.00)
No-split	0.964	0.0472	0.3469	1.3524	37.15 (3.27)	7,855 (2.93)

Appendix

Table 2: Comparison of the two-fold and no-split PPCI implementations. Wall time and peak GPU memory are averages over the 250 paired runs; parenthetical values are ratios relative to the two-fold implementation.

Component	Configuration	Coverage	Bias	RMSE	Width	Fallback
$c_{\rm bias}$	$0.06$	0.952 (0.930)	0.0124	0.3876	1.4510	1.000
	$0.08$	0.952 (0.931)	0.0126	0.3862	1.4465	0.897
	$0.10$	0.950 (0.925)	0.0222	0.3629	1.3640	0.000
	$0.12$	0.945 (0.909)	0.0402	0.3392	1.2721	0.000
	$0.14$	0.939 (0.891)	0.0577	0.3245	1.2076	0.000
Bandwidth cap	$\{0.8,1.0,1.15,1.2\}h_{0}$	0.950 (0.925)	0.0222	0.3629	1.3640	0.000

Appendix

Table 3: Sensitivity to the bias threshold, bandwidth cap, and regularization-factor envelope. Coverage is the mean (minimum across target points); fallback is the proportion of tuning decisions invoking the least-violation rule.

Data	Reference	Shift	Coverage		RMSE		Width
			PPCI	LO	PPCI	LO	PPCI	LO
Income	Matérn, $0.8h_{0}$	0.1976	0.9329	0.9329	0.5181	0.5782	1.9079	2.1528
	Matérn, $h_{0}$	0	0.9429	0.9375	0.5571	0.6134	1.9079	2.1528
	Matérn, $1.2h_{0}$	0.1681	0.9233	0.9242	0.6404	0.6898	1.9079	2.1528
	Gaussian, ESS matched	0.0630	0.9400	0.9408	0.5364	0.5946	1.9079	2.1528
BlogFeedback	Matérn, $0.8h_{0}$	0.0200	0.9175	0.9404	0.0425	0.0603	0.1442	0.2227

Appendix

Table 4: Sensitivity to the full-data reference smoother. Shift is the mean absolute target change from the main Matérn-

5/2

reference at

h_{0}

. Coverage uses eight targets and 300 intervals per target and method; RMSE and width use the formal 1000-replication summaries at all targets.

$b$	Method	Coverage	Bias	RMSE	Width
0	NW-CV	0.964	0.0302	0.1818	0.7110
0	NW-US	0.970	0.0145	0.2027	0.7958
0	RKHS localization	0.970	-0.0005	0.1674	0.6808
2	NW-CV	0.960	0.0288	0.1833	0.7171
2	NW-US	0.970	0.0135	0.2044	0.8014
2	RKHS localization	0.970	-0.0005	0.1674	0.6813

Appendix

Table 5: Data-driven comparison of NW-PPCI and PPCI with the RKHS localization weight over 500 paired replications.

$h$	Method	Coverage	Bias	RMSE	$\widehat{\Delta}_{2}$
0.08	NW	0.956	0.0260	0.1766	0.007690
0.08	RKHS localization	0.980	0.0023	0.3806	0.000047
0.20	NW	0.652	0.1851	0.2175	0.047563
0.20	RKHS localization	0.968	0.0019	0.2601	0.000127
0.40	NW	0.000	0.5879	0.5943	0.147705
0.40	RKHS localization	0.952	-0.0049	0.1960	0.000259

Appendix

Table 6: Same-bandwidth comparison for

b=0

over 500 paired replications. Here

\widehat{\Delta}_{2}

is the normalized empirical quadratic localization error.

Figure 8: Diabetes progression data illustration. Left: split conformal predictive intervals for the outcome

Y

at test covariates

X

. Right: a bootstrap uncertainty band for the conditional mean

\theta_{0}(x)=\mathbb{E}[Y|X=x]

Explore similar work

MEC: Machine-Learning-Assisted Generalized Entropy Calibration for Semi-Supervised Mean Estimation
Apr 7, 2026Se Yoon Lee, Jae Kwang KimPrediction-Powered Inference Semi-Supervised Learning
Texas A&M University · Iowa State University
Prediction-powered Inference by Mixture of Experts
Apr 30, 2026Yanwu Gu, Linglong Kong, Dong XiaUncertainty Quantification Mixture-of-Experts Models
Department of Mathematics, Hong Kong University of Science and Technology · Department of Mathematical and Statistical Sciences, University of Alberta
How Synthetic Labels Improve Conformal Prediction: A Perspective on Conditional Coverage
Sep 27, 2026Qianyi Chen, Bo LiQuantile Regression Conformal Prediction
School of Economics and Management, Tsinghua University

Apr 7, 2026J/K move · Enter open · S save

MEC: Machine-Learning-Assisted Generalized Entropy Calibration for Semi-Supervised Mean Estimation

Se Yoon Lee, Jae Kwang Kim

Texas A&M University · Iowa State University

Prediction-Powered Inference Semi-Supervised Learning Semiparametric Inference Confidence Region Estimation

Open paper AskPDF arXiv

Obtaining high-quality labels is costly, whereas unlabeled covariates are often abundant, motivating semi-supervised inference methods with reliable uncertainty quantification. Prediction-powered inference (PPI) leverages a machine-learning predictor trained on a small labeled sample to improve efficiency, but it can lose efficiency under model misspecification and suffer from coverage distortions due to label reuse. We introduce Machine-Learning-Assisted Generalized Entropy Calibration (MEC), a cross-fitted, calibration-weighted variant of PPI. MEC improves efficiency by reweighting labeled samples to better align with the target population, using a principled calibration framework based on Bregman projections. This yields robustness to affine transformations of the predictor and relaxes requirements for validity by replacing conditions on raw prediction error with weaker projection-error conditions. As a result, MEC attains the semiparametric efficiency bound under weaker assumptions than existing PPI variants. Across simulations and a real-data application, MEC achieves near-nominal coverage and tighter confidence intervals than CF-PPI and vanilla PPI.

Abstract

Explore similar work

MEC: Machine-Learning-Assisted Generalized Entropy Calibration for Semi-Supervised Mean Estimation

Prediction-powered Inference by Mixture of Experts

How Synthetic Labels Improve Conformal Prediction: A Perspective on Conditional Coverage