Apache Spark MLlib provides scalable implementation of popular machine learning algorithms, which lets users train models from big dataset and iterate fast. The existing implementations assume that the number of parameters is small enough to fit in the memory of a single machine. However, many applications require solving problems with billions of parameters on a huge amount of data such as Ads CTR prediction and deep neural network. This requirement far exceeds the capacity of exisiting MLlib algorithms many of who use L-BFGS as the underlying solver. In order to fill this gap, we developed Vector-free L-BFGS for MLlib. It can solve optimization problems with billions of parameters in the Spark SQL framework where the training data are often generated. The algorithm scales very well and enables a variety of MLlib algorithms to handle a massive number of parameters over large datasets. In this talk, we will illustrate the power of Vector-free L-BFGS via logistic regression with real-world dataset and requirement. We will also discuss how this approach could be applied to other ML algorithms.

1. Scaling Apache Spark MLlib to
billions of parameters
Yanbo Liang
(Apache Spark committer @ Hortonworks)
joint with Weichen Xu

2. Outline
• Background
• Vector-free L-BFGS on Spark
• Logistic regression on vector-free L-BFGS
• Performance
• Integrate with existing MLlib
• Future work

3. Outline
• Background
• Vector-free L-BFGS on Spark
• Logistic regression on vector-free L-BFGS
• Performance
• Integrate with existing MLlib
• Future work

4. Big data + big models = better accuracies

5. Spark: a unified platform
“Hidden Technical Debt in Machine Learning Systems”, Google

6. Spark: a unified platform
“Hidden Technical Debt in Machine Learning Systems”, Google

7. Optimization
• L-BFGS
• SGD
• WLS

8. MLlib logistic regression
Initialize
coefficients
Broadcast coefficients
to Executors
Comput loss and
gradient for each
instance, and sum
them up locally
Comput loss and
gradient for each
instance, and sum
them up locally
Comput loss and
gradient for each
instance, and sum
them up locally
Reduce from
executors to get
lossSum and
gradientSum
Handle regularization
and
use L-BFGS/OWL-QN
to find next step Final model
coefficients
Executors/Workers
Driver/Controller
loop untial converge

9. Bottleneck 1
Initialize
coefficients
Broadcast coefficients
to Executors
Comput loss and
gradient for each
instance, and sum
them up locally
Comput loss and
gradient for each
instance, and sum
them up locally
Comput loss and
gradient for each
instance, and sum
them up locally
Reduce from
executors to get
lossSum and
gradientSum
Handle regularization
and
use L-BFGS/OWL-QN
to find next step Final model
coefficients
Executors/Workers
Driver/Controller
loop untial converge

10. Solution 1
Comput loss and
gradient for each
instance, and sum
them up locally
Comput loss and
gradient for each
instance, and sum
them up locally
Comput loss and
gradient for each
instance, and sum
them up locally
Reduce from
executors to get
final lossSum
and gradientSum
Comput loss and
gradient for each
instance, and sum
them up locally
Reduce from
executors to get
partial lossSum
and gradientSum
Reduce from
executors to get
partial lossSum
and gradientSum

11. Bottleneck 2
Initialize
coefficients
Broadcast coefficients
to Executors
Comput loss and
gradient for each
instance, and sum
them up locally
Comput loss and
gradient for each
instance, and sum
them up locally
Comput loss and
gradient for each
instance, and sum
them up locally
Reduce from
executors to get
lossSum and
gradientSum
Handle regularization
and
use L-BFGS/OWL-QN
to find next step Final model
coefficients
Executors/Workers
Driver/Controller
loop untial converge

12. Outline
• Background
• Vector-free L-BFGS on Spark
• Logistic regression on vector-free L-BFGS
• Performance
• Integrate with existing MLlib
• Future work

13. Distributed Vector
Driver
[2.5, 6.8, 3.1, 9.0, 0.7, 1.9, 2.6, 8.7, 6.2, 1.1, 5.0, 0.0]
worker0 worker1 worker2 worker3

14. Distributed Vector
Driver
[2.5, 6.8, 3.1, 9.0, 0.7, 1.9, 2.6, 8.7, 6.2, 1.1, 5.0, 0.0]
[2.5, 6.8, 3.1]
worker0
[9.0, 0.7, 1.9]
worker1
[2.6, 8.7, 6.2]
worker2
[1.1, 5.0, 0.0]
worker3

15. Distributed Vector
• Definition:
• Linear algebra operations:
– a * x + y
– a * x + b * y + c * z + …
– x.dot(y)
– x.norm
– …
class DistribuedVector(
val values: RDD[Vector],
val sizePerPart: Int,
val numPartitions: Int,
val size: Long)

16. L-BFGS
Coefficients: X(k)
Choose descent direction
(Two loop recursion)
X(k+1) = alpha * dir(k) + X(k)
Calculate G(k+1) and F(k+1)

17. L-BFGS
Coefficients: X(k)
Choose descent direction
(Two loop recursion)
X(k+1) = alpha * dir(k) + X(k)
Calculate G(k+1) and F(k+1)
Worker Worker Worker Worker
WorkerWorkerWorkerWorker
Driver
Space: (2m+1)d
Time: 6md

18. Vector-free L-BFGS
Coefficients: X(k)
Choose descent direction
(Two loop recursion)
X(k+1) = alpha * dir(k) + X(k)
Calculate G(k+1) and F(k+1)

19. Vector-free L-BFGS
Coefficients: X(k)
Choose descent direction
(Two loop recursion)
X(k+1) = alpha * dir(k) + X(k)
Calculate G(k+1) and F(k+1)
Worker Worker Worker Worker
WorkerWorkerWorkerWorker
Driver
Space: (2m+1)*d
Time: 6md
Space: (2m+1)^2
Time: 8m^2

20. Outline
• Background
• Vector-free L-BFGS on Spark
• Logistic regression on vector-free L-BFGS
• Performance
• Integrate with existing MLlib
• Future work

21. numFeatures
numInstance
partition1
partition2
partition3

22. numFeatures
numInstance

23. coefficients: DistributedVectorcoeff0 coeff1 coeff2

24. coeff2coeff0
coeff0
coeff0
coeff1
coeff1
coeff1
coeff2
coeff2

25. multiplier02multiplier22multiplier12
multiplier01
multiplier00
multiplier11
multiplier10
multiplier21
multiplier20

26. multiplier02multiplier22multiplier12
multiplier01
multiplier00
multiplier11
multiplier10
multiplier21
multiplier20
label0label2label1

27. multiplier0multiplier2multiplier1

28. multiplier0multiplier2multiplier1
multiplier0
multiplier0
multiplier1
multiplier1
multiplier2
multiplier2

29. grad02grad00
grad10
grad20
grad01
grad11
grad21
grad12
grad22

30. grad0 grad1 grad2 gradient: DistributedVector

31. Regularization and OWL-QN
• Regularization:
– L2
– L1
– elasticNet
• Implement vector-free OWL-QN for objective with L1/elasticNet
regularization.

32. Checkpoint
Coefficients: X(k)
Choose descent direction
(Two loop recursion)
X(k+1) = alpha * dir(k) + X(k)
Calculate G(k+1) and F(k+1)
Worker Worker Worker Worker
WorkerWorkerWorkerWorker
Driver
Space: (2m+1)*d
Time: 6md
Space: (2m+1)^2
Time: 8m^2

33. Outline
• Background
• Vector-free L-BFGS on Spark
• Logistic regression on vector-free L-BFGS
• Performance
• Integrate with existing MLlib
• Future work

34. Performance(WIP)
0
20
40
60
80
100
120
1M 10M 50M 100M 250M 500M 750M 1B
Time for each epoch
L-BFGS VL-BFGS

35. Outline
• Background
• Vector-free L-BFGS on Spark
• Logistic regression on vector-free L-BFGS
• Performance
• Integrate with existing MLlib
• Future work

36. APIs
val dataset: Dataset[_] =
spark.read.format("libsvm").load("data/a9a")
val trainer = new VLogisticRegression()
.setColsPerBlock(100)
.setRowsPerBlock(10)
.setColPartitions(3)
.setRowPartitions(3)
.setRegParam(0.5)
val model = trainer.fit(dataset)
println(s"Vector-free logistic regression coefficients:
${model.coefficients}")

37. Adaptive logistic regression
Number of features Optimization method
less than 4096 WLS/IRLS
more than 4096,
but less than 10 million
L-BFGS
more than 10 million VL-BFGS

38. Whole picture of MLlib
Application layer
(Ads CTR, anti-fraud, data science, NLP, …)
WLS/IRLS
Optimization layer
Spark core
L-BFGS VL-BFGS
LinearRegression
Machine learning algorithm layer
LogisticRegression MLP …
SGD …

39. APIs
val dataset: Dataset[_] =
spark.read.format("libsvm").load("data/a9a")
val trainer = new LogisticRegression()
.setColsPerBlock(100)
.setRowsPerBlock(10)
.setColPartitions(3)
.setRowPartitions(3)
.setRegParam(0.5)
.setSolver(“vl-bfgs”)
val model = trainer.fit(dataset)
println(s"Vector-free logistic regression coefficients:
${model.coefficients}")

40. Outline
• Background
• Vector-free L-BFGS on Spark
• Logistic regression on vector-free L-BFGS
• Performance
• Integrate with existing MLlib
• Future work

41. Future work
• Reduce stages for each iteration.
• Performance improvements.
• Implement LinearRegression, SoftmaxRegression and
MultilayerPerceptronClassifier base on vector-free L-BFGS/OWL-QN.
• Real word use case for Ads CTR prediction with billions parameters
and will share experiences and lessons we learned:
– https://github.com/yanboliang/spark-vlbfgs

42. Key takeaways
• Full distributed calculation, full distributed model, can run
successfully without OOM.
• VL-BFGS API is consistent with breeze L-BFGS.
• VL-BFGS output exactly the same solution as breeze L-BFGS.
• Does not require special components such as parameter servers.
• Pure library and can be deployed and used easily.
• Can benefit from the Spark cluster operation and development
experience.
• Use the same platform as the data size increasing.

43. Reference
• https://papers.nips.cc/paper/5333-large-scale-l-bfgs-using-
mapreduce
• https://github.com/mengxr/spark-vl-bfgs
• https://spark-summit.org/2015/events/large-scale-lasso-and-elastic-
net-regularized-generalized-linear-models/

44. Thank You.
Yanbo Liang