Bài 10: Cơ chế Self-Attention & Kỷ nguyên Transformer

Năm 2017, nhóm nghiên cứu tại Google công bố bài báo khoa học lịch sử: "Attention Is All You Need", giới thiệu kiến trúc mạng Transformer. Sự kiện này đã kết liễu thời kỳ thống trị kéo dài hàng thập kỷ của các mạng nơ-ron tuần hoàn RNN/LSTM và mở ra kỷ nguyên của các mô hình ngôn ngữ lớn (LLM) thay đổi thế giới ngày nay.

Điểm mấu chốt của Transformer là việc loại bỏ hoàn toàn các vòng lặp tuần tự theo bước thời gian, cho phép song song hóa tính toán 100% trên phần cứng đồ họa GPU. Bài học này sẽ đưa kiến trúc Transformer dưới kính hiển vi: mổ xẻ toán học của phép toán Scaled Dot-Product Self-Attention (Q, K, V), cơ chế Multi-Head Attention, mã hóa vị trí (Positional Encoding), và tự tay lắp ráp một khối Transformer Block hoàn chỉnh bằng PyTorch.

10.1 Cuộc cách mạng song song hóa: Cơ chế Self-Attention

Như chúng ta đã tìm hiểu ở Bài 9, mạng nơ-ron tuần hoàn (RNN) xử lý câu văn một cách tuần tự từng từ một. Điều này đồng nghĩa với việc nơ-ron không thể tính toán từ thứ 10 nếu chưa tính xong trạng thái ẩn của từ thứ 9.

Ngược lại, Transformer xử lý tất cả các từ trong câu văn cùng một lúc. Để bắt được ngữ cảnh liên kết giữa các từ, nó giới thiệu phép toán Self-Attention (Tự chú ý). Mỗi từ đầu vào sẽ tự động tính toán mức độ liên quan ngữ nghĩa của chính nó với tất cả các từ khác trong câu thông qua ba vector đại diện:

• Query ($Q$ - Câu hỏi): Đại diện cho thông tin từ hiện tại đang đi tìm ngữ cảnh liên quan.
• Key ($K$ - Từ khóa): Đại diện cho nhãn định danh đặc trưng của các từ khác trong câu để đối chiếu.
• Value ($V$ - Giá trị): Đại diện cho nội dung thông tin thực tế của từ đó được nạp đi sau khi đã tìm thấy sự chú ý thích hợp.

10.2 Học đa chiều: Multi-Head Attention & Mã hóa vị trí Positional Encoding

Nếu chỉ sử dụng một luồng Attention đơn lẻ, mô hình chỉ có thể tập trung chú ý vào một mối liên kết ngữ nghĩa duy nhất tại một thời điểm (ví dụ: liên kết ngữ pháp danh từ - động từ).

Để khắc phục, Transformer sử dụng cơ chế Multi-Head Attention (Chú ý đa luồng). Thay vì tính Attention trực tiếp trên các vector có số chiều gốc $d_{\text{model}}$, mạng nơ-ron chia nhỏ các chiều này thành $h$ luồng song song (mỗi luồng có số chiều nhỏ hơn là $d_k = d_{\text{model}} / h$). Mỗi luồng "Head" sẽ được khởi tạo các ma trận trọng số độc lập để tự do học các mối quan hệ ngữ cảnh khác nhau (ví dụ: Head 1 học quan hệ cú pháp chủ vị; Head 2 học quan hệ địa lý; Head 3 học quan hệ thời gian). Sau đó, kết quả từ toàn bộ các Head được ghép nối (concatenate) lại và chiếu qua một lớp Linear để trả về số chiều gốc.

Vì Transformer xử lý song song toàn bộ các từ cùng một lúc, mô hình không hề biết thứ tự từ trước sau. Với Transformer, câu văn "Tôi yêu bạn" và câu "Bạn yêu tôi" có ma trận biểu diễn hoàn toàn giống hệt nhau (xem câu văn như một túi từ xáo trộn không thứ tự - Bag of Words).

Để đưa thông tin thứ tự thời gian vào, ta bắt buộc phải cộng thêm một vector Positional Encoding (Mã hóa vị trí) trực tiếp vào vector nhúng từ (Word Embedding) trước khi đưa vào khối Attention đầu tiên.

10.3 Kiến trúc Encoder-Decoder và mô hình GPT (Decoder-only)

Kiến trúc Transformer nguyên bản gồm hai phần lớn xếp chồng lên nhau:

• Encoder (Bộ mã hóa): Đọc toàn bộ câu gốc để trích xuất đặc trưng hai chiều (Bidirectional). Được sử dụng trong các mô hình hiểu ngôn ngữ như BERT.
• Decoder (Bộ giải mã): Tích hợp thêm cơ chế tự hồi quy (Autoregressive) để sinh từng từ tiếp theo.

Đối với phần Decoder, khi huấn luyện sinh văn bản, chúng ta bắt buộc phải sử dụng một bộ mặt nạ gọi là Causal Mask (Mặt nạ nhân quả / Look-ahead Mask). Lớp mặt nạ này che đi toàn bộ các từ ở tương lai bằng cách ghi đè điểm số tích vô hướng $Q K^T$ tại các vị trí đó thành $-\infty$ trước khi đưa vào hàm Softmax: $$\text{Softmax}(-\infty) = 0$$ Điều này đảm bảo khi dự đoán từ tiếp theo ở bước $t$, mô hình hoàn toàn không thể "ăn gian" nhìn thấy dữ liệu ở các bước $t+1, t+2$ trong quá trình huấn luyện.

10.4 Dự án thực hành bài 10: Xây dựng khối Transformer Block hoàn chỉnh bằng PyTorch

Dự án thực hành của bài học này là tự tay lập trình từ đầu một khối Transformer Block hoàn chỉnh sử dụng PyTorch.

Chúng ta sẽ xây dựng lớp tự chú ý đa luồng MultiHeadAttention, kết hợp với các kết nối tắt (Residual Connections), lớp chuẩn hóa LayerNorm và mạng thần kinh truyền thẳng FeedForward để đóng gói thành khối kiến trúc hoàn chỉnh.

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model phải chia hết cho num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Tạo các ma trận chiếu tuyến tính cho Query, Key, Value
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        
        # Lớp tuyến tính cuối sau khi ghép các luồng Attention
        self.W_o = nn.Linear(d_model, d_model)
        
    def forward(self, q, k, v, mask=None):
        batch_size, seq_len, _ = q.size()
        
        # 1. Chiếu tuyến tính đầu vào và phân tách thành các Heads
        # Kích thước chuyển đổi: (Batch, Seq, d_model) -> (Batch, Seq, heads, d_k) -> (Batch, heads, Seq, d_k)
        Q = self.W_q(q).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(k).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(v).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        # 2. Tính điểm tương đồng Scaled Dot-Product
        # Q K^T: (Batch, heads, Seq, d_k) x (Batch, heads, d_k, Seq) -> (Batch, heads, Seq, Seq)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        
        # Áp dụng mặt nạ (nếu có)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
            
        # 3. Softmax để tạo phân phối trọng số chú ý
        attention_weights = F.softmax(scores, dim=-1)
        
        # 4. Nhân với Value và chuyển đổi lại số chiều gốc
        # (Batch, heads, Seq, Seq) x (Batch, heads, Seq, d_k) -> (Batch, heads, Seq, d_k)
        context = torch.matmul(attention_weights, V)
        
        # Ghép (Concatenate) các heads lại với nhau: (Batch, Seq, d_model)
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        
        # Lớp chiếu tuyến tính đầu ra
        output = self.W_o(context)
        return output
 
class FeedForwardNetwork(nn.Module):
    def __init__(self, d_model, d_ff):
        super(FeedForwardNetwork, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(d_ff, d_model)
        
    def forward(self, x):
        return self.linear2(self.relu(self.linear1(x)))
 
class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(TransformerBlock, self).__init__()
        
        # Khối chú ý đa luồng
        self.attention = MultiHeadAttention(d_model, num_heads)
        
        # Lớp chuẩn hóa Layer Normalization
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        # Mạng truyền thẳng FFN
        self.feed_forward = FeedForwardNetwork(d_model, d_ff)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask=None):
        # 1. Khối Self-Attention + Kết nối tắt Residual Connection + LayerNorm
        attn_out = self.attention(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))
        
        # 2. Khối FeedForward + Kết nối tắt Residual Connection + LayerNorm
        ff_out = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_out))
        
        return x
 
if __name__ == "__main__":
    print("=== Khởi tạo khối Transformer Block ===")
    d_model = 64
    num_heads = 8
    d_ff = 256
    seq_len = 10
    batch_size = 2
    
    # Khởi tạo mô hình
    block = TransformerBlock(d_model=d_model, num_heads=num_heads, d_ff=d_ff)
    print(f"Tham số: d_model={d_model}, heads={num_heads}, d_ff={d_ff}")
    
    # Giả lập dữ liệu đầu vào (Batch, Seq Length, d_model)
    dummy_input = torch.randn(batch_size, seq_len, d_model)
    print(f"Kích thước tensor đầu vào: {dummy_input.shape}")
    
    # Khởi tạo mặt nạ Causal Mask giả lập cho bộ giải mã Decoder
    # Mặt nạ tam giác dưới (Lower Triangular Mask) chứa số 1, góc trên chứa số 0
    causal_mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).unsqueeze(1) # (1, 1, Seq, Seq)
    print(f"Kích thước mặt nạ Causal Mask: {causal_mask.shape}\n")
    
    # Lan truyền xuôi qua khối Transformer Block
    output = block(dummy_input, mask=causal_mask)
    print("=== Chạy Lan truyền xuôi qua mô hình thành công ===")
    print(f"Kích thước tensor đầu ra: {output.shape} (Bằng chính xác kích thước đầu vào)")
    print(f"Kiến trúc khối Transformer Block đã hoạt động đồng bộ.")

Tóm tắt bài học & Cầu nối kiến thức

Cầu nối bài tiếp theo: Sau khi đã hiểu cấu tạo phần cứng bên trong LLM, chúng ta sẽ bắt đầu học cách lập trình tương tác với các LLM đã được huấn luyện thông qua API và các kỹ thuật Prompt Engineering ở Bài số 11.