Transformer 位置编码代码解析

Transformer 的 Multi-Head-Attention 无法判断各个编码的位置信息。因此 Attention is all you need 中加入三角函数位置编码（sinusoidal position embedding），表达形式为：
$$\begin{array}{rl}& P{E}_{(\mathrm{pos},2i)}=\sin \left(\mathrm{pos}/10000^{2i/d_{\text{modal }}}\right)\\ & P{E}_{(pos,2i+1)}=\cos \left(\mathrm{pos}/10000^{2i/d_{\text{model }}}\right)\end{array}$$
其中 pos 是单词位置，i = (0,1,... d_model) 所以d_model为 512 情况下，第一个单词的位置编码可以表示为：
$$\begin{array}{rl}& PE(1)=[\sin \left(1/10000^{0/512}\right),\cos \left(1/10000^{0/512}\right),\sin \left(1/10000^{2/512}\right),\cos \\ & \left(1/10000^{2/512}\right),\dots ]\end{array}$$

代码

import numpy as np
import matplotlib.pyplot as pltdef get_angles(pos, i, d_model):angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))return pos * angle_ratesdef positional_encoding(position, d_model):angle_rads = get_angles(np.arange(position)[:, np.newaxis],np.arange(d_model)[np.newaxis, :],d_model)# apply sin to even indices in the array; 2iangle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])# apply cos to odd indices in the array; 2i+1angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])pos_encoding = angle_rads[np.newaxis, ...]return pos_encodingtokens = 10
dimensions = 64pos_encoding = positional_encoding(tokens, dimensions)
print (pos_encoding.shape)plt.figure(figsize=(12,8))
plt.pcolormesh(pos_encoding[0], cmap='viridis')
plt.xlabel('Embedding Dimensions')
plt.xlim((0, dimensions))
plt.ylim((tokens,0))
plt.ylabel('Token Position')
plt.colorbar()
plt.show()