機器之心報道
編輯:+0,冷貓
經曆了前段時間的雞飛狗跳,紮克伯格的投資似乎終於初見成效。
近期,Meta Superintelligence Labs 聯合提出了一個名為 REFRAG 的高效解碼框架,旨在解決 LLM 在處理長上下文輸入時麵臨的效率瓶頸,尤其是在 RAG 等應用場景下。
論文標題:REFRAG:Rethinking RAG based Decoding論文地址:http://arxiv.org/abs/2509.01092
為什麽長上下文處理如此困難?
在當前的 AI 應用中,利用 LLM 處理包含大量外部知識的長文本輸入,是提升問答、對話和智能體應用能力的關鍵。然而,這一過程也帶來了嚴峻的挑戰:在傳統 LLM 中,注意力機製的計算和內存開銷會隨著輸入長度的平方(N²)增長。
這意味著文本長度翻一倍,速度可能會慢 4 倍,這會導致顯著的係統延遲,並消耗大量內存用於存儲 KV Cache,進而降低係統吞吐量。這使得開發者不得不在知識豐富度與係統效率之間做出痛苦的權衡。
Meta 的研究指出,在 RAG 應用中,LLM 處理的上下文中包含了大量從外部知識庫檢索拚接而成的段落,但其中隻有一小部分與用戶查詢緊密相關。這些不相關的段落導致了計算資源的浪費。REFRAG 的核心思想正是基於這一觀察,通過識別並跳過對這些非相關上下文的無效計算,來優化解碼過程。
REFRAG 是如何解決問題的?
REFRAG 框架通過一個精巧的四步流程,利用注意力稀疏結構,實現了顯著的性能提升。它與傳統 RAG 的關鍵差異在於,它避免了讓 LLM 直接處理冗長的原始文本。
壓縮:首先,一個輕量級的編碼器會讀取檢索到的文檔,將每 16 個 token 壓縮成一個濃縮了語義精華的「塊向量」。縮短:接下來,主模型不再讀取原始的 token,而是直接處理這些塊向量。輸入序列的長度因此立刻縮短了 16 倍。加速:由於輸入變得極短,注意力機製的計算開銷大幅降低,同時作為顯存消耗大頭的 KV cache 也變得更小。這正是其能實現驚人速度提升的根本原因。選擇:為了防止在壓縮過程中丟失關鍵信息,框架引入了一個基於 RL 的策略充當「質檢員」,它能智能地挑出信息密度最高、與任務最相關的關鍵片段,確保它們不被壓縮,從而保留核心信息。
Meta 表示,該框架的有效性已在包括 RAG、多輪對話和長文檔摘要在內的多種長上下文任務中得到驗證,取得了突破性的成果:
速度提升: 將首個 token 生成時間(TTFT)加速高達 30.8 倍。在 16k tokens 的場景下,相比 CEPE 等基線方法,實現了超過 16 倍的 TTFT 加速。從性能圖表可以看出,文本越長,REFRAG 的優勢越明顯,其加速效果隨上下文規模增加呈指數級提升,而基線方法僅為線性增長。
上下文擴展: 能夠將現有 LLM 的有效上下文大小擴展 16 倍,使其能處理更海量的信息。精度不降反升: 在大幅提升速度和擴展上下文的同時,做到了模型的準確率沒有損失。更關鍵的是,在 GSM8K 基準測試上,REFRAG 不僅能處理 8 倍更長的上下文(80 個 chunk vs 10 個 chunk),運行速度還提升了一倍,最終成績更是幾乎翻倍,從 6.71 提升到 12.08。
簡而言之,REFRAG 讓「大上下文 RAG」從理想變成了現實。
雖然其效果聽起來非常不錯,但評論區也表示,它最終的價值仍需要在更廣泛的實際應用場景中進行檢驗。
還有人對該研究中的 RL 策略提出了質疑。
方法
為實現編碼器與解碼器的有效對齊,本研究遵循 Yen et al. (2024) 的工作,采用了一種基於「下一段落預測」任務的持續預訓練方法。
在訓練中,每個數據點包含總計 s+o=T 個詞元(token)。通過這一預訓練過程,模型能夠學習如何利用塊嵌入(chunk embeddings)來高效執行下遊任務。
為了進一步提升模型性能,該方法還引入了通過 RL 實現的選擇性壓縮機製。在完成 CPT 對齊後,模型會經過監督微調 ,以適應具體的下遊應用場景,例如 RAG 和多輪對話。
持續預訓練方案
為確保 CPT 階段的成功,研究者提出了一個包含重建任務課程學習方法的訓練方案。消融研究表明,該方案對於實現優異的 CPT 性能至關重要。
該任務主要實現兩個目標:
高效壓縮:訓練編碼器將 k 個詞元壓縮成一個塊嵌入,同時最大程度地保留原始信息。空間映射:訓練投影層有效地將編碼器輸出的塊嵌入映射到解碼器的詞元空間中,使解碼器能夠「理解」並準確重建原始信息。
設計重建任務的一個特定意圖是,鼓勵模型在訓練時更多地依賴其上下文記憶(即從輸入中獲取信息),而非其固有的參數化記憶(即模型自身已經學到的知識)。一旦通過此任務初步對齊了編碼器與解碼器,便會解凍解碼器,正式開始 CPT。
選擇性壓縮。為了進一步提升答案預測的準確性,該方法(REFRAG)引入了選擇性詞元壓縮機製。其核心思想是,對於上下文中特別重要的信息塊,可以不進行壓縮,而是以原始形式保留,從而避免關鍵信息丟失。
一個強化學習策略被用來決定哪些塊應當被保留。該策略以下一段落預測的困惑度作為負向獎勵信號進行指導(即困惑度越低,獎勵越高),從而學習識別並保留關鍵信息。編碼器和解碼器都經過微調,以適應這種壓縮塊與未壓縮塊混合的輸入形式。該策略網絡利用塊嵌入和掩碼技術來優化塊的擴展順序,既保留了解碼器的自回歸特性,又實現了壓縮位置的靈活安排。
更多技術細節請參看原論文。</p>