本文涵盖torch.dot、@、matmul等8大操作符区别 | 附批量计算性能优化技巧 | 大模型注意力机制源码解析。

1. 点积的数学本质

1.1 代数定义

对于两个n维向量 a= [a₁, a₂, ..., aₙ] 和 b= [b₁, b₂, ..., bₙ]：

a · b = Σ(aᵢ × bᵢ) = a₁b₁ + a₂b₂ + ... + aₙbₙ

1.2 几何定义

a · b = |a| × |b| × cos(θ)

其中：

• |a|, |b|：向量的模（长度）
• θ：两向量间的夹角【余弦】

1.3 几何意义解读

• θ = 0°（同向）：点积最大且为正
• θ = 90°（垂直）：点积为0
• θ = 180°（反向）：点积最小且为负
• 投影意义：a·b = (a在b上的投影长度) × |b|

2. PyTorch中的点积操作

先总结，后分条阐述

🎯 形状推断记忆口诀

• "去尾对齐"：(a,b) @ (b,c) → (a,c)
• "保留批次"：(batch,a,b) @ (batch,b,c) → (batch,a,c)
• "广播优先"：(a,b) @ (batch,b,c) → (batch,a,c)

🎯 选择指南

1. 日常使用：优先用 @运算符
2. 明确场景：用专用函数（如torch.mv用于矩阵×向量）
3. 复杂操作：考虑torch.einsum
4. 避免混淆：分清*（逐元素）和@（矩阵乘法）

运算符/函数	输入形状要求	输出形状	主要用途	等价写法	注意事项
*	任意相同形状或可广播	与输入同形	逐元素相乘	torch.mul()	不是矩阵乘法！
@	任意兼容形状	根据输入决定	通用矩阵乘法（推荐）	torch.matmul()	最简洁高效，支持广播
torch.matmul	任意兼容形状	根据输入决定	通用矩阵乘法	a @ b	功能最全，支持批量
torch.dot	必须都是1D	标量(0D)	严格向量点积	a @ b（仅1D时）	只适用于一维向量
torch.mv	矩阵(2D) × 向量(1D)	1D向量	矩阵与向量乘法	matrix @ vector	专门优化此场景
torch.mm	必须都是2D	2D矩阵	矩阵乘法（无批量）	a @ b（仅2D时）	已逐渐被@替代
torch.bmm	必须都是3D且批次相同	3D张量	批量矩阵乘法，Transformer 用到	a @ b（3D时）	专用批量版本
torch.einsum	任意形状	根据公式决定	爱因斯坦求和约定	无直接等价	最灵活但语法复杂

2.1 torch.dot() 严格的向量点积

📌 核心特性：只适用于一维张量

import torch

# 基本用法
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = torch.dot(a, b)  # tensor(32)
# 计算：1 * 4 + 2 * 5 + 3 * 6 = 32

# 错误示例：二维张量会报错
A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[5, 6], [7, 8]])
# torch.dot(A, B)  # 报错：RuntimeError

2.2 torch.matmul 或 @ 通用的矩阵乘法

📌 核心特性：支持各种维度的张量乘法

2.2.1 向量点积（与torch.dot等价）

# 1. 向量点积（与torch.dot等价）
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result1 = torch.matmul(a, b)  # tensor(32)
result2 = a @ b               # tensor(32) - 推荐写法

2.2.2 矩阵乘法

📌 重要规律总结

计算过程：矩阵乘法 A @ B的结果中每个元素都是 A的对应行向量与 B的对应列向量的点积

1. 维度匹配：A的列数必须等于B的行数
2. 结果维度：结果矩阵的行数 = A的行数，列数 = B的列数
3. 计算顺序：总是"行点积列"
4. 不可交换：A @ B≠ B @ A（矩阵乘法不满足交换律）

A = torch.tensor([[1, 2], [3, 4]])  # 2x2
B = torch.tensor([[5, 6], [7, 8]])  # 2x2
result = A @ B
# 结果
# tensor([[19, 22],
#         [43, 50]])


# A = [[1, 2],   B = [[5, 6],
#      [3, 4]]        [7, 8]]
# 计算过程如下：
计算结果矩阵的(0,0)位置（第一行第一列）
取A的第0行：[1, 2]
取B的第0列：[5, 7]
计算点积：1×5 + 2×7 = 5 + 14 = 19

计算结果矩阵的(0,1)位置（第一行第二列）
取A的第0行：[1, 2]
取B的第1列：[6, 8]
计算点积：1×6 + 2×8 = 6 + 16 = 22

计算结果矩阵的(1,0)位置（第二行第一列）
取A的第1行：[3, 4]
取B的第0列：[5, 7]
计算点积：3×5 + 4×7 = 15 + 28 = 43

计算结果矩阵的(1,1)位置（第二行第二列）
取A的第1行：[3, 4]
取B的第1列：[6, 8]
计算点积：3×6 + 4×8 = 18 + 32 = 50

2.2.3 批量矩阵乘法（大模型常用）

batch_A = torch.randn(10, 3, 4)  # 10个3x4矩阵
batch_B = torch.randn(10, 4, 5)  # 10个4x5矩阵
batch_result = batch_A @ batch_B  # 结果形状: (10, 3, 5)

🎯 批量矩阵乘法形状推断技巧

• 核心规则："去尾对齐，保留批次"
• 记忆口诀："批次不变，中间维度消失，头尾维度结合"
• 一句话掌握：批量矩阵乘法的结果形状 = (批次维度 + 第一个矩阵的行数 + 第二个矩阵的列数)

🎯 通用推断公式

• 对于 (batch, m, p) @ (batch, p, n)：
• 结果形状= (batch, m, n)

💡 快速推断技巧（3秒搞定）

1. 看最后两个维度：(3,4) @ (4,5) → (3,5)（普通矩阵乘法规则）
2. 看批次维度：(10,3,4) @ (10,4,5) → 批次维度不变 → (10,3,5)

batch_A = (10, 3, 4)  # 理解为：10个 3×4 的矩阵
batch_B = (10, 4, 5)  # 理解为：10个 4×5 的矩阵

# 计算过程：对应位置的矩阵相乘
# batch_A[0] @ batch_B[0] → (3,4) @ (4,5) = (3,5)
# batch_A[1] @ batch_B[1] → (3,4) @ (4,5) = (3,5)
# ...
# batch_A[9] @ batch_B[9] → (3,4) @ (4,5) = (3,5)

# 最终得到：10个 (3,5) 的矩阵 → 形状 (10, 3, 5)

2.3 torch.mv 矩阵与向量乘法

📌 核心特性：专门处理矩阵×向量的情况

M = torch.tensor([[1, 2], [3, 4]])  # 2x2矩阵
v = torch.tensor([5, 6])            # 1D向量
result = torch.mv(M, v)  # tensor([17, 39])
# 计算：[1 * 5+2 * 6, 3 * 5+4 * 6] = [17, 39]

2.4 逐元素相乘后求和（类似点积的逻辑）

🎯 理解要点：这实际上计算的是两个矩阵所有对应位置元素的"总体相似度"。

# 逐元素相乘后求和（类似点积的逻辑）
A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[5, 6], [7, 8]])
result = (A * B).sum()  # tensor(70)
# 相当于把所有元素当作一维向量求点积

2.4.1 计算过程拆解

A = [[1, 2],   B = [[5, 6],
     [3, 4]]        [7, 8]]

# 步骤1：逐元素相乘 (A * B)
element_wise = A * B = [[1 * 5, 2 * 6],   = [[5, 12],
                        [3 * 7, 4 * 8]]      [21, 32]]

# 步骤2：对所有元素求和
result = 5 + 12 + 21 + 32 = 70

2.4.2 为什么这相当于点积？

🎯 因为如果把两个矩阵展平成一维向量，这个计算就是标准的点积：

A_flat = [1, 2, 3, 4]  # 按行展平
B_flat = [5, 6, 7, 8]  # 按行展平

# 点积计算
dot_product = 1 * 5 + 2 * 6 + 3 * 7 + 4 * 8 = 5 + 12 + 21 + 32 = 70

2.5 批量点积（保持批次维度）

batch_a = torch.randn(10, 3)  # 10个3维向量
batch_b = torch.randn(10, 3)  # 10个3维向量

batch_dot = (batch_a * batch_b).sum(dim=1)  # 形状: (10,)

2.5.1 数据形状理解：

batch_a.shape = (10, 3)  # 10个样本，每个样本是3维向量
batch_b.shape = (10, 3)  # 10个样本，每个样本是3维向量

2.5.2 计算过程可视化：

假设具体数值为：

batch_a = [[a11, a12, a13],  # 样本1的3维向量
           [a21, a22, a23],  # 样本2的3维向量
           ...]              # 共10个样本

batch_b = [[b11, b12, b13],  # 样本1的3维向量  
           [b21, b22, b23],  # 样本2的3维向量
           ...]              # 共10个样本

2.5.3 逐步计算：

# 步骤1：逐元素相乘
element_wise = batch_a * batch_b = 
[[a11*b11, a12*b12, a13*b13],  # 样本1：对应元素相乘
 [a21*b21, a22*b22, a23*b23],  # 样本2：对应元素相乘
 ...]

# 步骤2：沿dim=1求和（对每个样本内部的3个元素求和）
batch_dot = [
    a11*b11 + a12*b12 + a13*b13,  # 样本1的点积结果
    a21*b21 + a22*b22 + a23*b23,  # 样本2的点积结果  
    ...
]

形状变化理解：

🎯 dim=1的含义：在第1个维度（列维度）上进行求和，但保留第0个维度（样本维度）

输入: (10, 3) * (10, 3) → 逐元素乘 → (10, 3)
然后: (10, 3).sum(dim=1) → 沿第1维求和 → (10,)

实际数值示例：

🎯 核心思想：dim=1表示"在保持批次维度的情况下，对每个样本内部进行点积计算"

# 具体数值示例
batch_a = torch.tensor([[1, 2, 3],    # 样本1
                        [4, 5, 6]])   # 样本2

batch_b = torch.tensor([[2, 2, 2],    # 样本1  
                        [3, 3, 3]])   # 样本2

# 计算过程
element_wise = batch_a * batch_b = [[1 * 2, 2 * 2, 3 * 2],   # [[2, 4, 6]
                                    [4 * 3, 5 * 3, 6 * 3]]   #  [12, 15, 18]]

batch_dot = element_wise.sum(dim=1) = [2+4+6, 12+15+18] = [12, 45]

2.5.4 关键区别总结

计算方式	输入形状	输出形状
(A * B).sum()	(2,2) * (2,2) → 标量	把所有元素当作一个向量求点积
(batch_a * batch_b).sum(dim=1)	(10,3) * (10,3) → (10,)	对每个样本独立计算点积

3. NumPy与PyTorch的差异对比

3.1 numpy.dot() 的多功能行为

import numpy as np

# 一维数组：点积（与torch.dot相同）
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
print(np.dot(a, b))  # 32

# 二维数组：矩阵乘法（与torch.matmul相同）
A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])
print(np.dot(A, B))
# [[19 22]
#  [43 50]]

# 高维数组：张量缩并
# 行为比torch.dot复杂得多

3.2 关键差异总结

操作	PyTorch	NumPy	说明
向量点积	torch.dot(a, b)	np.dot(a, b)	行为相同
矩阵乘法	A @ B或 torch.matmul(A, B)	np.dot(A, B)	关键差异
批量处理	torch.matmul支持批量	np.dot行为复杂	PyTorch更一致

4. 大模型中的点积应用场景

4.1 自注意力机制（Self-Attention）

# 简化的自注意力计算
def self_attention(Q, K, V):
    """
    Q: Query矩阵 (batch_size, seq_len, d_k)
    K: Key矩阵 (batch_size, seq_len, d_k)  
    V: Value矩阵 (batch_size, seq_len, d_v)
    """
    # 计算注意力分数：Q和K的点积
    scores = torch.matmul(Q, K.transpose(-2, -1))  # (batch_size, seq_len, seq_len)

    # 缩放
    d_k = Q.size(-1)
    scores = scores / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))

    # Softmax归一化
    attention_weights = torch.softmax(scores, dim=-1)

    # 加权求和
    output = torch.matmul(attention_weights, V)  # (batch_size, seq_len, d_v)
    return output

4.2 词向量相似度计算

# 计算词向量间的相似度
word_embeddings = torch.randn(10000, 300)  # 词表大小10000，维度300
query_vector = torch.randn(300)

# 高效计算所有词与查询词的相似度
similarities = torch.matmul(word_embeddings, query_vector)  # 形状: (10000,)
top_k_similar = torch.topk(similarities, k=10)

4.3 分类器输出计算

# 多分类任务的输出计算
features = torch.randn(32, 512)  # 批量大小32，特征维度512
classifier_weights = torch.randn(10, 512)  # 10个类别

# 计算每个样本属于每个类别的分数
logits = torch.matmul(features, classifier_weights.T)  # 形状: (32, 10)

5. 性能优化与最佳实践

5.1 选择正确的操作符

5.1.1 使用 @运算符（最简洁高效）

• 避免：不必要的转换
• 错误：先将矩阵展平再点积
• 正确：直接使用矩阵乘法

# 以下三种写法等价，但 @ 最推荐
result1 = tensor1 @ tensor2           # ✅ 最推荐
result2 = torch.matmul(tensor1, tensor2)  # ✅ 也可以
result3 = torch.mm(tensor1, tensor2)  # ⚠️ 仅适用于2D矩阵

❌ 错误：先将矩阵展平再点积

A = torch.randn(3, 4)  # 矩阵
B = torch.randn(3, 4)  # 矩阵

# ❌ 低效做法：展平后点积
result_bad = torch.dot(A.flatten(), B.flatten())

✅ 正确做法：直接计算Frobenius内积

result_good = (A * B).sum()
# 或者使用更数学的表达
result_better = torch.sum(A * B)

5.1.2 批量处理时使用矩阵乘法而不是循环

❌ 低效做法（循环）

batch_size, dim = 1000, 768
tensors_a = torch.randn(batch_size, dim)  # 1000个768维向量
tensors_b = torch.randn(batch_size, dim)  # 1000个768维向量

# ❌ 低效：Python循环
inefficient = torch.tensor([torch.dot(a, b) for a, b in zip(tensors_a, tensors_b)])

问题：

• Python循环每次迭代都有开销
• 无法利用GPU的并行计算能力
• 每次torch.dot()调用都有函数调用开销

✅ 高效做法（向量化）

# ✅ 高效：批量矩阵乘法
efficient = torch.matmul(tensors_a.unsqueeze(1), tensors_b.unsqueeze(2)).squeeze()

分解理解：

• 步骤1：改变形状以便批量矩阵乘法

• • tensors_a: (1000, 768) → unsqueeze(1) → (1000, 1, 768)
  • tensors_b: (1000, 768) → unsqueeze(2) → (1000, 768, 1)

• 步骤2：批量矩阵乘法

• • (1000, 1, 768) @ (1000, 768, 1) = (1000, 1, 1)
  • 每个 (1,768) 与 (768,1) 矩阵乘法得到 (1,1)

• 步骤3：压缩多余的维度

• • (1000, 1, 1) → squeeze() → (1000,)

5.1.3 性能对比

import torch
import time

batch_size, dim = 10000, 768
tensors_a = torch.randn(batch_size, dim).cuda()  # 使用GPU
tensors_b = torch.randn(batch_size, dim).cuda()

# 测试循环方法
start = time.time()
result_loop = torch.tensor([torch.dot(a, b) for a, b in zip(tensors_a, tensors_b)]).cuda()
time_loop = time.time() - start

# 测试向量化方法  
start = time.time()
result_vec = (tensors_a * tensors_b).sum(dim=1)
time_vec = time.time() - start

print(f"循环方法: {time_loop:.4f}秒")
print(f"向量化方法: {time_vec:.4f}秒") 
print(f"加速比: {time_loop/time_vec:.1f}x")

# 验证结果相同
print(f"结果是否一致: {torch.allclose(result_loop, result_vec)}")

5.2 内存布局优化

看另一篇文章👇

此处为语雀内容卡片，点击链接查看：https://www.yuque.com/lhyyh/ai/vw9pacthew5anp84?singleDoc#

6. 总结对比表

操作	输入维度	输出维度	主要用途	PyTorch函数	推荐写法
向量点积	1D, 1D	0D (标量)	向量相似度	torch.dot()	a @ b
矩阵乘法	2D, 2D	2D	线性变换	torch.mm()	A @ B
批量矩阵乘	3D+, 3D+	3D+	批量处理	torch.matmul()	A @ B
矩阵×向量	2D, 1D	1D	投影变换	torch.mv()	M @ v
广播乘法	任意	任意	灵活计算	torch.matmul()	A @ B

7. 实用技巧与调试建议

7.1 维度检查

def safe_matmul(a, b):
    # 检查维度兼容性
    assert a.dim() >= 1 and b.dim() >= 1, "输入至少是一维张量"

    if a.dim() == 1 and b.dim() == 1:
        # 向量点积
        assert a.shape[0] == b.shape[0], "向量长度必须相同"
        return a @ b
    else:
        # 矩阵乘法或批量乘法
        assert a.shape[-1] == b.shape[-2], "最后两个维度必须匹配"
        return a @ b

7.2 梯度计算验证

# 验证点积的梯度计算
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = torch.tensor([4.0, 5.0, 6.0], requires_grad=True)

z = x @ y  # 点积
z.backward()  # 反向传播

print(x.grad)  # 应该是 tensor([4., 5., 6.])
print(y.grad)  # 应该是 tensor([1., 2., 3.])

AI知识库将陆续推出大模型算法相关技术！

https://www.yuque.com/lhyyh/ai