Attention thực hành: kiểm tra kích thước và gỡ lỗi

Bài trước tập trung vào ý tưởng. Bài này chuyển attention thành tensor cụ thể và tập trung vào những lỗi sinh viên hay gặp nhất: kích thước tensor, chiều softmax, mặt nạ và cách tách/gộp nhiều head.

Mục tiêu học tập

Sau bài này, bạn cần:

• theo dõi được kích thước của $Q,K,V$, ma trận điểm, trọng số và đầu ra;
• tự cài scaled dot-product attention bằng PyTorch;
• thêm causal mask đúng chỗ;
• tách/gộp multi-head mà không làm lẫn batch, head và chiều chuỗi;
• viết được các kiểm tra nhỏ để phát hiện lỗi trước khi huấn luyện mô hình lớn.

1. Quy ước kích thước tensor

Dùng:

• $B$: kích thước batch;
• $T$: độ dài chuỗi;
• $D=d_{model}$: số chiều embedding;
• $H$: số head;
• $d_h=D/H$: chiều mỗi head.

Nếu chưa tách head:

Sau khi tách head:

Điều kiện bắt buộc:

2. Tính tay một query nhỏ

Giả sử:

Với $d_k=2$:

Softmax xấp xỉ:

Đầu ra:

Điểm cần nhìn thấy: đầu ra không phải một value có sẵn, mà là tổ hợp có trọng số của các value.

3. Scaled dot-product attention từ đầu

import math
import torch


def scaled_dot_product_attention(q, k, v, mask=None):
    """
    q: [..., Tq, d]
    k: [..., Tk, d]
    v: [..., Tk, dv]
    mask: broadcastable to [..., Tq, Tk]
          True = allowed, False = blocked
    """
    scores = q @ k.transpose(-2, -1)
    scores = scores / math.sqrt(q.size(-1))

    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))

    weights = torch.softmax(scores, dim=-1)
    output = weights @ v
    return output, weights

Vì sao dim=-1?

Mỗi query tương ứng với một hàng điểm và phải phân phối trọng số qua các key. Vì vậy tổng trên chiều cuối phải bằng 1:

assert torch.allclose(
    weights.sum(dim=-1),
    torch.ones_like(weights.sum(dim=-1)),
    atol=1e-6,
)

4. Causal mask

Với $T=4$:

query 0: ✓ × × ×
query 1: ✓ ✓ × ×
query 2: ✓ ✓ ✓ ×
query 3: ✓ ✓ ✓ ✓

PyTorch:

T = 4
mask = torch.tril(torch.ones(T, T, dtype=torch.bool))
print(mask)

Khi broadcast vào batch/head:

mask = mask.view(1, 1, T, T)

Kiểm tra bắt buộc

Sau attention:

future = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
assert torch.all(weights[..., future] == 0)

Trong số học dấu phẩy động, đôi khi nên kiểm bằng ngưỡng:

assert weights[..., future].abs().max() < 1e-7

5. Multi-head attention: tách/gộp head đúng cách

Giả sử các lớp projection trả về tensor [B, T, D].

def split_heads(x, num_heads):
    B, T, D = x.shape
    assert D % num_heads == 0
    d_head = D // num_heads

    x = x.view(B, T, num_heads, d_head)
    x = x.transpose(1, 2)  # [B, H, T, d_head]
    return x


def merge_heads(x):
    B, H, T, d_head = x.shape
    x = x.transpose(1, 2).contiguous()  # [B, T, H, d_head]
    return x.view(B, T, H * d_head)

Lỗi phổ biến là view sau transpose khi tensor không contiguous. Nếu dùng .view(...), hãy gọi .contiguous() trước; hoặc dùng .reshape(...) khi phù hợp.

6. Một lớp multi-head attention tối giản

import torch
import torch.nn as nn


class TinyMultiHeadAttention(nn.Module):
    def __init__(self, d_model=64, num_heads=4, causal=False):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads
        self.causal = causal

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def _split(self, x):
        B, T, D = x.shape
        return x.view(B, T, self.num_heads, self.d_head).transpose(1, 2)

    def forward(self, x):
        B, T, _ = x.shape

        q = self._split(self.q_proj(x))
        k = self._split(self.k_proj(x))
        v = self._split(self.v_proj(x))

        mask = None
        if self.causal:
            mask = torch.tril(
                torch.ones(T, T, dtype=torch.bool, device=x.device)
            ).view(1, 1, T, T)

        out, weights = scaled_dot_product_attention(q, k, v, mask)

        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.out_proj(out), weights

Kiểm tra nhanh:

torch.manual_seed(0)

x = torch.randn(2, 5, 64)
mha = TinyMultiHeadAttention(d_model=64, num_heads=4, causal=True)

y, attn = mha(x)

print(y.shape)      # [2, 5, 64]
print(attn.shape)   # [2, 4, 5, 5]

7. Dùng hàm tối ưu có sẵn của PyTorch

PyTorch hiện có torch.nn.functional.scaled_dot_product_attention, và nn.MultiheadAttention/TransformerEncoderLayer có thể tận dụng các kernel tối ưu bên dưới.

import torch.nn.functional as F

out = F.scaled_dot_product_attention(
    q,
    k,
    v,
    dropout_p=0.0,
    is_causal=True,
)

Cẩn thận về quy ước mask

Đây là lỗi rất dễ gặp khi đổi API:

• với F.scaled_dot_product_attention, boolean True có nghĩa là được tham gia attention;
• một số mask của nn.MultiheadAttention dùng quy ước ngược lại cho padding mask.

Không nên đoán. Hãy kiểm tra tài liệu đúng với API đang dùng.

8. Padding mask khác causal mask

Hai mask giải quyết hai vấn đề khác nhau:

Một mô hình có thể cần cả hai.

9. Checklist gỡ lỗi trước khi huấn luyện

9.1 Kích thước tensor

assert q.shape[-2] == T
assert q.shape[-1] == d_head
assert k.shape[-1] == d_head
assert weights.shape[-2:] == (T, T)

9.2 Tổng trọng số

row_sum = weights.sum(dim=-1)
assert torch.allclose(row_sum, torch.ones_like(row_sum), atol=1e-5)

9.3 Không có NaN/Inf

assert torch.isfinite(weights).all()
assert torch.isfinite(out).all()

9.4 Rò rỉ thông tin tương lai

Nếu là causal attention, toàn bộ vùng tam giác trên phải gần 0.

9.5 Gradient có chạy

loss = out.square().mean()
loss.backward()

assert mha.q_proj.weight.grad is not None
assert torch.isfinite(mha.q_proj.weight.grad).all()

10. Những lỗi hay gặp

11. Chi phí tính toán

Với full attention chuẩn, ma trận điểm có kích thước $T\times T$ cho mỗi head:

và cần lưu hoặc xử lý quan hệ theo từng cặp token, nên chuỗi rất dài sẽ nhanh tốn tài nguyên.

Các kernel như scaled-dot-product attention tối ưu hoặc FlashAttention có thể giảm đáng kể chi phí bộ nhớ và I/O, trong khi vẫn tính cùng phép attention toán học ở trường hợp exact attention. Tối ưu kernel không làm công thức cơ bản biến thành tuyến tính theo $T$.

12. Bài tập

Bài 1 — Tính tay

Dùng ví dụ Q/K/V đầu bài và thay query thành $q=[0,1]$. Tính điểm, softmax và đầu ra.

Bài 2 — Kiểm tra rò rỉ thông tin tương lai

Tạo đầu vào [B=1,T=6,D=16], chạy causal attention và viết assert chứng minh mọi trọng số phía trên đường chéo bằng 0.

Bài 3 — So sánh 1 head và 4 head

Giữ $D=64$, chạy:

• H=1 → $d_h=64$;
• H=4 → $d_h=16$.

So sánh kích thước tensor và số tham số projection. Không kết luận chất lượng chỉ từ một lần forward ngẫu nhiên.

Bài 4 — Debug cố ý

Đổi softmax(dim=-1) thành softmax(dim=-2). Viết test phát hiện lỗi mà không cần huấn luyện.

Tài liệu tham khảo

• Vaswani et al. Attention Is All You Need. https://arxiv.org/abs/1706.03762
• PyTorch: scaled_dot_product_attention. https://docs.pytorch.org/docs/main/generated/torch.nn.functional.scaled_dot_product_attention.html
• PyTorch tutorial: Implementing High-Performance Transformers with Scaled Dot Product Attention. https://docs.pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial
• Raschka. Build a Large Language Model (From Scratch), Ch. 3. Manning, 2025.

Bản tiếng Anh đang soạn.