spark-liblinear源码解析

spark-liblinear是一个可用于大规模LR的分布式库，其核心迭代算法为信赖域牛顿法(Trust-region Newton method 简称TRON)。可能很多人并没有听说过这种算法，也不知为何要使用它，它与那些常见的SGD、L-BFGS等优化方法有何不同，下面让我们一探究竟。

1、从LR模型说起

LR模型常被用于二分类问题，给定数据 $\text{x}$ 和权重$(\text{w},b)$，可用如下概率公式来表示：
$$P(y=\pm 1|\text{x,w})=\frac{1}{1+\exp (-y({\text{w}}^T\text{x}+b))}\text{\hspace{1em}(1)}$$
这里 $y$ 是 label 标签。对于训练样本 ${\text{x}}_i,i=1,...,l$，label w为 y i ∈ { − 1 , + 1 } y_i\in \{-1, +1\} yi​∈{−1,+1}，通常最小化负对数似然来对参数$(\text{w},b)$进行估计：
$$\underset{\text{w},b}{\min }\sum _{i=1}^l\log (1+e^{-y_i({\text{w}}^T{\text{x}}_i+b)})\text{\hspace{1em}(2)}$$.
为了推导方便，我们一般将标量$b$作为样本$\text{x}$的一个维度考虑：
$$[{\text{x}}_i^T,1]\to {\text{x}}_i^T[{\text{w}}^T,b]\to {\text{w}}^T\text{\hspace{1em}(3)}$$
同时为了获得更好的泛化能力，我们将加入L2正则化项${\text{w}}^T\text{w}/2$，因此本文将考虑解决如下正则化LR问题：

$$\underset{\text{w}}{\min }f(\text{w})\equiv \frac{1}{2}{\text{w}}^T\text{w}+C\sum _{i=1}^l\log (1+e^{-y_i{\text{w}}^T{\text{x}}_i})\text{\hspace{1em}(4)}$$

2、牛顿法&截断牛顿法

为了解决问题(4)，我们首先考虑牛顿法，它是一种二阶优化算法，需要计算$f(\text{w})$的梯度与Hessian矩阵：
$$\begin{gathered} \nabla f(\text{w})=\text{w}+C\sum _{i=1}^l(\sigma (y_i{\text{w}}^T{\text{x}}_i)-1)y_i{\text{x}}_i \\ {\nabla }^{2}f(\text{w})=I+C{X}^{T}DX,\text{\hspace{1em}(5)} \end{gathered}$$
牛顿法更新权重$\text{w}$： ${\text{w}}^{k+1}={\text{w}}^k+{\text{s}}^k$, 其中${\text{s}}^k$是牛顿方向, 是如下线性系统的解：
$${\nabla }^{2}f({\text{w}}^k){\text{s}}^k=-\nabla f({\text{w}}^k)\text{\hspace{1em}(6)}$$
但是，使用如上更新规则时会遇到两个问题：

1. 序列 ${\text{w}}^k$ 可能无法收敛到最优解，甚至不能确定函数的值会降低。
2. 我们假设$X$足够稀疏，但是$X^{T}DX$会仍是稠密的，当数据规模很大时，Hessian矩阵难以存储。

我们可以通过调整牛顿方向的长度来解决第一个问题，当前的技术主要有两种：线搜索与置信域方法。我们可以这么理解线搜索方法就像年轻人的梦想：沿着搜索方向，勇往直前的寻求下降而不计后果。与之截然相反的是，信赖域方法则像四十而不惑的人，回首朝气蓬勃的年轻时代，首先考虑到的是在当前模型下走一步所产生的预期，随后通过函数值是否下降或者超过预期的程度来调整模型状态。

对于第二个问题，同样有两类方法可供选择：直接法与迭代法。迭代法主要是计算Hessian矩阵与向量$\text{s}$的乘积：
$${\nabla }^{2}f(\text{w})\text{s}=(I+C{X}^{T}DX\text{s})=\text{s}+C\cdot X^T(D(X\text{s}))\text{\hspace{1em}(7)}$$
来避免对矩阵的直接存储，这种方法一般比直接计算法高效得多。在众多迭代算法中，共轭梯度法是使用最多的。因此我们看到对于求解问题(4)，外层是牛顿法进行迭代，内层是共轭梯度法迭代来求牛顿方向。然而，使用共轭梯度法会遇到迭代时间过长的问题，因此为了节省时间，我们一般在外层迭代时使用“近似”的牛顿方向，这种技术称为：截断牛顿法。

3、Trust Region Newton Method

首先让我们看看什么是信赖域方法：

在每次迭代时强制性地要求新的迭代点与当前的迭代点之间的距离不超过某一控制量. 引入控制步长是因为传统的线搜索方法常常由于步长过大而导致算法失败,特别是当问题是病态时. 控制步长实质上等价于在以当前迭代点为中心的一个邻域内对一个近似于原问题的简单模型求极值. 这种技巧可理解为只在一个邻域内对近似模型信赖,所以此邻域被称为信赖域(trust region)。

通过将置信域方法与截断牛顿法相结合，精细地设置内部共轭梯度法的迭代停止条件来挑选出好的牛顿方向，从而可以保证算法收敛到全局最小值，我们将之称为置信域牛顿法(Trust Region Newton Method)。接下来我们将通过公式来进一步阐述TRON算法。

首先，在利用tron算法最小化$f(\text{w})$时，第k步迭代存在一个大小为${\Delta }_k$的置信区间，此时用二阶泰勒展开：
$$q_k(\text{s})=\nabla ({\text{w}}^k{)}^T\text{s}+\frac{1}{2}{\text{s}}^T{\nabla }^{2}f({\text{w}}^k)\text{s}\text{\hspace{1em}(5)}$$
来作为$f({\text{w}}^k+\text{s})-f({\text{w}}^k)$的近似值。接下来，我们需要寻找步长${\text{s}}^k$来近似最小化$q_k(\text{s})$同时满足$||\text{s}||\le {\Delta }_k$。然后根据真实模型中loss的减少值与二阶近似模型中预期的减少值的比值${\rho }_k$，来更新${\text{w}}^k、{\Delta }_k$:
$${\rho }_k=\frac{f({\text{w}}^k+\text{s})-f({\text{w}}^k)}{q_k({\text{s}}^k)}\text{\hspace{1em}(8)}$$
${\text{w}}^k$的更新规则如下，其中${\eta }_0>0$是需要设置的参数

$${\text{w}}^{k+1}=\{\begin{array}{llc}{\text{w}}^k+{\text{s}}^k& {\rho }_k>{\eta }_0& \\ {\text{w}}^k& {\rho }_k\le {\eta }_0& \end{array}\text{\hspace{1em}(9)}$$

${\Delta }_k$的更新规则如下，其中参数${\eta }_1<{\eta }_2<1,0<{\sigma }_1<{\sigma }_2<1<{\sigma }_3$
(10) { Δ k + 1 ∈ [ σ 1 min ⁡ { ∣ ∣ s k ∣ ∣ , Δ k } , σ 2 Δ k ] ρ k ≤ η 1 Δ k + 1 ∈ [ σ 1 Δ k , σ 3 Δ k ] ρ k ∈ ( η 1 , η 2 ) Δ k + 1 ∈ [ Δ k , σ 3 Δ k ] ρ k ≥ η 2 \begin{cases} \Delta_{k+1}\in[\sigma_1\min\{||\textbf s^k||, \Delta_k\}, \sigma_2\Delta_k]&amp;&amp; \rho_k\le\eta_1\\ \Delta_{k+1}\in[\sigma_1\Delta_k, \sigma_3\Delta_k] &amp;&amp; \rho_k\in(\eta_1,\eta_2)\\ \Delta_{k+1}\in[\Delta_k, \sigma_3\Delta_k] &amp;&amp; \rho_k\ge\eta_2 \\ \end{cases} \tag {10} ⎩⎪⎨⎪⎧​Δk+1​∈[σ1​min{∣∣sk∣∣,Δk​},σ2​Δk​]Δk+1​∈[σ1​Δk​,σ3​Δk​]Δk+1​∈[Δk​,σ3​Δk​]​​ρk​≤η1​ρk​∈(η1​,η2​)ρk​≥η2​​(10)

因此，我们给出完整的tron算法如下：

其中子问题(11)由如下的共轭梯度法来近似求解：

4、分布式的spark-liblinear

了解tron算法之后，让我们再来看看在spark-liblinear中是如何做分布式的以及代码实现中的一些技巧。

在spark-liblinear中通过将梯度与Hessian矩阵的计算公式改写为分量形式，以此便于分布式计算。
$$\begin{gathered} f(\text{w})=\frac{1}{2}{\text{w}}^T\text{w}+C\sum _{k=1}^p{f}_k(\text{w}), \\ \nabla f(\text{w})=\text{w}+C\sum _{k=1}^p\nabla f_k(\text{w}), \\ {\nabla }^{2}f(\text{w})\text{v}=\text{v}+C\sum _{k=1}^p{\nabla }^2{f}_k(\text{w})\text{v}\text{\hspace{1em}(14)} \end{gathered}$$
这里：
$$\begin{gathered} f_k(\text{w})\equiv e_k^T\log (\sigma (Y_k{X}_k\text{w})), \\ \nabla f_k(\text{w})\equiv (Y_k{X}_k{)}^T(\sigma (Y_k{X}_k\text{w}{)}^{-1}-e_k), \\ {\nabla }^2{f}_k(\text{w})\text{v}\equiv X_k^T(D_k(X_k\text{v})), \\ {D}_k\equiv diag((\sigma (Y_k{X}_k\text{w})-e_k)/\sigma (Y_k{X}_k\text{w}{)}^2)\text{\hspace{1em}(15)} \end{gathered}$$

对于计算(14)，每个分量的计算只需要与之相关的数据分量$X_k、Y_k$，因此可以做分布式计算。

spark-liblinear 是基于scala语言开发的，为了降低开发scala程序的开销，需要涉及到一些复杂的设计问题。为实现高效的计算、通信和内存使用，需要分析以下不同的实现问题。

1、Loop Structure

算法中主要的计算量在于：矩阵与向量相乘， 需要用到循环结构。可以考虑三种方式： for-generator、for-range、while。 通过实验发现while 最高效。for表达式 在编译时先转译成高阶运算，开销较大

while(i < p.index.length)
		{
			z += p.value(i) * wB(p.index(i))
			i += 1
		}
		z = (1.0 / (1.0 + exp(-p.y * z)) - 1.0) * p.y
		i = 0
		while(i < p.index.length)
		{
			grad(p.index(i)) += z * p.value(i)
			i += 1
		}

2、Data Encapsulation

将数据表示成稀疏矩阵，我们只存储非零值处的索引与值。文中使用两个Array来存储实例对象的索引与特征值，来减少内存开销与调用时间。

class DataPoint(val index : Array[Int], val value : Array[Double], val y : Double) extends Serializable

3、Using mapPartitions Rather Than map

map的输入变换函数是应用于RDD中每个元素，而mapPartitions的输入函数是应用于每个分区。使用MapPartitions操作之后，一个task仅仅会执行一次function，function一次接收所有的partition数据。只要执行一次就可以了，性能比较高。 在计算函数值、梯度、hessian以及数据转换时用到了mapPartitions

4、not to cache $\sigma (Y_k{X}_k\text{w})$

对中间量的缓存是单机版liblinear使用的策略，对于数据量巨大的情形下，分布式下的带来的额外通信会抵消掉缓存带来的计算时间节省，因此不做缓存

5、 Using Broadcast Variables

算法迭代过程中，每个分区都使用来自master的相同信息的w与s，可以通过使用广播变量来减少通信间的开销

6、 The Cost of the reduce Function

使用coalesce函数在分区的输出发送到主服务器之前，将它们局部地组合在一起，减少通信间开销。

5、spark-liblinear性能调优

spark-liblinear库的使用十分简单方便，可以直接在spark-shell里调用jar包进行训练预测。

关于用spark-liblinear时的遇到的一些性能问题，可以参考如下：

另外关于spark-liblinear里的参数 -N 如何设置：

1、-N 用于将数据重新分区，使用coalesce函数使分区数减少，这个过程不会发生shuffle。

2、当样本量或特征量非常大时，运行spark-liblinear会出现这样的错误：
ERROR scheduler.TaskSetManager: Total size of serialized results of 342 tasks (8.3 GB) is bigger than spark.driver.maxResultSize (2.0 GB)。这是由于它在计算梯度时将所需计算的vector全部拉到driver上进行，会导致程序爆掉。通过合理的设置 -N 可以有效地避免这种错误。

3、例如原始数据分区3000，设置 -N 1600。在计算梯度时，会先在executor上把1600个task hash分成40*40，先计算40个梯度累加的结果，再将这40个拉到driver上进行运算。一般建议 -N 设置成 executor-num * core 的 1-2倍，且不超过原始分区数。

---

文章参考：
1、C.-J. Lin, R. C. Weng, and S. S. Keerthi, “Trust region Newton method for large-scale logistic regression,” Journal of Machine Learning Research, vol. 9, pp. 627–650, 2008.
[Online]. Available: [http://www.csie.ntu.edu.tw/ ∼ cjlin/papers/logistic.pdf]
2、Large-scale Logistic Regression and Linear Support Vector Machines Using Spark