Prediction-Augmented Quadrature (PAQ)

• Prediction-Augmented Quadrature (PAQ) is a learning-augmented estimator that merges ML predictions with quadrature for debiased inference in scientific analysis.
• It leverages the smoothness of prediction residuals to apply Taylor-based corrections, achieving faster error decay and tighter confidence intervals.
• PAQ is applied in fields like protein structure analysis, deforestation monitoring, and astronomy, demonstrating improved estimation accuracy under regular conditions.

Prediction-Augmented Quadrature (PAQ) is a learning-augmented estimator designed to combine the strengths of statistical prediction with numerical quadrature for reliable inference in scientific data analysis. PAQ emerges naturally as the infinite-depth limit of the Prediction-Augmented Residual Tree (PART) estimator, offering substantial improvements in variance and bias under regularity and smoothness assumptions on the residuals resulting from machine learning prediction corrections. The approach generalizes classical sample correction strategies by leveraging the structure of the residual function, leading to accelerated error decay rates and tighter confidence intervals in estimation tasks across diverse scientific applications.

1. Mathematical Foundation of PAQ

In its canonical form, PAQ exploits the smoothness of the residual function $r(x) = y(x) - f(x)$, where $f(x)$ is an ML prediction and $y(x)$ is the true label, to construct a debiased estimator for the population mean. With access to $N$ unlabeled samples $\tilde{x}_i$ and $n$ labeled samples $x_j$, the estimator is defined as:

$$\mu_{\text{PAQ}} = \frac{1}{N} \sum_{i=1}^{N} f(\tilde{x}_i) + r(h(\tilde{x}_i))$$

where $h(\tilde{x}_i)$ denotes the nearest neighbor of $\tilde{x}_i$ in the labeled set, and $r(h(\tilde{x}_i))$ is the corresponding residual.

In the regime where the residual $r$ is sufficiently smooth, PAQ aligns with classical quadrature—from a numerical analysis perspective, the empirical correction term approximates the integral $\int r(u) du$ over the feature space by the trapezoidal rule:

$$\int_0^1 r(u) du \approx x_1\, r(x_1) + \sum_{i=1}^{n-1} \frac{x_{i+1}-x_i}{2}\left[ r(x_i) + r(x_{i+1}) \right] + (1-x_n) r(x_n)$$

Taylor expansions show that with residual regularity $\sup_x |r''(x)| \leq L_2$, the approximation error decays as $O(1/n^2)$, while variance decays as $O(1/N + 1/n^4)$. Higher-degree polynomial interpolation yields further improvements:

• Bias: $O(1/n^{p+1})$
• Variance: $O(1/N + 1/n^{2p+2})$

2. Assumptions and Error Rates

PAQ’s theoretical guarantees are contingent on two primary conditions:

Assumption	Mathematical Statement	Implication
Residual Smoothness	$r \in C^2$ (trapezoidal), $C^{p+1}$ (degree-$p$)	Enables error rate $1/n^2$ or better
Data Distribution	$P_X = \text{Uniform}[0, 1]$ (basic case)	Quadrature nodes evenly spread

The smoothness of $r$ ensures that debiasing benefits from local accuracy, as PAQ's correction approaches the true integral for regular functions. For unknown marginals, a probability integral transform can be employed so that labeled nodes act as effective quadrature nodes.

If the residual is non-smooth or highly stochastic, PAQ's variance and bias benefits diminish, and it reverts to performance similar to classical estimators such as Prediction-Powered Inference (PPI).

3. Relationship to PART and Tree-based Estimators

PART adaptively partitions the feature space, assigning residual correction by averaging labeled sample residuals within each leaf:

• PART with zero splits: Reduces to PPI estimator
• Increasing tree depth: Localizes correction; residual estimation approaches true $r(x)$
• Infinite depth limit: Correction transitions to numerical integration (quadrature), yielding the PAQ estimator.

This formal equivalence demonstrates that PAQ can be viewed as the “atomic leaf” limit of tree-based estimators, capitalizing on residual smoothness to attain quadrature-level errors.

4. Comparison to Prior Learning-Augmented Estimators

PAQ notably improves over learning-augmented estimators that do not exploit structure in the residual function. For instance, the standard PPI estimator’s variance decays like $O(1/N + 1/n)$:

$$\text{Var}[\mu_{\text{PPI}}] \approx \frac{\text{Var}(Y - f(X))}{n}$$

In contrast,

$\text{Var}[\mu_{\text{PAQ}}] = O(1/N + 1/n^4)$

under regularity, and faster for higher-degree quadrature corrections. These rates are provably tighter (Kher et al., 19 Oct 2025).

5. Practical Applications and Empirical Results

PAQ and PART have broad applicability, especially in domains where ML prediction is very accurate and residuals are smooth:

• Protein structure analysis: When predictions (e.g., AlphaFold) are highly deterministic, PAQ yields precise estimation of odds ratios and other population attributes.
• Deforestation monitoring: Residuals from forest cover prediction models exhibit geographical smoothness; PAQ produces tighter confidence intervals for regional means.
• Astronomy (Galaxy Zoo): High accuracy in galaxy morphology prediction makes PAQ effective for estimating mean spiral fractions.
• Econometrics and environmental datasets: When ML predictions capture most variability, PAQ delivers improved estimates for means across wine, housing, or census data.

These empirical benefits stem from PAQ's ability to integrate prediction and correction in a way that adapts to the smoothness of the underlying signal.

6. Analytical Perspective and Extensions

The PAQ methodology unifies ideas from quadrature theory and modern statistical correction. Its analytical error bounds are derived via Taylor expansions and quadrature error decomposition, leading to the following generalization for degree-$p$ interpolation:

• If $r \in C^{p+1}$ and labeled points partition $[0,1]$ evenly,

  $$\left| E[\mu_{\text{PAQ}}^{(p)}] - E[Y] \right| = O\left(\frac{1}{n^{p+1}}\right)$$

  $$\operatorname{Var}(\mu_{\text{PAQ}}^{(p)}) = O\left( \frac{1}{N} + \frac{1}{n^{2p+2}} \right)$$

Extensions to higher dimensions are possible using tensor product quadrature rules, but curse-of-dimensionality effects can attenuate error improvement rates.

7. Limitations and Implications

PAQ is justified primarily under assumptions of residual smoothness and regularity. In cases where residuals are rough or noisy, PAQ’s advantages are diminished, reverting the variance rate to that of existing methods. A plausible implication is that performance improvements will be most pronounced in inference tasks where the ML predictor captures nearly all input-output variability, leaving only smooth corrections.

PAQ’s approach is closely aligned with developments in learning-augmented statistical inference, and its main contribution is to show that under appropriate regularity, quadrature-inspired statistical estimators yield demonstrably superior sample efficiency for tasks commonly encountered in scientific analysis (Kher et al., 19 Oct 2025).