[LLM] Attention is All You Need 의 Base Transformer 파라미터 수 계산

오랜만에 논문을 다시 읽다가 파라미터 수에 꽂혔다.

 

Attention is ALl You Need 논문에 나온 실험 결과 테이블

 

여기서 베이스 모델의 파라미터가 65M이라고 나와있길래, 재미삼아 GPT에게 물어봤는데, 당연히 위에 사진만 보여주면 환각 현상 때문에 65M이라고 대답한다. 그래서 구체적인 수치를 주고 다시 대답시켜봤는데 아래처럼 답변했다.

 

 

계산에 사용된 파이썬 코드를 보면 다음과 같다.

V = 37000  # Vocabulary Size
L = 512    # Sequence Length
d_model = 512  # Model dimension
h = 8  # Number of attention heads
d_ff = 2048  # Feedforward dimension
N = 6  # Number of layers in encoder and decoder

# 1. Embedding Layers
token_embedding = V * d_model
positional_embedding = L * d_model * 2  # Encoder + Decoder

# 2. Multi-Head Attention (Encoder + Decoder)
# Each attention block has Query, Key, Value projections + Output projection
attention_params_per_layer = 3 * d_model * d_model + d_model * d_model
encoder_attention = N * attention_params_per_layer
decoder_attention = N * attention_params_per_layer
cross_attention = N * attention_params_per_layer  # Cross-Attention in Decoder

# 3. Feedforward Layers (Encoder + Decoder)
feedforward_params_per_layer = 2 * d_model * d_ff
encoder_feedforward = N * feedforward_params_per_layer
decoder_feedforward = N * feedforward_params_per_layer

# 4. Final Linear Projection
final_projection = d_model * V

# 총 파라미터 수 계산
total_params = (
    token_embedding + positional_embedding +
    encoder_attention + decoder_attention + cross_attention +
    encoder_feedforward + decoder_feedforward +
    final_projection
)

print(total_params)

 

Vocab size, d_model 같은 파라미터는 모두 논문에 나와있는 수치를 그대로 인용, GPT에 제공했다. Vocab size의 경우엔 논문에서

We trained on the standard WMT 2014 English-German dataset consisting of about 4.5 million sentence pairs. Sentences were encoded using byte-pair encoding [3], which has a shared sourcetarget vocabulary of about 37000 tokens.

 

대략 37000개의 토큰을 얻었다는 말에서 차용했다. 이는 상기했던 테이블에서 WMT 2014 Englsih-German 데이터셋을 기준으로 작성됐다고 언급한 것에서 기인한다.

 

코드는 정상적인 것 같은데 수치 차이가 꽤 커서 내가 뭔가를 오해한건가, 싶어 이글 저글 뒤져봤는데 어떤 곳은 아래처럼 부정확한 소리를 하기도 했다.

 

또 아직까지도 GPT가 수식에서 약한 모습들을 보인다고 느껴지는게, 위 대화 이후에 재질문했을 때는

이런 식으로 숫자 연산을 할 때 틀린 답을 내놓기 일쑤다. (프롬프트에서 왜 65M이 안되는지, 뭘 실수한건지 다시 생각해보고 다시 계산해달라고 했다.)

 

아무튼 그래서 믿을 만한 것이 없다고 판단, 내가 직접 손과 계산기를 사용해 파라미터 계산을 진행해보고, 신뢰할 수 있는 글들 몇 개와 비교해봤다. 내가 계산한 내용에는 Layer Normalization, Positional Embedding, Bias 등 너무 사소한 값이라 계산에 큰 영향이 없는 값들은 제외했다. 또한 Notation은 논문에 나온 그대로 사용하도록 하겠다.

https://towardsdatascience.com/how-to-estimate-the-number-of-parameters-in-transformer-models-ca0f57d8dff0/

How to Estimate the Number of Parameters in Transformer models | Towards Data Science

 

실제로 계산한 내용

 

1. 임베딩 레이어

출처 : https://petuum.medium.com/embeddings-a-matrix-of-meaning-4de877c9aa27

 

임베딩 레이어의 파라미터는 단순히 V * d_model 행렬에 불과하다. 또한, 논문에선 인풋 임베딩 행렬과 아웃풋 임베딩 행렬을 공유한다고 했다. 따라서, 트랜스포머에서 임베딩 행렬에 사용된 파라미터는 V * d_model, 즉 V를 37000으로 가정한다면 37000 * 512 = 18,944,000 이 된다. 

 

2. Multi-Head Attention

출처 : 논문

 

Multi-Head Attention에서의 계산은 약간 복잡하다. 헤드의 개수로 Query, Key, Value로 나눠주는 부분이 있기 때문인데, 이는 어차피 나중에 Concatenate 하므로 파라미터 계산상에서는 굳이 나눠줄 이유가 없다. (512/8*3*8 이나 512*3이나 결과는 같기 때문.) 따라서 계산할 때 $Wq,Wk,Wv$ 모두 나누지 않고 사용하겠다. $Wq = Wk = Wv$ 이므로, 하나의 $W$라고 생각하면 이는 512*512 크기의 행렬을 갖는다. 그리고 마지막에 Output Projection을 해주는 행렬 $Wo$가 있으므로, 이 또한 512*512 크기의 파라미터가 발생한다. 따라서 하나의 Multi-Head Attention 단에선 4*512*512 = 1,048,576 개의 파라미터를 갖는다.

 

3. Feed Forward Neural Network

 

출처 : https://www.researchgate.net/figure/Feedforward-neural-network-FFNN-The-circles-represent-network-layers-and-the-solid_fig1_333119454

트랜스포머의 FFNN은 한 개의 히든레이어를 가지며, 이는 2048의 크기다. 또한 인풋과 아웃풋의 크기는 모두 512이므로, FFNN은 2*2048*512 = 2,097,512 개의 파라미터를 갖는다.

 

4. Total Multi-Head Attention + FFNN

 

위에서 계산한 것은 단순히 각각의 MHA와 FFNN 이었고, 전체 개수를 구해줘야한다. 인코더에 있던 디코더에 있던 하나의 블록이 갖는 파라미터 개수는 동일하므로, 이를 고려하지 않고 계산해보면 레이어 개수가 각각 6개이므로

Encoder : 6 * (#MHA + #FFNN) = 18874368

Decoder : 6 * (#MHA + #MHA + #FFNN) = 25165824

가 나오게 된다.

5. Linear Before Softmax

마지막으로 생성할 단어를 예측하는 단계에서 Linear Matrix가 한 번 더 사용된다. 이는 처음 임베딩 레이어와 같은 크기의 가중치를 가지므로, 18,944,000 을 그대로 사용할 수 있다.

 

Total Parameters

위에서 계산한 파라미터를 모두 더해보면 약 82M, 처음에 GPT가 계산해준 숫자가 나오게 된다. 어떻게 보면 안해도 될 짓을 한 셈이긴 한데, 뒤에 엉뚱한 말들을 해갖고 신뢰도가 확 떨어져 어쩔 수 없이 손으로 해봤다. 어쨌든 논문에서 말한 65M과는 차이가 꽤 크게 나서, 이 원인이 무엇인지 찾아보려고 했으나 이렇게 다 검증한 글을 찾지 못했는데, 중요한 키포인트를 마침내 찾았다.

 

논문의 3.4를 자세히 읽어보면

In our model, we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation

 

이럴수가, Linear Before Softmax 에서도 Input Embedding과 같은 가중치를 공유하고 있었다! 따라서 82M에서 대략 19M 정도를 빼면 63M 정도가 나오고, 내가 생략했던 Layer Normalization의 파라미터, bias 등을 모두 더하면 65M 정도가 나올 것이다.

 

이 부분을 분명히 읽었는데도 생각하지 못 했던 게, 내 상식으로는 임베딩 레이어와 최종 예측 단계에서 같은 가중치를 공유할 수가 없었기 때문이다. 임베딩은 단어 인덱스를 받아서 긴 표현 벡터로 만드는 것인데, 최종 단계 예측에선 마찬가지로 수많은 계산을 거쳐 나온 표현 벡터가 단어 인덱스로 변환되는 것이므로 이 두 표현벡터의 괴리가 있을 거라고 자연스럽게 인지했다. 하지만 이를 공유했다니.. 앞으로 공부할 다른 모델에서도 이를 유념하고 읽어봐야겠다.

 

PS. GPT가 이 사실을 좀 미리 알려줬다면 몇 시간 고생을 안 했어도 됐는데 킹받게 내가 말하니 그제서야 그렇다고 한다...

 

 

추가) 인풋 임베딩과 출력 직전의 Linear 행렬을 공유하는 방식은 https://arxiv.org/pdf/1608.05859 이 논문에서 가져온 아이디어라고 언급되어있다. 읽어보면 이 방식은 perplexity를 크게 감소시켰다고 저자들이 언급하고 있다.

추가2) https://nlp.seas.harvard.edu/2018/04/03/attention.html#embeddings-and-softmax

The Annotated Transformer

 

논문에 파이토치로 주석을 달아 설명해주는 하버드대의 칼럼이다. BERT에 인용될 정도로 굉장히 잘 구현되어 있는데, 이 부분의 Final Linear 부분을 보면 

class Generator(nn.Module):
    "Define standard linear + softmax generation step."

    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return log_softmax(self.proj(x), dim=-1)

 

이렇게 Generator라는 클래스를 만들고 Final Linear를 임베딩 레이어와는 별개로 모델을 만든 것을 볼 수 있다.

 

그리고 이후 마지막에 주석으로

이런 식으로 코딩하여 임베딩을 공유하게 설정할 수도 있다고 한다. 하지만 굳이 구현하지 않았다고 하는데, 이유까진 명시하지 않은 것으로 보아 이렇게 딥하게 생각할 필요는 없다고 본 것 같다. (이걸 공유할 때와 안 할 때 d_model * Vocab_size 만큼 파라미터 수가 차이가 나는데, 그렇게 크지 않다고 본 것 같다.)

따라서, Orignial Transformer는 이 Output Layer를 공유한다고 하지만, 실제로는 그렇지 않게 구현하는 사람들도 많은 것으로 보이고, 성능에 실제로 어느 정도 향상이 있었는지도 파악이 안 된다. 

https://datascience.stackexchange.com/questions/84930/weights-shared-by-different-parts-of-a-transformer-model

Weights shared by different parts of a transformer model

이 글의 답변을 보면 인풋 임베딩 / 아웃풋 임베딩 / 출력 임베딩 모두 공유하고 하지 않고는 개발자의 마음이라고 언급되어 있고, 독일어-영어가 아닌 중국어-영어 처럼 사용하는 문자가 아예 다른 경우엔 인풋 임베딩과 아웃풋 임베딩 모두 공유하지 않을 수 있다고 한다.