Research map / Probabilistic and scientific modeling / Molecules and proteins

Prediction-Augmented Trees for Reliable Statistical Inference

NeurIPS 2026 · Probabilistic and scientific modeling › Molecules and proteins · approach: Statistical analysis of ML practice · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes PART, a CART-style variance-greedy tree for prediction-powered mean and regression-coefficient inference, and PAQ, a quadrature-based limit with O(1/N + n^-4) variance under smooth deterministic residuals.

Problem

Prediction-powered inference: constructing tighter valid confidence intervals for a mean or linear regression coefficient using a small labeled set, a large unlabeled set, and a pre-trained ML predictor.

Approach

Adaptively partition the feature space with a greedy variance-based (VMS) regression tree, apply leaf-wise residual corrections weighted by unlabeled leaf masses (reducing to PPI at depth zero), and in the deterministic smooth limit replace the residual correction with nearest-neighbor quadrature (PAQ) achieving faster variance decay.

Key results

• PART consistently yields tighter confidence intervals than PPI++ across the real-world datasets while maintaining nominal coverage.
• PAQ variance is shown to scale as O(1/N + 1/n^4) versus O(1/N + 1/n) for PPI/PPI++ under smoothness assumptions on the residual.
• Theoretical coverage bound for PART is conservative; empirical coverage is close to the nominal 95% level.

Surprising findings

• The finite-sample coverage correction term grows with number of leaves, so deeper trees can overfit and reduce true coverage in practice.

Limitations

• PAQ's fast rate requires deterministic, smooth (C^2 or higher) residuals; not established for noisy labels
• Theoretical coverage guarantees for PART rely on fixed leaf count and asymptotic arguments; finite-sample correction term is loose
• Experiments use manually tuned depth limits and a pre-trained predictor assumed independent of the labeled data

Experimental setup

Datasets: Amazon deforestation satellite dataset (1596 observations) Galaxy Zoo 2 (16,743 observations) Vinho Verde wine quality (1280 samples) California census housing (16,512 instances) AlphaFold phosphorylation/IDR data (10,802 samples) California census income (380,091 observations) synthetic data

Benchmarks: Amazon deforestation Galaxy Zoo 2 spiral fraction Vinho Verde wine quality California census median house price AlphaFold odds ratio California census income regression synthetic PAQ experiment

Metrics: confidence interval width coverage probability

Baselines: PPI PPI++ empirical mean

Builds on: Prediction-powered inference (PPI) PPI++ CART Wald-type confidence intervals

Closest papers by meaning

• Task-Agnostic Machine-Learning-Assisted Inference (NeurIPS 2024 · similarity 0.87)
• Position: Data-driven Discovery with Large Generative Models (ICML 2024 · similarity 0.85)
• Calibration without labels in multiple testing (NeurIPS 2026 · similarity 0.85)
• TreeVI: Reparameterizable Tree-structured Variational Inference for Instance-level Correlation Capturing (NeurIPS 2024 · similarity 0.84)
• Common Task Framework For a Critical Evaluation of Scientific Machine Learning Algorithms (NeurIPS 2025 · similarity 0.84)
• Regression for the Mean: Auto-Evaluation and Inference with Few Labels through Post-hoc Regression (ICML 2025 · similarity 0.84)