透過多粒度自我注入擴展上下文窗口：雙劍合璧

Core Concepts

本文提出了一種名為 SharedLLM 的新型架構，旨在解決大型語言模型上下文窗口有限的問題，透過多粒度上下文壓縮和查詢感知信息檢索，在效率和性能之間取得平衡。

Abstract

論文資訊

標題：雙劍合璧：透過多粒度自我注入擴展上下文窗口 (Two are better than one: Context window extension with multi-grained self-injection)
作者：Wei Han, Pan Zhou, Soujanya Poria, Shuicheng Yan
發佈日期：2024 年 10 月 25 日
類型：研究論文
研究目標
本研究旨在解決當代大型語言模型 (LLM) 上下文窗口有限的問題，這個問題阻礙了它們在需要處理長文本的各種領域中的廣泛應用。
方法
SharedLLM 採用分層架構，由上層模型和下層模型組成，兩個模型都從相同的預訓練短上下文 LLM（例如 LLaMA-2）初始化。

下層模型（壓縮器）： 將長輸入上下文分解成較小的塊，並使用相同的 LLM 模型將每個上下文塊壓縮成緊湊的結構化表示形式，這些表示形式組織成一個稱為上下文樹的二叉樹，用於存儲不同級別的多粒度信息。
上層模型（解碼器）： 接收輸入文本的後一部分（運行上下文，例如問題），並根據深度優先搜索算法從上下文樹中檢索相關信息，然後通過交叉注意力將檢索到的信息與輸入信息整合，從而使模型能夠回答問題或執行語言建模。
主要發現

SharedLLM 在語言建模任務中表現出出色的外推能力，即使在訓練期間僅看到長達 8K 個詞元的序列，也能夠處理長達 128K 個詞元的序列。
在輸入長度從數千到數百萬個詞元的長上下文指令遵循任務中，SharedLLM 提供了與其他幾個強基準相當的結果。
SharedLLM 在速度和 GPU 內存利用率方面都表現出很高的計算效率，在保持相對較低內存消耗的同時，提供了比所有基準測試快幾倍的速度。
結論
SharedLLM 透過利用自我注入機制來適應一對短上下文 LLM 以進行高效的長上下文建模，提供了一種可擴展的方法，可以將上下文窗口擴展到任意長度。
局限性和未來研究方向

雖然 SharedLLM 在模型架構級別在效率和性能之間取得了相對較好的平衡，但可以通過在系統和硬件級別進行優化來實現進一步的改進。
雖然這項工作實現了一種簡單有效的檢索機制，但並未探索更高級的檢索技術，例如 BM25 和 Graph-RAG，這些技術可能會進一步提高性能。

Stats

SharedLLM 在語言建模任務中，使用 8,192 個詞元的文本進行訓練，並在長度從 4K 到 128K 的序列上進行評估。
在監督微調中，SharedLLM 使用了由 RedPajama 和 LongAlpaca 組成的混合數據集，輸入長度在 1200 到 8192 個詞元之間。
SharedLLM 的訓練是在一台配備 8 個 A800 GPU 的機器上進行的，語言建模的批次大小為每個 GPU 1 個樣本，梯度累積為 16 步，監督微調的批次大小為每個 GPU 1 個樣本，梯度累積為 1 步。
SharedLLM 的上下文塊大小設置為 1,024 個詞元（語言建模）或 512 個詞元（監督微調），樹高為 3，壓縮比為 8。

Key Insights Distilled From

Two are better than one: Context window extension with multi-grained self-injection

by Wei Han, Pan... at arxiv.org 10-28-2024

https://arxiv.org/pdf/2410.19318.pdf

Deeper Inquiries

SharedLLM 如何與其他上下文窗口擴展技術（例如位置插值和稀疏注意力）相結合，以進一步提高性能？

SharedLLM 可以與其他上下文窗口擴展技術結合，進一步提高性能，以下是一些可能的策略：

結合位置插值： 位置插值技術，如 YaRN，可以通過調整位置編碼來處理更長的上下文。SharedLLM 可以將其整合到上層模型中，以增強其對長距離依賴關係的建模能力。具體來說，可以在上層模型的 self-attention 模塊中使用位置插值，使其能夠更好地處理來自下層模型的長距離上下文信息。
結合稀疏注意力： 稀疏注意力機制，如 Longformer 使用的注意力模式，可以減少計算複雜度，同時允許模型關注更長的上下文。SharedLLM 可以通過在下層模型中採用稀疏注意力來提高效率，使其能夠處理更大的文本塊，並為上層模型提供更豐富的上下文信息。
混合架構： 可以構建一個混合架構，結合 SharedLLM、位置插值和稀疏注意力等技術的優勢。例如，可以使用位置插值來增強上層模型的長距離建模能力，同時使用稀疏注意力來提高下層模型的效率。
總之，SharedLLM 提供了一個靈活的框架，可以與其他上下文窗口擴展技術相結合，以進一步提高性能。通過探索這些組合的可能性，可以開發出更強大、更高效的長文本處理模型。

上下文樹結構的設計是否可以針對特定任務進行優化，例如通過使用更複雜的樹結構或不同的壓縮比？

是的，上下文樹結構的設計可以針對特定任務進行優化，以下是一些可行的方向：

使用更複雜的樹結構： 針對不同任務，可以探索更複雜的樹結構，例如：

N 叉樹： 可以使用 N 叉樹（N>2）來更細粒度地組織信息，例如，可以根據文本的語義結構將其分成不同的段落或章節，並將其存儲在不同的分支中。
有向無環圖 (DAG)： 對於需要捕捉更複雜關係的任務，可以使用 DAG 來表示文本單元之間的依賴關係，例如，在問答任務中，可以使用 DAG 來表示問題、答案和相關文本片段之間的關係。

使用不同的壓縮比： 可以根據任務需求調整不同層級節點的壓縮比，例如：

對於需要精確理解細節的任務： 可以降低底層節點的壓縮比，保留更多細節信息。
對於側重於提取概要信息的任務： 可以提高高層節點的壓縮比，側重於保留關鍵信息。

根據任務調整检索策略： 可以根據任務調整查詢感知的动态树构建和检索策略，例如：

對於需要查找特定信息的任務： 可以使用更精確的匹配算法來查找相關節點。
對於需要綜合多方面信息的任務： 可以使用更寬泛的检索策略，检索更多相關節點。
總之，通過針對特定任務調整上下文樹結構的設計，可以提高模型的性能和效率。

SharedLLM 的自我注入機制是否可以應用於其他需要處理長序列數據的領域，例如時間序列分析或生物信息學？

是的，SharedLLM 的自我注入機制可以應用於其他需要處理長序列數據的領域，例如：

時間序列分析： 在時間序列分析中，可以將長序列數據分成多个时间窗口，使用下層模型对每个窗口进行编码和压缩，构建上下文树。上層模型可以接收当前时间窗口的信息，并根据任务需求从上下文树中检索相关信息，进行预测或异常检测。
生物信息學： 在生物信息學中，可以将基因序列或蛋白质序列等长序列数据分成多个片段，使用下層模型对每个片段进行编码和压缩，构建上下文树。上層模型可以接收特定片段的信息，并根据任务需求从上下文树中检索相关信息，进行基因功能预测或蛋白质结构预测。
总的来说，SharedLLM 的自我注入机制可以应用于任何需要处理长序列数据的领域，并具有以下优势：

提高效率： 通过将长序列数据压缩成结构化表示，可以减少计算复杂度，提高模型的效率。
增强信息提取： 通过构建上下文树，可以更好地组织和提取长序列数据中的关键信息。
提高可解释性： 通过分析上下文树的结构和节点之间的关系，可以更好地理解模型的决策过程。
然而，在将 SharedLLM 应用于其他领域时，需要根据具体问题进行调整和优化，例如：

选择合适的编码器： 需要根据数据的特点选择合适的编码器，例如，对于时间序列数据，可以使用循环神经网络 (RNN) 或 Transformer 模型进行编码。
设计合适的树结构： 需要根据数据的结构和任务需求设计合适的树结构，例如，对于基因序列数据，可以使用系统发育树来表示不同基因之间的进化关系。
总而言之，SharedLLM 的自我注入机制为处理长序列数据提供了一种新的思路，具有广泛的应用前景。