在工業(yè)場(chǎng)景中，往往會(huì)利用檢索技術(shù)來為大語言模型添加一些來自外部數(shù)據(jù)庫的知識(shí)文檔，從而增強(qiáng)大語言模型的回復(fù)可信度。一般來說，RAG 被公認(rèn)是最有效的為 LLM 注入特定領(lǐng)域知識(shí)的方式。

然而，RAG 也有其不足之處。通常來說，在實(shí)際應(yīng)用中，為確保能召回包含正確知識(shí)的文檔，對(duì)于每個(gè)用戶的查詢，會(huì)檢索多個(gè)文檔（一般在 5 到 30 個(gè)之間），并把這些文檔整合到輸入提示中供大語言模型處理。這樣一來，輸入提示的序列長度增加，使得推理效率大幅降低。具體來講，以首次生成標(biāo)記的時(shí)間（TTFT）來衡量，RAG 大語言模型的推理延遲比非 RAG 大語言模型高很多。

由于數(shù)據(jù)庫中同一文檔經(jīng)常會(huì)被不同 query 召回，大家很自然的會(huì)想到：是否能夠把已經(jīng)算好的文檔表示（KV states）存在緩存中，以供二次使用？很遺憾， 由于自回歸注意力機(jī)制的限制，大語言模型中每個(gè)文檔的 KV States 都與上下文相關(guān)，所以遇到新的 query 時(shí)，模型必須重新編碼 KV states 才能確保準(zhǔn)確預(yù)測(cè)。

最近，論文《Block-Attention for Efficient RAG》為檢索增強(qiáng) (RAG) 場(chǎng)景實(shí)現(xiàn)了一種塊狀注意力機(jī)制，Block-Attention，通過分塊獨(dú)立編碼檢索到的文檔，使得模型無需重復(fù)編碼計(jì)算已經(jīng)在其他 query 中已經(jīng)見過的文檔，從而實(shí)現(xiàn)線上推理效率的有效提升。在實(shí)驗(yàn)中，該方法能夠讓使用 RAG 技術(shù)的模型與不使用 RAG 的模型有幾乎一樣的響應(yīng)速度。同時(shí)，該方法甚至還能略微提升在 RAG 場(chǎng)景下的模型準(zhǔn)確率。

• 論文標(biāo)題：Block-Attention for Efficient RAG
• 論文地址：https://arxiv.org/pdf/2409.15355

如下圖所示，該方法把整個(gè)輸入序列分成若干個(gè) block，每個(gè) block 獨(dú)立計(jì)算其 KV States，只有最后一個(gè) block 能夠關(guān)注其他 blocks（在 RAG 場(chǎng)景中，最后一個(gè) block 即用戶的輸入）。在 RAG 場(chǎng)景中，block-attention 讓模型不再需要重復(fù)計(jì)算已經(jīng)在其他 query 中見過的文檔。

Block-Attention 的實(shí)現(xiàn)并不復(fù)雜：1）獨(dú)立編碼除最后一個(gè) block 以外的所有 blocks；2）為每個(gè) blocks 重新計(jì)算位置編碼；3）將所有 blocks 拼接在一起，并計(jì)算最后一個(gè) block 的 KV State。然而直接把模型不加任何修改的從 self-attention 切換到 block-attention 會(huì)導(dǎo)致大語言模型懵圈，畢竟模型在訓(xùn)練階段從來沒見過 block-attention 方式編碼的輸入。一個(gè)量化的對(duì)比是，直接切換為 block-attention 會(huì)讓 Llama3-8B 在四個(gè) RAG 數(shù)據(jù)集上的平均準(zhǔn)確率由 67.9% 下降至 48.0%。

為了讓模型適應(yīng) block-attention，作者們對(duì)模型進(jìn)行了進(jìn)一步微調(diào)，作者們發(fā)現(xiàn)在 100-1000 步微調(diào)之后，模型就能快速適應(yīng) block-attention，在四個(gè) RAG 數(shù)據(jù)集上的平均準(zhǔn)確率恢復(fù)至 68.4%。另外，block-attention 方式的模型在 KV cache 技術(shù)的幫助下，能達(dá)到與無 RAG 模型相似的效率。在用戶輸入長度為 50 而 prompt 總長度為 32K 的極端情況下，block-attention model 的首字延時(shí)（Time To First Token, TTFT）和首字浮點(diǎn)運(yùn)算數(shù)（FLOPs To Frist Token, (FLOPs-TFT）分別能降低至 self-attention model 的 1.3% 和 0.2%，與無 RAG 模型的效率基本持平。

推理流程

關(guān)于 block-attention 的實(shí)現(xiàn)和詳細(xì)推導(dǎo)，讀者們請(qǐng)移步原文，這里主要介紹 block-attention 模型的推理流程。如下圖所示，首先從緩存中查詢并提取前 K 個(gè) block 的 KV states。然后，根據(jù)每個(gè) block 在輸入序列中的位置，作者們對(duì)每個(gè) block 的位置編碼進(jìn)行了重新計(jì)算。具體的操作過程詳見論文的公式 3。最后，根據(jù)前 k-1 個(gè) KV States 計(jì)算最后一個(gè)數(shù)據(jù)塊的鍵值狀態(tài)以及模型的輸出。

實(shí)驗(yàn)結(jié)果

在實(shí)驗(yàn)中，作者們主要想探究兩個(gè)問題的答案：1）在 RAG 場(chǎng)景中，block-attention 模型能否達(dá)到與自 self-attention 相同的準(zhǔn)確率？2）block-attention 對(duì)效率的提升有多大？

對(duì)于問題一，上圖給出了答案。作者們根據(jù)實(shí)驗(yàn)結(jié)果給出了三個(gè)結(jié)論：

1. 直接從 self-attention 切換到 block-attention 是不可取的，因?yàn)檫@會(huì)導(dǎo)致準(zhǔn)確率急劇下降。例如，對(duì)于 Llama3-8B 和 Mistral-7B 模型，去除微調(diào)過程會(huì)導(dǎo)致在所有四個(gè)基準(zhǔn)上平均絕對(duì)性能下降 21.99%。

2. 然而，如果作者們?cè)谖⒄{(diào)階段使用塊注意力機(jī)制，那么得到的模型與自注意力模型的性能幾乎相同，甚至在某些數(shù)據(jù)集上略好。例如，Mistral-7B-block-ft 在四個(gè)基準(zhǔn)上的性能優(yōu)于自回歸方式訓(xùn)練的模型，平均準(zhǔn)確率由 59.6% 上升至 62.3%。

3. 位置重新編碼操作對(duì)于 block-attention 模型至關(guān)重要。去除它會(huì)導(dǎo)致性能顯著下降 —— 在四個(gè)數(shù)據(jù)集上準(zhǔn)確率平均下降 4%。

對(duì)于效率的提升，作者們也通過另一組實(shí)驗(yàn)進(jìn)行了驗(yàn)證。他們將用戶的問題長度固定在 50 個(gè) token，然后逐漸增加被召回文檔的數(shù)量，讓輸入序列總長度從 50 一直增加到 32K。模型在不同 prompt 長度下的首字延時(shí)（Time To First Token, TTFT）和首字浮點(diǎn)運(yùn)算數(shù)（FLOPs To Frist Token, (FLOPs-TFT）如下圖所示。顯然，加速效果令人滿意：當(dāng)輸入序列的長度為 512 時(shí)，使用 block-attention 可以將 TTFT 減少 48%，將 FLOPs-TFT 減少 90.1%。隨著總長度的增加，block-attention 模型的 TTFT 和 FLOPs-TTF 保持基本不變的趨勢(shì)。當(dāng)總長度達(dá)到 32K 時(shí)，加速效果可以達(dá)到驚人的 98.7%，F(xiàn)LOPs-TFT 的消耗甚至減少了 99.8%。作者們將此實(shí)驗(yàn)結(jié)果總結(jié)為：“文本越長，block-attention 越重要”。

作者們最后還指出，block-attention 在很多場(chǎng)景中都有著重要作用，并不局限于 RAG。由于一些保密原因，作者們暫時(shí)無法透露在其他工業(yè)應(yīng)用中是如何使用它的。作者們期待社區(qū)的研究人員能夠進(jìn)一步探索 block-attention 的潛力，并將其應(yīng)用于合適的場(chǎng)景。