LLAMA3.1 임베딩 모델로 변환하기 via. LLM2VEC

Retrieval-Augmented Generation (RAG)는 대규모 언어 모델이 가진 환각 (Hallucination) 현상을 완화하고, 최신 정보나 특정 도메인의 지식을 기반으로 정확하고 신뢰성 높은 답변을 생성할 수 있게 돕는 강력한 프레임워크이다. 이러한 RAG의 성능을 좌우하는 가장 중요한 구성 요소 중 하나가 바로 임베딩 모델이라고 해도 무방하다. 임베딩 모델은 RAG 프레임워크에서 심장과 비슷한 역할을 수행하며, 전체 시스템의 효율성과 정확성에 결정적인 영향을 미치게 된다. 하지만, 일반적인 목적으로 학습된 임베딩 모델은 특정 전문 분야의 용어나 미묘한 문맥 차이를 제대로 파악하지 못하는 경우가 대부분이며, 이를 해결하고자 임베딩 모델을 찾거나 scratch 부터 학습 시키는 것은 굉장히 복잡할 수 있다. 이에 따라 본 게시물에서는 LLM2VEC 논문을 기반으로 LLAMA3.1:8B 모델을 임베딩 모델로 변환하는 과정에 대해 다루고자 한다.

---

거대 언어 모델과 임베딩 모델: 학습 방식의 근본적인 차이

앞서 LLM이 어떻게 학습되는 지에 대해 간략하게 알아보고자 한다.  LLM은 일반적으로 Causal Language Modeling Loss로 학습이 된다. 주어진 토큰 시퀸스를 바탕으로 다음 토큰을 예측하도록 학습을 하게 된다. 이러한 학습 방식을 Causal 이라고 하는 이유는 모델이 다음 단어를 예측할 때, 오직 이전에 등장한 단어들만을 참고하고 이후의 단어들은 참고할 수 없기 때문이다. 마치 시간은 과거에서 미래로만 흐르는 것처럼, 정보의 흐름 또한 한쪽으로만 흐르는 것으로 생각하면 좋을거 같다. 

 

위 내용에 대해 조금 더 구체적으로 설명을 하자면, 모델이 특정 토큰을 예측하는 것을 학습할 때 해당 토큰 뒤에 오는 모든 토큰들이 Masking 되어 어텐션 계산에 영향을 주지 않도록 시퀸스에 Causal Attention Mask가 적용이 되는 것이다. 예를 들어, 다음과 같은 시퀸스가 있다고 가정해보자. 

The dog is sleeping in the kitchen <eos>

 

학습 중에는 다음과 같은 어텐션 마스크가 순차적으로 적용이 되며, 0 같은 경우 해당 토큰이 어텐션 계산에 고려되지 않음을 의미한다. 

 

1 0 0 0 0 0 0 0 0
The

1 1 0 0 0 0 0 0
The cat

1 1 1 0 0 0 0 0
The cat is

1 1 1 1 0 0 0 0
The cat is sleeping

1 1 1 1 1 0 0 0
The cat is sleeping in

1 1 1 1 1 1 0 0
The cat is sleeping in the

1 1 1 1 1 1 1 0
The cat is sleeping in the kitchen

1 1 1 1 1 1 1 1
The cat is sleeping in the kitchen <eos>

 

 

이러한 과정에서 모델이 예측한 단어와 실제 정답 단어 사이의 차이를 계산하는 데, 이 차이를 손실(Loss)라고 칭한다. LLM의 학습 목표는 수많은 텍스트 데이터를 통해 이러한 손실 값을 지속적으로 줄여나가는 것이다. 어찌보면 굉장히 단순해 보이는 다음 단어  예측하기 과정을 수조 번 반복하면서, LLM은 단어와 단어 사이의 복잡한 관계, 문법, 문맥, 그리고 세상의 방대한 사실적 기반의 지식까지 스스로 터득하게 된다. 

 

이처럼 대부분의 LLM은 이전 단어를 바탕으로 다음 단어를 예측하는 단방향(Autoregressive) 구조로 작동하게 된다. 반면, 본 게시물의 핵심 주제인 임베딩 모델은 문장 전체의 의미를 파악/압축하기 위해 양방향(Bidirectional) 구조를 사용하게 된다. 즉, 문장을 왼쪽에서 오른쪽으로, 또 오른쪽에서 왼쪽으로 양쪽에서 모두 읽어 문맥을 종합적으로 이해하는 것을 목표로 하는 것이 임베딩 모델 (Embedding Model)이다.

---

LLM2VEC 논문의 제안 방법

LLM2VEC 세 가지의 방법론 (이미지 출처: LLM2VEC 논문)

 

1. Restrictions of Causal Attention 

LLM을 임베딩 모델로 변환하는 LLM2VEC 논문에서 제안한 가장 첫번 째 방법은 Decoder-Only LLM에서 Causal Attention Mask의 모든 값을 1인 행렬로 대체하는 것이다.

 

좀 더 구체적으로, Causal Attention Mask는 수학적으로 하삼각행렬 (Lower Triangular Matrix) 형태로 구현이 된다. 이러한 마스크는  어텐션 스코어를 계산할 때, 현재 위치 i의 토큰은 미래 위치 j > i 토큰을 참조하지 못하도록 해당 어텐션 스코어 값을 음의 무한대로 설정한다. 이후 Softmax 함수를 거치면서 이 값은 0이 되고, 결과적으로는 미래 토큰의 정보가 현재 토큰의 표현(Representation)이 어텐션 계산에 반영되는 것을 차단하는 것이 Causal Attention Mask의 원리이다. 

 

LLM2VEC 저자들은 Causal Attention Mask를 제거하고 모든 값이 1인 행렬(즉, 마스킹이 없는 상태)로 대체한다. 이로 인해 토큰 i는 시퀸스 내의 모든 토큰과 자유롭게 어텐션 스코어를 계산할 수 있게 된다. 이러한 구조 변경으로 인해 사실상 모델은 양방향 (Bidirectional) 기반의 LLM으로 변환이 된다. 하지만, 단방향 (Autoregressive)로 학습이 된 Decoder-Only LLM의 구조를 단숨에 변경하여, 미래의 토큰을 인코딩하게 된다면 성능이 크게 하락된다. 이에 대한 가장 큰 이유는 가중치(weight)의 불일치와 데이터 분포 이동 (Distribution Shift)으로 인해 발생한다. 

Distribution Shift (이미지 출처: NannyML)

• 가중치 불일치: LLM의 모든 가중치 (e.g. Attention Layer, Feedforward Network, etc.)는 수조 개의 단어를 통해 단방향 문맥(unidirectional context)을 입력받아 다음 단어를 예측하는 Causal Language Model Loss Function에 최적화가 되어 있는데, 갑자기 양방향(Bidirection)으로 변경이 된다면... "머지? 내가 알았던 내용이 아닌데?" 라는 생각을 하게 된다. 
• 데이터 분포 이동: 마스크를 제거하게 된다면, 각 레이어에 입력되는 Hidden State들의 분포가 완전히 달라지게 된다. 이전에는 과거 정보만으로 구성되었던 값들이 이제는 과거와 미래 정보가 혼합된, 즉 모델이 한번도 본 적이 없는 새로운 데이터 분포의 입력으로 바뀌게 되어, 모델이 "이거 아닌거 같은데"? 라는 생각을 하게 된다. 

이를 해결하기 위해 MNTP을 사용하는 것을 제안하였다. 

---

2. Masked Next Token Prediction Training (MNTP) 

 

이미지 출처: MNTP 논문

 

Masked Next Token Prediction (MNTP)는 단순하게 다음 토큰 예측과 BERT에서 사용되는 마스킹된 언어 모델을 결합한 것을 의미한다. 이를 구현하기 위해, 임의의 시퀸스를 가져와 입력 토큰의 일부를 마스킹한 다음, 과거와 미래의 문맥을 모두 사용하여 모델이 이 마스킹된 토큰들을 예측하도록 학습시킨다. 여기에서 중요한 점은 위치 i에서 마스크된 토큰을 예측할 때, BERT 모델에서 마스킹된 토큰 위치 자체의 로짓(Logit)을 사용하는 것이 아니라, 그 이전 위치 i-1의 토큰 표현으로부터 얻은 로짓을 사용하여 손실을 계산하는 것이다. 

 

위 문장을 읽게 되면, 약간 헷갈릴 수 있을거 같아 다음과 같은 예시를 준비하였습니다. 예를 들어, brown을 맞추는 상황을 가정해보면, 원래의 시퀸스에서의 brown은 [MASK]로 아래와 같이 변환이 된다. 

A quick brown fox jumps --> A quick [MASK] fox jumps

• BERT의 방식: BERT는 3번째 위치, 즉 [MASK] 토큰 자체의 최종적인 표현을 가지고, 이 빈칸에 들어갈 가장 적절한 단어가 무엇인지에 대해 풀게 된다. 
• MNTP의 방식: 반면. MNTP는 2번째, 즉 [quick]이라는 토큰의 최종적인 표현을 사용하여, 모델이 양쪽의 모든 문맥을 고려했을 때, [quick] 다음에 올 단어가 무엇인지에 대해 풀게 되는것과 같다. 쉽게 말하자면, quick의 표현으로 brown을 예측하는 것이다. 

이처럼 MNTP는 마스킹된 바로 이전 토큰의 표현에서 로짓을 계산하여 손실을 구하게 된다. 이러한 방식은 모델이 양방향 문맥 고려할 뿐만 아니라, 기존 Decoder-Only LLM이 훈련하였던 다음 토큰 예측 메커니즘을 유지하는 기법이다. 이러한 이점을 가진 MNTP을 통해 새로운 목표를 기준으로 손실을 계산하고, 역전파를 통해 손실을 줄이는 방향으로 그래디언트를 계산한다. 이에 최종적으로 기존 Causal LM에 최적화되어 있던 가중치들은 점차 bidirection에 적합한 값으로 수렴하게 되는 것이다. 

 

MNTP 학습을 마친 모델은 시퀸스의 양방향 기반으로 문맥을 잘 이해하게 되었지만, 문장 전체의 의미를 하나의 벡터로 압축 시키는 능력은 여전히 부족할 수 있다. 이에 본 논문의 저자들은 SimCSE 방식을 사용하여, 이와 같은 문제점을 해결했다고 한다. 

---

 

3. Unsupervised Constrastive Learning with Mean Pooling 

이 학습 과정의 핵심적인 목표는, 같은 의미의 문장은 가깝게, 다른 의미의 문장은 멀게 배치하여 벡터 공간을 의미적으로 정리하는 것과 같다. 구체적으로, 모델은 동일한 문장은 서로 다른 "Dropout Mask"를 적용하여 두 번 처리함으로써, 미세하게 다른 두 개의 표현 벡터를 생성하게 된다. 특이하게, 이 드롭아웃의 본래 목적인 과적합 방지를 넘어, 데이터 증강의 역할을 수행하게 된다. 예를 들어, 문장 "고양이가 귀엽다"가 모델에 입력이 되어, 모델 내부의 특정 레이어를 통과하게 되면서 드롭아웃 마스크가 적용되어 두 개의 벡터가 생성이 된다. 결과적으로 놓고 보면, 같은 문장을 입력하였지만 무작위로 선택된 뉴런에 의해 최종 결과물인 벡터 A와 벡터 B는 미세하게 다른 값을 가지게 되는 것이다. 이렇게 만들어진 A와 B를 같은 의미를 가진 데이터로 간주하여, 이 둘을 최대한 가깝게 만들도록 학습하게 된다. 즉, 학습의 목표는 두 벡터 간의 유사도를 높이는 동시에, 동일한 학습 배치 내에 있는 다른 시퀸스의 표현과의 유사도를 낮추는 것이다. 본 논문의 저자들은 이러한 과정을 Unsupervised Constrastive Learning (SimCSE)이라고 한다. 

 

---

 

LLAMA 3.1:8B 임베딩 모델로 변환시키는 법 

생각보다 간단하지만 A100-40GB 기준으로 거의 65시간이 소요되므로, 이에 적절하게 precision을 설정해야한다.

git clone https://github.com/McGill-NLP/llm2vec.git

import torch
from llm2vec import LLM2Vec

l2v = LLM2Vec.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    device_map="cuda" if torch.cuda.is_available() else "cpu",
    torch_dtype=torch.bfloat16,
)

l2v.save("Llama-3-8B-Emb")

JSON_CONFIG='''
{
    "model_name_or_path": "meta-llama/Meta-Llama-3-8B",
    "dataset_name": "wikitext",
    "dataset_config_name": "wikitext-103-raw-v1",
    "per_device_train_batch_size": 1,
    "per_device_eval_batch_size": 1,
    "gradient_accumulation_steps": 16,
    "do_train": true,
    "do_eval": true,
    "max_seq_length": 512,
    "mask_token_type": "blank",
    "data_collator_type": "all_mask",
    "mlm_probability": 0.8,
    "overwrite_output_dir": true,
    "output_dir": "Llama-3-8B-llm2vec-MNTP-Emb",
    "evaluation_strategy": "steps",
    "eval_steps": 100,
    "save_steps": 200,
    "stop_after_n_steps": 1000,
    "lora_r": 16,
    "gradient_checkpointing": true,
    "torch_dtype": "bfloat16",
    "attn_implementation": "flash_attention_2"
}
'''

with open("mntp_config.json", 'w') as f:
  f.write(JSON_CONFIG)

python llm2vec/experiments/run_mntp.py mntp_config.json