ChatGPT 的誕生，讓基于 Transformer 的大型語(yǔ)言模型 (LLM) 為通用人工智能（AGI）鋪開(kāi)了一條革命性的道路，并在知識(shí)庫(kù)、人機(jī)交互、機(jī)器人等多個(gè)領(lǐng)域得到應(yīng)用。然而，目前存在一個(gè)普遍的限制：由于資源受限，當(dāng)前大多 LLM 主要是在較短的文本上進(jìn)行預(yù)訓(xùn)練，導(dǎo)致它們?cè)谳^長(zhǎng)上下文方面的表現(xiàn)較差，而長(zhǎng)上下文在現(xiàn)實(shí)世界的環(huán)境中是更加常見(jiàn)的。
最近的一篇綜述論文對(duì)此進(jìn)行了全面的調(diào)研，作者重點(diǎn)關(guān)注了基于 Transformer 的 LLM 模型體系結(jié)構(gòu)在從預(yù)訓(xùn)練到推理的所有階段中優(yōu)化長(zhǎng)上下文能力的進(jìn)展。
論文鏈接：https://arxiv.org/pdf/2311.12351.pdf
論文首先分析了使用當(dāng)前基于 Transformer 的模型處理長(zhǎng)上下文輸入和輸出的問(wèn)題。然后，提供了一個(gè)全面的分類體系，以指導(dǎo) Transformer 架構(gòu)升級(jí)的領(lǐng)域，來(lái)解決這些問(wèn)題。作者對(duì)長(zhǎng)上下文 LLM 廣泛使用的評(píng)估需求進(jìn)行了調(diào)研，包括數(shù)據(jù)集、度量標(biāo)準(zhǔn)和基準(zhǔn)模型，以及一些令人驚奇的優(yōu)化工具包，如庫(kù)、系統(tǒng)和編譯器，以增強(qiáng) LLM 在不同階段的效率和功效。最后，文章進(jìn)一步討論了這一領(lǐng)域未來(lái)研究的主要挑戰(zhàn)和潛在方向。作者還建立了一個(gè)倉(cāng)庫(kù)，匯總了相關(guān)文獻(xiàn)，并提供實(shí)時(shí)更新 https://Github.com/Strivin0311/long-llms-learning。
綜述概覽
文章從基本的語(yǔ)言建模目標(biāo) (第 2.1 節(jié)) 開(kāi)始，內(nèi)容涵蓋從典型的建模階段到在基于 Transformer 的僅解碼 LLM 中找到的關(guān)鍵架構(gòu)模塊，如圖 1 (a) 所示。隨后，作者對(duì) LLM 在遇到擴(kuò)展上下文窗口時(shí)的架構(gòu)限制進(jìn)行了簡(jiǎn)要分析 (第 2.2 節(jié))。最后提出了一個(gè)全面的方法論分類法 (第 2.3 節(jié))，旨在通過(guò)架構(gòu)創(chuàng)新增強(qiáng) LLM 的長(zhǎng)上下文能力 (見(jiàn)圖 1 (b))。這個(gè)分類法作為文章的第 3、4、5、6、7 節(jié)的指南。
圖 1：文章核心概述：(a) 現(xiàn)代基于 Transformer 的僅解碼 LLMs 的典型架構(gòu)解剖圖，右上角有圖例；(b) 用于增強(qiáng) Transformer 架構(gòu)模塊的方法論分類法（與 (a) 相對(duì)應(yīng)的顏色）：高效注意力（注意力核心的子模塊），長(zhǎng)期記憶（針對(duì) KV 緩存），外推性 PEs（針對(duì)位置嵌入模塊），上下文處理（與上下文預(yù) / 后處理有關(guān)）和雜項(xiàng)（整個(gè)解碼器塊以及損失模塊通用）。
長(zhǎng)上下文，目前有哪些難點(diǎn)待突破？
注意力復(fù)雜度。在典型情況下 L ? d，MHA 的計(jì)算復(fù)雜性可以簡(jiǎn)潔總結(jié)如下：它涉及 O (L 2d) 的時(shí)間復(fù)雜度，包括 QKV 投影的 O (Ld2)，計(jì)算 P 的 O (L 2d)，通過(guò) softmax 運(yùn)算獲取 A 的 O (L 2 )，A 與 V 相乘的 O (L 2d)，以及輸出投影 O 的 O (Ld2)。它還產(chǎn)生 O (L 2) 的空間復(fù)雜度，包括 Q、K、V、O 的嵌入的 O (Ld)，以及額外的 O (L 2) 緩沖區(qū)用于存儲(chǔ)權(quán)重 P 和 A。因此，隨著序列長(zhǎng)度的增加，時(shí)間和空間計(jì)算成本都呈二次增加，這對(duì)于訓(xùn)練和推理可能都是繁重的。
上下文記憶。LLM 缺乏顯式的記憶機(jī)制，完全依賴 KV 緩存來(lái)存儲(chǔ)列表中所有先前 token 的表示。這種設(shè)計(jì)一旦在一個(gè)調(diào)用中完成查詢，Transformer 在后續(xù)的調(diào)用中不會(huì)保留或召回任何先前的狀態(tài)或序列，除非整個(gè)歷史記錄逐個(gè) token 重新加載到 KV 緩存中。因此，Transformer 在每個(gè)調(diào)用中僅具有一個(gè)上下文工作記憶，而不是像長(zhǎng)短時(shí)記憶 (LSTM) 這樣的內(nèi)在記憶機(jī)制。這種無(wú)狀態(tài)性在并行性方面提供了計(jì)算優(yōu)勢(shì)，但在聊天機(jī)器人應(yīng)用等需要長(zhǎng)期記憶保留的任務(wù)中的挑戰(zhàn)很明顯。
最大長(zhǎng)度約束。在訓(xùn)練階段，工程師通常需要確定一個(gè)關(guān)鍵的超參數(shù) max-length，本文中表示為 L_max。這個(gè)超參數(shù)代表了批次中任何訓(xùn)練樣本的序列長(zhǎng)度的上限，通常根據(jù)可用的計(jì)算資源設(shè)置為 1k、2k 或 4k，以避免在 GPU 上發(fā)生內(nèi)存溢出 (OOM) 錯(cuò)誤。在推理階段，LLM 的服務(wù)提供者還必須限制用戶提示的長(zhǎng)度或自動(dòng)截?cái)嗨鼈円耘c預(yù)定義的 L_max 對(duì)齊，即使推理資源通常比訓(xùn)練階段更豐富。需要注意的是 Transformer 的任何模塊在本質(zhì)上都不需要這樣的限制，因?yàn)樗袑W(xué)習(xí)的權(quán)重僅依賴于維度大小。因此，理論上只要資源足夠，Transformer 可以處理任意長(zhǎng)度的序列。然而，當(dāng)前的語(yǔ)言模型在處理超過(guò) L_max 的輸入序列時(shí)通常表現(xiàn)出明顯的性能下降，經(jīng)常導(dǎo)致重復(fù)和不切實(shí)際的輸出。
改進(jìn)的新方法
對(duì)于上述限制，有多種改進(jìn)方法可以探索，例如在訓(xùn)練過(guò)程中減少注意力復(fù)雜性、設(shè)計(jì)高效的記憶機(jī)制，以及增強(qiáng)長(zhǎng)度外推的能力，該模型在短序列上進(jìn)行訓(xùn)練，但在推理過(guò)程中對(duì)更長(zhǎng)的序列進(jìn)行測(cè)試。
因此，論文全面回顧了致力于改進(jìn) LLM 長(zhǎng)上下文能力的各個(gè)階段的最新方法，并將它們組織成一個(gè)統(tǒng)一的分類法，如圖 1 (b) 所示。具體而言，這些方法被分為五個(gè)主要的類別，如下：
高效注意力 (論文第 3 節(jié))：這些方法側(cè)重于實(shí)現(xiàn)具有降低計(jì)算要求的高效注意力機(jī)制，甚至實(shí)現(xiàn)了線性復(fù)雜度。通過(guò)這樣做，它們能夠通過(guò)直接在預(yù)訓(xùn)練階段增加 L_max 來(lái)推進(jìn) LLM 在推理期間的有效上下文長(zhǎng)度邊界。
長(zhǎng)期記憶 (論文第 4 節(jié))：為了解決上下文工作記憶的局限性，一些方法旨在設(shè)計(jì)明確的記憶機(jī)制，彌補(bǔ) LLM 中缺乏高效和有效的長(zhǎng)期記憶的不足。
外推性 PEs (論文第 5 節(jié))：最新的研究致力于通過(guò)改進(jìn)現(xiàn)有位置編碼方案的外推性能來(lái)增強(qiáng) LLM 的長(zhǎng)度泛化能力。
上下文處理 (論文第 6 節(jié))：除了增強(qiáng)特定低級(jí) Transformer 模塊的方法外，一些方法涉及對(duì)現(xiàn)成的 LLM 與額外的上下文預(yù) / 后處理。這些方法確保每次調(diào)用 LLM 時(shí)輸入始終滿足最大長(zhǎng)度要求，并通過(guò)引入多個(gè)調(diào)用開(kāi)銷打破上下文窗口限制。
雜項(xiàng) (論文第 7 節(jié))：探討了各種一般且有價(jià)值的方法，這些方法不容易歸入前面四類，為推進(jìn) LLM 的長(zhǎng)上下文能力提供了更廣泛的視角。
未來(lái)方向
論文的第 3、4、5、6 節(jié)中討論了該領(lǐng)域取得的顯著進(jìn)展，但仍然存在一些挑戰(zhàn)。下面是對(duì)一些關(guān)鍵挑戰(zhàn)的探討以及未來(lái)在增強(qiáng)基于 Transformer 的 LLM 的長(zhǎng)上下文能力方面進(jìn)行研究和開(kāi)發(fā)的潛在方向，重點(diǎn)關(guān)注架構(gòu)的增強(qiáng)。
注意力 Trade-off。在第 3 節(jié)，作者探討了高效注意方法往往涉及在保持全尺度注意力依賴性（例如局部注意力）或通過(guò)近似注意力提高注意力分?jǐn)?shù)精度以減輕標(biāo)準(zhǔn)注意內(nèi)核的計(jì)算需求之間的微妙權(quán)衡。然而，隨著上下文的延長(zhǎng)，話語(yǔ)結(jié)構(gòu)和相互關(guān)聯(lián)的信息變得越來(lái)越復(fù)雜，需要捕捉全局、長(zhǎng)距離的依賴性，同時(shí)保持精確的相關(guān)性。
解決這一挑戰(zhàn)需要在計(jì)算效率和盡可能保留注意模式精度之間找到最佳平衡。因此，在長(zhǎng)上下文 LLM 領(lǐng)域，這仍然是一個(gè)持續(xù)追求的目標(biāo)。最近的創(chuàng)新如 Flash Attention，探索了算法級(jí)別之外的 IO 感知解決方案，極大地提高了運(yùn)行時(shí)和記憶開(kāi)銷的效率，而不會(huì)喪失注意精度。這是在實(shí)際應(yīng)用中解決這個(gè)問(wèn)題的一個(gè)激動(dòng)人心的潛在途徑。此外，可以探索在「即插即用」替代方案中集成先前的高效策略，利用強(qiáng)大的 GPU 內(nèi)核編程工具 (如 CUDA) 或更輕量級(jí)的 Triton。
記憶效果和效率。正如在文章第 2.1、2.2 節(jié)中前面討論的，作者已經(jīng)概述了由于缺乏明確的記憶機(jī)制，僅依賴上下文內(nèi)工作記憶以及在延長(zhǎng)上下文交互期間 KV 緩存記憶消耗顯著增加而產(chǎn)生的限制。這些挑戰(zhàn)共同強(qiáng)調(diào)了在基于 Transformer 的 LLM 領(lǐng)域需要更有效和高效的記憶機(jī)制。雖然第 4 節(jié)中引入了各種長(zhǎng)期記憶機(jī)制，但它們受到其復(fù)雜啟發(fā)式設(shè)計(jì)引入的額外記憶開(kāi)銷的限制，因此隨著時(shí)間的推移可能導(dǎo)致性能下降。為了解決這一挑戰(zhàn)，研究人員可以從最近的進(jìn)展中汲取靈感，比如 Paged Attention，研發(fā)更有效的記憶存儲(chǔ)策略，增強(qiáng)讀 / 寫吞吐量。
長(zhǎng)度外推挖掘。在第 5 節(jié)中，作者對(duì)與基于 Transformer 的模型的長(zhǎng)度外推相關(guān)的挑戰(zhàn)進(jìn)行了徹底的分析，重點(diǎn)關(guān)注了位置嵌入的普遍設(shè)計(jì)。文章提供了對(duì)最近突破的全面概述，特別是應(yīng)用于 RoPE 的擴(kuò)展策略，作者相信這在解決外推限制方面具有重要的前景。值得注意的是，這些進(jìn)步往往依賴于對(duì)復(fù)雜高維位置嵌入屬性的簡(jiǎn)化觀察，并包含簡(jiǎn)單的啟發(fā)式調(diào)整。作者對(duì)使用高維嵌入來(lái)建模序列性的理論基礎(chǔ)提出質(zhì)疑，并探索在這些啟發(fā)式設(shè)計(jì)的指導(dǎo)下引導(dǎo)具有許多超參數(shù)的可學(xué)習(xí)嵌入的潛在復(fù)蘇。作者認(rèn)為未來(lái)的研究應(yīng)該深入探討這一領(lǐng)域，尤其是在 Transformer 設(shè)置下為建模序列性開(kāi)發(fā)健壯的理論框架方面，比如 CLEX 所實(shí)現(xiàn)的內(nèi)容。
特定但通用目標(biāo)。前文已經(jīng)為長(zhǎng)文本建模量身定制的具體目標(biāo)做了討論，但值得注意的是，許多目標(biāo)僅限于某些類型的任務(wù)，或者僅與 MLM 目標(biāo)兼容，而不是如今更普遍的 CLM 目標(biāo)。這突顯了需要特定但普遍適用的因果語(yǔ)言建模目標(biāo)，可以在模型訓(xùn)練的早期有效捕捉長(zhǎng)距離依賴性。通過(guò)與先前提到的目標(biāo)相一致，這是可能實(shí)現(xiàn)的。
可靠的度量需求。在評(píng)估度量方面，文章的第 8 節(jié)中研究了許多可選項(xiàng)。根據(jù)在評(píng)估中的先前經(jīng)驗(yàn)，常用的度量，如 ROUGE 分?jǐn)?shù)，與人類判斷分?jǐn)?shù)存在顯著差異，后者可以看作是「神諭」。隨著 LLM 在現(xiàn)實(shí)世界場(chǎng)景中的快速部署，越來(lái)越迫切地需要更可靠的度量來(lái)評(píng)估長(zhǎng)上下文能力，特別是在生成性任務(wù)中，其中精確的真實(shí)性難以捉摸。一個(gè)有希望的途徑涉及利用最先進(jìn)的 LLM (如 GPT4) 的魯棒性作為人類評(píng)審的替代，盡管相關(guān)的高成本仍然在更廣泛地在研究界中采用方面帶來(lái)挑戰(zhàn)。
更多研究細(xì)節(jié)，可參見(jiàn)原論文。
