从 QKᵀ 到加权聚合：用 3 个 Agent 直观理解 Attention

本文只解释两次批量矩阵乘法：$QK^T$ 与 $\alpha V$。示例中的 $\alpha$ 假定已经归一化，并不代表完整的 scaled dot-product attention 实现。

现有代码

score = torch.bmm(q, k.transpose(1, 2))

假设有 3 个 agent，每个元素都有一个查询向量 q 和一个键向量 k。为便于理解，这里设置查询向量和键向量的维度为 2。假设批次大小 batch_size = 1，只有一个批次的数据。

查询向量 q (形状为 (1, 3, 2))：

q = torch.tensor([[[1.0, 0.0],    # 第1个元素的查询向量
                   [0.0, 1.0],    # 第2个元素的查询向量
                   [1.0, 1.0]]])  # 第3个元素的查询向量

键向量 k (形状为 (1, 3, 2))：

k = torch.tensor([[[1.0, 0.0],    # 第1个元素的键向量
                   [0.0, 1.0],    # 第2个元素的键向量
                   [1.0, 1.0]]])  # 第3个元素的键向量

第 1 行：第 1 个元素与所有元素的相似性：

第 1 个元素和第 1 个元素的相似性：$1.0 \times 1.0 + 0.0 \times 0.0 = 1.0$

第 1 个元素和第 2 个元素的相似性：$1.0 \times 0.0 + 0.0 \times 1.0 = 0.0$

第 1 个元素和第 3 个元素的相似性：$1.0 \times 1.0 + 0.0 \times 1.0 = 1.0$

得到第 1 行的分数为 [1.0, 0.0, 1.0]。

第 2 行：第 2 个元素与所有元素的相似性：

第2 个元素和第1个元素的相似性：$0.0 \times 1.0 + 1.0 \times 0.0 = 0.0$

第 2 个元素和第 2 个元素的相似性：$0.0 \times 0.0 + 1.0 \times 1.0 = 1.0$

第 2 个元素和第 3 个元素的相似性：$0.0 \times 1.0 + 1.0 \times 1.0 = 1.0$

得到第 2 行的分数为 [0.0, 1.0, 1.0]。

第 3 行：第 3 个元素与所有元素的相似性：

第 3 个元素和第 1 个元素的相似性：$1.0 \times 1.0 + 1.0 \times 0.0 = 1.0$

第 3 个元素和第 2 个元素的相似性：$1.0 \times 0.0 + 1.0 \times 1.0 = 1.0$

第 3 个元素和第 3 个元素的相似性：$1.0 \times 1.0 + 1.0 \times 1.0 = 2.0$

得到第 3 行的分数为 [1.0, 1.0, 2.0]。

将这些结果组合起来得到 score：

score = torch.tensor([[[1.0, 0.0, 1.0],
                       [0.0, 1.0, 1.0],
                       [1.0, 1.0, 2.0]]])

现有代码

att = torch.bmm(alpha, v)

假设我们有 3 个 agent 在一个批次中观察其他 agent 的状态。为了方便说明，只考虑一个批次的数据，因此 batch_size = 1。假设每个 agent 的特征向量 v 是 2 维的，最终生成的注意力权重 alpha 和特征向量 v 的形状分别如下：

alpha (注意力权重矩阵)：形状为 (1, 3, 3)，其中每个 3x3 的矩阵表示每个 agent 对其他 agent 的关注度。

v (值向量矩阵)：形状为 (1, 3, 2)，其中每个 3x2 的矩阵表示每个 agent 的特征向量。

假设有以下输入数据：

# alpha 表示每个agent对其他agent的关注程度(归一化权重矩阵)
alpha = torch.tensor([[[0.1, 0.7, 0.2],  # agent1对agent1、agent2、agent3的权重
                       [0.5, 0.2, 0.3],  # agent2对agent1、agent2、agent3的权重
                       [0.3, 0.3, 0.4]]])  # agent3对agent1、agent2、agent3的权重

# v 表示每个agent的特征向量
v = torch.tensor([[[1.0, 2.0],   # agent1的特征
                   [3.0, 4.0],   # agent2的特征
                   [5.0, 6.0]]]) # agent3的特征

att = torch.bmm(alpha, v) 通过批量矩阵乘法来计算 alpha 和 v 的乘积：

agent1 的注意力聚合输出：

使用 alpha[0, 0]，即 [0.1, 0.7, 0.2]，对所有 agent 的 v 进行加权求和。

agent1 的输出 $= 0.1 \cdot [1.0, 2.0] + 0.7 \cdot [3.0, 4.0] + 0.2 \cdot [5.0, 6.0] = [3.2, 4.2]$

agent2 的注意力聚合输出：

使用 alpha[0, 1]，即 [0.5, 0.2, 0.3]，对所有 agent 的 v 进行加权求和。

agent2 的输出 $= 0.5 \cdot [1.0, 2.0] + 0.2 \cdot [3.0, 4.0] + 0.3 \cdot [5.0, 6.0] = [2.6, 3.6]$

agent3 的注意力聚合输出：

使用 alpha[0, 2]，即 [0.3, 0.3, 0.4]，对所有 agent 的 v 进行加权求和。计算如下：

agent3 的输出 $= 0.3 \cdot [1.0, 2.0] + 0.3 \cdot [3.0, 4.0] + 0.4 \cdot [5.0, 6.0] = [3.2, 4.2]$

将每个 agent 的注意力聚合输出组合起来形成 att：

att = torch.tensor([[[3.2, 4.2],
                     [2.6, 3.6],
                     [3.2, 4.2]]])