一種相對直接使用的distillation的變體方法

在提出備受矚目的“膠囊網(wǎng)絡(luò)”（Capsule networks）之后，深度學(xué)習(xí)領(lǐng)域的大牛、多倫多大學(xué)計算機(jī)科學(xué)教授Geoffrey Hinton近年在distillation這一想法做了一些前沿工作，包括Distill the Knowledge in a Neural Network等。今天我們介紹的是Hinton作為作者之一，谷歌大腦、DeepMind等的研究人員提交的distillation的更進(jìn)一步工作：通過online distillation進(jìn)行大規(guī)模分布式神經(jīng)網(wǎng)絡(luò)訓(xùn)練。該工作提出了Codistillation的概念，通過大規(guī)模實(shí)驗(yàn)，發(fā)現(xiàn)codistillation方法提高了準(zhǔn)確性并加快了訓(xùn)練速度，并且易于在實(shí)踐中使用。

與幾乎任何基本模型配合時，諸如集成（ensembling）和蒸餾（distillation）等技術(shù)都可以提升模型的質(zhì)量。但是，由于增加了測試時間成本（對于ensembling）和訓(xùn)練pipeline的復(fù)雜性（對于distillation），這些技術(shù)在工業(yè)環(huán)境中使用具有挑戰(zhàn)性。

來自谷歌、谷歌大腦、DeepMind的研究人員，包括Geoffrey Hinton等人，在他們提交給ICLR 2018的論文“Large scale distributed neural network training through online distillation”中，探討了一種相對直接使用的distillation的變體方法，該方法不需要復(fù)雜的多級設(shè)置或非常多的新超參數(shù)。

研究者提出的第一個主張是：在線蒸餾（online distillation）使我們能夠使用額外的并行性來適應(yīng)非常大的數(shù)據(jù)集，并且速度提高一倍。更重要的是，即使我們已經(jīng)達(dá)到額外的并行性對同步或異步隨機(jī)梯度下降沒有好處的程度，我們?nèi)匀豢梢约涌煊?xùn)練速度。在不相交的數(shù)據(jù)子集上訓(xùn)練的兩個神經(jīng)網(wǎng)絡(luò)可以通過鼓勵每個模型同意另一個模型的預(yù)測來共享知識。這些預(yù)測可能來自另一個模型的舊版本，因此可以使用很少被傳輸?shù)臋?quán)重來安全地計算它們。

論文提出的第二個主張是：online distillation是一種成本效益高的方法，可以使模型的精確預(yù)測更具可重復(fù)性。研究者通過在Criteo Display Ad Challenge數(shù)據(jù)集，ImageNet和用于神經(jīng)語言建模的最大數(shù)據(jù)集（包含6×1011個tokens）上進(jìn)行實(shí)驗(yàn)，支持了提出的這些主張。

Codistillation：優(yōu)于分布式SGD

對于大規(guī)模的、具有商業(yè)價值的神經(jīng)網(wǎng)絡(luò)訓(xùn)練問題，如果訓(xùn)練時間能夠大幅加快，或最終模型的質(zhì)量能夠大幅提高，從業(yè)者會愿意投入更多的機(jī)器用于訓(xùn)練。目前，分布式隨機(jī)梯度下降（SGD），包括其同步和異步形式（Chen et al.，2016）是在多個互聯(lián)機(jī)器上進(jìn)行大規(guī)模神經(jīng)網(wǎng)絡(luò)訓(xùn)練的主要算法。但是，隨著機(jī)器數(shù)量的增加，訓(xùn)練一個高質(zhì)量模型所需時間的改善程度會降低，直到繼續(xù)增加機(jī)器卻無法進(jìn)一步縮短訓(xùn)練時間。基礎(chǔ)架構(gòu)的限制，以及優(yōu)化上的障礙，一起限制了分布式 minibatch SGD的可擴(kuò)展性。

分布式SGD的精確可擴(kuò)展性限制將取決于算法的實(shí)現(xiàn)細(xì)節(jié)，基礎(chǔ)架構(gòu)的具體情況以及硬件的能力，但根據(jù)我們的經(jīng)驗(yàn)，在實(shí)際設(shè)置中，在超過100臺GPU workers上進(jìn)行有效擴(kuò)展可能非常困難。沒有任何訓(xùn)練神經(jīng)網(wǎng)絡(luò)的算法可以無限擴(kuò)展，但即使擴(kuò)展得比分布式SGD的限制多一點(diǎn)，也是非常有價值的。

一旦我們達(dá)到了向分布式SGD添加worker的限制，我們就可以使用額外的機(jī)器來訓(xùn)練模型的另一個副本，并創(chuàng)建一個集成（ensemble）以提高準(zhǔn)確性（或通過以更少的step訓(xùn)練ensemble中的成員來提高訓(xùn)練時的精度）。ensemble能夠做出更穩(wěn)定和可重復(fù)的預(yù)測，這在實(shí)際應(yīng)用中很有用。但是，ensemble增加了測試時的成本，可能會影響延遲或其他成本限制。

為了在不增加測試時間成本的情況下獲得幾乎與ensemble相同的好處，我們可以對一個n-way模型的ensemble進(jìn)行蒸餾（distill），得到一個單一模型，這包括兩個階段：首先我們使用nM機(jī)器來訓(xùn)練分布式SGD的n-way ensemble，然后使用M機(jī)器來訓(xùn)練student網(wǎng)絡(luò)，以模擬n-way ensemble。通過在訓(xùn)練過程中增加另一個階段并使用更多機(jī)器，distill一般會增加訓(xùn)練時間和復(fù)雜性，以換取接近更大的 teacher ensemble 模型的質(zhì)量改進(jìn)。

我們認(rèn)為，從時間和pipeline復(fù)雜性兩方面來看，額外的訓(xùn)練成本阻礙了從業(yè)者使用ensemble distillation，盡管這種方法基本上總是能夠改善結(jié)果。在這項(xiàng)新的工作中，我們描述了一個簡單的distillation的在線變體，我們稱之為codistillation。Codistillation通過向第i個模型的損失函數(shù)添加一個項(xiàng)來匹配其他模型的平均預(yù)測值，可以并行訓(xùn)練n個模型的副本。

通過大規(guī)模實(shí)驗(yàn)，我們發(fā)現(xiàn)，與分布式SGD相比，通過允許有效利用更多計算資源，codistillation提高了準(zhǔn)確性并加快了訓(xùn)練速度，甚至加速效果超過了給SGD方法添加更多worker。具體來說，codistillation提供了在不增加訓(xùn)練時間的情況下distill一個模型ensemble的好處。與Multi-phase的distillation訓(xùn)練過程相比，Codistillation在實(shí)踐中使用也相當(dāng)簡單。

這項(xiàng)工作的主要貢獻(xiàn)是codistillation的大規(guī)模實(shí)驗(yàn)驗(yàn)證。另一個貢獻(xiàn)是，我們探索了不同的設(shè)計選擇和codistillation的實(shí)現(xiàn)考慮因素，提出了實(shí)用的建議。

總的來說，我們認(rèn)為在實(shí)踐中，codistillation比精心調(diào)參的offline distillation得到的質(zhì)量提升是次要的，更有趣的研究方向是將codistillation作為一種分布式訓(xùn)練算法來研究。

在這篇論文中，我們使用codistillation來指代執(zhí)行的distillation：

所有模型使用相同的架構(gòu);

使用相同的數(shù)據(jù)集來訓(xùn)練所有模型;

在任何模型完全收斂之前使用訓(xùn)練期間的distillation loss。

codistillation算法

實(shí)驗(yàn)與結(jié)果

為了研究分布式訓(xùn)練的可擴(kuò)展性，我們需要一個代表重要的大規(guī)模神經(jīng)網(wǎng)絡(luò)訓(xùn)練問題的任務(wù)。神經(jīng)語言建模是一個理想的測試平臺，因?yàn)榫W(wǎng)絡(luò)上有大量的文本，并且也因?yàn)樯窠?jīng)語言模型的訓(xùn)練成本可能非常高。神經(jīng)語言模型是實(shí)現(xiàn)分布式SGD（如機(jī)器翻譯和語音識別）常用重要問題的代表，但是語言建模更容易評估和使用更簡單的管線。為了盡可能清晰地提高潛在的可擴(kuò)展性，我們選擇了一個足夠大的數(shù)據(jù)集，它完全不可能通過現(xiàn)有的SGD并行化策略來訓(xùn)練一個表達(dá)模型。為了證實(shí)我們的結(jié)果并不是特定于語言建模的某些特性，我們也驗(yàn)證了我們在ImageNet (Russakovsky et al.， 2015)上的一些大規(guī)模的codistillation結(jié)果。為了證明在減少預(yù)測過程中進(jìn)行編碼的好處，并研究算法的其他特性，我們可以使用更小的更便宜的實(shí)驗(yàn)，但在研究可擴(kuò)展性時，真正達(dá)到分布式SGD的極限是很重要的。

1. 達(dá)到分布式SGD的極限，用于在Common Crawl上訓(xùn)練RNN

在我們的第一組實(shí)驗(yàn)中，我們的目標(biāo)是大致確定在我們的Common Crawl神經(jīng)語言模型設(shè)置中可以有效使用SGD的GPU的最大數(shù)量。由于我們的數(shù)據(jù)集比英文維基百科數(shù)據(jù)量大兩個數(shù)量級，因此不用擔(dān)心重新訪問數(shù)據(jù)，即使在相對較大規(guī)模的實(shí)驗(yàn)中，這也會使獨(dú)立副本更加相似。我們嘗試使用32和128個workers的異步SGD，在必要的情況下，通過增加參數(shù)服務(wù)器的數(shù)量來分配權(quán)重，以確保訓(xùn)練速度被GPU計算時間瓶頸。我們發(fā)現(xiàn)很難保持訓(xùn)練的穩(wěn)定性，并防止RNNs與大量GPU的異步SGD出現(xiàn)差異。我們嘗試了一些workers提升計劃和不同的學(xué)習(xí)速率，但最終決定將重點(diǎn)放在同步算法上，使我們的結(jié)果更少地依賴于我們的基礎(chǔ)結(jié)構(gòu)和實(shí)現(xiàn)的特定特性。梯度的穩(wěn)定性很難獨(dú)立于特定的條件下進(jìn)行分析，而實(shí)現(xiàn)和基礎(chǔ)結(jié)構(gòu)的差異更容易抽象為同步SGD。雖然可能需要更多的努力使異步工作良好，但陳舊的漸變對學(xué)習(xí)進(jìn)度的削弱作用是一個眾所周知的問題，Mitliagkas等人(2016)認(rèn)為，異步可以有效地增加動量，這也是為什么它容易分化的原因。在初步實(shí)驗(yàn)中，從codistillation中獲得的收益似乎與選擇異步或同步SGD作為基本算法無關(guān)。

可同時用于同步SGD的最大GPU的數(shù)量取決于基礎(chǔ)架構(gòu)限制、尾延遲和批處理大小的影響。完全同步的SGD相當(dāng)于批量大得多的單機(jī)算法。增加有效的批量大小可減少梯度估計中的噪聲，從而允許更大的步長，并有望實(shí)現(xiàn)更高質(zhì)量的更新，從而實(shí)現(xiàn)更快的收斂速度。鑒于有效的無限訓(xùn)練數(shù)據(jù)（即使使用256個GPU，我們也不訪問所有常見爬網(wǎng)訓(xùn)練數(shù)據(jù)），我們直觀地預(yù)計會增加有效批量，最終增加步驟時間。我們采用完全同步的SGD在Common Crawl上訓(xùn)練了語言模型，并使用128和32,64,128和256GPU的批處理大小。因此，有效的批量大小范圍為4096至32768.一般來說，我們應(yīng)該期望需要增加學(xué)習(xí)率，因?yàn)槲覀冊黾恿擞行У呐看笮?，所以對于每個GPU我們嘗試的學(xué)習(xí)率為0.1,0.2和0.4。對于32和64GPU，0.1表現(xiàn)最好，因?yàn)樵鹊娜N學(xué)習(xí)率都沒有對256GPU表現(xiàn)良好，我們還嘗試了0.3的額外中級學(xué)習(xí)率，這是256的最佳表現(xiàn)學(xué)習(xí)率。

圖1a將驗(yàn)證誤差作為我們嘗試過的不同數(shù)量的GPU的全局步驟的函數(shù)，使用了每個GPU數(shù)量的最佳學(xué)習(xí)速率。增加數(shù)量(因此有效的批處理大小)減少了達(dá)到最佳驗(yàn)證錯誤所需的步驟數(shù)，直到128GPU，在這一點(diǎn)上沒有額外的改進(jìn)。圖1b對相同數(shù)量的同步GPU的驗(yàn)證誤差與Wall time進(jìn)行對比。在這些特殊的實(shí)驗(yàn)中，與128個GPU同步的SGD是訓(xùn)練時間和最終準(zhǔn)確度方面最強(qiáng)的基線。因此，我們將其余實(shí)驗(yàn)集中在與128位同步SGD進(jìn)行比較，并研究使用同步SGD作為子程序的codistillation，但它也適用于異步算法。

2. 帶同步SGD的CODISTILLATION

對于Common Crawl上的語言建模，具有128GPU的同步SGD實(shí)現(xiàn)了標(biāo)準(zhǔn)分布式訓(xùn)練的最佳結(jié)果，至少是我們嘗試過的配置，并且我們無法使用256個GPU來提高訓(xùn)練時間。雖然額外的GPU似乎不能幫助基本的同步SGD，但我們的假設(shè)是，如果我們使用兩組128GPU（使用同步SGD定期交換檢查點(diǎn)）雙向codistillation，那么額外的128GPU將提高訓(xùn)練時間。

一個問題是，codistillation僅僅是一種懲罰確信輸出分布（penalizing confident output distributions）或平滑標(biāo)簽的方法，所以我們也將其與兩個標(biāo)簽平滑基線進(jìn)行比較。第一個基線用符合統(tǒng)一分配的術(shù)語代替codistillation損失術(shù)語，第二個基線使用與使用一個與unigram分布相匹配的術(shù)語。在初步實(shí)驗(yàn)中手動調(diào)整權(quán)衡超參數(shù)。另一個重要的比較是兩個神經(jīng)網(wǎng)絡(luò)的集成，每個神經(jīng)網(wǎng)絡(luò)都有128GPU和同步SGD。

圖2a利用128個GPU的兩個組，以及同步SGD的訓(xùn)練曲線和標(biāo)簽平滑基線(每個使用128個GPU)和一個包含128個GPU基線的兩個實(shí)例的訓(xùn)練曲線，繪制了驗(yàn)證交叉信息和同步訓(xùn)練的步驟。

討論和未來的工作

Distillation是一種非常靈活的工具，尤其是在模型訓(xùn)練過程期間而不是之后進(jìn)行。它可用于加速培訓(xùn)、提高質(zhì)量，以新的、更有效的溝通方式下進(jìn)行訓(xùn)練，并減少預(yù)測的流失。然而，我們?nèi)杂性S多問題需要探討。例如，我們主要關(guān)注對彼此進(jìn)行codistilling的模型對上。如果成對是有用的，那么其他拓?fù)湟彩怯杏玫?。完全連接的圖形可能使模型過于相似，太快以至于環(huán)形結(jié)構(gòu)也可能變得有趣。我們也沒有探究teacher models預(yù)測的準(zhǔn)確性的局限性。也許可以對teacher models進(jìn)行積極的量化，使其在很大程度上與普通訓(xùn)練一樣廉價，即使是非常大的模型。

有些矛盾的是，壞模型之間codistilling可以比獨(dú)立的模型訓(xùn)練更快地學(xué)習(xí)。從某種程度上講，teacher models所犯的錯誤可以幫助student model做得更好，而且比僅僅看到數(shù)據(jù)中的實(shí)際標(biāo)簽要好。描述teacher models的理想特性是未來工作的一個令人興奮的途徑。在這項(xiàng)工作中，我們只從檢查點(diǎn)提取預(yù)測，因?yàn)轭A(yù)測是可識別的，而且不像網(wǎng)絡(luò)的內(nèi)部結(jié)構(gòu)，沒有虛假的對稱。也就是說，從一個檢查點(diǎn)提取更多信息可能比僅僅預(yù)測沒有觸及到工作人員交流梯度的相同問題，允許使用teacher models作為更強(qiáng)的調(diào)節(jié)者。也許可以用基于codistillation的方法來增強(qiáng)McMahan等人(2017)提出的在帶寬受限的設(shè)置環(huán)境下的聯(lián)合學(xué)習(xí)。