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

 

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

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

 

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

 

 

목표

손글씨 데이터를 text-instruct 모델과 결합하여 텍스트만 input이 가능한 모델에 손글씨를 인식하도록 하고자 한다.

 

준비물

• MNIST
• LLM(HyperClovaX SEED Text-Instruct 0.5B)

ViT Encoder

ViT Encoder를 구현하기 위해서는 아래 단계가 필요합니다.

 

1. 이미지를 토큰화 하는 임배딩 레이어

2. 토큰화된 이미지를 학습 & 추론하는 트랜스포머 모델

 

번외로, 트랜스포머는 라이브러리로 제공되어있어 간단하게 사용 가능하지만, 공부를 목적으로 pytorch로 직접 구현을 시도해봤습니다.

 

ImageEmbeddingLayer

코드

import torch
import torch.nn as nn


class ImageEmbeddingLayer(nn.Module):
    def __init__(
        self,
        in_channels: int,
        img_size: int,
        patch_size: int,
        embedding_size: int = 768,
    ):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2
        self.embedding_size = embedding_size

        self.div_img = nn.Conv2d(
            in_channels, self.embedding_size, kernel_size=patch_size, stride=patch_size
        )

        self.cls_token = nn.Parameter(torch.randn(1, 1, self.embedding_size))
        self.position_div_imgs = nn.Parameter(
            torch.randn(1, self.num_patches + 1, self.embedding_size)
        )

    def forward(self, x):
        x = self.div_img(x)
        x = x.flatten(2).transpose(1, 2)
        b = x.shape[0]
        cls_tokens = self.cls_token.expand(b, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x += self.position_div_imgs
        return x

 

설명

 

파라미터

• patch size : 이미지를 몇등분할지 정하는 수
• in channels : 입력되는 채널의 수. 보통 이미지는 [R, G, B] 이렇게 3이지만, MNIST 데이터는 흑백의 1차원이기에 1
• embedding size : 이미지를 토큰으로 임배딩을 했을때 그 결과물의 차원 크기
  • (ex : [1 x 28 x 28] -> [1 x 1 x 768]으로 변환 하면 768이 embedding size) 

• cls token
  • 맨 앞에 있는 토큰으로 원래 llm(transformers)에서는 다른 토큰들의 정보들을 담는 역할.
  • ViT는 분리(패치)된 이미지들의 정보를 담는 토큰. 학습 전에는 그냥 빈 값(더미 값)으로 있지만 학습을 통해 의미있는 값으로 변경됨

 

초기화 단계

1. 이미지를 몇등분할지 정하고 나눠진 하나의 이미지의 픽셀 크기가 몇일지 계산 (self.num_patches = ...)

2. 이미지를 분리 (self.div_img = nn.Conv2d...)

3. cls token 생성 (self.cls_token = nn.Parameter...)

4. 분리된 이미지의 순서(원래 위치) 정보를 생성 (self.position_div_imgs=nn.Parameter...)

 

출력 단계

1. 이미지를 분리(패치)

2. 분리된 데이터를 transformers 모델이 넣을 수 있도록 데이터(행렬, 텐서) 구조 변경

3. cls token을 가져오기

4. 앞서 변경한 이미지 데이터 맨 앞에 cls token을 붙이기

5. 순서 정보를 (cls token + 이미지) 데이터에 첨부

---

ViTModule (Transformers)

그 다음 Transformers를 이용한 ViT 모델은 이렇게 구성되어 있습니다.

 

• MultiHead Attention
• Feed-Forward Layer
• 정규화 계층(Norm Layer)

 

Norm Layer는 pytorch로 제공해서 넘어가고 MultiHead Attention, Feed-Forward Layer만 구현했습니다.

 

MultiHead Attention

사실 multihead attention도 pytorch에서 다 구현해줘서 크게 할건 없습니다.

이 모듈의 목적은 들어온 정보들 중 어떤 정보가 중요한지 다각도에서 살펴보는 역할입니다.

 

코드

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, embedding_size: int, num_heads: int) -> None:
        super().__init__()
        self.multihead_attention = nn.MultiheadAttention(
            embed_dim=embedding_size,
            num_heads=num_heads,
            batch_first=True,
            dropout=0.2,
        )
        self.query = nn.Linear(embedding_size, embedding_size)
        self.key = nn.Linear(embedding_size, embedding_size)
        self.value = nn.Linear(embedding_size, embedding_size)

    def forward(self, x):
        query = self.query(x)
        key = self.key(x)
        value = self.value(x)
        attention_output, attention = self.multihead_attention(query, key, value)
        return attention_output, attention

 

설명

• embedding size : 들어올 데이터의 규격 크기를 의미합니다.
• num_heads : 헤드의 수 입니다. 위에서 언급한 "다각도"에서 "몇가지 각도에서 볼지" 정하는 수 입니다.
• query, key, value : 데이터들의 query, key, value로 만들 가중치들입니다.

 

초기화 단계

1. MultiheadAttention 초기화

2. query, key, value 초기화

 

출력 단계

1. 데이터를 query, key, value로 변경

2. multihead attention 통과 후 나온 데이터와 attention(어디를 집중했는지)을 리턴

 

Feed-Forward Layer

Multihead attention에서 다각도로 집중한 데이터들을 신경망을 통해 추론하는 역할입니다.

 

코드

import torch
import torch.nn as nn

class FeedForwardLayer(nn.Sequential):
    def __init__(
        self, embedding_size: int, expansion: int = 4, drop_out: float = 0.2
    ) -> None:
        super().__init__(
            nn.Linear(embedding_size, expansion * embedding_size),
            nn.GELU(),
            nn.Dropout(p=drop_out),
            nn.Linear(expansion * embedding_size, embedding_size),
        )

 

설명

• GELU : LLM에서 효과가 기존 RELU(활성화 함수)보다 좋은 버전입니다.
• Dropout : 신경망 중 무작위로 일부를 비활성화하여 과적합 등을 방지하고 다양성을 유지합니다.

초기화 & 출력

 

1. 데이터를 기존 크기의 4배(상황에 따라 사용해야 할 수치가 다름) 확장합니다.

2. GELU 통과

3. Dropout 통과

4. 확장된 데이터를 원 크기로 변경

 

ViTModule

위에서 구현한 Multihead Attention + Feed-Forward Layer + Norm하면 끝입니다.

각 모듈을 활용하여 임배딩된 이미지 데이터를 transformers 방법으로 추론합니다.

 

코드

import torch
import torch.nn as nn

class ViTModule(nn.Module):
    def __init__(self, embedding_size: int, num_heads: int = 8):
        super().__init__()
        self.multihead_attention = MultiHeadAttention(
            embedding_size, num_heads=num_heads
        )
        self.FFL = FeedForwardLayer(embedding_size)
        self.norm = nn.LayerNorm(embedding_size)

    def forward(self, x):
        # Attention + Skip Connection
        norm_x = self.norm(x)
        attn_out, attension = self.multihead_attention(norm_x)
        x = x + attn_out

        # FFL + Skip Connection
        x = x + self.FFL(self.norm(x))
        return x, attension

 

설명

• Multihead Attention : 들어온 정보들 중 어떤 정보가 중요한지 다각도에서 살펴보는 역할
• Feed-Forward Layer : Multihead attention에서 다각도로 집중한 데이터들을 신경망을 통해 추론하는 역할
• Norm : 데이터 분포를 고르게 만들어 학습 & 추론이 잘 되도록 도와주는 역할

 

초기화 단계

1. Multihead Attention을 헤드 수, 데이터 차원 수 지정

2. Feed-Forward Layer를 데이터 차원 수 지정

3. 들어올 데이터 차원 수에 맞게 지정

 

출력 단계

1. 데이터를 정규화 진행

2. Multihead Attention에 정규화 된 데이터 통과

3. Multihead Attention를 통과한 데이터를 Feed-Forward Layer를 통과시켜서 추론

 

이렇게 ViT 모듈도 완성이 되었습니다.

이제 임배딩 모듈과 ViT 모듈을 합쳐서 MNIST 전용 ViT Encoder를 만들면 됩니다.

 

MNISTViTEncoder

MNIST 데이터를 받으면 이를

토큰화 → transformers 모델으로 추론 → 추론 결과 리턴

하는 인코더를 만들어 보겠습니다.

 

코드

class MNISTViTEncoder(nn.Module):
    def __init__(
        self,
        embedding_size: int,
        img_size: int,
        patch_size: int,
        return_token_type: bool,  # True: Origin Data Token, False: CLS Token
        num_heads: int = 8,
        in_channels: int = 1,
    ) -> None:
        super().__init__()
        self.embedding_size = embedding_size
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_heads = num_heads
        self.in_channels = in_channels

        self.return_token_type = return_token_type
        self.embedding = ImageEmbeddingLayer(
            in_channels=in_channels,
            img_size=img_size,
            patch_size=patch_size,
            embedding_size=embedding_size,
        )

        self.layers = nn.ModuleList(
            [
                ViTModule(embedding_size=embedding_size, num_heads=num_heads)
                for _ in range(5)
            ]
        )

    def forward(self, x):
        x = self.embedding(x)
        for layer in self.layers:
            x, _ = layer(x)

        if self.return_token_type:
            return x  # Origin Data Token
        else:
            return x[:, 0]  # CLS Token

 

설명

• return_token_type
  • 원래 ViT 모델은 cls 토큰을 classifier에 통과시켜서 어떤 데이터인지 알 수 있습니다.
  • 하지만 llm에 넘길때는 cls토큰 뿐만 아니라 전체 토큰을 다 줘야 하기 때문에 어떤 형태로 넘길지 정하는 옵션입니다.
  • True : 원본, False : CLS 토큰 만
• layers : ViTModule을 하나만 하는것이 아닌 여러 개를 나열 & 추론하여 데이터 추론 성능을 높힙니다.

 

초기화 단계

1. 필요한 정보들을 모듈 클래스 내부에 저장 (embedding_size ~ return_token_type)

2. 임배딩 레이어 등록

3. ViTModule을 5개 등록

 

출력 단계

1. 이미지 데이터를 임배딩

2. 임배딩 데이터를 5개의 ViTModule에 통과

3. 리턴 타입에 따라 추론된 데이터 리턴

 

이제 ViT Encoder 과정이 끝났고 LLM과 결합하여 학습 & 추론하는 단계만 남았습니다.

 

다음 포스팅에서 LLM과 결합하여 학습 & 추론하는 과정을 살펴보겠습니다.

 

 

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