當大家不斷升級迭代自家大模型的時候，LLM（大語言模型）對上下文窗口的處理能力，也成為一個重要評估指標。

比如明星大模型 GPT-4 支持 32k token，相當于 50 頁的文字；OpenAI 前成員創(chuàng)立的 Anthropic 更是將 Claude 處理 token 能力提升到 100k，約 75000 個單詞，大概相當于一鍵總結《哈利波特》第一部。

在微軟最新的一項研究中，他們這次直接將 Transformer 擴展到 10 億 token。這為建模非常長的序列開辟了新的可能性，例如將整個語料庫甚至整個互聯(lián)網(wǎng)視為一個序列。

作為比較，普通人可以在 5 小時左右的時間里閱讀 100,000 個 token，并可能需要更長的時間來消化、記憶和分析這些信息。Claude 可以在不到 1 分鐘的時間里完成這些。要是換算成微軟的這項研究，將會是一個驚人的數(shù)字。

圖片

• 論文地址：https://arxiv.org/pdf/2307.02486.pdf
• 項目地址：https://github.com/microsoft/unilm/tree/master

具體而言，該研究提出了 LONGNET，這是一種 Transformer 變體，可以將序列長度擴展到超過 10 億個 token，而不會犧牲對較短序列的性能。文中還提出了 dilated attention，它能指數(shù)級擴展模型感知范圍。

LONGNET 具有以下優(yōu)勢：

1）它具有線性計算復雜性；

2）它可以作為較長序列的分布式訓練器；

3）dilated attention 可以無縫替代標準注意力，并可以與現(xiàn)有基于 Transformer 的優(yōu)化方法無縫集成。

實驗結果表明，LONGNET 在長序列建模和一般語言任務上都表現(xiàn)出很強的性能。

在研究動機方面，論文表示，最近幾年，擴展神經(jīng)網(wǎng)絡已經(jīng)成為一種趨勢，許多性能良好的網(wǎng)絡被研究出來。在這當中，序列長度作為神經(jīng)網(wǎng)絡的一部分，理想情況下，其長度應該是無限的。但現(xiàn)實卻往往相反，因而打破序列長度的限制將會帶來顯著的優(yōu)勢：

• 首先，它為模型提供了大容量的記憶和感受野，使其能夠與人類和世界進行有效的交互。
• 其次，更長的上下文包含了更復雜的因果關系和推理路徑，模型可以在訓練數(shù)據(jù)中加以利用。相反，較短的依賴關系則會引入更多虛假的相關性，不利于模型的泛化性。
• 第三，更長的序列長度可以幫助模型探索更長的上下文，并且極長的上下文也可幫助模型緩解災難性遺忘問題。

然而，擴展序列長度面臨的主要挑戰(zhàn)是在計算復雜性和模型表達能力之間找到合適的平衡。

例如 RNN 風格的模型主要用于增加序列長度。然而，其序列特性限制了訓練過程中的并行化，而并行化在長序列建模中是至關重要的。

最近，狀態(tài)空間模型對序列建模非常有吸引力，它可以在訓練過程中作為 CNN 運行，并在測試時轉換為高效的 RNN。然而這類模型在常規(guī)長度上的表現(xiàn)不如 Transformer。

另一種擴展序列長度的方法是降低 Transformer 的復雜性，即自注意力的二次復雜性。現(xiàn)階段，一些高效的基于 Transformer 的變體被提出，包括低秩注意力、基于核的方法、下采樣方法、基于檢索的方法。然而，這些方法尚未將 Transformer 擴展到 10 億 token 的規(guī)模（參見圖 1）。

圖片

下表為不同計算方法的計算復雜度比較。N 為序列長度，d 為隱藏維數(shù)。

圖片

方法

該研究的解決方案 LONGNET 成功地將序列長度擴展到 10 億個 token。具體來說，該研究提出一種名為 dilated attention 的新組件，并用 dilated attention 取代了 Vanilla Transformer 的注意力機制。通用的設計原則是注意力的分配隨著 token 和 token 之間距離的增加而呈指數(shù)級下降。該研究表明這種設計方法獲得了線性計算復雜度和 token 之間的對數(shù)依賴性。這就解決了注意力資源有限和可訪問每個 token 之間的矛盾。

圖片

在實現(xiàn)過程中，LONGNET 可以轉化成一個密集 Transformer，以無縫地支持針對 Transformer 的現(xiàn)有優(yōu)化方法（例如內核融合（kernel fusion）、量化和分布式訓練）。利用線性復雜度的優(yōu)勢，LONGNET 可以跨節(jié)點并行訓練，用分布式算法打破計算和內存的約束。

最終，該研究有效地將序列長度擴大到 1B 個 token，而且運行時（runtime）幾乎是恒定的，如下圖所示。相比之下，Vanilla Transformer 的運行時則會受到二次復雜度的影響。

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

圖片

分布式訓練

雖然 dilated attention 的計算復雜度已經(jīng)大幅降低到，但由于計算和內存的限制，在單個 GPU 設備上將序列長度擴展到百萬級別是不可行的。有一些用于大規(guī)模模型訓練的分布式訓練算法，如模型并行 [SPP+19]、序列并行 [LXLY21, KCL+22] 和 pipeline 并行 [HCB+19]，然而這些方法對于 LONGNET 來說是不夠的，特別是當序列維度非常大時。

該研究利用 LONGNET 的線性計算復雜度來進行序列維度的分布式訓練。下圖 4 展示了在兩個 GPU 上的分布式算法，還可以進一步擴展到任意數(shù)量的設備。

實驗

該研究將 LONGNET 與 vanilla Transformer 和稀疏 Transformer 進行了比較。架構之間的差異是注意力層，而其他層保持不變。研究人員將這些模型的序列長度從 2K 擴展到 32K，與此同時減小 batch 大小，以保證每個 batch 的 token 數(shù)量不變。

表 2 總結了這些模型在 Stack 數(shù)據(jù)集上的結果。研究使用復雜度作為評估指標。這些模型使用不同的序列長度進行測試，范圍從 2k 到 32k 不等。當輸入長度超過模型支持的最大長度時，研究實現(xiàn)了分塊因果注意力（blockwise causal attention，BCA）[SDP+22]，這是一種最先進的用于語言模型推理的外推方法。

此外，研究刪除了絕對位置編碼。首先，結果表明，在訓練過程中增加序列長度一般會得到更好的語言模型。其次，在長度遠大于模型支持的情況下，推理中的序列長度外推法并不適用。最后，LONGNET 一直優(yōu)于基線模型，證明了其在語言建模中的有效性。

序列長度的擴展曲線

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

擴展模型規(guī)模

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

長上下文 prompt

Prompt 是引導語言模型并為其提供額外信息的重要方法。該研究通過實驗來驗證 LONGNET 是否能從較長的上下文提示窗口中獲益。

該研究保留了一段前綴（prefixes）作為 prompt，并測試其后綴（suffixes）的困惑度。并且，研究過程中，逐漸將 prompt 從 2K 擴展到 32K。為了進行公平的比較，保持后綴的長度不變，而將前綴的長度增加到模型的最大長度。圖 7 (b) 報告了測試集上的結果。它表明，隨著上下文窗口的增加，LONGNET 的測試損失逐漸減少。這證明了 LONGNET 在充分利用長語境來改進語言模型方面的優(yōu)越性。