【深度学习】爱因斯坦求和约定

你正在上数学课，教授突然停下来，神秘地微笑着说：“今天，我要教你们一种数学界的终极偷懒秘诀，连爱因斯坦都用它来省事儿。”全班立刻竖起了耳朵。教授接着说：“你们知道，那些看起来复杂的求和符号，其实就是在浪费我们宝贵的脑细胞。于是，爱因斯坦决定，干脆直接忽略掉这些累赘的符号——凡是公式里上下标相同的，自动就给它们求和！没错，这就是爱因斯坦求和约定，一个聪明得让人心情愉快的偷懒小技巧，让我们的公式变得优雅简洁。”于是，全班同学顿时茅塞顿开，纷纷感叹：原来，连爱因斯坦也有这么“偷懒”的一面！

1. 介绍

爱因斯坦求和约定（Einstein summation convention）是爱因斯坦提出的一种简化数学表达的记号方法，尤其在张量运算中非常常见。这种约定可以让复杂的线性代数运算变得更加简洁明了。

1.1 基本概念

在传统的数学表达式中，当我们进行张量运算（例如矩阵乘法、内积等）时，通常会明确地写出求和符号（$\sum$）。而在爱因斯坦求和约定中，省略了显式的求和符号，约定对于那些在表达式中重复出现的下标，会自动进行求和。

爱因斯坦求和约定的核心是：当一个表达式中同一个指标（下标或上标）在同一个项内出现两次时，假定要对该指标的所有可能取值进行求和，而不需要显式地写出求和符号。

1.2 约定的规则

• 重复出现的指标：在一个张量表达式中，如果某个指标（通常是下标或上标）在同一个项中出现了两次，就默认对这个指标求和。这个指标被称为“哑标”（dummy index）。哑标的求和值的范围是该维度的全范围。
• 单次出现的指标：如果一个指标在一个表达式中只出现一次，它表示该指标是自由的，不需要求和。这个指标被称为“自由标”（free index）。
• 简化表达：使用爱因斯坦求和约定，可以省去求和符号，使得表达式更加简洁。

1.3 示例解释

向量点积

考虑两个向量 𝑎 和 𝑏 的内积：

$$c=\sum _i{\mathbf{a}}_i{\mathbf{b}}_i$$

使用爱因斯坦求和约定，可以写成：

$$c={\mathbf{a}}_i{\mathbf{b}}_i$$

在这个表达式中，$i$ 是一个哑标，表示对它进行求和。

矩阵乘法

考虑两个矩阵 $\mathbf{A}$ 和 $\mathbf{B}$ 的乘积：

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

使用爱因斯坦求和约定，可以写成：

$${\mathbf{C}}_{ij}={\mathbf{A}}_{ik}{\mathbf{B}}_{kj}$$

这里，$k$ 是哑标，表示对它进行求和；$i$ 和 $j$ 是自由标，表示矩阵 $\mathbf{C}$ 的行和列。

2. torch.einsum()函数

在 PyTorch 中，torch.einsum() 函数允许使用爱因斯坦求和约定来高效、简洁地执行复杂的张量操作。通过 torch.einsum()，可以指定张量的维度排列和操作，而无需显式地编写多个矩阵乘法或转置操作。

2.1 使用步骤

• 指定爱因斯坦求和表达式：表达式包含输入张量的索引和输出张量的索引。
• 提供输入张量：将要操作的张量传入函数。

2.2 代码示例

示例1. 向量点积

假设我们有两个向量 𝑎 和 𝑏，长度为 3。根据爱因斯坦求和约定，点积可以表示为：
$$c=\sum _i{\mathbf{a}}_i{\mathbf{b}}_i$$

# -*- coding: utf-8 -*-
# @time: 2024/8/23 17:15

import torch

a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

# Einsum for dot product
c = torch.einsum('i,i->', a, b)

print(c)  # 输出: tensor(32)

在这个例子中，'i,i->' 表示对索引 i 进行求和，并且输出是一个标量（没有索引）。

示例2. 矩阵乘法

考虑两个矩阵 $\mathbf{A}$ 和 $\mathbf{B}$ 的矩阵乘法$\mathbf{C}=\mathbf{A}\times \mathbf{B}$：

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

# -*- coding: utf-8 -*-
# @time: 2024/8/23 17:15

import torch

A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[5, 6], [7, 8]])

# Einsum for matrix multiplication
C = torch.einsum('ik,kj->ij', A, B)

print(C)  # 输出: tensor([[19, 22], [43, 50]])

在这个例子中，'ik,kj->ij' 表示输入的两个张量分别按 ik 和 kj 维度进行操作，并输出按 ij 维度排列的结果。

示例3. 批次矩阵乘法

假设我们有一组批次矩阵 $\mathbf{A}$和 $\mathbf{B}$，我们想要计算每个批次的矩阵乘法：

# -*- coding: utf-8 -*-
# @time: 2024/8/23 17:15

import torch

A = torch.randn(10, 3, 4)  # 10 个 3x4 的矩阵
B = torch.randn(10, 4, 5)  # 10 个 4x5 的矩阵

# Einsum for batch matrix multiplication
C = torch.einsum('bij,bjk->bik', A, B)

print(C.shape)  # 输出: torch.Size([10, 3, 5])

在这个例子中，'bij,bjk->bik' 中的 b 表示批次维度，不进行求和操作，仅在结果中保留。

3.总结

爱因斯坦求和约定通过省略显式的求和符号，使得张量运算的表达更加简洁。它在处理高维数据时尤为方便，特别是在深度学习中，许多复杂的线性代数操作可以通过这种方式简化。torch.einsum() 函数正是利用这一约定，允许我们以直观的方式定义复杂的张量操作。