一種相對(duì)直接使用的distillation的變體方法

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

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

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

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

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

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

對(duì)于大規(guī)模的、具有商業(yè)價(jià)值的神經(jīng)網(wǎng)絡(luò)訓(xùn)練問題，如果訓(xùn)練時(shí)間能夠大幅加快，或最終模型的質(zhì)量能夠大幅提高，從業(yè)者會(huì)愿意投入更多的機(jī)器用于訓(xùn)練。目前，分布式隨機(jī)梯度下降（SGD），包括其同步和異步形式（Chen et al.，2016）是在多個(gè)互聯(lián)機(jī)器上進(jìn)行大規(guī)模神經(jīng)網(wǎng)絡(luò)訓(xùn)練的主要算法。但是，隨著機(jī)器數(shù)量的增加，訓(xùn)練一個(gè)高質(zhì)量模型所需時(shí)間的改善程度會(huì)降低，直到繼續(xù)增加機(jī)器卻無(wú)法進(jìn)一步縮短訓(xùn)練時(shí)間?；A(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è)置中，在超過(guò)100臺(tái)GPU workers上進(jìn)行有效擴(kuò)展可能非常困難。沒有任何訓(xùn)練神經(jīng)網(wǎng)絡(luò)的算法可以無(wú)限擴(kuò)展，但即使擴(kuò)展得比分布式SGD的限制多一點(diǎn)，也是非常有價(jià)值的。

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

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

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

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

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

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

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

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

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

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

codistillation算法

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

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

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

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

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

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

2. 帶同步SGD的CODISTILLATION

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

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

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

討論和未來(lái)的工作

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

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