Prediction-Augmented Trees for Reliable Statistical Inference

Machine learning (ML) is increasingly central to scientific discovery, but translating ML predictions into rigorous scientific claims still requires statistical inference on gold-standard labeled data. A setting that has emerged across the sciences provides the analyst with a small labeled sample, a much larger unlabeled sample, and a pre-trained predictor, with the goal of constructing a confidence interval for a quantity of interest, as was formalized in Angelopoulos et al. (2023). In this ML-driven scientific discovery literature, labeled data is typically the binding constraint on precision while computation is comparatively unconstrained, motivating the design question of how additional computation can extract more inferential value from a fixed labeled set, a fixed unlabeled set, and a fixed predictor. To address this question, we propose the *Prediction-Augmented Residual Tree* (PART) estimator. PART takes the same inputs as the original PPI estimator and replaces their single global rectifier with locally computed rectifiers obtained from an adaptive partition of the feature space. We show that PART outperforms existing methods by producing tighter confidence intervals across real-world datasets from ecology, astronomy, and census reports, among other domains. We then describe and analyze PAQ, an estimator that arises when considering the limit of PART when the depth of its tree grows to infinity. Under appropriate assumptions in the input data, we show that the variance of PAQ shrinks at a rate of $O(N^{-1} + n^{-4})$, showing that there are settings where the rate of $O(N^{-1}+n^{-1})$ of existing methods can be provably broken.