今天小編分享的科學經驗：陳丹琦團隊新作：Llama-2上下文擴展至128k，10倍吞吐量僅需1/6内存，歡迎閲讀。
陳丹琦團隊剛剛發布了一種新的 LLM上下文視窗擴展方法：
它僅用 8k 大小的 token 文檔進行訓練，就能将 Llama-2 視窗擴展至 128k。
最重要的是，在這個過程中，只需要原來1/6 的内存，模型就獲得了10 倍吞吐量。
除此之外，它還能大大降低訓練成本：
用該方法對 7B 大小的羊駝 2 進行改造，只需要一塊 A100就能搞定。
團隊表示：
希望這個方法有用、好用，為未來的 LLM 們提供廉價又有效的長上下文能力。
目前，模型和代碼都已在 HuggingFace 和 GitHub 上發布。
只需添加兩個組件
這個方法名叫CEPE，全稱 "并行編碼上下文擴展（Context Expansion with Parallel Encoding）"。
作為輕量級框架，它可用于擴展任何預訓練和指令微調模型的上下文視窗。
對于任何預訓練的僅解碼器語言模型，CEPE 通過添加兩個小組件來實現擴展：
一個是小型編碼器，用于對長上下文進行塊編碼；
一個是交叉注意力模塊，插入到解碼器的每一層，用于關注編碼器表示。
完整架構如下：
在這個示意圖中，編碼器模型并行編碼上下文的 3 個額外塊，并與最終隐藏表示進行連接，然後作為解碼器交叉注意力層的輸入。
在此，交叉注意力層主要關注解碼器模型中自注意力層和前饋層之間的編碼器表示。
通過仔細選擇無需标記的訓練數據，CEPE 就幫助模型具備了長上下文能力，并且也擅長文檔檢索。
作者介紹，這樣的 CEPE 主要包含 3 大優勢：
（1）長度可泛化
因為它不受位置編碼的約束，相反，它的上下文是分段編碼的，每一段都有自己的位置編碼。
（2）效率高
使用小型編碼器和并行編碼來處理上下文可以降低計算成本。
同時，由于交叉注意力僅關注編碼器最後一層的表示，而僅使用解碼器的語言模型則需要緩存每個層每個 token 的鍵 - 值對，所以對比起來，CEPE 需要的内存大大減少。
（3）降低訓練成本
與完全微調方法不同，CEPE 只調整編碼器和交叉注意力，同時保持大型解碼器模型凍結。
作者介紹，通過将 7B 解碼器擴充為具有 400M 編碼器和交叉注意力層的模型（總計 14 億參數），用一塊 80GB 的 A100 GPU 就可以完成。
困惑度持續降低
團隊将 CEPE 應用于 Llama-2，并在 200 億 token 的 RedPajama 過濾版本上進行訓練（僅為 Llama-2 預訓練預算的 1%）。
首先，與 LLAMA2-32K 和 YARN-64K 這兩種完全微調的模型相比，CEPE 在所有數據集上都實現了更低或相當的困惑度，同時具有更低的内存使用率和更高的吞吐量。
在将上下文提升到 128k 時（遠超其 8k 訓練長度），CEPE 的困惑度更是持續保持降低，同時保持低内存狀态。
相比之下，Llama-2-32K 和 YARN-64K 不僅不能推廣到其訓練長度之外，還伴随着内存成本顯著增加。
其次，檢索能力增強。
如下表所示：
通過使用檢索到的上下文，CEPE 可以有效改善模型困惑度，性能優于 RePlug。
值得注意的是，即使讓段落 k=50 （訓練是 60），CEPE 仍會繼續改善困惑度。
這表明 CEPE 可以很好地轉移到檢索增強設定，而全上下文解碼器模型在這個能力上卻退化了。
第三，開放網域問答能力顯著超越。
如下圖所示，CEPE 在所有數據集和段落 k 參數上都大幅優于其他模型，且不像别的模型那樣，k 值越來越大之後性能明顯下降。
這也表明，CEPE 對大量冗餘或不相關的段落并不敏感。
所以總結一下就是，與大多數其他解決辦法相比，CEPE 在上述所有任務上都能以低得多的内存和計算成本勝出。
最後，作者在這些基礎上，提出了專門用于指令調優模型的 CEPE-Distilled ( CEPED ) 。
它僅使用未标記的數據來擴展模型的上下文視窗，通過輔助 KL 散度損失将原始指令調整模型的行為提煉為新架構，由此無需管理昂貴的長上下文指令跟蹤數據。
最終，CEPED 可以在保留指令理解能力的同時，擴展 Llama-2 的上下文視窗，提高模型長文本性能。
團隊介紹
CEPE 一共 3 位作者。
一作為顏和光（Howard Yen），普林斯頓大學計算機科學專業碩士生在讀。
二作為高天宇，同校博士生在讀，清華本科畢業。
他們都是通訊作者陳丹琦的學生。
論文原文：
https://arxiv.org/abs/2402.16617
參考鏈接：
https://twitter.com/HowardYen1/status/1762474556101661158
