[코드 리뷰] mlx로 Vision Transformer 만들어보기

8 분 소요

이 글이 맘에 들거나 도움이 드셨다면 아래 레포에 스타 하나 부탁드립니다!!

서론

6월 3일 선거가 있던 날 나의 구현 능력이 궁금해서 뭐라도 손으로 구현해보고 싶었다. 그러던 중 “나는 Vision Transformer를 얼마나 이해하고 있을까?” 라는 질문이 문득 떠올랐다.

Q1: 왜 ViT?
A1: 연구에서 가장 자주 사용한게 Vision Transformer(ViT)이기도 하고 Sequential과 ModuleList라는 개념을 최근에 공부했는데 둘 다 사용해보고 싶어서 구현하고 싶었다. 본 구현에선 ViT 논문만을 참고해서 구현하는 것을 목표로 했다.

Q2: 왜 mlx?
A2: macbook과 mac mini를 사용하는 한명의 apple의 팬으로써, Mac slicon에서 동작하는 mlx라는 프레임워크에 관심이 있었다. mlx는 numpy-like, torch-like라 별로 어렵지 않겠지… 라는 생각과 함께 구현을 했다.

Q3: 참고한 사이트?
A3: timm 라이브러리를 주로 참고하였다. 코딩 에이전트/LLM은 cls 토큰을 배치 단위로 expand하는 법을 몰라서 찾아본 것 외에는 사용하지 않았다.

아래는 ViT 논문에서 제공하는 ViT의 구조도이다.

png

ViT를 구현하기 위해 고려해야 하는 사항은 크게 3가지가 있었다.

patch embedding을 어떻게 구현할지?
Multi head Attention을 어떻게 구현할지?
forward()를 어떻게 구현할지?

이제 차근차근 고려사항을 어떻게 구현했는지 알아보자.

구현

1. Patch Embedding

png ViT는 Transformer 기반의 입력을 받기 위해 이미지를 고정된 크기의 패치로 잘라서 flatten한 후 Linear Projection을 한다. 이 과정이 굉장히 복잡할 것 같지만 실제로는 nn.Conv2d를 통해 간단하게 구현이 가능하다.

patch size와 stride를 같게 주면 겹치는 부분없이 patchfy 하는 꼴이 되이 된다.
convolution 연산은 결국 패치들에 대하여 Linear projection을 하는 꼴이 된다.
해당 피처는 NHWC의 모양의 피처맵이 나오게 된다. H와 W에 대하여 flatten 해주면 1D의 토큰 나열이 나오게 된다.

class PatchEmbedding(nn.Module):
    def __init__(self, is_rgb=True, patch_size=16, embed_dim=768):
        super().__init__()
        self.patch_size = patch_size
        self.embed_dim = embed_dim
        self.in_channels = 3 if is_rgb else 1
        self.proj = nn.Conv2d(
                in_channels=self.in_channels,
                out_channels=embed_dim,
                kernel_size=patch_size,
                stride=patch_size
            )

    def __call__(self, x):
        # B, C, H, W = x.shape
        # print(f"Input Shape: {x.shape}")
        x = self.proj(x).flatten(1, 2) # 두번째 차원인 H/patch_size와 W/patch_size를 하나의 차원으로 합침
        return x

2. Block - Multi Head Self-Attention(MHSA)

png

MHSA는 Transformer block(layer)에서 필수적으로 구현을 해야한다. 비록 요즘에는 flash attention이라는게 나와서 단순히 호출을 하면 되지만, 본 글에서는 직접 수식을 구현하는 형태로 진행한다.

우선 self-attention의 아주 기본적인 수식은 다음과 같다.

2-1. QKV 계산

$Q = XW^Q, \quad K = XW^K, \quad V = XW^V$

2-2. Single-Head Scaled Dot-Product Attention

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

2-3. Multi-Head

$\text{head}_i = \text{Attention}(QW_i^Q,\ KW_i^K,\ VW_i^V)$

\[\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O\]

위 식을 코드로 구현하기 위해

우리는 projection 행렬을 만들어 Q, K, V를 얻어야 한다.
self-attention 수식을 구현해야 하고
그 와중에 head를 고려해서 구현해야 한다.

class MHSA(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads

        assert embed_dim % num_heads == 0, "embed_dim % num_heads != 0 !!!!" 

        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.qkv = nn.Linear(embed_dim, embed_dim * 3, bias=False)

        self.proj = nn.Linear(embed_dim, embed_dim)


    def __call__(self, x):
        # BND -> BN(3*D)
        x = self.qkv(x)
        # Q, K, V 3개로 chunk -> [BND, BND, BND]
        qkv = mx.split(x, axis=2, indices_or_sections=3) 
        # print(qkv[0].shape, qkv[1].shape, qkv[2].shape)
        
        # 1. Q,K,V 각각 BHND로 reshape (관례)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.num_heads), qkv) 
        # 2. Mat. Mul. -> BHN(Q)N(K)
        attn = q @ k.transpose(0, 1, 3, 2) 
        # 3. Scale: sqrt(dk)
        attn_score = attn * self.scale 
        # 4. attn_w: K에 대하여 softmax하기 위해서 axis=-1로 설정
        attn_weight = mx.softmax(attn_score, axis=-1) 
        # 5. Mat. Mul. -> BHND
        out = attn_weight @ v 
        # 6. 각 헤드 concat
        out = rearrange(out, 'b h n d -> b n (h d)', h = self.num_heads) # MHSA concat
        # 7. concat 후 정보 섞어주기 위해 같은 차원으로 projection
        out = self.proj(out)
        return out, attn_weight # 최종 출력과 attn_weight 같이 출력

구현체의 포인트는 다음과 같다.

projection 행렬은 아래처럼 nn.Linear로 QKV를 한번에 projection한 후

self.qkv = nn.Linear(embed_dim, embed_dim * 3, bias=False)

forward에서 각각 QKV로 2번째 차원인 D를 3개로 split해서 각각을 Q, K, V로 사용한다.

        x = self.qkv(x)
        # Q, K, V 3개로 chunk -> [BND, BND, BND]
        qkv = mx.split(x, axis=2, indices_or_sections=3) 

그 후에 Q,K,V 각각 BHND로 reshape을 한 후 Q, K를 행렬곱한다. 이때 K.transpose를(0,1,3,2)를 통해 K의 N과 D차원을 바꾸어 N $\times$ N 모양을 만들어준 후 scaling을 해주어 attention score가 만들어지도록 한다.

        # 1. Q,K,V 각각 BHND로 reshape (관례)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.num_heads), qkv) 
        # 2. Mat. Mul. -> BHN(Q)N(K)
        attn = q @ k.transpose(0, 1, 3, 2) 
        attn_score = attn * self.scale 

이후에 K에 대하여 softmax를 하는데 이때 Q에 대하여 해야할지 K에 대하여 해야할지 헷갈렸다. 어쨋든 Query에 대한 Key값의 확률을 구하는 과정이므로 axis=-1로 하여 K에 대하여 softmax를 하도록 했다.

         # 4. attn_w: K에 대하여 softmax하기 위해서 axis=-1로 설정
        attn_weight = mx.softmax(attn_score, axis=-1) 

그 이후에 V를 곱하고 head를 concat한 후 head 별 정보를 섞어주기 위해 동일한 차원으로 projection한다. forward 과정에서는 out만 사용하겠지만, 여러 분석 과정에서 attn_weight가 필요할 수 있으므로 동시에 반환하도록 설정했다.

        out = attn_weight @ v 
        # 6. 각 헤드 concat
        out = rearrange(out, 'b h n d -> b n (h d)', h = self.num_heads) # MHSA concat
        # 7. concat 후 정보 섞어주기 위해 같은 차원으로 projection
        out = self.proj(out)
        return out, attn_weight # 최종 출력과 attn_weight 같이 반환

3. Block - MLP

3-1. MLP

Transformer block의 또 다른 요소는 MLP인데 논문에 그렇게 자세히 나와있지 않아서 MLP는 timm을 참고하여 다음과 같이 간단하게 구현했다. 이 글에서는 짧게 작성하지만 실제로는 비선형성이 추가가 되는 부분이라 굉장히 중요한 부분이다. 절대 무시해선 안된다!

class MLP(nn.Module):
    def __init__(self, embed_dim=768, mlp_ratio=4, dropout_rate=0.1):
        super().__init__() 
        # 아래 timm-like MLP 참조함
        # https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/mlp.py
        self.net = nn.Sequential(
                nn.Linear(embed_dim, embed_dim * mlp_ratio),
                nn.GELU(),
                nn.Dropout(dropout_rate),
                nn.LayerNorm(embed_dim * mlp_ratio),
                nn.Linear(embed_dim * mlp_ratio, embed_dim ),
                nn.Dropout(dropout_rate),
            )
        
    def __call__(self, x):
        return self.net(x)

3-2. Block

Block도 여기서는 짧게 넘어가고자 한다. x를 normalization 하고 각 모듈을 통과한 후에 그 값에 x를 더하는 식의 Residual connection을 적용한다. 이는 논문 figure에 자세히 설명이 되어 있다.

class Block(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4, dropout_rate=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(dims=embed_dim) # 별도의 norm1
        self.norm2 = nn.LayerNorm(dims=embed_dim) # 별도의 norm1
        self.mhsa = MHSA(embed_dim=embed_dim, num_heads=num_heads)
        self.mlp = MLP(embed_dim=embed_dim, mlp_ratio=mlp_ratio, dropout_rate=dropout_rate)
    
    def __call__(self, x):
        # 1단계 - attention: 뭐가 더 중요한지 확인
        x_norm = self.norm1(x) # Layer normalization 1
        x_attn, _ = self.mhsa(x_norm) # MHSA
        x = x_attn + x # Residual connection
        
        # 2단계 - MLP: 비선형성 증가
        x_norm = self.norm2(x) # Layer normalization 2
        x_mlp = self.mlp(x_norm) # 비선형성 증가를 위한 MLP
        x = x_mlp + x

        return x

4. ViT

이제 입력을 받는 부분과 Transformer block을 전부 완성했다. 그럼에도 불구하고 아직 세가지가 더 완성이 되어야하는데 cls token과 positional embedding을 처리하는 부분과 간단한 classification head이다. 그리고 이를 한번에 실행시킬 forward 부분이 필요하다.

4-1. cls token과 positional embedding 그리고 classification head

cls token과 positional embedding은 결국 하나의 파라미터다. 따라서 mx.random.normal()을 사용한다. mlx는 명시적으로 parameter와 tensor를 구분하지 않는데 이는 model.parameters()를 호출하면 모듈 안의 모든 mx.array를 자동으로 pytree로 수집해버리기 때문이라고 한다.

classification head는 매우 간단하게 embed_dim에서 num_classes개로 projection하도록 만들었다.

mlx와 torch랑 다른 점은 nn.ModuleList를 사용하는 것이 아니라 그냥 리스트에 []에 각 모듈을 집어 넣어서 사용한다.

class ViT(nn.Module):
    """
        cls token based ViT 
    """
    def __init__(self, img_size=224, patch_size=16, embed_dim=768, num_heads=12, mlp_ratio=4, dropout_rate=0.1, depth=12, num_classes=1000):
        super().__init__()
        self.patch_size = patch_size
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio
        self.dropout_rate = dropout_rate
        self.depth = depth
        self.num_patches = int((img_size // patch_size) ** 2)
        
        self.norm = nn.LayerNorm(dims=embed_dim) # final norm


        # !주의! self.cls_token은 single batch 기준으로 만들어졌기 때문에 (0번째 차원이 1), _pos_embed 매서드에서 동적으로 배치 차원을 늘려줘야 함. 
        self.pos_embed = mx.random.normal((1, self.num_patches + 1, embed_dim), scale=0.02)  # 학습 가능한 파라미터로 position embedding 학습
        self.cls_token = mx.random.normal((1, 1, embed_dim), scale=0.02) # 학습 가능한 파라미터로 cls 토큰 학습 

        # FOR TEST 실제 실험 시 동작 변경!
        # self.cls_token = mx.zeros((1, 1, embed_dim))  # 동작 확인 용 cls token
        # self.pos_embed = mx.ones((1, self.num_patches + 1, embed_dim)) # 동작 확인 용 position embedding

        self.patch_embed = PatchEmbedding(
            patch_size=patch_size, 
            embed_dim=embed_dim
            )
        
        self.blocks = [
            Block(
                embed_dim=embed_dim, 
                num_heads=num_heads, 
                mlp_ratio=mlp_ratio, 
                dropout_rate=dropout_rate
                )  for _ in range(depth)
            ]
        
        self.head = nn.Linear(embed_dim, num_classes)

4-2. forward

forward는 단순하게 구현된다.

patch embedding을 한 후 _pos_embed 매서드를 이용하여 positional embedding과 cls token을 추가한다.
[]에 넣어놓은 Transformer block을 차례로 통과시켜 최종 출력을 얻는다.
최종 출력을 normalization한 후 cls token을 이용하여 head를 통과시켜 최종 결과를 얻는다. mlx는 torch와 다르게 forward()대신 __call__()를 이용하여 forward 기능을 제공한다. 이는 jax의 영향을 받은 것으로 생각한다.

_pos_embed를 만들면서 가장 신경 썼던 부분은 self.cls_token의 B 차원을 1로 정의했기 때문에 배치 단위로의 확장하는 것 이었다. torch에서는 보통 expand로 구현이 되는데 mlx에서는 broadcast_to라는 함수를 통해 구현을 해야했다.

또한 self.cls_token과 self.pos_embed이 정확히 동작하는지 확인하기 위해서 각각 zeros와 ones로 만들어 의도대로 정확히 동작하는지 확인했다.

그 결과 N차원이 1 늘어났으며 0번째에 0으로 이루어진 벡터가 추가된 것을 확인하여 self.cls_token이 제대로 추가되고 있는 것을 확인했다.
그 이후에 self.pos_embed이 1씩 더해지는 것을 확인하여 positional embedding이 잘 더해지는 것을 확인했다.

    def __call__(self, x):
        x = self.patch_embed(x)
        x = self._pos_embed(x)
        for i, block in enumerate(self.blocks):
            x = block(x)
            # print(f"{i} 번째 layer")
        x = self.norm(x)
        cls = x[:, 0]       # CLS token만 추출
        return self.head(cls)

    def _pos_embed(self, x):
        B, N, C = x.shape
        # print(B, N, C)
        # broadcast_to 함수로 cls token B 만큼 복제함. 
        # numpy-like 이므로 자세한 동작은 다음 문서 참고. https://numpy.org/doc/2.2/reference/generated/numpy.broadcast_to.html
        cls = mx.broadcast_to(self.cls_token, (B, 1, C))   
        x = mx.concatenate((cls, x), axis=1) # [cls] token N차원의 제일 앞에 concat, axis=1로 해줘야 N+1 됨. (BNC)
        # print(x)
        # array([[[0, 0, 0, ..., 0, 0, 0], > zeros로 설정 해놓고 제대로 추가 됐는지 테스트 완료
        #         [-0.625268, 0.0861489, 0.387924, ..., 0.565768, 1.14554, 0.420791],
        #         [1.4328, 0.820314, -0.0266257, ..., 0.0697592, -1.02677, 0.830738],
        #         ...,
        #         [0.164638, 0.335743, 0.710777, ..., 0.172818, -0.326656, 0.0479117],
        #         [-0.466198, 0.0355091, -0.264295, ..., -0.378135, 0.381905, -0.481186],
        #         [0.474137, 1.21557, -0.281954, ..., 0.562486, -0.0671904, 0.0877942]]], dtype=float32)
        
        x += self.pos_embed # position 정보 postion wise 하게 추가
        # print(x)
        # array([[[1, 1, 1, ..., 1, 1, 1], > ones로 설정 해놓고 제대로 추가 됐는지 테스트 완료
        #         [0.374732, 1.08615, 1.38792, ..., 1.56577, 2.14554, 1.42079],
        #         [2.4328, 1.82031, 0.973374, ..., 1.06976, -0.0267704, 1.83074],
        #         ...,
        #         [1.16464, 1.33574, 1.71078, ..., 1.17282, 0.673344, 1.04791],
        #         [0.533802, 1.03551, 0.735705, ..., 0.621866, 1.38191, 0.518814],
        #         [1.47414, 2.21557, 0.718046, ..., 1.56249, 0.93281, 1.08779]]], dtype=float32)
        return x
    
    def get_params_info(self):
        params = self.trainable_parameters()
        flat = mlx.utils.tree_flatten(params)
        # flat = [("layer.weight", array), ("layer.bias", array), ...]
        
        total = sum(v.size for _, v in flat)  # 언패킹 필요
        print(f"Total trainable parameters: {total:,}")
        return total

글을 마무리하며…

이번 구현을 통해 얻은 것은 다음과 같다.

softmax, MHSA 등에서 얻은 차원에 대한 이해
모델 설계 과정에서 동작 테스트하는 법
모델 중간 상태의 shape을 예상하는 능력
개발의 주도권을 AI에 뺏기지 않고 문서를 찾아보고 이해하는 능력

사전에 ViT와 torch에 대한 전반적인 지식이 있어서 구현에 큰 어려움을 겪지는 않았으나, 어려웠던 점은 다음과 같다.

cls token을 배치 단위로 확장하는 과정
MHSA에서 multi-head를 확장하는 과정
글에 적지는 않았지만 Block에서 MHSA 이전과 MLP 이전에 normalization 할 때 같은 norm을 사용했던 버그

앞으로도 AI를 사용하지 않고 문서 찾아보고 논문 찾아보며 직접 구현하여 AI modeling에 대한 이해를 높혀야겠다. 잘 알아야 AI도 잘쓰니까…

Twitter Facebook LinkedIn

YSPARK