谷歌的下一代架構 Pathways 已經用來訓練大模型了。
隨著規模的增加,模型在處理多個任務時的性能逐漸提高,而且還在不斷解鎖新的能力。
在探討現有 AI 模型的局限時,谷歌人工智能主管 Jeff Dean 曾經說過,今天的人工智能系統總是從頭開始學習新問題。最終,我們為數千個單獨的任務開發了數千個模型。以這種方式學習每項新任務不僅需要更長的時間,而且還需要更多的數據,效率非常低。
在 Jeff Dean 看來,理想的發展模式應該是訓練一個模型來做成千上萬件事情。為了實現這一愿景,他所在的團隊去年提出了一種名叫「Pathways」的通用 AI 架構。Jeff Dean 介紹說,Pathways 旨在用一個架構同時處理多項任務,并且擁有快速學習新任務、更好地理解世界的能力。前段時間,該團隊終于公布了 Pathways 的論文。
論文寫道,PATHWAYS 使用了一種新的異步分布式數據流設計。這種設計允許 PATHWAYS 采用單控制器模型,從而更容易表達復雜的新并行模式。實驗結果表明,當在 2048 個 TPU 上運行 SPMD(single program multiple data)計算時,PATHWAYS 的性能(加速器利用率接近 100%)可以媲美 SOTA 系統。
谷歌 Pathways 系統架構概覽。
有了強大的系統,接下來就是訓練模型了。
在剛剛公布的論文——「PaLM: Scaling Language Modeling with Pathways」中,谷歌宣布,他們用 Pathways 系統訓練了一個 5400 億參數的大型語言模型——PaLM(Pathways Language Model)。
圖片
論文鏈接://storage.googleapis.com/pathways-language-model/PaLM-paper.pdf
這是一個只有解碼器的密集 Transformer 模型。為了訓練這個模型,谷歌動用了 6144 塊 TPU,讓 Pathways 在兩個 Cloud TPU v4 Pods 上訓練 PaLM。
強大的系統和算力投入帶來了驚艷的結果。研究者在數百個語言理解和生成任務上評估了 PaLM,發現它在大多數任務上實現了 SOTA 少樣本學習性能,可以出色地完成笑話解讀、bug 修復、從表情符號中猜電影等語言、代碼任務。
原型驅動的文本生成使用非參數模型,該模型首先從句子庫中選擇“原型”,然后修改原型生成輸出文本。這些方法雖然有效,但測試時效率低下,因為需要對整個訓練語料庫進行存儲和索引。此外,現有的方法通常需要啟發式來確定在訓練時引用哪個原型。在本文中,我們提出了一種新的生成模型,它可以自動學習稀疏原型支持集,同時也可以獲得較強的語言建模性能。通過(1)在原型選擇分布上施加稀疏誘導先驗,(2)利用平攤變分推理學習原型檢索函數來實現。在實驗中,我們的模型優于以前的原型驅動的語言模型,同時實現了高達1000倍的內存減少,以及測試時1000倍的加速。更有趣的是,當我們改變原型選擇的稀疏性時,我們展示了學習的原型能夠在不同的粒度捕獲語義和語法,并且可以通過指定生成的原型來控制某些句子屬性。