Centralize row subsampling in GBTree
- Dominant language
- C++
- Stars
- 28.8k
- Forks
- 8.9k
- Avg merge
- 1d 12h
- Merged PRs (30d)
- 54
Description
## Summary
Row subsampling is currently implemented separately by multiple tree updaters:
- `grow_colmaker` marks sampled-out row positions.
- CPU `hist` and `approx` zero gradient pairs using the shared CPU sampler.
- CUDA histogram uses a separate device sampler.
- SYCL histogram builds a compact list of sampled row indices.
This duplicates sampling policy across backends and ties gradient-based sampling to histogram updaters.
We should move row subsampling into `GBTree` and represent sampled-out rows using zero gradient/Hessian pairs. Tree updaters would consume already-sampled gradients and no longer own sampling policy or random selection.
## Proposed design
For each tree constructed by `GBTree`:
1. Copy or prepare the original gradient container.
2. Generate a new row sample.
3. Apply the sample to the gradient pairs:
- Uniform sampling: zero rejected rows.
- Gradient-based sampling: zero rejected rows and reweight retained rows using inverse sampling probabilities.
4. Run the configured updater chain for exactly one tree.
5. Return final leaf positions for every row, including rows excluded from tree statistics.
Conceptually:
```text
for each output group:
for each parallel tree:
sampled_gradients = Sample(original_gradients)
for each configured updater:
updater.Update(sampled_gradients, one_tree)
obtain leaf positions for all rows
```
Moving the forest loop into `GBTree` preserves independent sampling for every tree constructed with `num_parallel_tree`.
## Gradient containers
Sampling should operate on the complete `GradientContainer`, including objectives that provide separate split and value gradients.
For these objectives:
- Select rows using the split gradients.
- Apply the same row-selection mask to the value gradients.
- Apply the appropriate gradient-based sampling weights to both representations.
The implementation should support host and device gradient storage without requiring unnecessary device-to-host transfers.
## Row positions and prediction caching
Whether a row contributes to tree statistics should be separate from whether it has a leaf position.
Sampled-out rows should:
- Contribute zero gradient and Hessian statistics.
- Still be routed through the completed tree.
- Receive a valid final leaf position.
- Remain usable by prediction-cache and adaptive-leaf functionality.
This could eliminate the current convention where complemented node positions encode sampled-out rows.
An explicit temporary sampling mask should remain available where necessary. Using only `hess == 0` as the semantic mask may be ambiguous for objectives that legitimately produce zero Hessians, although all-zero gradient rows can still be skipped as an optimization.
## Updater behavior
After this change, tree updaters would no longer perform random row sampling.
They may still optimize already-sampled inputs by:
- Skipping all-zero gradient rows when constructing histograms.
- Compacting active row indices.
- Avoiding partition work for inactive rows where final positions can be recovered afterward.
These should be backend implementation details rather than separate definitions of subsampling behavior.
Gradient-based sampling would consequently become available to all compatible tree methods rather than being restricted to histogram updaters.
## Benefits
- One consistent definition of row subsampling.
- Independent sampling remains available for parallel trees.
- Gradient-based sampling can support additional tree methods.
- Less duplicated CPU, CUDA, SYCL, and legacy updater logic.
- Clear separation between sampling policy and tree construction.
- Simpler semantics for row positions and prediction caching.
- Backend-specific compaction and zero-row skipping remain possible.
## Implementation outline
- [ ] Change the tree-building contract so an updater invocation builds one tree.
- [ ] Move the `num_parallel_tree` loop into `GBTree`.
- [ ] Introduce a booster-level sampler for `GradientContainer`.
- [ ] Implement uniform sampling by zeroing rejected gradient pairs.
- [ ] Move gradient-based selection and inverse-probability weighting into the centralized sampler.
- [ ] Handle separate split and value gradients.
- [ ] Make updaters return leaf positions for every row.
- [ ] Update prediction-cache and adaptive-leaf handling.
- [ ] Remove sampling from `colmaker`, CPU histogram, CUDA histogram, and SYCL histogram.
- [ ] Remove negative-Hessian/complemented-position sampling conventions where no longer needed.
- [ ] Add optional zero-row skipping or compaction optimizations to histogram backends.
- [ ] Add cross-updater tests for reproducibility, parallel trees, multi-output objectives, prediction caching, and distributed training.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by tracing the GBTree tree-building contract and the grow_colmaker, CPU hist/approx, CUDA histogram, and SYCL histogram updaters named in the issue. Define the centralized GradientContainer sampling flow, then review prediction-cache, adaptive-leaf, and distributed-training handling. Done means one-tree updater calls consume sampled gradients, all backends share sampling policy, and the listed cross-updater tests cover reproducibility, parallel trees, multi-output objectives, caching, and distribution.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, machine-learning
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100