本文經(jīng)AI新媒體量子位（公眾號ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。

不夠快！還不夠快？

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

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

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

論文地址：https://arxiv.org/abs/2204.06240開源地址：https://github.com/bytedance/LargeBatchCTR

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

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

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

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

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

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

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

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

文中提出的 CowClip 便希望解決上述問題，通過在嵌入層（Embedding layer）逐列進行的動態(tài)梯度裁剪，和一組簡單有效的設(shè)置不同批量大小下超參數(shù)值的方法，讓擴大 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)化過程。梯度裁剪是一種優(yōu)化更新過程的方法，它將范數(shù)值大于一定閾值的梯度裁剪到范數(shù)內(nèi)。給定一個固定的閾值 clip_t，梯度裁剪過程如下：

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

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

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

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

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

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

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

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

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

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

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

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