In-batch negatives Embedding模型介紹與實踐

語義索引(可通俗理解爲向量索引)技術是搜索引擎、推薦系統、廣告系統在召回階段的核心技術之一。語義索引模型的目標是:給定輸入文本,模型可以從海量候選召回庫中快速、準確地召回一批語義相關文本。語義索引模型的效果直接決定了語義相關的物料能否被成功召回進入系統參與上層排序,從基礎層面影響整個系統的效果。

In-batch negatives

我們採用百度paddleNLP裏提到的In-batch Negatives方案。

In-batch Negatives 策略的訓練數據爲語義相似的 Pair 對,策略核心是在 1 個 Batch 內同時基於 N 個負例進行梯度更新,將Batch 內除自身之外其它所有 Source Text 的相似文本 Target Text 作爲負例,例如: 上例中“我手機丟了,我想換個手機” 有 1 個正例(”我想買個新手機,求推薦“),3 個負例(1.求秋色之空全集漫畫,2.手機學日語的軟件,3.俠盜飛車罪惡都市怎麼改車)。

具體來說,In-batch negatives策略的實施步驟如下:

  1. 選擇正樣本:首先從當前批次中選擇出一個正樣本,這個樣本是模型需要正確識別的目標樣本。
  2. 選擇負樣本:然後從同一批次中隨機選擇或根據特定規則選擇一些負樣本。這些負樣本可以是與正樣本相似但被錯誤標記的樣本,也可以是完全不相關的樣本。
  3. 模型訓練:將正樣本和負樣本一起輸入模型進行訓練。模型需要學會區分正樣本和負樣本,從而提高推薦或檢索的準確性。

In-batch negatives策略的優勢在於:

  • 提高模型的區分能力:通過在每個批次中引入負樣本,模型被迫學習如何區分正樣本和負樣本,這有助於提高模型的泛化能力和區分度。
  • 利用現有數據:不需要額外的負樣本庫,可以直接利用當前批次中的數據作爲負樣本,這在數據有限的情況下尤其有用。
  • 減少計算資源消耗:與從全局樣本集中採樣負樣本相比,In-batch negatives可以減少計算資源的消耗,因爲它避免了在整個數據集上進行負採樣的需要。

然而,In-batch negatives策略也存在一些潛在的問題,例如:

  • 批次大小的限制:如果批次大小較小,可能無法提供足夠多樣化的負樣本,這可能影響模型的學習效果。
  • 偏差問題:由於負樣本是在同一個批次中選擇的,可能會出現某些樣本被頻繁選爲負樣本的情況,這可能導致模型學習到的表示存在偏差。

一般通過 Recall@1,Recall@5 ,Recall@10 ,Recall@20 和 Recall@50 指標來評估語義索引模型的召回效果。按照paddleNLP給出的基線:

策略 模型 Recall@1 Recall@5 Recall@10 Recall@20 Recall@50
In-batch Negatives ernie 1.0 51.301 65.309 69.878 73.996 78.881
In-batch Negatives rocketqa-zh-base-query-encoder 59.622 75.089 79.668 83.404 87.773

rocketqa作爲打底transformer模型效果更好。

總結,爲什麼爲採用In-batch negatives,一方面能充分利用現有數據,不用單獨準備負樣例,減少投入,另外一方面模型的區分能力也比較好。

模型數據方案

流傳一句話,用1億條數據,訓練10個epoch,不如用10億數據訓練一個epoch,也就是見多識廣,大力出奇跡。

我們要訓練一個給搜索用的向量召回模型,核心就是讓準備足夠多的正樣例數據。正樣例數據,一方面網上有較多的開源數據,可以直接利用。另外一方面,之間瞭解SimBERT 時,他們的數據很多也源自於搜索數據,所以可以通過搜索引擎將query和召回結果的doc作爲相似句對。

作爲試驗,我們構造了8000萬的一個小訓練集,用rocketqa-zh-mini-query-encoder作爲打底模型,訓練256維的embedding模型。

root_path=inbatch
python -u -m paddle.distributed.launch --gpus "0" \
    train_batch_neg.py \
    --device gpu \
    --save_dir ./checkpoints/${root_path} \
    --batch_size 64 \
    --learning_rate 5E-5 \
    --epochs 3 \
    --output_emb_size 256 \
    --model_name_or_path rocketqa-zh-mini-query-encoder \
    --save_steps 5000 \
    --max_seq_length 128 \
    --margin 0.2 \
    --train_set_file recall/train.csv \
    --recall_result_dir "recall_result_dir" \
    --recall_result_file "recall_result.txt" \
    --hnsw_m 100 \
    --hnsw_ef 100 \
    --recall_num 50 \
    --similar_text_pair_file "recall/dev.csv" \
    --corpus_file "recall/corpus.csv"

訓練完成導出onnx模型:

def convert_model(model_path):

    try:
        import onnx
        import onnxruntime as ort
        import paddle2onnx
        from onnxconverter_common import float16
    except ImportError:
        print(
            "The inference precision is change to 'fp32', please install the dependencies that required for 'fp16' inference, pip install onnxruntime-gpu onnx onnxconverter-common"
        )
    onnx_dir = os.path.join(model_path, "onnx")

    if not os.path.exists(onnx_dir):
        os.mkdir(onnx_dir)
    float_onnx_file = os.path.join(onnx_dir, "model.onnx")
    if not os.path.exists(float_onnx_file):
        onnx_model = paddle2onnx.command.c_paddle_to_onnx(
            model_file=os.path.join(model_path, "inference.pdmodel"),
            params_file=os.path.join(model_path, "inference.pdiparams"),
            opset_version=13,
            enable_onnx_checker=True,
        )
        with open(float_onnx_file, "wb") as f:
            f.write(onnx_model)
    fp16_model_file = os.path.join(onnx_dir, "fp16_model.onnx")
    if not os.path.exists(fp16_model_file):
        onnx_model = onnx.load_model(float_onnx_file)
        trans_model = float16.convert_float_to_float16(onnx_model, keep_io_types=True)
        onnx.save_model(trans_model, fp16_model_file)

加載測試:

class MiniRocketQAEmbedding():

    def __init__(self, model_file: str = model_file, use_gpu: bool = True):

        providers = ['CUDAExecutionProvider'] if use_gpu else ['CPUExecutionProvider']
        sess_options = ort.SessionOptions()
        self.predictor = ort.InferenceSession( 
            model_file, sess_options=sess_options, providers=providers)
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)


    def embeding(self, embeding_text):
        features = self.tokenizer(embeding_text, max_seq_len=128,
                                  pad_to_max_seq_len=True, truncation_strategy="longest_first")

        vecs = self.predictor.run(None, features.data)
        return vecs[0]


    def similarity(self, pairs):
        query = pairs[0][0]
        texts = [item[1] for item in pairs]
        emdbeding_text = [query]
        emdbeding_text.extend(texts)
        features = self.tokenizer(emdbeding_text, max_seq_len=128,
                                  pad_to_max_seq_len=True, truncation_strategy="longest_first")

        vecs = self.predictor.run(None, features.data)

        # print(vecs)

        query_embeding = vecs[0][0]
        vecs_text1 = query_embeding / (query_embeding**2).sum() ** 0.5

        result = []
        for i in range(1, len(vecs[0])):
            vecs_text2 = vecs[0][i]
            vecs_text2 = vecs_text2 / (vecs_text2**2).sum() ** 0.5
            similarity = (vecs_text1 * vecs_text2).sum()
            result.append({"similarity": float(similarity)})

        return result

if __name__ == "__main__":
    bert = MiniRocketQAEmbedding(use_gpu=False)
    import time
    start = time.time()
    bert.embeding(["雙魚座性格特點","雙魚座性格特點"])
    print((time.time() - start) * 1000)

通過MTEB框架來測試自建搜索測試集效果:

if __name__ == '__main__':

    model = MyModel()
    task_names = ["SSRetrieval"]

    for task in task_names:
        model.query_instruction_for_retrieval = None
        evaluation = MTEB(tasks=[task], task_langs=['zh', 'zh-CN'])
        evaluation.run(model, output_folder=f"zh_results/256_model", batch_size=64)

測試結果:

{
  "dataset_revision": null,
  "dev": {
    "evaluation_time": 251.86,
    "map_at_1": 0.13427,
    "map_at_10": 0.62859,
    "map_at_100": 0.72526,
    "map_at_1000": 0.72564,
    "map_at_3": 0.31398,
    "map_at_5": 0.45025,
    "mrr_at_1": 0.71863,
    "mrr_at_10": 0.81982,
    "mrr_at_100": 0.82077,
    "mrr_at_1000": 0.82078,
    "mrr_at_3": 0.80707,
    "mrr_at_5": 0.81587,
    "ndcg_at_1": 0.71803,
    "ndcg_at_10": 0.77357,
    "ndcg_at_100": 0.83634,
    "ndcg_at_1000": 0.83907,
    "ndcg_at_3": 0.72048,
    "ndcg_at_5": 0.73003,
    "precision_at_1": 0.71803,
    "precision_at_10": 0.53373,
    "precision_at_100": 0.07386,
    "precision_at_1000": 0.00747,
    "precision_at_3": 0.68889,
    "precision_at_5": 0.65699,
    "recall_at_1": 0.13427,
    "recall_at_10": 0.78675,
    "recall_at_100": 0.98082,
    "recall_at_1000": 0.99181,
    "recall_at_3": 0.35371,
    "recall_at_5": 0.53211
  },
  "mteb_dataset_name": "SSRetrieval",
  "mteb_version": "1.1.1"
}

同樣的數據集,用peg模型測試:

{
  "dataset_revision": null,
  "dev": {
    "evaluation_time": 1036.11,
    "map_at_1": 0.09911,
    "map_at_10": 0.42835,
    "map_at_100": 0.49497,
    "map_at_1000": 0.49681,
    "map_at_3": 0.2277,
    "map_at_5": 0.31901,
    "mrr_at_1": 0.56794,
    "mrr_at_10": 0.67111,
    "mrr_at_100": 0.6737,
    "mrr_at_1000": 0.67386,
    "mrr_at_3": 0.65495,
    "mrr_at_5": 0.66559,
    "ndcg_at_1": 0.56794,
    "ndcg_at_10": 0.56275,
    "ndcg_at_100": 0.62991,
    "ndcg_at_1000": 0.64939,
    "ndcg_at_3": 0.55564,
    "ndcg_at_5": 0.54815,
    "precision_at_1": 0.56794,
    "precision_at_10": 0.38468,
    "precision_at_100": 0.05755,
    "precision_at_1000": 0.00641,
    "precision_at_3": 0.53329,
    "precision_at_5": 0.49464,
    "recall_at_1": 0.09911,
    "recall_at_10": 0.55328,
    "recall_at_100": 0.7634,
    "recall_at_1000": 0.84758,
    "recall_at_3": 0.25931,
    "recall_at_5": 0.38263
  },
  "mteb_dataset_name": "SSRetrieval",
  "mteb_version": "1.1.1"
}
模型 Recall@1 Recall@10 Recall@100 Recall@1000
peg模型 9.911 55.328 76.34 84.758
微調256模型 13.427 78.675 98.082 99.181

可以看到,微調的模型,用更小的參數,見多識廣後,整體效果明顯優於未經歷大規模數據訓練的更大尺寸的模型。

參考

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章