
微軟新出熱乎論文:Transformer 擴展到 10 億 token

已經擴展到了 10 億 token,未來能否將整個互聯網作為一個序列處理?
比如明星大模型 GPT-4 支持 32ktoken,相當於 50 頁的文字;OpenAI 前成員創立的 Anthropic 更是將 Claude 處理 token 能力提升到 100k,約 75000 個單詞,大概相當於一鍵總結《哈利波特》第一部。在微軟最新的一項研究中,他們這次直接將 Transformer 擴展到 10 億 token。
這為建模非常長的序列開闢了新的可能性,例如將整個語料庫甚至整個互聯網視為一個序列。
作為比較,普通人可以在 5 小時左右的時間裏閲讀 100,000 個 token,並可能需要更長的時間來消化、記憶和分析這些信息。
Claude 可以在不到 1 分鐘的時間裏完成這些。要是換算成微軟的這項研究,將會是一個驚人的數字。
- 論文地址:
https://arxiv.org/pdf/2307.02486.pdf
- 項目地址:
https://github.com/microsoft/unilm/tree/master
具體而言,該研究提出了 LONGNET,這是一種 Transformer 變體,可以將序列長度擴展到超過 10 億個 token,而不會犧牲對較短序列的性能。
文中還提出了 dilatedattention,它能指數級擴展模型感知範圍。
LONGNET 具有以下優勢:
1)它具有線性計算複雜性;
2)它可以作為較長序列的分佈式訓練器;
3)dilatedattention 可以無縫替代標準注意力,並可以與現有基於 Transformer 的優化方法無縫集成。
實驗結果表明,LONGNET 在長序列建模和一般語言任務上都表現出很強的性能。在研究動機方面,論文表示,最近幾年,擴展神經網絡已經成為一種趨勢,許多性能良好的網絡被研究出來。
在這當中,序列長度作為神經網絡的一部分,理想情況下,其長度應該是無限的。但現實卻往往相反,因而打破序列長度的限制將會帶來顯著的優勢:
- 首先,它為模型提供了大容量的記憶和感受野,使其能夠與人類和世界進行有效的交互。
- 其次,更長的上下文包含了更復雜的因果關係和推理路徑,模型可以在訓練數據中加以利用。相反,較短的依賴關係則會引入更多虛假的相關性,不利於模型的泛化性。
- 第三,更長的序列長度可以幫助模型探索更長的上下文,並且極長的上下文也可幫助模型緩解災難性遺忘問題。
然而,擴展序列長度面臨的主要挑戰是在計算複雜性和模型表達能力之間找到合適的平衡。
例如 RNN 風格的模型主要用於增加序列長度。然而,其序列特性限制了訓練過程中的並行化,而並行化在長序列建模中是至關重要的。最近,狀態空間模型對序列建模非常有吸引力,它可以在訓練過程中作為 CNN 運行,並在測試時轉換為高效的 RNN。
然而這類模型在常規長度上的表現不如 Transformer。另一種擴展序列長度的方法是降低 Transformer 的複雜性,即自注意力的二次複雜性。現階段,一些高效的基於 Transformer 的變體被提出,包括低秩注意力、基於核的方法、下采樣方法、基於檢索的方法。
然而,這些方法尚未將 Transformer 擴展到 10 億 token 的規模(參見圖 1)。
下表為不同計算方法的計算複雜度比較。N 為序列長度,d 為隱藏維數。
方法
該研究的解決方案 LONGNET 成功地將序列長度擴展到 10 億個 token。具體來説,該研究提出一種名為 dilatedattention 的新組件,並用 dilatedattention 取代了 VanillaTransformer 的注意力機制。
通用的設計原則是注意力的分配隨着 token 和 token 之間距離的增加而呈指數級下降。該研究表明這種設計方法獲得了線性計算複雜度和 token 之間的對數依賴性。
這就解決了注意力資源有限和可訪問每個 token 之間的矛盾。
在實現過程中,LONGNET 可以轉化成一個密集 Transformer,以無縫地支持針對 Transformer 的現有優化方法(例如內核融合(kernelfusion)、量化和分佈式訓練)。
利用線性複雜度的優勢,LONGNET 可以跨節點並行訓練,用分佈式算法打破計算和內存的約束。最終,該研究有效地將序列長度擴大到 1B 個 token,而且運行時(runtime)幾乎是恆定的,如下圖所示。
相比之下,VanillaTransformer 的運行時則會受到二次複雜度的影響。

該研究進一步引入了多頭 dilatedattention 機制。
如下圖 3 所示,該研究通過對查詢 - 鍵 - 值對的不同部分進行稀疏化,在不同的頭之間進行不同的計算。

分佈式訓練
雖然 dilatedattention 的計算複雜度已經大幅降低到
,但由於計算和內存的限制,在單個 GPU 設備上將序列長度擴展到百萬級別是不可行的。有一些用於大規模模型訓練的分佈式訓練算法,如模型並行 [SPP+19]、序列並行 [LXLY21,KCL+22] 和 pipeline 並行 [HCB+19],然而這些方法對於 LONGNET 來説是不夠的,特別是當序列維度非常大時。
該研究利用 LONGNET 的線性計算複雜度來進行序列維度的分佈式訓練。
下圖 4 展示了在兩個 GPU 上的分佈式算法,還可以進一步擴展到任意數量的設備。

實驗
該研究將 LONGNET 與 vanillaTransformer 和稀疏 Transformer 進行了比較。架構之間的差異是注意力層,而其他層保持不變。
研究人員將這些模型的序列長度從 2K 擴展到 32K,與此同時減小 batch 大小,以保證每個 batch 的 token 數量不變。
表 2 總結了這些模型在 Stack 數據集上的結果。研究使用複雜度作為評估指標。
這些模型使用不同的序列長度進行測試,範圍從 2k 到 32k 不等。當輸入長度超過模型支持的最大長度時,研究實現了分塊因果注意力(blockwisecausalattention,BCA)[SDP+22],這是一種最先進的用於語言模型推理的外推方法。此外,研究刪除了絕對位置編碼。
首先,結果表明,在訓練過程中增加序列長度一般會得到更好的語言模型。其次,在長度遠大於模型支持的情況下,推理中的序列長度外推法並不適用。
最後,LONGNET 一直優於基線模型,證明了其在語言建模中的有效性。

序列長度的擴展曲線
圖 6 繪製了 vanillatransformer 和 LONGNET 的序列長度擴展曲線。該研究通過計算矩陣乘法的總 flops 來估計計算量。
結果表明,vanillatransformer 和 LONGNET 都能從訓練中獲得更大的上下文長度。
然而,LONGNET 可以更有效地擴展上下文長度,以較小的計算量實現較低的測試損失。這證明了較長的訓練輸入比外推法更具有優勢。實驗表明,LONGNET 是一種更有效的擴展語言模型中上下文長度的方法。這是因為 LONGNET 可以更有效地學習較長的依賴關係。

擴展模型規模
大型語言模型的一個重要屬性是:損失隨着計算量的增加呈冪律擴展。為了驗證 LONGNET 是否仍然遵循類似的擴展規律,該研究用不同的模型規模(從 1.25 億到 27 億個參數)訓練了一系列模型。
27 億的模型是用 300B 的 token 訓練的,而其餘的模型則用到了大約 400B 的 token。圖 7(a) 繪製了 LONGNET 關於計算的擴展曲線。該研究在相同的測試集上計算了複雜度。
這證明了 LONGNET 仍然可以遵循冪律。這也就意味着 denseTransformer 不是擴展語言模型的先決條件。此外,可擴展性和效率都是由 LONGNET 獲得的。

長上下文 prompt
Prompt 是引導語言模型併為其提供額外信息的重要方法。
該研究通過實驗來驗證 LONGNET 是否能從較長的上下文提示窗口中獲益。
該研究保留了一段前綴(prefixes)作為 prompt,並測試其後綴(suffixes)的困惑度。
並且,研究過程中,逐漸將 prompt 從 2K 擴展到 32K。為了進行公平的比較,保持後綴的長度不變,而將前綴的長度增加到模型的最大長度。
圖 7(b) 報告了測試集上的結果。它表明,隨着上下文窗口的增加,LONGNET 的測試損失逐漸減少。這證明了 LONGNET 在充分利用長語境來改進語言模型方面的優越性。
本文來源:人工智能學家,原標題:《微軟新出熱乎論文:Transformer 擴展到 10 億 token》
