字節(jié)跳動開源 CowClip :推薦模型單卡訓練最高加速72倍
本文經(jīng)AI新媒體量子位(公眾號ID:QbitAI)授權轉載,轉載請聯(lián)系出處。
不夠快!還不夠快?
在 NLP 和 CV 任務上,為了加速神經(jīng)網(wǎng)絡的訓練,借助 32K 的批量大小(batch size)和 8 塊 GPU,只需 14 分鐘就完成 ImageNet 的訓練,76 分鐘完成 Bert 的訓練。研究人員對訓練速度的渴望從來沒有停止過。
那,只用 1 塊 GPU 夠不夠?在推薦系統(tǒng)上,不僅可以,還能將批量大小繼續(xù)提升!
最近,字節(jié)跳動AML(應用機器學習團隊)和新加坡國立大學的研究人員提出了一個新的優(yōu)化方法 CowClip,在公開點擊率預測數(shù)據(jù)集 Criteo 上最高支持 128K 批量大小,在單張 GPU 上訓練 DeepFM,可以將 12 小時的訓練時間壓縮到 10 分鐘,性能不降反升!
論文地址:https://arxiv.org/abs/2204.06240開源地址:https://github.com/bytedance/LargeBatchCTR
為了支持如此大批量大小的訓練,保持模型的訓練精度,文中提出了一套全新的優(yōu)化訓練方法:通過參數(shù)轉換公式確定大批量大小下的參數(shù),對 embedding 層進行自適應梯度裁剪。
用了 CowClip 優(yōu)化方法的不同推薦模型(文中測試了 DeepFM 等四個模型),在兩個公開數(shù)據(jù)集上進行了理論和實驗驗證,證實了該方法的合理性和有效性。
作者表示,使用該優(yōu)化方法,任何人都可以很容易的分分鐘訓練一個中小規(guī)模的推薦模型。
CowClip 加速的理論基礎
用戶交互會成為推薦系統(tǒng)新的訓練數(shù)據(jù),模型在一次次的重新訓練中都學到最新的知識。目前的推薦系統(tǒng)面對著數(shù)以億計的用戶和數(shù)以千億計的訓練數(shù)據(jù),一次完整的訓練要花費大量的時間和計算成本。
為了加速推薦系統(tǒng)的訓練,目前推薦系統(tǒng)會利用 GPU 進行加速訓練。然而,隨著 GPU 計算能力和顯存的不斷增加,過去推薦系統(tǒng)的訓練過程沒有完全利用好目前 GPU 的性能。舉例而言,在 Criteo 數(shù)據(jù)集上,當批量大?。╞atch size)從 1K 提升到 8K 后,用一塊 V100 進行訓練每次迭代所需的時間只有少量增加。這說明在目前的高性能 GPU 上,以往使用的小批量大小不足以充分利用 GPU 的并行能力。
使用更大的批量大小可以更充分的挖掘 GPU 的性能,讓 GPU 真正物有所值。只要大批量大小下訓練的模型精度沒有損失,我們就可以不斷提高模型的批量大小,直到塞滿 GPU 的顯存。
然而防止模型精度損失并不是一件易事。一方面,更大的批量大小可能會使網(wǎng)絡訓練不穩(wěn)定,并減弱網(wǎng)絡的泛化能力;另一方面,如果沒有規(guī)則指導在更大的批量大小上進行超參選擇,那調參會浪費大量的資源。
文中提出的 CowClip 便希望解決上述問題,通過在嵌入層(Embedding layer)逐列進行的動態(tài)梯度裁剪,和一組簡單有效的設置不同批量大小下超參數(shù)值的方法,讓擴大 128 倍的批量大小成為可能。
CowClip 方法
為了讓大批量大小下網(wǎng)絡的訓練更加穩(wěn)定,研究者提出了自適應逐列梯度裁剪策略(Adaptive Column-wise Gradient Norm Clipping, CowClip)以穩(wěn)定網(wǎng)絡的優(yōu)化過程。梯度裁剪是一種優(yōu)化更新過程的方法,它將范數(shù)值大于一定閾值的梯度裁剪到范數(shù)內。給定一個固定的閾值 clip_t,梯度裁剪過程如下:
然而直接運用該方法到嵌入層的梯度上效果并不佳。該原因不僅在于難以確定一個有效的閾值,更在于訓練過程中,每個特征取值(ID 特征)對應的編碼向量(對應嵌入層中嵌入矩陣的一列)在訓練過程中的梯度值大小各不相同(如圖 4 所示),在全局應用梯度裁剪忽視了不同梯度值之間的差異。
因此,研究者提出對每個特征取值對應的編碼向量單獨應用裁剪閾值,并自適應地設置該閾值??紤]到如果梯度大小超過參數(shù)大小本身時訓練過程會很不穩(wěn)定,研究者提出用特征取值對應的編碼向量自身的范數(shù)值確定閾值。為了防止裁剪閾值過小,參數(shù) ζ 保證了裁剪值不會低于一定的值。
上式中,還需要考慮到由于特征 ID 的總頻次不同,在每個批次中出現(xiàn)的次數(shù)也不同。為了平衡出現(xiàn)的不同頻次,最后獲得的裁剪值還需乘以批次中對應頻次出現(xiàn)的次數(shù)。
基于以上分析,研究者提出的 CowClip 算法如下:
大批量大小下的參數(shù)轉換
在 CV 和 NLP 任務中,已經(jīng)有一套擴大批量大小時調整學習率和 L2 正則項參數(shù)的方法。常用的方法包括線性調整(Linear Scaling),即在擴大 s 倍批量大小時,擴大 s 倍的學習率;以及平方根調整(Sqrt Scaling),即在擴大 s 倍批量大小時,學習率和正則項參數(shù)均以根號下 s 的大小擴大。
研究者首先探索了應用這些調整方法是否能有效的在大批量大小下保持性能。表一左側的實驗表明,當批量大小擴大時,這些方法的性能都出現(xiàn)了一定程度的下降。
文中指出,以往方法的失敗的原因在于,輸入的特征 ID 具有不同的頻率,而這是 CV 和 NLP 模型輸入不具有的特點。舉例而言,熱榜上的視頻播放量高,出現(xiàn)在數(shù)據(jù)集中的次數(shù)也就多,則視頻 ID 特征中該視頻 ID 的出現(xiàn)頻次就遠高于一些冷門視頻。
為了驗證上述想法,研究者改造了一個只包含高頻特征的數(shù)據(jù)集。果不其然,以往的參數(shù)調整方法此時可以取得好的結果(表一右側)。該實驗說明了頻次分布不一致確實阻礙了之前的參數(shù)調整方法,
論文中對該現(xiàn)象還進行了進一步的理論分析。簡單而言,如果重新考慮線性調整方法,其背后思想在于當批量大小增大后,更新迭代的步數(shù)減小,所以要擴大學習率。但對于出現(xiàn)次數(shù)非常少的特征,擴大批量大小時不會減小其更新迭代的次數(shù)。
由于點擊率預測數(shù)據(jù)集中絕大部分數(shù)據(jù)是此類低頻的特征 ID,結合 CowClip 方法,對模型的嵌入層可以不做學習率調整,并同時線性增大 L2 參數(shù)。
通過最后的實驗結果可以看到,利用 CowClip 訓練的模型比其它方法不僅精度更高,訓練速度也大幅度提升。