torch.einsum详解

torch.einsum 是 PyTorch 中用于执行高效张量运算的函数，基于爱因斯坦求和约定（Einstein summation convention）。它能够处理复杂的张量操作，并简化代码书写。

基本语法

torch.einsum(subscripts, *operands)

• subscripts：一个字符串，用于描述输入张量的维度如何结合。
• *operands：待操作的张量。

爱因斯坦求和约定

爱因斯坦求和约定是一个简化张量运算的方式，省略了显式的求和符号。通过指定各维度的标签，可以直接描述复杂的张量运算。

语法结构

• "nqhd,nkhd->nhqk": 这个字符串描述了如何对两个张量进行操作，并生成输出张量的维度。

  • n：批次大小（batch size）
  • q：查询序列长度（query length）
  • k：键序列长度（key length）
  • h：注意力头的数量（number of heads）
  • d：每个注意力头的维度（dimension per head）

示例代码

以下是使用 torch.einsum 计算多头注意力机制中点积相似性的示例代码：

import torch

# 定义多头注意力机制的点积计算函数
def compute_attention_scores(queries, keys):
    # 计算点积相似性分数
    energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
    return energy

# 示例数据
N = 1            # 批次大小
q = 2            # 查询序列长度
k = 3            # 键序列长度
h = 2            # 注意力头数量
d = 4            # 每个注意力头的维度

# 随机生成 queries 和 keys
queries = torch.rand((N, q, h, d))  # Shape (1, 2, 2, 4)
keys = torch.rand((N, k, h, d))    # Shape (1, 3, 2, 4)

# 计算注意力分数
energy = compute_attention_scores(queries, keys)

print("Energy shape:", energy.shape)
print(energy)

计算过程

1. 维度解释：

   • queries 的维度为 (1, 2, 2, 4)：N = 1（批次大小），q = 2（查询序列长度），h = 2（注意力头数量），d = 4（每个头部的维度）。
   • keys 的维度为 (1, 3, 2, 4)：N = 1（批次大小），k = 3（键序列长度），h = 2（注意力头数量），d = 4（每个头部的维度）。

2. 点积计算：

   • 对每个批次和每个头部，计算 queries 和 keys 在 d 维度上的点积。
   • 结果的维度为 (N, h, q, k)，其中：
     • N 是批次大小
     • h 是注意力头的数量
     • q 是查询序列的长度
     • k 是键序列的长度

   点积计算的实际操作是：

   • 对于每个批次（n）和每个头部（h），对 queries 和 keys 张量在 d 维度上进行点积运算，得到形状为 (q, k) 的张量。

简单计算示例

假设我们有如下示例数据：

queries = torch.tensor([[[[1.0, 0.5, 0.2, 1.5], [0.3, 0.7, 0.6, 0.8]], [[0.9, 0.4, 1.2, 0.5], [0.2, 0.6, 0.8, 0.7]]]])
keys = torch.tensor([[[[0.1, 1.0, 0.3, 0.5], [0.2, 0.4, 0.6, 0.7], [0.8, 1.0, 0.9, 0.5]], [[0.1, 0.5, 0.2, 0.8], [0.3, 0.4, 0.7, 0.9], [0.6, 0.8, 1.0, 0.2]]]])

点积计算：

• 对于第一个批次和第一个头部：

  • queries[0, :, 0, :] 和 keys[0, :, 0, :] 的点积计算如下：

  计算:

  energy[0, 0, 0, 0] = (1.0*0.1 + 0.5*1.0 + 0.2*0.3 + 1.5*0.5) = 0.1 + 0.5 + 0.06 + 0.75 = 1.41
  energy[0, 0, 0, 1] = (1.0*0.2 + 0.5*0.4 + 0.2*0.6 + 1.5*0.7) = 0.2 + 0.2 + 0.12 + 1.05 = 1.59
  energy[0, 0, 0, 2] = (1.0*0.8 + 0.5*1.0 + 0.2*0.9 + 1.5*0.5) = 0.8 + 0.5 + 0.18 + 0.75 = 1.23
  energy[0, 0, 1, 0] = (0.3*0.1 + 0.7*1.0 + 0.6*0.3 + 0.8*0.5) = 0.03 + 0.7 + 0.18 + 0.4 = 1.31
  energy[0, 0, 1, 1] = (0.3*0.2 + 0.7*0.4 + 0.6*0.6 + 0.8*0.7) = 0.06 + 0.28 + 0.36 + 0.56 = 1.26
  energy[0, 0, 1, 2] = (0.3*0.8 + 0.7*1.0 + 0.6*0.9 + 0.8*0.5) = 0.24 + 0.7 + 0.54 + 0.4 = 1.88
  

总结

torch.einsum("nqhd,nkhd->nhqk", [queries, keys]) 用于计算 queries 和 keys 张量在注意力机制中的点积，相似性得分。它通过爱因斯坦求和约定指定了如何在多维张量上执行这些操作，使得代码更简洁、效率更高。

Code

AI_With_NumPy
此项目汇集了很多AI相关的代码实现，供大家学习使用，欢迎点赞收藏👏🏻

备注

个人水平有限，有问题随时交流~