class TinyTransformer(nn.Module):
def __init__(self):
super().__init__()
# setting the constructor for the initial values that we are every gonna need for the training of the data
self.char_embedding = nn.Embedding(65, 64)
self.pos_embedding = nn.Embedding(64, 64)
self.query = nn.Linear(64, 64)
self.key = nn.Linear(64, 64)
self.value = nn.Linear(64, 64)
self.mask = torch.tril(torch.ones(64, 64))
# these are for changing the dimensions we are doing this to enlarge the matrix as to make it of higher resolution so as to make the
# data and weights more refined
self.ff1 = nn.Linear(64, 128)
# this is to join them back again
self.ff2 = nn.Linear(128, 64)
self.output_head = nn.Linear(64, 65)
self.norm1 = nn.LayerNorm(64)
self.norm2 = nn.LayerNorm(64)
self.out_proj = nn.Linear(64, 64)
def forward(self, x):
# feed forward function
x = self.char_embedding(x) + self.pos_embedding(torch.arange(64))
# this is the start of the attention stuff i am writing this as a way to separate the code in section inside a functions
#
Q = self.query(x)
Q = Q.view(32, 64, 2, 32)
Q = Q.transpose(1, 2)
K = self.key(x)
K = K.view(32, 64, 2, 32)
K = K.transpose(1, 2)
V = self.value(x)
V = V.view(32, 64, 2, 32)
V = V.transpose(1, 2)
A = (Q @ K.transpose(-2, -1)) / 32**0.5
A = A.masked_fill(self.mask == 0, float("-inf"))
At = A.softmax(dim=-1)
# the -1 this is just to tell the
output = At @ V
output = output.transpose(1, 2).contiguous().view(32, 64, 64)
output = self.out_proj(output)
# this is where the attention ends and we start with the feed forward thing that will give us the predictions
# added another form of normalization below to improve accuracy the first time the loss function reached 1.8 max now after adding the
# below line it reached to like 1.5 something
x = x + output
x = self.norm1(x)
output = self.ff1(x)
output = torch.relu(output)
output = self.ff2(output)
x = x + output # ← merge back into main flow
x = self.norm2(x)
x = self.output_head(x)
return x
이 코드는 기본적으로 AI 분야의 누구에게나 트랜스포머를 훈련시키기 위한 표준 템플릿입니다.
저는 여기에 있는 한 줄을 이해하고 싶은데, 그 뒤에는 정말 흥미로운 역사가 있습니다.
x = x + output
우리는 왜 이것을 하고 있는가 - X=X+output?
신경망은 역전파 과정을 통해 학습하는데, 이는 기본적으로 필터를 변경함으로써 올바른 예측에 더 가까워지게 하는 변화를 찾는 것을 의미합니다. 그런데 여기에는 특별한 문제가 있습니다. 한 층에서 다른 층으로 이동할수록 기울기가 점점 작아진다는 점인데, 이는 계산이 점점 더 어려워지고 계산 비용이 증가하기 때문에 큰 문제입니다. 이는 기본적으로 편미분의 연쇄 법칙 때문에 발생합니다. 그런데 이 방식은 어떻게 이 문제를 해결할까요?
*이 문제의 역사 *
여기에서 이 논문을 읽으세요 --
이 논문은 Microsoft 연구팀에 의해 작성되었으며, 기본적으로 '더 많은 것이 항상 더 좋은 것은 아니다'라는 문제를 어떻게 해결했는지에 관한 것입니다. 이 논문 이전에 딥러닝 모델을 훈련하는 경우, 모델의 깊이가 깊어질수록 즉 레이어가 많을수록 더 많은 오류도 발생했습니다. 그리고 그것은 큰 문제였고 사람들은 그것을 해결하는 방법을 몰랐습니다. 왜냐하면 한편으로는 더 나은 이해의 깊이가 있었고, 다른 한편으로는 더 많은 오류를 얻게 되는 이 문제도 있었기 때문입니다.
솔루션
이제 우리는 모든 계산 후에 얻은 컨텍스트 행렬에 (내 경우에는) 원래 임베딩 벡터를 추가하는 것이 해결책이라고 생각할 수도 있습니다. 그리고 그렇게 생각하는 것은 맞지만, 여러분이 생각하는 이유는 아닙니다. 논문 자체에서도 이것이 이 문제의 이유가 아니라고 말하고 있습니다.
우리는 이 최적화 어려움이 다음과 같을 가능성이 낮다고 주장합니다.
기울기 소실(vanishing gradients)로 인해 발생합니다.
왜 그럴까요? - 기울기 소실 문제를 최소화하고 방지하기 위한 작업들이 이루어졌기 때문에 우리가 그 의심을 제거한 이유입니다. 이러한 작업은 배치 정규화(batch normalization)와 같은 도구와 여기서는 ReLu를 사용하여 수행됩니다. 다음은 우리 코드에서 이를 수행하는 방법입니다. -
x = self.norm1(x) # the batch normalization equivalent in transformers
output = self.ff1(x)
output = torch.relu(output) # another way to solve the vanishing gradient problem
output = self.ff2(output)
x = x + output # ← merge back into main flow
x = self.norm2(x)
x = self.output_head(x)
보시다시피, this, that, these는 기울기 소실 문제를 해결하지만, x=x+output을 제거하면 결과가 더 나빠질 것입니다. 자, 한번 해볼까요? 알겠습니다. --
이것은 정상적으로 진행하고 아무것도 변경하지 않은 경우입니다. 이제 한 가지만 변경해 보겠습니다. 바로 x=x+output 줄을 제거하는 것입니다. 그게 전부이며, 그것이 손실 함수에 어떤 영향을 미치는지 살펴보겠습니다.
그래서 손실 함수가 이 한 줄만으로 1.70에서 2.47로 점프했는데, 별거 아닌 것처럼 보일 수 있지만, 이것은 단순화를 위해 1층 모델에 불과하며 더 많은 층을 추가할수록 오류도 더 커진다는 점을 기억하세요. 제 요점을 확실히 하기 위해, 여기서 약간의 작은 조정을 통해 살아있는 그래디언트를 보여드리겠습니다 -
step 44500, loss: 2.5130
char_embedding.weight grad_norm: 0.006248
pos_embedding.weight grad_norm: 0.005838
ff1.weight grad_norm: 0.024721
ff2.weight grad_norm: 0.053932
output_head.weight grad_norm: 0.163109
norm1.weight grad_norm: 0.007271
norm2.weight grad_norm: 0.024594
step 45000, loss: 2.4751
char_embedding.weight grad_norm: 0.005574
pos_embedding.weight grad_norm: 0.005913
ff1.weight grad_norm: 0.023506
ff2.weight grad_norm: 0.056331
output_head.weight grad_norm: 0.161182
norm1.weight grad_norm: 0.007898
norm2.weight grad_norm: 0.020992
step 45500, loss: 2.4623
char_embedding.weight grad_norm: 0.006224
pos_embedding.weight grad_norm: 0.006075
ff1.weight grad_norm: 0.025461
ff2.weight grad_norm: 0.051210
output_head.weight grad_norm: 0.145062
norm1.weight grad_norm: 0.008452
norm2.weight grad_norm: 0.018521
step 46000, loss: 2.4764
char_embedding.weight grad_norm: 0.006709
pos_embedding.weight grad_norm: 0.006148
ff1.weight grad_norm: 0.026940
ff2.weight grad_norm: 0.057071
output_head.weight grad_norm: 0.163159
norm1.weight grad_norm: 0.008988
norm2.weight grad_norm: 0.025112
step 46500, loss: 2.4746
char_embedding.weight grad_norm: 0.006127
pos_embedding.weight grad_norm: 0.006181
ff1.weight grad_norm: 0.025931
ff2.weight grad_norm: 0.056799
output_head.weight grad_norm: 0.158272
norm1.weight grad_norm: 0.008369
norm2.weight grad_norm: 0.025981
그래서 제가 말하고 있던 것은, 이것이 기울기 소실 문제를 해결하는 것처럼 보이지만 실제로는 전혀 그렇지 않다는 점입니다.
이것이 실제로 하는 일은 훨씬 더 흥미로운 것입니다.-
모든 레이어와 비선형 함수는 어떤 변화를 일으키며, 이러한 변화는 매우 빠르게, 정말 빠르게 축적됩니다. 20개의 레이어 정도에서는 작동할 수도 있는데, 레이어 수가 많긴 하지만, 여기서 50개 레이어 정도로 넘어가면 복잡성이 급격히 증가합니다. 그리고 이러한 "작은" 변화들은 그 자체로는 아주 아주 작더라도 값에 큰 영향을 줄 수 있으며, 그 결과 생성되는 값은 원래 값과 완전히 다를 수 있습니다, 정말 많이 다를 수 있습니다. 예를 들어보겠습니다 -
5, 3, 8 --->5, 3, 8--->0.3, 0.01, 0.2 그리고 이것들이 전혀 기울기 소실 같은 것이 아니라 중간에 발생하는 작은 변화들 때문임을 알게 됩니다. 그렇다면 여기에 원본을 추가하면 어떻게 될까요?
[5, 3, 8] + [0.3, 0.01, 0.2] = [5.3, 3.01, 8.2]
그래서 이 결과는 원본과 매우 가깝죠? 그것이 잔차(residual)의 주요 아이디어이며, 잔차 알고리즘도 많지만 간단하게 하기 위해 우리는 좋은 옛날 덧셈을 고수할 것이고, 솔직히 이 방법이 더 낫습니다.