Augmenting Self-attention with Persistent Memory
Augmenting Self-attention with Persistent Memory
arXiv:\url{1907.01470 [cs.LG]}
自注意力机制中的 key、value 和 query 分别基于输入的信息，并由公式 W^KX 表示。
我们可以同时使用 N 个不依赖输入的关键值对。
前馈层实际上是一种特殊的注意力机制，并因而能够整合到自注意力机制中作为不依赖输入的关键值对的运算。
见下图.
上图左边是正常注意力机制, 右边蓝色部分是 M_V, 橙色部分是 M_U
编码层 由 注意力机制 & 前馈神经网络 & 残差归一化 组成.
自注意力
设序列 X:=(x_1,\cdots,x_T), 其中每个列向量 x_i 的维度为d。在自注意力机制中, 查询器、键和值分别对应于
q_i:=W^Qx_i, \qquad k_i:=W^Kx_i, \qquad v_i:=W^Vx_i.
W^Q,W^K,W^V 都是是 d_h \times d 矩阵, 其中 d_h:=\frac{d}{H}, H 为注意力的头数目.
由 N 组 key-value 对构成的关键-值对集与输入无关，并定义其对应的向量空间维度均为 d_h。通过扩展的方式生成对应的 key 和 value 矩阵分别为：
(k_1,\cdots,k_T,k_{T+1},\cdots,k_{T+N}) := (k_1,\cdots,k_T,M_V)
(v_1,\cdots,v_T,v_{T+1},\cdots,v_{T+N}) := (v_1,\cdots,v_T,M_U)
x_i 和 x_j 的相似度为
a_{i,j} 被定义为 q_i^T(k_j + p(i, j))。
其中位置信息由 P(i, j) 表示。
经过 Softmax 处理后得到矩阵 \hat{A} = (\hat{a}_{i, j})。
从而得到注意力权重矩阵 \hat{A}。
对于每个输入向量 \mathbf{x}_j = (x_1, x_2, \dots, x_n)^T，
我们有：
y_j = \sum_{k=1}^{n} \hat{a}_{j,k}(v_k + p'(j, k))
其中另一个位置编码由 p'(j, k) 给出。
关于前馈神经网络被取消的解释.
前馈神经网络
上文阐述了注意力机制的相关内容。随后介绍 feed 网络及其隐含层的激活函数特性。其具体形式为：对于任意时间步 t 的输入 x_t ， feed 网络的前向计算定义为 \operatorname{FF}(x_t):=U \operatorname{ReLU}(Vx_t +b) +b'
Persistent Memory
在Transformer架构中，前馈神经网络的计算开销相对较高。如果直接去除这一部分，则会导致整个Transformer模型性能下降。在这里，我们采用Softmax函数来替代前馈层的激活函数，并定义t时刻的输出变量为：
y_t = U \cdot \text{softmax}(V x_t + b) + b'
令b = b' = 0向量，则前馈层可简化为：
y_t = U \cdot \text{softmax}(V x_t)
其中，
- V_{i,*} 是第i个行向量，
- x_t 被视为查询，
- U_{*,j} 是第j个列向量。
基于此分析的基础之上