Details
-
Improvement
-
Status: Closed
-
Major
-
Resolution: Won't Fix
-
None
-
None
Description
The current implementation of (Stochastic) Gradient Descent performs one map-reduce shuffle per iteration. Moreover, when the sampling fraction gets smaller, the algorithm becomes shuffle-bound instead of CPU-bound.
(1 to numIterations or convergence) {
rdd
.sample(fraction)
.map(Gradient)
.reduce(Update)
}
A more performant variation requires only one map-reduce regardless from the number of iterations. A local mini-batch SGD could be run on each partition, then the results could be averaged. This is based on (Zinkevich, Martin, Markus Weimer, Lihong Li, and Alex J. Smola. "Parallelized stochastic gradient descent." In Advances in neural information processing systems, 2010, http://www.research.rutgers.edu/~lihong/pub/Zinkevich11Parallelized.pdf).
rdd
.shuffle()
.mapPartitions((1 to numIterations or convergence) {
iter.sample(fraction).map(Gradient).reduce(Update)
})
.reduce(Average)
A higher level iteration could enclose the above variation; shuffling the data before the local mini-batches and feeding back the average weights from the last iteration. This allows more variability in the sampling of the mini-batches with the possibility to cover the whole dataset. Here is a Spark based implementation https://github.com/mashin-io/rich-spark/blob/master/main/src/main/scala/org/apache/spark/mllib/optimization/ParallelSGD.scala
(1 to numIterations1 or convergence) {
rdd
.shuffle()
.mapPartitions((1 to numIterations2 or convergence) {
iter.sample(fraction).map(Gradient).reduce(Update)
})
.reduce(Average)
}