谷歌的下一代架構(gòu) Pathways 已經(jīng)用來(lái)訓(xùn)練大模型了。

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

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

有網(wǎng)友感嘆說(shuō)，「終于知道谷歌開發(fā) TPU 是用來(lái)干嘛的了。這種級(jí)別的自然語(yǔ)言理解，一旦被應(yīng)用程序所利用，并變得足夠高效并廣泛使用。這將徹底改變谷歌所做的一切。擁有能夠?qū)崿F(xiàn)這一目標(biāo)的芯片是非常有價(jià)值的，而那些陷入購(gòu)買或租用英偉達(dá)芯片的公司將處于不利地位。」

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

并行方法可以寫成

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

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

PaLM 540B 在 BIG-bench 任務(wù)上的 1-shot 性能展示：標(biāo)記因果關(guān)系、概念理解、從表情符號(hào)中猜測(cè)電影以及查找同義詞和反事實(shí)。

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

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

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

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

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