引言：知識密集型NLP任務中的挑戰與RaLM的潛力

在知識密集型自然語言處理（NLP）任務中，傳統的大語言模型面臨著將海量知識編碼進全參數化模型的巨大挑戰。這不僅在訓練和部署階段需要大量的努力，而且在模型需要適應新數據或不同的下遊任務時，問題更加嚴重。為了應對這些挑戰，近期的研究提出了檢索增強型語言模型（Retrieval-augmented Language Models, RaLM），它通過檢索增強將參數化的語言模型與非參數化的知識庫結合起來。

RaLM通過一次性（one-shot）或迭代（iterative）的檢索與語言模型的交互，來輔助生成過程。盡管迭代式RaLM在生成質量上表現更好，但它由於頻繁的檢索步驟而遭受高昂的開銷。因此，本文提出了一個問題：我們能否在不影響生成質量的情況下減少迭代式RaLM的開銷？

為了解決這一問題，我們提出了RaLMSpec框架，它采用推測性檢索（speculative retrieval）和批量驗證（batched verification）來減少迭代式RaLM的服務開銷，同時保證模型輸出的正確性。RaLMSpec利用了計算機體系結構中的推測執行概念，通過更高效但不那麼精確的推測性檢索步驟來替代昂貴的迭代檢索步驟。通過進一步引入預取（prefetching）、最優推測步長調度器（optimal speculation stride scheduler）和異步驗證（asynchronous verification），RaLMSpec能夠自動地充分利用加速潛力。

在對三種語言模型在四個下遊QA數據集上的廣泛評估中，RaLMSpec在保持與基線相同的模型輸出的同時，實現了顯著的速度提升。這些結果表明，RaLMSpec可以作為一個通用的加速框架，用於服務迭代式RaLM。

聲明：本期論文解讀非人類撰寫，全文由賽博馬良「AI論文解讀達人」智能體自主完成，經人工審核、配圖後發佈。

公眾號「夕小瑤科技說」後臺回復“智能體內測”獲取智能體內測邀請鏈接。

論文標題:

ACCELERATING RETRIEVAL-AUGMENTED LANGUAGE MODEL SERVING WITH SPECULATION

論文鏈接:

https://arxiv.org/pdf/2401.14021.pdf

RaLM與傳統語言模型的對比

傳統語言模型的局限性

傳統的大型語言模型（如GPT-3、LLaMA-2、PaLM）在多樣化的自然語言處理（NLP）任務中表現出色，但它們在訓練和部署時需要巨大的努力來編碼大量知識。當基礎模型需要適應新數據或不同的下遊任務時，這種情況會進一步惡化。例如，將知識編碼到完全參數化的模型中，如GPT-3，需要在訓練和部署中投入大量的努力。此外，當基礎模型需要適應新數據或多種下遊任務時，這種情況會變得更加嚴峻。

RaLM的優勢與應用

為了解決這一挑戰，最近的工作引入了檢索增強型語言模型（RaLM），它通過檢索增強將參數化語言模型與非參數化知識庫集成。RaLM在低成本適應最新數據和更好的源歸因機制方面表現出色。與傳統語言模型相比，RaLM通過結合非參數化知識庫與參數化語言模型，展現了解決知識密集型NLP任務的潛力。例如，迭代RaLM由於檢索器和語言模型之間更頻繁的交互，提供了更好的生成質量。然而，迭代RaLM通常由於頻繁的檢索步驟而遇到高開銷。

RaLMSpec框架的提出

迭代RaLM的效率瓶頸

現有的迭代RaLM方法的一個關鍵瓶頸是檢索的低效率。由於生成性語言模型的自回歸特性，檢索步驟通常是以單個查詢進行的，該查詢總結了當前的上下文。現有的迭代RaLM方法通過不斷地使用最新的上下文依賴查詢（例如q0、q1和q2）從知識庫中檢索，交錯執行檢索步驟和生成步驟。相應檢索到的內容（如A、B、C）然後通過提示或註意力級組合幫助生成過程，提供相關信息。然而，順序發出這些查詢以檢索知識庫在本質上是低效的。

RaLMSpec框架的核心思想

RaLMSpec框架采用了推測性檢索與批量驗證的方法，以減少迭代RaLM的服務開銷，同時在理論上保留模型輸出。RaLMSpec通過推測性檢索替代現有RaLM方法中昂貴、迭代的檢索步驟，采用更高效但準確度較低的推測性檢索步驟。因此，RaLMSpec使用批量驗證步驟來糾正任何不正確的推測結果，並保持模型的生成質量。更具體地說，經過一系列推測性檢索步驟後，RaLMSpec通過執行批量檢索（即圖1(b)中的➅），啟動驗證步驟，其中批次中的查詢對應於推測性檢索步驟中的查詢。如果推測的文檔與驗證步驟中檢索到的真實文檔之間存在不匹配，RaLMSpec會通過回滾到第一個錯誤推測的位置並使用真實文檔重新運行語言模型解碼來自動糾正不匹配。結果顯示，RaLMSpec通過高效的批量檢索，即使用n個查詢從知識庫檢索比順序執行n次檢索更高效，從而節省了延遲。

RaLMSpec的工作原理

推測性檢索與批量驗證

RaLMSpec框架通過推測性檢索（speculative retrieval）與批量驗證（batched verification）來減少迭代式檢索增強語言模型（iterative RaLM）的服務開銷，同時保證模型輸出的正確性。推測性檢索的概念源自計算機架構中的推測執行（speculative execution），其核心思想是用更高效但準確度較低的推測性檢索步驟替代昂貴的迭代檢索步驟。在進行了一定數量的推測性檢索後，RaLMSpec會啟動一個批量驗證步驟，其中批量查詢對應於推測性檢索步驟中的查詢。如果推測得到的文檔與驗證步驟中檢索到的真實文檔不匹配，RaLMSpec會自動回滾到第一個錯誤推測的位置，並使用真實文檔重新運行語言模型解碼。

本地緩存與預取技術

RaLMSpec利用本地緩存（local cache）來存儲每個請求的過去文檔，並在推測性檢索中使用本地緩存而非知識庫進行檢索。這種方法利用了檢索文檔的時空局部性（temporal and spatial locality），即在生成過程中可能會多次檢索到相同或連續的文檔。為了提高推測成功率，RaLMSpec在每次驗證步驟中更新本地緩存，直接添加從知識庫檢索到的相同或連續文檔。此外，RaLMSpec支持緩存預取（cache prefetching），通過在本地緩存中預先存儲來自知識庫的top-k檢索文檔，以提升推測性能。

異步驗證與最優推測步長調度器

RaLMSpec通過允許異步驗證（asynchronous verification）來提高並發性，使得額外的推測步驟可以與驗證步驟異步執行。這種技術在驗證延遲小於語言模型解碼延遲的情況下特別有益。此外，RaLMSpec引入了最優推測步長調度器（Optimal Speculation Stride Scheduler, OS3），它動態調整推測步長（speculation stride），即兩次驗證步驟之間連續推測步驟的數量，以最小化推測開銷。

實驗設置與評估方法

語言模型與數據集選擇

實驗中選用了三種標準的自然語言生成（NLG）模型類別：GPT-2、OPT和LLaMA-2。這些模型被廣泛用作RaLM的基礎語言模型，並覆蓋了不同的模型大小。實驗涉及的知識密集型開放域問答任務數據集包括Wiki-QA、Web Questions、Natural Questions和Trivia-QA。

檢索器類型與基線比較

為了展示RaLMSpec的一致性，實驗中測試了包括密集型檢索器（dense retrievers）和稀疏檢索器（sparse retrievers）在內的不同檢索器。對於密集型檢索器，進一步區分了精確（exact）和近似（approximate）方法。基線實現對於迭代式RaLM服務，直接遵循了現有的實現，其中檢索是在語言模型生成每四個令牌後觸發。對於KNN-LM服務，基線使用了Khandelwal等人（2019）的實現，其中檢索在生成每個令牌時執行。

實驗結果與分析

RaLMSpec在不同檢索器上的性能

RaLMSpec框架在不同類型的檢索器上表現出了顯著的加速效果。具體來說，當使用精確的密集檢索器（Exact Dense Retriever, EDR）時，RaLMSpec能夠實現最高2.39倍的加速比；而在使用近似的密集檢索器（Approximate Dense Retriever, ADR）時，加速比為1.04至1.39倍；對於稀疏檢索器（Sparse Retriever, SR），加速比介於1.31至1.77倍之間。這些結果表明，RaLMSpec在處理迭代式檢索增強語言模型（iterative RaLM）的服務時，能夠有效減少檢索步驟所帶來的高開銷，同時保持模型輸出的質量。

各組件對加速效果的貢獻

RaLMSpec的加速效果歸功於幾個關鍵組件的協同工作。首先，本地緩存的使用為推測性檢索提供了基礎，通過存儲過去的文檔來加速後續的檢索步驟。其次，批量驗證步驟通過並行處理多個查詢來校正任何不正確的推測結果，從而保證了模型輸出的正確性。此外，預取機制（prefetching）通過將知識庫中的前k個檢索文檔添加到本地緩存中，進一步提高了推測性能。最後，通過動態調整推測步長（speculation stride）和允許異步驗證，RaLMSpec能夠自動地發掘出最大的加速潛力。

KNN-LM服務的RaLMSpec評估

KNN-LM的挑戰與RaLMSpec的適用性

KNN-LM（K-Nearest Neighbour Language Models）是一種在生成下一個詞分佈時，會通過插值方式結合來自檢索到的k個最近鄰文檔的分佈和語言模型輸出的方法。盡管KNN-LM在提高基礎語言模型的困惑度（perplexity）方面非常有效，但它在推理過程中的開銷是極其高昂的，因為每一步生成都需要進行檢索。RaLMSpec通過修改緩存更新規則和驗證協議，能夠顯著提高KNN-LM的服務速度，即使在k值較大時（例如k=1024）也能實現高達3.88倍的加速比。

RaLMSpec在KNN-LM上的加速效果

在KNN-LM的服務中，RaLMSpec通過推測性檢索和批量驗證，顯著減少了每個令牌生成步驟所需的檢索開銷。實驗結果顯示，當使用精確的密集檢索器時，RaLMSpec能夠實現最高7.59倍的加速比；而在使用近似的密集檢索器時，加速比可達2.45倍。這些結果證明了RaLMSpec在處理檢索密集型工作負載時的有效性，並且通過啟用最優推測步長調度器（OS3），RaLMSpec能夠在不同的場景中一致地實現最佳性能。

討論與未來工作

RaLMSpec的潛在應用與擴展

RaLMSpec框架通過引入推測性檢索（speculative retrieval）和批量驗證（batched verification）的概念，為迭代式檢索增強語言模型（iterative RaLM）提供了顯著的服務加速，同時保持了模型輸出的質量。這一機制的潛在應用范圍廣泛，不僅限於當前的問答（QA）任務，也可能擴展到其他知識密集型自然語言處理（NLP）任務中，例如機器翻譯、文本摘要或對話系統。

在未來的工作中，RaLMSpec可以進一步與新興的大語言模型（如LLaMA-2, GPT-3, PaLM）集成，以提高這些模型在實際應用中的效率。此外，RaLMSpec的推測性檢索和批量驗證機制也可以與其他類型的檢索器（如TF-IDF或BM25）結合，以探索在不同檢索精度和效率權衡下的性能表現。

挑戰與未來研究方向

盡管RaLMSpec在多個數據集和語言模型上展示了其加速能力，但在實際部署時仍面臨一些挑戰。例如，推測性檢索的成功率高度依賴於檢索的時空局部性（temporal/spatial locality），這可能在某些任務或數據集中不那麼顯著。此外，批量驗證步驟的效率取決於並行處理能力，這可能受限於硬件資源。

未來的研究可以探索如何優化RaLMSpec的各個組件，例如改進本地緩存策略、調整預取（prefetching）大小或進一步優化異步驗證（asynchronous verification）機制。此外，研究者也可以探索自適應調整推測步長（speculation stride）的算法，以便在不同的運行時環境中動態平衡推測性能和驗證開銷。

總結：RaLMSpec在提升RaLM服務效率中的貢獻

RaLMSpec通過引入推測性檢索和批量驗證，顯著提高了迭代式檢索增強語言模型的服務效率。實驗結果表明，RaLMSpec能夠在保持模型輸出質量的同時，實現對不同檢索器（包括精確密集檢索器、近似密集檢索器和稀疏檢索器）的顯著加速。特別是在使用精確密集檢索器時，RaLMSpec PSA（結合預取、最優推測步長調度器和異步驗證）能夠與基線相比，在不同的語言模型和數據集上實現高達2.39倍的加速比。

此外，RaLMSpec的三個附加技術——預取、最優推測步長調度器（OS3）和異步驗證——進一步降低了RaLM服務的延遲。這些技術的組合使用，使得RaLMSpec能夠在各種場景中自動地、一致地實現最佳的加速比，證明了其作為迭代式RaLM服務的通用加速框架的潛力。