Trong chươngđầu của dự án này, tôi sẽ hướng dẫn bạn qua quá trình forward propagation và tính toán giá trị loss. Mục tiêu của chúng ta là hiểu sâu hơn về phần cốt lõi, yếu tố chính (“backbone”) của kiến trúc Transformers trong LLAMA2.
Về cơ bản, kiến trúc cốt lõi của LLAMA2 rất giống với kiến trúc Transformers (như hình ở trên).
Trong quá trình viết code, tôi sẽ tiếp tục trình bày các ý tưởng chính của từng phần trong kiến trúc này. Tuy nhiên, nếu bạn muốn có cái nhìn chi tiết và sâu sắc hơn, bạn có thể xem phần Appendix, trong đó tôi sẽ giải thích từng phần trong kiến trúc này với chi tiết ở mức độ Character Level để bạn có thể hiểu rõ hơn về nó.
from datasets import load_dataset
dataset = load_dataset("roneneldan/TinyStories")
datasetRepo card metadata block was not found. Setting CardData to empty.DatasetDict({
train: Dataset({
features: ['text'],
num_rows: 2119719
})
validation: Dataset({
features: ['text'],
num_rows: 21990
})
})Trong dự án này, tôi sử dụng bộ dữ liệu “Tiny Datasets”, một tập dữ liệu chứa các câu tiếng Anh đơn giản, thích hợp cho trẻ 3-4 tuổi có khả năng đọc dễ dàng. Bộ dữ liệu này bao gồm hơn 2 triệu câu cho phần huấn luyện (trainset) và gần 22.000 câu cho phần thử nghiệm (valid set).
sample = 20
subset_dataset = dataset['train'][:sample]['text']
for example in subset_dataset[:1]:
print(example)One day, a little girl named Lily found a needle in her room. She knew it was difficult to play with it because it was sharp. Lily wanted to share the needle with her mom, so she could sew a button on her shirt.
Lily went to her mom and said, "Mom, I found this needle. Can you share it with me and sew my shirt?" Her mom smiled and said, "Yes, Lily, we can share the needle and fix your shirt."
Together, they shared the needle and sewed the button on Lily's shirt. It was not difficult for them because they were sharing and helping each other. After they finished, Lily thanked her mom for sharing the needle and fixing her shirt. They both felt happy because they had shared and worked together.Ở đây, tôi sẽ chỉ sử dụng một lượng nhỏ dữ liệu, để bắt đầu xây dựng model của mình. Mục tiêu ban đầu là tạo ra một model hoàn chỉnh. Sau khi chúng ta đạt được một model đáng tin cậy, chúng ta có thể bắt đầu sử dụng toàn bộ hoặc một lượng lớn hơn của dữ liệu để train model.
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
tokenizer.pad_token = tokenizer.eos_token
# Tokenize the text data in the new subset dataset with padding and truncation
tokenized_dataset = tokenizer(
subset_dataset,
return_tensors='pt',
padding=True, # Enable padding
truncation=True # Enable truncation
)
tokenized_dataset['input_ids'][:1]tensor([[ 3198, 1110, 11, 257, 1310, 2576, 3706, 20037, 1043, 257,
17598, 287, 607, 2119, 13, 1375, 2993, 340, 373, 2408,
284, 711, 351, 340, 780, 340, 373, 7786, 13, 20037,
2227, 284, 2648, 262, 17598, 351, 607, 1995, 11, 523,
673, 714, 34249, 257, 4936, 319, 607, 10147, 13, 198,
198, 43, 813, 1816, 284, 607, 1995, 290, 531, 11,
366, 29252, 11, 314, 1043, 428, 17598, 13, 1680, 345,
2648, 340, 351, 502, 290, 34249, 616, 10147, 1701, 2332,
1995, 13541, 290, 531, 11, 366, 5297, 11, 20037, 11,
356, 460, 2648, 262, 17598, 290, 4259, 534, 10147, 526,
198, 198, 41631, 11, 484, 4888, 262, 17598, 290, 384,
19103, 262, 4936, 319, 20037, 338, 10147, 13, 632, 373,
407, 2408, 329, 606, 780, 484, 547, 7373, 290, 5742,
1123, 584, 13, 2293, 484, 5201, 11, 20037, 26280, 607,
1995, 329, 7373, 262, 17598, 290, 18682, 607, 10147, 13,
1119, 1111, 2936, 3772, 780, 484, 550, 4888, 290, 3111,
1978, 13, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256]])Vì dữ liệu hiện tại của chúng ta là văn bản, chúng ta cần sử dụng một tokenizer để chuyển đổi dữ liệu thành định dạng số để mô hình có thể học được.
Tôi đã quyết định sử dụng tokenizer từ mô hình EleutherAI/gpt-neo-125M vì tôi thấy nó phù hợp và tiện lợi cho giai đoạn bắt đầu. Sử dụng tokenizer từ LLAMA2 có thể đòi hỏi đăng nhập vào Hugging Face để truy cập, nhưng tôi cho rằng, ít nhất ở giai đoạn đầu, chúng ta nên giữ mọi thứ đơn giản. Trong các chương tiếp theo, tôi có thể xem xét cải tiến bằng cách sử dụng tokenizer từ LLAMA2 hoặc thậm chí tạo tokenizer từ đầu nếu cần.
data = tokenized_dataset['input_ids']
x = data[:, :-1].contiguous()
y = data[:, 1:].contiguous()
x.shape, y.shape(torch.Size([20, 218]), torch.Size([20, 218]))Tương tự như nhiều dự án khác, chúng ta cần hai thành phần chính: một là đầu vào (input) và hai là nhãn (label). Trong dự án này, đầu vào (input) sẽ bao gồm một chuỗi các từ được cung cấp, và nhãn (label) là từ tiếp theo sẽ xuất hiện trong chuỗi đó.
for i in range(5):
print(f"Input: {x[0, :i+1]} --> Labels: {y[0,i]}")Input: tensor([3198]) --> Labels: 1110
Input: tensor([3198, 1110]) --> Labels: 11
Input: tensor([3198, 1110, 11]) --> Labels: 257
Input: tensor([3198, 1110, 11, 257]) --> Labels: 1310
Input: tensor([3198, 1110, 11, 257, 1310]) --> Labels: 2576vocab_size = tokenizer.vocab_size
sequence_len = x.size(1)
print(f"Vocab Size: {vocab_size}")
print(f"Max Sequence Length: {sequence_len}")Vocab Size: 50257
Max Sequence Length: 218Bởi vì chúng ta đang sử dụng EleutherAI/gpt-neo-125M làm tokenizer, Vocab Size của chúng ta sẽ là 50257. Đồng thời, độ dài tối đa của các chuỗi (max sequence length) mà chúng ta đang xử lý sẽ tạm thời là độ dài của chuỗi “x” trong chiều thứ hai, tức là “x.size(1)”.
import torch.nn as nn
import torch
# Output Embedding
n_embd = 36
wte = nn.Embedding(vocab_size, n_embd) # word to embedding
token_embd = wte(x)
token_embd.shapetorch.Size([20, 218, 36])# Positional Encoding
position = nn.Embedding(sequence_len, n_embd)
position_embd = position(torch.arange(sequence_len))
position_embd.shapetorch.Size([218, 36])x_embd = token_embd + position_embd
x_embd.shapetorch.Size([20, 218, 36])Ý tưởng cốt lõi của chúng ta là sử dụng một vectơ đặc trưng để biểu diễn mỗi từ trong bộ từ vựng (vocab size) của chúng ta. Đồng thời, chúng ta cũng áp dụng nguyên tắc này cho việc biểu diễn từng vị trí của các từ trong câu.
Hãy xem xét từ “bàn” trong vocab size của chúng ta. Khi chúng ta thực hiện quá trình embedding, từ “bàn” này được chuyển đổi thành một vectơ đặc trưng duy nhất. Điều thú vị là vectơ này có khả năng biểu diễn từ “bàn” trong nhiều ngữ cảnh khác nhau - có thể là một cái bàn gỗ, một cuộc họp trên bàn, một cái bàn đạp xe, hoặc thậm chí một phần của chân. Điều này giúp model của chúng ta hiểu và biểu diễn nhiều sắc thái và ngữ nghĩa của từ “bàn” trong các tình huống khác nhau.
Ngoài ra, chúng ta cũng tạo ra các biểu diễn độc lập cho từng vị trí của các từ trong câu. Ví dụ, từ “bàn” trong câu “Tôi vừa mua một cái bàn làm việc” sẽ có một biểu diễn vector vị trí riêng biệt, khác với từ “bàn” trong câu “Cái bàn kia thật đẹp”. Mặc dù cả hai trường hợp này có cùng nghĩa cho từ “bàn,” nhưng do vị trí của nó trong câu khác nhau (vị trí thứ 6 vs thứ 2), nó sẽ được biểu diễn theo cách khác nhau.
Tóm lại, chúng ta sử dụng vectơ đặc trưng để biểu diễn từ (word embedding) và vị trí của từ trong câu (position embedding), cho phép model hiểu và biểu diễn ngữ nghĩa và ngữ cảnh của các từ một cách hiệu quả trong các ngữ cảnh khác nhau.
# Norm before calculate attention
norm = nn.LayerNorm(n_embd)
x_embd_norm = norm(x_embd)Trong kiến trúc Transformer, có một sự thay đổi quan trọng mà chúng ta đang áp dụng so với cách truyền thống. Thay vì tính toán output của attention trước rồi mới normalize output đó, chúng ta sẽ tiến hành ngược lại: chúng ta sẽ thực hiện normalization trước, sau đó sử dụng kết quả này như input để tính attention.
Sự thay đổi này không chỉ áp dụng trong kiến trúc LLAMA2 mà còn xuất hiện rộng rãi trong hầu hết các kiến trúc Transformers hiện đại.
# Multi-head Attention
n_head = 4
head_size = n_embd // n_head
opt_size = n_head * head_size # output size
head_size, opt_size(9, 36)Wqkv = nn.Linear(n_embd, 3 * opt_size)
qkv = Wqkv(x_embd_norm)
from einops import rearrange
qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, h = n_head)
q, k, v = qkv.unbind(dim=2)
q.shape, k.shape, v.shape(torch.Size([20, 218, 4, 9]),
torch.Size([20, 218, 4, 9]),
torch.Size([20, 218, 4, 9]))Mục tiêu của đoạn code trên là đi tính query, key và value, còn chúng có ý nghĩa là gì thì hãy cũng xem xét ví dụ dưới đây.
Hãy tưởng tượng rằng bạn là một nhà báo nổi tiếng đang thực hiện một cuộc phỏng vấn với một ngôi sao nổi tiếng, và bạn muốn thu thập thông tin quan trọng từ cuộc trò chuyện đó.
Key có thể coi như danh sách câu hỏi bạn chuẩn bị trước cuộc phỏng vấn. Mỗi câu hỏi là một Key, và mỗi câu hỏi sẽ tập trung vào một khía cạnh cụ thể của cuộc trò chuyện. Ví dụ, một Key có thể là “Bạn đã từng giành giải Oscar chưa?”
Value là câu trả lời mà ngôi sao đưa ra cho từng câu hỏi. Mỗi câu trả lời chứa thông tin quan trọng về cuộc trò chuyện, và nó sẽ được lưu trữ và sử dụng sau này khi bạn cần nắm bắt thông tin cụ thể từ cuộc phỏng vấn. Chúng ta có thể coi câu trả lời này là “value” của câu hỏi.
Query là cách bạn đặt câu hỏi hoặc tìm kiếm thông tin trong cuộc phỏng vấn. Khi bạn muốn biết điều gì đó cụ thể hoặc muốn nắm bắt một thông tin quan trọng từ cuộc trò chuyện, bạn sẽ đặt câu hỏi hoặc tạo một “Query” riêng. Ví dụ, “Giới thiệu về những vai diễn nổi bật nhất của bạn?” có thể là một Query.
Khi bạn đặt một câu hỏi (Query), model sẽ so sánh nó với danh sách các câu hỏi trước đó (Key) và quyết định câu trả lời nào (Value) chứa thông tin phù hợp nhất với câu hỏi của bạn. Điều này giống như việc bạn tập trung vào câu hỏi cụ thể nào trong cuộc trò chuyện để thu thập thông tin bạn cần.
# Masked multi-head
import math
softmax_scale = 1.0 / math.sqrt(q.shape[-1])
scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
# Masking
mask = torch.triu(torch.full((sequence_len, sequence_len), -10000.0), 1)
scores = scores + mask
attention = torch.softmax(scores, dim=-1)Ý tưởng chính ở đây là ta đang xây dựng một hệ thống dự đoán từ tiếp theo trong câu văn dựa trên các từ đã xuất hiện trước đó. Để thực hiện điều này, mỗi từ cần được dự đoán sẽ đóng vai trò là một Query, và các từ đã xuất hiện trước đó sẽ đóng vai trò là Key. Chúng ta sau đó so sánh tỉ lệ phù hợp giữa các query và các key để xác định những từ nào quan trọng hơn và sẽ được sử dụng để dự đoán từ tiếp theo.
Đặc điểm quan trọng là chúng ta sử dụng một cơ chế “masking” để che đi thông tin của các từ đứng sau từ cần dự đoán. Điều này giúp mô hình tập trung vào việc sử dụng thông tin từ các từ trước đó để dự đoán từ tiếp theo mà không bị ảnh hưởng bởi các từ sau đó trong câu.
attn_out = torch.einsum("bhts,bshd->bthd", attention, v)
attn_out = rearrange(attn_out, "... h d -> ... (h d)")
out_proj = nn.Linear(opt_size, n_embd)
attn_out = out_proj(attn_out)
# Add residual
residual = x_embd
attn_out += residual
attn_out.shapetorch.Size([20, 218, 36])Quá trình này được gọi là Self-Attention, vì điểm đặc biệt là giá trị của key và value được tạo ra từ chính bản thân câu văn hoặc chuỗi đầu vào.
Self-Attention là một khía cạnh quan trọng trong kiến trúc Transformer, vì nó cho phép mô hình xác định mức độ quan trọng của các từ hoặc phần tử trong câu văn đối với từ hoặc phần tử cụ thể khác trong câu văn đó. Điều này giúp mô hình xử lý ngôn ngữ tự nhiên một cách linh hoạt và hiểu ngữ cảnh một cách tốt hơn.
Khác với Self Attention, Cross Attention đặt ra sự khác biệt bằng cách mà giá trị key và value không đến từ bản thân câu văn, mà chúng đến từ nguồn thông tin bên ngoài.
Hãy cùng tưởng tượng một ví dụ để hiểu rõ hơn: Self Attention có thể được tưởng tượng như bạn đặt ra những câu hỏi cho chính bản thân mình (Key) và tự mình trả lời chúng (Value). Trong một ngày khác, bạn tiếp tục đặt ra những câu hỏi mới, nhưng lần này bạn không tự trả lời mà bạn xem xét những câu hỏi bạn đã đặt trước đó (Query) và xem câu hỏi nào (Key) phù hợp nhất với câu hỏi hiện tại (Query), sau đó bạn sẽ dựa vào câu trả lời đó (Value).
Cross Attention, ngược lại, có thể được tưởng tượng như bạn tham gia vào một cuộc phỏng vấn với một diễn viên nổi tiếng. Trong cuộc phỏng vấn, bạn đặt ra những câu hỏi (Key) và ghi chép lại câu trả lời của người diễn viên đó (Value). Sau đó, vào một thời điểm khác, bạn tự đặt ra những câu hỏi cho chính mình (Query) và xem xét xem câu hỏi nào bạn đã đặt trong cuộc phỏng vấn (Key) phù hợp nhất với câu hỏi của bạn (Query) và sử dụng câu trả lời từ cuộc phỏng vấn đó (Value).
Tuy nhiên trong dự án hiện tại của chúng ta, không cần thiết phải sử dụng Cross Attention vì chúng ta không cần sử dụng các key và value từ nguồn bên ngoài. Do đó, chúng ta có thể bỏ qua bước Cross Attention và tiếp tục với quá trình Feed Forward.
# Normalize before calc feed forward
norm = nn.LayerNorm(n_embd)
attn_out_norm = norm(attn_out)Tương tự như quá trình Attention, chúng ta sử dụng normalize đầu ra của Attention trước khi đưa nó vào lớp Feed Forward.
hidden_size = 4 * n_embd
linear_1 = nn.Linear(n_embd, hidden_size)
linear_2 = nn.Linear(hidden_size, n_embd)
act = nn.ReLU()
# Feed forward output
hidden_states = linear_1(attn_out_norm)
hidden_states = act(hidden_states)
ffwd_out = linear_2(hidden_states)
# Add residual
residual = attn_out
ffwd_out += residual
ffwd_out.shapetorch.Size([20, 218, 36])Về bản chất, phần Feed Forward trong kiến trúc Transformer không quá phức tạp. Nó bao gồm một hidden layer và một output layer, kèm theo một hàm activation. Mặc dù có một số cải tiến trong LLAMA2, nhưng nguyên tắc cơ bản vẫn thì giống như đã được mô tả ở trên.
# Self Attention
# Normalize
attention_norm = nn.LayerNorm(n_embd)
x_embd_norm = attention_norm(x_embd)
# Multi-head Attention
n_head = 4
head_size = n_embd // n_head
opt_size = n_head * head_size # output size
Wqkv = nn.Linear(n_embd, 3 * opt_size)
qkv = Wqkv(x_embd_norm)
qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, h = n_head)
q, k, v = qkv.unbind(dim=2)
# Masked multi-head
softmax_scale = 1.0 / math.sqrt(q.shape[-1])
scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
mask = torch.triu(torch.full((sequence_len, sequence_len), -10000.0), 1)
scores = scores + mask
attention = torch.softmax(scores, dim=-1)
attn_out = torch.einsum("bhts,bshd->bthd", attention, v)
attn_out = rearrange(attn_out, "... h d -> ... (h d)")
# Attention output
out_proj = nn.Linear(opt_size, n_embd)
attn_out = out_proj(attn_out)
# Add residual
residual = x_embd
attn_out += residual
# Feed Forward
hidden_size = 4 * n_embd
linear_1 = nn.Linear(n_embd, hidden_size)
linear_2 = nn.Linear(hidden_size, n_embd)
act = nn.ReLU()
# Normalize
ffwd_norm = nn.LayerNorm(n_embd)
attn_out_norm = ffwd_norm(attn_out)
# Feed forward output
hidden_states = linear_1(attn_out_norm)
hidden_states = act(hidden_states)
ffwd_out = linear_2(hidden_states)
# Add residual
residual = attn_out
ffwd_out += residual
ffwd_out.shapetorch.Size([20, 218, 36])Quá trình từ input đã được embedding đi qua lớp Attention và lớp Feed Forward mà chúng ta vừa hoàn thành được gọi là một “block” trong kiến trúc Transformers. Trong hình ảnh, việc có một khung bao quanh quá trình Attention và Feed Forward ngụ ý rằng chúng hoạt động cùng nhau như một block duy nhất. Chữ “Nx” chỉ ra rằng ta có thể có nhiều block tùy ý. Trong trường hợp này, tôi chỉ sử dụng 1 block làm ví dụ, và nếu bạn muốn sử dụng 2 block, bạn có thể đơn giản sao chép toàn bộ đoạn code ở trên và dán vào một ô mới để tạo ra 2 block.
# Normalize output feed forward
norm = nn.LayerNorm(n_embd)
output = norm(ffwd_out)
last_linear = nn.Linear(n_embd, vocab_size)
logits = last_linear(output)
logits.shapetorch.Size([20, 218, 50257])Bởi vì có sự thay đổi trong quá trình normalize trong Feed Forward, nơi chúng ta thực hiện việc normalize trước khi truyền vào lớp Feed Forward. Vì vậy, chúng ta cần điều chỉnh kiến trúc model bằng cách thêm lớp normalize cho output của lớp Feed Forward trước khi truyền nó vào lớp Linear cuối cùng.
loss_fct = nn.CrossEntropyLoss()
logits = logits.view(-1, logits.shape[-1])
labels = y.view(-1)
loss = loss_fct(logits, labels)
losstensor(10.9201, grad_fn=<NllLossBackward0>)
Vậy là chúng ta đã hoàn thành toàn bộ kiến trúc Transformers được mô tả trong hình ảnh ở trên, và điều này cũng đồng nghĩa với việc chúng ta đã hiểu “backbone” của kiến trúc Transformers trong LLAMA2. Tuy nhiên, trong chương này, chúng ta chỉ đang viết code mà chưa sử dụng bất kỳ class nào. Trong chương tiếp theo, tôi sẽ sắp xếp code vào các class để làm cho kiến trúc trở nên có cấu trúc và giống với LLAMA2 hơn.