[AAAI 18] Accelerated Training for Massive Classification via Dynamic Class Selection

Xingcheng Zhang, Lei Yang, Junjie Yan, Dahua Lin

from CUHK & SenseTime

Motivation

這篇文章研究當分類器分類個數非常大的時候，如何高效訓練分類器的問題。在網路的輸出層，Softmax分類器會對每一個類產生一個輸出。因此輸出層的引數量、佔用的空間和計算量是和分類個數呈正相關的。當在某些實際應用場景下，如人臉識別、自然語言處理等，其分類個數（如人的身份數量、單詞種類數）會非常大，以至於使網路的輸出層的引數多到無法儲存在顯示卡中，而且無法承受如此量級的計算。這篇文章針對這個問題提出一點觀察結果和一個解決該問題的方案。其分析和觀察的結果是：

1. 對於一個輸入樣本，其輸出概率只集中在少量的類別中，可以稱這些類別為活躍類(active class)；
2. 在同一迭代過程中，回傳梯度主要受這些活躍類影響，其他類別的貢獻很小。

因此一個解決方案便是前向和反向過程中只計算和這些活躍類相關的東西。本文提出的解決方案便是如何動態搜尋得到這些活躍類到底是哪些。下面就詳細介紹這篇文章的分析和觀察的細節，以及解決方案的內容。

An Empirical Study of Softmax

作者通過實驗來驗證其關於Softmax概率分佈的分析結論。為此提出了兩個新的概念，

1. top K累積概率：將輸出的概率向量裡最大的K個值相加(k=1,100,1000)，其結果往往接近於1；
2. 歸一化的top K梯度累積能量：雖然名字很拗口，但是其實質是輸出層所有梯度與top K類產生的梯度的餘弦相似度，用來衡量top K類產生的梯度佔所有梯度的比例。

作者在MS-Celeb-1M資料庫上訓練人臉識別網路，下圖為統計結果。可見活躍類確實在產生的概率(圖中(a))和回傳的梯度(圖中(b))上都佔主導作用。

Selective Softmax

在上個圖中，即便K=1000，其top K類仍佔總類別數很小的比例（如MS-Celeb-1M中只佔1.5%）。因此在訓練過程中能針對不同的樣本找到對應的活躍類是降低計算和儲存代價的關鍵。

對一個輸入特徵x，輸出層的權重矩陣為W。尋找top K類即是找到W中和這些活躍類關聯的k個列向量。因為每個活躍類的概率由

Hashing Forest

為了快速實現聚類，作者使用了構建雜湊樹的方法，即不停將各個列向量二分類，構建兩個葉子節點。之後在葉子節點上遞迴二分類，繼續二分葉子節點，直至每個葉子節點的向量個數小於某個值。

每次對W的列向量二分類的操作如下：

1. 隨機抽取兩個列向量wi,wj$w_i,w_j$，求兩個向量的平均作為分類基準h=wi+wj2$h=\frac{w_i+w_j}{2}$；
2. 將每個向量同該分類基準向量作點積wTih$w_i^{T}h$，依據點積結果大於0與否將其放到左、右子節點中；
3. 當待操作的葉子節點向量個數小於某個值的時候，停止操作。

因為構建每個雜湊樹存在隨機因素，實際使用中可以並行構建多個雜湊樹，組成森林以提高可靠性。

當給定輸入特徵x，查詢top K類對應的權重向量時，搜尋過程從雜湊樹的根節點開始。每次搜尋的時候，x都會和分類基準向量作點積，同樣根據結果大於0與否決定選擇左節點還是右節點。該過程直到搜尋抵達某個葉子節點位置。

Adaptive Allocation

在訓練過程中，所搜活躍類的過程中有三個變數需要進行控制：

1. 每次搜尋活躍類的個數M；
2. 構建雜湊樹的個數L；
3. 重新構建雜湊森林的間隔迭代次數T。

其中M和L決定了搜尋活躍類的可靠性，T的存在是因為權重W在訓練過程中也在不停變化。三個變數的選取也決定了演算法的速度。作者認為隨著網路迭代，模型逐漸趨穩，因此當迭代次數增加時：

1. M要滿足top M類的概率和大於一個門檻，該門檻線性增加；
2. T線性減少；
3. L線性增加。

對於一個訓練batch，其選取的活躍類是各個輸入選取得到的活躍類的整合。

Experiment

作者在幾個不同的大型人臉資料集上進行了訓練，來驗證本文提出方案的有效性。
同時也展示了幾個重要引數的敏感性。舉其中一例，作者在MS-Celeb-1M和Megaface合併的資料集上訓練ResNet-101模型進行人臉識別，總類別數達到了75萬。其對比結果如下圖所示，本文提出的方法較傳統訓練方法和隨機選取“偽”活躍類的方法相比都很大優勢：