Sample Splitting for Data-Selected Partitions

Why Separate Partition Selection from Leaf Estimation?

A tree partition is often chosen by searching over candidate splits using observed data. If the same observations are then used to estimate the predictions within the selected leaves, those predictions and the chosen partition are statistically entangled. Sample splitting separates the two tasks: one subsample selects the partition, and an independent subsample estimates its leaf predictions. This permits a risk statement that is valid conditional on the selected partition, even though the partition itself was data-selected.

The result is not a guarantee that the selected tree is optimal. It describes the estimation error for whichever partition the first subsample produces, and adds that error to the partition’s approximation error. The distinction matters: a stable estimate within a poor partition can still have substantial total risk.

Definition: Let $D_1$ be a discovery sample and $D_2$ an independent estimation sample, each consisting of independent observations from the same distribution of $(X,Y)$. A sample-split partition estimator selects a partition $\Pi$ using only $D_1$, then assigns a prediction to each cell of $\Pi$ using only observations in $D_2$. Conditional on $D_1$, the selected partition is fixed and $D_2$ remains an independent sample from the data distribution.

Risk for a Fixed Partition

We study squared-error regression. Assume $E[Y^2]<\infty$, and define the regression function $m(x)=E[Y\mid X=x]$. For any predictor $g$, its risk is $R(g)=E[(Y-g(X))^2]$, where the expectation is over a fresh observation. The Bayes predictor under squared-error loss is $m$; thus $R(g)-R(m)$ measures excess risk.

Suppose the selected partition has $K$ cells $A_1,\ldots,A_K$. For a cell with positive probability $p_j=P(X\in A_j)$, define its best constant prediction by $\mu_j=E[Y\mid X\in A_j]$. The partition’s piecewise-constant oracle predictor is $\mu_\Pi(x)=\mu_j$ when $x\in A_j$. Its approximation error is the excess risk it would have even if these cell means were known exactly.

Definition: The approximation error of a partition $\Pi$ is $A_\Pi=E[(m(X)-\mu_\Pi(X))^2]$. For estimation observations $(X_i,Y_i)$ in $D_2$, let $N_j$ be the number with $X_i\in A_j$. If $N_j>0$, the fitted leaf prediction is $\hat\mu_j=N_j^{-1}\sum_{i:X_i\in A_j}Y_i$. If $N_j=0$, set $\hat\mu_j=0$; any finite convention for an empty cell gives the same integrated risk contribution when $p_j=0$, but a positive-probability empty cell must be handled in the risk analysis.

Result: For any piecewise-constant predictor with leaf values $a_j$, its excess risk decomposes as $R(g)-R(m)=A_\Pi+\sum_{j:p_j>0}p_j(a_j-\mu_j)^2$.

To verify the decomposition, first condition on $X$. Since $E[Y-m(X)\mid X]=0$, expanding the square gives $R(g)-R(m)=E[(g(X)-m(X))^2]$. Within cell $A_j$, write $g(X)-m(X)=(a_j-\mu_j)+(\mu_j-m(X))$. The cross term has expectation zero within that cell because $E[m(X)-\mu_j\mid X\in A_j]=0$. Summing the resulting squared terms over cells gives the stated identity. This separates a property of the partition from the error in estimating its leaf values.

Worked Example: Approximation Error in Two Cells

Suppose a selected partition has two cells with probabilities $0.4$ and $0.6$. In the first cell, $m(X)$ takes values $0$ and $0.4$ with equal conditional probabilities; in the second, it takes values $0.5$ and $1$ with equal conditional probabilities. The cell means are $\mu_1=0.2$ and $\mu_2=0.75$.

In the first cell, the conditional mean squared approximation error is $(0-0.2)^2/2+(0.4-0.2)^2/2=0.04$. In the second cell it is $(0.5-0.75)^2/2+(1-0.75)^2/2=0.0625$. Therefore,

A_\Pi=0.4(0.04)+0.6(0.0625)=0.016+0.0375=0.0535.

Even with perfectly estimated cell means, the partition’s excess risk is $0.0535$. Any estimation error from the second subsample adds to this quantity; it cannot reduce the approximation error in the decomposition.

A Conditional Bound for the Leaf Estimates

The next result makes the role of independent estimation data precise. Assume $|Y|\le M$ almost surely for a known $M>0$. Conditional on $D_1$, suppose the selected partition has $K$ positive-probability cells. Write $p_j=P(X\in A_j)>0$ for those cells. The probabilities may vary from one selected partition to another, but conditional on $D_1$ they are fixed.

Theorem: Fix $0<\delta<1$. Conditional on $D_1$, with probability at least $1-\delta-\sum_{j=1}^K\exp(-n_2p_j/8)$ over the estimation sample, all cells have $N_j\ge n_2p_j/2$ and $$ R(\hat f)-R(m) \le A_\Pi+\frac{4M^2K\log(2K/\delta)}{n_2}, $$ where $\hat f(x)=\hat\mu_j$ for $x\in A_j$. Here the probability bound is informative when the displayed lower bound is positive. The result is conditional on the discovery sample: it holds for the partition it selected, not just for a partition fixed in advance.

Proof. Conditional on $D_1$, the partition and the cell probabilities are fixed, while the observations in $D_2$ are independent draws from the original distribution. Thus $N_j$ has a binomial distribution with parameters $n_2,p_j$. The binomial lower-tail bound gives $P(N_j<n_2p_j/2\mid D_1)\le \exp(-n_2p_j/8)$. A union bound shows that the probability any cell violates its count bound is at most the sum of these terms.

For a cell with $N_j=k>0$, conditional on which estimation observations fall in that cell, their responses are independent draws from the distribution of $Y$ given $X\in A_j$, with mean $\mu_j$ and range in $[-M,M]$. Hoeffding’s inequality for a mean of $k$ independent variables with range length $2M$ gives $P(|\hat\mu_j-\mu_j|>t\mid N_j=k,D_1)\le 2\exp(-kt^2/(2M^2))$. Set $L=\log(2K/\delta)$ and $t=M\sqrt{2L/k}$. The probability of exceeding this cell-specific threshold is at most $2e^{-L}=\delta/K$. A union bound over the $K$ cells shows that, apart from probability at most $\delta$, all these mean-error bounds hold. This remains true when conditioning on the counts being large: the concentration bound applies conditional on the cell memberships, and hence conditional on their counts.

On the event that every $N_j\ge n_2p_j/2$, the mean-error bound implies, for each cell, $|\hat\mu_j-\mu_j|\le M\sqrt{2L/N_j}\le 2M\sqrt{L/(n_2p_j)}$. Consequently,

\sum_{j=1}^Kp_j(\hat\mu_j-\mu_j)^2 \le \sum_{j=1}^Kp_j\frac{4M^2L}{n_2p_j} =\frac{4M^2KL}{n_2}.

Apply the risk decomposition to the fitted leaf values and add the two failure-probability bounds. This proves the theorem. $\square$

Worked Example: Evaluating the Bound for a Selected Partition

Suppose the discovery sample selects the two-cell partition in the first example. Let $n_2=200$, $M=1$, and $\delta=0.05$. The cell probabilities are $0.4$ and $0.6$, so the probability excluded for a small cell count is bounded by $e^{-200(0.4)/8}+e^{-200(0.6)/8}=e^{-10}+e^{-15}$, which is approximately $0.0000454$. The theorem’s total failure bound is therefore approximately $0.0500454$.

Here $K=2$ and $\log(2K/\delta)=\log(80)\approx4.382$. The estimation term is $4(1)^2(2)(4.382)/200\approx0.1753$. Since the approximation error is $0.0535$, the theorem gives the conditional bound $R(\hat f)-R(m)\le0.0535+0.1753=0.2288$ with probability at least about $0.94995$, conditional on the discovery sample. The bound can be conservative, but it explicitly separates approximation from estimation error.

Worked Example: Why the Estimation Sample Must Be Independent

Let $X$ take values $0$ and $1$ with probability $1/2$ each. Suppose $Y\mid X=0$ is equally likely to be $0$ or $2$, and $Y\mid X=1$ is equally likely to be $1$ or $3$. The regression means are $m(0)=1$ and $m(1)=2$. A discovery procedure selects the partition into the two singleton cells.

Now consider an independent estimation sample that happens to contain two observations with $X=0,Y=0$ and two with $X=1,Y=3$. The fitted leaf means are $0$ and $3$, so the integrated estimation error relative to the cell means is $(1/2)(0-1)^2+(1/2)(3-2)^2=1$. The partition approximation error is zero because each cell has a constant regression mean. The excess risk for this fitted tree is therefore $1$.

This realized error does not invalidate the conditional guarantee: the estimation sample is random, and the bound controls its error with a stated probability. Independence is what permits that probability statement to be made after conditioning on the partition chosen from the discovery data. If the same outcomes had also been used to choose the partition, conditioning on the selected partition could change the distribution of those outcomes; the theorem’s conditional argument would no longer apply as stated.

Why Cell Probabilities Matter

The bound’s estimation term scales with the number of positive-probability cells and inversely with the estimation-sample size. Its occupancy condition also depends on each $n_2p_j$. A large leaf count can make estimation noisy, while a rare cell may receive too few observations to estimate its mean reliably. Sample splitting resolves selection dependence; it does not guarantee that every selected cell has adequate data.

Worked Example: A Rare Cell Can Be Empty

Suppose a selected partition has a cell with probability $p=0.01$, and $n_2=100$ observations are used for estimation. The probability that this cell receives no observations is $(1-0.01)^{100}=0.99^{100}\approx0.366$. Thus it is not reasonable to expect this cell to be populated with high probability.

For a concrete risk consequence, let $Y=1$ in this cell and $Y=0$ outside it, and set the empty-cell prediction to $0$. The partition’s approximation error is zero, since the regression function is constant on each cell. When the rare cell is empty, its contribution to excess risk is $0.01(1-0)^2=0.01$. That event occurs with probability about $0.366$. This illustrates why a risk statement must account for cell probabilities, rather than claiming a useful $K/n_2$ bound regardless of occupancy.

Common Pitfall

Wrong reasoning: “The partition is data-selected, so a fixed-partition risk bound cannot apply.” Correction: Condition on the discovery sample. The selected partition is then fixed, and independence leaves the estimation sample distributed as a fresh sample. The resulting guarantee is conditional on the selected partition.
Wrong reasoning: “Independence ensures every selected leaf has enough observations.” Correction: Independence ensures the estimation sample is not selected using the discovery outcomes; it does not prevent rare leaves from being empty. Cell probabilities and sample size still govern occupancy.
Wrong reasoning: “The theorem proves that the selected tree is good.” Correction: It bounds estimation error around the selected partition’s cell means. Its approximation error remains, and the result does not compare that partition with alternative trees.

Key takeaway: Sample splitting makes leaf-mean estimation conditionally analyzable after a data-selected partition is fixed. The risk separates into partition approximation error and estimation error, while adequate cell probabilities are still necessary for a useful finite-sample guarantee.

Check Your Understanding

Try these questions before moving on.

Why does conditioning on the discovery sample leave the estimation sample independent of the selected partition?
In the squared-error decomposition, which term describes the partition itself, and which term describes estimation of its leaf values?
For $K$ positive-probability cells, how does the theorem’s estimation term depend on $K$, $M$, and $n_2$?
Why does a rare cell make the count condition difficult to satisfy, even when the estimation sample is independent?
Does the conditional bound establish that the selected partition is optimal among tree partitions? Explain.