토큰 임베딩 및 위치 임베딩 레이어: 토큰화된 텍스트는 토큰 임베딩 레이어와 위치 임베딩 레이어를 통과하여, 시퀀스에서 토큰의 위치를 캡처합니다. 이는 단어 순서를 이해하는 데 중요합니다.
트랜스포머 블록: 모델은 12개의 트랜스포머 블록을 포함하며, 각 블록은 여러 레이어로 구성됩니다. 이 블록은 다음 시퀀스를 반복합니다:
마스크드 멀티-헤드 어텐션: 모델이 입력 텍스트의 다양한 부분에 동시에 집중할 수 있게 합니다.
레이어 정규화: 훈련을 안정화하고 개선하기 위한 정규화 단계입니다.
피드 포워드 레이어: 어텐션 레이어에서 정보를 처리하고 다음 토큰에 대한 예측을 수행하는 역할을 합니다.
드롭아웃 레이어: 이 레이어는 훈련 중 무작위로 유닛을 드롭하여 과적합을 방지합니다.
최종 출력 레이어: 모델은 4x50,257 차원의 텐서를 출력하며, 여기서 50,257은 어휘의 크기를 나타냅니다. 이 텐서의 각 행은 모델이 시퀀스에서 다음 단어를 예측하는 데 사용하는 벡터에 해당합니다.
목표: 목표는 이러한 임베딩을 가져와 다시 텍스트로 변환하는 것입니다. 구체적으로, 출력의 마지막 행은 이 다이어그램에서 "forward"로 표시된 다음 단어를 생성하는 데 사용됩니다.
코드 표현
import torchimport torch.nn as nnimport tiktokenclassGELU(nn.Module):def__init__(self):super().__init__()defforward(self,x):return0.5* x * (1+ torch.tanh(torch.sqrt(torch.tensor(2.0/ torch.pi)) *(x +0.044715* torch.pow(x, 3))))classFeedForward(nn.Module):def__init__(self,cfg):super().__init__()self.layers = nn.Sequential(nn.Linear(cfg["emb_dim"], 4* cfg["emb_dim"]),GELU(),nn.Linear(4* cfg["emb_dim"], cfg["emb_dim"]),)defforward(self,x):return self.layers(x)classMultiHeadAttention(nn.Module):def__init__(self,d_in,d_out,context_length,dropout,num_heads,qkv_bias=False):super().__init__()assert d_out % num_heads ==0,"d_out must be divisible by num_heads"self.d_out = d_outself.num_heads = num_headsself.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dimself.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)self.out_proj = nn.Linear(d_out, d_out)# Linear layer to combine head outputsself.dropout = nn.Dropout(dropout)self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))defforward(self,x):b, num_tokens, d_in = x.shapekeys = self.W_key(x)# Shape: (b, num_tokens, d_out)queries = self.W_query(x)values = self.W_value(x)# We implicitly split the matrix by adding a `num_heads` dimension# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)values = values.view(b, num_tokens, self.num_heads, self.head_dim)queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)keys = keys.transpose(1, 2)queries = queries.transpose(1, 2)values = values.transpose(1, 2)# Compute scaled dot-product attention (aka self-attention) with a causal maskattn_scores = queries @ keys.transpose(2, 3)# Dot product for each head# Original mask truncated to the number of tokens and converted to booleanmask_bool = self.mask.bool()[:num_tokens,:num_tokens]# Use the mask to fill attention scoresattn_scores.masked_fill_(mask_bool, -torch.inf)attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)attn_weights = self.dropout(attn_weights)# Shape: (b, num_tokens, num_heads, head_dim)context_vec = (attn_weights @ values).transpose(1, 2)# Combine heads, where self.d_out = self.num_heads * self.head_dimcontext_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)context_vec = self.out_proj(context_vec)# optional projectionreturn context_vecclassLayerNorm(nn.Module):def__init__(self,emb_dim):super().__init__()self.eps =1e-5self.scale = nn.Parameter(torch.ones(emb_dim))self.shift = nn.Parameter(torch.zeros(emb_dim))defforward(self,x):mean = x.mean(dim=-1, keepdim=True)var = x.var(dim=-1, keepdim=True, unbiased=False)norm_x = (x - mean) / torch.sqrt(var + self.eps)return self.scale * norm_x + self.shiftclassTransformerBlock(nn.Module):def__init__(self,cfg):super().__init__()self.att =MultiHeadAttention(d_in=cfg["emb_dim"],d_out=cfg["emb_dim"],context_length=cfg["context_length"],num_heads=cfg["n_heads"],dropout=cfg["drop_rate"],qkv_bias=cfg["qkv_bias"])self.ff =FeedForward(cfg)self.norm1 =LayerNorm(cfg["emb_dim"])self.norm2 =LayerNorm(cfg["emb_dim"])self.drop_shortcut = nn.Dropout(cfg["drop_rate"])defforward(self,x):# Shortcut connection for attention blockshortcut = xx = self.norm1(x)x = self.att(x)# Shape [batch_size, num_tokens, emb_size]x = self.drop_shortcut(x)x = x + shortcut # Add the original input back# Shortcut connection for feed forward blockshortcut = xx = self.norm2(x)x = self.ff(x)x = self.drop_shortcut(x)x = x + shortcut # Add the original input backreturn xclassGPTModel(nn.Module):def__init__(self,cfg):super().__init__()self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])self.drop_emb = nn.Dropout(cfg["drop_rate"])self.trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ inrange(cfg["n_layers"])])self.final_norm =LayerNorm(cfg["emb_dim"])self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)defforward(self,in_idx):batch_size, seq_len = in_idx.shapetok_embeds = self.tok_emb(in_idx)pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))x = tok_embeds + pos_embeds # Shape [batch_size, num_tokens, emb_size]x = self.drop_emb(x)x = self.trf_blocks(x)x = self.final_norm(x)logits = self.out_head(x)return logitsGPT_CONFIG_124M ={"vocab_size":50257,# Vocabulary size"context_length":1024,# Context length"emb_dim":768,# Embedding dimension"n_heads":12,# Number of attention heads"n_layers":12,# Number of layers"drop_rate":0.1,# Dropout rate"qkv_bias":False# Query-Key-Value bias}torch.manual_seed(123)model =GPTModel(GPT_CONFIG_124M)out =model(batch)print("Input batch:\n", batch)print("\nOutput shape:", out.shape)print(out)
GELU 활성화 함수
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04classGELU(nn.Module):def__init__(self):super().__init__()defforward(self,x):return0.5* x * (1+ torch.tanh(torch.sqrt(torch.tensor(2.0/ torch.pi)) *(x +0.044715* torch.pow(x, 3))))
목적 및 기능
GELU (가우시안 오류 선형 유닛): 모델에 비선형성을 도입하는 활성화 함수입니다.
부드러운 활성화: 음수 입력을 0으로 만드는 ReLU와 달리, GELU는 입력을 출력으로 부드럽게 매핑하여 음수 입력에 대해 작은 비영 값도 허용합니다.
수학적 정의:
FeedForward 레이어 내부의 선형 레이어 이후 이 함수를 사용하는 목적은 선형 데이터를 비선형으로 변경하여 모델이 복잡하고 비선형적인 관계를 학습할 수 있도록 하는 것입니다.
FeedForward 신경망
행렬의 형태를 더 잘 이해하기 위해 주석으로 형태가 추가되었습니다:
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04classFeedForward(nn.Module):def__init__(self,cfg):super().__init__()self.layers = nn.Sequential(nn.Linear(cfg["emb_dim"], 4* cfg["emb_dim"]),GELU(),nn.Linear(4* cfg["emb_dim"], cfg["emb_dim"]),)defforward(self,x):# x shape: (batch_size, seq_len, emb_dim)x = self.layers[0](x)# x shape: (batch_size, seq_len, 4 * emb_dim)x = self.layers[1](x)# x shape remains: (batch_size, seq_len, 4 * emb_dim)x = self.layers[2](x)# x shape: (batch_size, seq_len, emb_dim)return x # Output shape: (batch_size, seq_len, emb_dim)
목적 및 기능
위치별 FeedForward 네트워크: 각 위치에 대해 별도로 동일하게 두 개의 완전 연결 네트워크를 적용합니다.
레이어 세부사항:
첫 번째 선형 레이어: 차원을 emb_dim에서 4 * emb_dim으로 확장합니다.
GELU 활성화: 비선형성을 적용합니다.
두 번째 선형 레이어: 차원을 다시 emb_dim으로 줄입니다.
보시다시피, Feed Forward 네트워크는 3개의 레이어를 사용합니다. 첫 번째는 선형 레이어로, 선형 가중치(모델 내부에서 훈련할 매개변수)를 사용하여 차원을 4배로 곱합니다. 그런 다음, GELU 함수가 모든 차원에서 사용되어 더 풍부한 표현을 포착하기 위한 비선형 변화를 적용하고, 마지막으로 또 다른 선형 레이어가 원래 차원 크기로 되돌리기 위해 사용됩니다.
다중 헤드 주의 메커니즘
이것은 이전 섹션에서 이미 설명되었습니다.
목적 및 기능
다중 헤드 자기 주의: 모델이 토큰을 인코딩할 때 입력 시퀀스 내의 다양한 위치에 집중할 수 있게 합니다.
주요 구성 요소:
쿼리, 키, 값: 입력의 선형 프로젝션으로, 주의 점수를 계산하는 데 사용됩니다.
헤드: 병렬로 실행되는 여러 주의 메커니즘(num_heads), 각기 축소된 차원(head_dim)을 가집니다.
주의 점수: 쿼리와 키의 내적을 계산하여 스케일링 및 마스킹합니다.
마스킹: 미래의 토큰에 주의를 기울이지 않도록 인과 마스크가 적용됩니다(자기 회귀 모델인 GPT에 중요).
주의 가중치: 마스킹되고 스케일된 주의 점수의 소프트맥스입니다.
컨텍스트 벡터: 주의 가중치에 따라 값의 가중 합입니다.
출력 프로젝션: 모든 헤드의 출력을 결합하는 선형 레이어입니다.
이 네트워크의 목표는 동일한 컨텍스트 내에서 토큰 간의 관계를 찾는 것입니다. 또한, 토큰은 과적합을 방지하기 위해 서로 다른 헤드로 나뉘며, 최종적으로 각 헤드에서 발견된 관계는 이 네트워크의 끝에서 결합됩니다.
또한, 훈련 중에 인과 마스크가 적용되어 특정 토큰에 대한 관계를 찾을 때 이후의 토큰이 고려되지 않으며, 드롭아웃도 적용되어 과적합을 방지합니다.
레이어 정규화
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04classLayerNorm(nn.Module):def__init__(self,emb_dim):super().__init__()self.eps =1e-5# Prevent division by zero during normalization.self.scale = nn.Parameter(torch.ones(emb_dim))self.shift = nn.Parameter(torch.zeros(emb_dim))defforward(self,x):mean = x.mean(dim=-1, keepdim=True)var = x.var(dim=-1, keepdim=True, unbiased=False)norm_x = (x - mean) / torch.sqrt(var + self.eps)return self.scale * norm_x + self.shift
목적 및 기능
레이어 정규화: 배치의 각 개별 예제에 대해 특징(임베딩 차원) 전반에 걸쳐 입력을 정규화하는 데 사용되는 기술입니다.
구성 요소:
eps: 정규화 중 0으로 나누는 것을 방지하기 위해 분산에 추가되는 작은 상수(1e-5)입니다.
scale 및 shift: 정규화된 출력을 스케일하고 이동할 수 있도록 모델이 학습할 수 있는 매개변수(nn.Parameter)입니다. 각각 1과 0으로 초기화됩니다.
정규화 과정:
평균 계산(mean): 임베딩 차원(dim=-1)에 걸쳐 입력 x의 평균을 계산하며, 브로드캐스팅을 위해 차원을 유지합니다(keepdim=True).
분산 계산(var): 임베딩 차원에 걸쳐 x의 분산을 계산하며, 차원을 유지합니다. unbiased=False 매개변수는 분산이 편향 추정기를 사용하여 계산되도록 보장합니다(샘플이 아닌 특징에 대해 정규화할 때 적합한 N으로 나누기).
정규화(norm_x):x에서 평균을 빼고 분산에 eps를 더한 값의 제곱근으로 나눕니다.
스케일 및 이동: 정규화된 출력에 학습 가능한 scale 및 shift 매개변수를 적용합니다.
목표는 동일한 토큰의 모든 차원에서 평균이 0이고 분산이 1이 되도록 하는 것입니다. 이는 딥 뉴럴 네트워크의 훈련을 안정화하기 위한 것으로, 훈련 중 매개변수 업데이트로 인한 네트워크 활성화의 분포 변화인 내부 공변량 이동을 줄이는 것을 의미합니다.
토큰 임베딩 (tok_emb): 토큰 인덱스를 임베딩으로 변환합니다. 상기 참고로, 이는 어휘의 각 토큰의 각 차원에 주어진 가중치입니다.
위치 임베딩 (pos_emb): 임베딩에 위치 정보를 추가하여 토큰의 순서를 캡처합니다. 상기 참고로, 이는 텍스트에서의 위치에 따라 토큰에 주어진 가중치입니다.
드롭아웃 (drop_emb): 정규화를 위해 임베딩에 적용됩니다.
트랜스포머 블록 (trf_blocks): 임베딩을 처리하기 위한 n_layers 개의 트랜스포머 블록 스택입니다.
최종 정규화 (final_norm): 출력 레이어 이전의 레이어 정규화입니다.
출력 레이어 (out_head): 최종 은닉 상태를 어휘 크기로 투영하여 예측을 위한 로짓을 생성합니다.
이 클래스의 목표는 시퀀스에서 다음 토큰을 예측하기 위해 언급된 모든 다른 네트워크를 사용하는 것입니다. 이는 텍스트 생성과 같은 작업에 기본적입니다.
얼마나 많은 트랜스포머 블록이 지정된 대로 사용될 것인지 주목하고, 각 트랜스포머 블록이 하나의 멀티 헤드 어텐션 네트워크, 하나의 피드 포워드 네트워크 및 여러 정규화를 사용하는지 주목하십시오. 따라서 12개의 트랜스포머 블록이 사용되면 이를 12로 곱합니다.
또한, 출력 이전에 정규화 레이어가 추가되고, 결과를 적절한 차원으로 얻기 위해 마지막에 선형 레이어가 적용됩니다. 각 최종 벡터가 사용된 어휘의 크기를 가지는 이유는 어휘 내 가능한 각 토큰에 대한 확률을 얻으려고 하기 때문입니다.
훈련할 매개변수 수
GPT 구조가 정의되면 훈련할 매개변수 수를 파악할 수 있습니다:
GPT_CONFIG_124M ={"vocab_size":50257,# Vocabulary size"context_length":1024,# Context length"emb_dim":768,# Embedding dimension"n_heads":12,# Number of attention heads"n_layers":12,# Number of layers"drop_rate":0.1,# Dropout rate"qkv_bias":False# Query-Key-Value bias}model =GPTModel(GPT_CONFIG_124M)total_params =sum(p.numel() for p in model.parameters())print(f"Total number of parameters: {total_params:,}")# Total number of parameters: 163,009,536
모델이 이전과 같은 다음 토큰을 예측하는 경우, 출력에서 마지막 토큰 값을 가져오기만 하면 됩니다(예측된 토큰의 값이 될 것이므로). 이는 어휘의 각 항목에 대한 값이 될 것이며, 그런 다음 softmax 함수를 사용하여 차원을 확률로 정규화하여 합이 1이 되도록 하고, 가장 큰 항목의 인덱스를 가져옵니다. 이 인덱스는 어휘 내의 단어 인덱스가 됩니다.
defgenerate_text_simple(model,idx,max_new_tokens,context_size):# idx is (batch, n_tokens) array of indices in the current contextfor _ inrange(max_new_tokens):# Crop current context if it exceeds the supported context size# E.g., if LLM supports only 5 tokens, and the context size is 10# then only the last 5 tokens are used as contextidx_cond = idx[:,-context_size:]# Get the predictionswith torch.no_grad():logits =model(idx_cond)# Focus only on the last time step# (batch, n_tokens, vocab_size) becomes (batch, vocab_size)logits = logits[:,-1,:]# Apply softmax to get probabilitiesprobas = torch.softmax(logits, dim=-1)# (batch, vocab_size)# Get the idx of the vocab entry with the highest probability valueidx_next = torch.argmax(probas, dim=-1, keepdim=True)# (batch, 1)# Append sampled index to the running sequenceidx = torch.cat((idx, idx_next), dim=1)# (batch, n_tokens+1)return idxstart_context ="Hello, I am"encoded = tokenizer.encode(start_context)print("encoded:", encoded)encoded_tensor = torch.tensor(encoded).unsqueeze(0)print("encoded_tensor.shape:", encoded_tensor.shape)model.eval()# disable dropoutout =generate_text_simple(model=model,idx=encoded_tensor,max_new_tokens=6,context_size=GPT_CONFIG_124M["context_length"])print("Output:", out)print("Output length:", len(out[0]))
