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

1. Problem Statement

Tree ensemble methods, such as Random Forests, Gradient-Boosted Decision Trees (GBDT), and XGBoost, are widely used for tabular data due to their strong predictive performance and interpretability. However, large ensembles incur high inference costs and memory footprints, complicating downstream tasks like robustness verification and fairness checks.

Post-hoc ensemble pruning is a standard solution, but existing methods face a trade-off between compression and fidelity (prediction equivalence):

• Accuracy-oriented pruning: Methods that optimize for accuracy-compression trade-offs often alter a subset of predictions. This lack of fidelity can disrupt decision consistency in high-stakes domains (e.g., medicine, finance) where downstream workflows rely on specific model outputs.
• Faithful pruning: Methods like FIPE (Emine et al., 2025) and Born-Again Tree Ensembles (Vidal & Schiffer, 2020) guarantee exact prediction equivalence for all inputs in the entire input space. While this ensures decision consistency, the requirement to preserve predictions even for rare or out-of-distribution (OOD) inputs severely limits the achievable compression ratio.

The core problem addressed is how to achieve higher compression rates while maintaining high fidelity, without the restrictive burden of guaranteeing equivalence on OOD inputs that are rarely observed in practice.

2. Methodology: PINE

The authors propose PINE (Pruning with In-distribution Equivalence), a pruning framework that guarantees prediction equivalence only within a calibrated in-distribution (ID) region, $X_{ID}(\alpha)$.

Core Concept

PINE relaxes the constraint of faithful pruning by restricting the equivalence guarantee to inputs that are "plausible" under the data distribution. It introduces a single hyperparameter, $\alpha$ (miscoverage level), which controls the size of this region. The method guarantees that for a future input sampled from the same distribution, the probability that the pruned model matches the original model's prediction is at least $1-\alpha$.

Algorithmic Pipeline

PINE builds upon the iterative Pruner-Oracle framework used by FIPE but modifies the Oracle's search space:

1. Plausible Score Definition: A score function $s(x)$ is defined such that lower scores indicate higher plausibility (likelihood) under the data distribution. The in-distribution region is defined as $X_{ID}(\alpha) = \{x \in \mathcal{X} \mid s(x) \leq \tau(\alpha)\}$.
2. Chow-Liu Tree Score: The paper adopts the negative log-likelihood of a Chow-Liu tree as the plausible score. Continuous features are discretized into bins, and the Chow-Liu tree approximates the joint distribution as a tree structure. This choice is critical because the negative log-likelihood decomposes into an additive form (root marginal + edge conditional terms), allowing the constraint $s(x) \leq \tau(\alpha)$ to be encoded efficiently as a Mixed-Integer Linear Programming (MILP) constraint.
3. Conformal Calibration: Using a calibration set $D_{cal}$, the threshold $\tau(\alpha)$ is determined via split conformal prediction. This ensures that, under the assumption of exchangeability, the region $X_{ID}(\alpha)$ covers future inputs with probability at least $1-\alpha$.
4. Iterative Pruning:
   • Pruner: Updates tree weights to satisfy prediction equivalence on a growing set of constraints.
   • Oracle: Searches for counterexamples (inputs where predictions differ) but restricts its search only to the region $X_{ID}(\alpha)$.
   • The process repeats until no counterexamples are found within the ID region.

Theoretical Guarantees

• Probabilistic Equivalence: Under the exchangeability assumption, PINE provides a lower bound of $1-\alpha$ on the probability that the pruned model's prediction matches the original model for a future input.
• Complexity: The number of discrete states the Oracle must explore is upper-bounded by $e^{\tau(\alpha)}$. Since $\tau(\alpha)$ is monotonic in $\alpha$, increasing $\alpha$ (allowing a larger miscoverage rate) exponentially reduces the search space, facilitating higher compression.

3. Key Contributions

1. PINE Framework: A novel pruning method that trades off strict global fidelity for high-probability in-distribution fidelity, enabling significantly higher compression ratios than faithful pruning baselines.
2. Conformal Integration: The first application of split conformal prediction to define a guaranteed in-distribution region for ensemble pruning, providing a rigorous probabilistic bound on prediction equivalence.
3. MILP-Embeddable Scoring: The proposal of using Chow-Liu tree negative log-likelihood as a plausible score, which is both expressive for capturing feature dependencies and computationally efficient for MILP encoding.
4. Systematic Trade-off Control: A mechanism to systematically tune the fidelity-compression trade-off via the single parameter $\alpha$.

4. Experimental Results

The authors evaluated PINE on 12 public tabular datasets (e.g., Adult, COMPAS, Pima-Diabetes) using XGBoost ensembles ($M=30$ trees, depth $D=2$).

• Compression vs. Fidelity: PINE significantly outperforms faithful pruning baselines (FIPE). As $\alpha$ increases from 0.05 to 0.8, the mean pruning rate increases from ~45% to ~68%, while mean fidelity remains high (99.15%–99.96%). In contrast, FIPE achieves a mean pruning rate of ~44% with 100% fidelity.
• Comparison to Accuracy-Oriented Methods: While accuracy-oriented methods (IC, DREP, MDEP) can maintain test accuracy, they exhibit substantial fidelity drops (often <90%) even within the in-distribution region. PINE maintains near-perfect fidelity within the calibrated region.
• Coverage Validation: Empirical test coverage $\hat{\pi}_{ID}$ closely tracks the target $1-\alpha$, validating the effectiveness of the split conformal calibration.
• Scalability: The method is sensitive to tree depth ($D$) more than the number of trees ($M$). Deeper trees reduce attainable compression due to more localized decision regions, while larger ensembles primarily increase optimization runtime.

5. Significance and Claims

The paper claims that PINE addresses a critical gap in model compression: the inability of existing faithful methods to compress models effectively without sacrificing decision consistency, and the inability of standard pruning methods to guarantee consistency.

• Practical Impact: PINE allows practitioners to compress models by up to 30% more than faithful baselines while preserving decision consistency on likely future inputs. This is particularly valuable for high-stakes domains where model outputs trigger downstream workflows or require auditability.
• Limitations: The authors modestly acknowledge that the guarantee relies on the exchangeability assumption (no distribution shift). If the data distribution shifts, the guarantee may weaken. Additionally, the computational cost is tied to solving MILPs, which can be expensive for multi-class settings or very deep trees.
• Conclusion: PINE offers a verifiable, probabilistic approach to pruning that balances the competing goals of efficiency and reliability, moving beyond the binary choice between "accurate but inconsistent" and "consistent but uncompressed."