[LLM] RAG 개선하기 (1) - 언어 모델을 임베딩 모델로 만들기

이전 포스팅에서는 두 문장 간 유사도를 비교하는 방법 두 가지인 바이 인코더와 교차 인코더 중 비교적 구현이 간단한 바이 인코더에 대해서 알아보았다. 하지만 교차 인코더의 정확도도 포기할 수 없으므로, 두 인코더의 장점을 살려 더 높은 성능의 검색을 구현하는 방법에 대해 알아보자. 또 기존의 문장 임베딩 모델을 가져와서 사용하기만 하는 것이 아니라, 이번에는 순수의 BERT 모델을 가져와 문장 임베딩 모델로 학습시켜 보자.

---

검색 성능을 높이기 위한 두 가지 방법

 

허깅페이스 모델 허브에서 사전 학습된 문장 임베딩 모델을 가져와 그대로 사용할 경우, 우리의 RAG가 가진 데이터셋으로 학습을 한 것이 아니므로 정확도가 떨어질 수 있다. 이를 극복하기 위해 할 수 있는 두 가지 방법을 알아보자.

 

데이터 추가 학습하기

말 그대로 추가 데이터셋으로 학습을 더 시키는 것이다. 기존의 데이터로 학습된 모델에 이렇게 추가로 학습을 시키면 우리가 타깃으로 하는 데이터에 맞는 성능으로 어느 정도 튜닝이 가능할 것이다.

 

교차 인코더 추가하기

바이 인코더로 검색된 유사도가 높은 문서나 문장의 상위 몇 가지를 추려, 이 데이터들에 대해 교차 인코더를 사용하는 것이다. 이렇게 한다면 유사도가 높은 데이터들로 1차 필터링이 가능하고, 그 뒤에 남아있는 데이터셋은 그 양이 많지 않아 모든 조합에 대한 유사도 계산을 하더라도 시간이나 속도가 현저히 줄지는 않는다.

 

언어 모델을 임베딩 모델로 만들기

 

대조 학습 Contrastive Learning

 

대조 학습은 관련이 있거나 유사한 데이터는 벡터 상에서 더 가까운 곳에 위치시키고, 관련이 없거나 유사하지 않은 데이터는 더 멀어지도록 학습하는 방식을 의미한다. 대조 학습에 넣을 데이터는 일반적으로 서로 유사한 데이터를 수집한 데이터셋이나 서로 이어지는 문장을 수집하여 학습에 사용한다.

 

기존의 사전 학습된 모델 사용하기

 

우리에게 최적화된 모델을 구축하기 전에 앞서 기존의 언어 모델을 가져와 문장 임베딩 모델을 만들어 보자.

from sentence_transformers import SentenceTransformer, models
transformer_model = models.Transformer('klue/roberta-base')

pooling_layer = models.Pooling(
    transformer_model.get_word_embedding_dimension(),
    pooling_mode_mean_tokens=True
)
embedding_model = SentenceTransformer(modules=[transformer_model, pooling_layer])

 

roberta 모델을 불러오고 평균 모드를 사용하는 풀링 층을 생성하였다. 이어서 이 둘을 SentenceTransformer에 넣어 문장 임베딩 모델을 만들었다.

 

from datasets import load_dataset
klue_sts_train = load_dataset('klue', 'sts', split='train')
klue_sts_test = load_dataset('klue', 'sts', split='validation')
klue_sts_train[0]

# {'guid': 'klue-sts-v1_train_00000',
#  'source': 'airbnb-rtt',
#  'sentence1': '숙소 위치는 찾기 쉽고 일반적인 한국의 반지하 숙소입니다.',
#  'sentence2': '숙박시설의 위치는 쉽게 찾을 수 있고 한국의 대표적인 반지하 숙박시설입니다.',
#  'labels': {'label': 3.7, 'real-label': 3.714285714285714, 'binary-label': 1}}

 

이렇게 만든 모델로 테스트를 해보자. KLUE의 STS 데이터셋을 가져와서 훈련, 테스트 데이터셋을 받아 온다.

 

# 학습 데이터셋의 10%를 검증 데이터셋으로 구성한다.
klue_sts_train = klue_sts_train.train_test_split(test_size=0.1, seed=42)
klue_sts_train, klue_sts_eval = klue_sts_train['train'], klue_sts_train['test']

 

그중 훈련 데이터셋의 10%를 검증 데이터셋으로 구성하고, 훈련 데이터와 검증 데이터를 새롭게 정의해 준다.

 

from sentence_transformers import InputExample
# 유사도 점수를 0~1 사이로 정규화 하고 InputExample 객체에 담는다.
def prepare_sts_examples(dataset):
    examples = []
    for data in dataset:
        examples.append(
            InputExample(
                texts=[data['sentence1'], data['sentence2']],
                label=data['labels']['label'] / 5.0)
            )
    return examples

train_examples = prepare_sts_examples(klue_sts_train)
eval_examples = prepare_sts_examples(klue_sts_eval)
test_examples = prepare_sts_examples(klue_sts_test)

 

다음으로 해줄 과정은 유사도 점수의 정규화이다. 이건 우리가 사용하는 데이터셋의 label 점수가 0 ~ 5까지의 값으로 이루어져 있기 때문에 모든 label 점수를 0 ~ 1로 정규화해 주는 것뿐이다.

 

from torch.utils.data import DataLoader
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)

 

이렇게 전처리된 훈련 데이터를 16의 배치 사이즈로 만든다.

 

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

eval_evaluator = EmbeddingSimilarityEvaluator.from_input_examples(eval_examples)
test_evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_examples)

 

이제 EmbeddingSimilarityEvaluator 라이브러리를 통해 우리가 전처리한 eval_examples와 test_examples를 평가할 수 있는 클래스를 선언해 준다. 이 값들에 앞서 생성한 embedding_model을 입력해 주면, 각각의 데이터셋을 가지고 임베딩 모델이 얼마나 문장의 의미를 잘 반영해 문장 임베딩을 생성하는지 0부터 1까지의 값으로 알 수 있다.

 

test_evaluator(embedding_model)
# 0.36460670798564826

 

이렇듯 기존의 사전 학습된 모델을 그냥 사용한다면 0.365의 낮은 성능을 지닌다는 것을 알 수 있다.

 

유사한 문장 데이터로 임베딩 모델 학습하기 (데이터 추가 학습하기)

이제 주어진 데이터셋과 사전 학습된 모델을 이용하여 튜닝을 진행해 보자.

from sentence_transformers import losses

num_epochs = 4
model_name = 'klue/roberta-base'
model_save_path = 'output/training_sts_' + model_name.replace("/", "-")
train_loss = losses.CosineSimilarityLoss(model=embedding_model)

# 임베딩 모델 학습
embedding_model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    evaluator=eval_evaluator,
    epochs=num_epochs,
    evaluation_steps=1000,
    warmup_steps=100,
    output_path=model_save_path
)

 

임베딩 모델에 .fit() 메서드를 통해 어떤 방식으로 학습을 진행할지 설정할 수 있다. 훈련 데이터로더, 손실 함수, 평가자, epoch 수 등등을 조절할 수 있다.

 

trained_embedding_model = SentenceTransformer(model_save_path)
test_evaluator(trained_embedding_model)
# 0.8965595666246748

 

결과는 다음과 같다. 위 학습을 통해 만들어진 모델을 불러오고, 테스트를 해보면 0.897의 높은 성능을 보임을 증명할 수 있었다!