MRC(7) : Retrieval - Dense Embedding

SeongGyun Hong·2024년 10월 3일
Dense EmbeddingMRCfaissnaverboostcamp

NaverBoostCamp

목록 보기
7/64
post-thumbnail

1. Dense Embedding

  • Sparse Embedding의 대표인 TF-IDF의 경우에 대부분의 벡터가 0이며, 차원의 수가 매우 크다는 단점이 있다.
  • 물론 차원의 문제는 Compressed format(Sparse 행렬은 대부분의 값이 0이므로, 0이 아닌 값만 저장하는 방식으로 데이터를 압축 가능) 으로 극복이 가능하지만, 그럼에도 유사성을 고려하지 못한다는 문제는 여전하다.
    이하 코드를 보면
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import pandas as pd

# 샘플 문서
documents = [
    "This is a sample document.",
    "This document is another sample.",
    "And this is yet another document, another sample.",
]

# TF-IDF 벡터화
vectorizer = TfidfVectorizer()
tfidf_matrix = vectorizer.fit_transform(documents)

# TF-IDF 행렬을 DataFrame으로 변환
df_tfidf = pd.DataFrame(tfidf_matrix.toarray(), columns=vectorizer.get_feature_names_out())
print("TF-IDF 행렬:")
print(df_tfidf)

# 코사인 유사도 계산
cosine_sim = cosine_similarity(tfidf_matrix)
similarity_df = pd.DataFrame(cosine_sim)
print("\n코사인 유사도 행렬:")
print(similarity_df)
  • TF - IDF 행렬에서 많은 값이 0이며
  • 각 단어가 한 차원을 차지하기에 어휘의 증가는 차원의 급격한 증가를 야기한다.
  • 코사인 유사도를 이용하여 뭐,,, 문서 간 유사성을 계산할 수는 있지만 이는 엄밀한 의미에서 '유사도'가 아니다. 단순히 단어 출현 빈도에 기반하여 측정되기 떄문이다. 예컨데, documnet와 sample이 공통으로 등장하는 문서 0과 1의 유사도는 0.8408이다.
  • 그리고 documnet와 text가 사전 정의적으로는 유사하더라도 다른 의미로 취급되는 경우를 반영하지 못한다. TF-IDF는 단어의 의미나 맥락을 고려하지 못하기 때문이다.

정확한 단어매칭에는 좋다. 하지만 의미적 유사성을 포착하기에는 어렵다 ! (법률, 의학등의 정확한 정보 전달에는 좋을지도..?)

어쨌든 이런 Sparse Embedding의 한계를 극복하기 위해 Dense Embedding이 등장했으며, 이는 더 작은 차원의 고밀도 벡터를 사용하여 단어의 의미와 맥락을 더 잘 표현할 수 있게 됨.

1.1 개념

  • 더 작은 차원의 고밀도 벡터(길이 50-1000)로 단어나 문서를 표현하는 방식
  • 각 차원이 특정 단어에 대응되지 않음
  • 대부분의 요소가 0이 아닌 값을 가짐

1.2 Sparse Embedding과의 비교

  • Sparse Embedding:
    • 차원이 매우 크고 대부분 0인 값으로 구성
    • 정확한 단어 일치에 유용
    • 엔지니어링 측면이 강하다. BoW 또는 TF-IDF를 통한 구현이기에 추가 학습이 불가능하다.
  • Dense Embedding:
    • 더 작은 차원의 밀집된 벡터
    • 단어의 의미와 맥락 파악에 유용
    • 추가 학습 가능
    • 각 차원이 특정 term에 대응되지 않는다.
    • 고밀도 벡터이기 때문에 대부분의 요소가 non-zero값이다.

코드 예시

from transformers import AutoTokenizer, BertModel

model_checkpoint = 'bert-base-multilingual-cased'
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = BertModel.from_pretrained(model_checkpoint)

text = "This is an example sentence"
inputs = tokenizer(text, return_tensors="pt")
outputs = model(**inputs)

dense_embedding = outputs.last_hidden_state[:, 0, :] # [CLS] token embedding

2. Dense Encoder 학습시키기

2.1 Dense Encoder란?

  • 문서나 질문을 dense embedding으로 변환하는 모델
  • 주로 BERT 등의 사전학습 언어모델을 fine-tuning하여 사용
  • BERT의 경우 [CLS] token의 output을 사용
    • [CLS]토큰은 모델의 입력 문장 맨 앞에 추가되는 특별한 토큰으로, 문장 전체의 의미를 요약하는 벡터로 사용된다. 그렇기에 [CLS]토큰은 해당 벡터의 값으로 문장 분류, 문서 검색의 역할을 할 수 있는 것

2.2 Dense Enocoder의 학습 목표

  • 각각 embedding된 query와 passage 간의 거리를 좁히거나 inner product(내적)를 높이는 것.
  • 연관된 query와 passage는 좁히고, 연관이 없는 qeury-passage 사이의 거리는 넓히는 것이 학습의 방향 !
  • MRC 데이터셋의 question/context 쌍을 활용하여 학습시킬 수 있다.
  • 학습은 항상 그 경계에 있는 데이터를 잘 다뤄야 한다. 높은 TF-IDF 스코어를 가지지만, 답을 포함하지 않는 데이터샘플을 negative 샘플로 뽑아서 적극적으로 활용해보자.

Dense Encoder 개선 방안

  • 학습 방법 개선 : DPR 적용
  • 인코더 모델 개선 (BERT보다 크고 정확한 Pretrained 모델 사용)
  • 데이터 분석 (증강, 전처리 등등)

2.4 NNL loss : negative log likelihood loss

주로 분류 문제에서 사용되는 손실 함수로, 모델이 예측한 확률 분포와 실제 레이블 간의 차이를 측정한다. NLL Loss는 모델이 예측한 확률이 실제 레이블과 일치할수록 낮아지며, 반대로 예측이 틀릴수록 높아진다.

  • 수학적 정의
    NLL Loss=−∑i=1Nlog⁡(p(yi∣xi))\text{NLL Loss} = -\sum_{i=1}^{N} \log(p(y_i | x_i))

    (N)( N )은 데이터 포인트의 수
    (yi)( y_i )는 데이터 포인트 (i)( i )의 실제 레이블
    (xi)( x_i )는 데이터 포인트 (i)( i )의 입력
    (p(yi∣xi))( p(y_i | x_i) )는 모델이 (xi)( x_i )를 입력으로 받아 (yi)( y_i )일 확률을 예측한 값

  • 양성 문서(positive passage)에 대한 손실은 다음과 같이 정의됨.
    L(qi;pi;pi1,…,pin)=−log⁡(e−(λ+μqi)∑j=1ne−(λ+μpj))

  • 실제 계산 예시

  import torch
from transformers import BertModel, BertTokenizer

# 모델과 토크나이저 로드
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 입력 예시
question = "What is the capital of France?"
positive_passage = "Paris is the capital of France."
negative_passage = "Berlin is the capital of Germany."

# 토큰화 및 인코딩
inputs_q = tokenizer(question, return_tensors='pt')
inputs_p_pos = tokenizer(positive_passage, return_tensors='pt')
inputs_p_neg = tokenizer(negative_passage, return_tensors='pt')

# 임베딩 생성
with torch.no_grad():
    outputs_q = model(**inputs_q)
    outputs_p_pos = model(**inputs_p_pos)
    outputs_p_neg = model(**inputs_p_neg)

# [CLS] 토큰의 출력 사용
h_q = outputs_q.last_hidden_state[:, 0, :]
h_p_pos = outputs_p_pos.last_hidden_state[:, 0, :]
h_p_neg = outputs_p_neg.last_hidden_state[:, 0, :]

# 유사도 계산 (내적)
sim_pos = torch.matmul(h_q, h_p_pos.T)
sim_neg = torch.matmul(h_q, h_p_neg.T)

# 손실 계산 (NLL)
loss_fn = torch.nn.CrossEntropyLoss()
logits = torch.cat([sim_pos, sim_neg], dim=1)
labels = torch.zeros(logits.size(0), dtype=torch.long)  # 양성 샘플의 인덱스는 0
loss = loss_fn(logits, labels)

print(f"Loss: {loss.item()}")
  • CrossEntropyLoss vs NLL Loss
    CrossEntropyLoss는 내부적으로 LogSoftmax와 NLLLoss를 결합한 것이다. 즉, CrossEntropyLoss를 사용하면 별도로 LogSoftmax를 적용할 필요가 없다.

3. Passage Retrieval with Dense Encoder

3.1 Dense Encoder로 문서 검색하는 과정

  1. SQuAD Dataset 클래스를 정의하여 SQuAD 데이터를 처리한다.
  2. 데이터 로더를 사용하여 배치 처리를 수행
  3. 학습 루프를 구현하여 DPR 모델(question encoder와 context encoder)을 파인튜닝 한다.
  4. 학습된 모델을 사용하여 문맥과 질문의 임베딩을 생성
  5. FAISS를 사용하여 효율적인 유사도 검색 수행
  • 본 예시 코드에서는 in-batch negative sampling을 사용하여 학습하는데, 이는 배치 내의 다른 문맥들을 negative example로 사용한다는 의미이다.
  • 실제 적용시에는 더 복잡한 negative sampling 전략을 사용할 수 있다.

3.2 코드 예시

import torch
from torch.utils.data import DataLoader, Dataset
from transformers import DPRContextEncoder, DPRQuestionEncoder, DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer, AdamW
from datasets import load_dataset
import faiss
import numpy as np

# 디바이스 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# SQuAD 데이터셋 로드
squad_dataset = load_dataset("squad", split="train[:5000]")  # 예시를 위해 5000개만 사용

# DPR 모델 및 토크나이저 로드
ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base").to(device)
ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")

q_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base").to(device)
q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")

# 데이터셋 클래스 정의
class SQuADDataset(Dataset):
    def __init__(self, dataset, q_tokenizer, ctx_tokenizer):
        self.dataset = dataset
        self.q_tokenizer = q_tokenizer
        self.ctx_tokenizer = ctx_tokenizer

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        item = self.dataset[idx]
        question = self.q_tokenizer(item['question'], truncation=True, padding='max_length', max_length=128, return_tensors='pt')
        context = self.ctx_tokenizer(item['context'], truncation=True, padding='max_length', max_length=512, return_tensors='pt')
        
        return {
            'question_input_ids': question['input_ids'].squeeze(),
            'question_attention_mask': question['attention_mask'].squeeze(),
            'context_input_ids': context['input_ids'].squeeze(),
            'context_attention_mask': context['attention_mask'].squeeze(),
        }

# 데이터셋 및 데이터로더 생성
train_dataset = SQuADDataset(squad_dataset, q_tokenizer, ctx_tokenizer)
train_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=True)

# 손실 함수 및 옵티마이저 정의
loss_fn = torch.nn.CrossEntropyLoss()
optimizer = AdamW(list(q_encoder.parameters()) + list(ctx_encoder.parameters()), lr=2e-5)

# 학습 함수
def train(q_encoder, ctx_encoder, train_dataloader, optimizer, epochs=3):
    for epoch in range(epochs):
        q_encoder.train()
        ctx_encoder.train()
        total_loss = 0
        for batch in train_dataloader:
            optimizer.zero_grad()
            
            q_inputs = {'input_ids': batch['question_input_ids'].to(device),
                        'attention_mask': batch['question_attention_mask'].to(device)}
            ctx_inputs = {'input_ids': batch['context_input_ids'].to(device),
                          'attention_mask': batch['context_attention_mask'].to(device)}
            
            q_embeddings = q_encoder(**q_inputs).pooler_output
            ctx_embeddings = ctx_encoder(**ctx_inputs).pooler_output
            
            sim_scores = torch.matmul(q_embeddings, ctx_embeddings.t())
            labels = torch.arange(sim_scores.size(0)).to(device)
            
            loss = loss_fn(sim_scores, labels)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_dataloader):.4f}")

# 모델 학습
train(q_encoder, ctx_encoder, train_dataloader, optimizer)

# 문맥 임베딩 생성 함수
def get_context_embedding(context):
    inputs = ctx_tokenizer(context, return_tensors="pt", max_length=512, truncation=True).to(device)
    with torch.no_grad():
        embeddings = ctx_encoder(**inputs).pooler_output
    return embeddings.cpu().numpy()

# 질문 임베딩 생성 함수
def get_question_embedding(question):
    inputs = q_tokenizer(question, return_tensors="pt", max_length=512, truncation=True).to(device)
    with torch.no_grad():
        embeddings = q_encoder(**inputs).pooler_output
    return embeddings.cpu().numpy()

# 문맥 임베딩 생성
context_embeddings = [get_context_embedding(ctx) for ctx in squad_dataset["context"]]
context_embeddings = np.vstack(context_embeddings)

# FAISS 인덱스 생성 및 추가
d = context_embeddings.shape[1]  # 임베딩 차원
index = faiss.IndexFlatIP(d)
index.add(context_embeddings.astype('float32'))

# 검색 함수
def search(query, top_k=5):
    question_embedding = get_question_embedding(query)
    scores, indices = index.search(question_embedding, top_k)
    return [(squad_dataset[int(i)]["context"], squad_dataset[int(i)]["question"], float(s)) for i, s in zip(indices[0], scores[0])]

# 테스트
test_query = "What is the capital of France?"
results = search(test_query)

print(f"Query: {test_query}\n")
for i, (context, original_question, score) in enumerate(results, 1):
    print(f"Result {i}:")
    print(f"Context: {context[:100]}...")
    print(f"Original Question: {original_question}")
    print(f"Similarity Score: {score:.4f}\n")

FAISS?

  • FAISS(Facebook AI Similarity Search)란, 대규모 벡터 데이터베이스에서 효율적인 유사성 검색을 수행하기 위한 라이브러리이다.
  • Inner Product 또는 Cosine Similarity를 사용하여 벡터 간 유사성을 계산할 수 있다.
  • 대규모 데이터셋에서도 빠른 검색이 가능하다.

이하 FAISS를 활용한 기본 예제 코드 참고

import numpy as np
import faiss

# 데이터 준비
d = 768  # BERT 임베딩의 차원
nb = len(document_embeddings)  # 문서 수
xb = np.array(document_embeddings).astype('float32')

# FAISS 인덱스 생성
index = faiss.IndexFlatIP(d)  # 내적(Inner Product) 유사도 사용
index.add(xb)  # 문서 임베딩을 인덱스에 추가

# 쿼리 처리 함수
def retrieve_faiss(query_or_dataset, topk=1):
  with torch.no_grad():
      query_vec = model.encode([query_or_dataset])[0].astype('float32')

  # FAISS를 사용한 유사도 검색
  D, I = index.search(query_vec.reshape(1, -1), topk)
  
  return I[0]  # 상위 k개의 문서 인덱스 반환

# 사용 예시
query = "What is the capital of France?"
top_k_indices = retrieve_faiss(query, topk=5)

# 결과 출력
for idx in top_k_indices:
  print(f"Document: {documents[idx]}")
profile
SeongGyun Hong
헤매는 만큼 자기 땅이다.
이전 포스트

MRC(6) : Retrieval

다음 포스트

MRC(8) : FAISS

0개의 댓글