字節(jié)跳動(dòng)開(kāi)源 CowClip：推薦模型單卡訓(xùn)練最高加速72倍

量子位的朋友們 2022-04-28 10:21:17 來(lái)源：量子位

不夠快！還不夠快？

在 NLP 和 CV 任務(wù)上，為了加速神經(jīng)網(wǎng)絡(luò)的訓(xùn)練，借助 32K 的批量大小（batch size）和 8 塊 GPU，只需 14 分鐘就完成 ImageNet 的訓(xùn)練，76 分鐘完成 Bert 的訓(xùn)練。研究人員對(duì)訓(xùn)練速度的渴望從來(lái)沒(méi)有停止過(guò)。

那，只用 1 塊 GPU 夠不夠？在推薦系統(tǒng)上，不僅可以，還能將批量大小繼續(xù)提升！

最近，字節(jié)跳動(dòng)AML（應(yīng)用機(jī)器學(xué)習(xí)團(tuán)隊(duì)）和新加坡國(guó)立大學(xué)的研究人員提出了一個(gè)新的優(yōu)化方法 CowClip，在公開(kāi)點(diǎn)擊率預(yù)測(cè)數(shù)據(jù)集 Criteo 上最高支持 128K 批量大小，在單張 GPU 上訓(xùn)練 DeepFM，可以將 12 小時(shí)的訓(xùn)練時(shí)間壓縮到 10 分鐘，性能不降反升！

論文地址：https://arxiv.org/abs/2204.06240
開(kāi)源地址：https://github.com/bytedance/LargeBatchCTR

為了支持如此大批量大小的訓(xùn)練，保持模型的訓(xùn)練精度，文中提出了一套全新的優(yōu)化訓(xùn)練方法：通過(guò)參數(shù)轉(zhuǎn)換公式確定大批量大小下的參數(shù)，對(duì) embedding 層進(jìn)行自適應(yīng)梯度裁剪。

用了 CowClip 優(yōu)化方法的不同推薦模型（文中測(cè)試了 DeepFM 等四個(gè)模型），在兩個(gè)公開(kāi)數(shù)據(jù)集上進(jìn)行了理論和實(shí)驗(yàn)驗(yàn)證，證實(shí)了該方法的合理性和有效性。

作者表示，使用該優(yōu)化方法，任何人都可以很容易的分分鐘訓(xùn)練一個(gè)中小規(guī)模的推薦模型。

CowClip 加速的理論基礎(chǔ)

用戶交互會(huì)成為推薦系統(tǒng)新的訓(xùn)練數(shù)據(jù)，模型在一次次的重新訓(xùn)練中都學(xué)到最新的知識(shí)。目前的推薦系統(tǒng)面對(duì)著數(shù)以億計(jì)的用戶和數(shù)以千億計(jì)的訓(xùn)練數(shù)據(jù)，一次完整的訓(xùn)練要花費(fèi)大量的時(shí)間和計(jì)算成本。

為了加速推薦系統(tǒng)的訓(xùn)練，目前推薦系統(tǒng)會(huì)利用 GPU 進(jìn)行加速訓(xùn)練。然而，隨著 GPU 計(jì)算能力和顯存的不斷增加，過(guò)去推薦系統(tǒng)的訓(xùn)練過(guò)程沒(méi)有完全利用好目前 GPU 的性能。舉例而言，在 Criteo 數(shù)據(jù)集上，當(dāng)批量大?。╞atch size）從 1K 提升到 8K 后，用一塊 V100 進(jìn)行訓(xùn)練每次迭代所需的時(shí)間只有少量增加。這說(shuō)明在目前的高性能 GPU 上，以往使用的小批量大小不足以充分利用 GPU 的并行能力。

使用更大的批量大小可以更充分的挖掘 GPU 的性能，讓 GPU 真正物有所值。只要大批量大小下訓(xùn)練的模型精度沒(méi)有損失，我們就可以不斷提高模型的批量大小，直到塞滿 GPU 的顯存。

然而防止模型精度損失并不是一件易事。一方面，更大的批量大小可能會(huì)使網(wǎng)絡(luò)訓(xùn)練不穩(wěn)定，并減弱網(wǎng)絡(luò)的泛化能力；另一方面，如果沒(méi)有規(guī)則指導(dǎo)在更大的批量大小上進(jìn)行超參選擇，那調(diào)參會(huì)浪費(fèi)大量的資源。

文中提出的 CowClip 便希望解決上述問(wèn)題，通過(guò)在嵌入層（Embedding layer）逐列進(jìn)行的動(dòng)態(tài)梯度裁剪，和一組簡(jiǎn)單有效的設(shè)置不同批量大小下超參數(shù)值的方法，讓擴(kuò)大 128 倍的批量大小成為可能。

CowClip 方法

為了讓大批量大小下網(wǎng)絡(luò)的訓(xùn)練更加穩(wěn)定，研究者提出了自適應(yīng)逐列梯度裁剪策略（Adaptive Column-wise Gradient Norm Clipping, CowClip）以穩(wěn)定網(wǎng)絡(luò)的優(yōu)化過(guò)程。梯度裁剪是一種優(yōu)化更新過(guò)程的方法，它將范數(shù)值大于一定閾值的梯度裁剪到范數(shù)內(nèi)。給定一個(gè)固定的閾值 clip_t，梯度裁剪過(guò)程如下：

然而直接運(yùn)用該方法到嵌入層的梯度上效果并不佳。該原因不僅在于難以確定一個(gè)有效的閾值，更在于訓(xùn)練過(guò)程中，每個(gè)特征取值（ID 特征）對(duì)應(yīng)的編碼向量（對(duì)應(yīng)嵌入層中嵌入矩陣的一列）在訓(xùn)練過(guò)程中的梯度值大小各不相同（如圖 4 所示），在全局應(yīng)用梯度裁剪忽視了不同梯度值之間的差異。

因此，研究者提出對(duì)每個(gè)特征取值對(duì)應(yīng)的編碼向量單獨(dú)應(yīng)用裁剪閾值，并自適應(yīng)地設(shè)置該閾值?？紤]到如果梯度大小超過(guò)參數(shù)大小本身時(shí)訓(xùn)練過(guò)程會(huì)很不穩(wěn)定，研究者提出用特征取值對(duì)應(yīng)的編碼向量自身的范數(shù)值確定閾值。為了防止裁剪閾值過(guò)小，參數(shù) ζ 保證了裁剪值不會(huì)低于一定的值。

上式中，還需要考慮到由于特征 ID 的總頻次不同，在每個(gè)批次中出現(xiàn)的次數(shù)也不同。為了平衡出現(xiàn)的不同頻次，最后獲得的裁剪值還需乘以批次中對(duì)應(yīng)頻次出現(xiàn)的次數(shù)。

基于以上分析，研究者提出的 CowClip 算法如下：

大批量大小下的參數(shù)轉(zhuǎn)換

在 CV 和 NLP 任務(wù)中，已經(jīng)有一套擴(kuò)大批量大小時(shí)調(diào)整學(xué)習(xí)率和 L2 正則項(xiàng)參數(shù)的方法。常用的方法包括線性調(diào)整（Linear Scaling），即在擴(kuò)大 s 倍批量大小時(shí)，擴(kuò)大 s 倍的學(xué)習(xí)率；以及平方根調(diào)整（Sqrt Scaling），即在擴(kuò)大 s 倍批量大小時(shí)，學(xué)習(xí)率和正則項(xiàng)參數(shù)均以根號(hào)下 s 的大小擴(kuò)大。

研究者首先探索了應(yīng)用這些調(diào)整方法是否能有效的在大批量大小下保持性能。表一左側(cè)的實(shí)驗(yàn)表明，當(dāng)批量大小擴(kuò)大時(shí)，這些方法的性能都出現(xiàn)了一定程度的下降。

文中指出，以往方法的失敗的原因在于，輸入的特征 ID 具有不同的頻率，而這是 CV 和 NLP 模型輸入不具有的特點(diǎn)。舉例而言，熱榜上的視頻播放量高，出現(xiàn)在數(shù)據(jù)集中的次數(shù)也就多，則視頻 ID 特征中該視頻 ID 的出現(xiàn)頻次就遠(yuǎn)高于一些冷門(mén)視頻。

為了驗(yàn)證上述想法，研究者改造了一個(gè)只包含高頻特征的數(shù)據(jù)集。果不其然，以往的參數(shù)調(diào)整方法此時(shí)可以取得好的結(jié)果（表一右側(cè)）。該實(shí)驗(yàn)說(shuō)明了頻次分布不一致確實(shí)阻礙了之前的參數(shù)調(diào)整方法，

論文中對(duì)該現(xiàn)象還進(jìn)行了進(jìn)一步的理論分析。簡(jiǎn)單而言，如果重新考慮線性調(diào)整方法，其背后思想在于當(dāng)批量大小增大后，更新迭代的步數(shù)減小，所以要擴(kuò)大學(xué)習(xí)率。但對(duì)于出現(xiàn)次數(shù)非常少的特征，擴(kuò)大批量大小時(shí)不會(huì)減小其更新迭代的次數(shù)。

由于點(diǎn)擊率預(yù)測(cè)數(shù)據(jù)集中絕大部分?jǐn)?shù)據(jù)是此類低頻的特征 ID，結(jié)合 CowClip 方法，對(duì)模型的嵌入層可以不做學(xué)習(xí)率調(diào)整，并同時(shí)線性增大 L2 參數(shù)。

通過(guò)最后的實(shí)驗(yàn)結(jié)果可以看到，利用 CowClip 訓(xùn)練的模型比其它方法不僅精度更高，訓(xùn)練速度也大幅度提升。