【论文精读】Self-Attentive Assocative Memory，2020

1 引言

这篇论文介绍了基于对象和对象关系的记忆模型，这对于设计类脑记忆模型有很大的启发作用。

2 Outer product attention (OPA)

标准transformer模型中定义的是内积注意力，即dot product attention：
$$A^{°}(q,K,V)=\sum _{i=1}^{n_{kv}}S(q\cdot k_i)v_i$$
其中，$A^{°}\in R^{d_v},q,k_i\in R^{d_{qk}},v_i\in R^{d_v}$，$\cdot$表示内积，计算结果是个标量，$S$是一个对向量元素的softmax计算函数。

作者定义了外积注意力命名为Outer product attention：
$$A^{\otimes }(q,K,V)=\sum _{i=1}^{n_{kv}}\text{F}(q\odot k_i)\otimes v_i$$
其中，$A^{\otimes }\in R^{d_{qk}\times d_v},q,k_i\in R^{d_{qk}},v\in R^{d_v}$，$\odot$表示对应位置元素的相乘，计算结果是个同维数向量，$\otimes$表示外积，$F$是一个对向量元素的tanh计算函数。

最好对照着标准注意力去理解。
差异：$A^{°}$是token序列中受注意力关注的token，$A^{\otimes }$是token序列中token之间的关系表征。

3 Self-attentive Associative Memory (SAM)

作者设计了一个关联记忆网络模块，命名为SAM，用来表征item及item之间的关系。

$$\begin{array}{rl}{\text{SAM}}_{\theta }(M)[s]& =A^{\otimes }(M_q[s],M_k,M_v)\\ & =\sum _{j=1}^{n_{kv}}\text{F}(M_q[s]\odot M_k[j])\otimes M_v[j]\end{array}$$
其中，
与注意力相关的q,k,v三个向量$M_q,M_k,M_v$：
$$\begin{array}{rl}{M}_q& =\mathcal{L}\mathcal{N}(W_{q}M)\\ M_k& =\mathcal{L}\mathcal{N}(W_{k}M)\\ M_v& =\mathcal{L}\mathcal{N}(W_{v}M)\end{array}$$
M是输入token序列组成的向量矩阵, $M\in R^{n\times d}$，n为token序列长度，d为token的维度；
$s$为M中第s行；
$W_q,W_k,W_v$是q,k,v对应线性变换层的参数矩阵；
$\mathcal{L}\mathcal{N}$是 layer normalization操作，而不是激活函数；
$\theta$代表SAM模块的内部参数是 { W q ∈ R n k v × n , W k ∈ R n k v × n , W v ∈ R n k v × n } \{W_q ∈ R^{n_{kv}×n},W_k ∈ R^{n_{kv}×n},W_v ∈ R^{n_{kv}×n}\} {Wq​∈Rnkv​×n,Wk​∈Rnkv​×n,Wv​∈Rnkv​×n}，$n_q$是query的个数，$n_{kv}$是key-value对的个数；

4 SAM-based Two-Memory Model (STM)

作者设计了2个记忆模块分别为$M_t^i\in R^{d\times d},M_t^r\in R^{n_q\times d\times d}$，都是基于SAM实现的，前者是用来记忆item，后者用来记忆item之间的关联关系。

4.1 $M^i$写操作

$$\begin{array}{rl}{X}_t& =f_1(x_t)\otimes f_2(x_t)\\ M_t^i& =F_t(M_{t-1}^i,x_t)\odot M_{t-1}^i+I_t(M_{t-1}^i,x_t)\odot X_t\end{array}$$
其中，
$x_t$是输入数据；
$f_1,f_2$是前馈神经网络，输出维度为d;
$F_t$为遗忘门，计算公式为$F_t(M_{t-1}^i,x_t)=W_F{x}_t+U_{F}tanh(M_{t-1}^i)+b_F$，其中$W_F,U_F\in R^{d\times d}$为网络参数；
$I_t$为输入的门控，计算公式为$I_t(M_{t-1}^i,x_t)=W_I{x}_t+U_{I}tanh(M_{t-1}^i)+b_I$，其中$W_I,U_I\in R^{d\times d}$为网络参数；

4.2 $M^r$读操作

$$\begin{array}{r}{v}_t^r=softmax(f_3(x_t{)}^{\top })M_{t-1}^r{f}_2(x_t)\end{array}$$
其中，
$v_t^r$为从关系记忆模块$M^r$中读出的值，将在下式（9）中使用；
$f_3$是前馈神经网络，输出维度为$n_q$;
$M_{t-1}^r$为$M^r$的前一个状态，其状态值由下式（9）计算得到；

4.3 $M^i$读操作和$M^r$写操作过程

$$\begin{array}{r}{M}_t^r=M_{t-1}^r+{\alpha }_1{\text{SAM}}_{\theta }(M_t^i+{\alpha }_2{v}_t^r\otimes f_2(x_t))\end{array}$$
其中，
${\alpha }_1,{\alpha }_2$是调和超参数，用于平衡量纲的，又类似于学习率；

4.4 用$M^r$实现item转移

$M^i$利用$M^r$实现更新，可以认为是hebbian更新，更新公式如下：
$$\begin{array}{r}{M}_t^i=M_t^i+{\alpha }_3{\mathcal{G}}_1◦{\mathcal{V}}_f◦M_t^r\end{array}$$
其中，
${\mathcal{V}}_f$是输入X(其shape为（batch_size, sequeue_length, dimension）)的前两维展开的向量；
${\mathcal{G}}_1$是前馈神经网络，负责维度变换$R^{(n_{q}d)\times d}\to R^{d\times d}$，其计算公式为${\mathcal{G}}_1(X)=W^g{\mathcal{V}}_f(X)$；
${\alpha }_3$是调和超参数；

4.5 模型输出$o_t$

$$\begin{array}{r}{o}_t={\mathcal{G}}_3◦{\mathcal{V}}_l◦{\mathcal{G}}_2◦{\mathcal{V}}_l◦M_t^r\end{array}$$
其中，
${\mathcal{V}}_l$是输入X(其shape为（batch_size, sequeue_length, dimension）)的后两维展开的向量；
${\mathcal{G}}_2，{\mathcal{G}}_3$是前馈神经网络，分别负责维度变换$R^{n_q\times dd}\to R^{d\times d}$，$R^{n_q{n}_r}\to R^{n_o}$，其计算公式为${\mathcal{G}}_2(X)=W^g{\mathcal{V}}_l(X)$，$n_q$是query的个数，$n_r$是超参数；

5 实验结果

源代码：https://github.com/thaihungle/SAM
作者做了消融实验，并在几何与图任务、强化学习任务、问答任务上做了测试。具体可以看论文附录和源码。

6 总结

该论文一个有趣的idea就是用两个前馈神经网络$M^i,M^r$分别表示对象与对象间关系，但是参数更新方法不是梯度下降而是赫布更新，后续可能是一个改进点。

7 参考资料

[1]. Self-Attentive Assocative Memory , 2020.