你還在糾結單個GPU怎麼訓練GPT-3嗎?快來看看HP調優新正規化吧!

語言: CN / TW / HK

編者按:偉大的科學成就不能僅靠反覆試驗取得。在構建大規模人工智慧系統時,基礎研究所形成的理論見解能夠幫助研究員大大減少試錯次數並提高成本效益。在今天的文章中,微軟研究院的研究員們將介紹基礎研究如何首次能夠調整龐大的神經網路。由於龐大的神經網路訓練十分昂貴,所以研究員們通過展示特定引數化在不同模型大小上保留最佳超引數來解決這一問題。通過與 OpenAI 合作,微軟研究院的研究員們在一系列現實場景中也驗證了該技術的實際優勢。本文轉載自公眾號“機器之心”編譯的微軟研究院部落格“µTransfer: A technique for hyperparameter tuning of enormous neural networks”,論文作者對部分內容做了調整。

偉大的科學成就不能僅靠反覆試驗取得。例如太空計劃中的每一次發射都是基於數百年的空氣動力學、推進和天體等基礎研究。同樣,在構建大規模人工智慧系統時,基礎研究大大減少了試錯次數,效益明顯。

超引數(Hyperparameter,HP)調優是深度學習的關鍵,但也是一個昂貴的過程,對於具有數十億引數的神經網路來說更是如此。假如 HP 選擇不當,會導致模型效能不佳、訓練不穩定。當訓練超大型深度學習模型時,這些問題更加嚴重。

最近,有研究表明不同的神經網路引數化會導致不同的無限寬度極限(infinite-width limits),他們提出了最大更新引數化(Maximal Update Parametrization,µP),該引數化的無限寬極限是實現“最大 特徵學習的。直觀地說,它確保每一層啟用在訓練期間的更新大小保持一致,而不管寬度如何。相比之下,雖然標準引數化 (standard parametrization,SP) 在初始化時保證了啟用的一致性,但實際上在訓練時,由於每層學習率的不平衡,導致啟用在寬模型中爆炸。

來自微軟和 OpenAI 的研究者首次提出了基礎研究如何調優大型神經網路(這些神經網路過於龐大而無法多次訓練)。他們通過展示特定引數化保留不同 大小 模型的最佳超引數來實現這一點。利用 µP 將 HP 從小型模型遷移到大型模型。也就是說,該研究在大型模型上獲得了接近最優的 HP。

本文的想法非常簡單,論文中引入了一種特殊引數化 µP,窄和寬的神經網路共享同一組最優超引數。即使寬度→∞也是如此。

論文作者之一、來自微軟的 Greg Yang 在推特上表示:“你不能在單個 GPU 上訓練 GPT-3,更不用說調優它的超引數了。但是利用新的理論進步,你竟可以在單個 GPU 上調優它的超引數!”

具體而言,該研究證明,在 µP 中,即使模型大小發生變化,許多最優的 HP 仍保持穩定。這導致一種新的 HP 調優正規化:µTransfer,即在 µP 中對目標模型進行引數化,並在較小的模型上間接調優 HP,將其零樣本遷移到全尺寸模型上,無需調優後者。該研究在 Transformer 和 ResNet 上驗證 µTransfer,例如,1)通過從 13M 引數的模型中遷移預訓練 HP,該研究優於 BERT-large (350M 引數),總調優成本相當於一次預訓練 BERT-large;2)通過從 40M 引數遷移,該研究的效能優於已公開的 6.7B GPT-3 模型,調優成本僅為總預訓練成本的7%。

  • 論文地址: https://arxiv.org/pdf/2203.03466.pdf

  • 專案地址: https://github.com/microsoft/mup

通過大大減少猜測訓練超引數的需要,這種技術可以加快對巨大神經網路的研究,例如 GPT-3 和未來可能更大的繼任者。

擴充套件初始化容易,但擴充套件訓練難

大型神經網路很難訓練,部分原因是不瞭解其行為如何隨著規模增加而變化。在深度學習的早期工作中,研究者採用啟發式演算法。一般來說,啟發式方法試圖在模型初始化時保持啟用大小一致,無論寬度如何。然而,隨著訓練的開始,這種一致性會在不同的模型寬度處中斷,如圖1左側所示。

與隨機初始化不同,模型訓練期間的行為更難進行數學分析。該研究用 µP 解決,如圖1右側所示,該圖顯示了網路啟用擴充套件(activation scales)在模型寬度增加的最初幾個訓練步驟中的穩定性。

圖1:在 PyTorch 的預設引數化中,左圖,在經過一次 step 訓練後,啟用擴充套件的寬度會出現差異。但是在右圖的 µP 中,無論訓練 step 寬度如何,啟用擴充套件都會發生一致的變化。

事實上,除了在整個訓練過程中保持啟用擴充套件一致之外,µP 還確保不同且足夠寬的神經網路在訓練過程中表現相似,以使它們收斂到一個理想的無窮寬極限,該研究稱之為特徵學習極限。

如圖所示,µP 是唯一在寬度上保持最佳學習率的引數化,在寬度為8192的模型中實現了最佳效能,並且對於給定的學習率,更寬的模型效能更好——即曲線不相交。

圖2:左側,該研究在 CIFAR10 上以不同的學習率(沿 x 軸顯示)訓練不同寬度(對應於不同顏色和圖案的曲線)的多層感知器 (MLP),並沿 y 軸繪製訓練損失。右側,引數化的 2D 平面由以下插值形成:1)PyTorch 預設值和 µP(x 軸)之間的初始化擴充套件,以及 2)PyTorch 預設值和 µP(y 軸)之間的學習率擴充套件。在這個平面上,PyTorch 預設用 (0,0) 表示,µP 預設用 (1,1) 表示。

基於張量程式(Tensor Programs)的理論基礎,µTransfer 自動適用於高階架構,例如 Transformer 和 ResNet。此外,它還可以同時遷移各種超引數。

以 Transformer 為例,圖3展示了關鍵超引數如何在寬度上保持穩定。超引數可以包括學習率、學習率 schedule、初始化、引數乘數等,甚至可以單獨針對每個引數張量。該研究在最大寬度為4096的 Transformer 上驗證了這一點。

圖3:在 µP 中引數化並在 WikiText-2 上訓練的不同寬度的 transformer。隨著模型寬度的增加,最優學習率、交叉熵溫度、初始化規模和學習率方案保持穩定。在右下角的圖中,該研究嘗試瞭如下學習率方案:(a) 線性衰減,(b) StepLR @ [5k, 8k],衰減因子為0.1,(c) StepLR @ [4k, 7k],衰減因子為0.3,(d) 餘弦退火,(e) 常數,(f) 逆平方根衰減。

模型深度的實驗擴充套件

現代神經網路擴充套件不止涉及寬度一個維度。該研究還探索瞭如何通過將 µP 與非寬度維度的簡單啟發式演算法相結合,將其應用於現實的訓練場景。下圖4使用相同的 transformer 設定來顯示最佳學習率如何在合理的非寬度維度範圍內保持穩定。

圖4:在 µP 中引數化並在 Wikitext-2 上訓練的不同大小的 transformer。如圖3所示,最優學習率不僅可以跨寬度遷移,還可在測試範圍內實驗性地跨其他擴充套件維度遷移,例如深度、批大小和序列長度。這意味著可以將理論上的跨寬度遷移與實驗驗證的跨其他擴充套件維度遷移相結合,以獲得能在小模型上間接調整超引數並遷移到大模型的 µTransfer。

除了學習率,其他超引數的情況如下圖所示:

測試 µTransfer

在驗證完單個超引數的遷移之後,研究者試圖將它們組合到更現實的場景中。下圖5對比了兩種情況,一種是 µTransfer 從一個小型 proxy 模型遷移調整過的超引數,另一種是直接調整大型目標模型。在這兩種情況下,調整都是通過隨機搜尋完成的。

圖5:µTransfer 大約將計算效率提高了一個數量級。

由於 proxy 模型即使很小也能有意義地預測最佳超引數(如圖3、圖4所示),因此隨著該研究用數十億個引數訓練更大的目標模型,研究者預計效能差距會擴大。

µP + GPT-3

在這項工作之前,模型越大,調優成本越高,預計調優效果就越差。研究者預計 µTransfer 將給最大的模型帶來最大的增益,因此該研究與 OpenAI 合作,在 GPT-3 上評估 µTransfer。

使用 µP 對 GPT-3 的一個相對位置編碼版本進行引數化後,該研究調整了一個具有4000萬個引數的小型 proxy 模型,然後按照 µTransfer 的方法將最佳超引數組合複製到 GPT-3 的67億引數變體中。在此調整階段使用的總計算量僅為67億模型預訓練使用計算量的7%。如下圖所示,這個使用 µTransfer 的模型優於 GPT-3 論文中相同大小的模型(絕對位置編碼),它的效能與 GPT-3 論文中引數數量翻倍的模型(絕對位置編碼)相當。

理論意義

µP 給出了一個擴充套件規則,該規則在訓練損失方面唯一地保留了跨不同寬度模型的最佳超引數組合。相反,其他擴充套件規則(如 PyTorch 中的預設初始化或 NTK 引數化),隨著網路變得越來越寬,超引數空間中的最優值卻越來越遠。研究者認為:實際使用特徵學習神經網路時,為了獲得適用的見解,µP 的特徵學習極限會比 NTK 極限更自然。因此,過引數化神經網路應該在大型寬度設定中重現 µP 的特徵學習極限。

過去幾年開發的張量程式(TP 理論使這項進展成為可能。TP 理論使研究人員能夠計算任何通用計算圖在其矩陣維數變大時的極限。TP 方法產生了基本的理論結果,例如神經網路 - 高斯過程對應的架構普遍性和動態二分定理,並通過推導 µP 和特徵學習極限形成了 µTransfer。研究者認為將 TP 理論擴充套件到深度、批大小等擴充套件維度是大型模型在寬度之外可靠擴充套件的關鍵。

研究者表示:基礎研究是對反覆試錯的一種高成本效益補充,該研究將繼續推匯出更具原則性的大規模機器學習方法。

你也許還想看