广播机制（涉及tf.multiply和pytorch.matmul）

广播机制与矩阵乘法详解

原创已于 2025-11-07 11:39:54 修改 · 457 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#lstm #机器学习 #人工智能 #数据结构 #算法

于 2024-07-11 10:19:06 首次发布

这里是逐元素运算的广播规则，最后还有矩阵乘法的广播规则介绍，注意区分概念。
广播机制（broadcasting）使得不同维度的张量在进行逐元素运算（如加减乘除）时能够自动扩展为兼容的形状。

从右到左对齐维度规则是指：在操作两个张量时，从它们的形状数组的最后一个维度开始比较，逐步向前进行匹配。这种匹配方式适用于任意数量的维度。
详细规则
在这里插入图片描述
广播的执行过程
1️⃣ 从最后一个维度开始对齐。
2️⃣ 按上面三条规则检查每个维度是否可广播。
3️⃣ 缺少的高维在前面补 1。
4️⃣ 广播的结果形状是对应维度的最大值。

规则解析

从右到左比较：

对两个张量的形状，从最后一维开始，逐个比较每一维的大小。

- 匹配条件：

维度相等
某个张量的当前维度为 1（可以通过广播扩展匹配）。
某个张量缺失该维度（会被认为是 1）。

扩展：
若某个张量在当前维度为 1，则扩展为另一个张量的大小。
若某个张量缺少某维度，则在操作中隐式补充 1 并扩展为另一个张量的大小。

在这里插入图片描述

在DeepFM的代码中，embedding向量与特征值相乘时用到广播机制, 其含义：
广播是指在不同形状的张量之间进行逐元素操作时，TensorFlow 自动扩展较小的张量以匹配较大张量的形状的过程。这样可以避免显式复制数据，提高计算效率。

            # model     embeddings = [batch_size, field_size, embedding_size]
            self.embeddings = tf.nn.embedding_lookup(self.weights['feature_embeddings'], self.feat_index)   # (batch_size、field_size、embedding_size)

            feat_value = tf.reshape(self.feat_value, shape=[-1, self.field_size, 1])    # -1 让 TensorFlow 自动确定第一个维度的大小，
                                                                                         # 以保证总的元素数量不变。-1(N)表示样本数量，
                                                                                        # self.field_size（F）表示特征域数量
            # print(feat_value.shape)
            # (batch_size、field_size、embedding_size)
            self.embeddings = tf.multiply(self.embeddings, feat_value)  # 逐元素相乘,內积 tf.multiply支持广播机制

示例:
self.feat_value 的形状为 [2, 3]，其中 batch_size=2，field_size=3：

self.feat_value = [
    [1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0]
]

self.embeddings 的形状为 [2, 3, 8]，其中 batch_size=2，field_size=3，embedding_size=8：

self.embeddings = [
    [
        [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8],
        [0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6],
        [1.7, 1.8, 1.9, 2.0, 2.1, 2.2, 2.3, 2.4]
    ],
    [
        [2.5, 2.6, 2.7, 2.8, 2.9, 3.0, 3.1, 3.2],
        [3.3, 3.4, 3.5, 3.6, 3.7, 3.8, 3.9, 4.0],
        [4.1, 4.2, 4.3, 4.4, 4.5, 4.6, 4.7, 4.8]
    ]
]

在将 self.feat_value 调整形状后：

self.feat_value = tf.reshape(self.feat_value, shape=[-1, self.field_size, 1])

调整后的 self.feat_value 的形状为 [2, 3, 1]：

self.feat_value = [
    [
        [1.0],
        [2.0],
        [3.0]
    ],
    [
        [4.0],
        [5.0],
        [6.0]
    ]
]

当与 self.embeddings 进行逐元素相乘时：

self.embeddings = tf.multiply(self.embeddings, self.feat_value)

结果将是：

self.embeddings = [
    [
        [0.1 * 1.0, 0.2 * 1.0, 0.3 * 1.0, 0.4 * 1.0, 0.5 * 1.0, 0.6 * 1.0, 0.7 * 1.0, 0.8 * 1.0],
        [0.9 * 2.0, 1.0 * 2.0, 1.1 * 2.0, 1.2 * 2.0, 1.3 * 2.0, 1.4 * 2.0, 1.5 * 2.0, 1.6 * 2.0],
        [1.7 * 3.0, 1.8 * 3.0, 1.9 * 3.0, 2.0 * 3.0, 2.1 * 3.0, 2.2 * 3.0, 2.3 * 3.0, 2.4 * 3.0]
    ],
    [
        [2.5 * 4.0, 2.6 * 4.0, 2.7 * 4.0, 2.8 * 4.0, 2.9 * 4.0, 3.0 * 4.0, 3.1 * 4.0, 3.2 * 4.0],
        [3.3 * 5.0, 3.4 * 5.0, 3.5 * 5.0, 3.6 * 5.0, 3.7 * 5.0, 3.8 * 5.0, 3.9 * 5.0, 4.0 * 5.0],
        [4.1 * 6.0, 4.2 * 6.0, 4.3 * 6.0, 4.4 * 6.0, 4.5 * 6.0, 4.6 * 6.0, 4.7 * 6.0, 4.8 * 6.0]
    ]
]

广播机制使得 self.feat_value 的形状 [batch_size, field_size, 1] 可以与 self.embeddings 的形状 [batch_size, field_size, embedding_size] 进行逐元素相乘，从而无需显式地复制 self.feat_value 的数据。这提高了计算的效率和内存的利用率。

矩阵乘法广播规则 pytorch.matmul
torch.matmul 的广播规则更复杂一些。
因为它不是逐元素操作，而是矩阵乘法，它会自动把前面的维度（即 batch 维度）广播。
记忆口诀：
在 matmul 中：
最后两维做矩阵乘法；
其余前面所有维度自动广播。

假设：

query.shape = [B, H, L, D]
key.shape   = [B, H, L, D]

执行：

scores = torch.matmul(query, key.transpose(-2, -1))

🔍 分析逐步广播过程：

1️⃣ 先看矩阵相乘的“核心两维”：

query 的最后两维是 [L, D]
key.transpose(-2, -1) 的最后两维是 [D, L]
→ 这两维完全匹配矩阵乘法规则：
(L, D) @ (D, L) → 结果是 (L, L)。

2️⃣ 再看前面的批次维度 [B, H]

两个张量在这两维上完全一致，PyTorch 会认为它们是batch维度。

对于每个 [B, H] 的组合，独立执行一次矩阵乘法。

3️⃣ 所以，最终 scores 的形状是：

[B, H, L, L]

也就是说：

👉 每个 batch、每个 head 都会得到一个自己的 [L, L] 注意力矩阵。

💡 四、如果前面的批次维不同怎么办？

假设：

query.shape = [2, 4, 10, 64]
key.shape   = [1, 4, 10, 64]

从右往左比较：
在这里插入图片描述