[딥러닝] LLaMA2 모델 정리(Few-Shot Learning, Instruction Tuning, RLHF) + Constitutional AI

- 이번 포스팅에서는 LLaMA2 모델을 분석하겠습니다.

기존 딥러닝 모델 분석 게시글들과 다른점으로는 모델 구조적인 내용보다는 Few-shot Learning, Instruction Tuning, RLHF 와 같은 LLM 모델의 효과적인 학습 기법들을 중점적으로 정리하고 이해할 것입니다.

해당 방식은 이후 나오는 다른 모델들의 학습 방식에서도 중요하게 활용되는 것이기에 제대로 이론을 정리하고 코드를 작성하도록 하겠습니다.

 

(LLaMA2 모델 설명)

- 앞서 정리한 LLaMA1 정리글을 기반으로 변경된 부분만을 빠르게 알아보겠습니다.

 

- 외부 변경사항

1. LLaMA1 의 파라미터는 7B, 13B, 33B, 65B 이지만 LLaMA2 에서는 7B, 13B, 70B 이렇게 최대 파라미터가 증량되었습니다.

2. 훈련 데이터는 LLaMA1 에서 사용한 훈련 데이터에 더하여 고품질 웹 데이터를 추가하였습니다.

3. 훈련 방식은 기존 방식과 더불어 Instruction tuning 기법, RLHF 기법이 추가되었습니다.

4. 훈련 데이터에 유해 콘텐츠 필터링을 강화하였습니다.

5. 성능은 GPT-3 보다는 높고, GPT-3.5 보다는 약간 낮은 정도입니다.

6. LLaMA1 에 비해 모델 가중치, 토크나이저 등의 상세 내용이 공개되었고, 파인 튜닝 및 커스텀 튜닝 가이드와 툴까지 제공하는 등 적극적인 오픈소스 정책을 펼쳤습니다.

 

- 구조 변경사항

1. 일부 레이어에서 드롭아웃을 제외하였습니다.

2. FFN Hidden Dim 비율을 기존의 2배(SwiGLU 를 적용하기 위해)에서 증가하여 4배 이상을 사용합니다.

 

- LLaMA2 의 변경사항은 위와 같습니다.

보시다시피 구조는 거의 변함이 없으며, 몇가지 파라미터나 세부사항에 대한 조정이 이루어졌고,

대부분은 학습 방식과 학습 기법의 변화입니다.

 

이로써 딥러닝 연구에 대한 흐름은 획기적인 모델이 발명되지 않는 이상 그 당시의 최신 기술들을 사용하여 학습 방식을 바꿔가며 시도하는 것을 반복하는 방식으로 이루어지는 것을 알 수 있습니다.

앞으로 첨단에 닿을 때 까지 지름길을 통해 계속해서 흐름을 따라가보도록 하죠.

 

(Few-Shot Learning)

- LLM 에서 Few-shot learning (소량 학습)은 모델이 학습된 이후 적은 수의 예시만으로 새로운 작업을 수행할 수 있는 능력을 말하는 것입니다.

 

GPT 시리즈와 같은 대규모 언어 모델에서 강력하게 나타나는 특징으로,

Few-shot 뿐만 아니라 학습 샘플이 몇개가 필요한지에 따라 종류가 나뉘는데,

클래스마다 보통 2~10 사이의 샘플이 필요하다면 Few-Shot learning,

클래스마다 1개의 샘플만이 필요하다면 One-shot learning,

예제 자체가 없이 설명만으로 일반화하는 것을 zero-shot learning 이라고 합니다.

 

- Few-Shot Learning 은 파인튜닝과 같은 파라미터 학습이 아니라 프롬프트엔지니어링의 영역입니다.

예를들어 설명하겠습니다.

 

"주어진 나라의 수도를 대답하라"

라는 태스크가 주어졌을 때,

Q: What is the capital of France?
A:

 

zero shot learning 은 위와 같이 질문만 던지는 것이고,

Q: What is the capital of Germany?
A: Berlin

Q: What is the capital of France?
A:

 

one shot learning 은 하나만,

Q: What is the capital of Germany?
A: Berlin

Q: What is the capital of Italy?
A: Rome

Q: What is the capital of France?
A:

 

few shot learning 은 여러개를 던져서 LLM 이 반환해야 하는 답변의 형식을 앞선 문장 속에서 파악하여 답변하게 하는 것입니다.

 

위와 같은 개념이 가능한 이유는,

사전학습을 거친 범용 LLM 모델은, 학습 당시 엄청난 양의 텍스트 데이터에서 확률 기반의 패턴(통계적 연관)을 학습한 모델로써,
프롬프트에 반복적으로 제시되는 형식과 의미의 관계를 보고 규칙처럼 작동하는 패턴을 따르며, 결과적으로 언어적 패턴을 인식할 능력이 있기 때문입니다.

 

즉, LLM 에 문맥 기반 예측 기능이 있기 때문이죠.

 

- LLM 의 Few-shot learning 은 그저 정확도 보정을 위해서뿐 아니라 사용자가 원하는 형식의 데이터를 LLM 모델이 생성해내도록 제약하는 기법으로도 사용할 수 있기 때문에 활용도가 높은 개념입니다.

 

(Instruction Tuning)

- Instruction Tuning 은 파인튜닝 기법입니다.

앞서 딥러닝 정리 게시글을 보셨다면 이해하실 수 있듯,

생성형 AI 모델의 사전 학습은 사람으로 치자면 책을 읽고 그 내용을 전부 외워서 다시 낭독함으로써 언어의 의미를 이해시키는 방식입니다.

 

자연어의 의미를 이해하고, 보다 출현 확률이 높은 단어들을 차례로 추론할 수 있는 것이죠.

하지만 이 상태만으로는 사용자의 '지시'에 대한 학습이 되어있지 않은 상태입니다.

 

예를들어,

"산넘고 물건너 만신창이로 도착한 곳에는 어느새 봄이 있었다."

대부분 이러한 문장으로 학습된 모델은,

"다음 뉴스를 읽고 요약해주세요 <뉴스 전문>. 답변 : <답변 전문>"

이러한 목적성을 지닌 지시와 답변에는 익숙하지 않다는 것입니다.

 

"요약해줘", "의견을 말해줘", "질문에 답해줘"

이러한 유형의 요청과, 그 요청에 대한 답변을 전문적으로 훈련받지 못하고, 외워둔 방대한 문장을 출력하는 모델에게 비교적 적은 양의 데이터셋으로 질문에 대한 해석과 답변 방식을 학습시키는 것이 Instruction Tuning 입니다.

 

- Instruction Tuning 은 파인튜닝 시점에서 지시문과 그 지시문에 대한 답변들 학습시켜서 목적을 수행할 수 있도록 훈련시키는 것을 의미합니다.

 

예를들어보겠습니다.

자주 사용되는 Super-NaturalInstructions 데이터셋의 경우,

{
  "instruction": "주어진 문장을 영어에서 한국어로 번역하세요.",
  "input": "The weather is nice today.",
  "output": "오늘 날씨가 좋습니다."
}

 

이런식으로 데이터가 존재합니다.

 

이 데이터를 가지고,

"
Instruction: 주어진 문장을 영어에서 한국어로 번역하세요.
Input: The weather is nice today.
Output: 오늘 날씨가 좋습니다.
"

 

이런식으로 입력하여 파인튜닝을 했을 때,

모델 사용 시점에,

prompt = f"Instruction: {instruction}\nInput: {input_text}\nOutput:"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs)

 

이런식으로 학습된 데이터의 형식에 맞게 프롬프트를 입력한다면 파인튜닝한 지시문 형식을 파악하여 그 법칙에 따라 응답값을 반환하게 하는 것입니다.

 

- 사전 학습된 성능 좋은 범용 언어 모델을 실용적으로 활용하기 위한 파인튜닝 시점에 꼭 필요한 학습 방식입니다.

대표적인 오픈 데이터셋으로는,

FLAN, Super-NaturalInstructions, Self-Instruct, TO, PromptSource 등이 있습니다.

 

(RLHF( Reinforcement Learning from Human Feedback, 인간 피드백 기반 강화학습))

- RLHF 는 기계 학습 모델, 특히 언어 모델을 더 사람처럼 유용하고 안전하게 만드는 데 사용되는 학습 기법입니다.

이 방법은 인간이 직접 제공한 피드백을 바탕으로 모델의 출력을 강화학습을 통해 조정하는 방식입니다.

쉽게 말하자면 사람이 직접 관여하는 파인튜닝 방법론을 의미합니다.

 

- RLHF 는 아래와 같은 순서로 진행됩니다.

1. 지도 학습

사람 전문가가 직접 작성한 고품질 응답 예시를 사용하여 파인튜닝을 진행합니다.

사전 학습이 완료된 직후 질문에 대한 답변을 배우는 단계이며, 모델의 기본적인 지식과 응답 형식을 배우는 과정입니다.

앞서 배운 Instruction Tuning 을 이곳에서 적용하는 것도 일반적인 방법입니다.

 

2. 보상 모델 훈련

RLHF 의 강화학습에 적용할 보상 모델을 훈련합니다.

보상 모델은 프롬프트 + 응답을 받아서 이것에 대한 실수 형태의 품질 점수를 출력하는 모델입니다.

품질이 높을수록 프롬프트에 대한 응답이 적절하다는 의미죠.

 

이 모델에 대한 학습 방법은,

(prompt, chosen_response, rejected_response) 쌍이 주어졌을 때,

 

r_chosen = reward_model(prompt, chosen_response)
r_rejected = reward_model(prompt, rejected_response)

 

이런식으로 각 응답별 점수를 예측합니다.

 

그리고,

 

-log(σ(r_chosen - r_rejected))

 

이런 형태의 Pairwise Ranking Loss(chosen 이 rejected 보다 더 높은 점수를 받도록 유도) 를 사용하여 학습을 진행합니다.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

from transformers import AutoTokenizer, AutoModel

class RewardDataset(Dataset):
    def __init__(self, data, tokenizer, max_length=512):
        self.data = data  # list of dicts with 'prompt', 'chosen', 'rejected'
        self.tokenizer = tokenizer
        self.max_length = max_length

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

    def __getitem__(self, idx):
        item = self.data[idx]
        prompt = item['prompt']
        chosen = item['chosen']
        rejected = item['rejected']

        # concatenate prompt + response
        chosen_input = self.tokenizer(prompt + chosen, truncation=True, padding="max_length", max_length=self.max_length, return_tensors="pt")
        rejected_input = self.tokenizer(prompt + rejected, truncation=True, padding="max_length", max_length=self.max_length, return_tensors="pt")

        return {
            "chosen_input": {k: v.squeeze(0) for k, v in chosen_input.items()},
            "rejected_input": {k: v.squeeze(0) for k, v in rejected_input.items()}
        }

class RewardModel(nn.Module):
    def __init__(self, base_model_name="bert-base-uncased"):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(base_model_name)
        self.value_head = nn.Linear(self.encoder.config.hidden_size, 1)  # outputs scalar

    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls_output = outputs.last_hidden_state[:, 0]  # [CLS] token
        reward = self.value_head(cls_output).squeeze(-1)
        return reward  # shape: (batch_size,)

def pairwise_ranking_loss(r_chosen, r_rejected):
    return -torch.nn.functional.logsigmoid(r_chosen - r_rejected).mean()

def train_reward_model(model, dataloader, epochs=3, lr=1e-5):
    optimizer = optim.AdamW(model.parameters(), lr=lr)
    model.train()

    for epoch in range(epochs):
        total_loss = 0.0
        for batch in dataloader:
            chosen_input = batch["chosen_input"]
            rejected_input = batch["rejected_input"]

            r_chosen = model(chosen_input["input_ids"].to(device), chosen_input["attention_mask"].to(device))
            r_rejected = model(rejected_input["input_ids"].to(device), rejected_input["attention_mask"].to(device))

            loss = pairwise_ranking_loss(r_chosen, r_rejected)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        print(f"Epoch {epoch+1}, Loss: {total_loss:.4f}")

if __name__ == "__main__":
    # 샘플 데이터
    dummy_data = [
        {
            "prompt": "What is the capital of France?",
            "chosen": " The capital of France is Paris.",
            "rejected": " I don't know, maybe Berlin?"
        },
        # 추가 데이터...
    ]

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    dataset = RewardDataset(dummy_data, tokenizer)
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

    model = RewardModel().to(device)
    train_reward_model(model, dataloader)

 

예시 코드는 위와 같습니다.

중요한 점만 짚어보자면,

Dataset 에 들어가는 데이터는 prompt, chosen, rejected 의 분류로 쌍을 이루고,

보상 모델은 자연어를 해석해서 인코딩 해야하므로 기본 모델을 BERT 로 하여 Scalar 값을 반환하는 선형 레이어를 붙인 형태이며,

앞서 설명한 방식으로 학습을 진항하도록 한 것입니다.

 

3. 강화 학습

위에서 학습한 보상 모델을 활용하여 언어 모델을 학습시키는 단계입니다.

옳고 그름을 판단할 수 있는 보상 모델이 준비된 상태이므로, 이 모델과 더불어 프롬프트만 존재한다면 따로 정답 라벨을 만들 필요 없이 보상 모델의 반환값을 정답값으로 하여 모델을 학습 시킬 수 있습니다.

고로 강화학습입니다.

 

위와 같은 컨셉만 이해한다면 구현은 어떤 방식이든 가능할텐데,

강화학습에서 널리 쓰이는 PPO 기법에 대해 알아보고 적용하는 방식을 정리하겠습니다.

 

PPO(Proximal Policy Optimization) 는 OpenAI 가 제안한 강화학습에서 매우 널리 사용되는 알고리즘 중 하나로,

간단히 설명하자면, 한번에 너무 과하게 움직이지 말고 천천히 바르게 조정해나가는 방식으로 이해하면 됩니다.

이에 대한 자세한 내용은 별도 게시글로 떼어내기로 하고, 예시 코드만 보겠습니다.

import torch
import torch.nn as nn
import torch.optim as optim
from transformers import AutoModelForCausalLM, AutoTokenizer

# 정책 모델 (GPT류 모델)
class PolicyModel(nn.Module):
    def __init__(self, model_name="gpt2"):
        super().__init__()
        self.model = AutoModelForCausalLM.from_pretrained(model_name)

    def forward(self, input_ids, attention_mask):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        return outputs.logits

    def generate(self, input_ids, max_length=50):
        return self.model.generate(input_ids, max_length=max_length, do_sample=True)

# PPO용 로그 확률 계산
def get_log_probs(logits, labels):
    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
    selected = log_probs.gather(2, labels.unsqueeze(-1)).squeeze(-1)
    return selected

# PPO 손실 함수 (clip 방식)
def ppo_loss(old_log_probs, new_log_probs, advantages, epsilon=0.2):
    ratio = torch.exp(new_log_probs - old_log_probs)
    clipped = torch.clamp(ratio, 1 - epsilon, 1 + epsilon)
    loss = -torch.min(ratio * advantages, clipped * advantages)
    return loss.mean()

# 샘플 강화학습 루프
def train_with_ppo(policy_model, reward_model, tokenizer, prompts, epochs=3, gamma=0.99):
    optimizer = optim.AdamW(policy_model.parameters(), lr=1e-5)
    policy_model.train()
    reward_model.eval()

    for epoch in range(epochs):
        for prompt in prompts:
            # 1. 인코딩 및 응답 생성
            inputs = tokenizer(prompt, return_tensors="pt").to(device)
            generated = policy_model.generate(inputs['input_ids'], max_length=50)
            
            # 2. 보상 평가
            with torch.no_grad():
                reward = reward_model(
                    input_ids=generated,
                    attention_mask=(generated != tokenizer.pad_token_id).long()
                )
                reward = reward.item()

            # 3. 이전 로그 확률 계산
            logits = policy_model(inputs['input_ids'], inputs['attention_mask'])
            old_log_probs = get_log_probs(logits[:, :-1], inputs['input_ids'][:, 1:])

            # 4. 현재 로그 확률 (for PPO 업데이트)
            new_logits = policy_model(inputs['input_ids'], inputs['attention_mask'])
            new_log_probs = get_log_probs(new_logits[:, :-1], inputs['input_ids'][:, 1:])

            # 5. Advantage 계산 (여기선 보상값만 사용)
            advantage = torch.tensor(reward).to(device)

            # 6. PPO 손실 계산 및 역전파
            loss = ppo_loss(old_log_probs, new_log_probs, advantage)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f"[Epoch {epoch+1}] Loss: {loss.item():.4f}")

# 실행 준비
if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    policy_model = PolicyModel().to(device)
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token  # GPT2는 pad_token이 없어서 설정 필요

    # RewardModel은 앞서 훈련한 모델 사용
    reward_model = RewardModel().to(device)
    reward_model.eval()

    # 간단한 프롬프트들
    prompts = [
        "Explain the process of photosynthesis.",
        "What is the capital of Japan?",
        "How do airplanes fly?"
    ]

    # PPO 기반 훈련
    train_with_ppo(policy_model, reward_model, tokenizer, prompts)

 

위 코드는 학습을 시킬 대상인 정책 모델을 앞서 작성했던 보상 모델을 이용하여 강화학습하는 예시 코드인데,

학습을 진행시 정책 모델에서 반환된 응답값을 보상 모델에 입력하여 reward 를 받고,

급격한 변화를 막는 처리가 되어있는 ppo_loss 로 모델을 학습시킨다고 이해하면 됩니다.

 

(Constitutional AI)

- 추가로, LLaMA 가 아닌 Claude 모델에서 RLHF 를 대신하여 적용한 Constitutional AI 기법을 정리하겠습니다.

Constitutional AI 란, 인간의 피드백을 사용하는 RLHF 와 달리, 사전에 정의된 윤리적 원칙과 규칙(=헌법, Constitution)에 따라 모델이 스스로 자신의 출력을 평가하고 수정하도록 유도하는 학습 방법입니다.

 

- 학습 프로세스는 아래와 같습니다.

1. Supervised Fine-Tuning(SFT)

RLHF 에서 본 지도 학습과 동일합니다.

 

2. AI Self-Critique

AI 의 자가 비판 단계입니다.

악의적인 답변을 요구하는 프롬프트를 준비하여 LLM 에게 요청하고,

그 답변을 샘플링합니다.

 

예를들어,

"사람을 다치게 하는 방법을 알려줘"

이렇게 입력했을 때,

아직 제약되지 않은 LLM 이라면 유해한 답변을 필터 없이 반환할 것입니다.

이러한 답변들을 모으는 것입니다.

 

다음으로는 위와 같은 프롬프트와 통제되지 않은 응답을 모델 스스로가 평가하게 합니다.

["사람을 다치게 하는 응답을 내면 안된다", "욕설을 적으면 안된다"]

위와 같이 LLM 이 지켜야하는 제약사항을 적은 헌법(Constitution)을 기준삼아,

 

"<프롬프트> 이런 요청이 있을 때, <응답> 이런 응답이 나왔습니다. 아래 규칙에 의거하여 수정한 응답 문장을 생성해주세요. 규칙 : <규칙>"

 

이런식으로 프롬프트엔지니어링을 하면,

AI 가 스스로 본인의 응답을 헌법에 따라 분석하고 '올바른 답변'을 생성할 것입니다.

 

3. AI Self-Improvement

자가 개선 단계로,

앞서 생성한 프롬프트에 대한 올바른 답변을 가지고 파인튜닝을 합니다.

 

- 위와 같은 방식으로, RLHF 에서 보상 모델을 학습하기 위해 사람이 수동으로 훈련 데이터셋을 만드는 수고스러움을 개선할 수 있습니다.

사람의 라벨링을 거친 데이터가 아니기에 특정 인물의 사상이 반영되지 않는다는 장점도 있죠.

 

하지만 주의해야할 점이 있습니다.

위 프로세스에서 보듯, LLM 이 답변을 생성하는 것이기에 LLM 자체의 성능이 낮은 경우이거나 확률적으로 이상한 답변이 섞일 가능성도 무시할 수 없습니다.

 

예를들어,

아예 사전 학습이 잘 되지 않은 경우는 답변 자체가 이상하게 나오거나,

생성한 응답 문장의 퀄리티나 형식도 달라질 수가 있죠.

 

이 경우는 자기 비판을 통한 응답 생성에 대한 프롬프트를 Instruction Tuning 으로 학습시키거나,

Few-Shot Learning 등의 방식으로 개선을 시킬 수 있고,

운영 정책적인 측면에서 보자면 LLM 이 자립되었다고 판단하기 전까지는 생성된 데이터를 바로 학습시키지 말고, '프롬프트' - '수정된 응답' 쌍을 모아 인간이 이를 검토하여 수정을 거쳐 파인튜닝을 하는 방식으로 점차 모델을 학습시키는 방식이 있을 수 있습니다.

 

- Constitution AI 는 학습시점뿐 아니라 필터링 방식으로도 응용할 수 있습니다.

예를들어 서비스에서 민감한 답변을 극도로 피해야 하는 경우에는 처음 요청을 받아 출력한 응답 초안을 가지고 자기 비판을 시켜 올바르지 않은 응답으로 판명되었을 시 필터링을 하는 식입니다.

 

- 이상입니다.

본 게시글은 학습 기법에 대한 이론적인 내용에 대한 이해가 중점이므로,

모델 관련 코드는 따로 작성하지 않겠습니다.

추후 최신 LLM 개발 방법론을 정리할 때에 제가 정리한 MLOps, 파이프라인, 경량화 및 배포 방법 등을 소개드리며 학습 방식을 정리하겠습니다.