• 您的位置:首頁 >聚焦 >

    NeurIPS 2021 Spotlight | 準確、快速、內存經濟,新框架MEST實現邊...

    2022-04-05 17:06:01    來源:程序員客棧

    點擊上方“邁微AI研習社”,選擇“星標★”公眾號

    重磅干貨,第一時間送達

    美國東北大學王言治教授、林雪教授研究組與威廉瑪麗學院任彬教授研究組共同提出了一種稀疏訓練新框架MEST,有望在邊緣設備上實現準確、快速以及內存經濟的稀疏訓練。

    在剪枝技術被成功應用于神經網絡的壓縮和加速之后,稀疏訓練在近年來受到了越來越多研究者的關注,即如何從零開始直接訓練一個高質量的稀疏神經網絡。稀疏訓練旨在有效降低神經網絡訓練過程中的計算和存儲開銷,從而加速訓練過程,為在資源有限的邊緣設備上的神經網絡訓練提供了更多可能性。

    多數現有稀疏訓練方法著力于設計更好的稀疏訓練算法來追求更高的網絡稀疏度同時保持高準確率。然而,稀疏訓練的關鍵,即稀疏訓練能否帶來實際的訓練加速以及計算和存儲資源的節省,卻往往被忽視了。

    為此,美國東北大學王言治教授、林雪教授研究組與威廉瑪麗學院任彬教授研究組共同提出了 MEST 稀疏訓練框架,有望在邊緣設備上實現準確、快速以及內存經濟的稀疏訓練。該論文?[1] 已被 NeurIPS 2021 會議收錄為 spotlight 論文。

    論文鏈接:https://arxiv.org/abs/2110.14032

    什么是稀疏訓練?

    我們知道通常的神經網絡剪枝(pruning)方法是從一個訓練好的密集模型(dense model)出發,通過應用剪枝算法去除模型中的冗余權重,來達到降低模型參數量和計算量的目的。而這個密集模型的預訓練過程依然會耗費大量的計算資源。所以,近年來有越來越多的研究試圖在預訓練過程中直接使用稀疏模型(sparse model)進行訓練,從而可以一步到位的獲得最終的稀疏模型,我們可以管這個直接訓練稀疏模型的過程叫做稀疏訓練(sparse training)。稀疏訓練可以有效的降低訓練過程的計算量,并有望降低設備的存儲開銷以及加速訓練過程。

    稀疏訓練大致可以分為兩個類別,即靜態稀疏訓練(static sparse training)和動態稀疏訓練(dynamic sparse training)。

    靜態稀疏訓練通常會在訓練的初始階段通過使用相應的算法來確定稀疏網絡的結構。之后,在整個的稀疏訓練過程中始終保持相同的稀疏網絡結構。靜態稀疏訓練的代表性工作有 SNIP [1] 與 GraSP [2] 等。此外,彩票假說(Lottery Ticket Hypothesis)[3] 中直接使用稀疏子網絡進行訓練的過程也可以被認為是靜態稀疏訓練。

    動態稀疏訓練一般會在訓練的初始階段隨機或簡單選擇一個稀疏網絡的結構作為起始,并在隨后的稀疏訓練過程通過中不斷變化稀疏網絡的結構,來達到搜尋更好稀疏結構的目的。動態稀疏訓練的代表性工作有 DSR [5] 和 RigL [6] 等。相較于靜態稀疏訓練,動態稀疏訓練往往更容易在高稀疏度下獲得更好的稀疏模型準確率。

    現有的先進稀疏訓練方法中存在的問題

    1. 非結構性稀疏方案 (Unstructured Sparsity Scheme)

    絕大多數現有的先進的稀疏訓練方法通常都會采用非結構性稀疏方案。這種稀疏方案允許任意位置的權重被移除。因此,稀疏模型的結構具有高靈活性,可以更好地在高稀疏度下維持稀疏模型的準確率。但是,與在剪枝技術中使用非結構性剪枝一樣,由于非結構性稀疏模型中非零權重分布的不規則性,導致了硬件并行計算的不友好,比如不良的數據局部性(data locality)、負載不平衡(load imbalance)、線程分歧(thread divergence)和繁重的控制流指示(heavy control-flow instructions)等。這最終使得非結構性稀疏模型即便在高稀疏度下也很難實現可觀的計算速度提升。

    2. 內存不經濟

    為了能夠找到更好的稀疏網絡結構,在一些先進的稀疏訓練方法中會使用到密集模型的信息,比如在靜態稀疏訓練中,一些方法會先進行少量次數迭代的密集模型訓練來在初始化模型中找到想要的稀疏結構。同樣,在動態稀疏訓練中,一些方法會頻繁地使用密集反向傳播來計算出密集的權重梯度,從而指引稀疏結構的動態選擇。雖然這些方法在大部分的稀疏訓練過程中維持了稀疏模型的存儲與計算,但是只要訓練過程中涉及到密集模型的計算,就會極大地增加計算與存儲資源的峰值使用量。這使得這些稀疏訓練方法難以真正的在資源受限的邊緣設備上進行端到端訓練。

    MEST 框架設計

    為了解決以上提到的現有稀疏訓練方法中存在的問題,我們認為在訓練過程應當避免使用任何的密集模型的信息,保持整個訓練過程的稀疏性。并且稀疏訓練應當使用更加硬件友好的稀疏方案,從而達到真正加速訓練的目的。此外,現有的稀疏訓練方法多著眼于更好的算法設計,在保持稀疏模型準確率的情況下盡可能的提升模型稀疏率(sparsity ratio),從而進一步降低訓練開銷。然而,我們認為這并不是降低訓練開銷的唯一途徑。提高對于訓練數據的使用效率可以從另一個角度來進一步降低訓練開銷和加速訓練過程,而這一角度被先前的稀疏訓練研究忽視了。

    1.MEST 稀疏訓練方法

    在我們的 MEST 稀疏訓練框架中我們首先強調的是算法的內存經濟性(Memory-economic)。我們采用傳統的動態稀疏訓練思想,通過在訓練過程中不斷地對稀疏網絡的結構進行變異(mutation)來最終獲得高準確率的稀疏網絡。具體來說,為了使全部訓練過程不涉及密集模型的存儲和計算,我們在模型初始化階段直接使用隨機生成的索引(indices)以及 CSR 形式來直接存儲隨機初始化的稀疏模型。稀疏模型結構的突變過程可以看作是模型剪枝(prune)和模型生長(grow)的結合。首先,我們從稀疏模型(非零權重)中去除具有相對較低重要性的權重,再從不在稀疏模型中的權重中隨機選擇相同數量的權重加回到稀疏模型中,從而完成稀疏模型結構的突變,如圖 1 中的 MEST(vanilla)方法所示。為了保持模型的稀疏性,我們直接將被去除權重的權重值設為零,并直接改變索引值,來實現模型結構的突變。

    我們還提出一種彈性突變(Elastic Mutation)策略,簡稱 MEST+EM。我們的彈性突變策略意在整個訓練過程中逐漸降低網絡結構的突變率,如圖 1 中 MEST+EM 方法所示。這樣即保證了在網絡訓練過程中擁有足夠大的搜索范圍,同時在訓練后期幫助稀疏網絡結構更好的收斂,增加訓練過程的穩定性。

    此外,如果在應用場景中,內存占用可能是一個軟約束,我們提出了一種稀疏訓練的增強方法,稱為軟內存約束的彈性突變(Soft-bounded Elastic Mutation),簡稱 MEST+EM&S。

    與 MEST+EM 方法不同的是,我們的 MEST+EM&S 允許將新增長的權重加入到現有的權重中,然后進行訓練,隨后從包括新增長的權重在內的權重中選擇重要性較低的權重進行移除,如圖 1 中 MEST+EM&S 方法所示。這可以避免在模型突變中強制移除比新增長的權重更正要的現有權重。所以我們的 MEST+EM&S 方法也可以被看做是在突變過程中增加了一個 "撤銷" 機制。需要注意的是,雖然使用軟內存約束會額外增加少量的內存和計算開銷,但是整個稀疏訓練過程仍處在高稀疏度下,且訓練結束時仍然達到目標稀疏率。

    圖 1. MEST 框架中的稀疏訓練方法

    2. 探究不同的稀疏方案

    在 MEST 框架中,我們研究了不同稀疏方案在稀疏訓練中的速度表現以及使用 MEST 方法的準確度。如圖 2 所示,他們分別為(a)非結構性稀疏(unstructured sparsity)、(b)結構性稀疏(structured sparsity)、(c)基于區塊的細粒度結構性稀疏(block-based sparsity)以及(d)基于模式的細粒度結構稀疏(pattern-based sparsity)。

    圖 2. 不同稀疏方案

    如前面提到的,非結構性稀疏被廣泛用于之前的稀疏訓練研究。作為一種細粒度的稀疏形式,其具有高靈活性,可以更好地維持網絡的準確率,但是很難達到可觀的加速。我們在 MEST 中也對非結構性稀疏進行測試,以便從算法角度與之前的工作進行公平比較。同時,在對比不同稀疏方案的準確率表現時將其視為準確率上限。

    我們還測試了結構性稀疏方案,即移除整個卷積濾波器(filter)或卷積通道(channel)。這種粗粒度的稀疏方案非常硬件友好,可以實現較高的加速性能,但是往往模型準確率會極大地降低。

    在基于區塊的細粒度結構性稀疏方案中,我們將權重矩陣分為若干個等大小的塊,然后在每一個塊中選擇整行或整列進行權重的稀疏化,在不同的塊中稀疏行列位置和數量可以不同,這樣提高了稀疏的靈活性,從而保持更好的模型準確率。由于受到硬件并行度的限制,即便將權重矩陣進行分割,每一個塊依然可以很好地滿足高計算并行度,所以這種相對細粒度的結構性稀疏可以達到較高的加速效果。

    基于模式的細粒度結構稀疏結合了卷積核模式稀疏和連通性稀疏。其中,卷積核模式稀疏從每一個濾波器的卷積核中將給定數量的權重設為零,使得剩余非零權重形成特殊的模式。圖 2(d)所示為在 3x3 卷積核中使用 4 非零權重模式。此外,連通性稀疏擇將整個卷積核移除,從而進一步提升基于模式的細粒度結構稀疏方案的整體稀疏度。

    這兩種稀疏方案通過細粒度的稀疏性在一定程度上保證了稀疏結構的靈活性,從而有效的提高模型準確率。同時,他們又滿足一定的結構約束,使得這種稀疏結構可以在編譯器優化的幫助下實現較高性能的加速。我們在稀疏訓練過程中,分別將不同的稀疏方案應用在我們提出的 MEST+EM 和 MEST+EM&S 方法中來探究不同稀疏方案在準確率和加速性能等方面的表現。

    3. 數據效率在稀疏訓練中的應用

    除了不斷增加模型的稀疏度,在稀疏訓練中使用更高效的訓練數據同樣可能有效的加速訓練過程。經過證明,對于網絡訓練來說,每個訓練樣本的學習難度是不同的,同時其能提供給網絡訓練的信息量也是不同的 [7]。一些訓練樣本在訓練的早期階段就可以被網絡學會(正確識別),其中一些樣本在被網絡學會之后就不會再被錯誤識別,換言之,他們永遠不會被網絡遺忘。這里我們定義當一個訓練樣本在被網絡正確識別后又被錯誤分類為一次 “遺忘事件”。

    如圖 3 所示,左手邊展示了兩個不同的青蛙訓練樣本。上面的訓練樣本在整個訓練過程中被遺忘了 0 次,所以我們也稱它為不被遺忘的樣本 (unforgettable example),而下面的樣本則被遺忘了 35 次。同時,通過觀察對比圖 3 右側的具有高遺忘次數和低遺忘次數的訓練樣本,我們可以看出一般具有較低遺忘次數的樣本通常比較簡單,既樣本具有鮮明的特征且背景較為簡單。因此,之前的工作使用遺忘次數作為衡量訓練樣本難易程度的標準,并且發現容易的樣本為網絡訓練提供的信息量較少,如果將不被遺忘的樣本從訓練數據集中移除,訓練出的網絡依然可以達到與完整數據集相當的準確率,但是訓練速度因訓練數據集的減小而加快。

    圖 3. 不同訓練樣本具有不同遺忘次數

    然而,對于數據效率在稀疏訓練場景下的探索仍然是缺失的。由于稀疏模型的容量較小,且模型結構是動態變化的,是否可以利用數據效率來進一步加速稀疏訓練仍是未知數。因此,我們希望研究稀疏率、稀疏方案和變異機制對可移除訓練樣本數量的影響,然后討論利用數據效率加速稀疏訓練的可能性。

    之前的工作使用了整個訓練過程來進行訓練樣本難易度的統計工作,而我們希望通過引入數據效率來進行訓練加速,所以我們提出了一種兩階段數據高效訓練方法。我們使用與常規訓練相同的訓練周期數,并將訓練過程分為數據統計階段與使用部分數據集的高效訓練階段,如圖 4 所示。我們在 MEST 框架中將這種數據高效方法與我們提出的稀疏訓練方法相結合,既在兩個階段中均使用 MEST+EM 或 MEST+EM&S 稀疏訓練。

    圖 4. 兩階段數據高效訓練方法

    實驗結果

    1. 稀疏算法對于數據效率的探究

    首先,我們對在動態稀疏訓練中引入數據高效訓練進行了多方面的探究。

    模型稀疏度、稀疏算法對不被遺忘的樣本數量的影響:

    我們首先測試了不同的稀疏訓練方法對于不被遺忘的樣本數量的影響,結果如圖 5(a)所示,隨著模型稀疏度的增加,不被遺忘的樣本數量在減少,并且與模型的準確率呈現正相關。原因是在高稀疏度下,模型的泛化性能下降,使得一些簡單的樣本更難記住。此外,我們還觀察到,更好的稀疏度訓練方法(如 MEST+EM&S)會增加不被遺忘的樣本數量,這表明在不影響模型準確率的情況下,有機會去除更多的訓練樣本,從而獲得更高的訓練加速。

    圖 5. 數據高效訓練的探究。(a) 在不同的稀疏訓練算法和稀疏率下,整個稀疏訓練過程后不被遺忘的樣本數量;(b) 在稀疏率為 90% 和 95% 下,移除被遺忘次數分別為 0 和 1 的樣本,所使用的樣本數據統計階段(第一階段)的周期數及其相應的模型最終準確率狀態;(c) 每個訓練周期中增加的遺忘例子的數量(在訓練周期 50 到 80 之間,稀疏結構突變頻率為 5 個周期)。結果為 ResNet-32 網絡在 CIFAR-10 數據集上的表現。

    兩階段數據高效訓練中數據統計所需周期數:

    為了在一個常規的訓練長度中實現數據統計以及數據高效訓練,我們需要在第一個階段統計樣本的難易程度,以便在第二階段將簡單樣本從數據集中移除。由于數據高效訓練僅在第二階段中擁有加速效果,所以我們希望使用盡可能少的訓練周期完成第一階段的數據統計工作。然而,在訓練初期,由于網絡沒有被良好的訓練,可能確保所統計出的難易程度的準確性,從而會影響最終數據高效訓練模型的準確率。因此,我們觀察了不同第一階段周期數與最終模型準確率的關系。

    如圖 5(b)所示,紅色代表在相應的第一階段周期數下,最終模型的準確率低于使用完整數據集的準確率,而黃色代表使用相應的第一階段周期數可以達到與使用完整數據集相同的準確率。同時,我們研究了不同稀疏訓練方法、稀疏結構方案、模型稀疏率、以及移除的樣本數量,所需要的第一階段的訓練周期。通過對比結果,我們可以在保持最終準確率的情況下,選擇最少的第一階段周期數量來獲得最高的訓練加速。

    模型結構變異是否會導致遺忘?

    不斷進行的模型結構變異作為動態稀疏訓練的一大特征,是否會導致模型對以學會樣本的嚴重遺忘是我們在稀疏訓練中引入數據高效訓練關注的問題。為此,我們觀察了每次模型結構突變前后的相鄰周期中被遺忘樣本數量的變化。

    圖 5(c)顯示了不同稀疏訓練方法從第 50 到第 80 個訓練周期,每個訓練周期中被遺忘的樣本數量的增多值,也就是兩個連續的訓練周期中被遺忘樣本的差異。我們可能會直觀的認為,模型結構的頻繁變異會導致明顯的遺忘現象發生。但是,令我們驚訝的是,被遺忘的樣本數量在模型結構變異的訓練周期與非變異的訓練周期并無明顯差異,這表明模型結構的變異并不會加劇遺忘現象。這是因為每次變異的權重均為相對最不重要的權重,因此對模型的性能影響很小。

    2.MEST 框架的準確率與加速性能表現

    訓練方法有效性檢驗:

    首先,我們對 MEST 框架中所提出的稀疏訓練方法的進行了準確率方面的測試,并與具有代表性的稀疏訓練方法進行了比較。為了進行更公平的比較,我們 MEST 在這里也使用了跟其他方法一樣的非結構性稀疏方案。

    我們在 CIFAR-10 和 CIFAR-100 數據集上對 MEST 框架中使用的 MEST+EM 和 MEST+EM&S 方法進行了稀疏訓練模型準確率的測試。我們對比了具有代表性的靜態稀疏訓練方法以及不需要涉及密集信息的動態稀疏訓練方法。可以看出,我們的 MEST+EM 方法在 CIFAR-10 上達到與參考方法相似或更高的準確率,在 CIFAR-100 上我們的 MEST+EM 方法具有更高的準確率優勢。此外,我們的 MEST+EM&S 方法可以顯著的進一步提高稀疏訓練準確率。

    表 1. 在 CIFAR-10 和 CIFAR-100 上使用 ResNet-32 與 SOTA 工作的準確率比較。

    雖然現今在端設備上使用龐大的 ImageNet 數據集進行訓練并不實際,但是為了進一步驗證我們的稀疏訓練方法的有效性,我們同樣在 ImageNet 數據集上進行了準確率的測試。

    如表 2 所示,與先進的 RigL 方法相比,我們的 MEST+EM 和 MEST+EM&S 方法可以在相同的訓練計算量(FLOPs)下,得到更高的稀疏訓練準確率,且我們的方法在訓練過程中不涉及任何密集模型信息,因此更加的硬件友好。此外,我們還測試了引入數據高效訓練的效果(MEST+EM&S+DE)。可以看到,我們的數據高效訓練方法可以在保持準確率的同時,進一步減少訓練開銷,從而實現訓練加速。

    表 2. 在 ImageNet 上使用 ResNet-50 與 SOTA 工作的準確率比較。

    不同稀疏方案的準確率與加速性能表現:

    我們在這篇工作中探究了不同稀疏方案在稀疏訓練中的準確率與加速性能的表現。如圖 6 (a)和(b)所示,粗粒度的結構性稀疏方案會產生較大的準確率下降,而兩種細粒度的結構性稀疏方案(基于區塊的細粒度結構性稀疏方案與基于模式的細粒度結構稀疏方案)則可以在高達 90% 的稀疏度下達到與非結構性稀疏方案相似的準確率。圖 6(c)展示了在 90% 稀疏度下,不同方法和稀疏方案在準確率、訓練加速率以及相對存儲資源的對比。我們可以看到,即便在如此高的稀疏度下,所有的非結構性稀疏方案都無法達到可觀的加速性能,而細粒度的稀疏方案可以在保持較高準確率的情況下,大大提升加速性能。同時,我們的數據高效訓練可以進一步有效地提升訓練的加速。

    圖 6. 不同稀疏方案在不同稀疏度下的模型精度對比。

    結語

    由于現有的稀疏訓練工作普遍忽視了稀疏方案對實際加速性能的影響,以及對稀疏訓練方法是否內存經濟且硬件友好的顧慮,我們提出了 MEST 稀疏訓練框架,通過簡單且有效的稀疏訓練方式和適當的稀疏方案,實現準確、快速以及對邊緣設備友好的稀疏訓練。我們研究了不同稀疏方案對于稀疏訓練準確率以及加速性能的影響。同時,我們提出引入數據高效訓練方法來作為模型稀疏度的補充。我們探究了數據高效訓練在稀疏訓練場景下的特性,并提出一種兩階段數據高效訓練方法,并在不影響準確率的情況下進一步加速了稀疏訓練過程。我們也希望這種新的稀疏訓練加速思路可以被更廣泛的研究。

    參考文獻:

    [1] Yuan, Geng, et al. "MEST: Accurate and Fast Memory-Economic Sparse Training Framework on the Edge." NeurIPS 2021.

    [2] Lee, Namhoon, et al.? "Snip: Single-shot network pruning based on connection sensitivity." ICLR 2019.

    [3] Wang, Chaoqi, et al. "Picking winning tickets before training by preserving gradient flow." ICLR 2020.

    [4] Frankle, Jonathan, et al. "The lottery ticket hypothesis: Finding sparse, trainable neural networks." ICLR 2019.

    [5] Mostafa, Hesham, et al. "Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization." ICML 2019.?

    [6] Evci, Utku, et al. "Rigging the lottery: Making all tickets winners." ICML 2020.

    [7] Toneva, Mariya, et al. "An empirical study of example forgetting during deep neural network learning." ICLR 2019.

    更多細節可參考論文原文,更多精彩內容請關注邁微AI研習社,每天晚上七點不見不散!

    ??THE END?

    投稿或尋求報道微信:MaiweiE_com

    GitHub中文開源項目《計算機視覺實戰演練:算法與應用》,“免費”“全面“”前沿”,以實戰為主,編寫詳細的文檔、可在線運行的notebook和源代碼。

    項目地址?https://github.com/Charmve/computer-vision-in-action

    項目主頁?https://charmve.github.io/L0CV-web/

    關鍵詞: 訓練樣本 模型結構

    相關閱讀

    BB电子