白话 Broadcasting

什么是 Broadcasting

在 《白话张量》 中我们讲过，张量之间进行运算需要满足一定的维数条件的，形状大小不满足条件的两个张量之间是无法进行运算的。为了能执行运算，我们通常需要将较小的张量被“拉伸”到与较大张量兼容的形状，这一步“拉伸”操作就是 Broadcasting。

图1. Broadcasting 示意图

Broadcasting 在不创建重复数据的情况下确保张量运算可以执行，从而确保算法可以更高效的执行。

Broadcasting 的规则

根据 PyTorch 的文档，张量在以下情况下可以 Broadcasting ：

每个张量至少有一个维度

迭代维度大小时，会从尾部维度开始，向前推进，维度大小要么相等，要么其中一个为1或不存在。

比较形状时，尾部维度指的是最右边的数字。我们以上面的图 1 为例，解释一下 Broadcasting 的通用过程：

1. 确定最右边的维度是否兼容
   • 每个张量是否至少有一个维度？
   • 尺寸是否相等？其中一个是 1 或不存在吗？
2. 将尺寸拉伸到合适的大小
3. 对下一个维度重复前面的步骤

逐元素操作

所有按元素操作的运算都要求张量具有相同或兼容的形状。

向量与标量运算

图2. 向量与标量运算中的 Broadcasting 示意图

在此示例中，标量的形状为 $(1,)$ ，向量的形状为 $(3,)$。 如图 2 所示，标量 $b$ 被 Broadcasting 为 $(3,)$ 的形状，并且按预期执行了 Hadamard 乘积。

上面的示例用 PyTorch 实现如下：

import torch

a = torch.tensor([1, 2, 3])
b = 2 # 会拉伸为 ([2, 2, 2])

a * b
# Output: tensor([2, 4, 6])

矩阵与向量运算

行向量

图3. 矩阵与行向量运算中的 Broadcasting 示意图

在上面示例中，矩阵 $A$ 的形状为 $(3,3)$，向量 $b$ 的形状为 $(3,)$。当运算时，向量 $b$ 被逐行拉伸成一个 $(3,3)$ 矩阵，如图 3 所示。 现在，A 和 b 的形状都是 $(3,3)$，可以进行逐元素操作。

上面的示例用 PyTorch 实现如下：

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

b = torch.tensor([1, 2, 3])

A * b
# Output: tensor([[ 1,  4,  9],
#         		  [ 4, 10, 18],
#         		  [ 7, 16, 27]])

列向量

图4. 矩阵与列向量运算中的 Broadcasting 示意图

在上面示例中，矩阵 $A$ 的形状为 $(3,3)$，列向量 $b$ 的形状为 $(3,1)$。当运算时， $b$ 被按列拉伸创建另外 2 列，如图 4 所示。 现在，A 和 b 的形状都是 $(3,3)$，可以进行逐元素操作。

上面的示例用 PyTorch 实现如下：

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

b = torch.tensor([[1], 
                  [2], 
                  [3]])
A * b
# Output: tensor([[ 1,  2,  3],
#         		  [ 8, 10, 12],
#         		  [21, 24, 27]])

张量与向量运算

$$\left[\begin{array}{c}\left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\\ \left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\end{array}\right]\odot \left[\begin{array}{c}[1]\\ [2]\\ [3]\end{array}\right]=\left[\begin{array}{c}\left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\\ \left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\end{array}\right]\odot \left[\begin{array}{c}\left[\begin{array}{c}[111]\\ [222]\\ [333]\end{array}\right]\\ \left[\begin{array}{c}[111]\\ [222]\\ [333]\end{array}\right]\end{array}\right]$$

3-D 及以上维度的张量很难用图形清晰地展示，这里我改用公式来表达。上面公式中，红字部分展示了向量如何 Broadcasting 到与张量兼容的尺寸。具体来说，上面示例中张量的形状为 $(2,3,3)$，向量的形状为 $(3,1)$，维度大小之间的对应关系为：
$$\begin{gathered} \mathcal{A}=(2,3,3) \\ b=(,3,1) \end{gathered}$$
从最右边的维度开始，每个元素按列拉伸以生成 $(3,3)$ 矩阵。中间维度尺寸相等，无需拉伸。$b$ 不存在最左边的维度，因此必须添加一个维度。因此，必须 Broadcasting $b$ 拉伸后的矩阵以创建大小为 $(2,3,3)$ 的张量，张量中包含两个 $(3,3)$ 矩阵。此时张量 $\mathcal{A}$ 与向量 $b$ Broadcasting 后得到的张量形状相同，可以计算 Hadamard 积。

上面的示例用 PyTorch 实现如下：

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

                  [[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9]]])

b = torch.tensor([[1], 
                  [2], 
                  [3]])

A * b
# Output: tensor([[[ 1,  2,  3],
# 			       [ 8, 10, 12],
#         		   [21, 24, 27]],
#
#        		  [[ 1,  2,  3],
#         		   [ 8, 10, 12],
#         		   [21, 24, 27]]])

张量与矩阵运算

$$\left[\begin{array}{c}\left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\\ \left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\end{array}\right]\odot \left[\begin{array}{c}[123]\\ [123]\\ [123]\end{array}\right]=\left[\begin{array}{c}\left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\\ \left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\end{array}\right]\odot \left[\begin{array}{c}\left[\begin{array}{c}[123]\\ [123]\\ [123]\end{array}\right]\\ \left[\begin{array}{c}[123]\\ [123]\\ [123]\end{array}\right]\end{array}\right]$$

理解了上面张量与向量的运算，张量与矩阵的运算会很简单。上面示例中张量的形状为 $(2,3,3)$，矩阵的形状为 $(3,3)$，维度大小之间的对应关系为：
$$\begin{gathered} \mathcal{A}=(2,3,3) \\ B=(,3,3) \end{gathered}$$
这个例子比上一个更容易，因为最右边的两个维度是相同的。这意味着矩阵只需在最左边的维度上 Broadcasting 即可变为 $(2,3,3)$ 的形状。而这仅意味着需要一个额外的矩阵。

上面的示例用 PyTorch 实现如下：

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

B = torch.tensor([[1, 2, 3], 
                  [1, 2, 3], 
                  [1, 2, 3]])

A * B
# Output: tensor([[[ 1,  4,  9],
#          		   [ 4, 10, 18],
#          		   [ 7, 16, 27]],
# 
#         		  [[ 1,  4,  9],
#          		   [ 4, 10, 18],
#          		   [ 7, 16, 27]]])

矩阵与张量的点积

前面所有的示例，目标都是最终得到相同的形状以允许逐元素操作。而点积并不要求两个张量形状一样，只需要第一个矩阵或张量的最后一个维度与第二个矩阵或张量的倒数第二个维度相同即可（详见 《白话张量》 ）。

矩阵点积要求
$$(m,n)\times (n,r)=(m,r)$$
3维张量点积要求
$$(c,m,n)\times (c,n,r)=(c,m,r)$$
4维张量点积要求
$$(z,c,m,n)\times (z,c,n,r)=(z,c,m,r)$$
以此类推，可以扩展到任意维张量。

举个例子
$$\left[\begin{array}{c}\left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\\ \left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\end{array}\right]\cdot \left[\begin{array}{c}[12]\\ [12]\\ [12]\end{array}\right]=\left[\begin{array}{c}\left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\\ \left[\begin{array}{c}[123]\\ [456]\\ [789]\end{array}\right]\end{array}\right]\cdot \left[\begin{array}{c}\left[\begin{array}{c}[12]\\ [12]\\ [12]\end{array}\right]\\ \left[\begin{array}{c}[12]\\ [12]\\ [12]\end{array}\right]\end{array}\right]$$
上面例子中，张量 $\mathcal{A}$ 的形状为 $(2,3,3)$，矩阵 $B$ 的形状为 $(3,2)$。截至目前，最后两个维度是符合点积乘法条件的。只需要将一个维度添加到 $B$，在该维度上 Broadcasting $(3,2)$ 矩阵到 $(2,3,2)$ 的形状。因此实例中点积结果为 $(2,3,3)\times (2,3,2)=(2,3,2)$。

上面的示例用 PyTorch 实现如下：

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

B = torch.tensor([[1, 2], 
                  [1, 2], 
                  [1, 2]])

A @ B # A.matmul(B)
# Output: tensor([[[ 6, 12],
#         		   [15, 30],
#         		   [24, 48]],
#
#        		  [[ 6, 12],
#         		   [15, 30],
#         		   [24, 48]]])

总结

Broadcasting 还是很好理解的，并且在实际开发中机器学习框架和库会自动处理，但了解 Broadcasting 的机制对理解代码实现非常有益，因此建议大家还是要了解一下 Broadcasting 。