6144塊TPU，5400億參數，會改bug、解讀笑話，谷歌剛剛用Pathways訓練了一個大模型
谷歌的下一代架構 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)。
論文鏈接：https://storage.googleapis.com/pathways-language-model/PaLM-paper.pdf
這是一個只有解碼器的密集 Transformer 模型。為了訓練這個模型，谷歌動用了 6144 塊 TPU，讓 Pathways 在兩個 Cloud TPU v4 Pods 上訓練 PaLM。
強大的系統和算力投入帶來了驚艷的結果。研究者在數百個語言理解和生成任務上評估了 PaLM，發現它在大多數任務上實現了 SOTA 少樣本學習性能，可以出色地完成笑話解讀、bug 修復、從表情符號中猜電影等語言、代碼任務。
有網友感嘆說，「終于知道谷歌開發 TPU 是用來干嘛的了。這種級別的自然語言理解，一旦被應用程序所利用，并變得足夠高效并廣泛使用。這將徹底改變谷歌所做的一切。擁有能夠實現這一目標的芯片是非常有價值的，而那些陷入購買或租用英偉達芯片的公司將處于不利地位。」
PaLM 架構概覽
PaLM 只包含解碼器(每個時間步只能關注自身和過去的時間步)，對一種標準的 Transformer 架構((Vaswani et al., 2017))做出了如下更改：
SwiGLU 激活
研究者使用 SwiGLU 激活 (Swish(xW) · xV) 用于 MLP 中間激活，因為研究表明，與標準 ReLU、GeLU 或 Swish 激活相比，SwiGLU 激活能顯著提高質量。注意，在 MLP 中，這確實需要三個矩陣乘法，而不是兩個，但 Shazeer (2020) 在計算等效實驗中證明了質量的提升。
并行層
研究者在每個 Transformer 模塊中使用「并行」方法，而不是標準的「串行」方法。具體來說，標準方法可以寫成：
并行方法可以寫成
由于 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 在許多非常困難的任務上顯示出了突破性的能力，包括語言理解、生成、推理和代碼等相關任務。
語言理解與生成
在 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 的表現也優于解決相同任務的人類平均表現。
PaLM 在 58 個 BIG-bench 任務的子集上的 scaling behavior。
PaLM 在幾個 BIG-bench 任務中展示了令人印象深刻的自然語言理解和生成能力。例如，該模型可以區分因果關系，理解上下文中的概念組合，甚至可以從表情符號中猜測電影。
PaLM 540B 在 BIG-bench 任務上的 1-shot 性能展示：標記因果關系、概念理解、從表情符號中猜測電影以及查找同義詞和反事實。
推理
通過將模型 scale 與 chain-of-thought prompting 相結合，PaLM 在需要多步驟算術或常識推理的推理任務上展示出了突破性的能力。以往諸如 Gopher 這樣的大型語言模型在提高性能方面從模型 scale 中獲益較少。
小學數學問題示例中的標準 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 甚至可以為需要多步驟邏輯推理、世界認知和深度語言理解的復雜組合的場景生成明確的解釋。例如，它可以為網絡上搜不到的新笑話提供高質量的解釋。
PaLM 用 two-shot prompts 解釋了一個原創笑話。
代碼生成
大型語言模型已被證明可以很好地推廣到編碼任務，比如在給定自然語言描述(文本到代碼)的情況下編寫代碼，將代碼從一種語言翻譯成另一種語言，以及修復編譯錯誤(代碼到代碼)。
PaLM 540B 在單個模型中顯示了橫跨編碼任務和自然語言任務的強大性能，即使它在預訓練數據集中只有 5% 的代碼。具體而言，PaLM 540B 的 few-shot 性能十分顯著，與經過微調的 Codex 12B 相當，同時使用的 Python 訓練代碼減少到了 50 分之一。這一結果印證了之前的發現，即較大的模型比較小的模型更高效，因為它們可以更好地從其他編程語言和自然語言數據中實現遷移學習。
PaLM 540B 微調模型在文本到代碼任務(例如 GSM8K - Python 和 HumanEval)和代碼到代碼任務(例如 Transcoder)上的示例。
此外，通過在純 Python 代碼數據集上微調 PaLM ，模型進一步提高了性能，團隊稱之為 PaLM-Coder。如下圖所示，PaLM-Coder 接到了一個名為 DeepFix 的示例代碼修復任務，目標是修改最初損壞的 C 程序直到它們編譯成功，PaLM-Coder 540B 展示了令人印象深刻的性能，實現了 82.1% 的編譯率，優于之前 71.7% 的 SOTA 結果。這為修復軟件開發過程中出現的更復雜的錯誤提供了機會。
DeepFix 代碼修復任務的示例。經過微調的 PaLM-Coder 540B 將編譯錯誤(左)修復為可編譯的代碼版本(右)。
從倫理方面考慮，最近的研究強調了受過網絡文本訓練的 LLM 相關的各種潛在風險。通過模型卡片和數據表等透明工件分析并記錄這些潛在的不良風險是至關重要的，其中還包括有關預期用途和測試的信息。為此，谷歌的論文提供了數據表、模型卡片和 Responsible AI 基準測試結果，并報告了對數據集和模型輸出的全面分析，以發現偏差和風險。雖然分析有助于概述模型的一些潛在風險，但特定領域和任務的分析對于真正校準、情境化和減輕可能的危害至關重要。進一步了解這些模型的風險和收益是正在進行的研究的主題，同時開發可擴展的解決方案，防止惡意使用語言模型。
