Causal attention theo lô bằng PyTorch

Khó

PyTorchTransformerattention4 test · 20 giây · sai số 1e-4

Cài đặt attention có mặt nạ nhân quả cho cả lô, đúng như trong mã nguồn GPT:

1. att = Q @ K.transpose(-2, -1) / math.sqrt(d) → (B, T, T)

2. Che tương lai: mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1), rồi att = att.masked_fill(mask, float('-inf'))

3. att = torch.softmax(att, dim=-1), trả về att @ V

(PyTorch có sẵn F.scaled_dot_product_attention(..., is_causal=True) — hãy tự viết để hiểu, rồi đối chiếu.)

Viết hàm

causal_attention(Q, K, V)

Tham số

Q: tensor (B, T, d)
K: tensor (B, T, d)
V: tensor (B, T, dᵥ)

Trả về

Tensor (B, T, dᵥ).

Ví dụ

Tham số

Q = [[[1,1],[1,1],[1,1]]]
K = [[[1,1],[1,1],[1,1]]]
V = [[[1],[2],[3]]]

Kết quả

[[[1.0],[1.5],[2.0]]]

Q·K bằng nhau → trung bình các vị trí được nhìn: [1, 1.5, 2]

Tham số

Q = [[[1.2715,1.6841,-1.1702,2.6258],[-0.1946,2.5894,2.4644,-0.5931]],[[-2.7379,-0.8192,0.2088,2.664],[0.4541,-2.7496,0.4742,2.8998]]]
K = [[[-0.6822,-2.0592,-2.1206,0.6947],[-0.6999,0.3447,0.54,0.0147]],[[1.8453,0.8471,-0.303,0.1706],[-1.1679,2.0037,1.7816,0.1101]]]
V = [[[-1.4566,-1.892,-0.6739],[-2.8601,2.9893,-2.5006]],[[1.8982,-0.9715,-1.5861],[-0.5439,-2.0704,-1.6701]]]

Kết quả

[[[-1.4566,-1.892,-0.6739],[-2.858182044,2.982629449,-2.498103719]],[[1.8982,-0.9715,-1.5861],[1.571427589,-1.118541564,-1.597339868]]]

Bài tiếp theo: Điểm attention bằng einsum →

Python · PyTorch

import math
import torch


def causal_attention(Q, K, V):
    """Causal attention theo lô bằng PyTorch"""
    # TODO: viết code của bạn ở đây
    pass

import math
import torch


def causal_attention(Q, K, V):
    """Causal attention theo lô bằng PyTorch"""
    # TODO: viết code của bạn ở đây
    pass

Bấm “Chạy thử” để gọi causal_attention() với 2 ví dụ trong đề và so sánh kết quả. Dùng print() thoải mái để gỡ lỗi.