6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

机器之心發表於2022-04-06

谷歌的下一代架構 Pathways 已經用來訓練大模型了。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型
隨著規模的增加,模型在處理多個任務時的效能逐漸提高,而且還在不斷解鎖新的能力。

在探討現有 AI 模型的侷限時,谷歌人工智慧主管 Jeff Dean 曾經說過,今天的人工智慧系統總是從頭開始學習新問題。最終,我們為數千個單獨的任務開發了數千個模型。以這種方式學習每項新任務不僅需要更長的時間,而且還需要更多的資料,效率非常低。

在 Jeff Dean 看來,理想的發展模式應該是訓練一個模型來做成千上萬件事情。為了實現這一願景,他所在的團隊去年提出了一種名叫「Pathways」的通用 AI 架構。Jeff Dean 介紹說,Pathways 旨在用一個架構同時處理多項任務,並且擁有快速學習新任務、更好地理解世界的能力。前段時間,該團隊終於公佈了 Pathways 的論文

論文寫道,PATHWAYS 使用了一種新的非同步分散式資料流設計。這種設計允許 PATHWAYS 採用單控制器模型,從而更容易表達複雜的新並行模式。實驗結果表明,當在 2048 個 TPU 上執行 SPMD(single program multiple data)計算時,PATHWAYS 的效能(加速器利用率接近 100%)可以媲美 SOTA 系統。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型
谷歌 Pathways 系統架構概覽。

有了強大的系統,接下來就是訓練模型了。

在剛剛公佈的論文——「PaLM: Scaling Language Modeling with Pathways」中,谷歌宣佈,他們用 Pathways 系統訓練了一個 5400 億引數的大型語言模型——PaLM(Pathways Language Model)。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

論文連結:https://storage.googleapis.com/pathways-language-model/PaLM-paper.pdf

這是一個只有解碼器的密集 Transformer 模型。為了訓練這個模型,谷歌動用了 6144 塊 TPU,讓 Pathways 在兩個 Cloud TPU v4 Pods 上訓練 PaLM。

強大的系統和算力投入帶來了驚豔的結果。研究者在數百個語言理解和生成任務上評估了 PaLM,發現它在大多數任務上實現了 SOTA 少樣本學習效能,可以出色地完成笑話解讀、bug 修復、從表情符號中猜電影等語言、程式碼任務。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

有網友感嘆說,「終於知道谷歌開發 TPU 是用來幹嘛的了。這種級別的自然語言理解,一旦被應用程式所利用,並變得足夠高效並廣泛使用。這將徹底改變谷歌所做的一切。擁有能夠實現這一目標的晶片是非常有價值的,而那些陷入購買或租用英偉達晶片的公司將處於不利地位。」

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

PaLM 架構概覽

PaLM 只包含解碼器(每個時間步只能關注自身和過去的時間步),對一種標準的 Transformer 架構((Vaswani et al., 2017))做出瞭如下更改:

SwiGLU 啟用

研究者使用 SwiGLU 啟用 (Swish(xW) · xV) 用於 MLP 中間啟用,因為研究表明,與標準 ReLU、GeLU 或 Swish 啟用相比,SwiGLU 啟用能顯著提高質量。注意,在 MLP 中,這確實需要三個矩陣乘法,而不是兩個,但 Shazeer (2020) 在計算等效實驗中證明了質量的提升。

並行層

研究者在每個 Transformer 模組中使用「並行」方法,而不是標準的「序列」方法。具體來說,標準方法可以寫成:

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

並行方法可以寫成

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

由於 MLP 和注意力輸入矩陣乘法可以融合,這裡的並行方法可以讓大規模訓練速度提升 15%。消融實驗顯示,在 8B 的規模下,質量下降很小,但在 62B 規模下,質量沒有下降,因此研究者推斷,並行層的影響會在 540B 規模下達到 quality neutral。

多查詢(Multi-Query)注意力

標準 Transformer 方法使用 k 個注意力頭,其中每個時間步長的輸入向量被線性投影成形狀 [k,h] 的「查詢」、「鍵」和「值」張量,其中 h 是注意力頭大小。這裡,鍵 / 值投影對於每個頭是共享的,即「鍵」和「值」被投影到[1,h],但是「查詢」仍然被投影到形狀[k,h]。此前有研究表明,這對模型質量和訓練速度的影響呈中性,但在自迴歸解碼時間上可以帶來顯著的成本節約。這是因為在自迴歸解碼過程中,標準多頭注意力在加速器硬體上的效率很低,因為鍵 / 值張量不在例項之間共享,並且一次只有單個 token 被解碼。

RoPE 嵌入

研究者使用了 RoPE 嵌入而不是絕對或相對位置嵌入,因為 RoPE 嵌入已被證明在長序列長度上具有更好的效能。

共享輸入 - 輸出嵌入

研究者共享了輸入和輸出嵌入矩陣,這是在過去的工作中經常做的(但不是普遍的)。

No Biases

研究者在任何密集核或層 norm 中都沒有使用 biases。他們發現,這可以增加大型模型的訓練穩定性。

詞彙表

研究者使用了具有 256k token 的 SentencePiece 詞彙表,選擇這個詞彙表是為了支援訓練語料庫中的多種語言(沒有過多的分詞)。詞彙表是從訓練資料中生成的,研究者發現這提高了訓練效率。

用 Pathways 訓練一個 5400 億引數的語言模型

PaLM 是谷歌首次大規模使用 Pathways 系統將訓練擴充套件到 6144 塊晶片,這是迄今為止用於訓練的基於 TPU 的最大系統配置。研究者在 Pod 級別上跨兩個 Cloud TPU v4 Pods 使用資料並行對訓練進行擴充套件,同時在每個 Pod 中使用標準資料和模型並行。與以前的大多數 LLM 相比,這是一個顯著的規模增長。

PaLM 實現了 57.8% 的硬體 FLOPs 利用率的訓練效率,這是 LLM 在這個規模上實現的最高效率。為了達到這一水平,研究者將並行策略和 Transformer 塊的重新設計結果相結合,這使得注意力層和前饋層平行計算成為可能,從而實現了 TPU 編譯器最佳化帶來的加速。

PaLM 使用英語和多語言資料集進行訓練,包括高質量的 web 文件、書籍、維基百科、對話和 GitHub 程式碼。研究者還建立了一個「無損(lossless)」詞彙表,它保留了所有空格(對於程式碼來說尤其重要),將詞彙表之外的 Unicode 字元拆分成位元組,並將數字拆分成單獨的 token,每個 token 對應一個數字。

突破性的語言、推理和程式碼生成能力

PaLM 在許多非常困難的任務上顯示出了突破性的能力,包括語言理解、生成、推理和程式碼等相關任務。

語言理解與生成

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

在 29 項基於英語的 NLP 任務上,PaLM 540B 的效能比之前的 SOTA 結果有所提高。

除了英語 NLP 任務外,PaLM 在包括翻譯在內的多語言 NLP 基準測試中也表現出強大的效能,儘管它只有 22% 的訓練語料庫是非英語的。

研究者還在 Beyond the Imitation Game Benchmark (BIG-bench) 上探索了 PaLM 新的和未來功能,這是一個近期釋出的套件,包含 150 多個新的語言建模任務。在這個過程中, PaLM 與 Gopher 和 Chinchilla 的效能進行了比較,涉及這些任務的 58 個公共子集。

有趣的是,研究者注意到,作為 scale 函式(function of scale)的 PaLM 的效能遵循與先前模型類似的對數線性表現,這表明 scale 的效能改進尚未趨於平穩。PaLM 540B 5-shot 的表現也優於解決相同任務的人類平均表現。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型
PaLM 在 58 個 BIG-bench 任務的子集上的 scaling behavior。

PaLM 在幾個 BIG-bench 任務中展示了令人印象深刻的自然語言理解和生成能力。例如,該模型可以區分因果關係,理解上下文中的概念組合,甚至可以從表情符號中猜測電影。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

PaLM 540B 在 BIG-bench 任務上的 1-shot 效能展示:標記因果關係、概念理解、從表情符號中猜測電影以及查詢同義詞和反事實。

推理

透過將模型 scale 與 chain-of-thought prompting 相結合,PaLM 在需要多步驟算術或常識推理的推理任務上展示出了突破性的能力。以往諸如 Gopher 這樣的大型語言模型在提高效能方面從模型 scale 中獲益較少。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型
小學數學問題示例中的標準 prompting 與 chain-of-thought prompting。Chain-of-thought prompting 將多步驟推理問題的 prompt 分解為中間步驟(黃色部分),類似於人類處理它的方式。

研究者在三個算術資料集和兩個常識推理資料集上觀察到了 PaLM 540B 在 chain-of-thought prompt 加持下的強大效能。例如,藉助 8-shot prompt,PaLM 解決了 GSM8K 中 58% 的問題,這是一個包含數千個具有挑戰性的小學水平數學問題的基準,超過了之前 GPT-3 175B 微調模型(訓練集包含 7500 個問題,並與外部計算器和驗證器相結合)獲得的 55% 的最高分。

這個新的得分值得關注,因為它接近 60% 的 9 到 12 歲兒童解決問題的水平,這些兒童正是問題集的目標受眾。研究者猜測,PaLM 詞彙表中數字的獨立編碼有助於實現這些效能改進。

值得注意的是,PaLM 甚至可以為需要多步驟邏輯推理、世界認知和深度語言理解的複雜組合的場景生成明確的解釋。例如,它可以為網路上搜不到的新笑話提供高質量的解釋。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

PaLM 用 two-shot prompts 解釋了一個原創笑話。

程式碼生成

大型語言模型已被證明可以很好地推廣到編碼任務,比如在給定自然語言描述(文字到程式碼)的情況下編寫程式碼,將程式碼從一種語言翻譯成另一種語言,以及修復編譯錯誤(程式碼到程式碼)。

PaLM 540B 在單個模型中顯示了橫跨編碼任務和自然語言任務的強大效能,即使它在預訓練資料集中只有 5% 的程式碼。具體而言,PaLM 540B 的 few-shot 效能十分顯著,與經過微調的 Codex 12B 相當,同時使用的 Python 訓練程式碼減少到了 50 分之一。這一結果印證了之前的發現,即較大的模型比較小的模型更高效,因為它們可以更好地從其他程式語言和自然語言資料中實現遷移學習。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型

PaLM 540B 微調模型在文字到程式碼任務(例如 GSM8K - Python 和 HumanEval)和程式碼到程式碼任務(例如 Transcoder)上的示例。

此外,透過在純 Python 程式碼資料集上微調 PaLM ,模型進一步提高了效能,團隊稱之為 PaLM-Coder。如下圖所示,PaLM-Coder 接到了一個名為 DeepFix 的示例程式碼修復任務,目標是修改最初損壞的 C 程式直到它們編譯成功,PaLM-Coder 540B 展示了令人印象深刻的效能,實現了 82.1% 的編譯率,優於之前 71.7% 的 SOTA 結果。這為修復軟體開發過程中出現的更復雜的錯誤提供了機會。

6144塊TPU,5400億引數,會改bug、解讀笑話,谷歌剛剛用Pathways訓練了一個大模型


DeepFix 程式碼修復任務的示例。經過微調的 PaLM-Coder 540B 將編譯錯誤(左)修復為可編譯的程式碼版本(右)。

從倫理方面考慮,最近的研究強調了受過網路文字訓練的 LLM 相關的各種潛在風險。透過模型卡片和資料表等透明工件分析並記錄這些潛在的不良風險是至關重要的,其中還包括有關預期用途和測試的資訊。為此,谷歌的論文提供了資料表、模型卡片和 Responsible AI 基準測試結果,並報告了對資料集和模型輸出的全面分析,以發現偏差和風險。雖然分析有助於概述模型的一些潛在風險,但特定領域和任務的分析對於真正校準、情境化和減輕可能的危害至關重要。進一步瞭解這些模型的風險和收益是正在進行的研究的主題,同時開發可擴充套件的解決方案,防止惡意使用語言模型。

參考連結:https://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html

相關文章