[Encoder Only Model with Random Data]

Transformer Model을 완전히 구현하기 위해서는 Encoder, Decoder를 모두 구현해야 한다. 이 포스팅에서는 그 중 Encoder만 우선적으로 구현한다.

 

우측의 Input을 받는 부분이 Encoder이다.

 

이전 포스팅에서 Config class를 만들었는데, 이번 포스팅에서 Config class를 다음과 같이 업그레이드한다.

@dataclass
class Config:
    vocab_size : int = 1000
    d_model : int = 512
    block_size : int = None
    num_heads : int = None
    key_dim : int = None
    dff : int = None
    dropout_rate : float = None
    num_layers : int = None

 

마찬가지로 이전 포스팅을 참고하여 TokenEmbedding class를 구현한다.

class TokenEmbedding(keras.layers.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.tok_embedding = keras.layers.Embedding(config.vocab_size, config.d_model)

    def call(self, x):
        x = self.tok_embedding(x)

        return x

 

Positional Embedding

Input Embedding을 진행하였으면 다시 거기에 Positional Embedding을 추가적으로 진행해야 한다.

여태까지 진행한 token embedding은, 들어온 input 값을 index로 하여 해당 row에 해당하는 값들을 반환하는 것이다. 앞으로 구현할 positional embedding은 0~block size-1 값을 index로 하여 해당 row에 해당하는 값들을 반환한다.

 

positional embedding을 구현하기 전에, 앞서 정의한 Config class를 다음과 같이 수정해야 한다.

@dataclass
class Config:
    vocab_size : int = 1000
    d_model : int = 512
    block_size : int = 10
    num_heads : int = None
    key_dim : int = None
    dff : int = None
    dropout_rate : float = None
    num_layers : int = None

block_size 값이 None에서 10으로 바뀌었다. 이는 한 문장이 10개의 token을 가지게 하겠다는 의미이다.

 

이후, 앞서 Token Embedding만 고려하여 구현했던 TokenEmbedding class를 Positional Embedding 내용까지 포함할 수 있도록 다음과 같이 개선한다.

class TokenAndPositionEmbedding(keras.layers.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.tok_embedding = keras.layers.Embedding(config.vocab_size, config.d_model)
        self.pos_embedding = keras.layers.Embedding(config.block_size, config.d_model)

    def call(self, x):
        b, t = x.shape
        tok_emb = self.tok_embedding(x)
        pos_emb = self.pos_embedding(keras.ops.arange(self.config.block_size))

        return tok_emb + pos_emb

vocabulary 1000개, 10개 token으로 이루어진 2개의 문장을 대상으로 테스트해 보면 다음과 같은 결과를 얻을 수 있다.

x = np.random.randint(0, 1000, (2, 10))
temp = TokenAndPositionEmbedding(Config)
temp(x)

이 값이 token embedding, positional embedding을 모두 완료한 결과값이 된다. 

 

The global self attention layer

참조 그림에서 Multi-Head Attention에 표시된 부분으로, tensorflow에서는 이 부분을 'global self attention layer'라고 부른다. 정확히는 Multi-Head Attention 중 자기 자신에게서만 input을 받는 부분을 'global self attention layer'라고 부르는 것이다. 

그림에서는 하단 두 개의 Multi-Head Attention이 이에 해당한다. 다른 곳에서 input을 받는 중단 Multi-Head Attention은 'cross attention layer'라고 부른다. 

 

self attention layer는 다시 global self attention layer, casual self attention layer로 나뉜다. Input의 모든 information을 참조해서 output을 만드는 것을 global self attention layer, 반면 n번째 ouput token이 1~n번째 input token만을 참조하는 것을 casual self attention layer라고 부른다. 

 

이번 포스팅에서는 global self attention layer를 구현하는 과정을 알아본다. 

 

먼저, global self attention layer에 맞게 Config class를 다음과 같이 개선한다. num_head 값을 None이 아닌 값(여기서는 우선적으로 1)로 설정하고, key_dim 값을 None이 아닌 d_model 값과 같도록 바꾼다.

@dataclass
class Config:
    vocab_size : int = 1000
    d_model : int = 512
    block_size : int = 10
    num_heads : int = 1 
    key_dim : int = 512 
    dff : int = None
    dropout_rate : float = None
    num_layers : int = None

 

감사하게도 keras에는 이미 Multi-head Attention이 구현되어 있고 우리는 이를 적절히 가져다 쓰면 된다. 이 때 매개변수로 넘겨줘야 할 값이 num_head, key_dim이기 때문에 Config class를 위와 같이 개선한 것이다. 

mha = keras.layers.MultiHeadAttention(Config.num_heads, Config.key_dim)

keras로부터 MultiHeadAttention instance를 생성하였으면, query, key, value 3개 값을 전달해야 한다.(앞서 참조 그림에서 Multi-Head Attention 블록으로 들어가는 input이 3개였음을 기억해 보자)

mha(temp(x), temp(x), temp(x))

temp(x)는 앞서 초기 input값을 Token, Positional Embedding을 거친 값이다. 여기까지 data 형태는 [2,10,512], 즉 512개 벡터로 이루어진 10개의 token으로 구성된 2개의 문장으로 이해하면 된다.

 

이  일련의 과정을 다음과 같은 class로 구현할 수 있다.

class GlobalSelfAttention(keras.layers.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.mha = keras.layers.MultiHeadAttention(config.num_heads, config.key_dim)
        self.layernorm = keras.layers.LayerNormalization()
        self.add = keras.layers.Add()

    def call(self, x):
        attn_output = self.mha(
            query=x,
            value=x,
            key=x,
            use_causal=False)
        x = self.add([x, attn_output])
        x = self.layernorm(x)
        return x

GlobalSelfAttention(Config)

 

The Feedforward Network

Feedforward Network를 구현하기 전에 또 다시 Config를 조금 개선한다.

@dataclass
class Config:
    vocab_size : int = 1000
    d_model : int = 512
    block_size : int = 10
    num_heads : int = 4
    key_dim : int = int(d_model/num_heads)
    dff : int = 2048
    dropout_rate : float = None
    num_layers : int = None

앞선 단락에서는 num_head를 일단 1로 설정하였지만, 여기서는 4로 지정하였다. num_head는 일종의 "관점"역할을 한다고 이해하면 좋다. 따라서 여기서는 주어진 데이터를 4개의 관점에서 attention 계산을 하겠다는 의미이다.

 

key_dim 값은 이전에는 num_head가 1이었기 때문에 d_model 값과 같았지만 이제는 d_model/num_head 값이 되었다. num_head만큼 나누어 연산한 후에 합쳐져서 d_model과 같은 값이 될 것이다.

 

새로이 dff(dimension of feedforward network)값을 2048로 설정하였다. 이 값은 꼭 2048이어야 할 필요는 없으나, 경험적으로 가장 효율적으로 model이 돌아가는 숫자이다. 

 

다음과 같이 feedforward network class를 구현한다.

class FeedForward(keras.layers.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.seq = keras.Sequential([
            keras.layers.Dense(config.dff, activation='relu'),
            keras.layers.Dense(config.d_model),
            keras.layers.Dropout(config.dropout_rate)
        ])
        self.add = keras.layers.Add()
        self.layer_norm = keras.layers.LayerNormalization()

    def call(self, x):
        x = self.add([x, self.seq(x)])
        x = self.layer_norm(x)
        return x

self.seq = keras.Sequential([
            keras.layers.Dense(config.dff, activation='relu'),
            keras.layers.Dense(config.d_model),
            keras.layers.Dropout(config.dropout_rate)
        ])

이 중 constructor 내의 위 코드 부분이 neural network가 구현된 부분이다. Dense layer 통과*2 + Dropout을 순서대로 통과한다. 앞서 Config에서 설정했던 값들을 복기해 보면, data가 Dense Layer를 통과하며 2048(dff)개가 되고, 다시 Dense Layer를 통과하며 512(d_model)개가 되며, 여기에 Dropout(앞선 포스팅 참조, dropout_rate만큼 랜덤한 값을 0으로 바꾸고 다른 값들을 조정함)이 적용되는 흐름이 구현된 것이다.

 

call 함수를 보면 이 class가 하는 역할(즉 feedforward network가 하는 역할)은 윗 문단에서 서술한 과정을 거친 data와  원본 data를 더하고, 이를 layer normalization해서 반환하는 역할임을 알 수 있다. 

 

The Encoder Layer

Input을 Multi-head Attention으로, 다시 그 값을 feedforward network로 보내는 이 일련의 과정의 encoder layer다. 그림 박스 맨 하단에 조그맣게 써진 글씨를 확인하자. 이 과정을 여러 번 거치는 것이 Encoder이다.

class EncoderLayer(keras.layers.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.gsa = GlobalSelfAttention(config)
        self.ff = FeedForward(config)

    def call(self, x):
        x = self.gsa(x)
        x = self.ff(x)
        return x

앞서 MultiHead Attention - 이 경우에는 global self attention - 과 feedforward network를 모두 구현하였고, 이를 그냥 class에 구현해 주기만 하면 된다. 

 

Encoder

encoder layer를 반복하는 것이 encoder라고 했다. 

 

먼저 encoder 구현에 들어가기 앞서 Config를 완전히 완성시켜야 한다. 

@dataclass
class Config:
    vocab_size = 1000
    d_model = 512
    block_size = 5
    num_heads = 4
    key_dim = int(d_model/num_heads)
    dff = 2048
    dropout_rate = 0.1
    num_layers = 1

dropout_rate와 num_layers가 추가되었다. dropout_rate는 dropout에 사용될 값, num_layers는 encoder layer를 반복시킬 횟수이다.

class Encoder(keras.layers.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embedding = TokenAndPositionEmbedding(config)
        self.enc_layers = [EncoderLayer(config) for _ in range(config.num_layers)]
        self.dropout = keras.layers.Dropout(config.dropout_rate)

    def call(self, x):
        x = self.embedding(x)  # Shape `(batch_size, seq_len, d_model)`.
        x = self.dropout(x)
        for i in range(self.config.num_layers):
            x = self.enc_layers[i](x)

        return x

위 class가 우리가 여태 했던 것들을 총집합하여 Encoder를 구현한 것이다. 

1. Input이 들어오면 그것을 (Token+Positioning) Embedding 하고

2. 그것을 dropout하고

3. 다시 그것을 만큼 num_layers 만큼 encoder layer를 통과시킨다.

 

이 과정을 모두 통과하면 input의 의미를 알게 되고, 이 값을 cross attention layer의 input으로 보내 주는 것이다.