Self-Attention

Self-attention là dạng attention trong đó query, key và value đều xuất phát từ cùng một chuỗi, vì vậy mỗi token xây dựng biểu diễn mới của nó từ mọi token khác trong chuỗi đó. Một lớp self-attention có thể kết nối từ đầu tiên với từ cuối cùng của đoạn văn, hoặc hai góc đối diện của ảnh, chỉ trong một bước; trong khi mạng nơ-ron hồi quy (RNN) cần một bước cho mỗi vị trí và mạng nơ-ron tích chập (CNN) cần nhiều lớp xếp chồng để mở rộng phạm vi quan sát. Khi được xếp chồng cùng các lớp feed-forward, self-attention tạo thành transformer được giới thiệu trong bài báo năm 2017 Attention Is All You Need, nền tảng của mô hình ngôn ngữ lớn (LLM), Vision Transformer (ViT) và detector dựa trên transformer.

Cách thức hoạt động của self-attention#

Mỗi token đầu vào, được biểu diễn dưới dạng embedding, được nhân với ba ma trận đã học để tạo ra query, key và value. Đầu ra của token i là trung bình có trọng số của vector value của mọi token; trọng số được xác định bằng softmax trên tích vô hướng đã co giãn giữa query của token i với tất cả key, tức cơ chế attention tích vô hướng đã co giãn được mô tả trên trang về cơ chế attention. Vì query, key và value là các phép chiếu của cùng những token, chuỗi gồm n token tạo ra ma trận trọng số n × n. Trong câu “The animal didn't cross the street because it was too tired,” hàng ứng với “it” có thể đặt phần lớn trọng số lên “animal” và đưa ý nghĩa của từ đó vào biểu diễn của chính nó.

Multi-head self-attention thực hiện song song nhiều phép tính như vậy trên các phép chiếu nhỏ hơn rồi nối kết quả, nhờ đó các head khác nhau có thể theo dõi những quan hệ khác nhau, chẳng hạn một head xử lý cú pháp và head khác xử lý đồng tham chiếu. Transformer ban đầu sử dụng 8 head, mỗi head có 64 chiều, trong model 512 chiều.

Mất thông tin thứ tự và causal mask#

Self-attention xem đầu vào như một tập hợp. Xáo trộn các token sẽ xáo trộn đầu ra theo cùng cách mà không làm thay đổi điều gì khác, vì vậy các model bổ sung thông tin vị trí, chẳng hạn embedding hình sin trong transformer ban đầu hoặc rotary position embedding (RoPE) trong hầu hết LLM hiện nay. Các encoder như BERT cho phép mỗi token attention theo cả hai chiều. LLM chỉ có decoder áp dụng causal mask để mỗi token chỉ attention tới chính nó và các token đứng trước, nhờ đó model có thể tạo văn bản từ trái sang phải và tái sử dụng key cùng value của các token trước đó từ KV cache. MultiheadAttention của PyTorch thể hiện cả hai đặc tính này khi cùng một tensor được truyền làm query, key và value:

import torch

torch.manual_seed(0)
tokens = torch.randn(1, 5, 16)  # một chuỗi gồm 5 token, mỗi token có 16 đặc trưng
attn = torch.nn.MultiheadAttention(embed_dim=16, num_heads=2, batch_first=True)

# Self-attention: query, key và value đều là các phép chiếu của cùng những token
out, weights = attn(tokens, tokens, tokens)
print("Attention weights shape:", tuple(weights.shape))

# Xáo trộn token chỉ làm xáo trộn đầu ra, vì vậy thứ tự từ không thể nhận biết nếu thiếu thông tin vị trí
perm = torch.tensor([3, 0, 4, 1, 2])
out_shuffled, _ = attn(tokens[:, perm], tokens[:, perm], tokens[:, perm])
print("Same outputs after shuffling:", torch.allclose(out_shuffled, out[:, perm], atol=1e-6))

# Causal mask chặn attention tới các token đứng sau, như trong LLM chỉ có decoder
mask = torch.triu(torch.ones(5, 5, dtype=torch.bool), diagonal=1)
_, causal_weights = attn(tokens, tokens, tokens, attn_mask=mask)
print(causal_weights[0].detach().round(decimals=2))

Kết quả in ra có shape trọng số (1, 5, 5), True cho phép thử xáo trộn và một ma trận tam giác dưới: token đầu tiên đặt trọng số 1.00 lên chính nó, token thứ hai chia trọng số 0.31 và 0.69 cho hai token đầu tiên, và mọi phần tử phía trên đường chéo đều bằng 0.00.

Self-attention và các thuật ngữ liên quan#

• Cơ chế attention: Phép toán tổng quát gán trọng số cho value dựa trên độ tương đồng giữa query và key. Self-attention là trường hợp cả ba thành phần đều xuất phát từ một chuỗi.
• Cross-attention: Query xuất phát từ một chuỗi, còn key và value xuất phát từ chuỗi khác. Decoder dịch máy attention tới câu nguồn đã mã hóa, object query của decoder DETR attention tới đặc trưng ảnh, và một model khuếch tán chuyển văn bản thành ảnh attention tới prompt.
• Tích chập: Áp dụng cùng một bộ lọc nhỏ tại mọi vị trí, vì vậy mỗi đầu ra nhìn thấy một cửa sổ cục bộ cố định với trọng số không phụ thuộc vào nội dung. Self-attention tính trọng số từ nội dung và bao phủ toàn bộ đầu vào.

Transformer encoder-decoder ban đầu kết hợp self-attention hai chiều, self-attention nhân quả và cross-attention:

Chi phí và các biến thể hiệu quả#

Tính điểm cho mọi cặp token khiến thời gian và bộ nhớ tăng theo bình phương độ dài chuỗi: 1.000 token tạo ra một triệu điểm cho mỗi head, còn 100.000 token tạo ra mười tỷ điểm. FlashAttention tính self-attention chính xác theo các tile mà không lưu toàn bộ ma trận, grouped-query attention (GQA) chia sẻ key và value giữa các head để thu nhỏ KV cache, còn sliding-window, sparse và linear attention giảm số lượng cặp cần tính. Trong thị giác máy tính, số lượng token phụ thuộc vào độ phân giải ảnh, vì vậy nhiều detector chỉ áp dụng self-attention tại những vị trí feature map nhỏ nhất. YOLO26 chỉ chạy multi-head self-attention trong khối C2PSA ở cuối backbone và một khối C3k2 có attention trong neck; cả hai đều hoạt động trên feature map 20×20 stride-32 của ảnh 640 pixel, tương ứng 400 token. Module attention của model tính query, key và value bằng một tích chập 1×1 trên feature map đó, đồng thời bổ sung thông tin vị trí bằng tích chập depthwise 3×3 áp dụng lên value.