Phân tích Kiến trúc Transformer

1. Tổng quan về Kiến trúc Transformer

Transformer là mô hình học sâu dựa trên cơ chế attention, không phụ thuộc vào các cấu trúc tuần tự như RNN hay CNN. Nó được thiết kế cho các tác vụ xử lý ngôn ngữ tự nhiên (NLP) như dịch máy, sinh văn bản và xây dựng mô hình ngôn ngữ tiền huấn luyện — sau đó có thể tinh chỉnh cho nhiều nhiệm vụ khác nhau.

Kiến trúc tổng thể gồm bốn thành phần chính:

• Phần đầu vào: Bao gồm lớp nhúng (embedding) và mã hóa vị trí cho cả chuỗi nguồn và chuỗi đích.
• Phần đầu ra: Gồm một lớp tuyến tính và hàm softmax để sinh xác suất từ vựng.
• Bộ mã hóa (Encoder): Gồm N tầng mã hóa đồng nhất, mỗi tầng chứa hai khối con: (1) chú ý tự động đa đầu kèm chuẩn hóa và kết nối dư; (2) mạng fully-connected kèm chuẩn hóa và kết nối dư.
• Bộ giải mã (Decoder): Cũng gồm N tầng giải mã đồng nhất, nhưng mỗi tầng có ba khối con: (1) chú ý tự động đa đầu (cho dữ liệu đích); (2) chú ý chéo đa đầu (giữa đích và nguồn); (3) mạng fully-connected — tất cả đều đi kèm chuẩn hóa và kết nối dư.

2. Triển khai Phần Đầu vào

2.1 Lớp Nhúng Từ (Word Embedding)

Mục tiêu của lớp nhúng là chuyển đổi biểu diễn rời rạc (mã số từ vựng) thành vector liên tục trong không gian chiều cao — nhằm nắm bắt mối quan hệ ngữ nghĩa giữa các từ.

Dưới đây là phiên bản đã tái cấu trúc với logic tương đương nhưng tên biến rõ ràng hơn và cải tiến cách khởi tạo:

import torch
import torch.nn as nn
import math

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, embedding_dim: int):
        super().__init__()
        self.embedding_table = nn.Embedding(vocab_size, embedding_dim)
        self.scale_factor = math.sqrt(embedding_dim)

    def forward(self, token_ids: torch.LongTensor) -> torch.Tensor:
        # Nhân với hệ số chuẩn hóa để ổn định độ lớn gradient
        return self.embedding_table(token_ids) * self.scale_factor

Ví dụ sử dụng:

# Khởi tạo với từ vựng 1000 từ, mỗi từ biểu diễn bởi vector 512 chiều
embedder = TokenEmbedding(vocab_size=1000, embedding_dim=512)
input_tokens = torch.tensor([[100, 2, 421, 508], [491, 998, 1, 221]])
output_vectors = embedder(input_tokens)  # Kích thước: [2, 4, 512]

2.2 Mã Hóa Vị Trí (Positional Encoding)

Vì Transformer không có cơ chế tuần tự, cần bổ sung thông tin vị trí bằng tín hiệu cố định dựa trên hàm sin/cos. Mỗi chiều của vector vị trí được gán một tần số khác nhau, đảm bảo khoảng cách tương đối giữa hai vị trí luôn có thể phân biệt.

class PositionalEncoder(nn.Module):
    def __init__(self, embedding_dim: int, dropout_rate: float = 0.1, max_seq_len: int = 5000):
        super().__init__()
        self.dropout = nn.Dropout(dropout_rate)
        # Tạo ma trận mã hóa vị trí tĩnh (không huấn luyện)
        position_enc = torch.zeros(max_seq_len, embedding_dim)
        position = torch.arange(0, max_seq_len).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, embedding_dim, 2).float() *
                             -(math.log(10000.0) / embedding_dim))
        position_enc[:, 0::2] = torch.sin(position * div_term)
        position_enc[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pos_enc', position_enc.unsqueeze(0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        seq_len = x.size(1)
        x += self.pos_enc[:, :seq_len]
        return self.dropout(x)

Để trực quan hóa, ta vẽ đường cong sin/cos cho một vài chiều đặc trưng:

import matplotlib.pyplot as plt
pe = PositionalEncoder(embedding_dim=20, dropout_rate=0)
x = torch.zeros(1, 100, 20)
encoded = pe(x).squeeze(0).data.numpy()

plt.figure(figsize=(12, 4))
for dim in [4, 5, 6, 7]:
    plt.plot(encoded[:, dim], label=f'Chiều {dim}')
plt.legend()
plt.title("Phân bố giá trị mã hóa vị trí theo chiều")
plt.xlabel("Vị trí trong chuỗi")
plt.ylabel("Giá trị mã hóa")
plt.show()

3. Triển khai Bộ Mã hóa

3.1 Mặt nạ Chuỗi (Causal Mask)

Khi huấn luyện bộ giải mã, ta phải ngăn việc "nhìn thấy" các ký tự tương lai — điều này được thực hiện bằng mặt nạ tam giác dưới (lower-triangular mask). Dưới đây là hàm tạo mặt nạ đã được viết lại gọn hơn:

def create_causal_mask(size: int) -> torch.Tensor:
    """Tạo mặt nạ tam giác dưới để che các vị trí tương lai"""
    mask = torch.tril(torch.ones(size, size, dtype=torch.bool))
    return mask.unsqueeze(0)  # Thêm batch dimension

Ví dụ: create_causal_mask(4) trả về ma trận 1×4×4 với các phần tử trên đường chéo chính và dưới đều là True.

3.2 Cơ Chế Chú Ý Đơn Giản

Cơ chế chú ý tính toán trọng số giữa các từ bằng cách so sánh query (câu hỏi), key (chìa khóa) và value (giá trị). Công thức chuẩn là "scaled dot-product attention":

import torch.nn.functional as F

def scaled_dot_attention(
    queries: torch.Tensor,
    keys: torch.Tensor,
    values: torch.Tensor,
    mask: torch.Tensor = None,
    dropout: nn.Dropout = None
) -> tuple[torch.Tensor, torch.Tensor]:
    d_k = queries.size(-1)
    scores = torch.matmul(queries, keys.transpose(-2, -1)) / math.sqrt(d_k)
    
    if mask is not None:
        scores = scores.masked_fill(~mask, float('-inf'))
    
    attention_weights = F.softmax(scores, dim=-1)
    if dropout is not None:
        attention_weights = dropout(attention_weights)
    
    output = torch.matmul(attention_weights, values)
    return output, attention_weights

3.3 Chú Ý Đa Đầu (Multi-Head Attention)

Thay vì dùng một cặp Q/K/V duy nhất, mô hình chia không gian nhúng thành H phần độc lập — mỗi phần học một loại mối quan hệ ngữ nghĩa riêng biệt. Kết quả cuối cùng là ghép nối và biến đổi tuyến tính.

def clone_module(module: nn.Module, n: int) -> nn.ModuleList:
    """Nhân bản module n lần để dùng chung kiến trúc"""
    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])

class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads: int, embedding_dim: int, dropout_rate: float = 0.1):
        super().__init__()
        assert embedding_dim % num_heads == 0
        self.head_dim = embedding_dim // num_heads
        self.num_heads = num_heads
        self.dropout = nn.Dropout(dropout_rate)
        
        # 4 lớp tuyến tính: 3 cho Q/K/V, 1 cho đầu ra sau khi ghép
        self.linears = clone_module(nn.Linear(embedding_dim, embedding_dim), 4)

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                mask: torch.Tensor = None) -> torch.Tensor:
        batch_size = q.size(0)
        
        # Ánh xạ tuyến tính và chia đầu
        qkv = [l(x).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
               for l, x in zip(self.linears, (q, k, v))]
        q, k, v = qkv
        
        # Tính chú ý cho từng đầu
        x, self.attn_weights = scaled_dot_attention(q, k, v, mask, self.dropout)
        
        # Ghép các đầu và phục hồi kích thước
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)
        return self.linears[-1](x)

3.4 Lớp Tuyến Tính Vị Trí (Position-wise Feed-Forward)

Mỗi tầng mã hóa/giải mã đều có một mạng fully-connected hai lớp — giúp tăng khả năng biểu diễn phi tuyến mà cơ chế chú ý đơn thuần chưa đủ:

class PositionWiseFFN(nn.Module):
    def __init__(self, embedding_dim: int, hidden_dim: int, dropout_rate: float = 0.1):
        super().__init__()
        self.linear1 = nn.Linear(embedding_dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, embedding_dim)
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear2(self.dropout(F.relu(self.linear1(x))))

3.5 Chuẩn Hóa Lớp (Layer Normalization)

Chuẩn hóa theo chiều đặc trưng (khác với BatchNorm) giúp ổn định quá trình huấn luyện ở các tầng sâu:

class LayerNorm(nn.Module):
    def __init__(self, feature_dim: int, eps: float = 1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(feature_dim))
        self.beta = nn.Parameter(torch.zeros(feature_dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

3.6 Kết Nối Phụ Trợ (Residual + Norm)

Mỗi khối con trong Encoder/Decoder đều tuân theo mẫu: Input → Norm → Sublayer → Dropout → Add(Input):

class ResidualConnection(nn.Module):
    def __init__(self, feature_dim: int, dropout_rate: float):
        super().__init__()
        self.norm = LayerNorm(feature_dim)
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, x: torch.Tensor, sublayer: callable) -> torch.Tensor:
        return x + self.dropout(sublayer(self.norm(x)))

3.7 Tầng Mã Hóa và Bộ Mã Hóa Hoàn Chỉnh

class EncoderLayer(nn.Module):
    def __init__(self, feature_dim: int, self_attn: MultiHeadAttention,
                 feed_forward: PositionWiseFFN, dropout_rate: float):
        super().__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayers = clone_module(ResidualConnection(feature_dim, dropout_rate), 2)

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        x = self.sublayers[0](x, lambda z: self.self_attn(z, z, z, mask))
        return self.sublayers[1](x, self.feed_forward)

class Encoder(nn.Module):
    def __init__(self, layer: EncoderLayer, num_layers: int):
        super().__init__()
        self.layers = clone_module(layer, num_layers)
        self.norm = LayerNorm(layer.norm.gamma.size(0))

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

4. Triển khai Bộ Giải Mã và Toàn Bộ Mô Hình

Tầng giải mã mở rộng thêm một khối chú ý chéo (cross-attention) để kết hợp thông tin từ bộ mã hóa:

class DecoderLayer(nn.Module):
    def __init__(self, feature_dim: int, self_attn: MultiHeadAttention,
                 cross_attn: MultiHeadAttention, feed_forward: PositionWiseFFN,
                 dropout_rate: float):
        super().__init__()
        self.self_attn = self_attn
        self.cross_attn = cross_attn
        self.feed_forward = feed_forward
        self.sublayers = clone_module(ResidualConnection(feature_dim, dropout_rate), 3)

    def forward(self, x: torch.Tensor, memory: torch.Tensor,
                src_mask: torch.Tensor, tgt_mask: torch.Tensor) -> torch.Tensor:
        x = self.sublayers[0](x, lambda z: self.self_attn(z, z, z, tgt_mask))
        x = self.sublayers[1](x, lambda z: self.cross_attn(z, memory, memory, src_mask))
        return self.sublayers[2](x, self.feed_forward)

class Decoder(nn.Module):
    def __init__(self, layer: DecoderLayer, num_layers: int):
        super().__init__()
        self.layers = clone_module(layer, num_layers)
        self.norm = LayerNorm(layer.norm.gamma.size(0))

    def forward(self, x: torch.Tensor, memory: torch.Tensor,
                src_mask: torch.Tensor, tgt_mask: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

class Transformer(nn.Module):
    def __init__(self, encoder: Encoder, decoder: Decoder,
                 src_embed: nn.Sequential, tgt_embed: nn.Sequential,
                 generator: nn.Linear):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def forward(self, src: torch.Tensor, tgt: torch.Tensor,
                src_mask: torch.Tensor, tgt_mask: torch.Tensor) -> torch.Tensor:
        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)

    def encode(self, src: torch.Tensor, src_mask: torch.Tensor) -> torch.Tensor:
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, memory: torch.Tensor, src_mask: torch.Tensor,
               tgt: torch.Tensor, tgt_mask: torch.Tensor) -> torch.Tensor:
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

Hàm Xây Dựng Mô Hình Đầy Đủ

def build_transformer(
    src_vocab_size: int,
    tgt_vocab_size: int,
    num_layers: int = 6,
    d_model: int = 512,
    d_ff: int = 2048,
    num_heads: int = 8,
    dropout_rate: float = 0.1
) -> Transformer:
    c = copy.deepcopy
    
    attn = MultiHeadAttention(num_heads, d_model, dropout_rate)
    ff = PositionWiseFFN(d_model, d_ff, dropout_rate)
    pos = PositionalEncoder(d_model, dropout_rate)
    
    model = Transformer(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout_rate), num_layers),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout_rate), num_layers),
        nn.Sequential(TokenEmbedding(src_vocab_size, d_model), c(pos)),
        nn.Sequential(TokenEmbedding(tgt_vocab_size, d_model), c(pos)),
        nn.Linear(d_model, tgt_vocab_size)
    )
    
    # Khởi tạo trọng số Xavier
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    
    return model

Sử dụng:

model = build_transformer(
    src_vocab_size=10000,
    tgt_vocab_size=10000,
    num_layers=6,
    d_model=512,
    d_ff=2048,
    num_heads=8
)
print(model)