一般來(lái)說(shuō)，大模型預(yù)訓(xùn)練時(shí)文本長(zhǎng)度是固定的，如果想要支持更長(zhǎng)文本，就需要對(duì)模型進(jìn)行微調(diào)。但是訓(xùn)練具有長(zhǎng)上下文的 LLM 計(jì)算成本很高，需要大量的訓(xùn)練時(shí)間和 GPU 資源。

比如，訓(xùn)練一個(gè)具有 8192 長(zhǎng)度上下文的模型，相比于 2048 長(zhǎng)度上下文，需要 16 倍的計(jì)算資源。就算如此，上下文長(zhǎng)度對(duì)模型性能至關(guān)重要，因?yàn)樗砹?LLM 回應(yīng)時(shí)對(duì)整個(gè)上下文清晰理解的能力。

近日，MIT 與香港中文大學(xué)聯(lián)合研究，提出了 LongLoRA。它是一種有效的微調(diào)方法，以有限的計(jì)算成本擴(kuò)展了預(yù)訓(xùn)練大型語(yǔ)言模型上下文大小。

論文地址：https://arxiv.org/pdf/2309.12307.pdf

項(xiàng)目地址：https://github.com/dvlab-research/LongLoRA

本文從兩個(gè)方面加快了 LLM 的上下文擴(kuò)展。

一方面，盡管在推理過(guò)程中需要密集的全局注意力，但通過(guò)稀疏的局部注意力可以有效且高效地對(duì)模型進(jìn)行微調(diào)。本文提出的 shift short attention 有效地實(shí)現(xiàn)了上下文擴(kuò)展，節(jié)省了大量的計(jì)算，與使用 vanilla attention 進(jìn)行微調(diào)的性能相似。

另一方面，用于上下文擴(kuò)展的 LoRA 在可訓(xùn)練嵌入和歸一化的前提下工作得很好。LongLoRA 在 LLaMA2 模型從 7B/13B 到 70B 的各種任務(wù)上都展現(xiàn)了很好的結(jié)果。在單臺(tái) 8x A100 設(shè)備上，LongLoRA 將 LLaMA2 7B 從 4k 上下文擴(kuò)展到 100k， LLaMA2 70B 擴(kuò)展到 32k。LongLoRA 擴(kuò)展了模型的上下文，同時(shí)保留了其原始架構(gòu)，并與大多數(shù)現(xiàn)有技術(shù)兼容，如 FlashAttention-2。為使 LongLoRA 實(shí)用，研究者收集了一個(gè)數(shù)據(jù)集 LongQA，用于監(jiān)督微調(diào)。該數(shù)據(jù)集包含超過(guò) 3k 個(gè)長(zhǎng)上下文問(wèn)題 - 答案對(duì)。

LongLoRA 的能夠在注意力水平和權(quán)重水平上加速預(yù)訓(xùn)練大型語(yǔ)言模型的上下文擴(kuò)展。亮點(diǎn)如下：

• Shift short attention 易于實(shí)現(xiàn)，與 Flash-Attention 兼容，且在推理過(guò)程中不需要使用。
• 發(fā)布了所有模型，包括從 7B 到 70B 的模型，上下文長(zhǎng)度從 8k 到 100k，包括 LLaMA2-LongLoRA-7B-100k、LLaMA2-LongLoRA-13B-64k 和 LLaMA2-LongLoRA-70B-32k。
• 建立了一個(gè)長(zhǎng)上下文 QA 數(shù)據(jù)集 LongQA，用于監(jiān)督微調(diào)。研究者已經(jīng)發(fā)布了 13B 和 70B 32k 型號(hào)的 SFT、Llama-2-13b-chat-longlora-32k-sft 和 Llama-2-70b-chat-longlora-32k-sft，并將在下個(gè)月發(fā)布數(shù)據(jù)集。

LongLoRA 技術(shù)細(xì)節(jié) 

Shift short attention

標(biāo)準(zhǔn)自注意力模式的計(jì)算開(kāi)銷為 O (n^2 )，使得長(zhǎng)序列上的 LLM 內(nèi)存開(kāi)銷高且速度慢。為了在訓(xùn)練中避免這個(gè)問(wèn)題，本文提出了 shift short attention（S^2 -Attn），如下圖 2 所示。

研究者驗(yàn)證了微調(diào)的重要性，如下表 1 所示。如果沒(méi)有微調(diào)，隨著上下文長(zhǎng)度的增長(zhǎng)，即使配備了適當(dāng)?shù)奈恢们度?，模型的表現(xiàn)也會(huì)變差。

現(xiàn)有的 efficient  attention 設(shè)計(jì)也可以提高長(zhǎng)上下文語(yǔ)言模型的效率。在下表 2 中，研究者將 S^2 -Attn 與幾種典型的 efficient  attention 進(jìn)行了比較，可以發(fā)現(xiàn)，前者不僅能夠?qū)崿F(xiàn)高效的微調(diào)，還支持 full attention 測(cè)試。

此外，S^2 -Attn 容易實(shí)現(xiàn)，它只涉及兩個(gè)步驟：(1) 轉(zhuǎn)換半注意力頭中的 token (2) 將 token 維度的特征移至批次維度。這個(gè)過(guò)程使用幾行代碼就夠了。

改進(jìn)長(zhǎng)上下文 LoRA

LoRA 是一種有效且流行的方法，可使 LLM 適應(yīng)其他數(shù)據(jù)集。與完全微調(diào)相比，它節(jié)省了很多可訓(xùn)練參數(shù)和內(nèi)存成本。然而，將 LLM 從短上下文長(zhǎng)度調(diào)整為長(zhǎng)上下文長(zhǎng)度并不容易。研究者觀察到 LoRA 和完全微調(diào)之間存在明顯的差距。如下表 3 所示，隨著目標(biāo)上下文長(zhǎng)度的增大，LoRA 和完全微調(diào)之間的差距也會(huì)增大。

為了彌補(bǔ)這一差距，研究者打開(kāi)嵌入層和歸一化層進(jìn)行訓(xùn)練。如表 3 所示，它們占用的參數(shù)有限，但對(duì)長(zhǎng)上下文適應(yīng)有影響。特別是歸一化層，在整個(gè) LLaMA2 7B 的參數(shù)占比僅為 0.004%。在實(shí)驗(yàn)中，研究者將這種改進(jìn)的 LoRA 表示為 LoRA+。

實(shí)驗(yàn)及結(jié)果

研究者擴(kuò)展了預(yù)訓(xùn)練的 7B、13B 和 70B LLaMA2 模型。7B 模型的最大擴(kuò)展上下文窗口大小為 100k，13B 模型的最大擴(kuò)展上下文窗口大小為 65536，70B 模型的最大擴(kuò)展上下文窗口大小為 32768。 

研究者沿用了 Position Interpolation 中的大部分訓(xùn)練超參數(shù)，不過(guò)批大小更小，因?yàn)橹皇窃谀承┣闆r下使用單臺(tái) 8×A100 GPU 設(shè)備。所有模型都通過(guò)下一個(gè) token 預(yù)測(cè)目標(biāo)進(jìn)行微調(diào)。研究者使用 AdamW，其中 β_1 = 0.9，β_2 = 0.95。7B 和 13B 模型的學(xué)習(xí)率設(shè)定為 2 × 10^?5，70B 模型的學(xué)習(xí)率設(shè)定為 10^?5。

他們還使用了線性學(xué)習(xí)率預(yù)熱。權(quán)重衰減為零。每臺(tái)設(shè)備的批大小設(shè)為 1，梯度累積步驟設(shè)為 8，這意味著使用 8 個(gè) GPU，全局批大小等于 64。模型進(jìn)行了 1000 步的訓(xùn)練。

研究者使用 Redpajama 數(shù)據(jù)集進(jìn)行訓(xùn)練，并構(gòu)建了一個(gè)長(zhǎng)上下文 QA 數(shù)據(jù)集 LongQA，用于監(jiān)督微調(diào)。Redpajama 微調(diào)的模型呈現(xiàn)了良好的困惑度，但它們的聊天能力是有限的。研究者收集了超過(guò) 3k 個(gè)問(wèn)題 - 答案對(duì)，它們都是與技術(shù)論文、科幻小說(shuō)和其他書(shū)籍等材料有關(guān)的。設(shè)計(jì)的問(wèn)題包括總結(jié)、關(guān)系、人物等。

從下表 4 中可以發(fā)現(xiàn)，在相同的訓(xùn)練和評(píng)估上下文長(zhǎng)度的情況下，困惑度隨著上下文大小的增加而降低。

在下表 5 中，研究者進(jìn)一步考察了在單臺(tái) 8×A100 設(shè)備上可微調(diào)的最大上下文長(zhǎng)度。他們分別將 LLaMA2 7B、13B 和 70B 擴(kuò)展到 100k、65536 和 32768 上下文長(zhǎng)度。LongLoRA 在這些超大設(shè)置上取得了令人滿意的結(jié)果。此外，實(shí)驗(yàn)還發(fā)現(xiàn)擴(kuò)展模型在較小的上下文長(zhǎng)度上會(huì)出現(xiàn)一些困惑度下降。

在下表 6 中，研究者將模型與其他開(kāi)放式 LLM 在 LongChat 中引入的主題檢索任務(wù)上進(jìn)行比較。這個(gè)任務(wù)是從很長(zhǎng)的對(duì)話中檢索目標(biāo)話題，對(duì)話長(zhǎng)度從 3k、6k、10k、13k 到 16k 不等。

消融實(shí)驗(yàn)

在下表 7 中，研究者將 LLaMA2 7B 細(xì)分為各種類型的層。他們分析了 FLOPs：對(duì)于 full attention，隨著上下文長(zhǎng)度的增加，Attn 的比例也急劇增加。例如，在上下文長(zhǎng)度為 8192 時(shí)，Attn 占總 FLOP 的 24.5%，而在上下文長(zhǎng)度為 65536 時(shí)，則增至 72.2%。當(dāng)使用 S^2 -Attn 時(shí)，則下降到 39.4%。

下表 8 展示了在 PG19 驗(yàn)證集上擴(kuò)展到 8192 上下文長(zhǎng)度時(shí)， LLaMA2 7B 模型的復(fù)雜度與微調(diào)步驟之間的關(guān)系?？梢园l(fā)現(xiàn)，如果不進(jìn)行微調(diào)，在第 0 步時(shí)，模型的長(zhǎng)上下文能力有限。完全微調(diào)比低階訓(xùn)練收斂得更快。兩者在 200 步后逐漸接近，最后沒(méi)有出現(xiàn)大的差距。

下表 2 顯示了微調(diào)過(guò)程中不同注意力模式的效果。

效果展示

模型在閱讀《哈利?波特》的內(nèi)容后，能夠告訴你斯內(nèi)普為什么看起來(lái)不喜歡哈利，甚至還能總結(jié)人物之間的關(guān)系。

不僅如此，給它一篇論文，還能幫助你立刻了解相關(guān)信息。

更多詳細(xì)內(nèi)容，請(qǐng)參閱原文。