Điểm attention bằng einsum

Trung bình

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

torch.einsum mô tả phép nhân tensor bằng chuỗi chỉ số — rất tiện cho attention nhiều đầu. Tính ma trận điểm S với S[b, h, i, j] = Σ_d Q[b, h, i, d] · K[b, h, j, d]:

torch.einsum('bhid,bhjd->bhij', Q, K) (tương đương Q @ K.transpose(-2, -1)).

Viết hàm

attention_scores(Q, K)

Tham số

Q: tensor (B, H, T, d)
K: tensor (B, H, S, d)

Trả về

Tensor (B, H, T, S).

Ví dụ

Tham số

Q = [[[[1,1,1],[1,1,1]]]]
K = [[[[1,1,1]]]]

Kết quả

[[[[3.0],[3.0]]]]

Tham số

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

Kết quả

[[[[0.0,1.0],[2.0,3.0]]]]

Python · PyTorch

import torch


def attention_scores(Q, K):
    """Điểm attention bằng einsum"""
    # TODO: viết code của bạn ở đây
    pass

import torch


def attention_scores(Q, K):
    """Điểm attention bằng einsum"""
    # TODO: viết code của bạn ở đây
    pass

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