为什么 Amplect 应该取代张量运算中的广播机制
所有张量算法的核心是矩阵乘法,而矩阵乘法在底层严重依赖于广播机制,这是一种允许不同形状张量之间进行操作的机制。虽然广播机制在许多情况下都有效,但其任意的规则既不优雅又容易出错。
有一种更好的方法。遵循爱因斯坦求和约定 (Einop notation) 的模式,Amplect 提供了一种更清晰、更一致的方式来指定这些逐元素张量操作。
为什么广播机制很奇怪
假设您想添加两个张量
T1 的形状为 [a b c]
,T2 的形状为 [c]
。
广播机制可以优雅地处理这种情况。它会自动拉伸 T2 以匹配 T1 的形状进行加法运算。
但是,如果您希望在 T2 的形状为 [b]
时执行相同的加法运算呢?现在广播机制失败了,声明这些张量“不可广播”。为了使其工作,您需要手动将 T2 重塑为 [b 1]
,但这又是为什么呢?问题在于广播机制的任意规则。
广播规则
为了理解广播机制为何表现得如此不一致,让我们看看它具体是如何工作的。
广播机制逐维度比较张量,从最右边的维度开始
- 如果两个维度大小相同 -> 继续
- 如果任一维度大小为 1 -> 将其拉伸以匹配另一个维度 -> 继续
- 如果张量耗尽维度 -> 将缺失维度视为大小为 1
- 否则 -> 广播失败
示例 1:T1 [a b c]
和 T2 [c]
- 比较
c
与c
-> 相等,继续 - 比较
b
与 (缺失) -> 将 T2 视为大小为 1,拉伸到b
,继续 - 比较
a
与 (缺失) -> 将 T2 视为大小为 1,拉伸到a
,继续
结果:广播成功。T2 实际上变为 [a b c]
。
示例 2:T1 [a b c]
和 T2 [b]
- 比较
c
与b
-> 大小不同,两者都不是 1 -> 广播失败
结果:广播失败。
示例 3:T1 [a b c]
和 T2 [b 1]
- 比较
c
与1
-> T2 的维度为 1,拉伸到c
,继续 - 比较
b
与b
-> 相等,继续 - 比较
a
与 (缺失) -> 将 T2 视为大小为 1,拉伸到a
,继续
结果:广播成功。T2 实际上变为 [a b c]
。
广播机制常见问题
假设您正在处理一批序列
data
形状为 [batch, seq_len, features]
,mask
形状为 [batch, seq_len]
您想通过将 mask 与数据相乘来将填充的 token 归零。
masked = data * mask
但由于上述规则,这两个张量不可广播。
为了使其工作,您必须重塑 mask
masked = data * mask[..., None]
对于有经验的用户来说,[..., None]
可能已经习以为常,但这是广播机制任意维度匹配规则的一种变通方法。
为什么要记住拉伸哪个轴?意图很明显,对齐 batch
和 seq_len
,并跨特征执行逐元素乘法。然而,广播机制却强制您手动操作维度,只是为了表达一些在数学上简单明了的东西。
广播机制将我们限制在一个框中
上述示例揭示了广播机制的根本缺陷:其规则是任意的。为什么 [a b c]
和 [c]
兼容,而 [a b c]
和 [b]
不兼容?这两个操作在数学上都说得通。
事实上,任何两个张量都可以通过它们的外积结合。
假设张量 A
的形状为 [a b c]
,张量 B
的形状为 [d e f]
。
广播机制将这些称为“不可广播”,但如果我们给 A
添加足够多的单例维度,使其形状变为 [a b c 1 1 1]
,它们就可以突然广播为形状 [a b c d e f]
,这就是外积。
问题不在于某些操作不可能,而在于广播机制强制我们根据任意的定位规则来猜测我们想要哪种特定操作。
我们需要一种明确表达意图的方法。
Amplect
Amplect 通过让我们明确指定哪些维度相互对应来解决广播机制的猜测问题。我们直接声明我们的意图,而不是依赖于基于位置的规则。这是一种理论方法,最终应该取代广播机制。
还记得我们有问题的情况吗?
我们想添加形状为 [a b c]
和 [b]
的张量。
使用 Amplect,我们只需写下
a b c, b -add> a b c
这告诉计算机:“取第二个张量的 b
维度,并将其与第一个张量的 b
维度对齐进行加法运算。”
对于输出中形状为 [a b c]
的每个位置 [i, j, k]
,操作变得清晰
Output[i, j, k] = A[i, j, k] + B[j]
我们之前示例中的外积变为
a b c, d e f -multiply> a b c d e f
对于输出中形状为 [a b c d e f]
的每个位置 [i, j, k, l, m, n]
Output[i, j, k, l, m, n] = A[i, j, k] * B[l, m, n]
我们甚至可以组合超过 2 个张量
a b, a c, b c, c d -multiply> a b c d
对于输出中形状为 [i, j, k, l]
的每个位置 [a b c d]
Output[i, j, k, l] = A[i, j] * B[i, k] * C[j, k] * D[k, l]
对于多张量操作,如果函数是可结合的,amplect 出错的可能性更小。
这是因为在并行硬件上,对于 n 个张量,我们可以执行 O(log(n)) 个并行步骤,而不是 O(n) 个顺序步骤。
结合律是函数的一种性质,其中求值顺序无关紧要。
所以
f(f(x, y), z) = f(x, f(y, z))
例如,乘法
(a * b) * c = a * (b * c)
让我们的常见问题不再那么令人沮丧
我们可以将当前和提议的方法与我之前分享的挫折感进行比较。
广播机制
masked = data * mask[..., None]
Amplect
masked = amplect("batch seq features, batch seq -multiply> batch seq features", data, mask)
您可以看到,我们不仅避免了 [..., None]
的变通方法,而且意图也更加清晰明了。
结论
广播机制基于位置的规则在张量编程中造成了不必要的摩擦。当 [a b c]
和 [c]
可以一起使用,而 [a b c]
和 [b]
不行时,我们处理的不是数学限制;我们正在与任意的语法作斗争。
Amplect 消除了这种猜测游戏。通过明确声明哪些维度对齐,我们编写的代码能够清晰地表达我们的意图。不再需要对维度顺序进行脑力体操,不再出现意外的“不可广播”错误,也不再仅仅为了满足任意规则而重塑张量。
结果是张量代码读起来就像它实际所做的那样,这是使张量编程更直观、更不容易出错的一步。
随着张量编程变得越来越复杂,转向像 Amplect 这样明确的、意图驱动的工具将减少错误,并使代码更易于阅读和维护。