PINE: Pruning Boosted Tree Ensembles with Conformal In-Distribution Prediction Equivalence

Abstract: Tree ensembles are machine learning models with strong predictive performance and interpretability, and remain widely used for tabular data. Standard pruning methods for tree ensembles typically optimize an accuracy-compression trade-off and may change a subset of predictions, potentially compromising decision consistency. Faithful pruning methods address this issue by preserving prediction equivalence over the entire input space, but this requirement leads to lower compression ratios. We propose PINE, a pruning method that provides strong guarantees within an in-distribution region. PINE preserves prediction equivalence within this region and controls the region size using a single parameter $α$ via conformal calibration. Experiments on 12 public tabular datasets show that PINE improves the compression ratio by up to 30% while preserving predictions at a comparable level to existing faithful pruning methods.

• The paper introduces PINE, a method that prunes boosted tree ensembles by enforcing prediction equivalence strictly within a conformally calibrated in-distribution region.
• It leverages MILP with Chow-Liu tree scoring to efficiently restrict counterexample search, achieving up to 67.8% pruning while preserving near-100% fidelity.
• The approach offers flexible tradeoffs between compression and decision consistency, making it valuable for safety-critical and regulated deployment contexts.

Pruning Tree Ensembles with In-Distribution Prediction Equivalence: The PINE Framework

Introduction and Motivation

Ensemble methods based on decision trees, including Random Forests and GBDTs (such as XGBoost), remain dominant in tabular data modeling due to their competitive predictive accuracy and interpretability. However, large ensembles incur substantial inference and memory costs, and complicate post-hoc analysis, robustness verification, and regulatory audits. Post-training pruning of ensembles is thus a central topic, with methods typically trading off model size versus predictive fidelity.

Standard pruning strategies optimize accuracy-compression, accepting some loss of fidelity to the original model's predictions. Faithful (or functionally-identical) pruning algorithms enforce prediction equivalence across the entire input domain, ensuring no downstream decision shifts, but at the cost of lower achievable compression, especially in high-dimensional regimes and with rare or out-of-distribution (OOD) feature configurations.

"PINE: Pruning Boosted Tree Ensembles with Conformal In-Distribution Prediction Equivalence" (2605.28068) proposes a middle ground: guarantee prediction equivalence—i.e., decision consistency—strictly within a calibrated in-distribution region, thus supporting higher model compression without risk of changed system behavior on plausible future data. The main innovation is the conformal-calibrated, data-driven definition of this in-distribution region, together with strong theoretical and empirical analyses of the resulting pruning-fidelity tradeoff.

Methodology: PINE Framework

PINE (Pruning with In-Distribution Equivalence) follows the faithful pruning template, but restricts the enforced equivalence region $X(\alpha)$ to cover likely future samples, with miscoverage at most $\alpha$. The approach comprises the following components:

• Region Definition: $X(\alpha)$ is $\{\bm{x} : s(\bm{x}) \leq \tau(\alpha)\}$, where $s$ is a plausible score (negative Chow-Liu log-likelihood), and $\tau(\alpha)$ is calibrated to yield $1-\alpha$ in-distribution covering probability via split conformal prediction.
• Constraint Enforcement: As in previous functionally identical pruning (FIPE), counterexamples are sought and rules enforced, but only in $X(\alpha)$. The resulting pruned ensemble has identical predictions to the original model within $X(\alpha)$.
• Plausible Score Construction: The Chow-Liu tree is used to approximate the training-data joint feature distribution, with its NLL as a sample's outlier/plausibility score. This admits efficient MILP encoding for constraint search.

  Figure 1: Overview of PINE—calibration of the in-distribution region by split conformal prediction, and counterexample search restricted to the certified region $X(\alpha)$.

This design enables a tunable tradeoff: smaller $\alpha$0 yields a stricter guarantee and smaller feasible compression; larger $\alpha$1 expands feasible compression while allowing increased but controlled miscoverage.

Theoretical Guarantees and Analysis

• Probabilistic Equivalence: For any $\alpha$2, PINE certifies that for an i.i.d. future input, the pruned ensemble is identical to the original on $\alpha$3, with probability at least $\alpha$4 under exchangeability assumptions.
• Solving Complexity: The Chow-Liu-based MILP restricts counterexample search to at most $\alpha$5 discrete input states, a strong exponential reduction compared to the full space $\alpha$6, with $\alpha$7 feature bins and $\alpha$8 features.
• Compression Bound: Since $\alpha$9, constraint relaxation can only improve achievable tree sparsity compared to whole-domain equivalence.
• Calibration Validity: The split conformal construction guarantees in-distribution region coverage holds distribution-free, even under arbitrary score functions.

Empirical Results

Comprehensive experiments on 12 public tabular datasets showcase the fidelity-compression advantages of PINE relative to both standard (accuracy-oriented) and functionally-identical pruning baselines.

• Test Accuracy and Fidelity: At identical or higher compression rates, PINE maintains fidelity near 100% on in-distribution data, unlike accuracy-oriented methods where fidelity may degrade sharply even as test accuracy appears stable (Figure 2).

Figure 2: Test accuracy versus pruning rate for PINE and baselines, with PINE preserving high accuracy and facilitating larger compression budgets.

• Control Over Coverage: Empirical test coverage $X(\alpha)$0 tracks $X(\alpha)$1 closely across datasets and seeds, validating the practical effectiveness of the conformal calibration (Figure 3).

  Figure 3: Empirical test coverage $X(\alpha)$2 as a function of miscoverage $X(\alpha)$3, confirming precise region calibration.

• Conditional Fidelity: Even within the certified $X(\alpha)$4 region, non-equivalence-guaranteed baseline methods continue to exhibit prediction changes, while PINE achieves perfect (or near-perfect) in-distribution fidelity (Figure 4).

Figure 4: In-distribution conditional fidelity $X(\alpha)$5—PINE uniquely delivers decision consistency in the calibrated region.

• Compression and Fidelity Gains: Across all datasets, PINE increases mean pruning rate from 44.6% (FIPE) up to 67.8% (for $X(\alpha)$6), while mean overall fidelity remains in the 99.1%–99.9% interval.
• Flexibility: PINE's $X(\alpha)$7 parameter enables practitioners to directly set target fidelity or control admitted compression dynamically, with provided analytical and empirical recipes for practical selection.

Practical and Theoretical Implications

The PINE framework offers a new axis of control in model compression for deployed tabular classifiers: practitioners no longer face the dichotomy between maximizing efficiency and retaining global functional compliance. Instead, the explicit calibration of guarantee regions, with formal error rates, enables rational quantification and communication of model update risk.

PINE is especially suited for safety- and regulation-critical pipelines, where decision shifts within plausible region of operation are unacceptable, but OOD prediction equivalence is not. High-stakes deployments in finance, healthcare, or legal domains with model-centric audits or recourse computations benefit from both the compression and certified decision consistency PINE offers.

Theoretically, the approach unifies conformal inference with model pruning—a notable synthesis—bringing statistical coverage guarantees to the domain of model reduction and interpretability.

Limitations and Future Directions

While PINE delivers strong guarantees under exchangeability and i.i.d. assumptions, its calibration and OOD demarcation rely on the training-data distribution. Under significant covariate shift, miscoverage may be underestimated. Robust extensions involving distribution shift detection or online re-calibration, possibly via adaptive or conditional conformal methods, constitute an important research direction.

Moreover, the necessity of repeated MILP solves, while mitigated by the Chow-Liu constraints, still introduces computational challenges for extremely large ensembles or deep trees. Algorithmic advances in counterexample search, structure-exploiting relaxations, and incremental pruning may enhance scalability further.

Conclusion

PINE establishes a principled, practical approach for ensemble pruning that combines verifiable, statistical prediction equivalence guarantees with efficient model compression. By calibrating the guarantee region through conformal inference and embedding plausible score constraints within the MILP-driven pruning loop, PINE enables both high test accuracy and fidelity at improved compression levels relative to previous faithful pruning approaches. Its theoretical guarantees, strong empirical results, and flexible control over the fidelity-compression tradeoff position it as a robust tool for high-stakes model deployment and maintenance (2605.28068).