Yggdrasil: An Optimized System for Training Deep Decision Trees at Scale
| Firas Abuzaid | Joseph Bradley |
| Feynman Liang | Andrew Feng |
| Lee Yang | Matei Zaharia |
| Ameet Talwalkar | |
Motivation
- Decision trees work well for lots of ML problems
- Easy to debug; easy to tune, easy to understand
- As data grows in size $(n)$ and dimensionality $(p)$, two issues arise:
- Training sets can’t fit on a single machine—more machines required for training
- Deeper trees (e.g., $D \geq$ 10) are needed; greater depth $\rightarrow$ higher model accuracy
How do we train deep distributed decision trees?
Planet: The Classic Approach
- Partition training set by row
- Each worker computes sufficient statistics on subset of training data
- Since data partitioned by row, workers must compute statistics for all features
The Problem with Planet
- Communication cost between workers scales poorly as $p$ and $D$ increase
- Meng et al., 2016—also at this year's NIPS!
- To limit communication, workers only consider $B$ thresholds
- Instead of $n-1$, like in the serial algorithm
- $B \ll n$
- Optimal split may not always be found!
Excerpt from page 5 of the Planet
paper
The Problem with Planet
Three problems:
- Another hyperparameter to tune: $B$
-
Important tradeoff:
-
finding optimal split
↕
runtime efficiency - (Meng et al. improve the tradeoff—but it's still there)
-
finding optimal split
- Large $p$ and $D$ $\rightarrow$ poor runtime (even if you choose optimal $B$!)
Yggdrasil: The New Approach
- Partition training set by column
- Each worker computes sufficient statistics on all local features
- Master performs no aggregation; only responsible for picking globally optimal
- No $B$ parameter; no tradeoff for finding optimal split
- Equivalent to serial algorithm on a single machine
Yggdrasil in Action
- Partition features across workers
- Workers sort each feature by value
- Compute best split for each feature
- Pick best feature for each worker, send feature's split to master
- Master selects best global split among the candidates
Tale of the Tape: Planet vs. Yggdrasil
Tale of the Tape: Planet vs. Yggdrasil
🔑 Yggdrasil is more efficient
than Planet for large $p$ and
$D$
But Column-Partitioning has been done before!
- Caragea et al., 2004; Ye et al., 2009; Svore and Burges, 2011
- If people have thought of this before...why is it not done in practice?
Ain't as easy as it looks
Excerpt from page 2 of Meng et al., 2016
Yggdrasil in Action
- Partition features across workers
- Workers sort each feature by value
- Compute best split for each feature
- Pick best feature for each worker, send feature's split + bit vector to master
- Master selects best global split among the candidates
- Master sends bit vector for best global split to each worker
- Worker sorts each feature by bit vector, then value
Itty-bitty Details
-
Bitvectors encode the partitioning of instances for each split
- 1 = left child, 0 = right child
- Lookups performed using label indices
- $O(Dn)$ term now isn't quite as scary
-
Training is done per-depth, not per-node
- Single bitvector for all active leaves in the tree
- All workers send split + bitvector, even if the split is not used
- Downside: Wasted bits
- Upside: Only one roundtrip needed for one iteration of training
Why sort?
- Must keep track of the split history for each feature
- New column per split $\rightarrow$ exponential memory footprint
- If we sort...
- Sorting based on bit vector is $O(n)$, since the data is already sorted
- Memory overhead is constant
- Computing splits across all leaves still only requires a single scan over all the features
- What about leaves that don't need to be split? Just skip the sub-array!
Additional Optimizations
Partitioning by column unlocks the potential for more optimizations:
- Sparse bitvectors
- Label encoding
- Columnar compression: Train on compressed features—without decompression!
Training on Compressed Features
-
Idea popularized by the databases community
- Sort column
- Compress column using run-length encoding (RLE), delta encoding, etc.
- That's it! Never have to fully materialize (i.e., decompress) the column again
- Works well for sparse features
Training on Compressed Features
- We already need to sort the features to perform greedy split-finding
- Feature values are visited in sequential order, so we never have to decompress!
- Downside: Can't use our sorting trick anymore; need data structure to keep track of split history
- Upside: Feature can now fit entirely in cache; no more DRAM accesses
Results: Real-world Datasets
Parameters:
| # instances | 8.1 $\times$ 106 |
| # features | 784 |
| Size | 18.2 GB (sparse) |
| Task | classification |
Results: Real-world Datasets
Parameters:
| # instances | 2 $\times$ 106 |
| # features | 3500 |
| Size | 52.2 GB (dense) |
| Task | regression |
Results: Scalability
Parameters:
| # instances | 2 $\times$ 106 |
Results: Scalability
Parameters:
| # instances | 2 $\times$ 106 |
| Task | regression |
🔑 Empirical results match the
expected tradeoffs
Results: Optimizations
Parameters:
| Dataset | MNIST 8M |
| Depth | 10 |
🔑 Optimizations improve runtime, particularly for sparse datasets
Using Yggdrasil
-
Yggdrasil is open-source:
https://github.com/fabuzaid21/yggdrasil - Pull requests welcomed!
-
Also available as a Spark package:
https://bit.ly/ygg-spark - Compatible with Spark 1.6+
- Code already written in Spark MLlib? Not a problem!
Future Work
- Merge Yggdrasil into Spark MLlib 2.1
- Add decision rule that automatically chooses best algorithm for you a priori
- Add support for approximate binning
- Single feature doesn't fit on a single node? Fix it!
Thanks!
Check out the full paper at NIPS 2016!
Any questions? Shoot me an email:
fabuzaid at cs dot stanford dot edu
