在深度学习，尤其是大模型（如 Transformer）的性能分析中，“计算量”是一个核心概念。它通常以 浮点运算次数（FLOPs, Floating Point Operations） 来衡量，直接影响模型的训练/推理速度、硬件资源需求以及能效。而几乎所有神经网络中的计算，最终都可以归结为矩阵乘法及其变体。因此，深入理解矩阵乘法的计算复杂度，是分析 Transformer 等现代架构计算开销的第一步。

本节将从最基础的向量和二维矩阵出发，逐步过渡到高维张量（即“高维矩阵”），系统讲解如何分析其浮点运算量，并引入三个关键维度概念：收缩维度、批处理维度 和 自由维度。

向量与二维矩阵

我们先回顾几个基本运算及其计算量。假设：

• 向量 \(x, y \in \mathbb{R}^K\)（形状为 \([K]\)）
• 矩阵 \(A \in \mathbb{R}^{M \times K}\)，\(B \in \mathbb{R}^{K \times N}\)

向量点积

计算 \(x \cdot y = \sum_{i=1}^K x_i y_i\)

• 需要 \(K\) 次乘法 + \(K - 1\) 次加法
  
• 总计约 \(2K - 1 \approx 2K\) 次浮点运算（FLOPs）

💡 为什么近似为 \(2K\)？

因为在大 \(K\) 时，减 1 可忽略；且现代硬件常将一次乘加（multiply-add）视为一个操作，但为统一标准，学术界通常按“乘法+加法”分别计数，故取 \(2K\)。

矩阵-向量乘法

计算 \(Ax \in \mathbb{R}^M\)

• 结果的每个元素是 \(A\) 的一行与 \(x\) 的点积
  
• 共 \(M\) 个点积，每个耗 \(2K\) FLOPs
  
• 总计算量：\(2MK\) FLOPs

矩阵-矩阵乘法

计算 \(C = AB \in \mathbb{R}^{M \times N}\)

• 可看作对 \(B\) 的每一列（共 \(N\) 列）执行一次 \(A \times (\text{列向量})\)
  
• 每列需 \(2MK\) FLOPs → 总计 \(2MKN\) FLOPs

所有这些运算的本质，都是多个向量点积的组合。因此，向量点积是矩阵乘法的基本计算单元。

换句话说，矩阵乘法就是批量的向量乘法。

高维张量乘法

在深度学习中，数据往往具有更高维度。例如，Transformer 中的注意力机制涉及三维甚至四维张量（如 [batch, head, seq_len, dim]）。此时，直接套用二维公式会出错。我们需要一套通用框架来分析任意维度下的矩阵乘法（更准确地说，是张量缩并，tensor contraction）。

为此，我们引入三种维度角色：

下面举一个实际例子来说明：

设：

• \(A \in \mathbb{R}^{\color{blue}{B} \times \color{green}{M} \times \color{red}{K}}\)
• \(B \in \mathbb{R}^{\color{blue}{B} \times \color{red}{K} \times \color{green}{N}}\)

我们计算 \(C = A \times B\)（沿最后一个维度对 \(A\) 和中间维度对 \(B\) 做矩阵乘）。

逐维度分析：

• 🔵 Batch 维度 \(B\)：两个张量都有，且不做求和 → 对每个 batch 独立做一次 \(M \times K\) 与 \(K \times N\) 的矩阵乘 → 输出保留 \(B\)
• 🔴 收缩维度 \(K\)：在 \(A\) 的最后一维、\(B\) 的中间维 → 相乘后沿 \(K\) 求和 → 不出现在输出中
• 🟢 自由维度 \(M\)（仅在 \(A\) 中）、\(N\)（仅在 \(B\) 中）→ 直接进入输出

因此，输出形状为：
\[ C \in \mathbb{R}^{\color{blue}{B} \times \color{green}{M} \times \color{green}{N}} \]

计算量分析：

• 每个 batch 执行一次 \(M \times K\) 与 \(K \times N\) 的矩阵乘 → 耗 \(2MKN\) FLOPs
  
• 共有 \(B\) 个 batch → 总计算量为
  \[ \boxed{2 \cdot \color{blue}{B} \cdot \color{green}{M} \cdot \color{green}{N} \cdot \color{red}{K}} \quad \text{FLOPs} \]

下面给出一个 通用规则：

对于任意两个张量的乘法，总 FLOPs =
\[ 2 \times (\text{所有批处理维度的乘积}) \times (\text{所有自由维度的乘积}) \times (\text{所有收缩维度的乘积}) \]

einsum

在 PyTorch 中，torch.einsum（爱因斯坦求和约定）是实现任意维度张量缩并的直观且强大的工具。它通过一个简洁的字符串公式，明确指定了输入张量的维度角色（收缩维度、批处理维度、自由维度）以及输出张量的形状。

基本语法

torch.einsum 的基本调用格式为：

torch.einsum('下标公式', 张量1, 张量2, ...)

其中下标公式是一个字符串，描述了输入张量的维度标签和输出张量的维度标签。

维度标签规则

1. 小写字母（如 i, j, k） 表示维度标签。
2. 逗号 , 分隔不同输入张量的维度标签。
3. 箭头 -> 后面是输出张量的维度标签。
4. 重复出现在输入中但未出现在输出中的标签，即为收缩维度（🔴），会沿该维度求和消去。
5. 同时出现在输入和输出中的标签，即为批处理维度（🔵），会保留并做逐元素对应。
6. 仅出现在一个输入中且出现在输出中的标签，即为自由维度（🟢），会传递到输出。

示例

矩阵乘法

A = torch.randn(M, K)  # 形状 [M, K]
B = torch.randn(K, N)  # 形状 [K, N]
# 计算 C = A @ B
C = torch.einsum('mk,kn->mn', A, B)

• k 在输入中出现两次，在输出中未出现 → 收缩维度（🔴）。
• m 仅出现在 A 和输出中 → 自由维度（🟢）。
• n 仅出现在 B 和输出中 → 自由维度（🟢）。
• 无批处理维度。
• 计算量：2 * M * N * K FLOPs。

带批处理的高维张量乘法

A = torch.randn(B, M, K)  # 形状 [B, M, K]
B = torch.randn(B, K, N)  # 形状 [B, K, N]
# 计算 C = A @ B（沿 K 维收缩）
C = torch.einsum('bmk,bkn->bmn', A, B)

• k 在输入中出现两次，输出中未出现 → 收缩维度（🔴）。
• b 在两个输入和输出中都出现 → 批处理维度（🔵）。
• m 仅出现在 A 和输出中 → 自由维度（🟢）。
• n 仅出现在 B 和输出中 → 自由维度（🟢）。
• 输出形状：[B, M, N]。
• 计算量：2 * B * M * N * K FLOPs。

更复杂的四维张量缩并

# Q: [batch, head, seq_len_q, dim]
# K: [batch, head, seq_len_k, dim]
Q = torch.randn(B, H, Lq, D)
K = torch.randn(B, H, Lk, D)
# 计算注意力分数 S = Q @ K^T，保持批和头维度并行
S = torch.einsum('bhqd,bhkd->bhqk', Q, K)

• d（dim）在输入中出现两次，输出中未出现 → 收缩维度（🔴）。
• b（batch）和 h（head）在两个输入和输出中都出现 → 批处理维度（🔵）。
• q（seq_len_q）和 k（seq_len_k）分别仅出现在一个输入和输出中 → 自由维度（🟢）。
• 输出形状：[B, H, Lq, Lk]。
• 计算量：2 * B * H * Lq * Lk * D FLOPs。

使用 einsum 的优势

1. 表达清晰：直接通过公式表达计算意图，避免复杂的 permute、view、matmul 组合。
2. 维度角色一目了然：从公式中可直接识别收缩、批处理和自由维度。
3. 便于计算量分析：根据公式中的重复标签和输出标签，可快速套用通用 FLOPs 公式： \[ \text{FLOPs} = 2 \times (\prod \text{批处理维度大小}) \times (\prod \text{自由维度大小}) \times (\prod \text{收缩维度大小}) \]
4. 支持任意高阶张量：无需将张量展平或重排，直接进行缩并。

小结

• 向量点积是所有矩阵运算的基础单元，耗约 \(2K\) FLOPs。
• 二维矩阵乘法 \(AB\) 的计算量为 \(2MKN\)。
• 高维张量乘法需区分三类维度：
  • 🔴 收缩维（求和消去）
  • 🔵 批处理维（并行独立计算）
  • 🟢 自由维（直接保留）
• 总 FLOPs = \(2 \times \prod(\text{batch}) \times \prod(\text{free}) \times \prod(\text{contracting})\)