torch.einsum —— 爱因斯坦标记法下的张量相乘的理解

$Att{n}_{ijbh}\times Value{s}_{jdbh}=Z_{idbh}$ （爱因斯坦标记法）

这里的$i,d,b,h$于‘->’左右都出现了，被称作自由标，而 $j$ 只在左边出现，被称为哑标，计算结果则是保留自由标，对哑标进行逐元素相乘并求和

展开为求和公式

${\Sigma }_{j}Att{n}_{ijbh}\times Value{s}_{jdbh}=Z_{idbh}$

例如，对于单个head和单个batch的Attention Matrix，可以表示如下：

$\left[\begin{array}{cccc}Att{n}_{11bh}& Att{n}_{12bh}& \cdots & Att{n}_{1jbh}\\ Att{n}_{21bh}& Att{n}_{22bh}& \cdots & Att{n}_{2jbh}\\ ⋮& ⋮& \ddots & ⋮\\ Att{n}_{i1bh}& Att{n}_{i2bh}& \cdots & Att{n}_{ijbh}\end{array}\right]$

相应的Value Matrix表示为

$\left[\begin{array}{c}Value{s}_{1dbh}\\ Value{s}_{2dbh}\\ \cdots \\ Value{s}_{jdbh}\end{array}\right]$

我们可以得到Z

$\left[\begin{array}{cccc}{Z}_{11bh}& Z_{12bh}& \cdots & Z_{1dbh}\\ ⋮& ⋮& \ddots & ⋮\\ Z_{i1bh}& Z_{i2bh}& \cdots & Z_{idbh}\end{array}\right]=\left[\begin{array}{cccc}Att{n}_{11bh}& Att{n}_{12bh}& \cdots & Att{n}_{1jbh}\\ Att{n}_{21bh}& Att{n}_{22bh}& \cdots & Att{n}_{2jbh}\\ ⋮& ⋮& \ddots & ⋮\\ Att{n}_{i1bh}& Att{n}_{i2bh}& \cdots & Att{n}_{ijbh}\end{array}\right]\cdot \left[\begin{array}{c}Value{s}_{1dbh}\\ Value{s}_{2dbh}\\ \cdots \\ Value{s}_{jdbh}\end{array}\right]$

其中

$Z_{idbh}{|}_{i=1,d=1}=\left[\begin{array}{cccc}Att{n}_{ijbh}{|}_{i=1,j=1}& Att{n}_{ijbh}{|}_{i=1,j=2}& \cdots & Att{n}_{1jbh}\end{array}\right]\cdot \left[\begin{array}{c}Value{s}_{jdbh}{|}_{j=1}\\ Value{s}_{jdbh}{|}_{j=2}\\ \cdots \\ Value{s}_{jdbh}\end{array}\right]$

从中我们可以看出所谓的对哑标 $j$ 进行逐元素相乘并求和。