語言模型的演算法進展與Scaling Laws

引言

近年來,語言模型(Language Models, LLMs)在自然語言處理領域取得了顯著的進展。LLM的預訓練成本高昂,但scaling laws使我們能夠從成本較低的訓練過程中準確預測更大規模訓練的性能。最近的研究大幅改善了我們對LLM scaling laws的實際理解,可以在模型參數數量和訓練tokens與參數之比兩方面進行外推。

LLM發展背景

先前的研究發現了語言模型的scaling laws,顯示隨著模型變大和預訓練資料增加,模型性能會提高。然而,這些研究主要關注計算最佳情況,且性能評估多以下一個token預測損失為主。在實務應用中,語言模型很少以最佳計算方式訓練,通常使用更多tokens訓練較小模型(過度訓練)以降低推理成本,且更關注下游任務實際表現。這引出疑問:scaling laws是否適用於這些更實用的情況?

演算法進展與規模擴展

2014至2024年間,演算法改進和運算能力提升顯著改善了語言模型性能。新演算法帶來的進步相當於每5至14個月運算能力翻倍。儘管近年來性能改進更多源自於大規模計算擴展,但Transformer架構和Chinchilla縮放規律等創新分別帶來了近兩年和8到16個月的算法進步。

過度訓練的LLM

在實務中,通常採用過度訓練模式,即以更多tokens訓練較小模型,以在達到有競爭力指標的同時降低推理成本。研究發現,過度訓練模型同樣遵循冪律scaling laws,token倍增器M值的改變只影響y截距,形成一組平行線。因此,我們可從scaling laws中預測更大過度訓練模型的性能。

下游任務性能預測 

利用scaling laws,我們可繪製LLM在下游評估任務上的平均top-1錯誤率,發現其隨著模型在C4資料集上的驗證損失呈指數下降。這揭示了下游任務性能的明確趨勢,據此提出了top-1錯誤的scaling laws,可用較少計算預測大模型在下游任務上的錯誤率。

實證評估與應用價值

通過訓練不同規模的LLM並使用scaling laws外推,實驗顯示我們可以較低成本地預測大規模訓練的驗證損失和下游任務錯誤率,相對誤差分別僅有0.7%和3.6%。這證實了scaling laws在實踐中預測大規模訓練性能的價值。

個人見解

LLM的快速發展得益於算法創新和計算規模擴張的共同推動,但後者的貢獻更為關鍵。Scaling laws作為一種低成本預測工具,在提高LLM研發效率方面有重要價值。隨著對scaling laws理解的加深,相信我們能更好地指導語言模型的發展方向與資源配置。不過,當前的scaling laws仍有局限,如對特定創新影響的估計能力不足、難以區分算法改進和計算擴展的貢獻等。未來還需要更多更好的數據,來建立更精準和全面的scaling laws模型。

此外,我認為LLM發展的關鍵是平衡效能提升與成本控制。過度訓練等策略雖以效率為代價換取效能,但對於實際部署至關重要。如何在scaling laws框架下找到最優的訓練配置,是一個值得深入探索的問題。總的來說,語言模型的進展令人振奮,但要實現更通用和強大的AI,仍需在算法和計算兩個維度協同發力,並以切實應用為導向,審慎權衡效能和成本。Scaling laws作為連接理論和實踐的紐帶,在其中扮演著不可或缺的角色。

LLaVA Llama-3 和 Phi-3 現已在 MLX 上推出

 LLaVA Llama-3 和 Phi-3 現已在 MLX 上推出 

現在可以在你的 Mac 上在本地運行推理了。

pip install -U mlx-vlm

在 M3 Max 上可以達到每秒約 50 個 token 的速度。

LLaVA (Large Language and Vision Assistant)是一個多模態的大型語言模型,可以同時處理文字和圖像。它是基於Meta的LLaMA架構訓練而成。現在LLaVA的Llama-3和Phi-3版本都已經在MLX平台上推出。

使用pip install -U mlx-vlm指令,就可以在Mac上安裝最新版的MLX Python套件,直接在本地端運行LLaVA模型進行推理。

就效能而言,在配備M3 Max晶片的Mac機器上,推理速度可以達到每秒約50個token,相當不錯。

感興趣的開發者可以到MLX平台上查看LLaVA Llama-3和Phi-3的模型說明頁(Model card),了解更多模型的細節、性能指標和使用案例等資訊。

LLaVA模型在MLX平台上的發布,為Mac用戶帶來了方便的本地端推理選擇。這對於某些對隱私和安全有較高要求的應用場景來說很有幫助。期待看到社群能夠利用LLaVA打造出更多有趣且實用的多模態AI應用!

語言模型推論的光速極限

在開發 calm 的過程中,這是一個從頭開始打造、精簡且快速的transformer語言模型 CUDA 推論實作,一個關鍵的考量是確立推論過程的光速極限,並衡量相對於該極限的進展。在本文中,我們將討論這個理論上限及其意涵。              --Arseny Kapoulkine--

如果您有興趣瞭解更多推導過程和一些圖表,這個 Python notebook 做了相同的建模。
https://github.com/zeux/calm

推論機制

當語言模型 (特指 decoder-only 的文字 transformer 模型,但本文其餘部分將統稱為語言模型) 生成字詞時,它是一次一個字詞地生成;語言模型可以被理解為一個函數,它接收一個字詞作為輸入,並產生一個包含詞彙表中所有字詞 (通常有 5 萬到 25 萬個字詞,每個字詞由幾個字母組成) 機率的陣列。接著,程式根據這些機率從所有字詞中取樣,產生下一個字詞,然後重複這個過程。這意味著在生成一段文字序列時,沒有辦法平行化 ─ 生成過程可以被視為一次一個字詞。

語言模型在處理一個字詞時,大致上會進行兩種運算:矩陣向量乘法,其中一個大矩陣 (如 8192x8192) 乘以一個向量產生另一個向量;以及注意力 (attention) 的計算。在生成過程中,模型不僅可以看到當前字詞的狀態,還可以看到序列中所有先前字詞的內部狀態 ─ 包括使用者在提示中寫的字詞,以及模型本身生成的字詞。這些狀態儲存在KV 快取(key-value cache) 的結構中,本質上是每個先前位置的一組 key 和 value 向量。注意力機制為當前字詞生成一個 query 向量,計算它與所有先前位置的所有 key 向量之間的內積,然後將結果標準化,並透過對所有先前位置的所有 value 向量進行加權求和來計算一個 value 向量,其中內積作為權重。

現在,矩陣向量乘法和注意力計算有一個重要的共通點:對於從矩陣或 KV 快取讀取的每個元素,我們只需要做非常少量的浮點運算。矩陣向量乘法對每個矩陣元素做一次乘法和加法 (2 個 FLOP);注意力計算對每個 key 元素做一次乘法和加法來計算內積,對每個 value 元素做一次乘法和加法來計算加權和。

現代 CPU 和 GPU 的 ALU 運算(乘法、加法)速度遠高於從記憶體讀取輸入的速度。例如:

- AMD Ryzen 7950X 有 67 GB/s 的記憶體頻寬和 2735 GFLOPS,FLOP與byte之比為 40:1。

- NVIDIA GeForce RTX 4090 有 1008 GB/s 的記憶體頻寬和 83 TFLOPS,FLOP與byte之比為 82:1。 

- NVIDIA H100 SXM(一張資料中心用的卡)有 3350 GB/s 的記憶體頻寬和 67 TFLOPS,FLOP與byte之比似乎只有 20:1;然而,對於類似矩陣乘法的問題,tensor core 在沒有稀疏性的情況下可提供約 494 TFLOPS,使得 FLOP 與 byte 之比達到 147:1。

對於 FP16 或 FP8 等更小的浮點數,情況會變得更糟:H100 的 tensor core 在密集 FP8 矩陣上的理論吞吐量為 1979 TFLOPS,將 FLOP 與 byte 之比提高到 590:1。不用說,在任何這些配置下,無論是否使用 tensor core,以及無論使用何種浮點格式,ALU 都是非常充足的。

因此,任何只需要對每個元素做兩次運算的問題,都一定會受到頻寬的限制,我們應該能夠根據模型配置、KV快取的大小和可用頻寬來估計推論過程所需的最小時間。

Mistral 的光速極限

在不過於深入討論確切公式和矩陣的情況下,讓我們看看像 Mistral 7B 這樣的模型,它有 72 億個參數 (因此所有矩陣元素的總數是 72 億)。

參數的組成如下:

- 4096 * 32000 = 1.31 億個嵌入矩陣的參數;這個矩陣不會用在矩陣向量乘法中,因為對每個字詞只會讀取矩陣的一行,所以我們不會將其納入頻寬計算。

- 32 * (4096 * (128 * 32 + 128 * 8 * 2) + 4096 * 128 * 32) = 13.42 億個用於計算注意力相關向量的參數。 

- 32 * (4096 * 14336 * 3) = 56.37 億個透過前饋網路轉換隱藏狀態的參數。

- 4096 * 32000 = 1.31 億個將隱藏狀態轉換為字詞機率的參數;這個矩陣會用在矩陣乘法中,與嵌入矩陣不同。

這加起來大約有 71.11 億個在矩陣乘法中使用的「活躍」參數。如果模型對矩陣元素使用 FP16,我們最終必須為每個字詞讀取約 14.2 GB 的資料。此外,雖然每個矩陣在為下一個字詞執行推論時都會再次使用,但快取的大小通常在幾十 MB 左右,因此我們可以假設這個過程無法比記憶體頻寬更快,因為權重在推論之間不會留在快取中。

這涵蓋了矩陣運算;注意力計算需要讀取直到當前字詞的 KV 快取,因此讀取的資料量取決於模型在生成新字詞時看到的字詞數量 ─ 包括系統提示 (通常對使用者隱藏)、使用者提示、先前的模型輸出,並且可以包括較長聊天階段的多個使用者提示。

對於 Mistral,KV 快取為每層的每個 key 儲存 8 個 128 維的向量,為每層的每個 value 儲存 8 個 128 維的向量,總共是每個字詞 32 * 128 * 8 * 2 = 65K 個元素;如果 KV 快取使用 FP16 儲存單個元素,那麼對於第 P 個字詞,我們需要讀取 P * 130 KB 的記憶體 ─ 例如,第 1000 個字詞需要從 KV 快取讀取 130 MB 的資料。

有了這些數字,現在很容易計算推論所需的最小時間。例如,在 NVIDIA RTX 4090 (1008 GB/s) 上,讀取 14.2 GB 大約需要 14.1 毫秒,因此我們可以預期對位置數較低的字詞,每個字詞大約需要 14.1 毫秒 (KV 快取的影響可以忽略不計)。如果我們使用 8 位元的權重,我們需要讀取 7.1 GB,這大約需要 7.0 毫秒。這些是下限 ─ 它們代表了理論上每個字詞可能的最小時間。

理論上限有用嗎?

我們做了一堆數學計算,得到了一些告訴我們推論速度無法超過給定閾值的數字 ─ 這有用嗎?讓我們看看為什麼它可能有用的幾個原因。

要真正達到該時間,你需要高品質的軟體實作,以及能夠達到理論峰值頻寬的硬體。這意味著,如果給定的實作遠低於最佳數字,就值得進行調查:效能可能遺留在軟體或硬體方面。例如,在 RTX 4090 上,使用 16 位元權重時,calm 可以達到每個字詞約 15.4 毫秒;使用 8 位元權重時,可以達到每個字詞約 7.8 毫秒 ─ 這大約是理論上可能效能的 90%。在使用 CPU 推論的 Apple M2 Air 上,calm 和 llama.cpp 都只達到了理論 100 GB/s 頻寬的約 65%,這表明引用的峰值頻寬只有在 iGPU 的幫助下才能完全利用。

頻寬與每個元素使用的位元組數呈線性關係;這意味著我們既可以從更小的權重格式 (量化) 估計理論效益,也可以透過將實際效能與理論極限進行比較來驗證實作的品質。例如,在 RTX 4090 上,使用 16 位元權重時,llama.cpp 可以達到每個字詞約 17.1 毫秒 (峰值的 82%);使用 8.5 位元權重時,可以達到每個字詞約 10.3 毫秒 (峰值的 71%);使用 4.5 位元權重時,可以達到每個字詞約 6.7 毫秒 (峰值的 58%),這表明對更小的格式還有很大的優化空間。

除了提供解碼時間的下限之外,上面的建模還表明,推論過程嚴重低估了 ALU 單元的利用率。要解決這個問題,需要改變 FLOP 與 byte 的平衡;投機解碼等技術試圖幫助解決這個問題,但對於多使用者的例子,我們可以注意到,當同時處理多個使用者請求時,我們可以在同一個矩陣上同時執行多個矩陣向量乘法(也就是矩陣矩陣乘法!)─ 矩陣矩陣乘法的最佳實作會在矩陣足夠大時變成ALU bound。這就是為什麼這種 ALU 與 byte 的不平衡對於生產環境的推論系統而言不是一個關鍵問題 ─ 當你要求 ChatGPT 幫助一項任務時,你的請求會與同一 GPU 上的許多其他請求同時進行評估,頻寬會被更有效地利用。至關重要的是,請求批次處理通常無法幫助減少 KV 快取的頻寬需求(除非這些請求共享非常大的前綴),因為 KV 快取的大小和頻寬會隨著請求數量增加,而權重矩陣保持不變。

像 Mixtral 這樣的專家混合模型有略微不同的縮放特性:批次處理最初只會增加所需的頻寬,但一旦專家利用率變得顯著,推論就會越來越受 ALU 的限制。

最後,如果無法進行批次處理,頻寬就成為預期推論效能的關鍵估算指標,它在不同的模型變體/設備類型或架構中保持不變,你可以用它來決定需要使用的硬體。例如,NVIDIA RTX 4080 有 716 GB/s 的頻寬,所以你可以預期它執行語言模型推論的速度大約是 RTX 4090 的 0.7 倍 ─ 這可能與其他工作負載(如遊戲、光線追蹤或其他類型神經網路的推論)的相對效能不同!

結論

對於像這樣計算量和記憶體訪問量已知的問題,將理論光速建模作為基準非常重要,因為它有助於驗證實作的品質並預測架構變化的影響。

理想情況下,你的推論實作應該仔細計算實現的有效頻寬,你應該在效能分析時將其作為主要指引 ─ 因為你知道這個值的極限!不過,一定要仔細計算 ─ calm 曾多次出現架構上的特殊情況,導致計算出的頻寬略有偏差:)

附錄:群組查詢注意力 (Group Query Attention)

Mistral-7B 是一個非常平衡的模型;在上面的所有計算中,KV 快取似乎並不是成本結構中不可或缺的部分。其中一個原因是相對較短的上下文(Mistral-7B 使用了視窗注意力,將頻寬消耗限制在 4096 個字詞的視窗中),但另一個可能更重要的原因,是使用了群組查詢注意力(Group Query Attention, GQA)。

在群組查詢注意力中(4倍比率),為了產生4個內積,我們不是使用4個查詢向量並與4個對應的鍵向量計算內積,而是取一個鍵向量但使用4個查詢向量,並執行4次內積運算。這使我們能夠減少KV快取的大小和所需的頻寬 ─ 我們不再是從KV快取中讀取每個元素並只對其進行一次乘加運算,而是現在進行4次運算,這在一定程度上重新平衡了ALU與頻寬的比率,對我們有利。

這對KV快取的記憶體大小也至關重要,但對於如此短的上下文可能並不明顯:4096個字詞的上下文在Mistral中佔用0.5 GiB,但沒有GQA的類似模型(如Llama 7B)「只」需要2 GiB。讓我們看看最近一個沒有使用GQA的模型,Cohere的Command-R。

該模型本身有約350億個參數,因此在每個權重16位元的情況下,我們在推論期間需要為每個字詞讀取70 GB的權重。對於每個字詞,它需要在KV快取中儲存40 * 128 * 64 * 2 = 655K個元素,在每個元素16位元的情況下,每個字詞是1.3 MB。

因此,4096個字詞的上下文將佔用約5.3 GB;與約70 GB的權重相比,這已經相當可觀了。然而,如果考慮到Cohere的模型號稱有20萬個字詞的上下文視窗,情況會變得更加可怕 ─ 要計算20萬個上下文視窗的最後一個字詞,你需要讀取260 GB!(讓我們暫時忽略你還需要260 GB的顯存來儲存它)

在典型的「生產」(仍然是單一使用者)環境中,情況會進一步改變。權重通常會使用4位元量化(通常實作為每個權重約4.5位元),而KV快取可能會使用8位元(FP8)值。如果我們「保守地」假設10萬個字詞的上下文(是宣稱最大值的一半),這將給我們帶來約19.7 GB的模型權重和約65 GB的KV快取,而要計算最後一個字詞,我們需要從記憶體中讀取所有這些資料。突然間,注意力計算從微不足道變成了佔用時間的約75%,假設兩者都以峰值頻寬執行!

儘管10萬個字詞的上下文可能看起來有點極端,但在多使用者的環境中,這也是預期工作負載的公平表示。批次處理允許我們使矩陣乘法變成ALU bound,並為每批值(即每64+個使用者請求)讀取一次模型權重,但每個使用者請求通常都有自己的KV快取,因此注意力計算仍然是頻寬bound ─ 並且需要大量記憶體才能將所有使用者的請求放在單個節點上!

如果這些模型使用了4倍GQA,KV快取所需的大小和頻寬就會小4倍;雖然對於數萬個字詞的上下文來說仍然很可觀,但會更容易管理。對於Cohere的預期使用情境,使用GQA可能會有一些品質下降 ─ 看看技術報告會很有趣,因為它可能包含相關的消融研究,但純粹從成本/效能的角度來看,每個基於transformer的語言模型都需要評估GQA,因為其優勢太顯著了,不容忽視。

個人見解:

在生成文字時,語言模型是逐字逐句依序產生,沒有辦法平行化處理。模型在處理每個字詞時,主要進行矩陣向量乘法和注意力(attention)計算。現代 CPU 和 GPU 的算術運算能力遠高於記憶體讀取速度,以至於推論過程的效能幾乎完全取決於記憶體頻寬。

藉由分析模型參數的組成和注意力機制的快取需求,我們可以預估推論每個字詞所需的最小時間。這個理論下限對於評估軟硬體效能非常有用:

1. 如果實際效能遠低於理論值,代表軟硬體可能有優化空間。
2. 不同精度(如FP16、INT8)的理論效能可以預估量化的潛在效益。
3. 即使硬體架構不同,記憶體頻寬仍是一致的效能指標。

我認為作者提出了一個實用且具啟發性的效能模型。不只可以用來診斷效能瓶頸,也能輔助軟硬體設計的取捨。唯一的缺點是沒有涵蓋多使用者的情境,不過文中也簡單討論了批次處理(batching)的影響。

此外,作者特別強調了群組查詢注意力(Group Query Attention)的重要性。藉由增加每個注意力快取元素的運算量,GQA 顯著降低了快取占用的記憶體容量和頻寬。以 Cohere 的 Command-R 模型為例,如果使用 4x GQA,快取所需的記憶體和頻寬可以減少 4 倍之多。

我認為速度極限分析是效能工程不可或缺的一環。透過將實際效能與理論極限比較,我們可以系統性地找出優化方向,而不是憑直覺或運氣。當然,模型品質的權衡也不能忽視。畢竟再怎麼快,生成愚蠢的回覆也是沒有意義的。在效能與品質之間取得平衡,才是打造實用語言模型的不二法門。

不過,我也想補充幾點:

1. 文中假設讀取權重是效能的主要限制,但對某些硬體(如AI加速晶片)來說,這可能並不完全成立。晶片內部通常有專門的高速記憶體來儲存權重,因此讀取速度可能不是最關鍵的因素。

2. 推論速度只是語言模型實用性的一個面向。即使透過種種優化來逼近理論極限,品質和穩定性的問題仍然存在。生成不連貫或有害的內容,對使用者體驗的傷害可能大過延遲幾毫秒。品質和效能需要一起考量。

3. 文章聚焦於單一使用者的情境,但實際應用常常需要同時服務大量使用者。這時批次處理和記憶體效率的權衡會變得更加複雜。找出最佳的批次大小和排程策略,是一個值得深入研究的課題。

作者提出了一個實用的效能分析框架,揭示了當前語言模型推論的核心挑戰。這些洞見不僅有助於改進現有系統,也為未來的架構設計提供了方向。期待看到更多在效能、品質、應用場景等多方面尋求平衡的創新方案。

Google的"大力出奇蹟"!!

這是我第三次評論這篇論文, 重要的報告我講三次, 很重要!很重要!很重要!我認為這個模型會超越Mamba!

https://arxiv.org/pdf/2402.19427.pdf

近年來,Transformer語言模型在多個自然語言處理任務上取得了巨大成功,但其全局注意力機制在處理超長序列時面臨計算瓶頸。另一方面,循環神經網絡(RNN)長期以來被認為更適合對長距離依賴進行建模,但其串行計算特性限制了訓練和推理效率。為了兼顧長程建模能力和計算效率,Google提出了Griffin——一種融合了門控線性遞迴單元(RG-LRU)和局部多頭注意力(Local Multi-head Attention)的混合語言模型。

Griffin模型

整體架構

Griffin的主體是一個堆疊的Transformer結構,包含L個殘差塊(L一般取12~40)。每個殘差塊包含兩個子結構:多層感知器(MLP)和時序混合塊。其中,MLP塊在所有殘差塊中共享,而時序混合塊在不同層中交替使用RG-LRU塊和局部注意力(Local Attention)塊。模型的輸入是一個長度為N的token序列,通過詞嵌入(Embedding)層映射為H維稠密向量。最終輸出通過嵌入矩陣的轉置進行復原,並計算softmax交叉熵損失。

MLP塊

Griffin的MLP塊基於門控線性單元(GLU),包含兩個并列的全連接層(尺寸為H->H)和一個GLU激活函數。計算公式為:
$$ f(x) = (W_1 x + b_1) \otimes \sigma(W_2 x + b_2) $$
其中$W_1,W_2 \in R^{H \times H}, b_1,b_2 \in R^H$分別為兩個全連接層的權重矩陣和偏置項,$\sigma$為Sigmoid函數,$\otimes$為按元素乘法。相比ReLU等激活函數,GLU可以更好地建模輸入之間的非線性交互作用。

時序混合塊

RG-LRU

RG-LRU是Griffin的核心創新之一,旨在改進傳統RNN捕捉長距離依賴的能力。給定輸入序列$x_1,\dots,x_N$,RG-LRU逐步計算隱狀態序列$h_1,\dots,h_N$。在時刻t,RG-LRU執行以下運算:
$$r_t = \sigma(W_{xr} x_t + b_r)$$
$$i_t = \sigma(W_{xi} x_t + b_i)$$
$$\tilde{A}_t = \text{diag}(r_t) A^T \text{diag}(r_t)$$
$$h_t = \tilde{A}t h{t-1} + \sqrt{1-\tilde{A}_t^2} \odot (i_t \odot x_t)$$
其中$r_t,i_t \in R^H$分別為遞迴門和輸入門,$W_{xr},W_{xi} \in R^{H \times H}$和$b_r,b_i \in R^H$ 為門控單元的參數,$A \in R^{H \times H}$為一個對角矩陣,對角線元素在(0,1)範圍內。運算$\odot$表示按元素乘法。
RG-LRU引入了兩個關鍵的改進:
(1) 遞迴門$r_t$控制了過去狀態$h_{t-1}$的保留程度,其中$\tilde{A}_t$是對$A$應用$r_t$進行插值的結果。當$r_t$接近1時,模型傾向於保留過去信息;當$r_t$接近0時,模型傾向於遺忘過去信息,只關注當前輸入$x_t$。這種自適應的記憶機制使RG-LRU能靈活地應對不同時間尺度的依賴關係。
(2) 輸入門$i_t$控制了當前輸入$x_t$對隱狀態$h_t$的貢獻。與LSTM等傳統RNN不同,RG-LRU的輸入門不依賴上一時刻的隱狀態,從而實現了$O(1)$的順序計算復雜度。

局部注意力

除了RG-LRU塊,Griffin還周期性地插入局部多頭注意力(Local Multi-head Attention)塊,用於捕捉中短距離的token級別交互。給定隱狀態序列$h_1,\dots,h_N$,局部注意力在每個位置$t$處計算加權平均表示$o_t$:
$$o_t = \sum_{i=t-M}^{t+M} \alpha_{ti} (W_V h_i)$$
其中$M$為注意力窗口的半徑,$W_V \in R^{H \times H}$為值(Value)映射矩陣。注意力權重$\alpha_{ti}$通過查詢向量$q_t$和鍵向量$k_i$的內積計算:
$$\alpha_{ti} = \frac{\exp(q_t^T k_i)}{\sum_{j=t-M}^{t+M} \exp(q_t^T k_j)}$$
其中$q_t = W_Q h_t, k_i = W_K h_i$,對應的投影矩陣為$W_Q,W_K \in R^{H \times H}$。為了建模多種類型的依賴關係,Griffin使用多頭注意力機制,計算公式為:
$$o_t = W_O[o_t^{(1)},\dots,o_t^{(K)}] + b_o$$
其中$o_t^{(k)}$是第$k$個注意力頭的輸出,$W_O \in R^{H \times KH},b_o \in R^H$為最終的線性變換參數。
在實踐中,我們發現設置注意力窗口大小$2M+1=1024$,注意力頭數$K=32$可以取得最佳的性能。這種局部注意力機制顯著降低了計算複雜度(從$O(N^2)$減少到$O(NM)$),且與RG-LRU形成了很好的互補。

作為堅定支持者，Google 的 Griffin 論文完美地展示了Scaling laws。當他們將模型參數擴大 7 倍時,在各項任務上的表現大約提升了 10%。所有模型都是在相同的 3000 億 tokens 數據上訓練的。隨著參數的增加,模型變得更加樣本有效,外推能力也更強。

位置編碼

為了將序列中的位置信息引入模型,Griffin沿用了Transformer中的相對位置編碼方案。具體來說,在局部注意力模塊中,我們將位置編碼信息添加到鍵向量和查詢向量內積中:
$$\alpha_{ti} = \frac{\exp(q_t^T k_i + q_t^T r_{i-t})}{\sum_{j=t-M}^{t+M} \exp(q_t^T k_j + q_t^T r_{j-t})}$$
其中$r_{\Delta t} \in R^H$是一個可學習的相對位置編碼向量,表示時間步$t$和$i$之間的距離$\Delta t = i-t$的影響。

這是 Google 的一項突破性研究成果。

Google 發布了具有新 Griffin 架構的模型,其性能優於 Transformer。

訓練與推理優化

參數高效性

得益於RG-LRU的線性計算特性,Griffin的參數量和計算復雜度與層數L呈線性關係。具體來說,設模型隱藏層維度為H,詞表大小為V,則Griffin的參數量近似為:

$$\text{Params} = 4LH^2 + 2VH$$

其中$4LH^2$項對應MLP塊、RG-LRU和注意力層的權重矩陣參數,$2VH$對應詞嵌入矩陣及其轉置。考慮到$H << V$,模型的參數量主要由詞嵌入矩陣主導。這意味著Griffin可以通過增加深度L來提高模型容量,而不會導致參數量過度膨脹。通過控制隱藏層維度H,在保持參數量不變的情況下比較Griffin和Transformer在不同層數L下的性能。

GPU/TPU並行化

雖然RG-LRU易於實現廉價的串行計算,但為了充分利用現代加速器(如GPU和TPU)的並行能力,還需要對其進行並行化改造。受益於RG-LRU的簡潔性,可以輕松地將其計算過程表示為一系列矩陣乘法和逐元素操作,從而實現高效的批量化(Batching)和張量化(Tensorization)加速。

在訓練時,將一個批次的輸入序列表示為形狀為$[B, N, H]$的三維張量,其中$B$為批次大小。通過將第二維(長度維)的計算映射到加速器的不同線程/內核上,Griffin可以實現與Transformer相當的訓練吞吐量。對於局部注意力層,使用了快速注意力(Fast Attention)算法,通過計算局部注意力權重的前綴和(Prefix Sum),將總體複雜度降低到$O(BNH)$。

在推理時,需要逐步生成輸出序列。與Transformer需要維護一個隨生成長度增長的鍵值緩存(Key-Value Cache)不同,RG-LRU只需要維護一個固定大小的隱狀態向量$h_t$,從而顯著節省了內存佔用。此外,Griffin還可以利用內存化(Memorization)技術,即將每一層的輸出都保存到一個大小固定的循環緩衝區中,供下一層計算時復用。這種做法避免了在深層模型中重複計算前幾層的結果,進一步提升了推理速度。

接下來,研究者還選取了三個專門用於評測長距離依賴建模能力的任務:LAMBADA、ListOps和Pathfinder。其中,LAMBADA是一個基於上下文的單詞預測任務,ListOps需要模型執行算術運算,而Pathfinder則考察模型在網格圖上進行推理的能力。實驗結果顯示,Griffin在這些任務上的表現也一致地超過了基線模型。值得注意的是,隨著序列長度的增加,Griffin的優勢變得更加明顯,這歸因於RG-LRU層強大的遞歸歸納能力。

在推理效率方面,Griffin展現出了明顯的優勢。得益於RG-LRU層恆定的內存佔用和局部注意力的線性計算複雜度,Griffin在生成超長文本時的推理速度比Transformer-XL等模型快2~3倍,且內存佔用減少了50%以上。這使得Griffin非常適合應用於資源受限的場景,如移動設備或實時系統。

混合精度訓練

為了進一步提高訓練效率並節省顯存,採用了混合精度(Mixed Precision)訓練策略。具體來說,模型的前向和反向傳播過程使用半精度浮點數(FP16),而模型權重的更新則在單精度(FP32)下進行。通過這種設置,可以將訓練所需的顯存減少近50%,同時保持模型的收斂性和最終性能。我們還發現,對梯度應用動態缩放(Dynamic Scaling)技術可以進一步提高混合精度訓練的穩定性。

為了深入分析Griffin的行為特徵,研究者還進行了一系列消融實驗。首先,他們探究了不同超參數(如層數、隱藏狀態維度、注意力頭數等)對模型性能的影響。結果表明,增加Griffin的層數和寬度都能帶來性能的提升,且效果優於同等大小的Transformer模型。其次,研究者還考察了不同注意力範圍對Griffin的影響。他們發現,當注意力窗口大小在1024左右時,模型在各個任務上的表現最優,且即使在更長的序列長度下也能保持優勢。

他們在WikiText-103、One Billion Word和PG-19等語言建模數據集上測試了Griffin的泛化能力。這些數據集包含了不同長度、主題和風格的文本,可以全面評估模型在長距離依賴關係上的建模能力。實驗中,研究者控制了模型的參數量,以公平地比較Griffin與其他基線模型的性能。結果表明,Griffin在所有數據集上都取得了最優的perplexity(PPL)分數,顯著優於Transformer-XL、Compressive Transformer等現有方法。

接下來,研究者還選取了三個專門用於評測長距離依賴建模能力的任務:LAMBADA、ListOps和Pathfinder。其中,LAMBADA是一個基於上下文的單詞預測任務,ListOps需要模型執行算術運算,而Pathfinder則考察模型在網格圖上進行推理的能力。實驗結果顯示,Griffin在這些任務上的表現也一致地超過了基線模型。值得注意的是,隨著序列長度的增加,Griffin的優勢變得更加明顯,這歸因於RG-LRU層強大的遞歸歸納能力。

為了深入分析Griffin的行為特徵,研究者還進行了一系列消融實驗。首先,他們探究了不同超參數(如層數、隱藏狀態維度、注意力頭數等)對模型性能的影響。結果表明,增加Griffin的層數和寬度都能帶來性能的提升,且效果優於同等大小的Transformer模型。其次,研究者還考察了不同注意力範圍對Griffin的影響。他們發現,當注意力窗口大小在1024左右時,模型在各個任務上的表現最優,且即使在更長的序列長度下也能保持優勢。

在推理效率方面,Griffin展現出了明顯的優勢。得益於RG-LRU層恆定的內存佔用和局部注意力的線性計算複雜度,Griffin在生成超長文本時的推理速度比Transformer-XL等模型快2~3倍,且內存佔用減少了50%以上。這使得Griffin非常適合應用於資源受限的場景,如移動設備或實時系統。

結論

1. RG-LRU(Real-Gated Linear Recurrent Unit):

RG-LRU是Griffin模型的核心組件之一。它是一種新型的門控線性遞迴層,其設計靈感來自於標準的LRU(Linear Recurrent Unit)和LSTM/GRU中使用的門控機制。 

具體來說,RG-LRU引入了兩個門:遞迴門r_t和輸入門i_t,它們分別控制前一時刻隱狀態h_{t-1}和當前輸入x_t對新隱狀態h_t的貢獻。與LSTM/GRU不同的是,RG-LRU的門不依賴於前一時刻的隱狀態,這樣可以顯著提高計算效率。

遞迴門r_t進一步調節一個對角矩陣A_t,使其在保留歷史信息(r_t->1時A_t->I)和遺忘歷史信息(r_t->0時A_t->0)之間進行非線性插值。這種獨特的門控方式使RG-LRU能學會在長序列上選擇性地汲取信息。

RG-LRU還有一個復數版本CG-LRU(Complex-Gated Linear Recurrent Unit),它將輸入和隱狀態表示為復數,並使用復數矩陣參數,以提高层的表示能力。但實驗發現,在語言建模任務上,實數版本的RG-LRU性能並不亞於CG-LRU。

2. 時序混合(Temporal Mixing): 

Griffin採用分層的時序混合方式,在12-40層的網路中,每3層就混合一個RNN塊(含3個RG-LRU層)和一個局部MQA(Multi-Query Attention)層。這樣的結構讓模型在捕捉長距離依賴的同時,也能很好地對局部範圍內的信息進行匹配。

實驗顯示,對於序列長度為2048的預訓練,局部MQA的最優窗口大小在1024左右。當序列長度增加到4096、8192時,1024的窗口大小仍然能取得比全局注意力Transformer更好的效果。這表明Griffin利用RNN積累長程信息的能力可以很好地彌補局部注意力的不足。

3. 硬體加速:

為了高效地在TPU-v3上訓練RG-LRU,研究者實現了定製的Pallas內核,將線性掃描運算中的內存讀寫次數降到最低。這使得RG-LRU在訓練加速比上可以達到Transformer的水平(記憶體帶寬約為900GB/s)。

同時,由於RG-LRU的隱狀態大小和局部MQA的KV cache大小遠小於Transformer的KV cache,Griffin在推理階段表現出了顯著的延遲優勢(低20%以上)和吞吐量優勢(高3倍以上)。這在生成超長序列時尤為明顯。

4. 外推和零樣本學習:

研究者在一個書籍語料(Books)和arXiv語料上評估了Griffin的外推(extrapolation)能力,即在遠長於訓練序列的長度上生成文本的能力。實驗表明,Griffin可以在長達訓練序列4倍的序列上穩定地利用更長上下文改進預測。而Transformer受限於其位置編碼,很難在超出訓練序列長度時維持較好表現。

此外,Griffin還在幾個合成的復制/檢索任務(Selective Copying、Induction Heads、Phone Number Lookup)上接受了測試。結果顯示,Griffin可以在監督訓練下快速學會這些任務所需的復制檢索技能。但在零樣本(zero-shot)設定下,預訓練的Griffin模型在這些任務上的外推能力仍然不及Transformer。這可能與局部注意力和RNN對字面匹配(verbatim match)的能力有限有關。

5. 訓練規模與下游任務性能:

研究者在100M到14B參數規模下訓練了Griffin,Hawk和Transformer模型,並在MMLU、HellaSwag、PIQA等7個下游任務上進行了評估。結果顯示,在只用300B tokens訓練的情況下:

- Griffin和Hawk在所有任務上的平均得分隨模型規模增大而穩定上升,其中Griffin在所有規模下都超過Transformer。

- Hawk-3B顯著超過了用600B tokens訓練的Mamba-3B。

- Griffin-7B/14B則在多數任務上達到甚至超過了用2T tokens訓練的Transformer模型Llama-2的水平。

這充分展現了Griffin超越Transformer的樣本效率。

綜上所述,Griffin通過巧妙融合RNN和局部注意力,在計算效率、樣本效率、長程建模等多方面取得了超越Transformer的效果,是一種非常有前景的新型語言模型架構。未來還可以進一步探索如何改進其在超長序列上的注意力機制,以及在零樣本語言任務上的外推和泛化能力。在工程實踐中,Griffin有望憑藉其出色的推理性能在長文本應用場景大放異彩。

Transformer的「二次障礙」

 近期,在多位AI領域大佬的演講中,都提到了一個共同的問題:電力短缺。隨著AI技術的快速發展,尤其是大語言模型的出現,對算力和能源的需求呈指數級增長。然而,現有的電網基礎設施難以滿足如此龐大的電力需求。正如馬斯克所預測,在晶片短缺緩解後,電力短缺問題將接踵而至。

OpenAI執行長Sam Altman認為,能源問題是AI發展的第一性原理,而AI則是能源的無底洞。在可控核融合實現之前,能源短缺可能成為AI發展的瓶頸。

當前主流的AI大模型,如ChatGPT、Gemini、Claude等,大多基於Transformer架構。Transformer的強大之處在於其自注意力機制(Self-Attention)和前饋神經網路(FFN)能夠捕捉文字資料中的複雜模式和長距離依賴關係。

然而,Transformer在處理長序列時存在計算複雜度的問題。隨著輸入序列長度的增加,Self-Attention的計算量呈現二次增長。意味著輸入序列長度越長,模型的計算資源消耗就越高,從而導致能源消耗的急劇上升。儘管當前已有一些針對注意力機制的最佳化方法,如Flash Attention、Grouped-Query Attention等,但它們並未從根本上解決Transformer的二次複雜度問題。

潛在的解決方案：

1. 混合體架構

為了克服Transformer的「二次障礙」,研究者提出了各種注意力機制的變體和混合體架構。其中,Jamba架構透過將Mamba結構化狀態空間模型(SSM)與Transformer相結合,在保持模型性能的同時,大幅降低了計算成本。這種權衡的方式為優化模型效能提供了新的思路。

2. 混合專家(MoE)技術

混合專家(Mixture of Experts, MoE)技術透過將模型分解為多個專家子模型,並在推論過程中只啟動部分專家,從而減少不必要的計算,提高模型的計算效率。目前,MoE已成為新發布的大型語言模型中的常見設計。例如,GPT-4、Gemini 1.5、Mixtral 8x7B和Jamba都採用了MoE架構。

傳統的MoE模型仍面臨知識混合和冗餘的問題。Deepseek提出了一種新型的MoE架構,透過增加專家數量、引入共享專家等方式,進一步提升了模型的性能和效率。實驗結果表明,DeepSeekMoE在保持優異性能的同時,計算成本僅為原先的約40%。

3. 資料工程的進步

資料品質對於模型性能的影響至關重要。Meta最新發布的Llama-3模型,在架構不變的情況下,透過增加資料量和提高資料品質,實現了顯著的效果提升。這凸顯了資料工程在大模型訓練中的重要性。

透過資料篩選、清洗、增強等手段,可以在相同計算資源下獲得更好的模型效果。此外,對不同來源資料的混合比例進行探索和優化,有助於提升模型的泛化能力和適用性。因此,加強資料工程,提高資料品質,是提升模型效能和減少資源浪費的重要途徑。

對人類來說,思考模式跟 Attention 機制不太一樣。Attention 是在一個很長的上下文裡,去找出之前出現過的某些關鍵詞 (Token)。但人類的思考可沒那麼制式化。以目前的模型架構來看,其實還有很大的進步空間。

像Transformer 就很難去實現人類擁有的某些能力。舉例來說,工作記憶就是一個很好的例子。人類做某件事情久了,就會越做越上手,但 Transformer 就很難學會這一點。

空間記憶也是一樣。人類第一次到一個地方可能會暈頭轉向,但去多了之後,就會知道哪條路最近、最方便。但 Transformer 對這種記憶好像就有點無從下手。

所以,Transformer 再厲害,終究還是有些先天上的缺陷和侷限性。

Apple的逆襲: OpenELM大模型

Apple的新報告提出了開源的高效語言模型及訓練推理架"OpenELM: An Efficient Language Model Family with Open-source Training and Inference Framework"論文,跟之前Meta的OLMo一樣令人耳目一新。因為它詳細分享了模型架構、訓練方法和訓練數據,這在業界並不多見。

蘋果公司最近開源了一個名為OpenELM的大型語言模型(LLM)家族,旨在通過公開模型架構、訓練方法和資料集,提高LLM研究的可重複性和透明度。這不僅有助於學界驗證結果、研究數據和模型偏差,也為探索潛在風險提供了機會。   

OpenELM採用了一種逐層縮放(layer-wise scaling)策略,在transformer模型的每一層中高效分配參數,從而在給定參數預算下實現更高的精度。例如,在約10億參數規模下,OpenELM的精度比最近開源的OLMo高出2.36%,而訓練token數量只有後者的一半。

有別於此前只提供模型權重和推理程式碼,或在私有資料集上預訓練的做法,Apple的開源內容包括了在公開資料集上進行訓練和評估的完整框架,涵蓋訓練日誌、多個檢查點和預訓練配置。他們還發布了將模型轉換為MLX庫的程式碼,方便在蘋果設備上進行推理和微調。這種全面的開源有望賦能和加強開放研究社區,為未來的開放研究鋪平道路。

論文連結: https://arxiv.org/pdf/2404.14619

模型架構

OpenELM採用了decoder-only的transformer架構。參考SOTA的LLM,它在設計中融入了以下特點:  

1. 所有全連接層(線性層)不使用可學習的偏置參數

2. 使用RMSNorm進行預規範化,使用RoPE編碼位置信息  

3. 用分組查詢注意力(GQA)取代多頭注意力(MHA)

4. 用SwiGLU FFN取代前饋網絡(FFN)  

5. 使用flash attention計算點積注意力

6. 沿用LLama的分詞器

與現有LLM在每層採用相同配置、均勻分配參數不同,OpenELM根據逐層縮放策略,讓每層transformer的配置(如注意力頭數、前饋維度)各不相同,從而在整個模型中實現參數的非均勻分配。這種做法讓OpenELM能在相同參數量下,實現更高的精度。

訓練資料和方法

OpenELM在約1.8萬億token的公開資料集上進行了預訓練,包括RefinedWeb、去重的PILE、RedPajama和Dolma的子集。與之前使用預分詞數據的方法不同,OpenELM採用了即時過濾和分詞。這大大簡化了使用不同分詞器進行實驗的流程。 

訓練採用AdamW優化器,使用餘弦退火學習率調度,訓練35萬步。他們訓練了270M、450M、1.1B和3B四個不同規模的變體,對於某些模型還啟用了FSDP和activation checkpointing技術。

評估結果

研究者在三個評估框架下測試了OpenELM:標準zero-shot任務、OpenLLM排行榜任務和LLM360排行榜任務。結果顯示,不論在推理、知識理解還是識別錯誤信息和偏見方面,OpenELM都優於現有的開源LLM。尤其是在規模相當的情況下,OpenELM在精度上顯著領先OLMo,且只用了一半的訓練資料。  

進一步的指令微調實驗表明,在6萬個指令樣本上微調,可使OpenELM在各種任務上的平均精度提高1-2個百分點。此外,LoRA和DoRA等參數高效微調方法也被證明可用於OpenELM,在維持模型大小不變的情況下顯著提升下游任務性能。

OpenELM的源碼、預訓練權重和訓練配方已在GitHub上開源。模型本身也已上傳到HuggingFace,方便研究者使用。值得一提的是,蘋果還開源了將模型轉換為MLX的代碼,MLX是一個在Apple矽晶片上進行高效機器學習的庫。這為在蘋果設備上部署OpenELM掃清了障礙。  

OpenELM的發布代表了Apple在開放研究道路上又邁出的一大步。逐層縮放策略和全面的開源內容是其兩大亮點。儘管目前的推理速度還有待優化,但OpenELM在精度和資料效率上已經展現了巨大的潛力。未來,隨著更多的研究者參與改進和分析,OpenELM有望在各類自然語言任務上樹立新的標杆,為LLM的發展注入新的活力。

心得:

1. OpenELM有4種相對較小且方便使用的規模:270M、450M、1.1B和3B。

2. 儘管OpenELM的訓練token數量只有OLMo的一半,但性能略優於後者。

3. 主要的架構創新是一種逐層縮放策略。

這種逐層縮放策略非常有意思。希望他們做一個消融實驗,在同樣的資料集上訓練兩個版本的模型,一個用這種策略,一個不用,看看效果有何不同。但這樣的實驗成本很高。

另外, 他們也比較了LoRA和DoRA在參數高效微調上的表現!不過結果顯示,這兩種方法並沒有明顯差異。