Transformer の位置エンコーディング

import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure()
ax = fig.add_subplot(projection='3d')

ax.set_xlabel('i', size=15, color='black')
ax.set_ylabel('pos', size=15, color='black')
ax.set_zlabel('PE', size=15, color='black')

x, y = np.meshgrid(np.linspace(0,40,100),np.linspace(0,40,100))
z = np.sin(y/(10**(x/60)))

ax.plot_surface(x, y, z)
plt.show()

埋め込み空間の偶数インデックスのケースについて 3D 空間で可視化。インデックスが小さいうちは波長が短く、インデックスが大きくなると波長が長くなっていく。
$pos \simeq pos^\prime$ の時に $PE_{pos} \simeq PE_{pos^\prime}$ となるのかはいまいちピンと来ないが、そういう狙いではあるらしい・・・。

$\alpha = 10000^{1/d_{model}}$ として、

\begin{align*}
PE_{(pos + k,2i)} &= \cos(k/\alpha^{2i}) \sin(pos/\alpha^{2i}) + \sin(k/\alpha^{2i}) \cos(pos/\alpha^{2i}) \\
&= \cos(k/\alpha^{2i}) PE_{(pos,2i)} + \sin(k/\alpha^{2i}) PE_{(pos,2i+1)}
\end{align*}
\begin{align*}
PE_{(pos + k,2i+1)} &= \cos(k/\alpha^{2i}) \cos(pos/\alpha^{2i}) - \sin(k/\alpha^{2i}) \sin(pos/\alpha^{2i}) \\
&= \cos(k/\alpha^{2i}) PE_{(pos,2i+1)} - \sin(k/\alpha^{2i}) PE_{(pos,2i)}
\end{align*}

となる。よって、

\begin{align*}
\begin{pmatrix} PE_{(pos + k,2i)} \\ PE_{(pos + k,2i+1)} \end{pmatrix} = \begin{pmatrix}
\cos(k/\alpha^{2i}) & \sin(k/\alpha^{2i}) \\
- \sin(k/\alpha^{2i}) & \cos(k/\alpha^{2i})
\end{pmatrix}
\begin{pmatrix} PE_{(pos,2i)} \\ PE_{(pos,2i+1)} \end{pmatrix}
\end{align*}

となるので、$PE_{pos + k}$ は $PE_{pos}$ の一次変換として記述できる*1。

*1:対角成分に 2x2 行列が並んだ大きな行列をかけこむ。