傳統(tǒng)的注意力機(jī)制由于時(shí)間和空間復(fù)雜度的二次方增長(zhǎng)，以及在生成過程中鍵值緩存的內(nèi)存消耗不斷增加，限制了模型處理長(zhǎng)文本的能力。相關(guān)的解決方案包括減少計(jì)算復(fù)雜度、改進(jìn)記憶選擇和引入檢索增強(qiáng)語言建模。

檢索增強(qiáng)生成（RAG）和MemLong的記憶檢索流程。 (a) 當(dāng)檢索到的信息長(zhǎng)度超過模型的處理能力時(shí)，RAG甚至可能會(huì)降低生成性能（黃色）(b) MemLong利用外部檢索器來獲取歷史信息，然后將這些信息以鍵值對(duì)（K-V）的形式而不是文本形式傳遞給模型。

圖片

提出一種新方案MemLong，結(jié)合一個(gè)非可微分的檢索-記憶模塊和一個(gè)部分可訓(xùn)練的解碼器語言模型，來增強(qiáng)長(zhǎng)文本上下文的語言建模能力。

MemLong利用外部檢索器來檢索歷史信息，并通過細(xì)粒度、可控的檢索注意力機(jī)制，將語義級(jí)別的相關(guān)信息塊整合到模型中。這種方法不僅提高了模型處理長(zhǎng)文本的能力，還保持了信息分布的一致性，避免了訓(xùn)練過程中的分布偏移問題。

MemLong的一個(gè)示例：在底層，模型保持靜態(tài)，對(duì)整個(gè)數(shù)據(jù)塊Ci進(jìn)行因果語言建模，隨后，Ci被緩存為嵌入和鍵值對(duì)（K-V）形式。最后，上層進(jìn)行微調(diào)，以協(xié)調(diào)檢索偏好并整合檢索到的內(nèi)容。

圖片

MemLong的核心原理包括以下幾個(gè)方面：

• 檢索-記憶模塊：MemLong通過檢索-記憶模塊來存儲(chǔ)過去的上下文和知識(shí)，并利用這些存儲(chǔ)的嵌入向量來檢索輸入模型的塊級(jí)關(guān)鍵值（K-V）對(duì)。
• 檢索注意力機(jī)制：MemLong引入了一種細(xì)粒度、可控的檢索注意力機(jī)制，允許模型在局部上下文和通過檢索獲得的塊級(jí)過去上下文之間進(jìn)行注意力分配。

檢索因果注意力的說明。局部因果注意力應(yīng)用于最近的上下文，而通過檢索方法獲得的塊級(jí)鍵值對(duì)（K-V）允許雙向注意力，由于它們的歷史性質(zhì)，不會(huì)導(dǎo)致信息泄露。

圖片

• 動(dòng)態(tài)記憶管理：當(dāng)內(nèi)存溢出時(shí)，MemLong使用計(jì)數(shù)器機(jī)制智能地更新內(nèi)存，保留最有價(jià)值的信息，刪除最不相關(guān)的數(shù)據(jù)，以控制內(nèi)存大小并提高檢索效率。
• 推理過程：在處理超出模型最大處理長(zhǎng)度的輸入時(shí)，MemLong將文本存儲(chǔ)為記憶庫(kù)中的上下文信息，并在給定最近生成的文本塊時(shí)，使用檢索器顯式檢索過去的信息，通過索引對(duì)齊獲得額外的上下文信息。

通過這些原理，MemLong在多個(gè)長(zhǎng)文本上下文語言建模基準(zhǔn)測(cè)試中表現(xiàn)出色，證明了其在處理長(zhǎng)文本方面的有效性和優(yōu)越性。MemLong可以將單個(gè)3090 GPU上的上下文長(zhǎng)度從4k擴(kuò)展到80k。

不同上下文窗口擴(kuò)展模型在PG19、Proof-pile、BookCorpus、Wikitext-103上的滑動(dòng)窗口困惑度。所有實(shí)驗(yàn)均在一塊3090 24GB GPU上進(jìn)行。LongLLaMA-3B和MemLong-3B帶有?標(biāo)記的表示在沒有內(nèi)存的情況下進(jìn)行評(píng)估，而帶有?標(biāo)記的LongLLaMA-3B表示在無限內(nèi)存的情況下進(jìn)行評(píng)估。還評(píng)估了MemLong在4K/32K內(nèi)存場(chǎng)景下的表現(xiàn)。"- / 6.95"表示模型在單個(gè)GPU上導(dǎo)致內(nèi)存不足（OOM）錯(cuò)誤，而在雙GPU上則產(chǎn)生了相應(yīng)的結(jié)果。

圖片

https://arxiv.org/pdf/2408.16967
https://github.com/Bui1dMySea/MemLong
MemLong: Memory-Augmented Retrieval for Long Text Modeling

本文轉(zhuǎn)載自??PaperAgent??