熱門資訊> 正文
2025-09-08 14:20
(來源:機器之心)
經歷了前段時間的雞飛狗跳,扎克伯格的投資似乎終於初見成效。
近期,Meta Superintelligence Labs 聯合提出了一個名為 REFRAG 的高效解碼框架,旨在解決 LLM 在處理長上下文輸入時面臨的效率瓶頸,尤其是在 RAG 等應用場景下。
論文標題:REFRAG:Rethinking RAG based Decoding
論文地址:https://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 的核心任務中,模型的工作流程如下:編碼器首先處理前 s 個詞元
,其輸出的壓縮信息將輔助解碼器預測接下來的 o 個詞元
這項任務旨在訓練模型利用上下文信息進行高效預測,為其在實際應用中的表現奠定基礎。其最終目標是讓任意的編碼器和解碼器組合都能協同工作,確保解碼器基於壓縮上下文生成的內容,與它在擁有完整、未壓縮上下文時生成的內容高度相似。
持續預訓練方案
為確保 CPT 階段的成功,研究者提出了一個包含重建任務和課程學習方法的訓練方案。消融研究表明,該方案對於實現優異的 CPT 性能至關重要。
重建任務。
此任務的目標是讓編碼器學習如何以最小的信息損失壓縮文本。具體操作是,將前 s 個詞元
輸入編碼器,然后訓練模型在解碼器中重建出完全相同的詞元
。在此過程中,解碼器模型本身保持「凍結」(即參數不更新),訓練重點完全集中在編碼器和用於連接兩者的投影層上。
該任務主要實現兩個目標:
高效壓縮:訓練編碼器將 k 個詞元壓縮成一個塊嵌入,同時最大程度地保留原始信息。
空間映射:訓練投影層有效地將編碼器輸出的塊嵌入映射到解碼器的詞元空間中,使解碼器能夠「理解」並準確重建原始信息。
設計重建任務的一個特定意圖是,鼓勵模型在訓練時更多地依賴其上下文記憶(即從輸入中獲取信息),而非其固有的參數化記憶(即模型自身已經學到的知識)。一旦通過此任務初步對齊了編碼器與解碼器,便會解凍解碼器,正式開始 CPT。
課程學習。
儘管上述訓練任務在概念上清晰,但在實踐中卻極具挑戰性。其難度在於,隨着塊長度 k 的增加,可能的詞元組合數量會以
個詞元,進一步加劇了任務的複雜性。
的速率呈指數級增長(其中 V 是詞匯表的大小)。將如此巨大的多樣性有效壓縮到一個固定長度的嵌入中,是一項重大的技術挑戰。此外,從 L 個塊嵌入中重建出
與直覺相反,直接繼續預訓練解碼器以利用編碼器輸出,即使是在重建任務中,也未能降低困惑度。為解決這一優化挑戰,研究者建議對這兩項任務均採用課程學習。課程學習通過逐步增加任務難度,使模型能夠漸進且有效地掌握複雜技能。對於重建任務,訓練從重建單個塊開始:編碼器接收用於
,以此類推。爲了持續調整任務難度,研究者隨時間改變數據混合比例,從以較簡單任務(例如,單個塊嵌入)為主的樣本開始,逐步過渡到以更困難任務(即 L 個塊嵌入)為主的樣本。圖 6 提供了課程學習期間數據混合的可視化展示。
中重建
和
,解碼器則使用投影后的塊嵌入 ecnk1 來重建這 k 個詞元。隨后,模型從
的一個塊嵌入
選擇性壓縮。爲了進一步提升答案預測的準確性,該方法(REFRAG)引入了選擇性詞元壓縮機制。其核心思想是,對於上下文中特別重要的信息塊,可以不進行壓縮,而是以原始形式保留,從而避免關鍵信息丟失。
一個強化學習策略被用來決定哪些塊應當被保留。該策略以下一段落預測的困惑度作為負向獎勵信號進行指導(即困惑度越低,獎勵越高),從而學習識別並保留關鍵信息。編碼器和解碼器都經過微調,以適應這種壓縮塊與未壓縮塊混合的輸入形式。該策略網絡利用塊嵌入和掩碼技術來優化塊的擴展順序,既保留了解碼器的自迴歸特性,又實現了壓縮位置的靈活安排。
更多技術細節請參看原論文。