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

隨著規(guī)模的增加，模型在處理多個任務(wù)時的性能逐漸提高，而且還在不斷解鎖新的能力。

在探討現(xiàn)有 AI 模型的局限時，谷歌人工智能主管 Jeff Dean 曾經(jīng)說過，今天的人工智能系統(tǒng)總是從頭開始學(xué)習(xí)新問題。最終，我們?yōu)閿?shù)千個單獨的任務(wù)開發(fā)了數(shù)千個模型。以這種方式學(xué)習(xí)每項新任務(wù)不僅需要更長的時間，而且還需要更多的數(shù)據(jù)，效率非常低。

在 Jeff Dean 看來，理想的發(fā)展模式應(yīng)該是訓(xùn)練一個模型來做成千上萬件事情。為了實現(xiàn)這一愿景，他所在的團(tuán)隊去年提出了一種名叫「Pathways」的通用 AI 架構(gòu)。Jeff Dean 介紹說，Pathways 旨在用一個架構(gòu)同時處理多項任務(wù)，并且擁有快速學(xué)習(xí)新任務(wù)、更好地理解世界的能力。前段時間，該團(tuán)隊終于公布了 Pathways 的論文。

論文寫道，PATHWAYS 使用了一種新的異步分布式數(shù)據(jù)流設(shè)計。這種設(shè)計允許 PATHWAYS 采用單控制器模型，從而更容易表達(dá)復(fù)雜的新并行模式。實驗結(jié)果表明，當(dāng)在 2048 個 TPU 上運行 SPMD(single program multiple data)計算時，PATHWAYS 的性能(加速器利用率接近 100%)可以媲美 SOTA 系統(tǒng)。

谷歌 Pathways 系統(tǒng)架構(gòu)概覽。

有了強(qiáng)大的系統(tǒng)，接下來就是訓(xùn)練模型了。

在剛剛公布的論文——「PaLM: Scaling Language Modeling with Pathways」中，谷歌宣布，他們用 Pathways 系統(tǒng)訓(xùn)練了一個 5400 億參數(shù)的大型語言模型——PaLM(Pathways Language Model)。

論文鏈接：https://storage.googleapis.com/pathways-language-model/PaLM-paper.pdf

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

強(qiáng)大的系統(tǒng)和算力投入帶來了驚艷的結(jié)果。研究者在數(shù)百個語言理解和生成任務(wù)上評估了 PaLM，發(fā)現(xiàn)它在大多數(shù)任務(wù)上實現(xiàn)了 SOTA 少樣本學(xué)習(xí)性能，可以出色地完成笑話解讀、bug 修復(fù)、從表情符號中猜電影等語言、代碼任務(wù)。

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

PaLM 架構(gòu)概覽

PaLM 只包含解碼器(每個時間步只能關(guān)注自身和過去的時間步)，對一種標(biāo)準(zhǔn)的 Transformer 架構(gòu)((Vaswani et al., 2017))做出了如下更改：

SwiGLU 激活

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

并行層

研究者在每個 Transformer 模塊中使用「并行」方法，而不是標(biāo)準(zhǔn)的「串行」方法。具體來說，標(biāo)準(zhǔn)方法可以寫成：

并行方法可以寫成

由于 MLP 和注意力輸入矩陣乘法可以融合，這里的并行方法可以讓大規(guī)模訓(xùn)練速度提升 15%。消融實驗顯示，在 8B 的規(guī)模下，質(zhì)量下降很小，但在 62B 規(guī)模下，質(zhì)量沒有下降，因此研究者推斷，并行層的影響會在 540B 規(guī)模下達(dá)到 quality neutral。

多查詢(Multi-Query)注意力

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

RoPE 嵌入

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

共享輸入 - 輸出嵌入

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

No Biases

研究者在任何密集核或?qū)?norm 中都沒有使用 biases。他們發(fā)現(xiàn)，這可以增加大型模型的訓(xùn)練穩(wěn)定性。

詞匯表

研究者使用了具有 256k token 的 SentencePiece 詞匯表，選擇這個詞匯表是為了支持訓(xùn)練語料庫中的多種語言(沒有過多的分詞)。詞匯表是從訓(xùn)練數(shù)據(jù)中生成的，研究者發(fā)現(xiàn)這提高了訓(xùn)練效率。

用 Pathways 訓(xùn)練一個 5400 億參數(shù)的語言模型

PaLM 是谷歌首次大規(guī)模使用 Pathways 系統(tǒng)將訓(xùn)練擴(kuò)展到 6144 塊芯片，這是迄今為止用于訓(xùn)練的基于 TPU 的最大系統(tǒng)配置。研究者在 Pod 級別上跨兩個 Cloud TPU v4 Pods 使用數(shù)據(jù)并行對訓(xùn)練進(jìn)行擴(kuò)展，同時在每個 Pod 中使用標(biāo)準(zhǔn)數(shù)據(jù)和模型并行。與以前的大多數(shù) LLM 相比，這是一個顯著的規(guī)模增長。

PaLM 實現(xiàn)了 57.8% 的硬件 FLOPs 利用率的訓(xùn)練效率，這是 LLM 在這個規(guī)模上實現(xiàn)的最高效率。為了達(dá)到這一水平，研究者將并行策略和 Transformer 塊的重新設(shè)計結(jié)果相結(jié)合，這使得注意力層和前饋層并行計算成為可能，從而實現(xiàn)了 TPU 編譯器優(yōu)化帶來的加速。

PaLM 使用英語和多語言數(shù)據(jù)集進(jìn)行訓(xùn)練，包括高質(zhì)量的 web 文檔、書籍、維基百科、對話和 GitHub 代碼。研究者還創(chuàng)建了一個「無損(lossless)」詞匯表，它保留了所有空格(對于代碼來說尤其重要)，將詞匯表之外的 Unicode 字符拆分成字節(jié)，并將數(shù)字拆分成單獨的 token，每個 token 對應(yīng)一個數(shù)字。

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

PaLM 在許多非常困難的任務(wù)上顯示出了突破性的能力，包括語言理解、生成、推理和代碼等相關(guān)任務(wù)。

語言理解與生成

在 29 項基于英語的 NLP 任務(wù)上，PaLM 540B 的性能比之前的 SOTA 結(jié)果有所提高。

除了英語 NLP 任務(wù)外，PaLM 在包括翻譯在內(nèi)的多語言 NLP 基準(zhǔn)測試中也表現(xiàn)出強(qiáng)大的性能，盡管它只有 22% 的訓(xùn)練語料庫是非英語的。

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

有趣的是，研究者注意到，作為 scale 函數(shù)(function of scale)的 PaLM 的性能遵循與先前模型類似的對數(shù)線性表現(xiàn)，這表明 scale 的性能改進(jìn)尚未趨于平穩(wěn)。PaLM 540B 5-shot 的表現(xiàn)也優(yōu)于解決相同任務(wù)的人類平均表現(xiàn)。

PaLM 在 58 個 BIG-bench 任務(wù)的子集上的 scaling behavior。

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

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

推理

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

小學(xué)數(shù)學(xué)問題示例中的標(biāo)準(zhǔn) prompting 與 chain-of-thought prompting。Chain-of-thought prompting 將多步驟推理問題的 prompt 分解為中間步驟(黃色部分)，類似于人類處理它的方式。

研究者在三個算術(shù)數(shù)據(jù)集和兩個常識推理數(shù)據(jù)集上觀察到了 PaLM 540B 在 chain-of-thought prompt 加持下的強(qiáng)大性能。例如，借助 8-shot prompt，PaLM 解決了 GSM8K 中 58% 的問題，這是一個包含數(shù)千個具有挑戰(zhàn)性的小學(xué)水平數(shù)學(xué)問題的基準(zhǔn)，超過了之前 GPT-3 175B 微調(diào)模型(訓(xùn)練集包含 7500 個問題，并與外部計算器和驗證器相結(jié)合)獲得的 55% 的最高分。

這個新的得分值得關(guān)注，因為它接近 60% 的 9 到 12 歲兒童解決問題的水平，這些兒童正是問題集的目標(biāo)受眾。研究者猜測，PaLM 詞匯表中數(shù)字的獨立編碼有助于實現(xiàn)這些性能改進(jìn)。

值得注意的是，PaLM 甚至可以為需要多步驟邏輯推理、世界認(rèn)知和深度語言理解的復(fù)雜組合的場景生成明確的解釋。例如，它可以為網(wǎng)絡(luò)上搜不到的新笑話提供高質(zhì)量的解釋。

PaLM 用 two-shot prompts 解釋了一個原創(chuàng)笑話。

代碼生成

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

PaLM 540B 在單個模型中顯示了橫跨編碼任務(wù)和自然語言任務(wù)的強(qiáng)大性能，即使它在預(yù)訓(xùn)練數(shù)據(jù)集中只有 5% 的代碼。具體而言，PaLM 540B 的 few-shot 性能十分顯著，與經(jīng)過微調(diào)的 Codex 12B 相當(dāng)，同時使用的 Python 訓(xùn)練代碼減少到了 50 分之一。這一結(jié)果印證了之前的發(fā)現(xiàn)，即較大的模型比較小的模型更高效，因為它們可以更好地從其他編程語言和自然語言數(shù)據(jù)中實現(xiàn)遷移學(xué)習(xí)。

PaLM 540B 微調(diào)模型在文本到代碼任務(wù)(例如 GSM8K - Python 和 HumanEval)和代碼到代碼任務(wù)(例如 Transcoder)上的示例。

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

DeepFix 代碼修復(fù)任務(wù)的示例。經(jīng)過微調(diào)的 PaLM-Coder 540B 將編譯錯誤(左)修復(fù)為可編譯的代碼版本(右)。

從倫理方面考慮，最近的研究強(qiáng)調(diào)了受過網(wǎng)絡(luò)文本訓(xùn)練的 LLM 相關(guān)的各種潛在風(fēng)險。通過模型卡片和數(shù)據(jù)表等透明工件分析并記錄這些潛在的不良風(fēng)險是至關(guān)重要的，其中還包括有關(guān)預(yù)期用途和測試的信息。為此，谷歌的論文提供了數(shù)據(jù)表、模型卡片和 Responsible AI 基準(zhǔn)測試結(jié)果，并報告了對數(shù)據(jù)集和模型輸出的全面分析，以發(fā)現(xiàn)偏差和風(fēng)險。雖然分析有助于概述模型的一些潛在風(fēng)險，但特定領(lǐng)域和任務(wù)的分析對于真正校準(zhǔn)、情境化和減輕可能的危害至關(guān)重要。進(jìn)一步了解這些模型的風(fēng)險和收益是正在進(jìn)行的研究的主題，同時開發(fā)可擴(kuò)展的解決方案，防止惡意使用語言模型。