核心概念
本文提出了一種名為 SharedLLM 的新型架構,旨在解決大型語言模型上下文窗口有限的問題,透過多粒度上下文壓縮和查詢感知信息檢索,在效率和性能之間取得平衡。
論文資訊
標題:雙劍合璧:透過多粒度自我注入擴展上下文窗口 (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,這些技術可能會進一步提高性能。
统计
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。