$\infty$-former: Infinite Memory Transformer

Martins P., Marinho Z. and Martins A. $\infty$-former: Infinite Memory Transformer. arXiv preprint arXiv:2109.00301, 2021.

概

在transformer中引入一种长期记忆机制.

主要内容

假设$X\in {\mathbb{R}}^{L\times d}$, 即每一行$x_i$代表一个token对应的特征.
Attention需要进行如下的步骤:
$$\begin{gathered} Q=X{W}^Q,K=X{W}^K,V=X{W}^V, \\ Z=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}(\frac{Q{K}^T}{\sqrt{d}})V. \end{gathered}$$
为了符号简易起见, 我们不考虑multi-head的情形, 下面的思想可以直接应用之.

我们知道, 可以通过径向基函数来逼近任意的连续函数:
$$\sum _k{b}_k{\psi }_k(t)\to f(t).$$
现在, 我们令$t_i=\frac{i}{L}$, 即对$L$个tokens冠以时序, $X$的每一列都可以看成一个特殊的$f_j(t)$的位于$t_i,i=0,1,\cdots ,L-1$处的值.
给定$N$个基函数${\psi }_k(t),k=0,1,\cdots ,N-1$, 我们要通过求解系数${\mathbf{b}}_j=[b_{j0},b_{j1},\cdots b_{j,N-1}{]}^T$来逼近$f_j$($X$的第$j$列).
设$\Psi \in {\mathbb{R}}^{N\times L},{\Psi }_{ki}={\psi }_k(t_i)$, $B\in {\mathbb{R}}^{d\times N},B_{jk}=b_{jk}$.
作者通过岭回归来求解系数$b$:
$$B=\arg \underset{B}{\min }\Vert B\Psi -X^T{\Vert }_F^2+\lambda \Vert B{\Vert }_F^2,$$
其显示表达式为:
$$B=X^T{\Psi }^T(\Psi {\Psi }^T+\lambda I{)}^{-1}.$$
故
$$X^T\approx B\Psi \to x_i\approx B\psi (t_i).$$
现在我们用$\tilde{X}:={\Psi }^T{B}^T$来代替$X$, 则
$$K=\tilde{X}{W}^K={\Psi }^T{B}^T{W}^K,\tilde{V}=\tilde{X}{W}^V={\Psi }^T{B}^T{W}^V.$$
注意, 我们并不对$Q$进行替换, 因为这个只是用作长期的记录用, Q每次重新计算.
对于每个$q_i$, 我们构建一个其关于$t$的密度函数$p_i(t)$, 文中假设其满足高斯分布:
$$\mathcal{N}(t;{\mu }_i;{\sigma }_i^2).$$
${\mu }_i,{\sigma }_i^2$分别通过如下估计:
$$\begin{gathered} {\mu }_i=\mathrm{s}\mathrm{i}\mathrm{g}\mathrm{m}\mathrm{o}\mathrm{i}\mathrm{d}(w_{\mu }^{T}K{q}_i)=\mathrm{s}\mathrm{i}\mathrm{g}\mathrm{m}\mathrm{o}\mathrm{i}\mathrm{d}(w_{\mu }^T{B}^T{W}^K{q}_i), \\ {\sigma }_i^2=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{p}\mathrm{l}\mathrm{u}\mathrm{s}(w_{\sigma }^{T}K{q}_i)=\mathrm{s}\mathrm{i}\mathrm{g}\mathrm{m}\mathrm{o}\mathrm{i}\mathrm{d}(w_{\sigma }^T{B}^T{W}^K{q}_i). \end{gathered}$$
注意最后令$w^T{\Psi }^T=w^T$既然$\Psi$是事先确定的.
我们知道
$$\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}(\frac{K{q}_i}{\sqrt{d}})$$
实际上求解的是一个离散化的$p_i(t)$, 即$q_i$和$k_j$的相合程度, 而
$$\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}(\frac{K{q}_i}{\sqrt{d}}{)}^{T}V$$
实际上就是求解期望
$${\mathbb{E}}_{p_i}[v(t)].$$
现在我们近似了一个连续的$p_i(t)$, 也可以通过这种方式得到最后的$z_i$:
$${\mathbb{E}}_{p_i}[v(t)]={\mathbb{E}}_{p_i}[{\psi }^T(t)B^T{W}^V]={\mathbb{E}}_{p_i}[{\psi }^T(t)]B^T{W}^V.$$
当我们取$\psi$为高斯径向基函数的时候, 上述是由显示解的.

现在来剖析一下, 好在哪里?
原本的$K$是$L\times d$的, 现在由于我们只需要计算$B^{T}W$, 故实际上只有$N\times d$, 我们可以选取很大的$L$但是选择较小的$N$来避免较高的复杂度.

如何扩展?

难不成每一次都要重新计算$B$? 倘若真的是这样就谈不上是长期记忆了.
作者采取了一种比较巧的方法, 实际上, 现在的$B\psi (t)$可以看成是一个$d$维的向量函数.
我们首先将其进行压缩至$[0,\tau ],\tau \in (0,1)$:
$$B\psi (t/\tau ),$$
如此一来, 整个函数的能量集中在$[0,\tau ]$中, 我们可以用剩下的$(\tau ,1]$来放置新的$X$.
我们首先从$[0,\tau ]$中采样$M$个点$t_0,\cdots ,t_{M-1}$, 并得到:
$$X_{past}=[x_0,\cdots ,x_{M-1}{]}^T\in {\mathbb{R}}^{M\times d},x_m={\psi }^T(t_m/\tau )B^T.$$
加上新的$X_{new}$, 我们有
$$X=[X_{past}^T,X_{new}^T{]}^T\in {\mathbb{R}}^{(M+L)\times d},$$
对$X$按照上面的逻辑重新估计$B$即可更新记忆.

关于如何采样这$M$个点, 作者提了一种sticky memories的方法, 将其与密度函数联系在一起, 便不细讲了.

实验细节

在看这篇论文的时候, 困扰我的就是这个径向基函数是怎么选的?
举一个作者在Language Modeling中的例子便可:
选取150个高斯径向基函数$\mathcal{N}(t;\mu ,{\sigma }^2)$, 其中
$\mu$从$[0,1]$中均匀采样, σ ∈ { 0.01 , 0.05 } \sigma \in \{0.01, 0.05\} σ∈{0.01,0.05}.

还有用KL散度防止一般化就不讲了. 感觉本文有趣的点就是压缩这个地方, 还有对$\Psi$的处理.