Einstein Summation 爱因斯坦求和 torch.einsum

Einstein Summation 爱因斯坦求和 torch.einsum

flyfish

理解爱因斯坦求和的基本概念和语法，这对初学者来说可能有一定难度。对于不熟悉该表示法的用户来说，可能不如直接的矩阵乘法表达式易于理解。

整个思路是

向量的点积 -》矩阵乘法-》einsum

向量之间的点积在几何上表示两个向量的投影和夹角，在代数上用于衡量向量的相似性，并且在物理学中用于计算力做的功。
矩阵乘法是由多个向量点积组成的，可以看作是多个向量点积的组合
einsum 操作可以用其他内置的矩阵运算函数来实现

使用 einsum 进行矩阵乘法

import torch# 定义两个矩阵
A = torch.randn(2, 3)
B = torch.randn(3, 4)# 使用 einsum 表示矩阵乘法
C = torch.einsum('ik,kj->ij', A, B)
print(C)

使用 matmul 进行矩阵乘法

import torch# 定义两个矩阵
A = torch.randn(2, 3)
B = torch.randn(3, 4)# 使用 matmul 表示矩阵乘法
C = torch.matmul(A, B)
print(C)

开始解释

向量之间的点积（也称为内积或标量积）在数学、物理学和计算中有着重要的意义。点积是两个向量乘积的一种特殊形式，其结果是一个标量。点积在许多领域中都有广泛的应用，包括向量的投影、计算角度、物理学中的功、机器学习中的相似性度量等。

点积的定义

给定两个n维向量 $\mathbf{a}$ 和 $\mathbf{b}$，它们的点积定义如下：
$\mathbf{a}\cdot \mathbf{b}={\sum }_{i=1}^n{a}_i{b}_i$

点积的几何意义

1. 计算向量间的夹角：
   点积可以用来计算两个向量之间的夹角 $\theta$。根据点积的定义，可以得到：
   $\mathbf{a}\cdot \mathbf{b}=\Vert \mathbf{a}\Vert \Vert \mathbf{b}\Vert \cos (\theta )$
   其中 $\Vert \mathbf{a}\Vert$ 和 $\Vert \mathbf{b}\Vert$ 分别是向量 $\mathbf{a}$ 和 $\mathbf{b}$ 的模（或长度）。因此，可以通过点积计算两个向量之间的夹角：
   $\cos (\theta )=\frac{\mathbf{a}\cdot \mathbf{b}}{\Vert \mathbf{a}\Vert \Vert \mathbf{b}\Vert }$
2. 投影：
   点积可以用来计算一个向量在另一个向量方向上的投影。例如，向量 $\mathbf{a}$ 在向量 $\mathbf{b}$ 方向上的投影长度为：
   ${\text{proj}}_{\mathbf{b}}\mathbf{a}=\frac{\mathbf{a}\cdot \mathbf{b}}{\Vert \mathbf{b}\Vert }$

点积的代数意义

1. 向量相似性：
   在机器学习和数据分析中，点积可以用来衡量两个向量之间的相似性。如果两个向量的方向相同，它们的点积为正；如果两个向量的方向相反，它们的点积为负；如果两个向量正交，它们的点积为零。
2. 功的计算：
   在物理学中，点积用于计算力和位移的乘积，即功。例如，如果一个物体在力 $\mathbf{F}$ 的作用下移动了位移 $\mathbf{d}$，则做的功为：
   $W=\mathbf{F}\cdot \mathbf{d}$

例子

计算两个二维向量的点积

假设 $\mathbf{a}=[a_1,a_2]$ 和 $\mathbf{b}=[b_1,b_2]$，它们的点积为：
$\mathbf{a}\cdot \mathbf{b}=a_1{b}_1+a_2{b}_2$

import numpy as np# 定义两个向量
a = np.array([1, 2])
b = np.array([3, 4])# 计算点积
dot_product = np.dot(a, b)
print(dot_product)  # 输出: 11

计算两个三维向量的夹角

假设 $\mathbf{a}=[a_1,a_2,a_3]$ 和 $\mathbf{b}=[b_1,b_2,b_3]$，它们的点积和夹角计算如下：

import numpy as np# 定义两个向量
a = np.array([1, 0, 0])
b = np.array([0, 1, 0])# 计算点积
dot_product = np.dot(a, b)# 计算向量的模
norm_a = np.linalg.norm(a)
norm_b = np.linalg.norm(b)# 计算夹角的余弦值
cos_theta = dot_product / (norm_a * norm_b)
theta = np.arccos(cos_theta)print(f"夹角: {np.degrees(theta)} 度")  # 输出: 90.0 度

矩阵乘法

矩阵乘法是两个矩阵 $A$ 和 $B$ 的乘积 $C$，其中：

• 矩阵 $A$ 的形状为 $m\times n$
• 矩阵 $B$ 的形状为 $n\times p$
• 矩阵 $C$ 的形状为 $m\times p$
  矩阵乘法的定义是：
  $C_{ij}={\sum }_{k=1}^n{A}_{ik}{B}_{kj}$

换句话说，矩阵 $C$ 的元素 $C_{ij}$ 是矩阵 $A$ 的第 $i$ 行和矩阵 $B$ 的第 $j$ 列的点积。

从矩阵乘法到向量点积

考虑两个矩阵 $A$ 和 $B$，我们可以将矩阵乘法分解为一系列的向量点积：

1. 提取行向量和列向量：

• 矩阵 $A$ 的第 $i$ 行可以表示为向量 ${\mathbf{a}}_{\mathbf{i}}$。
• 矩阵 $B$ 的第 $j$ 列可以表示为向量 ${\mathbf{b}}_{\mathbf{j}}$。

2. 计算点积：

• $C_{ij}$ 是向量 ${\mathbf{a}}_{\mathbf{i}}$ 和向量 ${\mathbf{b}}_{\mathbf{j}}$ 的点积，即：
  $C_{ij}={\mathbf{a}}_{\mathbf{i}}\cdot {\mathbf{b}}_{\mathbf{j}}$
  例如，考虑矩阵 $A$ 和 $B$：
  $$A=\left(\begin{array}{cc}1& 2\\ 3& 4\end{array}\right)$$
  $$B=\left(\begin{array}{cc}5& 6\\ 7& 8\end{array}\right)$$
  它们的乘积 $C=AB$ 为：
  $$C=\left(\begin{array}{cc}1\cdot 5+2\cdot 7& 1\cdot 6+2\cdot 8\\ 3\cdot 5+4\cdot 7& 3\cdot 6+4\cdot 8\end{array}\right)=\left(\begin{array}{cc}19& 22\\ 43& 50\end{array}\right)$$
  在这里：
  $$\begin{gathered} C_{11}=1\cdot 5+2\cdot 7=19 \\ {C}_{12}=1\cdot 6+2\cdot 8=22 \\ {C}_{21}=3\cdot 5+4\cdot 7=43 \\ {C}_{22}=3\cdot 6+4\cdot 8=50 \end{gathered}$$

使用 PyTorch 进行矩阵乘法和点积

以下是如何在 PyTorch 中实现矩阵乘法和向量点积：

import torch
# 定义两个矩阵
A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[5, 6], [7, 8]])
# 使用 torch.matmul 进行矩阵乘法
C = torch.matmul(A, B)
print(C)
# 输出:
# tensor([[19, 22],
# [43, 50]])
# 提取行向量和列向量
a1 = A[0, :] # A 的第一行
b1 = B[:, 0] # B 的第一列
# 计算点积
dot_product = torch.dot(a1, b1)
print(dot_product)
# 输出: tensor(19)

爱因斯坦求和

爱因斯坦求和约定（Einstein Summation Convention）是一种在物理学和数学中简化张量运算表示的方法。它由阿尔伯特·爱因斯坦在他的广义相对论论文中引入。这个约定的核心思想是通过省略求和符号（∑），简化公式的书写，增强表达的简洁性和可读性。

背景与起源

在物理学中，尤其是在处理广义相对论和量子力学中的张量时，常常需要进行大量的求和运算。为了简化这些计算的书写，爱因斯坦提出了一种简便的表示法：对于任何重复出现的指标，默认对其进行求和。

具体规则

1. 求和隐含性：在一个表达式中，如果一个指标（下标或上标）在一个单项式中出现两次，则认为对该指标求和。例如：$a_i{b}_i={\sum }_i{a}_i{b}_i$在这种表示法中，i 被称为“哑指标”或“虚指标”。
2. 自由指标：如果一个指标在表达式中仅出现一次，则称其为自由指标，这个指标代表一个范围的所有可能值。例如：$c_i=a_{ij}{b}_j$这里的 i 是自由指标，而 j 是哑指标。
3. 多重求和：可以在一个表达式中使用多个哑指标进行多重求和。例如：$d=a_{ij}{b}_{jk}{c}_k$在这个例子中，j 和 k 都是哑指标，意味着：$d={\sum }_j{\sum }_k{a}_{ij}{b}_{jk}{c}_k$

在使用 torch.einsum 时，我们可以利用爱因斯坦求和约定来简洁地表示矩阵乘法、张量收缩等操作：

import torch# 矩阵乘法
A = torch.randn(3, 4)
B = torch.randn(4, 5)
C = torch.einsum('ik,kj->ij', A, B)

在上面的例子中，‘ik,kj->ij’ 表示矩阵乘法，其中 $k$ 是求和下标，最终结果的维度由 $i$ 和 $j$ 确定。

‘ik,kj->ij’ 是爱因斯坦求和约定在 torch.einsum 中的一个具体应用，表示矩阵乘法。让我们详细解析一下这个表示：

表达式解析

1. 输入张量：

• 假设我们有两个矩阵 A 和 B。
• A 的形状为 (i, k)，即 A 有 i 行和 k 列。
• B 的形状为 (k, j)，即 B 有 k 行和 j 列。

2. 爱因斯坦求和约定：

• ‘ik,kj->ij’ 中的 ik 和 kj 分别对应输入张量 A 和 B 的维度标签。
• 中间的 , 用于分隔多个输入张量的维度标签。
• 箭头 -> 左侧表示输入张量的维度标签，右侧表示输出张量的维度标签。
• ‘ik,kj’ 表示对两个输入张量 A 和 B 进行操作，其中 k 是求和下标。

3. 求和与输出：

• 在 ‘ik,kj’ 中，k 是求和下标，表示我们要对 k 维度进行求和。
• i 和 j 出现在箭头 -> 右侧，表示输出张量的维度标签。

举例说明

假设我们有两个矩阵：

$A=\left(\begin{array}{cc}{a}_{11}& a_{12}\\ a_{21}& a_{22}\\ a_{31}& a_{32}\end{array}\right),B=\left(\begin{array}{ccc}{b}_{11}& b_{12}& b_{13}\\ b_{21}& b_{22}& b_{23}\end{array}\right)$

其中：

• $A$ 是一个 $3\times 2$ 的矩阵，对应维度标签 ‘ik’（即 $i=3$, $k=2$）。
• $B$ 是一个 $2\times 3$ 的矩阵，对应维度标签 ‘kj’（即 $k=2$, $j=3$）。
  使用 torch.einsum 表示矩阵乘法：

import torchA = torch.tensor([[a11, a12],[a21, a22],[a31, a32]])B = torch.tensor([[b11, b12, b13],[b21, b22, b23]])C = torch.einsum('ik,kj->ij', A, B)

矩阵乘法过程

‘ik,kj->ij’ 表示通过对 k 维度进行求和，得到输出矩阵 C：

$C=A\cdot B$

其中：

$C_{ij}={\sum }_k{A}_{ik}{B}_{kj}$

即：

$C_{ij}=A_{i1}{B}_{1j}+A_{i2}{B}_{2j}$

结果

根据上面的定义，最终的结果 C 是一个 $3\times 3$ 的矩阵：

$C=\left(\begin{array}{ccc}{c}_{11}& c_{12}& c_{13}\\ c_{21}& c_{22}& c_{23}\\ c_{31}& c_{32}& c_{33}\end{array}\right)$

每个元素 $c_{ij}$ 由对应的矩阵乘法和求和计算得到。

注意力机制

在编写注意力的时候有这样的代码
scores = torch.einsum(“blhe,bshe->bhls”, queries, keys)

从向量内积的角度解释

假设 queries 和 keys 的形状分别为 $(B,L,H,E)$ 和 $(B,S,H,E)$，其中:

• $B$ 是批次大小 (Batch Size)
• $L$ 是查询序列的长度 (Length of queries)
• $S$ 是键序列的长度 (Length of keys)
• $H$ 是注意力头的数量 (Number of heads)
• $E$ 是嵌入维度 (Embedding Dimension)
  我们计算 queries 和 keys 在嵌入维度 $E$ 上的内积，即通过 torch.einsum(“blhe,bshe->bhls”, queries, keys) 来计算注意力得分。

示例

假设我们有以下输入：

import torchqueries = torch.tensor([[[[0.5, 1.2], [0.3, 0.7]],  # 第一个 query 序列，两个头，每个头两个维度[[1.5, 2.2], [1.3, 1.7]],  # 第二个 query 序列，两个头，每个头两个维度]
])  # 形状 (1, 2, 2, 2)keys = torch.tensor([[[[0.8, 1.5], [0.4, 0.9]],  # 第一个 key 序列，两个头，每个头两个维度[[1.1, 2.3], [1.0, 1.5]],  # 第二个 key 序列，两个头，每个头两个维度]
])  # 形状 (1, 2, 2, 2)

这里 queries 和 keys 的形状都是 (1, 2, 2, 2)，表示 1 个批次，2 个序列，2 个头，每个头 2 个维度。

我们希望计算注意力得分矩阵 scores，其形状为 (1, 2, 2, 2)。

计算步骤

使用 torch.einsum 计算 scores：

scores = torch.einsum("blhe,bshe->bhls", queries, keys)
print(scores)

手动计算

头 1：

• 第一个 query 序列和第一个 key 序列的内积：
  $0.5\times 0.8+1.2\times 1.5=0.4+1.8=2.2$
• 第一个 query 序列和第二个 key 序列的内积：
  $0.5\times 1.1+1.2\times 2.3=0.55+2.76=3.31$
• 第二个 query 序列和第一个 key 序列的内积：
  $1.5\times 0.8+2.2\times 1.5=1.2+3.3=4.5$
• 第二个 query 序列和第二个 key 序列的内积：
  $1.5\times 1.1+2.2\times 2.3=1.65+5.06=6.71$

头 2：

• 第一个 query 序列和第一个 key 序列的内积：
  $0.3\times 0.4+0.7\times 0.9=0.12+0.63=0.75$
• 第一个 query 序列和第二个 key 序列的内积：
  $0.3\times 1.0+0.7\times 1.5=0.3+1.05=1.35$
• 第二个 query 序列和第一个 key 序列的内积：
  $1.3\times 0.4+1.7\times 0.9=0.52+1.53=2.05$
• 第二个 query 序列和第二个 key 序列的内积：
  $1.3\times 1.0+1.7\times 1.5=1.3+2.55=3.85$

最终结果

根据上述计算，我们可以得到：

scores = torch.tensor([[[[2.2, 3.31], [4.5, 6.71]],  # 第一个头的得分[[0.75, 1.35], [2.05, 3.85]]  # 第二个头的得分]
])

使用 PyTorch 计算

运行以下代码验证手动计算结果：

import torchqueries = torch.tensor([[[[0.5, 1.2], [0.3, 0.7]],  [[1.5, 2.2], [1.3, 1.7]],  ]
]) keys = torch.tensor([[[[0.8, 1.5], [0.4, 0.9]],  [[1.1, 2.3], [1.0, 1.5]],  ]
])scores = torch.einsum("blhe,bshe->bhls", queries, keys)
print(scores)

输出：

tensor([[[[2.2000, 3.3100],[4.5000, 6.7100]],[[0.7500, 1.3500],[2.0500, 3.8500]]]])

从矩阵乘法的角度解释

使用矩阵乘法计算

为了将 queries 和 keys 的计算表示成矩阵乘法，我们可以按以下步骤操作：

1. 调整形状：

• 将 queries 和 keys 调整形状，使每个头的查询和键序列变成矩阵。

2. 矩阵乘法：

• 对每个头分别进行矩阵乘法。

调整形状并进行矩阵乘法

我们将 queries 和 keys 形状从 (B, L, H, E) 和 (B, S, H, E) 调整为 (B, H, L, E) 和 (B, H, E, S)，以便进行矩阵乘法。

queries_reshaped = queries.permute(0, 2, 1, 3)  # (B, H, L, E)
keys_reshaped = keys.permute(0, 2, 3, 1)       # (B, H, E, S)# 使用矩阵乘法
scores_matmul = torch.matmul(queries_reshaped, keys_reshaped)  # (B, H, L, S)
print(scores_matmul)

输出：

tensor([[[[2.2000, 3.3100],[4.5000, 6.7100]],[[0.7500, 1.3500],[2.0500, 3.8500]]]])

这里的矩阵乘法使用 torch.matmul ，没有使用 torch.mm
torch.matmul 和 torch.mm 是 PyTorch 中用于矩阵乘法的两个函数，但它们在适用的张量维度上有一些不同。具体来说：

torch.mm

• 用途：专门用于两个二维矩阵（矩阵）之间的乘法。
• 输入：必须是两个二维张量，形状分别为 (m, n) 和 (n, p)。
• 输出：结果是一个二维张量，形状为 (m, p)。
  示例：

import torch# 定义两个二维矩阵
A = torch.randn(2, 3)
B = torch.randn(3, 4)# 使用 torch.mm 进行矩阵乘法
C = torch.mm(A, B)
print(C.shape)  # 输出: torch.Size([2, 4])

torch.matmul

• 用途：更通用的矩阵乘法函数，可以处理二维及以上的张量。
• 输入：可以是二维矩阵，也可以是具有更多维度的张量。
• 输出：根据输入张量的维度，输出可能是一个矩阵或更高维度的张量。
• 广播：torch.matmul 可以处理广播（broadcasting），即输入张量的形状可以不完全匹配，但需要满足广播规则。
  示例：

import torch# 定义两个二维矩阵
A = torch.randn(2, 3)
B = torch.randn(3, 4)# 使用 torch.matmul 进行矩阵乘法
C = torch.matmul(A, B)
print(C.shape)  # 输出: torch.Size([2, 4])# 定义两个三维张量
A_3d = torch.randn(5, 2, 3)
B_3d = torch.randn(5, 3, 4)# 使用 torch.matmul 进行三维张量的矩阵乘法
C_3d = torch.matmul(A_3d, B_3d)
print(C_3d.shape)  # 输出: torch.Size([5, 2, 4])# 广播示例
A_broadcast = torch.randn(2, 3)
B_broadcast = torch.randn(5, 3, 4)# A_broadcast 的形状将广播成 (5, 2, 3)
C_broadcast = torch.matmul(A_broadcast, B_broadcast)
print(C_broadcast.shape)  # 输出: torch.Size([5, 2, 4])

主要区别

• 适用维度：torch.mm 只适用于二维矩阵；torch.matmul 则适用于二维及以上维度的张量。
• 广播：torch.matmul 支持广播，而 torch.mm 不支持。

permute

在 PyTorch 中，permute 是一个张量（tensor）的方法，用于改变张量的维度顺序。这个操作不会改变张量的数据，只是重新排列它的维度。这对于需要改变数据的形状以适应不同操作的需求非常有用。

举例来说，如果你有一个形状为 (batch_size, height, width, channels) 的图像张量，而你的模型需要输入形状为 (batch_size, channels, height, width) 的张量，你可以使用 permute 方法来重新排列维度。

以下是一个简单的例子：

import torch# 创建一个形状为 (2, 3, 4, 5) 的随机张量
x = torch.randn(2, 3, 4, 5)# 使用 permute 方法改变维度顺序
x_permuted = x.permute(0, 3, 1, 2)# 打印新张量的形状
print(x_permuted.shape)  # 输出: torch.Size([2, 5, 3, 4])

在这个例子中：

• 原始张量 x 的形状为 (2, 3, 4, 5)。
• 调用 x.permute(0, 3, 1, 2) 后，新张量 x_permuted 的形状变为 (2, 5, 3, 4)。
  permute 方法的参数是新维度顺序的索引。例如，x.permute(0, 3, 1, 2) 意味着将第 0 维保持不变，将原第 3 维移到第 1 位置，将原第 1 维移到第 2 位置，将原第 2 维移到第 3 位置。

https://pytorch.org/docs/stable/generated/torch.transpose.html
https://pytorch.org/docs/stable/generated/torch.einsum.html