用機器學習構建O(N)複雜度的排序演算法,可在GPU和TPU上加速計算

機器之心發表於2018-05-16

排序一直是電腦科學中最為基礎的演算法之一,從簡單的氣泡排序到高效的桶排序,我們已經開發了非常多的優秀方法。但隨著機器學習的興起與大資料的應用,簡單的排序方法要求在大規模場景中有更高的穩定性與效率。中國科技大學和蘭州大學等研究者提出了一種基於機器學習的排序演算法,它能實現 O(N) 的時間複雜度,且可以在 GPU 和 TPU 上高效地實現平行計算。這篇論文在 Reddit 上也有所爭議,我們也希望機器學習能在更多的基礎演算法上展現出更優秀的效能。

排序,作為資料上的基礎運算,從計算伊始就有著極大的吸引力。雖然當前已有大量的卓越演算法,但基於比較的排序演算法對Ω(N log N) 比較有著根本的需求,也就是 O(N log N) 時間複雜度。近年來,隨著大資料的興起(甚至萬億位元組的資料),效率對資料處理而言愈為重要,研究者們也做了許多努力來提高排序演算法的效率。

大部分頂尖的排序演算法採用平行計算來處理大資料集,也取得了卓越的成果。例如,2015 年阿里巴巴開發的 FuxiSort,就是在 Apsara 上的分散式排序實現。FuxiSort 能夠在隨機非偏態(non-skewed)資料集上用 377 秒完成 100TB 的 Daytona GraySort 基準,在偏態資料集上的耗時是 510 秒,而在 Indy GraySort 基準上的耗時是 329 秒。到了 2016 年,在 Indy GraySort 基準上,Tencent Sort 排序 100TB 資料時達到了 60.7TB/min 的速度,使用的是為超大資料中心優化過的包含 512 個 OpenPOWER 伺服器叢集。然而,這些演算法仍舊受下邊界複雜度和網路耗時的限制。

另一方面,機器學習在近年來發展迅速,已經在多個領域中得到廣泛應用。在 2012 年,使用深度卷積神經網路實現 ImageNet 影象的接近誤差減半的分類是一項重大突破,並使深度學習迅速被計算機視覺社群所接受。在 2016 年 3 月,AlphaGo 使用神經網路人工智慧的重大挑戰即圍棋中打敗了世界冠軍李世石。機器學習的巨大成功表明計算機 AI 可以在複雜任務中超越人類知識,即使是從零開始。在這之後,機器學習演算法被廣泛應用到了多種領域例如人類視覺、自然語言理解、醫學影象處理等,並取得了很高的成就。

由人類大腦結構啟發而來的神經網路方法擁有輸入層、輸出層和隱藏層。隱藏層由多個連結人工神經元構成。這些神經元連線強度根據輸入和輸出資料進行調整,以精確地反映資料之間的關聯。神經網路的本質是從輸入資料到輸出資料的對映。一旦訓練階段完成,我們可以應用該神經網路來對未知資料進行預測。這就是所謂的推理階段。推理階段的精度和效率啟發研究者應用機器學習技術到排序問題上。在某種程度上,可以將排序問題看成是從資料到其在資料集位置的對映

在本文中,研究者提出了一個複雜度為 O(N·M)的使用機器學習的排序演算法,其在大資料上表現得尤其好。這裡 M 是表示神經網路隱藏層中的神經元數量的較小常數。我們首先使用一個 3 層神經網路在一個小規模訓練資料集上訓練來逼近大規模資料集的分佈。然後使用該網路來評估每個位置資料在未來排序序列中的位置。在推理階段,我們不需要對兩個資料之間進行比較運算,因為我們已經有了近似分佈。在推理階段完成之後,我們得到了幾乎排序好的序列。因此,我們僅需要應用 O(N) 時間複雜度的運算來得到完全排序的資料序列。此外,該演算法還可以應用到稀疏雜湊表上。

演算法

若假定我們有一個實數序列 S,它的長度為 N、上邊界和下邊界分別為 x_max 和 x_min。對於一個有效的排序演算法,我們需要交換 x_i 的位置來確保新的序列 S' 是經過排序的。假設一個實數 x_i 在序列 S' 中的位置為 r_i,那麼我們可以將排序問題視為一個雙對映函式 G(x_i)=r_i。如果我們可以預先求得這個函式,那麼排序演算法的複雜度就為 O(N)。實際上,如果序列 S 中所有的實數都來自同一分佈 f(x),且當 N 足夠大時,那麼 x_i 在新序列 S' 中的排位 r_i 將近似等於:

用機器學習構建O(N)複雜度的排序演算法,可在GPU和TPU上加速計算

其中 F 為資料的概率分佈函式,且當 N 趨向於無窮大時,表示式左右兩邊取等號。

這樣形式化排序問題的困難時函式 G(x) 通常是很難推導的,概率密度函式 f(x) 同樣也如此。然而當我們處理大資料序列時,N 會足夠大以令序列保持一些統計屬性。因此如果我們能推出概率密度函式 f(x),那麼就有機會根據上面所示的方程 1 降低排序演算法的複雜度到 O(N)。

在這一篇論文中,作者們應用了廣義支援向量機(General Vector Machine,GVM)來逼近概率密度函式 f(x)。這種 GVM 是帶有一個隱藏層的三層神經網路,且它的結構展示在以下圖 1 中。GVM 的學習過程基於蒙特卡洛演算法而不是反向傳播,作者們也發現 GVM 非常適合擬合函式。

用機器學習構建O(N)複雜度的排序演算法,可在GPU和TPU上加速計算

圖 1:GVM 的簡單圖示。研究者在每個實驗中固定 M 為 100 個隱藏層神經元

在該神經網路中,輸入層僅有一個神經元且輸入是用於擬合函式的 x_i,輸出層也只有一個神經元,輸出為 y_i。研究者修改了隱藏層的神經元數量為 M=100。實際在某種程度上,隱藏層的神經元越多擬合的精度就越大,但同時也伴隨著過擬合問題,以及計算效率降低的問題。

N 個實數的排序估計過程僅需要 O(N·M) 的時間。M 與 N 是互相獨立的,且在理論分析上 M 是沒有下界的。例如如果資料序列服從高斯分佈且我們只使用一個隱藏神經元,那麼計算複雜度就為 log(N)。特別地,我們也可以用多個神經元擬合高斯分佈神經元的數量依賴於機器學習方法。

在預測過程中,這種演算法不需要比較和交換運算,並且每個資料的排序估計都是互相獨立的,這使得平行計算變得高效且網路負載小。除了高效平行計算之外,由於機器學習需要矩陣運算,它還適用於在 GPU 或 TPU 上工作以實現加速 [19]。

實驗

如圖 2 所示,我們選擇兩種分佈進行實驗:均勻分佈和截尾正態分佈。

用機器學習構建O(N)複雜度的排序演算法,可在GPU和TPU上加速計算

圖 2:資料分佈。(a)截尾正態分佈和(b)均勻分佈的 107 個資料點。(c)截尾正態分佈和(d)均勻分佈的訓練序列分佈的 103 個資料點。紫色實線是解析分佈,粉色點線是實驗資料。

圖 3 對比了 Tim Sorting 和 Machine Learning Sorting 的執行時間。

用機器學習構建O(N)複雜度的排序演算法,可在GPU和TPU上加速計算

圖 3:(a)截尾正態分佈的資料數量和時間複雜度的關係。(b)截尾正態分佈的資料數量和時間複雜度離均差的關係。(c)均勻分佈的資料數量和時間複雜度的關係。(d)均勻分佈的資料數量和時間複雜度離均差的關係,研究者使用了 102 次實現的總體均值來獲得結果。

論文:An O(N) Sorting Algorithm: Machine Learning Sorting

用機器學習構建O(N)複雜度的排序演算法,可在GPU和TPU上加速計算

論文地址:https://arxiv.org/pdf/1805.04272.pdf

我們提出了一種基於機器學習方法的 O(N) 排序演算法,其在大資料排序應用上有巨大的潛力。該排序演算法可以應用到並行排序,且適用於 GPU 或 TPU 加速。此外,我們還將該演算法應用到了稀疏雜湊表上。

相關文章