스크래치 nanoGPT

최

최윤진

Jan 25, 20242y ago

들어가며

네이처는 매년 이슈가 된 과학자 10인, nature’s 10을 뽑습니다. 2023년 nature’s 10에 ChatGPT가 명단을 올렸습니다. 네이처는 2023년 ChatGPT 가 세상 전반에 큰 영향력을 끼쳤다고 했습니다.

더 성능이 좋고 가벼운 모델들을 개발하려는 움직임에 막대한 자본과 인력이 투입되고 있습니다. 동시에 점점 더 커지고 복잡해지는 LLM을 이해하는 것이 어려워 지고 있습니다. 이번 글에서는 GPT를 아주 간소하게 만든 nanoGPT를 만들어봄으로써 LLM의 내부 메커니즘을 파악해보겠습니다.

출처 : nature

먼저 GPT를 잘 알기 위해서는 Transformer 모델을 알아야 합니다. 이 모델을 공식적으로 발간한 ‘Attention is all you need(NeurIPS, 2017)’ 논문은 인용수는 10만회를 넘었습니다. 이 논문을 기반으로 GPT 시리즈가 만들어졌습니다. GPT-1(2018년), GPT-2(2019년), GPT-3(2020년), InstructGPT(2022년), 그리고 2023년에 GPT-4가 나왔습니다.

[Transformer] Attention Is All You Need, 2017

[GPT-1] Improving Language Understanding by Generative Pre-Training, 2018

[GPT-2] Language Models are Unsupervised Multitask Learners, 2019

[GPT-3] Language Models are Few-Shot Learners, 2020

[Instruct GPT] Training language models to follow instructions with human feedback, 2022

[GPT-4] GPT-4 Technical Report, 2023

때문에 먼저 Transformer 의 아키텍쳐를 살펴보고 nanoGPT 를 만드는 순으로 진행하겠습니다. 코드는 Andrej Karpathy 의 Let's build GPT: from scratch, in code, spelled out. 를 참고했습니다. 셰익스피어의 문체를 학습하고 생성하는 모델을 만들어보겠습니다.

Before Transformer

출처 : Attention Is All You Need, 2017

Transformer도 딥러닝 모델 중 하나

딥러닝 모델을 만든다는 것은 input 과 output 이 있고 그것을 일반성있게 예측할 수 있는 함수를 만들어내는 것으로 볼 수 있습니다. 수 많은 가중치(weight)들을 조정하며 학습합니다.

Transformer 또한 딥러닝 모델 중 하나로, 수많은 가중치(weight)를 조정하며 input 과 output 을 연결합니다.

출처 : Alammar, J (2018). The Illustrated Transformer [Blog post]. Retrieved from https://jalammar.github.io/illustrated-transformer/

Recurrent 모델의 단점

저희는 처리하고 싶은 것은 언어입니다. 언어는 순차적으로 사용되기 때문에 대표적인 시퀀스 데이터입니다.

기존 딥러닝 방법들은 시퀀스 데이터를 처리하기 위해 순환적 구조를 가지도록 모델 아키텍쳐를 구성하는 Recurrent 방법을 사용했습니다. 대표적으로 RNN, LSTM, GRU 와 같은 알고리즘들이 있습니다.

출처 : Recurrent Neural Network Architectures, 2017

출처 : Deep Reinforcement Learning for Sequence-to-Sequence Models, 2018

사람은 언어를 사용 할 때 대명사를 사용하기도 하고, 비유적인 표현을 쓰기도 하며 매우 긴 문서를 작성하기도 합니다. 때문에 단어 간의 거리가 멀거나 복잡한 문장이더라도 단어 간의 연관 관계를 제대로 파악하는 모델을 만들어야 언어를 제대로 이해하는 모델을 만든다고 할 수 있습니다.

하지만 이런 Recurrent 방법들은 구조상 이전의 상태값 $h_{k-1}$의 값이 새로운 input 값 $X_{k}$ 와 결합하여 $H_{k}$ 를 만들어 내는 방식입니다. 이것은 hidden state 값들이 희석되어 거리가 먼 토큰 사이의 연관관계를 파악하지 못하는 long term dependency problem 이 발생시키게 됩니다. 더해, 이전 상태가 완료되어야 다음 상태를 만들어낼 수 있는 병목현상을 구조적으로 가지고 있습니다.

Encoder-Decoder 모델에서는 하나의 context vector로 encoder의 출력을 압축합니다. 이 과정에서 마찬가지로 정보가 손실된다는 문제점이 있습니다.

문맥 내에서 단어들간의 연관관계를 파악해 의미 손실 없이 전달하려는 시도가 계속 되었습니다. 이러한 흐름 속에서 나온 개념이 Attention 입니다. Attention은 말 그대로 ‘주의’ 이며, 얼마나 주의를 기울일 것이냐를 나타내는 값 입니다. 다음 문장을 살펴보겠습니다.

”The animal didn't cross the street because it was too tired”

여기서 ‘it’는 ‘animal’ 입니다. 이것을 어떻게 더 잘 학습시킬 수 있을까요?

아래 그림 처럼 Attention 을 통해서 hidden 에 대한 가중 합(weighted sum)을 사용하는 기법을 사용할 수 있습니다. Neural Machine Translation by Jointly Learning to Align and Translate, 2014 논문에서 제안한 이 방법은 hidden state 를 전달하는 과정에서 어떤 hidden state 에 영향을 줄 것인지는 attention을 줌으로써 더 중요한 것에 주목도를 높히는 방식으로 성능을 개선했습니다.

An attention-based seq2seq model 개요 출처 : Neural Abstractive Text Summarization with Sequence-to-Sequence Models, 2018

Transformer

생각보다 직관적입니다. 이러한 attention 의 잠재력을 완전히 이끌어낸 것이 Transformer 모델 입니다. Transformer 는 Encoder Decoder 모델을 취하고 오직 Attention 만을 사용합니다.

Encoder-Decoder 구조를 가지고 있으며 각각 Encoder block, Decoder block 으로 이루어져 있습니다.

출처 : Attention Is All You Need, 2017

Transformer 구성요소

Encoder block

Encoder block 은 2 가지 요소로 구성되어 있습니다.

Multi-Head Attention

Feed Forward

Decoder block

Decoder block 은 3 가지 요소로 구성되어 있습니다.

Masked Multi-Head Attention

Cross Multi-Head Attention

Feed Forward

Feed Forward 는 2번의 linear 와 1번의 activation function을 거치도록 구성되어 있습니다.

이제 남은 것은 Multi-Head Attention, Masked Multi-Head Attention, Cross Multi-Head Attention 이 있습니다. 이 세 부분의 작동 원리는 self-attention 메커니즘 입니다.

Embedding, Positional Encoding

그에 앞서 self-attention 입력단 전까지 작업되는 것들을 살펴보도록 하겠습니다.

다시 ”The animal didn't cross the street because it was too tired” 라는 문장을 생각해보겠습니다.

이 문장 자체는 컴퓨터가 이해할 수 없기 때문에 벡터로 매핑하는 임베딩을 해야 합니다.

이것은 사전에 정의된 Vocab 에 의해 토큰화가 된 상태에서 출발합니다. vocab을 만들 때는 BPE, sentencepiece 등 다양한 토큰화 알고리즘들이 사용될 수 있습니다. 편의상 character 단위로 설정하였습니다.

Vocab 만들기

text를 set 으로 만들고 list로 변환하는 과정을 통해 간단하게 vocab 을 만들수 있습니다. 아래 문장은 총 19개의 vocab을 가지게 됩니다. 이렇게 vocab을 만듦으로써 text 의 모든 character는 특정 index 로 변환할 수 있습니다.

text = "The animal didn't cross the street because it was too tired"

chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"chars : {''.join(chars)}")
print(f"vocab_size : {vocab_size}")

chars :  'Tabcdehilmnorstuw
vocab_size : 19

Embedding

각 토큰들은 아래와 같은 과정을 거쳐 벡터로 매핑 될 수 있습니다.

문자를 인덱스로, 인덱스를 문자로 매핑하는 dictionary 를 생성하고 파이토치의 벡터를 생성해주는 함수 nn.Embedding을 통해 아래와 같이 로직을 거치게 되면 토큰 별로 임베딩 값을 얻을 수 있습니다.

import torch
import torch.nn as nn

# 문자를 인덱스로, 인덱스를 문자로 매핑하는 사전 생성
**char_to_index = {char: index for index, char in enumerate(chars)}
index_to_char = {index: char for index, char in enumerate(chars)}**

# Embedding layer 생성 (임베딩 차원을 예로 vocab_size의 절반으로 설정)
embedding_dim = 4  # 혹은 원하는 다른 차원 수
**token_embedding_table = nn.Embedding(vocab_size, embedding_dim)**

# 전체 텍스트를 인덱스로 변환하여 임베딩 조회
text_indices = torch.tensor([char_to_index[char] for char in text])
**text_embedding = token_embedding_table(text_indices)**

# 결과
print(f"text_embedding.shape : {text_embedding.shape}")
print(f"text_embedding : \n {text_embedding[:10]}")

text_embedding.shape : torch.Size([59, 4])

text_embedding : 
 tensor([[ 0.6686, -0.4579,  0.5651,  0.5184],
        [-1.1360,  0.3221,  0.0946, -1.4244],
        [ 0.6121,  1.3113,  0.4926,  0.4148],
        [-1.3359,  0.4271,  0.9899,  1.0274],
        [-0.1111,  0.2487,  2.7271, -1.7861],
        [ 0.1996, -1.5881,  0.5240, -0.4852],
        [ 1.7475, -1.3174, -1.9797, -0.2429],
        [-0.0726,  0.2717, -1.5650, -0.1542],
        [-0.1111,  0.2487,  2.7271, -1.7861],
        [-0.9550,  0.8071, -1.2684, -1.8388]], grad_fn=<SliceBackward0>)

이런식 각 토큰들을 벡터로 임베딩 하게 되면 컴퓨터가 연산할 수 있습니다. 추가적으로 위치정보를 반영하는 positional encoding 도 더해 줘야 합니다. Transformer 에서는 사인, 코사인 함수를 이용하여 위치에 따른 고유값을 부여해줍니다.

이것까지 마무리하게 되면 이제 self-attention 을 위한 사전 준비가 마무리가 되었습니다.

Self-Attention

출처 : Attention Is All You Need, 2017

각 토큰들이 다른 토큰들과 어떻게 연관이 있는지에 대한 정보를 얻어내는 것이 이 self-attention의 역할이며 이러한 Encoder block 이 여러번 거치게 되면 토큰 간의 관계를 풍부한 정보를 얻어낼 수 있습니다.

주의해야 할 점은 input의 shape 과 self-attention을 통과한 결과 값의 shape이 같다는 것입니다.

self-attention 은 Query, Key, Value 3 가지 요소로 구성되어 있습니다.

단어 의미 그대로 Query로 Key를 조회하고 Value를 구한다 라고 생각하면 직관적으로 이해를 할 수 있습니다.

Query를 알고 싶은 대상이며, Key와 Value는 조회가 되는 대상입니다.

출처 : _Alammar, J (2018). The Illustrated Transformer [Blog post]. Retrieved from __[https://jalammar.github.io/illustrated-transformer/](https://jalammar.github.io/illustrated-transformer/)_

먼저 Query, Key, Value 행렬을 초기화 해야 합니다.

이때

W^Q, W^k, W^v

가중치 행렬이 사용됩니다.

이 행렬들은 transformer 를 통해 학습됩니다.

input 값의 shape을 (T,d) 라 해보겠습니다.

그렇다면 각 가중치 행렬들을 아래와 같은 shape 으로 초기화를 하고 input과 연산을 함으로써 각 Query, Key, Value를 만들어낼 수 있습니다.

shape\, of\, W^Q = d * d_q

shape\, of\, W^k = d * d_k

shape\, of\, W^v = d * d_v

shape\, of\, query = n*d_q

shape\, of\, key = n*d_k

shape\, of\, value = n*d_v

일반적으로 d_q, d_k, d_v 값은 값을 사용합니다.

Query와 Key를 내적 한 값이 attention score 입니다. 각 토큰 임베딩 값의 유사도를 계산했기 때문에 두 토큰의 관계 정도를 얻어 낼 수 있는 값이기 때문입니다.

아래 그림을 살펴보겠습니다. ‘Thinking’, ‘Machine’ 의 Query, Key, Value 를 구한 후, Query 와 Key 를 내적하여 attention score를 구하게 됩니다.

출처 : Alammar, J (2018). The Illustrated Transformer [Blog post]. Retrieved from https://jalammar.github.io/illustrated-transformer/

이제 Query 와 Key의 행렬 곱을 하려 합니다. 모든 Query 와 모든 Key의 attention score 로 초기화된 행렬이 만들어지게 되고 이것을 attention matrix 라 합니다. 이 행렬은 query 와 key 의 연관관계를 알려주는 행렬입니다. $\sqrt{d_k}$ 값으로 정규화를 하고 softmax 를 적용해 확률 분포로 만들어 줍니다.

Attention Matrix 출처 : Visualizing and Explaining Transformer Models From the Ground Up, 2023

encoder, decoder 에는 여러 encoder block, decoder block 이 있기 때문에 매 block 마다 self-attention 이 계산이 됩니다. 아래 그림을 통해 매 레이어마다 모든 토큰의 query, key 들이 서로 서로 attention score 를 연산하는 것을 확인할 수 있습니다.

출처 : Interpretable Multi-Head Self-Attention Architecture for Sarcasm Detection in Social Media, 2021

이제 attention matrix 에 Value 를 곱하면 과정만 남았습니다. 이것이 의미하는 바는 attention score 에 따른 weighted sum 입니다. 각 토큰들의 Value 를 attention score 에 비례해서 더해주는 것입니다.

self-attention 을 통해 Query 는 자기 자신을 포함해서 모든 Key 들 과의 중요도를 알 수 있게 됩니다. 그 것에 비례하여 각 토큰에 대응하는 value 벡터들을 weighted sum을 하게 되면 Query(=개별 토큰) 에 대응하는 새로운 의미 정보를 얻는 벡터를 얻게 되는 것입니다. 이것이 바로 self-attention 의 핵심입니다.

출처 : Attention Is All You Need, 2017

아래 그림을 다시 살펴 보겠습니다. Thinking과 Machines 라는 토큰들에 대해서 attention score 들을 구한 후 그 값을 확률값으로 바꾸어주고 그 것에 비례하여 value 값들을 혼합해주면 됩니다.

이것이 의미하는 바는 각 token 들이 문맥 내에서 다른 토큰들과 얼마만큼의 관계를 모두 고려하여(QK^T), 그것에 비례하여 새로운 표현 값(V) 을 얻어내겠다는 것 입니다.

출처 : Alammar, J (2018). The Illustrated Transformer [Blog post]. Retrieved from https://jalammar.github.io/illustrated-transformer/

self-attention 코드

아래와 같이 작업할 수 있습니다.

class Head(nn.Module):
    """ one head of self-attention """

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False) # weight matrix
        self.query = nn.Linear(n_embd, head_size, bias=False) # weight matrix
        self.value = nn.Linear(n_embd, head_size, bias=False) # weight matrix

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):

        B,T,C = x.shape
        k = self.key(x)   # (B,T,C)
        q = self.query(x) # (B,T,C)

        # compute attention scores ("affinities")
        **wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
        wei = F.softmax(wei, dim=-1) # (B, T, T)**
        wei = self.dropout(wei)

        # perform the weighted aggregation of the values
        v = self.value(x) # (B,T,C)
        **out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)**
        return out

Multi-Head Attention

마지막으로 Query, Key, Value 를 하나의 벡터로 표현하는 것보다는 여러 벡터로 표현을해서 self-attention을 하게 되면 여러 의미 정보를 담을 수 있게 됩니다. 이것이 Multi-Head Attention ****입니다. 이렇게 각각 생성된 Value 행렬을 Concat 한 후 linear 층($W^O)$을 거치게 되면 마찬가지로 input 과 같은 동등한 output shape을 얻을 수 있으며, 그 이후의 프로세스는 위와 일치합니다.

출처 : Attention Is All You Need, 2017

출처 : Attention Is All You Need, 2017

Multi-Head Attention 코드

num_heads 수 만큼 self-attention 을 거친후 proj Linear$(W^O)$ 를 통과하는 것을 확인할 수 있습니다.

class MultiHeadAttention(nn.Module):
    """ multiple heads of self-attention in parallel """

    def __init__(self, num_heads, head_size):
        super().__init__()
        **self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)**
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        **out = torch.cat([h(x) for h in self.heads], dim=-1)**
        out = self.dropout(**self.proj(out)**)
        return out

앞서 말씀 드렸듯, input과 output의 shape이 일치하기 때문에 동등한 Encoder block 에서의 연산을 계속해서 수행할 수 있습니다. self-attention 외에도 FFN, DropOut, Add & Norm, Residual 부분이 있습니다. Attention Is All You Need 논문에서는 총 6번의 Encoder block과 6번의 Decoder block을 거칩니다.

Masked Multi-Head Attention

이제 decoder 부분으로 넘어가겠습니다. Masked Multi-Head Attention 이 기존 Multi-Head Attention 과 다른 점은 문자 그대로 Mask가 붙었다는 점입니다.

다음 단어를 맞추는 방식으로 학습을 하기 위해서는 미래에 나올 토큰들을 고려하지 않고 지금까지 나온 토큰들을 토대로 생성을 해야 합니다. 이 때문에 attention matrix 에서 생성 시점 이후의 토큰들에 대해서 Mask를 씌워 attention이 고려되지 않도록 해야만 올바른 계산을 할 수 있습니다.

Masked Multi-Head Attention 코드

softmax 이전에 0으로 초기화된 부분에 -inf 로 변경하여 확률 값이 0이 되도록 구현할 수 있습니다.

class Head(nn.Module):
    """ one head of self-attention """

    def __init__(self, head_size):
        super().__init__()

                ...
            self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
                ...	
            

        def forward(self, x):
            ...
            wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
            **wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)**
            wei = F.softmax(wei, dim=-1) # (B, T, T)
      wei = self.dropout(wei)

****      # perform the weighted aggregation of the values
      v = self.value(x) # (B,T,C)
      out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
      return out

Cross Multi-Head Attention

Cross Multi-Head Attention은 self-attention 을 계산을 하되, Query 부분은 decoder 단에 들어오는 값을 사용하고, Key 와 Value 부분은 Encode를 통과한 것을 사용하는 방식입니다. Masked Multi-Head Attention 을 거쳤기 때문에 앞서 생성된 토큰들 만을 고려하고 encoder 를 통과한 출력값을 모두 활용하는 방식으로 새로운 출력을 만들어낼 수 있습니다.

출처 : Alammar, J (2018). The Illustrated Transformer [Blog post]. Retrieved from https://jalammar.github.io/illustrated-transformer/

여기까지 해서 간단하게 Transformer 를 살펴보았습니다.

이제 nanoGPT 를 만들어 보겠습니다.

nanoGPT

GPT는 Generative Pre-trained Transformer 의 약자입니다.

생성에 특화된 Decoder 부분만을 이용해 만든 Pre-trained transformer 라 할 수 있습니다.

기존 Transformer 의 Decoder block 구성요소 중에서 Cross Multi-Head Attention ****을 제거하여 구성합니다.

출처 : Intrusion Detection Method Using Bi-Directional GPT for In-Vehicle Controller Area Networks, 2021

load data

셰익스피어의 글과 문체를 학습시켜 보겠습니다. 아래와 같은 데이터가 준비되었다고 해보겠습니다. 단지 txt 파일을 불러오기만 하면 됩니다.

# read it in to inspect it
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

First Citizen:
You are all resolved rather to die than to famish?

All:
Resolved. resolved.

First Citizen:
First, you know Caius Marcius is chief enemy to the people.

...

vocab

vocab 을 만들어보겠습니다. 총 61 개의 character 로 구성이 되어 있습니다.

chars = sorted(list(set(text)))
vocab_size = len(chars)
print(''.join(chars))
print(vocab_size)

!&',-.:;?ABCDEFGHIJKLMNOPQRSTUVWYabcdefghijklmnopqrstuvwxyz
61

문자를 인덱스로, 인덱스를 문자로 매핑하는 dict 를 통해 encode 하고 decode를 할 수 있습니다.

# create a mapping from characters to integers
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }

encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string

print(f"encode('hii there') : {encode('hii there')}")
print(f"decode(encode('hii there')) : {decode(encode('hii there'))}")

encode('hii there') : [8, 9, 9, 0, 16, 8, 7, 14, 7]
decode(encode('hii there')) : hii there

train data setting

학습 데이터는 이전 입력값들을 바탕으로 다음 토큰을 예측하는 형태로 구성합니다.

x = train_data[:block_size]
y = train_data[1:block_size+1]
for t in range(block_size):
    context = x[:t+1]
    target = y[t]
    print(f"when input is {context} the target: {target}")

when input is tensor([788]) the target: 149
when input is tensor([788, 149]) the target: 140
when input is tensor([788, 149, 140]) the target: 1
when input is tensor([788, 149, 140,   1]) the target: 726
when input is tensor([788, 149, 140,   1, 726]) the target: 370
when input is tensor([788, 149, 140,   1, 726, 370]) the target: 680
when input is tensor([788, 149, 140,   1, 726, 370, 680]) the target: 996
when input is tensor([788, 149, 140,   1, 726, 370, 680, 996]) the target: 6

출처 : Intrusion Detection Method Using Bi-Directional GPT for In-Vehicle Controller Area Networks, 2021

BigramLanguageModel (overview)

먼저 decoder block 을 구현하지 않은 상태에서 모델을 만들어 보겠습니다. BigramLanguageModel 모델은 token_embedding_table 로 부터 임베딩 값을 얻습니다. 그리고 idx 값은 (B, T) shape 의 텐서로서 forward 를 통해 통과한 로짓이 다음 토큰의 예측값에 대한 로짓이 됩니다.

import torch
import torch.nn as nn
from torch.nn import functional as F
torch.manual_seed(1337)

class BigramLanguageModel(nn.Module):

    def __init__(self, vocab_size):
        super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):

        # idx and targets are both (B,T) tensor of integers
        logits = self.token_embedding_table(idx) # (B,T,C)
        
        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss
    
        
    def generate(self, idx, max_new_tokens):
        # idx is (B, T) array of indices in the current context
        for _ in range(max_new_tokens):
            # get the predictions
            logits, loss = self(idx)

            **# focus only on the last time step**
            logits = logits[:, -1, :] # becomes (B, C)

            # apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1) # (B, C)

            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)

            # append sampled index to the running sequence
            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)

        return idx

m = BigramLanguageModel(vocab_size)
logits, loss = m(xb, yb)
print(logits.shape)
print(loss)

print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist()))

torch.Size([256, 19])
tensor(2.9300, grad_fn=<NllLossBackward0>

 du u ts'iwonabbbbbacbbbbu clsbbsohdsbbbbnteln'nlbbssbnlblnnnlcrlconnuaacrmnnTtoiiocrnilwdmuu tolnnni

forward 함수의 경우 decoder 내부 요소를 구현하지 않은 상태입니다. 예측된 logit 은 vocab 에 있는 토큰 들 중 하나로 분류하기 때문에 손실함수는 cross entropy 를 사용합니다.

generate 함수의 경우 시퀀스의 마지막 토큰만을 이용해 logit 을 계산해 생성합니다. 아직은 무의미한 정보를 생성합니다.

train

아래처럼 배치를 두고 학습할 수 있습니다.

batch_size = 32
for steps in range(100): # increase number of steps for good results... 
    
    # sample a batch of data
    xb, yb = get_batch('train')

    # evaluate the loss
    logits, loss = m(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

BigramLanguageModel (Decoder Blocks)

Decoder 부분은 Decoder block 을 여러개 쌓아서 구성하면 됩니다.

먼저 하이퍼 파라메터를 선언해줍니다.

Decoder block 은 4개를 쌓아줍니다. embedding은 64 차원, head 수는 4 입니다.

# hyperparameters
batch_size = 16 # how many independent sequences will we process in parallel?
block_size = 32 # what is the maximum context length for predictions?
max_iters = 5000
eval_interval = 100
learning_rate = 1e-3
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
dropout = 0.0 
n_embd = 64 # embedding 차원  
n_head = 4 # head number
n_layer = 4 # Decoder block number

BigramLanguageModel 내부에 Decoder block 을 num_layer 수 만큼 쌓습니다.

# super simple bigram model
class BigramLanguageModel(nn.Module):

    def __init__(self):
        super().__init__()
                ...
                self.position_embedding_table = nn.Embedding(block_size, n_embd)
                self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
                ...

        ...

각 Block 내부엔 Multi-Head Attention 이 존재합니다.

class Block(nn.Module):
    """ Transformer block: communication followed by computation """

    def __init__(self, n_embd, n_head):
        # n_embd: embedding dimension, n_head: the number of heads we'd like
        super().__init__()
        head_size = n_embd // n_head
        **self.sa = MultiHeadAttention(n_head, head_size)**
        self.ffwd = FeedFoward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x)) # residual
        x = x + self.ffwd(self.ln2(x)) # residual
        return x

MultiHeadAttention 클래스는 num_heads 만큼 각각 Head(self-attention) 을 통과합니다.

class MultiHeadAttention(nn.Module):
    """ multiple heads of self-attention in parallel """

    def __init__(self, num_heads, head_size):
        super().__init__()
        **self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])**
        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        **out = torch.cat([h(x) for h in self.heads], dim=-1)**
        out = self.dropout(**self.proj(out)**)
        return out

각 Head 는 masked 로직을 포함하여 self-attention 을 수행하는 걸 확인할 수 있습니다.

class Head(nn.Module):
    """ one head of self-attention """

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        **self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))**

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B,T,C = x.shape
        k = self.key(x)   # (B,T,C)
        q = self.query(x) # (B,T,C)

        # compute attention scores ("affinities")
        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
        **wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)**
        wei = F.softmax(wei, dim=-1) # (B, T, T)
        wei = self.dropout(wei)

        # perform the weighted aggregation of the values
        v = self.value(x) # (B,T,C)
        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
        return out

Train

모델을 선언하고 미리 구성해놓은 데이터를 학습합니다.

model = BigramLanguageModel()
m = model.to(device)
# print the number of parameters in the model
print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')

# create a PyTorch optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

for iter in range(max_iters):

    # every once in a while evaluate the loss on train and val sets
    if iter % eval_interval == 0 or iter == max_iters - 1:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    # sample a batch of data
    xb, yb = get_batch('train')

    **# evaluate the loss
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()**

생성 결과

얼추 셰익스피어와 비슷한 텍스트를 생성한 것을 확인할 수가 있습니다.

context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=2000)[0].tolist()))

BUCKINGHAM:
Thou, his lost to betchsed ingron
So you not me, slate ine, but that peerd
But connurdererisHards a wall that gleed:
Op, weive? But not bsorn of
Good, imis it to death! God's but the optrange your your many
his med and ne'er, no not, care
Annds my someds not; my wints go yoursay.

NORDIZANUO:
Yet while;
So marry upon il his been, live is me, wonst da my son,
that do he doung I ress patius;
Your enague should movence aste-timal
Trought that God, hor densless not;
As most; queech seeme
Misel, lible imbraid, to yet stand,
theirs is trurness away, and thou drue
And my chooly being Roman, for that shall to: where his norber rump matter you,
Who sweet not forth goty me would sgeen,
To last needs them set it! this you your bearth,
And there is me that bity the face here buty gates
I was have and Bacleingmery, is Dolk they graveight:
I did conteremp; I benam you.

마치며

이렇게 해서 nanoGPT으로 만들어봤습니다. GPT-1, GPT-2, GPT-3, InstructGPT 모델들도 어느정도 차이가 있지만 큰 틀에서 이러한 방식으로 학습하게 됩니다.

추가로 Transformer 에 관심있으신 분들은 LLM Visualization (링크 : https://bbycroft.net/llm) 에서 실제 self-attention이 어떻게 이뤄지고 영향을 주고 받는지 살펴보시면 큰 도움을 받으실 수 있을거라 생각합니다.

출처 : https://bbycroft.net/llm

이 글을 읽는 분들께 조금이나마 도움이 되길 바라며 이상으로 마치도록 하겠습니다.

DALL·E 3

참고자료

•

Let's build GPT: from scratch, in code, spelled out.(링크 : https://www.youtube.com/watch?v=kCc8FmEb1nY&t=5329s)

•

The Illustrated GPT-2 (Visualizing Transformer Language Models)(링크 : https://jalammar.github.io/illustrated-gpt2/)

•

LLM Visualization(링크 : https://bbycroft.net/llm)

•

[NLP 논문 구현] pytorch로 구현하는 Transformer (Attention is All You Need) (링크 : https://cpm0722.github.io/pytorch-implementation/transformer)

•

Visualizing and Explaining Transformer Models From the Ground Up, 2023 (링크 : https://deepgram.com/learn/visualizing-and-explaining-transformer-models-from-the-ground-up)

Subscribe to 'KPMG Lighthouse'

Subscribe to my site to be the first to receive notifications and emails about the latest updates, including new posts.
Join Slashpage and subscribe to 'KPMG Lighthouse'!