首個線性注意力Transformer大模型！1750億參數，速度和精度更優

GPT 等大型語言模型（LLM）的成功離不開 Softmax 注意力機制，但這一機制也存在著成本高等一些缺點。

近日，上海人工智能實驗室和 OpenNLPLab 的一個研究團隊提出了一種新的大型語言模型 TransNormerLLM，其中完全拋棄了基于 Softmax 的注意力機制，而是使用了新提出的線性注意力。據介紹，TransNormerLLM 是首個基于線性注意力的大型語言模型（LLM），其在準確度和效率方面的表現優于傳統的基于 Softmax 注意力的模型。研究者也將發布其預訓練模型的開源版本。

論文：https://arxiv.org/abs/2307.14995

模型：https://github.com/OpenNLPLab/TransnormerLLM

大型語言模型已經為自然語言處理（NLP）領域帶來了變革。它們能出色地處理許多不同類型的任務，提升了在計算框架中使用人類語言進行理解、生成和交互的能力。之前語言建模的發展主要圍繞 Transformer 架構，其中堪稱支柱的模型包括基本的 Transformer、GPT 系列、BERT 和 BART 等開創性的模型。Transformer 架構的成功立足于 softmax 注意力機制，該機制可以識別出數據驅動的模型方案中輸入 token 之間的依賴關系，其還能感知全局位置，讓模型能有效處理自然語言的長程動態情況。

即便如此，傳統的 Transformer 依然存在局限。首要的一點，它們有著對于序列長度的二次時間復雜度，這會限制它們的可擴展性并拖累訓練和推理階段的計算資源和時間效率。為了將這種二次時間復雜性簡化至線性復雜性，已經有不少研究者提出了多種不同的序列建模方法。但是，這些方法卻難以用于 LLM，原因有二：1) 它們在語言建模任務上的表現往往不如人意；2) 它們在真實世界場景中并沒有展現出速度優勢。

這篇論文提出的 TransNormerLLM 是首個基于線性注意力的 LLM，其在準確度和效率方面都優于傳統的 softmax 注意力。TransNormerLLM 的構建基于之前的線性注意力架構 TransNormer，同時也做了一些修改讓性能更優。TransNormerLLM 中關鍵性的改進包括位置嵌入、線性注意力加速、門控機制、張量歸一化和推理加速。

其中值得格外注意的一項改進是將 TransNormer 的 DiagAttention 替換成線性注意力，從而可提升全局的互動性能。研究者還引入了帶指數衰減的 LRPE 來解決 dilution 問題。此外，研究者還引入了 Lightning Attention（閃電注意力）這種全新技術，并表示其可以將線性注意力在訓練時的速度提升兩倍，并且其還能通過感知 IO 將內存用量減少 4 倍。不僅如此，他們還簡化了 GLU 和歸一化方法，而后者將整體的速度提升了 20%。他們還提出了一種穩健的推理算法，可以在不同的序列長度下保證數值穩定和恒定的推理速度，由此能提升模型在訓練和推理階段的效率。

為了驗證 TransNormerLLM 的效果，研究者精心收集了一個大型語料庫，其大小超過 6TB，token 數更是超過 2 萬億。為了確保數據的質量，他們還開發了一種用于過濾所收集語料庫的自清理（self-cleaning）策略。如表 1 所示，研究者對原始 TransNormer 模型進行了擴展，得到了參數量從 3.85 億到 1750 億的多個 TransNormerLLM 模型。然后他們基于新的大型語料庫進行了全面實驗和控制變量研究，結果表明新方法的性能優于基于 softmax 注意力的方法并且還有更快的訓練和推理速度。

表 1：TransNormerLLM 模型的多種變體

為了促進 LLM 領域的研究，上海人工智能實驗室和 OpenNLPLab 的這些研究者也將開源自己的預訓練模型，他們表示這是為了「讓研究者和實踐者基于我們的成果構建應用以及探索用于 LLM 的高效 transformer 結構」。

TransNormerLLM

架構改進

下面將簡單介紹 TransNormerLLM 的各個模塊以及研究者提出的一些改進措施。

改進一：位置編碼

TransNormer 中的較低層使用了 DiagAttention 來避免 dilution 問題。但是，這會導致 token 之間缺乏全局互動能力。為了解決這個問題，研究者為 TransNormerLLM 使用了帶指數衰減的 LRPE（線性化相對位置編碼），從而可在較低層保留完整的注意力。研究者把這種方法稱為 LRPE-d。

改進二：門控機制

門控可以增強模型的性能并使訓練過程平滑。研究者為 TransNormerLLM 使用了來自論文《Transformer quality in linear time》的 Flash 方法并在 token 混合中使用了門控式線性注意力（GLA）的結構。

為了進一步提升模型速度，他們還提出了 Simple GLU（SGLU），其去除了原始 GLU 結構的激活函數，因為門本身就能引入非線性。

改進三：張量歸一化

研究者使用了 TransNormer 中引入的 NormAttention。在 TransNormerLLM 中，他們使用一種新的簡單歸一化函數 SimpleRMSNorm（簡寫為 SRMSNorm）替換了 RMSNorm。

整體結構

圖 1 展示了 TransNormerLLM 的整體結構。

圖 1：新提出模型的整體架構

在該結構中，輸入 X 的更新通過兩個連續步驟完成：首先，其通過使用了 SRMSNorm 歸一化的門控式線性注意力（GLA）模塊。然后，再次通過使用了 SRMSNorm 歸一化的簡單門控式線性單元（SGLU）模塊。這種整體架構有助于提升模型的性能表現。下方給出了這個整體流程的偽代碼：

訓練優化

閃電注意力

為了加快注意力計算速度，研究者引入了閃電注意力（Lightning Attention）算法，這能讓新提出的線性注意力更適合 IO（輸入和輸出）處理。

算法 1 展示了閃電注意力的前向通過的實現細節，算法 2 則是后向通過的。研究者表示，他們還有一個可以更快計算梯度的實現版本，這會在未來發布。

模型并行化

為了在計算機集群上分散所有模型參數、梯度和優化器狀態張量，研究者使用了全分片數據并行（FSDP/Fully Sharded Data Parallel）。這種策略性分區方法可減少對每個 GPU 的內存占用，從而優化了內存利用率。為了進一步提高效率，他們使用了激活檢查點（Activation Checkpointing），這可減少后向通過過程中緩存在內存中的激活數量。相反，當計算這些梯度時，這些梯度會被移除并重新計算。該技術有助于提升計算效率和節省資源。此外，為了在減少 GPU 內存消耗的同時加快計算速度，研究者還使用了自動混合精度（AMP）。

除了上述成果外，研究者還更進一步通過對線性 transformer 執行模型并行化而進行了系統工程優化，其靈感很大程度上來自于英偉達的 Megatron-LM 模型并行化，在傳統的 Transformer 模型中，每個 transformer 層都有一個自注意力模塊，其后跟著一個兩層多層感知器（MLP）模塊。當使用 Megatron-LM 模型并行性時，是在這兩個模塊上獨立使用。類似地，TransNormerLLM 結構也是由兩個主要模塊構成：SGLU 和 GLA；這兩者的模型并行化分開執行。

穩健推理

這讓 TransNormerLLM 能以 RNN 的形式執行推理。算法 3 給出了這個過程的細節。但其中存在數值精度問題。

為了避免這些問題，研究者提出了穩健推理算法，見算法 4。

原推理算法和穩健推理算法得到的結果是一樣的。

語料庫

研究者從互聯網收集了大量可公開使用的文本，總大小超過 700TB。收集到的數據經由他們的數據預處理程序進行處理，如圖 2 所示，留下 6TB 的干凈語料庫，其中包含大約 2 萬億 token。為了提供更好的透明度，幫助用戶更好理解，他們對數據源進行了分門別類。表 2 給出了具體的類別情況。

圖 2：數據預處理流程

表 2：語料庫統計數據

實驗

研究者在 Metaseq 框架中使用 PyTorch 和 Trition 實現了 TransNormerLLM。模型的訓練使用了 Adam 優化器，FSDP 也被用于高效地將模型擴展到 NVIDIA A100 80G 集群。他們也適當地使用了模型并行技術來優化性能。

架構消融實驗

表 3：Transformer vs TransNormerLLM。在相同的配置下，當模型參數數量為 385M 和 1B 時，TransNormerLLM 的性能比 Transformer 分別好 5% 和 9%。

表 4：TransNormer vs TransNormerLLM。TransNormerLLM 的實驗結果是最好的。

表 5：位置編碼組合使用 LRPE+LRPE-d 得到的結果最優。

表 6：衰減溫度方面的消融實驗結果。結果表明新方法更優。

表 7：門控機制方面的消融實驗結果。使用該門控機制的模型表現更好。

表 8：GLA 激活函數的消融實驗結果。用不同激活函數得到結果差不多。

表 9：GLU 激活函數的消融實驗結果。去除激活函數對結果沒有負面影響。

表 10：歸一化函數。使用以下歸一化函數的結果差異不大。

圖 3：對 SRMSNorm 實現的性能評估

圖 4：線性注意力和閃電注意力之間的內存和速度比較

圖 5：推理時間和內存占用情況

系統優化

表 11：模型并行性性能

表 12：訓練不同大小的模型的效率

表 13：訓練 Transformer 和 TransNormerLLM 的最大上下文長度