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

隨著規(guī)模的增加，模型在處理多個任務時的性能逐漸提高，而且還在不斷解鎖新的能力。
在探討現(xiàn)有 AI 模型的局限時，谷歌人工智能主管 Jeff Dean 曾經(jīng)說過，今天的人工智能系統(tǒng)總是從頭開始學習新問題。最終，我們?yōu)閿?shù)千個單獨的任務開發(fā)了數(shù)千個模型。以這種方式學習每項新任務不僅需要更長的時間，而且還需要更多的數(shù)據(jù)，效率非常低。
在 Jeff Dean 看來，理想的發(fā)展模式應該是訓練一個模型來做成千上萬件事情。為了實現(xiàn)這一愿景，他所在的團隊去年提出了一種名叫「Pathways」的通用 AI 架構。Jeff Dean 介紹說，Pathways 旨在用一個架構同時處理多項任務，并且擁有快速學習新任務、更好地理解世界的能力。前段時間，該團隊終于公布了 Pathways 的論文。
論文寫道，PATHWAYS 使用了一種新的異步分布式數(shù)據(jù)流設計。這種設計允許 PATHWAYS 采用單控制器模型，從而更容易表達復雜的新并行模式。實驗結果表明，當在 2048 個 TPU 上運行 SPMD（single program multiple data）計算時，PATHWAYS 的性能（加速器利用率接近 100%）可以媲美 SOTA 系統(tǒng)。
谷歌 Pathways 系統(tǒng)架構概覽。
有了強大的系統(tǒng)，接下來就是訓練模型了。
在剛剛公布的論文——「PaLM: Scaling Language Modeling with Pathways」中，谷歌宣布，他們用 Pathways 系統(tǒng)訓練了一個 5400 億參數(shù)的大型語言模型——PaLM（Pathways Language Model）。

論文鏈接：https://storage.googleapis.com/pathways-language-model/PaLM-paper.pdf
這是一個只有****的密集 Transformer 模型。為了訓練這個模型，谷歌動用了 6144 塊 TPU，讓 Pathways 在兩個 Cloud TPU v4 Pods 上訓練 PaLM。
強大的系統(tǒng)和算力投入帶來了驚艷的結果。研究者在數(shù)百個語言理解和生成任務上評估了 PaLM，發(fā)現(xiàn)它在大多數(shù)任務上實現(xiàn)了 SOTA 少樣本學習性能，可以出色地完成笑話解讀、bug 修復、從表情符號中猜電影等語言、代碼任務。

有網(wǎng)友感嘆說，「終于知道谷歌開發(fā) TPU 是用來干嘛的了。這種級別的自然語言理解，一旦被應用程序所利用，并變得足夠高效并廣泛使用。這將徹底改變谷歌所做的一切。擁有能夠實現(xiàn)這一目標的芯片是非常有價值的，而那些陷入購買或租用英偉達芯片的公司將處于不利地位?！?/span>

PaLM 架構概覽
PaLM 只包含****（每個時間步只能關注自身和過去的時間步），對一種標準的 Transformer 架構（(Vaswani et al., 2017)）做出了如下更改：
SwiGLU 激活
研究者使用 SwiGLU 激活 (Swish(xW) · xV) 用于 MLP 中間激活，因為研究表明，與標準 ReLU、GeLU 或 Swish 激活相比，SwiGLU 激活能顯著提高質量。注意，在 MLP 中，這確實需要三個矩陣乘法，而不是兩個，但 Shazeer (2020) 在計算等效實驗中證明了質量的提升。
并行層
研究者在每個 Transformer 模塊中使用「并行」方法，而不是標準的「串行」方法。具體來說，標準方法可以寫成：

并行方法可以寫成

由于 MLP 和注意力輸入矩陣乘法可以融合，這里的并行方法可以讓大規(guī)模訓練速度提升 15%。消融實驗顯示，在 8B 的規(guī)模下，質量下降很小，但在 62B 規(guī)模下，質量沒有下降，因此研究者推斷，并行層的影響會在 540B 規(guī)模下達到 quality neutral。
多查詢（Multi-Query）注意力
標準 Transformer 方法使用 k 個注意力頭，其中每個時間步長的輸入向量被線性投影成形狀 [k，h] 的「查詢」、「鍵」和「值」張量，其中 h 是注意力頭大小。這里，鍵 / 值投影對于每個頭是共享的，即「鍵」和「值」被投影到[1，h]，但是「查詢」仍然被投影到形狀[k，h]。此前有研究表明，這對模型質量和訓練速度的影響呈中性，但在自回歸解碼時間上可以帶來顯著的成本節(jié)約。這是因為在自回歸解碼過程中，標準多頭注意力在加速器硬件上的效率很低，因為鍵 / 值張量不在實例之間共享，并且一次只有單個 token 被解碼。
RoPE 嵌入
研究者使用了 RoPE 嵌入而不是絕對或相對位置嵌入，因為 RoPE 嵌入已被證明在長序列長度上具有更好的性能。
共享輸入 - 輸出嵌入
研究者共享了輸入和輸出嵌入矩陣，這是在過去的工作中經(jīng)常做的（但不是普遍的）。
No Biases
研究者在任何密集核或層 norm 中都沒有使用 biases。他們發(fā)現(xiàn)，這可以增加大型模型的訓練穩(wěn)定性。
詞匯表
研究者使用了具有 256k token 的 SentencePiece 詞匯表，選擇這個詞匯表是為了支持訓練語料庫中的多種語言（沒有過多的分詞）。詞匯表是從訓練數(shù)據(jù)中生成的，研究者發(fā)現(xiàn)這提高了訓練效率。
用 Pathways 訓練一個 5400 億參數(shù)的語言模型
PaLM 是谷歌首次大規(guī)模使用 Pathways 系統(tǒng)將訓練擴展到 6144 塊芯片，這是迄今為止用于訓練的基于 TPU 的最大系統(tǒng)配置。研究者在 Pod 級別上跨兩個 Cloud TPU v4 Pods 使用數(shù)據(jù)并行對訓練進行擴展，同時在每個 Pod 中使用標準數(shù)據(jù)和模型并行。與以前的大多數(shù) LLM 相比，這是一個顯著的規(guī)模增長。
PaLM 實現(xiàn)了 57.8% 的硬件 FLOPs 利用率的訓練效率，這是 LLM 在這個規(guī)模上實現(xiàn)的最高效率。為了達到這一水平，研究者將并行策略和 Transformer 塊的重新設計結果相結合，這使得注意力層和前饋層并行計算成為可能，從而實現(xiàn)了 TPU 編譯器優(yōu)化帶來的加速。
PaLM 使用英語和多語言數(shù)據(jù)集進行訓練，包括高質量的 web 文檔、書籍、維基百科、對話和 GitHub 代碼。研究者還創(chuàng)建了一個「無損（lossless）」詞匯表，它保留了所有空格（對于代碼來說尤其重要），將詞匯表之外的 Unicode 字符拆分成字節(jié)，并將數(shù)字拆分成單獨的 token，每個 token 對應一個數(shù)字。
突破性的語言、推理和代碼生成能力
PaLM 在許多非常困難的任務上顯示出了突破性的能力，包括語言理解、生成、推理和代碼等相關任務。
語言理解與生成

在 29 項基于英語的 NLP 任務上，PaLM 540B 的性能比之前的 SOTA 結果有所提高。
除了英語 NLP 任務外，PaLM 在包括翻譯在內(nèi)的多語言 NLP 基準測試中也表現(xiàn)出強大的性能，盡管它只有 22% 的訓練語料庫是非英語的。
研究者還在 Beyond the Imitation Game Benchmark (BIG-bench) 上探索了 PaLM 新的和未來功能，這是一個近期發(fā)布的套件，包含 150 多個新的語言建模任務。在這個過程中， PaLM 與 Gopher 和 Chinchilla 的性能進行了比較，涉及這些任務的 58 個公共子集。
有趣的是，研究者注意到，作為 scale 函數(shù)（function of scale）的 PaLM 的性能遵循與先前模型類似的對數(shù)線性表現(xiàn)，這表明 scale 的性能改進尚未趨于平穩(wěn)。PaLM 540B 5-shot 的表現(xiàn)也優(yōu)于解決相同任務的人類平均表現(xiàn)。
PaLM 在 58 個 BIG-bench 任務的子集上的 scaling behavior。
PaLM 在幾個 BIG-bench 任務中展示了令人印象深刻的自然語言理解和生成能力。例如，該模型可以區(qū)分因果關系，理解上下文中的概念組合，甚至可以從表情符號中猜測電影。

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

推理
通過將模型 scale 與 chain-of-thought prompting 相結合，PaLM 在需要多步驟算術或常識推理的推理任務上展示出了突破性的能力。以往諸如 Gopher 這樣的大型語言模型在提高性能方面從模型 scale 中獲益較少。
小學數(shù)學問題示例中的標準 prompting 與 chain-of-thought prompting。Chain-of-thought prompting 將多步驟推理問題的 prompt 分解為中間步驟（黃色部分），類似于人類處理它的方式。

研究者在三個算術數(shù)據(jù)集和兩個常識推理數(shù)據(jù)集上觀察到了 PaLM 540B 在 chain-of-thought prompt 加持下的強大性能。例如，借助 8-shot prompt，PaLM 解決了 GSM8K 中 58% 的問題，這是一個包含數(shù)千個具有挑戰(zhàn)性的小學水平數(shù)學問題的基準，超過了之前 GPT-3 175B 微調(diào)模型（訓練集包含 7500 個問題，并與外部計算器和驗證器相結合）獲得的 55% 的最高分。
這個新的得分值得關注，因為它接近 60% 的 9 到 12 歲兒童解決問題的水平，這些兒童正是問題集的目標受眾。研究者猜測，PaLM 詞匯表中數(shù)字的獨立編碼有助于實現(xiàn)這些性能改進。
值得注意的是，PaLM 甚至可以為需要多步驟邏輯推理、世界認知和深度語言理解的復雜組合的場景生成明確的解釋。例如，它可以為網(wǎng)絡上搜不到的新笑話提供高質量的解釋。

PaLM 用 two-shot prompts 解釋了一個原創(chuàng)笑話。
代碼生成
大型語言模型已被證明可以很好地推廣到編碼任務，比如在給定自然語言描述（文本到代碼）的情況下編寫代碼，將代碼從一種語言翻譯成另一種語言，以及修復編譯錯誤（代碼到代碼）。
PaLM 540B 在單個模型中顯示了橫跨編碼任務和自然語言任務的強大性能，即使它在預訓練數(shù)據(jù)集中只有 5% 的代碼。具體而言，PaLM 540B 的 few-shot 性能十分顯著，與經(jīng)過微調(diào)的 Codex 12B 相當，同時使用的 Python 訓練代碼減少到了 50 分之一。這一結果印證了之前的發(fā)現(xiàn)，即較大的模型比較小的模型更高效，因為它們可以更好地從其他編程語言和自然語言數(shù)據(jù)中實現(xiàn)遷移學習。

PaLM 540B 微調(diào)模型在文本到代碼任務（例如 GSM8K - Python 和 HumanEval）和代碼到代碼任務（例如 Transcoder）上的示例。
此外，通過在純 Python 代碼數(shù)據(jù)集上微調(diào) PaLM ，模型進一步提高了性能，團隊稱之為 PaLM-Coder。如下圖所示，PaLM-Coder 接到了一個名為 DeepFix 的示例代碼修復任務，目標是修改最初損壞的 C 程序直到它們編譯成功，PaLM-Coder 540B 展示了令人印象深刻的性能，實現(xiàn)了 82.1% 的編譯率，優(yōu)于之前 71.7% 的 SOTA 結果。這為修復軟件開發(fā)過程中出現(xiàn)的更復雜的錯誤提供了機會。

DeepFix 代碼修復任務的示例。經(jīng)過微調(diào)的 PaLM-Coder 540B 將編譯錯誤（左）修復為可編譯的代碼版本（右）。
從倫理方面考慮，最近的研究強調(diào)了受過網(wǎng)絡文本訓練的 LLM 相關的各種潛在風險。通過模型卡片和數(shù)據(jù)表等透明工件分析并記錄這些潛在的不良風險是至關重要的，其中還包括有關預期用途和測試的信息。為此，谷歌的論文提供了數(shù)據(jù)表、模型卡片和 Responsible AI 基準測試結果，并報告了對數(shù)據(jù)集和模型輸出的全面分析，以發(fā)現(xiàn)偏差和風險。雖然分析有助于概述模型的一些潛在風險，但特定領域和任務的分析對于真正校準、情境化和減輕可能的危害至關重要。進一步了解這些模型的風險和收益是正在進行的研究的主題，同時開發(fā)可擴展的解決方案，防止惡意使用語言模型。
參考鏈接：https://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html