变压器模型如何计算自我注意力？

Question

在变压器模型https://arxiv.org/pdf/1706.03762.pdf中，存在自我注意，它是通过对Query (Q)和Key (K)向量使用softmax来计算的：

我试图了解矩阵乘法：

Q = batch_size x seq_length x embed_size

K = batch_size x seq_length x embed_size

QK^T = batch_size x seq_length x seq_length

Softmax QK^T = Softmax (batch_size x seq_length x seq_length)

由于每个批处理元素有seq_length x seq_length个值，因此如何计算softmax？

参考Pytorch计算将非常有帮助。

干杯！

Answer 1

由于每个批处理元素有seq_length x seq_length个值，因此如何计算softmax？

softmax在最后一个轴上执行（torch.nn.Softmax(dim=-1)(tensor)，其中tensor的形状为batch_size x seq_length x seq_length），以获取参与输入序列中每个元素的每个元素的概率。

---

让我们假设，我们有一个文本序列“ Thinking Machines”，因此在执行seq_length = 2之后，我们有一个形状为“ 2 x 2”（其中QK^T）的矩阵。

我正在使用以下插图（reference）来说明自我注意的计算。如您所知，首先执行QK^T/square_root(d_k)的缩放点积，然后为每个序列元素计算softmax。

此处，对第一个序列元素“ Thinking”执行Softmax。通过执行softmax，将14 and 12的原始分数转换为0.88 and 0.12的概率。这些概率表明令牌“ Thinking”将以88％的概率参与其中，令牌“ Machines”的概率为12％。同样，也为令牌“机器”计算了注意力概率。

---

注意。我强烈建议您在Transformer上阅读出色的article。对于实现，您可以看一下OpenNMT。

Answer 2

QKᵀ乘法是一个批处理矩阵乘法-它是将seq_length x embed_size乘以embed_size x seq_length次batch_size来进行单独的seq_length x seq_length。每个结果都给出大小为batch_size x seq_length x seq_length的结果，这就是我们最终得到形状为torch.matmul(query, key.transpose(-2, -1)) 的QKᵀ的原因。

Gabriela Melo's suggested resource使用以下PyTorch代码进行此操作：

torch.matmul

之所以可行，是因为当输入至少具有3个维度时，let array = []; list.forEach((item_1) => { let variable = 0; fetch("data.json") .then((resp) => resp.json()) .then((data) => { data.forEach((item_2) => { item_2.forEach((item_3) => { item_3forEach((item_4) => { if (condition) { variable += 1; } }); }); }); }); array.push({a:variable, b: item_1}); }); 会进行批量矩阵乘法（请参见https://pytorch.org/docs/stable/torch.html#torch.matmul）。