香港中大-商湯科技聯合實驗室AAAI錄用論文詳解：超大規模分類加速算法

近日，香港中大-商湯科技聯合實驗室的新論文“Accelerated Training for Massive Classification via Dynamic Class Selection”被AAAI 2018錄用為Oral Presentation。該論文著力于解決超大規模分類問題(過百萬類)對模型訓練帶來的一系列挑戰，尤其是超大規模Softmax層造成的計算瓶頸。具體而言，該論文提出了一種新型的Dynamic Selective Softmax，能夠以較低的代價在每次迭代中根據類別空間的結構動態選擇一小部分類別參與計算，在保證精度基本一致的條件下，可以非常顯著地提升訓練速度，并大幅度減少訓練過程所占用的顯存。本文中我們將詳細介紹該論文中提出的方法，并介紹一些可能的后續工作等。

值得指出的是，這項工作是由商湯科技在城市級人臉識別的場景所驅動。而此項工作的成果也反映了在人工智能領域產業應用對于學術創新的重要意義。

作者：張行程，楊磊，閆俊杰，林達華

一個實際的問題—人臉識別

隨著數據的迅速積累，數據的類別數目也在不斷增加，大規模分類的問題受到越來越廣泛的重視，人臉識別就是其中一個重要的例子。在先前的一些人臉識別的工作中[2]，基于分類的訓練方式因其簡單且能提供高效監督信息的特點，為學界和工業界廣泛使用。

圖1 基于分類方法的人臉識別訓練流程

然而，如下圖所示，隨著類別數量急劇增加，兩個實際的問題浮出水面：1)計算Softmax所需要的計算量會大大增加，使得Softmax的計算成為制約網絡訓練速度的瓶頸。2)類別數量的增加會導致分類層的參數量大大增加，在如今GPU顯存有限的情況下，會極大制約能夠參與分類的類別數量。

圖2 類別數量增加時

對應的計算量和顯存的變化

針對類別數增加時帶來的計算增長和顯存增加的挑戰，本論文提出了一種全新的選擇性Softmax計算框架，我們稱之為“基于動態類別選擇的Softmax”(Dynamic Selective Softmax)。該框架可以顯著減少訓練過程中的計算量以及顯存消耗，并且取得近似完整Softmax訓練的性能。

對兩個指標的觀察

我們定義了兩個指標來幫助我們分析大規模分類的問題。第一個指標叫做CP_K，即前K累加概率和，表示Softmax計算出的前K大的概率的總和。CP_K的值越大表示概率越集中在前K個的類別中。

進一步地，從Softmax的梯度公式可知，除了真實標簽之外，其它類別的梯度和概率大小成正比。這說明如果概率是集中在前K個類別中，那么對應的梯度值也是集中在前K個類別中。

因此，我們進一步定義了一個指標NCG_K，即前K歸一化累計梯度。這個指標衡量了梯度在前K個類別上的集中程度。

根據定義，我們用一個在Ms1MCeleb上訓練好的Resnet-101模型，畫出了NCG_K隨著迭代輪數變化的曲線。

從上圖中我們可以得到幾點重要的觀察：1)在多數情況下，一個樣本只容易和一小部分的類別混淆。2)當使用Softmax作為損失函數時，從這一小部分類別產生的梯度主導了所有的梯度。3)隨著訓練的進行，容易混淆的類別的集合會越來越小。

圖3 Softmax計算中的“活躍類別”

如圖所示，這K個概率值最大的類別主導了Softmax計算，我們將其稱之為“活躍類別”。

選擇性Softmax框架

基于上述觀察，我們提出了一種選擇性的Softmax框架。該框架引入了一個選擇器，該選擇器通過結合當前特征和分類的權重矩陣，能夠高效且準確地挑選出分類上能夠產生較強監督信號的“活躍類別”。根據之前的分析，“活躍類別”可以近似提供和所有類別同樣有效的監督信息，因此，對于這占整體類別比例很小的“活躍類別”，我們可以利用一些稀疏性的操作來加速整體的計算流程。

如下圖所示，我們將大規模分類所需要的巨大的權重矩陣放置于CPU的內存中，根據選擇器挑選出的類別，可以從這個巨大的權重矩陣中挑選出需要的權重向量，放置到GPU的顯存中。當在GPU上計算出對應的梯度之后，再異步地將求出的梯度稀疏更新到存在于CPU上的權重矩陣中。

圖4 選擇性Softmax工作流程示意圖

To Be or not to be的困局

根據上文的介紹，選擇性Softmax框架的一大核心是“活躍類別”的選擇器，如何找到一種高效且準確的選擇器成為這個計算框架的關鍵。首先我們研究了兩種最為直觀的選擇器。

第一種是隨機選擇器。在每一輪迭代中，除了當前訓練樣本、的類別，還會通過隨機采樣得到其余的類別。該隨機選擇器不依賴于當前樣本的特征和分類層權重矩陣的信息，幾乎沒有額外的時間開銷，也因為是隨機選取，沒有引入任何額外信息，該方法的性能可以被視為是選擇器的性能下界。

第二種是最優選擇器。在每一輪迭代中，會根據當前樣本的特征和整個特征矩陣點乘的結果，從中選取響應最大的類別，用這些類別來作為“活躍類別”。這種方法保證了每次選取出來的活躍類別都是最優的，但是由于引入了整個權重矩陣參與計算，無法減少前傳的計算開銷，這有悖于我們減少計算量的初衷。但最優選擇器每次都選取了最活躍的類別，該方法的性能可以被視為是選擇器性能的上界。

如圖所示，隨機選擇器雖然速度快，但是損失了精度;最優選擇器雖然保證了精度，但是沒有減少計算開銷。這兩種選擇器都無法滿足我們的設計要求，因此，我們希望找到一種在速度和精度上都可以取得較好平衡的“活躍類別”選擇器。

哈希森林選擇器

為了有一個速度和精度上可以達到良好平衡的選擇器，我們提出了哈希森林選擇器。我們將從哈希森林的構建和哈希森林的使用兩個方面來介紹這一算法。

對于哈希森林的構建，我們首先來看一棵哈希樹的構建。對于N個類別分類的權重矩陣，我們可以將這個權重矩陣看作是N個類別的平均特征向量，那么我們就得到了D維空間的N個點。首先，我們從這N個點中隨機選擇2個點，計算出距離這兩個點距離最大的超平面，也即，該超平面把這兩個點劃分在了空間的兩側。那么，根據這個計算出的超平面，我們就可以把空間中的其它點分為兩部分，即在超平面的不同側。這樣，我們就把空間分成了兩個子空間。對于每一個子空間，我們會重復上面的過程，即每次隨機選取2個點，計算出超平面對空間進行劃分，當子空間中點的個數不超過B個點時停止繼續劃分。經過若干次操作后，當所有子空間中的點的個數都不超過B個點時，構建過程結束。如果將每個劃分的超平面視為中間節點，將最終劃分后的結果視為葉子結點，那么構建的這個結構就可以被當成是一棵哈希樹。

從一棵哈希樹的構建過程可以看出，我們構建哈希樹時采用的是隨機算法，因此，為了消除隨機選取帶來的不穩定性，我們會同時生成L棵哈希樹，這樣就組成了哈希森林。

圖5 哈希森林的構建

對于哈希森林的使用，我們也先來看每棵哈希樹對應的操作。對于每一顆哈希樹，樹上的中間節點都是一個超平面。因此，當一個新的特征向量來了之后，首先會被根結點的超平面分到左子樹或者右子樹，接著又會被相應子樹的根節點繼續往下分類，直到落入一個葉子結點中。那么葉子結點中所包含的類別就被認為是和這個特征向量相似的類別。

對于森林中的每一棵樹，都會重復這個過程，最終會得到一個類別的候選集合，如果候選集合中的類別超出了我們所需的類別數目，我們再通過計算其余弦距離，最終選取余弦距離最小的那些類別作為我們的“活躍類別”。

更好的平衡—動態選擇性Softmax

雖然哈希森林的選擇器在保證準確率的同時大大提升了我們的效率，但從之前對兩個指標的觀察中，我們可以看到，隨著訓練的進行，所需要的活躍類別數越來越少，因此我們可以針對選擇器中引入的參數做動態調整，進一步優化計算資源的配置。

為此，我們引入了一個閾值，每隔一段時間，在驗證集合上計算出CP_K(前K概率累加和)，我們選取的類別數K恰好使得算出的CP_K超過定義的閾值。相比于預先設定好采樣類別數，這種方式隨著訓練的進行，K的值會越來越小，也即需要采樣的類別數會越來越少，這樣可以進一步減少無用的計算資源消耗。

對于新引入的閾值和另外兩個參數，樹的數目和更新周期。基于“訓練初期多嘗試不同的可能，訓練穩定后尋求更精準的選擇”這一理念，我們發現簡單的線性更新的方式就可以取得不錯的效果，也即：閾值由小到大線性增長，樹的數目由少線性增多，更新周期由大線性減小。

實驗結果

我們在兩個分布迥異的人臉識別數據集上來驗證選擇性Softmax的性能。第一個數據集是Ms1M Celeb，該數據集包括10萬人共1000萬張圖，由于原始的數據集有比較多的噪聲，我們進行了一些清理，實際用于實驗的是8.7萬人約460萬張圖。第二個數據集是Megaface Challenge 2的訓練數據。該數據集共有67萬人約470萬人。我們在4個不同的人臉識別標準協議上測試了方法的性能，即LFW非受限的人臉驗證協議，人臉識別協議，IJB-A的1:1人臉驗證協議以及Megaface的人臉識別協議。

在這些不同的協議上，我們的提出的方法有著一致的性能，在指數級別地減少計算量的同時，可以取得逼近于完整Softmax的性能。從表圖1中可以看到，在取得相同性能的情況下，哈希森林的時間開銷比最優選擇器或者完整的Softmax要小一個數量級;在相同的時間消耗下，哈希森林的性能遠超其它的一些算法。并且可以從圖中看到，動態的哈希森林可以進一步地提升算法的性能。

表圖1. 選擇性Softmax和其它方法的比較

除了在小模型上的比較實驗，我們還將Ms1MCeleb和Megaface Challenge 2兩個數據集合并在一起，得到一個75萬分類的大數據集，用Resnet-101在這個數據集上進行了實驗。在這個大實驗中，我們在幾乎不損耗性能的情況下，可以將每輪的訓練迭代時間從3.5s縮短為1.5s，起到了60%的加速效果。并且由于我們將巨大的權重參數放到了CPU的內存中，我們幾乎移除了權重矩陣所占用的所有顯存。從整體網絡的顯存消耗上看，我們可以將顯存開銷減小24%。

表圖2. 在Resnet101上進行

75萬人臉分類的訓練

除了得到更好的性能，我們也詳細分析了哈希森林中不同決定因子對最后結果的影響。由下圖可以得到幾個觀察：1)不管用何種選擇器，隨著采樣類別數目的增加，模型的性能都會隨之增加。2)基于哈希森林的選擇器在采樣類別數目很小時，就可以逼近最優選擇器的性能。3)對于哈希森林中樹的數目不是特別敏感，在超過50顆樹時性能就趨于穩定。4)當哈希森林的更新周期在1個epoch內時，對性能的影響不顯著，如果更新周期超過了一個1 epoch，可能因為構建哈希森林的權重矩陣信息太過久遠而造成整體性能的下降。

表圖2. 采樣類別數對性能的影響

表圖3. 樹的數目對性能的影響

圖表4. 更新周期對性能的影響

思考與最后的話

回顧本文的算法，我們針對大規模分類時Softmax帶來的問題，提出了基于動態類別層級選擇的算法，高效且準確地找到了“活躍類別”，并且我們進一步提出了動態分配的方案來進一步充分利用計算資源。在不同測試集上的實驗結果，充分驗證了本文方法的有效性。另外，除了本文中提出的哈希森林選擇器，該動態選擇性Softmax的計算框架還可以靈活引入不同的采樣算法，在實際的使用情況中，可以根據速度和精度的不同需求來選擇不同的采樣算法。在未來，我們希望進一步探尋該方法的邊界，當數據規模再大一個甚至幾個量級時，會不可避免地引入更多的噪聲。當這些噪聲達到無法人工清理的規模時，可能會對構建的動態的層級結構造成很多干擾，使得選擇器無法準確區分“活躍類別”和“噪聲類別”等信息，如何在這種情況下改進我們的選擇性Softmax是一個值得繼續探究的問題。

相關文獻：

【1】“Accelerated Training for Massive Classification via Dynamic Class Selection”, Xingcheng Zhang, Lei Yang, Junjie Yan and Dahua Lin, AAAI 2018.

【2】“Deep learning face representation from predicting 10,000 classes”， Sun, Yi, Xiaogang Wang, and Xiaoou Tang, CVPR 2014.

Github代碼：

https://github.com/yl-1993/hfsoftmax