All Datasets V3 MiniLM L12

🚀 句子相似度模型

本項目旨在使用自監督對比學習目標，在超大型句子級數據集上訓練句子嵌入模型。通過該模型，能夠將輸入的句子轉化為向量，捕捉句子的語義信息，可廣泛應用於信息檢索、聚類或句子相似度任務等領域。

🚀 快速開始

本模型可作為句子編碼器使用。給定輸入句子，它將輸出一個捕獲句子語義信息的向量。句子向量可用於信息檢索、聚類或句子相似度任務。

✨ 主要特性

• 大規模數據訓練：使用超過10億個句子對的數據集進行微調，使模型能學習到豐富的語義信息。
• 高效架構：基於預訓練的MiniLM-L12模型進行微調，結合高效的對比學習目標，在保證性能的同時提高訓練效率。
• 多領域適用性：訓練數據涵蓋了多種不同類型的數據集，包括問答、圖像描述、代碼搜索等，使模型具有廣泛的適用性。

📦 安裝指南

使用此模型前，需安裝SentenceTransformers庫，可通過以下命令進行安裝：

pip install sentence-transformers

💻 使用示例

基礎用法

以下是如何使用SentenceTransformers庫來獲取給定文本特徵的示例代碼：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('flax-sentence-embeddings/all_datasets_v3_MiniLM-L12')
text = "Replace me by any text you'd like."
text_embbedding = model.encode(text)
# array([-0.01559514,  0.04046123,  0.1317083 ,  0.00085931,  0.04585106,
#        -0.05607086,  0.0138078 ,  0.03569756,  0.01420381,  0.04266302 ...],
#        dtype=float32)

📚 詳細文檔

訓練過程

預訓練

我們使用了預訓練的MiniLM-L12模型。有關預訓練過程的更多詳細信息，請參考該模型的卡片。

微調

我們使用對比目標對模型進行微調。具體來說，我們計算批次中每個可能句子對的餘弦相似度，然後通過與真實對進行比較來應用交叉熵損失。

超參數

我們在TPU v3 - 8上訓練模型。訓練共進行540k步，批次大小為1024（每個TPU核心128）。學習率預熱步數為500，序列長度限制為128個標記。我們使用AdamW優化器，學習率為2e - 5。完整的訓練腳本可在當前倉庫中獲取。

訓練數據

我們使用多個數據集的組合來微調模型。句子對的總數超過10億。每個數據集的採樣概率在data_config.json文件中有詳細配置。

🔧 技術細節

• 模型架構：基於MiniLM-L12預訓練模型，通過微調適應句子嵌入任務。
• 對比學習目標：計算句子對的餘弦相似度，並使用交叉熵損失進行優化，使相似句子的向量在嵌入空間中更接近。
• 硬件加速：使用7個TPU v3 - 8進行訓練，並得到了Google的Flax、JAX和雲團隊成員在高效深度學習框架方面的支持。