近日，上海人工智能實(shí)驗(yàn)室和 OpenNLPLab 的一個研究團(tuán)隊(duì)提出了一種新的大型語言模型 TransNormerLLM，其中完全拋棄了基于 Softmax 的注意力機(jī)制，而是使用了新提出的線性注意力。據(jù)介紹，TransNormerLLM 是首個基于線性注意力的大型語言模型（LLM），其在準(zhǔn)確度和效率方面的表現(xiàn)優(yōu)于傳統(tǒng)的基于 Softmax 注意力的模型。研究者也將發(fā)布其預(yù)訓(xùn)練模型的開源版本。

• 論文：https://arxiv.org/abs/2307.14995
• 模型：https://github.com/OpenNLPLab/TransnormerLLM

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

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

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

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

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

表 1：TransNormerLLM 模型的多種變體

為了促進(jìn) LLM 領(lǐng)域的研究，上海人工智能實(shí)驗(yàn)室和 OpenNLPLab 的這些研究者也將開源自己的預(yù)訓(xùn)練模型，他們表示這是為了「讓研究者和實(shí)踐者基于我們的成果構(gòu)建應(yīng)用以及探索用于 LLM 的高效 transformer 結(jié)構(gòu)」。

TransNormerLLM

架構(gòu)改進(jìn)

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

改進(jìn)一：位置編碼

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

改進(jìn)二：門控機(jī)制

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

為了進(jìn)一步提升模型速度，他們還提出了 Simple GLU（SGLU），其去除了原始 GLU 結(jié)構(gòu)的激活函數(shù)，因?yàn)殚T本身就能引入非線性。

改進(jìn)三：張量歸一化

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

整體結(jié)構(gòu)

圖 1 展示了 TransNormerLLM 的整體結(jié)構(gòu)。

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

訓(xùn)練優(yōu)化

閃電注意力

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

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

模型并行化

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

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

穩(wěn)健推理

這讓 TransNormerLLM 能以 RNN 的形式執(zhí)行推理。算法 3 給出了這個過程的細(xì)節(jié)。但其中存在數(shù)值精度問題。

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

原推理算法和穩(wěn)健推理算法得到的結(jié)果是一樣的。

語料庫

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

圖 2：數(shù)據(jù)預(yù)處理流程

表 2：語料庫統(tǒng)計數(shù)據(jù)

實(shí)驗(yàn)

研究者在 Metaseq 框架中使用 PyTorch 和 Trition 實(shí)現(xiàn)了 TransNormerLLM。模型的訓(xùn)練使用了 Adam 優(yōu)化器，F(xiàn)SDP 也被用于高效地將模型擴(kuò)展到 NVIDIA A100 80G 集群。他們也適當(dāng)?shù)厥褂昧四Ｐ筒⑿屑夹g(shù)來優(yōu)化性能。

架構(gòu)消融實(shí)驗(yàn)

圖片

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

圖片

表 4：TransNormer vs TransNormerLLM。TransNormerLLM 的實(shí)驗(yàn)結(jié)果是最好的。

表 5：位置編碼組合使用 LRPE+LRPE-d 得到的結(jié)果最優(yōu)。

表 6：衰減溫度方面的消融實(shí)驗(yàn)結(jié)果。結(jié)果表明新方法更優(yōu)。

表 7：門控機(jī)制方面的消融實(shí)驗(yàn)結(jié)果。使用該門控機(jī)制的模型表現(xiàn)更好。

圖片

表 8：GLA 激活函數(shù)的消融實(shí)驗(yàn)結(jié)果。用不同激活函數(shù)得到結(jié)果差不多。

表 9：GLU 激活函數(shù)的消融實(shí)驗(yàn)結(jié)果。去除激活函數(shù)對結(jié)果沒有負(fù)面影響。

表 10：歸一化函數(shù)。使用以下歸一化函數(shù)的結(jié)果差異不大。

圖 3：對 SRMSNorm 實(shí)現(xiàn)的性能評估

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

圖 5：推理時間和內(nèi)存占用情況

系統(tǒng)優(yōu)化

表 11：模型并行性性能

表 12：訓(xùn)練不同大小的模型的效率

表 13：訓(xùn)練 Transformer 和 TransNormerLLM 的最大上下文長度