Passage Retrieval: Re-rank

Re-rank: 재정렬

이 글은 Sparse Embedding과 Dense Embedding을 사용한 Reranking에 대해 설명한다. Reranking 작업은 정보 검색 및 추천 시스템에서 높은 정확도를 제공하기 위해 일반적으로 사용된다. 기본적인 검색 단계에서 다수의 후보 항목을 빠르게 선별한 다음, 이 후보들을 다시 평가하여 최종적인 순위를 매기는 방식이다.

Sparse와 Dense를 활용한 Rerank

기본 검색 단계에서 Sparse Embedding을 사용하여 후보 문서를 빠르게 찾은 후, Dense Embedding을 이용해 Reranking하는 방식은 효과적이다. 이 방식은 다음과 같은 단계로 이루어진다.

Sparse Retrieval (기본 검색):

초기 검색 단계에서는 빠른 검색이 필요하기 때문에 주로 Sparse Embedding을 사용한다.

예를 들어, 사용자가 입력한 쿼리에 대해 BM25 같은 모델을 이용하여 초기 후보 문서 집합을 빠르게 검색한다.

이 단계의 결과물은 수백에서 수천 개의 문서가 될 수 있으며, 이들은 고전적인 키워드 매칭을 통해 필터링된다.

Dense Rerank (재정렬):

초기 단계에서 추출된 후보 문서를 대상으로 Dense Embedding을 사용해 Reranking한다.

이 과정에서 BERT나 Siamese Network와 같은 딥러닝 기반의 임베딩 모델을 이용해 각 문서와 쿼리의 의미적 유사도를 계산한다.

Dense Embedding은 문서의 의미적 내용을 잘 포착할 수 있기 때문에 더 정교한 순위를 매길 수 있다.

예를 들어, BERT를 사용해 쿼리와 문서 간의 상호작용을 인코딩하고, 문서의 순위를 정교하게 재조정하는 방식이다.

Rerank의 방식

1. Bi-encoder 방식

구조: Bi-encoder는 두 개의 문장을 독립적으로 인코딩하여 고차원의 벡터로 표현한다. 일반적으로 쿼리와 후보 문서를 각각 독립된 신경망(대부분은 동일한 파라미터를 공유하는 BERT나 다른 트랜스포머 모델)을 통해 임베딩하고, 이 두 임베딩의 유사도를 계산하여 관련성을 측정한다.

사용 시나리오: 이 방식은 대량의 후보 문서를 빠르게 처리해야 하는 상황에서 유용하다. 쿼리와 문서의 벡터를 독립적으로 계산한 후, 벡터 간의 내적이나 코사인 유사도를 계산하는 방식이므로 효율적이다.

장점: 효율적인 계산과 저장이 가능하다. 사전에 후보 문서들을 임베딩하여 데이터베이스에 저장해 두고, 쿼리가 들어오면 빠르게 유사도를 계산할 수 있다.

단점: 쿼리와 문서를 독립적으로 인코딩하기 때문에 문맥에 따라 상호작용하는 정보가 충분히 반영되지 않을 수 있다. 특히, 더 미묘한 상호작용이 필요한 복잡한 질문에 대해선 정확도가 떨어질 수 있다.

2. Cross-encoder 방식

구조: Cross-encoder는 쿼리와 문서를 하나의 입력으로 합쳐 트랜스포머 모델에 넣어 인코딩한다. 예를 들어 [CLS] 쿼리 [SEP] 문서 [SEP] 형식으로 쿼리와 문서를 동시에 인코딩하고, 이를 통해 두 텍스트 간의 관련성을 예측한다.

사용 시나리오: 정확도가 매우 중요한 경우, 특히 최종 순위를 매기는 재정렬 작업에서 자주 사용된다. 모델이 두 문장 간의 상호작용을 충분히 반영할 수 있기 때문에 더 정밀한 관련성 판단이 가능하다.

장점: 쿼리와 문서를 하나의 입력으로 합쳐 모델에 넣기 때문에 문맥적인 상호작용이 잘 반영되며, 더 높은 정확도를 얻을 수 있다.

단점: 계산 비용이 높다. 쿼리가 들어올 때마다 쿼리와 모든 후보 문서를 조합하여 모델을 통해 인코딩해야 하므로, 대규모 데이터셋에서는 비효율적이다.

Hybrid Approach (결합 접근)

Sparse와 Dense를 결합하는 하이브리드 방법을 사용하는 경우도 있다. 이 방식은 각 Embedding의 장점을 결합하여 최적의 성능을 달성한다.

Sparse Embedding은 키워드 매칭을 잘 처리하기 때문에 특정 단어가 반드시 포함된 결과를 찾는 데 유리하다.

Dense Embedding은 의미적 유사도를 잘 반영하기 때문에 관련성은 있지만 정확히 같은 단어를 사용하지 않은 문서도 잘 찾을 수 있다.

이 두 Embedding을 Concat하거나 가중 평균을 구해 쿼리와 문서 간의 유사도를 측정하고, 최종적으로 순위를 매기는 방식으로 활용할 수 있다.

Bi-encoder Rerank

✅ 이번 실습에서는 하나의 query가 아닌 복수의 queries를 한번에 처리한다.

Dataset 로드하기

[Input]

from datasets import load_dataset

dataset = load_dataset("squad_kor_v1")

Tokenizer 준비하기

[Input]

from transformers import AutoTokenizer

model_checkpoint = "BM-K/KoSimCSE-roberta-multitask"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

데이터 샘플 준비하기

[Input]

import numpy as np
import random

random.seed(42)

num_samples = 128
sample_idx = np.random.choice(range(len(dataset['train'])), num_samples)
contexts = dataset['train'][sample_idx]['context']
queries = dataset['train'][sample_idx]['question']
ground_truth = dataset['train'][sample_idx]['context']

print(len(contexts), len(queries), len(ground_truth))

[output]

>>  128 128 128

✅ 128개의 랜덤 샘플만 사용하여 데이터를 구성한다. KorQuAD 데이터셋의 context, question feature를 사용해 Retrieval를 실습해 볼 수 있다.
▶ KorQuAD란 무엇일까?

First Step. Sparse Embedding (TF-IDF)

[Input]

def tokenizer_for_retrieve(text):
    tokens = tokenizer.tokenize(text, max_length=512, truncation=True)
    return tokens
    
from sklearn.feature_extraction.text import TfidfVectorizer
vectorizer = TfidfVectorizer(tokenizer=tokenizer_for_retrieve)

vectorizer.fit(contexts)
contexts_vec = vectorizer.transform(contexts)
queries_vec = vectorizer.transform(queries)

print(contexts_vec.shape, queries_vec.shape)

[output]

>>  (128, 6285) (128, 6285)

✅ contexts와 query를 TF-IDF를 사용해 임베딩한다.
▶ TF-IDF가 기억나지 않는다면?

Sparse Embedding 유사도 측정 (TF-IDF)

[Input]

result = queries_vec * contexts_vec.T
print(result.shape)

[output]

>>  (128, 128)

✅ 128개의 쿼리와 128개의 문장간의 유사도 점수

Sparse Embedding Top-K 추출하기 (TF-IDF)

[Input]

doc_scores = []
doc_indices = []
k = 50

if not isinstance(result, np.ndarray):
  result = result.toarray()
for i in range(result.shape[0]):
  sorted_result = np.argsort(result[i, :])[::-1]
  doc_scores.append(result[i, :][sorted_result].tolist()[:k])
  doc_indices.append(sorted_result.tolist()[:k])
           
tfidf_output = {'question': queries, 
                'contexts': [[contexts[i] for i in doc_ids] for doc_ids in doc_indices], 
                'answer': ground_truth,
                'scores': doc_scores}

tfidf_output_df = pd.DataFrame(tfidf_output)
tfidf_output_df["correct"] = tfidf_output_df.apply(lambda row: row["answer"] in row["contexts"], axis=1)

✅ 128개의 각각의 쿼리에 대해 유사도가 높은 50개의 contexts와 doc_scores(유사도 점수)를 저장한다.

Sparse Embedding Ranking 확인하기 (TF-IDF)

[Input]

example_idx = 0
example_question = queries[example_idx]
example_answer = ground_truth[example_idx]
doc_scores_by_example = doc_scores[example_idx]
doc_indices_by_example = doc_indices[example_idx]


print("[Search query]\n", example_question, "\n")

print("[Ground truth passage]")
print(example_answer, "\n")

for i in range(k):
  print("Top-%d passage with score %.4f" % (i + 1, doc_scores_by_example[i]))
  doc_id = doc_indices_by_example[i]
  print(contexts[doc_id], "\n")

[output]

>>  [Search query]
	2010년 8월 17일 FFF의 징계 청문회 후 징계를 면한 사람은? 

    [Ground truth passage]
    니콜라 아넬카, 파트리스 에브라, 프랑크 리베리, 제레미 툴랄랑, 그리고 에리크 아비달이 FIFA 월드컵 기간동안 주요 사건에 연루된 5명이었고, 2010년 8월 17일에 FFF의 징계 청문회로 소환되었다. ...

    Top-1 passage with score 0.1007
    니콜라 아넬카, 파트리스 에브라, 프랑크 리베리, 제레미 툴랄랑, 그리고 에리크 아비달이 FIFA 월드컵 기간동안 주요 사건에 연루된 5명이었고, 2010년 8월 17일에 FFF의 징계 청문회로 소환되었다. ...

    Top-2 passage with score 0.0452
    죽음은 죽은 사람에게만 비참을 가져오는 것은 아니다. 사망자가 앞서가 남겨진 산 사람에는 고독이 남겨진다. 산 사람이 친밀한 사람의 죽음의 정산에 실패하면 큰 후회가 남겨져 대단한 고통에 괴롭혀지게 된다. ...

    Top-3 passage with score 0.0390
    《Prism》은 2013년 10월 18일 발매되었고, 빌보드 200 1위로 데뷔했다. 네 달 후 페리는 로스앤젤레스에서 열린 iHeartRadio에서 앨범의 수록곡들을 선보였다. ...

	...

✅ 예시로 첫번째 쿼리에 대해 잘 예측하는 지 확인

Sparse Embedding 결과 보기 (TF-IDF)

[Input]

print(len(tfidf_output['question']))
print(len(tfidf_output['answer']))
print(len(tfidf_output['contexts'][0]))
print(len(tfidf_output['scores'][0]))

[output]

>>  128
    128
    50
    50

✅ tfidf_output에는 128개의 question과 answer, 그리고 각각의 question에는 topk=50의 contexts와 scores가 저장되었다.

Dense Embedding (Bi-encoder)

✅ Dense Embedding 학습을 구성하는 과정은 Dense Embedding 과 동일하다.

[Input]

from tqdm import tqdm, trange
import argparse
import random
import torch
import torch.nn.functional as F
from transformers import BertModel, BertPreTrainedModel, AdamW, TrainingArguments, get_linear_schedule_with_warmup

torch.manual_seed(42)
torch.cuda.manual_seed(42)
np.random.seed(42)

학습용 데이터 셋 구성하기

[Input]

from torch.utils.data import (DataLoader, RandomSampler, TensorDataset)

q_seqs = tokenizer(queries, padding="max_length", truncation=True, return_tensors='pt')
p_seqs = tokenizer(contexts, padding="max_length", truncation=True, return_tensors='pt')

train_dataset = TensorDataset(p_seqs['input_ids'], p_seqs['attention_mask'], p_seqs['token_type_ids'],
                        	  q_seqs['input_ids'], q_seqs['attention_mask'], q_seqs['token_type_ids'])

BertEncoder 구성하기 (Bi-encoder)

[Input]

class BertEncoder(BertPreTrainedModel):
  def __init__(self, config):
    super(BertEncoder, self).__init__(config)

    self.bert = BertModel(config)
    self.init_weights()

  def forward(self, input_ids,
              attention_mask=None, token_type_ids=None):

      outputs = self.bert(input_ids,
                          attention_mask=attention_mask,
                          token_type_ids=token_type_ids)

      pooled_output = outputs[1]

      return pooled_output
      
p_encoder = BertEncoder.from_pretrained(model_checkpoint)
q_encoder = BertEncoder.from_pretrained(model_checkpoint)

if torch.cuda.is_available():
  p_encoder.cuda()
  q_encoder.cuda()

BertEncoder 학습하기

[Input]

def train(args, dataset, p_model, q_model):

  # Dataloader
  train_sampler = RandomSampler(dataset)
  train_dataloader = DataLoader(dataset, sampler=train_sampler, batch_size=args.per_device_train_batch_size)


  no_decay = ['bias', 'LayerNorm.weight']
  optimizer_grouped_parameters = [
        {'params': [p for n, p in p_model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in p_model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
        {'params': [p for n, p in q_model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in q_model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]

  optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)

  t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
  scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)

  # Start training!
  global_step = 0

  p_model.zero_grad()
  q_model.zero_grad()
  torch.cuda.empty_cache()

  train_iterator = trange(int(args.num_train_epochs), desc="Epoch")

  for _ in train_iterator:
    epoch_iterator = tqdm(train_dataloader, desc="Iteration")

    for step, batch in enumerate(epoch_iterator):
      q_encoder.train()
      p_encoder.train()

      if torch.cuda.is_available():
        batch = tuple(t.cuda() for t in batch)

      p_inputs = {'input_ids': batch[0],
                  'attention_mask': batch[1],
                  'token_type_ids': batch[2]
                  }

      q_inputs = {'input_ids': batch[3],
                  'attention_mask': batch[4],
                  'token_type_ids': batch[5]}

      p_outputs = p_model(**p_inputs)  # (batch_size, emb_dim)
      q_outputs = q_model(**q_inputs)  # (batch_size, emb_dim)


      # Calculate similarity score & loss
      sim_scores = torch.matmul(q_outputs, torch.transpose(p_outputs, 0, 1))  # (batch_size, emb_dim) x (emb_dim, batch_size) = (batch_size, batch_size)

      # target: position of positive samples = diagonal element
      targets = torch.arange(0, args.per_device_train_batch_size).long()
      if torch.cuda.is_available():
        targets = targets.to('cuda')

      sim_scores = F.log_softmax(sim_scores, dim=1)

      loss = F.nll_loss(sim_scores, targets)

      loss.backward()
      optimizer.step()
      scheduler.step()
      q_model.zero_grad()
      p_model.zero_grad()
      global_step += 1

      torch.cuda.empty_cache()


  return p_model, q_model

[Input]

args = TrainingArguments(
    output_dir="dense_retireval",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    num_train_epochs=2,
    weight_decay=0.01
)

p_encoder, q_encoder = train(args, train_dataset, p_encoder, q_encoder)

Sparse Embedding Top-K 결과를 바탕으로 Dense Embedding 하기

[Input]

import torch

batch_size = 64

with torch.no_grad():
    p_encoder.eval()
    q_encoder.eval()

    # Query를 배치로 토큰화하고 임베딩
    all_q_embs = []
    for i in range(0, len(queries), batch_size):
        batch_queries = queries[i:i + batch_size]
        q_seqs_val = tokenizer(batch_queries, padding="max_length", truncation=True, return_tensors='pt').to('cuda')
        q_emb_batch = q_encoder(**q_seqs_val)  # (batch_size, emb_dim)
        all_q_embs.append(q_emb_batch)

    # 배치로 처리한 임베딩들을 하나로 병합
    q_emb = torch.cat(all_q_embs, dim=0)  # (num_query, emb_dim)

    # contexts의 3차원 데이터를 2차원으로 풀어내서 배치 처리
    flattened_contexts = [p for context_list in tfidf_output['contexts'] for p in context_list]
    total_passages = len(flattened_contexts)  # 128 * 50 = 6400
    print(f'Total passages: {total_passages}')  # 6400

    # Passage 임베딩을 배치로 처리
    all_p_embs = []
    for i in range(0, total_passages, batch_size):
        batch_contexts = flattened_contexts[i:i + batch_size]
        p_seqs_val = tokenizer(batch_contexts, padding="max_length", truncation=True, return_tensors='pt').to('cuda')
        p_emb_batch = p_encoder(**p_seqs_val)  # (batch_size, emb_dim)
        all_p_embs.append(p_emb_batch)

    p_embs = torch.cat(all_p_embs, dim=0)  # (total_passages, emb_dim)
    p_embs = p_embs.view(len(tfidf_output['contexts']), -1, p_embs.size(-1))

    print(f'p_embs shape: {p_embs.shape}')  # (128, 50, emb_dim)

✅ 학습된 Encoder를 활용해서 question과 contexts를 임베딩한다. 병렬처리를 위해 contexts를 flatten 후 임베딩하고 원래의 상태로 반환하는 과정을 진행하였다. 임베딩 또한 배치단위로 진행된다.

Dense Embedding을 바탕으로 유사도 측정

[Input]

q_emb = q_emb.unsqueeze(1) # query_embeddings은 (128, 1, 768)이고 passage_embeddings는 (128, 50, 768)

# 1. Cosine Similarity
cosine_similarities = F.cosine_similarity(q_emb, p_embs, dim=-1)  # (128, 50)
# 2. Dot Product Similarity
dot_product_similarities = torch.matmul(q_emb, p_embs.transpose(-1, -2))  
dot_product_similarities = dot_product_similarities.squeeze(1) # (128, 50)

print(dot_product_similarities.shape)  

[Output]

>> (128, 50)

✅ 128개의 question과 TF-IDF로 필터링한 50개의 context사이의 유사도이다.

Dense Embedding Re-ranked Top-K 추출하기

[Input]

reranked_doc_scores = []
reranked_doc_indices = []
k_rerank = 5

# 각 쿼리마다 유사도 기반으로 랭킹을 계산
for i in range(dot_product_similarities.size(0)):  # 쿼리 개수(128)
    # 각 쿼리와 50개의 패시지 간의 유사도를 정렬 (내림차순으로)
    rank = torch.argsort(dot_product_similarities[i], dim=-1, descending=True)
    reranked_doc_scores.append(dot_product_similarities[i][rank][:k_rerank].tolist())
    reranked_doc_indices.append(rank[:k_rerank].tolist())

DPR_output = {
    'question': queries,
    'contexts': [
        [tfidf_output['contexts'][i][idx] for idx in doc_ids]
        for i, doc_ids in enumerate(reranked_doc_indices)
    ],
    'answer': ground_truth,
    'scores': reranked_doc_scores
}

DPR_output_df = pd.DataFrame(DPR_output)
DPR_output_df["correct"] = DPR_output_df.apply(lambda row: row["answer"] in row["contexts"], axis=1)

✅ 이전 스텝(TF-IDF)에서 진행한 것과 같이 128개의 question에 대해 50개의 context 중에 Dense Embedding Similarity로 재정렬된 5개의 context와 socers만 DPR_output에 저장한다.

Dense Embedding Re-ranked Top-K 확인하기

[Input]

k = 5  # 상위 k개의 패시지 출력
query_idx = 0  # 첫 번째 쿼리의 인덱스 

# 첫 번째 쿼리에 대해 상위 k개의 패시지 랭킹을 출력
query = tfidf_output['question'][query_idx]
rank = reranked_doc_indices[query_idx] # 첫 번째 쿼리의 랭킹

print("[Search query]\n", query, "\n")
print("[Ground truth passage]")
print(ground_truth[query_idx], "\n")

# dot_prod_scores는 torch 텐서이므로 CPU로 옮기고 점수 확인
dot_prod_scores_query = dot_product_similarities[query_idx].cpu().numpy()  # 첫 번째 쿼리의 유사도 점수

# 상위 k개의 패시지와 점수 출력
for i in range(k):
    top_idx = rank[i]
    print(f"Top-{i+1} passage with score {dot_prod_scores_query[top_idx]:.4f}")
    print(tfidf_output['contexts'][query_idx][top_idx])
    print()

[Output]

[Search query]
 토너먼트 팔씨름에서 우승한 팀이 영입한 선수는? 

[Ground truth passage]
팀원 트레이드는 더욱 더 강한 팀을 만들기 위한 취지로 도입된 제도로서 트레이드 우선권 순으로 상대팀 에이스 1명과 자신의 팀 1명을 교환할 수 있다. ...

Top-1 passage with score 17.8967
팀원 트레이드는 더욱 더 강한 팀을 만들기 위한 취지로 도입된 제도로서 트레이드 우선권 순으로 상대팀 에이스 1명과 자신의 팀 1명을 교환할 수 있다. ...

Top-2 passage with score 13.7799
2015년 2월 26일 서울 태릉빙상경기장에서 열린 제96회 전국동계체전 대학부 3000m에 출전한 김보름(대구)이 대학부 3000m 4연패를 달성했다. ...

Top-3 passage with score 13.6980
그러나 그는 여기서 물러서지 않았다. 이윽고 펼쳐진 CYON MBC게임 스타리그 승자결승까지 진출한 그는 다시 한 번 지난 리그와 같은 위치에서 마재윤을 만났다. ...

Top-4 passage with score 13.4651
2014년 7월 11일부터 2014년 8월 29일까지 온게임넷을 통해 방송된 리그이다. 진행방식은 토너먼트 방식으로 8강과 3-4위전은 3판 2선승제, 4강과 결승전은 5판 3선승제로 진행되었다. ...

Top-5 passage with score 12.5125
2010년 말 현역 시절 활약하였던 포항 스틸러스의 부름을 받아 포항의 감독으로 부임하였고, 2011 시즌부터 포항을 이끌었다. ...

Re-rank Retriver 정확도 측정

[Input]

print("correct Re-rank retrieval result by exhaustive search", DPR_output_df["correct"].sum() / len(DPR_output_df['contexts']))

[Output]

>> correct Re-rank retrieval result by exhaustive search 0.9306640625

✅ 전체 128개의 Query에 대한 5개의 Retrieved Passage(contexts) 중 정답 Passage(contexts)가 포함될 확률을 나타낸다.

Cross-Encoder Re-rank

기본적으로 STS(Semantic Textual Similarity) Task를 사용해 Query와 Passage사이의 유사도를 구한다.

🚧 추후 정리 해보도록 하자 🚧