在Spark上實現線性收斂的隨機優化演算法SVRG

01-24

線性收斂的隨機優化演算法SVRG

很多常見的機器學習模型（比如SVM、邏輯回歸）都是解決凸函數的有限和問題，可以概括成以下這種形式：

$min_{mathbf{win R^{d}}} f(w):=frac{1}{n}sum_{i=1}^{n}{f_{i}(mathbf{w})}$

其中 $mathbf{w}$ 是模型的參數， $f_{i}(mathbf{w})$ 代表樣本 $mathbf{x}_i$ 的損失函數，一般具有以下形式：

${f_{i}(mathbf{w})}:=f(mathbf{x}_i, mathbf{w}) + lambda h(mathbf{w})$

其中， $h(mathbf{w})$ 代表正則化項（用於控制模型複雜度或者模型稀疏度等等）， $lambda$ 是正則化超參數。正則化是一個懲罰函數，通過減小模型參數的數值 $mathbf{w}(j)$ ，有效緩解過擬合。

比如邏輯回歸的損失函數：

$f_{i}(mathbf{w}):=log(1+x^{-y_{i}mathbf{x}_{i}^{T}mathbf{w}})+frac{lambda}{2}||mathbf{w}||^{2}$

SVM的損失函數：

$f_{i}(mathbf{w}):=maxleft{ 0, -y_{i}mathbf{x}_{i}^{T}mathbf{w} ight}+frac{lambda}{2}||mathbf{w}||^{2}$

因為每個 $f_{i}(mathbf{w})$ 都有一階導，可以用梯度下降法（Gradient Descent）迭代求解：

$mathbf{w}_{t+1} = mathbf{w}_t - eta frac{1}{n}sum_{i=1}^n riangledown f_i(mathbf{w}_t)$

其中 $mathbf{w}_{t+1}$ 表示第 t+1 次更新後的參數， $eta$ 表示學習率。

梯度下降每次迭代需要求解所有樣本的梯度，樣本數比較多的時候，可能造成內存不夠用。而且梯度下降每次更新的計算量太大，沒法做online-learning。實際生產環境中，訓練數據量很大，所以往往採用隨機梯度下降演算法（Stochastic Gradient Descent），一般簡寫做SGD。

SGD每次迭代的時候均勻隨機得選擇一個樣本或者mini-batch做更新：

$mathbf{w}_{t+1} = mathbf{w}_t - eta riangledown f_i(mathbf{w}_t)$

相對於梯度下降，SGD可以選擇合適的batch size有效利用GPU顯存。而且SGD迭代的速度很快，可以算作是一種online-learning。但是SGD帶來的問題是收斂速度不如梯度下降，也就是說為了達到同樣的精度，SGD需要的總迭代次數要大於梯度下降。

SGD，GD，SVRG的收斂速度對比如下：

Algorithm Convex Gradient Dominated

SGD $O(1/epsilon^{2})$ $O(1/epsilon^{2})$

Gradient Descent $O(n/epsilon)$ $O(n aulog(n/epsilon))$

SVRG $O(n+(sqrt{n}/epsilon))$ $O((n+n^{2/3} au)log(1/epsilon))$

從收斂速度分析上看，SGD能夠在目標函數強凸並且遞減步長的情況下做到 $O(1/epsilon^{2})$ 的次線性收斂（sublinear convergence），而梯度下降則可以在目標函數強凸的情況下做到 $O(n/epsilon)$ 的線性收斂（linear convergence）。總結起來就是，如果想快速得到一個可以勉強接受的解，SGD比梯度下降更加合適，但是如果想得到一個精確度高的解，應當選擇梯度下降，但是GPU訓練可能顯存不夠用，如果內存夠用，選擇CPU訓練，訓練速度會慢很多。

SGD後來也衍生出了非常多的變種，比如Adagrad、Dual Averaging、FTRL等，但是並沒有在收斂速度上實現突破。直到2012和2013年，SAG[1]與SVRG[2]演算法發表在NIPS上。

SAG演算法

SAG演算法在內存中為每個樣本維護一個舊的梯度 $y_i$ ，每次迭代隨機選擇一個樣本i來更新d，並用d來更新參數 $mathbf{w}$ 。具體來說，就是用新的梯度 $riangledown f_i(x)$ 替換掉d中的舊梯度 $y_i$ ，d累積了樣本新舊梯度的偏差。如此，每次更新的時候僅僅需要計算一個樣本的梯度，而不是所有樣本的梯度，計算開銷與SGD無異，但是內存開銷要大得多。[1]中已經證明SAG是一種線性收斂演算法，這個速度比SGD快得多。

SAG實驗結果（結果摘自[1]的arxiv長文版）

實驗中使用的是L2正則化的邏輯回歸，左一是訓練誤差，左二和左三分別是測試集上的損失和誤差。注意左一的縱坐標是對數坐標，一般衡量優化演算法的速度都會採用對數坐標，因為在對數坐標中可以明顯看出一個演算法是線性收斂（近乎直線下降）還是次線性收斂（大體是一條向下凸的曲線）。通過左一可以看出SAG是一種線性收斂演算法，且相對於其他參與比較的演算法有很大的優勢。具體實驗配置數據集等可以參考原文。

SVRG演算法

SVRG的演算法思路是，每過一段時間（1～5個epoch）計算一次所有樣本的梯度 $ilde{mu}$ ，稱為全量梯度，並維護一個舊的模型參數。每個階段內部的單次更新採用 $riangledown f_{i_t} (w_{t-1}) - riangledown f_{i_t} ( ilde{w}) + ilde{mu}$ 來更新當前參數，具體來說，每次更新都是使用全量梯度，只不過這個全量梯度是用舊模型上算得的梯度和當前模型上算得的梯度差來修正過的。跟SGD相比，SVRG增加了計算全量梯度的開銷，由於m等於1～5個epoch，所以最多增加一倍的梯度計算。相對於SAG來說，SVRG不需要在內存中為每個樣本都維護一個梯度，極大地節省了內存資源。[2]中分析SGD更新中樣本梯度的的方差是有常數上界，導致其無法線性收斂。而SVRG通過 $riangledown f_{i_t} (w_{t-1}) - riangledown f_{i_t} ( ilde{w}) + ilde{mu}$ 這種特殊的更新項來讓方差有一個可以不斷減少的上界，因此也就做到了線性收斂。

SVRG實驗結果（結果摘自[2]）

實驗中使用的是邏輯回歸，左一是訓練誤差，左二和左三分別是訓練誤差和方差的對數。實驗中可以看出SVRG顯然是線性收斂演算法。

2. 在Spark上實現SVRG

綜上所述，SVRG是一種線性收斂演算法，比SGD更快，比SAG更省內存。但是目前主流的機器學習平台還未實現SVRG演算法。Spark是一個通用的數據分析平台，雖然Spark RDD的特性限制了Spark MLLib的性能和擴展性，但是由於Spark極大地方便了機器學習任務中的數據準備工作，Spark MLlib依然成為目前最火的機器學習平台。想要在Spark中實現SVRG，還得先了解一下Spark的相關特性。

在 Spark 中，計算被建模成有向無環圖（DAG：directed acyclic graph），其中每一個頂點都代表一個彈性分散式數據集（RDD：Resilient Distributed Dataset），每一條邊都代表對 RDD 的一個運算。RDD 是被分到了不同邏輯分區的對象的集合，這些邏輯分區是在內存中存儲和處理的，當內存不夠用時可能spill到磁碟。運算有兩種：變換（transformation）和動作（action）。變換（比如：映射、過濾、連接）是指在一個 RDD 上執行一種運算生成一個新的 RDD。

用戶將計算建模為DAG，在DAG調度中需要對計算過程劃分stage，每個stage作為一系列並行運行的任務執行（每個分區執行一個任務）。stage劃分的依據是RDD之間的依賴關係。RDD之間的依賴關係分為窄依賴和寬依賴。窄依賴是指父RDD的每個分區只被子RDD的一個分區所使用，如map，filter，union等操作，寬依賴是指父RDD的每個分區都可能被多個子RDD分區所使用，如groupByKey操作。窄依賴有利於高效執行，寬依賴則需要通信密集的shuffle運算。

Spark中的分散式執行是通過將這種 DAG stage分割到不同的機器上執行的。這張圖清晰地顯示了這種 master-worker 架構。Driver包含了DAG調度器和任務調度器，並且負責將任務分配到對應的Worker。在Spark MLlib中，模型參數存儲在Driver的內存中，每次迭代Worker通過廣播（broadcast）或閉包（closure）更新Driver中的模型參數。當模型參數大小超過Driver的內存，就需要把模型參數存儲在RDD中，每次迭代更新模型參數，都需要重新創建RDD。在Spark MLLib中，一般通過treeAggregate操作來執行每次迭代的。

for t=0,1,2,…,T do val (fullGradient, totalLoss) = data.treeAggregate(calculate gradient and loss) newWeight = data.mapPartitions(iter => { Randomly pick up an instance, update weight }, true).treeAggregate(locally updated weight) / numParts;end for

我們在SkyDiscovery（公司內部Spark的深度優化版本）中實現了SVRG演算法，並基於SVRG實現了邏輯回歸演算法，我們將其跟Spark中基於L-BFGS實現邏輯回歸演算法進行對比：

我們使用spark-bench 生成5-140G的測試數據。可以看出SVRG演算法在不同數據集上都能實現3-4倍的加速，並且通過MSE評估訓練好的模型，SVRG訓練出來模型精度更高。

3. 舉例說明SkyDiscovery中如何使用SVRG

SkyDiscovery在官網中可以免費下載使用，下面是個SkyDiscovery中基於SVRG的邏輯回歸演算法的例子。以根據致病因素預測胃病發生的概率為例，在數據載入後，轉化為Spark中的LabelPoint對象，然後調用LogisticRegressionWithDSVRG介面。

我們目前正在從消除瓶頸，非凸優化演算法，以及具體的深度學習演算法方面深度優化Tensorflow，公司福利好，可遠程辦公，歡迎有深度學習系統開發背景和工程能力強的平台工程師加入。

簡歷投遞：zuo.wang@sky-data.cn

參考文獻：

[1] Roux, Nicolas L., Mark Schmidt, and Francis R. Bach. "A stochastic gradient method with an exponential convergence rate for finite training sets." Advances in Neural Information Processing Systems. 2012.

[2] Johnson, Rie, and Tong Zhang. "Accelerating stochastic gradient descent using predictive variance reduction." Advances in Neural Information Processing Systems. 2013.

[3] Defazio, Aaron, Francis Bach, and Simon Lacoste-Julien. "Saga: A fast incremental gradient method with support for non-strongly convex composite objectives." Advances in Neural Information Processing Systems. 2014.

[4] Schmidt, Mark, Nicolas Le Roux, and Francis Bach. "Minimizing finite sums with the stochastic average gradient." arXiv preprint arXiv:1309.2388 (2013).

[5] S. J. Reddi, S. Sra, B. Poczos, and A. Smola. Fast incremental method for nonconvex optimization. arXiv:1603.06159, 2016b