cleanUrl: "/e5-mistral"
floatFirstTOC: right
title: "LLM 기반 최고성능 Text Embedding | AI Lab"
description: "GPT-4로 학습데이터를 생성하고 Mistral-7B를 튜닝하여 MTEB 최고 성능을 달성한 Text Embedding 연구를 소개합니다."

(이미지출처) Microsoft Designer - Stunning designs in a flash 를 사용하여 자동 생성 후 편집

도입

• LLM이 외부의 지식이나 정보를 동적으로 활용하는 방법으로 RAG(Retrieval-Augmented Generation) framework가 매우 각광 받고 있음
• Embedding model은 retrieval의 성능을 결정하는 가장 중요한 요소로 embedding model의 성능을 측정하기 위한 대규모 벤치마크 중 대표적인 것으로 MTEB(Massive Text Embedding Benchmark)이 존재 (이전 아티클 참고 : MTEB 상위권 방법론들)
• 이 아티클에서는 최근 MTEB에서 최고의 성능을 달성한 방법론에 대하여 소개함
• 리뷰 논문 : https://arxiv.org/abs/2401.00368 (by Microsoft)

개요

• 한줄 요약 : GPT-4를 활용하여 학습 데이터를 생성하고 그것으로 Mistral-7b 모델을 tuning 하여 MTEB에서 SOTA 성능을 달성
• 다국어 데이터도 생성하여 총 93개 언어에 대한 embedding 능력을 학습
• DeepNet, LongNet, E5, Kosmos 시리즈 등 꽤 실용적인 연구를 많이 하는 그룹에서 발표한 논문

방법론 상세

• Data Generation

  

  • 다양한 종류의 데이터를 생성하기 위하여 embedding task를 크게 몇 개의 그룹으로 나눈 후 각 그룹별로 다른 prompt template을 활용하여 데이터를 생성
  • Asymmetric tasks
    • Query와 document가 의미상 연관은 있지만 유사한 내용은 아닌 경우(질문-답변, 검색 query-관련 문서 등)
    • 이 경우는 short-short, short-long, long-short, long-long 4가지 subgroup으로 나눔
    • 위 그림 1에서와 같이 먼저 task의 instruction을 LLM이 생성하게 한 후 생성한 task instruction에 대하여 query 및 document를 생성하는 2단계로 진행
  • Symmetric tasks
    • Query와 document가 유사한 의미를 가지는 경우
    • STS와 bitext retrieval 2가지 subgroup으로 나눔
    • 이 경우는 task instruction이 직관적이므로 사람이 직접 생성한 prompt를 통하여 바로 query 및 document를 생성 (e.g., “Retrieve semantically similar text.”)
  • 프롬프트에서 query_length, query_type, num_words 등 주요 feature들에 대하여 여러 개의 옵션을 사전에 만든 후 생성할 때마다 random sampling하여 최대한 다양한 데이터를 생성하였음

• 모델

  • Backbone : Mistral-7b

  • LoRA with rank=16, DeepSpeed ZeRO-3

  • seq_length=512, batch_size=2048, LR=10^-4, linear decay, weight decay=0.1

  • Query에 task instruction 을 합친 형태를 입력으로 사용함

    

  • Query와 document 입력 맨 뒤에 [EOS] token을 삽입한 후 이 토큰에 대한 embedding을 각각의 embedding vector($\mathbf{e}_q, \mathbf{e}_d$)로 사용.

  • 입력 query와 positive document가 유사한 embedding이 되도록 InfoNCE loss로 학습. 여기서 $\mathbb{N}$은 in-batch negative 및 hard negative 모두 포함한 집합

  $$ \min \mathbb{L}=-\log\frac{\phi(q_{inst}^{+},d^+)}{\phi(q_{inst}^{+},d^+)+\sum_{n_i\in\mathbb{N}}\phi(q_{inst}^{+},n_i)}\\ \phi(q,d)=\exp(\frac{1}{\tau}\cos(\mathbf{e}_q,\mathbf{e}_d)) $$

  • 학습은 1 epoch 진행. V100 32장으로 18시간 소요
  • MTEB evaluation은 document에 대한 encoding에 시간이 많이 걸려서 V100 8장으로 3일 소요

결과

• 생성 데이터 유형 및 언어 종류

  

  • 총 150k개의 task로부터 500k개의 데이터를 생성. 데이터 생성에 180M 토큰 사용
  • 기존 labeled 데이터도 학습에 활용(아래 표에서 full data). 다음의 데이터로부터 샘플링하여 위에서 생성한 데이터 포함 총 1.8M 데이터를 수집
    • ELI5, HotpotQA, FEVER, MIRACL, MS-MARCO, NQ, NLI, SQuAD, TriviaQA, Quora Duplicate Questions, MrTyDi, DuReader, T2Ranking
  • Task 종류는 short-long 이 가장 많고 long-short, STS 등의 순서
  • 언어는 (당연하게도) 영어가 가장 많음
  • GPT-3.5-turbo로 약 25%, GPT-4로 75%를 생성. GPT-3.5-turbo는 지시문이나 가이드라인을 명확히 따르지 않는 경우가 일부 있으나 전체 품질은 나쁘지 않아서 학습 데이터에 포함시킴