視覺多語言：mcdse-2b 模型介紹

今天，我將介紹一種新的實驗性多語言嵌入模型，用於靈活的視覺文件檢索。mcdse-2b-v1 (🤗) 基於 MrLight/dse-qwen2-2b-mrl-v1 訓練，並使用 DSE 方法。

這個模型允許您嵌入頁面/幻燈片截圖，並使用自然語言進行查詢。無論是表格、圖表、示意圖、影像還是文字，mcdse-2b-v1 都能將所有內容編碼成一個單一的嵌入向量，從而無需傳統的 OCR、文件佈局分析、閱讀順序檢測、分塊、表格/公式提取等操作。

• 在 🇮🇹 義大利語、🇪🇸 西班牙語、🇬🇧 英語、🇫🇷 法語和 🇩🇪 德語上表現出色.

• 套娃表示學習：嵌入可以高效地從 1536 維縮放到 256 維。您可以將大小減少 6 倍，同時仍保持 95% 的嵌入質量。

• 二值化表現卓越：768 維二值向量保留了 1536 維浮點向量 99% 的檢索質量。使用二值向量，您可以在 10GB 空間中編碼 1 億個多語言頁面。

• 快速 vLLM 推理：在 vLLM 上執行推理，並高效地大規模提供嵌入服務，可用於生產環境。檢視部署瞭解更多資訊。

我的基準測試並非完美無缺，因此我鼓勵您在自己的資料上測試模型。這是一個早期版本，還有很大的改進空間。儘管如此，結果仍突顯了強大的多語言檢索器，它能出色地適應各種記憶體/速度要求。

訓練

mcdse-2b 基於 MrLight/dse-qwen2-2b-mrl-v1 使用低秩介面卡 (LoRA) 在多語言文件語料庫上進行訓練。我使用 8xRTX3090 顯示卡，採用 DSE 方法，並使用以下引數進行了訓練：

• 紀元 = 1
• 預熱比例 = 0.1
• 學習率 = 1e-5
• 最佳化器 = adamw_torch
• 排程 = 線性
• 總批次大小 = 16
• LoRA
  • Alpha = 64
  • R = 16
  • Dropout = 0.1
  • DoRA = True

資料集

該資料集包含從公共網際網路自動抓取的 24K PDF 頁面。從每個文件中提取隨機頁面，轉換為壓縮 JPEG 影像，並過濾以去除空白頁和重複項。最終的頁面截圖是獨一無二的，涵蓋了廣泛的主題。

我使用 gemini-flash-1.5-002 根據每張圖片生成查詢。Gemini 被指示生成三種類型的查詢：

• 一個寬泛的主題查詢：總結文件的整體主題。
• 一個具體的詳細問題：捕捉內容中的細微差別。
• 一個視覺查詢：側重於圖表、圖形、影像或簽名等視覺元素。

整個訓練和評估資料集僅花費 2 歐元生成（感謝 Gemini Flash！）

然後，每張影像根據其文字密度進行 0 到 2 的分類。我使用了在 DocLayNet 上微調的 omoured YOLOv10n 模型，用於檢測圖形與文字等區域。根據這些比例，我啟發式地計算文字密度。我計劃使用這種分類來提高模型在文字密集文件上的效能。

• 0 = 僅視覺元素
• 1 = 視覺元素和文字混合
• 2 = 僅文字

評估和訓練資料集尚未釋出。我非常願意開源它們，但我仍然不確定如何在不違反任何許可（如果有的話）的情況下正確地做到這一點。如果您知道如何幫助我，請聯絡我！

訓練執行

模型按以下順序依次訓練每種語言：

1. 法語：6k 樣本
2. 西班牙語：6k 樣本
3. 義大利語：6k 樣本
4. 德語：6k 樣本

這個順序是根據基礎模型在這些語言中的檢索效能確定的，第一個是表現最好的。我的直覺是，鑑於資料集較小，從更強的語言開始可能有助於平衡模型在整體上的改進。

在達到這個最終檢查點之前，我進行了多次執行，以測試各種策略並驗證我的一些直覺。

• 語言順序： 我交換了最後兩種語言的順序，發現最後訓練德語可以使其在評估中的效能提高 1.7%，同時保持其他語言相似的分數。

• 模型初始化： 我使用每種語言的 10k mmarco 對初始化模型。這導致所有語言的效能更差，特別是低維嵌入。例如，使用 512 維嵌入的法語 NDCG@5 在使用 mmarco 訓練時下降了 2%。

• 不同的影像調整演算法： 我開發了一個自定義調整函式（custom_resize），它在將影像縮放到 min_pixels 和 max_pixels 範圍內時嚴格保留影像的寬高比。所有評估都使用 qwen_vl_utils 中的標準調整函式。使用自定義調整函式訓練的模型優於標準方法，NDCG@5 平均提高了 +1.7%（1536 維）。探索使用此 custom_resize 函式訓練 ColQwen 模型會很有趣。

  調整函式	平均	英語	義大利語	西班牙語	法語	德語
  qwen2_vl_utils	80.8	80.2	80.5	79.6	81	82.6
  自定義調整	82.2	80.8	81.2	80.7	84.5	83.8
  	+1.7%	+0.7%	+0.9%	+1.4%	+4.0%	+1.4%

評估

由於公開可用的多語言文件影像檢索資料集有限，該模型已使用自定義資料集進行評估。此評估資料集專門設計用於基準測試模型在各種語言中的效能。

此評估資料集是使用與訓練資料集相同的方法和流程建立的。但是，文件主題通常不同，並且訓練和評估資料集之間沒有共享影像，以避免任何評估汙染。NDCG 分數是透過對每種語言的 1K 文件索引執行 100 個唯一查詢來計算的。

套娃表示學習

此模型使用套娃表示學習（MRL）進行訓練，維度分別為：1536、1024、768、512、384、256。訓練過程中使用的損失函式經過校準，以跟蹤所有這些維度上的效能，從而使模型能夠預先載入最重要的識別資訊。這有效地允許您根據自己的規模和預算縮小嵌入維度。

每個維度的平均 NDCG@5。有趣的是，該模型在英語方面顯示出改進，儘管該語言未包含在訓練集中。該模型在 256 維度上效能提高了 6%，並且在所有維度上平均整體提高了 4%。評估使用 FAISS 的 IndexFlatL2 進行。

NDCG@5（浮點數）

二值嵌入

mcdse-2b-v1 在二值化方面表現明顯更好，尤其是在較低維度上。該模型在 256 維度上提高了 23%，整體平均提高了 13%。評估使用 FAISS 的 IndexBinaryFlat 進行。但是，為什麼二值嵌入更優呢？

此表顯示，mcdse-2b-v1 的二進位制嵌入比基礎模型的 1536 維浮點向量效能高出 1%，同時記憶體消耗減少了 16 倍。除了這些優勢，二進位制嵌入還可以透過漢明距離進行 40 倍的快速搜尋，因為比較兩個二進位制向量只需 2 個 CPU 週期（異或、位計數）。

NDCG@5（二進位制）

ShiftProject

vidore/shiftproject_test 資料集是 ViDoRe 基準測試的一部分。它包含來自 Shift Project 關於環境的法語查詢和文件。查詢是使用 Claude-3 Sonnet 在與生成 vidore/colpali_train_set 抓取文件查詢相同的法語翻譯提示下生成的。

這是 ShiftProject 資料集上的 NDCG@5，使用 1536 浮點維度，最多評估 960 個影像補丁。

我原以為 mcdse-2b-v1 的分數會高於基礎模型，結果卻低了 3%。基礎模型是在 colpali 訓練集上訓練的，我曾以為它可能對“Claude-3 Sonnet 風格”的查詢過度優化了。為了調查這一點，我使用 gemini-1.5-flash-002 和我的提示重新生成了 ShiftProject 資料集查詢。

分數變化很大，但在這種情況下，mcdse-2b-v1 提高了 5%。這些結果傾向於表明兩種可能的原因：

1. 基礎模型對“Claude-3 Sonnet 風格”的查詢過度最佳化
2. 我的模型對“gemini-1.5-flash-002 風格”的查詢過度最佳化

在這兩種情況下，我相信 mcdse-2b-v1 透過理解更廣泛的查詢分佈，減輕了這些過度最佳化。

我生成的 Gemini 查詢有兩種格式：問題和查詢。colpali_train_set 生成的查詢只有問題。我還測試了只使用 Gemini 查詢和只使用 Gemini 問題的兩種模型。

基礎模型在 Gemini 問題上表現好 7%，在 Gemini 查詢上表現差 7%。查詢和問題之間的平均分數幾乎相同（66.7 和 66.5）。這表明我的模型已經緩解了之前提到的過度最佳化，並且通常能夠更好地理解更廣泛的查詢。在更多多語言資料上進行訓練可能會提高這個平均值，並最終提高 ShiftProject 上的效能。

Cohere Embed v3 影像

我使用最近釋出的 cohere embed-multilingual-v3.0 多模態 嵌入模型對較小版本的英語資料集進行了一些初步（且倉促）測試。該模型取得了 71 的 NDCG@5 分數，而 mcdse-2b-v1 的分數約為 84。我正在為該模型進行更全面的評估。

部署

在 HuggingFace Transformers 上，您可以使用 RTX3090 顯示卡以 32 的批處理大小（35TFLOPS）編碼約 3 張圖片/秒。像 RTX 4000 Ada 這樣的更常見的推理端 GPU 大致能提供相同的吞吐量。

vLLM

vLLM 官方僅支援 Qwen2VL 用於生成，我添加了一個新的模型類 Qwen2VLForEmbeddingGeneration 來支援嵌入任務。在 vLLM 上執行推理應該比 HuggingFace Transformers 快約 5 倍。

下載新的模型類

git clone https://github.com/marplex/mcdse && cd mcdse

下載 mcdse-2b-v1 用於本地推理

from huggingface_hub import snapshot_download
snapshot_download(repo_id="marco/mcdse-2b-v1", local_dir="/path/to/model/mcdse-2b-v1")

編輯 config.json

將 Qwen2VLForConditionalGeneration 替換為 Qwen2VLForEmbeddingGeneration

sed -i -e 's/Qwen2VLForConditionalGeneration/Qwen2VLForEmbeddingGeneration/g' /path/to/model/mcdse-2b-v1/config.json

檢查 vllm/main.py 進行本地推理

#vllm/main.py
from qwen2_vl_dse import Qwen2VLForEmbeddingGeneration, get_query_prompt, get_document_prompt
from vllm import ModelRegistry, LLM
from PIL import Image

ModelRegistry.register_model("Qwen2VLForEmbeddingGeneration", Qwen2VLForEmbeddingGeneration)

llm = LLM(
    model="/path/to/model/mcdse-2b-v1",
    limit_mm_per_prompt={
        "image": 1
    }
)

# Encode queries
query_prompt, image = get_query_prompt("Quali erano le passività totali al 31 dicembre 2017?")
outputs = llm.encode({"prompt": query_prompt, "multi_modal_data": {"image": [image]}})
outputs[0].outputs.embedding #1536 dimensional embedding

# Encode documents
dummy_document_image = Image.new('RGB', (256, 256))
document_prompt, image = get_document_prompt(dummy_document_image)
outputs = llm.encode({"prompt": document_prompt, "multi_modal_data": {"image": [image]}})
outputs[0].outputs.embedding #1536 dimensional embedding

結論

這是我第一次訓練模型，雖然充滿挑戰，但也樂趣無窮。如果沒有 HuggingFace 團隊和貢獻者的出色工作，我不可能做到這一點。我還要感謝 Manuel Faysse、Tony Wu 以及整個 Vidore 團隊在 ColPali 上的工作，感謝 Xueguang Ma 在 Tevatron 程式碼庫上的所有工作以及訓練了一個非常強大的基礎模型。我還受到了 Benjamin Clavié 及其令人印象深刻的模型公告的啟發。

我希望這個模型對您的檢索和 RAG 管道有用。正如開頭所述，我的基準測試遠非完美，實際場景中的結果可能會有所不同。我鼓勵您在自己的用例中進行測試。總的來說，視覺檢索的一個顯著優勢是，您可以透過簡單地嵌入頁面來拋棄複雜的索引管道。這就是未來！