圖片搜索場景，指的是輸入文本去檢索相關的圖片，是一種典型的跨模態檢索場景。在圖片搜索中，強烈依賴于圖片的視覺信息，多模態能力是必要的，典型的第一階段檢索方式（也即是召回階段）采用的是稠密檢索，如Fig 1所示，目前主流的管道是：
- 以圖文匹配為目標，采用基于CLIP等大規模對比學習的方式訓練圖片和文本塔
- 對索引庫中的所有待索引圖片進行圖片稠密特征刷取，在線上部署文本塔以提取query的稠密特征
- 將文本塔和圖片塔的稠密向量進行度量計算（如余弦相似度），得到圖文之間的相似度
當然，在真實的業務場景中，還會采用ANN（近似最近鄰）等技術對大規模的圖文相似度計算進行速度和效果的折中。我們能從以上的描述中，感受到稠密檢索的幾個缺陷：
- 儲存代價高：以一個512維度的特征向量，如果是
float32形式儲存，那么10億左右的圖片量級就需要接近2T的儲存。 - 檢索速度慢：對2個維度為d的稠密向量進行相似度計算需要d次浮點乘法運算和加法運算。
- 可解釋性差：采用稠密檢索的方式，無法提供足夠好的可解釋性，去解釋為何一對圖文之間為何打分高/低，這對于業務場景中需要定期解case的需求而言，是一個不利因素。
這讓我們想到為何不采用稀疏的檢索方式呢，比如在傳統的文本檢索任務中，會采用基于BM25的方式給文檔和query進行詞頻統計，然后基于詞的字面精準匹配去計算相似度打分，這種方式儲存代價低、檢索速度快且具有更好的可解釋性。 然而，圖片搜索中的圖片無法通過傳統的方法將其進行“詞頻統計”，而LexLIP（Lexicon-Bottlenecked Language Image Pre-Training， 基于詞典瓶頸的圖文預訓練） [1] 這篇文章就是考慮將圖片也像文本一樣進行“詞頻統計”，從而使得圖片搜索也可以采用稀疏檢索，如Fig 2所示。 讓我們進一步深究下這篇文章是怎么做的。
Fig 1. 稠密檢索方式 vs 稀疏檢索方式。
首先我們反過來，先看到當LexLIP整個模型訓練完后，是如何使用的。 如Fig 2所示，訓練好的視覺編碼器會對索引庫中的所有圖片進行特征計算（俗稱刷庫），然后每一個圖片中的關鍵實體元素，將會被分解到文本詞表中，并且詞會有對應的詞權（weight），筆者將其稱之為視覺詞（visual word）。待刷完庫后，可以構建倒排索引鏈表，此時可以知道詞表中的每個詞都在哪些圖片中出現過，并且詞權是多少。用戶在使用時候，對于輸入的query文本也進行詞權計算，用Lq表示query中出現的詞，然后去查詢構建好的圖片倒排索引庫Ls，那么l∈Lq∩Ls就是查詢到詞的倒排列表的歸并拉鏈 (歸并倒排拉鏈的介紹可見 https://blog.csdn.net/LoseInVain/article/details/116377189 中信息召回一節Fig 4.1。)，用Wx(l),x=q,s表示其詞權，那么可知基于詞精準匹配的相似度可由公式(1)計算。
(1)score=∑l∈Lq∩LsWq(l)×Ws(l)
可以發現，這個過程中，由于不需要儲存稠密特征向量，只需要儲存構建好的倒排列表字典，儲存代價是非常低的。在計算相似度這塊，由于只需要對少許的詞進行詞權乘法，計算代價也是很小的。在可解釋性這塊，由于通過文本去表達圖片的關鍵信息，因此通過觀察圖片對應的視覺詞以及詞權的大小，能夠解釋為何一個query能和該圖片構成/不構成相關關系（因為此時無論query還是圖片都是用文本表示了）。因此，LexLIP這個工作通過將圖片表示為視覺詞和詞權，實現了圖片搜索中的多模態稀疏檢索。
Fig 2. LexLIP中的檢索過程。
這個就是LexLIP的檢索過程，那么這個模型是怎么訓練的呢？ 從直觀上看，這個模型應該具有幾大塊：
- 文本編碼器和詞權計算器（lexicon-bottleneck）： 用于對文本側進行稠密編碼，然后通過詞權計算器對文本稠密特征處理，得到詞權。
- 視覺編碼器和視覺詞權計算器：用于對圖片進行稠密編碼，然后通過視覺詞權計算器對視覺稠密特征處理，得到視覺詞權。
因此，這個模型的訓練需要考慮文本、圖片雙塔的編碼器如何訓練，也需要考慮各自的詞權計算模塊怎么訓練。如Fig 3.所示，讓我們看下這個模型訓練是怎么做的。在圖片側的編碼器采用的是ViT，需要對圖片進行切塊，切塊后結果記為x=[x1,?,xm]，那么經過ViT處理后的圖片稠密向量可表示為公式(2-a)，然后用一個LM頭將其映射到詞表空間，詞表大小為|V|，如公式(2-b)所示。為了避免一些詞權特別大將其他詞權的作用掩蓋了，在公式(2-c)采用log?(1+x)的方式進行了對數尺度放縮，MaxPool(⋅)則是沿著序列軸方向的最大池化，這個操作也容易理解，就是將整個文本/圖片序列中最為“重要”的信息給提取出來了。
(2)Hv=Transv([CLSv;x])∈R(m+1)×d(a)Sx(enc)v=LM_headv(Hv)∈R(m+1)×|V|(b)pv=log?(1+MaxPool(max(Sx(enc)v,0)))∈R|V|(c)
相似的，對于文本側y的處理，也是和圖片側一樣的，如公式(3)所示。
(3)Hl=Transl([CLSl;y])∈R(m+1)×d(a)Sx(enc)l=LM_headl(Hl)∈R(m+1)×|V|(b)pl=log?(1+MaxPool(max(Sx(enc)l,0)))∈R|V|(c)
在訓練階段，有三種類型的損失函數：
- Text MLM： 文本的MLM預訓練任務，即是將輸入文本y進行一定概率的掩膜(即是以一定概率p將文本tokens用特殊token
[MASK]替代，或者是詞表V中的隨機一個token。) 得到y¯ ，然后通過上下文去預測被掩膜的token，如公式(4)所示。這個目標主要是建模文本編碼器本身的能力。(4)Lmlm=−∑D∑j∈M(enc)log?(P(wj=yj|y¯))(a)P(wj)=softmax(Sy¯(enc)l[i,:])(b) - LexMLM： 這個損失用于建模詞權計算的能力，也即是原文中提到的"lexicon-bottleneck"，其中的lexicon表示詞典，bottleneck表示瓶頸，也即是通過這個模塊，能將稠密向量的關鍵信息“封鎖”到稀疏的詞典表示中，就像是瓶頸一樣，算是很形象了。這塊的建模比較復雜， 首先需要進行規范化，如公式(5)所示，注意到文本側的輸入是掩膜后的文本y¯。(5)av=Norm(MaxPool(Sx(enc)v))∈R[0,1]|V|(a)al=Norm(MaxPool(Sy¯(enc)l))∈R[0,1]|V|(b)
由于這是一個跨模態模型，要求詞權也具有跨模態建模的能力，在傳統的多模態融合模型UNITER（見博文 [2] ）中，是采用跨模態的MLM去建模不同模態之間的語義關聯，在本工作中也是采取了類似的方法，考慮到LexLIP的底座是一個雙塔模型，因此也引入了一個弱化的Transformer（被稱之為Weakened Language Decoder）作為跨模態關聯。
在模型學習好了的情況下，現在的a(⋅)應該蘊含有充分的圖片或者文本的信息，足以用于推測出被掩膜掉的文本token。 考慮到文本側輸入的文本y¯已經是進行過掩膜的輸入（掩膜概率30%）， 也就是al本身的信息已經是帶噪的， 因此需要構建出帶噪程度更大的文本輸入，因此作者構建了掩膜程度更高（比如50%）的文本y^，打算以某種形式去建模，讓模型通過av或者al的信息就能去預估出y^中被掩膜的token，去用于預估掩膜token的模型是弱化Transformer（只有2層），而輸入應該是a(⋅)和y¯，這個是大概的思路。
不過我們需要考慮一下細節，a(⋅)的是一個維度為|V|的語義向量，而詞表大小通常是很大的（比如32000），直接將這個那么長的向量喂到弱化Transformer中不是一個好的做法，因此作者考慮采用了連續詞袋（Continuous Bag-of-Words， CBoW）的做法， 將長度為|V|，蔓延整個詞表大小的向量轉化為了維度固定為d的向量，如Fig 6所示，此處的sg(⋅)表示梯度停止，而W(te)∈R|V|×d是文本編碼器的token embedding矩陣。
(6)b(⋅)=a(⋅)sg(W(te))∈Rd
爾后將b(⋅)和y^拼接起來，送到弱化Transformer中進行解碼即可，如公式(7)所示
(7)Sy^(dec)v=Decoderv([bv;y^])∈R(n+1)×|V|(a)Sy^(dec)l=Decoderl([bl;y^])∈R(n+1)×|V|(b)
損失即是和公式(4)類似的MLM loss，可表示為Li2t和Lt2i，此處就不贅述了。
- BaCO（In-Batch lexicon-Contrastive Loss）：采用對比損失去保證進行過視覺稠密特征稀疏化后的pv同樣具有良好的語義匹配能力，如公式(8)所示，此處的F(⋅)是一個控制稀疏度的正則項，在SPLADE [3] 這篇文章中引入的，筆者感覺不是很關鍵，暫時先不關注，那么Lbaco=(Lbacoi2t+Lbacot2i)/2。(8)Lbacoi2t=−∑Dlog?(exp?(pv(pl)T)/τ∑j∈Bexp?(pv(pjl)T)/τ)+λF(pv)(a)Lbacot2i=−∑Dlog?(exp?(pv(pl)T)/τ∑j∈Bexp?(pjv(pl)T)/τ)+λF(pl)(a)
就筆者感知而言，其實不妨將LexMLM看成是ALBEF [4] 中的MLM loss，而BaCO則是ALBEF中的ITC loss，只不過為了適配稀疏特征的特點做了一些優化而已，這樣理解這個論文就會容易一些。最終的損失為：
(9)Lp1=Lmlm+Li2t+Lt2i+Lbaco
作者在文章中為了提高性能，在上面的基礎上還采用了動量學習 + 大規模負樣本隊列的方式（即是MoCo），去進行第二階段的預訓練，這個筆者在博文 [2] 中也曾經介紹過，在此不再累述。
Fig 3. LexLIP的預訓練階段框架，主要分為Text MLM、lexicon bottleneck MLM（LexMLM）、In-batch Contrastive loss這3大塊損失建模。
以上模型就訓練好了，然而為了實現梯度反傳，以上提到的相似度計算都是在實數域進行的，比如pv(pl)T， 此時詞的概率分布彌散在整個詞表中。這和我們在公式(1)提到的整數型的詞權計算不同，使得無法在目前開源的的基于詞的檢索系統（如Anserini）中使用。因此作者還進行了詞權的“量化”， 將實數型的詞權轉化為了離散的整數詞權。作者采用的是直接量化，即是 ⌊100×p(⋅)⌋，此時無論是圖片還是文本都可以轉化為若干個詞以及詞權，然后采用公式(1)高效地進行計算相關性打分了。
作者進行了一些實驗去對比LexLIP的模型性能和速度，如Fig 4 (a)所示，作者對比了一些圖文檢索的SOTA模型，能發現LexLIP在公平比較下，大部分時候都能達到SOTA的性能，具體的實驗細節就不累述了，讀者有興趣可以自行翻閱論文。 同時，作者還進行了大規模檢索實驗，去探索LexLIP的速度和模型性能的均衡表現，如Fig 4. (b)所示，其中的BM25，指的是用圖片的caption信息去表示圖片，使得可以采用BM25方法去進行圖文檢索，這個可視為是稀疏檢索的基線，稠密檢索的基線則是采用CLIP。從結果上看，LexLIP所需的索引儲存量（最多152M）明顯小于稠密索引（2G），甚至比BM25的索引儲存量還少。在保持同樣效果的前提下，QPS是稠密檢索的270倍，在最佳模型效果下，QPS是稠密檢索的5.4倍，并且有著模型效果上的絕對優勢。和純文本的稀疏檢索基線對比，在保持相似QPS的情況下，LexLIP的模型效果也是有著明顯優勢（22.3 > 16.8）。這說明LexLIP在模型效果，索引儲存量和檢索速度上都有著明顯優勢，是一個適合業務落地的方法。
Fig 4. LexLIP的SOTA實驗對比和大規模檢索實驗。
作者還進行了詞權的可視化，如Fig 5所示，其中詞權越大則可視化出來的詞尺寸也越大，可以發現LexLIP這種方法還能提供比較好的可解釋性，這對于線上應用來說也是很友好的。作者還對本文提到的各種損失函數進行了消融試驗，結論就是各種損失函數都是有效的，不過本文就不繼續累述了，讀者感興趣的請自行翻閱。
Fig 5. 對圖片的視覺詞和文本的詞，根據詞權大小進行可視化，尺寸越大的詞表示詞權越大。
筆者之前在 《萬字淺析視頻搜索系統中的多模態能力建設》 [5] 中曾經討論了在真實的視頻搜索業務場景中，多模態稀疏特征的重要作用。我在看完這篇論文后，感覺這個工作還是很適合在真實的圖片搜索場景落地的（甚至是視頻搜索場景也可以），可以作為一個獨特的多模態稀疏召回通路而存在，其特征甚至排序階段也有作用。不過這個倒底還是不能完全取代稠密檢索的作用，筆者在博文 [6] 中討論過基于CLIP特征的視覺短板問題，在一些需要視覺本身的結構信息，而不是語義信息的情況下，基于CLIP方式訓練出來的特征并不夠完備，因此需要引入視覺自監督模型的特征，如DINO等。 本文還是在建模跨模態的語義信息，只是將其做成了視覺的離散特征，因此筆者覺得，可以考慮再引入對自監督視覺特征進行稀疏化（比如dVAE、VQ-VAE等）的召回通路，也許這樣就能徹底拋棄稠密檢索的召回通路了哈哈哈。算是拋磚引玉吧，讀者有啥想法歡迎交流~
Reference
[1]. Luo, Ziyang, Pu Zhao, Can Xu, Xiubo Geng, Tao Shen, Chongyang Tao, Jing Ma, Qingwei Lin, and Daxin Jiang. "Lexlip: Lexicon-bottlenecked language-image pre-training for large-scale image-text sparse retrieval." In Proceedings of the IEEE/CVF international conference on computer vision, pp. 11206-11217. 2023. aka LexLIP
[2]. https://fesianxu.github.io/2023/03/04/story-of-multimodal-models-20230304/, 《視頻與圖片檢索中的多模態語義匹配模型：原理、啟示、應用與展望》
[3]. Thibault Formal, Benjamin Piwowarski, and Stephane Clin- ´ chant. SPLADE: sparse lexical and expansion model for first stage ranking. In Fernando Diaz, Chirag Shah, Torsten Suel, Pablo Castells, Rosie Jones, and Tetsuya Sakai, editors, SIGIR ’21: The 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, Virtual Event, Canada, July 11-15, 2021, pages 2288–2292. ACM, 2021 aka SPLADE
[4]. Li, Junnan, Ramprasaath Selvaraju, Akhilesh Gotmare, Shafiq Joty, Caiming Xiong, and Steven Chu Hong Hoi. "Align before fuse: Vision and language representation learning with momentum distillation." Advances in Neural Information Processing Systems 34 (2021). short for ALBEF
[5]. https://fesianxu.github.io/2024/06/30/video-retrieval-multimodal-20240630/， 《萬字淺析視頻搜索系統中的多模態能力建設》
[6]. https://fesianxu.github.io/2024/07/06/20240706-visual-shortcome-mllm/， 《基于CLIP特征的多模態大模型中的視覺短板問題》