MNIST 데이터로 VLM 만들어보기 (Feat. HyperClovaX) - LLM 결합, 학습 & 추론

Llava를 보면 text to text llm + ViT encoder 조합으로 VLM을 만들었다.

그 과정을 MNIST 데이터를 통해 간단하게 따라 해보고자 한다.

 

ViT Encoder는 논문 및 블로그&AI를 참고하여 구현하였으며,
LLM은 구현하기에 제한이 많아 가벼운 네이버의 HyperClovaX SEED Text-Instruct 0.5B 모델을 사용.

 

 

오늘은 저번 포스팅에서 구현한 ViT Encoder를 LLM과 결합하고 학습 & 추론까지 진행하겠습니다.

 

※ ViT Encoder 구현 포스트 (저번 포스트)

MNIST 데이터로 VLM 만들어보기 (Feat. HyperClovaX) - ViT Encoder

 

LLM

Huggincface 라이브러리를 이용하여 준비

이미 학습 & 구현까지 라이브러리로 다 존재하여 가져다 쓰기만 하면 됩니다.

 

MNISTViTHyperClovaX

코드

import torch
import torch.nn as nn

from transformers import PreTrainedModel
from transformers import AutoModelForCausalLM, AutoTokenizer

class MNISTViTHyperClovaX(nn.Module):
    def __init__(
        self,
        vit_encoder: MNISTViTEncoder,
        llm_model: PreTrainedModel,
        llm_hidden_size: int = 1024,
    ) -> None:
        super().__init__()
        self.vit = vit_encoder
        self.llm = llm_model
        self.projector = nn.Linear(self.vit.embedding_size, llm_hidden_size)

        # LLM은 학습에서 제외 (Frozen)
        for param in self.llm.parameters():
            param.requires_grad = False

        # ViT와 Projector만 학습 가능하도록 설정
        self.vit.train()
        self.projector.train()

    def forward(self, images, input_ids, attention_mask, labels=None):
        # ViT 특징 추출 (Trainable)
        visual_tokens = self.vit(images)
        # Projector 통과 (Trainable)
        image_embeds = self.projector(visual_tokens)

        # 텍스트 임베딩 및 결합
        text_embeds = self.llm.get_input_embeddings()(input_ids)
        inputs_embeds = torch.cat([image_embeds, text_embeds], dim=1)

        # Mask 및 Labels 설정
        visual_mask = torch.ones(
            (images.size(0), image_embeds.size(1)),
            device=images.device,
            dtype=attention_mask.dtype,
        )
        combined_mask = torch.cat([visual_mask, attention_mask], dim=1)

        if labels is not None:
            visual_labels = torch.full(
                (images.size(0), image_embeds.size(1)),
                -100,
                device=labels.device,
                dtype=labels.dtype,
            )
            combined_labels = torch.cat([visual_labels, labels], dim=1)
            return self.llm(
                inputs_embeds=inputs_embeds,
                attention_mask=combined_mask,
                labels=combined_labels,
            )

        return self.llm(inputs_embeds=inputs_embeds, attention_mask=combined_mask)

 

설명

동작은 아래 순서로 진행됩니다.

 

1. 이미지를 토큰화 (ViTEncoder)

2. 이미지 토큰 규격을 LLM에 맞게 변경 (projector)

3. 변경된 이미지 토큰과 텍스트 토큰을 합쳐서 LLM에 전달

4. LLM이 추론 후 결과 리턴

 

파라미터

• vit : 저번 포스팅에서 구현한 MNISTViTEncoder
• llm : LLM 모델(추후 HyperClovaX SEED Text-Instruct 0.5B를 여기에 연결)
• projector : ViTEncoder의 출력 크기와 LLM의 입력 크기를 가운데에서 변환

 

forward

• label
  • 입력된 이미지의 정답 데이터 (손글씨 7 이미지의 정답은? 당연히 7)
  • label이 None이면 추론, 존재하면 학습 상태
• mask란?
  • attention이 토큰의 어디를 집중해서 보는지 나타내는 행렬

 

 

초기화 단계

1. vit, llm, projector 정의

2. llm의 파라미터가 학습이 안되도록 고정

    - 지금 학습할 데이터를 학습할 경우, 이미 잘 학습된 가중치가 흐트러질 가능성이 존재합니다.

3. vit, projector를 학습 가능한 모드로 전환

 

출력 단계

1. 이미지를 ViTEncoder를 통과 시켜서 토큰화.

2. 이미지 토큰을 Projector를 통과 시켜서 LLM의 input 규격으로 통일.

•   Projector를 여기서는 규격 통일을 메인 목적으로만 사용했지만, 만약 VLM의 규모가 커지면,  이미지 토큰의 임배딩 값과 LLM의 텍스트 토큰 임배딩 값 간의 차이를 번역해주는 MLP등의 더 복잡한 구조로 사용됩니다.
•   Linear의 단순 구조라 하여도 위에서 언급한 임배딩 값 간 차이를 번역해주기 위해 이후 학습 단계에서 Porjector도 학습을 진행합니다.

3. 텍스트를 LLM의 임배딩 함수를 통과 시켜서 텍스트 토큰으로 전환

4. 이미지 토큰 + 텍스트 토큰 결합

5. 이미지 mask 생성

 

visual_mask = torch.ones(
          (images.size(0), image_embeds.size(1)),
          device=images.device,
          dtype=attention_mask.dtype,
      )

• 목적 : 이미지 토큰을 추후 텍스트 토큰과 같이 담기 위한 규격 변환 과정
• images.size(0) : 이미지가 몇장이 들어왔는지 알려주는 데이터
• image_embeds.size(1) : 이미지 토큰의 수가 몇 개인지 (patch * patch + 1(cls토큰)) 알려주는 데이터 
• torch.ones : 위에서 정의한 크기의 행렬을 전부 1로 채움

• ※ 왜 1로 다 채우나?
  • attention이 집중해서 보는 기준이 mask 데이터가 1이면 집중, 0이면 넘깁니다.
  • 글(텍스트)은 길이가 전부 다르기에 입력 토큰보다 작은 크기의 글이 들어오면 남은 부분을 0으로 채웁니다.
  • 하지만 이미지는 규격이 정해져있기에 전부 다 봐야 하는 중요한 데이터이기에, 전부 1로 채웁니다.

 

6. 이미지 mask과 텍스트 mask을 합쳐서 combined_mask 생성.

7. label 여부를 통해서 추론 or 학습 상태 파악

    7-1. label이 존재 (학습 상태)

        7-1-1. 이미지 라벨 생성

visual_labels = torch.full(
    (images.size(0), image_embeds.size(1)),
    -100,
    device=labels.device,
    dtype=labels.dtype,

)

• images.size(0) : 이미지가 몇 장이 들어왔는지 알려주는 데이터
• image_embeds.size(1) : 이미지 토큰의 수가 몇 개인지 (patch * patch + 1(cls토큰)) 알려주는 데이터 
• torch.full: -100으로 다 채움

• ※ 왜 -100 으로 다 채우나?
  • 위 함수가 동작한다는건 label이 존재하는 "학습 상황"이라는걸 의미.
  • 지금 이미지만 보는 게 아니라 뒤에 오는 텍스트 데이터도 같이 보기 위해 지금은 넘어가라는 걸 의미.

 

      7-1-2. 이미지 라벨 + 텍스트 라벨 결합 (label은 이미 텍스트 라벨을 담은 값이기에 별도 과정 없이 바로 결합)

      7-1-3. LLM에 [텍스트 임베딩 데이터, 결합한 마스크 데이터, 결합한 라벨 데이터]를 입력

 

    7-2. label이 None(추론 상태)

        7-2-1. LLM에 [텍스트 임베딩 데이터, 결합한 마스크 데이터]를 입력

 

좀 많이 복잡했습니다.

이제 학습으로 넘어가보겠습니다.

---

Train(학습)

코드

def train(epochs: int = 10):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model_name = "naver-hyperclovax/HyperCLOVAX-SEED-Text-Instruct-0.5B"

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # HyperCLOVA X의 경우 패딩 토큰 설정 필수
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # 모델 로드
    llm = AutoModelForCausalLM.from_pretrained(model_name).to(device)  # type: ignore

    vit = MNISTViTEncoder(
        embedding_size=768,
        img_size=28,
        patch_size=7,
        return_token_type=True,
        num_heads=12,
    )
    vlm_model = MNISTViTHyperClovaX(vit, llm).to(device)

    # 파라미터 최적화 대상: ViT + Projector
    optimizer = AdamW(
        [
            {"params": vlm_model.vit.parameters(), "lr": 1e-5},
            {"params": vlm_model.projector.parameters(), "lr": 1e-4},
        ]
    )

    # MNIST 데이터셋 준비
    # 1. 표준 MNIST 로드 (원본 데이터)
    data_path = os.path.join("datasets", "mnist")
    raw_train_dataset = datasets.MNIST(root=data_path, train=True, download=True)

    # 2. 커스텀 데이터 셋(MNISTVLMDataset)으로 감싸기
    # 여기서 텍스트 템플릿과 토크나이징이 적용됩니다.
    train_vlm_dataset = MNISTVLMDataset(raw_train_dataset, tokenizer)

    # 3. DataLoader에 vlm_collate_fn 적용
    # collate_fn을 넣어야 리스트 형태의 출력을 텐서 배치로 묶어줍니다.
    train_loader = DataLoader(
        train_vlm_dataset, batch_size=64, shuffle=True, collate_fn=vlm_collate_fn
    )

    # 학습 루프
    vlm_model.train()
    for epoch in range(epochs):
        epoch_loss = 0.0

        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}")
        for batch in pbar:
            optimizer.zero_grad()

            with torch.autocast(device_type=device, dtype=torch.bfloat16):
                outputs = vlm_model(
                    images=batch["images"].to(device),
                    input_ids=batch["input_ids"].to(device),
                    attention_mask=batch["attention_mask"].to(device),
                    labels=batch["input_ids"].to(device),
                )
                loss = outputs.loss

            loss.backward()
            optimizer.step()

            # 통계 업데이트
            current_loss = loss.item()
            epoch_loss += current_loss

            # tqdm 상태바에 현재 배치의 Loss 표시
            pbar.set_postfix(loss=f"{current_loss:.4f}")

        # 한 에포크가 끝나면 평균 Loss 출력
        avg_epoch_loss = epoch_loss / len(train_loader)
        print(
            f"\n>>> Epoch [{epoch+1}/5] Completed. Average Loss: {avg_epoch_loss:.4f}"
        )
        print("-" * 50)

    # 학습 완료 후 저장
    torch.save(vlm_model.vit.state_dict(), "vit_768_mnist.pth")
    torch.save(vlm_model.projector.state_dict(), "projector_768_to_1024.pth")

 

설명

정의

• pad_token : LLM에서 입력값으로 들어온 값이 기준보다 짧을 경우 빈칸을 채우는 용도의 토큰
• tokenizer : 텍스트 등의 데이터를 LLM이 이해할 수 있도록 토큰으로 변경해주는 도구
• optimizer : 학습 과정 중에 가중치 변경 값을 모델에 반영하려고 할때 해당 값을 얼마나 어떻게 반영할지 정하는 도구
• train_vlm_datasets : 이미지만 있던 MNIST에 VLM에 맞게 데이터에 대한 텍스트 데이터를 합쳐진 데이터 셋

동작

 

모델 정의

1. HyperCLOVAX-SEED-Text-Instruct-0.5B 모델의 토크나이저 불러오기

2. LLM(HyperCLOVAX-SEED-Text-Instruct-0.5B) 불러오기

3. ViTEncoder 불러오기

4. MNISTViTHyperClovaX에 LLM, ViTEncoder를 넘기면서 객체 생성

 

데이터 불러오기

1. MNIST 원본 데이터 불러오기

2. 커스텀 MNIST 데이터 셋에 MNIST 데이터와 토크나이저를 같이 넘기면서 VLM용 학습 데이터 생성

3. train loader에 학습 데이터를 넘겨 VLM이 학습할 수 있도록 형식 변경

 

학습

1. vlm model을 학습 모드로 지정

2. epochs 수 만큼 학습 진행

   2-1. train loader에서 학습할 batch만큼 데이터 가져오기

       2-1-1. optimizer의 zero_grad(역전파로 산출해낸 가중치 변경값 임시 저장소) 초기화

        2-1-2. 데이터를 VLM 모델에 통과 (순전파) 및 loss를 저장

        2-1-3. loss를 활용해서 계산(역전파) 및 optimizer의 zero_grad에 저장

            ※ zero grad에 저장하는 부분은 코드에 없지만 내부적으로 실행

        2-1-4. optimizer를 활용하여 zero_grad에 있던 값을 가중치에 반영

    2-2. 현재 batch에 대한 평균 loss 계산

3. 현재 epoch에 대한 평균 loss 계산

4. 학습을 epochs 수 만큼 진행하고 그 결과(가중치)를 저장

 

학습 결과

# Epoch 1/5: 100%|█████████████████████████████████████████████████████████████████████| 938/938 [10:07<00:00,  1.54it/s, loss=0.0060]

# >>> Epoch [1/5] Completed. Average Loss: 0.0827
# --------------------------------------------------
# Epoch 2/5: 100%|█████████████████████████████████████████████████████████████████████| 938/938 [10:07<00:00,  1.54it/s, loss=0.0006]

# >>> Epoch [2/5] Completed. Average Loss: 0.0035
# --------------------------------------------------
# Epoch 3/5: 100%|█████████████████████████████████████████████████████████████████████| 938/938 [10:07<00:00,  1.55it/s, loss=0.0023]

# >>> Epoch [3/5] Completed. Average Loss: 0.0022
# --------------------------------------------------
# Epoch 4/5: 100%|█████████████████████████████████████████████████████████████████████| 938/938 [10:07<00:00,  1.54it/s, loss=0.0003]

# >>> Epoch [4/5] Completed. Average Loss: 0.0016
# --------------------------------------------------
# Epoch 5/5: 100%|█████████████████████████████████████████████████████████████████████| 938/938 [10:07<00:00,  1.55it/s, loss=0.0003]

# >>> Epoch [5/5] Completed. Average Loss: 0.0013

 

궁금 사항

• ViT Encoder, Projector만 학습을 한 이유가 있나요?
  • 이미 CLIP처럼 훌륭한 ViTEncoder가 있다면 ViTEncoder도 학습을 하지 않고, ViTEncoder와 LLM이 서로 대화를 할 수 있도록 Projector만 진행을 합니다.
  • 하지만 여기서는 공부를 위해서 ViTEncoder를 직접 구현하였기에 ViTEncoder까지 MNIST로 학습하는 과정이 있습니다.
  • 이후 더 자연스럽고 LLM의 학습 능력을 사용하고자 할때 encoder를 고정하고 Projector, LLM이 두개를 파인튜닝 합니다.
• 데이터는 왜 MNIST로 했나요?
  • 데이터 크기, 규모가 작고 대중적인 데이터이며, 데이터를 인식한 결과가 명확하기에 해당 데이터로 선택했습니다.
• LLM은 고정인데 왜 학습할때 로드를 해야 하나요?
  • LLM의 가중치는 고정일지라도, Projector가 ViTEncoder와 LLM사이를 번역하는 방법을 학습할때 LLM까지 데이터를 보내서 그 결과를 토대로 학습하기 때문에 필수 입니다.

---

Valid(검증)

코드

def valid():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model_name = "naver-hyperclovax/HyperCLOVAX-SEED-Text-Instruct-0.5B"

    ##########DataSet###########
    # 1. 이미지 전처리 정의 (학습 시와 동일해야 함)
    transform = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
    )

    # 2. 데이터셋에서 샘플 하나 추출
    test_ds = datasets.MNIST(root="./datasets/mnist", train=False, download=True)
    sample_img, sample_label = test_ds[0]  # 첫 번째 데이터 (숫자 7)

    # 3. 텐서 변환 및 배치 차원 추가 [1, 1, 28, 28]
    input_tensor = transform(sample_img).unsqueeze(0)

    ##############Inference###################
    # 1. 모델 및 토크나이저 로드
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    llm = AutoModelForCausalLM.from_pretrained(model_name).to(device)

    # 학습 시와 동일한 구조의 ViT 선언
    vit = MNISTViTEncoder(
        embedding_size=768,
        img_size=28,
        patch_size=7,
        return_token_type=True,
        num_heads=12,
    ).to(device)

    # VLM 클래스
    vlm_model = MNISTViTHyperClovaX(vit, llm).to(device)

    # 2. 학습된 가중치 로드
    vlm_model.vit.load_state_dict(torch.load("vit_768_mnist.pth", map_location=device))
    vlm_model.projector.load_state_dict(
        torch.load("projector_768_to_1024.pth", map_location=device)
    )
    vlm_model.eval()

    # 3. 추론용 프롬프트 구성 (답변 직전까지만 입력)
    prompt = "질문: 이 이미지에 있는 숫자는 무엇인가요?\n답변: 이 숫자는"
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    # 4. 생성 (Inference)
    with torch.no_grad():
        target_device_type = "cpu" if device == "cpu" else "cuda"
        with torch.autocast(device_type=target_device_type, dtype=torch.bfloat16):
            # 이미지 특징 추출 및 프로젝션
            visual_tokens = vlm_model.vit(input_tensor.to(device))
            image_embeds = vlm_model.projector(visual_tokens)

            # 텍스트 임베딩 추출
            text_embeds = vlm_model.llm.get_input_embeddings()(inputs["input_ids"])

            # [이미지 임베딩 + 텍스트 임베딩] 결합
            combined_embeds = torch.cat([image_embeds, text_embeds], dim=1)

            output_ids = vlm_model.llm.generate(
                inputs_embeds=combined_embeds,
                max_new_tokens=10,  # " 5입니다." 정도만 생성하면 되므로 짧게 설정
                do_sample=False,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )

    # 5. 결과 디코딩
    result = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    print(f"--- 추론 결과 ---")
    if sample_label is not None:
        print(f"실제 정답: {sample_label}")
    print(f"모델 답변: {result}")

설명

정의

• transform : 이미지 데이터의 규격을 조정(정규화, 크기 조정, 행렬 변환 등)하는 도구 (Tokenizer가 텍스트 전문이라면 Transform은 이미지 특화)
• max_new_tokens : LLM이 출력으로 생성할 토큰의 최대 갯수 설정
• decode : 결과로 나온 데이터(숫자형태)를 자연어로 매칭해주는 도구
• skip_special_tokens : LLM에서 시작, 종료 등을 알리는 사전에 미리 정의된 특별한 의미를 가진 토큰들을 없애고 출력하는 옵션

 

동작

데이터 불러오기

1. transforms 정의

 - (0.1307,), (0.3081,)과 같은 숫자는 transforms에서 정규화 할때 사용할 MNIST 데이터의 평균, 표준 편차 값 입니다.

2. MNIST 데이터 셋에서 숫자 이미지 하나 가져오기

3. 이미지를 VLM에 넣기 위해 변환 

 

모델 정의

1. ViTEncoder, LLM 불러오기

2. MNISTViTHyperClovaX에 ViTEncoder, LLM 입력 및 객체 생성

3. VLM을 검증(추론) 모드 - eval() 로 정의

 

추론

1. 프롬프트를 토큰화

2. 이미지를 ViT를 통과해 토큰화

3. 이미지 토큰을 임배딩

4. 텍스트 토큰도 임배딩

5. [이미지 임배딩 + 텍스트 임배딩]으로 결합

6. VLM(안에 있는 LLM)에 입력 & 추론

 

결과 출력

1. 추론 결과를 decode

2. 실제 데이터와 비교

 

추론 결과

# --- 추론 결과 ---
# 실제 정답: 7
# 모델 답변:  7입니다.

 

궁금 사항

• 이 구조면 VLM이 하나로 움직이는게 아니라 모듈화 구조인건가요?
  • 네 맞습니다. 이미지를 토큰화 하는 ViTEncoder, 언어 능력을 가진 LLM, 둘 사이를 연결하는 Projector 이렇게 3파트가 독립적으로 존재합니다.
• 그러면 LLM을 굳이 학습할때 모델을 사용안해도 되는건가요?
  • 네 맞습니다. LLM은 학습할때 고정을 했기 때문에 굳이 동일 모델을 사용하지 않아도 문제 없습니다.
  • 지금까지의 과정으로는 LLM은 학습할 때, 길잡이 역할만 했을 뿐 수정이 있지 않습니다.
  • 단, LLM까지 파인튜닝 했다면 그때는 세트로 움직여야 합니다. 

---

그 외

커스텀 데이터 셋 코드

from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Dataset

import os
import torch


def vlm_collate_fn(batch):
    images = torch.stack([item["images"] for item in batch])
    input_ids = torch.stack([item["input_ids"] for item in batch])
    attention_mask = torch.stack([item["attention_mask"] for item in batch])

    return {"images": images, "input_ids": input_ids, "attention_mask": attention_mask}


class MNISTVLMDataset(Dataset):
    def __init__(self, mnist_dataset, tokenizer, max_length=128):
        self.mnist = mnist_dataset  # torchvision.datasets.MNIST 객체
        self.tokenizer = tokenizer
        self.max_length = max_length

        # MNIST는 28x28이므로 ViT 입력에 맞게 resize하거나 정규화합니다.
        self.transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        )

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

    def __getitem__(self, idx):
        image, label = self.mnist[idx]
        image = self.transform(image)

        # 1. 대화 형식 구성
        # <image> 토큰은 나중에 ViT 특징값이 들어갈 자리임을 표시하는 특수 문자열입니다.
        prompt = (
            f"질문: 이 이미지에 있는 숫자는 무엇인가요?\n답변: 이 숫자는 {label}입니다."
        )

        # 2. 토크나이징
        inputs = self.tokenizer(
            prompt,
            return_tensors="pt",
            padding="max_length",
            truncation=True,
            max_length=self.max_length,
        )

        return {
            "images": image,  # [1, 28, 28]
            "input_ids": inputs["input_ids"].squeeze(),  # [L]
            "attention_mask": inputs["attention_mask"].squeeze(),  # [L]
        }


def mnist_dataloader():

    batch_size: int = 64

    transform = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
    )

    data_path = os.path.join("datasets", "mnist")
    train_dataset = datasets.MNIST(
        root=data_path, train=True, transform=transform, download=True
    )
    test_dataset = datasets.MNIST(root=data_path, train=False, transform=transform)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    return train_loader, test_loader

• 본 포스팅은 HyperCLOVA X SEED 모델을 활용한 학습 과정을 담고 있습니다.
• HyperCLOVA X SEED Model is licensed under the HyperCLOVA X SEED Model License Agreement, Copyright © NAVER Corp. All Rights Reserved.