延遲交互模型,為什么是下一代RAG的標(biāo)配?
張穎峰:英飛流聯(lián)合創(chuàng)始人,多年搜索、AI、Infra基礎(chǔ)設(shè)施開發(fā)經(jīng)歷,目前正致力于下一代 RAG 核心產(chǎn)品建設(shè)。
在 RAG 系統(tǒng)開發(fā)中,良好的 Reranker 模型處于必不可少的環(huán)節(jié),也總是被拿來(lái)放到各類評(píng)測(cè)當(dāng)中,這是因?yàn)橐韵蛄克阉鳛榇淼牟樵?,?huì)面臨命中率低的問(wèn)題,因此需要高級(jí)的 Reranker 模型來(lái)補(bǔ)救,這樣就構(gòu)成了以向量搜索為粗篩,以 Reranker 模型作精排的兩階段排序架構(gòu)。
目前排序模型的架構(gòu)主要有兩類:
1. 雙編碼器。以 BERT 模型為例,它針對(duì)查詢和文檔分別編碼,最后再經(jīng)過(guò)一個(gè) Pooling 層,使得輸出僅包含一個(gè)向量。在查詢時(shí)的 Ranking 階段,只需要計(jì)算兩個(gè)向量相似度即可,如下圖所示。雙編碼器既可以用于 Ranking 也可以用于 Reranking 階段,向量搜索實(shí)際上就是這種排序模型。由于雙編碼器針對(duì)查詢和文檔分別編碼,因此無(wú)法捕獲查詢和文檔的 Token 之間的復(fù)雜交互關(guān)系,在語(yǔ)義上會(huì)有很多損耗,但由于只需要向量搜索即可完成排序打分計(jì)算,因此執(zhí)行效率非常高。
2. 交叉編碼器(Cross Encoder)。Cross-Encoder 使用單編碼器模型來(lái)同時(shí)編碼查詢和文檔,它能夠捕捉查詢和文檔之間的復(fù)雜交互關(guān)系,因此能夠提供更精準(zhǔn)的搜索排序結(jié)果。Cross-Encoder 并不輸出查詢和文檔的 Token 所對(duì)應(yīng)的向量,而是再添加一個(gè)分類器直接輸出查詢和文檔的相似度得分。它的缺點(diǎn)在于,由于需要在查詢時(shí)對(duì)每個(gè)文檔和查詢共同編碼,這使得排序的速度非常慢,因此 Cross-Encoder 只能用于最終結(jié)果的重排序。例如針對(duì)初篩結(jié)果的 Top 10 做重排序,仍然需要耗時(shí)秒級(jí)才可以完成。
今年以來(lái),另一類以 ColBERT【參考文獻(xiàn)1】 為代表的工作,在 RAG 開發(fā)社區(qū)引起了廣泛關(guān)注,如下圖所示,它具備一些顯著區(qū)分于以上兩類排序模型的特點(diǎn):
其一是相比于 Cross Encoder,ColBERT 仍采用雙編碼器策略,將查詢和文檔分別采用獨(dú)立的編碼器編碼,因此查詢的 Token 和文檔的 Token 在編碼時(shí)互不影響,這種分離使得文檔編碼可以離線處理,查詢時(shí)僅針對(duì) Query 編碼,因此處理的速度大大高于 Cross Encoder;
其二是相比于雙編碼器,ColBERT 輸出的是多向量而非單向量,這是從 Transformer 的最后輸出層直接獲得的,而雙編碼器則通過(guò)一個(gè) Pooling 層把多個(gè)向量轉(zhuǎn)成一個(gè)向量輸出,因此丟失了部分語(yǔ)義。
在排序計(jì)算時(shí),ColBERT 引入了延遲交互計(jì)算相似度函數(shù),并將其命名為最大相似性(MaxSim),計(jì)算方法如下:對(duì)于每個(gè)查詢 Token 的向量都要與所有文檔 Token 對(duì)應(yīng)的向量進(jìn)行相似度計(jì)算,并跟蹤每個(gè)查詢 Token 的最大得分。查詢和文檔的總分就是這些最大余弦分?jǐn)?shù)的總和。例如對(duì)于一個(gè)有 32 個(gè) Token 向量的查詢(最大查詢長(zhǎng)度為 32)和一個(gè)有 128 個(gè) Token 的文檔,需要執(zhí)行 32*128 次相似性操作,如下圖所示。
因此相比之下, Cross Encoder 可以稱作早期交互模型 (Early Interaction Model),而以 ColBERT 為代表的工作可稱為延遲交互模型(Late Interaction Model)。
下圖從性能和排序質(zhì)量上,分別對(duì)以上排序模型進(jìn)行對(duì)比。由于延遲交互模型滿足了對(duì)排序過(guò)程中查詢和文檔之間復(fù)雜交互的捕獲,同時(shí)也避免了對(duì)文檔 Token 編碼的開銷,因此既能保證良好的排序效果,也能實(shí)現(xiàn)較快的排序性能 —— 相同數(shù)據(jù)規(guī)模下, ColBERT 的效率可達(dá) Cross Encoder 的 100 倍以上。因此延遲交互模型是一種非常有前景的排序模型,一個(gè)天然的想法是:能否在 RAG 中直接采用延遲交互模型替代向量搜索 + 精排這樣的兩階段排序架構(gòu)?
為此,我們需要考慮 ColBERT 工程化的一些問(wèn)題:
1. ColBERT 的 MaxSim 延遲交互相似度函數(shù),計(jì)算效率大大高于 Cross Encoder,但相比普通向量搜索,計(jì)算開銷仍然很大:因?yàn)椴樵兒臀臋n之間的相似度,是多向量計(jì)算,因此 MaxSim 的開銷是普通向量相似度計(jì)算的 M * N 倍 (M 為查詢的 Token 數(shù), N 為 文檔的 Token 數(shù))。針對(duì)這些,ColBERT 作者在 2021 年推出了 ColBERT v2 【參考文獻(xiàn) 2】,通過(guò) Cross Encoder 和模型蒸餾,改進(jìn)了生成的 Embedding 質(zhì)量,并且采用壓縮技術(shù),對(duì)生成的文檔向量進(jìn)行量化,從而改善 MaxSim 的計(jì)算性能。基于 ColBERT v2 包裝的項(xiàng)目 RAGatouille 【參考文獻(xiàn) 3】成為高質(zhì)量 RAG 排序的解決方案。然而,ColBERT v2 只是一個(gè)算法庫(kù),端到端的讓它在企業(yè)級(jí) RAG 系統(tǒng)使用,仍然是一件困難的事情。
2. 由于 ColBERT 是預(yù)訓(xùn)練模型,而訓(xùn)練數(shù)據(jù)來(lái)自于搜索引擎的查詢和返回結(jié)果,這些文本數(shù)據(jù)并不大,例如查詢 Token 數(shù) 32 , 文檔 Token 數(shù) 128 是典型的長(zhǎng)度限制。因此將 ColBERT 用于真實(shí)數(shù)據(jù)時(shí), 超過(guò)限制的長(zhǎng)度會(huì)被截?cái)?,這對(duì)于長(zhǎng)文檔檢索并不友好。
基于以上問(wèn)題, 開源 AI 原生數(shù)據(jù)庫(kù) Infinity 在最新版本中提供了 Tensor 數(shù)據(jù)類型,并原生地提供端到端的 ColBERT 方案。當(dāng) Tensor 作為一種數(shù)據(jù)類型,ColBERT 編碼輸出的多個(gè)向量,就可以直接用一個(gè) Tensor 來(lái)存放,因此 Tensor 之間的相似度就可以直接得出 MaxSim 打分。針對(duì) MaxSim 計(jì)算量大的問(wèn)題,Infinity 給出了 2 個(gè)方案來(lái)優(yōu)化:其一種是 binary 量化,它可以讓原始 Tensor 的空間只需原始尺寸的 1/32 , 但并不改變 MaxSim 計(jì)算的相對(duì)排序結(jié)果。這種方案主要用于 Reranker,因?yàn)樾枰鶕?jù)前一階段粗篩的結(jié)果取出對(duì)應(yīng)的 Tensor 。另一種是 Tensor Index,ColBERTv2 實(shí)際上就是 ColBERT 作者推出的 Tensor Index 實(shí)現(xiàn),Infinity 采用的則是 EMVB【參考文獻(xiàn) 4】,它可以看作是 ColBERT v2 的改進(jìn),主要通過(guò)量化和預(yù)過(guò)濾技術(shù),并在關(guān)鍵操作上引入 SIMD 指令來(lái)加速實(shí)現(xiàn)。Tensor Index 只能用來(lái)服務(wù) Ranker 而非 Reranker。此外,針對(duì)超過(guò) Token 限制的長(zhǎng)文本,Infinity 引入了 Tensor Array 類型:
一篇超過(guò) ColBERT 限制的文檔,會(huì)被切分成多個(gè)段落,分別編碼生成 Tensor 后,都跟原始文檔保存在一行。計(jì)算 MaxSim 的時(shí)候,查詢跟這些段落分別計(jì)算,然后取最大值作為整個(gè)文檔的打分。如下圖所示:
因此,采用 Infinity,可以端到端地引入延遲交互模型高質(zhì)量地服務(wù) RAG。那么,應(yīng)該是采用 ColBERT 作為 Ranker ,還是 Reranker 呢?下邊我們采用 Infinity 來(lái)在真實(shí)數(shù)據(jù)集上進(jìn)行評(píng)測(cè)。由于 Infinity 的最新版本實(shí)現(xiàn)了有史以來(lái)最全的混合搜索方案,召回手段包含向量搜索、全文搜索、稀疏向量搜索,上文所述的 Tensor ,以及這些手段的任意組合,并且提供了多種 Reranker 手段,如 RRF,以及 ColBERT Reranker 等,因此我們?cè)谠u(píng)測(cè)中包含了各種混合搜索和 Reranker 的組合。
我們采用 MLDR 數(shù)據(jù)集進(jìn)行評(píng)測(cè)。MLDR 是 MTEB 【參考文獻(xiàn) 5】用來(lái)評(píng)測(cè) Embedding 模型質(zhì)量的 benchmark 集,其中 MLDR 是其中一個(gè)數(shù)據(jù)集,全稱為 Multi Long Document Retrieval,一共包含 20 萬(wàn)長(zhǎng)文本數(shù)據(jù)。評(píng)測(cè)采用 BGE-M3【參考文獻(xiàn) 6】作為 Embedding 模型,采用 Jina-ColBERT 【參考文獻(xiàn) 7】來(lái)生成 Tensor,評(píng)測(cè)腳本也放到了 Infinity 倉(cāng)庫(kù)【參考文獻(xiàn) 8】。
評(píng)測(cè)一:ColBERT 作為 Reranker 是否有效。將 20 萬(wàn) MLDR 數(shù)據(jù)分別用 BGE-M3 生成稠密向量和稀疏向量,并插入到 Infinity 數(shù)據(jù)庫(kù)中,數(shù)據(jù)庫(kù)包含 4 列,分別保存原始文本,向量,稀疏向量,以及 Tensor,并分別構(gòu)建相應(yīng)全文索引、向量索引、稀疏向量索引。評(píng)測(cè)包含所有的召回組合,包含單路召回、雙路召回,以及三路召回,如下所示:
評(píng)測(cè)指標(biāo)采用 nDCG@10。其他參數(shù):采用 RRF Reranker 時(shí)粗篩返回的 Top N = 1000 ,查詢累計(jì)共有 800 條,平均每條查詢長(zhǎng)度在 10 個(gè) token 左右。
從圖中看到,所有的召回方案,在采用了 ColBERT Reranker 之后,都有明顯的效果提升。ColBERT 作為一種延遲交互模型,它可以提供跟在 MTEB 的 Reranker 排行榜上位居前列相提并論的排序質(zhì)量,但是性能卻是它們的 100 倍,所以可以在更大的范圍內(nèi)進(jìn)行重排序。圖中給出的結(jié)果是針對(duì) Top 100 進(jìn)行 Reranker,而采用 Top 1000 進(jìn)行 ColBERT 重排序,數(shù)值沒有明顯變化,性能還有明顯下降,因此不推薦采用。傳統(tǒng)上采用基于 Cross Encoder 的外部 Reranker ,Top 10 就會(huì)有秒級(jí)的延遲,而 Infinity 內(nèi)部實(shí)現(xiàn)了高性能的 ColBERT Reranker,即使針對(duì) Top 100 甚至 Top 1000 做重排序,也不會(huì)影響用戶體驗(yàn),而召回的范圍卻大大增加,因此可以顯著改進(jìn)最終的排序效果。此外,這種 ColBERT Reranker 計(jì)算只需在純 CPU 架構(gòu)上即可運(yùn)行,這也大大降低了部署的成本。
評(píng)測(cè)二:對(duì)比基于 ColBERT 作為 Ranker 而不是 Reranker。因此,這時(shí)需要針對(duì) Tensor 這列數(shù)據(jù)構(gòu)建 Tensor Index。同時(shí),為了評(píng)估 Tensor Index 引入的精度損耗,還進(jìn)行了暴力搜索。
可以看到,相比 Reranker ,即使是采用沒有精度損失的暴力搜索,也沒有顯著的提升,而采用基于 Tensor Index 的排序質(zhì)量甚至低于采用 Reranker。然而,作為 Ranker 的查詢時(shí)間卻要慢得多:MLDR 數(shù)據(jù)集包含 20 萬(wàn)文檔數(shù)據(jù),大約 2GB 左右,采用 Jina-ColBERT 轉(zhuǎn)成 Tensor 數(shù)據(jù)后,高達(dá) 320 G,這是因?yàn)?Tensor 數(shù)據(jù)類型是把一篇文檔的每個(gè) Token 對(duì)應(yīng)的向量都要保存下來(lái), ColBERT 模型的維度是 128 維,因此默認(rèn)數(shù)據(jù)量會(huì)膨脹 2 個(gè)數(shù)量級(jí),即使構(gòu)建了 Tensor Index,在查詢這么多數(shù)據(jù)的時(shí)候,也需要平均 7s 才能返回一個(gè)查詢,但得到的結(jié)果卻并沒有更好。
因此,很顯然,ColBERT 作為 Reranker 的收益比作為 Ranker 要高得多。當(dāng)前最佳的 RAG 檢索方案,是在 3 路混合搜索(全文搜索 + 向量 + 稀疏向量)的基礎(chǔ)上加 ColBERT Reranker。有伙伴可能會(huì)問(wèn)了,為了采用 ColBERT Reranker,就需要增加單獨(dú)的 Tensor 列,并且該列會(huì)相比原始數(shù)據(jù)集膨脹 2 個(gè)數(shù)量級(jí),這樣做是否值得?首先:Infinity 針對(duì) Tensor 提供了 Binary 量化手段,作為 Reranker,它并不影響排序結(jié)果很多,但卻可以讓最終的數(shù)據(jù)僅有原始 Tensor 大小的 1/32。其次,即便如此,也會(huì)有人認(rèn)為這樣的開銷過(guò)高。然而站在使用者的視角,用更多的存儲(chǔ),來(lái)?yè)Q取更高的排序質(zhì)量和更廉價(jià)的成本(排序過(guò)程無(wú)需 GPU),這樣做依然是非常值得的。最后,相信很快就可以推出效果上略有下降,但存儲(chǔ)開銷大大降低的 Late Interaction 模型,作為一款 Data Infra 基礎(chǔ)設(shè)施, 對(duì)這些變化保持透明,把這些 Trade Off 交給用戶是明智的選擇。
以上是基于 Infinity 在 MLDR 數(shù)據(jù)集上的多路召回評(píng)測(cè),在其他數(shù)據(jù)集的評(píng)測(cè)結(jié)果,可能會(huì)有所不同,但整體上結(jié)論不會(huì)變 —— 3 路混合搜索 + 基于 Tensor 的重排序,是當(dāng)前搜索結(jié)果質(zhì)量最高的召回手段。
由此可以看到,ColBERT 及其延遲交互模型,在 RAG 場(chǎng)景具有很大的應(yīng)用價(jià)值,以上是在文本對(duì)話內(nèi)容生成的相關(guān)工作,近期,延遲交互模型在多模態(tài)場(chǎng)景,也得到了 SOTA 的結(jié)果。這就是 ColPali【參考文獻(xiàn) 9】,它改變了 RAG 的工作流程,如下圖所示:
RAG 在面臨復(fù)雜格式文檔時(shí),當(dāng)下的 SOTA ,是采用文檔識(shí)別模型,對(duì)文檔的布局做識(shí)別,并針對(duì)識(shí)別出的部分結(jié)構(gòu),例如圖表,圖片等,再分別調(diào)用相應(yīng)的模型,將它們轉(zhuǎn)化為對(duì)應(yīng)的文字,再用各種格式保存到 RAG 配套的數(shù)據(jù)庫(kù)中。而 ColPali 則省掉了這些步驟,它直接采用多模態(tài)模型生成 Embedding 內(nèi)容。提問(wèn)的時(shí)候,可以直接針對(duì)文檔中的圖表進(jìn)行回答:
ColPali 模型的訓(xùn)練跟 ColBERT 類似,也是采用查詢 - 文檔頁(yè)面對(duì)的形式,從而捕獲查詢和文檔多模態(tài)數(shù)據(jù)之間的語(yǔ)義關(guān)聯(lián),只是采用 PaliGemma 【參考文獻(xiàn) 10】用來(lái)生成多模態(tài) Embedding 。相比沒有采用 Late Interaction 機(jī)制但同樣采用 PaliGemma 生成 Embedding 的方案 BiPali,在 nDCG@5 的評(píng)測(cè)指標(biāo)對(duì)比是 81.3 vs 58.8,這種差距是就是 “極好” 和 “壓根不能工作” 的區(qū)別。
因此,盡管 ColBERT 出現(xiàn)至今已有 4 年時(shí)間,可是 Late Interaction 模型在 RAG 的應(yīng)用才剛剛開始,它必將擴(kuò)大 RAG 的使用場(chǎng)景,在包含多模態(tài)在內(nèi)的復(fù)雜 RAG 場(chǎng)景提供高質(zhì)量的語(yǔ)義召回。而 Infinity 已經(jīng)為它的端到端應(yīng)用做好了準(zhǔn)備,歡迎關(guān)注和 Star Infinity,https://github.com/infiniflow/infinity, 致力于成為最好的 AI 原生數(shù)據(jù)庫(kù)!