2024年5月5日 星期日

Matryoshka Representation Learning (MRL) 與 Vespa

這幾天經由 Vespa 的文章,讀到 OpenAI 在今年(2024 年)二月份的一篇在解說 MRL (Matryoshka Representation Learning) 的文章,看起來真的相當有趣!不過我並不是這個領域的專業,所以細部的內容我也沒有完全理解,這篇筆記大體只是我自己在應用上的基本理解。如果內容有誤也歡迎留言提醒。😇

MRL 是什麼?

Matryoshka doll 就是俄羅斯娃娃,而 Matryoshka Representation Learning 就是一種利用類似俄羅斯娃娃的特性(小的娃娃可以被裝進大的娃娃裡)做出來的 Embedding 方法,或者說是東西的表達手法。

那麼 MRL 跟其他 Embedding 方法最大的差異是什麼呢?就我目前的理解,最大的差異就是上述提到的,它具有俄羅斯娃娃的特性。具體來說,就是小的 Embedding 可以被放在大的 Embedding 裡,或者說是大的 Embedding 切割掉尾巴後,就幾乎等於是小的 Embedding。這個特性最主要的優勢,是在於我們可以不用針對不同維度的 Embedding 分別儲存不同的 Embedding,而是可以統一存一個 Embedding,在需要高維的時候使用 Embedding 裡的比較多內容、需要低維(更快更低成本)的時候則使用 Embedding 裡的一部份內容即可。

俄羅斯娃娃

Vespa 與 MRL

那麼回到 Vespa 的使用情境,MRL 可以帶來什麼效果呢?首先如同 MRL 文章在上圖中提到的,MRL 會很適合用在 multi-phase ranking 的場合,例如先使用低維度來找出較符合結果的 candidates,然後第二階段時再針對這些 candidates 使用高維度來做更精準的計算,找出最佳的結果。

Vespa 二月份的文章中,首先提到他們做了一點基本的實驗,發現 MRL 產生的 Embedding 在使用相同模型的情況下,產生 8-dimension 跟 3072-dimension 的 Embedding,確實中間只差了 scale 而已,本質上非常相似。接著他們實際實驗了把 3072-dimension 的 Embedding,裁掉尾巴並只留下最前面的 8-dimension,然後去跟一開始就用 8-dimension 產生的 Embedding 做 cosine similarity 比較,發現兩者的相似度高達 99.99999999999996%

最後,Vespa 團隊用 COVID-19 的資料集做了實際的搜尋實驗,分別有以下幾種:

  • query_exact - 使用完整的 3072-dimension 並且不使用 ANN,也就是做最精準的語義搜尋。
  • query_256 - 改為使用 256-dimension 的低維度 Embedding,同樣也不使用 ANN。
  • query_256_ann - 使用 256-dimension 的低維度 Embedding,同時使用 ANN 做語義搜尋。
  • query_rerank - 使用 256-dimension 的 Embedding 作為 1st-phase,選出 top 100,然後用 3072-dimension 作為 2nd-phase
    結果如下:
run query_exact
.................................................. avg query time 2.7918 s
run query_256
.................................................. avg query time 0.3040 s
run query_256_ann
.................................................. avg query time 0.0252 s
run query_rerank
.................................................. avg query time 0.0310 s

query_exact     ndcg_cut_10: 0.7870
query_256     ndcg_cut_10: 0.7574
query_256_ann     ndcg_cut_10: 0.7552
query_rerank     ndcg_cut_10: 0.7886

先看效能的部份,從 3072-dimensions 改為使用 256-dimensions 可以大幅加快搜尋速度,query time 只有原本的大約 11% 而已。而使用 ANN 的話,跟最開始的 query_exact 比較,甚至是只用了不到 1% 的時間!但效能提昇了,搜尋結果變差了多少呢?從 NDCG 來看,差異從 0.787 下降到 0.755 ,其實差異並沒有非常顯著。

如果我們直接比較 query_exactquery_rerank 的話,得到的結果是 query time 是原本的 1.1% ,而 NDCG 幾乎一樣!從數據來看,MRL 的潛力真的相當大 😍